From c69e88d0e9c1fcd07905276033e4c63fac04e4be Mon Sep 17 00:00:00 2001 From: svonava Date: Mon, 19 Oct 2026 17:43:53 -0700 Subject: [PATCH 1/2] feat(server): keep evicted models in host RAM and wake them instead of reloading When the registry evicts a model to free GPU memory (memory pressure, room for another load, OOM recovery, idle eviction, group placement), it now parks the model's weights in host RAM, within a budget, and the next request wakes them back onto the same GPU instead of loading from the local disk cache, the cluster cache or Hugging Face. The least recently parked model is dropped when the budget fills; a miss loads the ordinary way. Manual unloads, config changes, shutdown and crashed engines unload as before, and pinned models are never evicted. One mechanism lives in the registry (core/host_cache.py bookkeeping, park in the unload path, wake ahead of the cold load path) with three small adapter hooks: host_offload_bytes, offload_to_host and restore_from_host. - In-process PyTorch adapters (BaseAdapter) park generically on CUDA: every tensor reachable from the adapter is copied to pageable host memory with parameter identity preserved, CUDA graph holders are cleared first, and an adapter declines when the walk cannot see all its weights. - SGLang generation and embedding engines on one GPU park through SGLang's memory saver (launched with --enable-memory-saver --enable-weights-cpu-backup while the tier is on): flush_cache waits for idle, release_memory_occupation backs the weights up to pinned host memory and frees weights, KV cache and graph memory; resume_memory_occupation restores them at the same addresses, so the engine process, kernels and graphs survive. The engine keeps its host copy after a wake, which stays charged to the budget; a parked engine that still holds device memory is dropped when its GPU needs room and nothing resident can be evicted. SIE_HOST_CACHE_BYTES configures the budget: auto (a quarter of the container memory limit on CUDA workers, off elsewhere), an explicit size, or 0 to turn the tier off. Requests for a parked model wait up to 15 s for the wake instead of answering MODEL_LOADING. New metrics sie.worker.model.host_cache.lookups, .wake.duration and .memory are declared in the telemetry contract, allowed through the collector and shown on the worker dashboard. --- deploy/helm/sie-cluster/README.md | 19 +- .../files/dashboards/worker-health.json | 50 ++ .../templates/_otel-collector-config.tpl | 18 +- packages/sie_server/README.md | 68 +- .../src/sie_server/adapters/_base_adapter.py | 102 ++- .../src/sie_server/adapters/_host_offload.py | 244 +++++++ .../src/sie_server/adapters/base.py | 66 ++ .../src/sie_server/adapters/fake/adapter.py | 26 + .../sie_server/adapters/sglang/_parking.py | 82 +++ .../src/sie_server/adapters/sglang/_server.py | 109 ++++ .../sie_server/adapters/sglang/embedding.py | 13 +- .../sie_server/adapters/sglang/generation.py | 13 +- .../sie_server/src/sie_server/api/helpers.py | 5 + .../src/sie_server/config/engine.py | 24 +- .../src/sie_server/core/host_cache.py | 302 +++++++++ .../src/sie_server/core/model_loader.py | 22 + .../src/sie_server/core/registry.py | 613 ++++++++++++++---- .../observability/worker_telemetry.py | 85 +++ .../src/sie_server/queue_executor.py | 5 + .../tests/adapters/test_host_offload.py | 230 +++++++ .../tests/adapters/test_sglang_parking.py | 233 +++++++ .../tests/api/test_encode_endpoint.py | 19 + packages/sie_server/tests/api/test_extract.py | 2 + .../tests/api/test_openai_compat.py | 2 + packages/sie_server/tests/api/test_score.py | 2 + .../sie_server/tests/core/test_host_cache.py | 154 +++++ .../tests/core/test_registry_host_cache.py | 297 +++++++++ .../observability/test_worker_telemetry.py | 38 ++ .../sie_server/tests/test_queue_executor.py | 21 + telemetry/README.md | 13 +- telemetry/contract.yaml | 47 ++ 31 files changed, 2779 insertions(+), 145 deletions(-) create mode 100644 packages/sie_server/src/sie_server/adapters/_host_offload.py create mode 100644 packages/sie_server/src/sie_server/adapters/sglang/_parking.py create mode 100644 packages/sie_server/src/sie_server/core/host_cache.py create mode 100644 packages/sie_server/tests/adapters/test_host_offload.py create mode 100644 packages/sie_server/tests/adapters/test_sglang_parking.py create mode 100644 packages/sie_server/tests/core/test_host_cache.py create mode 100644 packages/sie_server/tests/core/test_registry_host_cache.py diff --git a/deploy/helm/sie-cluster/README.md b/deploy/helm/sie-cluster/README.md index 364612ccf..3822f49e2 100644 --- a/deploy/helm/sie-cluster/README.md +++ b/deploy/helm/sie-cluster/README.md @@ -78,6 +78,18 @@ The current GKE/EKS scale-from-zero guidance is: | **Model loading** | 10-120s | Download weights (if not cached), load to GPU | | **Total cold start** | 3-7 min | First request to a scaled-to-zero pool | +Model loading also applies when a running worker reloads a model it evicted +to make room for another. Workers keep evicted models in host RAM and wake them +back onto the GPU instead, which is estimated to take well under a second for +a typical encoder and a few seconds for a single-GPU SGLang model. The figures +are estimates, not measurements. A model loads the ordinary way when it was +never loaded on that worker, or when it was dropped from host RAM to stay +within the budget. Under the default `SIE_HOST_CACHE_BYTES=auto` each GPU +worker container uses a quarter of its memory limit, 8 GiB on the default L4 +pool. Override it with `workers.common.extraEnv` or a bundle's `extraEnv`, or +set it to `0` to turn the host tier off. See the server README section "Host-RAM tier for evicted +models" for what parks and how it is charged. + These figures are not an ACK performance claim; this repository has no ACK cold-start baseline. The Alibaba Terraform module defaults its shared NAT EIP to 100 Mbit/s, and observed ACK cold start depends on current GPU inventory and @@ -86,9 +98,10 @@ quota, effective bandwidth, image and model size, and cache state. ### Reducing Cold Start Time 1. **Use cluster cache**: Pre-populate object storage with model weights (`--cluster-cache`) -2. **Set minReplicas=1**: Keep one warm replica per critical GPU type -3. **Use reserved capacity**: Avoid spot for latency-sensitive workloads -4. **Pre-warm models**: Call `/v1/encode/{model}` on startup to load weights +2. **Size the host-RAM tier**: Raise `SIE_HOST_CACHE_BYTES` (with the pool's memory limit) so models that take turns on a GPU wake from host RAM instead of reloading +3. **Set minReplicas=1**: Keep one warm replica per critical GPU type +4. **Use reserved capacity**: Avoid spot for latency-sensitive workloads +5. **Pre-warm models**: Call `/v1/encode/{model}` on startup to load weights ### Client Handling diff --git a/deploy/helm/sie-cluster/files/dashboards/worker-health.json b/deploy/helm/sie-cluster/files/dashboards/worker-health.json index ae79571ad..79e6873aa 100644 --- a/deploy/helm/sie-cluster/files/dashboards/worker-health.json +++ b/deploy/helm/sie-cluster/files/dashboards/worker-health.json @@ -468,6 +468,56 @@ "refId": "B" } ] + }, + { + "gridPos": { "h": 10, "w": 12, "x": 0, "y": 78 }, + "id": 76, + "title": "Host-RAM Tier Lookups and Wake Time", + "description": "Model loads that woke a copy parked in host RAM (hit) or loaded from disk, the cluster cache or the Hub (miss), and the p95 time to wake a parked model back onto its device.", + "type": "timeseries", + "fieldConfig": { "defaults": { "unit": "short" } }, + "options": { + "legend": { "calcs": ["mean", "max"], "displayMode": "table", "placement": "bottom", "showLegend": true }, + "tooltip": { "mode": "multi", "sort": "desc" } + }, + "pluginVersion": "10.0.0", + "targets": [ + { + "expr": "sum by (result, lane) (rate(sie_worker_model_host_cache_lookups_total{namespace=\"$namespace\", service=\"$collector\", endpoint=\"prometheus\", producer_service=\"sie-worker\"}[5m]))", + "legendFormat": "loads/sec {{result}} / {{lane}}", + "refId": "A" + }, + { + "expr": "histogram_quantile(0.95, sum by (le, lane) (rate(sie_worker_model_host_cache_wake_duration_seconds_bucket{namespace=\"$namespace\", service=\"$collector\", endpoint=\"prometheus\", producer_service=\"sie-worker\", outcome=\"success\"}[5m])))", + "legendFormat": "wake p95 seconds / {{lane}}", + "refId": "B" + }, + { + "expr": "sum by (outcome, lane) (rate(sie_worker_model_host_cache_wake_duration_seconds_count{namespace=\"$namespace\", service=\"$collector\", endpoint=\"prometheus\", producer_service=\"sie-worker\"}[5m]))", + "legendFormat": "wakes/sec {{outcome}} / {{lane}}", + "refId": "C" + } + ] + }, + { + "gridPos": { "h": 10, "w": 12, "x": 12, "y": 78 }, + "id": 77, + "title": "Host-RAM Tier Memory", + "description": "Host RAM the host tier holds for parked models (and for SGLang engines that keep their host copy after waking), by model.", + "type": "timeseries", + "fieldConfig": { "defaults": { "unit": "bytes" } }, + "options": { + "legend": { "calcs": ["mean", "max"], "displayMode": "table", "placement": "bottom", "showLegend": true }, + "tooltip": { "mode": "multi", "sort": "desc" } + }, + "pluginVersion": "10.0.0", + "targets": [ + { + "expr": "max by (model, profile, lane) (sie_worker_model_host_cache_memory_bytes{namespace=\"$namespace\", service=\"$collector\", endpoint=\"prometheus\", producer_service=\"sie-worker\"}) > 0", + "legendFormat": "{{model}} / {{profile}} / {{lane}}", + "refId": "A" + } + ] } ], "refresh": "10s", diff --git a/deploy/helm/sie-cluster/templates/_otel-collector-config.tpl b/deploy/helm/sie-cluster/templates/_otel-collector-config.tpl index d822deb01..87a00e67f 100644 --- a/deploy/helm/sie-cluster/templates/_otel-collector-config.tpl +++ b/deploy/helm/sie-cluster/templates/_otel-collector-config.tpl @@ -95,12 +95,12 @@ processors: error_mode: propagate metrics: metric: - - 'not IsMatch(name, "^(sie[.]dispatcher[.]invocations|sie[.]dispatcher[.]invocation[.]duration|sie[.]dispatcher[.]inflight|sie[.]dispatcher[.]sealed[.]stage|sie[.]dispatcher[.]sealed[.]stage[.]duration|sie[.]dispatcher[.]sealed[.]resident|sie[.]dispatcher[.]sealed[.]sandbox[.]seconds|sie[.]dispatcher[.]sealed[.]cold_start[.]seconds|sie[.]config[.]requests|sie[.]config[.]request[.]duration|sie[.]config[.]epoch|sie[.]config[.]models|sie[.]config[.]publish|sie[.]config[.]store[.]writes|sie[.]config[.]messaging[.]ready|sie[.]worker[.]queue[.]duration|sie[.]worker[.]scheduler[.]request_batch[.]dispatch_wait|sie[.]worker[.]scheduler[.]request_batch[.]total|sie[.]worker[.]queue[.]depth|sie[.]worker[.]batch[.]size|sie[.]worker[.]batch[.]cost|sie[.]worker[.]batch[.]fill_ratio|sie[.]worker[.]queue[.]pending_at_dispatch|sie[.]worker[.]scheduler[.]adaptive[.]wait|sie[.]worker[.]scheduler[.]adaptive[.]cost|sie[.]worker[.]scheduler[.]adaptive[.]p50|sie[.]worker[.]scheduler[.]starvation[.]resets|sie[.]worker[.]ipc[.]requests|sie[.]worker[.]ipc[.]request[.]duration|sie[.]worker[.]ipc[.]response[.]chunks|sie[.]worker[.]ipc[.]response[.]reconstructed[.]size|sie[.]worker[.]ipc[.]response[.]chunk[.]count|sie[.]worker[.]ipc[.]response[.]chunk[.]reserved|sie[.]worker[.]config[.]applies|sie[.]worker[.]config[.]epoch|sie[.]worker[.]config[.]degraded|sie[.]worker[.]nats[.]operations|sie[.]worker[.]nats[.]delivery[.]attempts|sie[.]worker[.]result[.]transport[.]attempts|sie[.]worker[.]result[.]chunks[.]published|sie[.]worker[.]result[.]chunk[.]size|sie[.]worker[.]payload[.]fetches|sie[.]worker[.]payload[.]fetch[.]duration|sie[.]worker[.]payload[.]size|sie[.]worker[.]gpu[.]slots|sie[.]worker[.]pending[.]items|sie[.]worker[.]pending[.]cost|sie[.]worker[.]inflight[.]batches|sie[.]worker[.]saturated|sie[.]worker[.]ipc[.]capacity|sie[.]worker[.]ipc[.]inflight|sie[.]worker[.]ipc[.]acquire[.]duration|sie[.]worker[.]generation[.]model_loading[.]responses|sie[.]worker[.]shutdown[.]drain[.]duration|sie[.]worker[.]runtime[.]batch[.]size|sie[.]worker[.]runtime[.]batch[.]subgroups|sie[.]worker[.]runtime[.]subgroup[.]size|sie[.]worker[.]requests|sie[.]worker[.]request[.]duration|sie[.]worker[.]inference[.]duration|sie[.]worker[.]units|sie[.]worker[.]model[.]loaded|sie[.]worker[.]model[.]load[.]duration|sie[.]worker[.]model[.]memory|sie[.]worker[.]oom[.]recoveries|sie[.]worker[.]model[.]evictions|sie[.]worker[.]generation[.]worker_wait|sie[.]worker[.]generation[.]ttft|sie[.]worker[.]generation[.]tpot|sie[.]worker[.]generation[.]tokens|sie[.]worker[.]generation[.]inflight|sie[.]worker[.]generation[.]kv[.]reserved|sie[.]worker[.]generation[.]kv[.]budget|sie[.]worker[.]generation[.]admission[.]decisions|sie[.]worker[.]generation[.]duplicate_prevented|sie[.]worker[.]generation[.]grammar[.]compile[.]duration|sie[.]worker[.]generation[.]grammar[.]cache[.]lookups|sie[.]worker[.]generation[.]grammar[.]requests|sie[.]worker[.]upstream[.]refusals|sie[.]worker[.]upstream[.]breaker[.]open|sie[.]worker[.]runtime[.]forward[.]duration|sie[.]worker[.]runtime[.]forward[.]permit[.]wait|sie[.]worker[.]runtime[.]forward[.]concurrent|sie[.]worker[.]runtime[.]forward[.]limit|sie[.]worker[.]work_item[.]age|sie[.]worker[.]work_item[.]deadline_exceeded)$")' + - 'not IsMatch(name, "^(sie[.]dispatcher[.]invocations|sie[.]dispatcher[.]invocation[.]duration|sie[.]dispatcher[.]inflight|sie[.]dispatcher[.]sealed[.]stage|sie[.]dispatcher[.]sealed[.]stage[.]duration|sie[.]dispatcher[.]sealed[.]resident|sie[.]dispatcher[.]sealed[.]sandbox[.]seconds|sie[.]dispatcher[.]sealed[.]cold_start[.]seconds|sie[.]config[.]requests|sie[.]config[.]request[.]duration|sie[.]config[.]epoch|sie[.]config[.]models|sie[.]config[.]publish|sie[.]config[.]store[.]writes|sie[.]config[.]messaging[.]ready|sie[.]worker[.]queue[.]duration|sie[.]worker[.]scheduler[.]request_batch[.]dispatch_wait|sie[.]worker[.]scheduler[.]request_batch[.]total|sie[.]worker[.]queue[.]depth|sie[.]worker[.]batch[.]size|sie[.]worker[.]batch[.]cost|sie[.]worker[.]batch[.]fill_ratio|sie[.]worker[.]queue[.]pending_at_dispatch|sie[.]worker[.]scheduler[.]adaptive[.]wait|sie[.]worker[.]scheduler[.]adaptive[.]cost|sie[.]worker[.]scheduler[.]adaptive[.]p50|sie[.]worker[.]scheduler[.]starvation[.]resets|sie[.]worker[.]ipc[.]requests|sie[.]worker[.]ipc[.]request[.]duration|sie[.]worker[.]ipc[.]response[.]chunks|sie[.]worker[.]ipc[.]response[.]reconstructed[.]size|sie[.]worker[.]ipc[.]response[.]chunk[.]count|sie[.]worker[.]ipc[.]response[.]chunk[.]reserved|sie[.]worker[.]config[.]applies|sie[.]worker[.]config[.]epoch|sie[.]worker[.]config[.]degraded|sie[.]worker[.]nats[.]operations|sie[.]worker[.]nats[.]delivery[.]attempts|sie[.]worker[.]result[.]transport[.]attempts|sie[.]worker[.]result[.]chunks[.]published|sie[.]worker[.]result[.]chunk[.]size|sie[.]worker[.]payload[.]fetches|sie[.]worker[.]payload[.]fetch[.]duration|sie[.]worker[.]payload[.]size|sie[.]worker[.]gpu[.]slots|sie[.]worker[.]pending[.]items|sie[.]worker[.]pending[.]cost|sie[.]worker[.]inflight[.]batches|sie[.]worker[.]saturated|sie[.]worker[.]ipc[.]capacity|sie[.]worker[.]ipc[.]inflight|sie[.]worker[.]ipc[.]acquire[.]duration|sie[.]worker[.]generation[.]model_loading[.]responses|sie[.]worker[.]shutdown[.]drain[.]duration|sie[.]worker[.]runtime[.]batch[.]size|sie[.]worker[.]runtime[.]batch[.]subgroups|sie[.]worker[.]runtime[.]subgroup[.]size|sie[.]worker[.]requests|sie[.]worker[.]request[.]duration|sie[.]worker[.]inference[.]duration|sie[.]worker[.]units|sie[.]worker[.]model[.]loaded|sie[.]worker[.]model[.]load[.]duration|sie[.]worker[.]model[.]memory|sie[.]worker[.]oom[.]recoveries|sie[.]worker[.]model[.]evictions|sie[.]worker[.]model[.]host_cache[.]lookups|sie[.]worker[.]model[.]host_cache[.]wake[.]duration|sie[.]worker[.]model[.]host_cache[.]memory|sie[.]worker[.]generation[.]worker_wait|sie[.]worker[.]generation[.]ttft|sie[.]worker[.]generation[.]tpot|sie[.]worker[.]generation[.]tokens|sie[.]worker[.]generation[.]inflight|sie[.]worker[.]generation[.]kv[.]reserved|sie[.]worker[.]generation[.]kv[.]budget|sie[.]worker[.]generation[.]admission[.]decisions|sie[.]worker[.]generation[.]duplicate_prevented|sie[.]worker[.]generation[.]grammar[.]compile[.]duration|sie[.]worker[.]generation[.]grammar[.]cache[.]lookups|sie[.]worker[.]generation[.]grammar[.]requests|sie[.]worker[.]upstream[.]refusals|sie[.]worker[.]upstream[.]breaker[.]open|sie[.]worker[.]runtime[.]forward[.]duration|sie[.]worker[.]runtime[.]forward[.]permit[.]wait|sie[.]worker[.]runtime[.]forward[.]concurrent|sie[.]worker[.]runtime[.]forward[.]limit|sie[.]worker[.]work_item[.]age|sie[.]worker[.]work_item[.]deadline_exceeded)$")' - 'resource.attributes["service.name"] != "sie-config" and resource.attributes["service.name"] != "sie-dispatcher" and resource.attributes["service.name"] != "sie-worker" and resource.attributes["service.name"] != "sie-worker-sidecar"' - 'resource.attributes["service.name"] == "sie-dispatcher" and not IsMatch(name, "^(sie[.]dispatcher[.]invocations|sie[.]dispatcher[.]invocation[.]duration|sie[.]dispatcher[.]inflight|sie[.]dispatcher[.]sealed[.]stage|sie[.]dispatcher[.]sealed[.]stage[.]duration|sie[.]dispatcher[.]sealed[.]resident|sie[.]dispatcher[.]sealed[.]sandbox[.]seconds|sie[.]dispatcher[.]sealed[.]cold_start[.]seconds)$")' - 'resource.attributes["service.name"] == "sie-config" and not IsMatch(name, "^(sie[.]config[.]requests|sie[.]config[.]request[.]duration|sie[.]config[.]epoch|sie[.]config[.]models|sie[.]config[.]publish|sie[.]config[.]store[.]writes|sie[.]config[.]messaging[.]ready)$")' - 'resource.attributes["service.name"] == "sie-worker-sidecar" and not IsMatch(name, "^(sie[.]worker[.]queue[.]duration|sie[.]worker[.]scheduler[.]request_batch[.]dispatch_wait|sie[.]worker[.]scheduler[.]request_batch[.]total|sie[.]worker[.]queue[.]depth|sie[.]worker[.]batch[.]size|sie[.]worker[.]batch[.]cost|sie[.]worker[.]batch[.]fill_ratio|sie[.]worker[.]scheduler[.]adaptive[.]wait|sie[.]worker[.]scheduler[.]adaptive[.]cost|sie[.]worker[.]scheduler[.]adaptive[.]p50|sie[.]worker[.]scheduler[.]starvation[.]resets|sie[.]worker[.]ipc[.]requests|sie[.]worker[.]ipc[.]request[.]duration|sie[.]worker[.]ipc[.]response[.]chunks|sie[.]worker[.]ipc[.]response[.]reconstructed[.]size|sie[.]worker[.]ipc[.]response[.]chunk[.]count|sie[.]worker[.]ipc[.]response[.]chunk[.]reserved|sie[.]worker[.]config[.]applies|sie[.]worker[.]config[.]epoch|sie[.]worker[.]config[.]degraded|sie[.]worker[.]nats[.]operations|sie[.]worker[.]nats[.]delivery[.]attempts|sie[.]worker[.]result[.]transport[.]attempts|sie[.]worker[.]result[.]chunks[.]published|sie[.]worker[.]result[.]chunk[.]size|sie[.]worker[.]payload[.]fetches|sie[.]worker[.]payload[.]fetch[.]duration|sie[.]worker[.]payload[.]size|sie[.]worker[.]gpu[.]slots|sie[.]worker[.]pending[.]items|sie[.]worker[.]pending[.]cost|sie[.]worker[.]inflight[.]batches|sie[.]worker[.]saturated|sie[.]worker[.]ipc[.]capacity|sie[.]worker[.]ipc[.]inflight|sie[.]worker[.]ipc[.]acquire[.]duration|sie[.]worker[.]generation[.]model_loading[.]responses|sie[.]worker[.]shutdown[.]drain[.]duration|sie[.]worker[.]work_item[.]age|sie[.]worker[.]work_item[.]deadline_exceeded)$")' - - 'resource.attributes["service.name"] == "sie-worker" and not IsMatch(name, "^(sie[.]worker[.]queue[.]duration|sie[.]worker[.]queue[.]depth|sie[.]worker[.]batch[.]size|sie[.]worker[.]batch[.]cost|sie[.]worker[.]batch[.]fill_ratio|sie[.]worker[.]queue[.]pending_at_dispatch|sie[.]worker[.]scheduler[.]adaptive[.]wait|sie[.]worker[.]scheduler[.]adaptive[.]cost|sie[.]worker[.]scheduler[.]adaptive[.]p50|sie[.]worker[.]scheduler[.]starvation[.]resets|sie[.]worker[.]runtime[.]batch[.]size|sie[.]worker[.]runtime[.]batch[.]subgroups|sie[.]worker[.]runtime[.]subgroup[.]size|sie[.]worker[.]requests|sie[.]worker[.]request[.]duration|sie[.]worker[.]inference[.]duration|sie[.]worker[.]units|sie[.]worker[.]model[.]loaded|sie[.]worker[.]model[.]load[.]duration|sie[.]worker[.]model[.]memory|sie[.]worker[.]oom[.]recoveries|sie[.]worker[.]model[.]evictions|sie[.]worker[.]generation[.]worker_wait|sie[.]worker[.]generation[.]ttft|sie[.]worker[.]generation[.]tpot|sie[.]worker[.]generation[.]tokens|sie[.]worker[.]generation[.]inflight|sie[.]worker[.]generation[.]kv[.]reserved|sie[.]worker[.]generation[.]kv[.]budget|sie[.]worker[.]generation[.]admission[.]decisions|sie[.]worker[.]generation[.]duplicate_prevented|sie[.]worker[.]generation[.]grammar[.]compile[.]duration|sie[.]worker[.]generation[.]grammar[.]cache[.]lookups|sie[.]worker[.]generation[.]grammar[.]requests|sie[.]worker[.]upstream[.]refusals|sie[.]worker[.]upstream[.]breaker[.]open|sie[.]worker[.]runtime[.]forward[.]duration|sie[.]worker[.]runtime[.]forward[.]permit[.]wait|sie[.]worker[.]runtime[.]forward[.]concurrent|sie[.]worker[.]runtime[.]forward[.]limit)$")' + - 'resource.attributes["service.name"] == "sie-worker" and not IsMatch(name, "^(sie[.]worker[.]queue[.]duration|sie[.]worker[.]queue[.]depth|sie[.]worker[.]batch[.]size|sie[.]worker[.]batch[.]cost|sie[.]worker[.]batch[.]fill_ratio|sie[.]worker[.]queue[.]pending_at_dispatch|sie[.]worker[.]scheduler[.]adaptive[.]wait|sie[.]worker[.]scheduler[.]adaptive[.]cost|sie[.]worker[.]scheduler[.]adaptive[.]p50|sie[.]worker[.]scheduler[.]starvation[.]resets|sie[.]worker[.]runtime[.]batch[.]size|sie[.]worker[.]runtime[.]batch[.]subgroups|sie[.]worker[.]runtime[.]subgroup[.]size|sie[.]worker[.]requests|sie[.]worker[.]request[.]duration|sie[.]worker[.]inference[.]duration|sie[.]worker[.]units|sie[.]worker[.]model[.]loaded|sie[.]worker[.]model[.]load[.]duration|sie[.]worker[.]model[.]memory|sie[.]worker[.]oom[.]recoveries|sie[.]worker[.]model[.]evictions|sie[.]worker[.]model[.]host_cache[.]lookups|sie[.]worker[.]model[.]host_cache[.]wake[.]duration|sie[.]worker[.]model[.]host_cache[.]memory|sie[.]worker[.]generation[.]worker_wait|sie[.]worker[.]generation[.]ttft|sie[.]worker[.]generation[.]tpot|sie[.]worker[.]generation[.]tokens|sie[.]worker[.]generation[.]inflight|sie[.]worker[.]generation[.]kv[.]reserved|sie[.]worker[.]generation[.]kv[.]budget|sie[.]worker[.]generation[.]admission[.]decisions|sie[.]worker[.]generation[.]duplicate_prevented|sie[.]worker[.]generation[.]grammar[.]compile[.]duration|sie[.]worker[.]generation[.]grammar[.]cache[.]lookups|sie[.]worker[.]generation[.]grammar[.]requests|sie[.]worker[.]upstream[.]refusals|sie[.]worker[.]upstream[.]breaker[.]open|sie[.]worker[.]runtime[.]forward[.]duration|sie[.]worker[.]runtime[.]forward[.]permit[.]wait|sie[.]worker[.]runtime[.]forward[.]concurrent|sie[.]worker[.]runtime[.]forward[.]limit)$")' # This is the metric field firewall shared by Prometheus and OTLP # destinations. Every descriptor and point is reduced to the checked-in # contract; resource and instrumentation-scope extensions are discarded. @@ -121,7 +121,7 @@ processors: statements: - set(description, "") - 'set(unit, "{request}") where IsMatch(name, "^(sie[.]gateway[.]requests|sie[.]gateway[.]admission[.]decisions|sie[.]gateway[.]dispatches|sie[.]gateway[.]remote[.]fallbacks|sie[.]gateway[.]remote[.]numerical_admissions|sie[.]config[.]requests|sie[.]worker[.]ipc[.]requests|sie[.]worker[.]ipc[.]capacity|sie[.]worker[.]ipc[.]inflight|sie[.]worker[.]generation[.]inflight|sie[.]worker[.]generation[.]duplicate_prevented|sie[.]worker[.]generation[.]grammar[.]requests|sie[.]worker[.]upstream[.]refusals|sie[.]gateway[.]pending_demand|sie[.]gateway[.]rejected[.]requests)$")' - - 'set(unit, "s") where IsMatch(name, "^(sie[.]gateway[.]request[.]duration|sie[.]gateway[.]dispatch[.]duration|sie[.]gateway[.]remote[.]serving[.]duration|sie[.]gateway[.]queue[.]publish[.]duration|sie[.]gateway[.]queue[.]result_wait[.]duration|sie[.]gateway[.]generation[.]ttft|sie[.]gateway[.]generation[.]tpot|sie[.]dispatcher[.]invocation[.]duration|sie[.]dispatcher[.]sealed[.]stage[.]duration|sie[.]dispatcher[.]sealed[.]sandbox[.]seconds|sie[.]dispatcher[.]sealed[.]cold_start[.]seconds|sie[.]config[.]request[.]duration|sie[.]worker[.]queue[.]duration|sie[.]worker[.]scheduler[.]request_batch[.]dispatch_wait|sie[.]worker[.]scheduler[.]request_batch[.]total|sie[.]worker[.]scheduler[.]adaptive[.]wait|sie[.]worker[.]scheduler[.]adaptive[.]p50|sie[.]worker[.]ipc[.]request[.]duration|sie[.]worker[.]payload[.]fetch[.]duration|sie[.]worker[.]ipc[.]acquire[.]duration|sie[.]worker[.]shutdown[.]drain[.]duration|sie[.]worker[.]request[.]duration|sie[.]worker[.]inference[.]duration|sie[.]worker[.]model[.]load[.]duration|sie[.]worker[.]generation[.]worker_wait|sie[.]worker[.]generation[.]ttft|sie[.]worker[.]generation[.]tpot|sie[.]worker[.]generation[.]grammar[.]compile[.]duration|sie[.]worker[.]runtime[.]forward[.]duration|sie[.]worker[.]runtime[.]forward[.]permit[.]wait|sie[.]gateway[.]lane[.]queue[.]snapshot[.]timestamp|sie[.]gateway[.]capacity[.]snapshot[.]timestamp|sie[.]gateway[.]settlement[.]confirm[.]duration|sie[.]worker[.]work_item[.]age)$")' + - 'set(unit, "s") where IsMatch(name, "^(sie[.]gateway[.]request[.]duration|sie[.]gateway[.]dispatch[.]duration|sie[.]gateway[.]remote[.]serving[.]duration|sie[.]gateway[.]queue[.]publish[.]duration|sie[.]gateway[.]queue[.]result_wait[.]duration|sie[.]gateway[.]generation[.]ttft|sie[.]gateway[.]generation[.]tpot|sie[.]dispatcher[.]invocation[.]duration|sie[.]dispatcher[.]sealed[.]stage[.]duration|sie[.]dispatcher[.]sealed[.]sandbox[.]seconds|sie[.]dispatcher[.]sealed[.]cold_start[.]seconds|sie[.]config[.]request[.]duration|sie[.]worker[.]queue[.]duration|sie[.]worker[.]scheduler[.]request_batch[.]dispatch_wait|sie[.]worker[.]scheduler[.]request_batch[.]total|sie[.]worker[.]scheduler[.]adaptive[.]wait|sie[.]worker[.]scheduler[.]adaptive[.]p50|sie[.]worker[.]ipc[.]request[.]duration|sie[.]worker[.]payload[.]fetch[.]duration|sie[.]worker[.]ipc[.]acquire[.]duration|sie[.]worker[.]shutdown[.]drain[.]duration|sie[.]worker[.]request[.]duration|sie[.]worker[.]inference[.]duration|sie[.]worker[.]model[.]load[.]duration|sie[.]worker[.]model[.]host_cache[.]wake[.]duration|sie[.]worker[.]generation[.]worker_wait|sie[.]worker[.]generation[.]ttft|sie[.]worker[.]generation[.]tpot|sie[.]worker[.]generation[.]grammar[.]compile[.]duration|sie[.]worker[.]runtime[.]forward[.]duration|sie[.]worker[.]runtime[.]forward[.]permit[.]wait|sie[.]gateway[.]lane[.]queue[.]snapshot[.]timestamp|sie[.]gateway[.]capacity[.]snapshot[.]timestamp|sie[.]gateway[.]settlement[.]confirm[.]duration|sie[.]worker[.]work_item[.]age)$")' - 'set(unit, "{epoch}") where IsMatch(name, "^(sie[.]gateway[.]config[.]applied_epoch|sie[.]config[.]epoch|sie[.]worker[.]config[.]epoch)$")' - 'set(unit, "{operation}") where IsMatch(name, "^(sie[.]gateway[.]config[.]operations|sie[.]config[.]publish|sie[.]config[.]store[.]writes|sie[.]worker[.]nats[.]operations)$")' - 'set(unit, "1") where IsMatch(name, "^(sie[.]gateway[.]config[.]bootstrap[.]degraded|sie[.]gateway[.]messaging[.]client[.]ready|sie[.]config[.]messaging[.]ready|sie[.]worker[.]batch[.]fill_ratio|sie[.]worker[.]config[.]degraded|sie[.]worker[.]saturated|sie[.]gateway[.]pool[.]pinned_model[.]loaded|sie[.]gateway[.]model_loads_in_progress)$")' @@ -148,7 +148,7 @@ processors: - 'set(unit, "{apply}") where name == "sie.worker.config.applies"' - 'set(unit, "{attempt}") where name == "sie.worker.nats.delivery.attempts"' - 'set(unit, "{fetch}") where name == "sie.worker.payload.fetches"' - - 'set(unit, "By") where IsMatch(name, "^(sie[.]worker[.]payload[.]size|sie[.]worker[.]model[.]memory)$")' + - 'set(unit, "By") where IsMatch(name, "^(sie[.]worker[.]payload[.]size|sie[.]worker[.]model[.]memory|sie[.]worker[.]model[.]host_cache[.]memory)$")' - 'set(unit, "{slot}") where name == "sie.worker.gpu.slots"' - 'set(unit, "{batch}") where name == "sie.worker.inflight.batches"' - 'set(unit, "{subgroup}") where name == "sie.worker.runtime.batch.subgroups"' @@ -156,7 +156,7 @@ processors: - 'set(unit, "{recovery}") where name == "sie.worker.oom.recoveries"' - 'set(unit, "{decision}") where name == "sie.worker.generation.admission.decisions"' - 'set(unit, "{decision}") where name == "sie.gateway.queue.lane_admission.decisions"' - - 'set(unit, "{lookup}") where name == "sie.worker.generation.grammar.cache.lookups"' + - 'set(unit, "{lookup}") where IsMatch(name, "^(sie[.]worker[.]generation[.]grammar[.]cache[.]lookups|sie[.]worker[.]model[.]host_cache[.]lookups)$")' - 'set(unit, "{forward}") where IsMatch(name, "^(sie[.]worker[.]runtime[.]forward[.]concurrent|sie[.]worker[.]runtime[.]forward[.]limit)$")' - 'set(unit, "{gpu}") where name == "sie.gateway.active_lease.gpus"' - 'set(unit, "{worker}") where name == "sie.gateway.pool.warm_floor"' @@ -214,10 +214,12 @@ processors: - 'keep_keys(attributes, ["operation", "outcome", "backend", "lane", "model", "profile"]) where IsMatch(metric.name, "^(sie[.]worker[.]requests|sie[.]worker[.]request[.]duration)$")' - 'keep_keys(attributes, ["operation", "outcome", "phase", "backend", "lane", "model", "profile"]) where metric.name == "sie.worker.inference.duration"' - 'keep_keys(attributes, ["operation", "backend", "lane", "model", "profile", "unit.type"]) where metric.name == "sie.worker.units"' - - 'keep_keys(attributes, ["backend", "lane", "model", "profile"]) where IsMatch(metric.name, "^(sie[.]worker[.]model[.]loaded|sie[.]worker[.]model[.]memory|sie[.]worker[.]generation[.]inflight|sie[.]worker[.]generation[.]kv[.]reserved|sie[.]worker[.]generation[.]kv[.]budget|sie[.]worker[.]runtime[.]forward[.]limit)$")' + - 'keep_keys(attributes, ["backend", "lane", "model", "profile"]) where IsMatch(metric.name, "^(sie[.]worker[.]model[.]loaded|sie[.]worker[.]model[.]memory|sie[.]worker[.]model[.]host_cache[.]memory|sie[.]worker[.]generation[.]inflight|sie[.]worker[.]generation[.]kv[.]reserved|sie[.]worker[.]generation[.]kv[.]budget|sie[.]worker[.]runtime[.]forward[.]limit)$")' - 'keep_keys(attributes, ["outcome", "stage", "backend", "lane", "model", "profile"]) where metric.name == "sie.worker.model.load.duration"' - 'keep_keys(attributes, ["strategy", "outcome", "backend", "lane", "model", "profile"]) where metric.name == "sie.worker.oom.recoveries"' - 'keep_keys(attributes, ["reason", "backend", "lane", "model", "profile"]) where metric.name == "sie.worker.model.evictions"' + - 'keep_keys(attributes, ["result", "backend", "lane", "model", "profile"]) where metric.name == "sie.worker.model.host_cache.lookups"' + - 'keep_keys(attributes, ["outcome", "backend", "lane", "model", "profile"]) where metric.name == "sie.worker.model.host_cache.wake.duration"' - 'keep_keys(attributes, ["grammar", "backend", "lane", "model", "profile"]) where IsMatch(metric.name, "^(sie[.]worker[.]generation[.]worker_wait|sie[.]worker[.]generation[.]ttft|sie[.]worker[.]generation[.]tpot)$")' - 'keep_keys(attributes, ["token.type", "grammar", "backend", "lane", "model", "profile"]) where metric.name == "sie.worker.generation.tokens"' - 'keep_keys(attributes, ["outcome", "reason", "backend", "lane", "model", "profile"]) where metric.name == "sie.worker.generation.admission.decisions"' @@ -297,12 +299,12 @@ processors: error_mode: propagate metrics: metric: - - 'not IsMatch(name, "^(sie[.]dispatcher[.]invocations|sie[.]dispatcher[.]invocation[.]duration|sie[.]dispatcher[.]inflight|sie[.]dispatcher[.]sealed[.]stage|sie[.]dispatcher[.]sealed[.]stage[.]duration|sie[.]dispatcher[.]sealed[.]resident|sie[.]dispatcher[.]sealed[.]sandbox[.]seconds|sie[.]dispatcher[.]sealed[.]cold_start[.]seconds|sie[.]config[.]requests|sie[.]config[.]request[.]duration|sie[.]config[.]epoch|sie[.]config[.]models|sie[.]config[.]publish|sie[.]config[.]store[.]writes|sie[.]config[.]messaging[.]ready|sie[.]worker[.]queue[.]duration|sie[.]worker[.]scheduler[.]request_batch[.]dispatch_wait|sie[.]worker[.]scheduler[.]request_batch[.]total|sie[.]worker[.]queue[.]depth|sie[.]worker[.]batch[.]size|sie[.]worker[.]batch[.]cost|sie[.]worker[.]batch[.]fill_ratio|sie[.]worker[.]queue[.]pending_at_dispatch|sie[.]worker[.]scheduler[.]adaptive[.]wait|sie[.]worker[.]scheduler[.]adaptive[.]cost|sie[.]worker[.]scheduler[.]adaptive[.]p50|sie[.]worker[.]scheduler[.]starvation[.]resets|sie[.]worker[.]ipc[.]requests|sie[.]worker[.]ipc[.]request[.]duration|sie[.]worker[.]ipc[.]response[.]chunks|sie[.]worker[.]ipc[.]response[.]reconstructed[.]size|sie[.]worker[.]ipc[.]response[.]chunk[.]count|sie[.]worker[.]ipc[.]response[.]chunk[.]reserved|sie[.]worker[.]config[.]applies|sie[.]worker[.]config[.]epoch|sie[.]worker[.]config[.]degraded|sie[.]worker[.]nats[.]operations|sie[.]worker[.]nats[.]delivery[.]attempts|sie[.]worker[.]result[.]transport[.]attempts|sie[.]worker[.]result[.]chunks[.]published|sie[.]worker[.]result[.]chunk[.]size|sie[.]worker[.]payload[.]fetches|sie[.]worker[.]payload[.]fetch[.]duration|sie[.]worker[.]payload[.]size|sie[.]worker[.]gpu[.]slots|sie[.]worker[.]pending[.]items|sie[.]worker[.]pending[.]cost|sie[.]worker[.]inflight[.]batches|sie[.]worker[.]saturated|sie[.]worker[.]ipc[.]capacity|sie[.]worker[.]ipc[.]inflight|sie[.]worker[.]ipc[.]acquire[.]duration|sie[.]worker[.]generation[.]model_loading[.]responses|sie[.]worker[.]shutdown[.]drain[.]duration|sie[.]worker[.]runtime[.]batch[.]size|sie[.]worker[.]runtime[.]batch[.]subgroups|sie[.]worker[.]runtime[.]subgroup[.]size|sie[.]worker[.]requests|sie[.]worker[.]request[.]duration|sie[.]worker[.]inference[.]duration|sie[.]worker[.]units|sie[.]worker[.]model[.]loaded|sie[.]worker[.]model[.]load[.]duration|sie[.]worker[.]model[.]memory|sie[.]worker[.]oom[.]recoveries|sie[.]worker[.]model[.]evictions|sie[.]worker[.]generation[.]worker_wait|sie[.]worker[.]generation[.]ttft|sie[.]worker[.]generation[.]tpot|sie[.]worker[.]generation[.]tokens|sie[.]worker[.]generation[.]inflight|sie[.]worker[.]generation[.]kv[.]reserved|sie[.]worker[.]generation[.]kv[.]budget|sie[.]worker[.]generation[.]admission[.]decisions|sie[.]worker[.]generation[.]duplicate_prevented|sie[.]worker[.]generation[.]grammar[.]compile[.]duration|sie[.]worker[.]generation[.]grammar[.]cache[.]lookups|sie[.]worker[.]generation[.]grammar[.]requests|sie[.]worker[.]upstream[.]refusals|sie[.]worker[.]upstream[.]breaker[.]open|sie[.]worker[.]runtime[.]forward[.]duration|sie[.]worker[.]runtime[.]forward[.]permit[.]wait|sie[.]worker[.]runtime[.]forward[.]concurrent|sie[.]worker[.]runtime[.]forward[.]limit|sie[.]worker[.]work_item[.]age|sie[.]worker[.]work_item[.]deadline_exceeded)$")' + - 'not IsMatch(name, "^(sie[.]dispatcher[.]invocations|sie[.]dispatcher[.]invocation[.]duration|sie[.]dispatcher[.]inflight|sie[.]dispatcher[.]sealed[.]stage|sie[.]dispatcher[.]sealed[.]stage[.]duration|sie[.]dispatcher[.]sealed[.]resident|sie[.]dispatcher[.]sealed[.]sandbox[.]seconds|sie[.]dispatcher[.]sealed[.]cold_start[.]seconds|sie[.]config[.]requests|sie[.]config[.]request[.]duration|sie[.]config[.]epoch|sie[.]config[.]models|sie[.]config[.]publish|sie[.]config[.]store[.]writes|sie[.]config[.]messaging[.]ready|sie[.]worker[.]queue[.]duration|sie[.]worker[.]scheduler[.]request_batch[.]dispatch_wait|sie[.]worker[.]scheduler[.]request_batch[.]total|sie[.]worker[.]queue[.]depth|sie[.]worker[.]batch[.]size|sie[.]worker[.]batch[.]cost|sie[.]worker[.]batch[.]fill_ratio|sie[.]worker[.]queue[.]pending_at_dispatch|sie[.]worker[.]scheduler[.]adaptive[.]wait|sie[.]worker[.]scheduler[.]adaptive[.]cost|sie[.]worker[.]scheduler[.]adaptive[.]p50|sie[.]worker[.]scheduler[.]starvation[.]resets|sie[.]worker[.]ipc[.]requests|sie[.]worker[.]ipc[.]request[.]duration|sie[.]worker[.]ipc[.]response[.]chunks|sie[.]worker[.]ipc[.]response[.]reconstructed[.]size|sie[.]worker[.]ipc[.]response[.]chunk[.]count|sie[.]worker[.]ipc[.]response[.]chunk[.]reserved|sie[.]worker[.]config[.]applies|sie[.]worker[.]config[.]epoch|sie[.]worker[.]config[.]degraded|sie[.]worker[.]nats[.]operations|sie[.]worker[.]nats[.]delivery[.]attempts|sie[.]worker[.]result[.]transport[.]attempts|sie[.]worker[.]result[.]chunks[.]published|sie[.]worker[.]result[.]chunk[.]size|sie[.]worker[.]payload[.]fetches|sie[.]worker[.]payload[.]fetch[.]duration|sie[.]worker[.]payload[.]size|sie[.]worker[.]gpu[.]slots|sie[.]worker[.]pending[.]items|sie[.]worker[.]pending[.]cost|sie[.]worker[.]inflight[.]batches|sie[.]worker[.]saturated|sie[.]worker[.]ipc[.]capacity|sie[.]worker[.]ipc[.]inflight|sie[.]worker[.]ipc[.]acquire[.]duration|sie[.]worker[.]generation[.]model_loading[.]responses|sie[.]worker[.]shutdown[.]drain[.]duration|sie[.]worker[.]runtime[.]batch[.]size|sie[.]worker[.]runtime[.]batch[.]subgroups|sie[.]worker[.]runtime[.]subgroup[.]size|sie[.]worker[.]requests|sie[.]worker[.]request[.]duration|sie[.]worker[.]inference[.]duration|sie[.]worker[.]units|sie[.]worker[.]model[.]loaded|sie[.]worker[.]model[.]load[.]duration|sie[.]worker[.]model[.]memory|sie[.]worker[.]oom[.]recoveries|sie[.]worker[.]model[.]evictions|sie[.]worker[.]model[.]host_cache[.]lookups|sie[.]worker[.]model[.]host_cache[.]wake[.]duration|sie[.]worker[.]model[.]host_cache[.]memory|sie[.]worker[.]generation[.]worker_wait|sie[.]worker[.]generation[.]ttft|sie[.]worker[.]generation[.]tpot|sie[.]worker[.]generation[.]tokens|sie[.]worker[.]generation[.]inflight|sie[.]worker[.]generation[.]kv[.]reserved|sie[.]worker[.]generation[.]kv[.]budget|sie[.]worker[.]generation[.]admission[.]decisions|sie[.]worker[.]generation[.]duplicate_prevented|sie[.]worker[.]generation[.]grammar[.]compile[.]duration|sie[.]worker[.]generation[.]grammar[.]cache[.]lookups|sie[.]worker[.]generation[.]grammar[.]requests|sie[.]worker[.]upstream[.]refusals|sie[.]worker[.]upstream[.]breaker[.]open|sie[.]worker[.]runtime[.]forward[.]duration|sie[.]worker[.]runtime[.]forward[.]permit[.]wait|sie[.]worker[.]runtime[.]forward[.]concurrent|sie[.]worker[.]runtime[.]forward[.]limit|sie[.]worker[.]work_item[.]age|sie[.]worker[.]work_item[.]deadline_exceeded)$")' - 'resource.attributes["service.name"] != "sie-config" and resource.attributes["service.name"] != "sie-dispatcher" and resource.attributes["service.name"] != "sie-worker" and resource.attributes["service.name"] != "sie-worker-sidecar"' - 'resource.attributes["service.name"] == "sie-dispatcher" and not IsMatch(name, "^(sie[.]dispatcher[.]invocations|sie[.]dispatcher[.]invocation[.]duration|sie[.]dispatcher[.]inflight|sie[.]dispatcher[.]sealed[.]stage|sie[.]dispatcher[.]sealed[.]stage[.]duration|sie[.]dispatcher[.]sealed[.]resident|sie[.]dispatcher[.]sealed[.]sandbox[.]seconds|sie[.]dispatcher[.]sealed[.]cold_start[.]seconds)$")' - 'resource.attributes["service.name"] == "sie-config" and not IsMatch(name, "^(sie[.]config[.]requests|sie[.]config[.]request[.]duration|sie[.]config[.]epoch|sie[.]config[.]models|sie[.]config[.]publish|sie[.]config[.]store[.]writes|sie[.]config[.]messaging[.]ready)$")' - 'resource.attributes["service.name"] == "sie-worker-sidecar" and not IsMatch(name, "^(sie[.]worker[.]queue[.]duration|sie[.]worker[.]scheduler[.]request_batch[.]dispatch_wait|sie[.]worker[.]scheduler[.]request_batch[.]total|sie[.]worker[.]queue[.]depth|sie[.]worker[.]batch[.]size|sie[.]worker[.]batch[.]cost|sie[.]worker[.]batch[.]fill_ratio|sie[.]worker[.]scheduler[.]adaptive[.]wait|sie[.]worker[.]scheduler[.]adaptive[.]cost|sie[.]worker[.]scheduler[.]adaptive[.]p50|sie[.]worker[.]scheduler[.]starvation[.]resets|sie[.]worker[.]ipc[.]requests|sie[.]worker[.]ipc[.]request[.]duration|sie[.]worker[.]ipc[.]response[.]chunks|sie[.]worker[.]ipc[.]response[.]reconstructed[.]size|sie[.]worker[.]ipc[.]response[.]chunk[.]count|sie[.]worker[.]ipc[.]response[.]chunk[.]reserved|sie[.]worker[.]config[.]applies|sie[.]worker[.]config[.]epoch|sie[.]worker[.]config[.]degraded|sie[.]worker[.]nats[.]operations|sie[.]worker[.]nats[.]delivery[.]attempts|sie[.]worker[.]result[.]transport[.]attempts|sie[.]worker[.]result[.]chunks[.]published|sie[.]worker[.]result[.]chunk[.]size|sie[.]worker[.]payload[.]fetches|sie[.]worker[.]payload[.]fetch[.]duration|sie[.]worker[.]payload[.]size|sie[.]worker[.]gpu[.]slots|sie[.]worker[.]pending[.]items|sie[.]worker[.]pending[.]cost|sie[.]worker[.]inflight[.]batches|sie[.]worker[.]saturated|sie[.]worker[.]ipc[.]capacity|sie[.]worker[.]ipc[.]inflight|sie[.]worker[.]ipc[.]acquire[.]duration|sie[.]worker[.]generation[.]model_loading[.]responses|sie[.]worker[.]shutdown[.]drain[.]duration|sie[.]worker[.]work_item[.]age|sie[.]worker[.]work_item[.]deadline_exceeded)$")' - - 'resource.attributes["service.name"] == "sie-worker" and not IsMatch(name, "^(sie[.]worker[.]queue[.]duration|sie[.]worker[.]queue[.]depth|sie[.]worker[.]batch[.]size|sie[.]worker[.]batch[.]cost|sie[.]worker[.]batch[.]fill_ratio|sie[.]worker[.]queue[.]pending_at_dispatch|sie[.]worker[.]scheduler[.]adaptive[.]wait|sie[.]worker[.]scheduler[.]adaptive[.]cost|sie[.]worker[.]scheduler[.]adaptive[.]p50|sie[.]worker[.]scheduler[.]starvation[.]resets|sie[.]worker[.]runtime[.]batch[.]size|sie[.]worker[.]runtime[.]batch[.]subgroups|sie[.]worker[.]runtime[.]subgroup[.]size|sie[.]worker[.]requests|sie[.]worker[.]request[.]duration|sie[.]worker[.]inference[.]duration|sie[.]worker[.]units|sie[.]worker[.]model[.]loaded|sie[.]worker[.]model[.]load[.]duration|sie[.]worker[.]model[.]memory|sie[.]worker[.]oom[.]recoveries|sie[.]worker[.]model[.]evictions|sie[.]worker[.]generation[.]worker_wait|sie[.]worker[.]generation[.]ttft|sie[.]worker[.]generation[.]tpot|sie[.]worker[.]generation[.]tokens|sie[.]worker[.]generation[.]inflight|sie[.]worker[.]generation[.]kv[.]reserved|sie[.]worker[.]generation[.]kv[.]budget|sie[.]worker[.]generation[.]admission[.]decisions|sie[.]worker[.]generation[.]duplicate_prevented|sie[.]worker[.]generation[.]grammar[.]compile[.]duration|sie[.]worker[.]generation[.]grammar[.]cache[.]lookups|sie[.]worker[.]generation[.]grammar[.]requests|sie[.]worker[.]upstream[.]refusals|sie[.]worker[.]upstream[.]breaker[.]open|sie[.]worker[.]runtime[.]forward[.]duration|sie[.]worker[.]runtime[.]forward[.]permit[.]wait|sie[.]worker[.]runtime[.]forward[.]concurrent|sie[.]worker[.]runtime[.]forward[.]limit)$")' + - 'resource.attributes["service.name"] == "sie-worker" and not IsMatch(name, "^(sie[.]worker[.]queue[.]duration|sie[.]worker[.]queue[.]depth|sie[.]worker[.]batch[.]size|sie[.]worker[.]batch[.]cost|sie[.]worker[.]batch[.]fill_ratio|sie[.]worker[.]queue[.]pending_at_dispatch|sie[.]worker[.]scheduler[.]adaptive[.]wait|sie[.]worker[.]scheduler[.]adaptive[.]cost|sie[.]worker[.]scheduler[.]adaptive[.]p50|sie[.]worker[.]scheduler[.]starvation[.]resets|sie[.]worker[.]runtime[.]batch[.]size|sie[.]worker[.]runtime[.]batch[.]subgroups|sie[.]worker[.]runtime[.]subgroup[.]size|sie[.]worker[.]requests|sie[.]worker[.]request[.]duration|sie[.]worker[.]inference[.]duration|sie[.]worker[.]units|sie[.]worker[.]model[.]loaded|sie[.]worker[.]model[.]load[.]duration|sie[.]worker[.]model[.]memory|sie[.]worker[.]oom[.]recoveries|sie[.]worker[.]model[.]evictions|sie[.]worker[.]model[.]host_cache[.]lookups|sie[.]worker[.]model[.]host_cache[.]wake[.]duration|sie[.]worker[.]model[.]host_cache[.]memory|sie[.]worker[.]generation[.]worker_wait|sie[.]worker[.]generation[.]ttft|sie[.]worker[.]generation[.]tpot|sie[.]worker[.]generation[.]tokens|sie[.]worker[.]generation[.]inflight|sie[.]worker[.]generation[.]kv[.]reserved|sie[.]worker[.]generation[.]kv[.]budget|sie[.]worker[.]generation[.]admission[.]decisions|sie[.]worker[.]generation[.]duplicate_prevented|sie[.]worker[.]generation[.]grammar[.]compile[.]duration|sie[.]worker[.]generation[.]grammar[.]cache[.]lookups|sie[.]worker[.]generation[.]grammar[.]requests|sie[.]worker[.]upstream[.]refusals|sie[.]worker[.]upstream[.]breaker[.]open|sie[.]worker[.]runtime[.]forward[.]duration|sie[.]worker[.]runtime[.]forward[.]permit[.]wait|sie[.]worker[.]runtime[.]forward[.]concurrent|sie[.]worker[.]runtime[.]forward[.]limit)$")' {{- end }} {{- if $betterStack.enabled }} # Scalar lookups select one value; keep_keys alone retains duplicate OTLP diff --git a/packages/sie_server/README.md b/packages/sie_server/README.md index 391ad2a1a..7afb73ad5 100644 --- a/packages/sie_server/README.md +++ b/packages/sie_server/README.md @@ -4,7 +4,8 @@ GPU inference server for embeddings, reranking, and entity extraction. ## Features -- Multi-model serving with LRU eviction +- Multi-model serving with LRU eviction, and a host-RAM tier that wakes evicted + models back onto the GPU instead of reloading them - Token-based dynamic batching - Hot reload model configs without restart - Unified API: `encode()`, `score()`, `extract()` @@ -1228,6 +1229,71 @@ When OOM recovery is exhausted on a request, the server returns `HTTP 503 RESOURCE_EXHAUSTED` with `Retry-After`. The Python SDK auto-retries; see `packages/sie_sdk/README.md` for client-side controls. +### Host-RAM tier for evicted models + +When the worker evicts a model to free GPU memory, it keeps the model's weights +in host RAM instead of discarding them, up to a budget. The next request for +the model copies the weights back to the same GPU. It does not rebuild the +adapter, re-read the local disk cache, or fetch from the cluster cache or +Hugging Face. When the budget fills, the least recently parked model is +unloaded first. A model that is not parked loads the ordinary way. + +| Env var | Default | Effect | +|--|--|--| +| `SIE_HOST_CACHE_BYTES` | `auto` | Host RAM the tier may hold. `auto` uses a quarter of the container memory limit (cgroup v2 or v1), or of physical RAM when there is no limit, on CUDA workers, and is off on CPU and Apple silicon. A size such as `8GiB`, `8Gi` or `8000000000` sets it explicitly. `0` turns the tier off. | + +With `auto`, an L4 pool with the chart's 32Gi container limit gets 8 GiB, and +the 128Gi and 256Gi A100 and H200 pools get 32 GiB and 64 GiB. + +What parks: + +- Evictions that free memory for other work park the model: memory pressure, + making room for another load, OOM recovery, idle eviction + (`SIE_IDLE_EVICT_S`) and displacement by a multi-GPU model. A manual unload, + a config change, shutdown and a crashed engine unload the model as before. + Pinned models are never evicted, so they never park. +- Adapters that hold their PyTorch model in the worker process + (sentence-transformers, the flash encoders, cross-encoders and the other + in-process adapters) park on CUDA by copying every tensor the adapter + references to ordinary host memory and freeing the GPU copies. Parameters keep their identity, so LoRA adapters loaded through PEFT + park and wake with the base model. CUDA graphs are dropped before parking and + record again on demand. An adapter declines when it holds GPU tensors the + copy cannot rebind or quantized parameters. +- SGLang generation and embedding engines on one GPU park through SGLang's + memory saver. SIE launches these engines with `--enable-memory-saver + --enable-weights-cpu-backup` while the tier is on. Parking waits up to 10 + seconds for in-flight requests and aborts any that remain, then releases the + weights, KV cache and CUDA graph memory. The weights go to pinned host memory. + The engine process, its compiled kernels and its captured graphs stay alive, + so a wake maps memory back at the same addresses, copies the weights, and + checks that the engine generates. The memory saver in the pinned SGLang keeps + that host copy after a wake, so an engine that has parked once keeps its weights' + size charged to the budget until it is fully unloaded. A parked engine also + keeps its CUDA context and some workspace on the GPU. The worker drops it + when that GPU needs the memory and no resident model can be evicted. Set + `adapter_options.loadtime.host_offload: false` on a profile to launch its + engine without the memory saver. +- Tensor-parallel engines, TensorRT-LLM, MLX, CTranslate2, the vision + extraction engine, Candle workers and remote profiles unload as before. + +A request for a parked model waits up to 15 seconds for the wake instead of +receiving `503 MODEL_LOADING`. A queued request is held by the sidecar, which +checks readiness every second. A wake that fails falls back to an ordinary +load. A wake runs one at a time with loads, so it can wait behind another +model's cold load. + +Wake time is bounded by the host-to-GPU copy. These are estimates, not +measurements: a 1 GB encoder should wake in about 0.1 to 0.3 seconds from +ordinary host memory. An 8 GB LLM should take about 1 to 2 seconds from pinned +memory, plus the engine's one-token health check. A load from the local disk +cache takes seconds to minutes, because it constructs the model, warms it up +and, for SGLang, starts an engine and captures its CUDA graphs. + +The worker reports `sie.worker.model.host_cache.lookups` (each load, `hit` or +`miss`), `sie.worker.model.host_cache.wake.duration` and +`sie.worker.model.host_cache.memory` (host bytes per model); see +`telemetry/contract.yaml`. + ### Batching & request handling | Env var | Default | Effect | diff --git a/packages/sie_server/src/sie_server/adapters/_base_adapter.py b/packages/sie_server/src/sie_server/adapters/_base_adapter.py index 10b09f909..38aff8e1b 100644 --- a/packages/sie_server/src/sie_server/adapters/_base_adapter.py +++ b/packages/sie_server/src/sie_server/adapters/_base_adapter.py @@ -10,7 +10,7 @@ from sie_server.adapters._spec import AdapterSpec from sie_server.adapters._types import ERR_NOT_LOADED from sie_server.adapters._utils import grouped_score_pairs -from sie_server.adapters.base import ModelAdapter, ModelCapabilities, ModelDims +from sie_server.adapters.base import HostOffload, ModelAdapter, ModelCapabilities, ModelDims, released_bytes if TYPE_CHECKING: import torch @@ -36,6 +36,12 @@ class BaseAdapter(ModelAdapter): spec: ClassVar[AdapterSpec] + # Host-RAM parking (``ModelAdapter.offload_to_host``) for in-process + # weights. The generic implementation below moves every tensor reachable + # from the adapter's attributes (see ``_host_offload``). An adapter that + # keeps device state that walk cannot see sets this to False. + host_offload: ClassVar[bool] = True + # Guards the lazy creation of each instance's ``_tokenizer_lock`` (below). # A HuggingFace *fast* tokenizer wraps a Rust object behind a ``RefCell``: # calling it from two threads at once — or mutating shared state like @@ -127,6 +133,7 @@ def unload(self) -> None: setattr(self, attr, None) self._device = None + self._parked_tensors = None gc.collect() @@ -138,6 +145,99 @@ def unload(self) -> None: elif str(device) == "mps": _torch.mps.empty_cache() + # -- Host-RAM parking ------------------------------------------------------ + + def _host_offload_device(self) -> torch.device | None: + """The CUDA device this adapter's weights live on, or None if parking would free nothing. + + On CPU the weights already are in host RAM, and on Apple silicon the + accelerator shares it. + """ + device = getattr(self, "_device", None) + if device is None: + return None + import torch as _torch + + try: + resolved = _torch.device(device) + except (RuntimeError, TypeError, ValueError): + return None + return resolved if resolved.type == "cuda" else None + + def host_offload_bytes(self) -> int | None: + """Bytes of device tensors the generic walk would park, or None when it cannot park.""" + if not self.host_offload or getattr(self, "_parked_tensors", None) is not None: + return None + device = self._host_offload_device() + if device is None: + return None + if any(not callable(getattr(holder, "clear", None)) for holder in self._graph_holders()): + # A graph holder that cannot be cleared would replay against + # addresses the park frees. + return None + from sie_server.adapters import _host_offload + + parkable = _host_offload.parkable_bytes(self, device) + # The model ``memory_footprint`` measures must be among what the walk + # reached; a shortfall means weights sit where parking cannot see them + # and would stay on the device. The slack covers parameters that share + # storage, which the footprint counts once each. + if parkable is None or parkable < 0.9 * released_bytes(self.memory_footprint()): + return None + return parkable + + def offload_to_host(self) -> HostOffload: + """Copy the weights to host RAM and free their device memory.""" + device = self._host_offload_device() + if device is None or getattr(self, "_parked_tensors", None) is not None: + msg = f"{type(self).__name__} has nothing on a CUDA device to park" + raise RuntimeError(msg) + from sie_server.adapters import _host_offload + + self._drop_device_caches() + parked = _host_offload.park(self, device) + self._parked_tensors = parked + gc.collect() + _host_offload.release_cached_device_memory(device) + return HostOffload(host_bytes=parked.nbytes) + + def restore_from_host(self) -> int: + """Copy parked weights back to the device they came from.""" + parked = getattr(self, "_parked_tensors", None) + if parked is None: + msg = f"{type(self).__name__} is not parked" + raise RuntimeError(msg) + from sie_server.adapters import _host_offload + + _host_offload.unpark(parked) + self._parked_tensors = None + return 0 + + def _drop_device_caches(self) -> None: + """Drop device state recorded against the current weight addresses. + + A CUDA graph replays kernels against the exact addresses it recorded, + and parking moves every weight. Graph runners (``self._graphs``) are + cleared, and anything else the adapter releases as optional memory + goes with them; graphs record again on demand after the model wakes. + """ + for holder in self._graph_holders(): + holder.clear() + if self.has_releasable_memory(): + self.release_optional_memory() + + def _graph_holders(self) -> list[Any]: + """Attributes holding recorded CUDA graphs: runners (``_graphs``) and dicts of them. + + Found by name, as every graph runner in this package is held. Scalars + that merely describe graphs (``_cuda_graphs = "off"``) are skipped. + """ + return [ + value + for name, value in vars(self).items() + if "graph" in name.lower() and value is not None and not isinstance(value, (bool, int, float, str)) + ] + def get_preprocessor(self) -> Any: """Return ``CharCountPreprocessor`` for cost estimation.""" from sie_server.core.preprocessor import CharCountPreprocessor diff --git a/packages/sie_server/src/sie_server/adapters/_host_offload.py b/packages/sie_server/src/sie_server/adapters/_host_offload.py new file mode 100644 index 000000000..2fe702fa0 --- /dev/null +++ b/packages/sie_server/src/sie_server/adapters/_host_offload.py @@ -0,0 +1,244 @@ +"""Move an in-process adapter's tensors to host RAM and back. + +This is the generic half of host-RAM parking for adapters that hold their +weights in this process (``BaseAdapter`` subclasses). It finds the tensors an +adapter references on its device and copies them to host memory, freeing the +device copies, then copies them back on wake. + +What it finds, starting from the adapter's instance attributes: + +- ``torch.nn.Module`` objects: every parameter, buffer and plain tensor + attribute of every submodule. Parameters keep their identity (only + ``.data`` moves), so references to the parameter objects held elsewhere, + such as a PEFT wrapper around the same modules, follow along; +- loose tensors, and tensors or modules inside dicts and lists, a few levels + deep; +- wrapper objects that hold a module as their ``model`` attribute + (transformers pipelines, FlagEmbedding models). + +An adapter is not parkable when a device tensor sits where it cannot be +rebound (inside a tuple), when a parameter is a quantized subclass the copy +would not preserve, when accelerate dispatch hooks manage placement, or when +nothing on the device was found at all. Device state outside this reach, such +as a tensor captured in a closure, stays where it is: it remains valid, but +its memory is not freed while the model is parked. Adapters with such state +override the parking hooks or set ``host_offload = False``. + +Host copies are ordinary (pageable) CPU tensors. PyTorch's pinned-memory +allocator rounds every block up to a power of two and keeps freed blocks +cached, which would make the host tier's budget fiction; pageable copies are +accounted exactly and released to the OS when dropped. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, field +from typing import Any + +import torch + +logger = logging.getLogger(__name__) + +# How deep to follow dicts and lists (and wrappers) below an adapter attribute. +_MAX_CONTAINER_DEPTH = 4 + + +@dataclass(frozen=True) +class _Slot: + """Where a loose tensor lives, so it can be rebound.""" + + owner: Any + key: Any + attribute: bool + + def get(self) -> object: + if self.attribute: + return getattr(self.owner, self.key, None) + try: + return self.owner[self.key] + except (IndexError, KeyError): + return None + + def set(self, value: torch.Tensor) -> None: + if self.attribute: + setattr(self.owner, self.key, value) + else: + self.owner[self.key] = value + + +@dataclass +class _Reach: + modules: list[torch.nn.Module] = field(default_factory=list) + loose: list[tuple[_Slot, torch.Tensor]] = field(default_factory=list) + unmovable_bytes: int = 0 + + +@dataclass +class ParkedTensors: + """An adapter's tensors parked in host RAM, and where they came from.""" + + modules: list[torch.nn.Module] + loose: list[tuple[_Slot, torch.Tensor]] + # Host copies keyed by their storage address, with the device each came + # from. Holding the copies keeps the addresses unique until restore. + origins: dict[int, torch.device] + copies: list[torch.Tensor] + + @property + def nbytes(self) -> int: + return sum(copy.nbytes for copy in self.copies) + + +def _on(device: torch.device, tensor: torch.Tensor) -> bool: + """Whether ``tensor`` holds memory on ``device`` (empty tensors hold none).""" + if tensor.device.type != device.type or tensor.untyped_storage().nbytes() == 0: + return False + return device.index is None or tensor.device.index == device.index + + +def _collect(adapter: object, device: torch.device) -> _Reach: + reach = _Reach() + seen_modules: set[int] = set() + seen: set[int] = set() + + def add_module(module: torch.nn.Module) -> None: + if id(module) not in seen_modules: + seen_modules.add(id(module)) + reach.modules.append(module) + + def visit(value: object, slot: _Slot | None, depth: int) -> None: + if isinstance(value, torch.nn.Module): + add_module(value) + return + if isinstance(value, torch.Tensor): + if _on(device, value): + if slot is None: + reach.unmovable_bytes += value.nbytes + else: + reach.loose.append((slot, value)) + return + if depth >= _MAX_CONTAINER_DEPTH or id(value) in seen: + return + seen.add(id(value)) + if isinstance(value, dict): + for key, item in list(value.items()): + visit(item, _Slot(value, key, attribute=False), depth + 1) + elif isinstance(value, list): + for index, item in enumerate(value): + visit(item, _Slot(value, index, attribute=False), depth + 1) + elif isinstance(value, tuple): + for item in value: + visit(item, None, depth + 1) + else: + # A wrapper's own ``model`` attribute, read from its instance + # dictionary so no property runs. + wrapped = getattr(value, "__dict__", {}).get("model") + if isinstance(wrapped, torch.nn.Module): + add_module(wrapped) + + for name, value in list(vars(adapter).items()): + visit(value, _Slot(adapter, name, attribute=True), 0) + return reach + + +def _module_tensors(module: torch.nn.Module) -> list[torch.Tensor]: + tensors: list[torch.Tensor] = [] + for submodule in module.modules(): + tensors.extend(p for p in submodule._parameters.values() if p is not None) + tensors.extend(b for b in submodule._buffers.values() if b is not None) + tensors.extend(v for v in vars(submodule).values() if isinstance(v, torch.Tensor)) + return tensors + + +def parkable_bytes(adapter: object, device: torch.device) -> int | None: + """Device bytes ``park`` would move to host RAM, or None when parking cannot be trusted.""" + reach = _collect(adapter, device) + if reach.unmovable_bytes: + return None + storages: dict[int, int] = {} + for module in reach.modules: + for submodule in module.modules(): + if getattr(submodule, "_hf_hook", None) is not None: + return None + for tensor in _module_tensors(module): + if type(tensor) not in (torch.Tensor, torch.nn.Parameter): + return None + if _on(device, tensor): + storage = tensor.untyped_storage() + storages[storage.data_ptr()] = storage.nbytes() + for _, tensor in reach.loose: + storage = tensor.untyped_storage() + storages[storage.data_ptr()] = storage.nbytes() + total = sum(storages.values()) + return total or None + + +def park(adapter: object, device: torch.device) -> ParkedTensors: + """Copy the adapter's device tensors to host RAM and rebind them there.""" + reach = _collect(adapter, device) + origins: dict[int, torch.device] = {} + copies: list[torch.Tensor] = [] + + def to_host(tensor: torch.Tensor) -> torch.Tensor: + copy = tensor.detach().to("cpu", copy=True) + origins[copy.untyped_storage().data_ptr()] = tensor.device + copies.append(copy) + return copy + + def still_on_device(tensor: torch.Tensor) -> bool: + # A parameter shared by two modules, or a loose reference to one, + # has already been moved once by the time it is seen again. + return _on(device, tensor) and tensor.untyped_storage().data_ptr() not in origins + + with torch.no_grad(): + for module in reach.modules: + _move_module(module, still_on_device, to_host) + for slot, tensor in reach.loose: + if still_on_device(tensor): + slot.set(to_host(tensor)) + return ParkedTensors(modules=reach.modules, loose=reach.loose, origins=origins, copies=copies) + + +def unpark(parked: ParkedTensors) -> None: + """Copy parked tensors back to the devices they came from. + + Safe to call again after a failure part way through (an out-of-memory + error): tensors already back on their device are skipped. + """ + + def parked_here(tensor: torch.Tensor) -> bool: + storage = tensor.untyped_storage() + return tensor.device.type == "cpu" and storage.nbytes() > 0 and storage.data_ptr() in parked.origins + + def to_device(tensor: torch.Tensor) -> torch.Tensor: + return tensor.to(parked.origins[tensor.untyped_storage().data_ptr()], copy=True) + + with torch.no_grad(): + for module in parked.modules: + _move_module(module, parked_here, to_device) + for slot, _ in parked.loose: + current = slot.get() + if isinstance(current, torch.Tensor) and parked_here(current): + slot.set(to_device(current)) + for device in {origin for origin in parked.origins.values() if origin.type == "cuda"}: + torch.cuda.synchronize(device) + + +def _move_module(module: torch.nn.Module, select: Any, move: Any) -> None: + for submodule in module.modules(): + for key, param in submodule._parameters.items(): + if param is not None and select(param): + param.data = move(param.data) + for key, buffer in submodule._buffers.items(): + if buffer is not None and select(buffer): + submodule._buffers[key] = move(buffer) + for key, value in list(vars(submodule).items()): + if isinstance(value, torch.Tensor) and select(value): + setattr(submodule, key, move(value)) + + +def release_cached_device_memory(device: torch.device) -> None: + """Return freed blocks to the driver so device usage reflects the park.""" + if device.type == "cuda": + torch.cuda.empty_cache() diff --git a/packages/sie_server/src/sie_server/adapters/base.py b/packages/sie_server/src/sie_server/adapters/base.py index b37dff0d9..c60639648 100644 --- a/packages/sie_server/src/sie_server/adapters/base.py +++ b/packages/sie_server/src/sie_server/adapters/base.py @@ -2,6 +2,7 @@ import logging from abc import ABC, abstractmethod +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal from pydantic import BaseModel, ConfigDict, Field @@ -26,6 +27,20 @@ def released_bytes(value: object) -> int: return value +@dataclass(frozen=True, slots=True) +class HostOffload: + """What ``ModelAdapter.offload_to_host`` left behind. + + Attributes: + host_bytes: Host RAM now holding the model's weights. + retains_device_memory: Whether the parked model still holds some device + memory, such as an engine process's CUDA context and workspaces. + """ + + host_bytes: int + retains_device_memory: bool = False + + class ModelCapabilities(BaseModel): """Capabilities supported by a model adapter.""" @@ -79,6 +94,23 @@ class ModelAdapter(ABC): ``has_releasable_memory()`` is the cheap check callers use first, so an adapter with nothing cached is never locked just to free 0 bytes. + Host-RAM Parking: + When the registry evicts a model to free device memory and the host + tier has room, it parks the model instead of unloading it: + ``offload_to_host()`` moves the weights to host RAM and frees their + device memory, and ``restore_from_host()`` later moves them back to + the device the adapter was loaded on. ``host_offload_bytes()`` is the + cheap check the registry asks first; ``None`` (the default) means the + adapter cannot park and is unloaded as before. + + The registry calls ``offload_to_host()`` only after the model's worker + has drained and stopped, so no forward pass is in flight, and runs it + off the event loop. It calls ``restore_from_host()`` before creating a + new worker. If either raises, the registry calls ``unload()``, which + must then release whatever the adapter still holds, on the device and + in host RAM. A parked adapter may also be unloaded at any time to make + room in the host tier. + Main Thread Requirement: Some adapters may use parent-process signal handlers that only work in the main thread. Set requires_main_thread = True for these adapters. @@ -274,6 +306,40 @@ def release_optional_memory(self) -> int: """ return 0 + def host_offload_bytes(self) -> int | None: + """Host bytes ``offload_to_host`` would hold, or ``None`` if this adapter cannot park. + + Cheap and side-effect free: the registry asks before it chooses between + parking the model in host RAM and unloading it, and checks the answer + against the host tier's budget. The default is ``None``. + """ + return None + + def offload_to_host(self) -> HostOffload: + """Move the loaded weights to host RAM and free their device memory. + + Only called when ``host_offload_bytes()`` returned a size. On success + the adapter serves nothing until ``restore_from_host()``. + + Raises: + NotImplementedError: If the adapter cannot park. + """ + msg = f"{self.__class__.__name__} cannot park its weights in host RAM" + raise NotImplementedError(msg) + + def restore_from_host(self) -> int: + """Move parked weights back to the device the adapter was loaded on. + + Returns: + Host bytes the adapter still holds after restoring (for example an + engine that keeps its host copy for the next park), or 0. + + Raises: + NotImplementedError: If the adapter cannot park. + """ + msg = f"{self.__class__.__name__} cannot park its weights in host RAM" + raise NotImplementedError(msg) + def engine_exit_code(self) -> int | None: """Return the exit code of this adapter's engine process once it has exited. diff --git a/packages/sie_server/src/sie_server/adapters/fake/adapter.py b/packages/sie_server/src/sie_server/adapters/fake/adapter.py index a0dbecb8a..e9ef27670 100644 --- a/packages/sie_server/src/sie_server/adapters/fake/adapter.py +++ b/packages/sie_server/src/sie_server/adapters/fake/adapter.py @@ -56,6 +56,7 @@ ) from sie_server.adapters._spec import AdapterSpec from sie_server.adapters._types import ERR_NOT_LOADED +from sie_server.adapters.base import HostOffload from sie_server.core.inference_output import ( EncodeOutput, ExtractItemError, @@ -487,6 +488,7 @@ def __init__( synthetic_usage: bool = False, faults: dict[str, Any] | None = None, fault_key: str | None = None, + host_offload: bool = True, **kwargs: Any, ) -> None: _ = (model_name_or_path, compute_precision, kwargs) @@ -502,6 +504,8 @@ def __init__( self._fault_runtime = _FaultRuntime(_resolve_faults(faults, fault_key)) self._loaded = False self._device: str | None = None + self._host_offload_enabled = host_offload + self._parked = False # -- Lifecycle ----------------------------------------------------------- @@ -522,6 +526,28 @@ def load_required_memory_bytes(self, *, device_type: str, device_total_bytes: in _ = (device_type, device_total_bytes) return self._memory_footprint_bytes + # -- Host-RAM parking (simulated: the declared footprint moves) ---------- + + def host_offload_bytes(self) -> int | None: + if not self._host_offload_enabled or not self._loaded: + return None + return self._memory_footprint_bytes + + def offload_to_host(self) -> HostOffload: + if not self._loaded: + raise RuntimeError(ERR_NOT_LOADED) + self._loaded = False + self._parked = True + return HostOffload(host_bytes=self._memory_footprint_bytes) + + def restore_from_host(self) -> int: + if not self._parked: + msg = "sie-fake model is not parked" + raise RuntimeError(msg) + self._parked = False + self._loaded = True + return 0 + # -- Surfaces -------------------------------------------------------------- def encode( diff --git a/packages/sie_server/src/sie_server/adapters/sglang/_parking.py b/packages/sie_server/src/sie_server/adapters/sglang/_parking.py new file mode 100644 index 000000000..4e8607dee --- /dev/null +++ b/packages/sie_server/src/sie_server/adapters/sglang/_parking.py @@ -0,0 +1,82 @@ +"""Host-RAM parking for SGLang engines. + +Shared by the generation and embedding adapters. A parked engine keeps its +process, compiled kernels and captured CUDA graphs; its weights move to pinned +host memory and its KV cache and graph memory are released (see +``_server.park_engine``). Waking copies the weights back into the same device +addresses, so the engine serves again without a restart. + +Parking needs the engine's memory saver, which is a launch flag, so an engine +started while the host tier was off (``SIE_HOST_CACHE_BYTES=0``), started for +a multi-device group, or whose profile sets ``host_offload: false`` never +parks and is unloaded on eviction as before. +""" + +from __future__ import annotations + +import subprocess + +from sie_server.adapters.base import HostOffload +from sie_server.adapters.sglang import _server +from sie_server.core.host_cache import host_cache_configured + + +class SGLangParkingMixin: + """``ModelAdapter`` parking hooks for an adapter that runs one SGLang child process.""" + + _process: subprocess.Popen[bytes] | None + _server_url: str | None + _model_name_or_path: str + _revision: str | None + _tensor_parallel_size: int + + # Profile option (``adapter_options.loadtime.host_offload``); the adapter + # constructor overrides it. + _host_offload_requested: bool = True + # Whether this engine was launched able to park. + _memory_saver: bool = False + _parked: bool = False + # Host bytes the engine holds once it has parked; its backup buffer is kept + # and reused after a wake. + _host_copy_bytes: int = 0 + + def _memory_saver_launch_args(self) -> list[str]: + """Launch flags that let this engine park later; empty when it never will.""" + self._memory_saver = ( + self._host_offload_requested and self._tensor_parallel_size == 1 and host_cache_configured() + ) + return list(_server.MEMORY_SAVER_LAUNCH_ARGS) if self._memory_saver else [] + + def host_offload_bytes(self) -> int | None: + if ( + not self._memory_saver + or self._parked + or self._server_url is None + or _server.engine_exit_code(self._process) is not None + ): + return None + if self._host_copy_bytes: + return self._host_copy_bytes + return _server.checkpoint_weight_bytes(self._model_name_or_path, self._revision) + + def offload_to_host(self) -> HostOffload: + host_bytes = self.host_offload_bytes() + if host_bytes is None or self._server_url is None: + msg = f"SGLang engine for {self._model_name_or_path!r} cannot park" + raise RuntimeError(msg) + # Refuse new requests before the engine is drained, not after. + self._parked = True + _server.park_engine(self._server_url) + self._host_copy_bytes = host_bytes + # The process keeps its CUDA context and the workspaces allocated + # outside the memory saver's regions. + return HostOffload(host_bytes=host_bytes, retains_device_memory=True) + + def restore_from_host(self) -> int: + if not self._parked or self._server_url is None: + msg = f"SGLang engine for {self._model_name_or_path!r} is not parked" + raise RuntimeError(msg) + _server.raise_if_engine_exited(self._process, self._model_name_or_path) + _server.wake_engine(self._server_url) + self._parked = False + return self._host_copy_bytes diff --git a/packages/sie_server/src/sie_server/adapters/sglang/_server.py b/packages/sie_server/src/sie_server/adapters/sglang/_server.py index 6c98bfa58..c0b559ce9 100644 --- a/packages/sie_server/src/sie_server/adapters/sglang/_server.py +++ b/packages/sie_server/src/sie_server/adapters/sglang/_server.py @@ -673,3 +673,112 @@ def terminate_process(process: subprocess.Popen[bytes] | None) -> None: process.wait(timeout=5) except (ProcessLookupError, subprocess.TimeoutExpired): pass + + +# -- Host-RAM parking --------------------------------------------------------- +# +# With the memory saver on, the engine allocates its weights, KV cache and CUDA +# graph pools through torch-memory-saver, which can release their physical +# device memory while keeping their virtual addresses. ``release_memory_occupation`` +# copies the weights to pinned host memory (``--enable-weights-cpu-backup``), +# then releases all three; ``resume_memory_occupation`` maps fresh device memory +# at the same addresses and copies the weights back. Captured CUDA graphs stay +# valid because no address changes. The engine process, its compiled kernels and +# its graphs survive, so a wake skips everything a cold start pays for except the +# copy itself. torch-memory-saver 0.0.9 keeps the host buffer after a resume and +# reuses it for the next release. + +MEMORY_SAVER_LAUNCH_ARGS = ("--enable-memory-saver", "--enable-weights-cpu-backup") +# The engine refuses to release while it runs requests; flush_cache waits up to +# this long for it to go idle, and empties the allocator cache when it does. +PARK_IDLE_TIMEOUT_S = 10.0 +# Release copies every weight to host memory; resume copies it back. +PARK_RELEASE_TIMEOUT_S = 300.0 +WAKE_RESUME_TIMEOUT_S = 300.0 +# After resume, /health generates one token: proof the weights are back. +WAKE_HEALTH_TIMEOUT_S = 60.0 +_CONTROL_TIMEOUT_S = 30.0 + + +def checkpoint_weight_bytes(model_name_or_path: str, revision: str | None) -> int | None: + """Bytes of a model's cached checkpoint files, or None when they cannot be found. + + Sized from the safetensors (else ``.bin``) files beside the cached + ``config.json``: the weights the engine backs up to host memory when it + parks, at checkpoint precision. A checkpoint quantized at load time is + over-counted, which only makes the budget conservative. + """ + directory = Path(model_name_or_path) + if not directory.is_dir(): + try: + from huggingface_hub import try_to_load_from_cache + from sie_sdk.cache import get_cache_config + + cached = try_to_load_from_cache( + model_name_or_path, + "config.json", + revision=revision, + cache_dir=get_cache_config().local_cache, + ) + except (ImportError, OSError, ValueError): + return None + if not isinstance(cached, str): + return None + directory = Path(cached).parent + for pattern in ("*.safetensors", "*.bin"): + try: + sizes = [path.stat().st_size for path in directory.glob(pattern)] + except OSError: + return None + if sizes: + return sum(sizes) + return None + + +def _flush_when_idle(server_url: str) -> bool: + response = requests.post( + f"{server_url}/flush_cache", + params={"timeout": PARK_IDLE_TIMEOUT_S}, + timeout=PARK_IDLE_TIMEOUT_S + _CONTROL_TIMEOUT_S, + ) + return response.status_code == 200 + + +def _raise_for_engine_error(response: requests.Response, operation: str) -> None: + if response.status_code != 200: + msg = f"SGLang {operation} failed with HTTP {response.status_code}: {response.text[:500]}" + raise RuntimeError(msg) + + +def park_engine(server_url: str) -> None: + """Release an idle engine's device memory, keeping its weights in host RAM. + + Waits for in-flight requests first. Requests still running after the wait + are aborted, and the wait repeats; an engine that is still busy is not + released. + + Raises: + RuntimeError: When the engine stays busy or refuses the release. The + caller then unloads the model instead. + """ + if not _flush_when_idle(server_url): + logger.warning("SGLang engine at %s is still busy; aborting its requests to park it", server_url) + abort = requests.post(f"{server_url}/abort_request", json={"abort_all": True}, timeout=_CONTROL_TIMEOUT_S) + _raise_for_engine_error(abort, "abort_request") + if not _flush_when_idle(server_url): + msg = f"SGLang engine at {server_url} did not become idle; it cannot be parked" + raise RuntimeError(msg) + response = requests.post(f"{server_url}/release_memory_occupation", json={}, timeout=PARK_RELEASE_TIMEOUT_S) + _raise_for_engine_error(response, "release_memory_occupation") + + +def wake_engine(server_url: str) -> None: + """Bring a parked engine's weights, KV cache and graphs back onto the device. + + Raises: + RuntimeError: When the engine refuses the resume or cannot serve after it. + """ + response = requests.post(f"{server_url}/resume_memory_occupation", json={}, timeout=WAKE_RESUME_TIMEOUT_S) + _raise_for_engine_error(response, "resume_memory_occupation") + health = requests.get(f"{server_url}/health", timeout=WAKE_HEALTH_TIMEOUT_S) + _raise_for_engine_error(health, "health check after resume") diff --git a/packages/sie_server/src/sie_server/adapters/sglang/embedding.py b/packages/sie_server/src/sie_server/adapters/sglang/embedding.py index e28e21ce8..e0bde41b1 100644 --- a/packages/sie_server/src/sie_server/adapters/sglang/embedding.py +++ b/packages/sie_server/src/sie_server/adapters/sglang/embedding.py @@ -37,6 +37,7 @@ from sie_server.adapters.sglang import _server from sie_server.adapters.sglang._lora import SGLangLoraManager, SGLangLoraMixin, leased_encode from sie_server.adapters.sglang._memory import SGLangMemoryBudget +from sie_server.adapters.sglang._parking import SGLangParkingMixin from sie_server.core.inference_output import EncodeOutput from sie_server.core.tokenizer import load_tokenizer @@ -84,7 +85,7 @@ def _resolve_embed_concurrency(configured: int | None) -> int: return _DEFAULT_EMBED_CONCURRENCY -class SGLangEmbeddingAdapter(SGLangLoraMixin, BaseAdapter): +class SGLangEmbeddingAdapter(SGLangParkingMixin, SGLangLoraMixin, BaseAdapter): """Adapter for LLM embedding models using SGLang HTTP server backend. SGLang pre-allocates GPU memory for the KV cache, providing stable memory @@ -117,7 +118,7 @@ class SGLangEmbeddingAdapter(SGLangLoraMixin, BaseAdapter): manages_own_load_timeout: bool = True def _check_loaded(self) -> None: - if self._server_url is None: + if self._server_url is None or self._parked: raise RuntimeError(ERR_NOT_LOADED) # Without this the adapter never notices its engine exiting, and every # later request opens a connection to a port nobody is listening on, @@ -161,6 +162,7 @@ def __init__( request_read_timeout_s: float | None = None, watchdog_timeout_s: float | None = None, nccl_port: int | None = None, + host_offload: bool = True, **kwargs: Any, # Accept extra args from loader (e.g., pooling) ) -> None: r"""Initialize the adapter. @@ -202,9 +204,15 @@ def __init__( max_concurrent_dispatch: Batches the worker may have in flight to SGLang at once (see ``max_concurrent_dispatch()``). 1 keeps the worker's one-batch-at-a-time dispatch. + host_offload: Launch the engine able to park its weights in host + RAM when the registry evicts it (see ``_parking``). False keeps + the previous launch. **kwargs: Additional arguments (ignored, for compatibility). """ _ = kwargs # Unused, but accepted for loader compatibility + if not isinstance(host_offload, bool): + raise ValueError(f"host_offload must be true or false, got {host_offload!r}") + self._host_offload_requested = host_offload self._model_name_or_path = str(model_name_or_path) self._normalize = normalize self._max_seq_length = max_seq_length @@ -384,6 +392,7 @@ def load(self, device: str) -> None: # /tmp in open_output_log, a failed exec, or a cancelled load included. try: cmd.extend(self._lora.launch_args()) + cmd.extend(self._memory_saver_launch_args()) cmd.extend(self._memory_budget.launch_args(device_indices[0])) # Only a group rendezvouses, so width one reserves nothing and its # launch argv is unchanged. Left to the engine it picks a random diff --git a/packages/sie_server/src/sie_server/adapters/sglang/generation.py b/packages/sie_server/src/sie_server/adapters/sglang/generation.py index 9ded555f4..22243c71b 100644 --- a/packages/sie_server/src/sie_server/adapters/sglang/generation.py +++ b/packages/sie_server/src/sie_server/adapters/sglang/generation.py @@ -44,6 +44,7 @@ FinishReason, GenerationAdapter, GenerationChunk, + GenerationDrainingError, GenerationError, GenerationInvalidRequestError, GenerationResult, @@ -60,6 +61,7 @@ from sie_server.adapters.sglang._compat.sie_thinking_budget import PROCESSOR_REFERENCE as THINKING_BUDGET_PROCESSOR from sie_server.adapters.sglang._lora import SGLangLoraManager, SGLangLoraMixin, leased_generation from sie_server.adapters.sglang._memory import SGLangMemoryBudget +from sie_server.adapters.sglang._parking import SGLangParkingMixin from sie_server.observability.generation_diagnostics import GenerationStreamTimer from sie_server.types.grammar import ( OUTLINES_JSON_SCHEMA_TYPE_DIAGNOSTIC, @@ -453,7 +455,7 @@ def _translate_to_mlx_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: return {k: v for k, v in kwargs.items() if k in _MLX_PASSTHROUGH_KWARGS} -class SGLangGenerationAdapter(SGLangLoraMixin, GenerationAdapter): +class SGLangGenerationAdapter(SGLangParkingMixin, SGLangLoraMixin, GenerationAdapter): """Streaming SGLang generation adapter. Lifecycle is identical to :class:`SGLangEmbeddingAdapter`: one subprocess @@ -574,9 +576,15 @@ def __init__( # it, then follows as usual. None leaves reasoning unbounded and the # launch unchanged. Needs a reasoning parser with known control tokens. thinking_budget: int | None = None, + # Launch the engine able to park its weights in host RAM when the + # registry evicts it (see ``_parking``). Off keeps the previous launch. + host_offload: bool = True, **kwargs: Any, # accept extra args from loader for compatibility ) -> None: _ = kwargs + if not isinstance(host_offload, bool): + raise ValueError(f"host_offload must be true or false, got {host_offload!r}") + self._host_offload_requested = host_offload self._json_number_max_digits = _validate_json_number_max_digits(json_number_max_digits, grammar_backend) self._thinking_budget = _validate_thinking_budget(thinking_budget, reasoning_parser) self._thinking_budget_tokens: dict[str, int | None] | None = None @@ -980,6 +988,7 @@ def _start_server(self, *, port: int, device_indices: list[int], server_url: str # ``sampling_params["lora_path"]``. Only emitted when the profile # declared adapters, so non-LoRA deployments are byte-identical. cmd.extend(self._lora.launch_args()) + cmd.extend(self._memory_saver_launch_args()) cmd.extend(self._extra_launch_args) self._output_file = _server.open_output_log() @@ -1462,6 +1471,8 @@ def load_required_memory_bytes(self, *, device_type: str, device_total_bytes: in def _check_loaded(self) -> None: if self._server_url is None: raise RuntimeError(ERR_NOT_LOADED) + if self._parked: + raise GenerationDrainingError("The model is parked in host RAM; retry while it wakes") self._check_engine_alive() def _check_engine_alive(self) -> None: diff --git a/packages/sie_server/src/sie_server/api/helpers.py b/packages/sie_server/src/sie_server/api/helpers.py index b64548b30..98b318126 100644 --- a/packages/sie_server/src/sie_server/api/helpers.py +++ b/packages/sie_server/src/sie_server/api/helpers.py @@ -460,6 +460,11 @@ async def ensure_loaded(self, device: str) -> None: # configuration. self.check_not_failed() + # A model parked in host RAM wakes in seconds: wait for it rather than + # asking the client to come back. + if await self.registry.wake_parked(self.model, device): + return + logger.info("Starting background load for model %s on device %s", self.model, device) await self.registry.start_load_async(self.model, device=device) diff --git a/packages/sie_server/src/sie_server/config/engine.py b/packages/sie_server/src/sie_server/config/engine.py index ac3aede46..615750bd6 100644 --- a/packages/sie_server/src/sie_server/config/engine.py +++ b/packages/sie_server/src/sie_server/config/engine.py @@ -8,9 +8,11 @@ from pathlib import Path from typing import TYPE_CHECKING, Annotated, Literal -from pydantic import BaseModel, Field, model_validator +from pydantic import BaseModel, Field, field_validator, model_validator from pydantic_settings import BaseSettings, SettingsConfigDict +from sie_server.core.host_cache import parse_host_cache_setting + if TYPE_CHECKING: from sie_server.core.oom import OomRecoveryConfig @@ -337,6 +339,19 @@ class EngineConfig(BaseSettings): ), ] = None + # Host-RAM tier for evicted models + host_cache_bytes: Annotated[ + int | str, + Field( + description=( + "Host RAM that evicted models may keep their weights in, so the next request wakes " + "them onto the GPU instead of loading from disk. 'auto' (default) uses a quarter of " + "the container memory limit (or of physical RAM) on CUDA workers and is off " + "elsewhere; a size such as 8GiB sets it explicitly; 0 disables the tier." + ), + ), + ] = "auto" + # Disk cache disk_cache_enabled: Annotated[ bool, @@ -405,6 +420,13 @@ class EngineConfig(BaseSettings): Field(description="Directory containing model configs"), ] = Path("./models") + @field_validator("host_cache_bytes") + @classmethod + def _validate_host_cache_bytes(cls, value: int | str) -> int | str: + """Fail at startup on a malformed size rather than guessing a budget.""" + parse_host_cache_setting(value) + return value + @model_validator(mode="after") def _apply_oom_kill_switch(self) -> "EngineConfig": """Honour ``SIE_DISABLE_OOM_RECOVERY=1`` as a top-level kill switch. diff --git a/packages/sie_server/src/sie_server/core/host_cache.py b/packages/sie_server/src/sie_server/core/host_cache.py new file mode 100644 index 000000000..e14db1271 --- /dev/null +++ b/packages/sie_server/src/sie_server/core/host_cache.py @@ -0,0 +1,302 @@ +"""Host-RAM tier for evicted models. + +When the registry evicts a model to free accelerator memory, it can park the +model's weights in host RAM instead of dropping them. The next request for the +model copies the weights back to the accelerator, which avoids rebuilding the +adapter and re-reading weights from the local disk cache, the cluster cache or +Hugging Face. + +This module owns the bookkeeping and nothing else: the budget, the least +recently parked order, and which parked models still hold accelerator memory. +The registry decides when to park, wake or drop a model, and the adapters move +their own weights (``ModelAdapter.offload_to_host`` / ``restore_from_host``). + +Budget accounting is in host bytes the tier is responsible for: + +- a parked model is charged the host bytes its adapter reported when it parked; +- an adapter that keeps its host copy after waking (SGLang's memory saver reuses + the buffer for the next park) stays charged while it is resident, until the + model is fully unloaded. Those charges cannot be dropped to make room; only + parked models can. +""" + +from __future__ import annotations + +import logging +import math +import os +import time +from collections import OrderedDict +from collections.abc import Iterable +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +logger = logging.getLogger(__name__) + +HOST_CACHE_ENV = "SIE_HOST_CACHE_BYTES" +AUTO = "auto" +# ``auto`` budgets this share of the container memory limit (or of physical RAM +# when the container has none). The rest stays for the worker process, engine +# subprocesses, the page cache and /dev/shm. +AUTO_FRACTION = 0.25 + +_CGROUP_V2_LIMIT = Path("/sys/fs/cgroup/memory.max") +_CGROUP_V1_LIMIT = Path("/sys/fs/cgroup/memory/memory.limit_in_bytes") +# cgroup v1 reports "no limit" as a huge page-aligned number. +_CGROUP_V1_UNLIMITED = 1 << 60 + +_SIZE_SUFFIXES: dict[str, int] = { + "KIB": 1024, + "MIB": 1024**2, + "GIB": 1024**3, + "TIB": 1024**4, + "KB": 1000, + "MB": 1000**2, + "GB": 1000**3, + "TB": 1000**4, + "KI": 1024, + "MI": 1024**2, + "GI": 1024**3, + "TI": 1024**4, +} + + +def parse_host_cache_setting(value: object) -> int | None: + """Parse ``SIE_HOST_CACHE_BYTES``: ``None`` for ``auto``, else a byte count. + + Accepts ``auto``, a plain integer byte count, or a number with a binary + (``KiB``/``MiB``/``GiB``/``TiB``, also Kubernetes ``Ki``/``Mi``/``Gi``/``Ti``) + or decimal (``KB``/``MB``/``GB``/``TB``) suffix. ``0`` disables the tier. + + Raises: + ValueError: For anything else, so a malformed budget fails at startup + instead of silently disabling or over-committing the tier. + """ + if isinstance(value, bool): + raise ValueError(f"invalid {HOST_CACHE_ENV} value {value!r}") + if isinstance(value, int): + if value < 0: + raise ValueError(f"{HOST_CACHE_ENV} must not be negative, got {value}") + return value + if not isinstance(value, str): + raise ValueError(f"invalid {HOST_CACHE_ENV} value {value!r}") + text = value.strip() + if text.lower() == AUTO: + return None + number, multiplier = text, 1 + upper = text.upper() + for suffix, factor in _SIZE_SUFFIXES.items(): + if upper.endswith(suffix): + number, multiplier = text[: -len(suffix)].strip(), factor + break + try: + parsed = float(number) + except ValueError: + parsed = math.nan + if not math.isfinite(parsed) or parsed < 0 or (multiplier == 1 and not number.isdigit()): + msg = f"invalid {HOST_CACHE_ENV} value {value!r}: expected 'auto', bytes, or a size such as 8GiB" + raise ValueError(msg) + return int(parsed * multiplier) + + +def host_cache_configured() -> bool: + """Whether the process-level setting leaves the tier on (anything but ``0``). + + Engines that must decide at launch whether they can park later (SGLang's + memory saver) read this. A malformed value counts as off here; the registry + rejects it at startup. + """ + try: + return parse_host_cache_setting(os.environ.get(HOST_CACHE_ENV, AUTO)) != 0 + except ValueError: + return False + + +def _read_int(path: Path) -> int | None: + try: + raw = path.read_text(encoding="ascii").strip() + except (OSError, UnicodeDecodeError): + return None + if not raw or raw == "max": + return None + try: + return int(raw) + except ValueError: + return None + + +def physical_memory_bytes() -> int | None: + """Physical RAM of the machine, or None when the platform does not say.""" + try: + pages = os.sysconf("SC_PHYS_PAGES") + page_size = os.sysconf("SC_PAGE_SIZE") + except (AttributeError, OSError, ValueError): + return None + if pages <= 0 or page_size <= 0: + return None + return pages * page_size + + +def container_memory_limit_bytes( + *, + cgroup_v2: Path = _CGROUP_V2_LIMIT, + cgroup_v1: Path = _CGROUP_V1_LIMIT, +) -> int | None: + """The memory limit of this container (cgroup v2, then v1), or physical RAM. + + The smaller of the cgroup limit and physical RAM wins: a limit larger than + the machine is not memory the process can use. + """ + physical = physical_memory_bytes() + limit = _read_int(cgroup_v2) + if limit is None: + limit = _read_int(cgroup_v1) + if limit is not None and limit >= _CGROUP_V1_UNLIMITED: + limit = None + if limit is None or limit <= 0: + return physical + return min(limit, physical) if physical else limit + + +def resolve_host_cache_budget(setting: object, devices: Iterable[str], *, memory_limit: int | None = None) -> int: + """Resolve the configured setting to a budget in bytes; 0 disables the tier. + + ``auto`` enables the tier only for CUDA workers, at ``AUTO_FRACTION`` of the + container memory limit. On CPU there is nothing to move, and on Apple + silicon the accelerator shares host memory, so parking would free nothing. + An explicit size applies to any device, so tests and fake engines can + exercise the tier on CPU; real adapters still decline to park where it would + not free accelerator memory. + """ + explicit = parse_host_cache_setting(setting) + if explicit is not None: + return explicit + if not any(device.strip().lower().startswith("cuda") for device in devices): + return 0 + limit = memory_limit if memory_limit is not None else container_memory_limit_bytes() + if not limit: + return 0 + return int(limit * AUTO_FRACTION) + + +@dataclass +class ParkedModel: + """One model whose weights live in host RAM. + + Attributes: + name: Registry model id. + loaded: The registry's ``LoadedModel`` for it. Its adapter has offloaded + its weights; its worker is stopped and is replaced on wake. + host_bytes: Host bytes charged to the tier while it is parked. + retains_device_memory: Whether the parked adapter still holds some + accelerator memory (an engine process keeps its CUDA context and + workspaces). Such models are dropped when that device needs memory + and nothing resident can be evicted. + parked_at: ``time.monotonic()`` when it parked. + """ + + name: str + loaded: Any + host_bytes: int + retains_device_memory: bool = False + parked_at: float = field(default_factory=time.monotonic) + + @property + def device(self) -> str: + return str(self.loaded.device) + + +class HostCache: + """Budgeted, least-recently-parked set of models parked in host RAM. + + Not thread-safe: the registry mutates it on its lifecycle event loop. + """ + + def __init__(self, budget_bytes: int) -> None: + if budget_bytes < 0: + raise ValueError(f"host cache budget must not be negative, got {budget_bytes}") + self._budget = budget_bytes + self._parked: OrderedDict[str, ParkedModel] = OrderedDict() + self._retained: dict[str, int] = {} + + @property + def enabled(self) -> bool: + return self._budget > 0 + + @property + def budget_bytes(self) -> int: + return self._budget + + @property + def used_bytes(self) -> int: + return sum(entry.host_bytes for entry in self._parked.values()) + sum(self._retained.values()) + + @property + def names(self) -> list[str]: + """Parked model ids, least recently parked first.""" + return list(self._parked) + + def __contains__(self, name: object) -> bool: + return name in self._parked + + def __len__(self) -> int: + return len(self._parked) + + def get(self, name: str) -> ParkedModel | None: + return self._parked.get(name) + + def retained_bytes(self, name: str) -> int: + """Host bytes a resident model's adapter still holds for the tier.""" + return self._retained.get(name, 0) + + def admission_victims(self, name: str, host_bytes: int) -> list[str] | None: + """Parked models to drop, oldest first, so ``name`` can park ``host_bytes``. + + Returns an empty list when it fits as is, and None when it cannot fit + even after dropping every parked model (its size, plus host memory that + resident adapters keep, exceeds the budget). + """ + if host_bytes <= 0 or not self.enabled: + return None + others = self.used_bytes - self._retained.get(name, 0) + if name in self._parked: + others -= self._parked[name].host_bytes + overflow = others + host_bytes - self._budget + victims: list[str] = [] + for parked_name, entry in self._parked.items(): + if overflow <= 0: + break + if parked_name == name: + continue + victims.append(parked_name) + overflow -= entry.host_bytes + return victims if overflow <= 0 else None + + def put(self, entry: ParkedModel) -> None: + """Park ``entry`` as the most recently parked model. + + Its host bytes replace any charge the resident adapter was carrying. + """ + self._retained.pop(entry.name, None) + self._parked.pop(entry.name, None) + self._parked[entry.name] = entry + + def take(self, name: str) -> ParkedModel | None: + """Remove and return a parked model to wake it; its charge goes with it.""" + return self._parked.pop(name, None) + + def retain(self, name: str, host_bytes: int) -> None: + """Charge host bytes a resident adapter keeps after waking.""" + if host_bytes > 0: + self._retained[name] = host_bytes + else: + self._retained.pop(name, None) + + def release(self, name: str) -> None: + """Forget a resident model's retained charge (it was fully unloaded).""" + self._retained.pop(name, None) + + def holding_device_memory(self, device: str) -> list[str]: + """Parked models that still hold memory on ``device``, oldest first.""" + return [name for name, entry in self._parked.items() if entry.retains_device_memory and entry.device == device] diff --git a/packages/sie_server/src/sie_server/core/model_loader.py b/packages/sie_server/src/sie_server/core/model_loader.py index 09ef4f804..2f0dcf261 100644 --- a/packages/sie_server/src/sie_server/core/model_loader.py +++ b/packages/sie_server/src/sie_server/core/model_loader.py @@ -636,12 +636,24 @@ def _load_main_thread( raise return self._finish_load(name, device, adapter, config) + def register_woken(self, name: str, parked: LoadedModel) -> LoadedModel: + """Register an adapter restored from the host-RAM tier. + + The adapter kept its weights and LoRA adapters while parked, so its + LoRA bookkeeping carries over and nothing is preloaded again. Only the + worker (stopped when the model parked) and the pre/postprocessor + registrations are rebuilt. + """ + return self._finish_load(name, parked.device, parked.adapter, parked.config, restore_from=parked) + def _finish_load( self, name: str, device: str, adapter: ModelAdapter, config: ModelConfig, + *, + restore_from: LoadedModel | None = None, ) -> LoadedModel: """Complete loading - register preprocessor, create worker. @@ -650,6 +662,8 @@ def _finish_load( device: Device string. adapter: The loaded adapter. config: Model configuration. + restore_from: The parked ``LoadedModel`` when the adapter woke from + the host tier; its LoRA state is kept instead of preloaded. Returns: LoadedModel containing the loaded state. @@ -740,6 +754,14 @@ def _finish_load( max_loras=self._max_loras_per_model, ) + if restore_from is not None: + loaded_model.loras = restore_from.loras + loaded_model.max_loras = restore_from.max_loras + loaded_model.lora_failures = restore_from.lora_failures + loaded_model.lora_runtime = restore_from.lora_runtime + logger.info("Model '%s' restored from host RAM", name) + return loaded_model + # Engine-owned runtimes report their startup residency. Ordinary # PEFT adapters use the existing preload path below. lora_cap = adapter.lora_capability() diff --git a/packages/sie_server/src/sie_server/core/registry.py b/packages/sie_server/src/sie_server/core/registry.py index 00644660b..a8c8f7acb 100644 --- a/packages/sie_server/src/sie_server/core/registry.py +++ b/packages/sie_server/src/sie_server/core/registry.py @@ -36,6 +36,7 @@ from sie_server.config.routing import validate_model_routing from sie_server.config.upstreams import validate_profile_upstreams from sie_server.core.disk_cache import DiskCacheConfig, ModelDiskCacheManager +from sie_server.core.host_cache import HostCache, ParkedModel, resolve_host_cache_budget from sie_server.core.hot_reload import HotReloader from sie_server.core.load_errors import ( DevicePlacementError, @@ -84,6 +85,17 @@ _LORA_RETRY_BASE_S = 5.0 _LORA_RETRY_MAX_S = 300.0 +# Evictions that free device memory for other work park the model in the host +# tier when it has room. A manual unload, shutdown, a config change or an +# engine that exited unloads the model outright. +_PARKABLE_EVICTIONS = frozenset( + {"memory_pressure", "preload_pressure", "load_oom", "oom_recovery", "idle", "group_placement"} +) +# How long a request waits for a parked model to wake before it is told to +# retry. A wake normally takes seconds; this bounds the wait when the wake +# queues behind another model's cold load. +_HOST_WAKE_WAIT_S = 15.0 + class _ConfigChangedDuringLoadError(RuntimeError): """The downloaded config is no longer current; a later request may retry.""" @@ -320,6 +332,19 @@ def __init__( self._memory_manager = self._memory_managers[self._devices[0]] self._drain_timeout_s = drain_timeout_s + # Host-RAM tier: evicted models park their weights here and wake back + # onto their device instead of loading from disk. Off (budget 0) + # without an engine config, on CPU-only workers under ``auto``, and + # with ``SIE_HOST_CACHE_BYTES=0``. + host_cache_setting = engine_config.host_cache_bytes if engine_config is not None else 0 + self._host_cache = HostCache(resolve_host_cache_budget(host_cache_setting, self._devices)) + # Models whose parked copy has been taken for a wake still in progress. + self._waking: set[str] = set() + # Background load tasks by model, so a request can wait for a wake. + self._load_tasks: dict[str, asyncio.Task[None]] = {} + if self._host_cache.enabled: + logger.info("Host-RAM tier enabled: budget %.2f GiB", self._host_cache.budget_bytes / 1024**3) + # Disk cache management self._disk_cache_manager: ModelDiskCacheManager | None = None if engine_config is None or engine_config.disk_cache_enabled: @@ -1166,6 +1191,20 @@ def is_unloading(self, name: str) -> bool: """Check if a model is currently being unloaded.""" return name in self._unloading + def is_parked(self, name: str) -> bool: + """Whether ``name`` is parked in the host tier or waking from it.""" + return name in self._host_cache or name in self._waking + + @property + def parked_model_names(self) -> list[str]: + """Models parked in host RAM, least recently parked first.""" + return self._host_cache.names + + @property + def host_cache(self) -> HostCache: + """The host-RAM tier's bookkeeping (read-only use).""" + return self._host_cache + def is_failed(self, name: str) -> bool: """Check whether a model has a recorded load failure still in cooldown. @@ -1651,8 +1690,17 @@ async def load_async(self, name: str, device: str) -> ModelAdapter: load_start = time.monotonic() load_outcome = "error" load_stage = "total" + woke = False try: + if self._host_cache.enabled and not serves_remotely(config): + woken = await self._wake_parked(name) if name in self._host_cache else None + worker_telemetry().host_cache_lookup(model=name, hit=woken is not None) + if woken is not None: + woke = True + load_outcome = "success" + return woken + if serves_remotely(config): # No weights and no device memory: no placement, admission or # eviction, and never registered with a memory manager. @@ -1724,41 +1772,7 @@ async def load_async(self, name: str, device: str) -> ModelAdapter: # Pre-load eviction: evict LRU non-pinned models until current pressure # and any adapter-provided load headroom requirement are satisfied. - while True: - await self._wait_for_unloads() - async with lock: - if not memory_manager.should_evict_for_load(required_load_bytes): - break - lru_model = memory_manager.get_lru_model(exclude=self._pinned_models) - if lru_model is None: - break # No non-pinned models to evict, proceed with load attempt - - if lru_model not in self._loaded: - # Belt-and-braces: a MemoryManager entry with no matching - # ``_loaded`` model is a ghost (see #1600). ``_do_unload`` - # would no-op on it, so drop the stale accounting entry and - # re-evaluate rather than spin. The ``_do_unload`` finally - # normally prevents ghosts; this guards any other source. - logger.warning( - "Evictor found ghost '%s' (in memory manager, not loaded); dropping stale entry", - lru_model, - ) - memory_manager.unregister_model(lru_model) - continue - - stats = memory_manager.get_stats() - logger.info( - "Pre-load eviction: %s %.2f GB free, %.2f GB required, %.1f%% used " - "(threshold %.1f%%); evicting '%s' before loading '%s'", - load_device, - stats.available_gb, - (required_load_bytes or 0) / (1024**3), - stats.usage_ratio * 100, - memory_manager.pressure_threshold_pct, - lru_model, - name, - ) - await self._do_unload(lru_model, reason="preload_pressure") + await self._evict_for_load(name, memory_manager, required_load_bytes) try: # Load onto device (loader handles main thread vs executor) @@ -1768,20 +1782,10 @@ async def load_async(self, name: str, device: str) -> ModelAdapter: raise # OOM despite pre-load eviction: evict LRU non-pinned model and retry once - await self._wait_for_unloads() - async with lock: - lru_model = memory_manager.get_lru_model(exclude=self._pinned_models) - if lru_model is None: + if not await self._evict_one_for_oom(name, memory_manager): logger.error("OOM loading '%s' but no non-pinned models to evict", name) raise - logger.warning( - "OOM loading '%s' despite pre-eviction, evicting '%s' and retrying", - name, - lru_model, - ) - await self._do_unload(lru_model, reason="load_oom") - # Retry once after eviction. Weights are still cached # on disk so we skip ``ensure_weights_cached_async``. adapter = await self._loader.instantiate_adapter_async(name, config, model_dir, load_device) @@ -1817,15 +1821,307 @@ async def load_async(self, name: str, device: str) -> ModelAdapter: self._release_device_claims(name) else: self.clear_failure(name) - worker_telemetry().model_load_completed( - model=name, - duration_s=time.monotonic() - load_start, - outcome=load_outcome, - stage=load_stage, - ) + # A wake is not a load; it reports its own duration. + if not woke: + worker_telemetry().model_load_completed( + model=name, + duration_s=time.monotonic() - load_start, + outcome=load_outcome, + stage=load_stage, + ) # Always clear loading state when done (success or failure) self._loading.discard(name) + async def _evict_for_load(self, name: str, memory_manager: MemoryManager, required_bytes: int | None) -> None: + """Evict until the device fits a load of ``required_bytes`` (and pressure clears). + + Resident unpinned models go first, least recently used first; they + park in the host tier when it has room, so this is cheap to undo. Once + none is left, parked models that still hold memory on this device (an + engine process's context and workspaces) are dropped, oldest first. + Stops when the load fits or nothing else can be freed; the load then + proceeds and its own OOM handling takes over. + """ + lock = self._get_load_lock() + while True: + await self._wait_for_unloads() + parked_victim: str | None = None + async with lock: + if not memory_manager.should_evict_for_load(required_bytes): + return + lru_model = memory_manager.get_lru_model(exclude=self._pinned_models) + if lru_model is None: + holders = self._host_cache.holding_device_memory(memory_manager.device) + if not holders: + return # Nothing left to free, proceed with the load attempt + parked_victim = holders[0] + elif lru_model not in self._loaded: + # Belt-and-braces: a MemoryManager entry with no matching + # ``_loaded`` model is a ghost (see #1600). ``_do_unload`` + # would no-op on it, so drop the stale accounting entry and + # re-evaluate rather than spin. The ``_do_unload`` finally + # normally prevents ghosts; this guards any other source. + logger.warning( + "Evictor found ghost '%s' (in memory manager, not loaded); dropping stale entry", + lru_model, + ) + memory_manager.unregister_model(lru_model) + continue + stats = memory_manager.get_stats() + logger.info( + "Pre-load eviction: %s %.2f GB free, %.2f GB required, %.1f%% used " + "(threshold %.1f%%); %s '%s' before loading '%s'", + memory_manager.device, + stats.available_gb, + (required_bytes or 0) / (1024**3), + stats.usage_ratio * 100, + memory_manager.pressure_threshold_pct, + "dropping parked" if parked_victim is not None else "evicting", + parked_victim or lru_model, + name, + ) + if parked_victim is not None: + await self._drop_parked(parked_victim, reason="device memory") + elif lru_model is not None: + await self._do_unload(lru_model, reason="preload_pressure") + + async def _evict_one_for_oom(self, name: str, memory_manager: MemoryManager) -> bool: + """Free memory once after a load ran out of it; False when nothing could be freed.""" + await self._wait_for_unloads() + parked_victim: str | None = None + async with self._get_load_lock(): + lru_model = memory_manager.get_lru_model(exclude=self._pinned_models) + if lru_model is None: + holders = self._host_cache.holding_device_memory(memory_manager.device) + parked_victim = holders[0] if holders else None + if lru_model is not None: + logger.warning("OOM loading '%s' despite pre-eviction, evicting '%s' and retrying", name, lru_model) + await self._do_unload(lru_model, reason="load_oom") + return True + if parked_victim is not None: + logger.warning("OOM loading '%s', dropping parked '%s' and retrying", name, parked_victim) + await self._drop_parked(parked_victim, reason="device memory") + return True + return False + + async def _run_adapter_call[T](self, func: Callable[[], T], name: str, what: str) -> T: + """Run a blocking adapter call on a worker thread, waiting it out if cancelled. + + A weight move or teardown that has started must finish before the + caller acts on the adapter again, or a cancelled caller could free or + reuse memory the thread is still copying. + """ + loop = asyncio.get_running_loop() + future = loop.run_in_executor(None, func) + try: + return await asyncio.shield(future) + except asyncio.CancelledError: + logger.warning("%s of '%s' was cancelled midway; waiting for it to finish", what, name) + with contextlib.suppress(BaseException): + await asyncio.shield(future) + raise + + async def _park(self, name: str, loaded: LoadedModel, reason: str) -> bool: + """Park an evicted model's weights in host RAM; False to unload it instead. + + Called by the unload path after the model's worker has drained, with + the load lock held. Parks only evictions that free memory for other + work, single-device models, and adapters that report a size the + budget can hold, after dropping older parked models to make room. + """ + if reason not in _PARKABLE_EVICTIONS or not self._host_cache.enabled or serves_remotely(loaded.config): + return False + if len(self._claimed_members(name)) > 1: + # A multi-device engine would have to wake on the very same group. + return False + adapter = loaded.adapter + try: + host_bytes = adapter.host_offload_bytes() + except Exception: + logger.exception("host_offload_bytes() failed for '%s'; unloading it", name) + return False + if isinstance(host_bytes, bool) or not isinstance(host_bytes, int) or host_bytes <= 0: + return False + # A resident model is never parked as well; never leave an older copy behind. + await self._drop_parked(name, reason="superseded") + victims = self._host_cache.admission_victims(name, host_bytes) + if victims is None: + logger.info( + "Host tier: '%s' needs %.2f GiB, more than the %.2f GiB budget can hold; unloading it", + name, + host_bytes / 1024**3, + self._host_cache.budget_bytes / 1024**3, + ) + return False + for victim in victims: + await self._drop_parked(victim, reason="host budget") + started = time.monotonic() + try: + offload = await self._run_adapter_call(adapter.offload_to_host, name, "Parking") + except Exception: # noqa: BLE001 - any failure unloads the model as before + logger.warning("Host tier: parking '%s' failed; unloading it instead", name, exc_info=True) + return False + parked_bytes = offload.host_bytes if offload.host_bytes > 0 else host_bytes + self._host_cache.put( + ParkedModel( + name=name, + loaded=loaded, + host_bytes=parked_bytes, + retains_device_memory=offload.retains_device_memory, + ) + ) + worker_telemetry().host_cache_memory_changed(model=name, host_bytes=parked_bytes) + logger.info( + "Host tier: parked '%s' from %s in %.2fs (%.2f GiB; tier %.2f of %.2f GiB)", + name, + loaded.device, + time.monotonic() - started, + parked_bytes / 1024**3, + self._host_cache.used_bytes / 1024**3, + self._host_cache.budget_bytes / 1024**3, + ) + return True + + async def _drop_parked(self, name: str, *, reason: str) -> bool: + """Unload a parked model and free its host RAM; False when it was not parked.""" + entry = self._host_cache.take(name) + if entry is None: + return False + await self._discard_parked(entry, reason=reason) + return True + + async def _discard_parked(self, entry: ParkedModel, *, reason: str) -> None: + """Unload an adapter taken out of the host tier and drop its charges.""" + name = entry.name + self._host_cache.release(name) + worker_telemetry().host_cache_memory_changed(model=name, host_bytes=0) + logger.info("Host tier: dropping parked '%s' (%s)", name, reason) + try: + await self._close_and_unload(entry.loaded.adapter, name) + except Exception: + logger.exception("Unloading parked '%s' failed", name) + try: + # A failed wake may have registered it again before giving up. + self._loader.unregister(name, entry.device) + except Exception: # noqa: BLE001 - best-effort cleanup + logger.warning("loader.unregister failed while dropping parked '%s'", name, exc_info=True) + if entry.retains_device_memory: + # Device memory came back, so an earlier OOM may no longer apply. + self._clear_transient_failures((LoadErrorClass.OOM, LoadErrorClass.PLACEMENT)) + + async def _wake_parked(self, name: str) -> ModelAdapter | None: + """Bring a parked model back onto its device; None to load it from disk instead. + + Runs under the load admission lock like any load. A parked copy whose + config changed, or a wake that fails, is dropped and the caller loads + the model the ordinary way. + """ + lock = self._get_load_lock() + async with self._get_load_admission_lock(): + await self._wait_for_unloads() + async with lock: + if name in self._loaded: + self._memory_manager_for_model(name).touch(name) + return self._loaded[name].adapter + if name in self._unloading: + msg = f"Model '{name}' is currently being unloaded" + raise RuntimeError(msg) + entry = self._host_cache.take(name) + if entry is None: + return None + # The host copy exists until the restore finishes, so it stays + # charged while models evicted to make room for it are parked. + self._host_cache.retain(name, entry.host_bytes) + self._waking.add(name) + stale = not _model_configs_semantically_equal(self._configs.get(name), entry.loaded.config) + started = time.monotonic() + restoring = False + try: + if stale: + await self._discard_parked(entry, reason="config changed while parked") + return None + memory_manager = await self._make_room_to_wake(name, entry) + restoring = True + retained = await self._restore_parked(name, entry, memory_manager) + loaded = self._loader.register_woken(name, entry.loaded) + async with lock: + self._loaded[name] = loaded + self._register_across_group(name, loaded) + except asyncio.CancelledError: + if not restoring: + # Nothing about the parked copy changed; keep it. + self._host_cache.put(entry) + else: + # The restore ran to completion or failed, but nobody will + # register it: drop it rather than guess its state. + await asyncio.shield(self._discard_parked(entry, reason="wake cancelled")) + raise + except Exception as exc: # noqa: BLE001 - any failure falls back to an ordinary load + timed_out = isinstance(exc, (TimeoutError, ModelLoadTimeoutError)) + worker_telemetry().host_cache_wake_completed( + model=name, duration_s=time.monotonic() - started, outcome="timeout" if timed_out else "error" + ) + logger.warning( + "Host tier: waking '%s' failed (%s: %s); loading it from disk instead", + name, + type(exc).__name__, + exc, + ) + await self._discard_parked(entry, reason="wake failed") + return None + finally: + self._waking.discard(name) + self._host_cache.retain(name, retained) + duration = time.monotonic() - started + worker_telemetry().host_cache_memory_changed(model=name, host_bytes=retained) + worker_telemetry().host_cache_wake_completed(model=name, duration_s=duration, outcome="success") + logger.info("Host tier: woke '%s' onto %s in %.2fs", name, loaded.device, duration) + return loaded.adapter + + async def _make_room_to_wake(self, name: str, entry: ParkedModel) -> MemoryManager: + """Place a parked model back on its own device and evict until it fits. + + A parked model wakes on the device it was loaded on: its tensors are + copied back to that device, and an engine process stays bound to it. + """ + device = entry.device + load_device = await self._resolve_single_device_evicting(name, device) + if load_device != device: + msg = f"Parked model '{name}' must wake on {device}, but placement resolved {load_device}" + raise DevicePlacementError(msg) + memory_manager = self._memory_manager_for_device(device) + await self._evict_for_load(name, memory_manager, entry.loaded.memory_bytes or None) + return memory_manager + + async def _restore_parked(self, name: str, entry: ParkedModel, memory_manager: MemoryManager) -> int: + """Copy the parked weights back, retrying once after an OOM; returns host bytes still held.""" + adapter = entry.loaded.adapter + try: + retained = await self._run_adapter_call(adapter.restore_from_host, name, "Wake") + except RuntimeError as exc: + if not self._is_oom_error(exc) or not await self._evict_one_for_oom(name, memory_manager): + raise + retained = await self._run_adapter_call(adapter.restore_from_host, name, "Wake") + return released_bytes(retained) + + async def wake_parked(self, name: str, device: str, *, timeout_s: float = _HOST_WAKE_WAIT_S) -> bool: + """Start waking a parked model and wait up to ``timeout_s`` for it. + + For request paths that would otherwise answer "loading, retry": a + wake usually takes seconds, so waiting for it serves the request now. + Returns whether the model is loaded; False also when it is not parked. + """ + if name in self._loaded and name not in self._unloading: + return True + if not self.is_parked(name): + return False + await self.start_load_async(name, device) + task = self._load_tasks.get(name) + if task is not None: + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(asyncio.shield(task), timeout_s) + return name in self._loaded and name not in self._unloading + async def start_load_async(self, name: str, device: str) -> bool: """Start loading a model in the background (non-blocking). @@ -1878,8 +2174,14 @@ async def start_load_async(self, name: str, device: str) -> bool: ) self._background_tasks.add(task) task.add_done_callback(self._background_tasks.discard) + self._load_tasks[name] = task + task.add_done_callback(lambda done: self._forget_load_task(name, done)) return True + def _forget_load_task(self, name: str, task: asyncio.Task[None]) -> None: + if self._load_tasks.get(name) is task: + del self._load_tasks[name] + async def _load_model_background(self, name: str, device: str) -> None: """Load a model in the background (called via asyncio.create_task). @@ -2056,6 +2358,7 @@ def _unload_sync(self, name: str, *, reason: str) -> None: # the process. self._unregister_across_group(name, device) self._release_device_claims(name) + self._host_cache.release(name) logger.info("Model '%s' unloaded", name) @@ -2093,18 +2396,43 @@ async def _run_adapter_unload(self, adapter: ModelAdapter, name: str) -> None: completion above; ``unload()`` itself is a synchronous contract on every adapter and is safe off the loop. """ - loop = asyncio.get_running_loop() - unload_future = loop.run_in_executor(None, adapter.unload) - try: - await asyncio.shield(unload_future) - except asyncio.CancelledError: - logger.warning( - "Unload of '%s' was cancelled mid-teardown; waiting for it to finish before releasing the load lock", - name, - ) - with contextlib.suppress(BaseException): - await asyncio.shield(unload_future) - raise + await self._run_adapter_call(adapter.unload, name, "Unload") + + async def _close_and_unload(self, adapter: ModelAdapter, name: str) -> None: + """Close an adapter's async client, then run its blocking ``unload()`` off the loop.""" + # Some adapters (e.g. the SGLang generation adapter) hold an + # async HTTP client to a subprocess. ``unload()`` is sync and + # can only fire-and-forget the client close, which races the + # subprocess termination on the next line — the close could be + # cut off before its connections drain, leaking fds / wedging + # on a half-open socket. When the adapter exposes an awaitable + # ``aclose_client``, drive it to completion HERE (we're on the + # event loop) so the client is fully closed against the still- + # live subprocess before ``unload()`` terminates it. + aclose_client = getattr(adapter, "aclose_client", None) + if aclose_client is not None: + try: + await aclose_client() + except Exception: # noqa: BLE001 - close is best-effort + logger.warning("aclose_client() failed during unload of '%s'", name, exc_info=True) + + # Adapter.unload() handles gc.collect + empty_cache, and for + # subprocess-backed adapters a SIGTERM/SIGKILL wait of up to + # 15s. All of it is blocking, and running it inline froze the + # event loop for the duration: every other coroutine — + # health probes, in-flight responses on models that are not + # being evicted, the metrics endpoint — stopped, not merely + # the load path the lock already serialises (design proposal + # ``settlement-and-queue-scalability.md``, B6 part one; same + # class as the #1600 scar above). + # + # The load lock stays held across this on purpose. The + # ``finally`` below drops this model's MemoryManager entry, + # and until ``unload()`` returns the VRAM is not actually + # free, so a load admitted against the freed accounting + # would size itself against memory the dying adapter still + # holds. A slow unload beats an OOM. + await self._run_adapter_unload(adapter, name) async def _drain_generation_before_deadline( self, @@ -2228,40 +2556,14 @@ async def _finish_unload(self, name: str, loaded: LoadedModel, *, reason: str) - # counted in ``_models`` — that ``get_lru_model`` keeps # re-selecting while ``_do_unload`` no-ops on it: a no-yield spin # under ``_load_lock`` that starves the whole event loop. See #1600. + parked = False try: - # Some adapters (e.g. the SGLang generation adapter) hold an - # async HTTP client to a subprocess. ``unload()`` is sync and - # can only fire-and-forget the client close, which races the - # subprocess termination on the next line — the close could be - # cut off before its connections drain, leaking fds / wedging - # on a half-open socket. When the adapter exposes an awaitable - # ``aclose_client``, drive it to completion HERE (we're on the - # event loop) so the client is fully closed against the still- - # live subprocess before ``unload()`` terminates it. - aclose_client = getattr(loaded.adapter, "aclose_client", None) - if aclose_client is not None: - try: - await aclose_client() - except Exception: # noqa: BLE001 - close is best-effort - logger.warning("aclose_client() failed during unload of '%s'", name, exc_info=True) - - # Adapter.unload() handles gc.collect + empty_cache, and for - # subprocess-backed adapters a SIGTERM/SIGKILL wait of up to - # 15s. All of it is blocking, and running it inline froze the - # event loop for the duration: every other coroutine — - # health probes, in-flight responses on models that are not - # being evicted, the metrics endpoint — stopped, not merely - # the load path the lock already serialises (design proposal - # ``settlement-and-queue-scalability.md``, B6 part one; same - # class as the #1600 scar above). - # - # The load lock stays held across this on purpose. The - # ``finally`` below drops this model's MemoryManager entry, - # and until ``unload()`` returns the VRAM is not actually - # free, so a load admitted against the freed accounting - # would size itself against memory the dying adapter still - # holds. A slow unload beats an OOM. - await self._run_adapter_unload(loaded.adapter, name) + # An eviction that frees memory for other work parks the + # weights in host RAM when the tier has room; the next + # request wakes them instead of loading from disk. + parked = await self._park(name, loaded, reason) + if not parked: + await self._close_and_unload(loaded.adapter, name) finally: # Always run, even if adapter teardown above raised, so a # failed unload can never leave a ghost behind (#1600). @@ -2277,6 +2579,9 @@ async def _finish_unload(self, name: str, loaded: LoadedModel, *, reason: str) - # Hand back any exclusively held devices. A no-op for the # single-device models that are every profile today. self._release_device_claims(name) + # A fully unloaded adapter no longer holds a host copy. + if not parked: + self._host_cache.release(name) # Freeing memory makes any prior OOM-class load failure on a # *sibling* model retryable; clear those records so the next @@ -2289,7 +2594,7 @@ async def _finish_unload(self, name: str, loaded: LoadedModel, *, reason: str) - name, ) - logger.info("Model '%s' unloaded", name) + logger.info("Model '%s' %s", name, "parked in host RAM" if parked else "unloaded") finally: self._unloading.discard(name) if self._unload_done.get(name) is done: @@ -2313,7 +2618,7 @@ async def unload_async(self, name: str, *, reason: str = "manual") -> None: # that model still holds its memory and its whole device group. Gating # on the config alone would make it unloadable for the life of the # process. - if name not in self._configs and name not in self._loaded: + if name not in self._configs and name not in self._loaded and name not in self._host_cache: msg = _ERR_MODEL_NOT_FOUND.format(name=name) raise KeyError(msg) @@ -2322,6 +2627,9 @@ async def unload_async(self, name: str, *, reason: str = "manual") -> None: if pending_unload is not None: await pending_unload.wait() if name not in self._loaded: + # An explicit unload of a parked model frees its host copy. + if await self._drop_parked(name, reason=reason): + return msg = _ERR_MODEL_NOT_LOADED.format(name=name) raise KeyError(msg) @@ -2349,10 +2657,12 @@ def _unload_all_sync(self) -> None: self._unload_sync(name, reason="shutdown") async def unload_all_async(self) -> None: - """Unload all loaded models (async, concurrency-safe).""" + """Unload all loaded and parked models (async, concurrency-safe).""" async with self._get_load_admission_lock(): for name in list(self._loaded.keys()): await self._do_unload(name, reason="shutdown") + for name in self._host_cache.names: + await self._drop_parked(name, reason="shutdown") def get_configs_snapshot(self, bundle_id: str | None = None) -> dict[str, ModelConfig]: """Return a shallow copy of the current model configs. @@ -2465,6 +2775,7 @@ async def add_config_async(self, config: ModelConfig, model_dir: Path | None = N for model_id in sorted(invalidated_ids): if model_id in self._loaded: await self._do_unload(model_id, reason="config_change") + await self._drop_parked(model_id, reason="config_change") for model_id in sorted(removed_ids): self._drop_config_entry(model_id) self._apply_config_entry(config, model_dir) @@ -2489,6 +2800,7 @@ async def remove_config_async(self, model_id: str) -> set[str]: for removed_id in sorted(removed_ids): if removed_id in self._loaded: await self._do_unload(removed_id, reason="config_change") + await self._drop_parked(removed_id, reason="config_change") self._drop_config_entry(removed_id) self._loader.update_configs(self._configs) @@ -2682,6 +2994,7 @@ async def replace_configs_async( for name in sorted(invalidated): if name in self._loaded: await self._do_unload(name, reason="config_change") + await self._drop_parked(name, reason="config_change") self._configs = dict(new_configs) self._loader.update_configs(self._configs) @@ -2884,6 +3197,7 @@ async def evict_lru_excluding(self, exclude_name: str, *, timeout_s: float = 5.0 conflated failure modes. """ lock = self._get_load_lock() + parked_victim: str | None = None try: await asyncio.wait_for(lock.acquire(), timeout=timeout_s) except TimeoutError: @@ -2909,15 +3223,32 @@ async def evict_lru_excluding(self, exclude_name: str, *, timeout_s: float = 5.0 releasing=self._optional_release_inflight(), ) if candidate is None: - return EvictionResult.NO_CANDIDATE - logger.info( - "OOM recovery: evicting LRU sibling '%s' on %s (caller='%s')", - candidate, - manager.device, - exclude_name, - ) + # Nothing resident can go; a parked engine may still hold + # memory on this device. + holders = self._host_cache.holding_device_memory(manager.device) + if not holders: + return EvictionResult.NO_CANDIDATE + parked_victim = holders[0] + logger.info( + "OOM recovery: dropping parked '%s' on %s (caller='%s')", + parked_victim, + manager.device, + exclude_name, + ) + else: + logger.info( + "OOM recovery: evicting LRU sibling '%s' on %s (caller='%s')", + candidate, + manager.device, + exclude_name, + ) finally: lock.release() + if parked_victim is not None: + await self._drop_parked(parked_victim, reason="device memory") + return EvictionResult.EVICTED + if candidate is None: + return EvictionResult.NO_CANDIDATE try: await self._do_unload(candidate, reason="oom_recovery") except Exception: @@ -3224,6 +3555,15 @@ async def stop_hot_reload(self) -> None: await self._hot_reloader.stop() self._hot_reloader = None + async def _reap_exited_parked_engines(self) -> None: + """Drop parked models whose engine process exited while they were parked.""" + for name in self._host_cache.names: + entry = self._host_cache.get(name) + if entry is None or _engine_exit_code(entry.loaded.adapter) is None: + continue + logger.warning("Parked model '%s' engine process exited; dropping its parked copy", name) + await self._drop_parked(name, reason="engine exited") + async def _reap_exited_engines(self) -> None: """Unload every loaded model whose engine process has exited. @@ -3274,6 +3614,7 @@ async def _memory_monitor_loop(self) -> None: try: await asyncio.sleep(self._memory_config.memory_check_interval_s) await self._reap_exited_engines() + await self._reap_exited_parked_engines() # Quick check without lock pressured_managers = [manager for manager in self._memory_managers.values() if manager.check_pressure()] @@ -3293,10 +3634,13 @@ async def _memory_monitor_loop(self) -> None: # Adapters with nothing cached are not locked. A busy model is # skipped instead of draining its in-flight batches. await self._release_optional_memory_on(manager, exclude_name=None) + victim: str | None = None + drop_parked = False async with lock: # Re-check under lock (may have resolved) if not manager.check_pressure(): break + lru_model: str | None = None # Don't evict if only 1 model loaded - would cause immediate reload if manager.loaded_model_count <= 1: logger.debug( @@ -3304,18 +3648,28 @@ async def _memory_monitor_loop(self) -> None: manager.device, manager.loaded_model_count, ) - break + else: + # A timed-out release still holds the inference thread. + # Joining it from unload would pin the monitor, and + # the timeout itself is not a reason to unload. + lru_model = manager.get_lru_model( + exclude=self._pinned_models | self._optional_release_inflight() + ) - # A timed-out release still holds the inference thread. - # Joining it from unload would pin the monitor, and - # the timeout itself is not a reason to unload. - lru_model = manager.get_lru_model( - exclude=self._pinned_models | self._optional_release_inflight() - ) if lru_model is None: - break - - if lru_model not in self._loaded: + # A parked engine can still hold memory on this + # device; dropping it never forces a resident + # model to reload. + holders = self._host_cache.holding_device_memory(manager.device) + if not holders: + break + victim, drop_parked = holders[0], True + logger.info( + "Memory monitor: %s under pressure, dropping parked '%s'", + manager.device, + victim, + ) + elif lru_model not in self._loaded: # Belt-and-braces ghost guard (see #1600): drop a # stale MemoryManager entry with no ``_loaded`` # model instead of spinning on a no-op _do_unload. @@ -3325,16 +3679,21 @@ async def _memory_monitor_loop(self) -> None: ) manager.unregister_model(lru_model) continue - - stats = manager.get_stats() - logger.info( - "Memory monitor: %s pressure %.1f%% > %.1f%% threshold, evicting LRU model '%s'", - manager.device, - stats.usage_ratio * 100, - manager.pressure_threshold_pct, - lru_model, - ) - await self._do_unload(lru_model, reason="memory_pressure") + else: + victim = lru_model + stats = manager.get_stats() + logger.info( + "Memory monitor: %s pressure %.1f%% > %.1f%% threshold, evicting LRU model '%s'", + manager.device, + stats.usage_ratio * 100, + manager.pressure_threshold_pct, + lru_model, + ) + assert victim is not None + if drop_parked: + await self._drop_parked(victim, reason="device memory") + else: + await self._do_unload(victim, reason="memory_pressure") except asyncio.CancelledError: break diff --git a/packages/sie_server/src/sie_server/observability/worker_telemetry.py b/packages/sie_server/src/sie_server/observability/worker_telemetry.py index 4b75d9ff4..e34f20464 100644 --- a/packages/sie_server/src/sie_server/observability/worker_telemetry.py +++ b/packages/sie_server/src/sie_server/observability/worker_telemetry.py @@ -87,6 +87,9 @@ def _endpoint_origin_for_log(endpoint: str) -> str: MODEL_MEMORY_METRIC_NAME: Final = "sie.worker.model.memory" OOM_RECOVERIES_METRIC_NAME: Final = "sie.worker.oom.recoveries" MODEL_EVICTIONS_METRIC_NAME: Final = "sie.worker.model.evictions" +HOST_CACHE_LOOKUPS_METRIC_NAME: Final = "sie.worker.model.host_cache.lookups" +HOST_CACHE_WAKE_DURATION_METRIC_NAME: Final = "sie.worker.model.host_cache.wake.duration" +HOST_CACHE_MEMORY_METRIC_NAME: Final = "sie.worker.model.host_cache.memory" ADAPTIVE_WAIT_METRIC_NAME: Final = "sie.worker.scheduler.adaptive.wait" ADAPTIVE_COST_METRIC_NAME: Final = "sie.worker.scheduler.adaptive.cost" ADAPTIVE_P50_METRIC_NAME: Final = "sie.worker.scheduler.adaptive.p50" @@ -157,6 +160,21 @@ def _endpoint_origin_for_log(endpoint: str) -> str: 600.0, 900.0, ) +HOST_CACHE_WAKE_BUCKETS_S: Final = ( + 0.01, + 0.025, + 0.05, + 0.1, + 0.25, + 0.5, + 1.0, + 2.5, + 5.0, + 10.0, + 30.0, + 60.0, + 120.0, +) TTFT_TPOT_BUCKETS_S: Final = ( 0.005, 0.01, @@ -343,6 +361,19 @@ def model_evicted( profile: object = "default", ) -> None: ... + def host_cache_lookup(self, *, model: object, hit: bool, profile: object = "default") -> None: ... + + def host_cache_wake_completed( + self, + *, + model: object, + duration_s: object, + outcome: object, + profile: object = "default", + ) -> None: ... + + def host_cache_memory_changed(self, *, model: object, host_bytes: object, profile: object = "default") -> None: ... + def adaptive_snapshot( self, *, @@ -443,6 +474,15 @@ def oom_recovery_completed(self, **_: Any) -> None: def model_evicted(self, **_: Any) -> None: return + def host_cache_lookup(self, **_: Any) -> None: + return + + def host_cache_wake_completed(self, **_: Any) -> None: + return + + def host_cache_memory_changed(self, **_: Any) -> None: + return + def adaptive_snapshot(self, **_: Any) -> None: return @@ -576,6 +616,21 @@ def __init__(self, meter: Meter) -> None: unit="{model}", description="Resident models evicted from the engine by bounded cause", ) + self._host_cache_lookups = meter.create_counter( + HOST_CACHE_LOOKUPS_METRIC_NAME, + unit="{lookup}", + description="Model loads that found the model parked in host RAM (hit) or loaded it from disk (miss)", + ) + self._host_cache_wake_duration = meter.create_histogram( + HOST_CACHE_WAKE_DURATION_METRIC_NAME, + unit="s", + description="Time to bring a model parked in host RAM back onto its device", + ) + self._host_cache_memory = meter.create_gauge( + HOST_CACHE_MEMORY_METRIC_NAME, + unit="By", + description="Host RAM charged to the host tier for a model's weights", + ) self._adaptive_wait = meter.create_gauge( ADAPTIVE_WAIT_METRIC_NAME, unit="s", @@ -910,6 +965,32 @@ def model_evicted( {"reason": _enum(reason, _EVICTION_REASONS), **_model_attributes(model, profile)}, ) + def host_cache_lookup(self, *, model: object, hit: bool, profile: object = "default") -> None: + self._host_cache_lookups.add( + 1, + {"result": "hit" if hit is True else "miss", **_model_attributes(model, profile)}, + ) + + def host_cache_wake_completed( + self, + *, + model: object, + duration_s: object, + outcome: object, + profile: object = "default", + ) -> None: + duration = _nonnegative_float(duration_s) + if duration is None: + return + self._host_cache_wake_duration.record( + duration, + {"outcome": _enum(outcome, _LOAD_OUTCOMES), **_model_attributes(model, profile)}, + ) + + def host_cache_memory_changed(self, *, model: object, host_bytes: object, profile: object = "default") -> None: + if (memory := _nonnegative_int(host_bytes)) is not None: + self._host_cache_memory.set(memory, _model_attributes(model, profile)) + def adaptive_snapshot( self, *, @@ -1210,6 +1291,7 @@ def metric_views() -> list[View]: (REQUEST_DURATION_METRIC_NAME, REQUEST_DURATION_BUCKETS_S), (INFERENCE_DURATION_METRIC_NAME, INFERENCE_DURATION_BUCKETS_S), (MODEL_LOAD_DURATION_METRIC_NAME, MODEL_LOAD_DURATION_BUCKETS_S), + (HOST_CACHE_WAKE_DURATION_METRIC_NAME, HOST_CACHE_WAKE_BUCKETS_S), (GENERATION_TTFT_METRIC_NAME, TTFT_TPOT_BUCKETS_S), (GENERATION_WORKER_WAIT_METRIC_NAME, TTFT_TPOT_BUCKETS_S), (GENERATION_TPOT_METRIC_NAME, TTFT_TPOT_BUCKETS_S), @@ -1246,6 +1328,9 @@ def metric_names() -> frozenset[str]: MODEL_MEMORY_METRIC_NAME, OOM_RECOVERIES_METRIC_NAME, MODEL_EVICTIONS_METRIC_NAME, + HOST_CACHE_LOOKUPS_METRIC_NAME, + HOST_CACHE_WAKE_DURATION_METRIC_NAME, + HOST_CACHE_MEMORY_METRIC_NAME, ADAPTIVE_WAIT_METRIC_NAME, ADAPTIVE_COST_METRIC_NAME, ADAPTIVE_P50_METRIC_NAME, diff --git a/packages/sie_server/src/sie_server/queue_executor.py b/packages/sie_server/src/sie_server/queue_executor.py index 3d8e85870..edeb2ecfb 100644 --- a/packages/sie_server/src/sie_server/queue_executor.py +++ b/packages/sie_server/src/sie_server/queue_executor.py @@ -880,6 +880,11 @@ async def ensure_model_ready(self, model_id: str) -> ReadinessState: if failure is not None and failure.is_permanent: return "failed" + # A model parked in host RAM wakes in seconds: wait for it so the + # dispatcher can run this work now instead of re-polling. + if self._registry.is_parked(model_id) and await self._registry.wake_parked(model_id, self._registry.device): + return "ready" + if self._registry.is_loading(model_id): return "loading_in_progress" diff --git a/packages/sie_server/tests/adapters/test_host_offload.py b/packages/sie_server/tests/adapters/test_host_offload.py new file mode 100644 index 000000000..44957787f --- /dev/null +++ b/packages/sie_server/tests/adapters/test_host_offload.py @@ -0,0 +1,230 @@ +"""Generic host-RAM parking of in-process tensors (``_host_offload`` and ``BaseAdapter``). + +CI has no GPU, so these tests park from ``cpu`` to host memory: every +non-empty CPU tensor counts as "on the device", and each park and wake must +produce fresh storage while keeping values, parameter identity and structure. +""" + +from __future__ import annotations + +from typing import Any, ClassVar + +import pytest +import torch +from sie_server.adapters import _host_offload +from sie_server.adapters._base_adapter import BaseAdapter +from sie_server.adapters._spec import AdapterSpec + +CPU = torch.device("cpu") + + +class _Tied(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.embed = torch.nn.Embedding(8, 4) + self.head = torch.nn.Linear(4, 8, bias=False) + self.head.weight = self.embed.weight # tied + self.register_buffer("scale", torch.ones(4)) + self.cached = torch.arange(4.0) # plain tensor attribute + + +class _Pipeline: + """Stands in for a transformers pipeline: the module is its ``model``.""" + + def __init__(self) -> None: + self.model = torch.nn.Linear(2, 2) + + +class _Holder: + def __init__(self) -> None: + self.model = _Tied() + self.loose = torch.full((3,), 7.0) + self.table = {"a": torch.ones(2), "nested": [torch.zeros(2)]} + self.pipe = _Pipeline() + self.name = "not a tensor" + + +def _ptr(tensor: torch.Tensor) -> int: + return tensor.untyped_storage().data_ptr() + + +def test_park_and_unpark_move_every_reachable_tensor_and_keep_identity() -> None: + holder = _Holder() + weight = holder.model.embed.weight + before = { + "weight": weight.detach().clone(), + "scale": holder.model.scale.clone(), + "cached": holder.model.cached.clone(), + "loose": holder.loose.clone(), + "a": holder.table["a"].clone(), + "pipe": holder.pipe.model.weight.detach().clone(), + } + old_ptrs = { + "weight": _ptr(weight), + "scale": _ptr(holder.model.scale), + "cached": _ptr(holder.model.cached), + "loose": _ptr(holder.loose), + "a": _ptr(holder.table["a"]), + "nested": _ptr(holder.table["nested"][0]), + "pipe": _ptr(holder.pipe.model.weight), + } + + expected = _host_offload.parkable_bytes(holder, CPU) + assert expected is not None + parked = _host_offload.park(holder, CPU) + # The tied weight moves once. + assert parked.nbytes == expected + assert holder.model.head.weight is holder.model.embed.weight is weight + parked_ptrs = { + "weight": _ptr(weight), + "scale": _ptr(holder.model.scale), + "cached": _ptr(holder.model.cached), + "loose": _ptr(holder.loose), + "a": _ptr(holder.table["a"]), + "nested": _ptr(holder.table["nested"][0]), + "pipe": _ptr(holder.pipe.model.weight), + } + assert all(parked_ptrs[key] != old_ptrs[key] for key in old_ptrs) + + _host_offload.unpark(parked) + assert holder.model.head.weight is weight + assert all( + _ptr(tensor) != parked_ptrs[key] + for key, tensor in { + "weight": weight, + "scale": holder.model.scale, + "loose": holder.loose, + "a": holder.table["a"], + "pipe": holder.pipe.model.weight, + }.items() + ) + assert torch.equal(weight, before["weight"]) + assert torch.equal(holder.model.scale, before["scale"]) + assert torch.equal(holder.model.cached, before["cached"]) + assert torch.equal(holder.loose, before["loose"]) + assert torch.equal(holder.table["a"], before["a"]) + assert torch.equal(holder.pipe.model.weight, before["pipe"]) + assert holder.name == "not a tensor" + + +def test_unpark_again_is_a_no_op() -> None: + holder = _Holder() + parked = _host_offload.park(holder, CPU) + _host_offload.unpark(parked) + restored = _ptr(holder.model.embed.weight) + _host_offload.unpark(parked) + assert _ptr(holder.model.embed.weight) == restored + + +def test_a_tensor_in_a_tuple_makes_the_adapter_unparkable() -> None: + holder = _Holder() + holder.pair = (torch.ones(2), "label") + assert _host_offload.parkable_bytes(holder, CPU) is None + + +def test_a_quantized_parameter_subclass_makes_the_adapter_unparkable() -> None: + class _Quantized(torch.nn.Parameter): + pass + + holder = _Holder() + holder.model.extra = torch.nn.Linear(2, 2) + holder.model.extra.weight = _Quantized(torch.ones(2, 2)) + assert _host_offload.parkable_bytes(holder, CPU) is None + + +def test_accelerate_hooks_make_the_adapter_unparkable() -> None: + holder = _Holder() + holder.model.embed._hf_hook = object() + assert _host_offload.parkable_bytes(holder, CPU) is None + + +def test_nothing_on_the_device_is_unparkable() -> None: + holder = _Holder() + assert _host_offload.parkable_bytes(holder, torch.device("cuda", 0)) is None + + +class _GraphRunner: + def __init__(self) -> None: + self.cleared = 0 + + def clear(self) -> None: + self.cleared += 1 + + +class _Adapter(BaseAdapter): + spec: ClassVar[AdapterSpec] = AdapterSpec(inputs=("text",), outputs=("dense",), unload_fields=("_model",)) + + def __init__(self, device: str | None) -> None: + self._model = torch.nn.Linear(4, 4) + self._device = device + self._graphs: Any = _GraphRunner() + self._cuda_graphs = "bucketed" + + def load(self, device: str) -> None: + self._device = device + + def encode(self, *args: Any, **kwargs: Any) -> Any: + raise NotImplementedError + + +def test_base_adapter_does_not_park_off_cuda() -> None: + assert _Adapter("cpu").host_offload_bytes() is None + assert _Adapter("mps").host_offload_bytes() is None + assert _Adapter(None).host_offload_bytes() is None + + +def test_base_adapter_parks_clears_graphs_and_wakes(monkeypatch: pytest.MonkeyPatch) -> None: + adapter = _Adapter("cpu") + monkeypatch.setattr(adapter, "_host_offload_device", lambda: CPU) + weight = adapter._model.weight + values = weight.detach().clone() + size = adapter.host_offload_bytes() + assert size == sum(p.nbytes for p in adapter._model.parameters()) + + offload = adapter.offload_to_host() + assert offload.host_bytes == size + assert not offload.retains_device_memory + assert adapter._graphs.cleared == 1 + # Parked: no second park, and the size check declines. + assert adapter.host_offload_bytes() is None + with pytest.raises(RuntimeError, match="nothing on a CUDA device"): + adapter.offload_to_host() + + assert adapter.restore_from_host() == 0 + assert adapter._model.weight is weight + assert torch.equal(weight, values) + with pytest.raises(RuntimeError, match="not parked"): + adapter.restore_from_host() + + +def test_a_graph_holder_without_clear_blocks_parking(monkeypatch: pytest.MonkeyPatch) -> None: + adapter = _Adapter("cpu") + monkeypatch.setattr(adapter, "_host_offload_device", lambda: CPU) + adapter._prefix_graphs = object() + assert adapter.host_offload_bytes() is None + + +def test_opting_out_disables_parking(monkeypatch: pytest.MonkeyPatch) -> None: + adapter = _Adapter("cpu") + monkeypatch.setattr(adapter, "_host_offload_device", lambda: CPU) + monkeypatch.setattr(_Adapter, "host_offload", False) + assert adapter.host_offload_bytes() is None + + +def test_unload_while_parked_drops_the_host_copy(monkeypatch: pytest.MonkeyPatch) -> None: + adapter = _Adapter("cpu") + monkeypatch.setattr(adapter, "_host_offload_device", lambda: CPU) + adapter.offload_to_host() + adapter.unload() + assert adapter._parked_tensors is None + assert adapter._model is None + + +def test_weights_the_walk_cannot_reach_block_parking(monkeypatch: pytest.MonkeyPatch) -> None: + adapter = _Adapter("cpu") + monkeypatch.setattr(adapter, "_host_offload_device", lambda: CPU) + reachable = adapter.host_offload_bytes() + assert reachable is not None + # The adapter reports more weights than the walk found. + monkeypatch.setattr(adapter, "memory_footprint", lambda: reachable * 2) + assert adapter.host_offload_bytes() is None diff --git a/packages/sie_server/tests/adapters/test_sglang_parking.py b/packages/sie_server/tests/adapters/test_sglang_parking.py new file mode 100644 index 000000000..c5c1850a6 --- /dev/null +++ b/packages/sie_server/tests/adapters/test_sglang_parking.py @@ -0,0 +1,233 @@ +"""Host-RAM parking for SGLang engines: launch flags, park and wake over HTTP.""" + +from __future__ import annotations + +import re +import subprocess +from pathlib import Path +from typing import Any + +import pytest +import requests +from sie_server.adapters._generation_base import GenerationDrainingError +from sie_server.adapters._types import ERR_NOT_LOADED +from sie_server.adapters.sglang import _server +from sie_server.adapters.sglang.embedding import SGLangEmbeddingAdapter +from sie_server.adapters.sglang.generation import SGLangGenerationAdapter +from sie_server.core.host_cache import HOST_CACHE_ENV + +URL = "http://localhost:30001" + + +class _Process: + def __init__(self, exit_code: int | None = None) -> None: + self.exit_code = exit_code + + def poll(self) -> int | None: + return self.exit_code + + +class _Response: + def __init__(self, status_code: int = 200, text: str = "") -> None: + self.status_code = status_code + self.text = text + + +class _Engine: + """Records the control calls a park or wake makes.""" + + def __init__(self, flush_statuses: list[int] | None = None) -> None: + self.calls: list[tuple[str, str, Any]] = [] + self.flush_statuses = list(flush_statuses or [200]) + self.status: dict[str, int] = {} + + def post(self, url: str, *, json: Any = None, params: Any = None, timeout: float) -> _Response: + path = url.removeprefix(URL) + self.calls.append(("POST", path, json if json is not None else params)) + if path == "/flush_cache": + return _Response(self.flush_statuses.pop(0) if self.flush_statuses else 200) + return _Response(self.status.get(path, 200)) + + def get(self, url: str, *, timeout: float) -> _Response: + path = url.removeprefix(URL) + self.calls.append(("GET", path, None)) + return _Response(self.status.get(path, 200)) + + +@pytest.fixture +def engine(monkeypatch: pytest.MonkeyPatch) -> _Engine: + fake = _Engine() + monkeypatch.setattr(_server.requests, "post", fake.post) + monkeypatch.setattr(_server.requests, "get", fake.get) + return fake + + +def _started(adapter: Any) -> Any: + """An adapter as ``load`` leaves it, without starting an engine.""" + adapter._memory_saver_launch_args() + adapter._server_url = URL + adapter._process = _Process() + return adapter + + +@pytest.fixture +def generation(monkeypatch: pytest.MonkeyPatch) -> SGLangGenerationAdapter: + monkeypatch.delenv(HOST_CACHE_ENV, raising=False) + monkeypatch.setattr(_server, "checkpoint_weight_bytes", lambda *_: 8 * 1024**3) + return _started(SGLangGenerationAdapter("org/model", mem_fraction_static=0.8)) + + +def test_launch_flags_follow_the_host_tier_setting(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv(HOST_CACHE_ENV, raising=False) + assert SGLangGenerationAdapter("org/model")._memory_saver_launch_args() == [ + "--enable-memory-saver", + "--enable-weights-cpu-backup", + ] + assert SGLangEmbeddingAdapter("org/model")._memory_saver_launch_args() == list(_server.MEMORY_SAVER_LAUNCH_ARGS) + + monkeypatch.setenv(HOST_CACHE_ENV, "0") + assert SGLangGenerationAdapter("org/model")._memory_saver_launch_args() == [] + + monkeypatch.setenv(HOST_CACHE_ENV, "16GiB") + assert SGLangGenerationAdapter("org/model", host_offload=False)._memory_saver_launch_args() == [] + wide = SGLangGenerationAdapter( + "org/model", tensor_parallel_size=2, request_read_timeout_s=60.0, startup_timeout_s=600.0 + ) + assert wide._memory_saver_launch_args() == [] + + +def test_host_offload_must_be_a_boolean() -> None: + with pytest.raises(ValueError, match="host_offload"): + SGLangGenerationAdapter("org/model", host_offload="yes") # type: ignore[arg-type] + with pytest.raises(ValueError, match="host_offload"): + SGLangEmbeddingAdapter("org/model", host_offload=1) # type: ignore[arg-type] + + +def test_only_a_running_engine_launched_with_the_memory_saver_can_park( + generation: SGLangGenerationAdapter, +) -> None: + assert generation.host_offload_bytes() == 8 * 1024**3 + generation._process = _Process(exit_code=1) + assert generation.host_offload_bytes() is None + generation._process = _Process() + generation._memory_saver = False + assert generation.host_offload_bytes() is None + generation._memory_saver = True + generation._server_url = None + assert generation.host_offload_bytes() is None + + +def test_park_waits_for_idle_then_releases(generation: SGLangGenerationAdapter, engine: _Engine) -> None: + offload = generation.offload_to_host() + assert offload.host_bytes == 8 * 1024**3 + assert offload.retains_device_memory + assert [(method, path) for method, path, _ in engine.calls] == [ + ("POST", "/flush_cache"), + ("POST", "/release_memory_occupation"), + ] + assert engine.calls[0][2] == {"timeout": _server.PARK_IDLE_TIMEOUT_S} + # Parked: requests are told to retry, and it cannot park twice. + with pytest.raises(GenerationDrainingError): + generation._check_loaded() + assert generation.host_offload_bytes() is None + + +def test_park_aborts_requests_that_outlive_the_wait(generation: SGLangGenerationAdapter, engine: _Engine) -> None: + engine.flush_statuses = [400, 200] + generation.offload_to_host() + assert [path for _, path, _ in engine.calls] == [ + "/flush_cache", + "/abort_request", + "/flush_cache", + "/release_memory_occupation", + ] + assert engine.calls[1][2] == {"abort_all": True} + + +def test_park_fails_when_the_engine_stays_busy(generation: SGLangGenerationAdapter, engine: _Engine) -> None: + engine.flush_statuses = [400, 400] + with pytest.raises(RuntimeError, match="did not become idle"): + generation.offload_to_host() + assert "/release_memory_occupation" not in [path for _, path, _ in engine.calls] + + +def test_park_fails_when_the_release_is_refused(generation: SGLangGenerationAdapter, engine: _Engine) -> None: + engine.status["/release_memory_occupation"] = 400 + with pytest.raises(RuntimeError, match="release_memory_occupation failed"): + generation.offload_to_host() + + +def test_wake_resumes_checks_health_and_keeps_the_host_copy( + generation: SGLangGenerationAdapter, engine: _Engine +) -> None: + generation.offload_to_host() + engine.calls.clear() + assert generation.restore_from_host() == 8 * 1024**3 + assert [(method, path) for method, path, _ in engine.calls] == [ + ("POST", "/resume_memory_occupation"), + ("GET", "/health"), + ] + generation._check_loaded() + # The engine keeps its host buffer, so parking again costs what it holds. + assert generation.host_offload_bytes() == 8 * 1024**3 + + +def test_wake_fails_loudly(generation: SGLangGenerationAdapter, engine: _Engine) -> None: + generation.offload_to_host() + engine.status["/health"] = 503 + with pytest.raises(RuntimeError, match="health check after resume"): + generation.restore_from_host() + with pytest.raises(RuntimeError, match="not parked"): + SGLangGenerationAdapter("org/model").restore_from_host() + + +def test_wake_refuses_an_engine_that_exited(generation: SGLangGenerationAdapter, engine: _Engine) -> None: + generation.offload_to_host() + generation._process = _Process(exit_code=-9) + with pytest.raises(RuntimeError, match="not running"): + generation.restore_from_host() + + +def test_parked_embedding_engine_refuses_requests(monkeypatch: pytest.MonkeyPatch, engine: _Engine) -> None: + monkeypatch.delenv(HOST_CACHE_ENV, raising=False) + monkeypatch.setattr(_server, "checkpoint_weight_bytes", lambda *_: 1024) + adapter = _started(SGLangEmbeddingAdapter("org/model")) + assert adapter.offload_to_host().host_bytes == 1024 + with pytest.raises(RuntimeError, match=re.escape(ERR_NOT_LOADED)): + adapter._check_loaded() + adapter.restore_from_host() + adapter._check_loaded() + + +def test_checkpoint_weight_bytes_sums_local_safetensors(tmp_path: Path) -> None: + (tmp_path / "model-00001-of-00002.safetensors").write_bytes(b"x" * 100) + (tmp_path / "model-00002-of-00002.safetensors").write_bytes(b"x" * 50) + (tmp_path / "config.json").write_text("{}") + assert _server.checkpoint_weight_bytes(str(tmp_path), None) == 150 + + bin_only = tmp_path / "bin" + bin_only.mkdir() + (bin_only / "pytorch_model.bin").write_bytes(b"x" * 70) + assert _server.checkpoint_weight_bytes(str(bin_only), None) == 70 + + empty = tmp_path / "empty" + empty.mkdir() + assert _server.checkpoint_weight_bytes(str(empty), None) is None + + +def test_checkpoint_weight_bytes_is_none_for_an_uncached_repo(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + monkeypatch.setenv("SIE_LOCAL_CACHE", str(tmp_path)) + monkeypatch.setenv("HF_HOME", str(tmp_path)) + assert _server.checkpoint_weight_bytes("org/never-downloaded", None) is None + + +def test_requests_errors_propagate_so_the_registry_unloads( + generation: SGLangGenerationAdapter, monkeypatch: pytest.MonkeyPatch +) -> None: + def refused(*_: Any, **__: Any) -> None: + raise requests.ConnectionError("engine gone") + + monkeypatch.setattr(_server.requests, "post", refused) + with pytest.raises(requests.ConnectionError): + generation.offload_to_host() + assert isinstance(generation._process, _Process | subprocess.Popen) diff --git a/packages/sie_server/tests/api/test_encode_endpoint.py b/packages/sie_server/tests/api/test_encode_endpoint.py index 25435a7f9..95e34b8fd 100644 --- a/packages/sie_server/tests/api/test_encode_endpoint.py +++ b/packages/sie_server/tests/api/test_encode_endpoint.py @@ -82,6 +82,8 @@ def mock_registry(mock_adapter: MagicMock) -> MagicMock: from sie_server.core.postprocessor_registry import PostprocessorRegistry registry = MagicMock(spec=ModelRegistry) + # Nothing is parked in the host-RAM tier. + registry.wake_parked.return_value = False registry.has_model.return_value = True registry.is_loaded.return_value = True registry.is_loading.return_value = False @@ -307,6 +309,23 @@ async def start_load_async_success(*args: Any, **kwargs: Any) -> bool: assert response.status_code == 503 mock_registry.start_load_async.assert_called_once_with("test-model", device="cpu") + def test_encode_waits_for_a_parked_model_to_wake( + self, client: TestClient, mock_registry: MagicMock, mock_adapter: MagicMock + ) -> None: + """A model parked in host RAM is woken inline instead of answering MODEL_LOADING.""" + mock_registry.is_loaded.return_value = False + mock_registry.is_loading.return_value = False + mock_registry.wake_parked.return_value = True + mock_registry.start_load_async = MagicMock() + response = client.post( + "/v1/encode/test-model", + json={"items": [{"text": "Hello"}]}, + headers=JSON_HEADERS, + ) + assert response.status_code == 200 + mock_registry.wake_parked.assert_awaited_once_with("test-model", "cpu") + mock_registry.start_load_async.assert_not_called() + def test_encode_unsupported_output_type(self, client: TestClient, mock_registry: MagicMock) -> None: """Returns 400 for unsupported output type.""" mock_registry.get_config.return_value = ModelConfig( diff --git a/packages/sie_server/tests/api/test_extract.py b/packages/sie_server/tests/api/test_extract.py index ef38cebee..c244a43ff 100644 --- a/packages/sie_server/tests/api/test_extract.py +++ b/packages/sie_server/tests/api/test_extract.py @@ -145,6 +145,8 @@ async def mock_submit_extract( def mock_registry(mock_adapter: MagicMock) -> MagicMock: """Create a mock registry with an extraction model.""" registry = MagicMock(spec=ModelRegistry) + # Nothing is parked in the host-RAM tier. + registry.wake_parked.return_value = False registry.has_model.return_value = True registry.is_loaded.return_value = True registry.is_loading.return_value = False diff --git a/packages/sie_server/tests/api/test_openai_compat.py b/packages/sie_server/tests/api/test_openai_compat.py index a1e51a986..51b1ac90d 100644 --- a/packages/sie_server/tests/api/test_openai_compat.py +++ b/packages/sie_server/tests/api/test_openai_compat.py @@ -84,6 +84,8 @@ def mock_adapter() -> MagicMock: def mock_registry(mock_adapter: MagicMock) -> MagicMock: """Create a mock registry.""" registry = MagicMock(spec=ModelRegistry) + # Nothing is parked in the host-RAM tier. + registry.wake_parked.return_value = False registry.has_model.return_value = True registry.is_loaded.return_value = True registry.is_loading.return_value = False diff --git a/packages/sie_server/tests/api/test_score.py b/packages/sie_server/tests/api/test_score.py index dfc0a9b63..621d01170 100644 --- a/packages/sie_server/tests/api/test_score.py +++ b/packages/sie_server/tests/api/test_score.py @@ -119,6 +119,8 @@ def mock_encoder_adapter() -> MagicMock: def mock_registry(mock_adapter: MagicMock) -> MagicMock: """Create a mock registry with a scoring model.""" registry = MagicMock(spec=ModelRegistry) + # Nothing is parked in the host-RAM tier. + registry.wake_parked.return_value = False registry.has_model.return_value = True registry.is_loaded.return_value = True registry.is_loading.return_value = False diff --git a/packages/sie_server/tests/core/test_host_cache.py b/packages/sie_server/tests/core/test_host_cache.py new file mode 100644 index 000000000..a7a3d1fad --- /dev/null +++ b/packages/sie_server/tests/core/test_host_cache.py @@ -0,0 +1,154 @@ +"""Host-RAM tier bookkeeping: settings, budget resolution and LRU admission.""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +import pytest +from sie_server.core import host_cache as hc +from sie_server.core.host_cache import HostCache, ParkedModel + +GIB = 1024**3 +MIB = 1024**2 + + +def _entry(name: str, host_bytes: int, *, device: str = "cuda:0", holds_device: bool = False) -> ParkedModel: + return ParkedModel( + name=name, + loaded=SimpleNamespace(device=device), + host_bytes=host_bytes, + retains_device_memory=holds_device, + ) + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + ("auto", None), + (" AUTO ", None), + ("0", 0), + (0, 0), + (1048576, 1048576), + ("1048576", 1048576), + ("8GiB", 8 * GIB), + ("8Gi", 8 * GIB), + ("1.5GiB", int(1.5 * GIB)), + ("512 MiB", 512 * MIB), + ("2GB", 2 * 1000**3), + ], +) +def test_parse_host_cache_setting(raw: object, expected: int | None) -> None: + assert hc.parse_host_cache_setting(raw) == expected + + +@pytest.mark.parametrize("raw", ["", "lots", "1.5", "-1", "8XB", "GiB", -5, True, 2.5]) +def test_parse_host_cache_setting_rejects_garbage(raw: object) -> None: + with pytest.raises(ValueError, match=hc.HOST_CACHE_ENV): + hc.parse_host_cache_setting(raw) + + +def test_auto_is_off_without_a_cuda_device() -> None: + assert hc.resolve_host_cache_budget("auto", ["cpu"], memory_limit=64 * GIB) == 0 + assert hc.resolve_host_cache_budget("auto", ["mps"], memory_limit=64 * GIB) == 0 + + +def test_auto_is_a_quarter_of_the_memory_limit_on_cuda() -> None: + assert hc.resolve_host_cache_budget("auto", ["cuda:0"], memory_limit=32 * GIB) == 8 * GIB + assert hc.resolve_host_cache_budget("auto", ["cuda:0", "cuda:1"], memory_limit=128 * GIB) == 32 * GIB + + +def test_explicit_size_applies_to_any_device() -> None: + assert hc.resolve_host_cache_budget("1GiB", ["cpu"]) == GIB + assert hc.resolve_host_cache_budget(0, ["cuda:0"]) == 0 + + +def test_container_limit_prefers_cgroup_v2(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(hc, "physical_memory_bytes", lambda: 64 * GIB) + v2 = tmp_path / "memory.max" + v1 = tmp_path / "memory.limit_in_bytes" + v2.write_text(f"{32 * GIB}\n") + v1.write_text(f"{16 * GIB}\n") + assert hc.container_memory_limit_bytes(cgroup_v2=v2, cgroup_v1=v1) == 32 * GIB + + +def test_container_limit_falls_back_to_v1_then_physical(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(hc, "physical_memory_bytes", lambda: 64 * GIB) + v2 = tmp_path / "memory.max" + v1 = tmp_path / "memory.limit_in_bytes" + v2.write_text("max\n") + v1.write_text(f"{16 * GIB}\n") + assert hc.container_memory_limit_bytes(cgroup_v2=v2, cgroup_v1=v1) == 16 * GIB + # cgroup v1 reports "unlimited" as a huge number. + v1.write_text(f"{1 << 62}\n") + assert hc.container_memory_limit_bytes(cgroup_v2=v2, cgroup_v1=v1) == 64 * GIB + # A limit above the machine is capped at the machine. + v2.write_text(f"{128 * GIB}\n") + assert hc.container_memory_limit_bytes(cgroup_v2=v2, cgroup_v1=v1) == 64 * GIB + + +def test_host_cache_configured_reads_the_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv(hc.HOST_CACHE_ENV, raising=False) + assert hc.host_cache_configured() + monkeypatch.setenv(hc.HOST_CACHE_ENV, "0") + assert not hc.host_cache_configured() + monkeypatch.setenv(hc.HOST_CACHE_ENV, "4GiB") + assert hc.host_cache_configured() + monkeypatch.setenv(hc.HOST_CACHE_ENV, "garbage") + assert not hc.host_cache_configured() + + +def test_disabled_cache_admits_nothing() -> None: + cache = HostCache(0) + assert not cache.enabled + assert cache.admission_victims("a", MIB) is None + + +def test_admission_drops_least_recently_parked_first() -> None: + cache = HostCache(100) + cache.put(_entry("a", 40)) + cache.put(_entry("b", 40)) + assert cache.admission_victims("c", 20) == [] + assert cache.admission_victims("c", 50) == ["a"] + assert cache.admission_victims("c", 100) == ["a", "b"] + assert cache.admission_victims("c", 101) is None + assert cache.used_bytes == 80 + + +def test_retained_bytes_count_against_the_budget_but_are_never_victims() -> None: + cache = HostCache(100) + cache.retain("resident", 60) + cache.put(_entry("a", 30)) + assert cache.used_bytes == 90 + assert cache.admission_victims("b", 40) == ["a"] + assert cache.admission_victims("b", 41) is None + # Parking the resident model replaces its retained charge. + assert cache.admission_victims("resident", 70) == [] + cache.put(_entry("resident", 70)) + assert cache.retained_bytes("resident") == 0 + assert cache.used_bytes == 100 + + +def test_take_removes_the_charge_and_release_forgets_retention() -> None: + cache = HostCache(100) + cache.put(_entry("a", 30)) + entry = cache.take("a") + assert entry is not None + assert entry.host_bytes == 30 + assert "a" not in cache + assert cache.used_bytes == 0 + cache.retain("a", 30) + assert cache.used_bytes == 30 + cache.release("a") + assert cache.used_bytes == 0 + assert cache.take("a") is None + + +def test_holding_device_memory_filters_by_device_and_flag() -> None: + cache = HostCache(1000) + cache.put(_entry("torch", 10, device="cuda:0")) + cache.put(_entry("engine-0", 10, device="cuda:0", holds_device=True)) + cache.put(_entry("engine-1", 10, device="cuda:1", holds_device=True)) + assert cache.holding_device_memory("cuda:0") == ["engine-0"] + assert cache.holding_device_memory("cuda:1") == ["engine-1"] + assert cache.names == ["torch", "engine-0", "engine-1"] diff --git a/packages/sie_server/tests/core/test_registry_host_cache.py b/packages/sie_server/tests/core/test_registry_host_cache.py new file mode 100644 index 000000000..ecb7db505 --- /dev/null +++ b/packages/sie_server/tests/core/test_registry_host_cache.py @@ -0,0 +1,297 @@ +"""Registry park and wake through the host-RAM tier, with weightless fakes. + +The synthetic memory tracker (``SIE_FAKE_MEMORY_BUDGET``) makes device usage +the sum of the fakes' declared footprints, and the fake adapter simulates +parking by moving that footprint to the host tier. Budget 150 MiB on the +device; ``sie-fake:small-a``/``small-b`` declare 64 MiB, ``sie-fake`` 128 MiB. +""" + +from __future__ import annotations + +import asyncio +from collections import Counter +from pathlib import Path +from typing import Any + +import pytest +from sie_server.adapters.base import HostOffload +from sie_server.config.engine import EngineConfig +from sie_server.core import registry as registry_module +from sie_server.core.loader import load_model_configs +from sie_server.core.memory import SIE_FAKE_MEMORY_BUDGET_ENV +from sie_server.core.registry import ModelRegistry + +MODELS_DIR = Path(__file__).resolve().parents[2] / "models" +MIB = 1024**2 +SMALL_A = "sie-fake:small-a" +SMALL_B = "sie-fake:small-b" +BIG = "sie-fake" + + +class _RecordingTelemetry: + def __init__(self) -> None: + self.lookups: Counter[tuple[str, bool]] = Counter() + self.wakes: list[tuple[str, str]] = [] + self.host_bytes: dict[str, int] = {} + self.load_completions: list[str] = [] + + def host_cache_lookup(self, *, model: str, hit: bool, **_: Any) -> None: + self.lookups[(model, hit)] += 1 + + def host_cache_wake_completed(self, *, model: str, outcome: str, **_: Any) -> None: + self.wakes.append((model, outcome)) + + def host_cache_memory_changed(self, *, model: str, host_bytes: int, **_: Any) -> None: + self.host_bytes[model] = host_bytes + + def model_load_completed(self, *, model: str, **_: Any) -> None: + self.load_completions.append(model) + + def __getattr__(self, name: str) -> Any: + return lambda **_: None + + +@pytest.fixture +def telemetry(monkeypatch: pytest.MonkeyPatch) -> _RecordingTelemetry: + recorder = _RecordingTelemetry() + monkeypatch.setattr(registry_module, "worker_telemetry", lambda: recorder) + return recorder + + +def _registry(monkeypatch: pytest.MonkeyPatch, *, host_budget: str = "256MiB") -> ModelRegistry: + monkeypatch.setenv(SIE_FAKE_MEMORY_BUDGET_ENV, "150MiB") + configs = load_model_configs(MODELS_DIR) + registry = ModelRegistry(engine_config=EngineConfig(host_cache_bytes=host_budget, disk_cache_enabled=False)) + registry.add_config(configs["sie-fake"]) + return registry + + +async def test_eviction_parks_and_the_next_load_wakes_the_same_adapter( + monkeypatch: pytest.MonkeyPatch, telemetry: _RecordingTelemetry +) -> None: + registry = _registry(monkeypatch) + adapter_a = await registry.load_async(SMALL_A, device="cpu") + await registry.load_async(SMALL_B, device="cpu") + + # Loading the big fake needs both small ones out of the way: they park. + await registry.load_async(BIG, device="cpu") + assert registry.loaded_model_names == [BIG] + assert registry.parked_model_names == [SMALL_A, SMALL_B] + assert registry.is_parked(SMALL_A) + assert registry.host_cache.used_bytes == 128 * MIB + assert telemetry.host_bytes[SMALL_A] == 64 * MIB + + # Waking small-a evicts (parks) the big one and brings back the very same + # adapter object, with no new load. + load_count = len(telemetry.load_completions) + woken = await registry.load_async(SMALL_A, device="cpu") + assert woken is adapter_a + assert registry.is_loaded(SMALL_A) + assert not registry.is_parked(SMALL_A) + assert registry.parked_model_names == [SMALL_B, BIG] + assert len(telemetry.load_completions) == load_count + assert telemetry.wakes == [(SMALL_A, "success")] + assert telemetry.lookups[(SMALL_A, True)] == 1 + assert telemetry.lookups[(SMALL_A, False)] == 1 + assert telemetry.host_bytes[SMALL_A] == 0 + assert registry.host_cache.used_bytes == 192 * MIB + # The woken model serves. + worker = await registry.start_worker(SMALL_A) + assert worker.is_running + + +async def test_budget_drops_the_least_recently_parked( + monkeypatch: pytest.MonkeyPatch, telemetry: _RecordingTelemetry +) -> None: + registry = _registry(monkeypatch, host_budget="128MiB") + await registry.load_async(SMALL_A, device="cpu") + await registry.load_async(SMALL_B, device="cpu") + await registry.load_async(BIG, device="cpu") + assert registry.parked_model_names == [SMALL_A, SMALL_B] + + # Parking the big one (128 MiB) takes the whole budget: both small ones go. + await registry.load_async("sie-fake:chat", device="cpu") + assert registry.parked_model_names == [BIG] + assert registry.host_cache.used_bytes == 128 * MIB + assert telemetry.host_bytes[SMALL_A] == 0 + assert telemetry.host_bytes[SMALL_B] == 0 + + +async def test_a_waking_copy_stays_charged_until_restored(monkeypatch: pytest.MonkeyPatch) -> None: + registry = _registry(monkeypatch, host_budget="128MiB") + await registry.load_async(SMALL_A, device="cpu") + await registry.load_async(SMALL_B, device="cpu") + await registry.load_async(BIG, device="cpu") + + # small-a's host copy exists until its restore finishes, so the big model + # evicted to make room for it does not fit beside small-b and is unloaded. + await registry.load_async(SMALL_A, device="cpu") + assert registry.parked_model_names == [SMALL_B] + assert not registry.is_loaded(BIG) + assert registry.host_cache.used_bytes == 64 * MIB + + +async def test_a_model_larger_than_the_budget_is_unloaded(monkeypatch: pytest.MonkeyPatch) -> None: + registry = _registry(monkeypatch, host_budget="100MiB") + await registry.load_async(BIG, device="cpu") + await registry.load_async(SMALL_A, device="cpu") + await registry.load_async(SMALL_B, device="cpu") + assert BIG not in registry.parked_model_names + assert not registry.is_loaded(BIG) + + +async def test_disabled_tier_unloads_as_before(monkeypatch: pytest.MonkeyPatch, telemetry: _RecordingTelemetry) -> None: + registry = _registry(monkeypatch, host_budget="0") + await registry.load_async(SMALL_A, device="cpu") + await registry.load_async(SMALL_B, device="cpu") + await registry.load_async(BIG, device="cpu") + assert registry.parked_model_names == [] + assert not registry.host_cache.enabled + assert not telemetry.lookups + + +async def test_a_failed_wake_falls_back_to_a_cold_load( + monkeypatch: pytest.MonkeyPatch, telemetry: _RecordingTelemetry +) -> None: + registry = _registry(monkeypatch) + adapter_a = await registry.load_async(SMALL_A, device="cpu") + await registry.load_async(SMALL_B, device="cpu") + await registry.load_async(BIG, device="cpu") + + unloaded: list[bool] = [] + + def broken_restore() -> int: + raise RuntimeError("restore exploded") + + monkeypatch.setattr(adapter_a, "restore_from_host", broken_restore) + original_unload = adapter_a.unload + monkeypatch.setattr(adapter_a, "unload", lambda: (unloaded.append(True), original_unload())) + + reloaded = await registry.load_async(SMALL_A, device="cpu") + assert reloaded is not adapter_a + assert registry.is_loaded(SMALL_A) + assert unloaded == [True] + assert telemetry.wakes == [(SMALL_A, "error")] + assert SMALL_A not in registry.parked_model_names + + +async def test_a_failed_park_unloads_the_model(monkeypatch: pytest.MonkeyPatch) -> None: + registry = _registry(monkeypatch) + adapter_a = await registry.load_async(SMALL_A, device="cpu") + await registry.load_async(SMALL_B, device="cpu") + + def broken_offload() -> HostOffload: + raise RuntimeError("offload exploded") + + monkeypatch.setattr(adapter_a, "offload_to_host", broken_offload) + await registry.load_async(BIG, device="cpu") + assert registry.parked_model_names == [SMALL_B] + assert not registry.is_loaded(SMALL_A) + + +async def test_config_change_drops_the_parked_copy(monkeypatch: pytest.MonkeyPatch) -> None: + registry = _registry(monkeypatch) + await registry.load_async(SMALL_A, device="cpu") + await registry.load_async(SMALL_B, device="cpu") + await registry.load_async(BIG, device="cpu") + assert SMALL_A in registry.parked_model_names + + config = load_model_configs(MODELS_DIR)["sie-fake"] + changed = config.model_copy(update={"max_sequence_length": (config.max_sequence_length or 512) + 1}) + await registry.add_config_async(changed) + assert registry.parked_model_names == [] + assert registry.host_cache.used_bytes == 0 + + +async def test_manual_unload_of_a_parked_model_frees_it(monkeypatch: pytest.MonkeyPatch) -> None: + registry = _registry(monkeypatch) + await registry.load_async(SMALL_A, device="cpu") + await registry.load_async(SMALL_B, device="cpu") + await registry.load_async(BIG, device="cpu") + + await registry.unload_async(SMALL_A) + assert registry.parked_model_names == [SMALL_B] + with pytest.raises(KeyError, match="not loaded"): + await registry.unload_async(SMALL_A) + + # A manual unload of a resident model never parks it. + await registry.unload_async(BIG) + assert registry.parked_model_names == [SMALL_B] + + +async def test_shutdown_drops_every_parked_model(monkeypatch: pytest.MonkeyPatch) -> None: + registry = _registry(monkeypatch) + await registry.load_async(SMALL_A, device="cpu") + await registry.load_async(SMALL_B, device="cpu") + await registry.load_async(BIG, device="cpu") + await registry.unload_all_async() + assert registry.loaded_model_names == [] + assert registry.parked_model_names == [] + assert registry.host_cache.used_bytes == 0 + + +async def test_parked_engines_holding_device_memory_are_dropped_when_nothing_resident_remains( + monkeypatch: pytest.MonkeyPatch, +) -> None: + registry = _registry(monkeypatch) + adapter_a = await registry.load_async(SMALL_A, device="cpu") + original_offload = adapter_a.offload_to_host + + def engine_offload() -> HostOffload: + offload = original_offload() + return HostOffload(host_bytes=offload.host_bytes, retains_device_memory=True) + + monkeypatch.setattr(adapter_a, "offload_to_host", engine_offload) + await registry.load_async(SMALL_B, device="cpu") + await registry.load_async(BIG, device="cpu") + assert registry.host_cache.holding_device_memory("cpu") == [SMALL_A] + + # Pretend the device is still short after every resident model left: the + # parked engine is the only thing left that holds memory there. + manager = registry.memory_managers["cpu"] + calls = {"n": 0} + original = manager.should_evict_for_load + + def short_until_engine_dropped(required: int | None = None) -> bool: + calls["n"] += 1 + return SMALL_A in registry.parked_model_names or original(required) + + monkeypatch.setattr(manager, "should_evict_for_load", short_until_engine_dropped) + await registry.load_async(SMALL_B, device="cpu") + assert SMALL_A not in registry.parked_model_names + assert registry.is_loaded(SMALL_B) + + +async def test_wake_parked_waits_for_the_wake(monkeypatch: pytest.MonkeyPatch) -> None: + registry = _registry(monkeypatch) + await registry.load_async(SMALL_A, device="cpu") + await registry.load_async(SMALL_B, device="cpu") + await registry.load_async(BIG, device="cpu") + + assert not await registry.wake_parked("sie-fake:chat", "cpu") # never parked + assert await registry.wake_parked(SMALL_A, "cpu") + assert registry.is_loaded(SMALL_A) + assert await registry.wake_parked(SMALL_A, "cpu") # already loaded + + +async def test_cancelled_wake_before_restore_keeps_the_parked_copy(monkeypatch: pytest.MonkeyPatch) -> None: + registry = _registry(monkeypatch) + await registry.load_async(SMALL_A, device="cpu") + await registry.load_async(SMALL_B, device="cpu") + await registry.load_async(BIG, device="cpu") + + blocked = asyncio.Event() + + async def stuck_make_room(name: str, entry: Any) -> Any: + blocked.set() + await asyncio.Event().wait() + + monkeypatch.setattr(registry, "_make_room_to_wake", stuck_make_room) + task = asyncio.create_task(registry.load_async(SMALL_A, device="cpu")) + await blocked.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert SMALL_A in registry.parked_model_names + assert registry.host_cache.retained_bytes(SMALL_A) == 0 + assert registry.host_cache.used_bytes == 128 * MIB diff --git a/packages/sie_server/tests/observability/test_worker_telemetry.py b/packages/sie_server/tests/observability/test_worker_telemetry.py index 187305740..fa088d436 100644 --- a/packages/sie_server/tests/observability/test_worker_telemetry.py +++ b/packages/sie_server/tests/observability/test_worker_telemetry.py @@ -315,6 +315,41 @@ def test_lifecycle_native_topology_and_oom_attributes_match_contract( provider.shutdown() +def test_host_cache_instruments_match_contract( + active_telemetry: tuple[wt.WorkerTelemetry, InMemoryMetricReader, MeterProvider], +) -> None: + telemetry, reader, provider = active_telemetry + telemetry.host_cache_lookup(model="catalog/model", hit=True) + telemetry.host_cache_lookup(model="catalog/model", hit=False) + telemetry.host_cache_lookup(model="catalog/model", hit=False) + telemetry.host_cache_wake_completed(model="catalog/model", duration_s=0.4, outcome="success") + telemetry.host_cache_wake_completed(model="catalog/model", duration_s="not a number", outcome="success") + telemetry.host_cache_wake_completed(model="catalog/model", duration_s=2.0, outcome="exploded") + telemetry.host_cache_memory_changed(model="catalog/model", host_bytes=8192) + + by_name = _metric_map(reader.get_metrics_data()) + model_attributes = { + "backend": "python", + "lane": "realtime|l4|default", + "model": "catalog/model", + "profile": "default", + } + lookups = {point.attributes["result"]: point.value for point in _points(by_name[wt.HOST_CACHE_LOOKUPS_METRIC_NAME])} + assert lookups == {"hit": 1, "miss": 2} + assert dict(_points(by_name[wt.HOST_CACHE_LOOKUPS_METRIC_NAME])[0].attributes).keys() == { + "result", + *model_attributes, + } + wakes = {point.attributes["outcome"]: point for point in _points(by_name[wt.HOST_CACHE_WAKE_DURATION_METRIC_NAME])} + assert set(wakes) == {"success", "other"} + assert wakes["success"].count == 1 + assert tuple(wakes["success"].explicit_bounds) == wt.HOST_CACHE_WAKE_BUCKETS_S + memory = _points(by_name[wt.HOST_CACHE_MEMORY_METRIC_NAME])[0] + assert memory.value == 8192 + assert dict(memory.attributes) == model_attributes + provider.shutdown() + + def test_unknown_values_collapse_and_never_become_attributes( active_telemetry: tuple[wt.WorkerTelemetry, InMemoryMetricReader, MeterProvider], ) -> None: @@ -551,6 +586,9 @@ def test_metric_inventory_is_exact() -> None: "sie.worker.model.load.duration", "sie.worker.model.memory", "sie.worker.model.evictions", + "sie.worker.model.host_cache.lookups", + "sie.worker.model.host_cache.wake.duration", + "sie.worker.model.host_cache.memory", "sie.worker.oom.recoveries", "sie.worker.scheduler.adaptive.wait", "sie.worker.scheduler.adaptive.cost", diff --git a/packages/sie_server/tests/test_queue_executor.py b/packages/sie_server/tests/test_queue_executor.py index ba05e4f44..2d2ee1073 100644 --- a/packages/sie_server/tests/test_queue_executor.py +++ b/packages/sie_server/tests/test_queue_executor.py @@ -152,6 +152,8 @@ def _make_registry(*, loaded: bool = True, loading: bool = False) -> MagicMock: # which ``ensure_model_ready`` now reads to gate the terminal ``failed`` # state (#1786 fast-path). Tests that want the terminal path set this. reg.get_failure.return_value = None + # Nothing parked in the host-RAM tier unless a test parks it. + reg.is_parked.return_value = False return reg @@ -299,6 +301,25 @@ async def test_already_loading_returns_in_progress(self) -> None: assert await ex.ensure_model_ready("test/model") == "loading_in_progress" reg.start_load_async.assert_not_called() + @pytest.mark.asyncio + async def test_parked_model_is_woken_and_ready(self) -> None: + reg = _make_registry(loaded=False, loading=False) + reg.is_parked.return_value = True + reg.wake_parked = AsyncMock(return_value=True) + reg.start_load_async = AsyncMock(return_value=True) + ex = QueueExecutor(reg) + assert await ex.ensure_model_ready("test/model") == "ready" + reg.wake_parked.assert_awaited_once_with("test/model", "cpu") + reg.start_load_async.assert_not_called() + + @pytest.mark.asyncio + async def test_parked_model_still_waking_reports_loading(self) -> None: + reg = _make_registry(loaded=False, loading=True) + reg.is_parked.return_value = True + reg.wake_parked = AsyncMock(return_value=False) + ex = QueueExecutor(reg) + assert await ex.ensure_model_ready("test/model") == "loading_in_progress" + @pytest.mark.asyncio async def test_triggers_new_load(self) -> None: reg = _make_registry(loaded=False, loading=False) diff --git a/telemetry/README.md b/telemetry/README.md index 59087a245..b2468acd6 100644 --- a/telemetry/README.md +++ b/telemetry/README.md @@ -12,7 +12,7 @@ conditionally owned queue/adaptive-scheduler surfaces. Only authoritative lane lifecycle signals remain excluded until a deployed producer owns them; undeclared new instruments are not permitted. -The current inventory is 122 application families: 120 may reach the remote +The current inventory is 125 application families: 123 may reach the remote OTLP branch and two remain Prometheus-only controls. Ten additional, exact collector-self families form a separate operational allowlist without turning collector self-telemetry into a version-dependent wildcard. @@ -359,6 +359,17 @@ process and collapses the rest to `other`, so `sie.worker.upstream.refusals` retains at most `17 × 4 = 68` series (three refusal reasons and `other`) and `sie.worker.upstream.breaker.open` at most 17. +The Python worker's host-RAM tier instruments use the catalog pair like the +other model lifecycle instruments. `sie.worker.model.host_cache.lookups` +counts one lookup per model load, `hit` when the load woke a copy parked in +host RAM and `miss` when it loaded from the disk cache, cluster cache or Hub, +so it retains at most `257 × 3 = 771` series. +`sie.worker.model.host_cache.wake.duration` times each wake by load outcome +(`257 × 4` series), and `sie.worker.model.host_cache.memory` reports the host +bytes the tier charges each model (257 series). It returns to 0 when the copy +is dropped or woken, except for an SGLang engine, which keeps its host copy +after a wake and stays charged until the model is fully unloaded. + Every Rust-engine instrument has an explicit SDK view derived from its checked-in admission tier and finite domains, so contract-valid series do not become `otel.metric.overflow`. Forward output paths use one closed six-value diff --git a/telemetry/contract.yaml b/telemetry/contract.yaml index 597d535aa..a57f5bc1f 100644 --- a/telemetry/contract.yaml +++ b/telemetry/contract.yaml @@ -97,6 +97,8 @@ cardinality_budgets: - worker_oom - worker_runtime - worker_eviction + - worker_host_cache_lookup + - worker_host_cache_wake - worker_scheduler - worker_adaptive_p50 - generation_stream @@ -1015,6 +1017,7 @@ enums: grammar_phase: [request, prewarm, other] grammar_outcome: [success, error, timeout, other] grammar_cache_result: [hit, miss, other] + host_cache_result: [hit, miss, other] generation_admission_outcome: [admitted, rejected, error, other] generation_admission_reason: [none, kv_budget, resolver_error, missing_budget, other] generation_duplicate_path: [first_chunk_fallback, other] @@ -1285,6 +1288,18 @@ attribute_sets: lane: lane model: model profile: profile + worker_host_cache_lookup: + result: host_cache_result + backend: worker_backend + lane: lane + model: model + profile: profile + worker_host_cache_wake: + outcome: model_load_outcome + backend: worker_backend + lane: lane + model: model + profile: profile worker_scheduler: lane: lane model: model @@ -1450,6 +1465,9 @@ histogram_buckets: worker_work_item_age_seconds: [0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0, 600.0, 1800.0] worker_request_seconds: [0.0001, 0.00025, 0.0005, 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0] worker_model_load_seconds: [0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0, 600.0, 900.0] + # A wake copies weights already in host RAM back to the device: tens of + # milliseconds for a small encoder, seconds for an LLM. + worker_model_wake_seconds: [0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0] worker_ipc_seconds: [0.0001, 0.00025, 0.0005, 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 300.0, 900.0] worker_batch_items: [1, 2, 4, 8, 16, 32, 64, 128, 256, 512] worker_batch_cost: [128, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072, 262144] @@ -2482,6 +2500,35 @@ metrics: attributes: worker_eviction prometheus_name: sie_worker_model_evictions_total export: [prometheus, otlp] + # Host-RAM tier: an evicted model parks its weights in host RAM and the next + # load wakes it back onto its device. A lookup is one load: a hit woke a + # parked copy, a miss loaded from the disk cache, cluster cache or Hub. + - name: sie.worker.model.host_cache.lookups + owner: worker_python_engine + facade_event: worker.model_host_cache_lookup + type: counter + unit: "{lookup}" + attributes: worker_host_cache_lookup + prometheus_name: sie_worker_model_host_cache_lookups_total + export: [prometheus, otlp] + - name: sie.worker.model.host_cache.wake.duration + owner: worker_python_engine + facade_event: worker.model_host_cache_wake_completed + type: histogram + unit: s + attributes: worker_host_cache_wake + buckets: worker_model_wake_seconds + prometheus_name: sie_worker_model_host_cache_wake_duration_seconds + export: [prometheus, otlp] + - name: sie.worker.model.host_cache.memory + owner: worker_python_engine + facade_event: worker.model_host_cache_memory_changed + type: gauge + unit: By + attributes: worker_model + prometheus_name: sie_worker_model_host_cache_memory_bytes + export: [prometheus, otlp] + measurement: authoritative_only # Generation is currently implemented by the Python engine. It remains in # the same sie-worker service stream and uses the same bounded catalog From dcbd650fcb8d71ffd25de4d63538ba157f8b6832 Mon Sep 17 00:00:00 2001 From: svonava Date: Mon, 19 Oct 2026 18:03:48 -0700 Subject: [PATCH 2/2] fix(server): keep aliases and LoRA state intact across park and wake Copy each aliased device view to host RAM once and restore it to one device tensor, so the budget charge matches the estimate and aliasing survives. Drop LoRA load placeholders and clear eviction marks left by a park when the model wakes. Bound the inline wake wait in EnsureModelReady by half of SIE_MODEL_READY_TIMEOUT_S when the worker shares it. --- .../src/sie_server/adapters/_host_offload.py | 49 ++++++++++++++----- .../src/sie_server/core/model_loader.py | 9 +++- .../src/sie_server/core/registry.py | 4 +- .../src/sie_server/queue_executor.py | 24 ++++++++- .../tests/adapters/test_host_offload.py | 16 ++++++ .../tests/core/test_registry_host_cache.py | 19 +++++++ .../sie_server/tests/test_queue_executor.py | 12 ++++- 7 files changed, 115 insertions(+), 18 deletions(-) diff --git a/packages/sie_server/src/sie_server/adapters/_host_offload.py b/packages/sie_server/src/sie_server/adapters/_host_offload.py index 2fe702fa0..3e6517052 100644 --- a/packages/sie_server/src/sie_server/adapters/_host_offload.py +++ b/packages/sie_server/src/sie_server/adapters/_host_offload.py @@ -142,6 +142,17 @@ def visit(value: object, slot: _Slot | None, depth: int) -> None: return reach +def _view_key(tensor: torch.Tensor) -> tuple[Any, ...]: + """Identity of a tensor view: two tensor objects with one key alias the same bytes.""" + return ( + tensor.untyped_storage().data_ptr(), + tensor.storage_offset(), + tuple(tensor.shape), + tuple(tensor.stride()), + tensor.dtype, + ) + + def _module_tensors(module: torch.nn.Module) -> list[torch.Tensor]: tensors: list[torch.Tensor] = [] for submodule in module.modules(): @@ -156,7 +167,8 @@ def parkable_bytes(adapter: object, device: torch.device) -> int | None: reach = _collect(adapter, device) if reach.unmovable_bytes: return None - storages: dict[int, int] = {} + # Counted per view, as ``park`` copies them: aliases of one view once. + views: dict[tuple[Any, ...], int] = {} for module in reach.modules: for submodule in module.modules(): if getattr(submodule, "_hf_hook", None) is not None: @@ -165,12 +177,10 @@ def parkable_bytes(adapter: object, device: torch.device) -> int | None: if type(tensor) not in (torch.Tensor, torch.nn.Parameter): return None if _on(device, tensor): - storage = tensor.untyped_storage() - storages[storage.data_ptr()] = storage.nbytes() + views[_view_key(tensor)] = tensor.nbytes for _, tensor in reach.loose: - storage = tensor.untyped_storage() - storages[storage.data_ptr()] = storage.nbytes() - total = sum(storages.values()) + views[_view_key(tensor)] = tensor.nbytes + total = sum(views.values()) return total or None @@ -179,16 +189,25 @@ def park(adapter: object, device: torch.device) -> ParkedTensors: reach = _collect(adapter, device) origins: dict[int, torch.device] = {} copies: list[torch.Tensor] = [] + # Device views already copied. A second tensor object over the same view + # (``self.original_inv_freq = self.inv_freq``, a loose ``weight.data``) + # gets the same host copy, so the bytes are copied and charged once and + # the aliasing survives the wake. + moved: dict[tuple[Any, ...], torch.Tensor] = {} def to_host(tensor: torch.Tensor) -> torch.Tensor: - copy = tensor.detach().to("cpu", copy=True) - origins[copy.untyped_storage().data_ptr()] = tensor.device - copies.append(copy) + key = _view_key(tensor) + copy = moved.get(key) + if copy is None: + copy = tensor.detach().to("cpu", copy=True) + origins[copy.untyped_storage().data_ptr()] = tensor.device + copies.append(copy) + moved[key] = copy return copy def still_on_device(tensor: torch.Tensor) -> bool: - # A parameter shared by two modules, or a loose reference to one, - # has already been moved once by the time it is seen again. + # A parameter shared by two modules, or a loose reference to one, has + # already been rebound to its host copy by the time it is seen again. return _on(device, tensor) and tensor.untyped_storage().data_ptr() not in origins with torch.no_grad(): @@ -211,8 +230,14 @@ def parked_here(tensor: torch.Tensor) -> bool: storage = tensor.untyped_storage() return tensor.device.type == "cpu" and storage.nbytes() > 0 and storage.data_ptr() in parked.origins + # One device tensor per host copy, so aliases stay aliases. + restored: dict[int, torch.Tensor] = {} + def to_device(tensor: torch.Tensor) -> torch.Tensor: - return tensor.to(parked.origins[tensor.untyped_storage().data_ptr()], copy=True) + ptr = tensor.untyped_storage().data_ptr() + if ptr not in restored: + restored[ptr] = tensor.to(parked.origins[ptr], copy=True) + return restored[ptr] with torch.no_grad(): for module in parked.modules: diff --git a/packages/sie_server/src/sie_server/core/model_loader.py b/packages/sie_server/src/sie_server/core/model_loader.py index 2f0dcf261..8941c522d 100644 --- a/packages/sie_server/src/sie_server/core/model_loader.py +++ b/packages/sie_server/src/sie_server/core/model_loader.py @@ -755,7 +755,14 @@ def _finish_load( ) if restore_from is not None: - loaded_model.loras = restore_from.loras + # A LoRA load still in flight when the model parked saw it leave + # the registry and gave up without clearing its placeholder, and a + # LoRA marked for eviction whose unload never ran is still attached. + for source, state in restore_from.loras.items(): + if state.loading: + continue + state.unloading = False + loaded_model.loras[source] = state loaded_model.max_loras = restore_from.max_loras loaded_model.lora_failures = restore_from.lora_failures loaded_model.lora_runtime = restore_from.lora_runtime diff --git a/packages/sie_server/src/sie_server/core/registry.py b/packages/sie_server/src/sie_server/core/registry.py index a8c8f7acb..f9639b7cb 100644 --- a/packages/sie_server/src/sie_server/core/registry.py +++ b/packages/sie_server/src/sie_server/core/registry.py @@ -94,7 +94,7 @@ # How long a request waits for a parked model to wake before it is told to # retry. A wake normally takes seconds; this bounds the wait when the wake # queues behind another model's cold load. -_HOST_WAKE_WAIT_S = 15.0 +HOST_WAKE_WAIT_S = 15.0 class _ConfigChangedDuringLoadError(RuntimeError): @@ -2104,7 +2104,7 @@ async def _restore_parked(self, name: str, entry: ParkedModel, memory_manager: M retained = await self._run_adapter_call(adapter.restore_from_host, name, "Wake") return released_bytes(retained) - async def wake_parked(self, name: str, device: str, *, timeout_s: float = _HOST_WAKE_WAIT_S) -> bool: + async def wake_parked(self, name: str, device: str, *, timeout_s: float = HOST_WAKE_WAIT_S) -> bool: """Start waking a parked model and wait up to ``timeout_s`` for it. For request paths that would otherwise answer "loading, retry": a diff --git a/packages/sie_server/src/sie_server/queue_executor.py b/packages/sie_server/src/sie_server/queue_executor.py index edeb2ecfb..07c5ec101 100644 --- a/packages/sie_server/src/sie_server/queue_executor.py +++ b/packages/sie_server/src/sie_server/queue_executor.py @@ -2,6 +2,7 @@ import asyncio import logging +import math import os import tempfile import time @@ -35,7 +36,7 @@ from sie_server.core.oom import is_oom_error from sie_server.core.pool_isolation import validate_no_legacy_scalar_lora_id from sie_server.core.prepared import AudioPayload, AudioPreparedItem -from sie_server.core.registry import ModelRegistry +from sie_server.core.registry import HOST_WAKE_WAIT_S, ModelRegistry from sie_server.core.runtime_options import merge_runtime_options, merge_runtime_options_with_profile from sie_server.core.score_cost import build_score_prepared_items_timed from sie_server.core.timing import RequestTiming @@ -77,6 +78,23 @@ logger = logging.getLogger(__name__) +# The sidecar's EnsureModelReady deadline. When the worker shares it, a wake +# waited for inline stays well inside it, so the call answers "ready" or +# "loading" instead of timing out. +_MODEL_READY_TIMEOUT_ENV: Final = "SIE_MODEL_READY_TIMEOUT_S" + + +def _host_wake_wait_s() -> float: + raw = os.environ.get(_MODEL_READY_TIMEOUT_ENV, "").strip() + try: + deadline = float(raw) if raw else None + except ValueError: + deadline = None + if deadline is None or not math.isfinite(deadline) or deadline <= 0: + return HOST_WAKE_WAIT_S + return min(HOST_WAKE_WAIT_S, deadline / 2) + + # --------------------------------------------------------------------------- # Queue-path wire error code # --------------------------------------------------------------------------- @@ -882,7 +900,9 @@ async def ensure_model_ready(self, model_id: str) -> ReadinessState: # A model parked in host RAM wakes in seconds: wait for it so the # dispatcher can run this work now instead of re-polling. - if self._registry.is_parked(model_id) and await self._registry.wake_parked(model_id, self._registry.device): + if self._registry.is_parked(model_id) and await self._registry.wake_parked( + model_id, self._registry.device, timeout_s=_host_wake_wait_s() + ): return "ready" if self._registry.is_loading(model_id): diff --git a/packages/sie_server/tests/adapters/test_host_offload.py b/packages/sie_server/tests/adapters/test_host_offload.py index 44957787f..093745732 100644 --- a/packages/sie_server/tests/adapters/test_host_offload.py +++ b/packages/sie_server/tests/adapters/test_host_offload.py @@ -228,3 +228,19 @@ def test_weights_the_walk_cannot_reach_block_parking(monkeypatch: pytest.MonkeyP # The adapter reports more weights than the walk found. monkeypatch.setattr(adapter, "memory_footprint", lambda: reachable * 2) assert adapter.host_offload_bytes() is None + + +def test_aliases_of_one_tensor_are_copied_once_and_stay_aliases() -> None: + holder = _Holder() + holder.model.original_scale = holder.model.scale # a plain-attribute alias of a buffer + holder.weight_alias = holder.model.embed.weight.data # a loose alias of a parameter's data + + expected = _host_offload.parkable_bytes(holder, CPU) + parked = _host_offload.park(holder, CPU) + assert parked.nbytes == expected + assert _ptr(holder.model.original_scale) == _ptr(holder.model.scale) + assert _ptr(holder.weight_alias) == _ptr(holder.model.embed.weight) + + _host_offload.unpark(parked) + assert _ptr(holder.model.original_scale) == _ptr(holder.model.scale) + assert _ptr(holder.weight_alias) == _ptr(holder.model.embed.weight) diff --git a/packages/sie_server/tests/core/test_registry_host_cache.py b/packages/sie_server/tests/core/test_registry_host_cache.py index ecb7db505..9f3cabbc2 100644 --- a/packages/sie_server/tests/core/test_registry_host_cache.py +++ b/packages/sie_server/tests/core/test_registry_host_cache.py @@ -19,6 +19,7 @@ from sie_server.core import registry as registry_module from sie_server.core.loader import load_model_configs from sie_server.core.memory import SIE_FAKE_MEMORY_BUDGET_ENV +from sie_server.core.model_loader import LoadedLora from sie_server.core.registry import ModelRegistry MODELS_DIR = Path(__file__).resolve().parents[2] / "models" @@ -295,3 +296,21 @@ async def stuck_make_room(name: str, entry: Any) -> Any: assert SMALL_A in registry.parked_model_names assert registry.host_cache.retained_bytes(SMALL_A) == 0 assert registry.host_cache.used_bytes == 128 * MIB + + +async def test_wake_drops_lora_placeholders_left_by_the_park(monkeypatch: pytest.MonkeyPatch) -> None: + registry = _registry(monkeypatch) + await registry.load_async(SMALL_A, device="cpu") + await registry.load_async(SMALL_B, device="cpu") + await registry.load_async(BIG, device="cpu") + + parked = registry.host_cache.get(SMALL_A) + assert parked is not None + parked.loaded.loras["in-flight"] = LoadedLora(adapter_id="in-flight", loading=True) + parked.loaded.loras["victim"] = LoadedLora(adapter_id="victim", unloading=True) + parked.loaded.loras["resident"] = LoadedLora(adapter_id="resident") + + await registry.load_async(SMALL_A, device="cpu") + loras = registry._loaded[SMALL_A].loras + assert list(loras) == ["victim", "resident"] + assert not loras["victim"].unloading diff --git a/packages/sie_server/tests/test_queue_executor.py b/packages/sie_server/tests/test_queue_executor.py index 2d2ee1073..cc9d7fad4 100644 --- a/packages/sie_server/tests/test_queue_executor.py +++ b/packages/sie_server/tests/test_queue_executor.py @@ -309,9 +309,19 @@ async def test_parked_model_is_woken_and_ready(self) -> None: reg.start_load_async = AsyncMock(return_value=True) ex = QueueExecutor(reg) assert await ex.ensure_model_ready("test/model") == "ready" - reg.wake_parked.assert_awaited_once_with("test/model", "cpu") + reg.wake_parked.assert_awaited_once_with("test/model", "cpu", timeout_s=15.0) reg.start_load_async.assert_not_called() + @pytest.mark.asyncio + async def test_parked_wake_wait_stays_inside_the_readiness_deadline(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SIE_MODEL_READY_TIMEOUT_S", "8") + reg = _make_registry(loaded=False, loading=False) + reg.is_parked.return_value = True + reg.wake_parked = AsyncMock(return_value=True) + ex = QueueExecutor(reg) + assert await ex.ensure_model_ready("test/model") == "ready" + reg.wake_parked.assert_awaited_once_with("test/model", "cpu", timeout_s=4.0) + @pytest.mark.asyncio async def test_parked_model_still_waking_reports_loading(self) -> None: reg = _make_registry(loaded=False, loading=True)