Isolate Omni V1 IPC runtime directories

Ratish1 · Ratish1 · commit 8f063d6c4d8e · 2026-05-07T00:13:18.000+05:30
diff --git a/sglang_omni_router/serve.py b/sglang_omni_router/serve.py
@@ -43,6 +43,8 @@ def build_parser() -> argparse.ArgumentParser:
 def main(argv: Sequence[str] | None = None) -> None:
     args = build_parser().parse_args(argv)
     logging.basicConfig(level=getattr(logging, args.log_level.upper(), logging.INFO))
+    logging.getLogger("httpx").setLevel(logging.WARNING)
+    logging.getLogger("httpcore").setLevel(logging.WARNING)
     config = build_router_config(
         worker_urls=args.worker_urls,
         host=args.host,
diff --git a/sglang_omni_v1/config/__init__.py b/sglang_omni_v1/config/__init__.py
@@ -1,5 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
-from sglang_omni_v1.config.compiler import compile_pipeline
+from sglang_omni_v1.config.compiler import (
+    IpcRuntimeDir,
+    compile_pipeline,
+    compile_pipeline_core,
+    create_ipc_runtime_dir,
+    prepare_pipeline_runtime,
+)
 from sglang_omni_v1.config.schema import (
     EndpointsConfig,
     PipelineConfig,
@@ -8,7 +14,11 @@
 )
 
 __all__ = [
+    "IpcRuntimeDir",
     "compile_pipeline",
+    "compile_pipeline_core",
+    "create_ipc_runtime_dir",
+    "prepare_pipeline_runtime",
     "PipelineConfig",
     "StageConfig",
     "RelayConfig",
diff --git a/sglang_omni_v1/config/compiler.py b/sglang_omni_v1/config/compiler.py
@@ -4,7 +4,11 @@
 from __future__ import annotations
 
 import inspect
+import logging
+import re
+import shutil
 import socket
+import tempfile
 from pathlib import Path
 from typing import Any
 
@@ -14,43 +18,144 @@
 from sglang_omni_v1.pipeline.stage.input import InputHandler
 from sglang_omni_v1.utils import import_string
 
+logger = logging.getLogger(__name__)
 
-def compile_pipeline(config: PipelineConfig) -> tuple[Coordinator, list[Stage]]:
-    """Build the coordinator and stage objects from the pipeline configuration."""
-    stages_cfg, name_map, entry_stage = config.apply_fusion()
-    endpoints = _allocate_endpoints(config, stages=stages_cfg)
 
-    coordinator = Coordinator(
-        completion_endpoint=endpoints["completion"],
-        abort_endpoint=endpoints["abort"],
-        entry_stage=entry_stage,
-        terminal_stages=config.terminal_stages or None,
-    )
+class IpcRuntimeDir:
+    """Runtime-owned IPC directory for one pipeline instance."""
+
+    def __init__(self, path: Path):
+        self.path = path
+        self._closed = False
 
-    stage_endpoints = {s.name: endpoints[f"stage_{s.name}"] for s in stages_cfg}
+    def __enter__(self) -> IpcRuntimeDir:
+        return self
 
-    stages: list[Stage] = []
-    for stage_cfg in stages_cfg:
-        stage = _compile_stage(
-            stage_cfg, config, stage_endpoints, endpoints, name_map=name_map
+    def __exit__(self, exc_type, exc, tb) -> None:
+        self.close()
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        try:
+            shutil.rmtree(self.path)
+        except FileNotFoundError:
+            return
+        except OSError as exc:
+            logger.warning("Failed to remove IPC runtime dir %s: %s", self.path, exc)
+
+
+def create_ipc_runtime_dir(config: PipelineConfig) -> IpcRuntimeDir | None:
+    """Create a per-run IPC namespace for one pipeline instance."""
+    if config.endpoints.scheme != "ipc":
+        return None
+
+    base_root = Path(config.endpoints.base_path)
+    base_root.mkdir(parents=True, exist_ok=True)
+
+    namespace_prefix = re.sub(r"[^0-9a-z]+", "-", config.name.lower()).strip("-")
+    if not namespace_prefix:
+        namespace_prefix = "pipeline"
+    path = Path(tempfile.mkdtemp(prefix=f"{namespace_prefix}-", dir=base_root))
+    return IpcRuntimeDir(path)
+
+
+def prepare_pipeline_runtime(
+    config: PipelineConfig,
+    *,
+    ipc_runtime_dir: IpcRuntimeDir | None = None,
+) -> tuple[
+    list[StageConfig], dict[str, str], str, dict[str, str], IpcRuntimeDir | None, bool
+]:
+    """Prepare fused stages and endpoint allocation for one runtime."""
+    runtime_dir = ipc_runtime_dir
+    created_runtime_dir = None
+    if runtime_dir is None:
+        runtime_dir = create_ipc_runtime_dir(config)
+        created_runtime_dir = runtime_dir
+    owns_runtime_dir = created_runtime_dir is not None
+
+    try:
+        stages_cfg, name_map, entry_stage = config.apply_fusion()
+        endpoints = _allocate_endpoints(
+            config,
+            stages=stages_cfg,
+            ipc_base_dir=runtime_dir.path if runtime_dir else None,
         )
-        coordinator.register_stage(stage.name, stage.control_plane.recv_endpoint)
-        stages.append(stage)
-
-    # Wire streaming targets
-    stage_map = {stage.name: stage for stage in stages}
-    cfg_map = {s.name: s for s in stages_cfg}
-    for stage_cfg in stages_cfg:
-        stage = stage_map.get(stage_cfg.name)
-        if stage is not None:
+    except Exception:
+        if created_runtime_dir is not None:
+            created_runtime_dir.close()
+        raise
+
+    return stages_cfg, name_map, entry_stage, endpoints, runtime_dir, owns_runtime_dir
+
+
+def compile_pipeline_core(
+    config: PipelineConfig,
+    *,
+    ipc_runtime_dir: IpcRuntimeDir | None = None,
+) -> tuple[Coordinator, list[Stage], IpcRuntimeDir | None]:
+    """Build the coordinator and stage objects from the pipeline configuration."""
+    stages_cfg, name_map, entry_stage, endpoints, runtime_dir, owns_runtime_dir = (
+        prepare_pipeline_runtime(
+            config,
+            ipc_runtime_dir=ipc_runtime_dir,
+        )
+    )
+
+    try:
+        coordinator = Coordinator(
+            completion_endpoint=endpoints["completion"],
+            abort_endpoint=endpoints["abort"],
+            entry_stage=entry_stage,
+            terminal_stages=config.terminal_stages or None,
+        )
+
+        stage_endpoints = {s.name: endpoints[f"stage_{s.name}"] for s in stages_cfg}
+
+        stages: list[Stage] = []
+        for stage_cfg in stages_cfg:
+            stage = _compile_stage(
+                stage_cfg, config, stage_endpoints, endpoints, name_map=name_map
+            )
+            coordinator.register_stage(stage.name, stage.control_plane.recv_endpoint)
+            stages.append(stage)
+
+        stage_map = {stage.name: stage for stage in stages}
+        cfg_map = {s.name: s for s in stages_cfg}
+        for stage_cfg in stages_cfg:
+            stage = stage_map.get(stage_cfg.name)
+            if stage is None:
+                continue
             _wire_stream_targets(
                 stage,
                 stage_cfg,
                 stage_map,
                 gpu_placement=config.gpu_placement,
                 cfg_map=cfg_map,
             )
+    except Exception:
+        if owns_runtime_dir and runtime_dir is not None:
+            runtime_dir.close()
+        raise
+
+    return coordinator, stages, runtime_dir
+
+
+def compile_pipeline(config: PipelineConfig) -> tuple[Coordinator, list[Stage]]:
+    """Build coordinator and stages directly from a pipeline config.
+
+    IPC pipelines need explicit runtime-directory ownership so multiple
+    replicas cannot bind the same local sockets.
+    """
+    if config.endpoints.scheme == "ipc":
+        raise ValueError(
+            "compile_pipeline() does not manage IPC runtime-dir ownership. "
+            "Use compile_pipeline_core(...) or MultiProcessPipelineRunner."
+        )
 
+    coordinator, stages, _ = compile_pipeline_core(config)
     return coordinator, stages
 
 
@@ -252,6 +357,7 @@ def _allocate_endpoints(
     config: PipelineConfig,
     *,
     stages: list[StageConfig],
+    ipc_base_dir: Path | None = None,
 ) -> dict[str, str]:
     endpoints: dict[str, str] = {}
 
@@ -261,8 +367,9 @@ def _allocate_endpoints(
         endpoints["abort"] = config.abort_endpoint
 
     if config.endpoints.scheme == "ipc":
-        base_dir = Path(config.endpoints.base_path) / config.name
-        base_dir.mkdir(parents=True, exist_ok=True)
+        if ipc_base_dir is None:
+            raise ValueError("IPC endpoint allocation requires an IPC runtime dir")
+        base_dir = ipc_base_dir
         endpoints.setdefault("completion", f"ipc://{base_dir}/completion.sock")
         endpoints.setdefault("abort", f"ipc://{base_dir}/abort.sock")
         for s in stages:
diff --git a/sglang_omni_v1/pipeline/mp_runner.py b/sglang_omni_v1/pipeline/mp_runner.py
@@ -16,10 +16,11 @@
 from typing import Any
 
 from sglang_omni_v1.config.compiler import (
-    _allocate_endpoints,
+    IpcRuntimeDir,
     _build_relay_config,
     _detect_same_gpu_targets,
     _resolve_factory_args,
+    prepare_pipeline_runtime,
 )
 from sglang_omni_v1.config.schema import PipelineConfig, StageConfig
 from sglang_omni_v1.pipeline import Coordinator
@@ -32,6 +33,10 @@
 def _build_stage_groups(
     config: PipelineConfig,
     ctx: multiprocessing.context.BaseContext | None = None,
+    *,
+    stages_cfg: list[StageConfig] | None = None,
+    name_map: dict[str, str] | None = None,
+    endpoints: dict[str, str] | None = None,
 ) -> list[StageGroup]:
     """Compile *config* into one :class:`StageGroup` per logical stage.
 
@@ -41,8 +46,10 @@ def _build_stage_groups(
     if ctx is None:
         ctx = multiprocessing.get_context("spawn")
 
-    stages_cfg, name_map, _ = config.apply_fusion()
-    endpoints = _allocate_endpoints(config, stages=stages_cfg)
+    if stages_cfg is None or name_map is None or endpoints is None:
+        if config.endpoints.scheme == "ipc":
+            raise ValueError("_build_stage_groups requires prepared IPC endpoints")
+        stages_cfg, name_map, _, endpoints, _, _ = prepare_pipeline_runtime(config)
     stage_endpoints = {s.name: endpoints[f"stage_{s.name}"] for s in stages_cfg}
     cfg_map = {s.name: s for s in stages_cfg}
 
@@ -264,6 +271,7 @@ class MultiProcessPipelineRunner:
     def __init__(self, config: PipelineConfig):
         self._config = config
         self._coordinator: Coordinator | None = None
+        self._ipc_runtime_dir: IpcRuntimeDir | None = None
         self._groups: list[StageGroup] = []
         self._completion_task: asyncio.Task | None = None
         self._monitor_task: asyncio.Task | None = None
@@ -281,10 +289,24 @@ async def start(self, timeout: float = 120.0) -> None:
 
         try:
             ctx = multiprocessing.get_context("spawn")
-            groups = _build_stage_groups(self._config, ctx)
-
-            stages_cfg, _, entry_stage = self._config.apply_fusion()
-            endpoints = _allocate_endpoints(self._config, stages=stages_cfg)
+            (
+                stages_cfg,
+                name_map,
+                entry_stage,
+                endpoints,
+                self._ipc_runtime_dir,
+                _,
+            ) = prepare_pipeline_runtime(
+                self._config,
+                ipc_runtime_dir=self._ipc_runtime_dir,
+            )
+            groups = _build_stage_groups(
+                self._config,
+                ctx,
+                stages_cfg=stages_cfg,
+                name_map=name_map,
+                endpoints=endpoints,
+            )
 
             self._coordinator = Coordinator(
                 completion_endpoint=endpoints["completion"],
@@ -373,6 +395,11 @@ async def stop(self) -> None:
 
         await self._coordinator.stop()
         self._groups.clear()
+        self._coordinator = None
+
+        if self._ipc_runtime_dir is not None:
+            self._ipc_runtime_dir.close()
+            self._ipc_runtime_dir = None
 
     async def _cleanup_on_failure(self) -> None:
         """Best-effort cleanup after a failed start()."""
@@ -401,3 +428,7 @@ async def _cleanup_on_failure(self) -> None:
             except Exception:
                 pass
             self._coordinator = None
+
+        if self._ipc_runtime_dir is not None:
+            self._ipc_runtime_dir.close()
+            self._ipc_runtime_dir = None
diff --git a/sglang_omni_v1/serve/launcher.py b/sglang_omni_v1/serve/launcher.py
@@ -29,14 +29,15 @@
 import os
 import socket
 import time
+from contextlib import suppress
 from typing import Any
 
 import uvicorn
 from fastapi import APIRouter
 from pydantic import BaseModel
 
 from sglang_omni_v1.client import Client
-from sglang_omni_v1.config import PipelineConfig, compile_pipeline
+from sglang_omni_v1.config import PipelineConfig, compile_pipeline_core
 from sglang_omni_v1.profiler.profiler_control import ProfilerControlClient
 from sglang_omni_v1.serve.openai_api import create_app
 
@@ -189,20 +190,21 @@ async def _run_server(
             await mp_runner.stop()
             logger.info("Pipeline stopped.")
     else:
-        coordinator, stages = compile_pipeline(pipeline_config)
+        coordinator, stages, runtime_dir = compile_pipeline_core(pipeline_config)
         stage_endpoints = _collect_stage_control_endpoints(stages)
-
-        # Start coordinator + all stages as async tasks
-        await coordinator.start()
-        completion_task = asyncio.create_task(coordinator.run_completion_loop())
-        stage_tasks = [asyncio.create_task(s.run()) for s in stages]
-        logger.info(
-            "Pipeline '%s' started (%d stages)",
-            pipeline_config.name,
-            len(stages),
-        )
+        completion_task = None
+        stage_tasks = []
 
         try:
+            await coordinator.start()
+            completion_task = asyncio.create_task(coordinator.run_completion_loop())
+            stage_tasks = [asyncio.create_task(s.run()) for s in stages]
+            logger.info(
+                "Pipeline '%s' started (%d stages)",
+                pipeline_config.name,
+                len(stages),
+            )
+
             cl_kwargs = client_kwargs or {}
             client = Client(coordinator, **cl_kwargs)
             app = create_app(
@@ -227,8 +229,17 @@ async def _run_server(
             logger.info("Shutting down pipeline …")
             for t in stage_tasks:
                 t.cancel()
-            completion_task.cancel()
-            await coordinator.stop()
+            if completion_task is not None:
+                completion_task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await completion_task
+            if stage_tasks:
+                await asyncio.gather(*stage_tasks, return_exceptions=True)
+            try:
+                await coordinator.stop()
+            finally:
+                if runtime_dir is not None:
+                    runtime_dir.close()
             logger.info("Pipeline stopped.")
 
 
diff --git a/tests/test_v1_ipc_runtime_dir.py b/tests/test_v1_ipc_runtime_dir.py
