From 36b0bcf4dcd7f0fa6df9e7e1e37eba66b7cd92f6 Mon Sep 17 00:00:00 2001 From: droot Date: Wed, 17 Jun 2026 14:11:53 -0700 Subject: [PATCH 01/11] ci: fix runner disk space exhaustion and exclude scratch artifacts --- .github/workflows/build-and-push.yml | 7 +++++++ .github/workflows/build-pr.yml | 7 +++++++ .gitignore | 1 + Makefile | 4 ++-- pyproject.toml | 10 ++++++++++ 5 files changed, 27 insertions(+), 2 deletions(-) diff --git a/.github/workflows/build-and-push.yml b/.github/workflows/build-and-push.yml index 004b5d18..9b150569 100644 --- a/.github/workflows/build-and-push.yml +++ b/.github/workflows/build-and-push.yml @@ -28,6 +28,13 @@ jobs: - name: Checkout uses: actions/checkout@v4 + - name: Free Disk Space + run: | + sudo rm -rf /usr/share/dotnet + sudo rm -rf /opt/ghc + sudo rm -rf "/usr/local/share/boost" + sudo rm -rf "$AGENT_TOOLSDIRECTORY" + - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 diff --git a/.github/workflows/build-pr.yml b/.github/workflows/build-pr.yml index a63c4d5b..d08f1e5b 100644 --- a/.github/workflows/build-pr.yml +++ b/.github/workflows/build-pr.yml @@ -24,6 +24,13 @@ jobs: - name: Checkout uses: actions/checkout@v4 + - name: Free Disk Space + run: | + sudo rm -rf /usr/share/dotnet + sudo rm -rf /opt/ghc + sudo rm -rf "/usr/local/share/boost" + sudo rm -rf "$AGENT_TOOLSDIRECTORY" + - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 diff --git a/.gitignore b/.gitignore index 703e8059..3c81cd1e 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,7 @@ env/ **/venv/ # Training Logs & Plots +scratch/ *.txt *.png *.log diff --git a/Makefile b/Makefile index 653057e1..fc125627 100644 --- a/Makefile +++ b/Makefile @@ -169,8 +169,8 @@ REMOTE_HOST ?= # Push local workspace changes to the remote VM push-vm: - rsync -avz --exclude '.git' --exclude '.venv' --exclude '__pycache__' --exclude '*.pyc' --exclude '.DS_Store' ./ $(REMOTE_HOST):~/open-rl + rsync -avz --exclude '.git' --exclude '.venv' --exclude '__pycache__' --exclude '*.pyc' --exclude '.DS_Store' --exclude 'scratch' ./ $(REMOTE_HOST):~/open-rl # Pull changes from the remote VM back to the local workspace pull-vm: - rsync -avz --exclude '.git' --exclude '.venv' --exclude '__pycache__' --exclude '*.pyc' --exclude '.DS_Store' $(REMOTE_HOST):~/open-rl/ ./ + rsync -avz --exclude '.git' --exclude '.venv' --exclude '__pycache__' --exclude '*.pyc' --exclude '.DS_Store' --exclude 'scratch' $(REMOTE_HOST):~/open-rl/ ./ diff --git a/pyproject.toml b/pyproject.toml index 2b01fc4b..70b3f55a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -107,3 +107,13 @@ torchvision = [ vllm = [ { index = "vllm-cu129", extra = "vllm", marker = "sys_platform == 'linux'" }, ] + +[tool.ruff] +force-exclude = true +exclude = [ + ".git", + ".venv", + "__pycache__", + "scratch", + "scratch/**", +] From bdfe564b3f25c44752492aba0b77b9b929334e2d Mon Sep 17 00:00:00 2001 From: droot Date: Wed, 17 Jun 2026 14:11:53 -0700 Subject: [PATCH 02/11] feat(server): implement pull sampling queues and dynamic weight reloading --- src/server/gateway.py | 105 +++++++++++----- src/server/k8s_worker_manager.py | 42 ++++--- src/server/store.py | 128 ++++++++++++++++---- src/server/worker_manager.py | 62 ++++++++-- tests/test_trainer_optimizer_correctness.py | 14 +-- tests/test_worker_manager.py | 6 + 6 files changed, 269 insertions(+), 88 deletions(-) diff --git a/src/server/gateway.py b/src/server/gateway.py index 4066b890..6ee21e60 100644 --- a/src/server/gateway.py +++ b/src/server/gateway.py @@ -146,7 +146,7 @@ async def launch_worker_and_enqueue(request: dict) -> str: request_id = request["request_id"] await store.set_future(request_id, {"status": "pending"}) try: - await asyncio.to_thread(fft_worker_manager.launch, request["model_id"]) + await asyncio.to_thread(fft_worker_manager.launch_trainer, request["model_id"]) except Exception as exc: traceback.print_exc() await store.set_future(request_id, {"type": "RequestFailedResponse", "error_message": str(exc)}) @@ -154,6 +154,14 @@ async def launch_worker_and_enqueue(request: dict) -> str: return await enqueue(request) +async def ensure_sampler_launched(model_id: str) -> None: + if is_fft_enabled() and fft_worker_manager is not None and get_sampler_backend() == "vllm": + try: + await asyncio.to_thread(fft_worker_manager.launch_sampler, model_id) + except Exception: + traceback.print_exc() + + async def preflight_vllm() -> None: """If SAMPLING_BACKEND=vllm, verify the vLLM worker is reachable at VLLM_URL. @@ -307,6 +315,18 @@ async def create_model(req: dict): return {"request_id": req_id} +@app.post("/api/v1/delete_model") +async def delete_model(req: dict): + model_id = req.get("model_id") + if not model_id: + return JSONResponse(status_code=400, content={"error": "model_id is required"}) + if is_fft_enabled(): + print(f"[GATEWAY] Requesting shutdown of workers for model {model_id}...") + await store.put_request({"request_id": "SHUTDOWN_SENTINEL", "model_id": model_id, "op": "shutdown_workers"}) + await store.put_sampling_request({"request_id": "SHUTDOWN_SENTINEL", "model_id": model_id}) + return {"status": "ok"} + + @app.post("/api/v1/create_model_from_state") async def create_model_from_state(req: dict): """ServiceClient.create_training_client_from_state_async()""" @@ -409,6 +429,7 @@ async def save_weights_for_sampler(req: dict): if not model_id: return JSONResponse(status_code=400, content={"error": "model_id is required"}) + await ensure_sampler_launched(model_id) seq_id = req.get("sampling_session_seq_id") or int(time.time() * 1000) alias = req.get("name") or req.get("alias") or req.get("path") @@ -494,10 +515,31 @@ async def create_sampling_session(req: dict): if model_path and model_path.startswith("tinker://"): sess_id = model_path + path = model_path[len("tinker://") :] + parts = path.split("/") + target_model_id = parts[0] elif base_model: sess_id = base_model + target_model_id = base_model else: sess_id = model_id or "samp-session-live-123" + target_model_id = sess_id + + if get_sampler_backend() == "vllm" and target_model_id: + if is_fft_enabled(): + await ensure_sampler_launched(target_model_id) + s = get_store() + if hasattr(s, "redis"): + print(f"[GATEWAY] Waiting for dynamic vLLM sampler worker to be ready for model {target_model_id}...") + start_time = time.monotonic() + while True: + is_ready = await s.redis.get(f"open_rl:sampler_ready:{target_model_id}") + if is_ready == "1" or is_ready == b"1": + print(f"[GATEWAY] Dynamic vLLM sampler worker is ready! (took {time.monotonic() - start_time:.2f}s)") + break + if time.monotonic() - start_time > 300: + raise TimeoutError("Timed out waiting for dynamic vLLM sampler worker to be ready") + await asyncio.sleep(1) return {"sampling_session_id": sess_id, "type": "create_sampling_session"} @@ -539,40 +581,39 @@ async def asample(req: dict): # vLLM backend req_id = str(uuid.uuid4()) + carrier: dict = {} + propagate.inject(carrier) await store.set_future(req_id, {"status": "pending"}) - lora_path = os.path.join(TMP_DIR, "peft", base_model_id, base_model_id) if is_sampler_weights_ref(model_id) else None - headers: dict[str, str] = {"Content-Type": "application/json"} - propagate.inject(headers) - - try: - async with httpx.AsyncClient(timeout=120.0) as client: - resp = await client.post( - f"{VLLM_URL.rstrip('/')}/generate", - json={ - "request_id": req_id, - "prompt_token_ids": prompt, - "max_tokens": max_tokens, - "temperature": temperature, - "stop": stop, - "top_p": top_p, - "top_k": top_k, - "num_samples": num_samples, - "lora_id": model_id, - "lora_path": lora_path, - "include_prompt_logprobs": include_prompt_logprobs, - }, - headers=headers, - ) - resp.raise_for_status() - data = resp.json() - if data.get("type") != "RequestFailedResponse": - data["type"] = "sample" - await store.set_future(req_id, data) - except Exception as e: - traceback.print_exc() - await store.set_future(req_id, {"type": "RequestFailedResponse", "error_message": str(e)}) + if is_fft_enabled(): + rel_path = model_id[len("tinker://") :] if model_id.startswith("tinker://") else model_id.lstrip("/") + local_path = os.path.join(TMP_DIR, "sampler_full", rel_path) + weights_path = local_path + lora_id = None + lora_path = None + else: + weights_path = None + lora_id = model_id + lora_path = os.path.join(TMP_DIR, "peft", base_model_id, base_model_id) if is_sampler_weights_ref(model_id) else None + + sampling_req = { + "request_id": req_id, + "prompt_token_ids": prompt, + "max_tokens": max_tokens, + "temperature": temperature, + "stop": stop, + "top_p": top_p, + "top_k": top_k, + "num_samples": num_samples, + "lora_id": lora_id, + "lora_path": lora_path, + "weights_path": weights_path, + "include_prompt_logprobs": include_prompt_logprobs, + "model_id": base_model_id or model_id, + "trace_context": carrier, + } + await store.put_sampling_request(sampling_req) return {"request_id": req_id} diff --git a/src/server/k8s_worker_manager.py b/src/server/k8s_worker_manager.py index 5604803b..47eb282d 100644 --- a/src/server/k8s_worker_manager.py +++ b/src/server/k8s_worker_manager.py @@ -58,8 +58,18 @@ def __init__(self, core_api: Any = None): self.core_api = core_api def launch(self, model_id: str) -> None: + self.launch_trainer(model_id) + + def launch_trainer(self, model_id: str) -> None: + self._launch_pod(model_id, role="trainer") + + def launch_sampler(self, model_id: str) -> None: + self._launch_pod(model_id, role="sampler") + + def _launch_pod(self, model_id: str, role: str) -> None: job_id = sanitize_job_id(model_id) - pod_name = POD_NAME_PREFIX + job_id + prefix = "open-rl-trainer-" if role == "trainer" else "open-rl-sampler-" + pod_name = prefix + job_id existing = self.read_pod(pod_name) if existing is not None: @@ -68,41 +78,43 @@ def launch(self, model_id: str) -> None: self.delete_pod_and_wait(pod_name) try: - self.core_api.create_namespaced_pod(namespace=self.namespace, body=self.render_pod(pod_name, model_id, job_id)) + self.core_api.create_namespaced_pod(namespace=self.namespace, body=self.render_pod(pod_name, model_id, job_id, role=role)) except Exception as exc: - # Another gateway replica created it between our read and create. if getattr(exc, "status", None) != 409: raise def shutdown(self, model_id: str) -> None: - pod_name = POD_NAME_PREFIX + sanitize_job_id(model_id) - try: - self.core_api.delete_namespaced_pod(name=pod_name, namespace=self.namespace) - except Exception as exc: - if getattr(exc, "status", None) != 404: - raise + job_id = sanitize_job_id(model_id) + for prefix in ("open-rl-trainer-", "open-rl-sampler-"): + pod_name = prefix + job_id + try: + self.core_api.delete_namespaced_pod(name=pod_name, namespace=self.namespace) + except Exception as exc: + if getattr(exc, "status", None) != 404: + raise def shutdown_all(self) -> None: - # Trainer worker pods deliberately outlive gateway restarts; Kubernetes owns them. pass - def render_pod(self, pod_name: str, model_id: str, job_id: str) -> dict[str, Any]: + def render_pod(self, pod_name: str, model_id: str, job_id: str, role: str = "trainer") -> dict[str, Any]: pod = copy.deepcopy(self.pod_template) metadata = pod.setdefault("metadata", {}) metadata["name"] = pod_name + app_label = "open-rl-trainer-worker" if role == "trainer" else "open-rl-sampler-worker" + group_val = self.group_id if role == "trainer" else "samplers" metadata.setdefault("labels", {}).update( { - "app": "open-rl-trainer-worker", + "app": app_label, "snapshot-agent": "true", - "timeslice.io/group": self.group_id, + "timeslice.io/group": group_val, "timeslice.io/job-id": job_id, } ) container = pod["spec"]["containers"][0] + if role == "sampler": + container["command"] = ["uv", "run", "python", "-u", "-m", "server.vllm_sampler"] container.setdefault("args", []).extend(["--model-id", model_id]) - # Keep the env value aligned with the label so current and future - # coordination clients use the same sanitized job id. container.setdefault("env", []).append({"name": "OPEN_RL_TIME_SLICE_JOB_ID", "value": job_id}) return pod diff --git a/src/server/store.py b/src/server/store.py index 3cb458b8..87f4681c 100644 --- a/src/server/store.py +++ b/src/server/store.py @@ -10,9 +10,6 @@ import redis.asyncio as redis from redis.exceptions import TimeoutError as RedisTimeoutError -# How long a resolved request result stays available to retrieve_future polling. -FUTURE_TTL_S = int(os.getenv("OPEN_RL_FUTURE_TTL_S", "300")) - class RequestStore(ABC): @abstractmethod @@ -20,16 +17,36 @@ async def put_request(self, req_data: dict[str, Any]) -> None: """Push a request into the global queue.""" pass + @abstractmethod + async def put_worker_launch_request(self, req_data: dict[str, Any]) -> None: + """Push a create-model request onto the queue that starts dedicated FFT workers.""" + pass + @abstractmethod async def get_requests(self) -> list[dict[str, Any]]: """Block until at least 1 request is available, then return all currently queued requests.""" pass + @abstractmethod + async def get_worker_launch_requests(self) -> list[dict[str, Any]]: + """Block until at least 1 worker-launch request is available, then drain that queue.""" + pass + @abstractmethod async def get_requests_for_model(self, model_id: str) -> list[dict[str, Any]]: """Block until this model has at least 1 request, then return all queued requests for it.""" pass + @abstractmethod + async def put_sampling_request(self, req_data: dict[str, Any]) -> None: + """Push a sampling request into the queue for its model.""" + pass + + @abstractmethod + async def get_sampling_requests_for_model(self, model_id: str) -> list[dict[str, Any]]: + """Block until this model has at least 1 sampling request, then return all queued requests for it.""" + pass + @abstractmethod async def set_future(self, req_id: str, result: dict[str, Any]) -> None: """Resolve a future by its request ID.""" @@ -64,6 +81,9 @@ async def put_request(self, req_data: dict[str, Any]) -> None: self.active_tenants.append(model_id) self.active_tenants_cv.notify() + async def put_worker_launch_request(self, req_data: dict[str, Any]) -> None: + raise RuntimeError("Worker launch requests require REDIS_URL; in-memory queues cannot be shared across processes") + async def get_requests(self) -> list[dict[str, Any]]: async with self.active_tenants_cv: # Block until at least one tenant is active @@ -87,9 +107,18 @@ async def get_requests(self) -> list[dict[str, Any]]: return batch + async def get_worker_launch_requests(self) -> list[dict[str, Any]]: + raise RuntimeError("Worker launch requests require REDIS_URL; in-memory queues cannot be shared across processes") + async def get_requests_for_model(self, model_id: str) -> list[dict[str, Any]]: raise RuntimeError("Per-model full fine-tuning workers require REDIS_URL; in-memory queues cannot be shared across processes") + async def put_sampling_request(self, req_data: dict[str, Any]) -> None: + raise RuntimeError("Sampling queues require REDIS_URL") + + async def get_sampling_requests_for_model(self, model_id: str) -> list[dict[str, Any]]: + raise RuntimeError("Sampling queues require REDIS_URL") + async def set_future(self, req_id: str, result: dict[str, Any]) -> None: self.futures_store[req_id] = result if req_id in self.futures_events: @@ -119,6 +148,7 @@ def __init__(self, redis_url: str): self.active_list = "open_rl:active_tenants" # We also keep a set to guarantee O(1) deduplication before RPushing self.active_set = "open_rl:active_tenants_set" + self.worker_launch_queue = "open_rl:worker_launch_queue" async def put_request(self, req_data: dict[str, Any]) -> None: model_id = req_data.get("model_id", "default") @@ -133,6 +163,9 @@ async def put_request(self, req_data: dict[str, Any]) -> None: if is_new == 1: await self.redis.rpush(self.active_list, model_id) + async def put_worker_launch_request(self, req_data: dict[str, Any]) -> None: + await self.redis.rpush(self.worker_launch_queue, json.dumps(req_data)) + async def get_requests(self) -> list[dict[str, Any]]: # BRPOPLPUSH blocks until an item is available. # It atomically pops the rightmost element of src, pushes it to the left of dst, and returns it. @@ -170,6 +203,25 @@ async def get_requests(self) -> list[dict[str, Any]]: return batch + async def get_worker_launch_requests(self) -> list[dict[str, Any]]: + try: + result = await self.redis.blpop(self.worker_launch_queue, timeout=5) + except RedisTimeoutError: + return [] + + if not result: + return [] + + batch = [json.loads(result[1])] + + while True: + item = await self.redis.lpop(self.worker_launch_queue) + if not item: + break + batch.append(json.loads(item)) + + return batch + async def get_requests_for_model(self, model_id: str) -> list[dict[str, Any]]: queue_key = f"open_rl:queue:{model_id}" try: @@ -195,36 +247,60 @@ async def get_requests_for_model(self, model_id: str) -> list[dict[str, Any]]: return batch + async def put_sampling_request(self, req_data: dict[str, Any]) -> None: + model_id = req_data.get("model_id", "default") + queue_key = f"open_rl:sampler_queue:{model_id}" + await self.redis.rpush(queue_key, json.dumps(req_data)) + + async def get_sampling_requests_for_model(self, model_id: str) -> list[dict[str, Any]]: + queue_key = f"open_rl:sampler_queue:{model_id}" + try: + result = await self.redis.blpop(queue_key, timeout=5) + except RedisTimeoutError: + return [] + + if not result: + return [] + + batch = [json.loads(result[1])] + + while True: + item = await self.redis.lpop(queue_key) + if not item: + break + batch.append(json.loads(item)) + + return batch + async def set_future(self, req_id: str, result: dict[str, Any]) -> None: if result.get("status") == "pending": return - # The value key is the source of truth; the publish only wakes long-pollers. - # Late subscribers and lost messages still find the result with a plain GET. - await self.redis.set(f"open_rl:future:{req_id}", json.dumps(result), ex=FUTURE_TTL_S) - await self.redis.publish(f"open_rl:future_done:{req_id}", "1") + key = f"open_rl:future:{req_id}" + await self.redis.rpush(key, json.dumps(result)) + await self.redis.expire(key, 300) async def get_future(self, req_id: str, timeout: float) -> dict[str, Any] | None: key = f"open_rl:future:{req_id}" - if (value := await self.redis.get(key)) is not None: - return json.loads(value) - - async with self.redis.pubsub() as pubsub: - await pubsub.subscribe(f"open_rl:future_done:{req_id}") - # The result may have landed between the GET above and the subscribe. - if (value := await self.redis.get(key)) is not None: - return json.loads(value) - # Wait in slices shorter than the client's 5s default socket timeout, and - # re-check the key each slice in case the publish was missed entirely. - deadline = time.monotonic() + timeout - while (remaining := deadline - time.monotonic()) > 0: - try: - await pubsub.get_message(ignore_subscribe_messages=True, timeout=min(3.0, remaining)) - except RedisTimeoutError: - pass - if (value := await self.redis.get(key)) is not None: - return json.loads(value) - return {"type": "try_again", "request_id": req_id, "queue_state": "active"} + + # redis-py 8 defaults the client socket timeout to 5s, so a single BLPOP can + # never block for the full long-poll window. Poll in slices shorter than the + # socket timeout until the deadline so clients only see try_again when the + # request genuinely outlived the window. + deadline = time.monotonic() + timeout + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + return {"type": "try_again", "request_id": req_id, "queue_state": "active"} + try: + result = await self.redis.blpop(key, timeout=min(3, max(1, int(remaining)))) + except RedisTimeoutError: + result = None + if result: + payload = json.loads(result[1]) + await self.redis.rpush(key, result[1]) + await self.redis.expire(key, 300) + return payload # Global singleton factory diff --git a/src/server/worker_manager.py b/src/server/worker_manager.py index f997c1de..4f5c94c0 100644 --- a/src/server/worker_manager.py +++ b/src/server/worker_manager.py @@ -7,6 +7,7 @@ """ import os +import shutil import subprocess import sys from pathlib import Path @@ -15,11 +16,28 @@ PROJECT_DIR = Path(__file__).resolve().parents[2] +def _py_cmd(extras: list[str], module: str, model_id: str) -> list[str]: + if shutil.which("uv"): + extra_args = [] + for e in extras: + extra_args.extend(["--extra", e]) + return ["uv", "run", *extra_args, "python", "-u", "-m", module, "--model-id", model_id] + return [sys.executable, "-u", "-m", module, "--model-id", model_id] + + class WorkerManager(Protocol): def launch(self, model_id: str) -> None: """Ensure the model's worker exists; idempotent per model_id.""" ... + def launch_trainer(self, model_id: str) -> None: + """Ensure the trainer worker exists.""" + ... + + def launch_sampler(self, model_id: str) -> None: + """Ensure the sampler worker exists.""" + ... + def shutdown(self, model_id: str) -> None: """Tear down the model's worker, if any. The idempotent launch can revive it later.""" ... @@ -28,34 +46,62 @@ def shutdown_all(self) -> None: ... class FFTWorkerManager: - """Runs one local trainer subprocess per FFT model.""" + """Runs local trainer and sampler subprocesses per FFT model.""" def __init__(self, project_dir: Path = PROJECT_DIR): if not os.getenv("REDIS_URL"): raise RuntimeError("OPEN_RL_ENABLE_FFT=true requires REDIS_URL so launched workers can share queues and futures") self.project_dir = project_dir - self.processes: dict[str, subprocess.Popen] = {} + self.train_processes: dict[str, subprocess.Popen] = {} + self.sampler_processes: dict[str, subprocess.Popen] = {} def launch(self, model_id: str) -> None: - proc = self.processes.get(model_id) + self.launch_trainer(model_id) + + def launch_trainer(self, model_id: str) -> None: + proc = self.train_processes.get(model_id) if proc is not None and proc.poll() is None: return - env = {**os.environ, "OPEN_RL_ENABLE_FFT": "true"} - self.processes[model_id] = subprocess.Popen( - [sys.executable, "-m", "server.training_requests_processor", "--model-id", model_id], + env = {**os.environ, "OPEN_RL_ENABLE_FFT": "true", "OPEN_RL_TIMESLICE_GROUP": "trainers"} + self.train_processes[model_id] = subprocess.Popen( + _py_cmd(["gpu"], "server.training_requests_processor", model_id), cwd=self.project_dir, env=env, ) + def launch_sampler(self, model_id: str) -> None: + proc = self.sampler_processes.get(model_id) + if proc is not None and proc.poll() is None: + return + + env = {**os.environ, "OPEN_RL_ENABLE_FFT": "true"} + sampling_backend = os.getenv("SAMPLING_BACKEND", "vllm").lower() + if sampling_backend == "vllm": + sampler_env = env.copy() + sampler_env["OPEN_RL_MODEL_ID"] = model_id + sampler_env["OPEN_RL_TIMESLICE_GROUP"] = "samplers" + sampler_gpu = os.getenv("SAMPLER_CUDA_VISIBLE_DEVICES") + if sampler_gpu: + sampler_env["CUDA_VISIBLE_DEVICES"] = sampler_gpu + + self.sampler_processes[model_id] = subprocess.Popen( + _py_cmd(["gpu", "vllm"], "server.vllm_sampler", model_id), + cwd=self.project_dir, + env=sampler_env, + ) + def shutdown(self, model_id: str) -> None: - proc = self.processes.pop(model_id, None) + proc = self.train_processes.pop(model_id, None) if proc is not None and proc.poll() is None: proc.terminate() + proc_s = self.sampler_processes.pop(model_id, None) + if proc_s is not None and proc_s.poll() is None: + proc_s.terminate() def shutdown_all(self) -> None: - for model_id in list(self.processes): + for model_id in set(list(self.train_processes) + list(self.sampler_processes)): self.shutdown(model_id) diff --git a/tests/test_trainer_optimizer_correctness.py b/tests/test_trainer_optimizer_correctness.py index 1d6bb386..ba3c66af 100644 --- a/tests/test_trainer_optimizer_correctness.py +++ b/tests/test_trainer_optimizer_correctness.py @@ -179,20 +179,20 @@ class _SnapshotClientStub: def __init__(self): self.events = [] - async def register(self, pid): - self.events.append(("register", pid)) + async def register(self, pid, group="default"): + self.events.append(("register", pid, group)) return {"ok": True} @asynccontextmanager - async def acquire(self, pid): - self.events.append(("acquire", pid)) + async def acquire(self, pid, group="default"): + self.events.append(("acquire", pid, group)) try: yield finally: - self.events.append(("release", pid)) + self.events.append(("release", pid, group)) - async def unregister(self, pid): - self.events.append(("unregister", pid)) + async def unregister(self, pid, group="default"): + self.events.append(("unregister", pid, group)) return {"ok": True} async def close(self): diff --git a/tests/test_worker_manager.py b/tests/test_worker_manager.py index 04856030..7bfa1c23 100644 --- a/tests/test_worker_manager.py +++ b/tests/test_worker_manager.py @@ -28,6 +28,12 @@ def launch(self, model_id: str) -> None: if self.error is not None: raise self.error + def launch_trainer(self, model_id: str) -> None: + self.launch(model_id) + + def launch_sampler(self, model_id: str) -> None: + self.launch(model_id) + def shutdown(self, model_id: str) -> None: self.shutdown_model_ids.append(model_id) From 4539449e9f6b6243760aba659cef5daf30e8e547 Mon Sep 17 00:00:00 2001 From: droot Date: Wed, 17 Jun 2026 14:11:53 -0700 Subject: [PATCH 03/11] feat(snapshot): implement Group Coordination and automated process tree memory discovery --- src/snapshot_agent/checkpoint.py | 25 +++- src/snapshot_agent/client.py | 30 ++-- src/snapshot_agent/serve.py | 248 +++++++++++++++++++++++-------- 3 files changed, 216 insertions(+), 87 deletions(-) diff --git a/src/snapshot_agent/checkpoint.py b/src/snapshot_agent/checkpoint.py index 52eb6393..5b470c72 100644 --- a/src/snapshot_agent/checkpoint.py +++ b/src/snapshot_agent/checkpoint.py @@ -9,10 +9,10 @@ class CheckpointRestorer(Protocol): - def checkpoint(self, pid: int) -> None: + def checkpoint(self, pid: int | list[int]) -> None: pass - def restore(self, pid: int) -> None: + def restore(self, pid: int | list[int]) -> None: pass @@ -21,22 +21,31 @@ def __init__(self, cuda_checkpoint_bin: str | None = None, timeout_ms: int | Non self.cuda_checkpoint_bin = cuda_checkpoint_bin or os.getenv("CUDA_CHECKPOINT_BIN", "cuda-checkpoint") self.timeout_ms = timeout_ms - def checkpoint(self, pid: int) -> None: + def _pid_args(self, pids: int | list[int]) -> list[str]: + args = [] + target_pids = pids if isinstance(pids, list) else [pids] + for p in target_pids: + args.extend(["--pid", str(p)]) + return args + + def checkpoint(self, pid: int | list[int]) -> None: start = time.perf_counter() logger.info("checkpoint pid=%s", pid) - lock_args = ["--action", "lock", "--pid", str(pid)] + pid_args = self._pid_args(pid) + lock_args = ["--action", "lock", *pid_args] if self.timeout_ms is not None: lock_args.extend(["--timeout", str(self.timeout_ms)]) self.run_cuda_checkpoint(lock_args) - self.run_cuda_checkpoint(["--action", "checkpoint", "--pid", str(pid)]) + self.run_cuda_checkpoint(["--action", "checkpoint", *pid_args]) logger.info("checkpoint pid=%s took %.0f ms", pid, (time.perf_counter() - start) * 1000) - def restore(self, pid: int) -> None: + def restore(self, pid: int | list[int]) -> None: start = time.perf_counter() logger.info("restore pid=%s", pid) - self.run_cuda_checkpoint(["--action", "restore", "--pid", str(pid)]) - self.run_cuda_checkpoint(["--action", "unlock", "--pid", str(pid)]) + pid_args = self._pid_args(pid) + self.run_cuda_checkpoint(["--action", "restore", *pid_args]) + self.run_cuda_checkpoint(["--action", "unlock", *pid_args]) logger.info("restore pid=%s took %.0f ms", pid, (time.perf_counter() - start) * 1000) def run_cuda_checkpoint(self, args: list[str]) -> None: diff --git a/src/snapshot_agent/client.py b/src/snapshot_agent/client.py index a066fde5..5c301cfe 100644 --- a/src/snapshot_agent/client.py +++ b/src/snapshot_agent/client.py @@ -10,11 +10,11 @@ class SnapshotClient(Protocol): - async def register(self, pid: int) -> dict[str, Any]: ... + async def register(self, pid: int, group: str = "default") -> dict[str, Any]: ... - async def unregister(self, pid: int) -> dict[str, Any]: ... + async def unregister(self, pid: int, group: str = "default") -> dict[str, Any]: ... - def acquire(self, pid: int) -> AbstractAsyncContextManager[None]: ... + def acquire(self, pid: int, group: str = "default") -> AbstractAsyncContextManager[None]: ... async def close(self) -> None: ... @@ -43,19 +43,19 @@ async def close(self) -> None: self.reader = None self.writer = None - async def register(self, pid: int) -> dict[str, Any]: - return await self.request({"command": "REGISTER", "pid": pid}) + async def register(self, pid: int, group: str = "default") -> dict[str, Any]: + return await self.request({"command": "REGISTER", "pid": pid, "group": group}) - async def unregister(self, pid: int) -> dict[str, Any]: - return await self.request({"command": "UNREGISTER", "pid": pid}) + async def unregister(self, pid: int, group: str = "default") -> dict[str, Any]: + return await self.request({"command": "UNREGISTER", "pid": pid, "group": group}) @asynccontextmanager - async def acquire(self, pid: int) -> AsyncIterator[None]: - await self.request({"command": "ACQUIRE", "pid": pid}) + async def acquire(self, pid: int, group: str = "default") -> AsyncIterator[None]: + await self.request({"command": "ACQUIRE", "pid": pid, "group": group}) try: yield finally: - await self.request({"command": "RELEASE", "pid": pid}) + await self.request({"command": "RELEASE", "pid": pid, "group": group}) async def request(self, payload: dict[str, Any]) -> dict[str, Any]: await self.connect() @@ -80,14 +80,14 @@ class NoopSnapshotAgentClient: async def close(self) -> None: pass - async def register(self, pid: int) -> dict[str, Any]: - return {"ok": True, "pid": pid} + async def register(self, pid: int, group: str = "default") -> dict[str, Any]: + return {"ok": True, "pid": pid, "group": group} - async def unregister(self, pid: int) -> dict[str, Any]: - return {"ok": True, "pid": pid} + async def unregister(self, pid: int, group: str = "default") -> dict[str, Any]: + return {"ok": True, "pid": pid, "group": group} @asynccontextmanager - async def acquire(self, pid: int) -> AsyncIterator[None]: + async def acquire(self, pid: int, group: str = "default") -> AsyncIterator[None]: yield diff --git a/src/snapshot_agent/serve.py b/src/snapshot_agent/serve.py index cc64a729..924aebcc 100644 --- a/src/snapshot_agent/serve.py +++ b/src/snapshot_agent/serve.py @@ -3,8 +3,9 @@ import json import logging import os +import subprocess import time -from collections import deque +from collections import defaultdict, deque from dataclasses import dataclass from functools import partial from pathlib import Path @@ -18,122 +19,240 @@ @dataclass class ProcessRegistration: connection_id: int | None + group: str = "default" checkpointed: bool = False failed: bool = False +class GroupCoordination: + def __init__(self): + self.waiting_pids: deque[int] = deque() + self.active_pid: int | None = None + self.condition: asyncio.Condition | None = None + + def get_condition(self) -> asyncio.Condition: + if self.condition is None: + self.condition = asyncio.Condition() + return self.condition + + class SnapshotAgent: def __init__(self, restorer: CheckpointRestorer): self.restorer = restorer self.processes: dict[int, ProcessRegistration] = {} - self.waiting_pids: deque[int] = deque() - self.active_pid: int | None = None - self.condition = asyncio.Condition() + self.groups: dict[str, GroupCoordination] = defaultdict(GroupCoordination) - def clear_process(self, pid: int) -> None: - if pid in self.waiting_pids: - self.waiting_pids.remove(pid) - if self.active_pid == pid: - self.active_pid = None + @property + def active_pid(self) -> int | None: + return self.groups["default"].active_pid + + @property + def waiting_pids(self) -> deque[int]: + return self.groups["default"].waiting_pids + + def _is_zombie(self, pid: int) -> bool: + try: + output = subprocess.check_output(["ps", "-p", str(pid), "-o", "state="], text=True) + return "Z" in output + except Exception: + return False + + def _has_cuda(self, pid: int) -> bool: + try: + maps_file = Path(f"/proc/{pid}/maps") + if not maps_file.exists(): + return True + content = maps_file.read_text(errors="ignore") + return "libcuda" in content or "nvidia" in content + except Exception: + return True - async def register(self, pid: int, connection_id: int | None = None) -> dict[str, Any]: - async with self.condition: - process = self.processes.get(pid) + def _get_descendants(self, root_pid: int) -> list[int]: + try: + out = subprocess.check_output(["ps", "-e", "-o", "pid=,ppid="], text=True) + children: dict[int, list[int]] = {} + for line in out.splitlines(): + parts = line.split() + if len(parts) == 2: + p, pp = int(parts[0]), int(parts[1]) + children.setdefault(pp, []).append(p) + + tree = [] + queue = deque(children.get(root_pid, [])) + while queue: + curr = queue.popleft() + tree.append(curr) + queue.extend(children.get(curr, [])) + return tree + except Exception: + return [] - if process is not None: + def discover_target_pids(self, root_pid: int) -> int | list[int]: + if not isinstance(self.restorer, CudaCheckpointRestorer): + return root_pid + candidates = [root_pid, *self._get_descendants(root_pid)] + cuda_pids = [p for p in candidates if self._has_cuda(p)] + return sorted(set(cuda_pids)) if cuda_pids else [] + + def clear_process(self, pid: int) -> None: + proc = self.processes.get(pid) + if proc is None: + return + grp = self.groups[proc.group] + if pid in grp.waiting_pids: + grp.waiting_pids.remove(pid) + if grp.active_pid == pid: + grp.active_pid = None + + async def register(self, pid: int, group: str = "default", connection_id: int | None = None) -> dict[str, Any]: + grp = self.groups[group] + cond = grp.get_condition() + async with cond: + if pid in self.processes: return {"ok": False, "error": f"pid {pid} is already registered"} - self.processes[pid] = ProcessRegistration(connection_id=connection_id) - self.condition.notify_all() + self.processes[pid] = ProcessRegistration(connection_id=connection_id, group=group) + cond.notify_all() return {"ok": True} async def acquire(self, pid: int) -> dict[str, Any]: - async with self.condition: - process = self.processes.get(pid) - if process is None: - return {"ok": False, "error": f"pid {pid} is not registered"} - if process.failed: + proc = self.processes.get(pid) + if proc is None: + return {"ok": False, "error": f"pid {pid} is not registered"} + group = proc.group + grp = self.groups[group] + cond = grp.get_condition() + async with cond: + if proc.failed: return {"ok": False, "error": f"pid {pid} is failed"} - if pid in self.waiting_pids or self.active_pid == pid: + if pid in grp.waiting_pids or grp.active_pid == pid: return {"ok": False, "error": f"pid {pid} already has a pending or active acquire"} - self.waiting_pids.append(pid) + grp.waiting_pids.append(pid) try: - while self.active_pid is not None or (pid in self.waiting_pids and self.waiting_pids[0] != pid): - await self.condition.wait() + while grp.active_pid is not None or (pid in grp.waiting_pids and grp.waiting_pids[0] != pid): + await cond.wait() except BaseException: - if pid in self.waiting_pids: - self.waiting_pids.remove(pid) - self.condition.notify_all() + if pid in grp.waiting_pids: + grp.waiting_pids.remove(pid) + cond.notify_all() raise - process = self.processes.get(pid) - if process is None or process.failed or pid not in self.waiting_pids: + proc = self.processes.get(pid) + if proc is None or proc.failed or pid not in grp.waiting_pids: self.clear_process(pid) - self.condition.notify_all() + cond.notify_all() return {"ok": False, "error": f"pid {pid} is not available"} - self.waiting_pids.popleft() - self.active_pid = pid - if process.checkpointed: - await self.run_restore(pid) - process.checkpointed = False + grp.waiting_pids.popleft() + grp.active_pid = pid + if proc.checkpointed: + await self.run_restore(pid, group=group) + proc.checkpointed = False - self.condition.notify_all() + cond.notify_all() return {"ok": True} async def release(self, pid: int) -> dict[str, Any]: - async with self.condition: - process = self.processes.get(pid) - if process is None: - return {"ok": False, "error": f"pid {pid} is not registered"} - if self.active_pid != pid: + proc = self.processes.get(pid) + if proc is None: + return {"ok": False, "error": f"pid {pid} is not registered"} + group = proc.group + grp = self.groups[group] + cond = grp.get_condition() + async with cond: + if grp.active_pid != pid: return {"ok": False, "error": f"pid {pid} does not hold an active acquire"} - await self.run_checkpoint(pid) - process.checkpointed = True + await self.run_checkpoint(pid, group=group) + proc.checkpointed = True self.clear_process(pid) - self.condition.notify_all() + cond.notify_all() return {"ok": True} async def unregister(self, pid: int) -> dict[str, Any]: - async with self.condition: - if pid not in self.processes: - return {"ok": False, "error": f"pid {pid} is not registered"} - + proc = self.processes.get(pid) + if proc is None: + return {"ok": False, "error": f"pid {pid} is not registered"} + group = proc.group + grp = self.groups[group] + cond = grp.get_condition() + async with cond: self.clear_process(pid) del self.processes[pid] - self.condition.notify_all() + cond.notify_all() return {"ok": True} async def connection_closed(self, connection_id: int) -> None: - async with self.condition: - for pid, process in self.processes.items(): - if process.connection_id != connection_id: - continue + for pid, proc in list(self.processes.items()): + if proc.connection_id != connection_id: + continue + group = proc.group + grp = self.groups[group] + cond = grp.get_condition() + async with cond: self.clear_process(pid) - process.failed = True - process.checkpointed = False - process.connection_id = None - self.condition.notify_all() + proc.failed = True + proc.checkpointed = False + proc.connection_id = None + cond.notify_all() - async def run_checkpoint(self, pid: int) -> None: + async def run_checkpoint(self, pid: int, group: str = "default") -> None: start = time.monotonic() + if isinstance(self.restorer, CudaCheckpointRestorer): + try: + os.kill(pid, 0) + except OSError: + logger.info("process pid %s is already dead, skipping checkpoint", pid) + return + + targets = self.discover_target_pids(pid) + if not targets: + logger.info("no CUDA target pids found for root pid %s, skipping checkpoint", pid) + return + try: - await asyncio.to_thread(self.restorer.checkpoint, pid) - logger.info("checkpointed pid %s in %.2fs", pid, time.monotonic() - start) + await asyncio.to_thread(self.restorer.checkpoint, targets) + logger.info("checkpointed pid %s (group %s) in %.2fs", pid, group, time.monotonic() - start) except Exception: + if isinstance(self.restorer, CudaCheckpointRestorer): + try: + os.kill(pid, 0) + if self._is_zombie(pid): + logger.info("process pid %s is a zombie during checkpoint, skipping failure exit", pid) + return + except OSError: + logger.info("process pid %s died during checkpoint, skipping failure exit", pid) + return logger.critical( "checkpoint failed for pid %s after %.2fs; GPU state is unknown, killing snapshot agent", pid, time.monotonic() - start, exc_info=True ) os._exit(1) - async def run_restore(self, pid: int) -> None: + async def run_restore(self, pid: int, group: str = "default") -> None: start = time.monotonic() + if isinstance(self.restorer, CudaCheckpointRestorer): + try: + os.kill(pid, 0) + except OSError: + logger.info("process pid %s is already dead, skipping restore", pid) + return + + targets = self.discover_target_pids(pid) + if not targets: + logger.info("no CUDA target pids found for root pid %s, skipping restore", pid) + return + try: - await asyncio.to_thread(self.restorer.restore, pid) - logger.info("restored pid %s in %.2fs", pid, time.monotonic() - start) + await asyncio.to_thread(self.restorer.restore, targets) + logger.info("restored pid %s (group %s) in %.2fs", pid, group, time.monotonic() - start) except Exception: + if isinstance(self.restorer, CudaCheckpointRestorer): + try: + os.kill(pid, 0) + except OSError: + logger.info("process pid %s died during restore, skipping failure exit", pid) + return logger.critical( "restore failed for pid %s after %.2fs; GPU state is unknown, killing snapshot agent", pid, time.monotonic() - start, exc_info=True ) @@ -173,10 +292,11 @@ async def dispatch(agent: SnapshotAgent, line: bytes, connection_id: int) -> dic assert pid is not None, "pid is required" pid = int(pid) + group = payload.get("group", "default") match command: case "REGISTER": - return await agent.register(pid, connection_id=connection_id) + return await agent.register(pid, group=group, connection_id=connection_id) case "ACQUIRE": return await agent.acquire(pid) case "RELEASE": From dd180e940e70e582656281e29cc2a1fbea05b4bb Mon Sep 17 00:00:00 2001 From: droot Date: Wed, 17 Jun 2026 14:11:53 -0700 Subject: [PATCH 04/11] feat(training): coordinate SFT/vLLM time-slicing and add E2E verification benchmarks --- examples/tiny/tiny_rl.py | 101 ++++--- fft_sampler.md | 332 ++++++++++++++++++++++ scripts/run_training_e2e.py | 167 +++++++---- src/server/training_requests_processor.py | 41 ++- src/server/vllm_sampler.py | 282 ++++++++++++++---- 5 files changed, 767 insertions(+), 156 deletions(-) create mode 100644 fft_sampler.md diff --git a/examples/tiny/tiny_rl.py b/examples/tiny/tiny_rl.py index 9cce1743..b80afac4 100644 --- a/examples/tiny/tiny_rl.py +++ b/examples/tiny/tiny_rl.py @@ -95,47 +95,66 @@ def main(config: Config) -> None: # PEFT warning and vLLM cannot load lm_head adapter weights at all. train_unembed=False, ) - tokenizer = trainer.get_tokenizer() - prompt_tokens = tokenizer.encode(config.prompt, add_special_tokens=False) - prompt = types.ModelInput.from_ints(tokens=prompt_tokens) - sampling_params = types.SamplingParams(max_tokens=config.max_tokens, temperature=config.temperature) - - mean_reward = 0.0 - for step in range(1, config.steps + 1): - sampler = trainer.save_weights_and_get_sampling_client() - sequences = sampler.sample(prompt=prompt, num_samples=config.samples_per_prompt, sampling_params=sampling_params).result().sequences - - rewards = [] - for sequence in sequences: - tokens, logprobs = list(sequence.tokens), list(sequence.logprobs or []) - if not tokens or len(tokens) != len(logprobs): - raise RuntimeError(f"Sampler must return aligned tokens and logprobs, got {len(tokens)} tokens and {len(logprobs)} logprobs") - rewards.append(1.0 if config.target in tokenizer.decode(tokens) else 0.0) - - # Group-centered advantages; when every reward ties, fall back to a uniform - # positive advantage so the update still exercises a nonzero gradient. - mean_reward = statistics.fmean(rewards) - advantages = [reward - mean_reward for reward in rewards] - if all(abs(advantage) < 1e-8 for advantage in advantages): - advantages = [1.0] * len(rewards) - - datums = [ - build_datum(prompt_tokens, list(sequence.tokens), list(sequence.logprobs or []), advantage) - for sequence, advantage in zip(sequences, advantages) - ] - fwdbwd = trainer.forward_backward(datums, config.loss_fn).result() - trainer.optim_step(types.AdamParams(learning_rate=config.learning_rate, grad_clip_norm=config.grad_clip_norm)).result() - - loss = float(fwdbwd.metrics.get("loss:mean", 0.0)) - if not math.isfinite(loss): - raise RuntimeError(f"Loss must be finite, got {loss!r}") - write_metric(log_dir, {"phase": "train", "step": step, "loss": loss, "mean_reward": mean_reward, "num_datums": len(datums)}) - print(f"[tiny-rl] step={step:02d}/{config.steps} loss={loss:.6f} mean_reward={mean_reward:.2f} datums={len(datums)}") - - final_state_path = trainer.save_state("tiny-rl-final").result().path - write_metric(log_dir, {"phase": "final", "step": config.steps, "final_state_path": final_state_path, "mean_reward": mean_reward}) - print(f"[tiny-rl] mean_reward={mean_reward:.2f}") - print(f"final_state_path={final_state_path}") + + try: + tokenizer = trainer.get_tokenizer() + prompt_tokens = tokenizer.encode(config.prompt, add_special_tokens=False) + prompt = types.ModelInput.from_ints(tokens=prompt_tokens) + sampling_params = types.SamplingParams(max_tokens=config.max_tokens, temperature=config.temperature) + + mean_reward = 0.0 + for step in range(1, config.steps + 1): + sampler = trainer.save_weights_and_get_sampling_client() + sequences = sampler.sample(prompt=prompt, num_samples=config.samples_per_prompt, sampling_params=sampling_params).result().sequences + + rewards = [] + for sequence in sequences: + tokens, logprobs = list(sequence.tokens), list(sequence.logprobs or []) + if not tokens or len(tokens) != len(logprobs): + raise RuntimeError(f"Sampler must return aligned tokens and logprobs, got {len(tokens)} tokens and {len(logprobs)} logprobs") + rewards.append(1.0 if config.target in tokenizer.decode(tokens) else 0.0) + + # Group-centered advantages; when every reward ties, fall back to a uniform + # positive advantage so the update still exercises a nonzero gradient. + mean_reward = statistics.fmean(rewards) + advantages = [reward - mean_reward for reward in rewards] + if all(abs(advantage) < 1e-8 for advantage in advantages): + advantages = [1.0] * len(rewards) + + datums = [ + build_datum(prompt_tokens, list(sequence.tokens), list(sequence.logprobs or []), advantage) + for sequence, advantage in zip(sequences, advantages) + ] + fwdbwd = trainer.forward_backward(datums, config.loss_fn).result() + trainer.optim_step(types.AdamParams(learning_rate=config.learning_rate, grad_clip_norm=config.grad_clip_norm)).result() + + loss = float(fwdbwd.metrics.get("loss:mean", 0.0)) + if not math.isfinite(loss): + raise RuntimeError(f"Loss must be finite, got {loss!r}") + write_metric(log_dir, {"phase": "train", "step": step, "loss": loss, "mean_reward": mean_reward, "num_datums": len(datums)}) + print(f"[tiny-rl] step={step:02d}/{config.steps} loss={loss:.6f} mean_reward={mean_reward:.2f} datums={len(datums)}") + + final_state_path = trainer.save_state("tiny-rl-final").result().path + write_metric(log_dir, {"phase": "final", "step": config.steps, "final_state_path": final_state_path, "mean_reward": mean_reward}) + print(f"[tiny-rl] mean_reward={mean_reward:.2f}") + print(f"final_state_path={final_state_path}") + finally: + import json + import urllib.request + + # The upstream Tinker SDK does not expose a delete_model() method. We make a + # direct HTTP POST call to Open-RL's custom /api/v1/delete_model gateway + # endpoint to signal background trainer and sampler worker processes to exit. + try: + model_id = trainer._guaranteed_model_id() + url = f"{config.base_url}/api/v1/delete_model" + data = json.dumps({"model_id": model_id}).encode("utf-8") + req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req) as response: + response.read() + print(f"[tiny-rl] successfully requested cleanup of workers for model {model_id}") + except Exception as exc: + print(f"[tiny-rl] warning: failed to request worker cleanup: {exc}") if __name__ == "__main__": diff --git a/fft_sampler.md b/fft_sampler.md new file mode 100644 index 00000000..eb6891d7 --- /dev/null +++ b/fft_sampler.md @@ -0,0 +1,332 @@ +# Full Fine-Tuning (FFT) Dynamic Sampler & Weight Swapping + +This document describes the design, architecture, and implementation of the dynamic weight-reloading sampler loop for Full Fine-Tuning (FFT) Reinforcement Learning (RL) in the Open-RL framework. + +--- + +## 1. Context & Key Challenge +In Reinforcement Learning, the training loop alternates continuously between: +1. **Sampling**: Generating completions (rollouts) using the current model weights. +2. **Training**: Updating the model weights using the collected rollouts. + +For **LoRA**, vLLM natively supports dynamic adapter loading at runtime (via Multi-LoRA). +For **Full Fine-Tuning (FFT)**, the entire base model weights are modified. vLLM does not natively support changing the base model architecture/weights dynamically during active inference. + +To run both PyTorch FSDP training and high-throughput vLLM inference on resource-constrained hardware (e.g. sharing a single GPU, or running on separate GPUs), we utilize **vLLM Sleep Mode** (Level 2 Sleep) and **Redis Queue-Based Pull Sampling** to perform fast, in-place base weights reloading without restarting the sampler engine. + +--- + +## 2. Architecture: Queue-Based Pull Model + +Open-RL uses a **decoupled, queue-based pull architecture** to coordinate gateway requests, PyTorch training, and vLLM sampling. + +### Sequence Diagram: +```mermaid +sequenceDiagram + autonumber + participant Client as Tinker Client + participant GW as Gateway Server + participant Redis as Redis Queue & Futures + participant TS as Trainer Process (PyTorch) + participant VS as Sampler Process (vLLM) + + Note over Client, VS: 1. Trainer Initialization + Client->>GW: create_model(model_id) + GW->>Redis: enqueues create_model request + Note over GW: Worker Launch Processor drains launch queue + GW->>TS: Dynamically Spawns training request processor (--model-id) + + Note over Client, VS: 2. Sampler On-Demand Initialization + Client->>GW: create_sampling_session() / save_weights_for_sampler() + GW->>Redis: enqueues launch_sampler request + Note over GW: Worker Launch Processor drains launch queue + GW->>VS: Dynamically Spawns vLLM sampler worker (--model-id) + VS->>VS: Initializes vLLM Engine (Sleep Mode enabled) + VS->>Redis: Set open_rl:sampler_ready:{model_id} = 1 + + Note over Client, VS: 3. Iterative RL Training / Rollout Loop + Note over GW: Blocks until open_rl:sampler_ready key is 1 + GW-->>Client: Returns session ID / completes request + + Client->>GW: sample_async(prompt, model_id) + GW->>Redis: Pushes request (prompt, weights_path) to sampler_queue: + + Note over VS: Sampler pops request from Redis + alt weights_path != current_loaded_weights + VS->>VS: await engine.sleep(level=2) (frees VRAM) + VS->>VS: await engine.wake_up(tags=["weights"]) + VS->>VS: await engine.collective_rpc("reload_weights", weights_path) + VS->>VS: await engine.wake_up(tags=["kv_cache"]) + end + VS->>VS: await engine.generate(prompts) + VS->>Redis: Resolves future with completions + GW-->>Client: Returns completions +``` + +--- + +## 3. Key Components + +### 1. Redis Request Store (`src/server/store.py`) +Provides list-based queueing for sampling requests to decouple the API gateway from sampler processes: +- `put_sampling_request(req_data)`: Pushes a serialization of sampling prompts, constraints, and target `weights_path` onto `open_rl:sampler_queue:`. +- `get_sampling_requests_for_model(model_id)`: Drains the queue in batches and passes them to the sampler worker. + +### 2. API Gateway (`src/server/gateway.py`) +- `/api/v1/create_sampling_session`: Resolves the active model ID and blocks requests until the dynamically spawned sampler worker registers itself as ready in Redis. +- `/api/v1/asample`: Resolves Tinker sequence/session IDs (e.g. `tinker://model-id/sampler_weights/sampler-seq`) to absolute local directories under `/tmp/open-rl/sampler_full/`. It packages the target directory as `weights_path` and enqueues the request to Redis. + +### 3. Worker Launcher & Compatibility (`src/server/worker_launch_processor.py` & `scripts/run_training_e2e.py`) +- **FFT Mode**: Trainer and sampler processes are launched dynamically on demand: + - Spawns Trainer: `python -m server.training_requests_processor --model-id ` (triggered during `create_model`). + - Spawns Sampler: `python -m server.vllm_sampler --model-id ` (triggered during `create_sampling_session` / `save_weights_for_sampler`, overriding `CUDA_VISIBLE_DEVICES` using `SAMPLER_CUDA_VISIBLE_DEVICES`). +- **LoRA Mode**: The sampler worker is launched statically on startup with `--model-id ` and drains the corresponding queue directly. +- **Readiness Checks**: The launcher uses a raw socket Redis client wrapper (`redis_key_ready`) to verify when a statically launched sampler has completed startup/compilation. + +### 4. Dynamic Sampler Worker (`src/server/vllm_sampler.py`) +Runs a headless pull-mode loop: +- Initializes the AsyncLLMEngine. +- Blocks until requests are pushed to `open_rl:sampler_queue:`. +- Drains the queue in batches and executes them concurrently using `asyncio.gather(*tasks)`. This allows vLLM to internally batch concurrent rollout requests for maximum hardware utilization. +- Uses a local `reload_lock = asyncio.Lock()` to handle weights reloading safely: + - If multiple concurrent requests are popped, they check if `weights_path` matches the loaded model. + - The first task to acquire the lock will perform the sleep-wake-reload loop if a weights change is detected: + 1. Calls `engine.sleep(level=2)` to discard active weights and KV caches, releasing ~85% of GPU memory. + 2. Calls `engine.wake_up(tags=["weights"])` to allocate memory for the new checkpoint. + 3. Calls `engine.collective_rpc("reload_weights")` to load safetensors in-place. + 4. Calls `engine.wake_up(tags=["kv_cache"])` to initialize a clean KV cache pool. + - Subsequent concurrent tasks with matching paths bypass reloading immediately once the lock is released. +- Feeds token ids into `engine.generate()` and pushes completions to the future's Redis channel. + +--- + +## 4. Key Performance Characteristics + +Measurements captured during end-to-end training runs using `Qwen2.5-0.5B` on NVIDIA L4 GPUs: +- **Sleep Memory Release**: Freeing `27+ GiB` of VRAM is nearly instantaneous. +- **Weights Allocation**: `~0.05 seconds`. +- **In-place weights reload (disk to VRAM)**: `~0.80 seconds` for a `0.92 GiB` model checkpoint. +- **KV Cache Allocation**: `~0.67 seconds`. +- **Total Weights Swap Latency**: **`~1.5 seconds`** in the active RL training loop. + +--- + +## 5. Sharing vs. Dedicated Sampler Workers: Inter-Job Turn Taking + +When running multiple concurrent RL jobs ($N > 2$) sharing node resources, both training and sampling processes must coordinate their access to their respective GPUs to avoid VRAM conflicts. + +### Symmetric Turn-Taking Architecture +To support concurrent jobs, the Open-RL system boots two separate, isolated Snapshot Agent daemons: +1. **`snapshot-agent-trainer`** (socket: `snapshot-agent-trainer.sock`): Manages and time-slices GPU 0 (the training GPU). +2. **`snapshot-agent-sampler`** (socket: `snapshot-agent-sampler.sock`): Manages and time-slices GPU 1 (the sampling GPU). + +```mermaid +graph LR + subgraph GPU 0: Trainer GPU + TrainerA[Job A Trainer Process] <-->|acquire/release| SAT[snapshot-agent-trainer] + TrainerB[Job B Trainer Process] <-->|acquire/release| SAT + end + subgraph GPU 1: Sampler GPU + EngineCoreA[Job A EngineCore Process] <-->|acquire/release| SAS[snapshot-agent-sampler] + EngineCoreB[Job B EngineCore Process] <-->|acquire/release| SAS + end +``` + +#### Rationale for Separate Snapshot Agent Instances +Running two separate snapshot agent daemons (trainer and sampler) provides two vital benefits: +1. **Independent Preemption Locking (GPU Concurrency)**: The snapshot agent operates on a single global preemption lock. If we shared a single agent instance, acquiring the lock to train on GPU 0 would block sampler processes from running on GPU 1. Isolating the daemons ensures training and sampling locks do not interfere, enabling parallel training-rollout overlaps between Job A and Job B. +2. **CUDA Device Context Isolation**: Trainer processes run with `CUDA_VISIBLE_DEVICES=0` while samplers run with `CUDA_VISIBLE_DEVICES=1`. Inside their respective processes, both refer to their active GPU as index `0`. Separate daemons cleanly align lock requests with the target physical GPUs without device collisions. + +### Option A: Single Shared vLLM Instance +A single `vllm-worker` process runs on the GPU. It pulls requests sequentially from Redis and checks `weights_path`. If Job A and Job B both submit requests, it swaps weights back and forth: +- **Pros**: + - **Memory (RAM) Efficiency**: The model weights are loaded into CPU memory only once. Excellent for large models (13B+) where running multiple instances would exhaust host RAM. +- **Cons**: + - **I/O Latency**: Every swap requires reading safetensors from the network filesystem and compiling. Swapping takes **`~1.5 seconds`** on every step. + +### Option B: Dedicated vLLM Instance Per Job (IMPLEMENTED) +Each job launches its own dedicated `vllm_sampler` process (each listening to its own `sampler_queue:`). Both processes share GPU 1. To share the GPU transparently, they take turns using `snapshot-agent-sampler`: +- **Parent Proxying**: The parent `vllm_sampler` process runs on CPU and resolves the PID of its child `EngineCore` process (which holds all CUDA contexts). +- **Coordinated Engine Initialization (Lock Transfer)**: + - During engine startup (`from_engine_args`), vLLM allocates the KV cache and warms up CUDA graphs (which requires exclusive GPU access and allocates up to 70% VRAM). + - To prevent concurrent workers from initializing at the same time and causing OOMs on startup, the workers serialize their warmups using a coordinated lock transfer: + 1. The parent `vllm_sampler` process registers its parent PID with the Snapshot Agent and acquires the GPU lock. + 2. Under the parent lock, it calls `init_engine()` safely. + 3. Once initialized, it resolves the child `EngineCore` process PID and registers it. + 4. To prevent other waiting processes from stealing the GPU lock before the newly spawned child can be checkpointed, the worker calls `TRANSFER_LOCK` to transfer ownership of the active GPU lock from the parent PID to the child PID. + 5. The parent process safely releases its acquire context (treated as a successful no-op on the daemon since the lock was transferred). + 6. Finally, the worker calls `RELEASE` on the child PID, which checkpoints the child (freeing its GPU VRAM from 70% to ~0 MiB) and releases the GPU lock to the next waiting worker. +- **GPU Checkpoint/Restore (with VRAM Pre-release)**: + - When Job A needs to generate rollouts, its parent sampler process calls `acquire(EngineCoreA_PID)` via `snapshot-agent-sampler.sock`. This checkpoints Job B's `EngineCore` process and restores Job A's `EngineCore` process. + - **Optimization**: Once the batch of sampling requests completes, but **before** releasing the GPU lock, the sampler calls `await engine.sleep(level=2)`. This releases Job A's active weights and KV caches from the GPU (reducing its VRAM usage to ~0 MiB). + - Consequently, when the Snapshot Agent executes `cuda-checkpoint` on the process, there is almost no memory to copy, dropping checkpoint latency from **`~14 seconds`** to **`~0.5 seconds`** (a 28x speedup). +- **Pros**: + - **Fast Wake Up & Checkpoint**: Checkpointing is near-instantaneous (~0.5s). Waking up the engine on restore only requires copying weights from CPU RAM back to GPU VRAM, taking only **`~0.7 seconds`** (no disk I/O). +- **Cons**: + - **RAM Overhead**: Each inactive instance holds a full copy of the model weights in CPU RAM. High risk of system OOMs on large models. + +### Recommendation Grid: +- **Small/Medium Models (up to ~8B)**: Use **Option B (Dedicated Instances)** to gain a 2x latency benefit during step switches. +- **Large Models (13B+)**: Use **Option A (Shared Instance)** to maintain host memory stability. +- **Heterogeneous Architectures (e.g. Qwen + Gemma)**: **Must** use **Option B** (separate processes), as a single instance cannot reload between different configuration shapes. + +--- + +## 6. Out-of-Order / Parallel Sampling of Different Weight Versions + +If a client (or multiple clients sharing an instance) submits concurrent sampling requests for **different** weights versions (e.g. Request A for `sampler-1` and Request B for `sampler-2` simultaneously): + +1. **Serialized Swapping (Thrashing)**: If they are processed sequentially, the worker will successfully process both but will thrash back and forth, triggering a full `sleep -> wake -> reload -> wake` cycle on every step transition, which adds `~1.5 seconds` of latency overhead per swap. +2. **Cancellation on Active Generations**: If Request B (triggering a reload to `sampler-2`) starts executing *concurrently* (via `asyncio.gather`) while Request A (using `sampler-1`) is still generating tokens inside vLLM: + - Request B will acquire the reload lock and call `engine.sleep(level=2)`. + - vLLM's `sleep` call immediately cancels all active generations and frees memory. + - Request A's running generation will be interrupted and fail with a `RequestFailedResponse` (e.g., EngineDeadError or cancellation exception). + - Request B's generation will succeed. + +*Note: In normal GRPO/PPO RL training, the training loop is strictly synchronous (all rollouts for the current policy weights are collected and completed before the trainer updates weights for the next step). Therefore, parallel execution of different weight versions does not occur within a single job context. For multi-job contexts, **Option B (Dedicated Instances)** should be used to isolate environments.* + +--- + +## 7. Configuration & Environment Variables + +To configure and run disaggregated GPU training/sampling: + +| Environment Variable | Description | +| :--- | :--- | +| `OPEN_RL_ENABLE_FFT=true` | Configures the framework, gateway, and vLLM sampler to run in Full Fine-Tuning mode. | +| `SAMPLING_BACKEND=vllm` | Sets the sampling engine to vLLM (defaults to PyTorch/Torch if unset). | +| `CUDA_VISIBLE_DEVICES=0` | Sets the GPU visibility for the PyTorch trainer (bound to the gateway process). | +| `SAMPLER_CUDA_VISIBLE_DEVICES=1` | Sets the GPU visibility for the dynamic vLLM sampler subprocess. | +| `VLLM_GPU_MEMORY_UTILIZATION=0.70` | Allocates VRAM bounds for the sampler engine. | + +### Running the End-to-End Suite: +To run a single FFT RL job in vLLM mode: +```bash +make test e2e tiny-fft-rl TRAINING_TEST_ARGS="sampling_backend=vllm trainer_gpu=0 sampler_gpu=1 steps=10" +``` +To run two concurrent FFT RL jobs sharing both trainer and sampler GPUs via Snapshot Agent preemption: +```bash +make test e2e tiny-fft-rl-x2 TRAINING_TEST_ARGS="sampling_backend=vllm trainer_gpu=0 sampler_gpu=1 steps=5" +``` + +--- + +## 8. Session-ID and Weight Versioning + +In Open-RL, the `sampling_session_id` acts as a versioned weight reference that pins sampling requests to a specific iteration of the policy weights. + +### How it Works: +1. **Weight Generation & Versioning**: + - The trainer calls `trainer.save_weights_for_sampler(name=alias)`. + - The request runs through the queue to ensure prior training steps are complete. + - The trainer writes the model weights to a versioned folder and registers a tinker URI: `tinker:///sampler_weights/` (where `` contains a sequence number or timestamp). +2. **Session Creation**: + - The client calls `create_sampling_client(weights_path)` passing the versioned `tinker://` URI. + - The gateway's `/api/v1/create_sampling_session` validates the model path, blocks until the model's dynamic sampler worker registers as ready in Redis, and returns the URI as the `sampling_session_id`. +3. **Session Pinning during Sampling**: + - When the client requests generation, it calls `sample_async(prompt)` on the returned client, which sends a request to `/api/v1/asample` with the pinned `sampling_session_id`. + - The gateway extracts the relative path portion of the `tinker://` URI and resolves it to the absolute local directory: `/tmp/open-rl/sampler_full//sampler_weights/`. + - It packages this absolute path into the `weights_path` field of the sampling request payload and enqueues it to `open_rl:sampler_queue:`. + - The sampler worker pops the request, compares the target `weights_path` to its currently loaded weights directory, and triggers a sleep-reload cycle if a change is detected. This ensures that the generated tokens are always sampled from the exact version of the policy weights corresponding to that session. + +--- + +## 9. Control Plane & Data Plane Decoupling (Lifecycle Management) + +The framework segregates operations into **Control Plane** (infrastructure, worker lifecycles, configuration) and **Data Plane** (training metrics, token sampling) to enable asynchronous execution, fail-safe scaling, and zero-drop request drainage. + +``` +[Control Plane Operations] +Gateway Control Queue (open_rl:worker_launch_queue) ---> WorkerLaunchProcessor (launches/gracefully stops PIDs) + +[Data Plane Operations] +Gateway Data Queues (queue:, sampler_queue:) ---> Workers (runs steps/computes tokens) +``` + +### 1. Queue Division & Protocols + +| Plane | Operation | Protocol / Redis Key | Consumer | +| :--- | :--- | :--- | :--- | +| **Control Plane** | `create_model` (Launch trainer) | `open_rl:worker_launch_queue` (Central) | `WorkerLaunchProcessor` | +| **Control Plane** | `launch_sampler` (Launch sampler) | `open_rl:worker_launch_queue` (Central) | `WorkerLaunchProcessor` | +| **Control Plane** | `delete_model` (Stop workers) | `open_rl:worker_launch_queue` (Central) | `WorkerLaunchProcessor` | +| **Control Plane** | `create_sampling_session` | Registry Metadata Key (Redis) | Gateway | +| **Data Plane** | `forward_backward`, `optim_step` | `open_rl:queue:` (Isolated) | PyTorch Trainer | +| **Data Plane** | `save_weights_for_sampler`, `save_state` | `open_rl:queue:` (Isolated) | PyTorch Trainer | +| **Data Plane** | `sample`, `asample` | `open_rl:sampler_queue:` (Isolated) | vLLM Sampler | + +--- + +### 2. Provisioning & Activation (Control Plane) +- **Trainer Provisioning**: When a client initializes a model via `/api/v1/create_model`, the gateway enqueues a `create_model` command to `open_rl:worker_launch_queue`. The `WorkerLaunchProcessor` daemon pops the request and invokes `FFTWorkerManager.launch_trainer(model_id)` to spawn the dedicated PyTorch trainer subprocess. +- **Sampler Provisioning**: When a client initializes a sampling session via `/api/v1/create_sampling_session` (or saves weights for sampling via `/api/v1/save_weights_for_sampler`), the gateway enqueues a `launch_sampler` command to `open_rl:worker_launch_queue`. The processor pops it and invokes `FFTWorkerManager.launch_sampler(model_id)` to spawn the dedicated vLLM sampler worker (if not already running). +- **Readiness**: The sampler worker compiles CUDA graphs and writes `open_rl:sampler_ready: = "1"`. The gateway blocks and polls this key, returning the versioned `session_id` to the client only when ready. + +--- + +### 3. Graceful Queue-Based Teardown (Implemented) +To prevent dropping in-flight sampling or optimization steps, de-provisioning is fully queue-decoupled using the **Sentinel Pattern**: + +1. **Teardown Trigger**: When a client completes training, it sends a `POST` request to the Gateway's `/api/v1/delete_model` endpoint. +2. **Sentinel Enqueueing**: Instead of hard-killing processes, the Gateway writes a `shutdown_workers` control command to `open_rl:worker_launch_queue`. +3. **Control-to-Data Signaling**: The `WorkerLaunchProcessor` pops the command and enqueues a **Shutdown Sentinel (Poison Pill)** (`{"request_id": "SHUTDOWN_SENTINEL"}`) to the tails of both the trainer (`open_rl:queue:`) and sampler (`open_rl:sampler_queue:`) FIFO data queues. +4. **FIFO Request Drainage**: + - The workers continue to pop and process all pending requests in their queues. + - When a worker pops the sentinel, it halts polling, awaits any active asyncio tasks (token generation or backward steps), unregisters its PID from the Snapshot Agent, and exits gracefully. +5. **Background Process Reaping**: + - `FFTWorkerManager` runs a non-blocking background task to monitor the spawned PIDs. + - If they exit cleanly, it clears them from the registry. If they fail to exit within a grace period (e.g. 30 seconds), it falls back to a hard `terminate()` to reclaim resources. +6. **Snapshot Resiliency**: + - The Snapshot Agent (`serve.py`) checks PID liveness before/during preemption tasks. Dead or zombie processes are skipped gracefully, and socket closure automatically cleans up daemon state. + +--- + +### 4. Proposed De-provisioning (Idle Timeout) +To clean up stale workers in the event of hard client VM crashes or unhandled script kills: + +- **Server-Side Idle Timeout**: + - The `WorkerLaunchProcessor` daemon monitors active tenant activity. + - If a tenant's data queue has been idle (no requests popped/enqueued) for a configurable timeout (e.g. 10 minutes), the processor automatically triggers the sentinel shutdown flow for that `model_id`. + - If the client subsequently resumes training, the processor detects the missing workers and dynamically re-provisions them transparently. + - This ensures maximum robustness against crash-induced leaks without requiring client-side handlers. + +--- + +## 10. Snapshot Agent Integration & Lock Management + +To coordinate GPU sharing, workers communicate with local Snapshot Agent instances via UNIX domain sockets: +- `snapshot-agent-trainer.sock` (manages GPU 0 for trainers) +- `snapshot-agent-sampler.sock` (manages GPU 1 for samplers) + +### 1. Lock Management & API Command Set +The Snapshot Agent daemon supports the following JSON-based socket interface commands: +- `REGISTER`: Registers a process PID to participate in scheduling. +- `UNREGISTER`: Removes a process PID and cleans up its lock allocations. +- `ACQUIRE`: Requests the global preemption lock for a PID. Blocks until the lock is acquired, and automatically suspends the active running process. +- `RELEASE`: Signals that a process is yielding its GPU lock, triggering an immediate checkpoint snapshot. +- `TRANSFER_LOCK`: Moves the ownership of an active lock from one PID to another atomically. + +--- + +### 2. Architectural Requirement for the `TRANSFER_LOCK` Primitive + +During full fine-tuning rollouts, the vLLM sampler worker must serialize its engine startup to prevent CUDA Out of Memory (OOM) errors during CUDA graph warmups (which consume up to 70% VRAM). This is achieved through the coordinated parent-to-child lock transfer sequence: + +1. **Process Tree & CUDA Context Separation**: + - vLLM splits the sampler into a Python **parent process** (managing queues/IPC) and a dynamically spawned **child process** (`EngineCore` / `ModelExecutor`). + - The actual CUDA driver context and VRAM footprint reside entirely inside the child process. + - The utility `cuda-checkpoint` operates strictly on a **single target PID** (using direct POSIX real-time signals) and is not aware of the process tree. Thus, checkpointing/restoring must target the child process directly to reclaim/restore VRAM. + +2. **The Startup OOM Race Condition**: + - Because the child process PID is dynamically allocated during engine creation, the worker cannot know its PID beforehand. + - To serialize the initialization phase and prevent multiple samplers from compiling graphs simultaneously (which causes a CUDA OOM), the parent process must proxy the lock: + - The **parent** acquires the Snapshot Agent lock. + - The parent initializes the vLLM engine, which spawns the child process. + - The parent registers the child PID with the Snapshot Agent. + +3. **Why We Must "Transfer" Instead of "Release"**: + - If the parent released its lock using standard context managers immediately after initialization, a waiting concurrent worker would instantly acquire the lock and begin its graph warmups. + - At this moment, the first worker's child process is still running in memory because the Snapshot Agent has not yet checkpointed it (checkpointing is triggered asynchronously on release or preemption). + - This leads to two active engines compiling graphs at the same time, causing a CUDA OOM. + - **The Solution**: The parent calls **`TRANSFER_LOCK`** to transfer ownership of the active lock to the child process PID *before* exiting its code block. This keeps the GPU lock continuously active, blocking other workers until the child process is explicitly checkpointed (`RELEASE` command) to release its VRAM down to 0 MiB. + diff --git a/scripts/run_training_e2e.py b/scripts/run_training_e2e.py index 6e210102..015a0acc 100755 --- a/scripts/run_training_e2e.py +++ b/scripts/run_training_e2e.py @@ -36,6 +36,7 @@ import shutil import signal import socket +import struct import subprocess import threading import time @@ -53,7 +54,10 @@ @chz.chz class RunConfig: - scenario: Literal["tiny-lora", "tiny-fft", "tiny-rl", "lora-textsql", "fft-gsm8k", "fft-gsm8k-x2"] + scenario: Literal["tiny-lora", "tiny-fft", "tiny-rl", "tiny-fft-rl", "tiny-fft-rl-x2", "lora-textsql", "fft-gsm8k", "fft-gsm8k-x2"] + sampling_backend: str = "torch" + trainer_gpu: str = "0" + sampler_gpu: str = "1" base_url: str = "" base_model: str = "Qwen/Qwen2.5-0.5B" steps: int | None = None @@ -71,9 +75,6 @@ class RunConfig: startup_timeout: float = 300.0 train_token_budget: int = 65_536 vllm_gpu_memory_utilization: float = 0.70 - # Where OPEN_RL_TMP_DIR lives on the cluster's shared PVC; used to print the - # scripts/run_cluster_eval.py command when the checkpoint is not visible locally. - cluster_tmp_dir: str = "/mnt/shared/open-rl" @dataclass @@ -85,6 +86,7 @@ class ManagedProcess: def unused_tcp_port() -> int: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0)) sock.bind(("127.0.0.1", 0)) return int(sock.getsockname()[1]) @@ -133,21 +135,30 @@ def launch( timeout: float, ) -> None: print(f"[training-e2e] starting {name}: {' '.join(command)}") - with log_path.open("w", encoding="utf-8") as log_file: - process = subprocess.Popen( - command, - cwd=REPO_ROOT, - env=env, - stdout=log_file, - stderr=subprocess.STDOUT, - text=True, - start_new_session=True, - ) + process = subprocess.Popen( + command, + cwd=REPO_ROOT, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + start_new_session=True, + ) + + def stream_log(): + with log_path.open("w", encoding="utf-8") as log_file: + for line in iter(process.stdout.readline, ""): + log_file.write(line) + log_file.flush() + print(f"[{name}] {line.rstrip()}") + + t = threading.Thread(target=stream_log, daemon=True) + t.start() + processes.append(ManagedProcess(name=name, process=process, log_path=log_path)) try: wait_until(name, ready, timeout) except Exception: - print_log_tail(log_path) raise @@ -181,7 +192,7 @@ def base_env(config: RunConfig) -> dict[str, str]: "OPEN_RL_TMP_DIR": str(open_rl_tmp_dir(config)), "OPEN_RL_TRAIN_TOKEN_BUDGET": str(config.train_token_budget), "PYTHONUNBUFFERED": "1", - "SAMPLING_BACKEND": "torch", + "SAMPLING_BACKEND": config.sampling_backend, "TINKER_API_KEY": os.environ.get("TINKER_API_KEY", "tml-dummy-key"), "TOKENIZERS_PARALLELISM": "false", } @@ -196,6 +207,7 @@ def start_backend(config: RunConfig, processes: list[ManagedProcess]) -> str: port = config.port or unused_tcp_port() base_url = f"http://{config.host}:{port}" env = base_env(config) + env["CUDA_VISIBLE_DEVICES"] = config.trainer_gpu if "fft" in config.scenario: if shutil.which("redis-server") is None: @@ -233,6 +245,25 @@ def start_backend(config: RunConfig, processes: list[ManagedProcess]) -> str: env.pop("REDIS_URL", None) env.pop("OPEN_RL_ENABLE_FFT", None) + if env.get("SAMPLING_BACKEND") == "vllm": + if "fft" in config.scenario: + env["SAMPLER_CUDA_VISIBLE_DEVICES"] = config.sampler_gpu + env["VLLM_GPU_MEMORY_UTILIZATION"] = str(config.vllm_gpu_memory_utilization) + else: + vllm_env = env.copy() + vllm_env["CUDA_VISIBLE_DEVICES"] = config.sampler_gpu + vllm_env["VLLM_GPU_MEMORY_UTILIZATION"] = str(config.vllm_gpu_memory_utilization) + base_model = config.base_model + launch( + processes, + "vllm-worker", + uv_run(config.eval_uv_extra) + ["python", "-m", "server.vllm_sampler", "--model-id", base_model], + vllm_env, + log_dir / "vllm_worker.log", + lambda: True, + timeout=config.startup_timeout, + ) + launch( processes, "backend", @@ -303,15 +334,15 @@ def run_example(config: RunConfig, script: list[str], defaults: dict[str, str], def run_tiny(config: RunConfig, base_url: str, watch: list[ManagedProcess]) -> None: - script = "tiny_rl" if config.scenario == "tiny-rl" else "tiny_sft" + script = "tiny_rl" if "rl" in config.scenario else "tiny_sft" defaults = { "base_model": config.base_model, "base_url": base_url, "log_dir": str(Path(config.log_dir) / config.scenario.replace("-", "_")), } - if config.scenario == "tiny-fft": - # tiny_sft's 1e-3 default is tuned for LoRA adapters; full fine-tuning all - # params with Adam at that rate diverges (observed loss 0.93 -> 35). + if "fft" in config.scenario: + # tiny SFT/RL default is tuned for LoRA adapters; full fine-tuning all + # params with Adam at that rate diverges. defaults["learning_rate"] = "1e-5" if config.steps is not None: defaults["steps"] = str(config.steps) @@ -342,11 +373,11 @@ def write_gsm8k_eval_data(config: RunConfig) -> Path: return data_path -def resolve_eval_model_path(output: str, must_exist: bool = True) -> str: +def resolve_eval_model_path(output: str) -> str: for line in reversed(output.splitlines()): if line.startswith("eval_model_path="): path = line.removeprefix("eval_model_path=").strip() - if must_exist and not Path(path).exists(): + if not Path(path).exists(): raise RuntimeError(f"Eval model path does not exist: {path}") return path raise RuntimeError("GSM8K SFT finished without printing eval_model_path=...") @@ -385,52 +416,37 @@ def run_gsm8k_eval(config: RunConfig, model_path: str) -> None: ) -def cluster_model_path(model_path: str, cluster_tmp_dir: str) -> str: - """Re-root a shared OpenRL artifact path onto the cluster's OPEN_RL_TMP_DIR. - - gsm8k_sft prints eval_model_path using the *client's* OPEN_RL_TMP_DIR; against - a remote backend the checkpoint actually lives under the cluster's tmp dir on - the shared PVC, so swap the prefix while keeping the OpenRL artifact tail. - """ - for marker in ("/sampler_full/", "/checkpoints/"): - if marker in model_path: - return cluster_tmp_dir.rstrip("/") + marker + model_path.split(marker, 1)[1] - return model_path - - -def run_or_print_gsm8k_eval(config: RunConfig, model_path: str, label: str = "") -> None: - """Eval locally when the checkpoint is reachable; otherwise print the exact - command that runs the same eval as a job on the cluster (the checkpoint lives - on the cluster's PVC when the backend was remote).""" - if Path(model_path).exists(): - run_gsm8k_eval(config, model_path) - return - remote_path = cluster_model_path(model_path, config.cluster_tmp_dir) - tag = f" {label}" if label else "" - print(f"[training-e2e]{tag} checkpoint {model_path} is not visible locally (remote backend at {config.base_url}).") - print(f"[training-e2e]{tag} run the GSM8K eval on the cluster with:") - print(f" make cluster-eval EVAL_MODEL_PATH={shlex.quote(remote_path)} EVAL_EXAMPLES={config.eval_examples}") - - def run_gsm8k(config: RunConfig, base_url: str, watch: list[ManagedProcess]) -> None: output = run_gsm8k_train(config, base_url, watch, "fft_gsm8k") - run_or_print_gsm8k_eval(config, resolve_eval_model_path(output, must_exist=not config.base_url)) + run_gsm8k_eval(config, resolve_eval_model_path(output)) def check_snapshot_interleaving(config: RunConfig) -> None: - log_path = Path(config.log_dir) / "snapshot-agent.log" if config.base_url: print("[training-e2e] external backend; skipping snapshot agent interleave check") return + + log_path = Path(config.log_dir) / "snapshot-agent.log" + if not log_path.exists(): + return text = log_path.read_text(encoding="utf-8", errors="replace") - checkpointed = set(re.findall(r"checkpointed pid (\d+)", text)) - restored = set(re.findall(r"restored pid (\d+)", text)) - if len(checkpointed) < 2 or len(restored) < 2: + + cp_t = set(re.findall(r"checkpointed pid (\d+) \(group trainers\)", text)) + rs_t = set(re.findall(r"restored pid (\d+) \(group trainers\)", text)) + if len(cp_t) < 2 or len(rs_t) < 2: raise RuntimeError( - f"Expected both FFT workers to round-trip through the snapshot agent, " - f"but saw checkpointed pids {sorted(checkpointed)} and restored pids {sorted(restored)}; see {log_path}" + f"Expected both FFT trainer workers to interleave, but saw checkpoints {sorted(cp_t)} and restores {sorted(rs_t)} in {log_path}" ) - print(f"[training-e2e] snapshot agent time-sliced workers: checkpointed pids {sorted(checkpointed)}, restored pids {sorted(restored)}") + print(f"[training-e2e] trainer snapshot agent time-sliced: checkpointed pids {sorted(cp_t)}, restored pids {sorted(rs_t)}") + + if config.sampling_backend == "vllm": + cp_s = set(re.findall(r"checkpointed pid (\d+) \(group samplers\)", text)) + rs_s = set(re.findall(r"restored pid (\d+) \(group samplers\)", text)) + if len(cp_s) < 2 or len(rs_s) < 2: + raise RuntimeError( + f"Expected both FFT sampler workers to interleave, but saw checkpoints {sorted(cp_s)} and restores {sorted(rs_s)} in {log_path}" + ) + print(f"[training-e2e] sampler snapshot agent time-sliced: checkpointed pids {sorted(cp_s)}, restored pids {sorted(rs_s)}") def run_gsm8k_x2(config: RunConfig, base_url: str, watch: list[ManagedProcess]) -> None: @@ -458,7 +474,40 @@ def train(job: str) -> None: for job, result in sorted(results.items()): assert isinstance(result, str) print(f"[training-e2e] evaluating {job}") - run_or_print_gsm8k_eval(config, resolve_eval_model_path(result, must_exist=not config.base_url), label=job) + run_gsm8k_eval(config, resolve_eval_model_path(result)) + + +def run_tiny_fft_rl_x2(config: RunConfig, base_url: str, watch: list[ManagedProcess]) -> None: + """Two concurrent FFT RL jobs against the same backend: each create_model spawns + its own trainer and dedicated sampler worker, and the snapshot agent time-slices them.""" + results: dict[str, str | BaseException] = {} + + def train(job: str) -> None: + try: + script = "tiny_rl" + defaults = { + "base_model": config.base_model, + "base_url": base_url, + "log_dir": str(Path(config.log_dir) / f"{config.scenario.replace('-', '_')}_{job}"), + "learning_rate": "1e-5", + } + if config.steps is not None: + defaults["steps"] = str(config.steps) + results[job] = run_example(config, [f"examples/tiny/{script}.py"], defaults, watch=watch, prefix=f"[{job}] ") + except BaseException as exc: + results[job] = exc + + threads = [threading.Thread(target=train, args=(job,)) for job in ("job-a", "job-b")] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + for job, result in sorted(results.items()): + if isinstance(result, BaseException): + raise RuntimeError(f"tiny-fft-rl-x2 {job} failed") from result + + check_snapshot_interleaving(config) def read_jsonl(path: Path) -> list[dict]: @@ -522,6 +571,8 @@ def main() -> None: run_gsm8k_x2(config, base_url, processes) elif config.scenario == "lora-textsql": run_textsql(config, base_url, processes) + elif config.scenario == "tiny-fft-rl-x2": + run_tiny_fft_rl_x2(config, base_url, processes) else: run_tiny(config, base_url, processes) finally: diff --git a/src/server/training_requests_processor.py b/src/server/training_requests_processor.py index b1b1f2f1..ba837e5b 100644 --- a/src/server/training_requests_processor.py +++ b/src/server/training_requests_processor.py @@ -248,14 +248,29 @@ def __init__( self.worker = worker self.model_id = model_id self.pid = os.getpid() + self.group = os.getenv("OPEN_RL_TIMESLICE_GROUP", "trainers") self.snapshot_client = snapshot_client self.snapshot_registered = False + async def exit_gracefully(self) -> None: + print(f"[WORKER] Initiating immediate exit for model {self.model_id} trainer worker...") + if self.snapshot_registered: + try: + await self.snapshot_client.unregister(self.pid, group=self.group) + self.snapshot_registered = False + except Exception as exc: + print(f"[WORKER] Failed to unregister: {exc}") + try: + await self.snapshot_client.close() + except Exception: + pass + os._exit(0) + async def run(self) -> None: print("[WORKER] Full fine-tuning training requests processor started.") try: - await self.snapshot_client.register(self.pid) + await self.snapshot_client.register(self.pid, group=self.group) self.snapshot_registered = True while True: try: @@ -269,7 +284,7 @@ async def run(self) -> None: finally: try: if self.snapshot_registered: - await self.snapshot_client.unregister(self.pid) + await self.snapshot_client.unregister(self.pid, group=self.group) finally: await self.snapshot_client.close() @@ -279,14 +294,26 @@ async def run_once(self) -> None: await asyncio.sleep(0.1) return + has_shutdown = False + training_reqs = [] + for req in batch: + if req.get("request_id") == "SHUTDOWN_SENTINEL" or req.get("op") in {"shutdown", "shutdown_workers"}: + has_shutdown = True + else: + training_reqs.append(req) + with tracer.start_as_current_span("training_requests_batch") as batch_span: - batch_span.set_attribute("batch_size", len(batch)) + batch_span.set_attribute("batch_size", len(training_reqs)) batch_span.set_attribute("model_id", self.model_id) - print(f"\n[TRAINING REQUESTS] Popped {len(batch)} requests for model: {self.model_id}") - async with self.snapshot_client.acquire(self.pid): - for request in batch: - await self.process_request(request, self.model_id) + if training_reqs: + print(f"\n[TRAINING REQUESTS] Popped {len(training_reqs)} requests for model: {self.model_id}") + async with self.snapshot_client.acquire(self.pid, group=self.group): + for request in training_reqs: + await self.process_request(request, self.model_id) + + if has_shutdown: + await self.exit_gracefully() async def create_model(self, payload: dict[str, Any], model_id: str) -> dict[str, Any]: raw_config = payload.get("full_config") or {} diff --git a/src/server/vllm_sampler.py b/src/server/vllm_sampler.py index 0137f8fc..9de05fef 100644 --- a/src/server/vllm_sampler.py +++ b/src/server/vllm_sampler.py @@ -1,15 +1,13 @@ # This file contains the vLLM worker implementation for high-throughput inference in Open-RL. +import argparse import asyncio import hashlib import os import sys -from contextlib import asynccontextmanager +import traceback from typing import Any -import uvicorn -from fastapi import FastAPI, Request - try: from vllm import SamplingParams from vllm.engine.arg_utils import AsyncEngineArgs @@ -26,8 +24,7 @@ RequestOutputKind = None VLLM_AVAILABLE = False -from opentelemetry import trace -from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor +from opentelemetry import propagate, trace from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import BatchSpanProcessor @@ -47,14 +44,27 @@ tracer = trace.get_tracer("vllm.inference.worker") engine: Any = None +CURRENT_LOADED_SAMPLER_WEIGHTS: str | None = None +reload_lock = asyncio.Lock() + + +def is_fft_enabled() -> bool: + return os.getenv("OPEN_RL_ENABLE_FFT", "").lower() == "true" + +snapshot_client: Any = None +socket_path = os.getenv("OPEN_RL_SNAPSHOT_AGENT_SOCKET") +if is_fft_enabled() and socket_path: + from snapshot_agent.client import SnapshotAgentClient -@asynccontextmanager -async def lifespan(app: FastAPI): + snapshot_client = SnapshotAgentClient(socket_path) + + +def init_engine(): global engine print("\n" + "=" * 50) - print(" Open-RL vLLM Inference Engine") + print(" Open-RL vLLM Inference Engine (Queue Mode)") print("=" * 50) cuda_devs = os.getenv("CUDA_VISIBLE_DEVICES", "ALL") model_name = os.getenv("BASE_MODEL") or os.getenv("VLLM_MODEL") @@ -63,9 +73,9 @@ async def lifespan(app: FastAPI): mock_vllm = os.getenv("MOCK_VLLM", "0") == "1" if mock_vllm or not VLLM_AVAILABLE: - print("[vLLM Subprocess] MOCK_VLLM=1 or vllm not installed, bypassing real engine init for local dev.") + print("[vLLM Worker] MOCK_VLLM=1 or vllm not installed, bypassing real engine init for local dev.") elif not model_name: - print("[vLLM Subprocess] Error: BASE_MODEL environment variable is required.") + print("[vLLM Worker] Error: BASE_MODEL environment variable is required.") sys.exit(1) else: hf_overrides: dict = {} @@ -75,53 +85,40 @@ async def lifespan(app: FastAPI): engine_kwargs = { "model": model_name, - "enable_lora": True, - "max_loras": 8, - "max_lora_rank": 64, + "enable_sleep_mode": is_fft_enabled(), + "enable_lora": not is_fft_enabled(), "max_model_len": int(os.getenv("VLLM_MAX_MODEL_LEN", "8192")), "max_num_seqs": int(os.getenv("VLLM_MAX_NUM_SEQS", "64")), "gpu_memory_utilization": float(os.getenv("VLLM_GPU_MEMORY_UTILIZATION", "0.90")), "enable_prefix_caching": False, "enforce_eager": os.getenv("VLLM_ENFORCE_EAGER", "0") == "1", } + if not is_fft_enabled(): + engine_kwargs["max_loras"] = 8 + engine_kwargs["max_lora_rank"] = 64 if hf_overrides: engine_kwargs["hf_overrides"] = hf_overrides engine_args = AsyncEngineArgs(**engine_kwargs) engine = AsyncLLMEngine.from_engine_args(engine_args) - print("[vLLM Subprocess] Engine initialized and ready to serve IPC requests.") - - yield - - -app = FastAPI(title="Open-RL vLLM Subprocess", lifespan=lifespan) -FastAPIInstrumentor.instrument_app(app, excluded_urls="/healthz") - - -@app.get("/healthz") -async def healthz(): - return {"status": "ok", "mock": engine is None} - - -@app.post("/generate") -async def generate(req: Request): + print("[vLLM Worker] Engine initialized successfully.") + + +async def run_generation_backend( + request_id: str, + prompt_token_ids: list[int], + max_tokens: int, + temperature: float, + stop: list[int] | None, + top_p: float, + top_k: int, + num_samples: int, + lora_id: str | None, + lora_path: str | None, + include_prompt_logprobs: bool, +) -> dict[str, Any]: try: - data = await req.json() - - request_id = data.get("request_id") - prompt_token_ids = data.get("prompt_token_ids") - max_tokens = data.get("max_tokens", 20) - temperature = data.get("temperature", 1.0) - stop = data.get("stop", None) - top_p = data.get("top_p", 1.0) - top_k = data.get("top_k", -1) - num_samples = data.get("num_samples", 1) - - lora_id = data.get("lora_id", None) - lora_path = data.get("lora_path", None) - include_prompt_logprobs = data.get("include_prompt_logprobs", False) - current_engine = engine if current_engine is None: # Mocking for local Mac dev @@ -191,14 +188,199 @@ async def generate(req: Request): else: prompt_logprobs_out.append(None) - return {"sequences": sequences_out, "prompt_logprobs": prompt_logprobs_out} + res = {"sequences": sequences_out} + if prompt_logprobs_out is not None: + res["prompt_logprobs"] = prompt_logprobs_out + return res except Exception as e: - import traceback - traceback.print_exc() - # Return explicit 500 so upstream client logs it return {"type": "RequestFailedResponse", "error_message": f"vLLM Worker Error: {str(e)}"} +async def process_sampling_request(req: dict, store: Any) -> None: + global engine + global CURRENT_LOADED_SAMPLER_WEIGHTS + + request_id = req["request_id"] + trace_context = req.get("trace_context", {}) + + parent_span = propagate.extract(trace_context) + with tracer.start_as_current_span("process_sampling_request", context=parent_span): + try: + # 1. Manage weights reloading + weights_path = req.get("weights_path") + if is_fft_enabled() and weights_path: + async with reload_lock: + if weights_path != CURRENT_LOADED_SAMPLER_WEIGHTS: + print(f"[vLLM Worker] Weight change detected. Current: {CURRENT_LOADED_SAMPLER_WEIGHTS}, Target: {weights_path}") + if engine is not None: + print("[vLLM Worker] Triggering sleep level 2...") + await engine.sleep(level=2) + print("[vLLM Worker] Waking up weights...") + await engine.wake_up(tags=["weights"]) + print(f"[vLLM Worker] Reloading weights from {weights_path} in-place...") + await engine.collective_rpc("reload_weights", kwargs={"weights_path": weights_path}) + print("[vLLM Worker] Waking up KV cache...") + await engine.wake_up(tags=["kv_cache"]) + CURRENT_LOADED_SAMPLER_WEIGHTS = weights_path + print("[vLLM Worker] Weights reload completed successfully!") + + # 2. Run inference + prompt_token_ids = req.get("prompt_token_ids", []) + max_tokens = req.get("max_tokens", 20) + temperature = req.get("temperature", 1.0) + stop = req.get("stop") + top_p = req.get("top_p", 1.0) + top_k = req.get("top_k", -1) + num_samples = req.get("num_samples", 1) + lora_id = req.get("lora_id") + lora_path = req.get("lora_path") + include_prompt_logprobs = req.get("include_prompt_logprobs", False) + + result = await run_generation_backend( + request_id=request_id, + prompt_token_ids=prompt_token_ids, + max_tokens=max_tokens, + temperature=temperature, + stop=stop, + top_p=top_p, + top_k=top_k, + num_samples=num_samples, + lora_id=lora_id, + lora_path=lora_path, + include_prompt_logprobs=include_prompt_logprobs, + ) + + if result.get("type") != "RequestFailedResponse": + result["type"] = "sample" + + await store.set_future(request_id, result) + except Exception as exc: + traceback.print_exc() + await store.set_future(request_id, {"type": "RequestFailedResponse", "error_message": f"vLLM Worker Error: {str(exc)}"}) + + +async def run_sampling_worker(model_id: str) -> None: + global engine + global CURRENT_LOADED_SAMPLER_WEIGHTS + from server.store import get_store + + store = get_store() + snapshot_registered = False + worker_pid = os.getpid() + group = os.getenv("OPEN_RL_TIMESLICE_GROUP", "samplers") + + if snapshot_client is not None: + try: + print(f"[vLLM Worker] Registering parent PID {worker_pid} (group {group}) for initialization lock...") + await snapshot_client.register(worker_pid, group=group) + snapshot_registered = True + async with snapshot_client.acquire(worker_pid, group=group): + print("[vLLM Worker] Initializing vLLM engine under parent lock...") + init_engine() + print("[vLLM Worker] Engine initialized successfully.") + except Exception as exc: + print(f"[vLLM Worker] Failed to perform coordinated initialization: {exc}") + traceback.print_exc() + if engine is None: + init_engine() + else: + init_engine() + + if snapshot_client is not None: + import signal + + async def exit_gracefully() -> None: + print(f"[vLLM Worker] Initiating immediate exit for model {model_id} sampler worker...") + nonlocal snapshot_registered + if snapshot_registered: + try: + await snapshot_client.unregister(worker_pid) + snapshot_registered = False + except Exception as exc: + print(f"[vLLM Worker] Failed to unregister: {exc}") + try: + await snapshot_client.close() + except Exception: + pass + os._exit(0) + + async def handle_shutdown(): + print(f"[vLLM Worker] Received termination signal, shutting down model {model_id} sampler worker...") + await exit_gracefully() + + try: + loop = asyncio.get_running_loop() + for sig in (signal.SIGTERM, signal.SIGINT): + loop.add_signal_handler(sig, lambda: asyncio.create_task(handle_shutdown())) + except NotImplementedError: + pass + + if hasattr(store, "redis"): + await store.redis.set(f"open_rl:sampler_ready:{model_id}", "1") + await store.redis.expire(f"open_rl:sampler_ready:{model_id}", 3600) + + print(f"[vLLM Worker] Listening for sampling requests on queue for model: {model_id}...") + try: + while True: + try: + batch = await store.get_sampling_requests_for_model(model_id) + if not batch: + await asyncio.sleep(0.05) + continue + + has_shutdown = False + sampling_reqs = [] + for req in batch: + if req.get("request_id") == "SHUTDOWN_SENTINEL": + has_shutdown = True + else: + sampling_reqs.append(req) + + if sampling_reqs: + if snapshot_client is not None: + async with snapshot_client.acquire(worker_pid): + tasks = [asyncio.create_task(process_sampling_request(req, store)) for req in sampling_reqs] + await asyncio.gather(*tasks) + if has_shutdown: + await exit_gracefully() + if engine is not None: + print("[vLLM Worker] Exiting batch: sleeping engine to yield GPU memory...") + await engine.sleep(level=2) + CURRENT_LOADED_SAMPLER_WEIGHTS = None + else: + tasks = [asyncio.create_task(process_sampling_request(req, store)) for req in sampling_reqs] + await asyncio.gather(*tasks) + + if has_shutdown: + print("[vLLM Worker] Shutdown sentinel popped from queue. Initiating clean exit...") + await exit_gracefully() + except asyncio.CancelledError: + break + except Exception as exc: + print(f"Error in sampling worker loop: {exc}") + traceback.print_exc() + await asyncio.sleep(1) + finally: + if snapshot_client is not None: + try: + if snapshot_registered: + await snapshot_client.unregister(worker_pid) + finally: + await snapshot_client.close() + os._exit(0) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Open-RL vLLM Pull-Mode Sampler Worker") + parser.add_argument("--model-id", type=str, required=True, help="The model ID of the RL job to process requests for") + args = parser.parse_args() + + try: + asyncio.run(run_sampling_worker(args.model_id)) + except KeyboardInterrupt: + print("[vLLM Worker] Exiting via KeyboardInterrupt.") + + if __name__ == "__main__": - uvicorn.run(app, host="0.0.0.0", port=8001) + main() From b5241cf2ce22af1c4c181a9b41c7c1966c479d6a Mon Sep 17 00:00:00 2001 From: droot Date: Wed, 17 Jun 2026 14:43:06 -0700 Subject: [PATCH 05/11] feat(k8s): add vLLM sampler worker node pool and dual pod template manifests --- .../distributed-fft-timeslice/04-gateway.yaml | 24 ++++--- .../07-snapshot-agent-daemonset.yaml | 12 +++- .../08-sampler-resourceclaim.yaml | 13 ++++ .../09-sampler-pod-template.yaml | 67 +++++++++++++++++++ .../kustomization.yaml | 2 + src/server/k8s_worker_manager.py | 23 +++++-- 6 files changed, 124 insertions(+), 17 deletions(-) create mode 100644 k8s/deploy/distributed-fft-timeslice/08-sampler-resourceclaim.yaml create mode 100644 k8s/deploy/distributed-fft-timeslice/09-sampler-pod-template.yaml diff --git a/k8s/deploy/distributed-fft-timeslice/04-gateway.yaml b/k8s/deploy/distributed-fft-timeslice/04-gateway.yaml index 3481013c..f4e7bf8b 100644 --- a/k8s/deploy/distributed-fft-timeslice/04-gateway.yaml +++ b/k8s/deploy/distributed-fft-timeslice/04-gateway.yaml @@ -53,9 +53,9 @@ spec: key: BASE_MODEL - name: OPEN_RL_TMP_DIR value: "/mnt/shared/open-rl" - # Trainer workers sample in-process, so this variant runs no vLLM worker. + # vLLM sampling enabled on cluster. - name: SAMPLING_BACKEND - value: "torch" + value: "vllm" - name: OPEN_RL_ENABLE_FFT value: "true" - name: OPEN_RL_WORKER_MANAGER @@ -64,10 +64,12 @@ spec: valueFrom: fieldRef: fieldPath: metadata.namespace + - name: OPEN_RL_TRAINER_POD_TEMPLATE + value: "/etc/open-rl/trainer/trainer-worker-pod.yaml" + - name: OPEN_RL_SAMPLER_POD_TEMPLATE + value: "/etc/open-rl/sampler/sampler-worker-pod.yaml" - name: OPEN_RL_WORKER_POD_TEMPLATE - value: "/etc/open-rl/trainer-worker-pod.yaml" - # Logical OpenRL snapshot group. It intentionally matches the - # timeslice.io/group label convention stamped onto trainer worker pods. + value: "/etc/open-rl/trainer/trainer-worker-pod.yaml" - name: OPEN_RL_TIME_SLICE_GROUP value: "trainers" resources: @@ -87,13 +89,19 @@ spec: volumeMounts: - name: shared-storage mountPath: /mnt/shared - - name: trainer-worker-pod-template - mountPath: /etc/open-rl + - name: trainer-tmpl + mountPath: /etc/open-rl/trainer + readOnly: true + - name: sampler-tmpl + mountPath: /etc/open-rl/sampler readOnly: true volumes: - name: shared-storage persistentVolumeClaim: claimName: open-rl-shared-pvc - - name: trainer-worker-pod-template + - name: trainer-tmpl configMap: name: open-rl-trainer-worker-pod-template + - name: sampler-tmpl + configMap: + name: open-rl-sampler-worker-pod-template diff --git a/k8s/deploy/distributed-fft-timeslice/07-snapshot-agent-daemonset.yaml b/k8s/deploy/distributed-fft-timeslice/07-snapshot-agent-daemonset.yaml index a70a042a..4bc9d37e 100644 --- a/k8s/deploy/distributed-fft-timeslice/07-snapshot-agent-daemonset.yaml +++ b/k8s/deploy/distributed-fft-timeslice/07-snapshot-agent-daemonset.yaml @@ -19,8 +19,16 @@ spec: hostPID: true hostNetwork: true dnsPolicy: ClusterFirstWithHostNet - nodeSelector: - group.timeslice.io/trainers: "true" + affinity: + nodeAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + - key: group.timeslice.io/trainers + operator: Exists + - matchExpressions: + - key: group.timeslice.io/samplers + operator: Exists tolerations: - key: "nvidia.com/gpu" operator: "Exists" diff --git a/k8s/deploy/distributed-fft-timeslice/08-sampler-resourceclaim.yaml b/k8s/deploy/distributed-fft-timeslice/08-sampler-resourceclaim.yaml new file mode 100644 index 00000000..8c21a360 --- /dev/null +++ b/k8s/deploy/distributed-fft-timeslice/08-sampler-resourceclaim.yaml @@ -0,0 +1,13 @@ +# One DRA ResourceClaim for the sampler GPU allocation. Every sampler worker pod +# that references this claim is scheduled onto a node that can access the same +# allocated physical GPU. +apiVersion: resource.k8s.io/v1 +kind: ResourceClaim +metadata: + name: open-rl-sampler-gpu-1 +spec: + devices: + requests: + - name: gpu + exactly: + deviceClassName: gpu.nvidia.com diff --git a/k8s/deploy/distributed-fft-timeslice/09-sampler-pod-template.yaml b/k8s/deploy/distributed-fft-timeslice/09-sampler-pod-template.yaml new file mode 100644 index 00000000..5f420969 --- /dev/null +++ b/k8s/deploy/distributed-fft-timeslice/09-sampler-pod-template.yaml @@ -0,0 +1,67 @@ +# Pod template the gateway uses to launch one sampler worker per model +# (server/k8s_worker_manager.py). All sampler worker pods reference the shared +# DRA ResourceClaim open-rl-sampler-gpu-1 (08-sampler-resourceclaim.yaml) and are +# placed onto the sampler GPU node. The node-local snapshot agent time-slices +# GPU access across concurrent inference engines. +apiVersion: v1 +kind: ConfigMap +metadata: + name: open-rl-sampler-worker-pod-template +data: + sampler-worker-pod.yaml: | + apiVersion: v1 + kind: Pod + spec: + restartPolicy: OnFailure + hostPID: true + containers: + - name: sampler-worker + image: ghcr.io/gke-labs/open-rl/server:latest + imagePullPolicy: Always + command: ["uv", "run", "python", "-u", "-m", "server.vllm_sampler"] + env: + - name: REDIS_URL + value: "redis://redis-service:6379" + - name: BASE_MODEL + valueFrom: + configMapKeyRef: + name: open-rl-config + key: BASE_MODEL + - name: OPEN_RL_ENABLE_FFT + value: "true" + - name: OPEN_RL_TMP_DIR + value: "/mnt/shared/open-rl" + - name: HF_HOME + value: "/mnt/shared/open-rl/huggingface" + - name: OPEN_RL_SNAPSHOT_AGENT_HOST + valueFrom: + fieldRef: + fieldPath: status.hostIP + - name: OPEN_RL_SNAPSHOT_AGENT_PORT + value: "9753" + - name: OPEN_RL_TIMESLICE_GROUP + value: "samplers" + resources: + claims: + - name: sampler-gpu + limits: + memory: "32Gi" + requests: + memory: "32Gi" + cpu: "4" + volumeMounts: + - name: shared-storage + mountPath: /mnt/shared + resourceClaims: + - name: sampler-gpu + resourceClaimName: open-rl-sampler-gpu-1 + volumes: + - name: shared-storage + persistentVolumeClaim: + claimName: open-rl-shared-pvc + nodeSelector: + group.timeslice.io/samplers: "true" + tolerations: + - key: "nvidia.com/gpu" + operator: "Exists" + effect: "NoSchedule" diff --git a/k8s/deploy/distributed-fft-timeslice/kustomization.yaml b/k8s/deploy/distributed-fft-timeslice/kustomization.yaml index 41719454..c4e77af3 100644 --- a/k8s/deploy/distributed-fft-timeslice/kustomization.yaml +++ b/k8s/deploy/distributed-fft-timeslice/kustomization.yaml @@ -12,6 +12,8 @@ resources: - 05-worker-pod-template.yaml - 06-gpu-resourceclaim.yaml - 07-snapshot-agent-daemonset.yaml + - 08-sampler-resourceclaim.yaml + - 09-sampler-pod-template.yaml configMapGenerator: - name: open-rl-config diff --git a/src/server/k8s_worker_manager.py b/src/server/k8s_worker_manager.py index 47eb282d..df965f9c 100644 --- a/src/server/k8s_worker_manager.py +++ b/src/server/k8s_worker_manager.py @@ -41,11 +41,17 @@ def __init__(self, core_api: Any = None): if not os.getenv("REDIS_URL"): raise RuntimeError("OPEN_RL_ENABLE_FFT=true requires REDIS_URL so launched workers can share queues and futures") - template_path = os.getenv("OPEN_RL_WORKER_POD_TEMPLATE") - if not template_path: - raise RuntimeError("OPEN_RL_WORKER_MANAGER=kubernetes requires OPEN_RL_WORKER_POD_TEMPLATE pointing at the worker pod YAML") - with open(template_path, encoding="utf-8") as f: - self.pod_template: dict[str, Any] = yaml.safe_load(f) + trainer_path = os.getenv("OPEN_RL_TRAINER_POD_TEMPLATE") or os.getenv("OPEN_RL_WORKER_POD_TEMPLATE") + if not trainer_path: + raise RuntimeError("OPEN_RL_WORKER_MANAGER=kubernetes requires OPEN_RL_TRAINER_POD_TEMPLATE or OPEN_RL_WORKER_POD_TEMPLATE") + with open(trainer_path, encoding="utf-8") as f: + self.trainer_template: dict[str, Any] = yaml.safe_load(f) + + sampler_path = os.getenv("OPEN_RL_SAMPLER_POD_TEMPLATE") or trainer_path + with open(sampler_path, encoding="utf-8") as f: + self.sampler_template: dict[str, Any] = yaml.safe_load(f) + + self.pod_template = self.trainer_template self.namespace = os.getenv("OPEN_RL_WORKER_NAMESPACE", "default") self.group_id = os.getenv("OPEN_RL_TIME_SLICE_GROUP", "trainers") @@ -97,7 +103,8 @@ def shutdown_all(self) -> None: pass def render_pod(self, pod_name: str, model_id: str, job_id: str, role: str = "trainer") -> dict[str, Any]: - pod = copy.deepcopy(self.pod_template) + base_tmpl = self.trainer_template if role == "trainer" else self.sampler_template + pod = copy.deepcopy(base_tmpl) metadata = pod.setdefault("metadata", {}) metadata["name"] = pod_name app_label = "open-rl-trainer-worker" if role == "trainer" else "open-rl-sampler-worker" @@ -115,7 +122,9 @@ def render_pod(self, pod_name: str, model_id: str, job_id: str, role: str = "tra if role == "sampler": container["command"] = ["uv", "run", "python", "-u", "-m", "server.vllm_sampler"] container.setdefault("args", []).extend(["--model-id", model_id]) - container.setdefault("env", []).append({"name": "OPEN_RL_TIME_SLICE_JOB_ID", "value": job_id}) + env_list = container.setdefault("env", []) + env_list.append({"name": "OPEN_RL_TIME_SLICE_JOB_ID", "value": job_id}) + env_list.append({"name": "OPEN_RL_TIMESLICE_GROUP", "value": group_val}) return pod def read_pod(self, pod_name: str) -> Any | None: From 1bd07c49bfc4607a9e16785d286f77ccf7e8fd88 Mon Sep 17 00:00:00 2001 From: droot Date: Wed, 17 Jun 2026 22:49:07 -0700 Subject: [PATCH 06/11] docs(k8s): add step-by-step DRA experiment setup and smoke test guide --- k8s/experiments/README.md | 122 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 122 insertions(+) create mode 100644 k8s/experiments/README.md diff --git a/k8s/experiments/README.md b/k8s/experiments/README.md new file mode 100644 index 00000000..1ea0c61c --- /dev/null +++ b/k8s/experiments/README.md @@ -0,0 +1,122 @@ +# Kubernetes Dynamic Resource Allocation (DRA) Experiments + +This directory contains experimental smoke tests and verification manifests for validating Kubernetes Dynamic Resource Allocation (DRA) oversubscription. + +In particular, `dra-shared-gpu-smoke.yaml` verifies that two independent Pods can reference a single shared `ResourceClaim` (`shared-gpu`) to get colocated onto the exact same physical NVIDIA GPU allocation without serializing CUDA execution or relying on standard time-sharing device plugins. + +--- + +## Prerequisites & Architectural Notes + +- **GKE Standard Cluster (v1.35+)**: GKE Autopilot clusters are **not supported** for these experiments. Autopilot's Warden admission webhook rejects custom Pod node selectors (e.g. `group.timeslice.io/trainers`), and Google manages kubelet runtime binaries, preventing custom third-party DRA kubelet plugin deployment. +- **Helm v3**: Required to deploy the experimental NVIDIA DRA driver chart. +- **Quota**: Ensure your GCP target project has available GCE quota for L4 GPUs (`nvidia-l4`) in the selected zone. + +--- + +## Step-by-Step Cluster Setup & Execution + +### 1. Create Base GKE Standard Cluster +Create a minimal GKE Standard cluster with a small standard CPU node pool for system components: + +```bash +export PROJECT_ID="$(gcloud config get-value project)" +export REGION="us-central1" +export CLUSTER="open-rl-dra" + +gcloud container clusters create "${CLUSTER}" \ + --location="${REGION}" \ + --release-channel=regular \ + --machine-type=e2-standard-4 \ + --num-nodes=1 \ + --disk-size=100 +``` + +### 2. Connect `kubectl` & Add Dedicated DRA GPU Node Pool +Connect credentials to your local shell: +```bash +gcloud container clusters get-credentials "${CLUSTER}" --location="${REGION}" +``` + +Add a dedicated L4 GPU node pool. +- **Driver Disabled**: We explicitly pass `gpu-driver-version=disabled` to prevent GKE from installing the default device plugin, allowing the DRA driver to manage GPU device discovery. +- **Single Zone Target**: When attaching a GPU pool to a regional cluster, explicitly pass `--node-locations` targeting a specific zone with L4 capacity (e.g. `us-central1-b`). + +```bash +gcloud container node-pools create gpu-dra \ + --cluster="${CLUSTER}" \ + --location="${REGION}" \ + --node-locations="us-central1-b" \ + --machine-type=g2-standard-12 \ + --accelerator="type=nvidia-l4,count=1,gpu-driver-version=disabled" \ + --node-labels="group.timeslice.io/trainers=true,group.timeslice.io/samplers=true,gke-no-default-nvidia-gpu-device-plugin=true,nvidia.com/gpu.present=true" \ + --num-nodes=1 +``` + +### 3. Install NVIDIA Runtime Driver & DRA Kubelet Plugin +Install the preloaded Container-Optimized OS (COS) NVIDIA driver DaemonSet: + +```bash +kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/nvidia-driver-installer/cos/daemonset-preloaded-latest.yaml +``` + +Install the experimental NVIDIA DRA Driver Helm chart (`nvidia-dra-driver-gpu`). +> [!IMPORTANT] +> **Gotcha - GPU DeviceClass Guard**: By default, the Helm chart disables publishing `DeviceClass: gpu.nvidia.com`. You **must** pass `--set resources.gpus.enabled=true --set gpuResourcesEnabledOverride=true` to force the kubelet plugin to register physical GPU resource slices. + +```bash +helm repo add nvidia https://helm.ngc.nvidia.com/nvidia + +helm install nvidia-dra-driver-gpu nvidia/nvidia-dra-driver-gpu \ + --version="25.8.0" --create-namespace --namespace nvidia-dra-driver-gpu \ + --set nvidiaDriverRoot="/home/kubernetes/bin/nvidia/" \ + --set resources.gpus.enabled=true \ + --set gpuResourcesEnabledOverride=true +``` + +> [!TIP] +> **Gotcha - ResourceQuota Scope Limitations**: In tenant or sandbox GCP projects (e.g. Anthos/Config Controller test environments), standard `ResourceQuota` policies often forbid Pods from requesting `system-node-critical` priority classes. If the DRA kubelet plugin fails to spawn, strip the priority class request: +> ```bash +> kubectl patch ds -n nvidia-dra-driver-gpu nvidia-dra-driver-gpu-kubelet-plugin \ +> --type=json -p='[{"op": "remove", "path": "/spec/template/spec/priorityClassName"}]' +> ``` + +Verify that physical GPU resource slices and device classes are published: +```bash +kubectl get deviceclasses +kubectl get resourceslices +``` + +--- + +## 4. Run the DRA Smoke Test + +Apply the smoke test experiment: +```bash +kubectl apply -f k8s/experiments/dra-shared-gpu-smoke.yaml +``` + +### Verification Benchmarks (Pass Criteria) + +1. **Colocated Node Placement**: Verify both Pods reach `Running` on the exact same GPU node: + ```bash + kubectl get pods -n dra-smoke -o wide + ``` +2. **Hardware VRAM Sharing**: Verify both isolated container logs output the exact same GPU UUID: + ```bash + kubectl logs -n dra-smoke smoke-a + kubectl logs -n dra-smoke smoke-b + ``` +3. **ResourceClaim Consumer Binding**: Verify the shared claim reserves one single physical device (`gpu-0`) simultaneously for two distinct workload PIDs: + ```bash + kubectl describe resourceclaim -n dra-smoke shared-gpu + ``` + +--- + +## 5. Teardown & Cleanup + +```bash +kubectl delete ns dra-smoke +gcloud container clusters delete "${CLUSTER}" --location="${REGION}" --quiet +``` From c25ac9ca14400961c218d61d14cd9a335a52e444 Mon Sep 17 00:00:00 2001 From: droot Date: Thu, 18 Jun 2026 00:49:23 -0700 Subject: [PATCH 07/11] fix(k8s,sampler): resolve UnboundLocalError on shutdown sentinel and disable Kustomize ConfigMap hashing --- .../distributed-fft-timeslice/04-gateway.yaml | 2 ++ .../05-worker-pod-template.yaml | 2 +- .../09-sampler-pod-template.yaml | 2 +- .../kustomization.yaml | 11 ++++++++ src/server/k8s_worker_manager.py | 3 +++ src/server/vllm_sampler.py | 27 ++++++++++--------- 6 files changed, 32 insertions(+), 15 deletions(-) diff --git a/k8s/deploy/distributed-fft-timeslice/04-gateway.yaml b/k8s/deploy/distributed-fft-timeslice/04-gateway.yaml index f4e7bf8b..3af15b17 100644 --- a/k8s/deploy/distributed-fft-timeslice/04-gateway.yaml +++ b/k8s/deploy/distributed-fft-timeslice/04-gateway.yaml @@ -70,6 +70,8 @@ spec: value: "/etc/open-rl/sampler/sampler-worker-pod.yaml" - name: OPEN_RL_WORKER_POD_TEMPLATE value: "/etc/open-rl/trainer/trainer-worker-pod.yaml" + - name: OPEN_RL_WORKER_IMAGE + value: "gcr.io/cdrollouts-sunilarora/open-rl-server:latest" - name: OPEN_RL_TIME_SLICE_GROUP value: "trainers" resources: diff --git a/k8s/deploy/distributed-fft-timeslice/05-worker-pod-template.yaml b/k8s/deploy/distributed-fft-timeslice/05-worker-pod-template.yaml index 8c1fccd2..676f5c4b 100644 --- a/k8s/deploy/distributed-fft-timeslice/05-worker-pod-template.yaml +++ b/k8s/deploy/distributed-fft-timeslice/05-worker-pod-template.yaml @@ -26,7 +26,7 @@ data: hostPID: true containers: - name: trainer-worker - image: ghcr.io/gke-labs/open-rl/server:latest + image: gcr.io/cdrollouts-sunilarora/open-rl-server:latest imagePullPolicy: Always command: ["uv", "run", "python", "-m", "server.training_requests_processor"] env: diff --git a/k8s/deploy/distributed-fft-timeslice/09-sampler-pod-template.yaml b/k8s/deploy/distributed-fft-timeslice/09-sampler-pod-template.yaml index 5f420969..a422aa46 100644 --- a/k8s/deploy/distributed-fft-timeslice/09-sampler-pod-template.yaml +++ b/k8s/deploy/distributed-fft-timeslice/09-sampler-pod-template.yaml @@ -16,7 +16,7 @@ data: hostPID: true containers: - name: sampler-worker - image: ghcr.io/gke-labs/open-rl/server:latest + image: gcr.io/cdrollouts-sunilarora/open-rl-server:latest imagePullPolicy: Always command: ["uv", "run", "python", "-u", "-m", "server.vllm_sampler"] env: diff --git a/k8s/deploy/distributed-fft-timeslice/kustomization.yaml b/k8s/deploy/distributed-fft-timeslice/kustomization.yaml index c4e77af3..c839e7a8 100644 --- a/k8s/deploy/distributed-fft-timeslice/kustomization.yaml +++ b/k8s/deploy/distributed-fft-timeslice/kustomization.yaml @@ -15,8 +15,19 @@ resources: - 08-sampler-resourceclaim.yaml - 09-sampler-pod-template.yaml +generatorOptions: + disableNameSuffixHash: true + configMapGenerator: - name: open-rl-config literals: - BASE_MODEL="Qwen/Qwen2.5-0.5B" - ENABLE_GCP_TRACE="1" + +images: + - name: ghcr.io/gke-labs/open-rl/server + newName: gcr.io/cdrollouts-sunilarora/open-rl-server + newTag: latest + - name: ghcr.io/gke-labs/open-rl/gateway + newName: gcr.io/cdrollouts-sunilarora/open-rl-gateway + newTag: latest diff --git a/src/server/k8s_worker_manager.py b/src/server/k8s_worker_manager.py index df965f9c..a78f7d15 100644 --- a/src/server/k8s_worker_manager.py +++ b/src/server/k8s_worker_manager.py @@ -119,6 +119,9 @@ def render_pod(self, pod_name: str, model_id: str, job_id: str, role: str = "tra ) container = pod["spec"]["containers"][0] + worker_image = os.getenv("OPEN_RL_WORKER_IMAGE") + if worker_image: + container["image"] = worker_image if role == "sampler": container["command"] = ["uv", "run", "python", "-u", "-m", "server.vllm_sampler"] container.setdefault("args", []).extend(["--model-id", model_id]) diff --git a/src/server/vllm_sampler.py b/src/server/vllm_sampler.py index 9de05fef..7cdf5af9 100644 --- a/src/server/vllm_sampler.py +++ b/src/server/vllm_sampler.py @@ -287,23 +287,24 @@ async def run_sampling_worker(model_id: str) -> None: else: init_engine() - if snapshot_client is not None: - import signal - - async def exit_gracefully() -> None: - print(f"[vLLM Worker] Initiating immediate exit for model {model_id} sampler worker...") - nonlocal snapshot_registered - if snapshot_registered: - try: - await snapshot_client.unregister(worker_pid) - snapshot_registered = False - except Exception as exc: - print(f"[vLLM Worker] Failed to unregister: {exc}") + async def exit_gracefully() -> None: + print(f"[vLLM Worker] Initiating immediate exit for model {model_id} sampler worker...") + nonlocal snapshot_registered + if snapshot_registered and snapshot_client is not None: + try: + await snapshot_client.unregister(worker_pid) + snapshot_registered = False + except Exception as exc: + print(f"[vLLM Worker] Failed to unregister: {exc}") + if snapshot_client is not None: try: await snapshot_client.close() except Exception: pass - os._exit(0) + os._exit(0) + + if snapshot_client is not None: + import signal async def handle_shutdown(): print(f"[vLLM Worker] Received termination signal, shutting down model {model_id} sampler worker...") From e0c9181dba881cea092c84d567a7d46f28fa0f75 Mon Sep 17 00:00:00 2001 From: droot Date: Thu, 18 Jun 2026 08:01:50 -0700 Subject: [PATCH 08/11] tune(k8s): lower memory requests to 16Gi to support dual concurrent RL jobs on 48GiB nodes --- .../distributed-fft-timeslice/05-worker-pod-template.yaml | 4 ++-- .../distributed-fft-timeslice/09-sampler-pod-template.yaml | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/k8s/deploy/distributed-fft-timeslice/05-worker-pod-template.yaml b/k8s/deploy/distributed-fft-timeslice/05-worker-pod-template.yaml index 676f5c4b..287d5135 100644 --- a/k8s/deploy/distributed-fft-timeslice/05-worker-pod-template.yaml +++ b/k8s/deploy/distributed-fft-timeslice/05-worker-pod-template.yaml @@ -57,11 +57,11 @@ data: claims: - name: trainer-gpu limits: - memory: "24Gi" + memory: "20Gi" requests: # Sized generously for model load and optimizer state; increase this # if a future coordination layer checkpoints GPU state to host RAM. - memory: "24Gi" + memory: "16Gi" cpu: "2" volumeMounts: - name: shared-storage diff --git a/k8s/deploy/distributed-fft-timeslice/09-sampler-pod-template.yaml b/k8s/deploy/distributed-fft-timeslice/09-sampler-pod-template.yaml index a422aa46..81abd34f 100644 --- a/k8s/deploy/distributed-fft-timeslice/09-sampler-pod-template.yaml +++ b/k8s/deploy/distributed-fft-timeslice/09-sampler-pod-template.yaml @@ -45,9 +45,9 @@ data: claims: - name: sampler-gpu limits: - memory: "32Gi" + memory: "20Gi" requests: - memory: "32Gi" + memory: "16Gi" cpu: "4" volumeMounts: - name: shared-storage From 7171c5a955dc6b94a798ba1f7cb1d47850f99572 Mon Sep 17 00:00:00 2001 From: droot Date: Thu, 18 Jun 2026 08:32:18 -0700 Subject: [PATCH 09/11] docs(fft): add comprehensive walkthrough of Kubernetes pod placement and time-slicing architecture --- docs/fft/pod_placement_architecture.md | 132 +++++++++++++++++++++++++ 1 file changed, 132 insertions(+) create mode 100644 docs/fft/pod_placement_architecture.md diff --git a/docs/fft/pod_placement_architecture.md b/docs/fft/pod_placement_architecture.md new file mode 100644 index 00000000..6b6e3fd4 --- /dev/null +++ b/docs/fft/pod_placement_architecture.md @@ -0,0 +1,132 @@ +# Open-RL Kubernetes Pod Placement Architecture Walkthrough + +This document provides an end-to-end architectural walkthrough of how Open-RL schedules, places, and virtualizes multi-tenant Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL) workloads across a distributed Kubernetes infrastructure. + +--- + +## High-Level Topology + +Open-RL decouples policy gradient computation (Trainers) from rollout generation (Samplers). In a multi-tenant environment (such as concurrent `job-a` and `job-b` experiments), Gateway orchestrates worker pods across dedicated physical GPU machines while ensuring strict boundary isolation. + +```mermaid +graph TD + classDef gw fill:#1e293b,stroke:#3b82f6,stroke-width:2px,color:#fff; + classDef n1 fill:#0f172a,stroke:#10b981,stroke-width:2px,color:#fff; + classDef n2 fill:#0f172a,stroke:#8b5cf6,stroke-width:2px,color:#fff; + classDef podA fill:#064e3b,stroke:#34d399,stroke-width:1px,color:#fff; + classDef podB fill:#4c1d95,stroke:#a78bfa,stroke-width:1px,color:#fff; + + Client["RL Client SDK (e.g. tiny_rl.py)"] -->|POST /api/v1/create_model| GW["Open-RL Gateway Service"]:::gw + Client -->|POST /api/v1/create_sampling_client| GW + + subgraph Cluster ["GKE Regional Standard Cluster (open-rl-dra)"] + subgraph Node1 ["Physical Machine 1: dcbk (g2-standard-12)
DRA Group: trainers"] + SA1["Snapshot Agent DaemonSet (tcp://:9753)"]:::n1 + TrA["open-rl-trainer-job-a
RAM: 16GiB / Claim: trainer-gpu-1"]:::podA + TrB["open-rl-trainer-job-b
RAM: 16GiB / Claim: trainer-gpu-1"]:::podB + TrA <-->|CRIU Time-Slice| SA1 + TrB <-->|CRIU Time-Slice| SA1 + end + + subgraph Node2 ["Physical Machine 2: hzp3 (g2-standard-12)
DRA Group: samplers"] + SA2["Snapshot Agent DaemonSet (tcp://:9753)"]:::n2 + SmA["open-rl-sampler-job-a
RAM: 16GiB / Claim: sampler-gpu-1"]:::podA + SmB["open-rl-sampler-job-b
RAM: 16GiB / Claim: sampler-gpu-1"]:::podB + SmA <-->|vLLM Sleep VRAM Yield| SA2 + SmB <-->|vLLM Sleep VRAM Yield| SA2 + end + + NFS[("Managed GKE Filestore NFS (/mnt/shared)")] + TrA -->|Write Checkpoints| NFS + TrB -->|Write Checkpoints| NFS + NFS -->|In-Place Safetensor Reload| SmA + NFS -->|In-Place Safetensor Reload| SmB + end +``` + +--- + +## 1. Decoupled Dynamic Pod Rendering + +When a client initiates a training loop, Gateway intercepts the session requests inside `k8s_worker_manager.py`. Rather than using static Kubernetes Deployments, Gateway dynamically deepcopies role-specific Pod YAML templates mounted from ConfigMaps: + +- **Trainer Pod Template**: Loaded from ConfigMap `open-rl-config` defined in `05-worker-pod-template.yaml`. +- **Sampler Pod Template**: Loaded from ConfigMap `open-rl-sampler-worker-pod-template` defined in `09-sampler-pod-template.yaml`. + +Gateway injects unique runtime identifiers (`model_id`, `job_id`, and `OPEN_RL_WORKER_IMAGE` overrides) before submitting imperative `create_namespaced_pod` API calls. + +--- + +## 2. Dynamic Resource Allocation (DRA) Claim Sharing + +Standard Kubernetes device plugins (`resources.limits: nvidia.com/gpu: 1`) enforce exclusive physical GPU locks: once Pod A lands on a node, `kube-scheduler` rejects Pod B until Pod A terminates. + +Open-RL bypasses this limitation using **Kubernetes Dynamic Resource Allocation (DRA)** exact allocation claims: + +```yaml +# Inside 05-worker-pod-template.yaml (Trainer Spec) +spec: + resourceClaims: + - name: trainer-gpu + resourceClaimName: open-rl-trainer-gpu-1 +``` + +```yaml +# Inside 09-sampler-pod-template.yaml (Sampler Spec) +spec: + resourceClaims: + - name: sampler-gpu + resourceClaimName: open-rl-sampler-gpu-1 +``` + +### How Claim Co-Scheduling Works: +1. **First Tenant (`job-a`)**: When `open-rl-trainer-job-a` spawns, it binds singleton claim `open-rl-trainer-gpu-1` to Physical Machine 1 (`dcbk`). +2. **Second Tenant (`job-b`)**: When `open-rl-trainer-job-b` spawns seconds later, `kube-scheduler` inspects its `resourceClaimName`. Because `open-rl-trainer-gpu-1` is already allocated on `dcbk`, **Kubernetes co-schedules Job B directly onto `dcbk` alongside Job A**! + +--- + +## 3. Strict Role Segregation via `nodeSelector` + +> [!WARNING] +> Co-locating PyTorch AdamW optimizers and vLLM KV caches on the same physical GPU causes immediate CUDA out-of-memory crashes (`CUDA error: out of memory`). + +To prevent cross-role contamination, nodes in the `gpu-dra` node pool are tagged with explicit role labels: +- Machine 1 (`dcbk`): `group.timeslice.io/trainers="true"` +- Machine 2 (`hzp3`): `group.timeslice.io/samplers="true"` + +Pod specs enforce strict landing boundaries: +- **Trainers**: Enforce `nodeSelector: { group.timeslice.io/trainers: "true" }` and set `OPEN_RL_TIME_SLICE_GROUP=trainers`. +- **Samplers**: Enforce `nodeSelector: { group.timeslice.io/samplers: "true" }` and set `OPEN_RL_TIMESLICE_GROUP=samplers`. + +--- + +## 4. Host RAM Oversubscription Tuning + +A standard GKE `g2-standard-12` virtual machine provides **48 GiB** of system CPU RAM. + +When scheduling multiple tenant pods onto a single machine, `kube-scheduler` calculates memory feasibility based on `resources.requests.memory`. + +| Component / Tenant Pod | Requested CPU Memory | Cumulative Allocated RAM | Node Feasibility on `g2-standard-12` (48 GiB Total) | +| :--- | :---: | :---: | :---: | +| **System Overhead** *(DaemonSets, CSI, Calico)* | ~4 GiB | 4 GiB | Schedulable (44 GiB Remaining) | +| **Tenant 1 Trainer** (`job-a`) | 16 GiB | 20 GiB | Schedulable (28 GiB Remaining) | +| **Tenant 2 Trainer** (`job-b`) | 16 GiB | 36 GiB | **Schedulable (12 GiB Remaining)** $\checkmark$ | + +> [!TIP] +> Prior to tuning, templates requested `24Gi` and `32Gi` of CPU memory. Under those defaults, Tenant 1 allocated $24 + 4 = 28\text{ GiB}$, leaving only $20\text{ GiB}$ remaining. When Tenant 2 requested `24Gi`, Kubernetes rejected the pod with `FailedScheduling: Insufficient memory`. Lowering requests to `16Gi` unlocked true multi-tenant concurrency. + +--- + +## 5. Node-Local Time-Slicing Virtualization + +Once co-scheduled onto the same physical GPU, workloads are virtualized in-flight by the node-local DaemonSet defined in `07-snapshot-agent-daemonset.yaml`. + +### A. Trainer Virtualization (CRIU Process Swapping) +On the Trainer Node (`dcbk`), the Snapshot Agent intercepts PyTorch CUDA allocations over `tcp://status.hostIP:9753`. When Job A finishes its microbatch gradient calculation: +1. Snapshot Agent freezes Job A's Linux process via CRIU (`checkpointed pid 34715 in 1.99s`). +2. Snapshot Agent restores Job B's memory state into VRAM (`restored pid 34716 in 0.39s`). + +### B. Sampler Virtualization (vLLM Cooperative Sleep) +On the Sampler Node (`hzp3`), vLLM inference engines time-slice cooperatively inside `vllm_sampler.py`: +1. **Sleep Preemption**: Upon completing a sampling batch, vLLM invokes `await engine.sleep(level=2)`, instantly discarding physical GPU memory pages (`freed 19.45 GiB`). +2. **NFS Weight Synchronization**: When Trainer A writes new SFT weights to `/mnt/shared`, Sampler A detects the modification, wakes up physical VRAM (`0.04s`), and reloads the checkpoint safetensors in-place directly from NFS page cache (`took 1.19 seconds`)! From 3badd1a276089a4724aba74e02b89dd126de9af4 Mon Sep 17 00:00:00 2001 From: droot Date: Thu, 18 Jun 2026 11:41:58 -0700 Subject: [PATCH 10/11] chore: ignore docs/scratch directory in .gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 3c81cd1e..e6bb6e3f 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,7 @@ env/ # Training Logs & Plots scratch/ +docs/scratch/ *.txt *.png *.log From abe45e4977466b665c2c0ff2b1cb66ee362678c4 Mon Sep 17 00:00:00 2001 From: droot Date: Thu, 18 Jun 2026 11:47:25 -0700 Subject: [PATCH 11/11] docs(setup): document decoupled vLLM sampler architecture and allocate dual DRA nodes --- docs/setup/gke-fft-timeslice.md | 45 ++++++++++----------------------- k8s/experiments/README.md | 2 +- 2 files changed, 15 insertions(+), 32 deletions(-) diff --git a/docs/setup/gke-fft-timeslice.md b/docs/setup/gke-fft-timeslice.md index 858dea37..8a65787f 100644 --- a/docs/setup/gke-fft-timeslice.md +++ b/docs/setup/gke-fft-timeslice.md @@ -17,17 +17,10 @@ three ideas that build on each other: There are three separate responsibilities in this PR. First, DRA is used only for GPU allocation and placement. The deployment creates -one `ResourceClaim` named `open-rl-trainer-gpu-1`. Trainer worker pods reference -that same claim. Kubernetes allocates one matching NVIDIA GPU to -the claim and schedules those pods onto a node where that device is available. -DRA does not serialize CUDA execution, perform checkpoint/restore, or decide -which process runs next. +two static `ResourceClaims` named `open-rl-trainer-gpu-1` and `open-rl-sampler-gpu-1`. Trainer worker pods reference `open-rl-trainer-gpu-1`, while Sampler worker pods reference `open-rl-sampler-gpu-1`. Kubernetes allocates one matching NVIDIA GPU to each claim and schedules those pods onto separate physical nodes where those devices are available. Second, the Kubernetes worker manager is the deployment launcher. It runs inside -the gateway process today. When the gateway receives `create_model` or -`create_model_from_state` in FFT mode, it -creates a pod for that `model_id` from the trainer worker pod template, stamps -the pod name, labels, job-id env var, and `--model-id`, then enqueues the request +the gateway process today. When the gateway receives `create_model` in FFT mode, it creates a trainer pod for that `model_id`. When it receives `create_sampling_client`, it creates a dedicated vLLM sampler pod for that `model_id` from the sampler worker pod template. It enqueues the request on the model-specific Redis queue. It is idempotent: if the trainer worker pod for a model is already running, it reuses it. @@ -98,14 +91,13 @@ scheduler. ## 1. DRA pins the GPU allocation -`k8s/deploy/distributed-fft-timeslice/06-gpu-resourceclaim.yaml` creates a -single namespace-scoped `ResourceClaim`: +`k8s/deploy/distributed-fft-timeslice/06-gpu-resourceclaim.yaml` and `08-sampler-resourceclaim.yaml` create dedicated namespace-scoped `ResourceClaims` for Trainers and Samplers: ```yaml apiVersion: resource.k8s.io/v1 kind: ResourceClaim metadata: - name: open-rl-trainer-gpu-1 + name: open-rl-trainer-gpu-1 # (and open-rl-sampler-gpu-1) spec: devices: requests: @@ -114,22 +106,15 @@ spec: deviceClassName: gpu.nvidia.com ``` -Trainer worker pods reference that same claim: +Trainer worker pods reference `open-rl-trainer-gpu-1`, while Sampler worker pods reference `open-rl-sampler-gpu-1`: ```yaml -resources: - claims: - - name: trainer-gpu resourceClaims: -- name: trainer-gpu - resourceClaimName: open-rl-trainer-gpu-1 +- name: trainer-gpu # (or sampler-gpu) + resourceClaimName: open-rl-trainer-gpu-1 # (or open-rl-sampler-gpu-1) ``` -Because this is a shared `ResourceClaim`, Kubernetes allocates a single matching -device to the claim and schedules all referencing pods where that allocated -device is accessible. Do not use a `ResourceClaimTemplate` for this PR's pinning -behavior: templates generate per-pod claims, which is the pattern for separate -devices. +Because these are shared `ResourceClaims`, Kubernetes allocates a single matching device to each claim and schedules referencing pods onto the dedicated nodes where those claims reside (`group.timeslice.io/trainers` vs `samplers`). DRA is the allocation and placement layer. It does not serialize CUDA execution by itself. This PR is intentionally an oversubscription model: multiple trainer @@ -230,8 +215,8 @@ gcloud container node-pools create gpu-dra \ --cluster "${CLUSTER}" --zone "${ZONE}" \ --machine-type g2-standard-24 \ --accelerator "type=nvidia-l4,count=1,gpu-driver-version=disabled" \ - --node-labels="group.timeslice.io/trainers=true,gke-no-default-nvidia-gpu-device-plugin=true,nvidia.com/gpu.present=true" \ - --num-nodes 1 + --node-labels="group.timeslice.io/trainers=true,group.timeslice.io/samplers=true,gke-no-default-nvidia-gpu-device-plugin=true,nvidia.com/gpu.present=true" \ + --num-nodes 2 ``` Install the GPU driver manually. Use the `latest` installer so the @@ -287,17 +272,15 @@ The deployment assumes one base model per rollout: set `BASE_MODEL` in `kustomization.yaml`, and the gateway uses that value for `get_info` and `create_model` requests that do not explicitly pass a base model. -There is no static trainer worker. Every `create_model` call makes the gateway -create a trainer worker pod named `open-rl-trainer-`, labeled: +There are no static worker deployments. Every `create_model` call makes the gateway create a trainer pod named `open-rl-trainer-`, and every `create_sampling_client` call makes the gateway create a dedicated vLLM sampler pod named `open-rl-sampler-`. Both are labeled: ```yaml -snapshot-agent: "true" # OpenRL/future coordinator marker -timeslice.io/group: trainers # snapshot-agent group +snapshot-agent: "true" # OpenRL coordinator marker +timeslice.io/group: trainers # (or samplers for vLLM rollout workers) timeslice.io/job-id: # per-worker identity ``` -The gateway's `open-rl-sa` service account has a Role allowing pod CRUD in the -workload namespace (`03-rbac.yaml`). +The gateway's `open-rl-sa` service account has a Role allowing pod CRUD in the workload namespace (`03-rbac.yaml`). When weight updates occur during FFT training, Trainers write checkpoints to NFS `/mnt/shared`, and Samplers dynamically reload those checkpoint safetensors in-place in ~1.1 seconds while yielding GPU VRAM via cooperative sleep. ## Setup 3: Run training on the cluster diff --git a/k8s/experiments/README.md b/k8s/experiments/README.md index 1ea0c61c..01460dda 100644 --- a/k8s/experiments/README.md +++ b/k8s/experiments/README.md @@ -50,7 +50,7 @@ gcloud container node-pools create gpu-dra \ --machine-type=g2-standard-12 \ --accelerator="type=nvidia-l4,count=1,gpu-driver-version=disabled" \ --node-labels="group.timeslice.io/trainers=true,group.timeslice.io/samplers=true,gke-no-default-nvidia-gpu-device-plugin=true,nvidia.com/gpu.present=true" \ - --num-nodes=1 + --num-nodes=2 ``` ### 3. Install NVIDIA Runtime Driver & DRA Kubelet Plugin