From 72bd42f79f676411d32c0b1efe092f5e05fc29b1 Mon Sep 17 00:00:00 2001 From: Shuby Mishra Date: Wed, 21 Oct 2026 08:02:04 +0000 Subject: [PATCH] sampling: groundwork for sampling pools Nothing uses these changes until the switch to sampling pools. RequestStore.set_future and get_future become set_result and get_result, and SamplingStore.commit becomes set_result, so both stores name the same operation the same way. The SDK endpoint retrieve_future and the Redis key open_rl:future: keep their names. engine_args returns the vLLM engine settings a sampler uses, and serve_args renders them as vllm serve flags, so both launch styles get the same settings from the same env. The deploy config adds the sampler-router ConfigMap with Envoy and llm-d EPP config for each pool's CPU pod. RBAC lets the API server create a pool Deployment and NetworkPolicy, and lets the dispatcher list its replicas. The Workload CRD keeps pod template labels, which the EPP selects replicas by. --- k8s/deploy/base/api-server.yaml | 6 ++ k8s/deploy/base/images.yaml | 2 + k8s/deploy/base/kustomization.yaml | 2 + k8s/deploy/base/rbac.yaml | 39 ++++++++- k8s/deploy/base/sampler-router.yaml | 87 +++++++++++++++++++++ k8s/deploy/kind-dra/kustomization.yaml | 3 + scheduler/controller/Makefile | 4 +- scheduler/deploy/base/00-workload-crd.yaml | 34 ++++++++ src/server/api_server.py | 4 +- src/server/sampling_store.py | 8 +- src/server/store.py | 12 +-- src/server/training_requests_processor.py | 2 +- src/server/vllm_options.py | 39 ++++++++- src/server/vllm_sampler.py | 6 +- tests/test_api_server_paths.py | 4 +- tests/test_delete_model.py | 2 +- tests/test_exclusive_trainer.py | 2 +- tests/test_fft_batch_failure.py | 2 +- tests/test_redis_store.py | 24 +++--- tests/test_sampler_patch.py | 6 +- tests/test_sampling_store.py | 14 ++-- tests/test_trainer_optimizer_correctness.py | 6 +- tests/test_worker_manager.py | 2 +- 23 files changed, 259 insertions(+), 51 deletions(-) create mode 100644 k8s/deploy/base/sampler-router.yaml 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/src/server/api_server.py b/src/server/api_server.py index 85e7c17b..095e19ce 100644 --- a/src/server/api_server.py +++ b/src/server/api_server.py @@ -575,7 +575,7 @@ async def launch_worker_and_enqueue(command: Command) -> str: await asyncio.to_thread(worker_manager.ensure, command.model_id, "trainer") except Exception as exc: traceback.print_exc() - await store.set_future(request_id, {"type": "RequestFailedResponse", "error_message": str(exc)}) + await store.set_result(request_id, {"type": "RequestFailedResponse", "error_message": str(exc)}) return request_id return await enqueue(command) @@ -985,7 +985,7 @@ async def retrieve_future(req: RetrieveFutureRequest, accept: str = Header(defau pending, failed, and every other result stay JSON. """ request_id = req.request_id - result = await store.get_future(request_id, timeout=60.0) + result = await store.get_result(request_id, timeout=60.0) if result is None: return JSONResponse(status_code=400, content={"type": "RequestFailedResponse", "error_message": "Future not found"}) if isinstance(result, dict) and result.get("type") == "RequestFailedResponse": diff --git a/src/server/sampling_store.py b/src/server/sampling_store.py index d5c85815..29af6c33 100644 --- a/src/server/sampling_store.py +++ b/src/server/sampling_store.py @@ -7,7 +7,7 @@ import redis.asyncio as redis RESULT_TTL_SECONDS = 300 -# An accepted request may never be committed, so its idempotency key expires instead of staying until Redis is full. +# An accepted request may never get a result, so its idempotency key expires instead of staying until Redis is full. PENDING_TTL_SECONDS = 3 * 24 * 3600 @@ -49,8 +49,8 @@ async def pending(self, pool: str, limit: int | None = 1000) -> list[dict[str, A 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 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.""" + async def set_result(self, record: dict[str, Any], result: dict[str, Any]) -> bool: + """Save the first final result and drop the pending entry. False if it already has one.""" stream, future, entry_id = stream_key(record["pool"]), future_key(record["request_id"]), record["entry_id"] async def save(pipe) -> bool: @@ -64,7 +64,7 @@ async def save(pipe) -> bool: pipe.expire(idempotency_key(record["pool"], key), RESULT_TTL_SECONDS) return True - # Every commit writes the result key, so watching it lets only one through. + # Every set_result writes the result key, so watching it lets only one through. return await self.redis.transaction(save, future, value_from_callable=True) async def pinned_revisions(self, pool: str) -> set[str]: diff --git a/src/server/store.py b/src/server/store.py index 90896400..59c3d411 100644 --- a/src/server/store.py +++ b/src/server/store.py @@ -42,12 +42,12 @@ async def get_sampling_requests_for_model(self, model_id: str) -> list[dict[str, pass @abstractmethod - async def set_future(self, req_id: str, result: dict[str, Any]) -> None: + async def set_result(self, req_id: str, result: dict[str, Any]) -> None: """Resolve a future by its request ID.""" pass @abstractmethod - async def get_future(self, req_id: str, timeout: float) -> dict[str, Any] | None: + async def get_result(self, req_id: str, timeout: float) -> dict[str, Any] | None: """Block until the future resolves or the timeout is reached.""" pass @@ -129,14 +129,14 @@ async def get_sampling_requests_for_model(self, model_id: str) -> list[dict[str, batch.append(queue.get_nowait()) return batch - async def set_future(self, req_id: str, result: dict[str, Any]) -> None: + async def set_result(self, req_id: str, result: dict[str, Any]) -> None: if result.get("status") == "pending": return async with self.futures_cv: self.futures_store[req_id] = result self.futures_cv.notify_all() - async def get_future(self, req_id: str, timeout: float) -> dict[str, Any] | None: + async def get_result(self, req_id: str, timeout: float) -> dict[str, Any] | None: async with self.futures_cv: try: await asyncio.wait_for(self.futures_cv.wait_for(lambda: req_id in self.futures_store), timeout=timeout) @@ -259,7 +259,7 @@ async def get_sampling_requests_for_model(self, model_id: str) -> list[dict[str, return batch - async def set_future(self, req_id: str, result: dict[str, Any]) -> None: + async def set_result(self, req_id: str, result: dict[str, Any]) -> None: if result.get("status") == "pending": return key = f"open_rl:future:{req_id}" @@ -270,7 +270,7 @@ async def set_future(self, req_id: str, result: dict[str, Any]) -> None: pipeline.expire(key, 300) await pipeline.execute() - async def get_future(self, req_id: str, timeout: float) -> dict[str, Any] | None: + async def get_result(self, req_id: str, timeout: float) -> dict[str, Any] | None: key = f"open_rl:future:{req_id}" deadline = time.monotonic() + timeout while True: diff --git a/src/server/training_requests_processor.py b/src/server/training_requests_processor.py index 70fb1cf5..2ff259b0 100644 --- a/src/server/training_requests_processor.py +++ b/src/server/training_requests_processor.py @@ -201,7 +201,7 @@ async def fetch_requests(self) -> list[dict[str, Any]]: async def publish_result(self, request_id: str | None, result: dict[str, Any]) -> None: # Only rank 0 writes to the store; the other ranks computed the same result. if request_id is not None and is_primary(): - await self.store.set_future(request_id, result) + await self.store.set_result(request_id, result) async def run_once(self) -> None: batch = await self.next_batch() diff --git a/src/server/vllm_options.py b/src/server/vllm_options.py index a62a6dbb..5bb35689 100644 --- a/src/server/vllm_options.py +++ b/src/server/vllm_options.py @@ -1,7 +1,9 @@ """Engine and sampling options shared by the two vLLM sampler workers.""" +import json import os -from collections.abc import Sequence +from collections.abc import Mapping, Sequence +from typing import Any def text_only_engine_kwargs() -> dict: @@ -83,3 +85,38 @@ def detected_device_memory_bytes() -> int: except Exception: pass return 0 + + +def engine_args(base_model: str, is_lora: bool, env: Mapping[str, str | None], sleep_mode: bool = False) -> dict[str, Any]: + """vLLM engine settings for a sampler, the same for a `vllm serve` replica and an engine in the sampler's process.""" + args: dict[str, Any] = {"enable_prefix_caching": True} + # A checkpoint's generation_config.json could add sampling params the client never asked for, such as a repetition penalty. + args["generation_config"] = "vllm" + args["max_model_len"] = int(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"] = gpu_memory_utilization() + # These are set 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()] = int(os.environ[name]) + if os.getenv("VLLM_ENFORCE_EAGER", "0") == "1": + args["enforce_eager"] = True + args.update(text_only_engine_kwargs()) + if architecture := env.get("VLLM_ARCHITECTURE_OVERRIDE"): + args["hf_overrides"] = {"architectures": [architecture]} + if is_lora: + args.update(enable_lora=True, max_loras=int(os.getenv("VLLM_MAX_LORAS", "8")), max_lora_rank=int(os.getenv("VLLM_MAX_LORA_RANK", "64"))) + else: + args["weight_transfer_config"] = {"backend": "delta_snapshot"} + if sleep_mode: + args["enable_sleep_mode"] = True + return args + + +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"] + for name, value in engine_args(base_model, is_lora, env).items(): + flag = "--" + name.replace("_", "-") + args += [flag] if value is True else [flag, json.dumps(value) if isinstance(value, dict) else str(value)] + return args diff --git a/src/server/vllm_sampler.py b/src/server/vllm_sampler.py index bc0dba96..ef627cbb 100644 --- a/src/server/vllm_sampler.py +++ b/src/server/vllm_sampler.py @@ -233,7 +233,7 @@ async def process_batch(sampler: Sampler, store: RequestStore, requests: list[di 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}")) + await store.set_result(request["request_id"], failed_response(f"vLLM weight update failed: {exc}")) continue await serve_group(sampler, store, batch) @@ -259,7 +259,7 @@ async def fail_requests(store: RequestStore, requests: list[dict[str, Any]], err """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}")) + await store.set_result(request["request_id"], failed_response(f"vLLM Worker Error: {error}")) except Exception as exc: print(f"[vLLM Worker] Could not fail request {request['request_id']}: {exc}") @@ -271,7 +271,7 @@ async def process_request(sampler: Sampler, store: RequestStore, request: dict[s result["type"] = "sample" except Exception as exc: result = failed_response(f"vLLM Worker Error: {exc}") - await store.set_future(request["request_id"], result) + await store.set_result(request["request_id"], result) async def mark_ready(state: StateStore, model_id: str) -> None: diff --git a/tests/test_api_server_paths.py b/tests/test_api_server_paths.py index 85268b3a..51a7d9b5 100644 --- a/tests/test_api_server_paths.py +++ b/tests/test_api_server_paths.py @@ -257,11 +257,11 @@ def test_retrieve_future_answers_protobuf_only_when_asked_and_only_for_proto_typ from server.proto import tinker_public_pb2 as pb asyncio.run( - api_server.store.set_future( + api_server.store.set_result( "samp-1", {"type": "sample_completed", "sequences": [{"tokens": [1, 2], "logprobs": [-0.5, -1.0], "stop_reason": "stop"}]} ) ) - asyncio.run(api_server.store.set_future("optim-1", {"type": "optim_step_completed", "metrics": {"grad_norm:mean": 0.0}})) + asyncio.run(api_server.store.set_result("optim-1", {"type": "optim_step_completed", "metrics": {"grad_norm:mean": 0.0}})) as_json = self.client.post("/api/v1/retrieve_future", json={"request_id": "samp-1"}) self.assertEqual(as_json.status_code, 200) diff --git a/tests/test_delete_model.py b/tests/test_delete_model.py index 8f25e43a..f177f119 100644 --- a/tests/test_delete_model.py +++ b/tests/test_delete_model.py @@ -38,7 +38,7 @@ async def run_commands(self, processor: trp.TrainingRequestsProcessor, *steps: c results = [] for step in steps: await processor.process_request(commands.wire(step)) - result = await processor.store.get_future(step.request_id, timeout=5) + result = await processor.store.get_result(step.request_id, timeout=5) self.assertNotEqual(result.get("type"), "RequestFailedResponse", f"{step.request_id}: {result}") results.append(result) return results diff --git a/tests/test_exclusive_trainer.py b/tests/test_exclusive_trainer.py index 7eea31c6..836d83ec 100644 --- a/tests/test_exclusive_trainer.py +++ b/tests/test_exclusive_trainer.py @@ -71,7 +71,7 @@ async def test_a_full_step_runs_without_a_time_slicer(self) -> None: before = worker.model.lm_head.weight.detach().clone() processor = asyncio.create_task(trp.run_training_requests_processor(worker, "run-a", store=store)) try: - results = {step.request_id: await store.get_future(step.request_id, timeout=60) for step in steps} + results = {step.request_id: await store.get_result(step.request_id, timeout=60) for step in steps} finally: processor.cancel() diff --git a/tests/test_fft_batch_failure.py b/tests/test_fft_batch_failure.py index f91a0e31..96239316 100644 --- a/tests/test_fft_batch_failure.py +++ b/tests/test_fft_batch_failure.py @@ -45,7 +45,7 @@ async def get_requests_for_model(self, model_id): batch, self.batch = self.batch, [] return batch - async def set_future(self, request_id, result): + async def set_result(self, request_id, result): self.futures[request_id] = result diff --git a/tests/test_redis_store.py b/tests/test_redis_store.py index ee7ec943..92e905b1 100644 --- a/tests/test_redis_store.py +++ b/tests/test_redis_store.py @@ -70,18 +70,18 @@ async def asyncTearDown(self) -> None: await self.store.redis.aclose() self.state.sync_redis.close() - async def test_get_future_returns_already_resolved_result(self) -> None: - await self.store.set_future("req-1", {"type": "sample", "ok": True}) - self.assertEqual(await self.store.get_future("req-1", timeout=1.0), {"type": "sample", "ok": True}) + async def test_get_result_returns_already_resolved_result(self) -> None: + await self.store.set_result("req-1", {"type": "sample", "ok": True}) + self.assertEqual(await self.store.get_result("req-1", timeout=1.0), {"type": "sample", "ok": True}) - async def test_get_future_wakes_on_resolution(self) -> None: + async def test_get_result_wakes_on_resolution(self) -> None: async def resolve_later() -> None: await asyncio.sleep(0.2) - await self.store.set_future("req-1", {"type": "sample"}) + await self.store.set_result("req-1", {"type": "sample"}) resolver = asyncio.create_task(resolve_later()) started = time.monotonic() - result = await self.store.get_future("req-1", timeout=10.0) + result = await self.store.get_result("req-1", timeout=10.0) await resolver self.assertEqual(result, {"type": "sample"}) @@ -89,21 +89,21 @@ async def resolve_later() -> None: self.assertLess(time.monotonic() - started, 5.0) async def test_result_survives_repeated_and_concurrent_reads(self) -> None: - waiters = [asyncio.create_task(self.store.get_future("req-1", timeout=10.0)) for _ in range(5)] + waiters = [asyncio.create_task(self.store.get_result("req-1", timeout=10.0)) for _ in range(5)] await asyncio.sleep(0.2) - await self.store.set_future("req-1", {"type": "sample"}) + await self.store.set_result("req-1", {"type": "sample"}) for result in await asyncio.gather(*waiters): self.assertEqual(result, {"type": "sample"}) - self.assertEqual(await self.store.get_future("req-1", timeout=1.0), {"type": "sample"}) + self.assertEqual(await self.store.get_result("req-1", timeout=1.0), {"type": "sample"}) async def test_unresolved_future_times_out_with_try_again(self) -> None: - result = await self.store.get_future("req-never", timeout=0.3) + result = await self.store.get_result("req-never", timeout=0.3) self.assertEqual(result["type"], "try_again") async def test_pending_markers_are_not_stored(self) -> None: - await self.store.set_future("req-1", {"status": "pending"}) - self.assertEqual((await self.store.get_future("req-1", timeout=0.3))["type"], "try_again") + await self.store.set_result("req-1", {"status": "pending"}) + self.assertEqual((await self.store.get_result("req-1", timeout=0.3))["type"], "try_again") async def test_get_requests_rotates_between_tenants(self) -> None: for i in range(3): diff --git a/tests/test_sampler_patch.py b/tests/test_sampler_patch.py index e8d22cf9..a27bc16f 100644 --- a/tests/test_sampler_patch.py +++ b/tests/test_sampler_patch.py @@ -36,7 +36,7 @@ async def asyncSetUp(self): 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} + return {call.args[0]: call.args[1] for call in self.store.set_result.call_args_list} async def test_groups_drain_before_switching_weights(self): active = set() @@ -163,7 +163,7 @@ async def test_initialization_and_batches_own_slots_and_cleanup_once(self): 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") + self.assertEqual(self.store.set_result.call_args.args[0], "1") async def test_registration_failure_never_constructs_engine(self): self.slicer.register.side_effect = RuntimeError("registration failed") @@ -208,7 +208,7 @@ async def test_unshared_sampler_does_not_sleep_or_wake_engine(self): 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} + return {call.args[0]: call.args[1] for call in self.store.set_result.call_args_list} async def test_dead_engine_fails_hanging_requests_and_exits(self): async def hang(request_id): diff --git a/tests/test_sampling_store.py b/tests/test_sampling_store.py index ba746dc0..21ccfd16 100644 --- a/tests/test_sampling_store.py +++ b/tests/test_sampling_store.py @@ -17,20 +17,20 @@ def record(self, revision: str = "", key: str | None = None) -> dict: async def pending(self) -> dict[str, dict]: return {r["request_id"]: r for r in await self.store.pending(self.pool)} - async def test_the_same_key_returns_the_first_id_before_and_after_commit(self) -> None: + async def test_the_same_key_returns_the_first_id_before_and_after_its_result(self) -> None: first = self.record(key="session:7") ids = [await self.store.accept(first), await self.store.accept(self.record(key="session:7"))] pending = await self.pending() self.assertEqual((ids, list(pending)), ([first["request_id"]] * 2, [first["request_id"]])) - await self.store.commit(pending[first["request_id"]], {"type": "sample"}) + await self.store.set_result(pending[first["request_id"]], {"type": "sample"}) self.assertEqual(await self.store.accept(self.record(key="session:7")), first["request_id"]) - async def test_a_key_expires_days_after_accept_and_minutes_after_commit(self) -> None: + async def test_a_key_expires_days_after_accept_and_minutes_after_its_result(self) -> None: await self.store.accept(self.record(key="session:7")) key = idempotency_key(self.pool, "session:7") self.assertAlmostEqual(await self.redis.ttl(key), PENDING_TTL_SECONDS, delta=60) [pending] = (await self.pending()).values() - await self.store.commit(pending, {"type": "sample"}) + await self.store.set_result(pending, {"type": "sample"}) self.assertAlmostEqual(await self.redis.ttl(key), RESULT_TTL_SECONDS, delta=60) async def test_concurrent_commits_save_exactly_one_result_and_clear_pending(self) -> None: @@ -39,7 +39,7 @@ async def test_concurrent_commits_save_exactly_one_result_and_clear_pending(self [pending] = (await self.pending()).values() # Open the connections first, so the commits overlap instead of queueing behind connection setup. await asyncio.gather(*(self.redis.ping() for _ in range(5))) - won = await asyncio.gather(*(self.store.commit(pending, {"type": "sample", "attempt": i}) for i in range(5))) + won = await asyncio.gather(*(self.store.set_result(pending, {"type": "sample", "attempt": i}) for i in range(5))) saved = [json.loads(raw)["attempt"] for raw in await self.redis.lrange(future_key(r["request_id"]), 0, -1)] self.assertEqual((won.count(True), saved, await self.pending()), (1, [won.index(True)], {})) @@ -49,7 +49,7 @@ async def test_pending_keeps_accept_order_without_removal_and_a_revision_stays_p await self.store.accept(r) self.assertEqual([list(await self.pending()), list(await self.pending())], [[r["request_id"] for r in records]] * 2) r1, r1_again, _ = (await self.pending()).values() - await self.store.commit(r1, {"type": "sample"}) + await self.store.set_result(r1, {"type": "sample"}) self.assertEqual(await self.store.pinned_revisions(self.pool), {"/ckpt/r1", "/ckpt/r2"}) - await self.store.commit(r1_again, {"type": "sample"}) + await self.store.set_result(r1_again, {"type": "sample"}) self.assertEqual(await self.store.pinned_revisions(self.pool), {"/ckpt/r2"}) diff --git a/tests/test_trainer_optimizer_correctness.py b/tests/test_trainer_optimizer_correctness.py index f8ea1313..16b35684 100644 --- a/tests/test_trainer_optimizer_correctness.py +++ b/tests/test_trainer_optimizer_correctness.py @@ -175,9 +175,9 @@ def __init__(self, events=None): self.results = {} self.events = events - async def set_future(self, req_id, result): + async def set_result(self, req_id, result): if self.events is not None: - self.events.append(("set_future", req_id)) + self.events.append(("set_result", req_id)) self.results[req_id] = result @@ -602,7 +602,7 @@ async def test_full_processor_publishes_result_after_release(self) -> None: processor = training_requests_processor_module.TrainingRequestsProcessor(store, worker, "model-a", time_slicer=time_slicer) await processor.run_once() - self.assertEqual([event[0] for event in events], ["acquire", "release", "set_future"]) + self.assertEqual([event[0] for event in events], ["acquire", "release", "set_result"]) self.assertEqual(store.results["req-a"]["type"], "model_created") diff --git a/tests/test_worker_manager.py b/tests/test_worker_manager.py index 727e3658..fec69538 100644 --- a/tests/test_worker_manager.py +++ b/tests/test_worker_manager.py @@ -17,7 +17,7 @@ def __init__(self): async def put_request(self, req_data: dict, active_set_id: str | None = None) -> None: self.forwarded_requests.append(req_data) - async def set_future(self, req_id: str, result: dict) -> None: + async def set_result(self, req_id: str, result: dict) -> None: self.futures[req_id] = result async def set_value(self, key: str, value: str, ttl_seconds: float | None = None) -> None: