examples/disaggregated/slurm/benchmark/submit.py (4)

27-35: Use safe_dump and explicit YAML formatting

Switch to yaml.safe_dump and stable formatting for reproducibility.

Apply this diff:

-    with open(output_path, 'w') as f:
-        yaml.dump(worker_config, f, default_flow_style=False)
+    with open(output_path, 'w') as f:
+        yaml.safe_dump(worker_config, f, default_flow_style=False, sort_keys=True)

---

42-47: Resolve script path robustly

If script_file is relative, resolve it against work_dir (or the config file’s directory) to avoid “file not found.”

Apply this diff:

 def submit_job(config):
     # Extract configurations
     slurm_config = config['slurm']
     hw_config = config['hardware']
     env_config = config['environment']
+    # Resolve script path relative to work_dir if needed
+    script_path = slurm_config['script_file']
+    if not os.path.isabs(script_path):
+        script_path = os.path.join(env_config['work_dir'], script_path)
+    if not os.path.exists(script_path):
+        raise FileNotFoundError(f"SLURM script not found: {script_path}")

---

91-103: Set working directory and use resolved script path

Ensure job starts in work_dir and passes the correct script. Also consider using --chdir for sbatch.

Apply this diff:

     cmd = [
         'sbatch',
         f'--partition={slurm_config["partition"]}',
         f'--gres=gpu:{hw_config["gpus_per_node"]}',
         f'--account={slurm_config["account"]}',
         f'--time={slurm_config["job_time"]}',
         f'--job-name={slurm_config["job_name"]}',
         f'--nodes={total_nodes}',
         f'--ntasks={total_tasks}',
         f'--ntasks-per-node={hw_config["gpus_per_node"]}',
-        f'--segment={total_nodes}',
-        slurm_config['script_file'],
+        f'--chdir={env_config["work_dir"]}',
+        script_path,

---

141-147: Pass cwd to subprocess.run and surface stderr/stdout on failure

Improves reliability and debugging of submission failures.

Apply this diff:

-    try:
-        subprocess.run(cmd, check=True)
+    try:
+        subprocess.run(cmd, check=True, cwd=env_config['work_dir'])
     except subprocess.CalledProcessError as e:
-        print(f"Error submitting job: {e}", file=sys.stderr)
+        print(f"Error submitting job: {e}", file=sys.stderr)
         sys.exit(1)

examples/disaggregated/slurm/benchmark/disaggr_torch.slurm (3)

81-86: Add traps for reliable cleanup

Ensure cleanup_on_failure runs on ERR/INT/TERM and that scancel is issued on failures automatically.

Apply this diff:

 cleanup_on_failure() {
     echo "Error: $1"
     scancel ${SLURM_JOB_ID}
     exit 1
 }
 
+trap 'cleanup_on_failure "Script interrupted (INT)"' INT
+trap 'cleanup_on_failure "Script terminated (TERM)"' TERM
+trap 'cleanup_on_failure "Script failed (ERR)"' ERR

---

156-166: Pin srun placement or node lists if contention arises

You compute gen_nodes/ctx_nodes but don’t use them. If you observe step contention, pass explicit node lists via -w "${gen_nodes_csv}" and -w "${ctx_nodes_csv}".

---

90-96: Container launch: run a real no-op command within bash -c

A bare echo is fine, but using bash -c avoids surprises with shell builtins across environments; keep logging as-is.

Apply this diff:

-if ! srun -l --container-image=${container_image} \
+if ! srun -l --container-image=${container_image} \
         --container-name=${container_name} \
         --container-mounts=${container_mount} \
         --mpi=pmix \
-        echo "Container up." &> ${log_dir}/container_launch.log; then
+        bash -lc 'echo "Container up."' &> ${log_dir}/container_launch.log; then

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro