refactor: resolve workload container by name with Containers[0] fallback; wire GMS by pointer

[galletas1712]

Overview:

One of three follow-up PRs closing reviewer feedback from #8194. This is the main-container resolution refactor: every site that hardcoded podSpec.Containers[0] as "the workload" is now resolved via a shared helper that prefers-by-name with Containers[0] fallback.

Companion PRs:

• PR 1 — fix(snapshot): require exec-form command for checkpoint containers (fix(snapshot): require exec-form command for checkpoint containers #8318, already open) — CodeRabbit B, C
• PR 3 — fix(operator): validate pod-info volume; operator review nits (next) — CodeRabbit A, G, H, J

GMS Python changes (F, M, N) deferred to a separate GMS Python cleanup PR already in flight.

Closed predecessor: #8299.

Architectural decisions

1. Prefer-by-name with Containers[0] fallback. Seven sites used to hardcode podSpec.Containers[0] as "the workload". That assumption breaks whenever ExtraPodSpec.PodSpec.Containers or a webhook-injected sidecar (istio-proxy, etc.) lands ahead of the workload in the slice. New contract:

ResolveMainContainer(podSpec):
  1. if no containers -> nil
  2. scan for container named "main" -> return it
  3. fall back to Containers[0]    (preserves pre-convention pods and tests)

The fallback is deliberate — dropping it would break any pod rendered before the naming convention or produced by non-operator tooling (snapshotctl manifests authored by hand, old fixtures, manual pods). snapshotctl tightens the rule further: in multi-container manifests it requires the worker to be explicitly named main, because in CLI-authored input the ambiguity is user-visible.

2. Operator and snapshot each own their own copy. The "main" container name convention predates the snapshot subsystem (LWS multinode enforcement, PR #3390, Oct 2025) and is owned by the operator. Operator code must not import snapshotprotocol helpers that operate on pod-spec shape. After this PR:

• commonconsts.MainContainerName + common.ResolveMainContainer are the authoritative operator-side copies. Operator callers (graph.go, checkpoint_job.go, checkpoint/podspec.go) use these.
• snapshotprotocol.MainContainerName + snapshotprotocol.ResolveMainContainer / ResolveMainContainerName are the snapshot-side copies. Snapshot agent, snapshotctl, and snapshot protocol internals use these.
• Both sides carry cross-reference comments pointing at the other copy.

Why duplicate instead of share: deploy/operator and deploy/snapshot are separate go.mods and cannot cross-import internal/ packages. Creating a third Go module purely for one constant and one 10-line function would cost more than the drift risk of duplication, especially given the cross-reference comments.

3. GMS wires onto the main container by pointer, before it lands in the slice. GenerateBasePodSpec now holds the main container as a local variable at the GMS wire point and passes &container directly to dra.ApplyClaim and gms.EnsureServerSidecar. Rationale:

• The main container exists as a local; searching the slice is unnecessary work.
• Removes graph.go's import of snapshotprotocol. Non-snapshot DGDs (most of them) no longer touch that package.
• Closes a latent bug: checkpoint_job.go's post-Commit 1 main-container lookup could put the workload at a non-zero index, but the old dra.ApplyClaim(podSpec, name) still stripped nvidia.com/gpu from Containers[0] and wrote the claim there. So GMS and the DRA claim could land on different containers. The new dra.ApplyClaim(podSpec, target *corev1.Container, name) signature matches gms.EnsureServerSidecar(podSpec, mainContainer), and both receive the caller-resolved pointer.

4. GMS ↔ failover ordering invariant. buildFailoverPod deletes the single main container and replaces it with engine-0 / engine-1 (DeepCopy + engine-specific overrides). Ordering must be:

build local `container` (name="main")
     |
     v
apply GMS wiring (DRA claim + shared volume + env + mount)
     |
     v
podSpec.Containers = append(..., container)
     |
     v
backend.UpdatePodSpec  (multinode init containers, etc.)
     |
     v
append frontend sidecar (if any)
     |
     v
buildFailoverPod: DeepCopy "main" -> engine-0, engine-1
                  engines inherit GMS volume/mount/env/claim transitively

Pinned by the new TestGenerateBasePodSpec_GMSWiredToMainBeforeFailover integration test.

5. PrepareRestorePodSpecForCheckpoint owns worker resolution. The helper no longer accepts an arbitrary container *corev1.Container pointer; it resolves the worker internally via ResolveMainContainer. Callers cannot pass a sidecar pointer by mistake. Closes the part of CodeRabbit I that was about API footguns.

Details:

Commit	Summary
fix(snapshot): resolve main container by name with Containers[0] fallback	Introduce snapshotprotocol.ResolveMainContainer / ResolveMainContainerName. Replace hardcoded index-0 sites in the operator and snapshot agent/snapshotctl. Make PrepareRestorePodSpecForCheckpoint own its own worker resolution.
refactor(operator): wire GMS onto main container by pointer, not by search	Reorder graph.go so GMS runs against the local container variable before it is appended to the slice; update dra.ApplyClaim signature to take a target container pointer; this closes a latent "GMS and DRA claim land on different containers" bug in checkpoint_job.go.
refactor(operator): keep MainContainerName on operator side; document the duplicate	Correct the provenance of the constant: operator is the owner, snapshot keeps a local copy with cross-reference docs. Also adds the GMS+failover integration test.
refactor(operator): move main-container resolver to operator common package	Finish decoupling: operator code no longer imports snapshotprotocol.ResolveMainContainer. New common.ResolveMainContainer lives in deploy/operator/internal/common/ with its own tests.

Where should the reviewer start?

1. deploy/snapshot/protocol/maincontainer.go and deploy/operator/internal/common/maincontainer.go — the two identical resolvers, each with tests. The cross-reference comments on both files spell out the "keep in sync" invariant.
2. deploy/operator/internal/dynamo/graph.go GMS block — verify the in-scope container local is passed directly, no snapshotprotocol import. The regression test TestGenerateBasePodSpec_GMSWiredToMainBeforeFailover pins the failover interaction.
3. deploy/operator/internal/dra/dra.go — signature change from (podSpec, claimName) to (podSpec, target *corev1.Container, claimName). Both callers updated. Tests updated.
4. deploy/snapshot/protocol/restore.go PrepareRestorePodSpecForCheckpoint — now resolves the worker internally; the sole operator caller in checkpoint/podspec.go is updated to not pass the pointer in.
5. deploy/snapshot/cmd/snapshotctl/kube.go — the tightened multi-container rule (requires explicit name: main when the manifest has multiple containers).

Tests

• New unit tests:
  • snapshotprotocol.ResolveMainContainer + ResolveMainContainerName — nil spec, empty containers, single-container-is-main, main-at-non-zero-index, fallback.
  • common.ResolveMainContainer — same cases on the operator side.
  • TestGenerateBasePodSpec_GMSWiredToMainBeforeFailover — asserts that after buildFailoverPod replaces the "main" container with engine-0/engine-1, both engines inherit the GMS shared-socket volume mount, GMS_SOCKET_DIR env, and DRA claim (via DeepCopy in buildEngineContainer), and that the pod-level gms-server init sidecar and shared volume land exactly once.
  • TestApplyClaim_TargetsCallerSuppliedContainer (replaces the old _AlwaysTargetsFirstContainer) — verifies the sidecar-at-index-0 case.
  • TestInjectCheckpointIntoPodSpec — added "resolves main even when not at index zero" case.
• All deploy/operator unit tests pass (go test against envtest assets).
• All deploy/snapshot tests pass.
• go vet ./... clean in both modules.
• pre-commit run --from-ref origin/main --to-ref HEAD passes.

Risk notes for deploy-codeowners

• Behavior preservation via fallback: any pod rendered before the naming convention (or produced by non-operator tooling) still resolves via Containers[0], so this refactor is safe to cherry-pick without auditing every pod-emitting path. The regression surface is "pods that have a sidecar at index 0 and a container named main at a non-zero index" — the fallback would have been wrong here pre-refactor anyway.
• dra.ApplyClaim is an internal-package function (deploy/operator/internal/dra/), not part of any published API surface. Signature change has zero external impact.
• PrepareRestorePodSpecForCheckpoint is exported from deploy/snapshot/protocol, but I am not aware of any out-of-tree callers. Its one in-repo caller (checkpoint/podspec.go:InjectCheckpointIntoPodSpec) is updated.
• No API changes to CRDs, no Go module or Python dependency changes.
• This PR is safe to merge before or after PR 1 (fix(snapshot): require exec-form command for checkpoint containers #8318). The two are independent — PR 1 keeps Containers[0] for its validator, PR 2 replaces it. Whichever lands second picks up the other's work via rebase.

Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

• Relates to: fix(gms): rewrite GMS checkpoint/restore operator support #8194 (addresses resolved CodeRabbit threads D, E, I)
• Replaces: fix(gms): follow-up fixes for #8194 CodeRabbit and reviewer comments #8299 (closed; split into three single-concern PRs)

[galletas1712]

Superseded by #8631. The underlying problem — snapshot helpers keying off Containers[0] / "main" — is resolved structurally in #8631 by making target-container selection an explicit mandatory annotation read by the protocol helpers themselves, so there is no fallback to harden and no implicit main-container convention to thread through the operator.