From b64bb342045a070458849e4eb8e6a38d3dffed57 Mon Sep 17 00:00:00 2001 From: Shuby Mishra Date: Mon, 26 Oct 2026 23:24:13 +0000 Subject: [PATCH 1/2] sampling: sample through pools Each model's sampling requests go to a pool stream, open_rl:sampling:, instead of a Redis list that a sampler emptied on read. asample adds a request and returns. The pool's dispatcher saves one result per request, so a sampler that stops loses nothing. A pool that one job owns runs server.vllm_sampler, vLLM in its own process, when it has one replica, for LoRA and FFT alike. It takes time-slicer turns only when it shares its trainer GPU. A pool with two or more replicas gets vllm serve Workloads behind a dispatcher pod with Envoy and the llm-d EPP. Local pools always have one replica. A pool that shared LoRA jobs use is always routed, even at one replica. Each job's ensure then only adds replicas, so the pool grows to the largest sampler_replicas any of its jobs asks for, and a second job never meets an in-process sampler it would have to replace. A sampler on its trainer's GPU is its pool's only replica, so create_model refuses sampler_replicas above 1 for an FFT job that shares its trainer's GPU, and tells the user to set openrl.exclusive=true. A routed pool's dispatcher Deployment and NetworkPolicy carry an openrl.io/owner label. release_owner deletes them by that label as well as through the pool's Workloads, so they do not outlive Workloads that went first. The worker launcher's Role gains list on deployments and networkpolicies for that lookup. The queue sampler and its vLLM patch are deleted, along with the store's sampler queue reads. The in-process sampler keeps sizing vLLM's batch for its GPU with sampler_batch_limits. sample_params and sample_result build the vLLM params and the saved result for both samplers. The in-process sampler raises RevisionUnloadable for a delta that does not apply on its weights, so those requests fail at once there too. It also fails a request vLLM refuses as a client error, such as a prompt past max_model_len or a token out of vocabulary, and a request that fails with EngineGenerateError, instead of sending it again until the deadline. vLLM loads a LoRA adapter inside the engine step, and a missing adapter raises LoRAAdapterNotFoundError there, which kills EngineCore. So the in-process sampler checks for the snapshot's adapter_config.json before it sends a record. On a miss it lists the snapshot directory once, since an NFS client keeps a failed lookup until it does, and if the file is still missing the record stays pending and is sent again. When the engine dies, every request in flight fails with EngineDeadError, and the sampler saves a failed result for each, as main's sampler did. A request that kills the engine would otherwise be served again by each restarted engine until the deadline. A request sent after the engine died never reached it, so it stays pending for the restarted worker. A sampler that shares its GPU with the trainer runs its batches through TimeSlicedDispatcher in vllm_sampler, one time-slicer turn per batch with main's holding_gpu. It loads weights through the same Dispatcher.ready as the routed dispatcher. The shared Dispatcher knows nothing about turns. A sampling session opened on a bare base-model name keeps main's rule, FFT when the deployment enables FFT and LoRA otherwise. Its runtime was the base model's name, which is also the name of the base model's shared LoRA runtime, so an FFT sampler for it took the LoRA pool's owner and read the LoRA pool's stream. TrainingModelMetadata.runtime now names such a runtime base-. Cluster samplers keep Triton's kernel cache at /vllm-cache/triton on the shared volume, beside vLLM's compile cache. Local samplers keep Triton's default, which already persists on one machine. The scheduler worker manager refuses to start without OPEN_RL_TMP_DIR, since its default of /tmp/open-rl is not on the shared volume. Each record names the job that sent it as "job", which job_of reads, so later changes can tell a pool's jobs apart. --- Makefile | 4 +- docs/configuration.md | 25 +- docs/fft/fft.md | 2 +- k8s/deploy/base/rbac.yaml | 4 +- skills/openrl/references/lifecycle.md | 7 +- src/server/api_server.py | 133 +++--- src/server/lora_snapshots.py | 37 +- src/server/model_metadata.py | 22 +- src/server/openrl_status.py | 6 +- src/server/sampling_dispatcher.py | 18 +- src/server/sampling_store.py | 5 + src/server/scheduler_worker_manager.py | 274 +++++++++++- src/server/store.py | 53 --- src/server/training_requests_processor.py | 24 +- src/server/vllm_options.py | 15 +- src/server/vllm_sampler.py | 505 +++++++--------------- src/server/worker_manager.py | 40 +- tests/test_api_server_paths.py | 7 + tests/test_lora_snapshots.py | 2 + tests/test_openrl_status.py | 59 ++- tests/test_pool_accept.py | 92 ++-- tests/test_redis_store.py | 14 - tests/test_sampler_patch.py | 256 ----------- tests/test_sampler_weight_rotation.py | 25 +- tests/test_scheduler_worker_manager.py | 92 +++- tests/test_session_lifecycle.py | 6 +- tests/test_vllm_sampler.py | 207 +++++++++ tests/test_worker_manager.py | 52 ++- 28 files changed, 1072 insertions(+), 914 deletions(-) delete mode 100644 tests/test_sampler_patch.py create mode 100644 tests/test_vllm_sampler.py diff --git a/Makefile b/Makefile index be38e5581..eb1b4d82b 100644 --- a/Makefile +++ b/Makefile @@ -19,7 +19,7 @@ HOST ?= 127.0.0.1 PORT ?= 9003 # The fully qualified base URL used by local CLI tools and clients BASE_URL ?= http://$(HOST):$(PORT) -UNIT_TESTS ?= tests.test_session_lifecycle tests.test_proto_codec tests.test_forward_only tests.test_fft_batch_failure tests.test_api_server_paths tests.test_accel_timeslicer tests.test_trainer_optimizer_correctness tests.test_trainer_logprob_chunks tests.test_worker_manager tests.test_scheduler_worker_manager tests.test_estimator tests.test_redis_store tests.test_cluster_eval_script tests.test_delta_weight_sync tests.test_diffing_backends tests.test_sampler_weight_rotation tests.test_commands tests.test_weight_sync_config tests.test_exclusive_trainer tests.test_delete_model tests.test_automodel_worker tests.test_worker_thread_device tests.test_trainer_data_parallel tests.test_sampler_batch_limits tests.test_openrl_status tests.test_skill_docs tests.test_trainer_ready tests.test_sampling_store tests.test_lora_snapshots tests.test_pool_accept tests.test_sampling_dispatcher +UNIT_TESTS ?= tests.test_session_lifecycle tests.test_proto_codec tests.test_forward_only tests.test_fft_batch_failure tests.test_api_server_paths tests.test_accel_timeslicer tests.test_trainer_optimizer_correctness tests.test_trainer_logprob_chunks tests.test_worker_manager tests.test_scheduler_worker_manager tests.test_estimator tests.test_redis_store tests.test_cluster_eval_script tests.test_delta_weight_sync tests.test_diffing_backends tests.test_sampler_weight_rotation tests.test_sampler_batch_limits tests.test_commands tests.test_weight_sync_config tests.test_exclusive_trainer tests.test_delete_model tests.test_automodel_worker tests.test_worker_thread_device tests.test_trainer_data_parallel tests.test_openrl_status tests.test_skill_docs tests.test_trainer_ready tests.test_sampling_store tests.test_lora_snapshots tests.test_pool_accept tests.test_sampling_dispatcher tests.test_vllm_sampler # Only forward BASE_URL to e2e when the user supplied it. The Makefile default # is for local CLI usage; e2e should start its own backend by default. TRAINING_TEST_BASE_URL ?= $(if $(filter environment command line,$(origin BASE_URL)),$(BASE_URL),) @@ -114,7 +114,7 @@ test: .PHONY: test-weight-transfer test-weight-transfer: - uv run --frozen --extra vllm --extra gpu --extra cluster python -m unittest tests.test_delta_weight_transfer_engine tests.test_delta_weight_sync tests.test_weight_sync_config tests.test_sampler_patch tests.test_weight_transfer_gpu + uv run --frozen --extra vllm --extra gpu --extra cluster python -m unittest tests.test_delta_weight_transfer_engine tests.test_delta_weight_sync tests.test_weight_sync_config tests.test_weight_transfer_gpu lint: uv run --extra dev ruff check . diff --git a/docs/configuration.md b/docs/configuration.md index 5d16bedbe..e1293d944 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -48,7 +48,7 @@ REDIS_URL=redis://127.0.0.1:6379/0 VLLM_ARCHITECTURE_OVERRIDE=Gemma4ForCausalLM | Env var | Default | What it does | | --- | --- | --- | | `BASE_MODEL` | unset | Hugging Face model id loaded by the trainer and, when using vLLM, by the sampler. | -| `SAMPLING_BACKEND` | `torch` in one process, `vllm` when `REDIS_URL` is set | Sampling backend selector. `torch` samples in the training process. `vllm` queues sampling requests for a vLLM sampler worker and needs `REDIS_URL`. | +| `SAMPLING_BACKEND` | `torch` in one process, `vllm` when `REDIS_URL` is set | Sampling backend selector. `torch` samples in the training process. `vllm` samples through the sampling pools described below and needs `REDIS_URL`. | | `REDIS_URL` | unset | Enables distributed mode by switching the request store to Redis. Leave unset for a single-machine run. | | `OPEN_RL_FUTURE_TTL_S` | `300` | How long resolved request results stay readable by `retrieve_future` after a worker resolves them. | @@ -56,7 +56,7 @@ REDIS_URL=redis://127.0.0.1:6379/0 VLLM_ARCHITECTURE_OVERRIDE=Gemma4ForCausalLM | Env var | Default | What it does | | --- | --- | --- | -| `OPEN_RL_TMP_DIR` | `/tmp/open-rl` | Root directory for adapter snapshots under `peft/` and saved states under `checkpoints/`. | +| `OPEN_RL_TMP_DIR` | `/tmp/open-rl` | Root directory for adapter snapshots under `peft/` and saved states under `checkpoints/`. Required with `OPEN_RL_WORKER_MANAGER=scheduler`, set to the shared volume every worker pod mounts. | | `OPEN_RL_TRAIN_TOKEN_BUDGET` | `auto` | Maximum `batch_size * max_sequence_length` for padded trainer chunks inside one `forward_backward` request. `auto` uses the smaller of 16384 and the tokens whose activations fit in half of the free device memory, read at each request. A number is used as given; `0` runs one datum per pass. | | `TRAINER_CUDA_VISIBLE_DEVICES`, `SAMPLER_CUDA_VISIBLE_DEVICES` | unset | GPUs given to the trainer and sampler worker processes the local worker manager launches. | @@ -100,3 +100,24 @@ On Kubernetes the release bundles set these variables. The API server runs with each trainer and sampler, and the OpenRL scheduler starts their pods. The API server writes each worker's environment into the `Workload`'s pod template, so there are no worker deployments to configure by hand. + +## Sampling pools + +Sampling goes through sampling pools. A pool serves one FFT job, one exclusive +LoRA job, or every shared LoRA job on one base model. Pending requests stay in a +Redis stream until a result is saved, so the Redis in `k8s/deploy/base` keeps an +append-only file on a persistent volume and never evicts keys. + +A pool that one job owns runs vLLM in the sampler's own process when it has one +replica. A pool with two or more replicas (`openrl.sampler_replicas`), and every +pool that shared LoRA jobs use, runs `vllm serve` Workloads behind one CPU +Deployment, named after the pool, that runs the dispatcher, Envoy and the llm-d +endpoint picker (configured by the `open-rl-sampler-router` ConfigMap in +`k8s/deploy/base`). A shared pool grows to the most replicas any of its jobs asks +for. On one machine every pool has one replica, in process. See +[the sampling pools design](designs/017-sampling-pools.md). + +| Env var | Default | What it does | +| --- | --- | --- | +| `OPEN_RL_DISPATCH_CONCURRENCY` | `256` | Requests one pool keeps in flight. | +| `OPEN_RL_DISPATCHER_IMAGE` | the worker image | Image for each pool's dispatcher container. The release sets it to the API-server image, which carries the dispatcher without the GPU stack and pulls in seconds. | diff --git a/docs/fft/fft.md b/docs/fft/fft.md index cdf9a1c43..a64da0e5e 100644 --- a/docs/fft/fft.md +++ b/docs/fft/fft.md @@ -128,7 +128,7 @@ API Server Data Queues (queue:, sampler_queue:) ---> Workers ### Provisioning & Activation Mechanics * **Trainer Provisioning:** When `/api/v1/create_model` is invoked, the API Server enqueues a launch command. `WorkerLaunchProcessor` pops the request and invokes `FFTWorkerManager.launch_trainer(model_id)` to provision the pod or subprocess. -* **Sampler Provisioning:** When `/api/v1/create_sampling_session` is invoked, `WorkerLaunchProcessor` invokes `FFTWorkerManager.launch_sampler(model_id)`. The sampler compiles CUDA graphs and posts `open_rl:sampler_ready: = "1"` to Redis. The API Server polls this key, blocking session return until physical readiness is verified. +* **Sampler Provisioning:** When `/api/v1/create_sampling_session` is invoked, the API Server makes sure the model's sampling pool exists and returns at once. Samples wait in the pool's Redis stream until a sampler is up to serve them. ### Graceful Queue-Based Teardown (Sentinel Pattern) To prevent dropping in-flight rollouts during de-provisioning, teardown uses the **Sentinel Pattern**: diff --git a/k8s/deploy/base/rbac.yaml b/k8s/deploy/base/rbac.yaml index d0df1fac5..41b395e86 100644 --- a/k8s/deploy/base/rbac.yaml +++ b/k8s/deploy/base/rbac.yaml @@ -11,10 +11,10 @@ rules: verbs: ["create", "get", "list", "watch", "delete"] - apiGroups: ["apps"] resources: ["deployments"] - verbs: ["create", "get", "delete"] + verbs: ["create", "get", "list", "delete"] - apiGroups: ["networking.k8s.io"] resources: ["networkpolicies"] - verbs: ["create", "delete"] + verbs: ["create", "list", "delete"] - apiGroups: [""] resources: ["pods"] verbs: ["list"] diff --git a/skills/openrl/references/lifecycle.md b/skills/openrl/references/lifecycle.md index 852c4d226..a4a35dcbb 100644 --- a/skills/openrl/references/lifecycle.md +++ b/skills/openrl/references/lifecycle.md @@ -20,14 +20,17 @@ The workers that serve one job appear together in the job's model view, Each worker in the model view has a `state`, the same five values on every server, and a `detail` in the scheduler's own words (Workload phase and reason, node). Branch on `state`; show `detail` to a human. `kubectl -n openrl-system get workloads` shows the same facts -`detail` quotes; the pods behind them are the scheduler's and the view does not read them. +`detail` quotes. The pods behind them are the scheduler's; the view reads only whether a +running sampler's pod is Ready, and names the reason when its container is waiting, for +example `CrashLoopBackOff`. A sampler is also `starting` while its pool's dispatcher has no +pod available, since no sample is sent without it. | State | Meaning | What `detail` says | Normal duration | |---|---|---|---| | `waiting_for_placement` | no accelerator assigned yet | `Pending (NoCapacity: ...)`: nothing in the cluster fits; `Pending (WaitingForCapacity)`: something fits but is busy | seconds; longer only when the cluster is full | | `starting` | placed on a node; pulling the image, mounting volumes, or loading the model | `Placing (PodCreated) on ` is the image pull and volume mount; `Running ..., loading the model` is the trainer loading the base model or vLLM starting | under a minute on a warm node; 5 to 10 minutes on a cold node. Trainer: `ready` about 2 minutes after `create_model`; sampler: 3 to 4 minutes more while vLLM starts | | `ready` | serving requests: the worker has its model loaded and is draining its queue | node name | | -| `failed` | the process exited with an error, or the request cannot be satisfied | the scheduler's reason, for example `Failed (PodFailed)`. A container that cannot pull its image or crash-loops is not reported as failed yet; it stays `starting` or `ready` past the normal duration, see the runbook | | +| `failed` | the process exited with an error, or the request cannot be satisfied | the scheduler's reason, for example `Failed (PodFailed)`. A container that cannot pull its image or crash-loops is not reported as failed yet; it stays `starting` or `ready` past the normal duration (a crash-looping sampler is `starting` with the reason in `detail`), see the runbook | | | `stopped` | exited normally, or being released | `being released; ...` while the reaper or `delete_model` tears it down | | While a call blocks, ask the server about the model from another shell; the examples diff --git a/src/server/api_server.py b/src/server/api_server.py index b8369afc4..f596cd882 100644 --- a/src/server/api_server.py +++ b/src/server/api_server.py @@ -10,6 +10,7 @@ import uuid from collections import defaultdict from contextlib import asynccontextmanager +from pathlib import Path from typing import Annotated, Any, Literal from fastapi import Depends, FastAPI, Header, HTTPException, Request @@ -35,7 +36,7 @@ ) from server.sampling_store import SamplingStore from server.session_registry import SessionRegistry -from server.store import RedisStateStore, get_state_store, get_store +from server.store import get_state_store, get_store from server.vllm_options import split_stop from server.worker_manager import LocalWorkerManager, WorkerManager, create_worker_manager, default_metadata, owner_of from training import commands @@ -75,7 +76,7 @@ async def reap_owner(owner: str) -> None: return print(f"[API_SERVER] No live session uses {owner}; tearing its workers down") for model in await asyncio.to_thread(worker_manager.release_owner, owner): - await state.delete_values(f"open_rl:sampler_ready:{model}", f"open_rl:trainer_ready:{model}") + await state.delete_values(f"open_rl:trainer_ready:{model}") await session_registry.forget(owner) @@ -457,15 +458,6 @@ def base_model_id_from_sampling_ref(model_id: str | None) -> str | None: return model_id.split("-samp-")[0] -def is_sampler_weights_ref(model_id: str | None) -> bool: - if not model_id or not model_id.startswith("tinker://"): - return False - - path = model_id[len("tinker://") :] - parts = path.split("/") - return len(parts) >= 3 and parts[1] == "sampler_weights" - - async def _extract_and_persist_model_metadata( req: CreateModelRequest | CreateModelFromStateRequest, request: Request | None = None, @@ -530,6 +522,8 @@ async def _extract_and_persist_model_metadata( trainer_cp=settings.trainer_cp, recipe=settings.recipe, ) + if meta_obj.sampler_replicas > 1 and meta_obj.shares_gpu(): + raise ValueError("openrl.sampler_replicas above 1 needs samplers on GPUs of their own; set openrl.exclusive=true") await persist_model_metadata(state, model_id, meta_obj) return model_id, meta_obj @@ -601,13 +595,11 @@ async def ensure_sampler_launched(model_id: str) -> None: traceback.print_exc() -def check_single_process_backend() -> None: - """vLLM sampling runs in a worker the API server launches, and they share queues - through Redis, so one process can only sample with torch.""" - if get_sampler_backend() == "vllm": +def check_sampler_backend() -> None: + if get_sampler_backend() == "vllm" and not os.getenv("REDIS_URL"): raise RuntimeError( - "SAMPLING_BACKEND=vllm needs REDIS_URL: the API server launches the vLLM sampler as a separate worker. " - "Set REDIS_URL (for example redis://127.0.0.1:6379/0), or unset SAMPLING_BACKEND to sample with torch in one process." + "Sampling pools need REDIS_URL, because vLLM sampling requests wait in a Redis stream for each pool's dispatcher. " + "Set REDIS_URL (for example redis://127.0.0.1:6379/0), or set BASE_MODEL and unset SAMPLING_BACKEND to sample with torch in one process." ) @@ -668,6 +660,7 @@ async def reap_dead_sessions(): async def lifespan(_: FastAPI): global worker_manager task = None + check_sampler_backend() if is_fft_enabled() or os.getenv("REDIS_URL") or os.getenv("OPEN_RL_WORKER_MANAGER"): worker_manager = create_worker_manager() if is_single_process_mode(): @@ -679,7 +672,6 @@ async def lifespan(_: FastAPI): print(f"-> Sampling backend: {get_sampler_backend()}") print(f"-> FFT enabled : {is_fft_enabled()}") print("-> Server mode : API server + worker loop in one process\n") - check_single_process_backend() if not is_fft_enabled(): from server import training_requests_processor @@ -920,7 +912,6 @@ async def delete_model(req: ModelRequest): elif worker_manager is not None: print(f"[API_SERVER] Requesting shutdown of workers for model {model_id}...") await store.put_request(commands.wire(commands.Shutdown(model_id=model_id)), active_set_id=await _resolve_active_set_id(model_id)) - await store.put_sampling_request({"request_id": "SHUTDOWN_SENTINEL", "model_id": model_id}) await asyncio.to_thread(worker_manager.release, model_id) now = time.time() meta.status = "completed" @@ -1168,26 +1159,13 @@ async def create_sampling_session(req: CreateSamplingSessionRequest): sess_id = req.model_id or "samp-session-live-123" target_model_id = sess_id - model_meta = await get_model_metadata(state, target_model_id) if target_model_id else None - ready_check_id = (model_meta.runtime(target_model_id) or target_model_id) if model_meta else target_model_id - await bind_session(req.session_id, target_model_id) - if get_sampler_backend() == "vllm" and ready_check_id: - # Launch by model ID so the worker manager retains the training kind. - # LoRA readiness is still reported under the shared base-model runtime. + if get_sampler_backend() == "vllm": + # Requests wait in the pool's stream, so there is no sampler to wait for. The SDK + # numbers seq_id from 0 in every sampling client, so each session needs its own id. await ensure_sampler_launched(target_model_id) - if isinstance(state, RedisStateStore): - print(f"[API_SERVER] Waiting for dynamic vLLM sampler worker to be ready for model {ready_check_id}...") - start_time = time.monotonic() - while True: - is_ready = await state.get_value(f"open_rl:sampler_ready:{ready_check_id}") - if is_ready == "1" or is_ready == b"1": - print(f"[API_SERVER] 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": f"{sess_id}@{uuid.uuid4().hex}", "type": "create_sampling_session"} return {"sampling_session_id": sess_id, "type": "create_sampling_session"} @@ -1202,7 +1180,8 @@ async def get_sampler(sampler_id: str): resolve to the base model, which is all the client wants -- it loads the tokenizer from the Hub itself. """ - base_model_id = base_model_id_from_sampling_ref(sampler_id) + ref = sampling_ref(sampler_id) + base_model_id = base_model_id_from_sampling_ref(ref) model_meta = await get_model_metadata(state, base_model_id) if base_model_id else None base_model = (model_meta.base_model if model_meta is not None else None) or base_model_id or get_default_model_name() if not base_model: @@ -1210,7 +1189,7 @@ async def get_sampler(sampler_id: str): return { "sampler_id": sampler_id, "base_model": base_model, - "model_path": sampler_id if sampler_id.startswith("tinker://") else None, + "model_path": ref if ref.startswith("tinker://") else None, } @@ -1223,6 +1202,35 @@ def sample_sequence_ids(request_id: str, num_samples: int) -> list[str]: return [f"{request_id}:{index}" for index in range(max(1, int(num_samples)))] +# How long a sample waits for weights another node just saved to show up here. +SAVED_WEIGHTS_WAIT_SECONDS = 2.0 + + +def visible(path: Path) -> bool: + """Whether a file another node saved shows here. An NFS client keeps a failed + lookup until it revalidates the directory, and listing a directory makes it + revalidate, so on a miss this lists the nearest directory above the file that + exists and looks again.""" + if path.is_file(): + return True + for directory in path.parents: + try: + os.listdir(directory) + break + except OSError: + continue + return path.is_file() + + +async def wait_until_visible(path: Path) -> bool: + deadline = time.monotonic() + SAVED_WEIGHTS_WAIT_SECONDS + while not await asyncio.to_thread(visible, path): + if time.monotonic() >= deadline: + return False + await asyncio.sleep(0.2) + return True + + async def accept_pool_sample(req: AsampleRequest, prompt: list[int]) -> str: """Missing weights are a 400, never a fallback to other weights.""" ref = sampling_ref(req.model_id or req.sampling_session_id or "") @@ -1233,16 +1241,13 @@ async def accept_pool_sample(req: AsampleRequest, prompt: list[int]) -> str: elif meta.fine_tuning_type == "full": revision, model = resolve_sampler_weights_path(ref), meta.base_model # The trainer writes metadata.json last, so a save in progress is missing too. - if not os.path.isfile(os.path.join(revision, "metadata.json")): + if not await wait_until_visible(Path(revision, "metadata.json")): raise HTTPException(status_code=400, detail=f"No saved sampler weights for {ref}") else: revision, model = "", snapshot_name(ref) - if not (snapshot_root() / model / "adapter_config.json").is_file(): + if not await wait_until_visible(snapshot_root() / model / "adapter_config.json"): raise HTTPException(status_code=400, detail=f"No saved LoRA sampler weights for {ref}") - try: - pool = meta.pool(lookup_id) - except ValueError as exc: - raise HTTPException(status_code=400, detail=str(exc)) from exc + pool = meta.pool(lookup_id) params = req.sampling_params stop, stop_token_ids = split_stop(params.stop) @@ -1250,6 +1255,7 @@ async def accept_pool_sample(req: AsampleRequest, prompt: list[int]) -> str: { "request_id": new_request_id(), "pool": pool, + "job": lookup_id, "revision": revision, "model": model, "idempotency_key": f"{req.sampling_session_id}:{req.seq_id}" if req.sampling_session_id and req.seq_id is not None else None, @@ -1285,44 +1291,7 @@ async def asample(req: AsampleRequest): ) return {"request_id": req_id, "sample_sequence_ids": sample_sequence_ids(req_id, num_samples)} - # vLLM backend - req_id = str(uuid.uuid4()) - carrier = await open_future(req_id) - - model_meta = await get_model_metadata(state, lookup_id) - fine_tuning_type = model_meta.fine_tuning_type if model_meta else "lora" - - if fine_tuning_type == "lora": - weights_path = None - lora_id = model_id - peft_dir = os.path.join(TMP_DIR, "peft", lookup_id, lookup_id) - lora_path = peft_dir if os.path.exists(peft_dir) else None - queue_id = (model_meta.runtime(lookup_id) if model_meta else None) or lookup_id - else: - resolved_path = resolve_sampler_weights_path(model_id) if is_sampler_weights_ref(model_id) or is_fft_enabled() else None - weights_path = resolved_path - lora_id = None - lora_path = None - queue_id = lookup_id - - sampling_req = { - "request_id": req_id, - "prompt_token_ids": prompt, - "max_tokens": params.max_tokens, - "temperature": params.temperature, - "stop": params.stop, - "top_p": params.top_p, - "top_k": params.top_k, - "num_samples": num_samples, - "lora_id": lora_id, - "lora_path": lora_path, - "weights_path": weights_path, - "include_prompt_logprobs": req.prompt_logprobs, - "model_id": queue_id, - "trace_context": carrier, - } - - await store.put_sampling_request(sampling_req) + req_id = await accept_pool_sample(req, prompt) return {"request_id": req_id, "sample_sequence_ids": sample_sequence_ids(req_id, num_samples)} diff --git a/src/server/lora_snapshots.py b/src/server/lora_snapshots.py index 82fad34b0..15ad2fd23 100644 --- a/src/server/lora_snapshots.py +++ b/src/server/lora_snapshots.py @@ -1,6 +1,7 @@ """Each saved LoRA reference gets an immutable copy of the adapter in vLLM's VLLM_LORA_RESOLVER_CACHE_DIR, so its directory name is the served model name.""" +import filecmp import hashlib import json import os @@ -18,18 +19,30 @@ def snapshot_name(ref: str) -> str: return hashlib.sha256(ref.encode()).hexdigest() +def same_files(a: Path, b: Path) -> bool: + names = sorted(p.relative_to(a) for p in a.rglob("*") if p.is_file()) + return names == sorted(p.relative_to(b) for p in b.rglob("*") if p.is_file()) and all(filecmp.cmp(a / n, b / n, shallow=False) for n in names) + + def freeze_adapter(model_id: str, base_model: str, *refs: str | None) -> None: - # Never overwrite a snapshot, because vLLM caches adapters and prefixes by name. paths = {snapshot_root() / snapshot_name(ref) for ref in refs if ref} - if any(path.exists() for path in paths): - raise FileExistsError("Sampler weights reference already exists; save with a new name or sequence ID") source = Path(os.getenv("OPEN_RL_TMP_DIR", "/tmp/open-rl"), "peft", model_id, model_id) - for path in paths: - path.parent.mkdir(parents=True, exist_ok=True) - with tempfile.TemporaryDirectory(prefix=".saving-", dir=path.parent) as staging: - staged = shutil.copytree(source, Path(staging, "adapter")) - # vLLM's filesystem resolver loads an adapter only when this names the - # served model exactly, and a trainer may have loaded it from a local path. - config_path = staged / "adapter_config.json" - config_path.write_text(json.dumps({**json.loads(config_path.read_text()), "base_model_name_or_path": base_model})) - staged.rename(path) + snapshot_root().mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(prefix=".saving-", dir=snapshot_root()) as staging: + staged = Path(staging, "adapter") + shutil.copytree(source, staged) + # vLLM's filesystem resolver loads an adapter only when this names the + # served model exactly, and a trainer may have loaded it from a local path. + config_path = staged / "adapter_config.json" + config = json.loads(config_path.read_text()) + config["base_model_name_or_path"] = base_model + config_path.write_text(json.dumps(config)) + # Never overwrite a snapshot, because vLLM caches adapters and prefixes by name. + # Saving the same weights under a name again, as an eval before and after no training does, changes nothing. + if any(path.exists() and not same_files(staged, path) for path in paths): + raise FileExistsError("Sampler weights reference already exists; save with a new name or sequence ID") + for path in paths: + if not path.exists(): + copy = Path(staging, path.name) + shutil.copytree(staged, copy) + copy.rename(path) diff --git a/src/server/model_metadata.py b/src/server/model_metadata.py index 0c4efe613..9fb2ce48f 100644 --- a/src/server/model_metadata.py +++ b/src/server/model_metadata.py @@ -75,9 +75,10 @@ def trainer_image(self) -> str | None: def runtime(self, model_id: str) -> str: """The id of the workers that serve this job. LoRA jobs share workers only - with jobs on the same trainer backend.""" + with jobs on the same trainer backend. A sampling session on a bare base-model + name that does not share is kept apart from the base model's shared runtime.""" if not self.shares_runtime(): - return model_id + return f"base-{model_id}" if model_id == self.base_model else model_id if self.trainer_backend == "pytorch": return self.base_model if image := self.trainer_image(): @@ -86,18 +87,13 @@ def runtime(self, model_id: str) -> str: return f"{self.trainer_backend}-{self.base_model}" def pool(self, model_id: str) -> str: - """The sampling pool that serves this job. A runtime named after a base model - (shared LoRA, or a bare base-model session) needs a name from - OPEN_RL_SAMPLER_POOLS, because base model names are not valid Kubernetes names.""" + """The sampling pool that serves this job's runtime. Base model names are not + valid Kubernetes names, so those runtimes get a slug and a hash.""" runtime = self.runtime(model_id) - name = json.loads(os.getenv("OPEN_RL_SAMPLER_POOLS", "{}")).get(runtime) - if name is None and not self.shares_runtime(): - name = f"sampler-{model_id}" - if name is None: - raise ValueError(f"OPEN_RL_SAMPLER_POOLS has no pool for {runtime}") - if not POOL_NAME.fullmatch(name): - raise ValueError(f"Pool name {name!r} is not a DNS-1123 label of at most 63 characters") - return name + if POOL_NAME.fullmatch(name := f"sampler-{runtime}"): + return name + slug = re.sub(r"[^a-z0-9]+", "-", runtime.lower()).strip("-")[:40] + return f"sampler-{slug}-{hashlib.sha256(runtime.encode()).hexdigest()[:8]}" def decode_model_metadata(raw: str | None) -> TrainingModelMetadata | None: diff --git a/src/server/openrl_status.py b/src/server/openrl_status.py index 9866dd9d3..4ea677700 100644 --- a/src/server/openrl_status.py +++ b/src/server/openrl_status.py @@ -54,9 +54,9 @@ async def workers_of_owner(ctx: StatusContext, owner: str) -> tuple[list[dict[st out = [] for found in found_workers: worker = {field: found.get(field) for field in WORKER_FIELDS} - # The launcher sees a pod run minutes before the model is loaded; each - # worker announces itself when it can serve (open_rl:_ready:). - announced = await ctx.state.get_value(f"open_rl:{found['role']}_ready:{found['model_id']}") is not None + # The launcher sees a pod run minutes before the model is loaded, so a + # trainer announces itself when it can serve (open_rl:trainer_ready:). + announced = found["role"] != "trainer" or await ctx.state.get_value(f"open_rl:trainer_ready:{found['model_id']}") is not None if worker["state"] == READY and not announced: worker["state"] = STARTING worker["detail"] = f"{worker['detail']}, loading the model" diff --git a/src/server/sampling_dispatcher.py b/src/server/sampling_dispatcher.py index 1cb84fd68..3c696ed65 100644 --- a/src/server/sampling_dispatcher.py +++ b/src/server/sampling_dispatcher.py @@ -1,7 +1,7 @@ """Serves one sampling pool's pending records with a sampler and saves the first result for each. -The sampler is HttpSampler, which sends through the pool's router, Envoy in front of the llm-d EPP, to its vllm serve -replicas.""" +The sampler is either vLLM in this process (vllm_sampler.Sampler), for a pool of one replica that one job owns, or +HttpSampler, which sends through the pool's router, Envoy in front of the llm-d EPP, to its vllm serve replicas.""" import asyncio import contextlib @@ -18,6 +18,7 @@ import redis.asyncio as redis from server.sampling_store import SamplingStore +from server.vllm_options import sample_params, sample_result logger = logging.getLogger(__name__) @@ -68,14 +69,13 @@ def __init__(self, revision: str, held: list[str]) -> None: def to_result(body: dict[str, Any], record: dict[str, Any]) -> dict[str, Any]: - """The old sampler's result shape, from a vLLM /inference/v1/generate response.""" - logprobs = [[e["logprob"] for e in c["logprobs"]["content"]] for c in body["choices"]] - sequences = [{"tokens": c["token_ids"], "logprobs": lp, "stop_reason": c["finish_reason"]} for c, lp in zip(body["choices"], logprobs)] - result: dict[str, Any] = {"type": "sample", "sequences": sequences} + """The result for a vLLM /inference/v1/generate response.""" + sequences = [(c["token_ids"], [e["logprob"] for e in c["logprobs"]["content"]], c["finish_reason"]) for c in body["choices"]] + prompt_logprobs = None if record["include_prompt_logprobs"]: pairs = zip(record["prompt_token_ids"], body["prompt_logprobs"], strict=True) - result["prompt_logprobs"] = [None if e is None else e[str(token)]["logprob"] for token, e in pairs] - return result + prompt_logprobs = [None if e is None else e[str(token)]["logprob"] for token, e in pairs] + return sample_result(sequences, prompt_logprobs) class Dispatcher: @@ -212,7 +212,7 @@ async def ensure_weights(self, revision: str) -> bool: async def generate(self, r: dict[str, Any]) -> dict[str, Any] | None: """The result, or None when the answer did not come from the record's revision.""" target, sent = r["revision"] or "default", time.monotonic() - params = {**r["sampling_params"], "logprobs": 0, "prompt_logprobs": 0 if r["include_prompt_logprobs"] else None} + params = sample_params(r) attempt_id = f"{r['request_id']}-{uuid.uuid4().hex[:8]}" # The store hands prompts back as array('I'), which JSON cannot encode. body = {"request_id": attempt_id, "token_ids": list(r["prompt_token_ids"]), "model": r["model"], "sampling_params": params} diff --git a/src/server/sampling_store.py b/src/server/sampling_store.py index 6b044acf9..e5f8ce2d6 100644 --- a/src/server/sampling_store.py +++ b/src/server/sampling_store.py @@ -35,6 +35,11 @@ def future_key(request_id: str) -> str: return f"open_rl:future:{request_id}" +def job_of(r: dict[str, Any]) -> str: + """The job that asked for the record. Records accepted before the API server named the job group under their model.""" + return r.get("job") or r["model"] + + class SamplingStore: def __init__(self, client: redis.Redis): self.redis = client diff --git a/src/server/scheduler_worker_manager.py b/src/server/scheduler_worker_manager.py index 9da063f99..eec91070d 100644 --- a/src/server/scheduler_worker_manager.py +++ b/src/server/scheduler_worker_manager.py @@ -19,6 +19,7 @@ import dataclasses import datetime +import json import logging import os import time @@ -28,18 +29,21 @@ from kubernetes import client, config from server.estimator import Footprint, footprint +from server.lora_snapshots import snapshot_root +from server.sampling_dispatcher import revision_file, tmp_root +from server.vllm_options import serve_args from server.worker_manager import ( FAILED, READY, STARTING, STOPPED, + TRAINER_MODULE, WAITING_FOR_PLACEMENT, base_model_of, owner_id, runtime_of, worker_args, worker_env, - worker_module, ) logger = logging.getLogger(__name__) @@ -48,6 +52,12 @@ VERSION = "v1alpha1" PLURAL = "workloads" +EPP_IMAGE = "ghcr.io/llm-d/llm-d-router-endpoint-picker:v0.11.0" +ENVOY_IMAGE = "docker.io/envoyproxy/envoy:distroless-v1.33.2" +ROUTER_CONFIG = "/etc/sampler-router" +# On a pool's Deployment and NetworkPolicy, so release_owner finds them after the pool's Workloads are gone. +OWNER_LABEL = "openrl.io/owner" + def workload_name(role: str, owner: str, is_lora: bool, index: int = 0) -> str: # A second compatible request renders the same name, and the create's @@ -75,6 +85,22 @@ class Worker: # Which replica of its role this is. index: int = 0 + # The sampling pool a sampler serves. + pool: str | None = None + + @property + def replicas(self) -> int: + # A sampler on its trainer's GPU is its pool's only replica, since only one fits beside the trainer. + if self.role != "sampler" or self.meta.shares_gpu(): + return 1 + return self.meta.sampler_replicas + + @property + def routed(self) -> bool: + # A pool several jobs share is always routed, so a job that asks for more replicas only adds them. A pool of one + # replica that a job owns alone has nothing to route, so that replica serves the pool itself. + return self.role == "sampler" and (self.replicas > 1 or self.meta.shares_runtime()) + @property def owner(self) -> str: return owner_id(self.runtime) @@ -84,30 +110,30 @@ def name(self) -> str: return workload_name(self.role, self.owner, self.is_lora, self.index) -def replicas_of(model_id: str, role: str) -> int: - """How many workers of this role the model asked for. Each sampler replica - is its own Workload draining the shared sampling queue.""" - if role != "sampler": - return 1 - meta, _, _ = runtime_of(model_id) - return meta.sampler_replicas +def pool_labels(component: str, pool: str) -> dict[str, str]: + return {"app.kubernetes.io/component": component, "app.kubernetes.io/instance": pool} + + +def downward(name: str, path: str) -> dict[str, Any]: + return {"name": name, "valueFrom": {"fieldRef": {"fieldPath": path}}} def describe_worker(model_id: str, role: str) -> Worker: meta, runtime, is_lora = runtime_of(model_id) base_model = base_model_of(meta, runtime) + pool = meta.pool(runtime) if role == "sampler" else None exclusive = not meta.shares_gpu() devices = meta.trainer_gpus if role == "trainer" else 1 size = footprint(base_model, meta.fine_tuning_type, role) # Each torchrun rank is its own process with its own host memory. size = replace(size, host_request_bytes=size.host_request_bytes * devices, host_limit_bytes=size.host_limit_bytes * devices) - return Worker(role, runtime, base_model, is_lora, exclusive, meta, size, devices) + return Worker(role, runtime, base_model, is_lora, exclusive, meta, size, devices, pool=pool) def pod_env(worker: Worker) -> list[dict[str, Any]]: """The shared worker env plus what only the cluster knows. The time-slice group is placement's and the scheduler stamps it.""" - tmp_dir = os.getenv("OPEN_RL_TMP_DIR", "/mnt/shared/open-rl") + tmp_dir = tmp_root() values = { "REDIS_URL": os.environ["REDIS_URL"], "OPEN_RL_TMP_DIR": tmp_dir, @@ -129,11 +155,15 @@ def pod_env(worker: Worker) -> list[dict[str, Any]]: # rest load the result. if worker.role == "sampler": values["VLLM_CACHE_ROOT"] = f"{tmp_dir}/vllm-cache" + # Triton writes each kernel to a temp dir and renames it into place, so samplers can share one cache. + values["TRITON_CACHE_DIR"] = f"{tmp_dir}/vllm-cache/triton" # No other worker shares an exclusive worker's GPUs, so it never parks. if worker.exclusive: values["OPEN_RL_TIME_SLICING"] = "off" env: list[dict[str, Any]] = [{"name": name, "value": value} for name, value in values.items()] - env.append({"name": "OPEN_RL_ACCEL_TIMESLICER_HOST", "valueFrom": {"fieldRef": {"fieldPath": "status.hostIP"}}}) + env.append(downward("OPEN_RL_ACCEL_TIMESLICER_HOST", "status.hostIP")) + if worker.pool: + env.append(downward("POD_IP", "status.podIP")) return env @@ -144,10 +174,10 @@ def worker_container(worker: Worker) -> tuple[str, list[str]]: image = worker.meta.trainer_image() or os.getenv("OPEN_RL_AUTOMODEL_IMAGE", "ghcr.io/gke-labs/open-rl/automodel:latest") if worker.devices > 1: torchrun = ["-m", "torch.distributed.run", "--standalone", f"--nproc-per-node={worker.devices}"] - return image, ["python", "-u", *torchrun, "-m", worker_module(worker.role)] - return image, ["python", "-u", "-m", worker_module(worker.role)] + return image, ["python", "-u", *torchrun, "-m", TRAINER_MODULE] + return image, ["python", "-u", "-m", TRAINER_MODULE] image = os.getenv("OPEN_RL_WORKER_IMAGE", "ghcr.io/gke-labs/open-rl/server:latest") - return image, ["uv", "run", "python", "-u", "-m", worker_module(worker.role)] + return image, ["uv", "run", "python", "-u", "-m", TRAINER_MODULE] def pod_template(worker: Worker) -> dict[str, Any]: @@ -183,11 +213,118 @@ def pod_template(worker: Worker) -> dict[str, Any]: template["spec"]["containers"][0]["volumeMounts"].append({"name": "dshm", "mountPath": "/dev/shm"}) template["spec"]["volumes"].append({"name": "dshm", "emptyDir": {"medium": "Memory"}}) + container = template["spec"]["containers"][0] + if worker.pool and not worker.routed: + container["command"] = ["uv", "run", "python", "-u", "-m", "server.vllm_sampler"] + container["args"] = ["--pool", worker.pool] + container["readinessProbe"] = {"exec": {"command": ["test", "-f", "/tmp/open-rl-sampler-ready"]}, "periodSeconds": 2} + template["metadata"] = {"labels": pool_labels("sampler", worker.pool)} + # A pool sampler is plain `vllm serve`, ready only while it serves the version the pool's dispatcher expects. + elif worker.pool: + container["command"] = ["uv", "run", "vllm", "serve"] + # The pool's EPP subscribes to these events at llm-d's default port and topic prefix; the topic names the pod and the served model. + kv_events = {"enable_kv_cache_events": True, "publisher": "zmq", "endpoint": "tcp://*:5557", "topic": f"kv@$(POD_IP):8000@{worker.base_model}"} + env = {e["name"]: e.get("value") for e in container["env"]} + container["args"] = [*serve_args(worker.base_model, worker.is_lora, env), "--kv-events-config", json.dumps(kv_events)] + if worker.is_lora: + # vLLM loads each frozen adapter snapshot by name from this directory. + container["env"] += [ + {"name": "VLLM_ALLOW_RUNTIME_LORA_UPDATING", "value": "true"}, + {"name": "VLLM_LORA_RESOLVER_CACHE_DIR", "value": str(snapshot_root())}, + ] + probe = ( + "import json, sys, urllib.request\n" + 'version = json.load(urllib.request.urlopen("http://127.0.0.1:8000/weight_info", timeout=1))["weight_version"]\n' + f"sys.exit(version != open({str(revision_file(worker.pool))!r}).read())" + ) + container["readinessProbe"] = {"exec": {"command": ["/app/.venv/bin/python", "-c", probe]}, "periodSeconds": 2, "timeoutSeconds": 2} + template["metadata"] = {"labels": pool_labels("sampler", worker.pool)} + if pull_policy := os.getenv("OPEN_RL_WORKER_IMAGE_PULL_POLICY"): template["spec"]["containers"][0]["imagePullPolicy"] = pull_policy return template +def dispatcher_body(pool: str, owner: str) -> dict[str, Any]: + """The pool's CPU pod, named after the pool: the dispatcher, Envoy and the llm-d + EPP. Recreate stops the old pod before the new one starts, so two dispatchers + never send for one pool.""" + labels = pool_labels("dispatcher", pool) + env = [ + {"name": "REDIS_URL", "value": os.environ["REDIS_URL"]}, + {"name": "OPEN_RL_TMP_DIR", "value": tmp_root()}, + {"name": "OPEN_RL_SAMPLER_POOL", "value": pool}, + downward("OPEN_RL_WORKER_NAMESPACE", "metadata.namespace"), + ] + if concurrency := os.getenv("OPEN_RL_DISPATCH_CONCURRENCY"): + env.append({"name": "OPEN_RL_DISPATCH_CONCURRENCY", "value": concurrency}) + dispatcher = { + "name": "dispatcher", + # The dispatcher is a small CPU process; the API-server image carries it without the GPU stack. + "image": os.getenv("OPEN_RL_DISPATCHER_IMAGE") or os.getenv("OPEN_RL_WORKER_IMAGE", "ghcr.io/gke-labs/open-rl/server:latest"), + "command": ["uv", "run", "python", "-u", "-m", "server.sampling_dispatcher"], + "env": env, + "volumeMounts": [{"name": "shared-storage", "mountPath": "/mnt/shared"}], + "resources": {"requests": {"cpu": "250m", "memory": "512Mi"}, "limits": {"memory": "2Gi"}}, + } + if pull_policy := os.getenv("OPEN_RL_WORKER_IMAGE_PULL_POLICY"): + dispatcher["imagePullPolicy"] = pull_policy + envoy = { + "name": "envoy", + "image": ENVOY_IMAGE, + "args": ["-c", f"{ROUTER_CONFIG}/envoy.yaml", "--log-level", "warn"], + "volumeMounts": [{"name": "router", "mountPath": ROUTER_CONFIG}], + "resources": {"requests": {"cpu": "250m", "memory": "256Mi"}, "limits": {"memory": "1Gi"}}, + } + epp = { + "name": "epp", + "image": EPP_IMAGE, + "args": [ + *("--endpoint-selector", f"app.kubernetes.io/component=sampler,app.kubernetes.io/instance={pool}", "--endpoint-target-ports", "8000"), + *("--config-file", f"{ROUTER_CONFIG}/epp.yaml", "--metrics-endpoint-auth=false"), + ], + "env": [downward("POD_NAME", "metadata.name"), downward("NAMESPACE", "metadata.namespace")], + "livenessProbe": {"grpc": {"port": 9003, "service": "liveness"}, "periodSeconds": 10}, + "resources": {"requests": {"cpu": "500m", "memory": "1Gi"}, "limits": {"memory": "4Gi"}}, + "volumeMounts": [{"name": "router", "mountPath": ROUTER_CONFIG}], + } + return { + "apiVersion": "apps/v1", + "kind": "Deployment", + "metadata": {"name": pool, "labels": {**labels, "app.kubernetes.io/managed-by": "open-rl-api-server", OWNER_LABEL: owner}}, + "spec": { + "replicas": 1, + "strategy": {"type": "Recreate"}, + "selector": {"matchLabels": labels}, + "template": { + "metadata": {"labels": labels}, + "spec": { + "serviceAccountName": "open-rl-sampler-dispatcher", + "containers": [dispatcher, envoy, epp], + "volumes": [ + {"name": "router", "configMap": {"name": "open-rl-sampler-router"}}, + {"name": "shared-storage", "persistentVolumeClaim": {"claimName": os.getenv("OPEN_RL_SHARED_PVC", "open-rl-shared-pvc")}}, + ], + }, + }, + }, + } + + +def network_policy_body(pool: str, owner: str) -> dict[str, Any]: + """Only the pool's CPU pod reaches its samplers, for inference and KV events.""" + return { + "apiVersion": "networking.k8s.io/v1", + "kind": "NetworkPolicy", + "metadata": {"name": pool, "labels": {"app.kubernetes.io/managed-by": "open-rl-api-server", OWNER_LABEL: owner}}, + "spec": { + "podSelector": {"matchLabels": pool_labels("sampler", pool)}, + "policyTypes": ["Ingress"], + "ingress": [{"from": [{"podSelector": {"matchLabels": pool_labels("dispatcher", pool)}}], "ports": [{"port": 8000}, {"port": 5557}]}], + }, + } + + def accelerator_spec(worker: Worker) -> dict[str, Any]: """A torchrun group asks for its devices, each sized for a whole replica.""" if worker.devices > 1: @@ -196,10 +333,12 @@ def accelerator_spec(worker: Worker) -> dict[str, Any]: def workload_body(worker: Worker) -> dict[str, Any]: + # A pool label marks a Workload whose pool has a dispatcher, which release_owner deletes with it. + labels = {"app.kubernetes.io/managed-by": "open-rl-api-server", **(pool_labels("sampler", worker.pool) if worker.routed else {})} return { "apiVersion": f"{GROUP}/{VERSION}", "kind": "Workload", - "metadata": {"name": worker.name, "labels": {"app.kubernetes.io/managed-by": "open-rl-api-server"}}, + "metadata": {"name": worker.name, "labels": labels}, "spec": { "role": worker.role, "trainingKind": "lora" if worker.is_lora else "fft", @@ -216,21 +355,64 @@ def workload_body(worker: Worker) -> dict[str, Any]: class SchedulerWorkerManager: """Runs trainer and sampler workers by creating Workload objects.""" - def __init__(self, custom_api: Any = None): + def __init__(self, custom_api: Any = None, apps_api: Any = None, networking_api: Any = None, 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") + if not os.getenv("OPEN_RL_TMP_DIR"): + raise RuntimeError("The scheduler worker manager needs OPEN_RL_TMP_DIR on the shared volume every worker pod mounts") self.namespace = os.getenv("OPEN_RL_WORKER_NAMESPACE", "openrl-system") if custom_api is None: config.load_incluster_config() custom_api = client.CustomObjectsApi() + apps_api, networking_api, core_api = client.AppsV1Api(), client.NetworkingV1Api(), client.CoreV1Api() self.custom_api = custom_api + self.apps_api, self.networking_api, self.core_api = apps_api, networking_api, core_api # The model view may be polled; one read per few seconds is plenty. self._workers_cache: tuple[float, list[dict[str, Any]]] | None = None self.view_cache_seconds = 5.0 def ensure(self, model_id: str, role: str) -> None: - for index in range(replicas_of(model_id, role)): - self.ensure_workload(dataclasses.replace(describe_worker(model_id, role), index=index)) + worker = describe_worker(model_id, role) + if worker.routed and worker.is_lora: + # vLLM's filesystem LoRA resolver exits at startup without its directory, and no adapter may be frozen yet. + # The API server mounts the shared volume the replicas read it from. + snapshot_root().mkdir(parents=True, exist_ok=True) + for index in range(worker.replicas): + self.ensure_workload(dataclasses.replace(worker, index=index)) + if worker.routed: + self.ensure_dispatcher(worker.pool, worker.owner) + + def ensure_dispatcher(self, pool: str, owner: str) -> None: + # A released pool's Deployment stays until its pod is gone, so waiting + # here keeps an old and a new dispatcher from overlapping. The + # NetworkPolicy of a released pool is gone long before that. + started, waited = time.monotonic(), False + while not self.create_once(self.apps_api.create_namespaced_deployment, dispatcher_body(pool, owner)): + try: + if not self.apps_api.read_namespaced_deployment(pool, self.namespace).metadata.deletion_timestamp: + break + except Exception as exc: + if getattr(exc, "status", None) != 404: + raise + if time.monotonic() > started + 180: + raise RuntimeError(f"dispatcher {pool} has been terminating for over three minutes") + if not waited: + # A warning, because the API server configures no logging and drops info lines. + logger.warning("Dispatcher %s of a released pool is still terminating, so this sampling session waits up to 180s for it", pool) + waited = True + time.sleep(2) + if waited: + logger.warning("Waited %.0fs for the old dispatcher %s to terminate", time.monotonic() - started, pool) + self.create_once(self.networking_api.create_namespaced_network_policy, network_policy_body(pool, owner)) + + def create_once(self, create: Any, body: dict[str, Any]) -> bool: + try: + create(self.namespace, body) + return True + except Exception as exc: + if getattr(exc, "status", None) != 409: + raise + return False def ensure_workload(self, worker: Worker) -> None: role = worker.role @@ -270,10 +452,22 @@ def release(self, model_id: str) -> None: self.release_owner(owner_id(runtime)) def release_owner(self, owner: str) -> set[str]: - """Delete the owner's workloads. The scheduler's finalizer frees the seats.""" + """Delete the owner's workloads. The scheduler's finalizer frees the seats. + A pool's CPU pod goes first. It is found through the samplers, or by its owner label once they are gone.""" selector = "app.kubernetes.io/managed-by=open-rl-api-server" found = self.custom_api.list_namespaced_custom_object(GROUP, VERSION, self.namespace, PLURAL, label_selector=selector) ours = [item for item in found["items"] if item["spec"]["ownerID"] == owner] + pools = {item["metadata"]["labels"].get("app.kubernetes.io/instance") for item in ours} - {None} + owned = f"{OWNER_LABEL}={owner}" + pools |= {d.metadata.name for d in self.apps_api.list_namespaced_deployment(self.namespace, label_selector=owned).items} + pools |= {p.metadata.name for p in self.networking_api.list_namespaced_network_policy(self.namespace, label_selector=owned).items} + for pool in pools: + for delete in (self.networking_api.delete_namespaced_network_policy, self.apps_api.delete_namespaced_deployment): + try: + delete(pool, self.namespace, propagation_policy="Foreground") + except Exception as exc: + if getattr(exc, "status", None) != 404: + raise for item in ours: self.delete_workload(item["metadata"]["name"]) return {item["spec"]["modelID"] for item in ours} @@ -304,11 +498,14 @@ def workers(self) -> list[dict[str, Any]]: def load_workers(self) -> list[dict[str, Any]]: """Our Workloads, one entry per worker, in the client vocabulary. The pod is the scheduler's; what the API server knows about it is what the - scheduler wrote into Workload status.""" + scheduler wrote into Workload status, plus whether a running sampler's + pod is Ready and its pool's dispatcher has a pod available.""" selector = "app.kubernetes.io/managed-by=open-rl-api-server" found = self.custom_api.list_namespaced_custom_object(GROUP, VERSION, self.namespace, PLURAL, label_selector=selector) now = time.time() workers = [] + sampler_pods = None + dispatchers: dict[str, bool] = {} for item in found["items"]: spec, status, meta = item["spec"], item.get("status") or {}, item["metadata"] state, detail = worker_state( @@ -317,6 +514,17 @@ def load_workers(self) -> list[dict[str, Any]]: node=status.get("nodeName"), deleting=bool(meta.get("deletionTimestamp")), ) + if state == READY and spec["role"] == "sampler": + if sampler_pods is None: + pods = self.core_api.list_namespaced_pod(self.namespace, label_selector="app.kubernetes.io/component=sampler", _request_timeout=10) + sampler_pods = {pod.metadata.name: pod for pod in pods.items} + state, detail = sampler_state(detail, sampler_pods.get(status.get("podName"))) + # The readiness probe reads a revision file that outlives the pool's dispatcher, and without the dispatcher no sample is sent. + if state == READY and spec["role"] == "sampler" and (pool := meta.get("labels", {}).get("app.kubernetes.io/instance")): + if pool not in dispatchers: + dispatchers[pool] = self.dispatcher_available(pool) + if not dispatchers[pool]: + state, detail = STARTING, f"{detail}, dispatcher {pool} has no available pod" workers.append( { "workload": meta["name"], @@ -331,12 +539,36 @@ def load_workers(self) -> list[dict[str, Any]]: ) return workers + def dispatcher_available(self, pool: str) -> bool: + try: + deployment = self.apps_api.read_namespaced_deployment(pool, self.namespace) + except Exception as exc: + if getattr(exc, "status", None) != 404: + raise + return False + return bool(deployment.status and deployment.status.available_replicas) + + +def sampler_state(detail: str, pod: Any) -> tuple[str, str]: + """A pool sampler's Workload is Running as soon as its container starts, and + stays Running while the container crash-loops. Its readiness probe passes only + while vLLM serves the weights the pool expects, so the pod's Ready condition is + what makes it ready.""" + status = getattr(pod, "status", None) + if any(c.type == "Ready" and c.status == "True" for c in getattr(status, "conditions", None) or []): + return READY, detail + for container in getattr(status, "container_statuses", None) or []: + waiting = container.state.waiting if container.state else None + if waiting and waiting.reason: + return STARTING, f"{detail}, {waiting.reason}, restart count {container.restart_count}" + return STARTING, f"{detail}, loading the model" + def worker_state(phase: str, reason: str, node: str | None, deleting: bool) -> tuple[str, str]: """Map a Workload's status onto the client vocabulary. The raw facts go into `detail`, for a human; `state` is what a client branches on. A node assignment is the placement boundary; phase Running is the readiness - boundary (the model view refines it for samplers).""" + boundary (sampler_state refines it for samplers).""" facts = phase + (f" ({reason})" if reason else "") if node: facts += f" on {node}" diff --git a/src/server/store.py b/src/server/store.py index 59c3d411b..7f610a972 100644 --- a/src/server/store.py +++ b/src/server/store.py @@ -31,16 +31,6 @@ 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_result(self, req_id: str, result: dict[str, Any]) -> None: """Resolve a future by its request ID.""" @@ -61,7 +51,6 @@ def __init__(self): self.active_tenants_cv = asyncio.Condition() self.futures_store: dict[str, dict[str, Any]] = {} self.futures_cv = asyncio.Condition() - self.sampling_queues: dict[str, asyncio.Queue] = {} async def put_request(self, req_data: dict[str, Any], active_set_id: str | None = None) -> None: model_id = req_data.get("model_id", "default") @@ -112,23 +101,6 @@ async def get_requests(self, active_set_id: str | None = None) -> list[dict[str, 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: - model_id = req_data.get("model_id", "default") - if model_id not in self.sampling_queues: - self.sampling_queues[model_id] = asyncio.Queue() - await self.sampling_queues[model_id].put(req_data) - - async def get_sampling_requests_for_model(self, model_id: str) -> list[dict[str, Any]]: - if model_id not in self.sampling_queues: - return [] - queue = self.sampling_queues[model_id] - if queue.empty(): - return [] - batch = [queue.get_nowait()] - while not queue.empty(): - batch.append(queue.get_nowait()) - return batch - async def set_result(self, req_id: str, result: dict[str, Any]) -> None: if result.get("status") == "pending": return @@ -234,31 +206,6 @@ 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_result(self, req_id: str, result: dict[str, Any]) -> None: if result.get("status") == "pending": return diff --git a/src/server/training_requests_processor.py b/src/server/training_requests_processor.py index 2ff259b0d..5d764f568 100644 --- a/src/server/training_requests_processor.py +++ b/src/server/training_requests_processor.py @@ -2,7 +2,6 @@ import argparse import asyncio -import json import os import shutil import threading @@ -19,8 +18,10 @@ from accel_timeslicer.time_slicer import TimeSlicerClient, time_slicer_client_from_env, time_slicing_enabled, workload_from_env from accel_timeslicer.workload import TRAINER_CLAIM, local_workload_name +from server.lora_snapshots import freeze_adapter from server.model_metadata import get_model_metadata -from server.store import RequestStore, StateStore, get_state_store, get_store +from server.sampling_store import SamplingStore +from server.store import RedisStore, RequestStore, StateStore, get_state_store, get_store from training import commands from training.commands import parse_command from training.distributed import broadcast_object, is_distributed, is_primary, local_rank @@ -108,6 +109,8 @@ def __init__( raise RuntimeError("A dedicated trainer worker requires REDIS_URL so it can share queues and futures with the API server") self.store = store + # Sampling pools live in Redis, so the in-process trainer's in-memory store has none. + self.sampling = SamplingStore(store.redis) if isinstance(store, RedisStore) else None self.worker = worker self.model_id = model_id self.active_tenant_set_id = active_tenant_set_id @@ -369,7 +372,10 @@ async def dispatch_operation(self, command: commands.TrainingCommand) -> dict[st ref = command.path or command.sampling_session_id checkpoint = await asyncio.to_thread(self.worker.save_for_sampler, command.model_id, command.alias, ref) if checkpoint: - await self.publish_checkpoint(command.model_id, checkpoint) + await self.drop_old_checkpoints(command.model_id, checkpoint) + elif is_primary() and ref: + # A sampler serves a LoRA save by its reference, so the adapter is frozen under that name. + await asyncio.to_thread(freeze_adapter, command.model_id, self.worker.base_model_name, command.path, command.sampling_session_id) return {"path": command.path, "sampling_session_id": command.sampling_session_id, "type": "sampler_weights_saved"} case commands.DeleteModel(): await asyncio.to_thread(self.worker.delete_model, command.model_id) @@ -379,14 +385,16 @@ async def dispatch_operation(self, command: commands.TrainingCommand) -> dict[st case _: raise NotImplementedError(f"Training request op {command.op!r} is not supported") - async def publish_checkpoint(self, model_id: str, local_path: str) -> None: - """Tell the samplers about a new checkpoint and drop the versions they no longer need.""" + async def drop_old_checkpoints(self, model_id: str, local_path: str) -> None: + """Drop sampler versions beyond the newest few, except ones pending requests still name.""" if not is_primary(): return - if hasattr(self.store, "redis"): - num_subs = await self.store.redis.publish(f"open_rl:weight_update:{model_id}", json.dumps({"weights_path": local_path})) - print(f"[Trainer] Published weight update signal to {num_subs} subscribers for version path: {local_path}") older = older_versions(local_path, SAMPLER_VERSIONS_KEPT) + if self.sampling is not None: + meta = await get_model_metadata(get_state_store(), model_id) + # A job with no metadata has no pool, so no request pins its versions. + pinned = await self.sampling.pinned_revisions(meta.pool(model_id)) if meta is not None else set() + older = [path for path in older if path not in pinned] for path in older: shutil.rmtree(path, ignore_errors=True) if older: diff --git a/src/server/vllm_options.py b/src/server/vllm_options.py index 606ccbb4e..63bf53358 100644 --- a/src/server/vllm_options.py +++ b/src/server/vllm_options.py @@ -1,4 +1,4 @@ -"""Engine and sampling options shared by the two vLLM sampler workers.""" +"""Engine and sampling options for the vLLM samplers.""" import json import os @@ -97,6 +97,19 @@ def detected_device_memory_bytes() -> int: return 0 +def sample_params(r: Mapping[str, Any]) -> dict[str, Any]: + """vLLM sampling params for a pool record, asking for the logprobs every result carries.""" + return {**r["sampling_params"], "logprobs": 0, "prompt_logprobs": 0 if r["include_prompt_logprobs"] else None} + + +def sample_result(sequences: list[tuple[list[int], list[float], str]], prompt_logprobs: list[float | None] | None) -> dict[str, Any]: + """The result a pool saves for a record, from each sequence's tokens, logprobs and stop reason.""" + result: dict[str, Any] = {"type": "sample", "sequences": [{"tokens": t, "logprobs": lp, "stop_reason": s} for t, lp, s in sequences]} + if prompt_logprobs is not None: + result["prompt_logprobs"] = prompt_logprobs + return result + + def engine_args(base_model: str, is_lora: bool, env: Mapping[str, str | None], sleep_mode: bool = False) -> dict[str, Any]: """vLLM engine settings for a sampler, the same for a `vllm serve` replica and an engine in the sampler's process.""" args: dict[str, Any] = {"enable_prefix_caching": True} diff --git a/src/server/vllm_sampler.py b/src/server/vllm_sampler.py index 19fe43c9f..4373efbe1 100644 --- a/src/server/vllm_sampler.py +++ b/src/server/vllm_sampler.py @@ -1,123 +1,41 @@ -"""Queue-driven vLLM sampler. One engine serves a model's queue; a request names -either a LoRA adapter (lora_id) or a whole-model weights version (weights_path).""" +"""A sampling pool replica with vLLM in this process. It serves its pool's stream itself, with no router, and when its +job shares a GPU with the trainer it takes a time-slicer turn around each batch, the way the trainer does. + + python -m server.vllm_sampler --pool +""" import argparse import asyncio -import hashlib +import logging import os import signal -import time -from collections.abc import Callable +import uuid from contextlib import asynccontextmanager -from itertools import groupby +from pathlib import Path from typing import Any -from opentelemetry import propagate, trace -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import BatchSpanProcessor -from vllm import SamplingParams -from vllm.config.weight_transfer import WeightTransferConfig -from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest -from vllm.engine.arg_utils import AsyncEngineArgs -from vllm.engine.async_llm_engine import AsyncLLMEngine -from vllm.lora.request import LoRARequest -from vllm.sampling_params import RequestOutputKind - from accel_timeslicer.time_slicer import TimeSlicerClient, time_slicer_client_from_env, time_slicing_enabled, workload_from_env from accel_timeslicer.workload import SAMPLER_CLAIM, WorkloadRef, local_workload_name -from server.store import RequestStore, StateStore, get_state_store, get_store -from server.vllm_options import cudagraph_capture_sizes, gpu_memory_utilization, sampler_batch_limits, split_stop, text_only_engine_kwargs - -tracer = trace.get_tracer("vllm.inference.worker") -SHUTDOWN_SENTINEL = "SHUTDOWN_SENTINEL" -TMP_DIR = os.getenv("OPEN_RL_TMP_DIR", "/tmp/open-rl") -READY_TTL_SECONDS = 3600 -ENGINE_POLL_SECONDS = 5 - - -def failed_response(message: str) -> dict[str, Any]: - return {"type": "RequestFailedResponse", "error_message": message} - - -def engine_kwargs_from_env(fft_enabled: bool) -> dict[str, Any]: - model_name = os.getenv("BASE_MODEL") or os.getenv("VLLM_MODEL") - if not model_name: - raise ValueError("BASE_MODEL or VLLM_MODEL is required") - batch_limits = sampler_batch_limits() - engine_kwargs = { - "model": model_name, - "enable_sleep_mode": fft_enabled, - "enable_lora": not fft_enabled, - "max_model_len": int(os.getenv("VLLM_MAX_MODEL_LEN", "8192")), - **batch_limits, - "gpu_memory_utilization": gpu_memory_utilization(), - "enable_prefix_caching": False, - "enforce_eager": os.getenv("VLLM_ENFORCE_EAGER", "0") == "1", - **text_only_engine_kwargs(), - } - if not engine_kwargs["enforce_eager"]: - engine_kwargs["compilation_config"] = {"cudagraph_capture_sizes": cudagraph_capture_sizes(batch_limits["max_num_seqs"])} - architecture = os.getenv("VLLM_ARCHITECTURE_OVERRIDE") - if architecture: - engine_kwargs["hf_overrides"] = {"architectures": [architecture]} - if fft_enabled: - engine_kwargs["weight_transfer_config"] = WeightTransferConfig(backend="delta_snapshot") - else: - engine_kwargs["max_loras"] = int(os.getenv("VLLM_MAX_LORAS", "8")) - engine_kwargs["max_lora_rank"] = int(os.getenv("VLLM_MAX_LORA_RANK", "64")) - return engine_kwargs - +from server.lora_snapshots import snapshot_root +from server.sampling_dispatcher import Dispatcher, RevisionUnloadable, delta_applies, failed, redis_client +from server.sampling_store import SamplingStore +from server.vllm_options import cudagraph_capture_sizes, engine_args, sample_params, sample_result, sampler_batch_limits -def resolve_lora_path(lora_id: str, lora_path: str | None) -> str: - """Find the PEFT directory holding adapter_config.json for this adapter.""" - if lora_path and os.path.exists(os.path.join(lora_path, "adapter_config.json")): - return lora_path +logger = logging.getLogger(__name__) - if lora_path: - # Check subfolders created by PEFT save_pretrained - base_candidates = [ - lora_id, - lora_id.rsplit("/", 1)[-1], - lora_id.split("://")[-1].split("/")[0] if "://" in lora_id else lora_id, - ] - for candidate in base_candidates: - candidate_path = os.path.join(lora_path, candidate) - if os.path.exists(os.path.join(candidate_path, "adapter_config.json")): - return candidate_path - - # Check auto-saved PEFT directory: TMP_DIR/peft// - base_id = lora_id.split("://")[-1].split("/")[0] if "://" in lora_id else lora_id - peft_dir = os.path.join(TMP_DIR, "peft", base_id, base_id) - if os.path.exists(os.path.join(peft_dir, "adapter_config.json")): - return peft_dir - - return lora_path or peft_dir - - -def lora_request_for(request: dict[str, Any]) -> LoRARequest | None: - """A LoRA request when the adapter exists on disk. Before the first save a - LoRA job samples from the base model, so a missing adapter is not an error.""" - lora_id = request.get("lora_id") - if not lora_id: - return None - path = resolve_lora_path(lora_id, request.get("lora_path")) - if not os.path.exists(os.path.join(path, "adapter_config.json")): - return None - lora_int_id = int(hashlib.md5(lora_id.encode("utf-8")).hexdigest(), 16) % (2**31 - 1) + 1 - return LoRARequest(lora_id, lora_int_id, path) +# The readiness probe passes once vLLM is up, and asleep when it shares its GPU. +READY_FILE = os.getenv("OPEN_RL_SAMPLER_READY_FILE", "/tmp/open-rl-sampler-ready") class Sampler: - """A vLLM engine that serves sampling requests and swaps its weights between them. - - weights_path is the whole-model version the engine holds; LoRA requests leave it - alone and attach their adapter per request. A failed update leaves partial - weights behind, so the sampler refuses work until the process is restarted. - """ - - def __init__(self, engine: AsyncLLMEngine) -> None: - self.engine = engine - self.weights_path: str | None = None + """A vLLM engine that answers a pool's records and swaps its weights between batches. `revision` is the weights + it holds, "" for the base model. A failed update leaves partial weights behind, so the sampler refuses work until + the process is restarted.""" + + def __init__(self, engine: Any, base_model: str) -> None: + self.engine, self.base_model = engine, base_model + self.lora_ids: dict[str, int] = {} + self.revision = "" self.update_failed = False async def wake(self) -> None: @@ -126,95 +44,92 @@ async def wake(self) -> None: async def sleep(self) -> None: await self.engine.sleep(level=1) - async def ensure_weights(self, weights_path: str | None) -> None: - """Make the engine hold weights_path; None keeps whatever it holds.""" + async def ensure_weights(self, revision: str) -> bool: + """Whether the engine holds `revision` after this, loading it if it can. Raises RevisionUnloadable for a delta + that does not apply on the weights it holds.""" if self.update_failed: - raise RuntimeError("Weight update failed; restart the sampler before serving") - if weights_path is None or weights_path == self.weights_path: - return - # No finally/resume on failure: partially updated weights must not be served. + return False + if revision == self.revision: + return True + if not revision: + return False # base weights cannot come back in place + try: + applies = delta_applies(revision, self.revision or "default") + except FileNotFoundError: + return False # the dispatcher fails the revision's records + if not applies: + raise RevisionUnloadable(revision, [self.revision or "default"]) + from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest + + # No resume on failure: partially updated weights must not be served. try: await self.engine.pause_generation(mode="wait", clear_cache=True) await self.engine.start_weight_update() - await self.engine.update_weights(WeightTransferUpdateRequest(update_info={"target_weights_path": weights_path})) - await self.engine.finish_weight_update(weight_version=weights_path) + await self.engine.update_weights(WeightTransferUpdateRequest(update_info={"target_weights_path": revision})) + await self.engine.finish_weight_update(weight_version=revision) await self.engine.reset_encoder_cache() await self.engine.resume_generation() except BaseException: self.update_failed = True raise - self.weights_path = weights_path - - async def generate(self, request: dict[str, Any]) -> dict[str, Any]: - request_id = request["request_id"] - prompt_token_ids = request.get("prompt_token_ids", []) - max_tokens = request.get("max_tokens", 20) - lora_request = lora_request_for(request) - stop_strings, stop_token_ids = split_stop(request.get("stop")) - sampling_params = SamplingParams( - n=request.get("num_samples", 1), - temperature=request.get("temperature", 1.0), - max_tokens=max_tokens, - stop=stop_strings, - stop_token_ids=stop_token_ids, - top_p=request.get("top_p", 1.0), - top_k=request.get("top_k", -1), - logprobs=1, # return logprobs for TITO RL - prompt_logprobs=1 if request.get("include_prompt_logprobs", False) else None, - output_kind=RequestOutputKind.FINAL_ONLY, - ) - - results_generator = self.engine.generate( - prompt={"prompt_token_ids": prompt_token_ids}, sampling_params=sampling_params, request_id=request_id, lora_request=lora_request - ) - - final_output = None - with tracer.start_as_current_span("vllm_generate_tokens") as span: - span.set_attribute("vllm.prompt_len", len(prompt_token_ids) if prompt_token_ids else 0) - span.set_attribute("vllm.max_tokens", max_tokens) - if lora_request is not None: - span.set_attribute("vllm.lora_id", lora_request.lora_name) - async for request_output in results_generator: - final_output = request_output - - outputs = final_output.outputs if final_output else [] - sequences_out = [] - for output in outputs: - generated_token_ids = list(output.token_ids) - logprobs = [] - if output.logprobs: - for idx, token_logprobs in enumerate(output.logprobs): - # token_logprobs is a dict of {token_id: Logprob} - token_id = generated_token_ids[idx] - if token_logprobs and token_id in token_logprobs: - logprob = token_logprobs[token_id].logprob - else: - logprob = -9999.0 - logprobs.append(logprob) - sequences_out.append({"tokens": generated_token_ids, "logprobs": logprobs, "stop_reason": output.finish_reason}) - - prompt_logprobs_out = None - if final_output and final_output.prompt_logprobs: - prompt_logprobs_out = [] - for idx, token_logprobs in enumerate(final_output.prompt_logprobs): - if token_logprobs is None: - prompt_logprobs_out.append(None) - else: - token_id = prompt_token_ids[idx] - if token_id in token_logprobs: - prompt_logprobs_out.append(token_logprobs[token_id].logprob) - else: - prompt_logprobs_out.append(None) - - res = {"sequences": sequences_out} - if prompt_logprobs_out is not None: - res["prompt_logprobs"] = prompt_logprobs_out - return res + self.revision = revision + return True + + async def generate(self, r: dict[str, Any]) -> dict[str, Any]: + """The result. A request vLLM refuses, such as a prompt past max_model_len or a token out of vocabulary, fails + at once, since sending it again would only take another turn from the trainer. So does a request in flight when + the engine dies, since a request that kills the engine would kill each restarted one too.""" + from vllm import SamplingParams + from vllm.exceptions import VLLMClientError + from vllm.sampling_params import RequestOutputKind + from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError + + lora_request = self.adapter(r) + if self.engine.errored: + raise RuntimeError("The engine is dead") # left pending for the restarted engine, since this one never saw it + prompt = list(r["prompt_token_ids"]) + request_id = f"{r['request_id']}-{uuid.uuid4().hex[:8]}" + final = None + try: + params = SamplingParams(**sample_params(r), output_kind=RequestOutputKind.FINAL_ONLY) + stream = self.engine.generate(prompt={"prompt_token_ids": prompt}, sampling_params=params, request_id=request_id, lora_request=lora_request) + async for output in stream: + final = output + except (TypeError, ValueError, VLLMClientError) as exc: + return failed(f"Sampler rejected the request: {exc}") + except EngineDeadError: + return failed("The sampler's vLLM engine died while serving this request") + except EngineGenerateError as exc: + return failed(f"Sampler failed the request: {exc.__cause__!r}") + sequences = [] + for choice in final.outputs: + tokens = list(choice.token_ids) + sequences.append((tokens, [step[token].logprob for token, step in zip(tokens, choice.logprobs or [])], choice.finish_reason)) + prompt_logprobs = None + if r["include_prompt_logprobs"]: + prompt_logprobs = [None if step is None else step[token].logprob for token, step in zip(prompt, final.prompt_logprobs)] + return sample_result(sequences, prompt_logprobs) + + def adapter(self, r: dict[str, Any]) -> Any: + """The record's LoRA, or None for base weights. vLLM loads an adapter inside the engine step, where a missing one + kills the engine, so a snapshot this pod's view of the shared volume does not show yet raises here instead, and + the record is sent again.""" + if r["model"] == self.base_model: + return None + from vllm.lora.request import LoRARequest + + path = snapshot_root() / r["model"] + if not (path / "adapter_config.json").is_file(): + os.listdir(path.parent) # an NFS client keeps a failed lookup until it lists the directory + if not (path / "adapter_config.json").is_file(): + raise FileNotFoundError(f"No LoRA snapshot at {path}") + lora_id = self.lora_ids.setdefault(r["model"], len(self.lora_ids) + 1) + return LoRARequest(r["model"], lora_id, str(path)) @asynccontextmanager async def holding_gpu(time_slicer: TimeSlicerClient | None, workload: WorkloadRef | None, sampler: Sampler | None = None): - """Hold the time-slicer slot for the body, waking the sampler inside it. Without a slicer the GPU is ours.""" + """Hold the time-slicer turn for the body, waking the sampler inside it. Without a slicer the GPU is ours.""" if time_slicer is None: yield return @@ -228,174 +143,86 @@ async def holding_gpu(time_slicer: TimeSlicerClient | None, workload: WorkloadRe await sampler.sleep() -async def process_batch(sampler: Sampler, store: RequestStore, requests: list[dict[str, Any]]) -> None: - """Serve consecutive requests that share a weights_path together, updating weights between groups.""" - for weights_path, group in groupby(requests, key=lambda req: req.get("weights_path")): - batch = list(group) - try: - await sampler.ensure_weights(weights_path) - except Exception as exc: - for request in batch: - await store.set_result(request["request_id"], failed_response(f"vLLM weight update failed: {exc}")) - continue - await serve_group(sampler, store, batch) - - -async def serve_group(sampler: Sampler, store: RequestStore, batch: list[dict[str, Any]]) -> None: - """A dead EngineCore leaves generate() hanging, so the wait polls for it and fails what is left.""" - tasks = {asyncio.create_task(process_request(sampler, store, req)): req for req in batch} - pending = set(tasks) - while pending: - _, pending = await asyncio.wait(pending, timeout=ENGINE_POLL_SECONDS) - if pending and sampler.engine.errored: - for task in pending: - task.cancel() - await fail_requests(store, [tasks[task] for task in pending], "vLLM engine died during the batch") - break - if sampler.engine.errored: - raise RuntimeError("vLLM engine is dead") - for task in tasks: - task.result() - - -async def fail_requests(store: RequestStore, requests: list[dict[str, Any]], error: Any) -> None: - """Best effort, so clients see the failure instead of waiting forever.""" - for request in requests: - try: - await store.set_result(request["request_id"], failed_response(f"vLLM Worker Error: {error}")) - except Exception as exc: - print(f"[vLLM Worker] Could not fail request {request['request_id']}: {exc}") - - -async def process_request(sampler: Sampler, store: RequestStore, request: dict[str, Any]) -> None: - with tracer.start_as_current_span("process_sampling_request", context=propagate.extract(request.get("trace_context", {}))): - try: - result = await sampler.generate(request) - result["type"] = "sample" - except Exception as exc: - result = failed_response(f"vLLM Worker Error: {exc}") - await store.set_result(request["request_id"], result) - +class TimeSlicedDispatcher(Dispatcher): + """A pool's dispatcher when its sampler shares the GPU with the trainer. Each batch runs inside one time-slicer turn.""" -async def mark_ready(state: StateStore, model_id: str) -> None: - """The gateway waits on this key. A shared LoRA sampler outlives the TTL, so the loop refreshes it.""" - await state.set_value(f"open_rl:sampler_ready:{model_id}", "1", ttl_seconds=READY_TTL_SECONDS) + def __init__( + self, store: SamplingStore, pool: str, sampler: Sampler, concurrency: int, time_slicer: TimeSlicerClient, workload: WorkloadRef + ) -> None: + super().__init__(store, pool, sampler, concurrency) + self.time_slicer, self.workload = time_slicer, workload - -async def serve( - model_id: str, - store: RequestStore, - make_engine: Callable[[], AsyncLLMEngine], - time_slicer: TimeSlicerClient | None = None, - workload: WorkloadRef | None = None, -) -> None: - """Build the engine, then serve the model's queue until the shutdown sentinel arrives. - - Queue and store errors are retried. A dead engine or a process the time slicer - could not park raises, so the worker exits and is restarted. - """ - async with holding_gpu(time_slicer, workload): - sampler = Sampler(make_engine()) - if time_slicer is not None: - await sampler.sleep() # give the memory back before releasing the slot - state = get_state_store() - try: - try: - ready_at = float("-inf") - while True: - if time_slicer is not None and time_slicer.faulted: - raise RuntimeError(f"Time slicer could not park this process: {time_slicer.faulted}") - requests = [] - started = False - try: - if time.monotonic() - ready_at > 60: - await mark_ready(state, model_id) - ready_at = time.monotonic() - batch = await store.get_sampling_requests_for_model(model_id) - if not batch: - await asyncio.sleep(0.05) - continue - shutdown = any(req.get("request_id") == SHUTDOWN_SENTINEL for req in batch) - requests = [req for req in batch if req.get("request_id") != SHUTDOWN_SENTINEL] - if requests: - async with holding_gpu(time_slicer, workload, sampler): - started = True - await process_batch(sampler, store, requests) - if shutdown: - return - except Exception as exc: - if sampler.engine.errored: - raise - print(f"[vLLM Worker] Sampling loop error, retrying: {exc}") - if not started: - await fail_requests(store, requests, exc) - await asyncio.sleep(1) - finally: - await state.delete_values(f"open_rl:sampler_ready:{model_id}") - finally: - sampler.engine.shutdown() - - -async def serve_time_sliced(model_id: str, store: RequestStore, make_engine: Callable[[], AsyncLLMEngine]) -> None: - """An FFT sampler shares its GPU with the trainer, so it is registered with the time slicer while it serves.""" - time_slicer = time_slicer_client_from_env() - workload = workload_from_env(os.getpid(), name=local_workload_name("sampler", model_id), claim=SAMPLER_CLAIM) - try: + async def dispatch(self) -> None: + await self.read() + if not self.to_send(): + return + async with holding_gpu(self.time_slicer, self.workload, self.sampler): + await self.read() # requests that arrived while the turn was pending join it + if await self.ready(): + await asyncio.gather(*(self.send(r) for r in self.to_send()[: self.concurrency])) + + +def make_engine(sleep_mode: bool, is_lora: bool) -> Any: + from vllm.config.weight_transfer import WeightTransferConfig + from vllm.engine.arg_utils import AsyncEngineArgs + from vllm.engine.async_llm_engine import AsyncLLMEngine + + args = engine_args(os.environ["BASE_MODEL"], is_lora, os.environ, sleep_mode=sleep_mode) + # This process knows its GPU, so it sizes vLLM's batch for it. A vllm serve replica leaves that to vLLM. + args.update(sampler_batch_limits()) + if not args.get("enforce_eager"): + args["compilation_config"] = {"cudagraph_capture_sizes": cudagraph_capture_sizes(args["max_num_seqs"])} + if transfer := args.get("weight_transfer_config"): + args["weight_transfer_config"] = WeightTransferConfig(**transfer) + return AsyncLLMEngine.from_engine_args(AsyncEngineArgs(model=os.environ["BASE_MODEL"], **args)) + + +async def run_sampling_worker(pool: str) -> None: + """Boots vLLM, inside a turn when it shares its GPU, then serves the pool until the engine dies, a weight update + breaks off, or the slicer cannot park this process. Each of those exits with an error, so the worker restarts.""" + time_slicer = time_slicer_client_from_env() if time_slicing_enabled() else None + workload = workload_from_env(os.getpid(), name=local_workload_name("sampler", pool), claim=SAMPLER_CLAIM) if time_slicer else None + store = SamplingStore(redis_client(os.environ["REDIS_URL"])) + if time_slicer is not None: await time_slicer.register(workload) - try: - await serve(model_id, store, make_engine, time_slicer, workload) - finally: - # A faulted process still holds the GPU; the grant moves on once it exits. - if not time_slicer.faulted: - await time_slicer.unregister(workload) - finally: - await time_slicer.close() - - -async def run_sampling_worker(model_id: str) -> None: - fft_enabled = os.getenv("OPEN_RL_ENABLE_FFT", "").lower() == "true" - engine_kwargs = engine_kwargs_from_env(fft_enabled) - - def make_engine() -> AsyncLLMEngine: - return AsyncLLMEngine.from_engine_args(AsyncEngineArgs(**engine_kwargs)) - - loop = asyncio.get_running_loop() - task = asyncio.current_task() - assert task is not None - installed_signals = [] + sampler = None try: - for sig in (signal.SIGTERM, signal.SIGINT): - try: - loop.add_signal_handler(sig, task.cancel) - installed_signals.append(sig) - except NotImplementedError: - break - # An exclusive sampler is alone on its GPU, so it serves awake like a LoRA one. - if fft_enabled and time_slicing_enabled(): - await serve_time_sliced(model_id, get_store(), make_engine) + async with holding_gpu(time_slicer, workload): + is_lora = os.environ.get("OPEN_RL_FINE_TUNING_TYPE") == "lora" + sampler = Sampler(make_engine(sleep_mode=time_slicer is not None, is_lora=is_lora), os.environ["BASE_MODEL"]) + if time_slicer is not None: + await sampler.sleep() # give the memory back before releasing the turn + Path(READY_FILE).touch() + concurrency = int(os.getenv("OPEN_RL_DISPATCH_CONCURRENCY", "256")) + if time_slicer is not None: + dispatcher = TimeSlicedDispatcher(store, pool, sampler, concurrency, time_slicer, workload) else: - await serve(model_id, get_store(), make_engine) - except asyncio.CancelledError: - pass # serve has completed cleanup before cancellation reaches here. + dispatcher = Dispatcher(store, pool, sampler, concurrency) + task = asyncio.create_task(dispatcher.run()) + for sig in (signal.SIGTERM, signal.SIGINT): + asyncio.get_running_loop().add_signal_handler(sig, task.cancel) + while not task.done(): + sampler_failed = sampler.engine.errored or sampler.update_failed + time_slicer_failed = time_slicer is not None and time_slicer.faulted + if sampler_failed or time_slicer_failed: + task.cancel() + raise RuntimeError(f"The sampler for pool {pool} cannot serve: {sampler_failed=} {time_slicer_failed=}") + await asyncio.wait({task}, timeout=1) finally: - for sig in installed_signals: - loop.remove_signal_handler(sig) + # A faulted process still holds the GPU; the turn moves on once it exits and its connection closes. + if time_slicer is not None and not time_slicer.faulted: + await time_slicer.unregister(workload) + if sampler is not None: + sampler.engine.shutdown() + await store.redis.aclose() 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") + parser = argparse.ArgumentParser(description="Open-RL sampling pool replica with vLLM in this process") + parser.add_argument("--pool", required=True, help="The sampling pool whose stream this replica serves") args = parser.parse_args() - provider = TracerProvider() - trace.set_tracer_provider(provider) - try: - if os.getenv("ENABLE_GCP_TRACE", "0") == "1": - from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter - - provider.add_span_processor(BatchSpanProcessor(CloudTraceSpanExporter())) - asyncio.run(run_sampling_worker(args.model_id)) - finally: - provider.shutdown() + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") + asyncio.run(run_sampling_worker(args.pool)) if __name__ == "__main__": diff --git a/src/server/worker_manager.py b/src/server/worker_manager.py index 3741e80b5..a9ed455aa 100644 --- a/src/server/worker_manager.py +++ b/src/server/worker_manager.py @@ -13,9 +13,11 @@ from pathlib import Path from typing import Any, Protocol -from accel_timeslicer.workload import SAMPLER_TIME_SLICE_GROUP, TRAINER_TIME_SLICE_GROUP, workload_job_id +from accel_timeslicer.time_slicer import time_slicing_enabled +from accel_timeslicer.workload import TRAINER_TIME_SLICE_GROUP, workload_job_id from server.estimator import footprint from server.model_metadata import TrainingModelMetadata, decode_model_metadata +from server.sampling_dispatcher import tmp_root from server.store import get_state_store PROJECT_DIR = Path(__file__).resolve().parents[2] @@ -108,7 +110,6 @@ def worker_env(meta: TrainingModelMetadata, base_model: str, runtime: str, is_lo if meta.trainer_cp > 1: env["OPEN_RL_AUTOMODEL_CP"] = str(meta.trainer_cp) else: - env["OPEN_RL_MODEL_ID"] = runtime env["VLLM_SERVER_DEV_MODE"] = "1" env["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" # The text-only class skips the multimodal front end vLLM otherwise builds @@ -121,8 +122,7 @@ def worker_env(meta: TrainingModelMetadata, base_model: str, runtime: str, is_lo return env -def worker_module(role: str) -> str: - return "server.training_requests_processor" if role == "trainer" else "server.vllm_sampler" +TRAINER_MODULE = "server.training_requests_processor" def worker_args(runtime: str, role: str, is_lora: bool) -> list[str]: @@ -182,20 +182,30 @@ def ensure(self, model_id: str, role: str) -> None: return base_model = base_model_of(meta, runtime) env = {**os.environ, **worker_env(meta, base_model, runtime, is_lora, role)} - env["OPEN_RL_TIME_SLICE_JOB_ID"] = workload_job_id(role, runtime) - env["OPEN_RL_TIME_SLICE_GROUP"] = TRAINER_TIME_SLICE_GROUP if role == "trainer" else SAMPLER_TIME_SLICE_GROUP gpus = os.getenv("TRAINER_CUDA_VISIBLE_DEVICES" if role == "trainer" else "SAMPLER_CUDA_VISIBLE_DEVICES") if gpus: env["CUDA_VISIBLE_DEVICES"] = gpus - extras = ["gpu"] if role == "trainer" else ["gpu", "vllm"] - command = python_command(extras, worker_module(role), worker_args(runtime, role, is_lora)) - log_dir = Path(os.getenv("OPEN_RL_TMP_DIR", "/tmp")) - log_dir.mkdir(parents=True, exist_ok=True) - with open(log_dir / f"{role}_{runtime.replace('/', '_')}.log", "a") as log: - self.processes[(role, runtime)] = subprocess.Popen( - command, cwd=self.project_dir, env=env, stdout=log, stderr=subprocess.STDOUT, start_new_session=True - ) - self.started[(role, runtime)] = time.time() + if role == "trainer": + env["OPEN_RL_TIME_SLICE_JOB_ID"] = workload_job_id(role, runtime) + env["OPEN_RL_TIME_SLICE_GROUP"] = TRAINER_TIME_SLICE_GROUP + self.spawn(role, runtime, python_command(["gpu"], TRAINER_MODULE, worker_args(runtime, role, is_lora)), env) + return + env["OPEN_RL_TMP_DIR"] = tmp_root() + if meta.shares_gpu() and time_slicing_enabled(): + # The sampler takes turns on the one GPU with its trainer, so it is in the trainer's claim. + env["OPEN_RL_TIME_SLICE_JOB_ID"] = workload_job_id(role, runtime) + env["OPEN_RL_TIME_SLICE_GROUP"] = TRAINER_TIME_SLICE_GROUP + else: + env["OPEN_RL_TIME_SLICING"] = "off" + self.spawn(role, runtime, python_command(["gpu", "vllm"], "server.vllm_sampler", ["--pool", meta.pool(runtime)]), env) + + def spawn(self, role: str, runtime: str, command: list[str], env: dict[str, str]) -> None: + log_dir = Path(os.getenv("OPEN_RL_TMP_DIR", "/tmp")) + log_dir.mkdir(parents=True, exist_ok=True) + with open(log_dir / f"{role}_{runtime.replace('/', '_')}.log", "a") as log: + process = subprocess.Popen(command, cwd=self.project_dir, env=env, stdout=log, stderr=subprocess.STDOUT, start_new_session=True) + self.processes[(role, runtime)] = process + self.started[(role, runtime)] = time.time() def release(self, model_id: str) -> None: try: diff --git a/tests/test_api_server_paths.py b/tests/test_api_server_paths.py index 51a7d9b54..06283b117 100644 --- a/tests/test_api_server_paths.py +++ b/tests/test_api_server_paths.py @@ -570,6 +570,13 @@ def test_sampler_replicas_need_a_server_that_launches_pods(self) -> None: self.assertEqual(response.status_code, 400) self.assertIn("openrl.sampler_replicas above 1", response.json()["error"]) + def test_an_fft_job_sharing_its_trainers_gpu_cannot_ask_for_more_samplers(self) -> None: + fft = {"openrl.fine_tuning_type": "full", "openrl.sampler_replicas": "4"} + with patch.dict(os.environ, {"OPEN_RL_ENABLE_FFT": "true"}): + self.assertEqual(self.post("create_model", {"base_model": "m", "user_metadata": fft}).status_code, 400) + exclusive = self.post("create_model", {"base_model": "m", "user_metadata": {**fft, "openrl.exclusive": "true"}}) + self.assertEqual(self.metadata(exclusive.json()["request_id"])["sampler_replicas"], 4) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_lora_snapshots.py b/tests/test_lora_snapshots.py index 4543ac66a..cdf2d51b6 100644 --- a/tests/test_lora_snapshots.py +++ b/tests/test_lora_snapshots.py @@ -17,6 +17,8 @@ def test_a_frozen_snapshot_keeps_its_weights_is_never_overwritten_and_names_the_ (live / "adapter_model.safetensors").write_bytes(b"version-one") ref = "tinker://job/sampler_weights/sampler-1" freeze_adapter("job", "Qwen/Qwen2.5-0.5B", ref) + # The same weights saved under the same name again, with no training between, are already frozen. + freeze_adapter("job", "Qwen/Qwen2.5-0.5B", ref) (live / "adapter_model.safetensors").write_bytes(b"version-two") with self.assertRaisesRegex(FileExistsError, "new name or sequence ID"): freeze_adapter("job", "Qwen/Qwen2.5-0.5B", ref) diff --git a/tests/test_openrl_status.py b/tests/test_openrl_status.py index d23ec57d8..c5f70c122 100644 --- a/tests/test_openrl_status.py +++ b/tests/test_openrl_status.py @@ -10,6 +10,8 @@ from collections import defaultdict from unittest.mock import patch +from kubernetes import client + from server import api_server from server.scheduler_worker_manager import SchedulerWorkerManager, worker_state from server.session_registry import SessionRegistry @@ -166,18 +168,6 @@ async def test_a_trainer_is_starting_until_it_announces_itself(self): (trainer,) = (await self.get(f"openrl/models/{model_id}")).json()["workers"] self.assertEqual(trainer["state"], "ready") - async def test_a_sampler_is_starting_until_it_announces_itself(self): - _, model_id = await self.create_model() - self.manager.ensured.append(("test-base", "sampler")) - body = (await self.get(f"openrl/models/{model_id}")).json() - sampler = next(w for w in body["workers"] if w["role"] == "sampler") - self.assertEqual(sampler["state"], "starting") - self.assertIn("loading the model", sampler["detail"]) - await self.state.set_value("open_rl:sampler_ready:test-base", "1") - body = (await self.get(f"openrl/models/{model_id}")).json() - sampler = next(w for w in body["workers"] if w["role"] == "sampler") - self.assertEqual(sampler["state"], "ready") - async def test_a_released_model_says_so_in_its_summary(self): session, model_id = await self.create_model({"openrl.fine_tuning_type": "full"}) await self.state.delete_values(f"open_rl:session:{session}") @@ -237,11 +227,31 @@ def workload(name: str, role: str, owner: str, model_id: str, status: dict, crea return {"metadata": metadata, "spec": {"role": role, "ownerID": owner, "modelID": model_id}, "status": status} +class FakeCoreApi: + def __init__(self, *pods): + self.pods, self.selectors = list(pods), [] + + def list_namespaced_pod(self, namespace: str, label_selector: str = "", _request_timeout=None): + self.selectors.append(label_selector) + return client.V1PodList(items=self.pods) + + +def sampler_pod(name: str, ready: bool, waiting: str | None = None, restarts: int = 0): + if waiting: + state = client.V1ContainerState(waiting=client.V1ContainerStateWaiting(reason=waiting)) + else: + state = client.V1ContainerState(running=client.V1ContainerStateRunning()) + container = client.V1ContainerStatus(name="worker", image="server", image_id="", ready=ready, restart_count=restarts, state=state) + condition = client.V1PodCondition(type="Ready", status="True" if ready else "False") + status = client.V1PodStatus(phase="Running", conditions=[condition], container_statuses=[container]) + return client.V1Pod(metadata=client.V1ObjectMeta(name=name), status=status) + + class SchedulerViewTest(unittest.TestCase): - """The cluster side of the model view: Workloads only. The pod is the scheduler's.""" + """The cluster side of the model view, from Workloads and a running sampler's pod readiness.""" def setUp(self): - self.enterContext(patch.dict(os.environ, {"REDIS_URL": "redis://test"})) + self.enterContext(patch.dict(os.environ, {"REDIS_URL": "redis://test", "OPEN_RL_TMP_DIR": "/mnt/shared/open-rl"})) self.custom = FakeCustomObjectsApi() self.custom.existing = { "lora-base-0-trainer": workload( @@ -276,6 +286,27 @@ def test_the_view_is_cached_briefly(self): self.manager.view_cache_seconds = 0 self.assertEqual(self.manager.inspect_owner("base"), []) + def test_a_running_sampler_is_ready_only_once_its_pod_is_ready(self): + # The Workload reads Running while vLLM loads and while the container crash-loops. + running = {"phase": "Running", "podName": "orw-s", "nodeName": "node-b"} + self.custom.existing["lora-base-0-sampler"] = workload("lora-base-0-sampler", "sampler", "base", "base", running) + crash_looping = sampler_pod("orw-s", ready=False, waiting="CrashLoopBackOff", restarts=4) + cases = [ + ([sampler_pod("orw-s", ready=False)], "starting", "Running on node-b, loading the model"), + ([crash_looping], "starting", "Running on node-b, CrashLoopBackOff, restart count 4"), + ([sampler_pod("orw-s", ready=True)], "ready", "Running on node-b"), + ([], "starting", "Running on node-b, loading the model"), + ] + for pods, state, detail in cases: + with self.subTest(detail=detail): + core = FakeCoreApi(*pods) + manager = SchedulerWorkerManager(custom_api=self.custom, core_api=core) + by_role = {w["role"]: w for w in manager.inspect_owner("base")} + self.assertEqual((by_role["sampler"]["state"], by_role["sampler"]["detail"]), (state, detail)) + # One pod list per view, and the trainer is still read from its Workload alone. + self.assertEqual(core.selectors, ["app.kubernetes.io/component=sampler"]) + self.assertEqual(by_role["trainer"]["state"], "ready") + if __name__ == "__main__": unittest.main() diff --git a/tests/test_pool_accept.py b/tests/test_pool_accept.py index f2da4fa9a..f8c133fcd 100644 --- a/tests/test_pool_accept.py +++ b/tests/test_pool_accept.py @@ -1,43 +1,48 @@ -import json +import asyncio import os import tempfile +import unittest import uuid from pathlib import Path -from unittest.mock import patch +from unittest.mock import Mock, patch import redis -from fastapi import HTTPException from server import api_server from server.lora_snapshots import freeze_adapter, snapshot_name from server.model_metadata import TrainingModelMetadata, persist_model_metadata from server.sampling_store import SamplingStore from server.store import RedisStateStore, RedisStore +from server.worker_manager import default_metadata +from tests.api_client import asgi_client from tests.redis_server import RedisTestCase class PoolAcceptTest(RedisTestCase): async def asyncSetUp(self) -> None: await super().asyncSetUp() - self.tmp, self.base, self.pool = self.enterContext(tempfile.TemporaryDirectory()), f"org/m-{uuid.uuid4().hex[:8]}", f"pool-{uuid.uuid4().hex}" - self.enterContext(patch.dict(os.environ, {"OPEN_RL_TMP_DIR": self.tmp, "OPEN_RL_SAMPLER_POOLS": json.dumps({self.base: self.pool})})) - self.state = RedisStateStore(self.redis, redis.Redis.from_url(self.url, decode_responses=True)) - for name, value in [("store", RedisStore(self.redis)), ("state", self.state), ("TMP_DIR", self.tmp)]: + self.tmp, self.base = self.enterContext(tempfile.TemporaryDirectory()), f"org/m-{uuid.uuid4().hex[:8]}" + self.enterContext(patch.dict(os.environ, {"OPEN_RL_TMP_DIR": self.tmp, "SAMPLING_BACKEND": "vllm"})) + self.state, self.store = RedisStateStore(self.redis, redis.Redis.from_url(self.url, decode_responses=True)), RedisStore(self.redis) + patched = [("store", self.store), ("state", self.state), ("TMP_DIR", self.tmp), ("worker_manager", Mock()), ("SAVED_WEIGHTS_WAIT_SECONDS", 0.3)] + for name, value in patched: self.enterContext(patch.object(api_server, name, value)) + self.client = await self.enterAsyncContext(asgi_client()) - async def accept(self, model_id: str | None = None, session: str | None = None, seq_id: int = 0) -> str: - req = api_server.AsampleRequest(model_id=model_id, sampling_session_id=session, seq_id=seq_id, prompt={"chunks": [{"tokens": [1]}]}) - return await api_server.accept_pool_sample(req, [1]) + async def sample(self, seq_id: int = 0, session: str | None = None, **create) -> dict: + session = session or (await self.client.post("/api/v1/create_sampling_session", json=create)).json()["sampling_session_id"] + body = {"sampling_session_id": session, "seq_id": seq_id, "prompt": {"chunks": [{"tokens": [1]}]}} + response = await self.client.post("/api/v1/asample", json=body) + return {"session": session, "status": response.status_code, **response.json()} async def pending(self, pool: str, field: str) -> list: return [r[field] for r in await SamplingStore(self.redis).pending(pool)] async def test_a_retried_seq_id_is_one_request_and_another_session_gets_its_own(self) -> None: - session, other = f"{self.base}@{uuid.uuid4().hex}", f"{self.base}@{uuid.uuid4().hex}" - first = await self.accept(session=session, seq_id=7) - retry, second = await self.accept(session=session, seq_id=7), await self.accept(session=other, seq_id=7) - self.assertEqual(retry, first) - self.assertEqual(await self.pending(self.pool, "request_id"), [first, second]) + first = await self.sample(7, base_model=self.base) + retry, other = await self.sample(7, first["session"]), await self.sample(7, base_model=self.base) + self.assertEqual(retry["request_id"], first["request_id"]) + self.assertEqual(await self.pending(default_metadata(self.base).pool(self.base), "request_id"), [first["request_id"], other["request_id"]]) async def test_fft_samples_need_saved_weights_and_name_their_immutable_revision(self) -> None: model_id = uuid.uuid4().hex @@ -45,23 +50,56 @@ async def test_fft_samples_need_saved_weights_and_name_their_immutable_revision( saved = Path(self.tmp, "sampler_full", model_id, "sampler_weights", "sampler-1") saved.mkdir(parents=True) (saved / "metadata.json").write_text("{}") - with self.assertRaises(HTTPException) as missing: - await self.accept(f"tinker://{model_id}/sampler_weights/sampler-9") - self.assertEqual( - (missing.exception.status_code, missing.exception.detail), (400, f"No saved sampler weights for tinker://{model_id}/sampler_weights/sampler-9") - ) - await self.accept(f"tinker://{model_id}") + missing = await self.sample(model_path=f"tinker://{model_id}/sampler_weights/sampler-9") + self.assertEqual((missing["status"], missing["error"]), (400, f"No saved sampler weights for tinker://{model_id}/sampler_weights/sampler-9")) + await self.sample(model_path=f"tinker://{model_id}") + self.assertEqual(await self.pending(f"sampler-{model_id}", "revision"), [str(saved)]) + + async def test_weights_saved_while_the_sample_waits_are_accepted(self) -> None: + # A save from the trainer's node can show up on NFS a moment after the client hears it finished. + self.enterContext(patch.object(api_server, "SAVED_WEIGHTS_WAIT_SECONDS", 5.0)) + model_id = uuid.uuid4().hex + await persist_model_metadata(self.state, model_id, TrainingModelMetadata(base_model=self.base, fine_tuning_type="full")) + saved = Path(self.tmp, "sampler_full", model_id, "sampler_weights", "sampler-1") + + def finish_save() -> None: + saved.mkdir(parents=True) + (saved / "metadata.json").write_text("{}") + + asyncio.get_running_loop().call_later(0.5, finish_save) + accepted = await self.sample(model_path=f"tinker://{model_id}/sampler_weights/sampler-1") + self.assertEqual(accepted["status"], 200) self.assertEqual(await self.pending(f"sampler-{model_id}", "revision"), [str(saved)]) async def test_lora_samples_need_a_frozen_snapshot(self) -> None: model_id = uuid.uuid4().hex - await persist_model_metadata(self.state, model_id, TrainingModelMetadata(base_model=self.base, fine_tuning_type="lora")) + meta = TrainingModelMetadata(base_model=self.base, fine_tuning_type="lora") + await persist_model_metadata(self.state, model_id, meta) ref = f"tinker://{model_id}/sampler_weights/sampler-1" - with self.assertRaises(HTTPException) as refused: - await self.accept(ref) - self.assertEqual((refused.exception.status_code, refused.exception.detail), (400, f"No saved LoRA sampler weights for {ref}")) + refused = await self.sample(model_path=ref) + self.assertEqual((refused["status"], refused["error"]), (400, f"No saved LoRA sampler weights for {ref}")) Path(self.tmp, "peft", model_id, model_id).mkdir(parents=True) Path(self.tmp, "peft", model_id, model_id, "adapter_config.json").write_text("{}") freeze_adapter(model_id, self.base, ref) - await self.accept(ref) - self.assertEqual(await self.pending(self.pool, "model"), [snapshot_name(ref)]) + self.assertEqual((await self.sample(1, refused["session"]))["status"], 200) + self.assertEqual(await self.pending(meta.pool(model_id), "model"), [snapshot_name(ref)]) + + async def test_pool_names_are_valid_kubernetes_names(self) -> None: + pools = {default_metadata("Qwen/Qwen3.5-9B").pool("Qwen/Qwen3.5-9B"), TrainingModelMetadata(base_model="Qwen/Qwen3-8B").pool("job")} + self.assertTrue(all(len(p) <= 63 and p == p.lower() and "/" not in p and "." not in p for p in pools), pools) + + +class PoolsNeedRedisTest(unittest.IsolatedAsyncioTestCase): + def env(self, **values: str): + unset = ("REDIS_URL", "BASE_MODEL", "SAMPLING_BACKEND", "OPEN_RL_ENABLE_FFT", "OPEN_RL_WORKER_MANAGER") + return patch.dict(os.environ, {**{k: v for k, v in os.environ.items() if k not in unset}, **values}, clear=True) + + async def test_a_server_that_samples_through_pools_will_not_start_without_redis(self) -> None: + with self.env(), self.assertRaisesRegex(RuntimeError, "Sampling pools need REDIS_URL"): + async with api_server.lifespan(api_server.app): + pass + + def test_torch_sampling_in_one_process_and_pools_with_redis_pass_the_check(self) -> None: + for values in ({"BASE_MODEL": "org/m"}, {"REDIS_URL": "redis://127.0.0.1:6379/0"}): + with self.subTest(**values), self.env(**values): + api_server.check_sampler_backend() diff --git a/tests/test_redis_store.py b/tests/test_redis_store.py index 92e905b11..003024eaa 100644 --- a/tests/test_redis_store.py +++ b/tests/test_redis_store.py @@ -159,20 +159,6 @@ def setUp(self) -> None: self.store = InMemoryStore() self.state = InMemoryStateStore() - async def test_sampling_queue_put_and_get(self) -> None: - req1 = {"model_id": "base-m1", "request_id": "r1"} - req2 = {"model_id": "base-m1", "request_id": "r2"} - await self.store.put_sampling_request(req1) - await self.store.put_sampling_request(req2) - - batch = await self.store.get_sampling_requests_for_model("base-m1") - self.assertEqual(len(batch), 2) - self.assertEqual(batch[0]["request_id"], "r1") - self.assertEqual(batch[1]["request_id"], "r2") - - empty_batch = await self.store.get_sampling_requests_for_model("base-m1") - self.assertEqual(empty_batch, []) - async def test_get_requests_rotates_between_tenants(self) -> None: for i in range(3): await self.store.put_request({"model_id": "tenant-a", "request_id": f"a{i}"}, active_set_id="base-1") diff --git a/tests/test_sampler_patch.py b/tests/test_sampler_patch.py deleted file mode 100644 index a27bc16f0..000000000 --- a/tests/test_sampler_patch.py +++ /dev/null @@ -1,256 +0,0 @@ -"""Weight-version ordering, failure handling, and resource cleanup for the vLLM sampler.""" - -import asyncio -import os -import tempfile -import unittest -from contextlib import asynccontextmanager -from types import SimpleNamespace -from unittest.mock import AsyncMock, Mock, patch - -from accel_timeslicer.workload import WorkloadRef -from server.vllm_sampler import SHUTDOWN_SENTINEL, Sampler, process_batch, serve, serve_time_sliced - - -def make_engine(on_generate=None): - engine = AsyncMock() - engine.shutdown = Mock() - engine.errored = False - - def generate(prompt, sampling_params, request_id, lora_request): - async def stream(): - if on_generate is not None: - await on_generate(request_id) - yield SimpleNamespace(outputs=[SimpleNamespace(token_ids=[4], logprobs=None, finish_reason="length")], prompt_logprobs=None) - - return stream() - - engine.generate = Mock(side_effect=generate) - return engine - - -class SamplerBatchTest(unittest.IsolatedAsyncioTestCase): - async def asyncSetUp(self): - self.store = AsyncMock() - self.engine = make_engine() - self.sampler = Sampler(self.engine) - - def results(self): - return {call.args[0]: call.args[1] for call in self.store.set_result.call_args_list} - - async def test_groups_drain_before_switching_weights(self): - active = set() - seen = [] - - async def update(update_request): - self.assertFalse(active) - seen.append(("update", update_request.update_info["target_weights_path"])) - - async def generate(request_id): - active.add(request_id) - await asyncio.sleep(0) - seen.append(("generate", request_id)) - active.remove(request_id) - - self.engine.update_weights.side_effect = update - self.engine.generate.side_effect = make_engine(generate).generate - requests = [{"request_id": path + str(i), "weights_path": path} for i, path in enumerate(["a", "a", "b", "a"])] - await process_batch(self.sampler, self.store, requests) - self.assertEqual( - seen, [("update", "a"), ("generate", "a0"), ("generate", "a1"), ("update", "b"), ("generate", "b2"), ("update", "a"), ("generate", "a3")] - ) - self.assertEqual(self.sampler.weights_path, "a") - - async def test_unchanged_weights_skip_update(self): - for request_id in ("1", "2"): - await process_batch(self.sampler, self.store, [{"request_id": request_id, "weights_path": "a"}]) - self.engine.update_weights.assert_awaited_once() - self.engine.pause_generation.assert_awaited_once_with(mode="wait", clear_cache=True) - self.engine.finish_weight_update.assert_awaited_once_with(weight_version="a") - self.engine.resume_generation.assert_awaited_once() - self.engine.sleep.assert_not_called() - self.engine.wake_up.assert_not_called() - - async def test_failed_update_poisons_sampler_without_committing(self): - self.engine.update_weights.side_effect = RuntimeError("invalid patch") - await process_batch(self.sampler, self.store, [{"request_id": "bad", "weights_path": "a"}]) - self.assertIn("invalid patch", self.results()["bad"]["error_message"]) - self.engine.generate.assert_not_called() - self.engine.finish_weight_update.assert_not_called() - self.engine.resume_generation.assert_not_called() - self.assertIsNone(self.sampler.weights_path) - await process_batch(self.sampler, self.store, [{"request_id": "no-path"}]) - self.assertIn("restart", self.results()["no-path"]["error_message"]) - - async def test_generation_failure_is_reported_without_poisoning_weights(self): - self.engine.generate.side_effect = RuntimeError("generation error") - await process_batch(self.sampler, self.store, [{"request_id": "1", "weights_path": "a"}]) - self.assertIn("generation error", self.results()["1"]["error_message"]) - self.engine.generate.side_effect = make_engine().generate - await process_batch(self.sampler, self.store, [{"request_id": "2", "weights_path": "a"}]) - self.assertEqual(self.results()["2"]["type"], "sample") - - async def test_generation_preserves_tokens_logprobs_and_stop_options(self): - async def outputs(): - yield SimpleNamespace( - outputs=[SimpleNamespace(token_ids=[4], logprobs=[{4: SimpleNamespace(logprob=-0.25)}], finish_reason="length")], - prompt_logprobs=[None, {2: SimpleNamespace(logprob=-0.5)}], - ) - - self.engine.generate = Mock(return_value=outputs()) - result = await self.sampler.generate( - {"request_id": "req", "prompt_token_ids": [1, 2], "stop": [7], "max_tokens": 1, "include_prompt_logprobs": True} - ) - self.assertEqual(result["sequences"], [{"tokens": [4], "logprobs": [-0.25], "stop_reason": "length"}]) - self.assertEqual(result["prompt_logprobs"], [None, -0.5]) - params = self.engine.generate.call_args.kwargs["sampling_params"] - self.assertEqual(params.stop_token_ids, [7]) - self.assertEqual(params.max_tokens, 1) - self.assertEqual(params.prompt_logprobs, 1) - self.assertIsNone(self.engine.generate.call_args.kwargs["lora_request"]) - - async def test_lora_request_attached_only_when_adapter_exists(self): - with tempfile.TemporaryDirectory() as adapter: - await process_batch(self.sampler, self.store, [{"request_id": "1", "lora_id": "job-a", "lora_path": adapter}]) - self.assertIsNone(self.engine.generate.call_args.kwargs["lora_request"]) - open(os.path.join(adapter, "adapter_config.json"), "w").close() - await process_batch(self.sampler, self.store, [{"request_id": "2", "lora_id": "job-a", "lora_path": adapter}]) - lora_request = self.engine.generate.call_args.kwargs["lora_request"] - self.assertEqual((lora_request.lora_name, lora_request.lora_path), ("job-a", adapter)) - self.engine.update_weights.assert_not_called() - - -class SamplerLifecycleTest(unittest.IsolatedAsyncioTestCase): - async def asyncSetUp(self): - self.events = [] - self.store = AsyncMock() - self.engine = make_engine() - self.slicer = AsyncMock() - self.slicer.faulted = None - self.slicer.acquire = Mock(side_effect=self.slot) - self.workload = WorkloadRef("sampler-test") - self.factory = Mock(side_effect=self.create_engine) - self.patches = [ - patch("server.vllm_sampler.time_slicer_client_from_env", return_value=self.slicer), - patch("server.vllm_sampler.workload_from_env", return_value=self.workload), - patch("server.vllm_sampler.get_state_store", return_value=self.store), - patch("server.vllm_sampler.ENGINE_POLL_SECONDS", 0.01), - ] - for p in self.patches: - p.start() - self.addCleanup(p.stop) - - @asynccontextmanager - async def slot(self, workload): - self.events.append("acquire") - try: - yield - finally: - self.events.append("release") - - def create_engine(self): - self.events.append("create") - return self.engine - - async def test_initialization_and_batches_own_slots_and_cleanup_once(self): - self.store.get_sampling_requests_for_model.return_value = [{"request_id": "1"}, {"request_id": "SHUTDOWN_SENTINEL"}] - await serve_time_sliced("test", self.store, self.factory) - self.assertEqual(self.events, ["acquire", "create", "release", "acquire", "release"]) - self.assertEqual(self.engine.sleep.await_count, 2) - self.engine.wake_up.assert_awaited_once() - self.engine.shutdown.assert_called_once() - self.slicer.register.assert_awaited_once_with(self.workload) - self.slicer.unregister.assert_awaited_once_with(self.workload) - self.slicer.close.assert_awaited_once() - self.store.delete_values.assert_awaited_once_with("open_rl:sampler_ready:test") - self.assertEqual(self.store.set_result.call_args.args[0], "1") - - async def test_registration_failure_never_constructs_engine(self): - self.slicer.register.side_effect = RuntimeError("registration failed") - with self.assertRaisesRegex(RuntimeError, "registration failed"): - await serve_time_sliced("test", self.store, self.factory) - self.factory.assert_not_called() - self.slicer.unregister.assert_not_called() - self.slicer.close.assert_awaited_once() - - async def test_initialization_failure_unregisters_without_unlocked_retry(self): - self.factory.side_effect = RuntimeError("engine failed") - with self.assertRaisesRegex(RuntimeError, "engine failed"): - await serve_time_sliced("test", self.store, self.factory) - self.factory.assert_called_once() - self.assertEqual(self.events, ["acquire", "release"]) - self.slicer.unregister.assert_awaited_once() - self.slicer.close.assert_awaited_once() - - async def test_cancellation_shuts_down_engine_and_unregisters(self): - entered = asyncio.Event() - - async def get_batch(model_id): - entered.set() - await asyncio.Event().wait() - - self.store.get_sampling_requests_for_model.side_effect = get_batch - task = asyncio.create_task(serve_time_sliced("test", self.store, self.factory)) - await entered.wait() - task.cancel() - with self.assertRaises(asyncio.CancelledError): - await task - self.engine.shutdown.assert_called_once() - self.store.delete_values.assert_awaited_once() - self.slicer.unregister.assert_awaited_once() - self.slicer.close.assert_awaited_once() - - async def test_unshared_sampler_does_not_sleep_or_wake_engine(self): - self.store.get_sampling_requests_for_model.return_value = [{"request_id": "1"}, {"request_id": "SHUTDOWN_SENTINEL"}] - await serve("test", self.store, lambda: self.engine) - self.engine.wake_up.assert_not_called() - self.engine.sleep.assert_not_called() - self.engine.shutdown.assert_called_once() - - def results(self): - return {call.args[0]: call.args[1] for call in self.store.set_result.call_args_list} - - async def test_dead_engine_fails_hanging_requests_and_exits(self): - async def hang(request_id): - self.engine.errored = True - await asyncio.Event().wait() - - self.engine.generate.side_effect = make_engine(hang).generate - self.store.get_sampling_requests_for_model.return_value = [{"request_id": "1"}] - with self.assertRaisesRegex(RuntimeError, "engine is dead"): - await serve_time_sliced("test", self.store, self.factory) - self.assertIn("engine died", self.results()["1"]["error_message"]) - self.engine.shutdown.assert_called_once() - self.slicer.unregister.assert_awaited_once() - - async def test_faulted_release_exits_without_unregistering(self): - @asynccontextmanager - async def faulting_slot(workload): - yield - if self.slicer.acquire.call_count == 2: - self.slicer.faulted = "park failed" - - self.slicer.acquire.side_effect = faulting_slot - self.store.get_sampling_requests_for_model.return_value = [{"request_id": "1"}] - with self.assertRaisesRegex(RuntimeError, "park failed"): - await serve_time_sliced("test", self.store, self.factory) - self.assertEqual(self.results()["1"]["type"], "sample") - self.engine.shutdown.assert_called_once() - self.slicer.unregister.assert_not_called() - self.slicer.close.assert_awaited_once() - - async def test_loop_errors_fail_unserved_requests_and_retry(self): - self.store.get_sampling_requests_for_model.side_effect = [ - RuntimeError("redis down"), - [{"request_id": "1"}], - [{"request_id": "2"}, {"request_id": SHUTDOWN_SENTINEL}], - ] - self.engine.wake_up.side_effect = [RuntimeError("wake failed"), None] - await serve_time_sliced("test", self.store, self.factory) - self.assertIn("wake failed", self.results()["1"]["error_message"]) - self.assertEqual(self.results()["2"]["type"], "sample") - self.slicer.unregister.assert_awaited_once() - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_sampler_weight_rotation.py b/tests/test_sampler_weight_rotation.py index 43d136f4a..af2466e74 100644 --- a/tests/test_sampler_weight_rotation.py +++ b/tests/test_sampler_weight_rotation.py @@ -6,7 +6,10 @@ from unittest.mock import patch from server import training_requests_processor as trp -from server.store import InMemoryStore +from server.model_metadata import TrainingModelMetadata, persist_model_metadata +from server.sampling_store import SamplingStore +from server.store import InMemoryStateStore, InMemoryStore, RedisStore +from tests.redis_server import RedisTestCase from tests.test_fft_batch_failure import SlicerStub from training import commands @@ -41,5 +44,25 @@ def test_only_the_newest_versions_stay_on_the_volume(self) -> None: self.assertEqual(sorted(os.listdir(versions)), ["sampler-3", "sampler-4", "sampler-5"]) +class PinnedVersionTest(RedisTestCase): + async def test_a_version_a_pending_request_names_outlives_retention(self) -> None: + tmp = self.enterContext(tempfile.TemporaryDirectory()) + self.enterContext(patch.dict(os.environ, {"REDIS_URL": self.url, "OPEN_RL_TMP_DIR": tmp})) + self.enterContext(patch.object(trp, "SAMPLER_VERSIONS_KEPT", 3)) + state, meta = InMemoryStateStore(), TrainingModelMetadata(base_model="base", fine_tuning_type="full") + await persist_model_metadata(state, "run-a", meta) + self.enterContext(patch.object(trp, "get_state_store", return_value=state)) + versions = os.path.join(tmp, "sampler_full", "run-a", "sampler_weights") + pinned = {"request_id": "waiting", "pool": meta.pool("run-a"), "revision": os.path.join(versions, "sampler-1"), "idempotency_key": None} + await SamplingStore(self.redis).accept(pinned) + proc = trp.TrainingRequestsProcessor(RedisStore(self.redis), RecordingWorker(), "run-a", time_slicer=SlicerStub()) + for step in range(1, 6): + await proc.dispatch_operation( + commands.SaveWeightsForSampler(request_id=f"r{step}", model_id="run-a", path=f"tinker://run-a/sampler_weights/sampler-{step}") + ) + await asyncio.sleep(0.01) + self.assertEqual(sorted(os.listdir(versions)), ["sampler-1", "sampler-3", "sampler-4", "sampler-5"]) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_scheduler_worker_manager.py b/tests/test_scheduler_worker_manager.py index 977b10a58..8eca542ff 100644 --- a/tests/test_scheduler_worker_manager.py +++ b/tests/test_scheduler_worker_manager.py @@ -1,6 +1,8 @@ import json import os +import tempfile import unittest +from types import SimpleNamespace from typing import Any from unittest.mock import patch @@ -54,11 +56,38 @@ def list_namespaced_custom_object(self, group: str, version: str, namespace: str return {"items": items} +class FakeApi: + def __init__(self): + self.objects: dict[tuple[str, str], dict] = {} + + def __getattr__(self, method: str): + verb, kind = method.split("_namespaced_") + + def call(*args, **kwargs): + if verb == "list": + key, _, value = kwargs["label_selector"].partition("=") + found = [name for (k, name), body in self.objects.items() if k == kind and body["metadata"].get("labels", {}).get(key) == value] + return SimpleNamespace(items=[SimpleNamespace(metadata=SimpleNamespace(name=name)) for name in found]) + if verb in ("delete", "read") and (kind, args[0]) not in self.objects: + raise ApiError(404) + if verb == "delete": + return self.objects.pop((kind, args[0])) + if verb == "read": + return SimpleNamespace(metadata=SimpleNamespace(deletion_timestamp=None)) + if (kind, args[1]["metadata"]["name"]) in self.objects: + raise ApiError(409) + self.objects[(kind, args[1]["metadata"]["name"])] = args[1] + + return call + + class SchedulerWorkerManagerTest(unittest.TestCase): def setUp(self) -> None: - self.enterContext(patch.dict(os.environ, {"REDIS_URL": "redis://localhost:6379"})) - self.api = FakeCustomObjectsApi() - self.manager = SchedulerWorkerManager(custom_api=self.api) + # The API server Deployment sets OPEN_RL_TMP_DIR to the shared volume every worker pod mounts. A temp dir stands in for it. + self.tmp = self.enterContext(tempfile.TemporaryDirectory()) + self.enterContext(patch.dict(os.environ, {"REDIS_URL": "redis://localhost:6379", "OPEN_RL_TMP_DIR": self.tmp})) + self.api, self.k8s = FakeCustomObjectsApi(), FakeApi() + self.manager = SchedulerWorkerManager(custom_api=self.api, apps_api=self.k8s, networking_api=self.k8s) def store_with(self, model_id: str, meta: dict) -> InMemoryStore: s = InMemoryStateStore() @@ -83,7 +112,7 @@ def test_lora_trainer_and_sampler_share_an_owner(self) -> None: t_container = trainer["spec"]["template"]["spec"]["containers"][0] s_container = sampler["spec"]["template"]["spec"]["containers"][0] self.assertEqual(t_container["command"][-1], "server.training_requests_processor") - self.assertEqual(s_container["command"][-1], "server.vllm_sampler") + self.assertEqual(s_container["command"], ["uv", "run", "vllm", "serve"]) self.assertIn("--active-tenant-set-id", t_container["args"]) def test_an_exclusive_fft_worker_is_placed_alone_and_never_time_sliced(self) -> None: @@ -144,16 +173,28 @@ def test_workers_get_the_deployment_settings(self) -> None: def test_samplers_keep_the_vllm_compile_cache_on_the_shared_volume(self) -> None: s = self.store_with("job-lora-1", {"base_model": "Qwen/Qwen2.5-0.5B", "fine_tuning_type": "lora"}) - with patch.dict(os.environ, {"OPEN_RL_TMP_DIR": "/mnt/shared/open-rl"}), patch("server.worker_manager.get_state_store", return_value=s): + with patch("server.worker_manager.get_state_store", return_value=s): self.manager.ensure("job-lora-1", "sampler") self.manager.ensure("job-lora-1", "trainer") envs = {} for pod in self.api.created: envs[pod["spec"]["role"]] = {e["name"]: e.get("value") for e in pod["spec"]["template"]["spec"]["containers"][0]["env"]} - self.assertEqual(envs["sampler"]["VLLM_CACHE_ROOT"], "/mnt/shared/open-rl/vllm-cache") + self.assertEqual(envs["sampler"]["VLLM_CACHE_ROOT"], f"{self.tmp}/vllm-cache") + self.assertEqual(envs["sampler"]["TRITON_CACHE_DIR"], f"{self.tmp}/vllm-cache/triton") self.assertNotIn("VLLM_CACHE_ROOT", envs["trainer"]) + def test_a_bare_base_model_sampler_on_an_fft_deployment_is_fft_with_a_pool_apart_from_the_shared_lora_pool(self) -> None: + s = self.store_with("job-lora-1", {"base_model": "Qwen/Qwen2.5-0.5B", "fine_tuning_type": "lora"}) + with patch.dict(os.environ, {"OPEN_RL_ENABLE_FFT": "true"}), patch("server.worker_manager.get_state_store", return_value=s): + self.manager.ensure("job-lora-1", "sampler") + self.manager.ensure("Qwen/Qwen2.5-0.5B", "sampler") + self.assertEqual([w["spec"]["trainingKind"] for w in self.api.created], ["lora", "fft"]) + shared, bare = self.api.created + self.assertNotEqual(bare["spec"]["ownerID"], shared["spec"]["ownerID"]) + pools = [w["spec"]["template"]["metadata"]["labels"]["app.kubernetes.io/instance"] for w in (shared, bare)] + self.assertNotEqual(pools[0], pools[1]) + def test_launch_is_idempotent(self) -> None: s = self.store_with("job-lora-1", {"base_model": "Qwen/Qwen2.5-0.5B", "fine_tuning_type": "lora"}) with patch("server.worker_manager.get_state_store", return_value=s): @@ -255,7 +296,7 @@ def test_automodel_jobs_share_a_runtime_apart_from_pytorch(self) -> None: self.assertEqual(t_container["args"], ["--model-id", runtime, "--active-tenant-set-id", f"{runtime}-1"]) self.assertEqual({e["name"]: e.get("value") for e in t_container["env"]}["OPEN_RL_TRAINER_BACKEND"], "automodel") self.assertNotEqual(s_container["image"], "am:1") - self.assertEqual({e["name"]: e.get("value") for e in s_container["env"]}["OPEN_RL_MODEL_ID"], runtime) + self.assertEqual(s_container["command"], ["uv", "run", "vllm", "serve"]) # Like any shared LoRA runtime, deleting one job leaves it up. self.assertEqual(self.api.deleted, []) @@ -387,6 +428,34 @@ def test_sampler_replicas_are_single_gpu_workloads_released_together(self) -> No self.assertEqual(sampler["spec"]["ownerID"], "job-sr") self.assertEqual(sorted(self.api.deleted), sorted(names)) + def test_a_pool_of_one_replica_serves_itself_with_no_router(self) -> None: + s = self.store_with("job-lora-1", {"base_model": "Qwen/Qwen2.5-0.5B", "fine_tuning_type": "lora", "exclusive": True}) + with patch("server.worker_manager.get_state_store", return_value=s): + self.manager.ensure("job-lora-1", "sampler") + (sampler,) = self.api.created + container = sampler["spec"]["template"]["spec"]["containers"][0] + self.assertEqual((container["command"][-1], container["args"]), ("server.vllm_sampler", ["--pool", "sampler-job-lora-1"])) + self.assertEqual({e["name"]: e.get("value") for e in container["env"]}["OPEN_RL_TIME_SLICING"], "off") + self.assertEqual(self.k8s.objects, {}) + + def test_a_shared_lora_pool_is_routed_and_grows_to_the_most_replicas_any_job_asks_for(self) -> None: + for first, second in ((1, 2), (2, 1)): + with self.subTest(first=first, second=second): + api, k8s = FakeCustomObjectsApi(), FakeApi() + manager = SchedulerWorkerManager(custom_api=api, apps_api=k8s, networking_api=k8s) + s = InMemoryStateStore() + for model_id, replicas in (("job-a", first), ("job-b", second)): + meta = {"base_model": "Qwen/Qwen2.5-0.5B", "fine_tuning_type": "lora", "sampler_replicas": replicas} + s.kv_store[f"open_rl:model_meta:{model_id}"] = json.dumps(meta) + with patch("server.worker_manager.get_state_store", return_value=s): + manager.ensure("job-a", "sampler") + manager.ensure("job-b", "sampler") + commands = {w["metadata"]["name"]: w["spec"]["template"]["spec"]["containers"][0]["command"] for w in api.created} + serve = ["uv", "run", "vllm", "serve"] + self.assertEqual(commands, {"lora-qwen-qwen2-5-0-5b-0-sampler": serve, "lora-qwen-qwen2-5-0-5b-1-sampler": serve}) + pools = {w["metadata"]["labels"]["app.kubernetes.io/instance"] for w in api.created} + self.assertEqual([name for kind, name in k8s.objects if kind == "deployment"], list(pools)) + class MixedSamplingSessionTest(unittest.IsolatedAsyncioTestCase): async def test_lora_and_fft_sessions_launch_their_own_sampler_types(self) -> None: @@ -394,14 +463,15 @@ async def test_lora_and_fft_sessions_launch_their_own_sampler_types(self) -> Non state = InMemoryStateStore() for model_id, kind in (("lora-a", "lora"), ("lora-b", "lora"), ("fft-a", "full")): await state.set_value(f"open_rl:model_meta:{model_id}", json.dumps({"base_model": "Qwen/Qwen2.5-0.5B", "fine_tuning_type": kind})) - api = FakeCustomObjectsApi() + api, tmp = FakeCustomObjectsApi(), self.enterContext(tempfile.TemporaryDirectory()) + env = {"REDIS_URL": "redis://localhost:6379", "OPEN_RL_ENABLE_FFT": "true", "SAMPLING_BACKEND": "vllm", "OPEN_RL_TMP_DIR": tmp} with ( - patch.dict(os.environ, {"REDIS_URL": "redis://localhost:6379", "OPEN_RL_ENABLE_FFT": "true", "SAMPLING_BACKEND": "vllm"}), + patch.dict(os.environ, env), patch("server.worker_manager.get_state_store", return_value=state), patch.object(api_server, "store", store), patch.object(api_server, "state", state), patch.object(api_server, "get_store", return_value=store), - patch.object(api_server, "worker_manager", SchedulerWorkerManager(custom_api=api)), + patch.object(api_server, "worker_manager", SchedulerWorkerManager(custom_api=api, apps_api=FakeApi(), networking_api=FakeApi())), ): async with asgi_client() as client: for model_id in ("lora-a", "fft-a", "lora-b"): @@ -410,7 +480,7 @@ async def test_lora_and_fft_sessions_launch_their_own_sampler_types(self) -> Non self.assertEqual(len(api.created), 2, "LoRA sessions should reuse one sampler while FFT gets its own") lora, fft = api.created self.assertEqual(lora["spec"]["trainingKind"], "lora") - self.assertEqual(lora["spec"]["template"]["spec"]["containers"][0]["command"][-1], "server.vllm_sampler") + self.assertEqual(lora["spec"]["template"]["spec"]["containers"][0]["command"], ["uv", "run", "vllm", "serve"]) self.assertEqual(fft["spec"]["trainingKind"], "fft") self.assertEqual(fft["metadata"]["name"], "fft-fft-a-sampler") diff --git a/tests/test_session_lifecycle.py b/tests/test_session_lifecycle.py index d6fb8e349..2518e6a8c 100644 --- a/tests/test_session_lifecycle.py +++ b/tests/test_session_lifecycle.py @@ -79,7 +79,7 @@ async def test_shared_lora_owner_outlives_the_session_that_created_it(self): adapter = (await self.post("create_model", {"base_model": "test-base", "session_id": training}))["request_id"] fft_body = {"base_model": "fft-base", "session_id": training, "user_metadata": {"openrl.fine_tuning_type": "full"}} fft_model = (await self.post("create_model", fft_body))["request_id"] - await self.state.set_value("open_rl:sampler_ready:test-base", "1") + await self.state.set_value("open_rl:trainer_ready:test-base", "1") sampling = (await self.post("create_session", {}))["session_id"] await self.post("create_sampling_session", {"model_path": f"tinker://{adapter}/sampler_weights/test", "session_id": sampling}) self.assertEqual(self.manager.ensured, [(adapter, "trainer"), (fft_model, "trainer"), (adapter, "sampler")]) @@ -91,7 +91,7 @@ async def test_shared_lora_owner_outlives_the_session_that_created_it(self): await self.expire(sampling) await self.reap() self.assertEqual(self.manager.released, [fft_model.lower(), "test-base"]) - self.assertIsNone(await self.state.get_value("open_rl:sampler_ready:test-base")) + self.assertIsNone(await self.state.get_value("open_rl:trainer_ready:test-base")) self.assertEqual(await self.registry.owners(), []) async def test_an_exclusive_lora_model_has_its_own_runtime(self): @@ -136,6 +136,8 @@ async def test_an_owner_stays_listed_until_forgotten(self): self.assertEqual(await self.registry.owners(), []) async def test_a_session_attaching_during_teardown_waits_for_it(self): + # Without FFT, a session on the bare base model shares the base model's LoRA owner. + self.enterContext(patch.dict(os.environ, {"OPEN_RL_ENABLE_FFT": "false"})) await self.registry.attach("a", "test-base") await self.expire("a") loop = asyncio.get_running_loop() diff --git a/tests/test_vllm_sampler.py b/tests/test_vllm_sampler.py new file mode 100644 index 000000000..40c980638 --- /dev/null +++ b/tests/test_vllm_sampler.py @@ -0,0 +1,207 @@ +"""The in-process sampler's turns on a GPU it shares with its trainer, with the real time slicer.""" + +import asyncio +import json +import os +import sys +import tempfile +import types +import uuid +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +from accel_timeslicer.serve import start_time_slicer +from accel_timeslicer.single_node import SingleNodeTimeSlicer +from accel_timeslicer.time_slicer import SocketTimeSlicerClient +from accel_timeslicer.workload import WorkloadRef +from server.sampling_dispatcher import Dispatcher +from server.sampling_store import SamplingStore, future_key +from server.vllm_sampler import Sampler, TimeSlicedDispatcher +from tests.redis_server import RedisTestCase +from tests.test_sampling_dispatcher import SAMPLE, until + + +class RecordingRestorer: + """Parks and restores nothing, and logs each park and restore the time slicer asks for.""" + + def __init__(self, log: list) -> None: + self.log = log + + def checkpoint(self, workload: WorkloadRef) -> None: + self.log.append(("park", workload.name)) + + def restore(self, workload: WorkloadRef) -> None: + self.log.append(("restore", workload.name)) + + +class GpuSampler: + """Stands in for vllm_sampler.Sampler, whose vLLM engine needs a GPU. It logs each call and fails any that reaches + it off the GPU, or while the trainer holds it.""" + + def __init__(self, log: list, gpu: set) -> None: + self.log, self.gpu, self.revision = log, gpu, "" + + async def wake(self) -> None: + self.gpu.add("sampler") + self.log.append(("wake", "sampler")) + + async def sleep(self) -> None: + self.gpu.discard("sampler") + self.log.append(("sleep", "sampler")) + + async def ensure_weights(self, revision: str) -> bool: + assert self.gpu == {"sampler"}, self.gpu + if revision != self.revision: + self.log.append(("load", revision)) + self.revision = revision + return True + + async def generate(self, r: dict) -> dict: + assert self.gpu == {"sampler"}, self.gpu + self.log.append(("generate", r["revision"])) + await asyncio.sleep(0.01) + return SAMPLE + + +class EngineDeadError(Exception): + pass + + +def fake_vllm() -> dict[str, types.ModuleType]: + """The vLLM names Sampler imports. vLLM needs a GPU, and the unit tests run without one.""" + names = ["exceptions", "sampling_params", "lora", "lora.request", "v1", "v1.engine", "v1.engine.exceptions"] + modules = {name: types.ModuleType(name) for name in ["vllm", *(f"vllm.{name}" for name in names)]} + modules["vllm"].SamplingParams = dict + modules["vllm.exceptions"].VLLMClientError = type("VLLMClientError", (Exception,), {}) + modules["vllm.sampling_params"].RequestOutputKind = SimpleNamespace(FINAL_ONLY="final_only") + modules["vllm.lora.request"].LoRARequest = lambda name, lora_id, path: path + modules["vllm.v1.engine.exceptions"].EngineDeadError = EngineDeadError + modules["vllm.v1.engine.exceptions"].EngineGenerateError = type("EngineGenerateError", (Exception,), {}) + return modules + + +class FakeEngine: + """Stands in for vLLM's in-process engine. It dies on the `poison` prompt the way vLLM's does: that request and every + later one raise EngineDeadError, and `errored` turns true.""" + + def __init__(self, poison: list[int] | None = None) -> None: + self.poison, self.errored, self.seen = poison, False, [] + + async def generate(self, prompt: dict, sampling_params: dict, request_id: str, lora_request: str | None): + if self.errored: + raise EngineDeadError() + self.seen.append((prompt["prompt_token_ids"], lora_request)) + if prompt["prompt_token_ids"] == self.poison: + self.errored = True + raise EngineDeadError() + choice = SimpleNamespace(token_ids=[100], logprobs=[{100: SimpleNamespace(logprob=-0.5)}], finish_reason="length") + yield SimpleNamespace(outputs=[choice], prompt_logprobs=None) + + +class InProcessSamplerTest(RedisTestCase): + """vllm_sampler.Sampler with a fake engine, served by the real dispatcher over Redis.""" + + async def asyncSetUp(self) -> None: + await super().asyncSetUp() + self.store, self.pool = SamplingStore(self.redis), f"sampler-{uuid.uuid4()}" + self.enterContext(patch.dict(sys.modules, fake_vllm())) + self.tmp = self.enterContext(tempfile.TemporaryDirectory()) + self.enterContext(patch.dict(os.environ, {"OPEN_RL_TMP_DIR": self.tmp})) + + async def accepted(self, prompt: list[int], model: str = "base") -> dict: + r = {"request_id": str(uuid.uuid4()), "pool": self.pool, "revision": "", "model": model, "idempotency_key": None} + r |= {"prompt_token_ids": prompt, "sampling_params": {"n": 1, "max_tokens": 1}, "include_prompt_logprobs": False} + await self.store.accept(r) + return r + + def serve(self, engine: FakeEngine) -> asyncio.Task: + run = asyncio.create_task(Dispatcher(self.store, self.pool, Sampler(engine, "base"), 4).run()) + self.addAsyncCleanup(lambda: (run.cancel(), asyncio.gather(run, return_exceptions=True))[1]) + return run + + async def results(self, r: dict) -> list[dict]: + return [json.loads(raw) for raw in await self.redis.lrange(future_key(r["request_id"]), 0, -1)] + + async def test_a_request_that_kills_the_engine_fails_once_and_the_restarted_engine_never_gets_it(self): + poison = await self.accepted([6, 6, 6]) + dying = FakeEngine(poison=[6, 6, 6]) + first = self.serve(dying) + await until(lambda: self.results(poison)) + later = await self.accepted([1, 2, 3]) + await asyncio.sleep(0.3) + self.assertEqual(await self.results(later), []) + # The worker exits once its engine is dead, and the restarted one serves the pool. + first.cancel() + restarted = FakeEngine(poison=[6, 6, 6]) + self.serve(restarted) + await until(lambda: self.results(later)) + self.assertEqual([r["type"] for r in await self.results(poison)], ["RequestFailedResponse"]) + self.assertEqual((await self.results(later))[0]["type"], "sample") + self.assertEqual([prompt for prompt, _ in dying.seen + restarted.seen], [[6, 6, 6], [1, 2, 3]]) + + async def test_a_lora_snapshot_this_pod_does_not_see_yet_waits_and_never_reaches_the_engine(self): + r = await self.accepted([1, 2, 3], model="adapter-1") + engine = FakeEngine() + self.serve(engine) + await asyncio.sleep(0.3) + self.assertEqual((engine.seen, engine.errored, await self.results(r)), ([], False, [])) + snapshot = Path(self.tmp, "sampler_lora", "adapter-1") + snapshot.mkdir(parents=True) + (snapshot / "adapter_config.json").write_text("{}") + await until(lambda: self.results(r)) + self.assertEqual((await self.results(r))[0]["type"], "sample") + self.assertEqual(engine.seen, [([1, 2, 3], str(snapshot))]) + + +class TimeSlicedPoolTest(RedisTestCase): + """A pool of one sampler on its trainer's GPU. The dispatcher runs each batch inside one turn from the real time + slicer, over its socket, and the trainer steps between batches.""" + + async def asyncSetUp(self) -> None: + await super().asyncSetUp() + self.store, self.pool, self.log, self.gpu = SamplingStore(self.redis), f"sampler-{uuid.uuid4()}", [], set() + self.tmp = Path(self.enterContext(tempfile.TemporaryDirectory())) + socket_path = str(self.tmp / "slicer.sock") + server = await start_time_slicer(SingleNodeTimeSlicer(RecordingRestorer(self.log)), socket_path) + self.addAsyncCleanup(server.wait_closed) + self.addCleanup(server.close) + self.sampler_slicer, self.trainer_slicer = SocketTimeSlicerClient(socket_path=socket_path), SocketTimeSlicerClient(socket_path=socket_path) + self.addAsyncCleanup(self.sampler_slicer.close) + self.addAsyncCleanup(self.trainer_slicer.close) + self.sampler, self.trainer = WorkloadRef("sampler", "gpu-0"), WorkloadRef("trainer", "gpu-0") + await self.sampler_slicer.register(self.sampler) + await self.trainer_slicer.register(self.trainer) + + async def accepted(self, revision: str = "") -> dict: + r = {"request_id": str(uuid.uuid4()), "pool": self.pool, "revision": revision, "model": "base", "idempotency_key": None} + r |= {"prompt_token_ids": [1, 2, 3], "sampling_params": {"n": 1, "max_tokens": 2}, "include_prompt_logprobs": False} + await self.store.accept(r) + return r + + async def saved(self, *records: dict) -> None: + for r in records: + await until(lambda r=r: self.redis.llen(future_key(r["request_id"]))) + + async def test_each_batch_runs_inside_one_turn_and_the_trainer_steps_between_batches(self): + w1 = str(self.tmp / "w1") + Path(w1).mkdir() + Path(w1, "metadata.json").write_text("{}") + first = [await self.accepted() for _ in range(3)] + dispatcher = TimeSlicedDispatcher(self.store, self.pool, GpuSampler(self.log, self.gpu), 4, self.sampler_slicer, self.sampler) + run = asyncio.create_task(dispatcher.run()) + self.addAsyncCleanup(lambda: (run.cancel(), asyncio.gather(run, return_exceptions=True))[1]) + await self.saved(*first) + # The trainer steps on the GPU, and the next batch waits for its turn. + async with self.trainer_slicer.acquire(self.trainer): + self.log.append(("step", "trainer")) + second = [await self.accepted(w1) for _ in range(2)] + await asyncio.sleep(0.3) + self.assertNotIn(("load", w1), self.log) + await self.saved(*second) + await until(lambda: self.log[-1] == ("park", "sampler")) + batch = [("generate", "")] * 3 + turn = [("restore", "sampler"), ("wake", "sampler"), ("load", w1), ("generate", w1), ("generate", w1), ("sleep", "sampler"), ("park", "sampler")] + expected = [("wake", "sampler"), *batch, ("sleep", "sampler"), ("park", "sampler"), ("step", "trainer"), ("park", "trainer"), *turn] + self.assertEqual(self.log, expected) + self.assertEqual([json.loads(raw) for r in first + second for raw in await self.redis.lrange(future_key(r["request_id"]), 0, -1)], [SAMPLE] * 5) diff --git a/tests/test_worker_manager.py b/tests/test_worker_manager.py index ede3725c3..9b1f382eb 100644 --- a/tests/test_worker_manager.py +++ b/tests/test_worker_manager.py @@ -1,13 +1,23 @@ import json +import tempfile import unittest +from pathlib import Path from unittest.mock import patch from server import api_server from server.session_registry import SessionRegistry +from server.store import InMemoryStateStore from server.worker_manager import LocalWorkerManager from tests.api_client import asgi_client, post_json +def fft_job(model_id: str) -> InMemoryStateStore: + """A state store that knows `model_id` as an FFT job.""" + s = InMemoryStateStore() + s.kv_store[f"open_rl:model_meta:{model_id}"] = json.dumps({"base_model": "base-model-a", "fine_tuning_type": "full"}) + return s + + class StoreStub: def __init__(self): self.forwarded_requests = [] @@ -215,6 +225,7 @@ async def test_requires_redis(self) -> None: async def test_local_launch_stamps_workload_tags_and_process_group(self) -> None: with ( patch.dict("os.environ", {"REDIS_URL": "redis://localhost:6379", "OPEN_RL_ENABLE_FFT": "true"}, clear=True), + patch("server.worker_manager.get_state_store", return_value=fft_job("Model_A.1")), patch("server.worker_manager.subprocess.Popen") as popen, ): manager = LocalWorkerManager() @@ -226,20 +237,16 @@ async def test_local_launch_stamps_workload_tags_and_process_group(self) -> None self.assertEqual(kwargs["env"]["OPEN_RL_TIME_SLICE_JOB_ID"], "trainer-Model_A.1") self.assertEqual(kwargs["env"]["OPEN_RL_TIME_SLICE_GROUP"], "trainers") - async def test_local_sampler_launch_stamps_workload_tags_and_process_group(self) -> None: + async def test_a_local_fft_sampler_on_the_trainers_gpu_serves_its_pool_itself_in_the_trainers_claim(self) -> None: with ( patch.dict("os.environ", {"REDIS_URL": "redis://localhost:6379", "SAMPLING_BACKEND": "vllm", "OPEN_RL_ENABLE_FFT": "true"}, clear=True), + patch("server.worker_manager.get_state_store", return_value=fft_job("fft-job")), patch("server.worker_manager.subprocess.Popen") as popen, ): - manager = LocalWorkerManager() - manager.ensure("Model_A.1", "sampler") - - _, kwargs = popen.call_args - self.assertTrue(kwargs["start_new_session"]) - self.assertEqual(kwargs["env"]["OPEN_RL_ENABLE_FFT"], "true") - self.assertEqual(kwargs["env"]["OPEN_RL_MODEL_ID"], "Model_A.1") - self.assertEqual(kwargs["env"]["OPEN_RL_TIME_SLICE_JOB_ID"], "sampler-Model_A.1") - self.assertEqual(kwargs["env"]["OPEN_RL_TIME_SLICE_GROUP"], "samplers") + LocalWorkerManager().ensure("fft-job", "sampler") + ((command, kwargs),) = [(c[0][0], c[1]) for c in popen.call_args_list] + self.assertEqual(command[command.index("-m") + 1 :], ["server.vllm_sampler", "--pool", "sampler-fft-job"]) + self.assertEqual((kwargs["env"]["OPEN_RL_TIME_SLICE_JOB_ID"], kwargs["env"]["OPEN_RL_TIME_SLICE_GROUP"]), ("sampler-fft-job", "trainers")) async def test_launch_fetches_metadata_from_store(self) -> None: import json @@ -374,11 +381,12 @@ def test_internal_future_result_types_translate_to_tinker_types(self) -> None: class LocalWorkerManagerSamplerLaunchTest(unittest.TestCase): def setUp(self) -> None: - from pathlib import Path with patch.dict("os.environ", {"REDIS_URL": "redis://127.0.0.1:6379"}): self.manager = LocalWorkerManager(project_dir=Path("/tmp")) self.store = StoreStub() + # A LoRA launch creates its snapshot directory under OPEN_RL_TMP_DIR. + self.enterContext(patch.dict("os.environ", {"OPEN_RL_TMP_DIR": self.enterContext(tempfile.TemporaryDirectory())})) @patch("server.worker_manager.metadata_for") @patch("subprocess.Popen") @@ -395,14 +403,11 @@ def test_launch_sampler_lora_uses_base_model_and_reuses_process(self, mock_popen fine_tuning_type="lora", ) - # Launch for first LoRA model ID self.manager.ensure("model-lora-1", "sampler") - self.assertIn(("sampler", "Qwen/Qwen2.5-0.5B"), self.manager.processes) - self.assertEqual(mock_popen.call_count, 1) - - cmd_args = mock_popen.call_args[0][0] - self.assertIn("server.vllm_sampler", cmd_args) - self.assertIn("Qwen/Qwen2.5-0.5B", cmd_args) + self.assertEqual(set(self.manager.processes), {("sampler", "Qwen/Qwen2.5-0.5B")}) + ((command, serve_kwargs),) = [(c[0][0], c[1]) for c in mock_popen.call_args_list] + self.assertEqual((command[-2], command[-1].startswith("sampler-qwen-qwen2-5-0-5b-")), ("--pool", True)) + self.assertEqual((serve_kwargs["env"]["BASE_MODEL"], serve_kwargs["env"]["OPEN_RL_FINE_TUNING_TYPE"]), ("Qwen/Qwen2.5-0.5B", "lora")) # Launch for second LoRA model ID sharing the same base model self.manager.ensure("model-lora-2", "sampler") @@ -422,15 +427,14 @@ def test_launch_sampler_fft_uses_model_id(self, mock_popen, mock_fetch) -> None: base_model="Qwen/Qwen2.5-0.5B", created_at=100.0, fine_tuning_type="full", + exclusive=True, ) self.manager.ensure("model-fft-1", "sampler") - self.assertIn(("sampler", "model-fft-1"), self.manager.processes) - self.assertEqual(mock_popen.call_count, 1) - - cmd_args = mock_popen.call_args[0][0] - self.assertIn("server.vllm_sampler", cmd_args) - self.assertIn("model-fft-1", cmd_args) + self.assertEqual(set(self.manager.processes), {("sampler", "model-fft-1")}) + ((command, kwargs),) = [(c[0][0], c[1]) for c in mock_popen.call_args_list] + self.assertEqual(command[-2:], ["--pool", "sampler-model-fft-1"]) + self.assertEqual((kwargs["env"]["BASE_MODEL"], kwargs["env"]["OPEN_RL_FINE_TUNING_TYPE"]), ("Qwen/Qwen2.5-0.5B", "full")) if __name__ == "__main__": From 2aaae41eb985e1880d4a3e8032f4ffe86330ff3f Mon Sep 17 00:00:00 2001 From: Shuby Mishra Date: Mon, 26 Oct 2026 23:24:13 +0000 Subject: [PATCH 2/2] sampling: send each pool's requests round-robin by job A shared LoRA pool serves every job on its base model, and the dispatcher sent its records oldest first. One job's big batch then held every slot until it finished, and other jobs waited behind it. Dispatcher.to_send takes one record from each job, read from the record's "job" field, in turn, oldest first within a job, starting with the job served longest ago. A record accepted before this change has no job and groups under its model. --- src/server/sampling_dispatcher.py | 15 +++++++++++++-- tests/test_sampling_dispatcher.py | 23 ++++++++++++++++++++++- 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/src/server/sampling_dispatcher.py b/src/server/sampling_dispatcher.py index 3c696ed65..c25e34a40 100644 --- a/src/server/sampling_dispatcher.py +++ b/src/server/sampling_dispatcher.py @@ -5,6 +5,7 @@ import asyncio import contextlib +import itertools import json import logging import os @@ -17,7 +18,7 @@ import httpx import redis.asyncio as redis -from server.sampling_store import SamplingStore +from server.sampling_store import SamplingStore, job_of from server.vllm_options import sample_params, sample_result logger = logging.getLogger(__name__) @@ -88,6 +89,8 @@ def __init__(self, store: SamplingStore, pool: str, sampler: Any, concurrency: i self.records: dict[str, dict[str, Any]] = {} self.newest = "0-0" self.in_flight: dict[str, asyncio.Task] = {} + self.sends = itertools.count() + self.served: dict[str, int] = {} # the number of each job's latest send async def run(self) -> None: try: @@ -152,7 +155,14 @@ async def read(self) -> None: self.revision = revision def to_send(self) -> list[dict[str, Any]]: - return [r for r in self.records.values() if r["revision"] == self.revision and r["request_id"] not in self.in_flight] + """The current revision's records not in flight, one from each job in turn and oldest first within a job. The job + served longest ago goes first, so one job's big batch does not hold up the others.""" + jobs: dict[str, list[dict[str, Any]]] = {} + for r in self.records.values(): + if r["revision"] == self.revision and r["request_id"] not in self.in_flight: + jobs.setdefault(job_of(r), []).append(r) + turns = sorted(jobs.items(), key=lambda job: self.served.get(job[0], -1)) + return [r for turn in itertools.zip_longest(*(records for _, records in turns)) for r in turn if r is not None] async def set_result(self, r: dict[str, Any], result: dict[str, Any]) -> None: await self.store.set_result(r, result) @@ -160,6 +170,7 @@ async def set_result(self, r: dict[str, Any], result: dict[str, Any]) -> None: async def send(self, r: dict[str, Any]) -> None: """Saves the sampler's answer. None, or an error, leaves the record pending to be sent again.""" + self.served[job_of(r)] = next(self.sends) try: # A sampler that never answers must not hold the send past its record's deadline. timeout = max(0.0, accepted_at(r) + SAMPLE_DEADLINE_SECONDS - time.time()) diff --git a/tests/test_sampling_dispatcher.py b/tests/test_sampling_dispatcher.py index e3458f414..d9cf51eae 100644 --- a/tests/test_sampling_dispatcher.py +++ b/tests/test_sampling_dispatcher.py @@ -184,9 +184,10 @@ def revision(self, name: str, **metadata) -> str: def loads(self, *names: str) -> dict[str, list[str]]: return {name: [entry[2] for entry in self.log if entry[:2] == ("update", name)] for name in names} - async def accepted(self, revision: str = "", prompt: list | None = None) -> dict: + async def accepted(self, revision: str = "", prompt: list | None = None, **fields) -> dict: identity = {"request_id": str(uuid.uuid4()), "pool": self.pool, "revision": revision, "model": "base", "idempotency_key": None} r = {**identity, "prompt_token_ids": prompt or [1, 2, 3], "sampling_params": {"n": 1, "max_tokens": 2}, "include_prompt_logprobs": False} + r |= fields await self.store.accept(r) return r @@ -255,3 +256,23 @@ async def test_a_replica_on_any_save_at_the_weights_a_delta_starts_from_or_reach self.assertEqual(self.loads(early.name), {early.name: [second, third]}) self.assertTrue(same(replayed(at_first, [second, third]), weights(worker))) self.assertTrue(same(replayed(at_first, [third]), weights(worker))) + + async def test_one_jobs_big_batch_does_not_hold_up_another_jobs_request(self): + (replica,) = await self.replicas("default") + sent = [] + + # Records which job's request each send carries. Each job samples its own adapter, so the model names the job. + async def route(request) -> Response: + body = await request.body() + sent.append(json.loads(body)["model"]) + upstream = await self.http.post(f"{replica.url}/inference/v1/generate", content=body, timeout=None) + return Response(upstream.content, upstream.status_code, headers={"x-openrl-upstream": replica.url.removeprefix("http://")}) + + router_url = await self.serve(Starlette(routes=[Route("/inference/v1/generate", route, methods=["POST"])])) + big = [await self.accepted(job="job-a", model="adapter-a") for _ in range(6)] + small = await self.accepted(job="job-b", model="adapter-b") + run = asyncio.create_task(run_routed(routed_dispatcher(self.store, self.pool, self.http, router_url, [replica.url], concurrency=1))) + self.addAsyncCleanup(lambda: (run.cancel(), asyncio.gather(run, return_exceptions=True))[1]) + for r in [*big, small]: + await until(lambda r=r: self.redis.llen(future_key(r["request_id"]))) + self.assertEqual(sent[:2], ["adapter-a", "adapter-b"])