Skip to content

Commit d28f2f5

Browse files
mkuznet1coketaste
andauthored
Aicomnet dev slurm multi (#142)
* feat(preparation): amd-smi JSON tolerance, manifest runtime-context restore, secrets precedence, lightweight-CLI isolation Stage A of porting the aicomnet multi-node local-image work onto the slurm_multi pipeline (develop). Preparation-phase changes only: - core/context: parse amd-smi JSON via JSONDecoder.raw_decode so trailing non-JSON text (deprecation banners, double-emitted blocks under concurrent slurm tasks) no longer aborts run-phase init. - orchestration/run_orchestrator: restore host-level runtime context (docker_mounts, docker_build_arg, docker_gpus, gpu_vendor, guest_os) from the manifest on execution, keeping runtime-detected values as priority; docker_env_vars from the manifest no longer overwrite runtime values already populated by Context (e.g. MAD_SECRETS_* read from os.environ), so unexpanded "${VAR}" manifest placeholders cannot clobber the resolved secret. - core/constants: --version/--help lightweight CLI invocations skip MODEL_DIR setup and keep output clean, avoiding global-state pollution of subsequent runs in the same shell. * feat(build): multi-node local-image build-on-primary + shared-tar distribution Stage B of porting the aicomnet multi-node local-image work onto the slurm_multi pipeline. Adds local-image preparation to the per-node run path (run_models_from_manifest, which slurm_multi invokes on every node in docker target mode): - _ensure_local_image_available: primary node (rank 0) ensures the local image is present (building it from the manifest dockerfile if missing), and when MAD_DOCKER_BUILDS points at a shared dir, docker-saves it into a shared tar; worker nodes load that tar instead of pulling. Local images are not in any registry, so upstreams registry-only pull path cannot reach them on workers. - Atomic tar write (sibling .tmp + os.replace) so peers never load a half-written tar. - _build_local_image_from_manifest / _get_build_args / _build_or_pull_local_image with shlex.quote shell-escaping for all path/image/build-arg components (Console.sh runs with shell=True). - Strict NODE_RANK/RANK and NNODES/WORLD_SIZE parsing (malformed values raise instead of silently defaulting to 0/1 and diverging). - _sync_after_local_image_ready: Stage B uses a shared-filesystem readiness wait (workers without the image poll for the primary tar); Stage C replaces this with a deterministic TCP barrier. Replaces the previous verify-then-pull block, which could only pull registry images and left local multi-node images unreachable on worker nodes. * feat(running): deterministic TCP image-ready barrier + docker mount de-duplication Stage C of porting the aicomnet multi-node local-image work onto the slurm_multi pipeline. Replaces the Stage B shared-FS readiness poll with a robust TCP rendezvous and adds docker mount de-duplication. - _tcp_image_ready_barrier / _recv_line: NODE_RANK=0 listens on a port range derived from MASTER_PORT+SLURM_JOB_ID; workers send "READY <token> <rank>" and wait for "GO". No shared-FS visibility required (avoids NFS/Weka metadata lag). Hardening: MAD_BARRIER_TOKEN secret (defaults JOB<id>); bind to MASTER_ADDR resolved IP, skipping loopback (127/8, ::1) so an /etc/hosts loopback mapping cannot leave the listener on lo while workers connect over the routable IP; line-based recv (partial-read safe); worker-rank normalization so NODE_RANK="01" matches the int ACK; strict rank / nnodes parsing; on master timeout the ACKed worker sockets are closed (EOF, not half-open); accepted-peer list logged for deadlock diagnosis. - _sync_after_local_image_ready now drives the TCP barrier instead of polling the shared filesystem. - get_mount_arg de-duplicates against -v/--volume targets already present in additional_docker_run_options (via _extract_additional_mount_targets) so docker does not reject "Duplicate mount point". SLURM->container env pass-through and --version/--help MODEL_DIR isolation were already covered (develop slurm_env_vars list; Stage A constants). * feat(postprocessing): best-non-empty multi-node results CSV selection + on-failure container diagnostics Stage D of porting the aicomnet multi-node local-image work onto the slurm_multi pipeline. Covers result aggregation and failure triage across nodes: - slurm.py collect_results: instead of taking the first node_<rank> directory that happens to hold the multiple_results CSV, gather every per-node candidate (plus the job-dir copy) and pass them to _select_best_multiple_results_csv, which ranks by the count of non-empty "performance" rows and breaks ties on total row count. In multi-node runs the master CSV is frequently present but empty while a worker holds the real throughput numbers, so first-match could aggregate empty data; this defers empty-perf handling to the richest file. Header/row keys are stripped so a leading-space CSV header still matches. - container_runner.py run path: wrap the model-script exec in a try/except that, on RuntimeError, parses the failed container id from the error and snapshots bounded diagnostics inside it (process table, listening ports, and tails of /run_logs and the model dir logs) via Console.sh so they land in the run log next to the failure. All diagnostic calls are time-bounded and non-fatal; the original error is re-raised unchanged. * fix(running): defer empty-perf verdict to login-node aggregation in SLURM in-job runs In multi-node SLURM runs the per-node in-job status check marked the job FAILED when the designated collector node (MAD_COLLECT_METRICS=true) had an empty multiple_results CSV, even though Primus emits throughput only on the last global rank, which frequently lives on a different node than the collector. With skip_perf_collection set, the login-node collect_results already selects the richest per-node CSV (_select_best_multiple_results_csv) and writes the authoritative perf/status record, so an empty local perf is not authoritative and must not hard-fail the node (exit 3 -> whole job fails). Add a skip_perf_collection guard to the status determination (and the except fallback). Real failures remain caught first by the error-pattern scan; genuinely-empty multi-node runs still get a FAILURE row from the login-node aggregation. * fix(running): shell-escape image_name in local-image validate/pull (Copilot review) Resolves Copilot review (High): in _create_manifest_from_local_image, image_name was interpolated into "docker image inspect" / "docker pull" commands passed to Console.sh (shell=True) without quoting -- a command-injection risk that also breaks valid image names containing shell-special characters. Escape it once via shlex.quote and use the quoted form in both shell calls (display strings keep the raw name). Matches the shlex.quote escaping already used on the Stage B/C container_runner shell paths. * fix(slurm): reconcile local image across nodes and rebuild stale images Multi-node local-image runs trusted the Docker tag alone, allowing two failure modes that silently ran mismatched software across ranks: - Cross-node identity: a worker holding an old image under the same tag kept it while rank 0 (re)built a different image, because workers only acted when the tag was absent. Rank 0 now broadcasts the ensured image ID over the rendezvous barrier (extended GO line) and workers force- reload rank 0's tar from MAD_DOCKER_BUILDS whenever their local image ID differs, failing hard if the mismatch persists. - Content staleness: reusing a tag for a different build (e.g. a new RCCL commit) was never rebuilt. Local builds now bake a mad.build_fingerprint label (hash of Dockerfile + docker_build_arg); a present-but-differing label triggers a rebuild and tar refresh. Images without the label (pulled or pre-existing) are left untouched to avoid spurious rebuilds. When no shared tar is configured, workers fall back to fingerprint-based content correctness instead of byte-identical ID reconciliation. Adds unit tests for GO-line parsing, build fingerprinting, staleness detection, and primary/worker reconciliation paths. * fix(tcp-barrier): catch _recv_line errors on master side On the master side of the TCP barrier, _recv_line(conn) could raise socket.timeout or other socket errors after accept(); the outer except socket.timeout only covered server.accept() itself, so a slow or misbehaving peer could abort the entire barrier. Wrap _recv_line in its own try/except Exception: close the connection and continue accepting remaining peers instead of propagating. * fix(gpu-mapping): restore kfd_renderDs None guard for ROCm pre-6.4.1 The legacy unique_id path has no amd-smi fallback for renderD discovery; only the 6.4.1+ branch can recover when kfd_renderDs is empty. Without this guard an inaccessible KFD topology surfaces as an obscure len(None) TypeError instead of a clear, actionable error. Restores the guard dropped in 3ee0500. --------- Co-authored-by: Stephen Shao <yu.shao@amd.com>
1 parent ef87d05 commit d28f2f5

6 files changed

Lines changed: 1045 additions & 34 deletions

File tree

src/madengine/core/constants.py

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,17 +26,28 @@
2626
import os
2727
import json
2828
import logging
29+
import sys
2930

3031

3132
# Utility function for optional verbose logging of configuration
3233
def _log_config_info(message: str, force_print: bool = False):
3334
"""Log configuration information either to logger or print if specified."""
35+
# Keep --version/--help output clean even if MAD_VERBOSE_CONFIG=true.
36+
if any(arg in {"--version", "-V", "--help", "-h"} for arg in sys.argv[1:]):
37+
logging.debug(message)
38+
return
3439
if force_print or os.environ.get("MAD_VERBOSE_CONFIG", "").lower() == "true":
3540
print(message)
3641
else:
3742
logging.debug(message)
3843

3944

45+
def _is_lightweight_cli_invocation() -> bool:
46+
"""Return True for metadata/help invocations that should avoid side effects."""
47+
lightweight_flags = {"--version", "-V", "--help", "-h"}
48+
return any(arg in lightweight_flags for arg in sys.argv[1:])
49+
50+
4051
# third-party modules
4152
from madengine.core.console import Console
4253

@@ -65,9 +76,12 @@ def _setup_model_dir():
6576
_log_config_info(f"Model dir: {MODEL_DIR} copied to current dir: {cwd_abs}")
6677

6778

68-
# Only setup model directory if explicitly requested (when not just importing for constants)
79+
# Only setup model directory if explicitly requested and invocation is not metadata-only.
6980
if os.environ.get("MAD_SETUP_MODEL_DIR", "").lower() == "true":
70-
_setup_model_dir()
81+
if _is_lightweight_cli_invocation():
82+
_log_config_info("Skipping MODEL_DIR setup for lightweight CLI invocation (--version/--help).")
83+
else:
84+
_setup_model_dir()
7185

7286
# madengine credentials configuration
7387
CRED_FILE = "credential.json"

src/madengine/core/context.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -816,7 +816,10 @@ def get_gpu_renderD_nodes(self) -> typing.Optional[typing.List[int]]:
816816
json_output = output
817817

818818
try:
819-
data = json.loads(json_output)
819+
# Use raw_decode so we tolerate any trailing non-JSON
820+
# text (deprecation banners, double-emitted blocks under
821+
# concurrent slurm tasks, stray newlines, etc.).
822+
data, _end = json.JSONDecoder().raw_decode(json_output.lstrip())
820823
except json.JSONDecodeError as e:
821824
raise ValueError(f"Failed to parse amd-smi JSON output: {e}. Output was: {output[:200]}")
822825

src/madengine/deployment/slurm.py

Lines changed: 66 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1632,6 +1632,60 @@ def _build_common_info_dict(
16321632
flatten_tags(result)
16331633
return result
16341634

1635+
def _select_best_multiple_results_csv(self, candidates: List[Path]) -> Optional[Path]:
1636+
"""Pick the CSV with the most non-empty performance entries.
1637+
1638+
In multi-node SLURM runs every node copies its local multi-results CSV
1639+
into job_dir/node_<rank>/. Only some nodes observe the final throughput
1640+
and populate the performance column; others have the file but with empty
1641+
values (e.g. master perf empty while a worker has the real numbers).
1642+
Ranking candidates by the count of non-empty performance rows lets
1643+
downstream aggregation use the richest data instead of depending on
1644+
node-0 winning the race or being non-empty. Header/row keys are stripped
1645+
so a leading space in the CSV header (some Primus configs) still matches.
1646+
Ties break on total row count; candidates[0] is the ultimate fallback.
1647+
"""
1648+
if not candidates:
1649+
return None
1650+
if len(candidates) == 1:
1651+
return candidates[0]
1652+
import csv as _csv
1653+
best_candidate: Optional[Path] = None
1654+
best_score = -1
1655+
best_rows = -1
1656+
for candidate in candidates:
1657+
non_empty_perf = 0
1658+
total_rows = 0
1659+
has_perf_column = False
1660+
try:
1661+
with open(candidate, "r", encoding="utf-8", errors="ignore") as f:
1662+
reader = _csv.DictReader(f)
1663+
fieldnames = reader.fieldnames or []
1664+
stripped_fields = [fn.strip() for fn in fieldnames]
1665+
has_perf_column = "performance" in stripped_fields
1666+
for row in reader:
1667+
total_rows += 1
1668+
if has_perf_column:
1669+
normalized_row = {(k.strip() if isinstance(k, str) else k): v for k, v in row.items()}
1670+
value = (normalized_row.get("performance") or "").strip()
1671+
if value:
1672+
non_empty_perf += 1
1673+
except Exception:
1674+
continue
1675+
score = non_empty_perf if has_perf_column else 0
1676+
if score > best_score or (score == best_score and total_rows > best_rows):
1677+
best_score = score
1678+
best_rows = total_rows
1679+
best_candidate = candidate
1680+
if best_candidate is None:
1681+
return candidates[0]
1682+
if best_score > 0:
1683+
self.console.print(
1684+
f"[dim] Selected multiple_results CSV with {best_score} non-empty performance rows: {best_candidate}[/dim]"
1685+
)
1686+
return best_candidate
1687+
1688+
16351689
def collect_results(self, deployment_id: str) -> Dict[str, Any]:
16361690
"""Collect performance results from SLURM output files.
16371691
@@ -1775,14 +1829,19 @@ def collect_results(self, deployment_id: str) -> Dict[str, Any]:
17751829
mult_res = model_info_for_entry.get("multiple_results")
17761830
if mult_res:
17771831
resolved_csv: Optional[Path] = None
1832+
# Multi-node: gather all node CSVs and pick the one with the most
1833+
# non-empty performance rows (master CSV may be empty while a worker
1834+
# holds the real numbers) instead of taking the first node that has
1835+
# the file.
1836+
candidates: List[Path] = []
17781837
if (job_dir / mult_res).is_file():
1779-
resolved_csv = job_dir / mult_res
1780-
else:
1781-
for i in range(self.nodes):
1782-
candidate = job_dir / f"node_{i}" / mult_res
1783-
if candidate.is_file():
1784-
resolved_csv = candidate
1785-
break
1838+
candidates.append(job_dir / mult_res)
1839+
for i in range(self.nodes):
1840+
per_node_candidate = job_dir / f"node_{i}" / mult_res
1841+
if per_node_candidate.is_file():
1842+
candidates.append(per_node_candidate)
1843+
if candidates:
1844+
resolved_csv = self._select_best_multiple_results_csv(candidates)
17861845
if not resolved_csv and Path(mult_res).is_file():
17871846
resolved_csv = Path(mult_res)
17881847
if not resolved_csv and Path("run_directory", mult_res).is_file():

0 commit comments

Comments
 (0)