diff --git a/Makefile b/Makefile index 48d54c97..09d324ac 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_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 +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_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_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 # 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 da4e9cdb..fa8d5af6 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. | @@ -100,3 +100,25 @@ 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. Its samplers are plain +`vllm serve` replicas, and a dispatcher sends each request to the replica the +pool's router picks, then commits the result. Pending requests stay in a Redis +stream until a result is committed, so the Redis in `k8s/deploy/base` keeps an +append-only file on a persistent volume and never evicts keys. + +On Kubernetes each pool is `vllm serve` Workloads on exclusive GPUs plus 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 replica is ready only while vLLM reports the weight version +the dispatcher expects, and a NetworkPolicy lets only the pool's CPU pod reach it. +On one machine the local worker manager starts one `vllm serve` on +`SAMPLER_CUDA_VISIBLE_DEVICES` and a dispatcher that sends to it directly. + +| Env var | Default | What it does | +| --- | --- | --- | +| `OPEN_RL_DISPATCH_CONCURRENCY` | `256` | Requests one pool's dispatcher keeps in flight across its replicas. | +| `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/k8s/deploy/base/api-server.yaml b/k8s/deploy/base/api-server.yaml index de9f6c17..49f4fbe3 100644 --- a/k8s/deploy/base/api-server.yaml +++ b/k8s/deploy/base/api-server.yaml @@ -71,6 +71,12 @@ spec: configMapKeyRef: name: open-rl-images key: worker + - name: OPEN_RL_DISPATCHER_IMAGE + valueFrom: + configMapKeyRef: + name: open-rl-images + key: dispatcher + optional: true resources: limits: memory: "2Gi" diff --git a/k8s/deploy/base/images.yaml b/k8s/deploy/base/images.yaml index e6df6a32..72b19be0 100644 --- a/k8s/deploy/base/images.yaml +++ b/k8s/deploy/base/images.yaml @@ -1,3 +1,5 @@ images: - path: data/worker kind: ConfigMap + - path: data/dispatcher + kind: ConfigMap diff --git a/k8s/deploy/base/kustomization.yaml b/k8s/deploy/base/kustomization.yaml index 1c95858c..a5bbe7ea 100644 --- a/k8s/deploy/base/kustomization.yaml +++ b/k8s/deploy/base/kustomization.yaml @@ -7,6 +7,7 @@ namespace: openrl-system resources: - ../../../scheduler/deploy/base - rbac.yaml + - sampler-router.yaml - redis.yaml - api-server.yaml @@ -20,6 +21,7 @@ configMapGenerator: - name: open-rl-images literals: - worker=ghcr.io/gke-labs/open-rl/server:latest + - dispatcher=ghcr.io/gke-labs/open-rl/api-server:latest # Dynamic worker images receive the same image transformations as containers. configurations: diff --git a/k8s/deploy/base/rbac.yaml b/k8s/deploy/base/rbac.yaml index e23a8106..d0df1fac 100644 --- a/k8s/deploy/base/rbac.yaml +++ b/k8s/deploy/base/rbac.yaml @@ -1,5 +1,6 @@ # The API server requests one Workload per worker process; the scheduler owns the -# pods and claims from there, so this account needs only the Workload verbs. +# pods and claims from there. It also manages each sampling pool's CPU pod and +# lists sampler pods, because a sampler is ready only once its pod is Ready. apiVersion: rbac.authorization.k8s.io/v1 kind: Role metadata: @@ -8,6 +9,15 @@ rules: - apiGroups: ["openrl.io"] resources: ["workloads"] verbs: ["create", "get", "list", "watch", "delete"] +- apiGroups: ["apps"] + resources: ["deployments"] + verbs: ["create", "get", "delete"] +- apiGroups: ["networking.k8s.io"] + resources: ["networkpolicies"] + verbs: ["create", "delete"] +- apiGroups: [""] + resources: ["pods"] + verbs: ["list"] --- apiVersion: rbac.authorization.k8s.io/v1 kind: RoleBinding @@ -20,3 +30,30 @@ roleRef: kind: Role name: open-rl-worker-launcher apiGroup: rbac.authorization.k8s.io +--- +# A pool's dispatcher and EPP find its samplers by label, ready or not. +apiVersion: v1 +kind: ServiceAccount +metadata: + name: open-rl-sampler-dispatcher +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: Role +metadata: + name: open-rl-sampler-dispatcher +rules: +- apiGroups: [""] + resources: ["pods"] + verbs: ["get", "list", "watch"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: open-rl-sampler-dispatcher +subjects: +- kind: ServiceAccount + name: open-rl-sampler-dispatcher +roleRef: + kind: Role + name: open-rl-sampler-dispatcher + apiGroup: rbac.authorization.k8s.io diff --git a/k8s/deploy/base/sampler-router.yaml b/k8s/deploy/base/sampler-router.yaml new file mode 100644 index 00000000..e1aab085 --- /dev/null +++ b/k8s/deploy/base/sampler-router.yaml @@ -0,0 +1,87 @@ +# The router of every sampling pool's CPU pod. Envoy asks the pod's llm-d EPP +# (ext_proc on 9002) for a sampler and connects to the one the EPP names; it never +# routes without the EPP. The response names the sampler that answered, so the +# dispatcher can check its weight version before committing. The EPP keeps a request +# on the replicas that hold its prefix (precise matching from each pod's KV events) +# until their in-flight load would cost more time to first token than a cold replica, +# llm-d's "sticky until saturated". Prefix weighting alone piles long decodes onto one +# replica, and the slowest decode sets an RL step's time. No token producer is declared +# because the dispatcher sends token IDs to vLLM's /inference/v1/generate, and the one +# the EPP creates by default passes them through unchanged, so a vLLM renderer adds nothing. +apiVersion: v1 +kind: ConfigMap +metadata: + name: open-rl-sampler-router +data: + envoy.yaml: | + static_resources: + listeners: + - name: inference + address: {socket_address: {address: 127.0.0.1, port_value: 8081}} + filter_chains: + - filters: + - name: envoy.filters.network.http_connection_manager + typed_config: + "@type": type.googleapis.com/envoy.extensions.filters.network.http_connection_manager.v3.HttpConnectionManager + stat_prefix: inference + # A long generation sends nothing until it finishes, so no idle timeout may cut it off. + stream_idle_timeout: 0s + route_config: + virtual_hosts: + - name: samplers + domains: ["*"] + routes: + - match: {prefix: "/"} + route: {cluster: samplers, timeout: 0s, idle_timeout: 0s} + response_headers_to_add: + - header: {key: x-openrl-upstream, value: "%UPSTREAM_HOST%"} + http_filters: + - name: envoy.filters.http.ext_proc + typed_config: + "@type": type.googleapis.com/envoy.extensions.filters.http.ext_proc.v3.ExternalProcessor + grpc_service: + envoy_grpc: {cluster_name: epp, authority: localhost:9002} + timeout: 10s + processing_mode: + request_body_mode: FULL_DUPLEX_STREAMED + response_body_mode: FULL_DUPLEX_STREAMED + request_trailer_mode: SEND + response_trailer_mode: SEND + message_timeout: 1000s + - name: envoy.filters.http.router + typed_config: {"@type": type.googleapis.com/envoy.extensions.filters.http.router.v3.Router} + clusters: + - name: samplers + type: ORIGINAL_DST + lb_policy: CLUSTER_PROVIDED + connect_timeout: 10s + original_dst_lb_config: {use_http_header: true, http_header_name: x-gateway-destination-endpoint} + - name: epp + connect_timeout: 10s + transport_socket: + name: envoy.transport_sockets.tls + typed_config: + "@type": type.googleapis.com/envoy.extensions.transport_sockets.tls.v3.UpstreamTlsContext + common_tls_context: {validation_context: {}} + typed_extension_protocol_options: + envoy.extensions.upstreams.http.v3.HttpProtocolOptions: + "@type": type.googleapis.com/envoy.extensions.upstreams.http.v3.HttpProtocolOptions + explicit_http_config: {http2_protocol_options: {}} + load_assignment: + cluster_name: epp + endpoints: [{lb_endpoints: [{endpoint: {address: {socket_address: {address: 127.0.0.1, port_value: 9002}}}}]}] + epp.yaml: | + apiVersion: llm-d.ai/v1alpha1 + kind: EndpointPickerConfig + plugins: + - type: endpoint-notification-source + - type: precise-prefix-cache-producer + - type: inflight-load-producer + parameters: {prefixMatchInfoProducerName: precise-prefix-cache-producer} + - type: prefix-cache-affinity-filter + parameters: {prefixMatchInfoProducerName: precise-prefix-cache-producer} + - type: token-load-scorer + dataLayer: + sources: + - pluginRef: endpoint-notification-source + extractors: [{pluginRef: precise-prefix-cache-producer}] diff --git a/k8s/deploy/kind-dra/kustomization.yaml b/k8s/deploy/kind-dra/kustomization.yaml index fa302c9e..8680b650 100644 --- a/k8s/deploy/kind-dra/kustomization.yaml +++ b/k8s/deploy/kind-dra/kustomization.yaml @@ -143,6 +143,9 @@ patches: - name: OPEN_RL_WORKER_IMAGE valueFrom: null value: localhost:5001/open-rl-server:kind-dev + - name: OPEN_RL_DISPATCHER_IMAGE + valueFrom: null + value: localhost:5001/open-rl-api-server:kind-dev # The API server also stamps the pull policy into each Workload. - name: OPEN_RL_WORKER_IMAGE_PULL_POLICY value: Always diff --git a/scheduler/controller/Makefile b/scheduler/controller/Makefile index c6b2432f..3cbb651d 100644 --- a/scheduler/controller/Makefile +++ b/scheduler/controller/Makefile @@ -28,7 +28,7 @@ generate: $(CONTROLLER_GEN) # become the schema descriptions, so the prose lives in Go, not in the YAML. .PHONY: manifests manifests: $(CONTROLLER_GEN) - $(CONTROLLER_GEN) crd paths="./api/..." output:crd:dir=$(GEN_DIR)/crd + $(CONTROLLER_GEN) crd:generateEmbeddedObjectMeta=true paths="./api/..." output:crd:dir=$(GEN_DIR)/crd cp $(GEN_DIR)/crd/openrl.io_workloads.yaml $(CRD_MANIFEST) cp $(GEN_DIR)/crd/openrl.io_claimledgers.yaml $(CRD_GROUP_MANIFEST) @@ -36,7 +36,7 @@ manifests: $(CONTROLLER_GEN) # the case where the generator cannot be run on the machine holding the diff. .PHONY: manifests-check manifests-check: $(CONTROLLER_GEN) - $(CONTROLLER_GEN) crd paths="./api/..." output:crd:dir=$(GEN_DIR)/crd + $(CONTROLLER_GEN) crd:generateEmbeddedObjectMeta=true paths="./api/..." output:crd:dir=$(GEN_DIR)/crd diff -u $(CRD_MANIFEST) $(GEN_DIR)/crd/openrl.io_workloads.yaml diff -u $(CRD_GROUP_MANIFEST) $(GEN_DIR)/crd/openrl.io_claimledgers.yaml diff --git a/scheduler/deploy/base/00-workload-crd.yaml b/scheduler/deploy/base/00-workload-crd.yaml index 63edb92e..2042ce6b 100644 --- a/scheduler/deploy/base/00-workload-crd.yaml +++ b/scheduler/deploy/base/00-workload-crd.yaml @@ -156,6 +156,23 @@ spec: description: |- Standard object's metadata. More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#metadata + properties: + annotations: + additionalProperties: + type: string + type: object + finalizers: + items: + type: string + type: array + labels: + additionalProperties: + type: string + type: object + name: + type: string + namespace: + type: string type: object spec: description: |- @@ -7189,6 +7206,23 @@ spec: May contain labels and annotations that will be copied into the PVC when creating it. No other fields are allowed and will be rejected during validation. + properties: + annotations: + additionalProperties: + type: string + type: object + finalizers: + items: + type: string + type: array + labels: + additionalProperties: + type: string + type: object + name: + type: string + namespace: + type: string type: object spec: description: |- diff --git a/skills/openrl/references/lifecycle.md b/skills/openrl/references/lifecycle.md index d71eaf4f..4ed08776 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 85e7c17b..bc69ab1d 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) @@ -452,15 +453,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, @@ -588,13 +580,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." ) @@ -655,6 +645,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(): @@ -666,7 +657,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 @@ -907,7 +897,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" @@ -1155,26 +1144,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"} @@ -1189,7 +1165,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: @@ -1197,7 +1174,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, } @@ -1210,6 +1187,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 "") @@ -1220,16 +1226,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) @@ -1272,44 +1275,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 82fad34b..15ad2fd2 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 eb345a6e..3a1ddd45 100644 --- a/src/server/model_metadata.py +++ b/src/server/model_metadata.py @@ -85,18 +85,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 16acf057..9a1a45f6 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 new file mode 100644 index 00000000..6b69c0cb --- /dev/null +++ b/src/server/sampling_dispatcher.py @@ -0,0 +1,273 @@ +"""Sends one sampling pool's pending records through its router, Envoy in front of the llm-d EPP, and commits the first result.""" + +import asyncio +import contextlib +import json +import logging +import os +import random +import signal +import time +import uuid +from pathlib import Path +from typing import Any + +import httpx +import redis.asyncio as redis + +from server.sampling_store import SamplingStore + +logger = logging.getLogger(__name__) + +# 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 +# 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" + + +def tmp_root() -> str: + """OPEN_RL_TMP_DIR, or the documented default. The API server renders the readiness probe's path + from it and hands the same value to the dispatcher and the replicas, so they all agree.""" + return os.getenv("OPEN_RL_TMP_DIR", "/tmp/open-rl") + + +def revision_file(pool: str) -> Path: + """The readiness probe compares vLLM's weight version with the revision in this file.""" + return Path(tmp_root(), "sampling", pool, "revision") + + +def accepted_at(r: dict[str, Any]) -> float: + """When the record's request was accepted. A stream id starts with that time in milliseconds.""" + return int(r["entry_id"].split("-")[0]) / 1000 + + +def failed(message: str) -> dict[str, str]: + return {"type": "RequestFailedResponse", "error_message": message} + + +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} + 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 + + +class Dispatcher: + """Sends through `router`. A run on one machine has no router and lists its replicas in `urls`.""" + + def __init__(self, store: SamplingStore, pool: str, http: httpx.AsyncClient, router: str | None, urls: list[str], concurrency: int) -> None: + self.store, self.pool, self.http, self.router, self.urls, self.concurrency = store, pool, http, router, urls, concurrency + self.revision, self.k8s = "", None + # This dispatcher is the pool's only committer, so it reads each record once and drops it on commit. + self.records: dict[str, dict[str, Any]] = {} + self.newest = "0-0" + self.in_flight: dict[str, asyncio.Task] = {} + self.loading: dict[str, asyncio.Task] = {} + self.paused: dict[str, float] = {} # when each replica last paused to load weights + self.checked: set[str] = set() # replicas checked for a pause or an open update, checked again after a load of ours breaks off + + async def run(self) -> None: + self.write_revision() + try: + while True: + try: + await self.dispatch() + await self.store.wait_for_new(self.pool, self.newest, RETRY_MS) + except Exception: + # A Redis or Kubernetes hiccup must not end the pool's dispatching. + logger.exception("Dispatching pool %s failed; retrying", self.pool) + await asyncio.sleep(RETRY_MS / 1000) + finally: + for task in self.in_flight.values(): + task.cancel() + # A load cut off midway leaves its replica paused, so loads get a moment to finish and resume before the HTTP client closes. + if self.loading: + _, pending = await asyncio.wait(list(self.loading.values()), timeout=5) + for task in pending: + task.cancel() + + async def dispatch(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"] + 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) + await self.commit(r, failed(f"No sampler answered in {waited:.0f}s; pool {self.pool} serves {self.revision or 'base weights'}")) + # A send still waiting on a replica that never answers would keep its slot forever. + if task := self.in_flight.get(r["request_id"]): + task.cancel() + records = list(self.records.values()) + # Nothing pending names the current revision, so nothing is in flight on it. + if records and all(r["revision"] != self.revision for r in records): + logger.info("Pool %s switching from %r to %r", self.pool, self.revision, records[0]["revision"]) + self.revision = records[0]["revision"] + self.write_revision() + target = self.revision or "default" + urls = await self.list_replicas() + versions = await asyncio.gather(*(self.inspect(url) for url in urls), return_exceptions=True) + for url, version in zip(urls, versions): + # A replica still starting does not answer yet, and a pool on base weights never loads any. + if self.revision and isinstance(version, str) and version != target and url not in self.loading: + self.loading[url] = asyncio.create_task(self.load_weights(url, self.revision)) + if target in versions: + # Read after the awaits above, so a request that committed meanwhile is not sent again. + to_send = [r for r in self.records.values() if r["revision"] == self.revision and r["request_id"] not in self.in_flight] + for r in to_send[: self.concurrency - len(self.in_flight)]: + self.in_flight[r["request_id"]] = task = asyncio.create_task(self.send(r)) + task.add_done_callback(lambda _, request_id=r["request_id"]: self.in_flight.pop(request_id)) + + async def commit(self, r: dict[str, Any], result: dict[str, Any]) -> None: + await self.store.commit(r, result) + self.records.pop(r["entry_id"], None) + + async def send(self, r: dict[str, Any]) -> None: + target, sent = self.revision or "default", time.monotonic() + params = {**r["sampling_params"], "logprobs": 0, "prompt_logprobs": 0 if r["include_prompt_logprobs"] else None} + 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} + url = self.router or random.choice(self.urls) + try: + # A replica 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()) + response = await self.http.post(f"{url}/inference/v1/generate", json=body, timeout=timeout) + # Envoy names the replica that answered. Its weights must still be the target, with no pause to load others since this was sent. + replica = f"http://{name}" if (name := response.headers.get("x-openrl-upstream")) else url + if response.status_code in (400, 404, 422): + await self.commit(r, failed(f"Sampler rejected the request ({response.status_code}): {response.text[:500]}")) + elif response.status_code == 200 and await self.weight_version(replica) == target and self.paused.get(replica, 0.0) < sent: + try: + result = to_result(response.json(), r) + except (KeyError, TypeError, ValueError) as exc: + result = failed(f"Sampler response broke the contract: {exc!r}") + await self.commit(r, result) + except httpx.HTTPError as exc: + logger.info("Retrying %s after %r", r["request_id"], exc) + + async def load_weights(self, url: str, revision: str) -> None: + try: + if not (path := Path(revision, "metadata.json")).is_file(): + # Retention can delete a revision between the API server's check and its pin. + for r in [r for r in self.records.values() if r["revision"] == revision]: + await self.commit(r, failed(f"Sampler weights {revision} no longer exist")) + return + metadata, current = json.loads(path.read_text()), await self.weight_version(url) + # A delta applies on any save at the weights it starts from or reaches. Deltas written before the trainer listed them name one parent. + bases = [base or "default" for base in metadata.get("bases", [metadata.get("parent")])] + if metadata.get("format") == "sparse_delta" and current not in bases: + raise RuntimeError(f"{url} holds {current!r}, not a revision delta {revision!r} applies on") + self.paused[url] = time.monotonic() + # A load that breaks off may leave the replica paused or inside its update, so it is checked again unless this one ends cleanly. + self.checked.discard(url) + await self.post(url, "/pause?mode=wait&clear_cache=false") + updating = True + try: + await self.post(url, "/start_weight_update") + await self.post(url, "/update_weights", {"update_info": {"target_weights_path": revision}}) + await self.post(url, "/finish_weight_update", {"weight_version": revision}) + updating = False + except Exception: + await self.forget_weights(url) + updating = False + raise + else: + if not (await self.post(url, "/reset_prefix_cache")).json()["success"]: + raise RuntimeError("vLLM did not reset the prefix cache") + finally: + # A load cut off inside its update leaves the replica paused, so nothing is served from part of a revision + # before a dispatcher finds the open update and forgets the weights. + if not updating: + await self.post(url, "/resume") + self.checked.add(url) + except Exception as exc: + logger.warning("Loading %r on %s failed: %r", revision, url, exc) # the next dispatch tries again + finally: + self.loading.pop(url) + + async def inspect(self, url: str) -> str: + """The version a replica holds. On first sight any weight update an earlier load left open is closed and its weights + forgotten, paused or not, since a dispatcher stopped mid-update used to resume the replica. A dispatcher killed + mid-load, or a failed /resume, leaves a replica paused, where the router would queue requests forever, so one found + paused outside this dispatcher's loads is resumed.""" + if url not in self.checked and url not in self.loading: + await self.forget_weights(url, if_updating=True) + if (await self.http.get(f"{url}/is_paused", timeout=5)).raise_for_status().json()["is_paused"]: + logger.info("Resuming %s, left paused by an earlier load", url) + await self.post(url, "/resume") + self.checked.add(url) + return await self.weight_version(url) + + async def forget_weights(self, url: str, if_updating: bool = False) -> None: + """Closes an update a cut-off load left active, which refuses every later start, and has the replica report + UNKNOWN_WEIGHTS until a full load, since it may hold part of a revision. vLLM has no abort, but a bare finish closes + an active update and is refused when none is, so with `if_updating` a replica that was not updating keeps its version.""" + # It waits behind an update a dead dispatcher left running, but a wedged replica must not hold up the pool. + response = await self.http.post(f"{url}/finish_weight_update", timeout=300) + if if_updating and response.status_code != 200: + return + logger.warning("Weights on %s may be partial, so it reports %r until a full load replaces them", url, UNKNOWN_WEIGHTS) + await self.post(url, "/update_weight_version", {"new_version": UNKNOWN_WEIGHTS}) + + async def post(self, url: str, path: str, body: dict[str, Any] | None = None) -> httpx.Response: + return (await self.http.post(url + path, json=body, timeout=None)).raise_for_status() + + async def list_replicas(self) -> list[str]: + if self.urls: + return self.urls + if self.k8s is None: + from kubernetes import client, config + + config.load_incluster_config() + self.k8s = client.CoreV1Api() + selector = f"app.kubernetes.io/component=sampler,app.kubernetes.io/instance={self.pool}" + namespace = os.getenv("OPEN_RL_WORKER_NAMESPACE", "openrl-system") + pods = await asyncio.to_thread(self.k8s.list_namespaced_pod, namespace, label_selector=selector, _request_timeout=10) + return [f"http://{p.status.pod_ip}:8000" for p in pods.items if p.status.pod_ip and p.status.phase == "Running"] + + async def weight_version(self, url: str) -> str: + return (await self.http.get(f"{url}/weight_info", timeout=5)).raise_for_status().json()["weight_version"] + + def write_revision(self) -> None: + path = revision_file(self.pool) + path.parent.mkdir(parents=True, exist_ok=True) + path.with_suffix(".tmp").write_text(self.revision or "default") + path.with_suffix(".tmp").rename(path) + + +def redis_client(url: str) -> redis.Redis: + """When replicas come up, every in-flight send commits at once. Past max_connections a + commit waits for a free connection; redis-py's default pool of 100 raises instead.""" + return redis.Redis.from_pool(redis.BlockingConnectionPool.from_url(url, decode_responses=True, max_connections=64, timeout=30)) + + +def http_client() -> httpx.AsyncClient: + """Sends have no connection limit, so a burst goes out at once. httpx scans its whole pool for each idle + connection whenever a request starts or ends, so the hundreds a burst leaves idle stall this loop for + seconds. Few are kept.""" + return httpx.AsyncClient(limits=httpx.Limits(max_connections=None, max_keepalive_connections=32)) + + +async def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") + logging.getLogger("httpx").setLevel(logging.WARNING) + urls = [url for url in os.getenv("OPEN_RL_SAMPLER_ENDPOINTS", "").split(",") if url] + router = None if urls else os.getenv("OPEN_RL_SAMPLER_ROUTER", "http://127.0.0.1:8081") + store = SamplingStore(redis_client(os.environ["REDIS_URL"])) + async with http_client() as http: + dispatcher = Dispatcher(store, os.environ["OPEN_RL_SAMPLER_POOL"], http, router, urls, int(os.getenv("OPEN_RL_DISPATCH_CONCURRENCY", "256"))) + task = asyncio.create_task(dispatcher.run()) + for sig in (signal.SIGTERM, signal.SIGINT): + asyncio.get_running_loop().add_signal_handler(sig, task.cancel) + with contextlib.suppress(asyncio.CancelledError): + await task + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/src/server/sampling_store.py b/src/server/sampling_store.py index d5c85815..00c37e06 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 +from array import array from typing import Any import redis.asyncio as redis @@ -15,6 +16,17 @@ def stream_key(pool: str) -> str: return f"open_rl:sampling:{pool}" +def compact_prompt(record: dict[str, Any]) -> dict[str, Any]: + """A json.loads object_hook that holds a prompt as array('I'), four bytes a token against about 36 as a list of ints, + and drops each list as soon as its record is read. A prompt array('I') cannot hold stays the list, for vLLM to refuse.""" + if isinstance(tokens := record.get("prompt_token_ids"), list): + try: + record["prompt_token_ids"] = array("I", tokens) + except (OverflowError, TypeError): + pass + return record + + def idempotency_key(pool: str, key: str) -> str: return f"open_rl:sampling:{pool}:idem:{key}" @@ -45,9 +57,13 @@ async def add(pipe) -> str: return await self.redis.transaction(add, key, value_from_callable=True) - async def pending(self, pool: str, limit: int | None = 1000) -> list[dict[str, Any]]: - entries = await self.redis.xrange(stream_key(pool), count=limit) - return [{**json.loads(fields["record"]), "entry_id": entry_id} for entry_id, fields in entries] + 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 it commits, 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] + + 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) async def commit(self, record: dict[str, Any], result: dict[str, Any]) -> bool: """Save the first final result and drop the pending entry. False if it was already committed.""" diff --git a/src/server/scheduler_worker_manager.py b/src/server/scheduler_worker_manager.py index 29743bed..7ef8194e 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,6 +29,9 @@ 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, @@ -48,6 +52,10 @@ 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" + 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 +83,9 @@ class Worker: # Which replica of its role this is. index: int = 0 + # The sampling pool a sampler serves. + pool: str | None = None + @property def owner(self) -> str: return owner_id(self.runtime) @@ -86,28 +97,38 @@ def name(self) -> str: 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.""" + is its own Workload in the model's sampling pool.""" 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) - exclusive = not meta.shares_gpu() + pool = meta.pool(runtime) if role == "sampler" else None + # A pool sampler keeps its GPUs between batches; nothing time-slices it. + exclusive = not meta.shares_gpu() or role == "sampler" 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, @@ -127,7 +148,9 @@ def pod_env(worker: Worker) -> list[dict[str, Any]]: 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 @@ -177,11 +200,107 @@ 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"}}) + # A pool sampler is plain `vllm serve`, ready only while it serves the version the pool's dispatcher expects. + if worker.pool: + container = template["spec"]["containers"][0] + 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)] + 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) -> 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"}}, + "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) -> 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"}}, + "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: @@ -190,10 +309,11 @@ def accelerator_spec(worker: Worker) -> dict[str, Any]: def workload_body(worker: Worker) -> dict[str, Any]: + labels = {"app.kubernetes.io/managed-by": "open-rl-api-server", **(pool_labels("sampler", worker.pool) if worker.pool 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", @@ -210,21 +330,62 @@ 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") 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: + worker = describe_worker(model_id, role) + if worker.pool 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(replicas_of(model_id, role)): - self.ensure_workload(dataclasses.replace(describe_worker(model_id, role), index=index)) + self.ensure_workload(dataclasses.replace(worker, index=index)) + if worker.pool: + self.ensure_dispatcher(worker.pool) + + def ensure_dispatcher(self, pool: 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)): + 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)) + + 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 @@ -264,10 +425,18 @@ 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, so a retry still finds it through the samplers.""" 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] + for pool in {item["metadata"]["labels"].get("app.kubernetes.io/instance") for item in ours} - {None}: + 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} @@ -298,11 +467,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( @@ -311,6 +483,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"], @@ -325,12 +508,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 90896400..468baacc 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_future(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_future(self, req_id: str, result: dict[str, Any]) -> None: if result.get("status") == "pending": return @@ -234,18 +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 [] diff --git a/src/server/training_requests_processor.py b/src/server/training_requests_processor.py index 70fb1cf5..ed33ca40 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,15 @@ 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) + pinned = await self.sampling.pinned_revisions(meta.pool(model_id)) + 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 a62a6dbb..84540064 100644 --- a/src/server/vllm_options.py +++ b/src/server/vllm_options.py @@ -1,7 +1,8 @@ -"""Engine and sampling options shared by the two vLLM sampler workers.""" +"""Engine and sampling options for the vLLM samplers.""" +import json import os -from collections.abc import Sequence +from collections.abc import Mapping, Sequence def text_only_engine_kwargs() -> dict: @@ -52,34 +53,21 @@ def gpu_memory_utilization() -> float: return 0.90 -def sampler_batch_limits(device_memory_bytes: int | None = None) -> dict[str, int]: - """vLLM's concurrency for this device: sequences in flight and prefill tokens per step. - - A sweep's sampler serves many groups at once, and the defaults (64 - sequences, 2048 prefill tokens) leave most of an 80 GB card idle. Explicit - VLLM_MAX_NUM_SEQS / VLLM_MAX_NUM_BATCHED_TOKENS win; otherwise size by memory. - """ - if device_memory_bytes is None: - device_memory_bytes = detected_device_memory_bytes() - gib = device_memory_bytes / 2**30 - if gib >= 60: - seqs, tokens = 256, 16384 - elif gib >= 36: - seqs, tokens = 128, 8192 - else: - seqs, tokens = 64, 4096 - return { - "max_num_seqs": int(os.getenv("VLLM_MAX_NUM_SEQS", seqs)), - "max_num_batched_tokens": int(os.getenv("VLLM_MAX_NUM_BATCHED_TOKENS", tokens)), - } - - -def detected_device_memory_bytes() -> int: - try: - import torch - - if torch.cuda.is_available(): - return int(torch.cuda.get_device_properties(0).total_memory) - except Exception: - pass - return 0 +def serve_args(base_model: str, is_lora: bool, env: Mapping[str, str | None], port: int = 8000) -> list[str]: + """`vllm serve` flags for a sampling pool replica.""" + args = [base_model, "--port", str(port), "--api-server-count", "1", "--enable-scale-out", "--enable-prefix-caching"] + # A checkpoint's generation_config.json could add sampling params the client never asked for, such as a repetition penalty. + args += ["--generation-config", "vllm", "--max-model-len", os.getenv("VLLM_MAX_MODEL_LEN", "8192")] + # vLLM's own default of 0.92 has run FFT delta applies out of memory on an L4. + args += ["--gpu-memory-utilization", str(gpu_memory_utilization())] + # These args are rendered where the replica's device is unknown, so vLLM sizes the batch for the device it runs on unless told. + for name in ("VLLM_MAX_NUM_SEQS", "VLLM_MAX_NUM_BATCHED_TOKENS"): + if os.getenv(name): + args += ["--" + name.removeprefix("VLLM_").lower().replace("_", "-"), os.environ[name]] + if text_only := text_only_engine_kwargs(): + args += ["--limit-mm-per-prompt", json.dumps(text_only["limit_mm_per_prompt"])] + if architecture := env.get("VLLM_ARCHITECTURE_OVERRIDE"): + args += ["--hf-overrides", json.dumps({"architectures": [architecture]})] + if is_lora: + return args + ["--enable-lora", "--max-loras", os.getenv("VLLM_MAX_LORAS", "8"), "--max-lora-rank", os.getenv("VLLM_MAX_LORA_RANK", "64")] + return args + ["--weight-transfer-config", '{"backend":"delta_snapshot"}'] diff --git a/src/server/vllm_sampler.py b/src/server/vllm_sampler.py deleted file mode 100644 index bc0dba96..00000000 --- a/src/server/vllm_sampler.py +++ /dev/null @@ -1,399 +0,0 @@ -"""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).""" - -import argparse -import asyncio -import hashlib -import os -import signal -import time -from collections.abc import Callable -from contextlib import asynccontextmanager -from itertools import groupby -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 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") - 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")), - **sampler_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(), - } - 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 - - -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 - - 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) - - -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 - self.update_failed = False - - async def wake(self) -> None: - await self.engine.wake_up(tags=["weights", "kv_cache"]) - - 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.""" - 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. - 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.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 - - -@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.""" - if time_slicer is None: - yield - return - async with time_slicer.acquire(workload): - if sampler is not None: - await sampler.wake() - try: - yield - finally: - if sampler is not None: - 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_future(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_future(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_future(request["request_id"], result) - - -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) - - -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: - 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 = [] - 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) - else: - await serve(model_id, get_store(), make_engine) - except asyncio.CancelledError: - pass # serve has completed cleanup before cancellation reaches here. - finally: - for sig in installed_signals: - loop.remove_signal_handler(sig) - - -def main() -> None: - parser = argparse.ArgumentParser(description="Open-RL vLLM Pull-Mode Sampler Worker") - parser.add_argument("--model-id", type=str, required=True, help="The model ID of the RL job to process requests for") - args = parser.parse_args() - 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() - - -if __name__ == "__main__": - main() diff --git a/src/server/worker_manager.py b/src/server/worker_manager.py index 9ca9c8d4..f72d7f9a 100644 --- a/src/server/worker_manager.py +++ b/src/server/worker_manager.py @@ -6,6 +6,7 @@ import os import re import shutil +import socket import subprocess import sys import threading @@ -13,10 +14,12 @@ 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.workload import TRAINER_TIME_SLICE_GROUP, workload_job_id from server.estimator import footprint +from server.lora_snapshots import snapshot_root from server.model_metadata import TrainingModelMetadata, decode_model_metadata from server.store import get_state_store +from server.vllm_options import serve_args PROJECT_DIR = Path(__file__).resolve().parents[2] @@ -108,16 +111,24 @@ 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" + # Triton kernels persist on the shared volume, so later replicas skip compiling them, + # and Triton renames each kernel into place from a unique temp dir, so concurrent + # writers are safe. vLLM's compile cache is not shared, because vLLM names its AOT + # temp file by PID alone and replicas in different pods can run the same PID. + env["TRITON_CACHE_DIR"] = str(Path(os.getenv("OPEN_RL_TMP_DIR", "/tmp/open-rl"), "cache", "triton")) + if is_lora: + # vLLM loads each frozen adapter snapshot by name from this directory. + env["VLLM_ALLOW_RUNTIME_LORA_UPDATING"] = "true" + env["VLLM_LORA_RESOLVER_CACHE_DIR"] = str(snapshot_root()) if not is_lora and ("gemma-4" in base_model.lower() or "gemma4" in base_model.lower()): env["VLLM_ARCHITECTURE_OVERRIDE"] = "Gemma4ForCausalLM" return env def worker_module(role: str) -> str: - return "server.training_requests_processor" if role == "trainer" else "server.vllm_sampler" + return "server.training_requests_processor" def worker_args(runtime: str, role: str, is_lora: bool) -> list[str]: @@ -165,6 +176,8 @@ def __init__(self, project_dir: Path = PROJECT_DIR): self.project_dir = project_dir self.processes: dict[tuple[str, str], subprocess.Popen] = {} self.started: dict[tuple[str, str], float] = {} + # Each local pool's dispatcher command and env, to start it again while its vLLM runs on. + self.dispatchers: dict[str, tuple[list[str], dict[str, str]]] = {} self.lock = threading.Lock() def ensure(self, model_id: str, role: str) -> None: @@ -174,23 +187,51 @@ def ensure(self, model_id: str, role: str) -> None: with self.lock: proc = self.processes.get((role, runtime)) if proc is not None and proc.poll() is None: + dispatcher = self.processes.get(("dispatcher", runtime)) + if role == "sampler" and (dispatcher is None or dispatcher.poll() is not None): + # A pool without its dispatcher answers nothing, so a dead one starts again against the running vLLM. + self.spawn("dispatcher", runtime, *self.dispatchers[runtime]) 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"], worker_module(role), worker_args(runtime, role, is_lora)), env) + return + # A local pool is one vllm serve replica and the dispatcher that sends to it directly. The dispatcher of a vLLM + # that died is stopped before a new pair starts, so two dispatchers never read one stream. + if (old := self.processes.get(("dispatcher", runtime))) is not None: + self.terminate(("dispatcher", runtime)) + try: + old.wait(timeout=10) + except subprocess.TimeoutExpired: + old.kill() + old.wait() + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + port = probe.getsockname()[1] + vllm = ["uv", "run", "--extra", "gpu", "--extra", "vllm", "vllm"] if shutil.which("uv") else [str(Path(sys.executable).with_name("vllm"))] + if is_lora: + # vLLM's filesystem LoRA resolver exits at startup without its directory, and no adapter may be frozen yet. + snapshot_root().mkdir(parents=True, exist_ok=True) + self.spawn(role, runtime, [*vllm, "serve", *serve_args(base_model, is_lora, env, port)], env) + dispatcher_env = {**os.environ, "OPEN_RL_SAMPLER_POOL": meta.pool(runtime), "OPEN_RL_SAMPLER_ENDPOINTS": f"http://127.0.0.1:{port}"} + # The dispatcher keeps the pool's revision file under the same root the API server uses. + dispatcher_env["OPEN_RL_TMP_DIR"] = os.getenv("OPEN_RL_TMP_DIR", "/tmp/open-rl") + self.dispatchers[runtime] = (python_command([], "server.sampling_dispatcher", []), dispatcher_env) + self.spawn("dispatcher", runtime, *self.dispatchers[runtime]) + + 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/src/training/fft_trainer_worker.py b/src/training/fft_trainer_worker.py index c547ab09..6e69e150 100644 --- a/src/training/fft_trainer_worker.py +++ b/src/training/fft_trainer_worker.py @@ -50,6 +50,11 @@ def __init__(self): self._prev_weights_shadow: dict[str, torch.Tensor] = {} self.model_layer_shapes: dict[str, tuple[int, ...]] = {} self.total_model_elements: int = 0 + # Sampler saves at the current weights, oldest first. "" names the base weights a replica starts with. + self.sampler_saves: list[str] = [""] + # Sampler saves at the weights the last optimizer step started from, which its sparse delta applies on. + # Before any step the delta is empty, and the list is empty when no save holds those weights. + self.delta_from: list[str] = [""] def set_weight_sync_strategy(self, strategy: str) -> None: if strategy not in ("full", "delta"): @@ -178,7 +183,8 @@ def save_state(self, model_id: str, state_path: str, include_optimizer: bool = F # Under the delta strategy save_state writes the sparse delta the sampler # consumes. load_from_state cannot open it, so FFT is not resumable yet. - if self.weight_sync_cfg.strategy == "delta": + # A sampler save that no replica could apply the delta on is written in full. + if self.weight_sync_cfg.strategy == "delta" and (kind != "sampler" or self.delta_from): if kind != "sampler": logger.warning("save_state for %s under the delta strategy writes a delta, not a resumable checkpoint", model_id) return self.save_state_delta(model_id=model_id, state_path=state_path, kind=kind) @@ -271,6 +277,9 @@ def save_state_delta( "total_elements": total_elements, "layer_names": layer_names_list, "density_pct": round(100.0 * total_changed / max(1, total_elements), 3), + # The last sampler save before the last optimizer step. A replica holding any of the bases can apply this delta. + "parent": self.delta_from[-1] if self.delta_from else None, + "bases": list(dict.fromkeys(self.delta_from + self.sampler_saves)), "timestamp": time.time(), } with open(os.path.join(state_path, "metadata.json"), "w") as f: @@ -286,6 +295,7 @@ def save_for_sampler(self, model_id: str, alias: str | None, ref: str | None) -> rel_path = ref[len("tinker://") :] if ref.startswith("tinker://") else ref.lstrip("/") local_path = os.path.join(os.getenv("OPEN_RL_TMP_DIR", "/tmp/open-rl"), "sampler_full", rel_path) self.save_state(model_id, local_path, False, "sampler") + self.sampler_saves.append(local_path) return local_path def load_from_state(self, model_id: str, state_path: str, restore_optimizer: bool = False) -> dict[str, Any]: @@ -307,6 +317,8 @@ def load_from_state(self, model_id: str, state_path: str, restore_optimizer: boo target_device = "auto" if num_gpus > 1 else self.device self.model = AutoModelForCausalLM.from_pretrained(state_path, dtype=dtype, device_map=target_device) self.prepare_model_for_training() + # No sampler save holds the loaded weights. + self.sampler_saves, self.delta_from = [], [] if restore_optimizer and metadata.get("has_optimizer"): optimizer_path = os.path.join(state_path, "optimizer.pt") @@ -372,6 +384,8 @@ def optim_step(self, adam_params: dict[str, Any], model_id: str | None = None) - self.optimizer.zero_grad() t_step_end = time.perf_counter() step_time = t_step_end - t_step_start + # The weights moved, so the sampler saves made so far hold the weights this step's delta starts from. + self.delta_from, self.sampler_saves = self.sampler_saves, [] delta_compute_time = 0.0 if self.weight_sync_cfg.strategy == "delta" and self.model is not None and hasattr(self.model, "named_parameters"): diff --git a/tests/test_lora_snapshots.py b/tests/test_lora_snapshots.py index 4543ac66..cdf2d51b 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 bd4a7f9c..e52ba564 100644 --- a/tests/test_openrl_status.py +++ b/tests/test_openrl_status.py @@ -10,13 +10,15 @@ 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 from server.store import InMemoryStateStore, InMemoryStore from server.worker_manager import WORKER_STATES from tests.api_client import asgi_client, post_json -from tests.test_scheduler_worker_manager import FakeCustomObjectsApi +from tests.test_scheduler_worker_manager import ApiError, FakeCustomObjectsApi from tests.test_session_lifecycle import RuntimeManager @@ -165,16 +167,11 @@ 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): + async def test_a_sampler_is_ready_once_its_workload_runs(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): @@ -236,8 +233,28 @@ 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"})) @@ -275,6 +292,57 @@ 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") + + def test_a_ready_pool_sampler_is_starting_while_its_pools_dispatcher_has_no_available_pod(self): + # Its readiness probe compares vLLM's version with the revision file, which outlives a dead or OOM-killed + # dispatcher, and without the dispatcher no sample is sent. + running = {"phase": "Running", "podName": "orw-s", "nodeName": "node-b"} + for index in range(2): + sampler = workload(f"lora-base-{index}-sampler", "sampler", "base", "base", running) + sampler["metadata"]["labels"]["app.kubernetes.io/instance"] = "sampler-base" + self.custom.existing[sampler["metadata"]["name"]] = sampler + no_pod = "Running on node-b, dispatcher sampler-base has no available pod" + cases = [({}, "starting", no_pod), ({"sampler-base": 0}, "starting", no_pod), ({"sampler-base": 1}, "ready", "Running on node-b")] + for dispatchers, state, detail in cases: + with self.subTest(dispatchers=dispatchers): + apps = FakeAppsApi(dispatchers) + manager = SchedulerWorkerManager(custom_api=self.custom, core_api=FakeCoreApi(sampler_pod("orw-s", ready=True)), apps_api=apps) + samplers = [(w["state"], w["detail"]) for w in manager.inspect_owner("base") if w["role"] == "sampler"] + self.assertEqual((samplers, apps.reads), ([(state, detail)] * 2, ["sampler-base"])) + + +class FakeAppsApi: + """Each pool's dispatcher Deployment by name, with its available pod count. A missing one is a 404.""" + + def __init__(self, available: dict[str, int]): + self.available, self.reads = available, [] + + def read_namespaced_deployment(self, name: str, namespace: str): + self.reads.append(name) + if name not in self.available: + raise ApiError(404) + return client.V1Deployment(status=client.V1DeploymentStatus(available_replicas=self.available[name])) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_pool_accept.py b/tests/test_pool_accept.py index f2da4fa9..f8c133fc 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 ee7ec943..60e42b8f 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_batch_limits.py b/tests/test_sampler_batch_limits.py deleted file mode 100644 index 5f16fb78..00000000 --- a/tests/test_sampler_batch_limits.py +++ /dev/null @@ -1,27 +0,0 @@ -import os -import unittest -from unittest.mock import patch - -from server.vllm_options import sampler_batch_limits - -GIB = 2**30 - - -class SamplerBatchLimitsTest(unittest.TestCase): - def test_limits_follow_the_device(self) -> None: - with patch.dict(os.environ, {}, clear=False): - os.environ.pop("VLLM_MAX_NUM_SEQS", None) - os.environ.pop("VLLM_MAX_NUM_BATCHED_TOKENS", None) - self.assertEqual(sampler_batch_limits(80 * GIB), {"max_num_seqs": 256, "max_num_batched_tokens": 16384}) - self.assertEqual(sampler_batch_limits(40 * GIB), {"max_num_seqs": 128, "max_num_batched_tokens": 8192}) - self.assertEqual(sampler_batch_limits(24 * GIB), {"max_num_seqs": 64, "max_num_batched_tokens": 4096}) - # No device visible: the old defaults' tier. - self.assertEqual(sampler_batch_limits(0)["max_num_seqs"], 64) - - def test_explicit_settings_win(self) -> None: - with patch.dict(os.environ, {"VLLM_MAX_NUM_SEQS": "32", "VLLM_MAX_NUM_BATCHED_TOKENS": "1024"}): - self.assertEqual(sampler_batch_limits(80 * GIB), {"max_num_seqs": 32, "max_num_batched_tokens": 1024}) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_sampler_patch.py b/tests/test_sampler_patch.py deleted file mode 100644 index e8d22cf9..00000000 --- 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_future.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_future.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_future.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 43d136f4..af2466e7 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_sampling_dispatcher.py b/tests/test_sampling_dispatcher.py new file mode 100644 index 00000000..4871ad0a --- /dev/null +++ b/tests/test_sampling_dispatcher.py @@ -0,0 +1,590 @@ +import asyncio +import inspect +import json +import random +import tempfile +import tracemalloc +import uuid +from array import array +from pathlib import Path +from unittest.mock import patch + +import httpx +import redis.asyncio as redis +import safetensors.torch +import torch +import uvicorn +from redis.asyncio.retry import Retry +from redis.backoff import NoBackoff +from starlette.applications import Starlette +from starlette.responses import JSONResponse, Response +from starlette.routing import Route + +from server.sampling_dispatcher import UNKNOWN_WEIGHTS, Dispatcher, http_client, redis_client, revision_file, to_result +from server.sampling_store import SamplingStore, future_key +from tests.redis_server import RedisTestCase +from tests.test_exclusive_trainer import tiny_worker +from training.fft_trainer_worker import FFTTrainingWorker +from training.types import Datum, TensorData + +SAMPLE = {"type": "sample", "sequences": [{"tokens": [100, 101], "logprobs": [-0.5, -0.25], "stop_reason": "length"}]} +CHOICE = {"token_ids": [100, 101], "finish_reason": "length", "logprobs": {"content": [{"logprob": -0.5}, {"logprob": -0.25}]}} + + +async def until(condition, timeout: float = 10) -> None: + async with asyncio.timeout(timeout): + while not (await result if inspect.isawaitable(result := condition()) else result): + await asyncio.sleep(0.02) + + +def fft_worker() -> FFTTrainingWorker: + """A real FFT trainer on a tiny Llama, resident on the CPU, publishing sparse deltas.""" + worker = tiny_worker() + worker.cpu_offload = False + worker.set_weight_sync_strategy("delta") + worker.prepare_model_for_training() + return worker + + +def train_step(worker: FFTTrainingWorker) -> None: + row = [3, 5, 8, 13, 21, 34] + datum = Datum(model_input=row[:-1], loss_fn_inputs={"target_tokens": TensorData(data=row[1:]), "weights": TensorData(data=[1.0] * 5)}) + worker.forward_backward([datum], "cross_entropy") + worker.optim_step({"learning_rate": 1e-2}) + + +def nudge_step(worker: FFTTrainingWorker, index: int) -> None: + """An optimizer step that moves one weight. In bf16 at a small learning rate each step moves a different few, so a + delta of the last step misses what the step before moved.""" + with torch.no_grad(): + worker.model.lm_head.weight.view(-1)[index] += 1.0 + worker.optim_step({}) + + +def weights(worker: FFTTrainingWorker) -> dict[str, torch.Tensor]: + return {name: param.detach().clone() for name, param in worker.model.named_parameters()} + + +def replayed(start: dict[str, torch.Tensor], revisions: list[str]) -> dict[str, torch.Tensor]: + """A replica's weights after it loads `revisions` in order, written the way the delta_snapshot engine writes them.""" + out = {name: tensor.clone() for name, tensor in start.items()} + for revision in revisions: + metadata = json.loads(Path(revision, "metadata.json").read_text()) + if metadata.get("format") == "sparse_delta": + tensors = safetensors.torch.load_file(str(Path(revision, "delta.safetensors"))) + for i, name in enumerate(metadata["layer_names"]): + out[name].view(-1)[tensors[f"{i}.indices"].long()] = tensors[f"{i}.values"] + else: + for file in Path(revision).glob("*.safetensors"): + out.update((name, tensor) for name, tensor in safetensors.torch.load_file(str(file)).items() if name in out) + return out + + +def same(a: dict[str, torch.Tensor], b: dict[str, torch.Tensor]) -> bool: + return a.keys() == b.keys() and all(torch.equal(a[name], b[name]) for name in a) + + +class FakeReplica: + """vLLM's weight endpoints and generate with a real weight version. `held` pauses both; an update logs what the readiness probe expects. + Like vLLM, a paused replica queues generate until it is resumed, one weight update is active from start to finish, + and only a finish that names a version, or /update_weight_version, changes the version. `prompts` keeps the token ids sent.""" + + def __init__(self, name: str, log: list, version: str, expected: Path): + self.name, self.log, self.version, self.expected, self.held, self.url = name, log, version, expected, asyncio.Event(), "" + self.resumed, self.probes, self.paused, self.updating, self.fail_update, self.prompts = 0, 0, False, False, False, [] + gets = [Route("/weight_info", lambda r: JSONResponse({"weight_version": self.version}))] + gets.append(Route("/is_paused", lambda r: JSONResponse({"is_paused": self.paused}))) + posts = {"/pause": self.pause, "/resume": self.resume, "/start_weight_update": self.start, "/update_weight_version": self.set_version} + posts["/reset_prefix_cache"] = lambda r: JSONResponse({"success": True}) + posts |= {"/update_weights": self.update, "/finish_weight_update": self.finish, "/inference/v1/generate": self.generate} + self.app = Starlette(routes=gets + [Route(path, handler, methods=["POST"]) for path, handler in posts.items()]) + + async def start(self, request) -> Response: + if self.updating: + return Response("start_weight_update called while a weight update is already active. Call finish_weight_update first.", status_code=500) + self.updating = True + return JSONResponse({}) + + async def update(self, request) -> Response: + if not self.updating: + return Response("start_weight_update must be called before update_weights.", status_code=500) + self.log.append(("update", self.name, (await request.json())["update_info"]["target_weights_path"], self.expected.read_text())) + if self.fail_update: # like vLLM, a failed update_weights ends the update and keeps the version + self.fail_update, self.updating = False, False + return Response("Target weights path does not exist", status_code=500) + await until(lambda: not self.held.is_set()) + return JSONResponse({}) + + async def pause(self, request) -> Response: + self.paused = True + return JSONResponse({}) + + async def resume(self, request) -> Response: + self.resumed, self.paused = self.resumed + 1, False + return JSONResponse({}) + + async def finish(self, request) -> Response: + body = await request.body() + self.probes += not body # a bare finish, which vLLM refuses with a 500 traceback when no update is open + if not self.updating: + return Response("finish_weight_update called without a matching start_weight_update.", status_code=500) + self.updating = False + if body: + self.version = json.loads(body)["weight_version"] + return JSONResponse({}) + + async def set_version(self, request) -> Response: + self.version = (await request.json())["new_version"] + return JSONResponse({}) + + async def generate(self, request) -> Response: + self.log.append(("generate", self.name, self.version)) + self.prompts.append((await request.json())["token_ids"]) + await until(lambda: not self.held.is_set() and not self.paused) + return JSONResponse({"choices": [CHOICE]}) + + +def router(log: list, replicas: list[FakeReplica], http: httpx.AsyncClient) -> Starlette: + # Forwards generate to a ready replica, one whose version is the expected one, like the EPP. + async def route(request) -> Response: + ready = [r for r in replicas if r.version == r.expected.read_text()] + if not ready: + log.append(("unrouted", "", "")) + return Response(status_code=503) + upstream = await http.post(f"{ready[0].url}/inference/v1/generate", content=await request.body(), timeout=None) + return Response(upstream.content, upstream.status_code, headers={"x-openrl-upstream": ready[0].url.removeprefix("http://")}) + + return Starlette(routes=[Route("/inference/v1/generate", route, methods=["POST"])]) + + +class SamplingDispatcherTest(RedisTestCase): + async def asyncSetUp(self) -> None: + await super().asyncSetUp() + self.store, self.pool, self.log, self.http = SamplingStore(self.redis), f"sampler-{uuid.uuid4()}", [], httpx.AsyncClient() + self.addAsyncCleanup(self.http.aclose) + self.tmp = Path(self.enterContext(tempfile.TemporaryDirectory())) + self.enterContext(patch.dict("os.environ", {"OPEN_RL_TMP_DIR": str(self.tmp)})) + + async def serve(self, app: Starlette) -> str: + server = uvicorn.Server(uvicorn.Config(app, port=0, log_level="warning", lifespan="off")) + serving = asyncio.create_task(server.serve()) + self.addAsyncCleanup(lambda: (setattr(server, "should_exit", True), serving)[1]) + await until(lambda: server.started) + return f"http://127.0.0.1:{server.servers[0].sockets[0].getsockname()[1]}" + + async def replicas(self, *versions: str) -> list[FakeReplica]: + replicas = [FakeReplica(f"pod-{i}", self.log, version, revision_file(self.pool)) for i, version in enumerate(versions)] + for replica in replicas: + replica.url = await self.serve(replica.app) + self.addAsyncCleanup(replica.held.clear) + return replicas + + def revision(self, name: str, **metadata) -> str: + (self.tmp / name).mkdir() + (self.tmp / name / "metadata.json").write_text(json.dumps(metadata)) + return str(self.tmp / name) + + 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: + 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} + await self.store.accept(r) + return r + + async def dispatcher(self, replicas: list[FakeReplica], concurrency: int = 4) -> asyncio.Task: + router_url = await self.serve(router(self.log, replicas, self.http)) + dispatcher = Dispatcher(self.store, self.pool, self.http, router_url, [r.url for r in replicas], concurrency=concurrency) + run = asyncio.create_task(dispatcher.run()) + self.addAsyncCleanup(lambda: (run.cancel(), asyncio.gather(run, return_exceptions=True))[1]) + return run + + async def committed(self, request_id: str) -> list[dict]: + await until(lambda: self.redis.llen(future_key(request_id))) + await asyncio.sleep(0.3) # a second commit would show by now + return [json.loads(raw) for raw in await self.redis.lrange(future_key(request_id), 0, -1)] + + async def test_each_request_is_generated_once(self): + replicas = await self.replicas("default", "default") + records = [await self.accepted() for _ in range(40)] + await self.dispatcher(replicas, concurrency=256) + for r in records: + await until(lambda r=r: self.redis.llen(future_key(r["request_id"]))) + await asyncio.sleep(0.5) # a duplicate would be generated by now + self.assertEqual(sum(entry[0] == "generate" for entry in self.log), len(records)) + + async def test_a_burst_of_answers_commits_each_once(self): + replicas = await self.replicas("default", "default") + records = [await self.accepted() for _ in range(300)] + # The dispatcher's own clients, as main() builds them: 256 sends answered together commit together. + self.store = SamplingStore(await self.enterAsyncContext(redis_client(self.url))) + self.http = await self.enterAsyncContext(http_client()) + await self.dispatcher(replicas, concurrency=256) + for r in records: + await until(lambda r=r: self.redis.llen(future_key(r["request_id"])), timeout=60) + await asyncio.sleep(0.5) # a resend after a failed commit would generate again by now + # A send the test router drops is retried, so a few may generate twice. Commits that + # find no free connection fail and resend about half the burst. + self.assertLess(sum(entry[0] == "generate" for entry in self.log) - len(records), len(records) // 20) + + async def test_a_burst_over_many_replicas_leaves_at_most_32_idle_connections(self): + # httpx scans every pooled connection whenever a request starts or ends, so hundreds left idle stall the dispatcher. + replicas = await self.replicas(*["default"] * 8) + for replica in replicas: + replica.held.set() # every send stays in flight, each on a connection of its own, until all are out + records = [await self.accepted() for _ in range(100)] + # The dispatcher's own clients, as main() builds them. + self.store = SamplingStore(await self.enterAsyncContext(redis_client(self.url))) + self.http = await self.enterAsyncContext(http_client()) + dispatcher = Dispatcher(self.store, self.pool, self.http, None, [r.url for r in replicas], concurrency=256) + run = asyncio.create_task(dispatcher.run()) + self.addAsyncCleanup(lambda: (run.cancel(), asyncio.gather(run, return_exceptions=True))[1]) + await until(lambda: sum(entry[0] == "generate" for entry in self.log) == len(records), timeout=30) + for replica in replicas: + replica.held.clear() + for r in records: + await until(lambda r=r: self.redis.llen(future_key(r["request_id"])), timeout=60) + idle = [connection for connection in self.http._transport._pool.connections if connection.is_idle()] + self.assertLessEqual(len(idle), 32) + + async def test_a_request_no_replica_answers_fails_at_the_deadline(self): + replicas = await self.replicas("other-weights") # never ready for the base weights, so nothing is sent + self.enterContext(patch("server.sampling_dispatcher.SAMPLE_DEADLINE_SECONDS", 1)) + r = await self.accepted() + await self.dispatcher(replicas) + [result] = await self.committed(r["request_id"]) + self.assertEqual(result["type"], "RequestFailedResponse") + self.assertIn("No sampler answered", result["error_message"]) + self.assertEqual(await self.store.pending(self.pool), []) + + async def test_sends_to_a_replica_that_never_answers_give_back_their_slots_when_their_records_fail(self): + hung, healthy = await self.replicas("default", "default") + hung.held.set() # it takes generate and never answers, like a wedged GPU or a node gone without closing TCP + routed = [] + + # Like the EPP, the first two picks land on the hung replica and every later pick on the healthy one. + async def route(request) -> Response: + target = hung if len(routed) < 2 else healthy + routed.append(target.name) + upstream = await self.http.post(f"{target.url}/inference/v1/generate", content=await request.body(), timeout=None) + return Response(upstream.content, upstream.status_code, headers={"x-openrl-upstream": target.url.removeprefix("http://")}) + + router_url = await self.serve(Starlette(routes=[Route("/inference/v1/generate", route, methods=["POST"])])) + first = [await self.accepted() for _ in range(2)] + dispatcher = Dispatcher(self.store, self.pool, self.http, router_url, [hung.url, healthy.url], concurrency=2) + run = asyncio.create_task(dispatcher.run()) + self.addAsyncCleanup(lambda: (run.cancel(), asyncio.gather(run, return_exceptions=True))[1]) + await until(lambda: self.log.count(("generate", "pod-0", "default")) == 2) + await asyncio.sleep(1.1) + # The deadline passes while both sends still wait. + with patch("server.sampling_dispatcher.SAMPLE_DEADLINE_SECONDS", 1): + for r in first: + [result] = await self.committed(r["request_id"]) + self.assertEqual(result["type"], "RequestFailedResponse") + # The slots free at once, not when the hung replica lets go some seconds later. + await until(lambda: not dispatcher.in_flight, timeout=2) + later = await self.accepted() + self.assertEqual(await self.committed(later["request_id"]), [SAMPLE]) + self.assertEqual(routed, ["pod-0", "pod-0", "pod-1"]) + + async def test_a_send_gives_up_at_its_records_deadline(self): + (hung,) = await self.replicas("default") + hung.held.set() + self.enterContext(patch("server.sampling_dispatcher.SAMPLE_DEADLINE_SECONDS", 1)) + await self.accepted() + [r] = await self.store.pending(self.pool) + dispatcher = Dispatcher(self.store, self.pool, self.http, None, [hung.url], concurrency=1) + await asyncio.wait_for(dispatcher.send(r), 5) + self.assertEqual(self.log, [("generate", "pod-0", "default")]) + + async def test_the_dispatcher_holds_a_pending_prompt_in_a_few_bytes_a_token(self): + # It keeps every record until it commits, inside a 2Gi limit, so a long-context backlog must stay small once read + # and while it is read. A list of ints costs about 36 bytes a token. + prompt, count = [random.randrange(1000, 150000) for _ in range(20_000)], 50 + for _ in range(count): + await self.accepted(prompt=prompt) + dispatcher = Dispatcher(self.store, self.pool, self.http, None, ["http://127.0.0.1:9"], concurrency=4) # nothing answers + tracemalloc.start() + self.addCleanup(tracemalloc.stop) + before, _ = tracemalloc.get_traced_memory() + tracemalloc.reset_peak() + await dispatcher.dispatch() + after, peak = tracemalloc.get_traced_memory() + held, peak_held = (after - before) / (count * len(prompt)), (peak - before) / (count * len(prompt)) + self.assertEqual(len(dispatcher.records), count) + self.assertLess(held, 8, f"{held:.1f} bytes a token once read") + self.assertLess(peak_held, 20, f"{peak_held:.1f} bytes a token while read") + + async def test_a_prompt_kept_compact_is_sent_and_scored_as_the_client_sent_it(self): + (replica,) = replicas = await self.replicas("default") + # A negative token does not fit array('I'), so that prompt stays the list the client sent, for vLLM to refuse. + records = [await self.accepted(), await self.accepted(prompt=[-1, 2])] + compact, as_sent = await self.store.pending(self.pool) + self.assertEqual((type(compact["prompt_token_ids"]), as_sent["prompt_token_ids"]), (array, [-1, 2])) + body = {"choices": [CHOICE], "prompt_logprobs": [None, {"2": {"logprob": -0.5}}, {"3": {"logprob": -0.25}}]} + self.assertEqual(to_result(body, {**compact, "include_prompt_logprobs": True})["prompt_logprobs"], [None, -0.5, -0.25]) + await self.dispatcher(replicas) + for r in records: + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + self.assertEqual(sorted(replica.prompts), [[-1, 2], [1, 2, 3]]) + + async def test_a_redis_timeout_does_not_stop_dispatching(self): + replicas = await self.replicas("default") + self.enterContext(patch("server.sampling_dispatcher.RETRY_MS", 100)) + impatient = await self.enterAsyncContext(redis.from_url(self.url, decode_responses=True, socket_timeout=0.3, retry=Retry(NoBackoff(), 0))) + patient, self.store = self.store, SamplingStore(impatient) + run = await self.dispatcher(replicas) + self.store = patient + await self.redis.execute_command("CLIENT", "PAUSE", "1000", "ALL") + r = await self.accepted() # waits out the pause; the dispatcher's reads time out meanwhile + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + self.assertFalse(run.done()) + + async def test_a_replica_that_is_still_starting_fails_nothing(self): + (replica,) = replicas = await self.replicas("default") + starting = "http://127.0.0.1:9" # nothing listens here yet + r = await self.accepted() + router_url = await self.serve(router(self.log, replicas, self.http)) + dispatcher = Dispatcher(self.store, self.pool, self.http, router_url, [starting, replica.url], concurrency=4) + run = asyncio.create_task(dispatcher.run()) + self.addAsyncCleanup(lambda: (run.cancel(), asyncio.gather(run, return_exceptions=True))[1]) + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + + async def test_crash_mid_request_commits_exactly_one_result(self): + (replica,) = replicas = await self.replicas("default") + replica.held.set() + r = await self.accepted() + first = await self.dispatcher(replicas) + await until(lambda: self.log.count(("generate", "pod-0", "default")) == 1) + first.cancel() + await self.dispatcher(replicas) + replica.held.clear() + self.assertEqual((await self.committed(r["request_id"]), await self.store.pending(self.pool)), ([SAMPLE], [])) + + async def test_an_answer_from_the_wrong_weights_is_not_committed_and_the_pod_is_prepared_again(self): + (replica,) = replicas = await self.replicas(w1 := self.revision("w1")) + replica.held.set() + r = await self.accepted(w1) + await self.dispatcher(replicas) + await until(lambda: ("generate", "pod-0", w1) in self.log) + replica.version = "default" # restarted in place while the request ran + replica.held.clear() + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + self.assertEqual(self.log, [("generate", "pod-0", w1), ("update", "pod-0", w1, w1), ("generate", "pod-0", w1)]) + + async def test_fft_switch_expects_the_new_version_before_any_weight_update_and_sends_only_to_a_ready_pod(self): + w1, w2 = self.revision("w1"), self.revision("w2") + r = await self.accepted(w2) + await self.dispatcher(await self.replicas(w1, w1)) + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + self.assertEqual(sorted(self.log[:2]), [("update", "pod-0", w2, w2), ("update", "pod-1", w2, w2)]) + self.assertEqual([(entry[0], entry[2]) for entry in self.log[2:]], [("generate", w2)]) + + async def test_a_deleted_revision_fails_its_requests_and_the_pool_moves_on(self): + (replica,) = replicas = await self.replicas("default") + gone, later = await self.accepted(str(self.tmp / "gone")), await self.accepted(self.revision("w2", format="sparse_delta", parent="")) + await self.dispatcher(replicas) + failed = {"type": "RequestFailedResponse", "error_message": f"Sampler weights {self.tmp / 'gone'} no longer exist"} + self.assertEqual(await self.committed(gone["request_id"]), [failed]) + self.assertEqual((await self.committed(later["request_id"]), replica.version), ([SAMPLE], later["revision"])) + + async def test_a_replica_left_paused_at_the_pool_revision_is_resumed_on_startup_and_serves(self): + (replica,) = replicas = await self.replicas(w1 := self.revision("w1")) + replica.paused = True # the previous dispatcher died after this replica loaded w1 and before its /resume + r = await self.accepted(w1) + await self.dispatcher(replicas) + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + self.assertEqual((replica.resumed, replica.paused, self.loads("pod-0")), (1, False, {"pod-0": []})) + + async def test_a_replica_left_paused_at_an_older_revision_is_resumed_and_then_loaded(self): + w1 = self.revision("w1", format="sparse_delta", parent="") + w2 = self.revision("w2", format="sparse_delta", parent=w1) + (replica,) = replicas = await self.replicas(w1) + replica.paused = True + r = await self.accepted(w2) + await self.dispatcher(replicas) + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + self.assertEqual((replica.resumed, replica.paused, replica.version, self.loads("pod-0")), (2, False, w2, {"pod-0": [w2]})) + + async def test_a_replica_left_inside_a_weight_update_takes_a_full_load_before_it_serves(self): + w1, w2 = self.revision("w1"), self.revision("w2") + (replica,) = replicas = await self.replicas(w1) + # The previous dispatcher was killed after /start_weight_update for w2, so the replica may hold w1, w2 or part of w2. + replica.paused, replica.updating = True, True + r = await self.accepted(w2) + await self.dispatcher(replicas) + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + self.assertEqual((replica.version, replica.updating, replica.paused), (w2, False, False)) + self.assertEqual([entry[:3] for entry in self.log], [("update", "pod-0", w2), ("generate", "pod-0", w2)]) + + async def test_a_replica_left_inside_a_weight_update_at_the_pool_revision_does_not_serve_until_reloaded(self): + w1 = self.revision("w1") + self.revision("w2") + (replica,) = replicas = await self.replicas(w1) + # The pool still names w1 while the replica, cut off loading w2, may hold w2's changes under w1's version. + replica.paused, replica.updating = True, True + r = await self.accepted(w1) + await self.dispatcher(replicas) + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + self.assertEqual((replica.version, replica.updating), (w1, False)) + self.assertEqual([entry[:3] for entry in self.log], [("update", "pod-0", w1), ("generate", "pod-0", w1)]) + + async def test_a_replica_resumed_inside_a_weight_update_has_it_closed_before_any_load(self): + w1, w2 = self.revision("w1"), self.revision("w2") + (replica,) = replicas = await self.replicas(w1) + replica.updating = True # a dispatcher cut off between start and finish resumed the replica before it died, so nothing marks it paused + r = await self.accepted(w2) + with self.assertLogs("server.sampling_dispatcher", "WARNING") as logs: + await self.dispatcher(replicas) + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + self.assertNotIn("start_weight_update", "\n".join(logs.output)) # no load starts into the open update and is refused + self.assertEqual((self.loads("pod-0"), replica.version, replica.updating), ({"pod-0": [w2]}, w2, False)) + + async def test_a_replica_resumed_inside_a_weight_update_at_the_pool_revision_does_not_serve_until_reloaded(self): + w1 = self.revision("w1") + self.revision("w2") + (replica,) = replicas = await self.replicas(w1) + # A dispatcher stopped while loading w2 resumed the replica, which may hold part of w2 under w1's version. + replica.updating = True + r = await self.accepted(w1) + await self.dispatcher(replicas) + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + self.assertEqual([entry[:3] for entry in self.log], [("update", "pod-0", w1), ("generate", "pod-0", w1)]) + self.assertEqual((replica.version, replica.updating), (w1, False)) + + async def test_a_load_cut_off_inside_its_update_stays_paused_until_the_next_dispatcher_reloads_the_replica(self): + w1, w2 = self.revision("w1"), self.revision("w2") + (replica,) = replicas = await self.replicas(w1) + replica.held.set() # update_weights for w2 outlasts the shutdown grace, as a dense reload over NFS can + pending_w2 = await self.accepted(w2) + first = await self.dispatcher(replicas) + await until(lambda: ("update", "pod-0", w2, w2) in self.log) + first.cancel() # SIGTERM, from a pod deletion, a rollout or a pool release + await asyncio.gather(first, return_exceptions=True) + await asyncio.sleep(0.3) # a resume sent by the cut-off load would land by now + self.assertEqual((replica.paused, replica.updating, replica.version), (True, True, w1)) + replica.held.clear() # vLLM finishes writing w2, but nothing finishes the update, so the replica still says w1 + # The w2 request is answered elsewhere, and a late request for w1 arrives. + [r] = [r for r in await self.store.pending(self.pool) if r["request_id"] == pending_w2["request_id"]] + await self.store.commit(r, SAMPLE) + late = await self.accepted(w1) + await self.dispatcher(replicas) + self.assertEqual(await self.committed(late["request_id"]), [SAMPLE]) + # The next dispatcher closed the update, which left the replica unknown, and loaded w1 in full before serving from it. + self.assertEqual(self.log, [("update", "pod-0", w2, w2), ("update", "pod-0", w1, w1), ("generate", "pod-0", w1)]) + self.assertEqual((replica.version, replica.updating, replica.paused), (w1, False, False)) + + async def test_a_failed_update_leaves_the_replica_reporting_unknown_weights_until_a_full_load(self): + w1, w2 = self.revision("w1"), self.revision("w2") + (replica,) = replicas = await self.replicas(w1) + replica.fail_update = True # the first update_weights fails inside vLLM, which may have written part of w2 + replica.held.set() + r = await self.accepted(w2) + await self.dispatcher(replicas) + await until(lambda: self.loads("pod-0") == {"pod-0": [w2, w2]}) + self.assertEqual((replica.version, replica.updating, replica.resumed), (UNKNOWN_WEIGHTS, True, 1)) + replica.held.clear() + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + self.assertEqual((self.loads("pod-0"), replica.version), ({"pod-0": [w2, w2]}, w2)) + + async def test_a_replica_is_probed_for_an_open_update_on_first_sight_and_after_a_load_that_breaks_off_only(self): + # The probe is a bare finish_weight_update, and vLLM logs a 500 traceback for each one that finds no update open. + (replica,) = await self.replicas("default") + dispatcher = Dispatcher(self.store, self.pool, self.http, None, [replica.url], concurrency=1) + + async def load(revision: str) -> None: + dispatcher.revision = revision + dispatcher.write_revision() + dispatcher.loading[replica.url] = task = asyncio.create_task(dispatcher.load_weights(replica.url, revision)) + await task + + await dispatcher.inspect(replica.url) + w1, w2 = self.revision("w1"), self.revision("w2") + await load(w1) + await dispatcher.inspect(replica.url) + self.assertEqual((replica.probes, replica.version), (1, w1)) + replica.fail_update = True + await load(w2) + probes = replica.probes # the broken load closes its own update with a bare finish + await dispatcher.inspect(replica.url) + self.assertEqual((replica.probes, replica.version), (probes + 1, UNKNOWN_WEIGHTS)) + + async def test_under_the_delta_strategy_a_replica_with_unknown_weights_stays_out_of_service(self): + # No full revision is kept beside the deltas yet, so a replica that lost its place in the chain cannot be brought back. + w1 = self.revision("w1", format="sparse_delta", parent="") + self.enterContext(patch("server.sampling_dispatcher.SAMPLE_DEADLINE_SECONDS", 1)) + (replica,) = replicas = await self.replicas("default") + replica.paused, replica.updating = True, True # cut off loading w1 + r = await self.accepted(w1) + await self.dispatcher(replicas) + [result] = await self.committed(r["request_id"]) + self.assertEqual(result["type"], "RequestFailedResponse") + self.assertEqual((replica.version, replica.paused, self.loads("pod-0")), (UNKNOWN_WEIGHTS, False, {"pod-0": []})) + + async def test_a_stopping_dispatcher_lets_a_load_in_flight_finish_and_resume(self): + (replica,) = replicas = await self.replicas("default") + replica.held.set() + await self.accepted(w1 := self.revision("w1")) + run = await self.dispatcher(replicas) + await until(lambda: ("update", "pod-0", w1, w1) in self.log) + run.cancel() + asyncio.get_running_loop().call_later(0.2, replica.held.clear) + await asyncio.gather(run, return_exceptions=True) + self.assertEqual((replica.version, replica.resumed, replica.paused), (w1, 1, False)) + + async def test_a_delta_loads_on_the_save_its_optimizer_step_started_from_past_a_save_the_pool_never_sampled(self): + self.enterContext(patch("server.sampling_dispatcher.SAMPLE_DEADLINE_SECONDS", 5)) + worker = fft_worker() + train_step(worker) + w1, at_w1 = worker.save_for_sampler("job", None, "tinker://job/sampler_weights/sampler-0"), weights(worker) + train_step(worker) + worker.save_for_sampler("job", "ckpt", "tinker://job/sampler_weights/ckpt") # a named save nobody samples + w2 = worker.save_for_sampler("job", None, "tinker://job/sampler_weights/sampler-1") + (replica,) = replicas = await self.replicas(w1) + r = await self.accepted(w2) + await self.dispatcher(replicas) + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + self.assertEqual((self.loads("pod-0"), replica.version), ({"pod-0": [w2]}, w2)) + self.assertTrue(same(replayed(at_w1, [w2]), weights(worker))) + self.assertEqual(json.loads(Path(w2, "metadata.json").read_text())["parent"], w1) + + async def test_a_sampler_save_two_optimizer_steps_after_the_last_is_written_in_full_and_loads_the_trainers_weights(self): + self.enterContext(patch("server.sampling_dispatcher.SAMPLE_DEADLINE_SECONDS", 5)) + worker = fft_worker() + train_step(worker) + w1, at_w1 = worker.save_for_sampler("job", None, "tinker://job/sampler_weights/sampler-0"), weights(worker) + nudge_step(worker, 0) + nudge_step(worker, 1) # a delta holds only this step, so no replica on w1 could apply it + w2 = worker.save_for_sampler("job", None, "tinker://job/sampler_weights/sampler-1") + (replica,) = replicas = await self.replicas(w1) + r = await self.accepted(w2) + await self.dispatcher(replicas) + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + self.assertEqual((self.loads("pod-0"), replica.version), ({"pod-0": [w2]}, w2)) + self.assertTrue(same(replayed(at_w1, [w2]), weights(worker))) + metadata = json.loads(Path(w2, "metadata.json").read_text()) + self.assertEqual((metadata.get("format"), metadata.get("parent")), (None, None)) + + async def test_a_replica_on_any_save_at_the_weights_a_delta_starts_from_or_reaches_applies_it(self): + worker = fft_worker() + train_step(worker) + # An RL step with nothing to learn from takes no optimizer step, so its rollouts sample a second save at the same weights. + first, at_first = worker.save_for_sampler("job", None, "tinker://job/sampler_weights/sampler-0"), weights(worker) + second = worker.save_for_sampler("job", None, "tinker://job/sampler_weights/sampler-1") + # The late replica started on the first save and never loads the second. + early, late = await self.replicas(first, first) + r = await self.accepted(second) + run = await self.dispatcher([early]) + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + run.cancel() + await asyncio.gather(run, return_exceptions=True) + train_step(worker) + third = worker.save_for_sampler("job", None, "tinker://job/sampler_weights/sampler-2") + r = await self.accepted(third) + await self.dispatcher([early, late]) + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + await until(lambda: self.loads(late.name) == {late.name: [third]}) + 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))) diff --git a/tests/test_scheduler_worker_manager.py b/tests/test_scheduler_worker_manager.py index 40becf0e..ecb3a2c9 100644 --- a/tests/test_scheduler_worker_manager.py +++ b/tests/test_scheduler_worker_manager.py @@ -1,12 +1,16 @@ import json import os +import tempfile import unittest +from pathlib import Path +from types import SimpleNamespace from typing import Any from unittest.mock import patch from server import api_server from server.estimator import footprint -from server.scheduler_worker_manager import GROUP, PLURAL, VERSION, SchedulerWorkerManager +from server.sampling_dispatcher import revision_file +from server.scheduler_worker_manager import GROUP, PLURAL, VERSION, SchedulerWorkerManager, describe_worker, dispatcher_body, pod_template from server.store import InMemoryStateStore, InMemoryStore from tests.api_client import asgi_client, post_json @@ -54,11 +58,32 @@ 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 == "delete": + if (kind, args[0]) not in self.objects: + raise ApiError(404) + return self.objects.pop((kind, args[0])) + 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 +108,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: @@ -222,7 +247,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, []) @@ -329,6 +354,26 @@ def finalizer_finishes(_seconds: float) -> None: self.assertEqual(sleep.call_count, 1) self.assertEqual(len(self.api.created), 2) + def test_ensure_says_why_it_waits_for_a_released_pools_dispatcher(self) -> None: + s = self.store_with("fft-job", {"base_model": "Qwen/Qwen3-8B", "fine_tuning_type": "full"}) + self.k8s.objects[("deployment", "sampler-fft-job")] = {"metadata": {"name": "sampler-fft-job"}} + terminating = SimpleNamespace(metadata=SimpleNamespace(deletion_timestamp="2026-10-18T00:00:00Z")) + self.k8s.read_namespaced_deployment = lambda name, namespace: terminating + + def old_pod_gone(_seconds: float) -> None: + self.k8s.objects.pop(("deployment", "sampler-fft-job")) + + with ( + patch("server.worker_manager.get_state_store", return_value=s), + patch("server.scheduler_worker_manager.time.sleep", side_effect=old_pod_gone), + self.assertLogs("server.scheduler_worker_manager", "WARNING") as logs, + ): + self.manager.ensure("fft-job", "sampler") + + self.assertIn("sampler-fft-job of a released pool is still terminating", logs.output[0]) + self.assertIn("Waited", logs.output[1]) + self.assertEqual(self.k8s.objects[("deployment", "sampler-fft-job")]["kind"], "Deployment") + def test_create_worker_manager_selects_scheduler_mode(self) -> None: from server.worker_manager import create_worker_manager @@ -354,6 +399,72 @@ 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_lora_pool_has_its_snapshot_directory_before_any_replica_is_requested(self) -> None: + # vLLM's filesystem LoRA resolver exits at startup without it, and no trainer may have frozen an adapter yet. + snapshots, create = Path(self.tmp, "sampler_lora"), self.api.create_namespaced_custom_object + seen = [] + self.api.create_namespaced_custom_object = lambda *args: (seen.append(snapshots.is_dir()), create(*args))[1] + 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): + self.manager.ensure("job-lora-1", "sampler") + (sampler,) = self.api.created + env = {e["name"]: e.get("value") for e in sampler["spec"]["template"]["spec"]["containers"][0]["env"]} + self.assertEqual((seen, env["VLLM_LORA_RESOLVER_CACHE_DIR"]), ([True], str(snapshots))) + + def test_a_pool_replica_takes_ninety_percent_of_its_gpu_unless_the_env_says_otherwise(self) -> None: + # vLLM's own 0.92 has run FFT delta applies out of memory on an L4. Its batch limits follow the device it lands on. + s = self.store_with("fft-job", {"base_model": "Qwen/Qwen3-8B", "fine_tuning_type": "full"}) + with patch("server.worker_manager.get_state_store", return_value=s), patch.dict(os.environ): + os.environ.pop("VLLM_GPU_MEMORY_UTILIZATION", None) + default = pod_template(describe_worker("fft-job", "sampler"))["spec"]["containers"][0]["args"] + os.environ["VLLM_GPU_MEMORY_UTILIZATION"] = "0.7" + explicit = pod_template(describe_worker("fft-job", "sampler"))["spec"]["containers"][0]["args"] + fractions = [float(args[args.index("--gpu-memory-utilization") + 1]) for args in (default, explicit)] + self.assertEqual((fractions, "--max-num-seqs" in default), ([0.90, 0.7], False)) + + def test_without_open_rl_tmp_dir_the_probe_reads_the_revision_file_the_dispatcher_writes(self) -> None: + # The API server renders the probe's path itself, and the dispatcher writes under the OPEN_RL_TMP_DIR its pod gets. + os.environ.pop("OPEN_RL_TMP_DIR") + s = self.store_with("fft-job", {"base_model": "Qwen/Qwen3-8B", "fine_tuning_type": "full"}) + with patch("server.worker_manager.get_state_store", return_value=s): + worker = describe_worker("fft-job", "sampler") + probe = pod_template(worker)["spec"]["containers"][0]["readinessProbe"]["exec"]["command"][2] + dispatcher = dispatcher_body(worker.pool)["spec"]["template"]["spec"]["containers"][0] + with patch.dict(os.environ, {e["name"]: e["value"] for e in dispatcher["env"] if e["name"] == "OPEN_RL_TMP_DIR"}): + written = revision_file(worker.pool) + self.assertIn(f"open({str(written)!r})", probe) + + def test_a_pool_is_vllm_samplers_ready_on_the_expected_version_behind_one_routed_cpu_pod_until_released(self) -> None: + s = self.store_with("fft-job", {"base_model": "Qwen/Qwen3-8B", "fine_tuning_type": "full"}) + k8s = FakeApi() + manager = SchedulerWorkerManager(custom_api=self.api, apps_api=k8s, networking_api=k8s) + with patch("server.worker_manager.get_state_store", return_value=s): + manager.ensure("fft-job", "sampler") + (sampler,) = self.api.created + labels = {"app.kubernetes.io/component": "sampler", "app.kubernetes.io/instance": "sampler-fft-job"} + container = sampler["spec"]["template"]["spec"]["containers"][0] + self.assertNotIn("--middleware", container["args"]) + probe = container["readinessProbe"]["exec"]["command"] + self.assertEqual( + (probe[:2], "/weight_info" in probe[2], f"{self.tmp}/sampling/sampler-fft-job/revision" in probe[2]), + (["/app/.venv/bin/python", "-c"], True, True), + ) + self.assertEqual(sorted(k8s.objects), [("deployment", "sampler-fft-job"), ("network_policy", "sampler-fft-job")]) + dispatcher, _, epp = k8s.objects[("deployment", "sampler-fft-job")]["spec"]["template"]["spec"]["containers"] + self.assertEqual(dispatcher["volumeMounts"], [{"name": "shared-storage", "mountPath": "/mnt/shared"}]) + self.assertIn({"name": "OPEN_RL_TMP_DIR", "value": self.tmp}, dispatcher["env"]) # so it writes the file the probe reads + self.assertEqual(dispatcher["image"], "ghcr.io/gke-labs/open-rl/server:latest") # the worker image unless a dispatcher image is set + with patch.dict(os.environ, {"OPEN_RL_DISPATCHER_IMAGE": "registry/api-server:v1"}): + self.assertEqual(dispatcher_body("sampler-fft-job")["spec"]["template"]["spec"]["containers"][0]["image"], "registry/api-server:v1") + selector = "app.kubernetes.io/component=sampler,app.kubernetes.io/instance=sampler-fft-job" + self.assertEqual( + epp["args"][:6], ["--endpoint-selector", selector, "--endpoint-target-ports", "8000", "--config-file", "/etc/sampler-router/epp.yaml"] + ) + policy = k8s.objects[("network_policy", "sampler-fft-job")]["spec"] + self.assertEqual((policy["podSelector"]["matchLabels"], policy["ingress"][0]["ports"]), (labels, [{"port": 8000}, {"port": 5557}])) + self.assertEqual(manager.release_owner("fft-job"), {"fft-job"}) + self.assertEqual((k8s.objects, self.api.existing), ({}, {})) + class MixedSamplingSessionTest(unittest.IsolatedAsyncioTestCase): async def test_lora_and_fft_sessions_launch_their_own_sampler_types(self) -> None: @@ -361,14 +472,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"): @@ -377,7 +489,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 d6fb8e34..1620e159 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): diff --git a/tests/test_worker_manager.py b/tests/test_worker_manager.py index 727e3658..fd752028 100644 --- a/tests/test_worker_manager.py +++ b/tests/test_worker_manager.py @@ -1,11 +1,18 @@ import json +import os +import sys +import tempfile +import time import unittest +import uuid +from pathlib import Path from unittest.mock import patch from server import api_server from server.session_registry import SessionRegistry from server.worker_manager import LocalWorkerManager from tests.api_client import asgi_client, post_json +from tests.redis_server import TEST_REDIS_URL class StoreStub: @@ -201,7 +208,7 @@ 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_local_sampler_is_vllm_serve_and_its_dispatcher_each_in_its_own_process_group(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.subprocess.Popen") as popen, @@ -209,12 +216,11 @@ async def test_local_sampler_launch_stamps_workload_tags_and_process_group(self) 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") + (serve, serve_kwargs), (dispatch, dispatch_kwargs) = [(c[0][0], c[1]) for c in popen.call_args_list] + self.assertEqual((serve_kwargs["start_new_session"], dispatch_kwargs["start_new_session"]), (True, True)) + self.assertEqual((serve[serve.index("serve") + 1], serve_kwargs["env"]["OPEN_RL_ENABLE_FFT"]), ("Model_A.1", "true")) + self.assertNotIn("OPEN_RL_TIME_SLICE_JOB_ID", serve_kwargs["env"]) + self.assertEqual(dispatch[-1], "server.sampling_dispatcher") async def test_launch_fetches_metadata_from_store(self) -> None: import json @@ -242,7 +248,7 @@ async def test_launch_fetches_metadata_from_store(self) -> None: self.assertEqual(kwargs["env"].get("OPEN_RL_WEIGHT_SYNC_STRATEGY"), "delta") manager.ensure("Model_A.1", "sampler") - _, kwargs_s = popen.call_args + _, kwargs_s = popen.call_args_list[-2] # vllm serve; the last launch is its dispatcher self.assertEqual(kwargs_s["env"].get("BASE_MODEL"), "base-model-a") self.assertEqual(kwargs_s["env"].get("OPEN_RL_WEIGHT_SYNC_STRATEGY"), "delta") @@ -354,6 +360,8 @@ def setUp(self) -> None: 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") @@ -370,19 +378,25 @@ def test_launch_sampler_lora_uses_base_model_and_reuses_process(self, mock_popen fine_tuning_type="lora", ) - # Launch for first LoRA model ID + # Launch for first LoRA model ID: a vllm serve replica and its dispatcher 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"), ("dispatcher", "Qwen/Qwen2.5-0.5B")}) + (serve, serve_kwargs), (dispatch, dispatch_kwargs) = [(c[0][0], c[1]) for c in mock_popen.call_args_list] + port = serve[serve.index("--port") + 1] + self.assertEqual(serve[serve.index("serve") + 1 : serve.index("serve") + 2], ["Qwen/Qwen2.5-0.5B"]) + self.assertIn("--enable-lora", serve) + self.assertEqual(serve_kwargs["env"]["VLLM_ALLOW_RUNTIME_LORA_UPDATING"], "true") + # Triton kernels live under OPEN_RL_TMP_DIR; vLLM's own compile cache stays where vLLM puts it. + tmp = os.getenv("OPEN_RL_TMP_DIR", "/tmp/open-rl") + self.assertEqual(serve_kwargs["env"]["TRITON_CACHE_DIR"], f"{tmp}/cache/triton") + self.assertEqual(serve_kwargs["env"].get("VLLM_CACHE_ROOT"), os.environ.get("VLLM_CACHE_ROOT")) + self.assertIn("server.sampling_dispatcher", dispatch) + self.assertEqual(dispatch_kwargs["env"]["OPEN_RL_SAMPLER_ENDPOINTS"], f"http://127.0.0.1:{port}") # Launch for second LoRA model ID sharing the same base model self.manager.ensure("model-lora-2", "sampler") # Should reuse existing process and NOT call popen again! - self.assertEqual(mock_popen.call_count, 1) + self.assertEqual(mock_popen.call_count, 2) @patch("server.worker_manager.metadata_for") @patch("subprocess.Popen") @@ -400,12 +414,103 @@ def test_launch_sampler_fft_uses_model_id(self, mock_popen, mock_fetch) -> None: ) 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"), ("dispatcher", "model-fft-1")}) + (serve, _), (_, dispatch_kwargs) = [(c[0][0], c[1]) for c in mock_popen.call_args_list] + self.assertEqual(serve[serve.index("serve") + 1], "Qwen/Qwen2.5-0.5B") + self.assertIn("--weight-transfer-config", serve) + self.assertEqual(dispatch_kwargs["env"]["OPEN_RL_SAMPLER_POOL"], "sampler-model-fft-1") + + +# Stands in for vllm serve, which is not installed here. Like vLLM's filesystem LoRA resolver, it exits at startup when +# VLLM_LORA_RESOLVER_CACHE_DIR names a directory that does not exist. +VLLM_STAND_IN = ( + "import os, sys, time\nd = os.environ.get('VLLM_LORA_RESOLVER_CACHE_DIR')\nsys.exit(3) if d and not os.path.isdir(d) else time.sleep(60)" +) + + +def stand_in_spawn(spawned: dict[str, dict[str, str]]): + """LocalWorkerManager.spawn with VLLM_STAND_IN for vllm serve and the dispatcher module on this interpreter, since + `uv run` would sync the environment. The env of each role lands in `spawned`.""" + original = LocalWorkerManager.spawn + + def spawn(manager: LocalWorkerManager, role: str, runtime: str, command: list[str], env: dict[str, str]) -> None: + spawned[role] = env + stand_in = [sys.executable, *command[command.index("-u") :]] if role == "dispatcher" else [sys.executable, "-c", VLLM_STAND_IN] + original(manager, role, runtime, stand_in, env) + + return spawn + + +@unittest.skipUnless(TEST_REDIS_URL, "needs OPEN_RL_TEST_REDIS_URL, for example from docker run -p 6411:6379 redis:7") +class LocalPoolProcessTest(unittest.TestCase): + """A local pool's processes as LocalWorkerManager starts them, with the real dispatcher.""" + + def test_with_the_documented_local_env_the_dispatcher_runs_and_keeps_its_revision_under_the_default_tmp_dir(self) -> None: + # docs/setup/local-setup.md starts the server with REDIS_URL set and OPEN_RL_TMP_DIR unset. + env = {name: value for name, value in os.environ.items() if name != "OPEN_RL_TMP_DIR"} + self.enterContext(patch.dict(os.environ, {**env, "REDIS_URL": TEST_REDIS_URL, "OPEN_RL_ENABLE_FFT": "true"}, clear=True)) + spawned, model_id = {}, f"pool-{uuid.uuid4().hex[:8]}" + self.enterContext(patch.object(LocalWorkerManager, "spawn", stand_in_spawn(spawned))) + revision = Path("/tmp/open-rl/sampling", f"sampler-{model_id}", "revision") + for directory in [d for d in reversed(revision.parents[:3]) if not d.exists()]: + self.addCleanup(lambda d=directory: d.is_dir() and not any(d.iterdir()) and d.rmdir()) + log = Path(f"/tmp/dispatcher_{model_id}.log") + for path in (revision, log, Path(f"/tmp/sampler_{model_id}.log")): + self.addCleanup(path.unlink, missing_ok=True) + manager = LocalWorkerManager() + self.addCleanup(manager.close) + manager.ensure(model_id, "sampler") + dispatcher = manager.processes[("dispatcher", model_id)] + deadline = time.monotonic() + 30 + while not revision.exists() and dispatcher.poll() is None and time.monotonic() < deadline: + time.sleep(0.1) + time.sleep(1.5) # a dispatch pass, so a dispatcher that dies after writing the file shows + self.assertIsNone(dispatcher.poll(), log.read_text()[-1500:]) + self.assertEqual(revision.read_text(), "default") + self.assertEqual(spawned["sampler"]["TRITON_CACHE_DIR"], "/tmp/open-rl/cache/triton") + self.assertEqual(spawned["dispatcher"]["OPEN_RL_TMP_DIR"], "/tmp/open-rl") + + def test_a_lora_pool_started_before_any_adapter_is_frozen_finds_its_snapshot_directory(self) -> None: + # Base-model sampling before the first LoRA save, as examples/harvey_labs/train.py does. + tmp = self.enterContext(tempfile.TemporaryDirectory()) + self.enterContext(patch.dict(os.environ, {"REDIS_URL": TEST_REDIS_URL, "OPEN_RL_TMP_DIR": tmp, "OPEN_RL_ENABLE_FFT": "false"})) + spawned = {} + self.enterContext(patch.object(LocalWorkerManager, "spawn", stand_in_spawn(spawned))) + manager = LocalWorkerManager() + self.addCleanup(manager.close) + manager.ensure("Qwen/Qwen3-0.6B", "sampler") + time.sleep(1) + self.assertEqual(spawned["sampler"]["VLLM_LORA_RESOLVER_CACHE_DIR"], f"{tmp}/sampler_lora") + self.assertIsNone(manager.processes[("sampler", "Qwen/Qwen3-0.6B")].poll()) + + def test_ensure_restarts_a_dead_dispatcher_and_stops_a_live_one_before_it_starts_a_new_pair(self) -> None: + tmp = self.enterContext(tempfile.TemporaryDirectory()) + self.enterContext(patch.dict(os.environ, {"REDIS_URL": TEST_REDIS_URL, "OPEN_RL_TMP_DIR": tmp, "OPEN_RL_ENABLE_FFT": "true"})) + stand_ins = {"sampler": ["sleep", "300"], "dispatcher": [sys.executable, "-c", "import sys; sys.exit(1)"]} + original = LocalWorkerManager.spawn + + def spawn(manager: LocalWorkerManager, role: str, runtime: str, command: list[str], env: dict[str, str]) -> None: + original(manager, role, runtime, stand_ins[role], env) + + self.enterContext(patch.object(LocalWorkerManager, "spawn", spawn)) + manager, model_id = LocalWorkerManager(), f"pool-{uuid.uuid4().hex[:8]}" + self.addCleanup(manager.close) + manager.ensure(model_id, "sampler") + vllm, dead = manager.processes[("sampler", model_id)], manager.processes[("dispatcher", model_id)] + self.assertEqual(dead.wait(timeout=10), 1) + # The next save_weights_for_sampler or sampling session finds the dispatcher dead and starts it again. + stand_ins["dispatcher"] = ["sleep", "300"] + manager.ensure(model_id, "sampler") + live = manager.processes[("dispatcher", model_id)] + self.addCleanup(live.kill) + self.assertEqual((manager.processes[("sampler", model_id)] is vllm, live is dead, live.poll()), (True, False, None)) + # Once vLLM dies, its dispatcher is stopped before a new pair starts, so two dispatchers never read one stream. + vllm.kill() + vllm.wait() + manager.ensure(model_id, "sampler") + sampler, dispatcher = manager.processes[("sampler", model_id)], manager.processes[("dispatcher", model_id)] + self.assertIsNotNone(live.poll()) + self.assertEqual((sampler is vllm, sampler.poll(), dispatcher is live, dispatcher.poll()), (False, None, False, None)) if __name__ == "__main__":