diff --git a/areal/experimental/openai/proxy/proxy_rollout_server.py b/areal/experimental/openai/proxy/proxy_rollout_server.py index 877d7f3943..7a60946368 100644 --- a/areal/experimental/openai/proxy/proxy_rollout_server.py +++ b/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( diff --git a/areal/scheduler/slurm.py b/areal/scheduler/slurm.py index b72f094787..ccc2c67b4a 100644 --- a/areal/scheduler/slurm.py +++ b/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}", diff --git a/areal/tests/test_slurm_scheduler.py b/areal/tests/test_slurm_scheduler.py index a168c73918..a0411f61fb 100644 --- a/areal/tests/test_slurm_scheduler.py +++ b/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)