Arctic rl/integration

.github/CODEOWNERS

@@ -1 +1 @@
-* @sfc-gh-aqiao @sfc-gh-jrasley @sfc-gh-mhidayetoglu @sfc-gh-yewang @sfc-gh-goliaro @sfc-gh-reyazda
+* @sfc-gh-aqiao @jeff-rasley_snow @mert-hidayetoglu_snow @ye-wang_snow @gabriele-oliaro_snow @reza-yazdani_snow

arctic_inference/envs.py

@@ -18,6 +18,24 @@
 
 if TYPE_CHECKING:
     ARCTIC_INFERENCE_SKIP_SPEC_MODEL_CHECK: bool = False
+    ARCTIC_FP32_LM_HEAD: bool = False
+
+
+def arctic_inference_effective_enabled(
+    extra_env: dict[str, str] | None = None,
+) -> bool:
+    """True if the Arctic vLLM plugin should load for this process or worker env.
+
+    When ``extra_env`` is passed (e.g. ``ModelConfig.extra_env``), it is checked
+    in addition to ``os.environ`` so the driver can omit Arctic-only engine
+    kwargs when workers will not enable the plugin.
+    """
+    if os.getenv("ARCTIC_INFERENCE_ENABLED", "0") == "1":
+        return True
+    if extra_env and str(extra_env.get("ARCTIC_INFERENCE_ENABLED", "0")) == "1":
+        return True
+    return False
+
 
 environment_variables: dict[str, Callable[[], Any]] = {
     "ARCTIC_INFERENCE_ENABLED":
@@ -28,6 +46,14 @@
     lambda: os.getenv("ARCTIC_INFERENCE_SKIP_SPEC_MODEL_CHECK", "0") == "1",
     "ARCTIC_INFERENCE_SKIP_VERSION_CHECK":
     lambda: os.getenv("ARCTIC_INFERENCE_SKIP_VERSION_CHECK", "0") == "1",
+    # Run the lm_head matmul in fp32 (weights stay in their native
+    # dtype; we upcast on the fly). vLLM's V1 sampler already does
+    # softmax in fp32, so the full final stage is fp32. Needed for
+    # RL workloads that require precise log-probs / token
+    # probabilities for off-policy correction. Equivalent to the
+    # ``--fp32-lm-head`` CLI flag.
+    "ARCTIC_FP32_LM_HEAD":
+    lambda: os.getenv("ARCTIC_FP32_LM_HEAD", "0") == "1",
 }
 
 # temporary workaround for gpt-oss model

arctic_inference/server/__init__.py

@@ -0,0 +1,16 @@
+from arctic_inference.server.config import ModelConfig
+from arctic_inference.server.pipeline import Pipeline
+from arctic_inference.server.replica_pool import ReplicaPool, ensure_ray
+from arctic_inference.server.weight_sync import (
+    NCCLEngine, WeightInfo, build_weights_info,
+)
+
+__all__ = [
+    "ModelConfig",
+    "Pipeline",
+    "ReplicaPool",
+    "ensure_ray",
+    "NCCLEngine",
+    "WeightInfo",
+    "build_weights_info",
+]

arctic_inference/server/api.py

