diff --git a/Makefile b/Makefile index 48d54c97..5be49609 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_sampler_launcher tests.test_sampler_gpu_leases 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/docs/setup/gke-fft-timeslice.md b/docs/setup/gke-fft-timeslice.md index 46afc6cb..aea2056f 100644 --- a/docs/setup/gke-fft-timeslice.md +++ b/docs/setup/gke-fft-timeslice.md @@ -182,6 +182,32 @@ in-memory queue and active/checkpointed state for workloads sharing the physical GPU. The worker pod labels provide the workload identity llm-d uses to discover the relevant pod and process set. +An FFT job's pool sampler shares the +trainer's GPU and takes turns with it through the same coordinator, with no +proxy in the request path. Its pod runs +`python -m server.sampler_launcher vllm serve ... --enable-sleep-mode` with +`OPEN_RL_TIME_SLICED_SAMPLER=true`. The launcher registers the sampler's +workload, starts vLLM inside one turn, puts it to sleep and releases, so vLLM +never takes GPU memory while the trainer holds the device, then keeps the +registration open for as long as vLLM runs. The container counts as started only +then, and the pool's dispatcher, which reads the same variable from the pod, +leases the replica for each turn through the slicer on its node +(`status.hostIP` and `OPEN_RL_ACCEL_TIMESLICER_PORT`), wakes vLLM, loads any +pending weights as a whole, sends requests to it directly rather than through +the router, puts it to sleep and releases. A turn ends once the replica is idle, +or after `OPEN_RL_SAMPLER_TURN_SECONDS` (60 by default) with requests still +pending, so the trainer gets its turn. Requests still running then get ten +seconds more. After that they are cancelled, aborted inside vLLM by its sleep, +and sent again in a later turn. If the dispatcher dies holding a lease, +the slicer checkpoints the replica before anyone else is granted the device. If +that checkpoint, or the one at the end of a turn, fails, the slicer tells the +launcher, which stops vLLM and exits, so the device frees for the trainer and +kubelet restarts the replica. A +pod deleted while vLLM is still starting is killed by its launcher a few seconds +after the SIGTERM vLLM ignores at that point, so the turn is freed at once +rather than at kubelet's SIGKILL. Exclusive jobs' pool samplers and LoRA pool +samplers keep a GPU of their own and never talk to the slicer. + ## Requirements - GKE Standard cluster on **1.35 or newer** ([DRA for GPUs](https://docs.cloud.google.com/kubernetes-engine/docs/how-to/set-up-dra) 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/accel_timeslicer/README.md b/src/accel_timeslicer/README.md index cf050acc..a5cd849f 100644 --- a/src/accel_timeslicer/README.md +++ b/src/accel_timeslicer/README.md @@ -38,6 +38,23 @@ one DRA `ResourceClaim`. The cluster manifest runs this process with `--backend llmd` so llm-d's snapshot agent owns pod/PID discovery and physical checkpoint/restore. +The process that owns a workload's work queue holds its lease. A trainer +registers, acquires and releases from inside its own process. A pool sampler +that shares its GPU is `vllm serve` run by +`server.sampler_launcher`, which registers the workload, starts vLLM inside one +turn, puts it to sleep and releases over a connection it keeps open for as long +as vLLM runs; from then on the pool's dispatcher acquires and releases the same +workload over its own connection for each turn. A registration belongs to the +connection that made it, and a grant to the connection whose `ACQUIRE` holds it, +which alone can `RELEASE` it. When a connection closes while holding a grant on +a workload it did not register, the slicer releases that workload, checkpointing +it, before the next waiter is granted. A workload whose checkpoint fails keeps +the grant until the connection that registered it closes. When another +connection's release is the one that failed, as at the end of a dispatcher's +turn, the slicer tells the registering connection with a line it did not ask +for, `{"event": "FAULT", "name": ..., "claim": ..., "error": ...}`, and the +sampler launcher stops vLLM and exits on it. + ## llm-d snapshot backend The same OpenRL acquire/release protocol can use llm-d's snapshot agent by diff --git a/src/accel_timeslicer/serve.py b/src/accel_timeslicer/serve.py index 76151dac..8be79046 100644 --- a/src/accel_timeslicer/serve.py +++ b/src/accel_timeslicer/serve.py @@ -34,12 +34,19 @@ async def start_tcp_time_slicer(time_slicer: SingleNodeTimeSlicer, host: str, po async def handle_connection(time_slicer: SingleNodeTimeSlicer, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: connection_id = id(writer) + + def tell(message: dict[str, Any]) -> None: + if not writer.is_closing(): + writer.write(json.dumps(message).encode("utf-8") + b"\n") + + time_slicer.connections[connection_id] = tell try: while line := await reader.readline(): response = await dispatch(time_slicer, line, connection_id) writer.write(json.dumps(response).encode("utf-8") + b"\n") await writer.drain() finally: + time_slicer.connections.pop(connection_id, None) await time_slicer.connection_closed(connection_id) writer.close() await writer.wait_closed() @@ -55,9 +62,9 @@ async def dispatch(time_slicer: SingleNodeTimeSlicer, line: bytes, connection_id case "REGISTER": return await time_slicer.register(workload, connection_id=connection_id) case "ACQUIRE": - return await time_slicer.acquire(workload) + return await time_slicer.acquire(workload, connection_id=connection_id) case "RELEASE": - return await time_slicer.release(workload) + return await time_slicer.release(workload, connection_id=connection_id) case "UNREGISTER": return await time_slicer.unregister(workload) case _: diff --git a/src/accel_timeslicer/single_node.py b/src/accel_timeslicer/single_node.py index eacfcdb2..5868d7b8 100644 --- a/src/accel_timeslicer/single_node.py +++ b/src/accel_timeslicer/single_node.py @@ -2,6 +2,7 @@ import logging import time from collections import deque +from collections.abc import Callable from dataclasses import dataclass from typing import Any @@ -18,6 +19,8 @@ class WorkloadState: workload: WorkloadRef checkpointed: bool = False failed: bool = False + # The connection whose acquire holds the grant, while one does. + holder: int | None = None class SingleNodeTimeSlicer(TimeSlicer): @@ -34,6 +37,8 @@ def __init__(self, restorer: CheckpointRestorer, scheduling_policy: str = "lrs") self.running: dict[str, str] = {} self.condition = asyncio.Condition() self.last_release_time: dict[str, float] = {} + # How the server reaches each open connection with a line it did not ask for. + self.connections: dict[int, Callable[[dict[str, Any]], None]] = {} def next_waiter(self, claim: str) -> str | None: """The earliest waiter under fifo, the least recently served otherwise.""" @@ -58,7 +63,7 @@ async def register(self, workload: WorkloadRef, connection_id: int | None = None self.condition.notify_all() return {"ok": True} - async def acquire(self, workload: WorkloadRef) -> dict[str, Any]: + async def acquire(self, workload: WorkloadRef, connection_id: int | None = None) -> dict[str, Any]: async with self.condition: name, claim = workload.name, workload.claim state = self.workloads.get(name) @@ -88,6 +93,7 @@ async def acquire(self, workload: WorkloadRef) -> dict[str, Any]: self.waiting_workloads.remove(name) self.running[claim] = name + state.holder = connection_id self.condition.notify_all() if state.checkpointed: @@ -99,12 +105,16 @@ async def acquire(self, workload: WorkloadRef) -> dict[str, Any]: self.condition.notify_all() return {"ok": True} - async def release(self, workload: WorkloadRef) -> dict[str, Any]: + async def release(self, workload: WorkloadRef, connection_id: int | None = None) -> dict[str, Any]: async with self.condition: name, claim = workload.name, workload.claim state = self.workloads.get(name) if state is None or self.running.get(claim) != name: return {"ok": False, "error": f"workload {name} does not hold an active acquire"} + # A grant the workload took again over another connection, as a restarted launcher does for its boot turn, is + # not this connection's to end. A caller in this process names no connection and is not checked. + if connection_id is not None and state.holder != connection_id: + return {"ok": False, "error": f"workload {name} is held by another connection; only the connection that acquired it can release it"} # The grant is held through the checkpoint, so the next waiter is not # granted until this workload is off the devices. @@ -115,13 +125,18 @@ async def release(self, workload: WorkloadRef) -> dict[str, Any]: # the grant so nobody is restored on top of its memory; the seat frees # when its process exits and the connection closes. logger.error("checkpoint failed for workload %s claim %s; it keeps the grant until its process exits: %s", name, claim, exc) + error = f"checkpoint failed for workload {name}; it still holds the accelerator and must exit: {exc}" async with self.condition: state = self.workloads.get(name) if state is not None: state.failed = True state.checkpointed = False + # Only the process that registered the workload can free the devices, by exiting, so it is told when another + # connection's release failed, as a pool dispatcher's can for a sampler. One that released it itself has the reply. + if state.connection_id != connection_id and (tell := self.connections.get(state.connection_id)): + tell({"event": "FAULT", **workload.as_payload(), "error": error}) self.condition.notify_all() - return {"ok": False, "faulted": True, "error": f"checkpoint failed for workload {name}; it still holds the accelerator and must exit: {exc}"} + return {"ok": False, "faulted": True, "error": error} async with self.condition: state = self.workloads.get(name) @@ -144,6 +159,13 @@ async def unregister(self, workload: WorkloadRef) -> dict[str, Any]: return {"ok": True} async def connection_closed(self, connection_id: int) -> None: + # A grant held over a connection other than the registering one is a lease another process took, the pool + # dispatcher for a sampler. The workload is still on the devices, so it is released, which checkpoints it. + async with self.condition: + leased = [state.workload for state in self.workloads.values() if state.holder == connection_id and state.connection_id != connection_id] + for workload in leased: + logger.warning("connection holding workload %s claim %s closed; releasing it", workload.name, workload.claim) + await self.release(workload, connection_id) async with self.condition: for name, state in self.workloads.items(): if state.connection_id != connection_id: @@ -160,6 +182,8 @@ def clear_workload(self, name: str) -> None: for claim, holder in list(self.running.items()): if holder == name: del self.running[claim] + if (state := self.workloads.get(name)) is not None: + state.holder = None async def run_checkpoint(self, state: WorkloadState) -> bool | None: """Parks the workload. Raises when the restorer could not, since the workload then still holds the devices.""" diff --git a/src/accel_timeslicer/time_slicer.py b/src/accel_timeslicer/time_slicer.py index ed05af3c..f45fe2aa 100644 --- a/src/accel_timeslicer/time_slicer.py +++ b/src/accel_timeslicer/time_slicer.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import json import os from collections.abc import AsyncIterator @@ -50,8 +51,9 @@ def __init__( self.port = port self.reader: asyncio.StreamReader | None = None self.writer: asyncio.StreamWriter | None = None - # Set when a release could not park this process. The caller finishes - # publishing its results and exits; the grant moves on once it is gone. + # Set when a release could not park this process, whether this connection + # released it or another one leased the workload it registered. The caller + # finishes publishing its results and exits; the grant moves on once it is gone. self.faulted: str | None = None async def connect(self, retries: int = 10, backoff: float = 0.5) -> None: @@ -103,12 +105,9 @@ async def request(self, payload: dict[str, Any]) -> dict[str, Any]: try: self.writer.write(json.dumps(payload).encode("utf-8") + b"\n") await self.writer.drain() - line = await self.reader.readline() - if not line: - await self.close() - raise RuntimeError("time slicer connection closed") - - response = json.loads(line.decode("utf-8")) + response = await self.receive() + while response.get("event") == "FAULT": + response = await self.receive() if not response.get("ok"): if response.get("faulted"): # The connection stays open: closing it would hand the grant to the @@ -122,6 +121,26 @@ async def request(self, payload: dict[str, Any]) -> dict[str, Any]: await self.close() raise + async def receive(self) -> dict[str, Any]: + """The next line from the slicer. A FAULT, which the slicer sends unasked and never as a reply, is noted.""" + assert self.reader is not None + line = await self.reader.readline() + if not line: + await self.close() + raise RuntimeError("time slicer connection closed") + message = json.loads(line.decode("utf-8")) + if message.get("event") == "FAULT": + self.faulted = message.get("error", "the time slicer could not park this process") + return message + + async def fault(self) -> str | None: + """Waits, with no request in flight, until the slicer says another connection's release could not park the + workload this connection registered. None if the connection closes first.""" + with contextlib.suppress(OSError, RuntimeError): + while self.faulted is None and self.reader is not None: + await self.receive() + return self.faulted + def workload_from_env(pid: int | None = None, name: str | None = None, claim: str = DEFAULT_CLAIM) -> WorkloadRef: # The launcher's stamps win over the caller's defaults. OPEN_RL_TIME_SLICE_GROUP 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/sampler_launcher.py b/src/server/sampler_launcher.py new file mode 100644 index 00000000..f8c23af8 --- /dev/null +++ b/src/server/sampler_launcher.py @@ -0,0 +1,129 @@ +"""Runs a pool sampler's vLLM for a replica that shares its GPU with its trainer. + + python -m server.sampler_launcher vllm serve --port ... --enable-sleep-mode + +vLLM takes its GPU memory as it starts, so the launcher registers the replica's workload with +the node's time slicer and starts vLLM inside one turn, waits until it is healthy, puts it to +sleep, releases the turn and writes the awake marker the startup probe looks for. From then on +the pool's dispatcher leases the replica for each turn over its own connection. The launcher +keeps the registration, which lasts as long as its connection, and supervises vLLM. It forwards +SIGTERM and SIGINT and exits with vLLM's status. vLLM ignores a SIGTERM while its engine is +starting, so one it has not acted on within a few seconds is followed by SIGKILL, and the turn +goes with the launcher's connection. When the slicer reports on that connection that it could +not park vLLM at the end of a dispatcher's turn, the launcher stops vLLM and exits 1, so the +GPU frees, the grant moves on and the container restarts. +""" + +import asyncio +import contextlib +import os +import signal +import sys +from pathlib import Path + +import httpx + +from accel_timeslicer.time_slicer import time_slicer_client_from_env, workload_from_env +from accel_timeslicer.workload import SAMPLER_CLAIM + +AWAKE_FILE = os.getenv("OPEN_RL_SAMPLER_AWAKE_FILE", "/tmp/open-rl-sampler-awake") +# A vLLM that is up exits well under a second after a SIGTERM; one still starting ignores it. +TERM_GRACE_SECONDS = 5.0 + + +def log(message: str) -> None: + print(f"[sampler_launcher] {message}", file=sys.stderr, flush=True) + + +def server_url(command: list[str]) -> str: + port = command[command.index("--port") + 1] if "--port" in command else "8000" + return f"http://127.0.0.1:{port}" + + +def exit_status(returncode: int) -> int: + return returncode if returncode >= 0 else 128 - returncode + + +class Launcher: + def __init__(self, command: list[str]) -> None: + self.command, self.url = command, server_url(command) + self.boot: asyncio.Task | None = None + self.child: asyncio.subprocess.Process | None = None + + def forward(self, sig: signal.Signals) -> None: + """Hands a signal to vLLM. One that comes before vLLM starts, perhaps while the acquire waits behind + the trainer, ends the boot instead, and the turn goes with the connection when this process exits.""" + if self.child is None: + self.boot.cancel() + elif self.child.returncode is None: + with contextlib.suppress(ProcessLookupError): + self.child.send_signal(sig) + asyncio.get_running_loop().call_later(TERM_GRACE_SECONDS, self.kill) + + def kill(self) -> None: + if self.child.returncode is None: + log(f"vLLM has not exited {TERM_GRACE_SECONDS:g}s after the signal; killing it") + with contextlib.suppress(ProcessLookupError): + os.killpg(self.child.pid, signal.SIGKILL) + + async def stop(self, reason: str) -> int: + log(f"{reason}; stopping vLLM") + self.forward(signal.SIGTERM) + await self.child.wait() + return 1 + + async def run(self) -> int: + awake = Path(AWAKE_FILE) + awake.unlink(missing_ok=True) + self.boot = asyncio.current_task() + loop = asyncio.get_running_loop() + for sig in (signal.SIGTERM, signal.SIGINT): + loop.add_signal_handler(sig, self.forward, sig) + slicer, workload = time_slicer_client_from_env(), workload_from_env(claim=SAMPLER_CLAIM) + payload = workload.as_payload() + try: + await slicer.register(workload) + await slicer.request({"command": "ACQUIRE", **payload}) + # vLLM gets a process group of its own, so the kill takes its engine processes with it. + self.child = await asyncio.create_subprocess_exec(*self.command, start_new_session=True) + except asyncio.CancelledError: + log("stopped before vLLM started") + return 1 + async with httpx.AsyncClient(base_url=self.url) as http: + while not await self.healthy(http): + if self.child.returncode is not None: + log(f"vLLM exited with {self.child.returncode} before it was healthy") + return exit_status(self.child.returncode) + await asyncio.sleep(0.2) + try: + (await http.post("/sleep", params={"level": 1}, timeout=120)).raise_for_status() + except httpx.HTTPError as exc: + return await self.stop(f"vLLM did not sleep ({exc!r})") + try: + await slicer.request({"command": "RELEASE", **payload}) + except Exception as exc: + return await self.stop(f"the time slicer could not park vLLM ({exc!r})") + awake.touch() + log("vLLM asleep and released; the pool's dispatcher takes its turns from here") + # A dispatcher turn the slicer could not park leaves vLLM on the GPU, and the grant moves only once this process is + # gone. The slicer says so on this connection, and then vLLM and the launcher stop. + fault, exited = asyncio.create_task(slicer.fault()), asyncio.create_task(self.child.wait()) + await asyncio.wait({fault, exited}, return_when=asyncio.FIRST_COMPLETED) + if fault.done() and fault.result(): + return await self.stop(f"the time slicer could not park vLLM after a turn ({fault.result()})") + fault.cancel() + return exit_status(await exited) + + async def healthy(self, http: httpx.AsyncClient) -> bool: + try: + return (await http.get("/health", timeout=2)).status_code == 200 + except httpx.HTTPError: + return False + + +def main() -> None: + sys.exit(asyncio.run(Launcher(sys.argv[1:]).run())) + + +if __name__ == "__main__": + main() diff --git a/src/server/sampling_dispatcher.py b/src/server/sampling_dispatcher.py new file mode 100644 index 00000000..94fef181 --- /dev/null +++ b/src/server/sampling_dispatcher.py @@ -0,0 +1,399 @@ +"""Sends one sampling pool's pending records through its router, Envoy in front of the llm-d EPP, and commits the first result. + +A replica that shares its GPU with its trainer is off the GPU between turns, checkpointed by the node's time slicer, so +the dispatcher holds its lease for every call it makes to it. It acquires the workload, wakes vLLM, inspects it, loads +any weights, sends to that replica directly, puts vLLM to sleep and releases. Nothing goes through the router to such a +replica, since the router cannot tell whose turn it is.""" + +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 accel_timeslicer.time_slicer import DEFAULT_TCP_PORT, SocketTimeSlicerClient, TimeSlicerClient, time_slicer_client_from_env, workload_from_env +from accel_timeslicer.workload import SAMPLER_CLAIM, WorkloadRef +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" +# A replica that shares its GPU gives it back after this long with requests still pending, so its trainer gets a turn. +MAX_TURN_SECONDS = float(os.getenv("OPEN_RL_SAMPLER_TURN_SECONDS", "60")) +# A turn ends once the replica has had nothing to send for this long. +IDLE_SECONDS = 1.0 +# Requests still running when a turn ends get this long to answer. The rest are cancelled and sent again in a later turn. +TURN_GRACE_SECONDS = 10.0 +# A call that wakes, prepares or sleeps a replica gives up after this long, so a wedged replica cannot keep its lease, +# and with it the trainer's GPU, for good. A weight update reads a whole checkpoint from shared storage, so it gets longer. +CALL_TIMEOUT_SECONDS = 300.0 +UPDATE_TIMEOUT_SECONDS = 1800.0 +# The worker manager sets this in the env of a replica that shares its GPU, and of its dispatcher on one machine. +TIME_SLICED_SAMPLER = "OPEN_RL_TIME_SLICED_SAMPLER" +# A time-sliced replica is leased with its workload and a client to the time slicer on its node. +Lease = tuple[WorkloadRef, TimeSlicerClient] + + +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 + # The replicas that share their GPU, each with the turn in progress on it. + self.leases: dict[str, Lease] = {} + self.turns: dict[str, asyncio.Task] = {} + self.stopping = False + + 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: + self.stopping = True + for task in self.in_flight.values(): + task.cancel() + # A load cut off midway leaves its replica paused, and a turn cut off leaves it awake on the GPU, so both get a + # moment to finish, resume, sleep and release before the HTTP client closes. + if pending := [*self.loading.values(), *self.turns.values()]: + _, pending = await asyncio.wait(pending, 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" + replicas = await self.list_replicas() + # A time-sliced replica is off the GPU between turns, so it is only looked at inside one. + urls, leased = [u for u in replicas if u not in self.leases], [u for u in replicas if u in self.leases] + 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. + for r in self.to_send()[: self.concurrency - len(self.in_flight)]: + self.start_send(r) + if self.to_send(): + for url in leased: + if url not in self.turns and not self.leases[url][1].faulted: + self.turns[url] = task = asyncio.create_task(self.turn(url)) + task.add_done_callback(lambda _, url=url: self.turns.pop(url)) + + def to_send(self) -> list[dict[str, Any]]: + return [r for r in self.records.values() if r["revision"] == self.revision and r["request_id"] not in self.in_flight] + + def start_send(self, r: dict[str, Any], url: str | None = None) -> asyncio.Task: + self.in_flight[r["request_id"]] = task = asyncio.create_task(self.send(r, url)) + task.add_done_callback(lambda _, request_id=r["request_id"]: self.in_flight.pop(request_id)) + return task + + async def turn(self, url: str) -> None: + """Turns on a replica that shares its GPU, one after another while requests for it keep coming. A turn that + ends on the time bound with requests still pending is followed by the next at once; the slicer serves the + trainer first if it is waiting.""" + workload, slicer = self.leases[url] + while not self.stopping: + try: + backlog = await self.hold(url, workload, slicer) + if slicer.faulted: + raise RuntimeError(f"the time slicer could not park the replica, which still holds its GPU: {slicer.faulted}") + except Exception as exc: + logger.warning("Turn on %s failed: %r", url, exc) + return + if not backlog: + return + + async def hold(self, url: str, workload: WorkloadRef, slicer: TimeSlicerClient) -> bool: + """One turn. It holds the replica's lease throughout, with vLLM woken and inspected, any weight load done as a whole, + the pending requests sent to the replica directly, then vLLM asleep before the lease is released. The turn ends once + the replica has been idle for IDLE_SECONDS, when the dispatcher stops, or after MAX_TURN_SECONDS so the trainer gets + its turn, which is the one case it returns True, for requests still pending. Requests in flight get TURN_GRACE_SECONDS + to finish, and any still running then are cancelled, to be sent again in a later turn.""" + started, idle_since, sends = time.monotonic(), None, set() + async with slicer.acquire(workload): + await self.post(url, "/wake_up") + try: + # A replica an earlier dispatcher left paused or inside a weight update is put right here, inside its turn. + version = await self.inspect(url) + while not self.stopping: + target = self.revision or "default" + if self.revision and version != target: + self.loading[url] = load = asyncio.create_task(self.load_weights(url, self.revision)) + await load + version = await self.weight_version(url) + if version != target: + return False # the load failed and is logged; the next dispatch plans again + if (room := self.concurrency - len(self.in_flight)) > 0 and (to_send := self.to_send()[:room]): + sends |= {self.start_send(r, url) for r in to_send} + idle_since = None + sends = {task for task in sends if not task.done()} + if not sends and not self.to_send(): + idle_since = idle_since or time.monotonic() + if time.monotonic() - idle_since >= IDLE_SECONDS: + return False + if time.monotonic() - started >= MAX_TURN_SECONDS: + break + await asyncio.sleep(0.02) + return bool(self.to_send()) and not self.stopping + finally: + if sends: + # A generate that has not answered by now would keep the trainer off the GPU. Cancelled, its record stays + # pending, and the sleep that follows aborts it inside vLLM. + _, late = await asyncio.wait(sends, timeout=TURN_GRACE_SECONDS) + for task in late: + task.cancel() + if late: + await asyncio.wait(late) + await self.post(url, "/sleep?level=1&mode=abort") + + 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], url: str | None = None) -> None: + """Sends through the router, or to `url`, the replica whose turn it is.""" + 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 = url or 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}}, UPDATE_TIMEOUT_SECONDS) + 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, timeout: float = CALL_TIMEOUT_SECONDS) -> httpx.Response: + return (await self.http.post(url + path, json=body, timeout=timeout)).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) + running = [p for p in pods.items if p.status.pod_ip and p.status.phase == "Running"] + # A time-sliced replica's container counts as started once its launcher has registered it and put vLLM to + # sleep, so no turn is taken on one still booting. + started = [p for p in running if all(s.started for s in p.status.container_statuses or [])] + urls = [f"http://{p.status.pod_ip}:8000" for p in started] + for pod, url in zip(started, urls): + if url not in self.leases and (lease := pod_lease(pod)): + self.leases[url] = lease + # A pod that is gone takes its slicer connection with it. + for url in [u for u in self.leases if u not in urls and u not in self.turns]: + await self.leases.pop(url)[1].close() + return urls + + 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 pod_lease(pod: Any) -> Lease | None: + """For a time-sliced replica, the workload its launcher registered, read from the same env it read, and + a client to the time slicer on the replica's node, which listens on the host network.""" + env = {e.name: e.value for c in pod.spec.containers if c.name == "worker" for e in c.env or []} + if env.get(TIME_SLICED_SAMPLER) != "true": + return None + name = env.get("OPEN_RL_WORKLOAD_ID") or env.get("OPEN_RL_TIME_SLICE_JOB_ID") or pod.metadata.name + claim = env.get("OPEN_RL_CLAIM") or env.get("OPEN_RL_TIME_SLICE_GROUP") or SAMPLER_CLAIM + port = int(env.get("OPEN_RL_ACCEL_TIMESLICER_PORT") or DEFAULT_TCP_PORT) + return WorkloadRef(name, claim), SocketTimeSlicerClient(host=pod.status.host_ip, port=port) + + +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"))) + if os.getenv(TIME_SLICED_SAMPLER) == "true": + # The worker manager marks a local replica's dispatcher time-sliced too, with the replica's tags, when the replica shares its GPU. + dispatcher.leases = {url: (workload_from_env(claim=SAMPLER_CLAIM), time_slicer_client_from_env()) for url in urls} + 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..13403a84 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 TIME_SLICED_SAMPLER, 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,37 @@ 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) + pool = meta.pool(runtime) if role == "sampler" else None exclusive = not meta.shares_gpu() devices = meta.trainer_gpus if role == "trainer" else 1 size = footprint(base_model, meta.fine_tuning_type, role) # Each torchrun rank is its own process with its own host memory. size = replace(size, host_request_bytes=size.host_request_bytes * devices, host_limit_bytes=size.host_limit_bytes * devices) - return Worker(role, runtime, base_model, is_lora, exclusive, meta, size, devices) + return Worker(role, runtime, base_model, is_lora, exclusive, meta, size, devices, pool=pool) def pod_env(worker: Worker) -> list[dict[str, Any]]: """The shared worker env plus what only the cluster knows. The time-slice group is placement's and the scheduler stamps it.""" - tmp_dir = os.getenv("OPEN_RL_TMP_DIR", "/mnt/shared/open-rl") + tmp_dir = tmp_root() values = { "REDIS_URL": os.environ["REDIS_URL"], "OPEN_RL_TMP_DIR": tmp_dir, @@ -127,7 +147,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 +199,115 @@ 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 not worker.exclusive: + # The replica shares its GPU with its trainer, so the launcher boots vLLM, able to sleep, inside a time-slicer + # turn, and the pool's dispatcher, which finds the replica by this variable, takes its turns after that. The + # container starts once the launcher marks the replica up, so no turn is taken on one still booting. + container["command"] = ["uv", "run", "python", "-u", "-m", "server.sampler_launcher", "vllm", "serve"] + container["args"].append("--enable-sleep-mode") + container["env"].append({"name": TIME_SLICED_SAMPLER, "value": "true"}) + container["startupProbe"] = {"exec": {"command": ["test", "-f", "/tmp/open-rl-sampler-awake"]}, "periodSeconds": 2, "failureThreshold": 600} + 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 +316,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 +337,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 +432,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 +474,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 +490,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 +515,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..2324f766 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,14 @@ from pathlib import Path from typing import Any, Protocol -from accel_timeslicer.workload import SAMPLER_TIME_SLICE_GROUP, TRAINER_TIME_SLICE_GROUP, workload_job_id +from accel_timeslicer.time_slicer import time_slicing_enabled +from accel_timeslicer.workload import TRAINER_TIME_SLICE_GROUP, workload_job_id from server.estimator import footprint +from server.lora_snapshots import snapshot_root from server.model_metadata import TrainingModelMetadata, decode_model_metadata +from server.sampling_dispatcher import TIME_SLICED_SAMPLER from server.store import get_state_store +from server.vllm_options import serve_args PROJECT_DIR = Path(__file__).resolve().parents[2] @@ -108,16 +113,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 +178,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 +189,65 @@ 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] + runner = ["uv", "run", "--extra", "gpu", "--extra", "vllm"] if shutil.which("uv") else [] + vllm = ["vllm"] if runner else [str(Path(sys.executable).with_name("vllm"))] + serve = ["serve", *serve_args(base_model, is_lora, env, port)] + 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) + 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") + if meta.shares_gpu() and time_slicing_enabled(): + # The replica takes turns on the one GPU with its trainer, so it is in the trainer's claim. The launcher boots + # vLLM inside a turn, and the dispatcher, marked with the same tags, holds the lease for every turn after. + tags = { + TIME_SLICED_SAMPLER: "true", + "OPEN_RL_TIME_SLICE_JOB_ID": workload_job_id(role, runtime), + "OPEN_RL_TIME_SLICE_GROUP": TRAINER_TIME_SLICE_GROUP, + } + env.update(tags) + dispatcher_env.update(tags) + self.spawn(role, runtime, python_command(["gpu", "vllm"], "server.sampler_launcher", [*vllm, *serve, "--enable-sleep-mode"]), env) + else: + self.spawn(role, runtime, [*runner, *vllm, *serve], env) + 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_accel_timeslicer.py b/tests/test_accel_timeslicer.py index cba838ff..57630e7b 100644 --- a/tests/test_accel_timeslicer.py +++ b/tests/test_accel_timeslicer.py @@ -11,7 +11,7 @@ from accel_timeslicer.llmd import LlmDCheckpointRestorer from accel_timeslicer.serve import start_tcp_time_slicer, start_time_slicer from accel_timeslicer.single_node import SingleNodeTimeSlicer -from accel_timeslicer.time_slicer import SocketTimeSlicerClient, time_slicer_client_from_env, workload_from_env +from accel_timeslicer.time_slicer import SocketTimeSlicerClient, TimeSlicerFault, time_slicer_client_from_env, workload_from_env from accel_timeslicer.workload import WorkloadRef @@ -67,6 +67,13 @@ def checkpoint(self, target: WorkloadRef) -> None: raise RuntimeError("cuda-checkpoint refused the process") +class FirstCheckpointFailsRestorer(RecordingRestorer): + def checkpoint(self, target: WorkloadRef) -> None: + super().checkpoint(target) + if len(self.calls) == 1: + raise RuntimeError("cuda-checkpoint refused the process") + + class SingleNodeTimeSlicerTest(unittest.IsolatedAsyncioTestCase): async def test_agent_grants_only_one_active_process_at_a_time(self) -> None: restorer = RecordingRestorer() @@ -313,6 +320,161 @@ async def test_socket_client_remembers_the_fault_and_keeps_its_connection(self) await server.wait_closed() +class LeaseHolderTest(unittest.IsolatedAsyncioTestCase): + """A workload one connection registered and another holds, as a pool replica's launcher and its dispatcher do.""" + + async def test_closing_the_holders_connection_checkpoints_the_workload_before_the_next_grant(self) -> None: + restorer = BlockingRestorer() + agent = SingleNodeTimeSlicer(restorer) + await agent.register(WorkloadRef(name="replica"), connection_id=1) + await agent.register(WorkloadRef(name="trainer"), connection_id=3) + self.assertTrue((await agent.acquire(WorkloadRef(name="replica"), connection_id=2))["ok"]) + blocked = asyncio.create_task(agent.acquire(WorkloadRef(name="trainer"), connection_id=3)) + await asyncio.sleep(0.05) + restorer.block_checkpoint = True + closed = asyncio.create_task(agent.connection_closed(2)) + await asyncio.to_thread(restorer.checkpoint_started.wait, 1.0) + # The replica is still being checkpointed, so the trainer is not granted yet. + self.assertFalse(blocked.done()) + self.assertEqual(agent.running, {"shared-accelerator": "replica"}) + restorer.finish_checkpoint.set() + await closed + self.assertTrue((await asyncio.wait_for(blocked, timeout=1.0))["ok"]) + self.assertEqual(agent.running, {"shared-accelerator": "trainer"}) + # The replica's own process is still up behind its own connection, so the next dispatcher leases it again. + self.assertEqual((agent.workloads["replica"].failed, agent.workloads["replica"].checkpointed), (False, True)) + await agent.release(WorkloadRef(name="trainer")) + self.assertTrue((await agent.acquire(WorkloadRef(name="replica"), connection_id=4))["ok"]) + self.assertEqual(restorer.simple_labels(), [("checkpoint", "replica"), ("checkpoint", "trainer"), ("restore", "replica")]) + + async def test_only_the_connection_holding_a_grant_can_release_it(self) -> None: + restorer = RecordingRestorer() + agent = SingleNodeTimeSlicer(restorer) + await agent.register(WorkloadRef(name="replica"), connection_id=1) + self.assertTrue((await agent.acquire(WorkloadRef(name="replica"), connection_id=2))["ok"]) + refused = await agent.release(WorkloadRef(name="replica"), connection_id=3) + error = "workload replica is held by another connection; only the connection that acquired it can release it" + self.assertEqual(refused, {"ok": False, "error": error}) + self.assertEqual((agent.running, restorer.calls), ({"shared-accelerator": "replica"}, [])) + self.assertTrue((await agent.release(WorkloadRef(name="replica"), connection_id=2))["ok"]) + # A caller in the slicer's own process names no connection. + self.assertTrue((await agent.acquire(WorkloadRef(name="replica"), connection_id=4))["ok"]) + self.assertTrue((await agent.release(WorkloadRef(name="replica")))["ok"]) + + async def test_a_failed_checkpoint_after_the_holders_connection_closed_keeps_the_grant(self) -> None: + agent = SingleNodeTimeSlicer(FailingCheckpointRestorer()) + await agent.register(WorkloadRef(name="replica"), connection_id=1) + self.assertTrue((await agent.acquire(WorkloadRef(name="replica"), connection_id=2))["ok"]) + blocked = asyncio.create_task(agent.acquire(WorkloadRef(name="trainer"), connection_id=3)) + await asyncio.sleep(0.05) + await agent.connection_closed(2) + await asyncio.sleep(0.05) + # The replica is still on the devices, so it keeps the grant until its own process is gone. + self.assertFalse(blocked.done()) + self.assertEqual(agent.running, {"shared-accelerator": "replica"}) + self.assertTrue(agent.workloads["replica"].failed) + await agent.connection_closed(1) + self.assertTrue((await asyncio.wait_for(blocked, timeout=1.0))["ok"]) + + async def test_a_socket_client_that_drops_while_holding_anothers_workload_is_released(self) -> None: + restorer = RecordingRestorer() + agent = SingleNodeTimeSlicer(restorer) + with tempfile.TemporaryDirectory() as tmp: + socket_path = str(Path(tmp) / "accel-timeslicer.sock") + server = await start_time_slicer(agent, socket_path) + launcher, dispatcher, trainer = (SocketTimeSlicerClient(socket_path) for _ in range(3)) + try: + await launcher.register(WorkloadRef(name="replica")) + await trainer.register(WorkloadRef(name="trainer")) + await dispatcher.request({"command": "ACQUIRE", "name": "replica", "claim": "shared-accelerator"}) + blocked = asyncio.create_task(acquire_once(trainer, WorkloadRef(name="trainer"))) + await asyncio.sleep(0.05) + self.assertFalse(blocked.done()) + await dispatcher.close() + self.assertEqual(await asyncio.wait_for(blocked, timeout=1.0), "trainer") + self.assertEqual(restorer.simple_labels()[0], ("checkpoint", "replica")) + self.assertFalse(agent.workloads["replica"].failed) + finally: + for client in (launcher, dispatcher, trainer): + await client.close() + server.close() + await server.wait_closed() + + async def test_a_failed_park_is_told_to_the_connection_that_registered_the_workload(self) -> None: + agent = SingleNodeTimeSlicer(FailingCheckpointRestorer()) + with tempfile.TemporaryDirectory() as tmp: + socket_path = str(Path(tmp) / "accel-timeslicer.sock") + server = await start_time_slicer(agent, socket_path) + launcher, dispatcher, trainer = (SocketTimeSlicerClient(socket_path) for _ in range(3)) + try: + await launcher.register(WorkloadRef(name="replica")) + async with dispatcher.acquire(WorkloadRef(name="replica")): + pass + # The dispatcher hears it in the reply to its release. The launcher, the one process that can free the devices, is told unasked. + self.assertIn("must exit", dispatcher.faulted or "") + self.assertEqual(await asyncio.wait_for(launcher.fault(), 1.0), dispatcher.faulted) + # A process that registered and released its workload itself hears it once, in the reply. + await trainer.register(WorkloadRef(name="trainer", claim="other")) + async with trainer.acquire(WorkloadRef(name="trainer", claim="other")): + pass + self.assertIn("must exit", trainer.faulted or "") + with self.assertRaises(TimeoutError): + await asyncio.wait_for(trainer.reader.readline(), 0.2) + finally: + for client in (launcher, dispatcher, trainer): + await client.close() + server.close() + await server.wait_closed() + + async def test_when_a_dead_dispatchers_turn_cannot_be_parked_the_launcher_is_told_and_its_exit_frees_the_claim(self) -> None: + agent = SingleNodeTimeSlicer(FailingCheckpointRestorer()) + with tempfile.TemporaryDirectory() as tmp: + socket_path = str(Path(tmp) / "accel-timeslicer.sock") + server = await start_time_slicer(agent, socket_path) + launcher, dispatcher, trainer = (SocketTimeSlicerClient(socket_path) for _ in range(3)) + try: + await launcher.register(WorkloadRef(name="replica")) + await trainer.register(WorkloadRef(name="trainer")) + await dispatcher.request({"command": "ACQUIRE", "name": "replica", "claim": "shared-accelerator"}) + blocked = asyncio.create_task(acquire_once(trainer, WorkloadRef(name="trainer"))) + # The dispatcher dies holding the turn, and the park the slicer runs for it fails. + await dispatcher.close() + self.assertIn("must exit", await asyncio.wait_for(launcher.fault(), 1.0) or "") + self.assertFalse(blocked.done()) + # The launcher stops vLLM and exits, which closes its connection, and only then is the trainer granted. + await launcher.close() + self.assertEqual(await asyncio.wait_for(blocked, timeout=1.0), "trainer") + finally: + for client in (launcher, dispatcher, trainer): + await client.close() + server.close() + await server.wait_closed() + + async def test_a_fault_told_while_a_request_waits_is_noted_and_not_taken_for_its_reply(self) -> None: + agent = SingleNodeTimeSlicer(FirstCheckpointFailsRestorer()) + with tempfile.TemporaryDirectory() as tmp: + socket_path = str(Path(tmp) / "accel-timeslicer.sock") + server = await start_time_slicer(agent, socket_path) + launcher, dispatcher = SocketTimeSlicerClient(socket_path), SocketTimeSlicerClient(socket_path) + try: + await launcher.register(WorkloadRef(name="replica")) + await dispatcher.request({"command": "ACQUIRE", "name": "replica", "claim": "shared-accelerator"}) + waiting = asyncio.create_task(launcher.request({"command": "ACQUIRE", "name": "other", "claim": "shared-accelerator"})) + await asyncio.sleep(0.05) + # The release fails and the launcher is told while its acquire waits. Closing the dispatcher parks the replica after all. + with self.assertRaises(TimeSlicerFault): + await dispatcher.request({"command": "RELEASE", "name": "replica", "claim": "shared-accelerator"}) + await dispatcher.close() + self.assertTrue((await asyncio.wait_for(waiting, timeout=1.0))["ok"]) + self.assertIn("must exit", launcher.faulted or "") + self.assertEqual(agent.running, {"shared-accelerator": "other"}) + finally: + await launcher.close() + await dispatcher.close() + server.close() + await server.wait_closed() + + class SingleNodeTimeSlicerSocketTest(unittest.IsolatedAsyncioTestCase): async def test_persistent_socket_clients_alternate(self) -> None: restorer = RecordingRestorer() 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_gpu_leases.py b/tests/test_sampler_gpu_leases.py new file mode 100644 index 00000000..239417c1 --- /dev/null +++ b/tests/test_sampler_gpu_leases.py @@ -0,0 +1,379 @@ +"""A pool whose replicas share their GPU with their trainer. The dispatcher holds each replica's lease from +the real time slicer, over its socket, for every call it makes to the replica.""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import patch + +import httpx +from starlette.responses import Response +from starlette.routing import Route + +from accel_timeslicer.serve import start_time_slicer +from accel_timeslicer.single_node import SingleNodeTimeSlicer +from accel_timeslicer.time_slicer import SocketTimeSlicerClient +from accel_timeslicer.workload import WorkloadRef +from server import sampling_dispatcher +from server.sampling_dispatcher import Dispatcher, revision_file +from tests.test_sampling_dispatcher import SAMPLE, DispatcherTestCase, FakeReplica, router, until + +CLAIM = "gpu-0" +# What vLLM with sleep mode still answers while asleep. A checkpointed process answers nothing. +ASLEEP_OK = {"/health", "/weight_info", "/is_sleeping", "/wake_up", "/sleep"} + + +class LeasedReplica(FakeReplica): + """A FakeReplica with vLLM's sleep mode. It starts asleep, as its launcher leaves it, is `parked` while the + slicer holds it checkpointed, and logs any call that reaches it asleep or parked, which vLLM could not answer. + `gpu` holds the names of everything awake on the device, so an overlap shows up in the log.""" + + def __init__(self, name: str, log: list, version: str, expected, gpu: set): + super().__init__(name, log, version, expected) + self.awake, self.parked, self.gpu, self.wakes, self.delay = False, False, gpu, 0, 0.0 + self.app.router.routes += [Route("/sleep", self.sleep, methods=["POST"]), Route("/wake_up", self.wake_up, methods=["POST"])] + inner = self.app + + async def guarded(scope, receive, send): + path = scope.get("path") + if scope["type"] == "http" and (self.parked or (not self.awake and path not in ASLEEP_OK)): + self.log.append(("parked" if self.parked else "asleep", self.name, path)) + await Response(status_code=500)(scope, receive, send) + return + await inner(scope, receive, send) + + self.app = guarded + + async def sleep(self, request) -> Response: + self.awake = False + self.gpu.discard(self.name) + self.log.append(("sleep", self.name)) + return Response() + + async def wake_up(self, request) -> Response: + self.note_overlap() + self.awake, self.wakes = True, self.wakes + 1 + self.gpu.add(self.name) + self.log.append(("wake", self.name)) + return Response() + + async def generate(self, request) -> Response: + self.note_overlap() + await asyncio.sleep(self.delay) + return await super().generate(request) + + def note_overlap(self) -> None: + if others := self.gpu - {self.name}: + self.log.append(("overlap", self.name, others)) + + +class ParkingRestorer: + """Records checkpoints and restores in the shared log and parks the replica, since a checkpointed process answers nothing. + The checkpoint of a workload named in `fail` raises instead, as when the node's host memory is full.""" + + def __init__(self, log: list, replicas: dict[str, LeasedReplica]): + self.log, self.replicas, self.fail = log, replicas, set() + + def checkpoint(self, workload: WorkloadRef) -> None: + if workload.name in self.fail: + self.log.append(("checkpoint-failed", workload.name)) + raise RuntimeError("could not park the process") + self.log.append(("checkpoint", workload.name)) + if replica := self.replicas.get(workload.name): + replica.parked = True + replica.gpu.discard(replica.name) + + def restore(self, workload: WorkloadRef) -> None: + self.log.append(("restore", workload.name)) + if replica := self.replicas.get(workload.name): + replica.parked = False + + +class Trainer: + """A trainer on the same claim, holding the GPU for each step over its own slicer connection.""" + + def __init__(self, socket_path: str, log: list, gpu: set): + self.client, self.workload = SocketTimeSlicerClient(socket_path=socket_path), WorkloadRef("trainer", CLAIM) + self.log, self.gpu, self.steps = log, gpu, 0 + + async def run(self, steps: int, hold: float = 0.1) -> None: + await self.client.register(self.workload) + for _ in range(steps): + async with self.client.acquire(self.workload): + if self.gpu: + self.log.append(("overlap", "trainer", set(self.gpu))) + self.gpu.add("trainer") + self.log.append(("step", "trainer")) + await asyncio.sleep(hold) + self.gpu.discard("trainer") + self.steps += 1 + await asyncio.sleep(0.01) + await self.client.close() + + +class SamplerGpuLeaseTest(DispatcherTestCase): + async def asyncSetUp(self) -> None: + await super().asyncSetUp() + self.enterContext(patch.object(sampling_dispatcher, "IDLE_SECONDS", 0.1)) + self.socket = str(self.tmp / "slicer.sock") + self.gpu: set = set() + self.leased: dict[str, LeasedReplica] = {} + self.launchers: dict[str, SocketTimeSlicerClient] = {} + self.agent = SingleNodeTimeSlicer(ParkingRestorer(self.log, self.leased)) + server = await start_time_slicer(self.agent, self.socket) + self.addAsyncCleanup(server.wait_closed) + self.addCleanup(server.close) + + def workload(self, replica: FakeReplica) -> WorkloadRef: + return WorkloadRef(replica.name, CLAIM) + + async def replica(self, version: str = "default") -> LeasedReplica: + """A time-sliced replica as its launcher leaves it: registered over the launcher's own connection, + which stays open, asleep and checkpointed off the GPU.""" + replica = LeasedReplica(f"pod-{len(self.leased)}", self.log, version, revision_file(self.pool), self.gpu) + replica.url = await self.serve(replica.app) + self.addAsyncCleanup(replica.held.clear) + await self.boot(replica) + self.leased[replica.name] = replica + return replica + + async def boot(self, replica: LeasedReplica) -> None: + """What the replica's launcher does once vLLM is up. It registers over a connection of its own, kept in + `launchers` and left open, and its boot turn ends in a release.""" + launcher = SocketTimeSlicerClient(socket_path=self.socket) + self.addAsyncCleanup(launcher.close) + await launcher.register(self.workload(replica)) + async with launcher.acquire(self.workload(replica)): + pass + self.launchers[replica.name] = launcher + + async def leased_dispatcher(self, replicas: list[LeasedReplica], concurrency: int = 4) -> tuple[Dispatcher, asyncio.Task]: + router_url = await self.serve(router(self.log, replicas, self.http)) + leases = {} + for replica in replicas: + slicer = SocketTimeSlicerClient(socket_path=self.socket) + self.addAsyncCleanup(slicer.close) + leases[replica.url] = (self.workload(replica), slicer) + dispatcher = Dispatcher(self.store, self.pool, self.http, router_url, [r.url for r in replicas], concurrency=concurrency) + dispatcher.leases = leases + run = asyncio.create_task(dispatcher.run()) + self.addAsyncCleanup(lambda: (run.cancel(), asyncio.gather(run, return_exceptions=True))[1]) + return dispatcher, run + + def events(self, *kinds: str) -> list[tuple]: + return [entry for entry in self.log if entry[0] in kinds] + + def violations(self) -> list[tuple]: + return self.events("asleep", "parked", "overlap", "unrouted") + + async def test_a_turn_restores_wakes_serves_and_sleeps_the_replica_before_releasing_it(self): + replica = await self.replica() + records = [await self.accepted() for _ in range(3)] + await self.leased_dispatcher([replica]) + for r in records: + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + await until(lambda: len(self.events("checkpoint")) == 2) + generates = [("generate", "pod-0", "default")] * 3 + expected = [("checkpoint", "pod-0"), ("restore", "pod-0"), ("wake", "pod-0"), *generates, ("sleep", "pod-0"), ("checkpoint", "pod-0")] + self.assertEqual(self.events("checkpoint", "restore", "wake", "sleep", "generate"), expected) + self.assertEqual((self.violations(), self.agent.running, replica.wakes), ([], {}, 1)) + + async def test_a_busy_replica_gives_the_trainer_its_turn_and_comes_back(self): + self.enterContext(patch.object(sampling_dispatcher, "MAX_TURN_SECONDS", 0.2)) + replica = await self.replica() + replica.delay = 0.05 # a backlog of 20 outlasts the longest turn several times over + trainer = Trainer(self.socket, self.log, self.gpu) + records = [await self.accepted() for _ in range(20)] + await self.leased_dispatcher([replica], concurrency=2) + steps = asyncio.create_task(trainer.run(steps=3)) + for r in records: + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + await asyncio.wait_for(steps, timeout=20) + self.assertEqual((self.violations(), trainer.steps), ([], 3)) + # The replica gave the GPU up with requests still pending and came back once the trainer had stepped. + turns = [entry[0] for entry in self.events("step", "wake")] + self.assertGreaterEqual(replica.wakes, 2) + first, last = turns.index("wake"), len(turns) - 1 - turns[::-1].index("wake") + self.assertIn("step", turns[first:last]) + + async def test_two_replicas_on_one_gpu_take_turns(self): + replicas = [await self.replica(), await self.replica()] + records = [await self.accepted() for _ in range(8)] + await self.leased_dispatcher(replicas, concurrency=2) + for r in records: + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + self.assertEqual(self.violations(), []) + self.assertEqual(sum(entry[0] == "generate" for entry in self.log), 8) + + async def test_a_weight_load_never_splits_across_turns(self): + self.enterContext(patch.object(sampling_dispatcher, "MAX_TURN_SECONDS", 0.05)) + w1 = self.revision("w1", format="sparse_delta", parent="") + w2 = self.revision("w2", format="sparse_delta", parent=w1) + replica = await self.replica(w1) + replica.held.set() # the update waits to be let go, so the load outlasts the longest turn + trainer = Trainer(self.socket, self.log, self.gpu) + r = await self.accepted(w2) + await self.leased_dispatcher([replica]) + await until(lambda: ("update", "pod-0", w2, w2) in self.log) + step = asyncio.create_task(trainer.run(steps=1)) + await asyncio.sleep(0.3) + self.assertFalse(step.done()) + replica.held.clear() + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + await asyncio.wait_for(step, timeout=10) + sequence = [entry[:2] for entry in self.events("checkpoint", "restore", "update", "generate", "sleep", "step")] + turn = [("restore", "pod-0"), ("update", "pod-0"), ("generate", "pod-0"), ("sleep", "pod-0"), ("checkpoint", "pod-0")] + self.assertEqual(sequence, [("checkpoint", "pod-0"), *turn, ("step", "trainer"), ("checkpoint", "trainer")]) + self.assertEqual((self.violations(), replica.version), ([], w2)) + + async def test_a_turn_after_a_clean_load_sends_no_finish_probe_and_a_turn_after_a_broken_load_does(self): + # The probe is a bare finish_weight_update, and vLLM logs a 500 traceback for each one that finds no update open. + w1, w2, w3 = self.revision("w1"), self.revision("w2"), self.revision("w3") + replica = await self.replica() + r1 = await self.accepted(w1) + await self.leased_dispatcher([replica]) + self.assertEqual(await self.committed(r1["request_id"]), [SAMPLE]) + await until(lambda: len(self.events("checkpoint")) == 2) # the first turn is over + r2 = await self.accepted(w2) + self.assertEqual(await self.committed(r2["request_id"]), [SAMPLE]) + self.assertEqual((replica.wakes, replica.probes), (2, 1)) + # The third turn's load breaks off and closes its own update with a bare finish, so the fourth turn probes. + replica.fail_update = True + r3 = await self.accepted(w3) + self.assertEqual(await self.committed(r3["request_id"]), [SAMPLE]) + self.assertEqual((replica.wakes, replica.probes, self.violations()), (4, 3, [])) + + async def test_a_dispatcher_that_dies_mid_turn_leaves_the_replica_checkpointed_before_the_next_grant(self): + replica = await self.replica() + replica.held.set() + r = await self.accepted() + dispatcher, run = await self.leased_dispatcher([replica]) + await until(lambda: ("generate", "pod-0", "default") in self.log) + trainer = Trainer(self.socket, self.log, self.gpu) + step = asyncio.create_task(trainer.run(steps=1)) + await asyncio.sleep(0.1) + self.assertFalse(step.done()) + # The dispatcher dies: it makes no more calls and its slicer connection drops without a release. + dead = httpx.MockTransport(lambda request: (_ for _ in ()).throw(httpx.ConnectError("the dispatcher is dead"))) + dispatcher.http = await self.enterAsyncContext(httpx.AsyncClient(transport=dead)) + dispatcher.leases[replica.url][1].writer.transport.abort() + run.cancel() + await asyncio.gather(run, return_exceptions=True) + await asyncio.wait_for(step, timeout=10) + grants = [entry[:2] for entry in self.events("checkpoint", "restore", "wake", "step")] + turn = [("restore", "pod-0"), ("wake", "pod-0"), ("checkpoint", "pod-0")] + self.assertEqual(grants, [("checkpoint", "pod-0"), *turn, ("step", "trainer"), ("checkpoint", "trainer")]) + self.assertEqual(self.violations(), []) + # The replica is registered and well behind its launcher's connection, so the next dispatcher leases it again. + self.assertFalse(self.agent.workloads["pod-0"].failed) + await self.leased_dispatcher([replica]) + replica.held.clear() + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + self.assertEqual(self.violations(), []) + + async def test_a_failed_park_at_the_end_of_a_turn_stops_the_replica_so_the_trainer_goes_next(self): + replica = await self.replica() + self.agent.restorer.fail.add(replica.name) + first = await self.accepted() + dispatcher, _ = await self.leased_dispatcher([replica]) + self.assertEqual(await self.committed(first["request_id"]), [SAMPLE]) + trainer = Trainer(self.socket, self.log, self.gpu) + step = asyncio.create_task(trainer.run(steps=1)) + # The park at the end of the turn fails, so the replica keeps the GPU, and its launcher is told on its own connection. + launcher = self.launchers[replica.name] + self.assertIn("must exit", await asyncio.wait_for(launcher.fault(), 10) or "") + await asyncio.sleep(0.2) + self.assertFalse(step.done()) + # The launcher stops vLLM and exits, which closes its connection, and then the trainer is granted. + await launcher.close() + await asyncio.wait_for(step, timeout=5) + # Restarted, the replica boots in a turn of its own again, and a new lease on it serves, as list_replicas makes one + # once the restarted container has started. + self.agent.restorer.fail.clear() + await self.boot(replica) + stale, fresh = dispatcher.leases[replica.url][1], SocketTimeSlicerClient(socket_path=self.socket) + self.addAsyncCleanup(fresh.close) + dispatcher.leases[replica.url] = (self.workload(replica), fresh) + await stale.close() + second = await self.accepted() + self.assertEqual(await self.committed(second["request_id"]), [SAMPLE]) + self.assertEqual(self.violations(), []) + + async def test_a_turn_that_outlives_its_launcher_cannot_end_the_restarted_launchers_boot_turn(self): + self.enterContext(patch.object(sampling_dispatcher, "MAX_TURN_SECONDS", 2.0)) + replica = await self.replica() + await self.accepted() + replica.held.set() # a generate is in flight when vLLM dies + dispatcher, _ = await self.leased_dispatcher([replica]) + await until(lambda: ("generate", "pod-0", "default") in self.log) + with self.assertLogs("server.sampling_dispatcher", "WARNING") as logs: + # vLLM dies mid-turn. The replica stops answering, and its launcher exits with it, which closes its connection. + dead = httpx.MockTransport(lambda request: (_ for _ in ()).throw(httpx.ConnectError("vLLM is gone"))) + dispatcher.http = await self.enterAsyncContext(httpx.AsyncClient(transport=dead)) + replica.held.clear() + await self.launchers[replica.name].close() + await until(lambda: self.agent.workloads["pod-0"].failed) + # kubelet restarts the container in the same pod, and the new launcher boots vLLM in a turn of its own. + restarted = SocketTimeSlicerClient(socket_path=self.socket) + self.addAsyncCleanup(restarted.close) + await restarted.register(self.workload(replica)) + await restarted.request({"command": "ACQUIRE", **self.workload(replica).as_payload()}) + booting = len(self.log) + # The old turn reaches MAX_TURN_SECONDS and releases the same workload, and the slicer refuses. + await until(lambda: any("held by another connection" in line for line in logs.output)) + self.assertEqual(self.agent.running, {CLAIM: "pod-0"}) + self.assertNotIn(("checkpoint", "pod-0"), self.log[booting:]) + # The trainer waits for the boot turn, which the restarted launcher ends itself once vLLM sleeps. + trainer = Trainer(self.socket, self.log, self.gpu) + step = asyncio.create_task(trainer.run(steps=1)) + await asyncio.sleep(0.2) + self.assertFalse(step.done()) + self.assertTrue((await restarted.request({"command": "RELEASE", **self.workload(replica).as_payload()}))["ok"]) + await asyncio.wait_for(step, timeout=5) + + async def test_a_generate_that_hangs_past_the_turn_is_cancelled_so_the_trainer_steps_and_it_is_sent_again(self): + self.enterContext(patch.object(sampling_dispatcher, "MAX_TURN_SECONDS", 0.5)) + self.enterContext(patch.object(sampling_dispatcher, "TURN_GRACE_SECONDS", 0.5)) + replica = await self.replica() + replica.held.set() # a wedged replica, whose generate does not answer + r = await self.accepted() + await self.leased_dispatcher([replica]) + await until(lambda: ("generate", "pod-0", "default") in self.log) + trainer = Trainer(self.socket, self.log, self.gpu) + self.addAsyncCleanup(trainer.client.close) + await asyncio.wait_for(trainer.run(steps=1), timeout=5) + # Once the replica answers again, the record goes out in a later turn and is committed once. + replica.held.clear() + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + generates = [i for i, entry in enumerate(self.log) if entry[0] == "generate"] + self.assertLess(self.log.index(("step", "trainer")), generates[-1]) + self.assertEqual(self.violations(), []) + + async def test_an_exclusive_pool_makes_no_slicer_calls(self): + replicas = await self.replicas("default") + r = await self.accepted() + await self.dispatcher(replicas) + self.assertEqual(await self.committed(r["request_id"]), [SAMPLE]) + self.assertEqual((self.events("checkpoint", "restore"), self.agent.workloads), ([], {})) + + async def test_pods_marked_time_sliced_are_leased_through_the_slicer_on_their_node_once_started(self): + def pod(name: str, ip: str | None, env: dict, started: bool = True) -> SimpleNamespace: + container = SimpleNamespace(name="worker", env=[SimpleNamespace(name=k, value=v) for k, v in env.items()]) + status = SimpleNamespace(pod_ip=ip, host_ip="10.1.0.5", phase="Running", container_statuses=[SimpleNamespace(started=started)]) + return SimpleNamespace(metadata=SimpleNamespace(name=name), spec=SimpleNamespace(containers=[container]), status=status) + + identity = {"OPEN_RL_WORKLOAD_ID": "fft-job-sampler", "OPEN_RL_TIME_SLICE_GROUP": "claim-7", "OPEN_RL_ACCEL_TIMESLICER_PORT": "9753"} + shared = pod("shared", "10.0.0.1", {**identity, "OPEN_RL_TIME_SLICED_SAMPLER": "true"}) + alone = pod("alone", "10.0.0.2", {"OPEN_RL_WORKLOAD_ID": "lora-x-0-sampler", "OPEN_RL_TIME_SLICING": "off"}) + # A time-sliced replica still booting has not passed its startup probe, so no turn is taken on it yet. + booting = pod("booting", "10.0.0.3", {**identity, "OPEN_RL_TIME_SLICED_SAMPLER": "true"}, started=False) + dispatcher = Dispatcher(self.store, self.pool, self.http, "http://router", [], concurrency=4) + dispatcher.k8s = SimpleNamespace( + list_namespaced_pod=lambda *args, **kwargs: SimpleNamespace(items=[shared, alone, booting, pod("no-ip", None, {})]) + ) + self.assertEqual(await dispatcher.list_replicas(), ["http://10.0.0.1:8000", "http://10.0.0.2:8000"]) + ((url, (workload, slicer)),) = dispatcher.leases.items() + self.assertEqual((url, workload, slicer.host, slicer.port), ("http://10.0.0.1:8000", WorkloadRef("fft-job-sampler", "claim-7"), "10.1.0.5", 9753)) + # The lease goes when the pod does. + dispatcher.k8s = SimpleNamespace(list_namespaced_pod=lambda *args, **kwargs: SimpleNamespace(items=[alone])) + self.assertEqual(await dispatcher.list_replicas(), ["http://10.0.0.2:8000"]) + self.assertEqual(dispatcher.leases, {}) diff --git a/tests/test_sampler_launcher.py b/tests/test_sampler_launcher.py new file mode 100644 index 00000000..e4b4d3c4 --- /dev/null +++ b/tests/test_sampler_launcher.py @@ -0,0 +1,181 @@ +"""Drives the sampler launcher, as its own process, for a replica that shares its GPU, against the real +time slicer over its socket and tests/vllm_stub.py on PATH as vllm.""" + +import asyncio +import os +import signal +import socket +import subprocess +import sys +import tempfile +import time +import unittest +from pathlib import Path + +from accel_timeslicer.serve import start_time_slicer +from accel_timeslicer.single_node import SingleNodeTimeSlicer +from accel_timeslicer.time_slicer import SocketTimeSlicerClient +from accel_timeslicer.workload import WorkloadRef +from tests.test_accel_timeslicer import BlockingRestorer, FailingCheckpointRestorer + +REPO = Path(__file__).resolve().parents[1] +CLAIM = "gpu-0" + + +def free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +async def until(condition, timeout: float = 20.0) -> None: + async with asyncio.timeout(timeout): + while not condition(): + await asyncio.sleep(0.02) + + +class SamplerLauncherTest(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self) -> None: + self.dir = Path(self.enterContext(tempfile.TemporaryDirectory())) + (self.dir / "vllm").symlink_to(REPO / "tests" / "vllm_stub.py") + self.socket = str(self.dir / "slicer.sock") + self.restorer = BlockingRestorer() + self.agent = SingleNodeTimeSlicer(self.restorer) + server = await start_time_slicer(self.agent, self.socket) + self.addAsyncCleanup(server.wait_closed) + self.addCleanup(server.close) + + def launch(self, name: str, *options: str, command: tuple[str, ...] = ()) -> tuple[subprocess.Popen, Path, Path]: + """The launcher for replica `name` on CLAIM, running the stub with `options` unless given a command, with the stub's log and the marker.""" + log, awake, port = self.dir / f"{name}.log", self.dir / f"{name}.awake", free_port() + env = { + **os.environ, + "PATH": f"{self.dir}:{Path(sys.executable).parent}:{os.environ['PATH']}", + "PYTHONPATH": str(REPO / "src"), + "OPEN_RL_SAMPLER_AWAKE_FILE": str(awake), + "OPEN_RL_ACCEL_TIMESLICER_SOCKET": self.socket, + "OPEN_RL_TIME_SLICE_JOB_ID": name, + "OPEN_RL_TIME_SLICE_GROUP": CLAIM, + } + env.pop("OPEN_RL_ACCEL_TIMESLICER_HOST", None) + command = command or ("vllm", "serve", "stub-model", "--port", str(port), "--log", str(log), *options) + launcher = subprocess.Popen([sys.executable, "-m", "server.sampler_launcher", *command], env=env, start_new_session=True) + # The stub exits on its own once the launcher is gone. + self.addCleanup(subprocess.run, ["pkill", "-KILL", "-g", str(launcher.pid)]) + return launcher, log, awake + + def requests(self, log: Path) -> list[str]: + return [line for line in log.read_text().splitlines() if not line.startswith("GET /health")] + + async def test_the_replica_boots_asleep_inside_its_own_turn_then_a_dispatcher_leases_it(self) -> None: + launcher, log, awake = self.launch("sampler-x") + await until(awake.exists) + state = self.agent.workloads["sampler-x"] + self.assertEqual((state.workload.claim, state.failed, state.checkpointed, self.agent.running), (CLAIM, False, True, {})) + self.assertEqual((self.restorer.simple_labels(), self.requests(log)), ([("checkpoint", "sampler-x")], ["POST /sleep?level=1"])) + self.assertIsNone(launcher.poll()) + # A dispatcher leases the replica over its own connection while the launcher keeps the registration. + dispatcher = SocketTimeSlicerClient(socket_path=self.socket) + self.addAsyncCleanup(dispatcher.close) + async with dispatcher.acquire(WorkloadRef("sampler-x", CLAIM)): + self.assertEqual(self.restorer.simple_labels()[-1], ("restore", "sampler-x")) + self.assertEqual(self.restorer.simple_labels()[-1], ("checkpoint", "sampler-x")) + # SIGTERM reaches the stub, whose status comes back, and the registration goes with the launcher. + sent = time.monotonic() + launcher.send_signal(signal.SIGTERM) + self.assertEqual(await asyncio.to_thread(launcher.wait, 20), 3) + self.assertLess(time.monotonic() - sent, 3) + self.assertEqual(self.requests(log)[-1], "SIGTERM") + await until(lambda: self.agent.workloads["sampler-x"].failed) + + async def test_two_replicas_on_one_claim_boot_one_after_the_other(self) -> None: + self.restorer.block_checkpoint = True + replicas = {name: self.launch(name) for name in ("sampler-a", "sampler-b")} + # The first to boot is checkpointed with the grant held, so the other's vLLM has not been started yet. + await asyncio.to_thread(self.restorer.checkpoint_started.wait, 20) + (holder,) = self.agent.running.values() + (other,) = set(replicas) - {holder} + self.assertFalse(replicas[other][1].exists()) + self.restorer.finish_checkpoint.set() + await until(lambda: all(awake.exists() for _, _, awake in replicas.values())) + self.assertEqual(self.restorer.simple_labels(), [("checkpoint", holder), ("checkpoint", other)]) + + async def test_a_child_that_exits_before_it_is_healthy_gives_the_launcher_its_status_and_frees_the_claim(self) -> None: + launcher, _, awake = self.launch("sampler-x", command=(sys.executable, "-c", "import sys; sys.exit(7)")) + self.assertEqual(await asyncio.to_thread(launcher.wait, 20), 7) + # The launcher's exit closes the connection holding the turn, so the claim is free again. + await until(lambda: self.agent.running == {} and self.agent.workloads["sampler-x"].failed) + self.assertEqual((self.restorer.simple_labels(), awake.exists()), ([], False)) + + async def test_a_vllm_that_does_not_sleep_is_stopped(self) -> None: + launcher, log, awake = self.launch("sampler-x", "--fail-on", "/sleep?level=1") + self.assertEqual(await asyncio.to_thread(launcher.wait, 20), 1) + self.assertEqual((self.requests(log), awake.exists()), (["POST /sleep?level=1", "SIGTERM"], False)) + await until(lambda: self.agent.running == {}) + + async def test_a_replica_the_slicer_cannot_park_is_stopped(self) -> None: + self.agent.restorer = FailingCheckpointRestorer() + launcher, log, awake = self.launch("sampler-x") + self.assertEqual(await asyncio.to_thread(launcher.wait, 20), 1) + self.assertEqual((self.requests(log), awake.exists()), (["POST /sleep?level=1", "SIGTERM"], False)) + await until(lambda: self.agent.running == {}) + + async def test_a_turn_the_slicer_cannot_park_stops_vllm_and_the_launcher_so_the_trainer_goes_next(self) -> None: + launcher, log, awake = self.launch("sampler-x") + await until(awake.exists) + trainer = SocketTimeSlicerClient(socket_path=self.socket) + self.addAsyncCleanup(trainer.close) + await trainer.register(WorkloadRef("trainer", CLAIM)) + # A dispatcher's turn ends in a park that fails, so vLLM still holds the GPU and the trainer waits. + self.agent.restorer = FailingCheckpointRestorer() + dispatcher = SocketTimeSlicerClient(socket_path=self.socket) + self.addAsyncCleanup(dispatcher.close) + async with dispatcher.acquire(WorkloadRef("sampler-x", CLAIM)): + acquire = asyncio.create_task(trainer.request({"command": "ACQUIRE", "name": "trainer", "claim": CLAIM})) + await until(lambda: "trainer" in self.agent.waiting_workloads) + self.assertIn("must exit", dispatcher.faulted or "") + # The launcher stops vLLM and exits 1, so the container restarts, and its exit hands the GPU to the trainer. + self.assertEqual(await asyncio.to_thread(launcher.wait, 20), 1) + self.assertEqual(self.requests(log), ["POST /sleep?level=1", "SIGTERM"]) + self.assertTrue((await asyncio.wait_for(acquire, 5))["ok"]) + self.assertEqual((self.agent.running, self.agent.workloads["sampler-x"].failed), ({CLAIM: "trainer"}, True)) + # The restarted launcher registers again and boots vLLM in a turn of its own once the trainer releases. + self.agent.restorer = self.restorer + restarted, _, _ = self.launch("sampler-x") + await until(lambda: "sampler-x" in self.agent.waiting_workloads) + await trainer.request({"command": "RELEASE", "name": "trainer", "claim": CLAIM}) + await until(lambda: self.restorer.simple_labels()[-2:] == [("checkpoint", "trainer"), ("checkpoint", "sampler-x")]) + await until(awake.exists) + self.assertEqual((restarted.poll(), self.agent.workloads["sampler-x"].failed, self.agent.running), (None, False, {})) + + async def test_sigterm_while_the_acquire_waits_behind_the_trainer_ends_the_launcher_at_once(self) -> None: + trainer = SocketTimeSlicerClient(socket_path=self.socket) + self.addAsyncCleanup(trainer.close) + await trainer.register(WorkloadRef("trainer", CLAIM)) + async with trainer.acquire(WorkloadRef("trainer", CLAIM)): + launcher, log, awake = self.launch("sampler-x") + await until(lambda: "sampler-x" in self.agent.waiting_workloads) + sent = time.monotonic() + launcher.send_signal(signal.SIGTERM) + self.assertEqual(await asyncio.to_thread(launcher.wait, 20), 1) + self.assertLess(time.monotonic() - sent, 3) + self.assertEqual((log.exists(), awake.exists()), (False, False)) # vLLM was never started + # The slicer resolves a queued ACQUIRE at grant time, so once the trainer releases, the grant + # the dead launcher would have had is dropped with its connection and nothing holds the claim. + await until(lambda: self.agent.running == {} and "sampler-x" not in self.agent.waiting_workloads) + + async def test_a_sigterm_vllm_ignores_while_booting_is_followed_by_sigkill_and_frees_the_turn(self) -> None: + # vLLM ignores SIGTERM while its engine starts and would hold the GPU, and the turn, until kubelet's SIGKILL. + launcher, log, awake = self.launch("sampler-x", "--boot-seconds", "60") + await until(log.exists) # the stub is up but not serving + self.assertEqual(self.agent.running, {CLAIM: "sampler-x"}) + sent = time.monotonic() + launcher.send_signal(signal.SIGTERM) + self.assertEqual(await asyncio.to_thread(launcher.wait, 20), 137) + self.assertLess(time.monotonic() - sent, 10) + self.assertEqual((log.read_text().splitlines(), awake.exists()), (["SIGTERM ignored"], False)) + await until(lambda: self.agent.running == {} and self.agent.workloads["sampler-x"].failed) + + +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..2f8da4ce --- /dev/null +++ b/tests/test_sampling_dispatcher.py @@ -0,0 +1,594 @@ +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 DispatcherTestCase(RedisTestCase): + """The pool, its fake replicas and the dispatcher under test. The tests on leased replicas build on it too.""" + + 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)] + + +class SamplingDispatcherTest(DispatcherTestCase): + 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..963aa13e 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: @@ -91,11 +116,16 @@ def test_an_exclusive_fft_worker_is_placed_alone_and_never_time_sliced(self) -> with patch("server.worker_manager.get_state_store", return_value=s): self.manager.ensure("Model_A.1", "trainer") self.manager.ensure("Model_A.1", "sampler") - for worker in self.api.created: self.assertTrue(worker["spec"]["exclusive"]) - env = {e["name"]: e.get("value") for e in worker["spec"]["template"]["spec"]["containers"][0]["env"]} + container = worker["spec"]["template"]["spec"]["containers"][0] + env = {e["name"]: e.get("value") for e in container["env"]} self.assertEqual(env["OPEN_RL_TIME_SLICING"], "off") + self.assertNotIn("OPEN_RL_TIME_SLICED_SAMPLER", env) + # The sampler has a GPU of its own, so it is plain vllm serve with nothing to boot inside a turn. + sampler = self.api.created[1]["spec"]["template"]["spec"]["containers"][0] + self.assertEqual((sampler["command"], "startupProbe" in sampler), (["uv", "run", "vllm", "serve"], False)) + self.assertNotIn("--enable-sleep-mode", sampler["args"]) def test_fft_worker_is_its_own_owner(self) -> None: s = self.store_with("Model_A.1", {"base_model": "Qwen/Qwen3-8B", "fine_tuning_type": "full"}) @@ -222,7 +252,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 +359,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 +404,82 @@ 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"]) + # An FFT replica shares its GPU with its trainer: the launcher boots vLLM, able to sleep, inside a time-slicer + # turn, the env marks it for the pool's dispatcher to lease, and the container starts once it is up. + self.assertFalse(sampler["spec"]["exclusive"]) + self.assertEqual(container["command"], ["uv", "run", "python", "-u", "-m", "server.sampler_launcher", "vllm", "serve"]) + self.assertEqual((container["args"][0], container["args"][-1]), ("Qwen/Qwen3-8B", "--enable-sleep-mode")) + self.assertEqual(container["startupProbe"]["exec"]["command"], ["test", "-f", "/tmp/open-rl-sampler-awake"]) + self.assertEqual(sampler["spec"]["template"]["metadata"]["labels"], labels) + env = {e["name"]: e.get("value") for e in container["env"]} + self.assertEqual((env["OPEN_RL_TIME_SLICED_SAMPLER"], env["OPEN_RL_WORKLOAD_ID"]), ("true", "fft-fft-job-sampler")) + self.assertEqual((env["OPEN_RL_ACCEL_TIMESLICER_PORT"], env.get("OPEN_RL_TIME_SLICING")), ("9753", None)) + 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 +487,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 +504,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..a8622f94 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,20 +208,33 @@ 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: - 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, - ): - 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") + def launch_local_sampler(self, **env: str) -> tuple[list, dict, list, dict]: + base = {"REDIS_URL": "redis://localhost:6379", "SAMPLING_BACKEND": "vllm", "OPEN_RL_ENABLE_FFT": "true"} + with patch.dict("os.environ", {**base, **env}, clear=True), patch("server.worker_manager.subprocess.Popen") as popen: + LocalWorkerManager().ensure("Model_A.1", "sampler") + (serve, serve_kwargs), (dispatch, dispatch_kwargs) = [(c[0][0], c[1]) for c in popen.call_args_list] + return serve, serve_kwargs, dispatch, dispatch_kwargs + + async def test_a_local_pool_sampler_is_vllm_serve_and_its_dispatcher_each_in_its_own_process_group(self) -> None: + serve, serve_kwargs, dispatch, dispatch_kwargs = self.launch_local_sampler(OPEN_RL_TIME_SLICING="off") + 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("server.sampler_launcher", serve) + self.assertNotIn("--enable-sleep-mode", serve) + self.assertNotIn("OPEN_RL_TIME_SLICE_JOB_ID", serve_kwargs["env"]) + self.assertNotIn("OPEN_RL_TIME_SLICED_SAMPLER", dispatch_kwargs["env"]) + self.assertEqual(dispatch[-1], "server.sampling_dispatcher") + + async def test_a_local_fft_sampler_boots_under_a_lease_in_the_trainers_claim_and_its_dispatcher_carries_its_tags(self) -> None: + serve, serve_kwargs, dispatch, dispatch_kwargs = self.launch_local_sampler() + self.assertEqual(serve[serve.index("-m") + 1 : serve.index("serve")], ["server.sampler_launcher", serve[serve.index("serve") - 1]]) + self.assertTrue(serve[serve.index("serve") - 1].endswith("vllm")) + self.assertEqual((serve[serve.index("serve") + 1], serve[-1]), ("Model_A.1", "--enable-sleep-mode")) + tags = {"OPEN_RL_TIME_SLICED_SAMPLER": "true", "OPEN_RL_TIME_SLICE_JOB_ID": "sampler-Model_A.1", "OPEN_RL_TIME_SLICE_GROUP": "trainers"} + self.assertEqual({k: serve_kwargs["env"][k] for k in tags}, tags) + self.assertEqual({k: dispatch_kwargs["env"][k] for k in tags}, tags) + port = serve[serve.index("--port") + 1] + self.assertEqual(dispatch_kwargs["env"]["OPEN_RL_SAMPLER_ENDPOINTS"], f"http://127.0.0.1:{port}") async def test_launch_fetches_metadata_from_store(self) -> None: import json @@ -242,7 +262,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 +374,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 +392,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 +428,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__": diff --git a/tests/vllm_stub.py b/tests/vllm_stub.py new file mode 100755 index 00000000..5b017a86 --- /dev/null +++ b/tests/vllm_stub.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +"""A stand-in for `vllm serve` with the endpoints the sampler launcher calls. + + vllm serve --port --log [--boot-seconds S] [--fail-on PATH] + +The launcher tests put this file on PATH as `vllm`. Every request lands in the log as one +line, `GET /health` or `POST /sleep?level=1`. Like vLLM, it ignores SIGTERM and SIGINT while +it boots, for `--boot-seconds`, logging the signal as ignored. Once it serves, a signal is +logged and exits 3, so a test can tell a forwarded signal from a crash. `--fail-on` answers +that path with a 500. It exits when its parent is gone, as vLLM's engine does. +""" + +import argparse +import os +import signal +import sys +import threading +import time +from http.server import BaseHTTPRequestHandler, HTTPServer + + +def watch_parent(parent: int) -> None: + while os.getppid() == parent: + time.sleep(0.2) + os._exit(4) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("serve") + parser.add_argument("model") + parser.add_argument("--port", type=int, required=True) + parser.add_argument("--log", required=True) + parser.add_argument("--boot-seconds", type=float, default=0.3) + parser.add_argument("--fail-on") + options, _ = parser.parse_known_args() + log = open(options.log, "a", buffering=1) + serving = False + + class Handler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + log.write(f"GET {self.path}\n") + self.reply() + + def do_POST(self) -> None: + body = self.rfile.read(int(self.headers.get("Content-Length") or 0)).decode() + log.write(f"POST {self.path} {body}".rstrip() + "\n") + self.reply(500 if self.path == options.fail_on else 200) + + def reply(self, status: int = 200) -> None: + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(b"{}") + + def log_message(self, *args) -> None: + pass + + def on_signal(signum, frame) -> None: + name = signal.Signals(signum).name + if not serving: + log.write(f"{name} ignored\n") + return + log.write(f"{name}\n") + sys.exit(3) + + signal.signal(signal.SIGTERM, on_signal) + signal.signal(signal.SIGINT, on_signal) + threading.Thread(target=watch_parent, args=(os.getppid(),), daemon=True).start() + time.sleep(options.boot_seconds) + server = HTTPServer(("127.0.0.1", options.port), Handler) + serving = True + server.serve_forever() + + +if __name__ == "__main__": + main()