Fix/copilot concurrency limit

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.

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

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")

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)
+	})
+}

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)
+}

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),
 	}
 }

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")
+	}
+}