Repository navigation
feat(ess,reconciler): decouple artifact lifecycle management into ESS - #1440
Draft
janzak-dev wants to merge 15 commits into
Draft
janzak-dev wants to merge 15 commits into
janzak-dev wants to merge 15 commits into
Conversation
…nPhase and etag monotonicity fencing
…ation across ESS and Actuator
…ons to ess consumer
… id, and notify-when detection
There was a problem hiding this comment.
Code Review
This pull request introduces an artifact lifecycle notification framework that publishes state transitions of Kubernetes JobSets to Redis and processes them in the Experiment State Server (ESS) to manage external Hypercompute MLRun resources. It updates the task queue API, Kubernetes client, reconciler, and ESS handlers to propagate and synchronize the notify-when annotation. Feedback on the changes identifies a missing import alias for the task_queue package in createartifact.go that will cause a compilation error.
| xidmsgpb "xmanager_cloud/api/protos/xid" | ||
| "xmanager_cloud/internal/grpcerr" | ||
| "xmanager_cloud/internal/logging" | ||
| "xmanager_cloud/internal/task_queue" |
There was a problem hiding this comment.
The import of xmanager_cloud/internal/task_queue is missing the client alias, which is required because the package is referenced as client on line 207. This will cause a compilation error. Please add the client alias to the import.
Suggested change
| "xmanager_cloud/internal/task_queue" | |
| client "xmanager_cloud/internal/task_queue" |
…, and reuse standard modules TAG=agy CONV=e7c98da6-30d1-49fd-960a-981e85cfd781
- Scope NotFound retries to the notify-when JobSet patch only; restore global retry classification. Dispatch JobSet updates on the field mask via new internal/task_queue/updatemask constants. - Add internal/kubernetes/common/notifywhen as the single owner of the notify-when annotation format (Encode in ESS, Matches in Reconciler). - Move RedisLifecycleNotifier to cmd/reconciler/shared, build and close the Redis publisher in the reconciler composition root, and share one PublishLifecycleEvent between the batcher and deletion path. - Keep internal/storage free of reconciler notification hints by carrying them in a reconciler-side LifecycleHints envelope. - Consolidate MLRun identity, label, error and console URL helpers in internal/hypercomputeclient, with a shared client constructor; resolve MLRun project/location once in ESS and drop invented fallback IDs. - Move lifecycle glue into server/artifactlifecycle_hooks.go with a single superseded-execution finalizer and paginated artifact listing. - Split ArtifactLifecycleEvent into artifact_lifecycle_event.proto. - Add a shared Helm mlRun values block for ESS and Reconciler, plus docs. TAG=agy
- ManagedXprofRule returns Hypercompute errors so Redis redelivers the event; never persists a constructed URL after a failed create. - ESS records each execution's MLRun URL on WorkUnitExecution.metadata.mlrun_url; restarts close exactly the superseded execution's MLRun (failures are logged, restart succeeds); MLRuns created by late events for superseded executions are closed. - Reconciler publishes only on committed state changes or newly added notify-when; removes the deletion-path publish. notify-when tracking is keyed by JobSet UID and execution ID. - Resumption recomputes notify-when when replaying Create/Restart tasks. - With the feature flag off: artifact read path, ListArtifacts calls, and Reconciler Redis config/secret match main. - ESS calls Hypercompute with its own service account (no impersonation); CTK grants the Hypercompute role when the feature flag is on. TAG=agy
An MLRun deleted outside XManager has nothing left to close, so the terminal, superseded and superseded-guard close paths log it at info and succeed instead of retrying until the dead-letter queue. The active path still creates the run. TAG=agy
ESS sent artifacts keyed "xprof_gcs_path", which the control plane
rejects, and its request diverged from the legacy reconciler's. Build
the request in one place (hypercomputeclient.BuildManagedXprofCreateOptions)
for both producers: artifacts {"gcs_path": folder}, the xprof tool, GKE
JobSet workload details with the full cluster path and creation time,
and the legacy gating (GCS folder, display name and run group required;
display name and run group fall back to the ManagedXprof artifact).
TAG=agy
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Decouples work-unit artifact lifecycle management (
ManagedXprof/ HypercomputeMLRunsynchronization) from the KubernetesJobSetReconcilerinto the Experiment State Server (ESS) via an event-driven Redis notification stream and a pluggableLifecycleRuleregistry, gated behind theartifact_lifecycle_enabled(XMC_ARTIFACT_LIFECYCLE_ENABLED, defaultfalse) feature flag for safe dual-path rollout.Motivation & Context
Previously,
JobSetReconcilerdirectly queried Spanner forManagedXprofartifacts and called the Hypercompute Cluster Director API on every reconciliation loop, coupling cluster-level controller reconciliation with external artifact API calls and lacking transactional concurrency protection across work unit restarts. Furthermore,GetArtifactandListArtifactsdynamically constructedxprof_urlon every read rather than persisting it on theArtifactrow.This change implements the approved Artifact Lifecycle architecture behind
XMC_ARTIFACT_LIFECYCLE_ENABLED:xmc.google.com/notify-when): When enabled, ESS computes the union of execution states requested by attachedLifecycleRuleimplementations and stampsxmc.google.com/notify-whenonto the KubernetesJobSet(both at creation time and dynamically via Actuator when artifacts are attached to running work units).UpdateBatcherrecords committed(OldState, NewState)transitions from Spanner and publishesArtifactLifecycleEventmessages to Redis streamxmc-artifact-lifecycle-streamonly when a state transition (or newly patchednotify-whenannotation) matches theJobSet'snotify-whenfilter.LifecycleRuleRegistry & Fenced Persistence: ESS consumesArtifactLifecycleEventtasks, creates/updates HypercomputeMLRunresources withrunPhaseandetagmonotonicity fencing, and persistsxprof_urlonto the SpannerArtifactrow inside a read-write transaction fenced onevent.EID == CurrentWorkUnitExecution.EID.RestartWorkUnitinvokesOnExecutionSupersededfor the superseded execution ID before scheduling the new execution.XMC_ARTIFACT_LIFECYCLE_ENABLED): WhenXMC_ARTIFACT_LIFECYCLE_ENABLEDisfalse(default), Reconciler continues running legacysyncMLRunand ESS leaves itsLifecycleRuleregistry empty and falls back toenrichManagedXprof; whentrue, ESS owns theManagedXproflifecycle end-to-end and Reconciler skips legacysyncMLRunin favor of edge-triggered Redis notifications.Key Changes
internal/featureflags: AddedArtifactLifecycleFlagName(XMC_ARTIFACT_LIFECYCLE_ENABLED, defaultfalse) andFlags.ArtifactLifecycleEnabled.internal/hypercomputeclient: Relocated fromcmd/reconcilerand addedrunPhaseterminal-state monotonicity fencing plusetagpreservation onUpdateMLRun. It is now the single source of truth for MLRun identity hashing (ConstructMLRunID), labels,IsAlreadyExists, console/response URL handling, andNewControlPlaneClientFromConfig(with impersonation), shared by ESS and the Reconciler.internal/kubernetes/common/notifywhen: Owns thexmc.google.com/notify-whenwire format (Encodeused by ESS,Matchesused by the Reconciler).internal/task_queue/updatemask: Exported update-mask path constants shared by the ESS enqueue client and the Actuator.internal/storage: Added transactionalArtifactTxRepository(FetchArtifact,UpdateArtifact) tostorage.Transactionacross Spanner and PostgreSQL.WorkUnitExecutionUpdatecarries no notification metadata.api/protos/task_queue/artifact_lifecycle_event.proto: NewArtifactLifecycleEvent/JobSetRuntimeContextmessages for the Reconciler → ESS stream (same proto package astask.proto).cmd/experiment_state_server/artifactlifecycle:LifecycleRuleregistry,ManagedXprofRule, EID transaction fencing (fencing.go), paginated work-unit artifact listing, Redis stream consumer (xmc-artifact-lifecycle-stream) and worker.cmd/experiment_state_server/server: Lifecycle hooks (artifactlifecycle_hooks.go: notify-when computation, late-attach sync, single superseded-execution finalizer used byRestartWorkUnitandBatchRestartWorkUnits).GetArtifact/ListArtifactsreturn persistedxprof_urlwith fallback toenrichManagedXprof. MLRun project/location/environment are resolved once at startup; ESS uses its own service account (no impersonation). Each execution's MLRun URL is recorded onWorkUnitExecution.metadata.mlrun_url; restarts close exactly the superseded execution's MLRun (failures are logged and never fail the restart). Hypercompute errors are returned so Redis redelivers the event. Resumption recomputes notify-when when replaying Create/Restart tasks. With the flag off, the artifact read path matches main.cmd/actuator:notify_whenon create/restart and a mask-drivenUpdateWorkUnitpatch; a JobSetNotFoundis retried only for the notify-when patch (global retry classification unchanged).cmd/reconciler: Committed(OldState, NewState)transitions fromUpdateWorkUnitExecutionsWithTransitions; notify-when predicate watching and edge detection carried in a reconciler-sideLifecycleHintsenvelope;shared.RedisLifecycleNotifierbuilt and closed in the manager composition root; notify-when tracking keyed by JobSet UID and execution ID; publishes only fromUpdateBatcheron committed state changes or newly added notify-when (no deletion-time publish); legacysyncMLRun(includingmetadata.mlrun_url) kept when the flag is off.deploy/helm&deploy/ctk_v2:featureFlags.artifact_lifecycle_enabled, Redis settings for the Reconciler, and a shared top-levelmlRun:values block rendered into both the ESS and Reconciler ConfigMaps. With the flag off, the Reconciler needs no Redis address or secret. CTK grantsroles/hypercomputecluster.editorand enables the Hypercompute API when the flag is on.Assumptions
kubernetes/aiml/run/{location}/{id}), and the API response URL is read from the MLRun resource's top-levelurlfield. Previously persisted/runs/{id}URLs are still parsed.Verification & Testing
./.githooks/pre-commitbazel test //internal/featureflags/... //internal/hypercomputeclient/... //internal/storage/... //internal/workunit/... //cmd/experiment_state_server/... //cmd/actuator/... //cmd/reconciler/... //tools/release/... //test/helm/...118/118 PASSED), plus Helm chart tests (test/helm, including newmlrun_config_test.py).ManagedXprofArtifact): Verifiedxmc.google.com/notify-when: STARTING.RUNNING.COMPLETED.FAILED.STOPPEDstamped onJobSet, Reconciler Redis event publication, ESSMLRuncreation, andxprof_urlpersistence/readback on the SpannerArtifactrow.ManagedXprofArtifact): Verified liveJobSetannotation patching viaCreateArtifact->UpdateNotifyWhenAction->Actuator.UpdateWorkUnit, ReconcilerForceNotifyedge detection, andxprof_urlpopulation.RestartWorkUnitSuperseding & Fencing): Verified synchronousOnExecutionSupersededinvocation forEID 1, transactional fencing rejecting staleEID 1updates (eventEID=1, curEID=2), and clean transition toEID 2.Related Issues / Bugs
Bug: 578366632
TAG=agy
CONV=e7c98da6-30d1-49fd-960a-981e85cfd781