From 96c2711ccd5099823d21889b6da80541be8d7682 Mon Sep 17 00:00:00 2001 From: Chengjie Yao Date: Tue, 7 Jul 2026 22:33:11 -0700 Subject: [PATCH 1/8] [Core] Add OCI artifact reuse support for ome-agent --- config/ome-agent/ome-agent.yaml | 1 + internal/ome-agent/replica/config.go | 5 +- internal/ome-agent/replica/config_test.go | 22 ++ internal/ome-agent/replica/module_test.go | 2 + internal/ome-agent/replica/replica.go | 340 ++++++++++++++++- internal/ome-agent/replica/replica_test.go | 416 +++++++++++++++++++++ pkg/constants/constants.go | 35 +- pkg/constants/constants_test.go | 35 ++ pkg/ociobjectstore/os_data_store.go | 157 +++++++- 9 files changed, 998 insertions(+), 15 deletions(-) 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..c6c281b58 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,7 @@ func TestDefaultConfig(t *testing.T) { assert.Equal(t, 10, config.NumConnections) assert.Equal(t, 650, config.DownloadSizeLimitGB) assert.Equal(t, true, config.EnableSizeLimitCheck) + 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..9afaa1a04 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,28 @@ const ( SourceStorageConfigKeyName = "source" TargetStorageConfigKeyName = "target" + + targetArtifactLockPollInterval = 30 * 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) (bool, error) { + return dataStore.UploadIfAbsent(source, target) + } + deleteArtifactUploadLockFunc = 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 +52,14 @@ type ReplicaAgent struct { ReplicationInput common.ReplicationInput } +type targetArtifactState struct { + Complete bool + UploadLocked bool + UploadLockModifiedTime *time.Time + UploadLockETag string + ArtifactSizeBytes *int64 +} + // NewReplicaAgent constructs a new replica agent from the given configuration. func NewReplicaAgent(config *Config) (*ReplicaAgent, error) { sourceStorageType, err := storage.GetStorageType(config.Source.StorageURIStr) @@ -82,6 +117,24 @@ func NewReplicaAgent(config *Config) (*ReplicaAgent, error) { func (r *ReplicaAgent) Start() 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 + } + + lockAcquired, skipReplication, err := r.prepareTargetArtifactUpload() + if err != nil { + r.writeTerminationLog(err.Error()) + return err + } + if skipReplication { + return nil + } + if lockAcquired { + defer r.releaseTargetArtifactUploadLock() + } + sourceObjs, err := r.listSourceObjects() if err != nil { r.writeTerminationLog(err.Error()) @@ -90,7 +143,7 @@ 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 @@ -99,8 +152,274 @@ func (r *ReplicaAgent) Start() error { err = replicatorImp.Replicate(sourceObjs) if err != nil { r.writeTerminationLog(err.Error()) + return err } - 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() (bool, bool, error) { + if r.ReplicationInput.TargetStorageType != storage.StorageTypeOCI { + return false, false, nil + } + if r.Config.Target.OCIOSDataStore == nil { + return false, false, fmt.Errorf("target OCI object store data store is nil") + } + reuseAllowed := r.Config.TargetArtifactReuseAllowed + if !reuseAllowed { + r.Logger.Infof("Target artifact reuse disabled because HF access validation marker is missing; upload lock still applies") + } + + state, err := r.targetArtifactState() + if err != nil { + return false, false, fmt.Errorf("failed to inspect target artifact state: %w", err) + } + if state.Complete { + if reuseAllowed { + r.Logger.Infof("Target artifact is already complete; skipping replication") + r.logTargetArtifactSize(state) + return false, true, nil + } + r.Logger.Infof("Target artifact is already complete but reuse is disabled; continuing with upload") + } + + for { + acquired, err := r.acquireTargetArtifactUploadLock() + if err != nil { + return false, false, err + } + if acquired { + return true, false, nil + } + + r.Logger.Infof("Target artifact upload lock already exists; waiting for completion marker") + state, err = r.waitForTargetArtifactStateChange() + if err != nil { + return false, false, err + } + if state.Complete { + if reuseAllowed { + r.Logger.Infof("Target artifact completed while waiting for upload lock; skipping replication") + r.logTargetArtifactSize(state) + return false, true, nil + } + r.Logger.Infof("Target artifact completed while waiting for upload lock but reuse is disabled; continuing with upload") + } + if r.isTargetArtifactUploadLockStale(state) { + if err := r.deleteStaleTargetArtifactUploadLock(state); err != nil { + return false, 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() (bool, 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) + acquired, err := tryAcquireArtifactUploadLockFunc( + r.Config.Target.OCIOSDataStore, + constants.ArtifactUploadLockBody, + lockURI, + ) + if err != nil { + return false, fmt.Errorf("failed to acquire target artifact upload lock: %w", err) + } + return acquired, nil +} + +func (r *ReplicaAgent) releaseTargetArtifactUploadLock() { + lockURI := r.targetArtifactUploadLockURI() + if err := deleteArtifactUploadLockFunc(r.Config.Target.OCIOSDataStore, lockURI); err != nil { + r.Logger.Errorf("Failed to release target artifact upload lock at oci://n/%s/b/%s/o/%s: %v", lockURI.Namespace, lockURI.BucketName, lockURI.ObjectName, err) + } +} + +func (r *ReplicaAgent) waitForTargetArtifactStateChange() (targetArtifactState, error) { + deadline := nowFunc().Add(r.targetArtifactUploadLockTimeout()) + for { + if !nowFunc().Before(deadline) { + return targetArtifactState{}, fmt.Errorf("timed out waiting for target artifact completion marker") + } + sleepFunc(targetArtifactLockPollInterval) + + 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.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 nil +} + +func (r *ReplicaAgent) writeCompletionMarker() error { + 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 +451,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 +475,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..2180711fb 100644 --- a/internal/ome-agent/replica/replica_test.go +++ b/internal/ome-agent/replica/replica_test.go @@ -1,18 +1,25 @@ package replica import ( + "errors" + "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" "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" @@ -57,6 +64,16 @@ func createMockOCIOSDataStore() *ociobjectstore.OCIOSDataStore { } } +type fakeReplicator struct { + err error + objects []common.ReplicationObject +} + +func (f *fakeReplicator) Replicate(objects []common.ReplicationObject) error { + f.objects = objects + return f.err +} + func TestNewReplicaAgent(t *testing.T) { mockLogger := testingPkg.SetupMockLogger() @@ -614,3 +631,402 @@ 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) (bool, error) { + return false, nil + } + stateCalls := 0 + targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { + stateCalls++ + if stateCalls == 1 { + return targetArtifactState{}, nil + } + return targetArtifactState{Complete: true, UploadLocked: true}, nil + } + + err := agent.Start() + require.NoError(t, err) + assert.False(t, replicatorCalled) + 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) (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) (bool, error) { + assert.Equal(t, "target-models/"+constants.ArtifactUploadLockFileName, target.ObjectName) + acquireAttempts++ + return acquireAttempts > 1, 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_StartUploadsCompletedTargetArtifactWhenReuseNotAllowed(t *testing.T) { + agent, cleanup := newTestAgentForCompletionMarker(t) + defer cleanup() + agent.Config.TargetArtifactReuseAllowed = false + + fake := &fakeReplicator{} + newReplicatorFunc = func(_ *ReplicaAgent) (replicator.Replicator, error) { + return fake, nil + } + stateCalls := 0 + targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { + stateCalls++ + return targetArtifactState{Complete: true}, nil + } + lockCalled := false + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (bool, error) { + lockCalled = 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, 1, stateCalls) + assert.True(t, lockCalled) +} + +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 + deleteArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI) error { + released = true + assert.Equal(t, "target-models/"+constants.ArtifactUploadLockFileName, target.ObjectName) + return nil + } + + err := agent.Start() + require.Error(t, err) + assert.Contains(t, err.Error(), "replication failed") + 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 newTestAgentForCompletionMarker(t *testing.T) (*ReplicaAgent, func()) { + t.Helper() + + oldNewReplicatorFunc := newReplicatorFunc + oldUploadCompletionMarkerFunc := uploadCompletionMarkerFunc + oldTryAcquireArtifactUploadLockFunc := tryAcquireArtifactUploadLockFunc + oldDeleteArtifactUploadLockFunc := deleteArtifactUploadLockFunc + oldDeleteStaleArtifactUploadLockFunc := deleteStaleArtifactUploadLockFunc + oldTargetArtifactStateFunc := targetArtifactStateFunc + oldSleepFunc := sleepFunc + oldNowFunc := nowFunc + cleanup := func() { + newReplicatorFunc = oldNewReplicatorFunc + uploadCompletionMarkerFunc = oldUploadCompletionMarkerFunc + tryAcquireArtifactUploadLockFunc = oldTryAcquireArtifactUploadLockFunc + deleteArtifactUploadLockFunc = oldDeleteArtifactUploadLockFunc + deleteStaleArtifactUploadLockFunc = oldDeleteStaleArtifactUploadLockFunc + targetArtifactStateFunc = oldTargetArtifactStateFunc + sleepFunc = oldSleepFunc + nowFunc = oldNowFunc + } + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (bool, error) { + return true, nil + } + deleteArtifactUploadLockFunc = 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..20e104d92 100644 --- a/pkg/ociobjectstore/os_data_store.go +++ b/pkg/ociobjectstore/os_data_store.go @@ -328,6 +328,133 @@ 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) { + 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) + + var putObjectBody io.ReadCloser + var uploadObjectSize *int64 + + if sourceFile, err := os.Open(source); err == nil { + fileInfo, err := sourceFile.Stat() + if err != nil { + return false, fmt.Errorf( + "failed to get source file info %q: %+v", + source, + err) + } + putObjectBody = sourceFile + tmp := fileInfo.Size() + uploadObjectSize = &tmp + } else { + putObjectBody = io.NopCloser(strings.NewReader(source)) + tmp := int64(len(source)) + uploadObjectSize = &tmp + } + 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)) + } + return 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 +550,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 +649,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() +} From 20a292719ee9a4dd74df15a73236611404c2a4ab Mon Sep 17 00:00:00 2001 From: Chengjie Yao Date: Mon, 27 Jul 2026 01:16:48 -0700 Subject: [PATCH 2/8] Harden OCI artifact upload marker handling --- internal/ome-agent/replica/replica.go | 80 +++++++-- internal/ome-agent/replica/replica_test.go | 182 +++++++++++++++++++-- 2 files changed, 237 insertions(+), 25 deletions(-) diff --git a/internal/ome-agent/replica/replica.go b/internal/ome-agent/replica/replica.go index 9afaa1a04..2262ee0c5 100644 --- a/internal/ome-agent/replica/replica.go +++ b/internal/ome-agent/replica/replica.go @@ -38,6 +38,9 @@ var ( deleteArtifactUploadLockFunc = func(dataStore *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI) error { return dataStore.DeleteObject(target) } + 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) } @@ -54,6 +57,7 @@ type ReplicaAgent struct { type targetArtifactState struct { Complete bool + CompletionMarked bool UploadLocked bool UploadLockModifiedTime *time.Time UploadLockETag string @@ -128,11 +132,21 @@ func (r *ReplicaAgent) Start() error { r.writeTerminationLog(err.Error()) return err } + if lockAcquired { + defer r.releaseTargetArtifactUploadLock() + } if skipReplication { return nil } if lockAcquired { - defer r.releaseTargetArtifactUploadLock() + skipReplication, err = r.prepareTargetArtifactAfterLockAcquired() + if err != nil { + r.writeTerminationLog(err.Error()) + return err + } + if skipReplication { + return nil + } } sourceObjs, err := r.listSourceObjects() @@ -165,28 +179,28 @@ func (r *ReplicaAgent) Start() error { } func (r *ReplicaAgent) prepareTargetArtifactUpload() (bool, bool, error) { + if !r.Config.TargetArtifactReuseAllowed { + return false, false, nil + } if r.ReplicationInput.TargetStorageType != storage.StorageTypeOCI { return false, false, nil } if r.Config.Target.OCIOSDataStore == nil { return false, false, fmt.Errorf("target OCI object store data store is nil") } - reuseAllowed := r.Config.TargetArtifactReuseAllowed - if !reuseAllowed { - r.Logger.Infof("Target artifact reuse disabled because HF access validation marker is missing; upload lock still applies") - } + waitDeadline := nowFunc().Add(r.targetArtifactUploadLockTimeout()) state, err := r.targetArtifactState() if err != nil { return false, false, fmt.Errorf("failed to inspect target artifact state: %w", err) } if state.Complete { - if reuseAllowed { + if r.canReuseCompleteTargetArtifact(state) { r.Logger.Infof("Target artifact is already complete; skipping replication") r.logTargetArtifactSize(state) return false, true, nil } - r.Logger.Infof("Target artifact is already complete but reuse is disabled; continuing with upload") + r.Logger.Infof("Target artifact is complete but upload lock still exists; waiting for upload lock release") } for { @@ -199,17 +213,16 @@ func (r *ReplicaAgent) prepareTargetArtifactUpload() (bool, bool, error) { } r.Logger.Infof("Target artifact upload lock already exists; waiting for completion marker") - state, err = r.waitForTargetArtifactStateChange() + state, err = r.waitForTargetArtifactStateChange(waitDeadline) if err != nil { return false, false, err } if state.Complete { - if reuseAllowed { + if r.canReuseCompleteTargetArtifact(state) { r.Logger.Infof("Target artifact completed while waiting for upload lock; skipping replication") r.logTargetArtifactSize(state) return false, true, nil } - r.Logger.Infof("Target artifact completed while waiting for upload lock but reuse is disabled; continuing with upload") } if r.isTargetArtifactUploadLockStale(state) { if err := r.deleteStaleTargetArtifactUploadLock(state); err != nil { @@ -220,6 +233,12 @@ func (r *ReplicaAgent) prepareTargetArtifactUpload() (bool, bool, error) { } } +func (r *ReplicaAgent) canReuseCompleteTargetArtifact(state targetArtifactState) bool { + // A complete marker is reusable once no active writer owns the prefix. A + // stale lock can be left behind after completion and should not block reuse. + return state.Complete && (!state.UploadLocked || r.isTargetArtifactUploadLockStale(state)) +} + func (r *ReplicaAgent) logTargetArtifactSize(state targetArtifactState) { if state.ArtifactSizeBytes == nil { r.Logger.Infof("Target artifact is complete but artifact size is unavailable") @@ -249,8 +268,7 @@ func (r *ReplicaAgent) releaseTargetArtifactUploadLock() { } } -func (r *ReplicaAgent) waitForTargetArtifactStateChange() (targetArtifactState, error) { - deadline := nowFunc().Add(r.targetArtifactUploadLockTimeout()) +func (r *ReplicaAgent) waitForTargetArtifactStateChange(deadline time.Time) (targetArtifactState, error) { for { if !nowFunc().Before(deadline) { return targetArtifactState{}, fmt.Errorf("timed out waiting for target artifact completion marker") @@ -261,7 +279,7 @@ func (r *ReplicaAgent) waitForTargetArtifactStateChange() (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) { + if !state.UploadLocked || r.isTargetArtifactUploadLockStale(state) { return state, nil } } @@ -352,6 +370,7 @@ func defaultTargetArtifactState(dataStore *ociobjectstore.OCIOSDataStore, target } } + state.CompletionMarked = hasCompleteMarker state.Complete = hasCompleteMarker && hasArtifactObject state.UploadLocked = hasUploadLock if artifactSizeBytes > 0 { @@ -370,7 +389,42 @@ func objectSummaryTime(object objectstorage.ObjectSummary) *time.Time { 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 + } + if !state.CompletionMarked { + return false, 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 false, fmt.Errorf("failed to delete target artifact completion marker before upload: %w", err) + } + return false, 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 diff --git a/internal/ome-agent/replica/replica_test.go b/internal/ome-agent/replica/replica_test.go index 2180711fb..bec4d73ef 100644 --- a/internal/ome-agent/replica/replica_test.go +++ b/internal/ome-agent/replica/replica_test.go @@ -65,12 +65,16 @@ func createMockOCIOSDataStore() *ociobjectstore.OCIOSDataStore { } type fakeReplicator struct { - err error - objects []common.ReplicationObject + 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 } @@ -729,7 +733,7 @@ func TestReplicaAgent_StartSkipsReplicationWhenTargetArtifactUploadLockCompletes if stateCalls == 1 { return targetArtifactState{}, nil } - return targetArtifactState{Complete: true, UploadLocked: true}, nil + return targetArtifactState{Complete: true}, nil } err := agent.Start() @@ -738,6 +742,71 @@ func TestReplicaAgent_StartSkipsReplicationWhenTargetArtifactUploadLockCompletes assert.Equal(t, 2, stateCalls) } +func TestReplicaAgent_StartWaitsWhenCompletedTargetArtifactStillHasActiveUploadLock(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) (bool, error) { + return false, nil + } + stateCalls := 0 + targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { + stateCalls++ + if stateCalls < 3 { + return targetArtifactState{Complete: true, CompletionMarked: true, UploadLocked: true}, nil + } + return targetArtifactState{Complete: true, CompletionMarked: true}, nil + } + + err := agent.Start() + require.NoError(t, err) + assert.False(t, replicatorCalled) + assert.Equal(t, 3, stateCalls) +} + +func TestReplicaAgent_StartSkipsCompletedTargetArtifactWhenUploadLockClearsBeforeAcquire(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) (bool, error) { + return true, nil + } + released := false + deleteArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI) error { + released = true + assert.Equal(t, "target-models/"+constants.ArtifactUploadLockFileName, target.ObjectName) + return 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{Complete: true, CompletionMarked: true, UploadLocked: true}, 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() @@ -855,7 +924,7 @@ func TestReplicaAgent_StartLogsTargetArtifactSizeWhenSkippingCompletedTargetArti mockLogger.AssertCalled(t, "Infof", "Total model size: %d bytes", []interface{}{artifactSizeBytes}) } -func TestReplicaAgent_StartUploadsCompletedTargetArtifactWhenReuseNotAllowed(t *testing.T) { +func TestReplicaAgent_StartBypassesTargetArtifactCoordinationWhenReuseNotAllowed(t *testing.T) { agent, cleanup := newTestAgentForCompletionMarker(t) defer cleanup() agent.Config.TargetArtifactReuseAllowed = false @@ -864,25 +933,109 @@ func TestReplicaAgent_StartUploadsCompletedTargetArtifactWhenReuseNotAllowed(t * newReplicatorFunc = func(_ *ReplicaAgent) (replicator.Replicator, error) { return fake, nil } - stateCalls := 0 targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { - stateCalls++ - return targetArtifactState{Complete: true}, nil + t.Fatal("target artifact state should not be inspected when reuse is disabled") + return targetArtifactState{}, nil } - lockCalled := false tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (bool, error) { - lockCalled = true - return true, nil + t.Fatal("target artifact upload lock should not be acquired when reuse is disabled") + return false, nil + } + deleteArtifactCompletionMarkerFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) error { + t.Fatal("completion marker should not be deleted when reuse is disabled") + 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, 1, stateCalls) - assert.True(t, lockCalled) +} + +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) (bool, error) { + return 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) (bool, error) { + t.Fatal("target artifact upload lock should not be acquired when reuse is disabled") + return false, nil + } + + lockAcquired, skipReplication, err := agent.prepareTargetArtifactUpload() + require.NoError(t, err) + assert.False(t, lockAcquired) + 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) (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_StartReleasesTargetArtifactUploadLockWhenReplicationFails(t *testing.T) { @@ -964,6 +1117,7 @@ func newTestAgentForCompletionMarker(t *testing.T) (*ReplicaAgent, func()) { oldUploadCompletionMarkerFunc := uploadCompletionMarkerFunc oldTryAcquireArtifactUploadLockFunc := tryAcquireArtifactUploadLockFunc oldDeleteArtifactUploadLockFunc := deleteArtifactUploadLockFunc + oldDeleteArtifactCompletionMarkerFunc := deleteArtifactCompletionMarkerFunc oldDeleteStaleArtifactUploadLockFunc := deleteStaleArtifactUploadLockFunc oldTargetArtifactStateFunc := targetArtifactStateFunc oldSleepFunc := sleepFunc @@ -973,6 +1127,7 @@ func newTestAgentForCompletionMarker(t *testing.T) (*ReplicaAgent, func()) { uploadCompletionMarkerFunc = oldUploadCompletionMarkerFunc tryAcquireArtifactUploadLockFunc = oldTryAcquireArtifactUploadLockFunc deleteArtifactUploadLockFunc = oldDeleteArtifactUploadLockFunc + deleteArtifactCompletionMarkerFunc = oldDeleteArtifactCompletionMarkerFunc deleteStaleArtifactUploadLockFunc = oldDeleteStaleArtifactUploadLockFunc targetArtifactStateFunc = oldTargetArtifactStateFunc sleepFunc = oldSleepFunc @@ -984,6 +1139,9 @@ func newTestAgentForCompletionMarker(t *testing.T) (*ReplicaAgent, func()) { deleteArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) error { return nil } + deleteArtifactCompletionMarkerFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) error { + return nil + } targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { return targetArtifactState{}, nil } From 0423a771e608020c0a7e9e5d7996cb6757749d80 Mon Sep 17 00:00:00 2001 From: Chengjie Yao Date: Wed, 19 Aug 2026 19:06:41 -0700 Subject: [PATCH 3/8] Remove completion marker before OCI artifact overwrite --- internal/ome-agent/replica/replica.go | 22 ++++++++++++++++++---- internal/ome-agent/replica/replica_test.go | 16 ++++++++++++---- 2 files changed, 30 insertions(+), 8 deletions(-) diff --git a/internal/ome-agent/replica/replica.go b/internal/ome-agent/replica/replica.go index 2262ee0c5..75b1909e2 100644 --- a/internal/ome-agent/replica/replica.go +++ b/internal/ome-agent/replica/replica.go @@ -163,6 +163,11 @@ func (r *ReplicaAgent) Start() 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()) @@ -409,16 +414,25 @@ func (r *ReplicaAgent) prepareTargetArtifactAfterLockAcquired() (bool, error) { r.logTargetArtifactSize(state) return true, nil } - if !state.CompletionMarked { - return false, 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 false, fmt.Errorf("failed to delete target artifact completion marker before upload: %w", err) + return fmt.Errorf("failed to delete target artifact completion marker before upload: %w", err) } - return false, nil + return nil } func (r *ReplicaAgent) writeCompletionMarker() error { diff --git a/internal/ome-agent/replica/replica_test.go b/internal/ome-agent/replica/replica_test.go index bec4d73ef..37e0f5dfa 100644 --- a/internal/ome-agent/replica/replica_test.go +++ b/internal/ome-agent/replica/replica_test.go @@ -924,12 +924,18 @@ func TestReplicaAgent_StartLogsTargetArtifactSizeWhenSkippingCompletedTargetArti mockLogger.AssertCalled(t, "Infof", "Total model size: %d bytes", []interface{}{artifactSizeBytes}) } -func TestReplicaAgent_StartBypassesTargetArtifactCoordinationWhenReuseNotAllowed(t *testing.T) { +func TestReplicaAgent_StartDeletesCompletionMarkerBeforeOverwritingWhenReuseNotAllowed(t *testing.T) { agent, cleanup := newTestAgentForCompletionMarker(t) defer cleanup() agent.Config.TargetArtifactReuseAllowed = false - fake := &fakeReplicator{} + 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 } @@ -941,8 +947,9 @@ func TestReplicaAgent_StartBypassesTargetArtifactCoordinationWhenReuseNotAllowed t.Fatal("target artifact upload lock should not be acquired when reuse is disabled") return false, nil } - deleteArtifactCompletionMarkerFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) error { - t.Fatal("completion marker should not be deleted when reuse is disabled") + 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 { @@ -953,6 +960,7 @@ func TestReplicaAgent_StartBypassesTargetArtifactCoordinationWhenReuseNotAllowed 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) { From cab6375588eb390094e7bb789b14b076e76b2182 Mon Sep 17 00:00:00 2001 From: Chengjie Yao Date: Wed, 19 Aug 2026 19:10:40 -0700 Subject: [PATCH 4/8] Fence OCI artifact lock release with ETag --- internal/ome-agent/replica/replica.go | 71 +++++++++++------- internal/ome-agent/replica/replica_test.go | 84 ++++++++++++++-------- pkg/ociobjectstore/os_data_store.go | 21 ++++-- 3 files changed, 113 insertions(+), 63 deletions(-) diff --git a/internal/ome-agent/replica/replica.go b/internal/ome-agent/replica/replica.go index 75b1909e2..d03ce6fc1 100644 --- a/internal/ome-agent/replica/replica.go +++ b/internal/ome-agent/replica/replica.go @@ -32,11 +32,11 @@ var ( 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) (bool, error) { - return dataStore.UploadIfAbsent(source, target) + tryAcquireArtifactUploadLockFunc = func(dataStore *ociobjectstore.OCIOSDataStore, source string, target ociobjectstore.ObjectURI) (string, bool, error) { + return dataStore.UploadIfAbsentWithETag(source, target) } - deleteArtifactUploadLockFunc = func(dataStore *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI) error { - return dataStore.DeleteObject(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) @@ -64,6 +64,12 @@ type targetArtifactState struct { 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) @@ -127,18 +133,18 @@ func (r *ReplicaAgent) Start() error { return err } - lockAcquired, skipReplication, err := r.prepareTargetArtifactUpload() + uploadLock, skipReplication, err := r.prepareTargetArtifactUpload() if err != nil { r.writeTerminationLog(err.Error()) return err } - if lockAcquired { - defer r.releaseTargetArtifactUploadLock() + if uploadLock != nil { + defer r.releaseTargetArtifactUploadLock(*uploadLock) } if skipReplication { return nil } - if lockAcquired { + if uploadLock != nil { skipReplication, err = r.prepareTargetArtifactAfterLockAcquired() if err != nil { r.writeTerminationLog(err.Error()) @@ -183,55 +189,55 @@ func (r *ReplicaAgent) Start() error { return nil } -func (r *ReplicaAgent) prepareTargetArtifactUpload() (bool, bool, error) { +func (r *ReplicaAgent) prepareTargetArtifactUpload() (*targetArtifactUploadLock, bool, error) { if !r.Config.TargetArtifactReuseAllowed { - return false, false, nil + return nil, false, nil } if r.ReplicationInput.TargetStorageType != storage.StorageTypeOCI { - return false, false, nil + return nil, false, nil } if r.Config.Target.OCIOSDataStore == nil { - return false, false, fmt.Errorf("target OCI object store data store is 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 false, false, fmt.Errorf("failed to inspect target artifact state: %w", err) + return nil, false, fmt.Errorf("failed to inspect target artifact state: %w", err) } if state.Complete { if r.canReuseCompleteTargetArtifact(state) { r.Logger.Infof("Target artifact is already complete; skipping replication") r.logTargetArtifactSize(state) - return false, true, nil + return nil, true, nil } r.Logger.Infof("Target artifact is complete but upload lock still exists; waiting for upload lock release") } for { - acquired, err := r.acquireTargetArtifactUploadLock() + uploadLock, err := r.acquireTargetArtifactUploadLock() if err != nil { - return false, false, err + return nil, false, err } - if acquired { - return true, false, nil + 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 false, false, err + return nil, false, err } if state.Complete { if r.canReuseCompleteTargetArtifact(state) { r.Logger.Infof("Target artifact completed while waiting for upload lock; skipping replication") r.logTargetArtifactSize(state) - return false, true, nil + return nil, true, nil } } if r.isTargetArtifactUploadLockStale(state) { if err := r.deleteStaleTargetArtifactUploadLock(state); err != nil { - return false, false, err + return nil, false, err } continue } @@ -252,24 +258,35 @@ func (r *ReplicaAgent) logTargetArtifactSize(state targetArtifactState) { r.Logger.Infof("Total model size: %d bytes", *state.ArtifactSizeBytes) } -func (r *ReplicaAgent) acquireTargetArtifactUploadLock() (bool, error) { +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) - acquired, err := tryAcquireArtifactUploadLockFunc( + etag, acquired, err := tryAcquireArtifactUploadLockFunc( r.Config.Target.OCIOSDataStore, constants.ArtifactUploadLockBody, lockURI, ) if err != nil { - return false, fmt.Errorf("failed to acquire target artifact upload lock: %w", err) + return nil, fmt.Errorf("failed to acquire target artifact upload lock: %w", err) + } + if !acquired { + return nil, nil } - return acquired, nil + if etag == "" { + return nil, fmt.Errorf("acquired target artifact upload lock without an etag") + } + return &targetArtifactUploadLock{ETag: etag}, nil } -func (r *ReplicaAgent) releaseTargetArtifactUploadLock() { +func (r *ReplicaAgent) releaseTargetArtifactUploadLock(uploadLock targetArtifactUploadLock) { lockURI := r.targetArtifactUploadLockURI() - if err := deleteArtifactUploadLockFunc(r.Config.Target.OCIOSDataStore, lockURI); err != nil { + released, err := releaseArtifactUploadLockFunc(r.Config.Target.OCIOSDataStore, lockURI, uploadLock.ETag) + if err != nil { r.Logger.Errorf("Failed to release target artifact upload lock at oci://n/%s/b/%s/o/%s: %v", lockURI.Namespace, lockURI.BucketName, lockURI.ObjectName, err) + return + } + if !released { + r.Logger.Infof("Target artifact upload lock changed before release; leaving the current owner lock in place") } } diff --git a/internal/ome-agent/replica/replica_test.go b/internal/ome-agent/replica/replica_test.go index 37e0f5dfa..c92ad17ab 100644 --- a/internal/ome-agent/replica/replica_test.go +++ b/internal/ome-agent/replica/replica_test.go @@ -724,8 +724,8 @@ func TestReplicaAgent_StartSkipsReplicationWhenTargetArtifactUploadLockCompletes replicatorCalled = true return &fakeReplicator{}, nil } - tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (bool, error) { - return false, nil + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (string, bool, error) { + return "", false, nil } stateCalls := 0 targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { @@ -751,8 +751,8 @@ func TestReplicaAgent_StartWaitsWhenCompletedTargetArtifactStillHasActiveUploadL replicatorCalled = true return &fakeReplicator{}, nil } - tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (bool, error) { - return false, nil + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (string, bool, error) { + return "", false, nil } stateCalls := 0 targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { @@ -778,14 +778,15 @@ func TestReplicaAgent_StartSkipsCompletedTargetArtifactWhenUploadLockClearsBefor replicatorCalled = true return &fakeReplicator{}, nil } - tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (bool, error) { - return true, nil + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (string, bool, error) { + return "acquired-lock-etag", true, nil } released := false - deleteArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI) error { + releaseArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI, etag string) (bool, error) { released = true assert.Equal(t, "target-models/"+constants.ArtifactUploadLockFileName, target.ObjectName) - return nil + 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") @@ -820,9 +821,9 @@ func TestReplicaAgent_StartSkipsCompletedTargetArtifactEvenWhenUploadLockIsStale replicatorCalled = true return &fakeReplicator{}, nil } - tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, target ociobjectstore.ObjectURI) (bool, error) { + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, target ociobjectstore.ObjectURI) (string, bool, error) { assert.Equal(t, "target-models/"+constants.ArtifactUploadLockFileName, target.ObjectName) - return false, nil + return "", false, nil } staleLockModifiedTime := now.Add(-2 * time.Hour) @@ -862,10 +863,13 @@ func TestReplicaAgent_StartDeletesStaleTargetArtifactUploadLockAndReplicates(t * } acquireAttempts := 0 - tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, target ociobjectstore.ObjectURI) (bool, error) { + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, target ociobjectstore.ObjectURI) (string, bool, error) { assert.Equal(t, "target-models/"+constants.ArtifactUploadLockFileName, target.ObjectName) acquireAttempts++ - return acquireAttempts > 1, nil + if acquireAttempts > 1 { + return "replacement-lock-etag", true, nil + } + return "", false, nil } stateCalls := 0 @@ -943,9 +947,9 @@ func TestReplicaAgent_StartDeletesCompletionMarkerBeforeOverwritingWhenReuseNotA t.Fatal("target artifact state should not be inspected when reuse is disabled") return targetArtifactState{}, nil } - tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (bool, error) { + 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 + return "", false, nil } deleteArtifactCompletionMarkerFunc = func(_ *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI) error { assert.Equal(t, "target-models/"+constants.ArtifactCompleteMarkerFileName, target.ObjectName) @@ -979,8 +983,8 @@ func TestReplicaAgent_StartDeletesStaleCompletionMarkerBeforeUploadingTargetArti targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { return targetArtifactState{CompletionMarked: true}, nil } - tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (bool, error) { - return 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") @@ -1007,14 +1011,14 @@ func TestReplicaAgent_PrepareTargetArtifactUploadSkipsWhenReuseDisabled(t *testi t.Fatal("target artifact state should not be inspected when reuse is disabled") return targetArtifactState{}, nil } - tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (bool, error) { + 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 + return "", false, nil } - lockAcquired, skipReplication, err := agent.prepareTargetArtifactUpload() + uploadLock, skipReplication, err := agent.prepareTargetArtifactUpload() require.NoError(t, err) - assert.False(t, lockAcquired) + assert.Nil(t, uploadLock) assert.False(t, skipReplication) } @@ -1029,12 +1033,12 @@ func TestReplicaAgent_PrepareTargetArtifactUploadUsesOneWaitDeadline(t *testing. current = current.Add(time.Hour) } acquireAttempts := 0 - tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (bool, error) { + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (string, bool, error) { acquireAttempts++ if acquireAttempts > 2 { - return false, errors.New("wait deadline was reset") + return "", false, errors.New("wait deadline was reset") } - return false, nil + return "", false, nil } targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { return targetArtifactState{}, nil @@ -1055,10 +1059,11 @@ func TestReplicaAgent_StartReleasesTargetArtifactUploadLockWhenReplicationFails( } released := false - deleteArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI) error { + releaseArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, target ociobjectstore.ObjectURI, etag string) (bool, error) { released = true assert.Equal(t, "target-models/"+constants.ArtifactUploadLockFileName, target.ObjectName) - return nil + assert.Equal(t, "test-lock-etag", etag) + return true, nil } err := agent.Start() @@ -1067,6 +1072,23 @@ func TestReplicaAgent_StartReleasesTargetArtifactUploadLockWhenReplicationFails( assert.True(t, released) } +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 + } + + agent.releaseTargetArtifactUploadLock(targetArtifactUploadLock{ETag: "acquired-lock-etag"}) + + assert.True(t, released) +} + func TestReplicaAgent_StartSkipsCompletionMarkerForNonOCITarget(t *testing.T) { agent, cleanup := newTestAgentForCompletionMarker(t) defer cleanup() @@ -1124,7 +1146,7 @@ func newTestAgentForCompletionMarker(t *testing.T) (*ReplicaAgent, func()) { oldNewReplicatorFunc := newReplicatorFunc oldUploadCompletionMarkerFunc := uploadCompletionMarkerFunc oldTryAcquireArtifactUploadLockFunc := tryAcquireArtifactUploadLockFunc - oldDeleteArtifactUploadLockFunc := deleteArtifactUploadLockFunc + oldReleaseArtifactUploadLockFunc := releaseArtifactUploadLockFunc oldDeleteArtifactCompletionMarkerFunc := deleteArtifactCompletionMarkerFunc oldDeleteStaleArtifactUploadLockFunc := deleteStaleArtifactUploadLockFunc oldTargetArtifactStateFunc := targetArtifactStateFunc @@ -1134,18 +1156,18 @@ func newTestAgentForCompletionMarker(t *testing.T) (*ReplicaAgent, func()) { newReplicatorFunc = oldNewReplicatorFunc uploadCompletionMarkerFunc = oldUploadCompletionMarkerFunc tryAcquireArtifactUploadLockFunc = oldTryAcquireArtifactUploadLockFunc - deleteArtifactUploadLockFunc = oldDeleteArtifactUploadLockFunc + releaseArtifactUploadLockFunc = oldReleaseArtifactUploadLockFunc deleteArtifactCompletionMarkerFunc = oldDeleteArtifactCompletionMarkerFunc deleteStaleArtifactUploadLockFunc = oldDeleteStaleArtifactUploadLockFunc targetArtifactStateFunc = oldTargetArtifactStateFunc sleepFunc = oldSleepFunc nowFunc = oldNowFunc } - tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (bool, error) { - return true, nil + tryAcquireArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ string, _ ociobjectstore.ObjectURI) (string, bool, error) { + return "test-lock-etag", true, nil } - deleteArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) error { - return nil + releaseArtifactUploadLockFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI, _ string) (bool, error) { + return true, nil } deleteArtifactCompletionMarkerFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) error { return nil diff --git a/pkg/ociobjectstore/os_data_store.go b/pkg/ociobjectstore/os_data_store.go index 20e104d92..8e8cf4011 100644 --- a/pkg/ociobjectstore/os_data_store.go +++ b/pkg/ociobjectstore/os_data_store.go @@ -332,10 +332,18 @@ func (cds *OCIOSDataStore) Upload(source string, target ObjectURI) error { // 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) + return "", false, fmt.Errorf("error upload object due to no namespace found: %+v", err) } target.Namespace = *namespace } @@ -349,7 +357,7 @@ func (cds *OCIOSDataStore) UploadIfAbsent(source string, target ObjectURI) (bool if sourceFile, err := os.Open(source); err == nil { fileInfo, err := sourceFile.Stat() if err != nil { - return false, fmt.Errorf( + return "", false, fmt.Errorf( "failed to get source file info %q: %+v", source, err) @@ -374,16 +382,19 @@ func (cds *OCIOSDataStore) UploadIfAbsent(source string, target ObjectURI) (bool } response, err := cds.Client.PutObject(context.Background(), putObjectRequest) if isObjectAlreadyPresent(response.RawResponse, err) { - return false, nil + return "", false, nil } if err != nil || response.RawResponse == nil || response.RawResponse.StatusCode != http.StatusOK { - return false, fmt.Errorf( + return "", false, fmt.Errorf( "failed to put object %q with response %+v: %s", objectFullName, response, errorMessage(err)) } - return true, nil + 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 From a6d226b7278ef57f989d6683b4f610e2a7d49594 Mon Sep 17 00:00:00 2001 From: Chengjie Yao Date: Fri, 21 Aug 2026 17:28:45 -0700 Subject: [PATCH 5/8] Harden OCI artifact upload lock cleanup --- internal/ome-agent/replica/replica.go | 48 ++++++--- internal/ome-agent/replica/replica_test.go | 69 ++++++++++++- pkg/ociobjectstore/os_data_store_test.go | 108 +++++++++++++++++++++ 3 files changed, 212 insertions(+), 13 deletions(-) diff --git a/internal/ome-agent/replica/replica.go b/internal/ome-agent/replica/replica.go index d03ce6fc1..dc709c8c0 100644 --- a/internal/ome-agent/replica/replica.go +++ b/internal/ome-agent/replica/replica.go @@ -23,8 +23,10 @@ const ( SourceStorageConfigKeyName = "source" TargetStorageConfigKeyName = "target" - targetArtifactLockPollInterval = 30 * time.Second - defaultArtifactUploadLockWaitTimeout = 120 * time.Hour + targetArtifactLockPollInterval = 30 * time.Second + targetArtifactLockReleaseMaxAttempts = 3 + targetArtifactLockReleaseRetryInterval = 2 * time.Second + defaultArtifactUploadLockWaitTimeout = 120 * time.Hour ) var ( @@ -124,7 +126,7 @@ 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 { @@ -139,7 +141,17 @@ func (r *ReplicaAgent) Start() error { return err } if uploadLock != nil { - defer r.releaseTargetArtifactUploadLock(*uploadLock) + 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 @@ -278,16 +290,28 @@ func (r *ReplicaAgent) acquireTargetArtifactUploadLock() (*targetArtifactUploadL return &targetArtifactUploadLock{ETag: etag}, nil } -func (r *ReplicaAgent) releaseTargetArtifactUploadLock(uploadLock targetArtifactUploadLock) { +func (r *ReplicaAgent) releaseTargetArtifactUploadLock(uploadLock targetArtifactUploadLock) error { lockURI := r.targetArtifactUploadLockURI() - released, err := releaseArtifactUploadLockFunc(r.Config.Target.OCIOSDataStore, lockURI, uploadLock.ETag) - if err != nil { - r.Logger.Errorf("Failed to release target artifact upload lock at oci://n/%s/b/%s/o/%s: %v", lockURI.Namespace, lockURI.BucketName, lockURI.ObjectName, err) - return - } - if !released { - r.Logger.Infof("Target artifact upload lock changed before release; leaving the current owner lock in place") + 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) { diff --git a/internal/ome-agent/replica/replica_test.go b/internal/ome-agent/replica/replica_test.go index c92ad17ab..fcc23e5ed 100644 --- a/internal/ome-agent/replica/replica_test.go +++ b/internal/ome-agent/replica/replica_test.go @@ -1072,6 +1072,73 @@ func TestReplicaAgent_StartReleasesTargetArtifactUploadLockWhenReplicationFails( 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() @@ -1084,7 +1151,7 @@ func TestReplicaAgent_ReleaseTargetArtifactUploadLockOnlyDeletesAcquiredLock(t * return false, nil } - agent.releaseTargetArtifactUploadLock(targetArtifactUploadLock{ETag: "acquired-lock-etag"}) + require.NoError(t, agent.releaseTargetArtifactUploadLock(targetArtifactUploadLock{ETag: "acquired-lock-etag"})) assert.True(t, released) } diff --git a/pkg/ociobjectstore/os_data_store_test.go b/pkg/ociobjectstore/os_data_store_test.go index fb4858577..e1a7ced56 100644 --- a/pkg/ociobjectstore/os_data_store_test.go +++ b/pkg/ociobjectstore/os_data_store_test.go @@ -2,12 +2,15 @@ package ociobjectstore import ( "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 +18,111 @@ 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, + } +} + +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) From a78907e124bf9994869a2e801a2e58b48f71ce9c Mon Sep 17 00:00:00 2001 From: Chengjie Yao Date: Fri, 21 Aug 2026 22:01:12 -0700 Subject: [PATCH 6/8] Close OCI upload source files reliably --- pkg/ociobjectstore/os_data_store.go | 76 ++++++++++++------------ pkg/ociobjectstore/os_data_store_test.go | 32 ++++++++++ 2 files changed, 70 insertions(+), 38 deletions(-) diff --git a/pkg/ociobjectstore/os_data_store.go b/pkg/ociobjectstore/os_data_store.go index 8e8cf4011..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, @@ -351,24 +366,9 @@ func (cds *OCIOSDataStore) UploadIfAbsentWithETag(source string, target ObjectUR objectFullName := fmt.Sprintf( "%s/%s/%s", target.Namespace, target.BucketName, target.ObjectName) - var putObjectBody io.ReadCloser - var uploadObjectSize *int64 - - if sourceFile, err := os.Open(source); err == nil { - fileInfo, err := sourceFile.Stat() - if err != nil { - return "", false, fmt.Errorf( - "failed to get source file info %q: %+v", - source, - err) - } - putObjectBody = sourceFile - tmp := fileInfo.Size() - uploadObjectSize = &tmp - } else { - putObjectBody = io.NopCloser(strings.NewReader(source)) - tmp := int64(len(source)) - uploadObjectSize = &tmp + putObjectBody, uploadObjectSize, err := newUploadBody(source) + if err != nil { + return "", false, err } defer putObjectBody.Close() diff --git a/pkg/ociobjectstore/os_data_store_test.go b/pkg/ociobjectstore/os_data_store_test.go index e1a7ced56..79cce3f06 100644 --- a/pkg/ociobjectstore/os_data_store_test.go +++ b/pkg/ociobjectstore/os_data_store_test.go @@ -1,6 +1,7 @@ package ociobjectstore import ( + "errors" "fmt" "net/http" "os" @@ -54,6 +55,37 @@ func objectStorageTestResponse(request *http.Request, statusCode int, headers ht } } +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"} From 9010288300c2aa97b682c8d8dd4961b214733d3f Mon Sep 17 00:00:00 2001 From: Chengjie Yao Date: Sat, 22 Aug 2026 17:58:00 -0700 Subject: [PATCH 7/8] Reuse completed OCI artifacts despite retained upload lock and add test coverage --- cmd/ome-agent/config_test.go | 79 ++++++++++ internal/ome-agent/replica/config_test.go | 1 + internal/ome-agent/replica/replica.go | 28 ++-- internal/ome-agent/replica/replica_test.go | 168 +++++++++++++++++++-- 4 files changed, 247 insertions(+), 29 deletions(-) create mode 100644 cmd/ome-agent/config_test.go 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/internal/ome-agent/replica/config_test.go b/internal/ome-agent/replica/config_test.go index c6c281b58..53aa3c115 100644 --- a/internal/ome-agent/replica/config_test.go +++ b/internal/ome-agent/replica/config_test.go @@ -446,6 +446,7 @@ 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) } diff --git a/internal/ome-agent/replica/replica.go b/internal/ome-agent/replica/replica.go index dc709c8c0..32451f45d 100644 --- a/internal/ome-agent/replica/replica.go +++ b/internal/ome-agent/replica/replica.go @@ -218,12 +218,12 @@ func (r *ReplicaAgent) prepareTargetArtifactUpload() (*targetArtifactUploadLock, return nil, false, fmt.Errorf("failed to inspect target artifact state: %w", err) } if state.Complete { - if r.canReuseCompleteTargetArtifact(state) { - r.Logger.Infof("Target artifact is already complete; skipping replication") - r.logTargetArtifactSize(state) - return nil, true, nil - } - r.Logger.Infof("Target artifact is complete but upload lock still exists; waiting for upload lock release") + // 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 { @@ -241,11 +241,9 @@ func (r *ReplicaAgent) prepareTargetArtifactUpload() (*targetArtifactUploadLock, return nil, false, err } if state.Complete { - if r.canReuseCompleteTargetArtifact(state) { - r.Logger.Infof("Target artifact completed while waiting for upload lock; skipping replication") - r.logTargetArtifactSize(state) - return nil, true, nil - } + 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 { @@ -256,12 +254,6 @@ func (r *ReplicaAgent) prepareTargetArtifactUpload() (*targetArtifactUploadLock, } } -func (r *ReplicaAgent) canReuseCompleteTargetArtifact(state targetArtifactState) bool { - // A complete marker is reusable once no active writer owns the prefix. A - // stale lock can be left behind after completion and should not block reuse. - return state.Complete && (!state.UploadLocked || r.isTargetArtifactUploadLockStale(state)) -} - func (r *ReplicaAgent) logTargetArtifactSize(state targetArtifactState) { if state.ArtifactSizeBytes == nil { r.Logger.Infof("Target artifact is complete but artifact size is unavailable") @@ -325,7 +317,7 @@ func (r *ReplicaAgent) waitForTargetArtifactStateChange(deadline time.Time) (tar if err != nil { return targetArtifactState{}, fmt.Errorf("failed to inspect target artifact state while waiting for upload lock: %w", err) } - if !state.UploadLocked || r.isTargetArtifactUploadLockStale(state) { + if state.Complete || !state.UploadLocked || r.isTargetArtifactUploadLockStale(state) { return state, nil } } diff --git a/internal/ome-agent/replica/replica_test.go b/internal/ome-agent/replica/replica_test.go index fcc23e5ed..6cbb9e2f8 100644 --- a/internal/ome-agent/replica/replica_test.go +++ b/internal/ome-agent/replica/replica_test.go @@ -2,6 +2,8 @@ package replica import ( "errors" + "io" + "net/http" "os" "path/filepath" "strings" @@ -13,6 +15,7 @@ import ( "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" @@ -32,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) @@ -64,6 +86,19 @@ 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 @@ -733,7 +768,11 @@ func TestReplicaAgent_StartSkipsReplicationWhenTargetArtifactUploadLockCompletes if stateCalls == 1 { return targetArtifactState{}, nil } - return targetArtifactState{Complete: true}, 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() @@ -742,7 +781,7 @@ func TestReplicaAgent_StartSkipsReplicationWhenTargetArtifactUploadLockCompletes assert.Equal(t, 2, stateCalls) } -func TestReplicaAgent_StartWaitsWhenCompletedTargetArtifactStillHasActiveUploadLock(t *testing.T) { +func TestReplicaAgent_StartSkipsCompletedTargetArtifactWhenUploadLockRemains(t *testing.T) { agent, cleanup := newTestAgentForCompletionMarker(t) defer cleanup() @@ -752,24 +791,19 @@ func TestReplicaAgent_StartWaitsWhenCompletedTargetArtifactStillHasActiveUploadL 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 } - stateCalls := 0 targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { - stateCalls++ - if stateCalls < 3 { - return targetArtifactState{Complete: true, CompletionMarked: true, UploadLocked: true}, nil - } - return targetArtifactState{Complete: true, CompletionMarked: true}, nil + return targetArtifactState{Complete: true, CompletionMarked: true, UploadLocked: true}, nil } err := agent.Start() require.NoError(t, err) assert.False(t, replicatorCalled) - assert.Equal(t, 3, stateCalls) } -func TestReplicaAgent_StartSkipsCompletedTargetArtifactWhenUploadLockClearsBeforeAcquire(t *testing.T) { +func TestReplicaAgent_StartSkipsCompletedTargetArtifactAfterLockAcquired(t *testing.T) { agent, cleanup := newTestAgentForCompletionMarker(t) defer cleanup() @@ -796,7 +830,7 @@ func TestReplicaAgent_StartSkipsCompletedTargetArtifactWhenUploadLockClearsBefor targetArtifactStateFunc = func(_ *ociobjectstore.OCIOSDataStore, _ ociobjectstore.ObjectURI) (targetArtifactState, error) { stateCalls++ if stateCalls == 1 { - return targetArtifactState{Complete: true, CompletionMarked: true, UploadLocked: true}, nil + return targetArtifactState{}, nil } return targetArtifactState{Complete: true, CompletionMarked: true}, nil } @@ -1050,6 +1084,52 @@ func TestReplicaAgent_PrepareTargetArtifactUploadUsesOneWaitDeadline(t *testing. assert.Equal(t, 2, acquireAttempts) } +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() @@ -1207,6 +1287,72 @@ func TestFilterInternalArtifactReplicationObjectsSkipsCompletionMarker(t *testin 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() From 117d4c50198b02bf797f468c9df44976314f3253 Mon Sep 17 00:00:00 2001 From: Chengjie Yao Date: Sun, 23 Aug 2026 05:30:31 -0700 Subject: [PATCH 8/8] Limit OCI artifact lock polling to configured timeout --- internal/ome-agent/replica/replica.go | 9 +++++++-- internal/ome-agent/replica/replica_test.go | 23 ++++++++++++++++++++++ 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/internal/ome-agent/replica/replica.go b/internal/ome-agent/replica/replica.go index 32451f45d..54428a76a 100644 --- a/internal/ome-agent/replica/replica.go +++ b/internal/ome-agent/replica/replica.go @@ -308,10 +308,15 @@ func (r *ReplicaAgent) releaseTargetArtifactUploadLock(uploadLock targetArtifact func (r *ReplicaAgent) waitForTargetArtifactStateChange(deadline time.Time) (targetArtifactState, error) { for { - if !nowFunc().Before(deadline) { + remaining := deadline.Sub(nowFunc()) + if remaining <= 0 { return targetArtifactState{}, fmt.Errorf("timed out waiting for target artifact completion marker") } - sleepFunc(targetArtifactLockPollInterval) + pollInterval := targetArtifactLockPollInterval + if remaining < pollInterval { + pollInterval = remaining + } + sleepFunc(pollInterval) state, err := r.targetArtifactState() if err != nil { diff --git a/internal/ome-agent/replica/replica_test.go b/internal/ome-agent/replica/replica_test.go index 6cbb9e2f8..626a4c2d3 100644 --- a/internal/ome-agent/replica/replica_test.go +++ b/internal/ome-agent/replica/replica_test.go @@ -1084,6 +1084,29 @@ func TestReplicaAgent_PrepareTargetArtifactUploadUsesOneWaitDeadline(t *testing. 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()