diff --git a/flytecopilot/README.md b/flytecopilot/README.md index 92304e947f8..a8d272ea93b 100644 --- a/flytecopilot/README.md +++ b/flytecopilot/README.md @@ -12,6 +12,9 @@ This is achieved using `flyte-copilot` a binary that runs in 2 modes, $ flyte-copilot downloader ``` +Useful flag: +- `--concurrency-per-cpu`: Limits parallel download work. Effective concurrency is `GOMAXPROCS * concurrency-per-cpu` (with a max cap). + In K8s `flyte-copilot downloader` can be run as part of the init containers with the download volume mounted. This guarantees that the metadata and any data (if configured) is downloaded before the main container starts up. @@ -27,6 +30,9 @@ is downloaded before the main container starts up. $ flyte-copilot sidecar ``` +Useful flag: +- `--concurrency-per-cpu`: Limits parallel upload work. Effective concurrency is `GOMAXPROCS * concurrency-per-cpu` (with a max cap). + ### Raw notes Solution 1: poll Kubeapi. diff --git a/flytecopilot/cmd/download.go b/flytecopilot/cmd/download.go index 3827308c335..bd51d321d21 100644 --- a/flytecopilot/cmd/download.go +++ b/flytecopilot/cmd/download.go @@ -21,6 +21,7 @@ type DownloadOptions struct { inputInterface []byte metadataFormat string downloadMode string + concurrencyPerCPU int timeout time.Duration } @@ -70,7 +71,7 @@ func (d *DownloadOptions) Download(ctx context.Context) error { if !ok { return fmt.Errorf("incorrect input download mode specified, given [%s], possible values [%+v]", d.downloadMode, GetDownloadModeVals()) } - dl := data.NewDownloader(ctx, d.Store, core.DataLoadingConfig_LiteralMapFormat(f), core.IOStrategy_DownloadMode(m)) + dl := data.NewDownloader(ctx, d.Store, core.DataLoadingConfig_LiteralMapFormat(f), core.IOStrategy_DownloadMode(m), d.concurrencyPerCPU) childCtx := ctx cancelFn := func() {} if d.timeout > 0 { @@ -115,6 +116,7 @@ func NewDownloadCommand(opts *RootOptions) *cobra.Command { downloadCmd.Flags().StringVarP(&downloadOpts.localDirectoryPath, "to-local-dir", "o", "", "The local directory on disk where data should be downloaded.") downloadCmd.Flags().StringVarP(&downloadOpts.metadataFormat, "format", "m", core.DataLoadingConfig_JSON.String(), fmt.Sprintf("What should be the output format for the primitive and structured types. Options [%v]", GetFormatVals())) downloadCmd.Flags().StringVarP(&downloadOpts.downloadMode, "download-mode", "d", core.IOStrategy_DOWNLOAD_EAGER.String(), fmt.Sprintf("Download mode to use. Options [%v]", GetDownloadModeVals())) + downloadCmd.Flags().IntVar(&downloadOpts.concurrencyPerCPU, "concurrency-per-cpu", data.DefaultConcurrencyPerCPU, "Concurrency multiplier used to cap concurrent download workers. Effective concurrency is CPU*concurrency-per-cpu.") downloadCmd.Flags().DurationVarP(&downloadOpts.timeout, "timeout", "t", time.Hour*1, "Max time to allow for downloads to complete, default is 1H") downloadCmd.Flags().BytesBase64VarP(&downloadOpts.inputInterface, "input-interface", "i", nil, "Input interface proto message - core.VariableMap, base64 encoced string") return downloadCmd diff --git a/flytecopilot/cmd/sidecar.go b/flytecopilot/cmd/sidecar.go index 073e3c98f3c..89ec15156d5 100644 --- a/flytecopilot/cmd/sidecar.go +++ b/flytecopilot/cmd/sidecar.go @@ -33,12 +33,13 @@ type UploadOptions struct { // Local directory path where the sidecar should look for outputs. localDirectoryPath string // Non primitive types will be dumped in this output format - metadataFormat string - uploadMode string - timeout time.Duration - typedInterface []byte - startWatcherType containerwatcher.WatcherType - exitWatcherType containerwatcher.WatcherType + metadataFormat string + uploadMode string + concurrencyPerCPU int + timeout time.Duration + typedInterface []byte + startWatcherType containerwatcher.WatcherType + exitWatcherType containerwatcher.WatcherType } func (u *UploadOptions) createWatcher(_ context.Context, w containerwatcher.WatcherType) (containerwatcher.Watcher, error) { @@ -99,7 +100,7 @@ func (u *UploadOptions) uploader(ctx context.Context) error { return err } - dl := data.NewUploader(ctx, u.Store, core.DataLoadingConfig_LiteralMapFormat(f), core.IOStrategy_UploadMode(m), ErrorFile) + dl := data.NewUploader(ctx, u.Store, core.DataLoadingConfig_LiteralMapFormat(f), core.IOStrategy_UploadMode(m), ErrorFile, u.concurrencyPerCPU) childCtx, cancelFn := context.WithTimeout(ctx, u.timeout) defer cancelFn() @@ -145,6 +146,7 @@ func NewUploadCommand(opts *RootOptions) *cobra.Command { uploadCmd.Flags().StringVarP(&uploadOptions.localDirectoryPath, "from-local-dir", "f", "", "The local directory on disk where data will be available for upload.") uploadCmd.Flags().StringVarP(&uploadOptions.metadataFormat, "format", "m", core.DataLoadingConfig_JSON.String(), fmt.Sprintf("What should be the output format for the primitive and structured types. Options [%v]", GetFormatVals())) uploadCmd.Flags().StringVarP(&uploadOptions.uploadMode, "upload-mode", "u", core.IOStrategy_UPLOAD_ON_EXIT.String(), fmt.Sprintf("When should upload start/upload mode. Options [%v]", GetUploadModeVals())) + uploadCmd.Flags().IntVar(&uploadOptions.concurrencyPerCPU, "concurrency-per-cpu", data.DefaultConcurrencyPerCPU, "Concurrency multiplier used to cap concurrent upload workers. Effective concurrency is CPU*concurrency-per-cpu.") uploadCmd.Flags().StringVarP(&uploadOptions.metaOutputName, "meta-output-name", "", "outputs.pb", "The key name under the remoteOutputPrefix that should be return to provide meta information about the outputs on successful execution") uploadCmd.Flags().DurationVarP(&uploadOptions.timeout, "timeout", "t", time.Hour*1, "Max time to allow for uploads to complete, default is 1H") uploadCmd.Flags().BytesBase64VarP(&uploadOptions.typedInterface, "interface", "i", nil, "Typed Interface - core.TypedInterface, base64 encoded string of the serialized protobuf") diff --git a/flytecopilot/data/concurrency_test.go b/flytecopilot/data/concurrency_test.go new file mode 100644 index 00000000000..ada373f23f1 --- /dev/null +++ b/flytecopilot/data/concurrency_test.go @@ -0,0 +1,68 @@ +package data + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +func TestSanitizeConcurrencyPerCPU(t *testing.T) { + assert.Equal(t, DefaultConcurrencyPerCPU, sanitizeConcurrencyPerCPU(0)) + assert.Equal(t, DefaultConcurrencyPerCPU, sanitizeConcurrencyPerCPU(-1)) + assert.Equal(t, DefaultConcurrencyPerCPU, sanitizeConcurrencyPerCPU(-100)) + assert.Equal(t, 1, sanitizeConcurrencyPerCPU(1)) + assert.Equal(t, 2, sanitizeConcurrencyPerCPU(2)) + assert.Equal(t, 10, sanitizeConcurrencyPerCPU(10)) +} + +func TestGetConcurrency(t *testing.T) { + result := getConcurrency(DefaultConcurrencyPerCPU) + assert.GreaterOrEqual(t, result, 1) + + result = getConcurrency(0) + assert.GreaterOrEqual(t, result, 1) + + result = getConcurrency(1) + assert.GreaterOrEqual(t, result, 1) + assert.LessOrEqual(t, result, MaxConcurrency) + + result = getConcurrency(100000) + assert.LessOrEqual(t, result, MaxConcurrency) +} + +func TestAcquireSemaphore(t *testing.T) { + t.Run("acquires successfully", func(t *testing.T) { + sem := make(chan struct{}, 2) + ctx := context.Background() + + assert.NoError(t, acquireSemaphore(ctx, sem)) + assert.NoError(t, acquireSemaphore(ctx, sem)) + assert.Len(t, sem, 2) + }) + + t.Run("respects context cancellation", func(t *testing.T) { + sem := make(chan struct{}, 1) + sem <- struct{}{} + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + + err := acquireSemaphore(ctx, sem) + assert.Error(t, err) + assert.ErrorIs(t, err, context.DeadlineExceeded) + }) + + t.Run("returns immediately on already canceled context", func(t *testing.T) { + sem := make(chan struct{}, 1) + sem <- struct{}{} + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err := acquireSemaphore(ctx, sem) + assert.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) + }) +} diff --git a/flytecopilot/data/constructor_test.go b/flytecopilot/data/constructor_test.go new file mode 100644 index 00000000000..30fea5712de --- /dev/null +++ b/flytecopilot/data/constructor_test.go @@ -0,0 +1,57 @@ +package data + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core" + "github.com/flyteorg/flyte/flytestdlib/promutils" + "github.com/flyteorg/flyte/flytestdlib/storage" +) + +func TestNewDownloaderSanitizesConcurrencyPerCPU(t *testing.T) { + store, err := storage.NewDataStore(&storage.Config{Type: storage.TypeMemory}, promutils.NewTestScope()) + assert.NoError(t, err) + + d := NewDownloader(context.Background(), store, core.DataLoadingConfig_YAML, core.IOStrategy_DOWNLOAD_EAGER, 0) + assert.Equal(t, core.DataLoadingConfig_YAML, d.format) + assert.Equal(t, store, d.store) + assert.Equal(t, core.IOStrategy_DOWNLOAD_EAGER, d.mode) + assert.Equal(t, DefaultConcurrencyPerCPU, d.concurrencyPerCPU) + + d = NewDownloader(context.Background(), store, core.DataLoadingConfig_JSON, core.IOStrategy_DOWNLOAD_EAGER, 2) + assert.Equal(t, 2, d.concurrencyPerCPU) +} + +func TestNewUploaderSanitizesConcurrencyPerCPU(t *testing.T) { + store, err := storage.NewDataStore(&storage.Config{Type: storage.TypeMemory}, promutils.NewTestScope()) + assert.NoError(t, err) + + u := NewUploader(context.Background(), store, core.DataLoadingConfig_YAML, core.IOStrategy_UPLOAD_ON_EXIT, "errors.pb", -10) + assert.Equal(t, core.DataLoadingConfig_YAML, u.format) + assert.Equal(t, store, u.store) + assert.Equal(t, core.IOStrategy_UPLOAD_ON_EXIT, u.mode) + assert.Equal(t, "errors.pb", u.errorFileName) + assert.Equal(t, DefaultConcurrencyPerCPU, u.concurrencyPerCPU) + + u = NewUploader(context.Background(), store, core.DataLoadingConfig_JSON, core.IOStrategy_UPLOAD_ON_EXIT, "errors.pb", 3) + assert.Equal(t, 3, u.concurrencyPerCPU) +} + +func TestRecursiveDownloadReturnsEmptyMapForNilOrEmptyInputs(t *testing.T) { + d := Downloader{} + + varMap, literalMap, err := d.RecursiveDownload(context.Background(), nil, t.TempDir(), true) + assert.NoError(t, err) + assert.NotNil(t, varMap) + assert.Empty(t, varMap) + assert.Nil(t, literalMap) + + varMap, literalMap, err = d.RecursiveDownload(context.Background(), &core.LiteralMap{}, t.TempDir(), true) + assert.NoError(t, err) + assert.NotNil(t, varMap) + assert.Empty(t, varMap) + assert.Nil(t, literalMap) +} diff --git a/flytecopilot/data/download.go b/flytecopilot/data/download.go index 084469faed2..766c106f828 100644 --- a/flytecopilot/data/download.go +++ b/flytecopilot/data/download.go @@ -9,6 +9,7 @@ import ( "path" "path/filepath" "reflect" + "runtime" "strconv" "strings" "sync" @@ -30,7 +31,43 @@ type Downloader struct { format core.DataLoadingConfig_LiteralMapFormat store *storage.DataStore // TODO support download mode - mode core.IOStrategy_DownloadMode + mode core.IOStrategy_DownloadMode + concurrencyPerCPU int +} + +const DefaultConcurrencyPerCPU = 4 +const MaxConcurrency = 1024 + +func sanitizeConcurrencyPerCPU(concurrencyPerCPU int) int { + if concurrencyPerCPU <= 0 { + return DefaultConcurrencyPerCPU + } + return concurrencyPerCPU +} + +func getConcurrency(concurrencyPerCPU int) int { + cpus := runtime.GOMAXPROCS(0) + factor := sanitizeConcurrencyPerCPU(concurrencyPerCPU) + + concurrency := cpus * factor + if concurrency < 1 { + concurrency = 1 + } + if concurrency > MaxConcurrency { + concurrency = MaxConcurrency + } + return concurrency +} + +// acquireSemaphore attempts to acquire a slot from the semaphore channel. +// It respects context cancellation so callers are not blocked indefinitely. +func acquireSemaphore(ctx context.Context, sem chan struct{}) error { + select { + case sem <- struct{}{}: + return nil + case <-ctx.Done(): + return ctx.Err() + } } // TODO add timeout and rate limit @@ -101,10 +138,19 @@ func (d Downloader) handleBlob(ctx context.Context, blob *core.Blob, toPath stri // We use Mutex to avoid race conditions when updating counters and creating directories var mu sync.Mutex var wg sync.WaitGroup + concurrency := getConcurrency(d.concurrencyPerCPU) + sem := make(chan struct{}, concurrency) + var acquireErr error for _, absPath := range absPaths { + absPath := absPath + if err := acquireSemaphore(ctx, sem); err != nil { + acquireErr = err + break + } wg.Add(1) go func() { + defer func() { <-sem }() defer wg.Done() defer func() { if err := recover(); err != nil { @@ -216,6 +262,9 @@ func (d Downloader) handleBlob(ctx context.Context, blob *core.Blob, toPath stri } // Go routines are synchronized with a WaitGroup to prevent goroutine leaks. wg.Wait() + if acquireErr != nil { + return nil, errors.Wrapf(acquireErr, "context canceled while downloading multipart blob from [%s]", blobRef) + } if downloadSuccess != itemCount || readerCloseSuccessCount != itemCount || writerCloseSuccessCount != itemCount { return nil, errors.Errorf( "Failed to copy %d out of %d remote files from [%s] to local [%s].\n"+ @@ -475,6 +524,8 @@ func (d Downloader) RecursiveDownload(ctx context.Context, inputs *core.LiteralM return VarMap{}, nil, nil } f := make(FutureMap, len(inputs.GetLiterals())) + concurrency := getConcurrency(d.concurrencyPerCPU) + sem := make(chan struct{}, concurrency) for variable, literal := range inputs.GetLiterals() { if literal.GetOffloadedMetadata() != nil { offloadedMetadataURI := literal.GetOffloadedMetadata().GetUri() @@ -488,7 +539,11 @@ func (d Downloader) RecursiveDownload(ctx context.Context, inputs *core.LiteralM } varPath := path.Join(dir, variable) lit := literal + if err := acquireSemaphore(childCtx, sem); err != nil { + return nil, nil, errors.Wrapf(err, "context canceled while scheduling download for variable [%s]", variable) + } f[variable] = futures.NewAsyncFuture(childCtx, func(ctx2 context.Context) (interface{}, error) { + defer func() { <-sem }() v, lit, err := d.handleLiteral(ctx2, lit, varPath, writePrimitiveToFile) if err != nil { return nil, err @@ -565,10 +620,11 @@ func (d Downloader) DownloadInputs(ctx context.Context, inputRef storage.DataRef return nil } -func NewDownloader(_ context.Context, store *storage.DataStore, format core.DataLoadingConfig_LiteralMapFormat, mode core.IOStrategy_DownloadMode) Downloader { +func NewDownloader(_ context.Context, store *storage.DataStore, format core.DataLoadingConfig_LiteralMapFormat, mode core.IOStrategy_DownloadMode, concurrencyPerCPU int) Downloader { return Downloader{ - format: format, - store: store, - mode: mode, + format: format, + store: store, + mode: mode, + concurrencyPerCPU: sanitizeConcurrencyPerCPU(concurrencyPerCPU), } } diff --git a/flytecopilot/data/download_error_test.go b/flytecopilot/data/download_error_test.go new file mode 100644 index 00000000000..a018ce859fa --- /dev/null +++ b/flytecopilot/data/download_error_test.go @@ -0,0 +1,44 @@ +package data + +import ( + "bytes" + "context" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core" + "github.com/flyteorg/flyte/flytestdlib/promutils" + "github.com/flyteorg/flyte/flytestdlib/storage" +) + +func TestHandleBlobMultipartReportsFailedCopies(t *testing.T) { + store, err := storage.NewDataStore(&storage.Config{Type: storage.TypeMemory}, promutils.NewTestScope()) + assert.NoError(t, err) + + ref := storage.DataReference("s3://container/prefix/file.txt") + payload := []byte("payload") + err = store.WriteRaw(context.Background(), ref, int64(len(payload)), storage.Options{}, bytes.NewReader(payload)) + assert.NoError(t, err) + + downloader := Downloader{store: store} + blob := &core.Blob{ + Uri: "s3://container/prefix", + Metadata: &core.BlobMetadata{ + Type: &core.BlobType{Dimensionality: core.BlobType_MULTIPART}, + }, + } + + blockedPath := filepath.Join(t.TempDir(), "blocked") + err = os.WriteFile(blockedPath, []byte("x"), 0o600) + assert.NoError(t, err) + + result, err := downloader.handleBlob(context.Background(), blob, blockedPath) + assert.Nil(t, result) + if assert.Error(t, err) { + assert.Contains(t, err.Error(), "Failed to copy 1 out of 1 remote files") + assert.Contains(t, err.Error(), "Failed to close 1 writers") + } +} diff --git a/flytecopilot/data/upload.go b/flytecopilot/data/upload.go index dc0c69db64d..5b2e9cfa613 100644 --- a/flytecopilot/data/upload.go +++ b/flytecopilot/data/upload.go @@ -30,6 +30,7 @@ type Uploader struct { store *storage.DataStore aggregateOutputFileName string errorFileName string + concurrencyPerCPU int } type dirFile struct { @@ -99,18 +100,24 @@ func (u Uploader) handleBlobTypeMultipart(ctx context.Context, localPath string, childCtx, cancel := context.WithCancel(ctx) defer cancel() + concurrency := getConcurrency(u.concurrencyPerCPU) + sem := make(chan struct{}, concurrency) fileUploader := make([]futures.Future, 0, len(files)) for _, f := range files { pth := f.path ref := f.ref size := f.info.Size() + if err := acquireSemaphore(childCtx, sem); err != nil { + cancel() + return nil, errors.Wrapf(err, "context canceled while scheduling upload from [%s]", localPath) + } fileUploader = append(fileUploader, futures.NewAsyncFuture(childCtx, func(i2 context.Context) (i interface{}, e error) { + defer func() { <-sem }() return nil, UploadFileToStorage(i2, pth, ref, size, u.store) })) } for _, f := range fileUploader { - // TODO maybe we should have timeouts, or we can have a global timeout at the top level _, err := f.Get(ctx) if err != nil { return nil, err @@ -198,11 +205,12 @@ func (u Uploader) RecursiveUpload(ctx context.Context, vars *core.VariableMap, f return nil } -func NewUploader(_ context.Context, store *storage.DataStore, format core.DataLoadingConfig_LiteralMapFormat, mode core.IOStrategy_UploadMode, errorFileName string) Uploader { +func NewUploader(_ context.Context, store *storage.DataStore, format core.DataLoadingConfig_LiteralMapFormat, mode core.IOStrategy_UploadMode, errorFileName string, concurrencyPerCPU int) Uploader { return Uploader{ - format: format, - store: store, - errorFileName: errorFileName, - mode: mode, + format: format, + store: store, + errorFileName: errorFileName, + mode: mode, + concurrencyPerCPU: sanitizeConcurrencyPerCPU(concurrencyPerCPU), } } diff --git a/flytecopilot/data/upload_test.go b/flytecopilot/data/upload_test.go index 9131b9ab916..7e2212f35aa 100644 --- a/flytecopilot/data/upload_test.go +++ b/flytecopilot/data/upload_test.go @@ -44,7 +44,7 @@ func TestUploader_RecursiveUpload(t *testing.T) { outputRef := storage.DataReference("output") rawRef := storage.DataReference("raw") - u := NewUploader(context.TODO(), store, core.DataLoadingConfig_JSON, core.IOStrategy_UPLOAD_ON_EXIT, "error") + u := NewUploader(context.TODO(), store, core.DataLoadingConfig_JSON, core.IOStrategy_UPLOAD_ON_EXIT, "error", DefaultConcurrencyPerCPU) assert.NoError(t, u.RecursiveUpload(context.TODO(), vmap, tmpDir, outputRef, rawRef)) outputs := &core.LiteralMap{} @@ -95,7 +95,7 @@ func TestUploader_RecursiveUpload(t *testing.T) { outputRef := storage.DataReference("output") rawRef := storage.DataReference("raw") - u := NewUploader(context.TODO(), store, core.DataLoadingConfig_JSON, core.IOStrategy_UPLOAD_ON_EXIT, "error") + u := NewUploader(context.TODO(), store, core.DataLoadingConfig_JSON, core.IOStrategy_UPLOAD_ON_EXIT, "error", DefaultConcurrencyPerCPU) assert.NoError(t, u.RecursiveUpload(context.TODO(), vmap, tmpDir, outputRef, rawRef)) outputs := &core.LiteralMap{} diff --git a/flytecopilot/data/utils.go b/flytecopilot/data/utils.go index 923c75f908b..cde105e7807 100644 --- a/flytecopilot/data/utils.go +++ b/flytecopilot/data/utils.go @@ -87,5 +87,10 @@ func DownloadFileFromHTTP(ctx context.Context, ref storage.DataReference) (io.Re if err != nil { return nil, errors.Wrapf(err, "Failed to download from url :%s", ref) } + if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { + _ = resp.Body.Close() + return nil, fmt.Errorf("failed to download from url %s, received HTTP status %d %s", ref, resp.StatusCode, resp.Status) + } + logger.Infof(ctx, "HTTP download started: [%s], status=%d, content-length=%s", ref, resp.StatusCode, resp.Header.Get("Content-Length")) return resp.Body, nil } diff --git a/flytecopilot/data/utils_test.go b/flytecopilot/data/utils_test.go index 517f9d30ef2..bc5e34eb684 100644 --- a/flytecopilot/data/utils_test.go +++ b/flytecopilot/data/utils_test.go @@ -3,6 +3,8 @@ package data import ( "bytes" "context" + "net/http" + "net/http/httptest" "os" "path" "testing" @@ -94,6 +96,18 @@ func TestDownloadFromHttp(t *testing.T) { assert.Error(t, err) } +func TestDownloadFromHttpReturnsErrorForNonSuccessStatus(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, "boom", http.StatusBadGateway) + })) + defer server.Close() + + _, err := DownloadFileFromHTTP(context.TODO(), storage.DataReference(server.URL)) + if assert.Error(t, err) { + assert.Contains(t, err.Error(), "received HTTP status 502") + } +} + func TestDownloadFromStorage(t *testing.T) { store, err := storage.NewDataStore(&storage.Config{Type: storage.TypeMemory}, promutils.NewTestScope()) assert.NoError(t, err)