Repository navigation
Conversation
Add a dedicated architecture reference document (`docs/architecture/jobset_and_work_unit_states.md`) detailing how Kubernetes `JobSet`, Kueue `Workload`, and owned `Pod` signals map to XMC `WorkUnitExecution` and `WorkUnit` states, how state transitions are guarded in storage, how the intent and observation loops synchronize state, and the end-to-end sequences of events across lifecycle transitions. Also link the new reference from the architecture overview, execution model, and root README documentation index. TAG=agy CONV=76617b8d-6898-4554-b6b1-758e3ba90593 Change-Id: I798bca6d7469f2c3c5acee33a6136e8ce41f699f
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive new documentation file, jobset_and_work_unit_states.md, detailing the mapping, transition guards, and synchronization mechanics between Kubernetes JobSet states and XMC WorkUnit states. It also updates existing documentation files (README.md, execution_model.md, and overview.md) to reference this new guide. The feedback highlights a broken relative link in the new document that should be updated to point to the correct Spanner storage path.
| Every call to [`UpdateWorkUnitExecutions`](../../internal/workunit/updateworkunitexecutions.go) | ||
| runs inside a read-write database transaction (`RunInTransaction`) and evaluates | ||
| [`ShouldUpdateWorkUnitExecutionState`](../../internal/storage/state/state.go) and | ||
| [`MergeWorkUnitExecution`](../../internal/storage/data/workunitexecutions.go) before persisting |
There was a problem hiding this comment.
The relative link to MergeWorkUnitExecution points to ../../internal/storage/data/workunitexecutions.go, but the actual Spanner storage data files are located under internal/storage/spanner/data/. Please update the path to ../../internal/storage/spanner/data/workunitexecutions.go to ensure the link is correct and consistent with the references on lines 51-53.
| [`MergeWorkUnitExecution`](../../internal/storage/data/workunitexecutions.go) before persisting | |
| [`MergeWorkUnitExecution`](../../internal/storage/spanner/data/workunitexecutions.go) before persisting |
There was a problem hiding this comment.
Keeping ../../internal/storage/data/workunitexecutions.go as written (Push Back — Technical Verification):
MergeWorkUnitExecution is defined in the database-agnostic package internal/storage/data/workunitexecutions.go (line 34, package data), which is shared across storage backends (UpdateWorkUnitExecutions in internal/workunit/updateworkunitexecutions.go calls data.MergeWorkUnitExecution). By contrast, internal/storage/spanner/data/workunitexecutions.go (package spannerdata, referenced on lines 51–53) handles Spanner row-to-proto scanning (WorkUnitExecutionRowToProto / WorkUnitCurrentExecutionRowToProto) and does not define MergeWorkUnitExecution.
| Kueue sets `JobSet.spec.suspend = true` and terminates active pods; as `totalReady` and | ||
| `totalActive` drop to `0`, `JobSetReconciler` transitions the execution directly from `RUNNING` (or | ||
| `STARTING`) back to `SCHEDULING` via Rule 9 on the **same `eid` and same `JobSet`**. |
There was a problem hiding this comment.
[Priority: Actionable] docs/architecture/jobset_and_work_unit_states.md:223-225
Observation: The documentation states that preemption always drops totalReady and totalActive to 0 and transitions the execution to SCHEDULING via Rule 9. However, ToWorkUnitState() evaluates Rule 8b (totalSucceeded > 0) before Rule 9, meaning if a workload had successful pods prior to preemption, it remains in RUNNING.
Risk / Rule: Accuracy of edge cases and boundary conditions in architecture documentation.
Suggestion:
| Kueue sets `JobSet.spec.suspend = true` and terminates active pods; as `totalReady` and | |
| `totalActive` drop to `0`, `JobSetReconciler` transitions the execution directly from `RUNNING` (or | |
| `STARTING`) back to `SCHEDULING` via Rule 9 on the **same `eid` and same `JobSet`**. | |
| Kueue sets `JobSet.spec.suspend = true` and terminates active pods; as `totalReady` and | |
| `totalActive` drop to `0` (with `Failed == 0` and `Succeeded == 0`), `JobSetReconciler` transitions | |
| the execution directly from `RUNNING` (or `STARTING`) back to `SCHEDULING` via Rule 9 on the **same | |
| `eid` and same `JobSet`**. (If the evicted workload already had successes, it remains `RUNNING` via Rule 8b). |
There was a problem hiding this comment.
Addressed in 53c84bb: qualified the preemption mapping in §3.2, §4.1, and §6.2 to note that Rule 9 (SCHEDULING) applies when Failed == 0 && Succeeded == 0, whereas evicted workloads with prior completed jobs stay RUNNING via Rule 8b (or PARTIALLY_FAILED via Rule 8a when totalFailed > 0 without target SuccessPolicy satisfied).
| - Cluster-backed execution: `{INITIALIZING | SCHEDULING | STARTING | RUNNING | PARTIALLY_FAILED | RESTARTING} → STOPPING → STOPPED` | ||
| - Paused shortcut (no `JobSet` on cluster): `SCHEDULING_PAUSED → STOPPED` (state committed directly in ESS; `Delete` task still closes `Operation.done = true` upon `NotFound`) | ||
|
|
||
| ```mermaid |
There was a problem hiding this comment.
This diagram is currently broken
There was a problem hiding this comment.
fixed. GitHub seems to get confused when there is a ; in description
There was a problem hiding this comment.
🤖 Incremental PR Approval
PR author is a code owner and changes since last approval (2 lines) are within the 50-line grace delta.
Preserving prior approval by @janzak-dev on commit 6e34de9.
|
|
||
| | Writer Component | Batcher Source | Mutates `status.state`? | Masked Fields Written | | ||
| | :--- | :--- | :---: | :--- | | ||
| | **ESS** (`Create` / `Stop` / `Restart`) | Direct transaction | **Yes** | `INITIALIZING`, `STOPPING`, `STOPPED` (only from `SCHEDULING_PAUSED`), `RESTARTING` (on `eid+1`). | |
There was a problem hiding this comment.
[Priority: Actionable] ESS Create and Restart write full execution rows. They never call the masked UpdateWorkUnitExecutions (cmd/experiment_state_server/server/utils.go:99-115, restartworkunitexecutions.go:99-147). Only Stop uses a mask, ["status.state"] (batchstopworkunits.go:163-172). ESS also writes status.system_message on enqueue errors (utils.go:173).
There was a problem hiding this comment.
Addressed in 53c84bb: split the ESS entry in §3.3 into full-row writes (CreateWorkUnit via InsertWorkUnit and RestartWorkUnit via SaveWorkUnitExecution + UpdateCurrentWorkUnitExecution) versus masked UpdateWorkUnitExecutions calls (Stop with ["status.state"] and handleEnqueueActionError with ["status.system_message"] and ForceStatusMessageUpdate: true).
There was a problem hiding this comment.
Thanks. The split between full-row writes for Create/Restart and masked writes for Stop is correct now. One new inaccuracy came in with the fix:
handleEnqueueActionErrorpassesForceStatusMessageUpdate: false, nottrue(cmd/experiment_state_server/server/utils.go:189-193), and no production caller sets it totrue.- It's also only reached from the Stop paths. Create and BatchRestart enqueue failures write a FATAL message by calling
updateWorkUnitExecutionStatusMessagedirectly (createworkunit.go:164,batchrestartworkunits.go:142).
Suggested wording: "Enqueue failures write ["status.system_message"] (no force): via handleEnqueueActionError on Stop, and via updateWorkUnitExecutionStatusMessage on Create/BatchRestart."
| ESS->>Redis: Publish trigger on xmc-scheduling-resume-stream | ||
|
|
||
| ESS->>DB: Resumption worker reads PausedTasks batches | ||
| ESS->>ESS: Filter out missing WUs, curEID > taskEID, or state != SCHEDULING_PAUSED |
There was a problem hiding this comment.
[Priority: Actionable] The resume filter has no state != SCHEDULING_PAUSED check. It skips a task when (a) the WU has a newer EID (> taskEID, or > taskEID+1 for Restart), (b) the WU is COMPLETED/STOPPED/STOPPING/FAILED, or (c) it's an Update that is stale or the WU is still paused. Missing WUs drop out through the SQL JOIN (filter.go:46-110). Also L697: PausedTasks rows are soft-marked via MarkProcessed (setting ResolutionStatus/ProcessedTime), not deleted.
There was a problem hiding this comment.
Addressed in 53c84bb: updated §6.5 to reflect the JOIN WorkUnits lookup for missing work units, the exact ShouldSkipTask skip conditions (newer curEID, COMPLETED/STOPPED/STOPPING/FAILED, and stale or still-paused Update tasks), and soft-marking via PausedTasks.MarkProcessed (ResolutionStatus, ProcessedTime, ResumptionRunId, 90-day TTL).
There was a problem hiding this comment.
Thanks, the JOIN, ShouldSkipTask and MarkProcessed parts are correct now. One error remains: the PausedTasks TTL is 7 days after ProcessedTime, not 90 days. Sources:
internal/storage/spanner/schema/ess-schema.sql:624:TIMESTAMP_ADD(ProcessedTime, INTERVAL 7 DAY)- the row-deletion policy at
:628 - migration
013_pause_scheduling.up.sql:73 - the Postgres schema (
INTERVAL '7 days')
Please change L795 to PausedTasks.MarkProcessed (ResolutionStatus, ProcessedTime, ResumptionRunId; row TTL 7d after ProcessedTime).
| | `JobSet` | `metadata.deletionTimestamp` & `metadata.finalizers` | Non-nil timestamp + `jobset.xmanager.cloud/finalizer` | Triggers synchronous terminal state persistence (`STOPPED`, unless `Completed`/`Failed` is already `True`) before finalizer removal. | | ||
| | `JobSet` | `spec.suspend` | `true` / `false` | Set `true` on creation; flipped `false` by Kueue upon quota admission and flipped back to `true` on Kueue eviction/preemption. | | ||
| | `JobSet` | `status.conditions` | `Completed` (`JobSetCompleted`), `Failed` (`JobSetFailed`), `Suspended` (`JobSetSuspended`) | `Completed=True` → `COMPLETED`; `Failed=True` → `FAILED`. `Suspended=False` `LastTransitionTime` anchors the startup pod error timer. | | ||
| | `JobSet` | `status.replicatedJobsStatus[]` | Summed across jobs into `JobTotals`: `Ready`, `Active`, `Succeeded`, `Failed`, `Replicas` | Drives `SCHEDULING`, `STARTING`, `RUNNING`, and `PARTIALLY_FAILED`, and populates `WorkUnitStatus.detailed_status.job_statuses`. | |
There was a problem hiding this comment.
[Priority: Nit] Bundled nits (one comment): L144 the field is detailed_status.jobset_status, not job_statuses, and Replicas comes from the spec. L146/L164/L539 the image-pull check covers all regular containers (init containers skipped), and the 2m is measured from admission. L147 Kueue Admitted/Evicted are never read; the inputs are QuotaReserved=False/Inadmissible, status.clusterName and podSetAssignments. L148 the field is positionInClusterQueue (stored +1). L36-40/L128 the experiment rollup also depends on launch_state. L166 ERROR messages defer to an existing message of the same severity. L238 the Inadmissible logic is in workload/reconciler.go. L314 the EID fence is a silent continue, not logged. L331 the function is the unexported restartWorkUnitExecutions. L334 a status.state mask without state_update_time fails the whole batch with InvalidArgument. L366 complete_time is time.Now(), not the transaction timestamp. L491 the constant is maxStoppedWriteRetryAge. L494 permanent errors also include FailedPrecondition and PermissionDenied. L451 done=true is also written on failure (task goes to the DLQ). L466 "3-level coalescing" is really latest-wins per slot inside a 3-level index.
There was a problem hiding this comment.
Addressed all 15 items in 53c84bb: (1) detailed_status.jobset_status and Replicas from spec, (2) startup image-pull check across all regular containers (skipping initContainers) after >= 2m from admission (JobSetSuspended=False), (3) Kueue inputs (QuotaReserved=False/Inadmissible, status.clusterName, podSetAssignments), (4) positionInClusterQueue (stored +1), (5) launch_state in UpdateExperimentActiveState (§1 & §2.2), (6) DeferToExistingMessage on ERROR messages, (7) workload/reconciler.go link in §3.3, (8) silent continue on EID fence, (9) unexported restartWorkUnitExecutions, (10) InvalidArgument batch failure when status.state omits state_update_time, (11) complete_time set via time.Now(), (12) maxStoppedWriteRetryAge, (13) all 4 permanent error codes (NotFound, InvalidArgument, FailedPrecondition, PermissionDenied), (14) Operations.done = true with error on failure and DLQ xmc-tasks-dlq, and (15) UpdateBatcher latest-wins replacement per (xid, wid, source) slot in its 3-level index.
There was a problem hiding this comment.
Thanks, 14 of the 15 nits now look correct. A few new small issues came in with 53c84bb1:
- L526 (wrong name): the dead-letter stream is
xmc-dlq-stream(RedisDLQ,internal/kubernetes/common/constants/constants.go:119, wired into the actuator consumer atcmd/actuator/prod.go:116), notxmc-tasks-dlq. - L507: Create does run
preconditionsValidator.ValidateJobSetManifest(cmd/actuator/k8s_client/jobset.go:71,93). It only skips the EID and last-edit guards. Suggest "Create skips the EID/last-edit guards" instead of "skipspreconditions_validator". - L380-382: only a nil
state_update_timewith astatus.statemask fails the batch withInvalidArgument(internal/storage/data/workunitexecutions.go:199). A non-nil zero value drops the state change silently (internal/storage/state/state.go:267). - L748:
k8sEid == Nreads like a precondition. The actual guard isk8sEid <= N(preconditions_validator.go:209) plus last-edit ordering, the same as §6.3 L692. SuggestDelete old JobSet (Preconditions: UID, k8sEid <= N, last-edit). - L155 (optional): Kueue's admission webhook only suspends at creation. Unsuspending on admission and re-suspending on eviction are done by Kueue's job reconciler (jobframework), not the webhook.
|
[Priority: Actionable] Request: run XMC Iterative Reflection on the authoring session A fact-check of this doc against
Since this doc was agent-authored ( |
| | :--- | :--- | :---: | :--- | | ||
| | **ESS** (`Create` / `Stop` / `Restart`) | Direct transaction | **Yes** | `INITIALIZING`, `STOPPING`, `STOPPED` (only from `SCHEDULING_PAUSED`), `RESTARTING` (on `eid+1`). | | ||
| | **Actuator** (`TaskPauseHandler`) | Direct transaction | **Yes** | `status.state = SCHEDULING_PAUSED`, `status.state_update_time` (only when global scheduling is paused). | | ||
| | **Actuator** (`WorkUnitTaskHandler`) | Direct (`Operations` repo) | **No** | Marks `Operations.done = true` (or error); never writes `RUNNING` or `STOPPED`. | |
There was a problem hiding this comment.
[Priority: Actionable] WorkUnitTaskHandler does not exist in the codebase (cmd/actuator/workunittaskhandler defines Handler, which has no storage.Storage dependency). Operations updates (done, response/error on Create/Delete/Restart) are written by k8sclient.Client.updateOperation (cmd/actuator/k8s_client/client.go:291-318), not via UpdateWorkUnitExecutions. Consider moving this to a note below the table or referencing k8sclient.Client.updateOperation.
There was a problem hiding this comment.
Addressed in 53c84bb: removed WorkUnitTaskHandler and added a note under §3.3 clarifying that workunittaskhandler.Handler has no WorkUnitExecution client and that k8sclient.Client.updateOperation updates Operations (done, response/error) directly.
| ### 5.3 Synchronous Finalizer Bypass on `JobSet` Deletion | ||
|
|
||
| Every XMC `JobSet` is created with finalizer `jobset.xmanager.cloud/finalizer` | ||
| ([`k8s_client/jobset.go`](../../cmd/actuator/k8s_client/jobset.go)). When `deletionTimestamp != nil`, |
There was a problem hiding this comment.
[Priority: Nit] jobset.xmanager.cloud/finalizer (constants.JobSetFinalizer) is attached in buildBaseJobSet inside cmd/actuator/k8s_client/jobset_builder.go:270 rather than k8s_client/jobset.go.
There was a problem hiding this comment.
Addressed in 53c84bb: updated the §5.3 finalizer reference to buildBaseJobSet in cmd/actuator/k8s_client/jobset_builder.go.
| | Writer Component | Batcher Source | Mutates `status.state`? | Masked Fields Written | | ||
| | :--- | :--- | :---: | :--- | | ||
| | **ESS** (`Create` / `Stop` / `Restart`) | Direct transaction | **Yes** | `INITIALIZING`, `STOPPING`, `STOPPED` (only from `SCHEDULING_PAUSED`), `RESTARTING` (on `eid+1`). | | ||
| | **Actuator** (`TaskPauseHandler`) | Direct transaction | **Yes** | `status.state = SCHEDULING_PAUSED`, `status.state_update_time` (only when global scheduling is paused). | |
There was a problem hiding this comment.
[Priority: Actionable] TaskPauseHandler (cmd/actuator/pausescheduling/pause_handler.go:208-221) masks ["status.state", "status.system_message"] (writing an INFO system message alongside SCHEDULING_PAUSED, with state_update_time carried implicitly by status.state), and runs this execution update only on Create and Restart (paused Update tasks only write PausedTasks). Also mention status.system_message among the shared mask paths above.
There was a problem hiding this comment.
Addressed in 53c84bb: updated §3.3 to list ["status.state", "status.system_message"] for TaskPauseHandler (executed on Create and Restart, while paused Update tasks only write PausedTasks) and included status.system_message among the shared mask paths.
| 3. `newState` is in `WorkUnitStatusRetainingStates` (`{STOPPED, FAILED}`), preserving the diagnostic | ||
| error or user log context that preceded failure or stop. | ||
| - **Audit & Message Tables**: Any change to `WorkUnitExecution.status` appends a snapshot row to | ||
| `StatusHistories` ([`internal/workunit/history.go`](../../internal/workunit/history.go)), and new |
There was a problem hiding this comment.
[Priority: Nit] recordStatusMessages (internal/workunit/history.go:129-140) mirrors only system_message and user_message (excluding scheduling_information.message) on any accepted execution change—even without a state transition—and Experiment.status.message_severity_counts tracks only the single highest-severity {system_message, user_message} per work unit (GetBestMessageFromStatus).
There was a problem hiding this comment.
Addressed in 53c84bb: clarified in §4.3 that recordStatusMessages mirrors only system_message and user_message on any accepted execution change (even when oldState == newState) and that message_severity_counts tallies each work unit's single highest-severity message via GetBestMessageFromStatus.
… doc Align `docs/architecture/jobset_and_work_unit_states.md` with exact Go, Protobuf, and SQL contracts across all open review threads and self-review verification findings: - Scope inline/interleaved `WorkUnitExecutions`, stored generated `ExecutionState`, and read-time `ACTIVE`/`INACTIVE` stamping to Cloud Spanner, noting PostgreSQL DDL and `Unimplemented` repository status. - Remove non-existent `GetWorkUnitExecution`, `WorkUnitTaskHandler`, `metadata.job_manifest`, `status.state_update_time` mask path, `StatusHistories`, `ResumeWorkloadsScheduling`, `max_concurrent_flushes`, and `finalizerMaxRetryDuration` references, replacing them with exact symbols, flags, and RPCs. - Clarify `JobSet.spec.suspend` creation/webhook ownership, startup pod diagnostics container scope and admission timer anchor, Kueue watched fields (`QuotaReserved`, `clusterName`, `podSetAssignments`, `positionInClusterQueue`), preemption Rule 8b/8a vs. Rule 9 edge cases, shared `UpdateMask` conflict resolution, `ShouldResetStatusMessage` replacement semantics, `WorkUnitStatusHistory` vs. `StatusMessages` triggers, `UpdateBatcher` latest-wins slot indexing, `Stop` vs. `Create`/`Restart` persistence ordering, `Delete`/`Restart` precondition guards, `STOPPING` + `Completed`/`Failed` finalizer interaction, and Health Service `ResumeScheduling` / `PausedTasks.MarkProcessed` flows. Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) Comment link: #1411 (comment) TAG=agy CONV=76617b8d-6898-4554-b6b1-758e3ba90593 Change-Id: If32c87bbb2b0290eabdb0bf379acb14c98c136f6
53c84bb
| | Writer Component | Write Path / Batcher Source | Mutates `status.state`? | Fields / `UpdateMask` Paths Written | | ||
| | :--- | :--- | :---: | :--- | | ||
| | **ESS** (`Create` / `Restart`) | Direct full-row write (not `UpdateWorkUnitExecutions`) | **Yes** | `InsertWorkUnit` writes `INITIALIZING` on `eid=1` ([`server/utils.go`](../../cmd/experiment_state_server/server/utils.go)); `SaveWorkUnitExecution` + `UpdateCurrentWorkUnitExecution` archive `eid=N` and write `RESTARTING` on `eid=N+1` ([`restartworkunitexecutions.go`](../../cmd/experiment_state_server/server/restartworkunitexecutions.go)). | | ||
| | **ESS** (`Stop` / enqueue error) | Direct `UpdateWorkUnitExecutions` | **Yes** (`Stop` only) | `StopWorkUnit` / `BatchStopWorkUnits` masks `["status.state"]` (`STOPPING`, or `STOPPED` from `SCHEDULING_PAUSED`, [`batchstopworkunits.go`](../../cmd/experiment_state_server/server/batchstopworkunits.go)); Redis enqueue failure (`handleEnqueueActionError`, [`server/utils.go`](../../cmd/experiment_state_server/server/utils.go)) masks `["status.system_message"]` (`ForceStatusMessageUpdate: true`). | |
There was a problem hiding this comment.
[Priority: Actionable] docs/architecture/jobset_and_work_unit_states.md:260
Observation: The ESS (Stop / enqueue error) row states that handleEnqueueActionError sets ForceStatusMessageUpdate: true and implies it is the sole enqueue-failure writer. In cmd/experiment_state_server/server/utils.go:160-196, updateWorkUnitExecutionStatusMessage sets ForceStatusMessageUpdate: forceUpdate and handleEnqueueActionError passes false (called by StopWorkUnit, BatchStopWorkUnits, and RestartWorkUnit), while CreateWorkUnit (createworkunit.go:164-168) and BatchRestartWorkUnits (batchrestartworkunits.go:142-146) call updateWorkUnitExecutionStatusMessage(..., false) directly.
Risk / Rule: Documentation Alignment (.agents/rules/languages/markdown.md — Accurate Contract Representation).
Suggestion:
| | **ESS** (`Stop` / enqueue error) | Direct `UpdateWorkUnitExecutions` | **Yes** (`Stop` only) | `StopWorkUnit` / `BatchStopWorkUnits` masks `["status.state"]` (`STOPPING`, or `STOPPED` from `SCHEDULING_PAUSED`, [`batchstopworkunits.go`](../../cmd/experiment_state_server/server/batchstopworkunits.go)); Redis enqueue failure (`handleEnqueueActionError`, [`server/utils.go`](../../cmd/experiment_state_server/server/utils.go)) masks `["status.system_message"]` (`ForceStatusMessageUpdate: true`). | | |
| | **ESS** (`Stop` / enqueue error) | Direct `UpdateWorkUnitExecutions` | **Yes** (`Stop` only) | `StopWorkUnit` / `BatchStopWorkUnits` masks `["status.state"]` (`STOPPING`, or `STOPPED` from `SCHEDULING_PAUSED`, [`batchstopworkunits.go`](../../cmd/experiment_state_server/server/batchstopworkunits.go)); Redis enqueue failures (`updateWorkUnitExecutionStatusMessage` / `handleEnqueueActionError`, [`server/utils.go`](../../cmd/experiment_state_server/server/utils.go)) mask `["status.system_message"]` (`ForceStatusMessageUpdate: false`). | |
| - Whenever `"status.state"` is included in `UpdateMask`, | ||
| `newExecution.status.state_update_time` must be non-nil and valid — omitting it fails the | ||
| **entire batch transaction** with `InvalidArgument` in `MergeWorkUnitExecution`. When present, | ||
| it must satisfy `newTime >= currentTime` (`!currentTime.After(newTime)`). | ||
| - **Exception (`WorkUnitTimeIgnoredStates = {STOPPING}`)**: `STOPPING` bypasses the timestamp | ||
| ordering check (though a valid `state_update_time` is still required) so a user `StopWorkUnit` | ||
| request on ESS is never rejected if a GKE Reconciler node's clock was slightly ahead when | ||
| writing `RUNNING`. |
There was a problem hiding this comment.
[Priority: Actionable] docs/architecture/jobset_and_work_unit_states.md:380-387
Observation: Section 4.2 item 5 states that any missing or invalid state_update_time fails the entire batch transaction with InvalidArgument. In internal/storage/data/workunitexecutions.go:197-202, mergeWorkUnitExecutionStatusPath only returns InvalidArgument when state_update_time == nil; a non-nil zero timestamp (AsTime().IsZero()) causes ShouldUpdateWorkUnitExecutionState (internal/storage/state/state.go:267-270) to return false, nil, silently dropping the state change without failing the batch transaction.
Risk / Rule: Documentation Alignment (.agents/rules/languages/markdown.md — Accurate Contract Representation).
Suggestion:
| - Whenever `"status.state"` is included in `UpdateMask`, | |
| `newExecution.status.state_update_time` must be non-nil and valid — omitting it fails the | |
| **entire batch transaction** with `InvalidArgument` in `MergeWorkUnitExecution`. When present, | |
| it must satisfy `newTime >= currentTime` (`!currentTime.After(newTime)`). | |
| - **Exception (`WorkUnitTimeIgnoredStates = {STOPPING}`)**: `STOPPING` bypasses the timestamp | |
| ordering check (though a valid `state_update_time` is still required) so a user `StopWorkUnit` | |
| request on ESS is never rejected if a GKE Reconciler node's clock was slightly ahead when | |
| writing `RUNNING`. | |
| - Whenever `"status.state"` is included in `UpdateMask`, a `nil` | |
| `newExecution.status.state_update_time` fails the **entire batch transaction** with | |
| `InvalidArgument` in `MergeWorkUnitExecution`, whereas a non-nil zero timestamp | |
| (`AsTime().IsZero()`) causes `ShouldUpdateWorkUnitExecutionState` to return `false, nil` | |
| (dropping the state change without failing the batch). When non-zero, it must satisfy | |
| `newTime >= currentTime` (`!currentTime.After(newTime)`). | |
| - **Exception (`WorkUnitTimeIgnoredStates = {STOPPING}`)**: `STOPPING` bypasses the timestamp | |
| ordering check (though a non-nil, non-zero `state_update_time` is still required) so a user | |
| `StopWorkUnit` request on ESS is never rejected if a GKE Reconciler node's clock was slightly | |
| ahead when writing `RUNNING`. |
| - **Parent Experiment Aggregation (on `state` change)**: `DenormalizeWorkUnitStatus` decrements | ||
| `oldState` and increments `newState` in `Experiment.status.child_state_counts`. | ||
| `UpdateExperimentActiveState` recomputes `Experiment.status.active_state` (`ACTIVE` while | ||
| `LaunchState == LAUNCHING` or at least one work unit is non-terminal; `INACTIVE` once | ||
| `LaunchState == FAILED` or all child work units are in `{STOPPED, COMPLETED, FAILED}`, setting | ||
| `Experiment.complete_time`). |
There was a problem hiding this comment.
[Priority: Clarity] docs/architecture/jobset_and_work_unit_states.md:438-443
Observation: Section 4.3 attributes setting Experiment.complete_time to UpdateExperimentActiveState (including when LaunchState == FAILED). In internal/storage/state/state.go:174-193, UpdateExperimentActiveState only updates exp.Status.ActiveState; Experiment.complete_time is set in updateExperimentFromExecutions (internal/workunit/updateworkunitexecutions.go:255-264) when an execution update transitions active_state from ACTIVE to INACTIVE (whereas UpdateExperimentLaunchState in internal/experiment/updateexperimentlaunchstate.go:66-91 does not set complete_time).
Risk / Rule: Documentation Alignment (.agents/rules/languages/markdown.md — Accurate Contract Representation).
Suggestion:
| - **Parent Experiment Aggregation (on `state` change)**: `DenormalizeWorkUnitStatus` decrements | |
| `oldState` and increments `newState` in `Experiment.status.child_state_counts`. | |
| `UpdateExperimentActiveState` recomputes `Experiment.status.active_state` (`ACTIVE` while | |
| `LaunchState == LAUNCHING` or at least one work unit is non-terminal; `INACTIVE` once | |
| `LaunchState == FAILED` or all child work units are in `{STOPPED, COMPLETED, FAILED}`, setting | |
| `Experiment.complete_time`). | |
| - **Parent Experiment Aggregation (on `state` change)**: `DenormalizeWorkUnitStatus` decrements | |
| `oldState` and increments `newState` in `Experiment.status.child_state_counts` and calls | |
| `UpdateExperimentActiveState` to recompute `Experiment.status.active_state` (`ACTIVE` while | |
| `LaunchState == LAUNCHING` or at least one work unit is non-terminal; `INACTIVE` once | |
| `LaunchState == FAILED` or all child work units are in `{STOPPED, COMPLETED, FAILED}`). When | |
| `updateExperimentFromExecutions` observes `active_state` transition from `ACTIVE` to `INACTIVE`, | |
| it sets `Experiment.complete_time` to `timestamppb.New(time.Now())`. |
| - **`Create` (`AlreadyExists` verification)**: `Create` skips `preconditions_validator` and | ||
| submits the `JobSet` directly; if Kubernetes returns `AlreadyExists`, the Actuator invokes | ||
| `verifyDuplicateJobSet` ([`k8s_client/client.go`](../../cmd/actuator/k8s_client/client.go)) and | ||
| treats a matching `JobSet` as an idempotent success. | ||
| - **`Delete` and `Restart` (`xmc.google.com/execution-id` & `xmc.google.com/last-edit` guards)**: | ||
| Evaluated by | ||
| [`preconditions_validator`](../../cmd/actuator/k8s_client/preconditions_validator/preconditions_validator.go) | ||
| **only on `Delete` and `Restart`**. If `k8sEid > targetEid`, the task is skipped as superseded; | ||
| `last-edit` rejects tasks whose `operationTime` predates the `JobSet`'s `xmc.google.com/last-edit` | ||
| label. On `Restart`, `targetEid` is the **outgoing `eid = N`** (while `N + 1` travels inside | ||
| `payload.WorkUnit`): if `k8sEid == N`, the Actuator deletes `eid=N` and creates `eid=N+1`; on a | ||
| redelivered retry after `eid=N+1` exists, `k8sEid (N+1) > targetEid (N)` skips idempotently. |
There was a problem hiding this comment.
[Priority: Actionable] docs/architecture/jobset_and_work_unit_states.md:507-518
Observation: Section 5.1 step 2 states that Create skips preconditions_validator and that Restart deletes the old JobSet when k8sEid == N. In cmd/actuator/k8s_client/jobset.go:71,87-98, createJobSet still runs c.preconditionsValidator.ValidateJobSetManifest(jobSetUnstructured) (hostNetwork exclusive-topology validation) and only skips the live-cluster EID and last-edit guards (IsActuationRequestValid / IsExecutionIDAligned), while Restart (preconditions_validator.go:209) permits deletion whenever k8sEid <= N.
Risk / Rule: Documentation Alignment (.agents/rules/languages/markdown.md — Accurate Contract Representation).
Suggestion:
| - **`Create` (`AlreadyExists` verification)**: `Create` skips `preconditions_validator` and | |
| submits the `JobSet` directly; if Kubernetes returns `AlreadyExists`, the Actuator invokes | |
| `verifyDuplicateJobSet` ([`k8s_client/client.go`](../../cmd/actuator/k8s_client/client.go)) and | |
| treats a matching `JobSet` as an idempotent success. | |
| - **`Delete` and `Restart` (`xmc.google.com/execution-id` & `xmc.google.com/last-edit` guards)**: | |
| Evaluated by | |
| [`preconditions_validator`](../../cmd/actuator/k8s_client/preconditions_validator/preconditions_validator.go) | |
| **only on `Delete` and `Restart`**. If `k8sEid > targetEid`, the task is skipped as superseded; | |
| `last-edit` rejects tasks whose `operationTime` predates the `JobSet`'s `xmc.google.com/last-edit` | |
| label. On `Restart`, `targetEid` is the **outgoing `eid = N`** (while `N + 1` travels inside | |
| `payload.WorkUnit`): if `k8sEid == N`, the Actuator deletes `eid=N` and creates `eid=N+1`; on a | |
| redelivered retry after `eid=N+1` exists, `k8sEid (N+1) > targetEid (N)` skips idempotently. | |
| - **`Create` (manifest validation & `AlreadyExists` verification)**: `Create` validates the | |
| constructed manifest via `preconditionsValidator.ValidateJobSetManifest` (`hostNetwork` | |
| exclusive-topology check) while skipping the live-cluster EID and `last-edit` guards; if | |
| Kubernetes returns `AlreadyExists`, the Actuator invokes `verifyDuplicateJobSet` | |
| ([`k8s_client/client.go`](../../cmd/actuator/k8s_client/client.go)) and treats a matching | |
| `JobSet` as an idempotent success. | |
| - **`Delete` and `Restart` (`xmc.google.com/execution-id` & `xmc.google.com/last-edit` guards)**: | |
| `IsActuationRequestValid` and `IsExecutionIDAligned` in | |
| [`preconditions_validator`](../../cmd/actuator/k8s_client/preconditions_validator/preconditions_validator.go) | |
| run **only on `Delete` and `Restart`**. If `k8sEid > targetEid`, the task is skipped as superseded; | |
| `last-edit` rejects tasks whose `operationTime` predates the `JobSet`'s `xmc.google.com/last-edit` | |
| label. On `Restart`, `targetEid` is the **outgoing `eid = N`** (while `N + 1` travels inside | |
| `payload.WorkUnit`): if `k8sEid <= N`, the Actuator deletes the old `JobSet` and creates `eid=N+1`; on a | |
| redelivered retry after `eid=N+1` exists, `k8sEid (N+1) > targetEid (N)` skips idempotently. |
| `error` on failure—and the consumer `XACK`s the message (routing exhausted/failed tasks to | ||
| dead-letter stream `xmc-tasks-dlq`). Cluster execution states (`SCHEDULING`, `STARTING`, |
There was a problem hiding this comment.
[Priority: Actionable] docs/architecture/jobset_and_work_unit_states.md:525-526
Observation: Section 5.1 step 3 refers to the Redis dead-letter stream as xmc-tasks-dlq, whereas constants.RedisDLQ wired into redisclient.NewConsumer (internal/kubernetes/common/constants/constants.go:119, cmd/actuator/prod.go:123) is xmc-dlq-stream.
Risk / Rule: Documentation Alignment (.agents/rules/languages/markdown.md — Accurate Contract Representation).
Suggestion:
| `error` on failure—and the consumer `XACK`s the message (routing exhausted/failed tasks to | |
| dead-letter stream `xmc-tasks-dlq`). Cluster execution states (`SCHEDULING`, `STARTING`, | |
| `error` on failure—and the consumer `XACK`s the message (routing exhausted/failed tasks to | |
| dead-letter stream `xmc-dlq-stream`). Cluster execution states (`SCHEDULING`, `STARTING`, |
| ESS-->>Client: Operation (done=false) | ||
|
|
||
| Redis->>Act: XREADGROUP Restart task (target eid=N) | ||
| Act->>K8s: Delete old JobSet (k8sEid == N, UID precondition) |
There was a problem hiding this comment.
[Priority: Clarity] docs/architecture/jobset_and_work_unit_states.md:748
Observation: In the §6.4 sequence diagram, step 8 writes k8sEid == N, UID precondition. validateAndDeleteJobSetForRestart and deleteJobSet (cmd/actuator/k8s_client/jobset.go:123-143,217-238, preconditions_validator.go:209) enforce k8sEid <= N alongside last-edit ordering and the UID precondition, matching §6.3 line 692.
Risk / Rule: Documentation Alignment (.agents/rules/languages/markdown.md — Accurate Contract Representation).
Suggestion:
| Act->>K8s: Delete old JobSet (k8sEid == N, UID precondition) | |
| Act->>K8s: Delete old JobSet ( |
| ESS->>DB: Read unprocessed PausedTasks joined with WorkUnits (drops missing WUs) | ||
| ESS->>ESS: ShouldSkipTask skips newer curEID, COMPLETED/STOPPED/STOPPING/FAILED, or stale/paused Update | ||
| ESS->>Redis: Re-sign fresh JWT/OIDC and XADD non-skipped tasks to xmc-tasks-stream | ||
| ESS->>DB: PausedTasks.MarkProcessed (ResolutionStatus, ProcessedTime, ResumptionRunId, 90d TTL) |
There was a problem hiding this comment.
[Priority: Actionable] docs/architecture/jobset_and_work_unit_states.md:795
Observation: The §6.5 sequence diagram states that PausedTasks.MarkProcessed sets a 90d TTL. Both internal/storage/spanner/schema/ess-schema.sql:623-628 and internal/storage/postgres/schema/ess-schema.sql:359-360 define PausedTasks.ExpirationTime as ProcessedTime + 7 days (INTERVAL 7 DAY / INTERVAL '7 days'), whereas 90 days is the TTL for Operations.
Risk / Rule: Documentation Alignment (.agents/rules/languages/markdown.md — Accurate Contract Representation).
Suggestion:
| ESS->>DB: PausedTasks.MarkProcessed (ResolutionStatus, ProcessedTime, ResumptionRunId, 90d TTL) | |
| ESS->>DB: PausedTasks.MarkProcessed (ResolutionStatus, ProcessedTime, ResumptionRunId, 7d TTL) |
There was a problem hiding this comment.
Please address the comments
Summary
Adds a technical architecture document (
docs/architecture/jobset_and_work_unit_states.md) specifying how KubernetesJobSet,WorkUnitExecution, andWorkUnitstates relate, map, synchronize, and transition across the XMC control and execution planes. Updates the root and architecture documentation indices to cross-reference the new guide.Motivation & Context
XMC tracks execution state across three distinct layers—cluster
JobSet/Workloadresources, numberedWorkUnitExecutionattempts (WorkUnitExecutionStateslot activity +WorkUnitStatus.WorkUnitState), and parentWorkUnit/Experimentprojections—synchronized asynchronously through Redis task streams, the Actuator, Reconciler controllers,UpdateBatcher, and synchronous deletion finalizers. Consolidating the exact evaluation priority order, storage monotonicity/transition guards, field-level conflict resolution, and end-to-end event sequences in one reference makes debugging and extending workload lifecycle behavior straightforward without reverse-engineering multiple binaries.Key Changes
docs/architecture/jobset_and_work_unit_states.mdcovering:WorkUnitsinlineACTIVEexecution vs.WorkUnitExecutionsinterleavedINACTIVEhistory, plus current PostgreSQL DDL/repository status).WorkUnitExecutionState, all 13WorkUnitStatus.WorkUnitStatevalues, writer ownership, execution-terminal semantics vs.RestartWorkUnitrestartability, and observed Kubernetes/JobSet/Kueue/Podsignals.JobSet→WorkUnitStatemapping decision tree, control-plane-only states,DESCHEDULINGproto vs.RUNNING → SCHEDULINGpreemption behavior, and execution writerUpdateMaskpaths and conflict resolution.ShouldUpdateWorkUnitExecutionState/MergeWorkUnitExecution(EID fencing, forward/backward rules,StateUpdateTimemonotonicity withSTOPPINGskew bypass, partial merges, and automatic reset/retention side effects).UpdateBatcherlatest-wins slot indexing, and synchronousjobset.xmanager.cloud/finalizerbypass) plus end-to-end walkthroughs and sequence diagrams for create, preemption/eviction, stop, restart, and global scheduling pause/resume.docs/architecture/jobset_and_work_unit_states.mdindocs/architecture/overview.md,docs/architecture/execution_model.md, andREADME.md.Assumptions
docs/architecture/overview.md,docs/architecture/execution_model.md, andREADME.mdalongsidedocs/architecture/jobset_and_work_unit_states.md— repository documentation alignment guidelines require keeping architecture index tables synchronized when adding new docs — if a strictly single-file PR is preferred, the three index link additions can be reverted without affecting the new document.Verification & Testing
./.githooks/pre-commit,python3 ./tools/run_affected_tests.py --dry-run,python3 ./tools/analyze_diff_size.py --base-ref origin/main...Related Issues / Bugs
N/A
TAG=agy
CONV=76617b8d-6898-4554-b6b1-758e3ba90593