From b64bb342045a070458849e4eb8e6a38d3dffed57 Mon Sep 17 00:00:00 2001 From: Shuby Mishra Date: Mon, 26 Oct 2026 23:24:13 +0000 Subject: [PATCH 1/3] 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/3] 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"]) From 9093a861ecbc17fd30d32d01c316ae7654decc13 Mon Sep 17 00:00:00 2001 From: Shuby Mishra Date: Mon, 26 Oct 2026 23:24:13 +0000 Subject: [PATCH 3/3] sampling: recreate missing pools and fail requests no sampler reads A request could wait forever in two ways. A pool whose Workloads or dispatcher were deleted while its jobs still had requests pending never came back, since only create_sampling_session and save_weights_for_sampler ensured it. And the deadline was enforced only by the pool's dispatcher, so with no dispatcher nothing failed the pool's requests. Two siblings of the session reaper, each a task of its own on the same 30 s interval, now tend the pool streams. The deadline sweep fails every entry older than SAMPLE_DEADLINE_SECONDS through SamplingStore.fail_older_than. The reconcile looks at each job with pending entries in each pool, under the owner lock the reaper takes. They run apart, so a reconcile stuck in an ensure never holds up the sweep. The reconcile handles a pass's jobs concurrently, since an ensure can wait minutes for an old pool to terminate, but at most RECONCILE_CONCURRENCY (4) at a time, so stuck ensures cannot take every thread the API server runs blocking calls on. A job whose pool still has a sampler Workload, or a local sampler process, is left alone, since a client may sample without a session. When the pool is gone, a job whose owner has a live session gets worker_manager.ensure(job, "sampler"), which creates only what is missing, so a deleted pool comes back within 30 s even when the client sends no new sample and only polls retrieve_future. A job with no live session has its pending entries failed instead, so a later pool on the same base model does not replay them. The lock keeps a reaper release from landing between the checks and the ensure, which would recreate a pool that nothing removes again. With no worker manager, standing worker deployments own the pools and nothing is reconciled. delete_model takes the same owner lock. Under it, delete_model marks the job completed, fails the job's pending pool requests, and releases its workers. The reconcile reads the job's status again under the lock, and fails a completed job's requests instead of bringing its pool back. set_result's XDEL leaves an empty stream behind, and Redis never evicts it, so each FFT or exclusive job left one key forever. SamplingStore.delete_if_empty now deletes a pool's stream when nothing is pending in it. The reaper calls it for each runtime release_owner returns, delete_model after it releases a job's workers, and the reconcile after it fails the requests of a job whose pool is gone. It checks XLEN and deletes in one transaction that watches the stream, so a request accepted at that moment stays. pool_name, split out of TrainingModelMetadata.pool, names the pool of a runtime. Stream entries carry the job and the revision as fields of their own. The deadline sweep, pending_jobs, fail_job and pinned_revisions read each stream 100 entries at a time, so a backlog of long prompts never sits in memory at once, and the last three read the job and revision fields without decoding prompts. An entry written before those fields existed decodes its record. A record from before records named their job counts under its model, and is skipped when that model is no job of the record's pool. Results are saved through set_result, which keeps only the first, so a dispatcher answering the same entry is harmless. Each dispatcher pass drops the records it holds whose result was saved elsewhere, found with one pipelined EXISTS on their result keys, so it does not keep sending them or let them hold the pool on their revision. OPEN_RL_SAMPLE_DEADLINE_SECONDS sets the deadline, default 3600. Worker pods and pool dispatchers get it from the API server's environment. The dispatcher_body docstring no longer claims two dispatchers never overlap. An eviction or a partition can still overlap them, and set_result's transaction keeps each request to one result. Out of scope: retrieve_future's "try again" answer still says nothing about the pool's state. --- docs/configuration.md | 1 + src/server/api_server.py | 104 ++++++++++++++++++++++--- src/server/model_metadata.py | 18 +++-- src/server/sampling_dispatcher.py | 10 ++- src/server/sampling_store.py | 78 ++++++++++++++++++- src/server/scheduler_worker_manager.py | 40 ++++++++-- src/server/worker_manager.py | 10 +++ tests/test_pool_accept.py | 78 ++++++++++++++++++- tests/test_sampler_weight_rotation.py | 3 +- tests/test_sampling_store.py | 3 +- tests/test_session_lifecycle.py | 13 +++- 11 files changed, 323 insertions(+), 35 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index e1293d944..1654ab6da 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -120,4 +120,5 @@ for. On one machine every pool has one replica, in process. See | Env var | Default | What it does | | --- | --- | --- | | `OPEN_RL_DISPATCH_CONCURRENCY` | `256` | Requests one pool keeps in flight. | +| `OPEN_RL_SAMPLE_DEADLINE_SECONDS` | `3600` | How long a sampling request waits for a result before it fails. Set it on the API server. Worker pods and pool dispatchers get the same value. | | `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/src/server/api_server.py b/src/server/api_server.py index f596cd882..25684d61a 100644 --- a/src/server/api_server.py +++ b/src/server/api_server.py @@ -33,12 +33,14 @@ WeightSyncConfig, get_model_metadata, persist_model_metadata, + pool_name, ) +from server.sampling_dispatcher import SAMPLE_DEADLINE_SECONDS from server.sampling_store import SamplingStore from server.session_registry import SessionRegistry 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 server.worker_manager import LocalWorkerManager, WorkerManager, create_worker_manager, default_metadata, owner_id, owner_of, runtime_of from training import commands from training.commands import Command from training.types import TRAINER_BACKENDS, Datum, FFTConfig, FineTuningType, LoraConfig @@ -49,6 +51,8 @@ session_registry = SessionRegistry(state) SESSION_REAP_INTERVAL_SEC = 30 +# Jobs whose pools the reconcile may ensure at once. +RECONCILE_CONCURRENCY = 4 # Attaching a session to an owner and reaping that owner take turns, so a # session cannot attach between the reaper deciding an owner is unused and # deleting its workers. In-process, which is why there is one API server replica. @@ -77,6 +81,8 @@ async def reap_owner(owner: str) -> None: 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:trainer_ready:{model}") + if get_sampler_backend() == "vllm": + await SamplingStore(store.redis).delete_if_empty(pool_name(model)) await session_registry.forget(owner) @@ -640,6 +646,16 @@ def translate_future_result(result: dict) -> dict: return result +async def every_interval(step) -> None: + """Runs `step` every SESSION_REAP_INTERVAL_SEC. A failed step is retried next time.""" + while True: + await asyncio.sleep(SESSION_REAP_INTERVAL_SEC) + try: + await step() + except Exception: + traceback.print_exc() + + async def reap_dead_sessions(): while True: await asyncio.sleep(SESSION_REAP_INTERVAL_SEC) @@ -656,6 +672,63 @@ async def reap_dead_sessions(): traceback.print_exc() +async def reconcile_pools() -> None: + """For each job with pending requests whose pool is gone, recreates the pool when the job has a live session and fails + the requests when it has none, so no request waits on a pool nothing will bring back. Standing worker deployments + run without a worker manager and own their pools.""" + if worker_manager is None: + return + sampling = SamplingStore(store.redis) + pairs = [] + for pool in await sampling.pools(): + try: + pairs += [(pool, job) for job in await sampling.pending_jobs(pool)] + except Exception: + traceback.print_exc() + # Concurrently, since an ensure can wait minutes for an old pool to terminate, but a few at a time so stuck ensures + # cannot take every thread the API server runs blocking calls on. + slots = asyncio.Semaphore(RECONCILE_CONCURRENCY) + + async def reconcile(pool: str, job: str) -> None: + async with slots: + await reconcile_job(sampling, pool, job) + + for outcome in await asyncio.gather(*(reconcile(pool, job) for pool, job in pairs), return_exceptions=True): + if isinstance(outcome, Exception): + traceback.print_exception(outcome) + + +async def reconcile_job(sampling: SamplingStore, pool: str, job: str) -> None: + meta, runtime, _ = await asyncio.to_thread(runtime_of, job) + if meta.pool(job) != pool: + return # a record from before records named their job, whose model is no job on this pool + owner = owner_id(runtime) + # Under the owner lock, so a reaper release cannot land between the checks and the ensure and leave a pool no one removes. + async with owner_locks[owner]: + # Read again under the lock, since delete_model may have completed the job meanwhile. + if (current := await get_model_metadata(state, job)) is not None and current.status == "completed": + if failed := await sampling.fail_job(pool, job, {"type": "RequestFailedResponse", "error_message": f"{job} was deleted"}): + print(f"[API_SERVER] Failed {failed} sampling requests of {job}, which was deleted") + return + # A pool that is up serves its requests, since a client may sample without a session. + if await asyncio.to_thread(worker_manager.sampler_exists, job): + return + if await session_registry.in_use(owner): + await asyncio.to_thread(worker_manager.ensure, job, "sampler") + return + message = f"The pool of {job} is gone and its session ended before a sampler read the request" + if failed := await sampling.fail_job(pool, job, {"type": "RequestFailedResponse", "error_message": message}): + print(f"[API_SERVER] Failed {failed} sampling requests of {job}, whose pool is gone and whose session ended") + await sampling.delete_if_empty(pool) + + +async def sweep_sample_deadlines() -> None: + """Fails pool requests past the deadline, which no dispatcher does while the pool is gone.""" + result = {"type": "RequestFailedResponse", "error_message": f"No sampler read the request within {SAMPLE_DEADLINE_SECONDS}s"} + if failed := await SamplingStore(store.redis).fail_older_than(SAMPLE_DEADLINE_SECONDS, result): + print(f"[API_SERVER] Failed {failed} sampling requests no sampler read in time") + + @asynccontextmanager async def lifespan(_: FastAPI): global worker_manager @@ -680,11 +753,16 @@ async def lifespan(_: FastAPI): await asyncio.to_thread(worker.load_base_model, base_model) task = asyncio.create_task(training_requests_processor.run_training_requests_processor(worker)) reap_task = asyncio.create_task(reap_dead_sessions()) if worker_manager is not None else None + pool_tasks = [] + if get_sampler_backend() == "vllm": + # Separate tasks, so a reconcile stuck in an ensure never holds up the deadline sweep. + pool_tasks = [asyncio.create_task(every_interval(sweep_sample_deadlines)), asyncio.create_task(every_interval(reconcile_pools))] try: yield finally: - if reap_task is not None: - reap_task.cancel() + for background in (reap_task, *pool_tasks): + if background is not None: + background.cancel() if task is not None: task.cancel() if worker_manager is not None: @@ -912,12 +990,20 @@ 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 asyncio.to_thread(worker_manager.release, model_id) - now = time.time() - meta.status = "completed" - meta.completed_at = now - meta.updated_at = now - await persist_model_metadata(state, model_id, meta) + # Under the owner lock the reconcile takes, so it cannot bring back a pool for the job between these steps. + async with owner_locks[await asyncio.to_thread(owner_of, model_id)]: + now = time.time() + meta.status = "completed" + meta.completed_at = now + meta.updated_at = now + await persist_model_metadata(state, model_id, meta) + if get_sampler_backend() == "vllm": + result = {"type": "RequestFailedResponse", "error_message": f"{model_id} was deleted"} + await SamplingStore(store.redis).fail_job(meta.pool(model_id), model_id, result) + if worker_manager is not None and not meta.shares_runtime(): + await asyncio.to_thread(worker_manager.release, model_id) + if get_sampler_backend() == "vllm": + await SamplingStore(store.redis).delete_if_empty(meta.pool(model_id)) return {"status": "ok"} diff --git a/src/server/model_metadata.py b/src/server/model_metadata.py index 9fb2ce48f..0e57c8c45 100644 --- a/src/server/model_metadata.py +++ b/src/server/model_metadata.py @@ -87,13 +87,17 @@ 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's runtime. Base model names are not - valid Kubernetes names, so those runtimes get a slug and a hash.""" - runtime = self.runtime(model_id) - 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]}" + """The sampling pool that serves this job's runtime.""" + return pool_name(self.runtime(model_id)) + + +def pool_name(runtime: str) -> str: + """The sampling pool that serves a runtime. Base model names are not valid + Kubernetes names, so those runtimes get a slug and a hash.""" + 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/sampling_dispatcher.py b/src/server/sampling_dispatcher.py index c25e34a40..141fcc0ea 100644 --- a/src/server/sampling_dispatcher.py +++ b/src/server/sampling_dispatcher.py @@ -26,7 +26,7 @@ # The dispatcher wakes for each new request, or after this long to retry a failed send and look at its replicas. RETRY_MS = 1000 # A request no replica has answered for this long fails, so its client stops waiting. -SAMPLE_DEADLINE_SECONDS = 3600 +SAMPLE_DEADLINE_SECONDS = int(os.getenv("OPEN_RL_SAMPLE_DEADLINE_SECONDS", "3600")) # The version a replica reports once its weights may not match any revision. Never a path, so no pool revision equals it. UNKNOWN_WEIGHTS = "unknown" # How long one /weight_info call may take. @@ -85,7 +85,7 @@ class Dispatcher: def __init__(self, store: SamplingStore, pool: str, sampler: Any, concurrency: int) -> None: self.store, self.pool, self.sampler, self.concurrency = store, pool, sampler, concurrency self.revision = "" - # This dispatcher is the pool's only reader, so it reads each record once and drops it when its result is saved. + # Each record is read once and dropped when its result is saved. self.records: dict[str, dict[str, Any]] = {} self.newest = "0-0" self.in_flight: dict[str, asyncio.Task] = {} @@ -140,6 +140,12 @@ async def read(self) -> None: for r in await self.store.pending(self.pool, after=self.newest): self.records[r["entry_id"]] = r self.newest = r["entry_id"] + # The API server's deadline sweep and reconcile save results too, so a record they answered is dropped. + held = list(self.records.values()) + answered = await self.store.answered([r["request_id"] for r in held]) + for r in held: + if r["request_id"] in answered: + self.records.pop(r["entry_id"]) for r in list(self.records.values()): if (waited := time.time() - accepted_at(r)) > SAMPLE_DEADLINE_SECONDS: logger.warning("Failing %s after %.0fs without a result", r["request_id"], waited) diff --git a/src/server/sampling_store.py b/src/server/sampling_store.py index e5f8ce2d6..7335eb6e2 100644 --- a/src/server/sampling_store.py +++ b/src/server/sampling_store.py @@ -2,6 +2,7 @@ Results land in the `open_rl:future:` list that `retrieve_future` already reads.""" import json +import time from array import array from typing import Any @@ -10,6 +11,8 @@ RESULT_TTL_SECONDS = 300 # An accepted request may never get a result, so its idempotency key expires instead of staying until Redis is full. PENDING_TTL_SECONDS = 3 * 24 * 3600 +# The API server reads streams this many entries at a time, so a backlog of long prompts never sits in memory at once. +SWEEP_PAGE = 100 def stream_key(pool: str) -> str: @@ -40,13 +43,28 @@ def job_of(r: dict[str, Any]) -> str: return r.get("job") or r["model"] +def job_and_revision(fields: dict[str, str]) -> tuple[str, str]: + """An entry's job and revision without decoding its prompt. Entries accepted before the stream held both fields + decode the record.""" + if "job" in fields: + return fields["job"], fields["revision"] + record = json.loads(fields["record"], object_hook=compact_prompt) + return job_of(record), record["revision"] + + +def entry_record(entry_id: str, fields: dict[str, str]) -> dict[str, Any]: + return {**json.loads(fields["record"], object_hook=compact_prompt), "entry_id": entry_id} + + class SamplingStore: def __init__(self, client: redis.Redis): self.redis = client async def accept(self, record: dict[str, Any]) -> str: """The new request's id, or the id already stored under its idempotency key.""" - stream, fields = stream_key(record["pool"]), {"record": json.dumps(record)} + # The job and revision get fields of their own, so the API server reads them without decoding prompts. + fields = {"record": json.dumps(record), "job": job_of(record), "revision": record["revision"]} + stream = stream_key(record["pool"]) if not record.get("idempotency_key"): await self.redis.xadd(stream, fields) return record["request_id"] @@ -65,7 +83,7 @@ async def add(pipe) -> str: async def pending(self, pool: str, limit: int | None = 1000, after: str = "0-0") -> list[dict[str, Any]]: """Pending records in accept order. The dispatcher holds every one until its result is set, so prompts come back compact.""" entries = await self.redis.xrange(stream_key(pool), min=f"({after}", count=limit) - return [{**json.loads(fields["record"], object_hook=compact_prompt), "entry_id": entry_id} for entry_id, fields in entries] + return [entry_record(entry_id, fields) for entry_id, fields in entries] async def wait_for_new(self, pool: str, after: str, timeout_ms: int) -> None: await self.redis.xread({stream_key(pool): after}, count=1, block=timeout_ms) @@ -88,6 +106,60 @@ async def save(pipe) -> bool: # Every set_result writes the result key, so watching it lets only one through. return await self.redis.transaction(save, future, value_from_callable=True) + async def entries(self, stream: str, until: str = "+"): + """The stream's entries up to `until`, a page at a time.""" + after = "-" + while page := await self.redis.xrange(stream, min=after, max=until, count=SWEEP_PAGE): + for entry in page: + yield entry + after = f"({page[-1][0]}" + + async def fail_older_than(self, seconds: float, result: dict[str, Any]) -> int: + """Saves `result` for every pending request, in any pool, accepted more than `seconds` ago. Returns how many.""" + cutoff, failed = str(int((time.time() - seconds) * 1000)), 0 + async for stream in self.redis.scan_iter(match=stream_key("*"), _type="stream"): + async for entry_id, fields in self.entries(stream, cutoff): + failed += await self.set_result(entry_record(entry_id, fields), result) + return failed + + async def answered(self, request_ids: list[str]) -> set[str]: + """The requests among `request_ids` that have a result.""" + async with self.redis.pipeline(transaction=False) as pipe: + for request_id in request_ids: + pipe.exists(future_key(request_id)) + found = await pipe.execute() + return {request_id for request_id, exists in zip(request_ids, found) if exists} + + async def pools(self) -> list[str]: + """Every pool with a stream, whether or not it has pending requests.""" + return [stream.removeprefix(stream_key("")) async for stream in self.redis.scan_iter(match=stream_key("*"), _type="stream")] + + async def pending_jobs(self, pool: str) -> set[str]: + """The jobs with a pending request in the pool.""" + return {job_and_revision(fields)[0] async for _, fields in self.entries(stream_key(pool))} + + async def fail_job(self, pool: str, job: str, result: dict[str, Any]) -> int: + """Saves `result` for every pending request of `job` in `pool`. Returns how many.""" + failed = 0 + async for entry_id, fields in self.entries(stream_key(pool)): + if job_and_revision(fields)[0] == job: + failed += await self.set_result(entry_record(entry_id, fields), result) + return failed + + async def delete_if_empty(self, pool: str) -> bool: + """Deletes the pool's stream when nothing is pending in it, since set_result's XDEL leaves the empty stream behind + and Redis never evicts it. Watching the stream keeps a request accepted meanwhile. Returns whether it deleted.""" + stream = stream_key(pool) + + async def delete(pipe) -> bool: + if await pipe.xlen(stream): + return False + pipe.multi() + pipe.delete(stream) + return True + + return await self.redis.transaction(delete, stream, value_from_callable=True) + async def pinned_revisions(self, pool: str) -> set[str]: """Revisions that pending requests still name, so their weights must stay.""" - return {record["revision"] for record in await self.pending(pool, limit=None)} - {""} + return {job_and_revision(fields)[1] async for _, fields in self.entries(stream_key(pool))} - {""} diff --git a/src/server/scheduler_worker_manager.py b/src/server/scheduler_worker_manager.py index eec91070d..046ca59a8 100644 --- a/src/server/scheduler_worker_manager.py +++ b/src/server/scheduler_worker_manager.py @@ -146,7 +146,8 @@ def pod_env(worker: Worker) -> list[dict[str, Any]]: } # MAX_JOBS caps FlashInfer's JIT build, which otherwise runs one ~3GB # compiler per core and blows through the pod's host memory limit. - for name in ("VLLM_GPU_MEMORY_UTILIZATION", "VLLM_MAX_MODEL_LEN", "VLLM_ENABLE_MULTIMODAL", "OPEN_RL_TRAIN_TOKEN_BUDGET", "MAX_JOBS"): + passed = ("VLLM_GPU_MEMORY_UTILIZATION", "VLLM_MAX_MODEL_LEN", "VLLM_ENABLE_MULTIMODAL", "OPEN_RL_TRAIN_TOKEN_BUDGET", "MAX_JOBS") + for name in (*passed, "OPEN_RL_SAMPLE_DEADLINE_SECONDS"): if os.getenv(name): values[name] = os.environ[name] # vLLM keys its torch.compile artifacts by model, GPU type and engine config. @@ -247,8 +248,9 @@ def pod_template(worker: Worker) -> dict[str, Any]: 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.""" + EPP. Recreate stops the old pod before the new one starts. An eviction or a + partition can still overlap two dispatchers, and set_result's transaction keeps + each request to one result.""" labels = pool_labels("dispatcher", pool) env = [ {"name": "REDIS_URL", "value": os.environ["REDIS_URL"]}, @@ -256,8 +258,9 @@ def dispatcher_body(pool: str, owner: str) -> dict[str, Any]: {"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}) + for name in ("OPEN_RL_DISPATCH_CONCURRENCY", "OPEN_RL_SAMPLE_DEADLINE_SECONDS"): + if value := os.getenv(name): + env.append({"name": name, "value": value}) dispatcher = { "name": "dispatcher", # The dispatcher is a small CPU process; the API-server image carries it without the GPU stack. @@ -384,8 +387,8 @@ def ensure(self, model_id: str, role: str) -> None: 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. + # here keeps an old and a new dispatcher from overlapping in the usual case. + # 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: @@ -472,6 +475,29 @@ def release_owner(self, owner: str) -> set[str]: self.delete_workload(item["metadata"]["name"]) return {item["spec"]["modelID"] for item in ours} + def sampler_exists(self, model_id: str) -> bool: + """Whether the model's pool is up: a sampler Workload of its runtime not being deleted and, when the pool is routed, + its dispatcher Deployment.""" + _, runtime, _ = runtime_of(model_id) + 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) + samplers = [ + item + for item in found["items"] + if item["spec"].get("role") == "sampler" and item["spec"]["modelID"] == runtime and not item["metadata"].get("deletionTimestamp") + ] + if not samplers: + return False + pool = samplers[0]["metadata"].get("labels", {}).get("app.kubernetes.io/instance") + if pool is None: + return True + try: + return not self.apps_api.read_namespaced_deployment(pool, self.namespace).metadata.deletion_timestamp + except Exception as exc: + if getattr(exc, "status", None) != 404: + raise + return False + def close(self) -> None: pass # Workloads outlive the API server; the scheduler owns them from here diff --git a/src/server/worker_manager.py b/src/server/worker_manager.py index a9ed455aa..2f9c0777a 100644 --- a/src/server/worker_manager.py +++ b/src/server/worker_manager.py @@ -149,6 +149,10 @@ def release_owner(self, owner: str) -> set[str]: model ids they served, so the caller can clear their state.""" ... + def sampler_exists(self, model_id: str) -> bool: + """Whether any sampler of this model's pool exists, ready or not.""" + ... + def close(self) -> None: """The API server is exiting.""" ... @@ -216,6 +220,12 @@ def release(self, model_id: str) -> None: for key in [key for key in self.processes if key[1] in {runtime, model_id}]: self.terminate(key) + def sampler_exists(self, model_id: str) -> bool: + _, runtime, _ = runtime_of(model_id) + with self.lock: + proc = self.processes.get(("sampler", runtime)) + return proc is not None and proc.poll() is None + def release_owner(self, owner: str) -> set[str]: with self.lock: keys = [key for key in self.processes if owner_id(key[1]) == owner] diff --git a/tests/test_pool_accept.py b/tests/test_pool_accept.py index f8c133fcd..b6bb19476 100644 --- a/tests/test_pool_accept.py +++ b/tests/test_pool_accept.py @@ -11,11 +11,14 @@ 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.sampling_store import SamplingStore, stream_key +from server.scheduler_worker_manager import SchedulerWorkerManager +from server.session_registry import SessionRegistry from server.store import RedisStateStore, RedisStore -from server.worker_manager import default_metadata +from server.worker_manager import default_metadata, owner_of from tests.api_client import asgi_client from tests.redis_server import RedisTestCase +from tests.test_scheduler_worker_manager import FakeApi, FakeCustomObjectsApi class PoolAcceptTest(RedisTestCase): @@ -88,6 +91,77 @@ 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) + def on_kubernetes(self, workloads: FakeCustomObjectsApi | None = None) -> None: + """The API server with the scheduler's worker manager over fake Kubernetes APIs.""" + self.workloads, self.apps = workloads or FakeCustomObjectsApi(), FakeApi() + self.enterContext(patch.dict(os.environ, {"REDIS_URL": self.url})) + self.manager = SchedulerWorkerManager(custom_api=self.workloads, apps_api=self.apps, networking_api=FakeApi()) + self.enterContext(patch.object(api_server, "worker_manager", self.manager)) + self.enterContext(patch.object(api_server, "session_registry", SessionRegistry(self.state))) + self.enterContext(patch("server.worker_manager.get_state_store", return_value=self.state)) + + async def sampling_session(self, live: bool, **create) -> str: + """A sampling session on a pool the scheduler launched, opened by a client session that is still live or has ended.""" + client_session = (await self.client.post("/api/v1/create_session", json={})).json()["session_id"] + body = {"session_id": client_session, **create} + sampling = (await self.client.post("/api/v1/create_sampling_session", json=body)).json()["sampling_session_id"] + self.assertTrue(self.launched()) + if not live: + await self.state.delete_values(f"open_rl:session:{client_session}") + return sampling + + async def sampling_on_kubernetes(self, live: bool) -> str: + self.on_kubernetes() + return await self.sampling_session(live, base_model=self.base) + + def launched(self) -> set: + return set(self.workloads.existing) | set(self.apps.objects) + + def delete_pool(self) -> None: + # By hand, by a node pool going away, or by the reaper. A pass also serves other tests' leftover requests in the + # shared Redis, so checks look only at this test's pool. + self.workloads.existing.clear() + self.apps.objects.clear() + + async def answer(self, pool: str, count: int = 1) -> None: + sampling = SamplingStore(self.redis) + for r in (await sampling.pending(pool))[:count]: + await sampling.set_result(r, {"type": "sample", "sequences": []}) + + async def test_the_reaper_deletes_a_released_pools_stream_once_nothing_is_pending_in_it(self) -> None: + pool = default_metadata(self.base).pool(self.base) + session = await self.sampling_on_kubernetes(live=False) + await self.sample(0, session) + await self.sample(1, session) + await self.answer(pool) + await api_server.reap_owner(owner_of(self.base)) + self.assertEqual(len(await self.pending(pool, "request_id")), 1) + # The pool comes back for the next session, and once its last request is answered the reaper deletes the stream. + session = await self.sampling_session(live=False, base_model=self.base) + await self.answer(pool) + await api_server.reap_owner(owner_of(self.base)) + self.assertFalse(await self.redis.exists(stream_key(pool))) + + async def test_deleting_an_fft_job_deletes_its_pools_empty_stream(self) -> None: + self.on_kubernetes() + model_id = uuid.uuid4().hex + meta = TrainingModelMetadata(base_model=self.base, fine_tuning_type="full") + await persist_model_metadata(self.state, model_id, meta) + record = {"request_id": str(uuid.uuid4()), "pool": meta.pool(model_id), "job": model_id, "revision": "", "model": self.base} + await SamplingStore(self.redis).accept({**record, "idempotency_key": None, "prompt_token_ids": [1]}) + await self.answer(meta.pool(model_id)) + self.assertTrue(await self.redis.exists(stream_key(meta.pool(model_id)))) + self.assertEqual((await self.client.post("/api/v1/delete_model", json={"model_id": model_id})).status_code, 200) + self.assertFalse(await self.redis.exists(stream_key(meta.pool(model_id)))) + + async def test_a_pool_deleted_under_a_live_jobs_pending_requests_comes_back_without_another_sample(self) -> None: + session = await self.sampling_on_kubernetes(live=True) + launched = self.launched() + await self.sample(0, session) + self.delete_pool() + await api_server.reconcile_pools() + self.assertLessEqual(launched, self.launched()) + class PoolsNeedRedisTest(unittest.IsolatedAsyncioTestCase): def env(self, **values: str): diff --git a/tests/test_sampler_weight_rotation.py b/tests/test_sampler_weight_rotation.py index af2466e74..a6d25954a 100644 --- a/tests/test_sampler_weight_rotation.py +++ b/tests/test_sampler_weight_rotation.py @@ -53,7 +53,8 @@ async def test_a_version_a_pending_request_names_outlives_retention(self) -> Non 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} + pinned = {"request_id": "waiting", "pool": meta.pool("run-a"), "job": "run-a", "model": "base", "idempotency_key": None} + pinned["revision"] = os.path.join(versions, "sampler-1") await SamplingStore(self.redis).accept(pinned) proc = trp.TrainingRequestsProcessor(RedisStore(self.redis), RecordingWorker(), "run-a", time_slicer=SlicerStub()) for step in range(1, 6): diff --git a/tests/test_sampling_store.py b/tests/test_sampling_store.py index 21ccfd166..b9bdfeba1 100644 --- a/tests/test_sampling_store.py +++ b/tests/test_sampling_store.py @@ -12,7 +12,8 @@ async def asyncSetUp(self) -> None: self.store, self.pool = SamplingStore(self.redis), f"sampler-{uuid.uuid4()}" def record(self, revision: str = "", key: str | None = None) -> dict: - return {"request_id": str(uuid.uuid4()), "pool": self.pool, "revision": revision, "idempotency_key": key, "prompt_token_ids": [1]} + r = {"request_id": str(uuid.uuid4()), "pool": self.pool, "revision": revision, "model": "base", "idempotency_key": key} + return {**r, "prompt_token_ids": [1]} async def pending(self) -> dict[str, dict]: return {r["request_id"]: r for r in await self.store.pending(self.pool)} diff --git a/tests/test_session_lifecycle.py b/tests/test_session_lifecycle.py index 2518e6a8c..dd84f8b6d 100644 --- a/tests/test_session_lifecycle.py +++ b/tests/test_session_lifecycle.py @@ -75,6 +75,8 @@ async def reap(self): await api_server.reap_owner(owner) async def test_shared_lora_owner_outlives_the_session_that_created_it(self): + # Pools keep requests in Redis, which this test has none of, and reaping an owner deletes its pools' empty streams. + self.enterContext(patch.dict(os.environ, {"SAMPLING_BACKEND": "torch"})) training = (await self.post("create_session", {}))["session_id"] 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"}} @@ -82,7 +84,7 @@ async def test_shared_lora_owner_outlives_the_session_that_created_it(self): 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")]) + self.assertEqual(self.manager.ensured, [(adapter, "trainer"), (fft_model, "trainer")]) await self.expire(training) await self.reap() @@ -95,6 +97,8 @@ async def test_shared_lora_owner_outlives_the_session_that_created_it(self): self.assertEqual(await self.registry.owners(), []) async def test_an_exclusive_lora_model_has_its_own_runtime(self): + # Pools keep requests in Redis, which this test has none of, and deleting a job fails its pool requests. + self.enterContext(patch.dict(os.environ, {"SAMPLING_BACKEND": "torch"})) session = (await self.post("create_session", {}))["session_id"] shared = (await self.post("create_model", {"base_model": "test-base", "session_id": session}))["request_id"] exclusive_metadata = {"openrl.exclusive": "true"} @@ -111,6 +115,8 @@ async def test_an_exclusive_lora_model_has_its_own_runtime(self): self.assertEqual(self.manager.released_models, [exclusive]) async def test_deleting_a_job_frees_its_adapter_or_its_workers(self): + # Pools keep requests in Redis, which this test has none of, and deleting a job fails its pool requests. + self.enterContext(patch.dict(os.environ, {"SAMPLING_BACKEND": "torch"})) session = (await self.post("create_session", {}))["session_id"] shared = (await self.post("create_model", {"base_model": "test-base", "session_id": session}))["request_id"] body = {"base_model": "test-base", "session_id": session, "user_metadata": {"openrl.exclusive": "true"}} @@ -136,8 +142,9 @@ 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"})) + # Without FFT, a session on the bare base model shares the base model's LoRA owner. Pools keep requests in Redis, + # which this test has none of, and reaping an owner deletes its pools' empty streams. + self.enterContext(patch.dict(os.environ, {"OPEN_RL_ENABLE_FFT": "false", "SAMPLING_BACKEND": "torch"})) await self.registry.attach("a", "test-base") await self.expire("a") loop = asyncio.get_running_loop()