#!/usr/bin/env python3 """Positive preflight and completion checks for the bounded two-Worker verl GRPO run.""" from __future__ import annotations import argparse import hashlib import importlib import importlib.metadata import json import math import os import re import subprocess import sys from datetime import datetime, timezone from pathlib import Path from typing import Any EXPECTED_COMMIT = "ddd86f527a4af75095e4677b02b5aa272913a088" EXPECTED_DATASET_REVISION = "740312add88f781978c0658806c59bc2815b9866" EXPECTED_DATA_CHECKSUMS = "0150c256df1500d5edf5f5462434b287ada29a9674ddf1d59ef779629d6e8a4e" EXPECTED_MODEL_MANIFEST = "0ea1d330342b4b9efbd1c3648360fbc4c2e3b1d5abd6120e49ef837649001f7b" EXPECTED_GPU_NAME = "NVIDIA A100-SXM4-80GB" MIN_CUDA_12_DRIVER = (525, 60, 13) REQUIRED_SCALARS = ( "training/global_step", "actor/pg_loss", "actor/grad_norm", "actor/lr", "critic/score/mean", ) FATAL_PATTERNS = ( r"Traceback \(most recent call last\)", r"Error executing job with overrides", r"CUDA out of memory", r"OutOfMemoryError", r"RayTaskError", r"ActorDiedError", r"WorkerCrashedError", r"\bNCCL (?:WARN|ERROR)\b", r"\bnccl(?:UnhandledCudaError|SystemError|InternalError|InvalidArgument|InvalidUsage|RemoteError)\b", r"\bProcessGroupNCCL\b[^\n]*(?:abort|watchdog|timeout|timed out)", r"Segmentation fault", r"Killed(?:\s|$)", ) def sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as stream: for chunk in iter(lambda: stream.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def run(*command: str) -> str: return subprocess.run(command, check=True, text=True, capture_output=True).stdout.strip() def package_record(distribution: str, module_name: str | None = None) -> dict[str, object]: module = importlib.import_module(module_name or distribution.replace("-", "_")) try: distribution_version = importlib.metadata.version(distribution) except importlib.metadata.PackageNotFoundError: distribution_version = None return { "distribution_version": distribution_version, "module_version": getattr(module, "__version__", None), "module_file": getattr(module, "__file__", None), } def atomic_json(path: Path, data: dict[str, object]) -> None: temporary = path.with_suffix(path.suffix + ".tmp") temporary.write_text(json.dumps(data, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") os.replace(temporary, path) def driver_version_tuple(version: str) -> tuple[int, int, int]: parts = version.split(".") if not parts or any(not part.isdigit() for part in parts): raise RuntimeError(f"Unrecognized NVIDIA driver version: {version}") numbers = [int(part) for part in parts[:3]] numbers.extend([0] * (3 - len(numbers))) return tuple(numbers) # type: ignore[return-value] def ordered_nodes(ray: Any) -> list[dict[str, Any]]: nodes = [node for node in ray.nodes() if node.get("Alive")] current_node_id = str(ray.get_runtime_context().get_node_id()) head = [node for node in nodes if str(node.get("NodeID")) == current_node_id] if len(head) != 1: raise RuntimeError(f"Cannot identify the Ray Head: current={current_node_id} nodes={nodes}") workers = sorted( (node for node in nodes if node is not head[0]), key=lambda node: str(node.get("NodeManagerAddress", "")), ) return head + workers def preflight(args: argparse.Namespace) -> int: import ray from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy output_dir = args.output_dir.resolve() output_dir.mkdir(parents=True, exist_ok=True) sentinel = output_dir / "shared-storage-sentinel" sentinel.write_text(f"{args.run_id}\n", encoding="utf-8") ray.init(address="auto", ignore_reinit_error=True) nodes = ordered_nodes(ray) if len(nodes) != args.expected_nodes: raise RuntimeError(f"Expected {args.expected_nodes} live Ray nodes, found {len(nodes)}") cluster_resources = ray.cluster_resources() expected_total_gpus = args.expected_nodes * args.expected_gpus_per_node if int(cluster_resources.get("GPU", 0)) != expected_total_gpus: raise RuntimeError( f"Expected {expected_total_gpus} cluster GPUs, found {cluster_resources}" ) for index, node in enumerate(nodes): node_gpus = int(node.get("Resources", {}).get("GPU", 0)) if node_gpus != args.expected_gpus_per_node: raise RuntimeError(f"Node {index} exposes {node_gpus} GPUs: {node}") @ray.remote(num_cpus=1) def probe_node( ordinal: int, storage_root: str, verl_root: str, model_root: str, data_root: str, output_root: str, expected_node_id: str, ) -> dict[str, object]: import json as remote_json import socket from pathlib import Path as RemotePath import pyarrow.parquet as pq import ray as remote_ray def remote_run(*command: str) -> str: return subprocess.run(command, check=True, text=True, capture_output=True).stdout.strip() def read_one(path: RemotePath) -> int: with path.open("rb") as stream: return len(stream.read(1)) storage = RemotePath(storage_root) source = RemotePath(verl_root) model = RemotePath(model_root) data = RemotePath(data_root) output = RemotePath(output_root) actual_node_id = str(remote_ray.get_runtime_context().get_node_id()) if actual_node_id != expected_node_id: raise RuntimeError( f"Node {ordinal}: affinity target {expected_node_id}, ran on {actual_node_id}" ) sentinel_path = output / "shared-storage-sentinel" if sentinel_path.read_text(encoding="utf-8") != f"{args.run_id}\n": raise RuntimeError(f"Node {ordinal}: cannot read the shared-storage sentinel") marker_root = output / "preflight-nodes" marker_root.mkdir(parents=True, exist_ok=True) marker_path = marker_root / f"node-{ordinal}.json" marker_path.write_text( remote_json.dumps({"ordinal": ordinal, "node_id": actual_node_id}) + "\n", encoding="utf-8", ) source_commit = remote_run("git", "-C", str(source), "rev-parse", "HEAD") if source_commit != EXPECTED_COMMIT: raise RuntimeError(f"Node {ordinal}: unexpected verl commit {source_commit}") if remote_run("git", "-C", str(source), "status", "--porcelain"): raise RuntimeError(f"Node {ordinal}: the verl checkout is not clean") model_manifest = sha256(model / "SHA256SUMS") if model_manifest != EXPECTED_MODEL_MANIFEST: raise RuntimeError(f"Node {ordinal}: unexpected model manifest {model_manifest}") dataset_manifest_path = data / "dataset-manifest.json" dataset_manifest = remote_json.loads(dataset_manifest_path.read_text(encoding="utf-8")) if dataset_manifest.get("revision") != EXPECTED_DATASET_REVISION: raise RuntimeError(f"Node {ordinal}: unexpected dataset revision") data_checksums = sha256(data / "SHA256SUMS") if data_checksums != EXPECTED_DATA_CHECKSUMS: raise RuntimeError(f"Node {ordinal}: unexpected dataset checksums") datasets: dict[str, dict[str, object]] = {} for split, expected_rows in (("train", 7473), ("test", 1319)): path = data / f"{split}.parquet" rows = pq.read_metadata(path).num_rows if rows != expected_rows or read_one(path) != 1: raise RuntimeError(f"Node {ordinal}: invalid {split} parquet") datasets[split] = {"path": str(path), "rows": rows, "bytes": path.stat().st_size} if read_one(model / "model.safetensors") != 1: raise RuntimeError(f"Node {ordinal}: model weights are not readable") gpu_rows = remote_run( "nvidia-smi", "--query-gpu=driver_version,name,memory.total,uuid", "--format=csv,noheader,nounits", ).splitlines() parsed_gpus = [] for row in gpu_rows: driver, name, memory, uuid = (part.strip() for part in row.split(",", 3)) parsed_gpus.append( {"driver": driver, "name": name, "memory_mib": int(memory), "uuid": uuid} ) if len(parsed_gpus) != args.expected_gpus_per_node: raise RuntimeError(f"Node {ordinal}: unexpected GPUs {parsed_gpus}") if any(gpu["name"] != EXPECTED_GPU_NAME for gpu in parsed_gpus): raise RuntimeError(f"Node {ordinal}: unexpected GPU model {parsed_gpus}") mount = remote_json.loads( remote_run("findmnt", "-J", "-T", str(storage), "-o", "TARGET,SOURCE,FSTYPE,OPTIONS") )["filesystems"][0] if mount["fstype"] == "overlay" or "rw" not in mount["options"].split(","): raise RuntimeError(f"Node {ordinal}: storage is not a writable shared mount {mount}") if os.environ.get("NCCL_IB_DISABLE") != "1": raise RuntimeError(f"Node {ordinal}: NCCL_IB_DISABLE must be 1") packages = { "torch": package_record("torch"), "ray": package_record("ray"), "transformers": package_record("transformers"), "sglang": package_record("sglang"), "pyarrow": package_record("pyarrow"), "tensorboard": package_record("tensorboard"), "verl": package_record("verl"), } return { "ordinal": ordinal, "node_id": actual_node_id, "node_ip": remote_ray.util.get_node_ip_address(), "hostname": socket.gethostname(), "gpus": parsed_gpus, "mount": mount, "python": {"version": sys.version, "executable": sys.executable}, "packages": packages, "verl_source": {"path": str(source), "commit": source_commit, "clean": True}, "model": {"path": str(model), "manifest_sha256": model_manifest}, "dataset": { "path": str(data), "checksums_sha256": data_checksums, "manifest": dataset_manifest, "splits": datasets, }, } probes = [] for ordinal, node in enumerate(nodes): strategy = NodeAffinitySchedulingStrategy(node_id=str(node["NodeID"]), soft=False) probes.append( probe_node.options(scheduling_strategy=strategy).remote( ordinal, str(args.storage_root), str(args.verl_root), str(args.model_root), str(args.data_root), str(output_dir), str(node["NodeID"]), ) ) workers = ray.get(probes) drivers = {gpu["driver"] for worker in workers for gpu in worker["gpus"]} model_manifests = {worker["model"]["manifest_sha256"] for worker in workers} dataset_manifests = {worker["dataset"]["checksums_sha256"] for worker in workers} package_identities = { json.dumps( {"python": worker["python"], "packages": worker["packages"]}, sort_keys=True, default=str, ) for worker in workers } if min(driver_version_tuple(version) for version in drivers) < MIN_CUDA_12_DRIVER: raise RuntimeError(f"Unsupported NVIDIA driver versions: {sorted(drivers)}") if len(model_manifests) != 1 or len(dataset_manifests) != 1: raise RuntimeError("Input identity differs across Ray nodes") if len(package_identities) != 1: raise RuntimeError("Python package identities differ across Ray nodes") marker_root = output_dir / "preflight-nodes" marker_records = [ json.loads((marker_root / f"node-{ordinal}.json").read_text(encoding="utf-8")) for ordinal in range(args.expected_nodes) ] if {record["node_id"] for record in marker_records} != { worker["node_id"] for worker in workers }: raise RuntimeError("The Driver cannot read both per-node shared-storage markers") for worker in workers: print( "WORKER_TOPOLOGY " + json.dumps( { "ordinal": worker["ordinal"], "node_id": worker["node_id"], "node_ip": worker["node_ip"], "hostname": worker["hostname"], "gpu_names": [gpu["name"] for gpu in worker["gpus"]], "driver": worker["gpus"][0]["driver"], "mount_source": worker["mount"]["source"], }, sort_keys=True, ) ) record = { "status": "passed", "phase": "preflight", "checked_at": datetime.now(timezone.utc).isoformat(), "run_id": args.run_id, "target_steps": args.target_steps, "train_batch_size": args.train_batch_size, "rollout_n": args.rollout_n, "expected_nodes": args.expected_nodes, "expected_gpus_per_node": args.expected_gpus_per_node, "cluster_resources": cluster_resources, "driver_versions": sorted(drivers), "ray_context": { "driver_node_id": str(ray.get_runtime_context().get_node_id()), "address": os.environ.get("RAY_ADDRESS", "auto"), }, "workers": workers, "shared_storage_markers": marker_records, } atomic_json(output_dir / "preflight.json", record) print(json.dumps(record, indent=2, ensure_ascii=False)) ray.shutdown() return 0 def load_scalars(tensorboard_root: Path) -> tuple[list[Path], dict[str, list[dict[str, float | int]]]]: from tensorboard.backend.event_processing.event_accumulator import EventAccumulator event_files = sorted(tensorboard_root.rglob("events.out.tfevents.*")) if not event_files or any(path.stat().st_size == 0 for path in event_files): raise RuntimeError("No non-empty TensorBoard event file was produced") scalars: dict[str, list[dict[str, float | int]]] = {} for event_file in event_files: accumulator = EventAccumulator(str(event_file), size_guidance={"scalars": 0}) accumulator.Reload() for tag in accumulator.Tags().get("scalars", []): scalars.setdefault(tag, []).extend( {"step": event.step, "value": event.value, "wall_time": event.wall_time} for event in accumulator.Scalars(tag) ) return event_files, scalars def actor_placement_from_log( log_text: str, preflight_record: dict[str, object], world_size: int, run_id: str, ) -> dict[str, object]: workers = preflight_record.get("workers", []) if not isinstance(workers, list): raise RuntimeError("Preflight worker records are invalid") worker_by_hostname = { str(worker["hostname"]): worker for worker in workers if isinstance(worker, dict) and worker.get("hostname") } plain_log = re.sub(r"\x1b\[[0-9;]*m", "", log_text) nccl_pattern = re.compile( r"\(WorkerDict pid=(?P\d+)(?:, ip=[^)]+)?\)\s+" r"(?P[^:\s]+):(?P\d+):\d+\s+\[\d+\] NCCL INFO " r"ncclCommInitRankConfig comm \S+ rank (?P\d+) nranks (?P\d+) " r"cudaDev \d+ nvmlDev (?P\d+) busId (?P\S+) " r"commId (?P\S+) - Init (?PSTART|COMPLETE)" ) communicators: dict[str, dict[int, dict[str, dict[str, object]]]] = {} for match in nccl_pattern.finditer(plain_log): if int(match.group("world")) != world_size: continue prefix_pid = int(match.group("prefix_pid")) pid = int(match.group("pid")) if prefix_pid != pid: raise RuntimeError(f"WorkerDict/NCCL PID mismatch: {prefix_pid} != {pid}") event: dict[str, object] = { "rank": int(match.group("rank")), "hostname": match.group("hostname"), "pid": pid, "nvml_dev": int(match.group("nvml_dev")), "bus_id": match.group("bus_id"), } comm_id = match.group("comm_id") rank = int(event["rank"]) phase = match.group("phase") phases = communicators.setdefault(comm_id, {}).setdefault(rank, {}) existing = phases.get(phase) if existing is not None and existing != event: raise RuntimeError(f"Conflicting NCCL {phase} evidence for rank {rank}") phases[phase] = event candidates: list[tuple[str, list[dict[str, object]]]] = [] for comm_id, ranks in communicators.items(): if sorted(ranks) != list(range(world_size)): continue actors = [] valid = True for rank in range(world_size): phases = ranks[rank] if set(phases) != {"START", "COMPLETE"} or phases["START"] != phases["COMPLETE"]: valid = False break actors.append(dict(phases["START"])) if valid and len({(actor["hostname"], actor["pid"]) for actor in actors}) == world_size: candidates.append((comm_id, actors)) if not candidates: raise RuntimeError("No complete world-size ActorRollout NCCL communicator was found") baseline = [ (actor["rank"], actor["hostname"], actor["pid"], actor["nvml_dev"], actor["bus_id"]) for actor in candidates[0][1] ] if any( [ (actor["rank"], actor["hostname"], actor["pid"], actor["nvml_dev"], actor["bus_id"]) for actor in actors ] != baseline for _, actors in candidates[1:] ): raise RuntimeError("Multiple full-world NCCL communicators have conflicting rank placement") comm_ids = sorted(comm_id for comm_id, _ in candidates) actors = candidates[0][1] gloo_pattern = re.compile( r"\(WorkerDict pid=(?P\d+)(?:, ip=[^)]+)?\).*?\[Gloo\] " r"Rank (?P\d+) is connected to (?P\d+) peer ranks\. " r"Expected number of connected peer ranks is : (?P\d+)" ) gloo_evidence: dict[tuple[int, int], set[tuple[int, int]]] = {} for match in gloo_pattern.finditer(plain_log): key = (int(match.group("pid")), int(match.group("rank"))) gloo_evidence.setdefault(key, set()).add( (int(match.group("peers")), int(match.group("expected"))) ) node_counts: dict[str, int] = {} for actor in actors: rank = int(actor["rank"]) pid = int(actor["pid"]) hostname = str(actor["hostname"]) worker = worker_by_hostname.get(hostname) if worker is None: raise RuntimeError(f"ActorRollout rank {rank} ran on unknown host {hostname}") peer_evidence = gloo_evidence.get((pid, rank), set()) expected_peers = (world_size - 1, world_size - 1) if expected_peers not in peer_evidence: raise RuntimeError( f"Missing full-mesh Gloo evidence for rank {rank}: {sorted(peer_evidence)}" ) node_ip = str(worker["node_ip"]) prefix = rf"\(WorkerDict pid={pid}(?:, ip=[^)]+)?\)\s+{re.escape(hostname)}:{pid}:\d+" required_transport = ( rf"{prefix}.*?NCCL INFO NCCL_IB_DISABLE set by environment to 1\.", rf"{prefix}.*?NCCL INFO NET/Socket : Using .*?{re.escape(node_ip)}<\d+>", rf"{prefix}.*?NCCL INFO Using network Socket", ) if any(not re.search(pattern, plain_log) for pattern in required_transport): raise RuntimeError(f"Missing socket-only NCCL evidence for rank {rank}") actor.update( { "class_name": "WorkerDict", "node_id": str(worker["node_id"]), "node_ip": node_ip, "connected_peer_ranks": expected_peers[0], "nccl_init": "complete", "nccl_transport": "Socket", "rdma_disabled": True, } ) node_id = str(actor["node_id"]) node_counts[node_id] = node_counts.get(node_id, 0) + 1 expected_node_ids = { str(worker["node_id"]) for worker in workers if isinstance(worker, dict) and worker.get("node_id") } if set(node_counts) != expected_node_ids or sorted(node_counts.values()) != [2, 2]: raise RuntimeError(f"Unexpected ActorRollout node distribution: {node_counts}") return { "status": "passed", "phase": "distributed_rank_placement", "source": "training_log", "checked_at": datetime.now(timezone.utc).isoformat(), "run_id": run_id, "expected_actors": world_size, "comm_ids": comm_ids, "node_counts": node_counts, "actors": actors, } def postflight(args: argparse.Namespace) -> int: output_dir = args.output_dir.resolve() log_text = args.training_log.read_text(encoding="utf-8", errors="replace") for pattern in FATAL_PATTERNS: if re.search(pattern, log_text, flags=re.IGNORECASE): raise RuntimeError(f"Fatal training marker found: {pattern}") required_log_markers = ( "[validate_config] All configuration checks passed successfully!", f"Total training steps: {args.target_steps}", f"global_step_{args.target_steps}", ) missing_markers = [marker for marker in required_log_markers if marker not in log_text] if missing_markers: raise RuntimeError(f"Training log is missing completion markers: {missing_markers}") preflight_record = json.loads((output_dir / "preflight.json").read_text(encoding="utf-8")) workers = preflight_record.get("workers", []) if len(workers) != args.expected_nodes: raise RuntimeError("Preflight topology evidence is missing") actor_placement = actor_placement_from_log( log_text, preflight_record, args.world_size, args.run_id ) atomic_json(output_dir / "actor-placement.json", actor_placement) resolved_command = (output_dir / "resolved-command.txt").read_text(encoding="utf-8") required_overrides = ( f"trainer.n_gpus_per_node={args.expected_gpus_per_node}", f"trainer.nnodes={args.expected_nodes}", "+ray_kwargs.ray_init.address=auto", "+ray_kwargs.ray_init.runtime_env.env_vars.TENSORBOARD_DIR=", "actor_rollout_ref.rollout.name=sglang", "actor_rollout_ref.actor.strategy=fsdp2", ) missing_overrides = [item for item in required_overrides if item not in resolved_command] if missing_overrides: raise RuntimeError(f"Resolved command is missing topology/runtime overrides: {missing_overrides}") event_files, scalars = load_scalars(output_dir / "tensorboard") metric_summary = {} for tag in REQUIRED_SCALARS: events = scalars.get(tag, []) if not events: raise RuntimeError(f"TensorBoard scalar is missing: {tag}") latest = max(events, key=lambda event: int(event["step"])) if int(latest["step"]) < args.target_steps or not math.isfinite(float(latest["value"])): raise RuntimeError(f"TensorBoard scalar did not reach a finite target step: {tag}={latest}") metric_summary[tag] = latest expected_rollouts = args.train_batch_size * args.rollout_n rollout_summary = {} for step in range(1, args.target_steps + 1): path = output_dir / "rollouts" / f"{step}.jsonl" rows = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()] if len(rows) != expected_rollouts: raise RuntimeError(f"Unexpected rollout count at step {step}: {len(rows)}") for row in rows: if row.get("step") != step or not row.get("input") or not isinstance(row.get("output"), str): raise RuntimeError(f"Invalid rollout row at step {step}") if not math.isfinite(float(row["score"])): raise RuntimeError(f"Non-finite rollout score at step {step}") rollout_summary[str(step)] = { "path": str(path), "rows": len(rows), "bytes": path.stat().st_size, "score_min": min(float(row["score"]) for row in rows), "score_max": max(float(row["score"]) for row in rows), } checkpoint_root = output_dir / "checkpoints" checkpoint = checkpoint_root / f"global_step_{args.target_steps}" actor = checkpoint / "actor" fsdp_config_path = actor / "fsdp_config.json" fsdp_config = json.loads(fsdp_config_path.read_text(encoding="utf-8")) if fsdp_config != {"FSDP_version": 2, "world_size": args.world_size}: raise RuntimeError(f"Unexpected FSDP checkpoint configuration: {fsdp_config}") checkpoint_files = {} for kind in ("model", "optim", "extra_state"): for rank in range(args.world_size): path = actor / f"{kind}_world_size_{args.world_size}_rank_{rank}.pt" if not path.is_file() or path.stat().st_size == 0: raise RuntimeError(f"Checkpoint shard is missing or empty: {path}") checkpoint_files[str(path.relative_to(output_dir))] = path.stat().st_size for path in ( checkpoint / "data.pt", checkpoint_root / "latest_checkpointed_iteration.txt", fsdp_config_path, ): if not path.is_file() or path.stat().st_size == 0: raise RuntimeError(f"Checkpoint file is missing or empty: {path}") checkpoint_files[str(path.relative_to(output_dir))] = path.stat().st_size if (checkpoint_root / "latest_checkpointed_iteration.txt").read_text().strip() != str(args.target_steps): raise RuntimeError("Checkpoint tracker does not match the target step") record = { "status": "passed", "phase": "postflight", "verified_at": datetime.now(timezone.utc).isoformat(), "run_id": args.run_id, "target_steps": args.target_steps, "train_batch_size": args.train_batch_size, "rollout_n": args.rollout_n, "world_size": args.world_size, "topology": { "nodes": len(workers), "gpus_per_node": preflight_record["expected_gpus_per_node"], "workers": workers, "actor_placement": actor_placement, }, "tensorboard": { "event_files": [{"path": str(path), "bytes": path.stat().st_size} for path in event_files], "metrics": metric_summary, }, "rollouts": rollout_summary, "checkpoint": { "path": str(checkpoint), "fsdp_config": fsdp_config, "files": checkpoint_files, }, } atomic_json(output_dir / "verification.json", record) (output_dir / "SUCCESS").write_text(datetime.now(timezone.utc).isoformat() + "\n", encoding="utf-8") print(json.dumps(record, indent=2, ensure_ascii=False)) return 0 def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() subparsers = parser.add_subparsers(dest="phase", required=True) common = argparse.ArgumentParser(add_help=False) common.add_argument("--run-id", required=True) common.add_argument("--output-dir", type=Path, required=True) common.add_argument("--target-steps", type=int, required=True) common.add_argument("--train-batch-size", type=int, required=True) common.add_argument("--rollout-n", type=int, required=True) before = subparsers.add_parser("preflight", parents=[common]) before.add_argument("--storage-root", type=Path, required=True) before.add_argument("--verl-root", type=Path, required=True) before.add_argument("--model-root", type=Path, required=True) before.add_argument("--data-root", type=Path, required=True) before.add_argument("--expected-nodes", type=int, required=True) before.add_argument("--expected-gpus-per-node", type=int, required=True) after = subparsers.add_parser("postflight", parents=[common]) after.add_argument("--training-log", type=Path, required=True) after.add_argument("--expected-nodes", type=int, required=True) after.add_argument("--expected-gpus-per-node", type=int, required=True) after.add_argument("--world-size", type=int, required=True) return parser.parse_args() if __name__ == "__main__": parsed = parse_args() handlers = {"preflight": preflight, "postflight": postflight} selected_handler = handlers[parsed.phase] raise SystemExit(selected_handler(parsed))