@@ -0,0 +1,213 @@
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+from typing import Any
+
+from arctic_inference.server.replica_pool import ReplicaPool
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field, model_validator
+
+from arctic_inference.server.config import ModelConfig
+
+
+# ---------------------------------------------------------------------------
+# Request models
+# ---------------------------------------------------------------------------
+
+
+class InitRequest(BaseModel):
+    config: ModelConfig
+    model_id: str | None = None
+    num_replicas: int | None = None
+
+
+class GenerateRequest(BaseModel):
+    model_id: str | None = None
+    prompts: list[str | list[int]]
+    sampling_params: dict[str, Any] = Field(default_factory=dict)
+    routing_key: str | list[str | None] | None = None
+    strict: bool = False
+
+
+class GroupConfig(BaseModel):
+    group_id: int
+    master_addr: str
+    master_port: int
+    world_size: int
+    replica_ids: list[int]
+
+
+class SleepRequest(BaseModel):
+    model_id: str
+    level: int = 1
+
+
+class WakeUpRequest(BaseModel):
+    model_id: str
+    tags: list[str] | None = None
+
+
+class SyncWeightsRequest(BaseModel):
+    model_id: str | None = None
+    groups: list[GroupConfig] | None = None
+    bucket_size: int = 256 * 1024 * 1024
+    strategy: str = "hotswap"
+    engine_only: bool = False
+    direct_mode: bool = False
+    reverse: bool = False
+
+    # Legacy flat fields
+    master_addr: str | None = None
+    master_port: int | None = None
+    world_size: int | None = None
+
+
+class SyncSpecWeightsRequest(BaseModel):
+    model_id: str | None = None
+    groups: list[GroupConfig] | None = None
+    bucket_size: int = 256 * 1024 * 1024
+    strategy: str = "hotswap"
+    engine_only: bool = False
+    reverse: bool = False
+
+    # Legacy flat fields
+    master_addr: str | None = None
+    master_port: int | None = None
+    world_size: int | None = None
+
+
+# ---------------------------------------------------------------------------
+# Backend – a bare ReplicaPool by default, swapped to Driver by multi_model.py
+# ---------------------------------------------------------------------------
+
+backend: Any = ReplicaPool()
+
+
+# ---------------------------------------------------------------------------
+# FastAPI app
+# ---------------------------------------------------------------------------
+
+
+@asynccontextmanager
+async def lifespan(_app: FastAPI):
+    yield
+    await backend.shutdown()
+
+
+app = FastAPI(title="Arctic Inference", lifespan=lifespan)
+
+
+@app.post("/init")
+async def init_endpoint(request: InitRequest):
+    try:
+        n = await backend.initialize(
+            request.config, model_id=request.model_id,
+            num_replicas=request.num_replicas,
+        )
+        return {"status": "ready", "model_id": request.model_id, "num_replicas": n}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/generate")
+async def generate_endpoint(request: GenerateRequest):
+    try:
+        results = await backend.generate(**request.model_dump())
+        return {"results": results}
+    except RuntimeError as e:
+        msg = str(e).lower()
+        if "paused" in msg or "cancelled" in msg:
+            raise HTTPException(status_code=503, detail=str(e))
+        raise HTTPException(status_code=500, detail=str(e))
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/sleep")
+async def sleep_endpoint(request: SleepRequest):
+    """Free GPU memory for a model (drain in-flight requests first)."""
+    try:
+        return await backend.sleep(request.model_id, level=request.level)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/wake_up")
+async def wake_up_endpoint(request: WakeUpRequest):
+    """Restore GPU memory for a model and resume serving."""
+    try:
+        return await backend.wake_up(request.model_id, tags=request.tags)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.get("/weights_info")
+async def weights_info_endpoint(model_id: str | None = None):
+    try:
+        infos = backend.get_weights_info(model_id=model_id)
+        return {"weights_info": infos, "count": len(infos)}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/sync_weights")
+async def sync_weights_endpoint(request: SyncWeightsRequest):
+    try:
+        return await backend.sync_weights(**request.model_dump())
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.get("/spec_weights_info")
+async def spec_weights_info_endpoint(model_id: str | None = None):
+    try:
+        infos = backend.get_spec_weights_info(model_id=model_id)
+        return {"weights_info": infos, "count": len(infos)}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/sync_spec_weights")
+async def sync_spec_weights_endpoint(request: SyncSpecWeightsRequest):
+    try:
+        return await backend.sync_spec_weights(**request.model_dump())
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/close_weight_sync")
+async def close_weight_sync_endpoint(model_id: str | None = None):
+    try:
+        return await backend.close_weight_sync(model_id=model_id)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.get("/status")
+async def status_endpoint():
+    try:
+        return await backend.get_status()
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.get("/metrics")
+async def metrics_endpoint(model_id: str | None = None):
+    """Drain per-replica snapshots and per-request records.
+
+    Each call empties the server-side ring buffers, so consumers see only
+    metrics produced since their previous call.
+    """
+    try:
+        return await backend.drain_metrics(model_id=model_id)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/shutdown")
+async def shutdown_endpoint(model_id: str | None = None):
+    try:
+        await backend.shutdown(model_id=model_id)
+        return {"status": "shutdown", "model_id": model_id}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))

arctic_inference/server/cli.py

@@ -0,0 +1,30 @@
+"""CLI entry point: ``arctic-inference-server``."""
+from __future__ import annotations
+
+import argparse
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        prog="arctic-inference-server",
+        description="Launch the Arctic Inference server",
+    )
+    parser.add_argument("--host", default="0.0.0.0", help="Bind address (default: 0.0.0.0)")
+    parser.add_argument("--port", type=int, default=8000, help="Bind port (default: 8000)")
+    parser.add_argument("--workers", type=int, default=1, help="Number of uvicorn workers (default: 1)")
+    parser.add_argument("--log-level", default="info", help="Log level (default: info)")
+    args = parser.parse_args()
+
+    import uvicorn
+
+    uvicorn.run(
+        "arctic_inference.server.api:app",
+        host=args.host,
+        port=args.port,
+        workers=args.workers,
+        log_level=args.log_level,
+    )
+
+
+if __name__ == "__main__":
+    main()