Aicomnet dev slurm multi

[mkuznet1]

Summary

Closes the multi-node coordination gap in the SLURM slurm_multi path. The
upstream refactor made madengine multi-node observable (per-rank log staging,
login-node collect_results, partial manifest restore) but did not yet
coordinate which node owns the local image, when workers may enter
docker run, whose CSV the aggregator should trust, and which manifest fields a
worker needs to reconstruct the builder's environment.

This PR ports that coordination onto slurm_multi as four staged commits
(preparation → build → running → postprocessing) plus one follow-up fix.
Single-node behavior is unchanged — every new path is gated on
MAD_DOCKER_BUILDS, NNODES > 1, skip_perf_collection, manifest presence, or
the failure flow.

What's new

Multi-node coordination

• Shared local-image tar cache (MAD_DOCKER_BUILDS): rank 0 docker saves
  the locally-built image into a shared dir (atomic sibling-.tmp + os.replace);
  workers docker load it instead of pulling. Local images are not in any
  registry, so upstream's registry-only pull path cannot reach them on workers.
• Rebuild-from-manifest fallback on a worker where the local image is missing
  and no tar is available (_build_local_image_from_manifest, with shlex.quote
  escaping for every path/image/build-arg component since Console.sh runs with
  shell=True).
• Deterministic TCP image-ready barrier (_tcp_image_ready_barrier): rank 0
  listens on a port derived from MASTER_PORT+SLURM_JOB_ID; workers send
  READY <token> <rank> and wait for GO, so no one enters docker run against
  a half-loaded image. No shared FS required (avoids NFS/Weka metadata lag).
  Hardening: MAD_BARRIER_TOKEN secret, bind to the routable MASTER_ADDR IP
  (skips loopback so an /etc/hosts mapping can't strand the listener on lo),
  line-based partial-read-safe recv, worker-rank normalization, strict
  rank/nnodes parsing, ACKed-socket cleanup + peer logging on master timeout.
• Best-non-empty multi-node CSV picker in collect_results
  (_select_best_multiple_results_csv): gathers every per-node candidate and
  ranks by count of non-empty performance rows (tie-break on total rows), so
  the aggregator uses the worker that actually holds the throughput numbers
  instead of a present-but-empty master CSV. Header/row keys are stripped so a
  leading-space CSV header still matches.
• Manifest runtime-context restore on workers now covers docker_mounts,
  docker_build_arg, docker_gpus, gpu_vendor, guest_os — runtime-detected
  values keep priority.

Bug fixes on the restore / verdict paths

• docker_env_vars restore no longer overwrites os.environ-sourced values
  (e.g. MAD_SECRETS_HFTOKEN) with unexpanded ${VAR} literals from an old
  manifest. Fixes HF 401/403 on Primus / gated models on rerun.
• Empty-perf verdict deferred to the login node (skip_perf_collection):
  Primus emits throughput only on the last global rank, which often lives on a
  different node than the designated collector (MAD_COLLECT_METRICS=true). The
  per-node in-job check used to mark the whole job FAILED (exit 3) on the
  collector's empty local CSV; it now defers to the authoritative login-node
  aggregation. Real failures are still caught first by the error-pattern scan,
  and a genuinely-empty multi-node run still gets a FAILURE row from the login
  step.

Robustness exposed by scale-out

• amd-smi JSON tolerance: parse via JSONDecoder.raw_decode so trailing
  non-JSON (deprecation banners, double-emitted blocks under concurrent SLURM
  tasks) no longer aborts run-phase init.
• Side-effect-free lightweight CLI: madengine --version/--help skip
  MODEL_DIR setup, avoiding global-state pollution of later runs in the same
  shell.
• Duplicate-mount protection: get_mount_arg de-duplicates against
  -v/--volume targets already present in additional_docker_run_options so
  Docker doesn't reject "Duplicate mount point".
• Container-failure diagnostics: on a docker exec RuntimeError, a bounded,
  non-fatal snapshot (process table, listening sockets via ss/netstat/lsof,
  tails of /run_logs and the model-dir logs) is captured inside the failed
  container and lands in the run log next to the failure; the original error is
  re-raised unchanged.

Validation

• Per-stage 2-node validation on a real SLURM cluster: shared-tar save/load,
  TCP barrier release with the image removed on the worker, and throughput parity
  vs the single-image baseline.
• End-to-end 2-node Primus Megatron-LM Llama-3.1-8B (2 × 8 GPU, IB): worker
  rank correctly reports ~10.1k tok/s/GPU, ~522 TFLOPS/GPU, and the run is now
  marked SUCCESS (previously a false FAILURE on the collector's empty CSV).
• pytest tests/unit/test_slurm_multi.py — 20 passed. Integration tests cover the
  shared-tar paths (primary-saves, existing-tar-is-loaded, worker-waits-for-primary).

Test plan

• 2-node Primus Megatron-LM 8B with image absent on the worker (exercises
  build-on-primary → save → worker load → TCP barrier → SUCCESS verdict).
• Rerun on a gated/Primus model to confirm MAD_SECRETS_HFTOKEN is not
  clobbered by manifest ${VAR} placeholders (no 401/403).
• Single-node SLURM run to confirm unchanged behavior (gated paths inactive).
• Confirm collect_results selects the worker CSV and writes the aggregated
  perf record.

[Copilot]

Pull request overview

This PR adds missing multi-node coordination and robustness improvements across the SLURM slurm_multi execution path, ensuring workers can reliably obtain locally-built images, synchronize before docker run, and aggregate performance results correctly across nodes.

Changes:

• Restore additional runtime/build context from manifests (mounts/build args/env vars) while preserving runtime-detected values.
• Add multi-node support utilities in the container runner (shared image tar cache, TCP barrier, duplicate mount avoidance, and failure diagnostics).
• Improve SLURM result aggregation by selecting the “best” per-node multiple_results CSV and harden GPU discovery JSON parsing.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
src/madengine/orchestration/run_orchestrator.py	Adjusts local-image manifest creation and improves manifest context restore behavior.
src/madengine/execution/container_runner.py	Adds multi-node local-image coordination helpers, mount de-dup, diagnostics, and verdict behavior for deferred perf collection.
src/madengine/deployment/slurm.py	Picks the richest per-node multiple_results CSV during login-node aggregation.
src/madengine/core/context.py	Makes amd-smi JSON parsing tolerant of trailing non-JSON output.
src/madengine/core/constants.py	Avoids side effects (MODEL_DIR setup / verbose config prints) for lightweight CLI invocations like --help/--version.

Comments suppressed due to low confidence (1)

src/madengine/core/context.py:765

• When ROCm < 6.4.1 and KFD topology is not accessible, kfd_renderDs can remain None (the code explicitly tolerates KFD permission errors). The subsequent len(kfd_renderDs) will then raise TypeError, losing the clearer error and making the fallback path brittle. Add an explicit None guard before using kfd_renderDs in the ROCm<6.4.1 mapping branch.

                if len(kfd_unique_ids) != len(kfd_renderDs):
                    raise RuntimeError(
                        f"Mismatch between unique_ids count ({len(kfd_unique_ids)}) "
                        f"and renderDs count ({len(kfd_renderDs)})"
                    )

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.