fix(scheduler): Fix SLURM scheduler spec.image and spec.additional_bash_cmds arguments.

areal/experimental/openai/proxy/proxy_rollout_server.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import argparse
+import asyncio
 import inspect
 import os
 import threading
@@ -74,6 +75,10 @@
 _server_host: str = "0.0.0.0"
 _server_port: int = 8000
 
+# Port allocation tracking
+_allocated_ports: set[int] = set()
+_port_alloc_lock = asyncio.Lock()
+
 # Server config (needed for name_resolve registration)
 _experiment_name: str | None = None
 _trial_name: str | None = None
@@ -117,6 +122,47 @@ def health():
     return {"status": "ok", "initialized": _engine is not None}
 
 
+@app.post("/alloc_ports")
+async def alloc_ports(raw_request: Request):
+    """Allocate multiple free ports.
+
+    Expected JSON payload:
+    {
+        "count": 5  # Number of ports to allocate
+    }
+
+    This endpoint is required for compatibility with SlurmScheduler's
+    fork_workers mechanism, which may request additional ports after
+    forking a new worker process.
+    """
+    global _allocated_ports
+
+    try:
+        data = await raw_request.json()
+        count = data.get("count")
+        if count is None:
+            raise HTTPException(
+                status_code=400, detail="Missing 'count' field in request"
+            )
+
+        if not isinstance(count, int) or count <= 0:
+            raise HTTPException(
+                status_code=400, detail="'count' must be a positive integer"
+            )
+
+        async with _port_alloc_lock:
+            ports = find_free_ports(count, exclude_ports=_allocated_ports)
+        _allocated_ports.update(ports)
+
+        return {"status": "success", "ports": ports, "host": _server_host}
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Error in alloc_ports: {e}")
+        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
+
+
 def _setup_openai_client():
     global _openai_client, _session_timeout_seconds
     config = _engine.config
@@ -616,6 +662,7 @@ def main():
 
     # Determine port
     _server_port = args.port if args.port != 0 else find_free_ports(1)[0]
+    _allocated_ports.add(_server_port)
 
     # Configure name_resolve and register this server
     name_resolve.reconfigure(

areal/scheduler/slurm.py

@@ -2,6 +2,7 @@
 import getpass
 import os
 import re
+import shlex
 import subprocess
 import time
 from dataclasses import dataclass
@@ -72,7 +73,6 @@ def __init__(
         fileroot: str | None = None,
         cluster_name: str | None = None,
         container_type: str = "apptainer",
-        container_image: str | None = "/storage/openpsi/images/areal-latest.sif",
         container_mounts: str | None = "/storage:/storage",
         srun_additional_args: str = "--unbuffered --mpi=pmi2 -K --chdir $PWD",
         startup_timeout: float = 300.0,
@@ -120,7 +120,6 @@ def __init__(
             )
 
         self.container_type = container_type
-        self.container_image = container_image
         self.container_mounts = container_mounts
         self.srun_additional_args = srun_additional_args
         self.startup_timeout = startup_timeout
@@ -798,6 +797,11 @@ def _generate_sbatch_script(
         # Build environment variables (common to all workers)
         env_vars_dict = spec.env_vars.copy() if spec.env_vars else {}
 
+        bash_cmds = (spec.additional_bash_cmds or []).copy()
+        bash_cmds.append(rpc_cmd)
+        bash_cmds_str = ";\n".join(bash_cmds)
+        cmd = f"bash -c {shlex.quote(bash_cmds_str)}"
+
         # Build final command and export string
         if self.container_type == "apptainer":
             # For apptainer, pass env vars to singularity
@@ -806,10 +810,10 @@ def _generate_sbatch_script(
             if self.container_mounts:
                 final_cmd += f" --bind {self.container_mounts}"
             final_cmd += f" {env_string}"
-            final_cmd += f" {self.container_image}"
-            final_cmd += f" {rpc_cmd}"
+            final_cmd += f" {spec.image}"
+            final_cmd += f" {cmd}"
         else:  # native
-            final_cmd = rpc_cmd
+            final_cmd = cmd
 
         srun_flags = [
             f"--nodes={nodes}",

areal/tests/test_slurm_scheduler.py

@@ -364,7 +364,7 @@ def test_create_workers_with_colocate_strategy(scheduler):
         replicas=2,
         tasks=[SchedulingSpec(cpu=2, gpu=4, mem=4, port_count=2)],
     )
-    actor_ids = scheduler.create_workers(actor_job)
+    scheduler.create_workers(actor_job)
     scheduler.get_workers("actor", timeout=180)
 
     # Create colocated role
@@ -378,20 +378,22 @@ def test_create_workers_with_colocate_strategy(scheduler):
     )
     ref_ids = scheduler.create_workers(ref_job)
 
-    # Verify colocated role returns the SAME worker IDs as target role
-    assert ref_ids == actor_ids
+    # Verify colocated role with fork=True returns NEW worker IDs (not same as target)
+    # Forked workers get their own IDs like "ref/0", "ref/1"
+    assert ref_ids == ["ref/0", "ref/1"]
 
     # Verify colocation tracking is set up correctly
     assert "ref" in scheduler._colocated_roles
     assert scheduler._colocated_roles["ref"] == "actor"
-    assert scheduler._role_to_workers["ref"] == actor_ids
+    # Forked workers have their own entries in _workers
+    assert "ref" in scheduler._workers
 
     # Cleanup
     scheduler.delete_workers()
 
 
 def test_get_workers_for_colocated_role_delegates_to_target(scheduler):
-    """Test that get_workers for colocated role returns target's workers."""
+    """Test that get_workers for forked colocated role returns forked workers."""
     # Create target role
     actor_job = Job(
         role="actor",
@@ -401,7 +403,7 @@ def test_get_workers_for_colocated_role_delegates_to_target(scheduler):
     scheduler.create_workers(actor_job)
     actor_workers = scheduler.get_workers("actor", timeout=180)
 
-    # Create colocated role
+    # Create colocated role (with fork=True by default)
     ref_job = Job(
         role="ref",
         replicas=2,
@@ -413,12 +415,16 @@ def test_get_workers_for_colocated_role_delegates_to_target(scheduler):
     scheduler.create_workers(ref_job)
     ref_workers = scheduler.get_workers("ref", timeout=60)
 
-    # Verify same worker info is returned
+    # Forked workers have their own IDs and ports, but same IP (same node)
     assert len(ref_workers) == len(actor_workers)
     for ref_w, actor_w in zip(ref_workers, actor_workers):
-        assert ref_w.id == actor_w.id
+        # Forked workers have different IDs (ref/0 vs actor/0)
+        assert ref_w.id != actor_w.id
+        assert ref_w.id.startswith("ref/")
+        # Same IP since they're on the same node
         assert ref_w.ip == actor_w.ip
-        assert ref_w.worker_ports == actor_w.worker_ports
+        # Different ports since they're separate processes
+        assert ref_w.worker_ports != actor_w.worker_ports
 
     # Cleanup
     scheduler.delete_workers()
@@ -451,7 +457,7 @@ def test_delete_colocated_role_does_not_kill_processes(scheduler):
 
     # Verify colocated role is removed from tracking
     assert "ref" not in scheduler._colocated_roles
-    assert "ref" not in scheduler._role_to_workers
+    assert "ref" not in scheduler._workers
 
     # Verify target role workers are still available
     actor_workers = scheduler.get_workers("actor", timeout=60)