diff --git a/cmd/ome-agent/config_test.go b/cmd/ome-agent/config_test.go new file mode 100644 index 000000000..43eefccd8 --- /dev/null +++ b/cmd/ome-agent/config_test.go @@ -0,0 +1,79 @@ +package main + +import ( + "os" + "path/filepath" + "testing" + + "github.com/spf13/cobra" + "github.com/spf13/viper" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/fx" + + "sigs.k8s.io/ome/internal/ome-agent/replica" + "sigs.k8s.io/ome/pkg/constants" +) + +func TestConfigProviderPassesArtifactReuseEnvironmentOverrideToReplicaConfig(t *testing.T) { + viper.Reset() + t.Cleanup(viper.Reset) + + configPath := filepath.Join(t.TempDir(), "ome-agent.yaml") + require.NoError(t, os.WriteFile(configPath, []byte("local_path: /tmp/models\n"), 0600)) + + previousConfigFilePath := configFilePath + configFilePath = configPath + t.Cleanup(func() { configFilePath = previousConfigFilePath }) + t.Setenv(constants.AgentTargetArtifactReuseAllowedEnvVarKey, "true") + + command := &cobra.Command{} + command.Flags().Bool("debug", false, "") + + var configuredViper *viper.Viper + app := fx.New( + configProvider(command, nil), + fx.Populate(&configuredViper), + ) + require.NoError(t, app.Err()) + + config, err := replica.NewReplicaConfig(replica.WithViper(configuredViper)) + require.NoError(t, err) + assert.True(t, config.TargetArtifactReuseAllowed) +} + +func TestConfigProviderLeavesArtifactReuseDisabledWithoutEnvironmentOverride(t *testing.T) { + viper.Reset() + t.Cleanup(viper.Reset) + + previousValue, wasSet := os.LookupEnv(constants.AgentTargetArtifactReuseAllowedEnvVarKey) + require.NoError(t, os.Unsetenv(constants.AgentTargetArtifactReuseAllowedEnvVarKey)) + t.Cleanup(func() { + if wasSet { + _ = os.Setenv(constants.AgentTargetArtifactReuseAllowedEnvVarKey, previousValue) + return + } + _ = os.Unsetenv(constants.AgentTargetArtifactReuseAllowedEnvVarKey) + }) + + configPath := filepath.Join(t.TempDir(), "ome-agent.yaml") + require.NoError(t, os.WriteFile(configPath, []byte("local_path: /tmp/models\n"), 0600)) + + previousConfigFilePath := configFilePath + configFilePath = configPath + t.Cleanup(func() { configFilePath = previousConfigFilePath }) + + command := &cobra.Command{} + command.Flags().Bool("debug", false, "") + + var configuredViper *viper.Viper + app := fx.New( + configProvider(command, nil), + fx.Populate(&configuredViper), + ) + require.NoError(t, app.Err()) + + config, err := replica.NewReplicaConfig(replica.WithViper(configuredViper)) + require.NoError(t, err) + assert.False(t, config.TargetArtifactReuseAllowed) +} diff --git a/config/ome-agent/ome-agent.yaml b/config/ome-agent/ome-agent.yaml index 1f696d982..79e49ef59 100644 --- a/config/ome-agent/ome-agent.yaml +++ b/config/ome-agent/ome-agent.yaml @@ -33,6 +33,7 @@ download_size_limit_gb: 650 enable_size_limit_check: true hf_download_timeout: "72h" hf_download_stale_progress_timeout: "30m" +artifact_upload_lock_timeout: "120h" source: storage_uri: "oci://n//b//o/" diff --git a/internal/ome-agent/replica/config.go b/internal/ome-agent/replica/config.go index ce4050837..040aa17c5 100644 --- a/internal/ome-agent/replica/config.go +++ b/internal/ome-agent/replica/config.go @@ -24,9 +24,11 @@ type Config struct { LocalPath string `mapstructure:"local_path" validate:"required"` DownloadSizeLimitGB int `mapstructure:"download_size_limit_gb"` EnableSizeLimitCheck bool `mapstructure:"enable_size_limit_check"` - NumConnections int `mapstructure:"num_connections"` + TargetArtifactReuseAllowed bool `mapstructure:"target_artifact_reuse_allowed"` + NumConnections int `mapstructure:"num_connections" validate:"gt=0"` HFDownloadTimeout time.Duration `mapstructure:"hf_download_timeout"` HFDownloadStaleProgressTimeout time.Duration `mapstructure:"hf_download_stale_progress_timeout"` + ArtifactUploadLockTimeout time.Duration `mapstructure:"artifact_upload_lock_timeout"` Source struct { StorageURIStr string `mapstructure:"storage_uri" validate:"required"` @@ -67,6 +69,7 @@ func defaultConfig() *Config { EnableSizeLimitCheck: true, HFDownloadTimeout: 72 * time.Hour, HFDownloadStaleProgressTimeout: 30 * time.Minute, + ArtifactUploadLockTimeout: 120 * time.Hour, } } diff --git a/internal/ome-agent/replica/config_test.go b/internal/ome-agent/replica/config_test.go index b53ac73f4..53aa3c115 100644 --- a/internal/ome-agent/replica/config_test.go +++ b/internal/ome-agent/replica/config_test.go @@ -164,6 +164,7 @@ func TestWithViper(t *testing.T) { v.Set("enable_size_limit_check", true) v.Set("hf_download_timeout", "2h") v.Set("hf_download_stale_progress_timeout", "15m") + v.Set("artifact_upload_lock_timeout", "3h") v.Set("source.storage_uri", "oci://n/test-src-namespace/b/test-src-bucket/o/models") v.Set("target.storage_uri", "oci://n/test-tgt-namespace/b/test-tgt-bucket/o/models") return v @@ -176,6 +177,7 @@ func TestWithViper(t *testing.T) { assert.Equal(t, true, c.EnableSizeLimitCheck) assert.Equal(t, 2*time.Hour, c.HFDownloadTimeout) assert.Equal(t, 15*time.Minute, c.HFDownloadStaleProgressTimeout) + assert.Equal(t, 3*time.Hour, c.ArtifactUploadLockTimeout) assert.Equal(t, "oci://n/test-src-namespace/b/test-src-bucket/o/models", c.Source.StorageURIStr) assert.Equal(t, "oci://n/test-tgt-namespace/b/test-tgt-bucket/o/models", c.Target.StorageURIStr) }, @@ -193,6 +195,7 @@ func TestWithViper(t *testing.T) { assert.Equal(t, true, c.EnableSizeLimitCheck) assert.Equal(t, 72*time.Hour, c.HFDownloadTimeout) assert.Equal(t, 30*time.Minute, c.HFDownloadStaleProgressTimeout) + assert.Equal(t, 120*time.Hour, c.ArtifactUploadLockTimeout) }, }, { @@ -384,6 +387,24 @@ func TestConfig_Validate(t *testing.T) { }, expectError: true, }, + { + name: "invalid NumConnections", + setupConfig: func() *Config { + return &Config{ + LocalPath: "/test/path", + DownloadSizeLimitGB: 100, + EnableSizeLimitCheck: true, + NumConnections: 0, + Source: SourceStruct{ + StorageURIStr: validSourceURI, + }, + Target: TargetStruct{ + StorageURIStr: validTargetURI, + }, + } + }, + expectError: true, + }, { name: "invalid source storage URI", setupConfig: func() *Config { @@ -425,6 +446,8 @@ func TestDefaultConfig(t *testing.T) { assert.Equal(t, 10, config.NumConnections) assert.Equal(t, 650, config.DownloadSizeLimitGB) assert.Equal(t, true, config.EnableSizeLimitCheck) + assert.False(t, config.TargetArtifactReuseAllowed) + assert.Equal(t, 120*time.Hour, config.ArtifactUploadLockTimeout) } func TestConfig_ValidateRequiredDependencies(t *testing.T) { diff --git a/internal/ome-agent/replica/module_test.go b/internal/ome-agent/replica/module_test.go index 1086b7020..d5c8373eb 100644 --- a/internal/ome-agent/replica/module_test.go +++ b/internal/ome-agent/replica/module_test.go @@ -192,6 +192,7 @@ func TestModuleIntegration(t *testing.T) { v.Set("num_connections", 5) v.Set("download_size_limit_gb", 100) v.Set("enable_size_limit_check", true) + v.Set("target_artifact_reuse_allowed", true) v.Set("source.storage_uri", "oci://n/test-src-namespace/b/test-src-bucket/o/models") v.Set("target.storage_uri", "oci://n/test-tgt-namespace/b/test-tgt-bucket/o/models") @@ -235,6 +236,7 @@ func TestModuleIntegration(t *testing.T) { require.NoError(t, err) assert.NotNil(t, config) assert.Equal(t, "/test/path", config.LocalPath) + assert.True(t, config.TargetArtifactReuseAllowed) assert.Equal(t, "oci://n/test-src-namespace/b/test-src-bucket/o/models", config.Source.StorageURIStr) assert.Equal(t, mockDataStores[0], config.Source.OCIOSDataStore) assert.Equal(t, mockDataStores[1], config.Target.OCIOSDataStore) diff --git a/internal/ome-agent/replica/replica.go b/internal/ome-agent/replica/replica.go index 00bb52c12..54428a76a 100644 --- a/internal/ome-agent/replica/replica.go +++ b/internal/ome-agent/replica/replica.go @@ -5,10 +5,15 @@ import ( "os" "path/filepath" "strings" + "time" + + "github.com/oracle/oci-go-sdk/v65/objectstorage" "sigs.k8s.io/ome/internal/ome-agent/replica/common" + "sigs.k8s.io/ome/pkg/constants" "sigs.k8s.io/ome/pkg/logging" + "sigs.k8s.io/ome/pkg/ociobjectstore" "sigs.k8s.io/ome/pkg/utils/storage" ) @@ -17,6 +22,33 @@ const ( SourceStorageConfigKeyName = "source" TargetStorageConfigKeyName = "target" + + targetArtifactLockPollInterval = 30 * time.Second + targetArtifactLockReleaseMaxAttempts = 3 + targetArtifactLockReleaseRetryInterval = 2 * time.Second + defaultArtifactUploadLockWaitTimeout = 120 * time.Hour +) + +var ( + newReplicatorFunc = NewReplicator + uploadCompletionMarkerFunc = func(dataStore *ociobjectstore.OCIOSDataStore, source string, target ociobjectstore.ObjectURI) error { + return dataStore.Upload(source, target) + } + tryAcquireArtifactUploadLockFunc = func(dataStore *ociobjectstore.OCIOSDataStore, source string, target ociobjectstore.ObjectURI) (string, bool, error) { + return dataStore.UploadIfAbsentWithETag(source, target) + } + releaseArtifactUploadLockFunc = func(dataStore *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI, etag string) (bool, error) { + return dataStore.DeleteObjectIfMatch(target, etag) + } + deleteArtifactCompletionMarkerFunc = func(dataStore *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI) error { + return dataStore.DeleteObject(target) + } + deleteStaleArtifactUploadLockFunc = func(dataStore *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI, etag string) (bool, error) { + return dataStore.DeleteObjectIfMatch(target, etag) + } + targetArtifactStateFunc = defaultTargetArtifactState + sleepFunc = time.Sleep + nowFunc = time.Now ) type ReplicaAgent struct { @@ -25,6 +57,21 @@ type ReplicaAgent struct { ReplicationInput common.ReplicationInput } +type targetArtifactState struct { + Complete bool + CompletionMarked bool + UploadLocked bool + UploadLockModifiedTime *time.Time + UploadLockETag string + ArtifactSizeBytes *int64 +} + +// targetArtifactUploadLock identifies the specific Object Storage generation +// acquired by this replica agent. Release must be conditional on this ETag. +type targetArtifactUploadLock struct { + ETag string +} + // NewReplicaAgent constructs a new replica agent from the given configuration. func NewReplicaAgent(config *Config) (*ReplicaAgent, error) { sourceStorageType, err := storage.GetStorageType(config.Source.StorageURIStr) @@ -79,9 +126,47 @@ func NewReplicaAgent(config *Config) (*ReplicaAgent, error) { } // Start initiates the replication process. -func (r *ReplicaAgent) Start() error { +func (r *ReplicaAgent) Start() (returnErr error) { r.Logger.Infof("Start replication from %s %v to %s %v with checksum config %+v", r.ReplicationInput.SourceStorageType, r.ReplicationInput.Source, r.ReplicationInput.TargetStorageType, r.ReplicationInput.Target, r.Config.Target.ChecksumConfig) + if r.Config.NumConnections <= 0 { + err := fmt.Errorf("num_connections must be greater than 0") + r.writeTerminationLog(err.Error()) + return err + } + + uploadLock, skipReplication, err := r.prepareTargetArtifactUpload() + if err != nil { + r.writeTerminationLog(err.Error()) + return err + } + if uploadLock != nil { + defer func() { + if releaseErr := r.releaseTargetArtifactUploadLock(*uploadLock); releaseErr != nil { + if returnErr != nil { + r.Logger.Errorf("Failed to release target artifact upload lock after replication error: %v", releaseErr) + return + } + + returnErr = releaseErr + r.writeTerminationLog(returnErr.Error()) + } + }() + } + if skipReplication { + return nil + } + if uploadLock != nil { + skipReplication, err = r.prepareTargetArtifactAfterLockAcquired() + if err != nil { + r.writeTerminationLog(err.Error()) + return err + } + if skipReplication { + return nil + } + } + sourceObjs, err := r.listSourceObjects() if err != nil { r.writeTerminationLog(err.Error()) @@ -90,17 +175,357 @@ func (r *ReplicaAgent) Start() error { r.validateModelSize(sourceObjs) - replicatorImp, err := NewReplicator(r) + replicatorImp, err := newReplicatorFunc(r) if err != nil { r.writeTerminationLog(err.Error()) return err } + if err = r.removeTargetArtifactCompletionMarkerBeforeOverwrite(); err != nil { + r.writeTerminationLog(err.Error()) + return err + } + err = replicatorImp.Replicate(sourceObjs) if err != nil { r.writeTerminationLog(err.Error()) + return err + } + + if err = r.writeCompletionMarker(); err != nil { + err = fmt.Errorf("failed to write target artifact completion marker: %w", err) + r.writeTerminationLog(err.Error()) + return err + } + + return nil +} + +func (r *ReplicaAgent) prepareTargetArtifactUpload() (*targetArtifactUploadLock, bool, error) { + if !r.Config.TargetArtifactReuseAllowed { + return nil, false, nil + } + if r.ReplicationInput.TargetStorageType != storage.StorageTypeOCI { + return nil, false, nil + } + if r.Config.Target.OCIOSDataStore == nil { + return nil, false, fmt.Errorf("target OCI object store data store is nil") + } + waitDeadline := nowFunc().Add(r.targetArtifactUploadLockTimeout()) + + state, err := r.targetArtifactState() + if err != nil { + return nil, false, fmt.Errorf("failed to inspect target artifact state: %w", err) + } + if state.Complete { + // The completion marker is written only after every artifact object has + // uploaded successfully. A retained lock from cleanup must not delay + // reuse of that complete artifact. + r.Logger.Infof("Target artifact is already complete; skipping replication") + r.logTargetArtifactSize(state) + return nil, true, nil + } + + for { + uploadLock, err := r.acquireTargetArtifactUploadLock() + if err != nil { + return nil, false, err + } + if uploadLock != nil { + return uploadLock, false, nil + } + + r.Logger.Infof("Target artifact upload lock already exists; waiting for completion marker") + state, err = r.waitForTargetArtifactStateChange(waitDeadline) + if err != nil { + return nil, false, err + } + if state.Complete { + r.Logger.Infof("Target artifact completed while waiting for upload lock; skipping replication") + r.logTargetArtifactSize(state) + return nil, true, nil + } + if r.isTargetArtifactUploadLockStale(state) { + if err := r.deleteStaleTargetArtifactUploadLock(state); err != nil { + return nil, false, err + } + continue + } + } +} + +func (r *ReplicaAgent) logTargetArtifactSize(state targetArtifactState) { + if state.ArtifactSizeBytes == nil { + r.Logger.Infof("Target artifact is complete but artifact size is unavailable") + return + } + r.Logger.Infof("Total model size: %d bytes", *state.ArtifactSizeBytes) +} + +func (r *ReplicaAgent) acquireTargetArtifactUploadLock() (*targetArtifactUploadLock, error) { + lockURI := r.targetArtifactUploadLockURI() + r.Logger.Infof("Acquiring target artifact upload lock at oci://n/%s/b/%s/o/%s", lockURI.Namespace, lockURI.BucketName, lockURI.ObjectName) + etag, acquired, err := tryAcquireArtifactUploadLockFunc( + r.Config.Target.OCIOSDataStore, + constants.ArtifactUploadLockBody, + lockURI, + ) + if err != nil { + return nil, fmt.Errorf("failed to acquire target artifact upload lock: %w", err) + } + if !acquired { + return nil, nil + } + if etag == "" { + return nil, fmt.Errorf("acquired target artifact upload lock without an etag") + } + return &targetArtifactUploadLock{ETag: etag}, nil +} + +func (r *ReplicaAgent) releaseTargetArtifactUploadLock(uploadLock targetArtifactUploadLock) error { + lockURI := r.targetArtifactUploadLockURI() + for attempt := 1; attempt <= targetArtifactLockReleaseMaxAttempts; attempt++ { + released, err := releaseArtifactUploadLockFunc(r.Config.Target.OCIOSDataStore, lockURI, uploadLock.ETag) + if err == nil { + if !released { + r.Logger.Infof("Target artifact upload lock changed before release; leaving the current owner lock in place") + } + return nil + } + + if attempt == targetArtifactLockReleaseMaxAttempts { + releaseErr := fmt.Errorf("failed to release target artifact upload lock at oci://n/%s/b/%s/o/%s after %d attempts: %w", lockURI.Namespace, lockURI.BucketName, lockURI.ObjectName, attempt, err) + r.Logger.Errorf("%v", releaseErr) + return releaseErr + } + + r.Logger.Errorf("Failed to release target artifact upload lock at oci://n/%s/b/%s/o/%s (attempt %d/%d): %v; retrying", lockURI.Namespace, lockURI.BucketName, lockURI.ObjectName, attempt, targetArtifactLockReleaseMaxAttempts, err) + sleepFunc(targetArtifactLockReleaseRetryInterval) + } + + return nil +} + +func (r *ReplicaAgent) waitForTargetArtifactStateChange(deadline time.Time) (targetArtifactState, error) { + for { + remaining := deadline.Sub(nowFunc()) + if remaining <= 0 { + return targetArtifactState{}, fmt.Errorf("timed out waiting for target artifact completion marker") + } + pollInterval := targetArtifactLockPollInterval + if remaining < pollInterval { + pollInterval = remaining + } + sleepFunc(pollInterval) + + state, err := r.targetArtifactState() + if err != nil { + return targetArtifactState{}, fmt.Errorf("failed to inspect target artifact state while waiting for upload lock: %w", err) + } + if state.Complete || !state.UploadLocked || r.isTargetArtifactUploadLockStale(state) { + return state, nil + } + } +} + +func (r *ReplicaAgent) targetArtifactUploadLockTimeout() time.Duration { + if r.Config.ArtifactUploadLockTimeout > 0 { + return r.Config.ArtifactUploadLockTimeout + } + return defaultArtifactUploadLockWaitTimeout +} + +func (r *ReplicaAgent) isTargetArtifactUploadLockStale(state targetArtifactState) bool { + if !state.UploadLocked || state.UploadLockModifiedTime == nil { + return false + } + return !nowFunc().Before(state.UploadLockModifiedTime.Add(r.targetArtifactUploadLockTimeout())) +} + +func (r *ReplicaAgent) deleteStaleTargetArtifactUploadLock(state targetArtifactState) error { + lockURI := r.targetArtifactUploadLockURI() + modifiedAt := "unknown" + if state.UploadLockModifiedTime != nil { + modifiedAt = state.UploadLockModifiedTime.Format(time.RFC3339) + } + if state.UploadLockETag == "" { + return fmt.Errorf("cannot delete stale target artifact upload lock without etag") + } + r.Logger.Infof( + "Deleting stale target artifact upload lock at oci://n/%s/b/%s/o/%s; modifiedAt=%s timeout=%s", + lockURI.Namespace, + lockURI.BucketName, + lockURI.ObjectName, + modifiedAt, + r.targetArtifactUploadLockTimeout(), + ) + deleted, err := deleteStaleArtifactUploadLockFunc(r.Config.Target.OCIOSDataStore, lockURI, state.UploadLockETag) + if err != nil { + return fmt.Errorf("failed to delete stale target artifact upload lock: %w", err) + } + if !deleted { + r.Logger.Infof("Stale target artifact upload lock changed before deletion; retrying lock acquisition") + } + return nil +} + +func (r *ReplicaAgent) targetArtifactState() (targetArtifactState, error) { + return targetArtifactStateFunc(r.Config.Target.OCIOSDataStore, r.targetArtifactPrefixURI()) +} + +func defaultTargetArtifactState(dataStore *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI) (targetArtifactState, error) { + objects, err := dataStore.ListObjects(target) + if err != nil { + return targetArtifactState{}, err + } + + completeMarkerName := normalizeObjectPrefix(target.Prefix) + constants.ArtifactCompleteMarkerFileName + uploadLockName := normalizeObjectPrefix(target.Prefix) + constants.ArtifactUploadLockFileName + var hasCompleteMarker bool + var hasArtifactObject bool + var hasUploadLock bool + var artifactSizeBytes int64 + state := targetArtifactState{} + for _, object := range objects { + if object.Name == nil { + continue + } + switch *object.Name { + case completeMarkerName: + hasCompleteMarker = true + case uploadLockName: + hasUploadLock = true + stateTime := objectSummaryTime(object) + if stateTime != nil { + lockModifiedTime := *stateTime + state.UploadLockModifiedTime = &lockModifiedTime + } + if object.Etag != nil { + state.UploadLockETag = *object.Etag + } + default: + if !constants.IsInternalArtifactObjectName(*object.Name) { + hasArtifactObject = true + if object.Size != nil { + artifactSizeBytes += *object.Size + } + } + } + } + + state.CompletionMarked = hasCompleteMarker + state.Complete = hasCompleteMarker && hasArtifactObject + state.UploadLocked = hasUploadLock + if artifactSizeBytes > 0 { + state.ArtifactSizeBytes = &artifactSizeBytes + } + return state, nil +} + +func objectSummaryTime(object objectstorage.ObjectSummary) *time.Time { + if object.TimeModified != nil { + return &object.TimeModified.Time + } + if object.TimeCreated != nil { + return &object.TimeCreated.Time } - return err + return nil +} + +func (r *ReplicaAgent) prepareTargetArtifactAfterLockAcquired() (bool, error) { + if !r.Config.TargetArtifactReuseAllowed { + return false, nil + } + if r.ReplicationInput.TargetStorageType != storage.StorageTypeOCI { + return false, nil + } + if r.Config.Target.OCIOSDataStore == nil { + return false, fmt.Errorf("target OCI object store data store is nil") + } + + state, err := r.targetArtifactState() + if err != nil { + return false, fmt.Errorf("failed to inspect target artifact state before upload: %w", err) + } + if state.Complete { + r.Logger.Infof("Target artifact completed before upload started; skipping replication") + r.logTargetArtifactSize(state) + return true, nil + } + return false, nil +} + +// removeTargetArtifactCompletionMarkerBeforeOverwrite ensures readers cannot +// treat a target prefix as complete while replication overwrites its objects. +func (r *ReplicaAgent) removeTargetArtifactCompletionMarkerBeforeOverwrite() error { + if r.ReplicationInput.TargetStorageType != storage.StorageTypeOCI { + return nil + } + if r.Config.Target.OCIOSDataStore == nil { + return fmt.Errorf("target OCI object store data store is nil") + } + + markerURI := r.targetArtifactCompleteMarkerURI() + r.Logger.Infof("Deleting target artifact completion marker before upload at oci://n/%s/b/%s/o/%s", markerURI.Namespace, markerURI.BucketName, markerURI.ObjectName) + if err := deleteArtifactCompletionMarkerFunc(r.Config.Target.OCIOSDataStore, markerURI); err != nil { + return fmt.Errorf("failed to delete target artifact completion marker before upload: %w", err) + } + return nil +} + +func (r *ReplicaAgent) writeCompletionMarker() error { + if !r.Config.TargetArtifactReuseAllowed { + return nil + } + if r.ReplicationInput.TargetStorageType != storage.StorageTypeOCI { + r.Logger.Infof("Skipping target artifact completion marker for non-OCI target storage type %s", r.ReplicationInput.TargetStorageType) + return nil + } + if r.Config.Target.OCIOSDataStore == nil { + return fmt.Errorf("target OCI object store data store is nil") + } + + markerURI := r.targetArtifactCompleteMarkerURI() + r.Logger.Infof("Writing target artifact completion marker to oci://n/%s/b/%s/o/%s", markerURI.Namespace, markerURI.BucketName, markerURI.ObjectName) + return uploadCompletionMarkerFunc( + r.Config.Target.OCIOSDataStore, + constants.ArtifactCompleteMarkerBody, + markerURI, + ) +} + +func (r *ReplicaAgent) targetArtifactCompleteMarkerURI() ociobjectstore.ObjectURI { + return ociobjectstore.ObjectURI{ + Namespace: r.ReplicationInput.Target.Namespace, + BucketName: r.ReplicationInput.Target.BucketName, + ObjectName: normalizeObjectPrefix(r.ReplicationInput.Target.Prefix) + constants.ArtifactCompleteMarkerFileName, + Region: r.ReplicationInput.Target.Region, + } +} + +func (r *ReplicaAgent) targetArtifactPrefixURI() ociobjectstore.ObjectURI { + return ociobjectstore.ObjectURI{ + Namespace: r.ReplicationInput.Target.Namespace, + BucketName: r.ReplicationInput.Target.BucketName, + Prefix: normalizeObjectPrefix(r.ReplicationInput.Target.Prefix), + Region: r.ReplicationInput.Target.Region, + } +} + +func (r *ReplicaAgent) targetArtifactUploadLockURI() ociobjectstore.ObjectURI { + return ociobjectstore.ObjectURI{ + Namespace: r.ReplicationInput.Target.Namespace, + BucketName: r.ReplicationInput.Target.BucketName, + ObjectName: normalizeObjectPrefix(r.ReplicationInput.Target.Prefix) + constants.ArtifactUploadLockFileName, + Region: r.ReplicationInput.Target.Region, + } +} + +func normalizeObjectPrefix(prefix string) string { + if prefix == "" || strings.HasSuffix(prefix, "/") { + return prefix + } + return prefix + "/" } func (r *ReplicaAgent) writeTerminationLog(message string) { @@ -132,8 +557,10 @@ func (r *ReplicaAgent) listSourceObjects() ([]common.ReplicationObject, error) { if err != nil { return nil, err } - r.Logger.Infof("Listed %d model weight objects under prefix %s", len(listOfObjectSummary), r.ReplicationInput.Source.Prefix) - return common.ConvertToReplicationObjectsFromObjectSummary(listOfObjectSummary), nil + sourceObjects := common.ConvertToReplicationObjectsFromObjectSummary(listOfObjectSummary) + sourceObjects = filterInternalArtifactReplicationObjects(sourceObjects) + r.Logger.Infof("Listed %d model weight objects under prefix %s", len(sourceObjects), r.ReplicationInput.Source.Prefix) + return sourceObjects, nil case storage.StorageTypeHuggingFace: repoFiles, err := r.Config.Source.HubClient.ListFiles(r.ReplicationInput.Source.BucketName, r.ReplicationInput.Source.Prefix) if err != nil { @@ -154,6 +581,17 @@ func (r *ReplicaAgent) listSourceObjects() ([]common.ReplicationObject, error) { } } +func filterInternalArtifactReplicationObjects(objects []common.ReplicationObject) []common.ReplicationObject { + filtered := make([]common.ReplicationObject, 0, len(objects)) + for _, object := range objects { + if constants.IsInternalArtifactObjectName(object.GetName()) { + continue + } + filtered = append(filtered, object) + } + return filtered +} + func (r *ReplicaAgent) validateModelSize(objects []common.ReplicationObject) { r.Logger.Info("Calculating model size from source") diff --git a/internal/ome-agent/replica/replica_test.go b/internal/ome-agent/replica/replica_test.go index a010f0776..626a4c2d3 100644 --- a/internal/ome-agent/replica/replica_test.go +++ b/internal/ome-agent/replica/replica_test.go @@ -1,18 +1,28 @@ package replica import ( + "errors" + "io" + "net/http" + "os" + "path/filepath" "strings" "testing" + "time" "sigs.k8s.io/ome/pkg/xet" "sigs.k8s.io/ome/internal/ome-agent/replica/common" + "sigs.k8s.io/ome/internal/ome-agent/replica/replicator" + ociCommon "github.com/oracle/oci-go-sdk/v65/common" "github.com/oracle/oci-go-sdk/v65/objectstorage" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" "sigs.k8s.io/ome/pkg/afero" + "sigs.k8s.io/ome/pkg/constants" "sigs.k8s.io/ome/pkg/ociobjectstore" "sigs.k8s.io/ome/pkg/principals" testingPkg "sigs.k8s.io/ome/pkg/testing" @@ -25,6 +35,25 @@ type TestReplicaAgent struct { mockValidateModelSize func(objects []common.ReplicationObject) } +type targetArtifactStateRequestDispatcher struct { + listResponse string +} + +func (d targetArtifactStateRequestDispatcher) Do(request *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(d.listResponse)), + Request: request, + }, nil +} + +type targetArtifactStateRequestSigner struct{} + +func (targetArtifactStateRequestSigner) Sign(*http.Request) error { + return nil +} + // Override Start method to use the mock func (t *TestReplicaAgent) Start() error { t.Logger.Infof("Start replication from %+v to %+v", t.ReplicationInput.Source, t.ReplicationInput.Target) @@ -57,6 +86,33 @@ func createMockOCIOSDataStore() *ociobjectstore.OCIOSDataStore { } } +func newTargetArtifactStateDataStore(listResponse string) *ociobjectstore.OCIOSDataStore { + return &ociobjectstore.OCIOSDataStore{ + Client: &objectstorage.ObjectStorageClient{ + BaseClient: ociCommon.BaseClient{ + HTTPClient: targetArtifactStateRequestDispatcher{listResponse: listResponse}, + Signer: targetArtifactStateRequestSigner{}, + Host: "https://objectstorage.test", + UserAgent: "replica-test", + }, + }, + } +} + +type fakeReplicator struct { + err error + objects []common.ReplicationObject + onReplicate func(objects []common.ReplicationObject) error +} + +func (f *fakeReplicator) Replicate(objects []common.ReplicationObject) error { + f.objects = objects + if f.onReplicate != nil { + return f.onReplicate(objects) + } + return f.err +} + func TestNewReplicaAgent(t *testing.T) { mockLogger := testingPkg.SetupMockLogger() @@ -614,3 +670,787 @@ func TestReplicaAgent_Start(t *testing.T) { err := testAgent.Start() assert.NoError(t, err) } + +func TestReplicaAgent_StartReturnsErrorWhenNumConnectionsInvalid(t *testing.T) { + mockLogger := testingPkg.SetupMockLogger() + agent := &ReplicaAgent{ + Logger: mockLogger, + Config: Config{ + AnotherLogger: mockLogger, + NumConnections: 0, + }, + } + + err := agent.Start() + require.Error(t, err) + assert.Contains(t, err.Error(), "num_connections") +} + +func TestReplicaAgent_StartWritesCompletionMarkerAfterSuccessfulOCITargetReplication(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + + fake := &fakeReplicator{} + newReplicatorFunc = func(_ *ReplicaAgent) (replicator.Replicator, error) { + return fake, nil + } + + var markerSource string + var markerTarget ociobjectstore.ObjectURI + uploadCompletionMarkerFunc = func(_ *ociobjectstore.OCIOSDataStore, source string, target ociobjectstore.ObjectURI) error { + markerSource = source + markerTarget = target + return nil + } + + err := agent.Start() + require.NoError(t, err) + require.Len(t, fake.objects, 1) + assert.Equal(t, constants.ArtifactCompleteMarkerBody, markerSource) + assert.Equal(t, "tgt-ns", markerTarget.Namespace) + assert.Equal(t, "tgt-bucket", markerTarget.BucketName) + assert.Equal(t, "target-models/"+constants.ArtifactCompleteMarkerFileName, markerTarget.ObjectName) + assert.Equal(t, "us-ashburn-1", markerTarget.Region) +} + +func TestReplicaAgent_StartSkipsCompletionMarkerWhenReplicationFails(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + + newReplicatorFunc = func(_ *ReplicaAgent) (replicator.Replicator, error) { + return &fakeReplicator{err: errors.New("replication failed")}, nil + } + + markerWritten := false + uploadCompletionMarkerFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) error { + markerWritten = true + return nil + } + + err := agent.Start() + require.Error(t, err) + assert.Contains(t, err.Error(), "replication failed") + assert.False(t, markerWritten) +} + +func TestReplicaAgent_StartReturnsErrorWhenCompletionMarkerUploadFails(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + + newReplicatorFunc = func(_ *ReplicaAgent) (replicator.Replicator, error) { + return &fakeReplicator{}, nil + } + uploadCompletionMarkerFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) error { + return errors.New("marker upload failed") + } + + err := agent.Start() + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to write target artifact completion marker") + assert.Contains(t, err.Error(), "marker upload failed") +} + +func TestReplicaAgent_StartSkipsReplicationWhenTargetArtifactUploadLockCompletes(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + + replicatorCalled := false + newReplicatorFunc = func(_ *ReplicaAgent) (replicator.Replicator, error) { + replicatorCalled = true + return &fakeReplicator{}, nil + } + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (string, bool, error) { + return "", false, nil + } + stateCalls := 0 + targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { + stateCalls++ + if stateCalls == 1 { + return targetArtifactState{}, nil + } + if stateCalls == 2 { + return targetArtifactState{Complete: true, CompletionMarked: true, UploadLocked: true}, nil + } + t.Fatal("waiter should reuse immediately after the completion marker appears") + return targetArtifactState{}, nil + } + + err := agent.Start() + require.NoError(t, err) + assert.False(t, replicatorCalled) + assert.Equal(t, 2, stateCalls) +} + +func TestReplicaAgent_StartSkipsCompletedTargetArtifactWhenUploadLockRemains(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + + replicatorCalled := false + newReplicatorFunc = func(_ *ReplicaAgent) (replicator.Replicator, error) { + replicatorCalled = true + return &fakeReplicator{}, nil + } + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (string, bool, error) { + t.Fatal("completed target artifact should be reused without acquiring the retained upload lock") + return "", false, nil + } + targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { + return targetArtifactState{Complete: true, CompletionMarked: true, UploadLocked: true}, nil + } + + err := agent.Start() + require.NoError(t, err) + assert.False(t, replicatorCalled) +} + +func TestReplicaAgent_StartSkipsCompletedTargetArtifactAfterLockAcquired(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + + replicatorCalled := false + newReplicatorFunc = func(_ *ReplicaAgent) (replicator.Replicator, error) { + replicatorCalled = true + return &fakeReplicator{}, nil + } + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (string, bool, error) { + return "acquired-lock-etag", true, nil + } + released := false + releaseArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI, etag string) (bool, error) { + released = true + assert.Equal(t, "target-models/"+constants.ArtifactUploadLockFileName, target.ObjectName) + assert.Equal(t, "acquired-lock-etag", etag) + return true, nil + } + deleteArtifactCompletionMarkerFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) error { + t.Fatal("completion marker should not be deleted when reuse is allowed and target is complete") + return nil + } + stateCalls := 0 + targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { + stateCalls++ + if stateCalls == 1 { + return targetArtifactState{}, nil + } + return targetArtifactState{Complete: true, CompletionMarked: true}, nil + } + + err := agent.Start() + require.NoError(t, err) + assert.False(t, replicatorCalled) + assert.True(t, released) + assert.Equal(t, 2, stateCalls) +} + +func TestReplicaAgent_StartSkipsCompletedTargetArtifactEvenWhenUploadLockIsStale(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + agent.Config.ArtifactUploadLockTimeout = time.Hour + + now := time.Date(2026, 7, 10, 1, 0, 0, 0, time.UTC) + nowFunc = func() time.Time { return now } + + replicatorCalled := false + newReplicatorFunc = func(_ *ReplicaAgent) (replicator.Replicator, error) { + replicatorCalled = true + return &fakeReplicator{}, nil + } + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, target ociobjectstore.ObjectURI) (string, bool, error) { + assert.Equal(t, "target-models/"+constants.ArtifactUploadLockFileName, target.ObjectName) + return "", false, nil + } + + staleLockModifiedTime := now.Add(-2 * time.Hour) + stateCalls := 0 + targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { + stateCalls++ + if stateCalls == 1 { + return targetArtifactState{}, nil + } + return targetArtifactState{ + Complete: true, + UploadLocked: true, + UploadLockModifiedTime: &staleLockModifiedTime, + }, nil + } + deleteStaleArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI, _ string) (bool, error) { + t.Fatal("stale upload lock should not be deleted before skipping a complete target artifact") + return false, nil + } + + err := agent.Start() + require.NoError(t, err) + assert.False(t, replicatorCalled) +} + +func TestReplicaAgent_StartDeletesStaleTargetArtifactUploadLockAndReplicates(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + agent.Config.ArtifactUploadLockTimeout = time.Hour + + now := time.Date(2026, 7, 10, 1, 0, 0, 0, time.UTC) + nowFunc = func() time.Time { return now } + + fake := &fakeReplicator{} + newReplicatorFunc = func(_ *ReplicaAgent) (replicator.Replicator, error) { + return fake, nil + } + + acquireAttempts := 0 + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, target ociobjectstore.ObjectURI) (string, bool, error) { + assert.Equal(t, "target-models/"+constants.ArtifactUploadLockFileName, target.ObjectName) + acquireAttempts++ + if acquireAttempts > 1 { + return "replacement-lock-etag", true, nil + } + return "", false, nil + } + + stateCalls := 0 + staleLockModifiedTime := now.Add(-2 * time.Hour) + targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { + stateCalls++ + if stateCalls == 1 { + return targetArtifactState{}, nil + } + return targetArtifactState{ + UploadLocked: true, + UploadLockModifiedTime: &staleLockModifiedTime, + UploadLockETag: "stale-lock-etag", + }, nil + } + + deletedLock := false + deleteStaleArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI, etag string) (bool, error) { + assert.Equal(t, "target-models/"+constants.ArtifactUploadLockFileName, target.ObjectName) + assert.Equal(t, "stale-lock-etag", etag) + deletedLock = true + return true, nil + } + uploadCompletionMarkerFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) error { + return nil + } + + err := agent.Start() + require.NoError(t, err) + require.Len(t, fake.objects, 1) + assert.Equal(t, 2, acquireAttempts) + assert.True(t, deletedLock) +} + +func TestReplicaAgent_StartLogsTargetArtifactSizeWhenSkippingCompletedTargetArtifact(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + + replicatorCalled := false + newReplicatorFunc = func(_ *ReplicaAgent) (replicator.Replicator, error) { + replicatorCalled = true + return &fakeReplicator{}, nil + } + artifactSizeBytes := int64(123456789) + targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { + return targetArtifactState{ + Complete: true, + ArtifactSizeBytes: &artifactSizeBytes, + }, nil + } + + err := agent.Start() + require.NoError(t, err) + assert.False(t, replicatorCalled) + mockLogger := agent.Logger.(*testingPkg.MockLogger) + mockLogger.AssertCalled(t, "Infof", "Total model size: %d bytes", []interface{}{artifactSizeBytes}) +} + +func TestReplicaAgent_StartDeletesCompletionMarkerBeforeOverwritingWhenReuseNotAllowed(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + agent.Config.TargetArtifactReuseAllowed = false + + operations := make([]string, 0, 2) + fake := &fakeReplicator{ + onReplicate: func(objects []common.ReplicationObject) error { + operations = append(operations, "replicate") + return nil + }, + } + newReplicatorFunc = func(_ *ReplicaAgent) (replicator.Replicator, error) { + return fake, nil + } + targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { + t.Fatal("target artifact state should not be inspected when reuse is disabled") + return targetArtifactState{}, nil + } + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (string, bool, error) { + t.Fatal("target artifact upload lock should not be acquired when reuse is disabled") + return "", false, nil + } + deleteArtifactCompletionMarkerFunc = func(_ *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI) error { + assert.Equal(t, "target-models/"+constants.ArtifactCompleteMarkerFileName, target.ObjectName) + operations = append(operations, "delete-completion-marker") + return nil + } + uploadCompletionMarkerFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) error { + t.Fatal("completion marker should not be written when reuse is disabled") + return nil + } + + err := agent.Start() + require.NoError(t, err) + require.Len(t, fake.objects, 1) + assert.Equal(t, []string{"delete-completion-marker", "replicate"}, operations) +} + +func TestReplicaAgent_StartDeletesStaleCompletionMarkerBeforeUploadingTargetArtifact(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + + operations := make([]string, 0, 3) + newReplicatorFunc = func(_ *ReplicaAgent) (replicator.Replicator, error) { + return &fakeReplicator{ + onReplicate: func(objects []common.ReplicationObject) error { + operations = append(operations, "replicate") + return nil + }, + }, nil + } + targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { + return targetArtifactState{CompletionMarked: true}, nil + } + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (string, bool, error) { + return "acquired-lock-etag", true, nil + } + deleteArtifactCompletionMarkerFunc = func(_ *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI) error { + operations = append(operations, "delete-completion-marker") + assert.Equal(t, "target-models/"+constants.ArtifactCompleteMarkerFileName, target.ObjectName) + return nil + } + uploadCompletionMarkerFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, target ociobjectstore.ObjectURI) error { + operations = append(operations, "write-completion-marker") + assert.Equal(t, "target-models/"+constants.ArtifactCompleteMarkerFileName, target.ObjectName) + return nil + } + + err := agent.Start() + require.NoError(t, err) + assert.Equal(t, []string{"delete-completion-marker", "replicate", "write-completion-marker"}, operations) +} + +func TestReplicaAgent_PrepareTargetArtifactUploadSkipsWhenReuseDisabled(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + agent.Config.TargetArtifactReuseAllowed = false + + targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { + t.Fatal("target artifact state should not be inspected when reuse is disabled") + return targetArtifactState{}, nil + } + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (string, bool, error) { + t.Fatal("target artifact upload lock should not be acquired when reuse is disabled") + return "", false, nil + } + + uploadLock, skipReplication, err := agent.prepareTargetArtifactUpload() + require.NoError(t, err) + assert.Nil(t, uploadLock) + assert.False(t, skipReplication) +} + +func TestReplicaAgent_PrepareTargetArtifactUploadUsesOneWaitDeadline(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + agent.Config.ArtifactUploadLockTimeout = time.Hour + + current := time.Date(2026, 7, 10, 1, 0, 0, 0, time.UTC) + nowFunc = func() time.Time { return current } + sleepFunc = func(time.Duration) { + current = current.Add(time.Hour) + } + acquireAttempts := 0 + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (string, bool, error) { + acquireAttempts++ + if acquireAttempts > 2 { + return "", false, errors.New("wait deadline was reset") + } + return "", false, nil + } + targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { + return targetArtifactState{}, nil + } + + _, _, err := agent.prepareTargetArtifactUpload() + require.Error(t, err) + assert.Contains(t, err.Error(), "timed out waiting for target artifact completion marker") + assert.Equal(t, 2, acquireAttempts) +} + +func TestReplicaAgent_WaitForTargetArtifactStateChangeDoesNotSleepPastDeadline(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + agent.Config.ArtifactUploadLockTimeout = 5 * time.Second + + current := time.Date(2026, 7, 10, 1, 0, 0, 0, time.UTC) + nowFunc = func() time.Time { return current } + sleepDurations := make([]time.Duration, 0, 1) + sleepFunc = func(duration time.Duration) { + sleepDurations = append(sleepDurations, duration) + current = current.Add(duration) + } + targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { + return targetArtifactState{UploadLocked: true}, nil + } + + _, err := agent.waitForTargetArtifactStateChange(nowFunc().Add(agent.targetArtifactUploadLockTimeout())) + + require.Error(t, err) + assert.Contains(t, err.Error(), "timed out waiting for target artifact completion marker") + assert.Equal(t, []time.Duration{5 * time.Second}, sleepDurations) +} + +func TestReplicaAgent_PrepareTargetArtifactUploadReclaimsLockThatBecomesStaleAtWaitDeadline(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + agent.Config.ArtifactUploadLockTimeout = time.Hour + + current := time.Date(2026, 7, 10, 1, 0, 0, 0, time.UTC) + nowFunc = func() time.Time { return current } + sleepFunc = func(time.Duration) { + current = current.Add(time.Hour) + } + + acquireAttempts := 0 + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (string, bool, error) { + acquireAttempts++ + if acquireAttempts == 1 { + return "", false, nil + } + return "replacement-lock-etag", true, nil + } + + lockCreatedAt := current + targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { + return targetArtifactState{ + UploadLocked: true, + UploadLockModifiedTime: &lockCreatedAt, + UploadLockETag: "stale-lock-etag", + }, nil + } + + deleted := false + deleteStaleArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI, etag string) (bool, error) { + deleted = true + assert.Equal(t, "stale-lock-etag", etag) + return true, nil + } + + uploadLock, skipReplication, err := agent.prepareTargetArtifactUpload() + + require.NoError(t, err) + require.NotNil(t, uploadLock) + assert.Equal(t, "replacement-lock-etag", uploadLock.ETag) + assert.False(t, skipReplication) + assert.True(t, deleted) + assert.Equal(t, 2, acquireAttempts) +} + +func TestReplicaAgent_StartReleasesTargetArtifactUploadLockWhenReplicationFails(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + + newReplicatorFunc = func(_ *ReplicaAgent) (replicator.Replicator, error) { + return &fakeReplicator{err: errors.New("replication failed")}, nil + } + + released := false + releaseArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI, etag string) (bool, error) { + released = true + assert.Equal(t, "target-models/"+constants.ArtifactUploadLockFileName, target.ObjectName) + assert.Equal(t, "test-lock-etag", etag) + return true, nil + } + + err := agent.Start() + require.Error(t, err) + assert.Contains(t, err.Error(), "replication failed") + assert.True(t, released) +} + +func TestReplicaAgent_StartRetriesTargetArtifactUploadLockRelease(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + + newReplicatorFunc = func(_ *ReplicaAgent) (replicator.Replicator, error) { + return &fakeReplicator{}, nil + } + uploadCompletionMarkerFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) error { + return nil + } + + releaseAttempts := 0 + releaseArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI, _ string) (bool, error) { + releaseAttempts++ + if releaseAttempts == 1 { + return false, errors.New("temporary delete failure") + } + return true, nil + } + + require.NoError(t, agent.Start()) + assert.Equal(t, 2, releaseAttempts) +} + +func TestReplicaAgent_StartReturnsErrorWhenTargetArtifactUploadLockReleaseFails(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + + newReplicatorFunc = func(_ *ReplicaAgent) (replicator.Replicator, error) { + return &fakeReplicator{}, nil + } + uploadCompletionMarkerFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) error { + return nil + } + + releaseAttempts := 0 + releaseArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI, _ string) (bool, error) { + releaseAttempts++ + return false, errors.New("temporary delete failure") + } + + err := agent.Start() + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to release target artifact upload lock") + assert.Equal(t, 3, releaseAttempts) +} + +func TestReplicaAgent_StartPreservesReplicationErrorWhenTargetArtifactUploadLockReleaseFails(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + + newReplicatorFunc = func(_ *ReplicaAgent) (replicator.Replicator, error) { + return &fakeReplicator{err: errors.New("replication failed")}, nil + } + + releaseAttempts := 0 + releaseArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI, _ string) (bool, error) { + releaseAttempts++ + return false, errors.New("temporary delete failure") + } + + err := agent.Start() + require.Error(t, err) + assert.Contains(t, err.Error(), "replication failed") + assert.Equal(t, 3, releaseAttempts) +} + +func TestReplicaAgent_ReleaseTargetArtifactUploadLockOnlyDeletesAcquiredLock(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + + released := false + releaseArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI, etag string) (bool, error) { + released = true + assert.Equal(t, "target-models/"+constants.ArtifactUploadLockFileName, target.ObjectName) + assert.Equal(t, "acquired-lock-etag", etag) + return false, nil + } + + require.NoError(t, agent.releaseTargetArtifactUploadLock(targetArtifactUploadLock{ETag: "acquired-lock-etag"})) + + assert.True(t, released) +} + +func TestReplicaAgent_StartSkipsCompletionMarkerForNonOCITarget(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + + agent.ReplicationInput.TargetStorageType = storage.StorageTypePVC + agent.ReplicationInput.Target = ociobjectstore.ObjectURI{ + BucketName: "target-pvc", + Prefix: "target-model", + } + agent.Config.Target = TargetStruct{ + StorageURIStr: "pvc://target-pvc/target-model", + PVCFileSystem: afero.NewOsFs().(*afero.OsFs), + } + + newReplicatorFunc = func(_ *ReplicaAgent) (replicator.Replicator, error) { + return &fakeReplicator{}, nil + } + markerWritten := false + uploadCompletionMarkerFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) error { + markerWritten = true + return nil + } + + err := agent.Start() + require.NoError(t, err) + assert.False(t, markerWritten) +} + +func TestFilterInternalArtifactReplicationObjectsSkipsCompletionMarker(t *testing.T) { + configName := "models/config.json" + weightName := "models/model.safetensors" + markerName := "models/" + constants.ArtifactCompleteMarkerFileName + lockName := "models/" + constants.ArtifactUploadLockFileName + rootMarkerName := constants.ArtifactCompleteMarkerFileName + size := int64(1) + + objects := []common.ReplicationObject{ + common.ObjectSummaryReplicationObject{ObjectSummary: objectstorage.ObjectSummary{Name: &configName, Size: &size}}, + common.ObjectSummaryReplicationObject{ObjectSummary: objectstorage.ObjectSummary{Name: &markerName, Size: &size}}, + common.ObjectSummaryReplicationObject{ObjectSummary: objectstorage.ObjectSummary{Name: &lockName, Size: &size}}, + common.ObjectSummaryReplicationObject{ObjectSummary: objectstorage.ObjectSummary{Name: &weightName, Size: &size}}, + common.ObjectSummaryReplicationObject{ObjectSummary: objectstorage.ObjectSummary{Name: &rootMarkerName, Size: &size}}, + } + + filtered := filterInternalArtifactReplicationObjects(objects) + + require.Len(t, filtered, 2) + assert.Equal(t, configName, filtered[0].GetName()) + assert.Equal(t, weightName, filtered[1].GetName()) +} + +func TestDefaultTargetArtifactState(t *testing.T) { + target := ociobjectstore.ObjectURI{ + Namespace: "target-namespace", + BucketName: "target-bucket", + Prefix: "models/", + } + + t.Run("complete artifact excludes replication metadata from size", func(t *testing.T) { + dataStore := newTargetArtifactStateDataStore(`{ + "objects": [ + {"name":"models/.ome-artifact-complete","size":1}, + {"name":"models/config.json","size":12}, + {"name":"models/model.safetensors","size":100}, + {"name":"models/.ome-artifact-upload.lock","size":1,"etag":"lock-etag","timeModified":"2026-08-20T12:00:00.000Z"} + ] + }`) + + state, err := defaultTargetArtifactState(dataStore, target) + + require.NoError(t, err) + assert.True(t, state.CompletionMarked) + assert.True(t, state.Complete) + assert.True(t, state.UploadLocked) + assert.Equal(t, "lock-etag", state.UploadLockETag) + require.NotNil(t, state.UploadLockModifiedTime) + assert.Equal(t, time.Date(2026, time.August, 20, 12, 0, 0, 0, time.UTC), *state.UploadLockModifiedTime) + require.NotNil(t, state.ArtifactSizeBytes) + assert.Equal(t, int64(112), *state.ArtifactSizeBytes) + }) + + t.Run("completion marker without model files is incomplete", func(t *testing.T) { + dataStore := newTargetArtifactStateDataStore(`{ + "objects": [ + {"name":"models/.ome-artifact-complete","size":1} + ] + }`) + + state, err := defaultTargetArtifactState(dataStore, target) + + require.NoError(t, err) + assert.True(t, state.CompletionMarked) + assert.False(t, state.Complete) + assert.False(t, state.UploadLocked) + assert.Nil(t, state.ArtifactSizeBytes) + }) + + t.Run("active lock uses creation time when modification time is unavailable", func(t *testing.T) { + dataStore := newTargetArtifactStateDataStore(`{ + "objects": [ + {"name":"models/.ome-artifact-upload.lock","size":1,"etag":"lock-etag","timeCreated":"2026-08-20T12:30:00.000Z"} + ] + }`) + + state, err := defaultTargetArtifactState(dataStore, target) + + require.NoError(t, err) + assert.False(t, state.CompletionMarked) + assert.False(t, state.Complete) + assert.True(t, state.UploadLocked) + assert.Equal(t, "lock-etag", state.UploadLockETag) + require.NotNil(t, state.UploadLockModifiedTime) + assert.Equal(t, time.Date(2026, time.August, 20, 12, 30, 0, 0, time.UTC), *state.UploadLockModifiedTime) + assert.Nil(t, state.ArtifactSizeBytes) + }) +} + +func newTestAgentForCompletionMarker(t *testing.T) (*ReplicaAgent, func()) { + t.Helper() + + oldNewReplicatorFunc := newReplicatorFunc + oldUploadCompletionMarkerFunc := uploadCompletionMarkerFunc + oldTryAcquireArtifactUploadLockFunc := tryAcquireArtifactUploadLockFunc + oldReleaseArtifactUploadLockFunc := releaseArtifactUploadLockFunc + oldDeleteArtifactCompletionMarkerFunc := deleteArtifactCompletionMarkerFunc + oldDeleteStaleArtifactUploadLockFunc := deleteStaleArtifactUploadLockFunc + oldTargetArtifactStateFunc := targetArtifactStateFunc + oldSleepFunc := sleepFunc + oldNowFunc := nowFunc + cleanup := func() { + newReplicatorFunc = oldNewReplicatorFunc + uploadCompletionMarkerFunc = oldUploadCompletionMarkerFunc + tryAcquireArtifactUploadLockFunc = oldTryAcquireArtifactUploadLockFunc + releaseArtifactUploadLockFunc = oldReleaseArtifactUploadLockFunc + deleteArtifactCompletionMarkerFunc = oldDeleteArtifactCompletionMarkerFunc + deleteStaleArtifactUploadLockFunc = oldDeleteStaleArtifactUploadLockFunc + targetArtifactStateFunc = oldTargetArtifactStateFunc + sleepFunc = oldSleepFunc + nowFunc = oldNowFunc + } + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (string, bool, error) { + return "test-lock-etag", true, nil + } + releaseArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI, _ string) (bool, error) { + return true, nil + } + deleteArtifactCompletionMarkerFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) error { + return nil + } + targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { + return targetArtifactState{}, nil + } + sleepFunc = func(time.Duration) {} + nowFunc = time.Now + + localPath := t.TempDir() + sourceDir := filepath.Join(localPath, "source-model") + require.NoError(t, os.MkdirAll(sourceDir, 0755)) + require.NoError(t, os.WriteFile(filepath.Join(sourceDir, "config.json"), []byte("model config"), 0644)) + + mockLogger := testingPkg.SetupMockLogger() + return &ReplicaAgent{ + Logger: mockLogger, + Config: Config{ + AnotherLogger: mockLogger, + LocalPath: localPath, + NumConnections: 1, + DownloadSizeLimitGB: 100, + EnableSizeLimitCheck: true, + TargetArtifactReuseAllowed: true, + Source: SourceStruct{ + StorageURIStr: "pvc://source-pvc/source-model", + PVCFileSystem: afero.NewOsFs().(*afero.OsFs), + }, + Target: TargetStruct{ + StorageURIStr: "oci://n/tgt-ns/b/tgt-bucket/o/target-models", + OCIOSDataStore: createMockOCIOSDataStore(), + }, + }, + ReplicationInput: common.ReplicationInput{ + SourceStorageType: storage.StorageTypePVC, + TargetStorageType: storage.StorageTypeOCI, + Source: ociobjectstore.ObjectURI{ + BucketName: "source-pvc", + Prefix: "source-model", + }, + Target: ociobjectstore.ObjectURI{ + Namespace: "tgt-ns", + BucketName: "tgt-bucket", + Prefix: "target-models/", + Region: "us-ashburn-1", + }, + }, + }, cleanup +} diff --git a/pkg/constants/constants.go b/pkg/constants/constants.go index baef245c2..aca118a6d 100644 --- a/pkg/constants/constants.go +++ b/pkg/constants/constants.go @@ -88,12 +88,13 @@ var ( AgentBaseModelTypeEnvVarKey = AgentAppName + "_" + "MODEL_TYPE" // General Configuration - AgentLocalPathEnvVarKey = AgentAppName + "_" + "LOCAL_PATH" - AgentNumOfGPUEnvVarKey = AgentAppName + "_" + "NUM_OF_GPU" - AgentDisableModelDecryptionEnvVarKey = AgentAppName + "_" + "DISABLE_MODEL_DECRYPTION" - AgentModelBucketNameEnvVarKey = AgentAppName + "_" + "MODEL_BUCKET_NAME" - AgentModelNamespaceEnvVarKey = AgentAppName + "_" + "MODEL_NAMESPACE" - AgentModelObjectName = AgentAppName + "_" + "MODEL_OBJECT_NAME" + AgentLocalPathEnvVarKey = AgentAppName + "_" + "LOCAL_PATH" + AgentNumOfGPUEnvVarKey = AgentAppName + "_" + "NUM_OF_GPU" + AgentDisableModelDecryptionEnvVarKey = AgentAppName + "_" + "DISABLE_MODEL_DECRYPTION" + AgentModelBucketNameEnvVarKey = AgentAppName + "_" + "MODEL_BUCKET_NAME" + AgentModelNamespaceEnvVarKey = AgentAppName + "_" + "MODEL_NAMESPACE" + AgentModelObjectName = AgentAppName + "_" + "MODEL_OBJECT_NAME" + AgentTargetArtifactReuseAllowedEnvVarKey = AgentAppName + "_" + "TARGET_ARTIFACT_REUSE_ALLOWED" // OCI Vault and Security AgentCompartmentIDEnvVarKey = AgentAppName + "_" + "COMPARTMENT_ID" @@ -117,11 +118,27 @@ var ( // Model agent Constants const ( - AgentConfigMapKeyName = "agent" - TensorRTLLM = "tensorrtllm" - HfArtifactConfigMapKeyPrefix = "artifact.huggingface." + AgentConfigMapKeyName = "agent" + TensorRTLLM = "tensorrtllm" + HfArtifactConfigMapKeyPrefix = "artifact.huggingface." + ArtifactCompleteMarkerFileName = ".ome-artifact-complete" + ArtifactCompleteMarkerBody = "complete\n" + ArtifactUploadLockFileName = ".ome-artifact-upload.lock" + ArtifactUploadLockBody = "uploading\n" ) +func IsArtifactCompleteMarkerObjectName(objectName string) bool { + return objectName == ArtifactCompleteMarkerFileName || strings.HasSuffix(objectName, "/"+ArtifactCompleteMarkerFileName) +} + +func IsArtifactUploadLockObjectName(objectName string) bool { + return objectName == ArtifactUploadLockFileName || strings.HasSuffix(objectName, "/"+ArtifactUploadLockFileName) +} + +func IsInternalArtifactObjectName(objectName string) bool { + return IsArtifactCompleteMarkerObjectName(objectName) || IsArtifactUploadLockObjectName(objectName) +} + // InferenceService Annotations var ( DeploymentMode = OMEAPIGroupName + "/deploymentMode" diff --git a/pkg/constants/constants_test.go b/pkg/constants/constants_test.go index fc2afe3be..e9e013a9f 100644 --- a/pkg/constants/constants_test.go +++ b/pkg/constants/constants_test.go @@ -62,3 +62,38 @@ func TestLWSNameTruncates(t *testing.T) { }) } } + +func TestIsArtifactCompleteMarkerObjectName(t *testing.T) { + tests := []struct { + name string + objectName string + want bool + }{ + { + name: "marker at root", + objectName: ArtifactCompleteMarkerFileName, + want: true, + }, + { + name: "marker under model prefix", + objectName: "customer-imported-basemodels/deepseek-ai/DeepSeek-V4-Pro/abc123/" + ArtifactCompleteMarkerFileName, + want: true, + }, + { + name: "regular model file", + objectName: "customer-imported-basemodels/deepseek-ai/DeepSeek-V4-Pro/abc123/config.json", + }, + { + name: "similar suffix without path separator", + objectName: "customer-imported-basemodels/deepseek-ai/DeepSeek-V4-Pro/abc123/not-" + ArtifactCompleteMarkerFileName, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := IsArtifactCompleteMarkerObjectName(tt.objectName); got != tt.want { + t.Fatalf("IsArtifactCompleteMarkerObjectName(%q) = %v, want %v", tt.objectName, got, tt.want) + } + }) + } +} diff --git a/pkg/ociobjectstore/os_data_store.go b/pkg/ociobjectstore/os_data_store.go index 8f69ff645..7802375d3 100644 --- a/pkg/ociobjectstore/os_data_store.go +++ b/pkg/ociobjectstore/os_data_store.go @@ -271,6 +271,37 @@ func (cds *OCIOSDataStore) Download(source ObjectURI, target string, opts ...Dow return nil } +type uploadSourceFile interface { + io.ReadCloser + Stat() (os.FileInfo, error) +} + +type uploadSourceFileOpener func(string) (uploadSourceFile, error) + +func newUploadBody(source string) (io.ReadCloser, *int64, error) { + return newUploadBodyWithOpener(source, func(source string) (uploadSourceFile, error) { + return os.Open(source) + }) +} + +func newUploadBodyWithOpener(source string, openSourceFile uploadSourceFileOpener) (io.ReadCloser, *int64, error) { + sourceFile, err := openSourceFile(source) + if err != nil { + body := io.NopCloser(strings.NewReader(source)) + size := int64(len(source)) + return body, &size, nil + } + + fileInfo, err := sourceFile.Stat() + if err != nil { + _ = sourceFile.Close() + return nil, nil, fmt.Errorf("failed to get source file info %q: %+v", source, err) + } + + size := fileInfo.Size() + return sourceFile, &size, nil +} + // Upload uploads a file (or string content) to OCI Object Storage. // // If the `source` is a file path, the file is read and uploaded. @@ -287,27 +318,11 @@ func (cds *OCIOSDataStore) Upload(source string, target ObjectURI) error { objectFullName := fmt.Sprintf( "%s/%s/%s", target.Namespace, target.BucketName, target.ObjectName) - var putObjectBody io.ReadCloser - var uploadObjectSize *int64 - - // When source is the path of the file which needs to be uploaded - if sourceFile, err := os.Open(source); err == nil { - fileInfo, err := sourceFile.Stat() - if err != nil { - return fmt.Errorf( - "failed to get source file info %q: %+v", - source, - err) - } - putObjectBody = io.NopCloser(sourceFile) - tmp := fileInfo.Size() - uploadObjectSize = &tmp - } else { - // When the source is pure string content that needs to be uploaded - putObjectBody = io.NopCloser(strings.NewReader(source)) - tmp := int64(len(source)) - uploadObjectSize = &tmp + putObjectBody, uploadObjectSize, err := newUploadBody(source) + if err != nil { + return err } + defer putObjectBody.Close() putObjectRequest := objectstorage.PutObjectRequest{ NamespaceName: &target.Namespace, @@ -328,6 +343,129 @@ func (cds *OCIOSDataStore) Upload(source string, target ObjectURI) error { return nil } +// UploadIfAbsent uploads a file or string only when the target object does not +// already exist. It returns false when Object Storage rejects the write because +// another writer has already created the object. +func (cds *OCIOSDataStore) UploadIfAbsent(source string, target ObjectURI) (bool, error) { + _, uploaded, err := cds.UploadIfAbsentWithETag(source, target) + return uploaded, err +} + +// UploadIfAbsentWithETag uploads a file or string only when the target object +// does not already exist. On success it returns the ETag of the object this +// caller created, which can be used to conditionally update or delete it. +func (cds *OCIOSDataStore) UploadIfAbsentWithETag(source string, target ObjectURI) (string, bool, error) { + if target.Namespace == "" { + namespace, err := cds.GetNamespace() + if err != nil { + return "", false, fmt.Errorf("error upload object due to no namespace found: %+v", err) + } + target.Namespace = *namespace + } + + objectFullName := fmt.Sprintf( + "%s/%s/%s", target.Namespace, target.BucketName, target.ObjectName) + + putObjectBody, uploadObjectSize, err := newUploadBody(source) + if err != nil { + return "", false, err + } + defer putObjectBody.Close() + + putObjectRequest := objectstorage.PutObjectRequest{ + NamespaceName: &target.Namespace, + BucketName: &target.BucketName, + ObjectName: &target.ObjectName, + ContentLength: uploadObjectSize, + PutObjectBody: putObjectBody, + IfNoneMatch: common.String("*"), + } + response, err := cds.Client.PutObject(context.Background(), putObjectRequest) + if isObjectAlreadyPresent(response.RawResponse, err) { + return "", false, nil + } + if err != nil || response.RawResponse == nil || response.RawResponse.StatusCode != http.StatusOK { + return "", false, fmt.Errorf( + "failed to put object %q with response %+v: %s", + objectFullName, + response, + errorMessage(err)) + } + if response.ETag == nil || *response.ETag == "" { + return "", false, fmt.Errorf("put object %q succeeded without an ETag", objectFullName) + } + return *response.ETag, true, nil +} + +// DeleteObject removes an object from OCI Object Storage. Missing objects are +// treated as already cleaned up so callers can use this for best-effort cleanup. +func (cds *OCIOSDataStore) DeleteObject(target ObjectURI) error { + if target.Namespace == "" { + namespace, err := cds.GetNamespace() + if err != nil { + return fmt.Errorf("error delete object due to no namespace found: %+v", err) + } + target.Namespace = *namespace + } + + objectFullName := fmt.Sprintf( + "%s/%s/%s", target.Namespace, target.BucketName, target.ObjectName) + deleteObjectRequest := objectstorage.DeleteObjectRequest{ + NamespaceName: &target.Namespace, + BucketName: &target.BucketName, + ObjectName: &target.ObjectName, + } + response, err := cds.Client.DeleteObject(context.Background(), deleteObjectRequest) + if isObjectMissing(err) { + return nil + } + if err != nil || response.RawResponse == nil || response.RawResponse.StatusCode != http.StatusNoContent { + return fmt.Errorf( + "failed to delete object %q with response %+v: %s", + objectFullName, + response, + errorMessage(err)) + } + return nil +} + +// DeleteObjectIfMatch removes an object only when its current ETag matches the +// caller's observed ETag. It returns false when the object is already gone or +// has changed since observation. +func (cds *OCIOSDataStore) DeleteObjectIfMatch(target ObjectURI, etag string) (bool, error) { + if etag == "" { + return false, fmt.Errorf("etag cannot be empty") + } + if target.Namespace == "" { + namespace, err := cds.GetNamespace() + if err != nil { + return false, fmt.Errorf("error delete object due to no namespace found: %+v", err) + } + target.Namespace = *namespace + } + + objectFullName := fmt.Sprintf( + "%s/%s/%s", target.Namespace, target.BucketName, target.ObjectName) + deleteObjectRequest := objectstorage.DeleteObjectRequest{ + NamespaceName: &target.Namespace, + BucketName: &target.BucketName, + ObjectName: &target.ObjectName, + IfMatch: common.String(etag), + } + response, err := cds.Client.DeleteObject(context.Background(), deleteObjectRequest) + if isObjectMissing(err) || isPreconditionFailed(response.RawResponse, err) { + return false, nil + } + if err != nil || response.RawResponse == nil || response.RawResponse.StatusCode != http.StatusNoContent { + return false, fmt.Errorf( + "failed to delete object %q with response %+v: %s", + objectFullName, + response, + errorMessage(err)) + } + return true, nil +} + // HeadObject fetches metadata headers for an object in OCI Object Storage. // // It returns an OCI HeadObjectResponse which contains fields such as size, ETag, and MD5 checksum. @@ -423,7 +561,7 @@ func (cds *OCIOSDataStore) ListObjects(target ObjectURI) ([]objectstorage.Object NamespaceName: &target.Namespace, BucketName: &target.BucketName, Prefix: &target.Prefix, //Virtual folder name within bucket - Fields: common.String("name,size,md5"), + Fields: common.String("name,size,md5,etag,timeCreated,timeModified"), } var allObjects []objectstorage.ObjectSummary @@ -522,3 +660,31 @@ func isMultipartMd5(md5 string) bool { _, err := strconv.Atoi(parts[1]) return err == nil } + +func isObjectAlreadyPresent(response *http.Response, err error) bool { + return isPreconditionFailed(response, err) +} + +func isPreconditionFailed(response *http.Response, err error) bool { + if response != nil && response.StatusCode == http.StatusPreconditionFailed { + return true + } + if serviceErr, ok := common.IsServiceError(err); ok { + return serviceErr.GetHTTPStatusCode() == http.StatusPreconditionFailed + } + return false +} + +func isObjectMissing(err error) bool { + if serviceErr, ok := common.IsServiceError(err); ok { + return serviceErr.GetHTTPStatusCode() == http.StatusNotFound + } + return false +} + +func errorMessage(err error) string { + if err == nil { + return "" + } + return err.Error() +} diff --git a/pkg/ociobjectstore/os_data_store_test.go b/pkg/ociobjectstore/os_data_store_test.go index fb4858577..79cce3f06 100644 --- a/pkg/ociobjectstore/os_data_store_test.go +++ b/pkg/ociobjectstore/os_data_store_test.go @@ -1,13 +1,17 @@ package ociobjectstore import ( + "errors" "fmt" + "net/http" "os" "path/filepath" "strings" "testing" "time" + "github.com/oracle/oci-go-sdk/v65/common" + "github.com/oracle/oci-go-sdk/v65/objectstorage" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -15,6 +19,142 @@ import ( testingPkg "sigs.k8s.io/ome/pkg/testing" ) +type objectStorageRequestDispatcher struct { + do func(*http.Request) (*http.Response, error) +} + +func (d objectStorageRequestDispatcher) Do(request *http.Request) (*http.Response, error) { + return d.do(request) +} + +type objectStorageRequestSigner struct{} + +func (objectStorageRequestSigner) Sign(*http.Request) error { + return nil +} + +func newTestOCIOSDataStore(dispatch func(*http.Request) (*http.Response, error)) *OCIOSDataStore { + return &OCIOSDataStore{ + Client: &objectstorage.ObjectStorageClient{ + BaseClient: common.BaseClient{ + HTTPClient: objectStorageRequestDispatcher{do: dispatch}, + Signer: objectStorageRequestSigner{}, + Host: "https://objectstorage.test", + UserAgent: "ociobjectstore-test", + }, + }, + } +} + +func objectStorageTestResponse(request *http.Request, statusCode int, headers http.Header) *http.Response { + return &http.Response{ + StatusCode: statusCode, + Header: headers, + Body: http.NoBody, + Request: request, + } +} + +type uploadSourceFileStub struct { + *strings.Reader + statErr error + closed bool +} + +func (f *uploadSourceFileStub) Close() error { + f.closed = true + return nil +} + +func (f *uploadSourceFileStub) Stat() (os.FileInfo, error) { + return nil, f.statErr +} + +func TestNewUploadBodyWithOpenerClosesFileWhenStatFails(t *testing.T) { + sourceFile := &uploadSourceFileStub{ + Reader: strings.NewReader("upload body"), + statErr: errors.New("stat failed"), + } + + body, size, err := newUploadBodyWithOpener("/source/file", func(string) (uploadSourceFile, error) { + return sourceFile, nil + }) + + require.Error(t, err) + assert.Nil(t, body) + assert.Nil(t, size) + assert.True(t, sourceFile.closed) +} + +func TestUploadIfAbsentWithETag(t *testing.T) { + target := ObjectURI{Namespace: "namespace", BucketName: "bucket", ObjectName: "prefix/lock"} + + t.Run("uses If-None-Match and returns created ETag", func(t *testing.T) { + dataStore := newTestOCIOSDataStore(func(request *http.Request) (*http.Response, error) { + assert.Equal(t, http.MethodPut, request.Method) + assert.Equal(t, "*", request.Header.Get("If-None-Match")) + return objectStorageTestResponse(request, http.StatusOK, http.Header{"Etag": []string{"lock-etag"}}), nil + }) + + etag, created, err := dataStore.UploadIfAbsentWithETag("lock body", target) + + require.NoError(t, err) + assert.True(t, created) + assert.Equal(t, "lock-etag", etag) + }) + + t.Run("treats precondition failure as already present", func(t *testing.T) { + dataStore := newTestOCIOSDataStore(func(request *http.Request) (*http.Response, error) { + return objectStorageTestResponse(request, http.StatusPreconditionFailed, http.Header{}), nil + }) + + etag, created, err := dataStore.UploadIfAbsentWithETag("lock body", target) + + require.NoError(t, err) + assert.False(t, created) + assert.Empty(t, etag) + }) +} + +func TestDeleteObjectIfMatch(t *testing.T) { + target := ObjectURI{Namespace: "namespace", BucketName: "bucket", ObjectName: "prefix/lock"} + + t.Run("uses If-Match and reports a successful delete", func(t *testing.T) { + dataStore := newTestOCIOSDataStore(func(request *http.Request) (*http.Response, error) { + assert.Equal(t, http.MethodDelete, request.Method) + assert.Equal(t, "lock-etag", request.Header.Get("If-Match")) + return objectStorageTestResponse(request, http.StatusNoContent, http.Header{}), nil + }) + + deleted, err := dataStore.DeleteObjectIfMatch(target, "lock-etag") + + require.NoError(t, err) + assert.True(t, deleted) + }) + + for _, statusCode := range []int{http.StatusNotFound, http.StatusPreconditionFailed} { + t.Run(fmt.Sprintf("treats %d as not deleted", statusCode), func(t *testing.T) { + dataStore := newTestOCIOSDataStore(func(request *http.Request) (*http.Response, error) { + return objectStorageTestResponse(request, statusCode, http.Header{}), nil + }) + + deleted, err := dataStore.DeleteObjectIfMatch(target, "lock-etag") + + require.NoError(t, err) + assert.False(t, deleted) + }) + } +} + +func TestDeleteObjectTreatsMissingObjectAsSuccess(t *testing.T) { + target := ObjectURI{Namespace: "namespace", BucketName: "bucket", ObjectName: "prefix/marker"} + dataStore := newTestOCIOSDataStore(func(request *http.Request) (*http.Response, error) { + return objectStorageTestResponse(request, http.StatusNotFound, http.Header{}), nil + }) + + require.NoError(t, dataStore.DeleteObject(target)) +} + func TestNewOCIOSDataStore(t *testing.T) { t.Run("Nil config", func(t *testing.T) { cds, err := NewOCIOSDataStore(nil)