diff --git a/docker/Dockerfile b/docker/Dockerfile index 989f378cac4..1b9250d406e 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -165,10 +165,7 @@ RUN pip install "git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@v0.3 RUN TMS_CUDA_MAJOR=$(python3 -c "import torch; print(torch.version.cuda.split('.')[0])") \ pip install git+https://github.com/fzyzcjy/torch_memory_saver.git@f05a8754daf68238d54e4cf31cb3ba866684bbaf --no-cache-dir --force-reinstall RUN pip install "nvidia-modelopt[torch]>=0.37.0" --no-build-isolation -# radixark/Megatron-Bridge#24 (TE 2.17 grouped-linear contract), merged to @bridge. -# Pinned by SHA, not branch: buildkit caches this layer on the instruction text, so a -# branch that moves leaves the old revision baked into the image with nothing to show it. -RUN pip install git+https://github.com/radixark/Megatron-Bridge.git@7f0fb3456f8ffe47599b5fd167b454605d85f932 --no-deps --no-build-isolation +RUN pip install git+https://github.com/radixark/Megatron-Bridge.git@bridge --no-deps --no-build-isolation RUN pip install megatron-energon --no-deps RUN pip install multi-storage-client --no-deps diff --git a/docs/advanced/lora.md b/docs/advanced/lora.md index a91e32afe3b..269e5f683ee 100644 --- a/docs/advanced/lora.md +++ b/docs/advanced/lora.md @@ -300,32 +300,32 @@ its own Adam state and independently clocked scheduler. The trainer coalesces ready prompt-group slices or partial adapter batches and selectively upserts only changed adapters into SGLang. -Set the slot capacity with `--multi-lora-n-adapters N`. A bounded run registers -repeatable `--multi-lora-adapter NAME PATH` entries at startup; service mode can -start with empty slots and register adapters through the controller HTTP API. +Set the slot capacity with `--multi-lora-n-adapters N`. The operation backend +starts as a long-running service with empty slots and registers adapters at +runtime through the controller HTTP API. This path currently forces Megatron-Bridge LoRA and requires disaggregated NCCL broadcast, PP1, THD, Adam, and no train offload. Shared-outer expert adapters are unsupported, and MoE expert adapters cannot use FP8/FP4 experts. Native multi-LoRA is not implied by the native single-adapter work: both current -`main` and the Tinker-oriented branch below still build multi-LoRA through +`main` and the operation-backend branch below still build multi-LoRA through Megatron-Bridge. Native multi-LoRA is tracked separately in [issue #2141](https://github.com/radixark/miles/issues/2141). -### Future Tinker-compatible operation backend +### Multi-LoRA operation backend and Tinker compatibility [PR #2273](https://github.com/radixark/miles/pull/2273) is the active -Tinker-oriented backend proposal. It changes ownership of the training loop: +Multi-LoRA operation-backend proposal. It changes ownership of the training loop: instead of the server owning a dataset, reward function, and one-step schedule, clients submit explicit operations against a registered adapter. Its primary -intended consumer is a Tinker-compatible training service rather than a generic -server-owned dataset scheduler. +consumer today is a Tinker-compatible protocol adapter, but the trainer-side +operation contract is named independently from that client protocol. ```text -Tinker-style client - | register + ordered operations +Tinker client -> Tinker protocol/frontend adapter + | normalized register + ordered operations v -controller / operation ledger +MultiLoraOperationBackend / operation ledger | bind one fixed LoRA slot v Megatron-Bridge multi-LoRA trainer @@ -335,6 +335,12 @@ Megatron-Bridge multi-LoRA trainer SGLang router + registration-scoped adapter identity ``` +The current concrete is `MultiLoraOperationBackend`, with +`MultiLoraOperationBatchFn` and `MultiLoraParameterExecutor` handling adapter +batching and slot execution. `Tinker` remains the wire/SDK compatibility name. +A future full-parameter executor may reuse the normalized operation contract, +but this PR does not implement or claim full-parameter training. + The operation surface separates compute, optimization, and publication: | Operation | Contract | @@ -349,7 +355,8 @@ are strictly serialized per registration, while idempotent retries, gap-buffered arrival, acknowledgements, and backpressure make execution retry-safe and order-safe. A registration-scoped serving identity prevents an old request from using a slot after that slot has been reassigned. Authenticated remote access is -the responsibility of the future frontend, not the Ray operation API in #2273. +provided by the stacked Tinker frontend in #2346; it is not part of the Ray +operation API in #2273. The v1 scope in the PR is deliberately narrow: text-only synchronous training, one shared base model, shifted 1-D targets, `cross_entropy`, importance-sampling, @@ -362,9 +369,10 @@ This backend is implemented in an open PR, not released on `main`; the PR reports H200 validation. PR #2273 provides the operation backend, but its v1 training operations are still exposed through the controller's Ray API. The stacked [PR #2346](https://github.com/radixark/miles/pull/2346) adds a REST -frontend compatible with the official `tinker==0.24.1` client; its GPU frontend -E2E is still pending. If #2273 lands as proposed, it replaces the current -dataset-driven driver. +frontend compatible with the official `tinker==0.24.1` client. The stacked +full system has passed both RL and pure-SFT 2xH200 acceptance; those results do +not expand #2273 beyond its fixed-slot Multi-LoRA scope. If #2273 lands as +proposed, it replaces the current dataset-driven driver. ## Compatibility and limitations @@ -404,5 +412,5 @@ dataset-driven driver. - `miles/rollout/session/core.py` attaches the single adapter to agentic session requests. - `miles/ray/multi_lora/`, `miles/rollout/multi_lora/`, and - `miles/backends/megatron_utils/multi_lora_*.py` implement the multi-adapter + `miles/backends/megatron_utils/api_backends/multi_lora/` implement the multi-adapter controller, routing, scheduling, optimization, and checkpoint path. diff --git a/docs/docs.json b/docs/docs.json index 9216e8008d0..bf25af3a1cb 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -242,7 +242,7 @@ "pages": [ "examples/geo3k-vlm", "examples/geo3k-vlm/multi-turn", - "examples/multi-lora", + "examples/multi-lora-operations", "examples/on-policy-distillation", "examples/on-policy-distillation/qwen3-5-35b-selfdistill", "examples/ppo", diff --git a/docs/examples/index.md b/docs/examples/index.md index 016fc9d5184..032069ffdf1 100644 --- a/docs/examples/index.md +++ b/docs/examples/index.md @@ -13,7 +13,7 @@ End-to-end training workflows — the place to start. - **[geo3k_vlm](/examples/geo3k-vlm)**: Training VLMs with FSDP using GRPO on the GEO3K dataset. - **[multi_turn](/examples/geo3k-vlm/multi-turn)**: The same dataset over multiple turns, with the model cropping images through an interactive environment. - **[lora](https://github.com/radixark/miles/tree/main/examples/lora)**: LoRA fine-tuning with the Megatron backend. -- **[multi_lora](/examples/multi-lora)**: Fully-async multi-adapter LoRA training with a slot-keyed adapter page table. +- **[multi_lora_operations](/examples/multi-lora-operations)**: Multi-adapter LoRA trained through explicit operations with Tinker REST/SDK compatibility. - **[on_policy_distillation](/examples/on-policy-distillation)**: Teacher–student distillation on the student's own rollouts, run inside the on-policy training loop. - **[qwen3_5_35b_selfdistill](/examples/on-policy-distillation/qwen3-5-35b-selfdistill)**: Two-phase self-distillation of Qwen3.5-35B-A3B on one 8xH200 node, with an in-process Megatron teacher. - **[ppo](/examples/ppo)**: Actor-critic PPO with GAE advantages, where the critic shares the actor's train GPUs. diff --git a/docs/examples/multi-lora-operations.md b/docs/examples/multi-lora-operations.md new file mode 100644 index 00000000000..4c32055ea5e --- /dev/null +++ b/docs/examples/multi-lora-operations.md @@ -0,0 +1,369 @@ +--- +title: "Multi-LoRA operation backend with Tinker compatibility" +description: "Multi-adapter LoRA trained through explicit operations with Tinker REST/SDK compatibility." +# Generated from examples/multi_lora_operations/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. +--- +Serve many LoRA training runs on one shared base model through the +`MultiLoraOperationBackend`. The +[tinker](https://tinker-docs.thinkingmachines.ai/)-compatible frontend maps the +official SDK onto explicit `forward_backward` / `optim_step` operations and +shared-engine sampling — no dataset, reward function, or batch schedule on the +server. + +``` +official Tinker SDK ──HTTP──> Tinker protocol/frontend adapter + │ +internal caller ──Ray operations───────┘ + ▼ + MultiLoraOperationBackend (head node) + ├─ registration + operation ledger + ├─ adapter-slot execution + └─ serving plane ─────────> SGLang router +trainer ranks <──Ray── driver loop (train_multi_lora_operations.py) +``` + +## Start the Miles engine + +For the documented SDK flow, start both the operation backend and the Tinker +frontend. The helper starts the shared training and sampling engines in +service mode; add `--tinker-frontend` through `--extra-args` so that the +official SDK can use the controller's `/api/v1` endpoint: + +```bash +# Once per node: download the example checkpoint. +python examples/multi_lora_operations/run_multi_lora_operations.py prepare + +# Start Miles in service mode, with both the backend and frontend enabled. +python examples/multi_lora_operations/run_multi_lora_operations.py serve \ + --extra-args "--tinker-frontend" +``` + +The following lower-level command is useful when deploying with custom +Megatron and SGLang flags: + +```bash +python train_multi_lora_operations.py \ + --tinker-backend \ + --tinker-frontend \ + --multi-lora-n-adapters 4 \ + --lora-rank 32 --lora-alpha 64 \ + --target-modules all-linear \ + --hf-checkpoint Qwen/Qwen3-0.6B \ + ... # the usual megatron/sglang flags; see run_multi_lora_operations.py +``` + +Key flags: + +| flag | meaning | +|------|---------| +| `--tinker-backend` | enable the Tinker protocol adapter for the Multi-LoRA operation backend (requires `--multi-lora-n-adapters > 0`) | +| `--multi-lora-n-adapters N` | fixed slot count; a registration binds a slot for life (queue when full) | +| `--lora-rank` / `--lora-alpha` | deployment-wide ceiling / fixed alpha — clients may lower `rank`, never set `alpha` | +| `--multi-lora-api-port` | control-plane API port for runtime adapter registration | +| `--tinker-max-coalesce-wait-s` | how long one train call coalesces additional ready client batches | +| `--tinker-max-empty-wait-s` | idle-queue yield back to the control phase (keep this small) | +| `--tinker-frontend` | serve the official tinker SDK REST protocol (`/api/v1`) on the controller HTTP server (requires `--tinker-backend`) | +| `--tinker-api-key` | X-API-Key the frontend requires (prefer `$MILES_TINKER_API_KEY` — a CLI flag shows in the process list); mandatory for a non-loopback bind | + +The operator plane (`/adapter_runs*`, `/info`) accepts loopback peers only, +whatever the bind: the SDK key is a client credential and never grants the +routes that read server-local YAML files, choose save paths, or deregister +tenants. `/health` is liveness (the socket is up); `/api/v1/healthz` is +readiness and answers 503 until the driver reports the trainer exists. + +### Activation recompute (memory saving) + +`--recompute-granularity selective` is always supported (default +`--recompute-modules core_attn`; add `moe_act` to also recompute the MoE +activation with grouped GEMM). `--recompute-granularity full` — and `moe` in +`--recompute-modules` when expert modules are targeted — is supported as +well: the deployment's Megatron-Bridge (branch `bridge`) recognizes +multi-LoRA `.adapters..` params in its PEFT recompute patch, so +checkpointed regions replay grad-enabled during adapter-only training. + +## Operation contract + +`Tinker` names the compatibility boundary, not the trainer implementation. +The current concrete is `MultiLoraOperationBackend`; its queue-backed +`MultiLoraOperationBatchFn` batches already-tokenized operations, and the +Megatron `MultiLoraParameterExecutor` applies them to adapter slots. A future +full-parameter composition can reuse the same operation contract and the +unwired `FullParameterExecutor` sibling; full-parameter launch, data-path, +checkpoint, and publish integration are not implemented by this stack today. + +``` +Tinker protocol frontend + │ + ▼ +generic training-operation contract + │ + ├── MultiLoraParameterExecutor (current wired target) + └── FullParameterExecutor (implemented seam; not wired) +``` + +`enqueue_operation(name, operation_id, ordinal, kind, payload)` — ordinals are +consecutive per registration starting at 1; arrival may be out of order +(gap-buffered, and a hole-filling ordinal is always admitted), execution is +strictly ordinal-ordered; retries with the same `operation_id`, same ordinal, +and identical payload return the original operation — anything else is a +typed conflict. + +A stream stalled on a never-arriving ordinal (the 0.24.1 SDK consumes a +seq_id and can then fail BEFORE HTTP — see the SDK limitations below) expires +after `--tinker-operation-gap-timeout` (default 600 s, `<= 0` disables): the +blocked, never-claimed operations terminal-fail `FAILED(user)` naming the +missing ordinal, and the hole is sealed — the missing identity never executes +(a late arrival is a typed conflict), nothing overtakes it, and the client +resubmits as new operations. Stalls are observable before expiry: +`service_info()` reports `gap_stalls`, and a blocked operation's +`get_operation` view carries `waiting_on_ordinal` / `gap_stalled_for`. + +| kind | payload | success result | +|------|---------|----------------| +| `forward_backward` | `{samples: [Datum...], loss: {loss_fn, loss_fn_config?}}` | `{logprobs: [[...]], metrics: {"loss:sum", "unmasked_tokens:sum", "loss_weight:sum" (CE only)}}` | +| `forward` | `{samples: [Datum...]}` | `{logprobs: [[...]]}` (zero gradient, structurally) | +| `optim_step` | `{adam_params: {learning_rate, beta1, beta2, eps, weight_decay, grad_clip_norm}}` | `{grad_norm, learning_rate}` | +| `save_weights_for_sampler` | `{}` | `{serving_version, serving_name}` — completes only after the weights are live | +| `save_state` | `{tag?, ttl_seconds?}` | `{path, step}` (named states are immutable) | +| `load_state` | `{path}` | `{step, path}` (re-publishes on the next push) | + +`Datum = {tokens, response_length, loss_mask, loss_weights?, advantages?, rollout_log_probs?}` +— per-token channels align with the response span. Losses reduce as plain +token sums (`Σ(-logp·w)` for `cross_entropy`), so K chunked forward_backward +calls accumulate exactly like one; `loss_weights` own the scale and no server +normalization or scheduler ever touches a tinker slot. Result `metrics` use +the SDK combiner's `name:reduction` keys. + +For an SFT-style per-token loss, divide `loss:sum` by `loss_weight:sum` +(cross-entropy only: Σ weight·mask, chunk-additive like the loss) — NOT by +`unmasked_tokens:sum`, which counts every loss_mask-active position and so +includes the weight-0 prompt tokens of a teacher-forced datum, silently +diluting the displayed loss. Guard the division: weights are arbitrary +finite floats, so the sum can be zero or negative. + +Operation states: `QUEUED → CLAIMED → SUCCEEDED | FAILED(user|server) | CANCELLED`; +poll `get_operation`, then `ack_operation` to release the record. These verbs +are the controller actor's Ray API; the tinker frontend drives them over +HTTP. Backpressure raises a retryable `OperationBackpressure` — the frontend +maps it to 429 + Retry-After, never to a 4xx the SDK treats as fatal. +Deregistering fences every open operation of that registration as +`FAILED(user)`. + +Gradient-window poison: `optim_step` delimits a window of `forward_backward` +operations. If any of them reached a terminal state without succeeding (a +rejected chunk, an execution failure, a cancel), the window holds PARTIAL +gradients — the window's `optim_step` executes as a discard (all ranks clear +the slot's gradient sum), terminal-fails `FAILED(user)`, and moves neither +the step clock nor the serving version. The consumed poison resets the +window; resubmit the batch and step again. + +## Tinker SDK frontend (tinker==0.24.1 JSON subset) + +With `--tinker-frontend` the controller's HTTP server also speaks the REST +protocol of the official [`tinker`](https://pypi.org/project/tinker/) SDK — +exactly the **`tinker==0.24.1` JSON core-loop subset** (wheel source and +captured traffic; pure JSON, no protobuf: `/api/v1/client/config` pins the +SDK to its own default JSON path). Other SDK versions are rejected at +bootstrap (`/client/config` and `create_session` fail fast on the reported +`sdk_version`): 0.25+ switches `forward_backward` to protobuf, and the +current cookbook's canonical final checkpoint needs named sampler +checkpoints — neither is served here, so this is NOT "current +Tinker/cookbook compatible". An unmodified 0.24.1 client drives training +and sampling: + +```python +import tinker +sc = tinker.ServiceClient(base_url="http://127.0.0.1:8068", api_key="tml-...") +tc = sc.create_lora_training_client(base_model=..., rank=32) +tc.forward_backward(data, "cross_entropy") +tc.optim_step(tinker.types.AdamParams(learning_rate=1e-4)).result() +sampler = tc.save_weights_and_get_sampling_client() +future = sampler.sample( # sample()/sample_async() submit /api/v1/asample; + prompt=tinker.types.ModelInput.from_ints(prompt_tokens), + num_samples=4, + sampling_params=tinker.types.SamplingParams(max_tokens=128, temperature=0.7), +) +response = future.result() # .sequences[i].tokens / .logprobs / .stop_reason +``` + +### Client-owned RL loop + +After the engine reports ready, connect the official SDK client to the +frontend endpoint and run the loop below. The backend executes each requested +operation; rollout generation, scoring, and `Datum` construction remain in +the client. + +Start the driver with both `--tinker-backend` and `--tinker-frontend`. The +backend then owns execution and serving, while the client owns data +preparation and the training loop. In particular, the client can run the +same pattern as the [target-flow example](https://github.com/radixark/miles/issues/2258): + +```python +import tinker +from transformers import AutoTokenizer + +service = tinker.ServiceClient(base_url="http://127.0.0.1:8068", api_key="tml-...") +base_model = service.get_server_capabilities().supported_models[0].model_name +training = service.create_lora_training_client(base_model=base_model, rank=16) +tokenizer = AutoTokenizer.from_pretrained(base_model) + +# Publish the initial LoRA so the first rollout has a policy to sample. +sampler = training.save_weights_and_get_sampling_client() + +rl_prompts = ["Solve: If a train travels 60 km in 2 hours, what is its speed?"] +prompt_ids = [tokenizer(p).input_ids for p in rl_prompts] + +for update_idx in range(num_rl_updates): + # Option 1 -- SFT data preparation (client-owned; replace the RL batch + # below and train with loss_fn="cross_entropy"). + # batch = [ + # datum_from_sft_example(example["prompt"], example["completion"]) + # for example in sft_examples + # ] + + # Option 2 -- RL data preparation (client-owned). sample() returns a + # future; .result() carries sequences with tokens and logprobs. + futures = [ + sampler.sample( + prompt=tinker.types.ModelInput.from_ints(ids), + num_samples=4, + sampling_params=tinker.types.SamplingParams(max_tokens=256, temperature=1.0), + ) + for ids in prompt_ids + ] + rollouts = [future.result() for future in futures] + scored = score_rollouts(rl_prompts, rollouts) # rewards -> advantages, client-owned + batch = [ + datum_from_scored_rollout(ids, sequence, advantage) + for ids, response, advantages in zip(prompt_ids, rollouts, scored) + for sequence, advantage in zip(response.sequences, advantages) + ] + + fb = training.forward_backward(batch, "importance_sampling") + step = training.optim_step(tinker.types.AdamParams(learning_rate=1e-4)) + fb.result() + step.result() + + # Publish explicitly so the next rollout samples the new policy. + # Serving is latest-only: the publish supersedes the previous sampling + # client, so re-acquire it here every update. + sampler = training.save_weights_and_get_sampling_client() +``` + +`datum_from_sft_example`, `score_rollouts`, and `datum_from_scored_rollout` +are application code: they define the task data, rollout scoring, and the +per-token loss channels. An RL datum pairs `model_input` (prompt + sampled +tokens, shifted) with `loss_fn_inputs` `target_tokens`, the sampler's +returned `logprobs`, and per-token `advantages`; an SFT datum needs +`target_tokens` plus 0/1 `weights`. The frontend translates the resulting +SDK requests to operations; the backend executes them in order and only +changes the sampler's policy on the explicit publish. The complete runnable +version of this loop is `tests/e2e/tinker_frontend/tinker_sdk_rl_quality.py` +(GRPO on GSM8K, four concurrent adapters through one deployment). + +Mapping: one training client = one registration (`create_model` registers, +`unload_model` deregisters), and every operation is pinned to its +`(name, registration_id)` — a stale handle fences instead of binding to a +same-name successor; every training verb forwards its SDK `seq_id` as the +registration ordinal (chunks posted out of order gap-buffer); futures poll +`/api/v1/retrieve_future` and terminal bodies replay until delivered (an +evicted delivered result leaves a fingerprint tombstone that answers a typed +410 — the 0.24.1 SDK surfaces it as a retryable "promise expired", it does +not re-run the original request); `save_state` mints `tinker://` paths +(resolved from an in-memory catalog; failures echo the public URI, not the +trainer filesystem); the ephemeral `save_weights_and_get_sampling_client` +publish binds `(name, registration_id, serving_version)` and samples through +the sglang router — a republish makes older sampling clients fail loud, and +the version is re-checked after generation so a publish landing mid-flight +fails the in-flight sample instead of returning cross-version output (the +identity is versioned, not leased: a publish committing between that check +and delivery is a documented residual race). Frontend rejections on a spent +`seq_id` become terminal `FAILED(user)` futures so the ordinal is still +consumed — bounded by the same unacked-results budget as every other record +(429 past it). + +Frontend-level v1 rejections (beyond the backend matrix): non-0.24.x SDK +versions, LoRA `seed` and per-module `train_*` flags (deployment-wide), +weights-only restore (`load_state` / `create_training_client_from_state` — +the backend restores the full training state; use the `_with_optimizer` +variants), named persistent sampler checkpoints +(`save_weights_for_sampler(name)` / `create_sampling_client(model_path=...)`), +`ttl_seconds` (checkpoint/sampler TTL expiry is not implemented), +`topk_prompt_logprobs`, sparse-CSR tensors, and negative +token ids anywhere (targets, inputs, prompts, stop tokens). A sampling +`seed` maps to sglang `sampling_seed`, offset per sample so +`num_samples > 1` stays diverse. `prompt_logprobs` maps to sglang +`logprob_start_len=0` on the same generate (the engine scores the prompt +natively; position 0 has no context and returns null) — this serves both +`sample(include_prompt_logprobs=True)` and the SDK's `compute_logprobs()`, +which the 0.24.1 wheel sends as a 1-sample, 1-token generation. + +Sampling architecture: `/asample` returns its future immediately and a +background task posts one router `/generate` per sample, carrying the +server-derived serving identity (`rid`/`lora_path`/`extra_key` are never +client-controllable — the wire models drop unknown fields and the sglang +params are rebuilt from an allowlist). SGLang's continuous batching is the +only sampling batcher: the frontend never coalesces prompts, and the +training-operation scheduler (`MultiLoraOperationBatchFn`) never sees a sampling +request. The legacy datasource rollout pipeline +(`RolloutManager.generate()`: datasets, rewards, training-data conversion) +is not on this path — the frontend shares only the router the rollout +engines already serve. + +Trust boundary (v1): the frontend authenticates clients and bounds aggregate +active sub-generations, rejects one request whose `num_samples` exceeds that +capacity, and preflights `prompt + max_tokens` against the discovered engine +limit. It still does not validate token ids against the vocabulary upper bound +or enforce request-body/output-byte quotas. Run it loopback/VPN-facing for +trusted clients; per-tenant quotas are future work. + +## v1 compatibility matrix + +Supported: text-only input; the synchronous training loop; 1-D shifted +targets; `loss_fn ∈ {cross_entropy, importance_sampling, ppo}` (per-op clip +config); per-call AdamParams; multi-chunk gradient accumulation with +independent `optim_step`; latest-only sampler weights behind the publish +barrier; prompt logprobs (`compute_logprobs()` / +`sample(include_prompt_logprobs=True)`, one sub-generation of admission +weight); named immutable `save_state` / `load_state` (create-from-checkpoint +included, shape-fenced); optional `num_step` auto-retirement. + +Explicitly rejected (boundary error, never a silent fallback): multimodal +inputs; nested `(N, K)` top-K targets; other loss functions (CISPO, DRO, ...); +client-set `alpha`; non-finite/out-of-domain AdamParams; a loss's required +per-token channels missing; `response_length == len(tokens)` (targets are +shifted); async/off-policy sampling against pinned snapshots; +cross-world-size state restore; state restore into a slot whose per-rank +optimizer ownership differs from the save (cross-slot restore requires an +identical dense-and-expert ownership signature); idle slot GC. + +## Known tinker SDK (0.24.1) client-side limitations + +The official `tinker==0.24.1` TrainingClient takes its per-model seq counter +BEFORE it serializes and POSTs a request, so a submission can die client-side +with the ordinal already spent (verified against the live stack, +codex-0817-sft-fix §4-§6): + +- **Pre-HTTP serialization failure** — e.g. `AdamParams(learning_rate=nan)` + raises a local JSON `ValueError`; the request never reaches Miles and later + operations of the same client queue behind the hole. The gap timeout above + terminal-fails them typed, and the SAME TrainingClient can resubmit + afterwards (its turn counter did advance). Validate that Adam params and + custom scalars are finite before calling the SDK to avoid the stall. +- **`.future().cancel()` on an SDK future** can spend the request id without + advancing the SDK's internal turn counter: later operations of that client + wait forever CLIENT-side and Miles receives nothing it could terminalize — + no server-side mitigation exists. Do not cancel underlying SDK futures; + `.result(timeout=...)` is safe (non-destructive, the future stays + retrievable). After an immediate cancel, discard the TrainingClient and + create a new one (a fresh registration). When some submissions did reach + the server, the gap timeout converts the surviving stall into typed + failures instead of a hang. +- The server never skips a missing ordinal and never guesses what it would + have been: the gap timeout only fails what is blocked and seals the hole, + so strict per-registration ordering, idempotent retries, and anti-replay + all hold. + +## Files + +- `run_multi_lora_operations.py` — disaggregated service launch (`prepare` / `serve`) diff --git a/docs/examples/multi-lora.md b/docs/examples/multi-lora.md deleted file mode 100644 index 91778cb21cb..00000000000 --- a/docs/examples/multi-lora.md +++ /dev/null @@ -1,149 +0,0 @@ ---- -title: "Multi-LoRA Training Example (fully-async)" -description: "Fully-async multi-adapter LoRA training with a slot-keyed adapter page table." -# Generated from examples/multi_lora/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. ---- -Train multiple LoRA adapters concurrently against a shared base model, using a -fully-async rollout (continuous producer) + a slot-keyed LoRA page table on the -SGLang engines (in-place upsert, no unload, no drain). - -This example trains two adapters on Qwen3-4B: - -- **gsm8k** — grade-school math, `rm_type: math` -- **dapo_math** — competition math (DAPO-Math-17k), `rm_type: deepscaler` - -## Layout - -``` -run_multi_lora.py # launcher: prepare / train / full-train / serve -service_smoke.py # register/deregister smoke test against the API -adapters/ - gsm8k.yaml - dapo_math.yaml -``` - -The implementation lives in the library: the driver is `train_multi_lora_async.py` -at the repo root (next to `train.py`/`train_async.py`), the rollout fn and data -source are `miles/rollout/multi_lora/`, and the controller is -`miles/ray/multi_lora/` (registry + backend + HTTP API, plus the named Ray -actor pinned to the head node). - -## Design (decoupled per-adapter optimizers) - -- **Controller** (Ray actor + control-plane HTTP API) is the source of truth: - `POST/GET/DELETE /adapter_runs` plus `GET /adapter_runs/state`. The data source - reads it; the trainer reads it. Generation traffic goes straight to the router; - on deregister the controller aborts the adapter's in-flight requests - engine-side by rid prefix (`rid = {adapter}::{uuid}`, set in `generate`). -- **Per-adapter gradient accumulation.** Each adapter has its own batch shape: - `rollout_batch_size` prompt groups per optimizer step, each group holding - `n_samples_per_prompt` responses (`adapter_global_batch_size = - rollout_batch_size x n_samples_per_prompt` samples per step). Completed - prompt groups flow into training continuously in multiples of the - adapter's `min_groups_per_dp_split` (the smallest group count whose samples - split evenly across data-parallel ranks), gradients - accumulate in the DDP buffers across train batches, and an adapter's - optimizer steps exactly when its adapter batch fills — independent of every other - adapter. The controller tracks adapter batch progress (`accumulated_groups`) and commits - it only after a successful train call. -- **Per-slot optimizers.** One Adam per adapter slot under Megatron's - `LayerWiseDistributedOptimizer` (whole-parameter ZeRO-1): per-slot state, - step counts, and gradient clipping; optimizer state sharded across DP ranks; - plain DDP all-reduce (no distributed optimizer) makes cross-batch gradient - retention idempotent. -- **Batch collection.** The collection loop (same shape as fully_async's) - pops groups from the per-adapter buffers round-robin, one - `min_groups_per_dp_split` at a time, capped at each adapter's remaining - batch, until the batch reaches `--global-batch-size` samples or a non-empty - batch makes no progress for `--multi-lora-max-coalesce-wait-s` (the target - can be permanently unreachable, so it trains on whatever is ready) — a - single adapter with a small batch trains alone without waiting for - anyone. Samples enter the gradient buffers with weight 1; at step time the - slot's accumulated gradient is scaled by `1/adapter_global_batch_size` - (a constant known in advance), so an adapter's update is identical to what - it would get training alone. -- **Selective weight sync.** Only adapters whose optimizer stepped are pushed - to the engines (upsert into the slot-keyed page table); only their slot - versions bump, keeping staleness filtering per-adapter accurate. -- Adapters deregister on committed optimizer-step count (`num_step`) in the - controller's train-commit path (`mark_batch_trained`), so stop checks happen - exactly when steps advance. `num_step` is relative to the adapter's - start/resume step. When an adapter doesn't set `num_step`, it is derived - from `num_epoch` (default 1) as `num_epoch x len(dataset) // - rollout_batch_size` once the data source loads the dataset (post-filter - length). The trainer's - `reconcile_adapters` (before each generate) retires it at the next sync - point and cleans up (save ckpt + clear Megatron slot + zero its optimizer - state and retained gradients). The adapter's untrained tail — buffered - groups and any partially accumulated gradients — is discarded. -- **Batch ⊆ loaded property:** `reconcile_adapters` runs before `generate`, so the - batch is fetched with loaded = active; active only shrinks during generate, so every - adapter in the batch is live on the trainer. - -## Provision (once) - -```bash -python examples/multi_lora/run_multi_lora.py prepare -``` - -Downloads `Qwen/Qwen3-4B` (to `/root/models`), `zhuzilin/dapo-math-17k`, and -`zhuzilin/gsm8k` (to `/root/datasets`). - -## Run - -```bash -python examples/multi_lora/run_multi_lora.py train # or: full-train (prepare + train) -``` - -Registers the two adapters from CLI flags and trains until each hits its `num_step`, -then exits. - -## Service mode - -```bash -python examples/multi_lora/run_multi_lora.py serve -``` - -Starts with no adapters and idles; register/deregister at runtime through the -control-plane API (port 8068): - -```bash -python examples/multi_lora/service_smoke.py --api-url http://127.0.0.1:8068 \ - --data /root/datasets/gsm8k/train.parquet --input-key messages --label-key label --rm-type math -``` - -## Multi-LoRA CLI flags - -| Flag | Purpose | -| --- | --- | -| `--multi-lora-n-adapters N` | Max concurrent adapter slots. `0` disables (default); `> 0` enables. | -| `--multi-lora-adapter NAME PATH` | Register an adapter at startup. Repeatable. `PATH` → an `adapter.yaml`. | - -Per-adapter `rank` in `adapter.yaml` must be `<= --lora-rank`. - -## adapter.yaml - -```yaml -rank: 16 -alpha: 16 -rollout_batch_size: 32 # prompt groups per optimizer step (defaults to --rollout-batch-size) -n_samples_per_prompt: 4 # group shape (defaults to --n-samples-per-prompt) -data: /root/datasets/gsm8k/train.parquet -input_key: messages -label_key: label -rm_type: math -num_step: 400 # stop adapter after N optimizer steps - # (default: derived from num_epoch, itself default 1) -# optional: save, num_epoch, custom_rm_path, ... -``` - -The derived `adapter_global_batch_size = rollout_batch_size x -n_samples_per_prompt` is the adapter's samples-per-optimizer-step (the -per-adapter analog of `--global-batch-size`). - -Batch-shape constraints (validated at registration, not at runtime): -`n_samples_per_prompt` must be a divisor or multiple of the trainer's -data-parallel size; `rollout_batch_size` must be a multiple of the adapter's -`min_groups_per_dp_split`; -`adapter_global_batch_size` is capped by -`--multi-lora-max-adapter-global-batch-size` (default 4x `--global-batch-size`). diff --git a/examples/README.md b/examples/README.md index c8ef2d7cdbe..33a340603b9 100644 --- a/examples/README.md +++ b/examples/README.md @@ -9,7 +9,7 @@ End-to-end training workflows — the place to start. - **[geo3k_vlm](./geo3k_vlm)**: Training VLMs with FSDP using GRPO on the GEO3K dataset. - **[multi_turn](./geo3k_vlm/multi_turn)**: The same dataset over multiple turns, with the model cropping images through an interactive environment. - **[lora](./lora)**: LoRA fine-tuning with the Megatron backend. -- **[multi_lora](./multi_lora)**: Fully-async multi-adapter LoRA training with a slot-keyed adapter page table. +- **[multi_lora_operations](./multi_lora_operations)**: Multi-adapter LoRA trained through explicit operations with Tinker REST/SDK compatibility. - **[on_policy_distillation](./on_policy_distillation)**: Teacher–student distillation on the student's own rollouts, run inside the on-policy training loop. - **[qwen3_5_35b_selfdistill](./on_policy_distillation/qwen3_5_35b_selfdistill)**: Two-phase self-distillation of Qwen3.5-35B-A3B on one 8xH200 node, with an in-process Megatron teacher. - **[ppo](./ppo)**: Actor-critic PPO with GAE advantages, where the critic shares the actor's train GPUs. diff --git a/examples/multi_lora/README.md b/examples/multi_lora/README.md deleted file mode 100644 index 8423ddae1d6..00000000000 --- a/examples/multi_lora/README.md +++ /dev/null @@ -1,146 +0,0 @@ -# Multi-LoRA Training Example (fully-async) - -Train multiple LoRA adapters concurrently against a shared base model, using a -fully-async rollout (continuous producer) + a slot-keyed LoRA page table on the -SGLang engines (in-place upsert, no unload, no drain). - -This example trains two adapters on Qwen3-4B: - -- **gsm8k** — grade-school math, `rm_type: math` -- **dapo_math** — competition math (DAPO-Math-17k), `rm_type: deepscaler` - -## Layout - -``` -run_multi_lora.py # launcher: prepare / train / full-train / serve -service_smoke.py # register/deregister smoke test against the API -adapters/ - gsm8k.yaml - dapo_math.yaml -``` - -The implementation lives in the library: the driver is `train_multi_lora_async.py` -at the repo root (next to `train.py`/`train_async.py`), the rollout fn and data -source are `miles/rollout/multi_lora/`, and the controller is -`miles/ray/multi_lora/` (registry + backend + HTTP API, plus the named Ray -actor pinned to the head node). - -## Design (decoupled per-adapter optimizers) - -- **Controller** (Ray actor + control-plane HTTP API) is the source of truth: - `POST/GET/DELETE /adapter_runs` plus `GET /adapter_runs/state`. The data source - reads it; the trainer reads it. Generation traffic goes straight to the router; - on deregister the controller aborts the adapter's in-flight requests - engine-side by rid prefix (`rid = {adapter}::{uuid}`, set in `generate`). -- **Per-adapter gradient accumulation.** Each adapter has its own batch shape: - `rollout_batch_size` prompt groups per optimizer step, each group holding - `n_samples_per_prompt` responses (`adapter_global_batch_size = - rollout_batch_size x n_samples_per_prompt` samples per step). Completed - prompt groups flow into training continuously in multiples of the - adapter's `min_groups_per_dp_split` (the smallest group count whose samples - split evenly across data-parallel ranks), gradients - accumulate in the DDP buffers across train batches, and an adapter's - optimizer steps exactly when its adapter batch fills — independent of every other - adapter. The controller tracks adapter batch progress (`accumulated_groups`) and commits - it only after a successful train call. -- **Per-slot optimizers.** One Adam per adapter slot under Megatron's - `LayerWiseDistributedOptimizer` (whole-parameter ZeRO-1): per-slot state, - step counts, and gradient clipping; optimizer state sharded across DP ranks; - plain DDP all-reduce (no distributed optimizer) makes cross-batch gradient - retention idempotent. -- **Batch collection.** The collection loop (same shape as fully_async's) - pops groups from the per-adapter buffers round-robin, one - `min_groups_per_dp_split` at a time, capped at each adapter's remaining - batch, until the batch reaches `--global-batch-size` samples or a non-empty - batch makes no progress for `--multi-lora-max-coalesce-wait-s` (the target - can be permanently unreachable, so it trains on whatever is ready) — a - single adapter with a small batch trains alone without waiting for - anyone. Samples enter the gradient buffers with weight 1; at step time the - slot's accumulated gradient is scaled by `1/adapter_global_batch_size` - (a constant known in advance), so an adapter's update is identical to what - it would get training alone. -- **Selective weight sync.** Only adapters whose optimizer stepped are pushed - to the engines (upsert into the slot-keyed page table); only their slot - versions bump, keeping staleness filtering per-adapter accurate. -- Adapters deregister on committed optimizer-step count (`num_step`) in the - controller's train-commit path (`mark_batch_trained`), so stop checks happen - exactly when steps advance. `num_step` is relative to the adapter's - start/resume step. When an adapter doesn't set `num_step`, it is derived - from `num_epoch` (default 1) as `num_epoch x len(dataset) // - rollout_batch_size` once the data source loads the dataset (post-filter - length). The trainer's - `reconcile_adapters` (before each generate) retires it at the next sync - point and cleans up (save ckpt + clear Megatron slot + zero its optimizer - state and retained gradients). The adapter's untrained tail — buffered - groups and any partially accumulated gradients — is discarded. -- **Batch ⊆ loaded property:** `reconcile_adapters` runs before `generate`, so the - batch is fetched with loaded = active; active only shrinks during generate, so every - adapter in the batch is live on the trainer. - -## Provision (once) - -```bash -python examples/multi_lora/run_multi_lora.py prepare -``` - -Downloads `Qwen/Qwen3-4B` (to `/root/models`), `zhuzilin/dapo-math-17k`, and -`zhuzilin/gsm8k` (to `/root/datasets`). - -## Run - -```bash -python examples/multi_lora/run_multi_lora.py train # or: full-train (prepare + train) -``` - -Registers the two adapters from CLI flags and trains until each hits its `num_step`, -then exits. - -## Service mode - -```bash -python examples/multi_lora/run_multi_lora.py serve -``` - -Starts with no adapters and idles; register/deregister at runtime through the -control-plane API (port 8068): - -```bash -python examples/multi_lora/service_smoke.py --api-url http://127.0.0.1:8068 \ - --data /root/datasets/gsm8k/train.parquet --input-key messages --label-key label --rm-type math -``` - -## Multi-LoRA CLI flags - -| Flag | Purpose | -| --- | --- | -| `--multi-lora-n-adapters N` | Max concurrent adapter slots. `0` disables (default); `> 0` enables. | -| `--multi-lora-adapter NAME PATH` | Register an adapter at startup. Repeatable. `PATH` → an `adapter.yaml`. | - -Per-adapter `rank` in `adapter.yaml` must be `<= --lora-rank`. - -## adapter.yaml - -```yaml -rank: 16 -alpha: 16 -rollout_batch_size: 32 # prompt groups per optimizer step (defaults to --rollout-batch-size) -n_samples_per_prompt: 4 # group shape (defaults to --n-samples-per-prompt) -data: /root/datasets/gsm8k/train.parquet -input_key: messages -label_key: label -rm_type: math -num_step: 400 # stop adapter after N optimizer steps - # (default: derived from num_epoch, itself default 1) -# optional: save, num_epoch, custom_rm_path, ... -``` - -The derived `adapter_global_batch_size = rollout_batch_size x -n_samples_per_prompt` is the adapter's samples-per-optimizer-step (the -per-adapter analog of `--global-batch-size`). - -Batch-shape constraints (validated at registration, not at runtime): -`n_samples_per_prompt` must be a divisor or multiple of the trainer's -data-parallel size; `rollout_batch_size` must be a multiple of the adapter's -`min_groups_per_dp_split`; -`adapter_global_batch_size` is capped by -`--multi-lora-max-adapter-global-batch-size` (default 4x `--global-batch-size`). diff --git a/examples/multi_lora/adapters/dapo_math.yaml b/examples/multi_lora/adapters/dapo_math.yaml deleted file mode 100644 index 3a1a3ff8a8b..00000000000 --- a/examples/multi_lora/adapters/dapo_math.yaml +++ /dev/null @@ -1,9 +0,0 @@ -rank: 32 -alpha: 32 -rollout_batch_size: 8 # prompt groups per optimizer step -n_samples_per_prompt: 8 # -> 64 samples per step -data: /root/datasets/dapo-math-17k/dapo-math-17k.jsonl -input_key: prompt -label_key: label -rm_type: deepscaler -num_step: 500 diff --git a/examples/multi_lora/adapters/gsm8k.yaml b/examples/multi_lora/adapters/gsm8k.yaml deleted file mode 100644 index 22906114647..00000000000 --- a/examples/multi_lora/adapters/gsm8k.yaml +++ /dev/null @@ -1,9 +0,0 @@ -rank: 16 -alpha: 16 -rollout_batch_size: 32 # prompt groups per optimizer step -n_samples_per_prompt: 4 # -> 128 samples per step -data: /root/datasets/gsm8k/train.parquet -input_key: messages -label_key: label -rm_type: math -num_step: 400 diff --git a/examples/multi_lora/run_multi_lora.py b/examples/multi_lora/run_multi_lora.py deleted file mode 100644 index 417d14af65e..00000000000 --- a/examples/multi_lora/run_multi_lora.py +++ /dev/null @@ -1,204 +0,0 @@ -"""Multi-LoRA fully-async GRPO example (Qwen3-4B, disaggregated 4 train + 4 rollout GPUs). - -Trains multiple LoRA adapters concurrently on a shared base model. Two example -adapters ship in ``adapters/``: gsm8k (rm_type=math) and dapo_math -(rm_type=deepscaler); each carries its own rank/alpha, batch shape, dataset, -reward, and ``num_step`` stop condition. The driver is -``train_multi_lora_async.py`` at the repo root; fully-async training forbids -``--colocate`` (generation needs continuous GPU). - -Usage: - python examples/multi_lora/run_multi_lora.py prepare # download Qwen3-4B + both datasets (once per node) - python examples/multi_lora/run_multi_lora.py train # bounded run: registers the two adapters, exits when each hits num_step - python examples/multi_lora/run_multi_lora.py full-train # prepare + train - python examples/multi_lora/run_multi_lora.py serve # service mode: no adapters preloaded, idles for registrations (API on :8068) - -Service mode pairs with the smoke client: - python examples/multi_lora/service_smoke.py --api-url http://127.0.0.1:8068 \\ - --data /root/datasets/gsm8k/train.parquet --input-key messages --label-key label --rm-type math -""" - -from dataclasses import dataclass - -import typer - -import miles.utils.external_utils.command_utils as U - -app = typer.Typer() - -_ADAPTER_DIR = f"{U.repo_base_dir}/examples/multi_lora/adapters" - - -@dataclass -class ScriptArgs(U.ExecuteTrainConfig): - run_id: str = U.create_run_id() - - hf_checkpoint: str | None = None - model_dir: str = "/root/models" - data_dir: str = "/root/datasets" - save_dir: str = "/tmp/multi_lora" - megatron_path: str = "/root/Megatron-LM" - - # Disaggregated split (fully-async forbids colocate). - num_gpus_per_node: int = 8 - actor_num_gpus: int = 4 - rollout_num_gpus: int = 4 - tp: int = 2 - - # LoRA slot pool. Per-adapter rank/alpha come from adapter.yaml, capped by lora_rank. - lora_rank: int = 32 - lora_alpha: int = 32 - lora_dropout: float = 0.0 - target_modules: str = "all-linear" - n_adapters: int = 4 - # Comma-separated adapter names; each resolves to adapters/{name}.yaml (train mode only). - adapters: str = "dapo_math,gsm8k" - - # Global rollout defaults; the per-adapter batch shapes live in the yamls. - num_rollout: int = 50 - rollout_batch_size: int = 32 - n_samples_per_prompt: int = 8 - rollout_max_response_len: int = 4096 - global_batch_size: int = 256 - max_weight_staleness: int = 3 - - # Service mode. - api_port: int = 8068 - - save_interval: int = 5 - enable_wandb: bool = False - extra_args: str = "" - - def __post_init__(self): - if self.hf_checkpoint is None: - self.hf_checkpoint = f"{self.model_dir}/Qwen3-4B" - - -def _prepare_download(args: ScriptArgs): - U.exec_command_cpu(f"mkdir -p {args.data_dir} {args.model_dir}") - U.exec_command_cpu(f"hf download Qwen/Qwen3-4B --local-dir {args.model_dir}/Qwen3-4B") - U.hf_download_dataset("zhuzilin/dapo-math-17k", data_dir=args.data_dir) - U.hf_download_dataset("zhuzilin/gsm8k", data_dir=args.data_dir) - - -def _train(args: ScriptArgs, service: bool): - mode = "service" if service else "bounded" - print( - f"[run] multi-LoRA ({mode}): {args.actor_num_gpus} train + {args.rollout_num_gpus} rollout GPUs, tp={args.tp}" - ) - - ckpt_args = f"--hf-checkpoint {args.hf_checkpoint} --megatron-to-hf-mode bridge " - - lora_args = ( - f"--lora-rank {args.lora_rank} --lora-alpha {args.lora_alpha} " - f'--lora-dropout {args.lora_dropout} --target-modules "{args.target_modules}" ' - ) - - multi_lora_args = f"--multi-lora-n-adapters {args.n_adapters} --multi-lora-idle-poll-s 5 " - if service: - # No adapters preloaded; the control-plane API accepts registrations at runtime. - multi_lora_args += f"--multi-lora-api-port {args.api_port} " - else: - for name in args.adapters.split(","): - multi_lora_args += f'--multi-lora-adapter "{name}" "{_ADAPTER_DIR}/{name}.yaml" ' - multi_lora_args += "--multi-lora-disable-service-mode " - - # in_place pause + upsert weight push is what lets adapters refresh without - # unloading (an unload would deadlock behind paused in-flight requests). - sync_args = f"--pause-generation-mode in_place --max-weight-staleness {args.max_weight_staleness} --use-tis " - - rollout_args = ( - "--apply-chat-template --rollout-shuffle " - f"--num-rollout {args.num_rollout} " - f"--rollout-batch-size {args.rollout_batch_size} " - f"--n-samples-per-prompt {args.n_samples_per_prompt} " - f"--rollout-max-response-len {args.rollout_max_response_len} " - "--rollout-temperature 1 " - f"--global-batch-size {args.global_batch_size} " - ) - - grpo_args = ( - "--advantage-estimator grpo --kl-loss-coef 0.00 --kl-coef 0.00 " - "--entropy-coef 0.00 --eps-clip 0.2 --eps-clip-high 0.28 " - ) - - optimizer_args = ( - "--optimizer adam --lr 1e-5 --lr-decay-style constant --weight-decay 0.1 " - "--adam-beta1 0.9 --adam-beta2 0.98 " - ) - - perf_args = ( - f"--tensor-model-parallel-size {args.tp} --sequence-parallel " - "--pipeline-model-parallel-size 1 --context-parallel-size 1 " - "--expert-model-parallel-size 1 --expert-tensor-parallel-size 1 " - "--use-dynamic-batch-size --max-tokens-per-gpu 9216 " - ) - - sglang_args = "--rollout-num-gpus-per-engine 1 --sglang-mem-fraction-static 0.8 " - - topology_args = ( - f"--actor-num-nodes 1 --actor-num-gpus-per-node {args.actor_num_gpus} " - f"--rollout-num-gpus {args.rollout_num_gpus} " - ) - - save_args = f"--save {args.save_dir} --save-interval {args.save_interval} " - - misc_args = ( - "--attention-dropout 0.0 --hidden-dropout 0.0 --accumulate-allreduce-grads-in-fp32 " - "--attention-softmax-in-fp32 --attention-backend flash " - ) - - wandb_args = U.get_default_wandb_args(__file__, run_id=args.run_id) if args.enable_wandb else "" - - train_args = ( - f"{ckpt_args} {lora_args} {multi_lora_args} {sync_args} {rollout_args} {grpo_args} " - f"{optimizer_args} {perf_args} {sglang_args} {topology_args} {save_args} {misc_args} " - f"{wandb_args} {args.extra_args} " - ) - - U.execute_train( - train_args=train_args, - config=args, - num_gpus_per_node=args.num_gpus_per_node, - megatron_model_type="qwen3-4B", - train_script="train_multi_lora_async.py", - megatron_path=args.megatron_path, - ) - - -@app.command() -@U.dataclass_cli -def prepare(args: ScriptArgs): - """Download Qwen3-4B and both task datasets. Run once per node before training.""" - _prepare_download(args) - - -@app.command() -@U.dataclass_cli -def train(args: ScriptArgs): - """Bounded run: register the adapters from adapters/, train until each hits num_step, exit.""" - _train(args, service=False) - - -@app.command() -@U.dataclass_cli -def full_train(args: ScriptArgs): - """Download model + datasets, then run the bounded training.""" - _prepare_download(args) - _train(args, service=False) - - -@app.command() -@U.dataclass_cli -def serve(args: ScriptArgs): - """Service mode: no adapters preloaded; register/deregister via the HTTP API while it idles.""" - _train(args, service=True) - - -@app.callback() -def _callback() -> None: - pass - - -if __name__ == "__main__": - app() diff --git a/examples/multi_lora/service_smoke.py b/examples/multi_lora/service_smoke.py deleted file mode 100644 index fb6685e34f2..00000000000 --- a/examples/multi_lora/service_smoke.py +++ /dev/null @@ -1,170 +0,0 @@ -"""Smoke test for multi-LoRA service mode: register/deregister against a running -trainer, using step counts as the race-free progress signal. - -Usage: python examples/multi_lora/service_smoke.py --api-url http://HOST:8068 \\ - --data /root/datasets/gsm8k/train.parquet --input-key messages --label-key label --rm-type math -""" - -import argparse -import sys -import time - -import httpx - -POLL_INTERVAL_S = 5.0 - - -class SmokeFailure(Exception): - pass - - -class ServiceClient: - def __init__(self, api_url: str, timeout_s: float): - self.api_url = api_url.rstrip("/") - self.timeout_s = timeout_s - self.http = httpx.Client(timeout=30.0) - - def adapters(self, states: set[str] | None = None) -> dict: - response = self.http.get(f"{self.api_url}/adapter_runs") - response.raise_for_status() - wanted_states = states if states is not None else {"ACTIVE"} - return { - status["name"]: { - "slot": status["slot"], - "version": status["version"], - "step": status["step"], - "state": status["state"], - } - for status in response.json()["adapters"] - if status["state"] in wanted_states - } - - def active_adapters(self) -> dict: - return self.adapters(states={"ACTIVE"}) - - def register(self, name: str, config: dict) -> httpx.Response: - return self.http.post(f"{self.api_url}/adapter_runs", json={"name": name, "config": config}) - - def deregister(self, name: str) -> None: - response = self.http.delete(f"{self.api_url}/adapter_runs/{name}") - response.raise_for_status() - - def wait_for(self, description: str, predicate) -> dict: - deadline = time.time() + self.timeout_s - while time.time() < deadline: - try: - adapters = self.active_adapters() - except httpx.HTTPError as e: - print(f" ... api not reachable yet ({e})") - adapters = None - if adapters is not None: - if predicate(adapters): - print(f" ok: {description} (active={adapters})") - return adapters - print(f" waiting for {description} (active={adapters})") - time.sleep(POLL_INTERVAL_S) - raise SmokeFailure(f"timed out after {self.timeout_s}s waiting for: {description}") - - def wait_for_step(self, name: str, min_step: int) -> None: - # Step-triggered deregistration can move an adapter to RETIRING quickly; - # count both ACTIVE and RETIRING for progress waits. - self.wait_for( - f"'{name}' to reach step {min_step}", - lambda _active: ( - (adapters := self.adapters(states={"ACTIVE", "RETIRING"})) - and name in adapters - and adapters[name]["step"] >= min_step - ), - ) - - def register_when_allowed(self, name: str, config: dict) -> None: - """Registration is rejected while a same-named adapter is cleaning up; - retry until the name frees.""" - deadline = time.time() + self.timeout_s - while time.time() < deadline: - response = self.register(name, config) - if response.status_code == 200: - print(f" ok: registered '{name}'") - return - print(f" register '{name}' rejected ({response.status_code}): {response.text[:200]}") - time.sleep(POLL_INTERVAL_S) - raise SmokeFailure(f"timed out registering '{name}'") - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--api-url", required=True, help="controller API listener, e.g. http://host:8068") - parser.add_argument("--data", required=True, help="prompt dataset path for the test adapters") - parser.add_argument("--input-key", default="text") - parser.add_argument("--label-key", default="label") - parser.add_argument("--rm-type", default="math") - parser.add_argument("--rank", type=int, default=16) - parser.add_argument("--alpha", type=int, default=16) - parser.add_argument("--save", default=None, help="per-adapter save dir root override (default: trainer --save)") - parser.add_argument("--steps", type=int, default=2, help="training steps to wait for per phase") - parser.add_argument( - "--num-step-smoke", - type=int, - default=1, - help="num_step used by the auto-deregister smoke adapter", - ) - parser.add_argument("--timeout", type=float, default=1800.0, help="per-phase timeout in seconds") - args = parser.parse_args() - - def config(name: str) -> dict: - cfg = { - "rank": args.rank, - "alpha": args.alpha, - "data": args.data, - "input_key": args.input_key, - "label_key": args.label_key, - "rm_type": args.rm_type, - } - if args.save: - cfg["save"] = f"{args.save}/{name}" - return cfg - - client = ServiceClient(args.api_url, args.timeout) - try: - print("phase 1: api reachable, no active adapters expected") - client.wait_for("api reachable", lambda adapters: True) - - print("phase 2: register smoke_auto with num_step; expect auto-deregister after committed steps") - auto_cfg = config("smoke_auto") - auto_cfg["num_step"] = args.num_step_smoke - client.register_when_allowed("smoke_auto", auto_cfg) - client.wait_for_step("smoke_auto", args.num_step_smoke) - client.wait_for("'smoke_auto' auto-deregistered", lambda adapters: "smoke_auto" not in adapters) - - print("phase 3: register smoke_a; expect promotion + training progress") - client.register_when_allowed("smoke_a", config("smoke_a")) - client.wait_for_step("smoke_a", args.steps) - - print("phase 4: register smoke_b mid-run; both must train") - client.register_when_allowed("smoke_b", config("smoke_b")) - client.wait_for_step("smoke_b", args.steps) - - print("phase 5: deregister smoke_a mid-run; smoke_b must keep training") - step_b = client.active_adapters()["smoke_b"]["step"] - client.deregister("smoke_a") - client.wait_for("'smoke_a' gone from active set", lambda adapters: "smoke_a" not in adapters) - client.wait_for_step("smoke_b", step_b + 1) - - print("phase 6: re-register the name smoke_a (waits out cleanup, reuses slot)") - client.register_when_allowed("smoke_a", config("smoke_a")) - client.wait_for_step("smoke_a", 1) - - print("phase 7: deregister everything; service should drain to idle") - client.deregister("smoke_a") - client.deregister("smoke_b") - client.wait_for("no active adapters", lambda adapters: not adapters) - - print("SMOKE TEST PASSED") - return 0 - except SmokeFailure as failure: - print(f"SMOKE TEST FAILED: {failure}", file=sys.stderr) - return 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/examples/multi_lora_operations/README.md b/examples/multi_lora_operations/README.md new file mode 100644 index 00000000000..eb783df6dc3 --- /dev/null +++ b/examples/multi_lora_operations/README.md @@ -0,0 +1,366 @@ +# Multi-LoRA operation backend with Tinker compatibility + +Serve many LoRA training runs on one shared base model through the +`MultiLoraOperationBackend`. The +[tinker](https://tinker-docs.thinkingmachines.ai/)-compatible frontend maps the +official SDK onto explicit `forward_backward` / `optim_step` operations and +shared-engine sampling — no dataset, reward function, or batch schedule on the +server. + +``` +official Tinker SDK ──HTTP──> Tinker protocol/frontend adapter + │ +internal caller ──Ray operations───────┘ + ▼ + MultiLoraOperationBackend (head node) + ├─ registration + operation ledger + ├─ adapter-slot execution + └─ serving plane ─────────> SGLang router +trainer ranks <──Ray── driver loop (train_multi_lora_operations.py) +``` + +## Start the Miles engine + +For the documented SDK flow, start both the operation backend and the Tinker +frontend. The helper starts the shared training and sampling engines in +service mode; add `--tinker-frontend` through `--extra-args` so that the +official SDK can use the controller's `/api/v1` endpoint: + +```bash +# Once per node: download the example checkpoint. +python examples/multi_lora_operations/run_multi_lora_operations.py prepare + +# Start Miles in service mode, with both the backend and frontend enabled. +python examples/multi_lora_operations/run_multi_lora_operations.py serve \ + --extra-args "--tinker-frontend" +``` + +The following lower-level command is useful when deploying with custom +Megatron and SGLang flags: + +```bash +python train_multi_lora_operations.py \ + --tinker-backend \ + --tinker-frontend \ + --multi-lora-n-adapters 4 \ + --lora-rank 32 --lora-alpha 64 \ + --target-modules all-linear \ + --hf-checkpoint Qwen/Qwen3-0.6B \ + ... # the usual megatron/sglang flags; see run_multi_lora_operations.py +``` + +Key flags: + +| flag | meaning | +|------|---------| +| `--tinker-backend` | enable the Tinker protocol adapter for the Multi-LoRA operation backend (requires `--multi-lora-n-adapters > 0`) | +| `--multi-lora-n-adapters N` | fixed slot count; a registration binds a slot for life (queue when full) | +| `--lora-rank` / `--lora-alpha` | deployment-wide ceiling / fixed alpha — clients may lower `rank`, never set `alpha` | +| `--multi-lora-api-port` | control-plane API port for runtime adapter registration | +| `--tinker-max-coalesce-wait-s` | how long one train call coalesces additional ready client batches | +| `--tinker-max-empty-wait-s` | idle-queue yield back to the control phase (keep this small) | +| `--tinker-frontend` | serve the official tinker SDK REST protocol (`/api/v1`) on the controller HTTP server (requires `--tinker-backend`) | +| `--tinker-api-key` | X-API-Key the frontend requires (prefer `$MILES_TINKER_API_KEY` — a CLI flag shows in the process list); mandatory for a non-loopback bind | + +The operator plane (`/adapter_runs*`, `/info`) accepts loopback peers only, +whatever the bind: the SDK key is a client credential and never grants the +routes that read server-local YAML files, choose save paths, or deregister +tenants. `/health` is liveness (the socket is up); `/api/v1/healthz` is +readiness and answers 503 until the driver reports the trainer exists. + +### Activation recompute (memory saving) + +`--recompute-granularity selective` is always supported (default +`--recompute-modules core_attn`; add `moe_act` to also recompute the MoE +activation with grouped GEMM). `--recompute-granularity full` — and `moe` in +`--recompute-modules` when expert modules are targeted — is supported as +well: the deployment's Megatron-Bridge (branch `bridge`) recognizes +multi-LoRA `.adapters..` params in its PEFT recompute patch, so +checkpointed regions replay grad-enabled during adapter-only training. + +## Operation contract + +`Tinker` names the compatibility boundary, not the trainer implementation. +The current concrete is `MultiLoraOperationBackend`; its queue-backed +`MultiLoraOperationBatchFn` batches already-tokenized operations, and the +Megatron `MultiLoraParameterExecutor` applies them to adapter slots. A future +full-parameter composition can reuse the same operation contract and the +unwired `FullParameterExecutor` sibling; full-parameter launch, data-path, +checkpoint, and publish integration are not implemented by this stack today. + +``` +Tinker protocol frontend + │ + ▼ +generic training-operation contract + │ + ├── MultiLoraParameterExecutor (current wired target) + └── FullParameterExecutor (implemented seam; not wired) +``` + +`enqueue_operation(name, operation_id, ordinal, kind, payload)` — ordinals are +consecutive per registration starting at 1; arrival may be out of order +(gap-buffered, and a hole-filling ordinal is always admitted), execution is +strictly ordinal-ordered; retries with the same `operation_id`, same ordinal, +and identical payload return the original operation — anything else is a +typed conflict. + +A stream stalled on a never-arriving ordinal (the 0.24.1 SDK consumes a +seq_id and can then fail BEFORE HTTP — see the SDK limitations below) expires +after `--tinker-operation-gap-timeout` (default 600 s, `<= 0` disables): the +blocked, never-claimed operations terminal-fail `FAILED(user)` naming the +missing ordinal, and the hole is sealed — the missing identity never executes +(a late arrival is a typed conflict), nothing overtakes it, and the client +resubmits as new operations. Stalls are observable before expiry: +`service_info()` reports `gap_stalls`, and a blocked operation's +`get_operation` view carries `waiting_on_ordinal` / `gap_stalled_for`. + +| kind | payload | success result | +|------|---------|----------------| +| `forward_backward` | `{samples: [Datum...], loss: {loss_fn, loss_fn_config?}}` | `{logprobs: [[...]], metrics: {"loss:sum", "unmasked_tokens:sum", "loss_weight:sum" (CE only)}}` | +| `forward` | `{samples: [Datum...]}` | `{logprobs: [[...]]}` (zero gradient, structurally) | +| `optim_step` | `{adam_params: {learning_rate, beta1, beta2, eps, weight_decay, grad_clip_norm}}` | `{grad_norm, learning_rate}` | +| `save_weights_for_sampler` | `{}` | `{serving_version, serving_name}` — completes only after the weights are live | +| `save_state` | `{tag?, ttl_seconds?}` | `{path, step}` (named states are immutable) | +| `load_state` | `{path}` | `{step, path}` (re-publishes on the next push) | + +`Datum = {tokens, response_length, loss_mask, loss_weights?, advantages?, rollout_log_probs?}` +— per-token channels align with the response span. Losses reduce as plain +token sums (`Σ(-logp·w)` for `cross_entropy`), so K chunked forward_backward +calls accumulate exactly like one; `loss_weights` own the scale and no server +normalization or scheduler ever touches a tinker slot. Result `metrics` use +the SDK combiner's `name:reduction` keys. + +For an SFT-style per-token loss, divide `loss:sum` by `loss_weight:sum` +(cross-entropy only: Σ weight·mask, chunk-additive like the loss) — NOT by +`unmasked_tokens:sum`, which counts every loss_mask-active position and so +includes the weight-0 prompt tokens of a teacher-forced datum, silently +diluting the displayed loss. Guard the division: weights are arbitrary +finite floats, so the sum can be zero or negative. + +Operation states: `QUEUED → CLAIMED → SUCCEEDED | FAILED(user|server) | CANCELLED`; +poll `get_operation`, then `ack_operation` to release the record. These verbs +are the controller actor's Ray API; the tinker frontend drives them over +HTTP. Backpressure raises a retryable `OperationBackpressure` — the frontend +maps it to 429 + Retry-After, never to a 4xx the SDK treats as fatal. +Deregistering fences every open operation of that registration as +`FAILED(user)`. + +Gradient-window poison: `optim_step` delimits a window of `forward_backward` +operations. If any of them reached a terminal state without succeeding (a +rejected chunk, an execution failure, a cancel), the window holds PARTIAL +gradients — the window's `optim_step` executes as a discard (all ranks clear +the slot's gradient sum), terminal-fails `FAILED(user)`, and moves neither +the step clock nor the serving version. The consumed poison resets the +window; resubmit the batch and step again. + +## Tinker SDK frontend (tinker==0.24.1 JSON subset) + +With `--tinker-frontend` the controller's HTTP server also speaks the REST +protocol of the official [`tinker`](https://pypi.org/project/tinker/) SDK — +exactly the **`tinker==0.24.1` JSON core-loop subset** (wheel source and +captured traffic; pure JSON, no protobuf: `/api/v1/client/config` pins the +SDK to its own default JSON path). Other SDK versions are rejected at +bootstrap (`/client/config` and `create_session` fail fast on the reported +`sdk_version`): 0.25+ switches `forward_backward` to protobuf, and the +current cookbook's canonical final checkpoint needs named sampler +checkpoints — neither is served here, so this is NOT "current +Tinker/cookbook compatible". An unmodified 0.24.1 client drives training +and sampling: + +```python +import tinker +sc = tinker.ServiceClient(base_url="http://127.0.0.1:8068", api_key="tml-...") +tc = sc.create_lora_training_client(base_model=..., rank=32) +tc.forward_backward(data, "cross_entropy") +tc.optim_step(tinker.types.AdamParams(learning_rate=1e-4)).result() +sampler = tc.save_weights_and_get_sampling_client() +future = sampler.sample( # sample()/sample_async() submit /api/v1/asample; + prompt=tinker.types.ModelInput.from_ints(prompt_tokens), + num_samples=4, + sampling_params=tinker.types.SamplingParams(max_tokens=128, temperature=0.7), +) +response = future.result() # .sequences[i].tokens / .logprobs / .stop_reason +``` + +### Client-owned RL loop + +After the engine reports ready, connect the official SDK client to the +frontend endpoint and run the loop below. The backend executes each requested +operation; rollout generation, scoring, and `Datum` construction remain in +the client. + +Start the driver with both `--tinker-backend` and `--tinker-frontend`. The +backend then owns execution and serving, while the client owns data +preparation and the training loop. In particular, the client can run the +same pattern as the [target-flow example](https://github.com/radixark/miles/issues/2258): + +```python +import tinker +from transformers import AutoTokenizer + +service = tinker.ServiceClient(base_url="http://127.0.0.1:8068", api_key="tml-...") +base_model = service.get_server_capabilities().supported_models[0].model_name +training = service.create_lora_training_client(base_model=base_model, rank=16) +tokenizer = AutoTokenizer.from_pretrained(base_model) + +# Publish the initial LoRA so the first rollout has a policy to sample. +sampler = training.save_weights_and_get_sampling_client() + +rl_prompts = ["Solve: If a train travels 60 km in 2 hours, what is its speed?"] +prompt_ids = [tokenizer(p).input_ids for p in rl_prompts] + +for update_idx in range(num_rl_updates): + # Option 1 -- SFT data preparation (client-owned; replace the RL batch + # below and train with loss_fn="cross_entropy"). + # batch = [ + # datum_from_sft_example(example["prompt"], example["completion"]) + # for example in sft_examples + # ] + + # Option 2 -- RL data preparation (client-owned). sample() returns a + # future; .result() carries sequences with tokens and logprobs. + futures = [ + sampler.sample( + prompt=tinker.types.ModelInput.from_ints(ids), + num_samples=4, + sampling_params=tinker.types.SamplingParams(max_tokens=256, temperature=1.0), + ) + for ids in prompt_ids + ] + rollouts = [future.result() for future in futures] + scored = score_rollouts(rl_prompts, rollouts) # rewards -> advantages, client-owned + batch = [ + datum_from_scored_rollout(ids, sequence, advantage) + for ids, response, advantages in zip(prompt_ids, rollouts, scored) + for sequence, advantage in zip(response.sequences, advantages) + ] + + fb = training.forward_backward(batch, "importance_sampling") + step = training.optim_step(tinker.types.AdamParams(learning_rate=1e-4)) + fb.result() + step.result() + + # Publish explicitly so the next rollout samples the new policy. + # Serving is latest-only: the publish supersedes the previous sampling + # client, so re-acquire it here every update. + sampler = training.save_weights_and_get_sampling_client() +``` + +`datum_from_sft_example`, `score_rollouts`, and `datum_from_scored_rollout` +are application code: they define the task data, rollout scoring, and the +per-token loss channels. An RL datum pairs `model_input` (prompt + sampled +tokens, shifted) with `loss_fn_inputs` `target_tokens`, the sampler's +returned `logprobs`, and per-token `advantages`; an SFT datum needs +`target_tokens` plus 0/1 `weights`. The frontend translates the resulting +SDK requests to operations; the backend executes them in order and only +changes the sampler's policy on the explicit publish. The complete runnable +version of this loop is `tests/e2e/tinker_frontend/tinker_sdk_rl_quality.py` +(GRPO on GSM8K, four concurrent adapters through one deployment). + +Mapping: one training client = one registration (`create_model` registers, +`unload_model` deregisters), and every operation is pinned to its +`(name, registration_id)` — a stale handle fences instead of binding to a +same-name successor; every training verb forwards its SDK `seq_id` as the +registration ordinal (chunks posted out of order gap-buffer); futures poll +`/api/v1/retrieve_future` and terminal bodies replay until delivered (an +evicted delivered result leaves a fingerprint tombstone that answers a typed +410 — the 0.24.1 SDK surfaces it as a retryable "promise expired", it does +not re-run the original request); `save_state` mints `tinker://` paths +(resolved from an in-memory catalog; failures echo the public URI, not the +trainer filesystem); the ephemeral `save_weights_and_get_sampling_client` +publish binds `(name, registration_id, serving_version)` and samples through +the sglang router — a republish makes older sampling clients fail loud, and +the version is re-checked after generation so a publish landing mid-flight +fails the in-flight sample instead of returning cross-version output (the +identity is versioned, not leased: a publish committing between that check +and delivery is a documented residual race). Frontend rejections on a spent +`seq_id` become terminal `FAILED(user)` futures so the ordinal is still +consumed — bounded by the same unacked-results budget as every other record +(429 past it). + +Frontend-level v1 rejections (beyond the backend matrix): non-0.24.x SDK +versions, LoRA `seed` and per-module `train_*` flags (deployment-wide), +weights-only restore (`load_state` / `create_training_client_from_state` — +the backend restores the full training state; use the `_with_optimizer` +variants), named persistent sampler checkpoints +(`save_weights_for_sampler(name)` / `create_sampling_client(model_path=...)`), +`ttl_seconds` (checkpoint/sampler TTL expiry is not implemented), +`topk_prompt_logprobs`, sparse-CSR tensors, and negative +token ids anywhere (targets, inputs, prompts, stop tokens). A sampling +`seed` maps to sglang `sampling_seed`, offset per sample so +`num_samples > 1` stays diverse. `prompt_logprobs` maps to sglang +`logprob_start_len=0` on the same generate (the engine scores the prompt +natively; position 0 has no context and returns null) — this serves both +`sample(include_prompt_logprobs=True)` and the SDK's `compute_logprobs()`, +which the 0.24.1 wheel sends as a 1-sample, 1-token generation. + +Sampling architecture: `/asample` returns its future immediately and a +background task posts one router `/generate` per sample, carrying the +server-derived serving identity (`rid`/`lora_path`/`extra_key` are never +client-controllable — the wire models drop unknown fields and the sglang +params are rebuilt from an allowlist). SGLang's continuous batching is the +only sampling batcher: the frontend never coalesces prompts, and the +training-operation scheduler (`MultiLoraOperationBatchFn`) never sees a sampling +request. The legacy datasource rollout pipeline +(`RolloutManager.generate()`: datasets, rewards, training-data conversion) +is not on this path — the frontend shares only the router the rollout +engines already serve. + +Trust boundary (v1): the frontend authenticates clients and bounds aggregate +active sub-generations, rejects one request whose `num_samples` exceeds that +capacity, and preflights `prompt + max_tokens` against the discovered engine +limit. It still does not validate token ids against the vocabulary upper bound +or enforce request-body/output-byte quotas. Run it loopback/VPN-facing for +trusted clients; per-tenant quotas are future work. + +## v1 compatibility matrix + +Supported: text-only input; the synchronous training loop; 1-D shifted +targets; `loss_fn ∈ {cross_entropy, importance_sampling, ppo}` (per-op clip +config); per-call AdamParams; multi-chunk gradient accumulation with +independent `optim_step`; latest-only sampler weights behind the publish +barrier; prompt logprobs (`compute_logprobs()` / +`sample(include_prompt_logprobs=True)`, one sub-generation of admission +weight); named immutable `save_state` / `load_state` (create-from-checkpoint +included, shape-fenced); optional `num_step` auto-retirement. + +Explicitly rejected (boundary error, never a silent fallback): multimodal +inputs; nested `(N, K)` top-K targets; other loss functions (CISPO, DRO, ...); +client-set `alpha`; non-finite/out-of-domain AdamParams; a loss's required +per-token channels missing; `response_length == len(tokens)` (targets are +shifted); async/off-policy sampling against pinned snapshots; +cross-world-size state restore; state restore into a slot whose per-rank +optimizer ownership differs from the save (cross-slot restore requires an +identical dense-and-expert ownership signature); idle slot GC. + +## Known tinker SDK (0.24.1) client-side limitations + +The official `tinker==0.24.1` TrainingClient takes its per-model seq counter +BEFORE it serializes and POSTs a request, so a submission can die client-side +with the ordinal already spent (verified against the live stack, +codex-0817-sft-fix §4-§6): + +- **Pre-HTTP serialization failure** — e.g. `AdamParams(learning_rate=nan)` + raises a local JSON `ValueError`; the request never reaches Miles and later + operations of the same client queue behind the hole. The gap timeout above + terminal-fails them typed, and the SAME TrainingClient can resubmit + afterwards (its turn counter did advance). Validate that Adam params and + custom scalars are finite before calling the SDK to avoid the stall. +- **`.future().cancel()` on an SDK future** can spend the request id without + advancing the SDK's internal turn counter: later operations of that client + wait forever CLIENT-side and Miles receives nothing it could terminalize — + no server-side mitigation exists. Do not cancel underlying SDK futures; + `.result(timeout=...)` is safe (non-destructive, the future stays + retrievable). After an immediate cancel, discard the TrainingClient and + create a new one (a fresh registration). When some submissions did reach + the server, the gap timeout converts the surviving stall into typed + failures instead of a hang. +- The server never skips a missing ordinal and never guesses what it would + have been: the gap timeout only fails what is blocked and seals the hole, + so strict per-registration ordering, idempotent retries, and anti-replay + all hold. + +## Files + +- `run_multi_lora_operations.py` — disaggregated service launch (`prepare` / `serve`) diff --git a/examples/multi_lora_operations/run_multi_lora_operations.py b/examples/multi_lora_operations/run_multi_lora_operations.py new file mode 100644 index 00000000000..7e78cadc72f --- /dev/null +++ b/examples/multi_lora_operations/run_multi_lora_operations.py @@ -0,0 +1,134 @@ +from dataclasses import dataclass + +import typer + +import miles.utils.external_utils.command_utils as U + +app = typer.Typer() + + +@dataclass +class ScriptArgs(U.ExecuteTrainConfig): + """Launch configuration for the Multi-LoRA operation backend example. + + Full-parameter targets can reuse the protocol-neutral operation contract, + not the LoRA slot and adapter settings defined by this launcher. + """ + + run_id: str = U.create_run_id() + + hf_checkpoint: str | None = None + model_dir: str = "/root/models" + save_dir: str = "/tmp/multi_lora_operations" + megatron_path: str = "/root/Megatron-LM" + + # Disaggregated split (the operation backend forbids colocate). + num_gpus_per_node: int = 8 + actor_num_gpus: int = 4 + rollout_num_gpus: int = 4 + tp: int = 2 + + # Deployment-wide LoRA slot constraints. + max_lora_rank: int = 32 + backend_lora_alpha: int = 64 + backend_target_modules: str = "all-linear" + max_adapters: int = 4 + + # Soft coalescing target for one train call (whole client batches only). + backend_batch_size: int = 32 + + api_port: int = 8068 + enable_wandb: bool = False + extra_args: str = "" + + def __post_init__(self): + if self.hf_checkpoint is None: + self.hf_checkpoint = f"{self.model_dir}/Qwen3-4B" + + +@app.command() +@U.dataclass_cli +def prepare(args: ScriptArgs): + """Download Qwen3-4B. Run once per node before serving.""" + U.exec_command_cpu(f"mkdir -p {args.model_dir}") + U.exec_command_cpu(f"hf download Qwen/Qwen3-4B --local-dir {args.model_dir}/Qwen3-4B") + + +def _serve(args: ScriptArgs): + print( + f"[run] Multi-LoRA operations (service): " + f"{args.actor_num_gpus} train + {args.rollout_num_gpus} rollout GPUs" + ) + + ckpt_args = f"--hf-checkpoint {args.hf_checkpoint} --megatron-to-hf-mode bridge " + lora_args = ( + f"--lora-rank {args.max_lora_rank} --lora-alpha {args.backend_lora_alpha} " + f'--lora-dropout 0.0 --target-modules "{args.backend_target_modules}" ' + ) + tinker_args = ( + f"--tinker-backend --multi-lora-n-adapters {args.max_adapters} " + f"--multi-lora-idle-poll-s 5 --multi-lora-api-port {args.api_port} " + ) + + # in_place pause + upsert push: adapters publish without unloading. + sync_args = "--pause-generation-mode in_place " + + rollout_args = ( + f"--rollout-batch-size {args.backend_batch_size} " + f"--n-samples-per-prompt 1 --global-batch-size {args.backend_batch_size} " + "--num-rollout 1000000 " + ) + + optimizer_args = "--optimizer adam --lr 1e-4 --lr-decay-style constant " + + perf_args = ( + f"--tensor-model-parallel-size {args.tp} --sequence-parallel " + "--pipeline-model-parallel-size 1 --context-parallel-size 1 " + "--expert-model-parallel-size 1 --expert-tensor-parallel-size 1 " + "--use-dynamic-batch-size --max-tokens-per-gpu 9216 " + ) + + sglang_args = "--rollout-num-gpus-per-engine 1 --sglang-mem-fraction-static 0.8 " + topology_args = ( + f"--actor-num-nodes 1 --actor-num-gpus-per-node {args.actor_num_gpus} " + f"--rollout-num-gpus {args.rollout_num_gpus} " + ) + # Tinker checkpoints move only through save_state operations, but megatron + # arg validation requires a save interval whenever --save is set. + save_args = f"--save {args.save_dir} --save-interval 1000000 " + misc_args = ( + "--attention-dropout 0.0 --hidden-dropout 0.0 --accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 --attention-backend flash " + ) + wandb_args = U.get_default_wandb_args(__file__, run_id=args.run_id) if args.enable_wandb else "" + + train_args = ( + f"{ckpt_args} {lora_args} {tinker_args} {sync_args} {rollout_args} " + f"{optimizer_args} {perf_args} {sglang_args} {topology_args} {save_args} {misc_args} " + f"{wandb_args} {args.extra_args} " + ) + + U.execute_train( + train_args=train_args, + config=args, + num_gpus_per_node=args.num_gpus_per_node, + megatron_model_type="qwen3-4B", + train_script="train_multi_lora_operations.py", + megatron_path=args.megatron_path, + ) + + +@app.command() +@U.dataclass_cli +def serve(args: ScriptArgs): + """Service mode: no adapters preloaded; register via the HTTP API while it idles.""" + _serve(args) + + +@app.callback() +def _callback() -> None: + pass + + +if __name__ == "__main__": + app() diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 4681c67ebc9..aa5cc41b016 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -24,13 +24,13 @@ from miles.utils.ft_utils.indep_dp import IndepDPInfo from miles.utils.hf_config import load_hf_config from miles.utils.memory_utils import clear_memory, print_memory -from miles.utils.multi_lora import is_multi_lora_enabled from miles.utils.processing_utils import load_tokenizer from miles.utils.ray_utils import Box from miles.utils.reloadable_process_group import destroy_process_groups, monkey_patch_torch_dist, reload_process_groups from miles.utils.replay_base import all_replay_managers, routing_replay_manager from miles.utils.test_utils.ft_test_actions import FTTestActionActorExecutor from miles.utils.timer import Timer, inverse_timer, timer +from miles.utils.tinker import is_tinker_enabled from miles.utils.tracking_utils.structured_log import with_logs from miles.utils.tracking_utils.tracking import init_tracking from miles.utils.types import RolloutBatch @@ -474,6 +474,12 @@ def train_actor( witness_info: WitnessInfo | None, attempt: int, ) -> TrainStepOutcome: + if rollout_data.get("batch_kind") == "tinker": + from miles.backends.megatron_utils.api_backends.multi_lora.trainer import validate_batch_lease + + validate_batch_lease(rollout_data, self.loaded_adapters) + rollout_data["tinker_logprob_collector"] = {} + # Create data iterator for log_probs and train. data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data) num_optimizer_steps = len(num_microbatches) @@ -499,7 +505,7 @@ def train_actor( ) with inverse_timer("train_wait"), timer("train"): - if self.args.compute_advantages_and_returns: + if self.args.compute_advantages_and_returns and rollout_data.get("batch_kind") != "tinker": if "ref" in self.weights_backuper.backup_tags: self._set_replay_stage("fallthrough") self._switch_model("ref") @@ -584,6 +590,7 @@ def train_actor( witness_info=witness_info, attempt=attempt, ft_test_action_executor=self._ft_test_action_executor, + forward_only=bool(rollout_data.get("tinker_forward_only")), ) self.prof.step(rollout_id=rollout_id) @@ -612,10 +619,10 @@ def train_actor( logger.info(f"Updating ref model at rollout_id {rollout_id}") self.weights_backuper.backup("ref") - if train_step_outcome == TrainStepOutcome.NORMAL and is_multi_lora_enabled(self.args): - from miles.backends.megatron_utils.multi_lora_utils import commit_trained_batch + if train_step_outcome == TrainStepOutcome.NORMAL and rollout_data.get("batch_kind") == "tinker": + from miles.backends.megatron_utils.api_backends.multi_lora.trainer import commit_batch - commit_trained_batch(rollout_data, rollout_id, self._multi_lora_pending_push) + commit_batch(rollout_data, self._multi_lora_pending_push) log_perf_data(rollout_id, self.args, extra_metrics=self.weight_updater.pop_metrics()) @@ -624,53 +631,35 @@ def train_actor( @with_logs @timer - def reconcile_adapters(self) -> None: - """Load adapters the controller wants served; retire deregistered ones, dropping their untrained tail.""" - if not is_multi_lora_enabled(self.args): - return - from miles.backends.megatron_utils.multi_lora_utils import cleanup_adapters as _cleanup_adapters - from miles.backends.megatron_utils.multi_lora_utils import load_adapters as _load_adapters - from miles.ray.multi_lora.controller import get_multi_lora_controller - - broadcast_buffer = [None] - if is_first_replica_megatron_main_rank(): - controller = get_multi_lora_controller() - ray.get(controller.retire_adapters.remote()) - broadcast_buffer[0] = ray.get(controller.snapshot.remote()) - if dist.is_initialized(): - dist.broadcast_object_list(broadcast_buffer, src=0, group=get_gloo_group()) - snapshot = broadcast_buffer[0] - should_be_loaded = {**snapshot["active"], **snapshot["pending"], **snapshot["retiring"]} - cleanup_names = set(snapshot["cleanup"]) - - loaded_names = set(self.loaded_adapters) - # Sorted so per-adapter collectives (checkpoint export) run in the same - # order on every rank; set iteration order is process-specific. - adapters_to_load = sorted( - (adapter for name, adapter in should_be_loaded.items() if name not in loaded_names), - key=lambda adapter: adapter.name, - ) - adapters_to_clean_up = sorted( - (self.loaded_adapters[n] for n in loaded_names if n in cleanup_names or n not in should_be_loaded), - key=lambda adapter: adapter.name, + def execute_tinker_controls(self, operations: list[dict], lease_metadata: dict) -> dict: + from miles.backends.megatron_utils.api_backends.multi_lora.trainer import execute_controls + + return execute_controls( + self.args, + self.model, + self.optimizer, + self.loaded_adapters, + self._multi_lora_pending_push, + self.weights_backuper, + operations, + lease_metadata, ) - if adapters_to_load: - _load_adapters(self.args, self.model, self.optimizer, adapters_to_load) - for adapter in adapters_to_load: - self.loaded_adapters[adapter.name] = adapter - self._multi_lora_pending_push.add(adapter.name) - self.weights_backuper.backup("actor") - if adapters_to_clean_up: - _cleanup_adapters(self.args, self.model, self.optimizer, adapters_to_clean_up) - for adapter in adapters_to_clean_up: - self.loaded_adapters.pop(adapter.name, None) - self._multi_lora_pending_push.discard(adapter.name) - self.weights_backuper.backup("actor") - # Deregistered before ever being loaded: nothing to save or clear. - if is_first_replica_megatron_main_rank(): - for name in cleanup_names - loaded_names: - ray.get(get_multi_lora_controller().free_slot.remote(name)) + @with_logs + @timer + def reconcile_tinker_adapters(self) -> None: + if not is_tinker_enabled(self.args): + return + from miles.backends.megatron_utils.api_backends.multi_lora.trainer import reconcile_adapters + + reconcile_adapters( + self.args, + self.model, + self.optimizer, + self.loaded_adapters, + self._multi_lora_pending_push, + self.weights_backuper, + ) @timer def save_model(self, rollout_id: int, force_sync: bool = False) -> None: @@ -683,13 +672,10 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None: maybe_finalize_async_save(blocking=True) - if is_multi_lora_enabled(self.args): - from miles.backends.megatron_utils.multi_lora_utils import save_due_adapter_checkpoints + if is_tinker_enabled(self.args): + return - if not save_due_adapter_checkpoints(self.args, self.model): - return - else: - save(rollout_id, self.model, self.optimizer, self.opt_param_scheduler) + save(rollout_id, self.model, self.optimizer, self.opt_param_scheduler) if force_sync and self.args.async_save: maybe_finalize_async_save(blocking=True) @@ -770,12 +756,16 @@ def update_weights(self, info: "EnginesAndLock") -> None: return version_update_names: list[str] = [] - if is_multi_lora_enabled(self.args): - from miles.backends.megatron_utils.multi_lora_utils import select_adapters_to_push + if is_tinker_enabled(self.args): + from miles.backends.megatron_utils.api_backends.multi_lora.trainer import select_adapters_to_push self.weight_updater.multi_lora_adapters, version_update_names = select_adapters_to_push( self.loaded_adapters, self._multi_lora_pending_push, has_new_engines ) + if not self.weight_updater.multi_lora_adapters: + if process_groups_are_temporary: + destroy_process_groups() + return with torch_memory_saver.disable() if self.args.offload_train else nullcontext(): print_memory("before update_weights") @@ -784,8 +774,8 @@ def update_weights(self, info: "EnginesAndLock") -> None: if dist.get_rank() == 0: ray.get(self.rollout_manager.set_weight_version.remote(self.weight_updater.weight_version)) - if is_multi_lora_enabled(self.args): - from miles.backends.megatron_utils.multi_lora_utils import commit_weight_push + if is_tinker_enabled(self.args): + from miles.backends.megatron_utils.api_backends.multi_lora.trainer import commit_weight_push self._multi_lora_pending_push.clear() commit_weight_push(version_update_names, self._is_first_replica_megatron_main_rank) diff --git a/miles/backends/megatron_utils/api_backends/__init__.py b/miles/backends/megatron_utils/api_backends/__init__.py new file mode 100644 index 00000000000..f7c4430e18f --- /dev/null +++ b/miles/backends/megatron_utils/api_backends/__init__.py @@ -0,0 +1 @@ +"""Megatron implementations behind the protocol-neutral training-operation API.""" diff --git a/miles/backends/megatron_utils/api_backends/full_parameter/__init__.py b/miles/backends/megatron_utils/api_backends/full_parameter/__init__.py new file mode 100644 index 00000000000..5c86f93d356 --- /dev/null +++ b/miles/backends/megatron_utils/api_backends/full_parameter/__init__.py @@ -0,0 +1,5 @@ +"""Full-parameter implementation of the generic operation executor port.""" + +from .executor import FullParameterBinding, FullParameterExecutor + +__all__ = ["FullParameterBinding", "FullParameterExecutor"] diff --git a/miles/backends/megatron_utils/api_backends/full_parameter/executor.py b/miles/backends/megatron_utils/api_backends/full_parameter/executor.py new file mode 100644 index 00000000000..792991b3964 --- /dev/null +++ b/miles/backends/megatron_utils/api_backends/full_parameter/executor.py @@ -0,0 +1,272 @@ +"""Execute protocol-neutral operations against one whole-model optimizer. +Each dispatch lease contains exactly one operation for the target.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed as dist + +from miles.backends.training_utils.operation_execution import StepRequest, resolve_adam_params +from miles.utils.operation_contract import BatchExecutionLease + + +@dataclass(frozen=True) +class FullParameterBinding: + """Immutable binding for one executor-owned whole-model target. + ``target_id`` identifies the deployment rather than an adapter slot.""" + + target_id: str + + +def _server_error(message: str, *, consumed: bool = False) -> dict: + outcome = dict(ok=False, error=message, category="server") + if consumed: + outcome["gradient_window_consumed"] = True + return outcome + + +@dataclass +class FullParameterExecutor: + """Execute one operation at a time against a stock Megatron optimizer. + Validate the request and singleton lease before mutating model state.""" + + model_chunks: Sequence[Any] + optimizer: Any + binding: FullParameterBinding + + def discard_many( + self, + lease: BatchExecutionLease[FullParameterBinding], + operation_ids: list[str], + ) -> dict[str, dict]: + if not operation_ids: + return {} + refusal = self._validate_singleton_lease(lease, operation_ids) + if refusal is not None: + return {operation_id: _server_error(refusal) for operation_id in operation_ids} + + runtime_error = self._validate_clear_runtime() + if runtime_error is not None: + return {operation_ids[0]: _server_error(runtime_error)} + + self._clear_gradient_window() + return {operation_ids[0]: dict(ok=True, gradient_window_consumed=True)} + + def step_many( + self, + lease: BatchExecutionLease[FullParameterBinding], + requests: list[StepRequest], + ) -> dict[str, dict]: + if not requests: + return {} + operation_ids = [request.operation_id for request in requests] + refusal = self._validate_singleton_lease(lease, operation_ids) + if refusal is not None: + return {operation_id: _server_error(refusal) for operation_id in operation_ids} + + request = requests[0] + runtime_error = self._validate_step_runtime() + if runtime_error is not None: + return {request.operation_id: _server_error(runtime_error)} + + try: + adam = resolve_adam_params(request.adam_params) + except Exception as exc: + return {request.operation_id: _server_error(f"invalid Adam parameters: {exc}")} + + update_successful = False + grad_norm: float | None = None + nonfinite_veto = False + primary_error: BaseException | None = None + primary_traceback = None + finalization_errors: list[str] = [] + config = self.optimizer.config + previous_clip = config.clip_grad + try: + self._apply_adam_to_param_groups(adam) + # MCore measures grad_norm only in its clip branch; infinity keeps + # ``0 = no clipping`` while still requesting the measurement. + config.clip_grad = adam["grad_clip_norm"] if adam["grad_clip_norm"] > 0.0 else float("inf") + nonfinite_veto = self._has_nonfinite_gradient_norm() + if not nonfinite_veto: + raw_outcome = self.optimizer.step() + if not isinstance(raw_outcome, tuple) or len(raw_outcome) != 3: + raise RuntimeError( + "stock optimizer.step() did not return (update_successful, grad_norm, num_zeros)" + ) + update_successful, raw_grad_norm, _ = raw_outcome + update_successful = bool(update_successful) + if update_successful and raw_grad_norm is None: + raise RuntimeError("stock optimizer.step() did not report a gradient norm") + if raw_grad_norm is not None: + grad_norm = float(raw_grad_norm) + except BaseException as exc: + # Execution failures may follow a partial physical update, so keep + # them fatal rather than returning a recoverable operation result. + primary_error = exc + primary_traceback = exc.__traceback__ + finally: + try: + config.clip_grad = previous_clip + except Exception as exc: + finalization_errors.append(f"failed to restore optimizer clip_grad: {exc}") + try: + self._clear_gradient_window() + except Exception as exc: + finalization_errors.append(str(exc)) + + if primary_error is not None: + if finalization_errors: + raise RuntimeError( + f"full-parameter optimizer execution failed ({primary_error}); " + f"finalization also failed: {'; '.join(finalization_errors)}" + ) from primary_error + raise primary_error.with_traceback(primary_traceback) + if finalization_errors: + raise RuntimeError("; ".join(finalization_errors)) + if nonfinite_veto: + return { + request.operation_id: _server_error( + "non-finite gradient norm; step vetoed and gradients cleared", + consumed=True, + ) + } + if not update_successful: + return { + request.operation_id: _server_error( + "stock optimizer vetoed the step; gradients cleared", + consumed=True, + ) + } + return { + request.operation_id: dict( + ok=True, + gradient_window_consumed=True, + result=dict(grad_norm=grad_norm, learning_rate=adam["learning_rate"]), + ) + } + + def _validate_singleton_lease( + self, + lease: BatchExecutionLease[FullParameterBinding], + operation_ids: list[str], + ) -> str | None: + if len(operation_ids) != 1: + return ( + f"full-parameter execution requires exactly one operation per dispatch; received {len(operation_ids)}" + ) + try: + bindings = lease.bindings_by_operation + if len(bindings) != 1: + return f"full-parameter execution requires a singleton whole-model lease; received {len(bindings)} bindings" + leased_operation_id, leased_binding = bindings[0] + operation_id = operation_ids[0] + if leased_operation_id != operation_id: + return f"operation '{operation_id}' is not the singleton operation in dispatch '{lease.dispatch_id}'" + if leased_binding != self.binding: + return f"operation '{operation_id}' is not bound to this executor's whole-model target" + except Exception as exc: + return f"invalid full-parameter batch lease: {exc}" + return None + + def _validate_clear_runtime(self) -> str | None: + if not isinstance(self.model_chunks, Sequence): + return "full-parameter executor model must be a sequence of model chunks" + if not self.model_chunks: + return "full-parameter executor requires at least one model chunk" + for index, model_chunk in enumerate(self.model_chunks): + if not callable(getattr(model_chunk, "zero_grad_buffer", None)): + return f"model chunk {index} does not provide zero_grad_buffer()" + if not callable(getattr(model_chunk, "parameters", None)): + return f"model chunk {index} does not provide parameters()" + if not callable(getattr(self.optimizer, "zero_grad", None)): + return "stock optimizer does not provide zero_grad()" + return None + + def _validate_step_runtime(self) -> str | None: + clear_error = self._validate_clear_runtime() + if clear_error is not None: + return clear_error + if not callable(getattr(self.optimizer, "step", None)): + return "stock optimizer does not provide step()" + config = getattr(self.optimizer, "config", None) + if config is None or not hasattr(config, "clip_grad"): + return "stock optimizer config does not provide clip_grad" + if str(getattr(config, "optimizer", "")).lower() != "adam": + return "full-parameter explicit operations require an Adam optimizer" + try: + param_groups = self.optimizer.param_groups + except Exception as exc: + return f"stock optimizer param_groups are unavailable: {exc}" + if not isinstance(param_groups, Sequence) or not param_groups: + return "stock optimizer must expose at least one parameter group" + if any(not isinstance(group, dict) for group in param_groups): + return "stock optimizer parameter groups must be dictionaries" + return None + + def _apply_adam_to_param_groups(self, adam: dict[str, float]) -> None: + for group in self.optimizer.param_groups: + group["lr"] = adam["learning_rate"] + group["betas"] = (adam["beta1"], adam["beta2"]) + group["eps"] = adam["eps"] + group["weight_decay"] = adam["weight_decay"] + + def _has_nonfinite_gradient_norm(self) -> bool: + """Veto NaN/Inf before the stock BF16 optimizer mutates parameters. + Scan all visible gradients and reduce their squared norm across ranks.""" + + gradients: list[torch.Tensor] = [] + seen: set[int] = set() + + def append_gradient(candidate: Any) -> None: + if candidate is None: + return + candidate = getattr(candidate, "_local_tensor", candidate) + if not isinstance(candidate, torch.Tensor) or id(candidate) in seen: + return + seen.add(id(candidate)) + gradients.append(candidate.coalesce().values() if candidate.is_sparse else candidate) + + for model_chunk in self.model_chunks: + for parameter in model_chunk.parameters(): + append_gradient(getattr(parameter, "main_grad", None)) + append_gradient(getattr(parameter, "grad", None)) + append_gradient(getattr(parameter, "decoupled_grad", None)) + for group in self.optimizer.param_groups: + for parameter in group.get("params", ()): + append_gradient(getattr(parameter, "main_grad", None)) + append_gradient(getattr(parameter, "grad", None)) + append_gradient(getattr(parameter, "decoupled_grad", None)) + + if gradients: + reduction_device = gradients[0].device + elif dist.is_initialized() and dist.get_backend() == "nccl": + reduction_device = torch.device("cuda", torch.cuda.current_device()) + else: + reduction_device = torch.device("cpu") + + squared_norm = torch.zeros(1, dtype=torch.float32, device=reduction_device) + for gradient in gradients: + local_norm = torch.linalg.vector_norm(gradient.detach().float()) + squared_norm.add_(local_norm.to(reduction_device).square()) + if dist.is_initialized(): + dist.all_reduce(squared_norm, op=dist.ReduceOp.SUM) + return not bool(torch.isfinite(squared_norm).item()) + + def _clear_gradient_window(self) -> None: + errors: list[str] = [] + for index, model_chunk in enumerate(self.model_chunks): + try: + model_chunk.zero_grad_buffer() + except Exception as exc: + errors.append(f"model chunk {index} zero_grad_buffer() failed: {exc}") + try: + self.optimizer.zero_grad() + except Exception as exc: + errors.append(f"optimizer zero_grad() failed: {exc}") + if errors: + raise RuntimeError("; ".join(errors)) diff --git a/miles/backends/megatron_utils/api_backends/multi_lora/__init__.py b/miles/backends/megatron_utils/api_backends/multi_lora/__init__.py new file mode 100644 index 00000000000..2e9a1234fd3 --- /dev/null +++ b/miles/backends/megatron_utils/api_backends/multi_lora/__init__.py @@ -0,0 +1 @@ +"""Megatron execution for client-driven Multi-LoRA training operations.""" diff --git a/miles/backends/megatron_utils/api_backends/multi_lora/checkpoint.py b/miles/backends/megatron_utils/api_backends/multi_lora/checkpoint.py new file mode 100644 index 00000000000..f2de3b7c3fe --- /dev/null +++ b/miles/backends/megatron_utils/api_backends/multi_lora/checkpoint.py @@ -0,0 +1,220 @@ +"""Serialize per-slot Multi-LoRA weights, optimizer state, and clocks. +Atomic rank shards and a committed manifest fence restores against torn saves.""" + +import hashlib +import logging +import os +import re +from pathlib import Path + +import torch +import torch.distributed as dist + +from miles.utils.distributed_utils import get_gloo_group + +logger = logging.getLogger(__name__) + +FORMAT = "miles-tinker-slot-v1" +_SLOT_INDEX = re.compile(r"\.adapters\.(\d+)\.") + + +def stable_slot_param_name(name: str, slot: int) -> str: + return _SLOT_INDEX.sub(lambda m: ".adapter." if int(m.group(1)) == slot else m.group(0), name) + + +def named_adapter_slot_parameters(model, slot: int): + from megatron.bridge.peft.multi_lora_layers import MultiLoRALinear + + marker = f".adapters.{slot}." + seen: set[int] = set() + model_chunks = model if isinstance(model, (list, tuple)) else [model] + for model_chunk in model_chunks: + for module_name, module in model_chunk.named_modules(): + if not isinstance(module, MultiLoRALinear): + continue + for param_name, param in module.named_parameters(prefix=module_name): + if marker in param_name and id(param) not in seen: + seen.add(id(param)) + yield stable_slot_param_name(param_name, slot), param + + +def _slot_children(optimizer, slot: int): + return [optimizer.chained_optimizers[i] for i in optimizer.miles_slot_child_indices[slot]] + + +def _slot_child_param_names(model, optimizer, slot: int) -> list[list[str | None]]: + names_by_param: dict[int, str] = {} + for name, param in named_adapter_slot_parameters(model, slot): + names_by_param[id(param)] = name + # fp16/bf16 children hold the fp32 masters in their param groups. + if (main := getattr(param, "main_param", None)) is not None: + names_by_param[id(main)] = name + return [ + [names_by_param.get(id(param)) for group in child.param_groups for param in group["params"]] + for child in _slot_children(optimizer, slot) + ] + + +def _save_token(adapter, reason: str) -> str: + return hashlib.sha256(f"{adapter.registration_id}:{adapter.step}:{reason}".encode()).hexdigest()[:16] + + +def sidecar_dir(adapter) -> Path | None: + """Default state location (retirement final state and resume).""" + save = adapter.config.save + return Path(save) / "slot_state" if save is not None else None + + +def named_state_dir(adapter, tag: str) -> Path | None: + save = adapter.config.save + return Path(save) / "states" / tag if save is not None else None + + +def _shard_path(base: Path, rank: int) -> Path: + return base / f"shard_rank{rank:05d}.pt" + + +def save_slot_state( + args, + model, + optimizer, + adapter, + *, + reason: str = "state", + base: Path | None = None, + ttl_seconds: int | None = None, +) -> Path | None: + base = base if base is not None else sidecar_dir(adapter) + if base is None: + logger.warning(f"[tinker] ({adapter.name}) no save dir; slot state NOT persisted ({reason})") + return None + base.mkdir(parents=True, exist_ok=True) + + slot = adapter.slot + weights = {name: param.detach().cpu() for name, param in named_adapter_slot_parameters(model, slot)} + optimizer_state = [child.state_dict() for child in _slot_children(optimizer, slot)] + + rank = dist.get_rank() if dist.is_initialized() else 0 + save_id = _save_token(adapter, reason) + payload = { + "format": FORMAT, + "save_id": save_id, + "name": adapter.name, + "registration_id": adapter.registration_id, + "rank_lora": adapter.config.rank, + "alpha": adapter.config.alpha, + "weights": weights, + "optimizer_state": optimizer_state, + "optimizer_param_names": _slot_child_param_names(model, optimizer, slot), + "clocks": {"optimizer_step": adapter.step, "serving_version": adapter.version}, + "topology": { + "rank": rank, + "world_size": dist.get_world_size() if dist.is_initialized() else 1, + }, + "reason": reason, + } + shard = _shard_path(base, rank) + tmp = shard.with_suffix(".tmp") + torch.save(payload, tmp) + os.replace(tmp, shard) # atomic per shard: a crash never leaves a torn file + + if dist.is_initialized(): + dist.barrier() + manifest = base / "manifest.pt" + if rank == 0: + tmp_manifest = manifest.with_suffix(".tmp") + torch.save( + { + "format": FORMAT, + "save_id": save_id, + "name": adapter.name, + "rank_lora": adapter.config.rank, + "alpha": adapter.config.alpha, + "optimizer_step": adapter.step, + "world_size": payload["topology"]["world_size"], + "ttl_seconds": ttl_seconds, + }, + tmp_manifest, + ) + os.replace(tmp_manifest, manifest) + if dist.is_initialized(): + dist.barrier() + logger.info(f"[tinker] ({adapter.name}) slot state saved at step {adapter.step} ({reason}) -> {base}") + return manifest if rank == 0 else shard + + +def find_slot_state(adapter, base: Path | None = None) -> Path | None: + base = base if base is not None else sidecar_dir(adapter) + if base is None or not (base / "manifest.pt").exists(): + return None + manifest = torch.load(base / "manifest.pt", map_location="cpu", weights_only=True) + if manifest.get("format") != FORMAT: + return None + world = dist.get_world_size() if dist.is_initialized() else 1 + if manifest.get("world_size") != world: + logger.warning(f"[tinker] ({adapter.name}) state world_size {manifest.get('world_size')} != {world}; ignoring") + return None + if manifest.get("rank_lora") != adapter.config.rank or manifest.get("alpha") != adapter.config.alpha: + logger.warning( + f"[tinker] ({adapter.name}) state shape rank/alpha " + f"{manifest.get('rank_lora')}/{manifest.get('alpha')} != " + f"{adapter.config.rank}/{adapter.config.alpha}; ignoring" + ) + return None + return base + + +def load_slot_state(args, model, optimizer, adapter, *, base: Path | None = None) -> int | None: + from megatron.bridge.peft.multi_lora_layers import init_adapter_slot, load_adapter + + base = find_slot_state(adapter, base) + if base is None: + return None + rank = dist.get_rank() if dist.is_initialized() else 0 + shard = _shard_path(base, rank) + payload = torch.load(shard, map_location="cpu", weights_only=True) + manifest = torch.load(base / "manifest.pt", map_location="cpu", weights_only=True) + + slot = adapter.slot + children = _slot_children(optimizer, slot) + saved_states = payload.get("optimizer_state") or [] + problem = None + if payload.get("format") != FORMAT: + problem = f"[tinker] ({adapter.name}) state shard format mismatch at {shard}" + elif payload.get("rank_lora") != adapter.config.rank or payload.get("alpha") != adapter.config.alpha: + problem = f"[tinker] ({adapter.name}) state shard shape mismatch at {shard}" + elif payload.get("save_id") != manifest.get("save_id"): + problem = ( + f"[tinker] ({adapter.name}) state at {base} is torn: shard and manifest come from " + "different saves (interrupted write); refusing to restore a mixed generation" + ) + elif len(saved_states) != len(children): + problem = ( + f"[tinker] ({adapter.name}) state has {len(saved_states)} optimizer children " + f"but slot {slot} has {len(children)}; refusing partial restore" + ) + elif payload.get("optimizer_param_names") != _slot_child_param_names(model, optimizer, slot): + problem = ( + f"[tinker] ({adapter.name}) state at {base} was sharded with a different per-rank " + f"parameter ownership than slot {slot} (mismatch on rank {rank}); cross-slot restore " + "requires an identical ownership signature" + ) + if dist.is_initialized(): + problems = [None] * dist.get_world_size(get_gloo_group()) + dist.all_gather_object(problems, problem, group=get_gloo_group()) + problem = next((p for p in problems if p is not None), None) + if problem is not None: + raise ValueError(problem) + + loaded = load_adapter(model, slot, payload["weights"]) + assert loaded > 0, f"[tinker] ({adapter.name}) state restored 0 weight tensors" + init_adapter_slot(model, slot, rank=payload["rank_lora"], alpha=payload["alpha"]) + + for child, state in zip(children, saved_states, strict=True): + child.load_state_dict(state) + for group in child.param_groups: + group["miles_multi_lora_slot"] = slot # the save carries the SOURCE slot's tag + + restored_step = int(payload["clocks"]["optimizer_step"]) + logger.info(f"[tinker] ({adapter.name}) slot state restored at step {restored_step} from {base}") + return restored_step diff --git a/miles/backends/megatron_utils/api_backends/multi_lora/executor.py b/miles/backends/megatron_utils/api_backends/multi_lora/executor.py new file mode 100644 index 00000000000..ad74be685d0 --- /dev/null +++ b/miles/backends/megatron_utils/api_backends/multi_lora/executor.py @@ -0,0 +1,110 @@ +"""Execute Multi-LoRA optimizer operations in slot-sorted collective order. +Lease bindings are validated against local residency before any mutation.""" + +import logging +from dataclasses import dataclass +from typing import Any + +from miles.backends.megatron_utils.api_backends.multi_lora.optimizer import step_adapter_slots, zero_adapter_slot_grads +from miles.backends.training_utils.operation_execution import StepRequest +from miles.ray.multi_lora.residency import ResidentBinding +from miles.utils.operation_contract import BatchExecutionLease + +logger = logging.getLogger(__name__) + + +@dataclass +class MultiLoraParameterExecutor: + model: Any + optimizer: Any + loaded_adapters: dict + + def discard_many(self, lease: BatchExecutionLease[ResidentBinding], operation_ids: list[str]) -> dict[str, dict]: + outcomes: dict[str, dict] = {} + targets: list[tuple[int, str]] = [] + for operation_id in operation_ids: + slot, refusal = self._resolve_slot(lease, operation_id) + if refusal is not None: + outcomes[operation_id] = refusal + continue + targets.append((slot, operation_id)) + for slot, operation_id in sorted(targets): + zero_adapter_slot_grads(self.model, slot) + outcomes[operation_id] = dict(ok=True, gradient_window_consumed=True) + return outcomes + + def step_many(self, lease: BatchExecutionLease[ResidentBinding], requests: list[StepRequest]) -> dict[str, dict]: + outcomes: dict[str, dict] = {} + adam_by_slot: dict[int, dict] = {} + operation_by_slot: dict[int, str] = {} + duplicate_slots: set[int] = set() + for request in requests: + slot, refusal = self._resolve_slot(lease, request.operation_id) + if refusal is not None: + outcomes[request.operation_id] = refusal + continue + if slot in operation_by_slot: + duplicate_slots.add(slot) + continue + adam_by_slot[slot] = request.adam_params + operation_by_slot[slot] = request.operation_id + if duplicate_slots: + for slot in duplicate_slots: + adam_by_slot.pop(slot, None) + operation_by_slot.pop(slot, None) + for request in requests: + binding = lease.binding_of(request.operation_id) + if binding is not None and binding.training_slot in duplicate_slots: + outcomes[request.operation_id] = dict( + ok=False, + error=( + f"operation '{request.operation_id}' shares physical slot " + f"{binding.training_slot} with another operation in this batch; " + "refusing every operation on that slot" + ), + category="server", + ) + if adam_by_slot: + grad_norms, vetoed, norm_blind = step_adapter_slots(self.optimizer, self.model, adam_by_slot) + for slot, operation_id in operation_by_slot.items(): + if slot in vetoed: + outcomes[operation_id] = dict( + ok=False, + error="non-finite gradients; step vetoed and gradients cleared", + category="server", + gradient_window_consumed=True, + ) + elif slot in norm_blind: + outcomes[operation_id] = dict( + ok=False, + error="grads exist but no grad-norm sources (param-flagging bug); step refused, grads cleared", + category="server", + gradient_window_consumed=True, + ) + else: + outcomes[operation_id] = dict( + ok=True, + gradient_window_consumed=True, + result=dict( + grad_norm=grad_norms.get(slot), + learning_rate=adam_by_slot[slot].get("learning_rate", 1e-4), + ), + ) + return outcomes + + def _resolve_slot(self, lease, operation_id: str) -> tuple[int | None, dict | None]: + """Lease -> local residency validation; (slot, None) or (None, outcome).""" + binding = lease.binding_of(operation_id) + if binding is None: + return None, dict( + ok=False, error=f"operation '{operation_id}' has no binding in the batch lease", category="server" + ) + name, registration_id = binding.registration_key + run = self.loaded_adapters.get(name) + if run is None or run.registration_id != registration_id or run.slot != binding.training_slot: + return None, dict( + ok=False, + error=f"adapter '{name}' is not resident in slot {binding.training_slot}", + category="server", + ) + return binding.training_slot, None diff --git a/miles/backends/megatron_utils/api_backends/multi_lora/model.py b/miles/backends/megatron_utils/api_backends/multi_lora/model.py new file mode 100644 index 00000000000..4d1ddbb0776 --- /dev/null +++ b/miles/backends/megatron_utils/api_backends/multi_lora/model.py @@ -0,0 +1,55 @@ +from argparse import Namespace + +import torch + + +def create_multi_lora_instance(args: Namespace): + """Create a MultiLoRA instance from training args.""" + from megatron.bridge.peft.multi_lora import MultiLoRA + + from miles.backends.megatron_utils.lora_utils import convert_target_modules_to_megatron + + lora_type_name = getattr(args, "lora_type", "lora").lower() + if lora_type_name == "canonical_lora": + from megatron.bridge.peft.canonical_lora import CanonicalLoRA + + lora_cls = CanonicalLoRA + else: + from megatron.bridge.peft.lora import LoRA + + lora_cls = LoRA + + # exclude_modules was already folded into target_modules during arg validation. + return MultiLoRA( + target_modules=convert_target_modules_to_megatron(args.target_modules, lora_type=lora_cls), + n_adapters=args.multi_lora_n_adapters, + dim=args.lora_rank, + alpha=args.lora_alpha, + dropout=getattr(args, "lora_dropout", 0.0), + lora_A_init_method=getattr(args, "lora_A_init_method", "xavier"), + lora_B_init_method=getattr(args, "lora_B_init_method", "zero"), + ) + + +def slice_lora_to_rank(hf_name: str, tensor: torch.Tensor, adapter_rank: int) -> torch.Tensor: + if "lora_A" in hf_name: + rank_dim = tensor.ndim - 2 + if adapter_rank < tensor.shape[rank_dim]: + remainder = tensor.narrow(rank_dim, adapter_rank, tensor.shape[rank_dim] - adapter_rank) + assert remainder.abs().max() == 0, ( + f"lora_A padded dims are non-zero: {hf_name}, " + f"max={remainder.abs().max().item():.6e}, shape={tensor.shape}, rank={adapter_rank}" + ) + return tensor.narrow(rank_dim, 0, adapter_rank) + return tensor + if "lora_B" in hf_name: + rank_dim = tensor.ndim - 1 + if adapter_rank < tensor.shape[rank_dim]: + remainder = tensor.narrow(rank_dim, adapter_rank, tensor.shape[rank_dim] - adapter_rank) + assert remainder.abs().max() == 0, ( + f"lora_B padded dims are non-zero: {hf_name}, " + f"max={remainder.abs().max().item():.6e}, shape={tensor.shape}, rank={adapter_rank}" + ) + return tensor.narrow(rank_dim, 0, adapter_rank) + return tensor + return tensor diff --git a/miles/backends/megatron_utils/api_backends/multi_lora/optimizer.py b/miles/backends/megatron_utils/api_backends/multi_lora/optimizer.py new file mode 100644 index 00000000000..0dbb815c6fe --- /dev/null +++ b/miles/backends/megatron_utils/api_backends/multi_lora/optimizer.py @@ -0,0 +1,218 @@ +import logging +import math +from argparse import Namespace +from collections.abc import Sequence +from contextlib import contextmanager + +import torch +import torch.distributed as dist + +from miles.backends.megatron_utils.api_backends.multi_lora.checkpoint import ( + _slot_children, + named_adapter_slot_parameters, +) +from miles.backends.training_utils.operation_execution import resolve_adam_params + +logger = logging.getLogger(__name__) + + +def adapter_slot_parameters(model, slot: int) -> list[torch.nn.Parameter]: + """All parameters belonging to one adapter slot, across model chunks.""" + return [param for _, param in named_adapter_slot_parameters(model, slot)] + + +def _adam_init_state_fn(opt, config=None): + for group in opt.param_groups: + for p in group["params"]: + if len(opt.state[p]) == 0: + opt.state[p]["exp_avg"] = torch.zeros_like(p.data) + opt.state[p]["exp_avg_sq"] = torch.zeros_like(p.data) + + +@contextmanager +def _only_slot_trainable(model_chunks, slot_params: list[torch.nn.Parameter]): + slot_ids = {id(p) for p in slot_params} + frozen = [] + for model_chunk in model_chunks: + for param in model_chunk.parameters(): + if param.requires_grad and id(param) not in slot_ids: + param.requires_grad = False + frozen.append(param) + try: + yield + finally: + for param in frozen: + param.requires_grad = True + + +def build_multi_lora_operation_optimizer(args: Namespace, config, model_chunks: Sequence): + assert ( + not config.use_distributed_optimizer + ), "per-slot optimizers require use_distributed_optimizer=False (LayerWise shards; grad retention all-reduces)" + assert not config.fp16, "tinker per-slot optimizers require bf16 (no dynamic loss scaler)" + assert ( + config.optimizer or "" + ).lower() == "adam", ( + f"tinker per-slot optimizers only implement Adam semantics; got optimizer={config.optimizer!r}" + ) + + from megatron.core.optimizer import get_megatron_optimizer + from megatron.core.optimizer.layer_wise_optimizer import LayerWiseDistributedOptimizer + from megatron.core.process_groups_config import ProcessGroupCollection + + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + + # Defer bf16 master-weight creation into LayerWise (post-sharding) so fp32 masters exist only for owned params. + reset_bf16 = config.bf16 + config.bf16 = False + + base_optimizers: list = [] + init_fns: list = [] + slot_child_indices: dict[int, list[int]] = {} + try: + for slot in range(args.multi_lora_n_adapters): + slot_params = adapter_slot_parameters(model_chunks, slot) + assert slot_params, f"adapter slot {slot} has no parameters; is this a multi-LoRA model?" + with _only_slot_trainable(model_chunks, slot_params): + chained = get_megatron_optimizer( + config, + list(model_chunks), + use_gloo_process_groups=args.use_gloo_process_groups, + ) + children = [ + child + for child in chained.chained_optimizers + if getattr(child, "optimizer", None) is not None and child.get_parameters() + ] + assert children, f"adapter slot {slot} produced no optimizer children" + slot_child_indices[slot] = list(range(len(base_optimizers), len(base_optimizers) + len(children))) + for child in children: + for group in child.param_groups: + group["miles_multi_lora_slot"] = slot + # LayerWise wraps raw torch optimizers itself; the pinned MCore + # rejects pre-wrapped children (slot tags survive via the proxy). + base_optimizers.append(child.optimizer) + init_fns.append(_adam_init_state_fn) + finally: + config.bf16 = reset_bf16 + + optimizer = LayerWiseDistributedOptimizer(base_optimizers, config, pg_collection, init_state_fn_list=init_fns) + + # Dense and expert params use independent ownership groups; norm/clip reductions must span the world. + for child in optimizer.chained_optimizers: + child.grad_stats_parallel_group = None + + optimizer.miles_slot_child_indices = slot_child_indices + logger.info( + f"[tinker] built LayerWise optimizer: {args.multi_lora_n_adapters} slots, " + f"{len(optimizer.chained_optimizers)} chained children" + ) + return optimizer + + +def reload_adapter_slot_model_params(optimizer, slot: int) -> None: + for child in _slot_children(optimizer, slot): + child.reload_model_params() + + +def zero_adapter_slot_grads(model, slot: int) -> None: + for param in adapter_slot_parameters(model, slot): + if (main_grad := getattr(param, "main_grad", None)) is not None: + main_grad.zero_() + param.grad = None + if (main_param := getattr(param, "main_param", None)) is not None: + main_param.grad = None + + +def _found_inf_anywhere(found_inf: bool) -> bool: + """The veto must agree on every rank, or the collective step order diverges.""" + if not dist.is_initialized(): + return found_inf + flag = torch.tensor([1.0 if found_inf else 0.0], device=torch.cuda.current_device()) + dist.all_reduce(flag, op=dist.ReduceOp.MAX) + return flag.item() > 0 + + +def _norm_source_flags_anywhere(has_norm_source: bool, has_grads: bool) -> tuple[bool, bool]: + if not dist.is_initialized(): + return has_norm_source, has_grads + flags = torch.tensor( + [1.0 if has_norm_source else 0.0, 1.0 if has_grads else 0.0], device=torch.cuda.current_device() + ) + dist.all_reduce(flags, op=dist.ReduceOp.MAX) + return bool(flags[0].item() > 0), bool(flags[1].item() > 0) + + +def apply_adam_params_to_slot(optimizer, slot: int, adam_params: dict | None) -> dict: + resolved = resolve_adam_params(adam_params) + for child in _slot_children(optimizer, slot): + for group in child.param_groups: + group["lr"] = resolved["learning_rate"] + group["betas"] = (resolved["beta1"], resolved["beta2"]) + group["eps"] = resolved["eps"] + group["weight_decay"] = resolved["weight_decay"] + return resolved + + +def step_adapter_slots( + optimizer, + model, + adam_params_by_slot: dict[int, dict | None], +) -> tuple[dict[int, float], set[int], set[int]]: + from megatron.core.optimizer.clip_grads import clip_grad_by_total_norm_fp32, get_grad_norm_fp32 + + grad_norms: dict[int, float] = {} + vetoed: set[int] = set() + norm_blind: set[int] = set() + + for slot in sorted(adam_params_by_slot): + children = _slot_children(optimizer, slot) + adam = apply_adam_params_to_slot(optimizer, slot, adam_params_by_slot[slot]) + + # Copy accumulated main_grads into the owned masters' grads, untouched. + found_inf = False + for child in children: + found_inf = bool(child.prepare_grads()) or found_inf + + # Reduce per-slot norms across the world to combine dense and expert ownership groups. + grads_for_norm = [] + slot_params = [] + for child in children: + grads_for_norm += child.get_main_grads_for_grad_norm() + slot_params += child.get_parameters() + slot_norm = get_grad_norm_fp32(grads_for_norm, grad_stats_parallel_group=None) + + # A non-finite step would otherwise be applied AND live-published to + # every engine; the veto must be unanimous across ranks. + if _found_inf_anywhere(found_inf) or not math.isfinite(float(slot_norm)): + logger.error( + f"[tinker] slot {slot}: non-finite gradients " + f"(found_inf={found_inf}, grad_norm={float(slot_norm)}); step vetoed, grads cleared" + ) + vetoed.add(slot) + zero_adapter_slot_grads(model, slot) + continue + + has_norm_source, has_grads = _norm_source_flags_anywhere( + bool(grads_for_norm), + any(param.grad is not None and bool((param.grad != 0).any().item()) for param in slot_params), + ) + if has_grads and not has_norm_source: + logger.error(f"[tinker] slot {slot}: no grad-norm source despite grads (mis-flagged params); step refused") + norm_blind.add(slot) + zero_adapter_slot_grads(model, slot) + continue + + if adam["grad_clip_norm"] > 0.0 and slot_params: + clip_grad_by_total_norm_fp32(slot_params, adam["grad_clip_norm"], slot_norm, False) + grad_norms[slot] = float(slot_norm) + + for child in children: + child.step_with_ready_grads() + + zero_adapter_slot_grads(model, slot) + + if grad_norms: + optimizer.allgather_params() + + return grad_norms, vetoed, norm_blind diff --git a/miles/backends/megatron_utils/api_backends/multi_lora/trainer.py b/miles/backends/megatron_utils/api_backends/multi_lora/trainer.py new file mode 100644 index 00000000000..c8d8869bcc1 --- /dev/null +++ b/miles/backends/megatron_utils/api_backends/multi_lora/trainer.py @@ -0,0 +1,336 @@ +import logging +import re +from dataclasses import replace as dataclass_replace +from pathlib import Path + +import ray +import torch +import torch.distributed as dist + +from miles.backends.megatron_utils.api_backends.multi_lora.checkpoint import ( + load_slot_state, + named_state_dir, + save_slot_state, +) +from miles.backends.megatron_utils.api_backends.multi_lora.executor import MultiLoraParameterExecutor +from miles.backends.megatron_utils.api_backends.multi_lora.optimizer import ( + reload_adapter_slot_model_params, + zero_adapter_slot_grads, +) +from miles.backends.training_utils.operation_execution import run_optim_controls +from miles.ray.multi_lora.controller import get_multi_lora_controller +from miles.ray.multi_lora.residency import lease_from_metadata +from miles.utils.distributed_utils import get_gloo_group + +logger = logging.getLogger(__name__) + +_STATE_TAG = re.compile(r"[A-Za-z0-9._-]+") + + +def zero_optimizer_state_for_adapter(optimizer, model, slot: int) -> None: + from megatron.bridge.peft.multi_lora_layers import MultiLoRALinear, _iter_multi_lora_modules + + target_main_params = set() + for module in _iter_multi_lora_modules(model): + if not isinstance(module, MultiLoRALinear): + continue + for param in module.adapters[slot].parameters(): + main = getattr(param, "main_param", None) + target_main_params.add(id(main if main is not None else param)) + + chained = getattr(optimizer, "chained_optimizers", [optimizer]) + for chained_optimizer in chained: + inner = getattr(chained_optimizer, "optimizer", chained_optimizer) + if inner is None: + continue + # TE/apex FusedAdam tracks the Adam step per param GROUP, not per param. + for group in inner.param_groups: + if group.get("miles_multi_lora_slot") == slot and "step" in group: + if isinstance(group["step"], torch.Tensor): + group["step"].zero_() + else: + group["step"] = 0 + for param, state in inner.state.items(): + if id(param) not in target_main_params: + continue + if "exp_avg" in state: + state["exp_avg"].zero_() + if "exp_avg_sq" in state: + state["exp_avg_sq"].zero_() + if "step" in state: + if isinstance(state["step"], torch.Tensor): + state["step"].zero_() + else: + state["step"] = 0 + + +def _install_adapter(adapter, args, model, optimizer) -> int | None: + from megatron.bridge.peft.multi_lora_layers import init_adapter_slot + + log_prefix = f"[tinker] ({adapter.name})" + try: + restored_step = load_slot_state(args, model, optimizer, adapter) + except ValueError as e: + logger.warning(f"{log_prefix} sidecar state not restorable into slot {adapter.slot} ({e}); fresh init") + restored_step = None + if restored_step is not None: + logger.info(f"{log_prefix} resumed slot {adapter.slot} from sidecar at step {restored_step}") + return restored_step + init_adapter_slot(model, adapter.slot, rank=adapter.config.rank, alpha=adapter.config.alpha) + logger.info(f"{log_prefix} fresh init at slot {adapter.slot}") + return None + + +def load_adapters(args, model, optimizer, adapters) -> int: + from miles.backends.megatron_utils.initialize import is_first_replica_megatron_main_rank + + if dist.is_initialized(): + dist.barrier(group=get_gloo_group()) + if not adapters: + return 0 + installed_steps: dict[str, int | None] = {} + for adapter in adapters: + installed_steps[adapter.name] = _install_adapter(adapter, args, model, optimizer) + if dist.is_initialized(): + dist.barrier(group=get_gloo_group()) + for adapter in adapters: + if installed_steps[adapter.name] is None: + reload_adapter_slot_model_params(optimizer, adapter.slot) + if is_first_replica_megatron_main_rank(): + controller = get_multi_lora_controller() + for name, step in installed_steps.items(): + if step: + ray.get(controller.set_adapter_step.remote(name, step)) + ray.get(controller.mark_ready.remote(sorted(installed_steps))) + return len(adapters) + + +def cleanup_adapters(args, model, optimizer, adapters) -> int: + from megatron.bridge.peft.multi_lora_layers import clear_adapter_slot + + from miles.backends.megatron_utils.initialize import is_first_replica_megatron_main_rank + + if dist.is_initialized(): + dist.barrier(group=get_gloo_group()) + if not adapters: + return 0 + for adapter in adapters: + save_slot_state(args, model, optimizer, adapter, reason="final") + clear_adapter_slot(model, adapter.slot) + zero_optimizer_state_for_adapter(optimizer, model, adapter.slot) + zero_adapter_slot_grads(model, adapter.slot) + reload_adapter_slot_model_params(optimizer, adapter.slot) + logger.info(f"[tinker] ({adapter.name}) slot {adapter.slot} retired and scrubbed") + if dist.is_initialized(): + dist.barrier(group=get_gloo_group()) + if is_first_replica_megatron_main_rank(): + for adapter in adapters: + ray.get(get_multi_lora_controller().free_slot.remote(adapter.name)) + return len(adapters) + + +def reconcile_adapters(args, model, optimizer, loaded_adapters: dict, pending_push: set, weights_backuper) -> None: + from miles.backends.megatron_utils.initialize import is_first_replica_megatron_main_rank + + broadcast_buffer = [None] + if is_first_replica_megatron_main_rank(): + controller = get_multi_lora_controller() + ray.get(controller.retire_adapters.remote()) + # Queued registrations take freed slots so this reconcile loads them. + ray.get(controller.bootstrap_pending.remote()) + snapshot = ray.get(controller.snapshot.remote()) + cleanup_steps = {name: ray.get(controller.adapter_step.remote(name)) for name in snapshot["cleanup"]} + broadcast_buffer[0] = (snapshot, cleanup_steps) + if dist.is_initialized(): + dist.broadcast_object_list(broadcast_buffer, src=0, group=get_gloo_group()) + snapshot, cleanup_steps = broadcast_buffer[0] + should_be_loaded = { + name: run + for name, run in {**snapshot["pending"], **snapshot["ready"], **snapshot["retiring"]}.items() + # Queued-but-unbound registrations have no residency to reconcile yet. + if run.slot is not None + } + cleanup_names = set(snapshot["cleanup"]) + + loaded_names = set(loaded_adapters) + # Sorted so per-adapter collectives run in the same order on every rank; + # set iteration order is process-specific. + adapters_to_load = sorted( + (adapter for name, adapter in should_be_loaded.items() if name not in loaded_names), + key=lambda adapter: adapter.name, + ) + adapters_to_clean_up = sorted( + (loaded_adapters[n] for n in loaded_names if n in cleanup_names or n not in should_be_loaded), + key=lambda adapter: adapter.name, + ) + if adapters_to_load: + load_adapters(args, model, optimizer, adapters_to_load) + for adapter in adapters_to_load: + loaded_adapters[adapter.name] = adapter + weights_backuper.backup("actor") + if adapters_to_clean_up: + # The registry's step clock is authoritative for the final state; the + # loaded views were captured at load time and lag it. + refreshed = [ + dataclass_replace(adapter, step=cleanup_steps.get(adapter.name, adapter.step)) + for adapter in adapters_to_clean_up + ] + cleanup_adapters(args, model, optimizer, refreshed) + for adapter in adapters_to_clean_up: + loaded_adapters.pop(adapter.name, None) + pending_push.discard(adapter.name) + weights_backuper.backup("actor") + + # Deregistered before ever being loaded: nothing to save or clear. + if is_first_replica_megatron_main_rank(): + for name in cleanup_names - loaded_names: + ray.get(get_multi_lora_controller().free_slot.remote(name)) + + +def execute_controls( + args, model, optimizer, loaded_adapters, pending_push, weights_backuper, operations, lease_metadata +) -> dict: + lease = lease_from_metadata(lease_metadata) + executor = MultiLoraParameterExecutor(model=model, optimizer=optimizer, loaded_adapters=loaded_adapters) + results = run_optim_controls(operations, lease, executor) + + def state_order(op: dict): + binding = lease.binding_of(op["operation_id"]) + return (op["kind"], binding.training_slot if binding is not None else -1) + + for op in sorted( + (op for op in operations if op["kind"] in ("save_weights_for_sampler", "save_state", "load_state")), + key=state_order, + ): + results[op["operation_id"]] = _execute_state_op( + op, lease, args, model, optimizer, loaded_adapters, pending_push + ) + if results[op["operation_id"]].get("ok") and op["kind"] == "load_state": + weights_backuper.backup("actor") + + for op in operations: + if op["operation_id"] not in results: + results[op["operation_id"]] = dict( + ok=False, error=f"operation kind '{op['kind']}' has no executor", category="server" + ) + return results + + +def _execute_state_op(op: dict, lease, args, model, optimizer, loaded_adapters, pending_push) -> dict: + name, kind = op["name"], op["kind"] + binding = lease.binding_of(op["operation_id"]) + if binding is None: + return dict( + ok=False, error=f"operation '{op['operation_id']}' has no binding in the batch lease", category="server" + ) + bound_name, bound_registration_id = binding.registration_key + if bound_name != name: + return dict( + ok=False, + error=f"operation '{op['operation_id']}' names adapter '{name}' but its lease binding " + f"names '{bound_name}'", + category="server", + ) + run = loaded_adapters.get(name) + if run is None or run.registration_id != bound_registration_id or run.slot != binding.training_slot: + return dict( + ok=False, error=f"adapter '{name}' is not resident in slot {binding.training_slot}", category="server" + ) + # The registry's clocks are authoritative; the loaded view can lag. + run = dataclass_replace(run, step=op.get("step", run.step), version=op.get("serving_version", run.version)) + + if kind == "save_weights_for_sampler": + pending_push.add(name) + return dict(ok=True, deferred="publish") + + payload = op.get("payload") or {} + if kind == "save_state": + tag = str(payload.get("tag") or f"step_{run.step}") + if not _STATE_TAG.fullmatch(tag) or tag in (".", ".."): + return dict(ok=False, error=f"invalid state tag '{tag}'", category="user") + base = named_state_dir(run, tag) + if base is None: + return dict(ok=False, error=f"adapter '{name}' has no save dir", category="user") + if (base / "manifest.pt").exists(): + return dict(ok=False, error=f"state '{tag}' already exists; states are immutable", category="user") + save_slot_state( + args, model, optimizer, run, reason=f"state:{tag}", base=base, ttl_seconds=payload.get("ttl_seconds") + ) + return dict(ok=True, result=dict(path=str(base), step=run.step)) + + assert kind == "load_state" + path = payload.get("path") + try: + restored_step = load_slot_state(args, model, optimizer, run, base=Path(path)) + except ValueError as e: + return dict(ok=False, error=str(e), category="user") + if restored_step is None: + return dict(ok=False, error=f"no loadable state at '{path}' for adapter '{name}'", category="user") + pending_push.add(name) + return dict(ok=True, deferred="publish", result=dict(step=restored_step, path=str(path))) + + +def validate_batch_lease(rollout_data, loaded_adapters: dict) -> None: + lease = rollout_data.get("batch_execution_lease") + if lease is None: + raise RuntimeError("tinker batch carries no execution lease") + for op_id, (name, registration_id, slot) in lease["bindings_by_operation"]: + run = loaded_adapters.get(name) + if run is None or run.registration_id != registration_id or run.slot != slot: + raise RuntimeError( + f"operation '{op_id}': lease binding ('{name}', {registration_id[:8]}, slot {slot}) " + "does not match this rank's loaded adapters; refusing to mutate" + ) + + +def commit_batch(rollout_data, pending_push: set) -> None: + from miles.backends.megatron_utils.initialize import is_first_replica_megatron_main_rank + + logprobs_by_op = _gather_logprobs(rollout_data) + if is_first_replica_megatron_main_rank(): + try: + registration_by_lane = rollout_data.get("registration_by_lane", {}) + # Forward batches accumulate nothing: no dirty streams. + accumulated = ( + [] + if rollout_data.get("tinker_forward_only") + else sorted({tuple(key) for key in registration_by_lane.values()}) + ) + operation_ids = [op_id for op_id in rollout_data.get("operation_by_lane", {}).values() if op_id] + ray.get(get_multi_lora_controller().commit_tinker_batch.remote(accumulated, operation_ids, logprobs_by_op)) + finally: + if (lease := rollout_data.get("batch_execution_lease")) is not None: + ray.get(get_multi_lora_controller().release_batch_lease.remote(lease)) + + +def _gather_logprobs(rollout_data) -> dict[str, list[list[float]]]: + collector = rollout_data.get("tinker_logprob_collector") or {} + if dist.is_initialized(): + shards = [None] * dist.get_world_size(get_gloo_group()) + dist.all_gather_object(shards, collector, group=get_gloo_group()) + merged: dict = {} + for shard in shards: + merged.update(shard or {}) + else: + merged = dict(collector) + + op_by_lane = rollout_data.get("operation_by_lane", {}) + logprobs_by_op: dict[str, list[list[float]]] = {} + for op_lane, op_id in op_by_lane.items(): + if op_id is None: + continue + # row -1 is DP padding: never part of the operation's result plane. + rows = sorted((row, lp) for (lane, row), lp in merged.items() if lane == op_lane and row >= 0) + logprobs_by_op[op_id] = [lp for _, lp in rows] + return logprobs_by_op + + +def select_adapters_to_push(loaded_adapters: dict, pending_push: set, has_new_engines: bool) -> tuple[dict, list]: + pending = pending_push & set(loaded_adapters) + push_names = set(loaded_adapters) if has_new_engines else pending + return {name: loaded_adapters[name] for name in sorted(push_names)}, sorted(pending) + + +def commit_weight_push(version_update_names: list, is_main_rank: bool) -> None: + if version_update_names and is_main_rank: + ray.get(get_multi_lora_controller().record_weight_update.remote(version_update_names)) diff --git a/miles/backends/megatron_utils/bridge_lora_helpers.py b/miles/backends/megatron_utils/bridge_lora_helpers.py index b0746cc7627..8b6c02df18d 100644 --- a/miles/backends/megatron_utils/bridge_lora_helpers.py +++ b/miles/backends/megatron_utils/bridge_lora_helpers.py @@ -168,7 +168,7 @@ def _setup_lora_model_via_bridge(args: Namespace) -> list: if is_multi_lora_enabled(args): _validate_multi_lora_moe_support(args, provider) - from miles.backends.megatron_utils.multi_lora_utils import create_multi_lora_instance + from miles.backends.megatron_utils.api_backends.multi_lora.model import create_multi_lora_instance lora = create_multi_lora_instance(args) else: diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index bd526e76ba0..dee907d48c9 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -32,8 +32,9 @@ from miles.utils.audit_utils.witness.module import witness_dump_and_clear_stale from miles.utils.dumper_utils import DumperMegatronUtil, DumperPhase from miles.utils.memory_utils import clear_memory -from miles.utils.multi_lora import is_multi_lora_enabled +from miles.utils.multi_lora import is_multi_lora_enabled, uses_multi_lora_operation_executor from miles.utils.test_utils.ft_test_actions import FTTestActionActorExecutor +from miles.utils.tinker import uses_explicit_training_operations from miles.utils.tracking_utils.structured_log import log_structured from ...utils.misc import filter_keys @@ -190,10 +191,12 @@ def setup_model_and_optimizer( use_gloo_process_groups=args.use_gloo_process_groups, layer_wise_distributed_optimizer="dist" in config.optimizer.lower(), ) - elif is_multi_lora_enabled(args): - from miles.backends.megatron_utils.multi_lora_optimizer import build_multi_lora_optimizer + elif uses_multi_lora_operation_executor(args): + from miles.backends.megatron_utils.api_backends.multi_lora.optimizer import ( + build_multi_lora_operation_optimizer, + ) - optimizer = build_multi_lora_optimizer(args, config, model) + optimizer = build_multi_lora_operation_optimizer(args, config, model) else: optimizer = get_megatron_optimizer( config=config, @@ -415,6 +418,7 @@ def train_one_step( witness_info: WitnessInfo | None, attempt: int, ft_test_action_executor: FTTestActionActorExecutor | None = None, + forward_only: bool = False, ) -> tuple[dict[str, float], float, TrainStepOutcome]: """Execute a single pipeline-parallel training step. @@ -422,8 +426,8 @@ def train_one_step( one scheduler step when gradients are valid. Multi-LoRA: gradients are retained across train calls (per-adapter - gradient accumulation); only the slots in the batch's ``step_slots`` step, - and only their gradients are zeroed. + gradient accumulation); slots step only when the client's optim_step + operation executes, and only their gradients are zeroed. Args: args: Runtime arguments. @@ -443,10 +447,10 @@ def train_one_step( parallel_state = get_parallel_state() dumper_phase_util = DumperMegatronUtil(args, model, DumperPhase.FWD_BWD, rollout_id=rollout_id) disable_optimizer = args.debug_disable_optimizer or optimizer is None - multi_lora = is_multi_lora_enabled(args) + explicit_optim_step = uses_explicit_training_operations(args) - if multi_lora: - from miles.backends.megatron_utils.multi_lora_optimizer import reset_grad_metadata_keep_grads + if explicit_optim_step: + from miles.backends.training_utils.operation_execution import reset_grad_metadata_keep_grads # Retain accumulated per-adapter gradients; reset only the per-iteration # DDP bookkeeping. Slot grads are zeroed selectively at step time. @@ -493,6 +497,8 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p "advantages", "returns", "rollout_log_probs", + "loss_weights", + "sample_indices", "max_seq_lens", "witness_ids", "opd_reverse_kl", @@ -553,7 +559,6 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p num_rollouts=num_rollouts, ) - # Forward pass. forward_backward_func = get_forward_backward_func() losses_reduced = forward_backward_func( forward_step_func=forward_step, @@ -563,7 +568,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p seq_length=args.seq_length, micro_batch_size=args.micro_batch_size, decoder_seq_length=args.decoder_seq_length, - forward_only=False, + forward_only=forward_only, ) outcome = TrainStepOutcome.NORMAL @@ -585,7 +590,11 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p outcome = TrainStepOutcome.DISCARDED_SHOULD_RETRY valid_step = False - if (not disable_optimizer) and (not multi_lora) and (not getattr(args, "check_for_nan_in_loss_and_grad", True)): + if ( + (not disable_optimizer) + and (not explicit_optim_step) + and (not getattr(args, "check_for_nan_in_loss_and_grad", True)) + ): found_inf_flag = optimizer.prepare_grads() if found_inf_flag: valid_step = False @@ -610,12 +619,8 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p dumper_phase_util.finalize(model) if not disable_optimizer and valid_step: - if multi_lora: - from miles.backends.megatron_utils.multi_lora_utils import step_stepped_adapter_slots - - grad_norm = step_stepped_adapter_slots( - args, model, optimizer, data_iterator[0].rollout_data, rollout_id, step_id - ) + if explicit_optim_step: + grad_norm = 0.0 else: # Update parameters. update_successful, grad_norm, num_zeros_in_grad = optimizer.step() @@ -624,9 +629,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p assert update_successful opt_param_scheduler.step(increment=num_rollouts) - # release grad (multi-LoRA retains accumulated grads; stepped slots were - # zeroed selectively inside step_adapter_slots) - if not multi_lora: + if not explicit_optim_step: _zero_grads(model, optimizer, disable_optimizer) log_structured( @@ -679,6 +682,7 @@ def train( witness_info: WitnessInfo | None, attempt: int, ft_test_action_executor: FTTestActionActorExecutor | None = None, + forward_only: bool = False, ) -> TrainStepOutcome: """Run training over a rollout consisting of multiple steps. @@ -693,6 +697,8 @@ def train( data_iterator (Sequence[DataIterator]): Iterable(s) yielding training batches. num_microbatches (Sequence[int]): Microbatches per step in the rollout. num_rollouts (Sequence[int]): Rollout count per step (total across DP). + forward_only (bool): Run the schedule without backward (tinker + ``forward`` operations: logprobs only, gradients untouched). """ parallel_state = get_parallel_state() args = get_args() @@ -786,6 +792,7 @@ def train( witness_info=witness_info, attempt=attempt, ft_test_action_executor=ft_test_action_executor, + forward_only=forward_only, ) if step_id == 0: diff --git a/miles/backends/megatron_utils/multi_lora_optimizer.py b/miles/backends/megatron_utils/multi_lora_optimizer.py deleted file mode 100644 index 8cd96b2a29c..00000000000 --- a/miles/backends/megatron_utils/multi_lora_optimizer.py +++ /dev/null @@ -1,195 +0,0 @@ -"""Per-slot decoupled Adam optimizers for multi-LoRA, chained under Megatron's LayerWiseDistributedOptimizer; -requires plain DDP all-reduce (use_distributed_optimizer OFF) so cross-batch gradient retention stays idempotent.""" - -import logging -from argparse import Namespace -from collections.abc import Sequence -from contextlib import contextmanager - -import torch -from megatron.core.optimizer import get_megatron_optimizer -from megatron.core.optimizer.clip_grads import clip_grad_by_total_norm_fp32, get_grad_norm_fp32 -from megatron.core.optimizer.layer_wise_optimizer import LayerWiseDistributedOptimizer -from megatron.core.optimizer.optimizer import MegatronOptimizer -from megatron.core.optimizer.optimizer_config import OptimizerConfig -from megatron.core.process_groups_config import ProcessGroupCollection - -logger = logging.getLogger(__name__) - - -def adapter_slot_parameters(model, slot: int) -> list[torch.nn.Parameter]: - """All parameters belonging to one adapter slot, across model chunks.""" - from megatron.bridge.peft.multi_lora_layers import MultiLoRALinear - - parameters = [] - seen = set() - model_chunks = model if isinstance(model, (list, tuple)) else [model] - for model_chunk in model_chunks: - for module in model_chunk.modules(): - if not isinstance(module, MultiLoRALinear): - continue - for param in module.adapters[slot].parameters(): - if id(param) not in seen: - parameters.append(param) - seen.add(id(param)) - return parameters - - -def _adam_init_state_fn(opt, config=None): - for group in opt.param_groups: - for p in group["params"]: - if len(opt.state[p]) == 0: - opt.state[p]["exp_avg"] = torch.zeros_like(p.data) - opt.state[p]["exp_avg_sq"] = torch.zeros_like(p.data) - - -@contextmanager -def _only_slot_trainable(model_chunks, slot_params: list[torch.nn.Parameter]): - """Temporarily freeze every trainable param outside ``slot_params`` so the - stock param-group builder sees exactly one slot (the Muon construction - pattern from megatron's ``get_megatron_muon_optimizer``).""" - slot_ids = {id(p) for p in slot_params} - frozen = [] - for model_chunk in model_chunks: - for param in model_chunk.parameters(): - if param.requires_grad and id(param) not in slot_ids: - param.requires_grad = False - frozen.append(param) - try: - yield - finally: - for param in frozen: - param.requires_grad = True - - -def build_multi_lora_optimizer( - args: Namespace, - config: OptimizerConfig, - model_chunks: Sequence, -) -> MegatronOptimizer: - """Build one Float16-wrapped Adam per adapter slot under a LayerWiseDistributedOptimizer (ChainedOptimizer); - each child's param groups are tagged with ``miles_multi_lora_slot`` and narrowed to this rank's shard.""" - assert not config.use_distributed_optimizer, ( - "multi-LoRA per-slot optimizers require use_distributed_optimizer=False: " - "gradient retention relies on all-reduce idempotency, and LayerWise " - "sharding replaces byte-level ZeRO" - ) - assert not config.fp16, "multi-LoRA per-slot optimizers require bf16 (no dynamic loss scaler)" - assert (config.optimizer or "").lower() == "adam", ( - "multi-LoRA per-slot optimizers only implement Adam semantics (state init, " - f"slot retirement cleanup, step clocks); got optimizer={config.optimizer!r}" - ) - - pg_collection = ProcessGroupCollection.use_mpu_process_groups() - - # Defer bf16 master-weight creation into LayerWise (post-sharding) so fp32 masters exist only for owned params. - reset_bf16 = config.bf16 - config.bf16 = False - - base_optimizers: list = [] - init_fns: list = [] - slot_child_indices: dict[int, list[int]] = {} - try: - for slot in range(args.multi_lora_n_adapters): - slot_params = adapter_slot_parameters(model_chunks, slot) - assert slot_params, f"adapter slot {slot} has no parameters; is this a multi-LoRA model?" - with _only_slot_trainable(model_chunks, slot_params): - chained = get_megatron_optimizer( - config, - list(model_chunks), - use_gloo_process_groups=args.use_gloo_process_groups, - ) - children = [ - child - for child in chained.chained_optimizers - if getattr(child, "optimizer", None) is not None and child.get_parameters() - ] - assert children, f"adapter slot {slot} produced no optimizer children" - slot_child_indices[slot] = list(range(len(base_optimizers), len(base_optimizers) + len(children))) - for child in children: - for group in child.param_groups: - group["miles_multi_lora_slot"] = slot - base_optimizers.append(child) - init_fns.append(_adam_init_state_fn) - finally: - config.bf16 = reset_bf16 - - optimizer = LayerWiseDistributedOptimizer(base_optimizers, config, pg_collection, init_state_fn_list=init_fns) - - # Params are scattered whole across DP ranks, so per-child norm/clip reductions must span the world. - for child in optimizer.chained_optimizers: - child.grad_stats_parallel_group = None - - optimizer.miles_slot_child_indices = slot_child_indices - logger.info( - f"Built multi-LoRA LayerWise optimizer: {args.multi_lora_n_adapters} slots, " - f"{len(optimizer.chained_optimizers)} chained children" - ) - return optimizer - - -def _slot_children(optimizer, slot: int): - return [optimizer.chained_optimizers[i] for i in optimizer.miles_slot_child_indices[slot]] - - -def reset_grad_metadata_keep_grads(model_chunks) -> None: - """Reset DDP per-iteration grad bookkeeping WITHOUT zeroing grad buffers, so per-adapter accumulation - survives across train batches (replaces ``DistributedDataParallel.zero_grad_buffer``).""" - for model_chunk in model_chunks: - if getattr(model_chunk.config, "cuda_graph_impl", "none") != "transformer_engine": - for param in model_chunk.params_with_grad: - param.grad_added_to_main_grad = False - for bucket_group in model_chunk.bucket_groups + model_chunk.expert_parallel_bucket_groups: - bucket_group.reset() - - -def zero_adapter_slot_grads(model, slot: int) -> None: - """Zero one slot's gradients everywhere they live: the DDP ``main_grad`` buffer views - and any lingering ``grad``/``main_param.grad`` references.""" - for param in adapter_slot_parameters(model, slot): - if (main_grad := getattr(param, "main_grad", None)) is not None: - main_grad.zero_() - param.grad = None - if (main_param := getattr(param, "main_param", None)) is not None: - main_param.grad = None - - -def step_adapter_slots( - optimizer, - model, - step_batch_sizes: dict[int, int], - clip_grad: float, -) -> dict[int, float]: - """Step exactly the slots in ``step_batch_sizes`` (slot -> batch size), retaining all other slots' gradients; - scales each slot's accumulated grad sum by 1/batch_size and returns the grad norm per stepped slot.""" - grad_norms: dict[int, float] = {} - - for slot, batch_size in step_batch_sizes.items(): - children = _slot_children(optimizer, slot) - # Copy accumulated main_grads into the owned masters' grads, then scale the sum to the adapter-batch mean. - for child in children: - child.prepare_grads() - for main_param in child.get_parameters(): - if main_param.grad is not None: - main_param.grad.mul_(1.0 / batch_size) - - # Per-slot grad norm over the slot's children, reduced across the whole world (whole-param DP scatter). - grads_for_norm = [] - slot_params = [] - for child in children: - grads_for_norm += child.get_main_grads_for_grad_norm() - slot_params += child.get_parameters() - slot_norm = get_grad_norm_fp32(grads_for_norm, grad_stats_parallel_group=None) - if clip_grad > 0.0 and slot_params: - clip_grad_by_total_norm_fp32(slot_params, clip_grad, slot_norm, False) - grad_norms[slot] = float(slot_norm) - - for child in children: - child.step_with_ready_grads() - - zero_adapter_slot_grads(model, slot) - - if step_batch_sizes: - optimizer.allgather_params() - - return grad_norms diff --git a/miles/backends/megatron_utils/multi_lora_scheduler.py b/miles/backends/megatron_utils/multi_lora_scheduler.py deleted file mode 100644 index 4f148eb599e..00000000000 --- a/miles/backends/megatron_utils/multi_lora_scheduler.py +++ /dev/null @@ -1,88 +0,0 @@ -"""Per-adapter LR/WD schedules for multi-LoRA: one ``OptimizerParamScheduler`` per adapter slot, positioned by -the adapter's own trained samples. Adapters without a known ``num_step`` warm up, then hold ``--lr`` constant.""" - -import logging -from argparse import Namespace - -from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler - -logger = logging.getLogger(__name__) - - -class _SlotParamGroups: - """Minimal optimizer facade: the scheduler only reads ``.param_groups``.""" - - def __init__(self, param_groups: list[dict]): - self.param_groups = param_groups - - -def build_slot_scheduler(args: Namespace, optimizer, adapter, resume_step: int) -> OptimizerParamScheduler: - """Build the slot's scheduler and position it at the adapter's committed - samples. Rebuilt on every adapter load, so slot reuse starts fresh.""" - from miles.backends.megatron_utils.multi_lora_optimizer import _slot_children - - groups = [group for child in _slot_children(optimizer, adapter.slot) for group in child.param_groups] - samples_per_step = adapter.config.adapter_global_batch_size - num_step = adapter.config.num_step - - decay_steps = num_step * samples_per_step if num_step is not None else None - if args.lr_warmup_fraction is not None and decay_steps is not None: - lr_warmup_steps = args.lr_warmup_fraction * decay_steps - else: - lr_warmup_steps = args.lr_warmup_iters * samples_per_step - if decay_steps is None: - # No horizon: warm up, then hold constant. The decay steps only need - # to satisfy the scheduler's warmup < decay invariant. - lr_decay_style = "constant" - decay_steps = int(lr_warmup_steps) + 1 - else: - lr_decay_style = args.lr_decay_style - - scheduler = OptimizerParamScheduler( - _SlotParamGroups(groups), - init_lr=args.lr_warmup_init, - max_lr=args.lr, - min_lr=args.min_lr, - lr_warmup_steps=lr_warmup_steps, - lr_decay_steps=decay_steps, - lr_decay_style=lr_decay_style, - start_wd=args.start_weight_decay, - end_wd=args.end_weight_decay, - wd_incr_steps=decay_steps, - wd_incr_style=args.weight_decay_incr_style, - use_checkpoint_opt_param_scheduler=False, - override_opt_param_scheduler=False, - wsd_decay_steps=( - args.lr_wsd_decay_iters * samples_per_step - if lr_decay_style == "WSD" and args.lr_wsd_decay_iters is not None - else None - ), - lr_wsd_decay_style=args.lr_wsd_decay_style, - ) - if resume_step: - scheduler.step(increment=resume_step * samples_per_step) - return scheduler - - -def install_slot_scheduler(args: Namespace, optimizer, adapter, resume_step: int) -> None: - """Attach the adapter's scheduler to the optimizer, keyed by slot.""" - if not hasattr(optimizer, "miles_slot_schedulers"): - optimizer.miles_slot_schedulers = {} - optimizer.miles_slot_schedulers[adapter.slot] = build_slot_scheduler(args, optimizer, adapter, resume_step) - - -def drop_slot_scheduler(optimizer, slot: int) -> None: - """Detach a retired slot's scheduler (the next tenant installs its own).""" - getattr(optimizer, "miles_slot_schedulers", {}).pop(slot, None) - - -def step_slot_schedulers(optimizer, step_batch_sizes: dict[int, int]) -> dict[int, float]: - """Advance exactly the stepped slots' schedules by their batch samples. - Returns slot -> new learning rate, for logging.""" - lr_by_slot: dict[int, float] = {} - for slot, batch_size in step_batch_sizes.items(): - scheduler = optimizer.miles_slot_schedulers[slot] - scheduler.step(increment=batch_size) - if scheduler.optimizer.param_groups: # empty on ranks owning none of the slot's params - lr_by_slot[slot] = scheduler.optimizer.param_groups[0]["lr"] - return lr_by_slot diff --git a/miles/backends/megatron_utils/multi_lora_utils.py b/miles/backends/megatron_utils/multi_lora_utils.py deleted file mode 100644 index 4d6d7689098..00000000000 --- a/miles/backends/megatron_utils/multi_lora_utils.py +++ /dev/null @@ -1,489 +0,0 @@ -import json -import logging -import os -from argparse import Namespace -from collections.abc import Mapping -from pathlib import Path - -import ray -import torch -import torch.distributed as dist - -from miles.backends.training_utils.parallel import get_parallel_state -from miles.ray.multi_lora.controller import get_multi_lora_controller -from miles.utils.adapter_config import AdapterRun -from miles.utils.distributed_utils import get_gloo_group - -logger = logging.getLogger(__name__) - -# Cached by adapter_shard_topology(); the topology is fixed for the run. -_shard_topology: tuple[bool, tuple[tuple[int, int, int], ...]] | None = None - - -def create_multi_lora_instance(args: Namespace): - """Create a MultiLoRA instance from training args.""" - from megatron.bridge.peft.multi_lora import MultiLoRA - - from miles.backends.megatron_utils.lora_utils import convert_target_modules_to_megatron - - lora_type_name = getattr(args, "lora_type", "lora").lower() - if lora_type_name == "canonical_lora": - from megatron.bridge.peft.canonical_lora import CanonicalLoRA - - lora_cls = CanonicalLoRA - else: - from megatron.bridge.peft.lora import LoRA - - lora_cls = LoRA - - # exclude_modules was already folded into target_modules during arg validation. - return MultiLoRA( - target_modules=convert_target_modules_to_megatron(args.target_modules, lora_type=lora_cls), - n_adapters=args.multi_lora_n_adapters, - dim=args.lora_rank, - alpha=args.lora_alpha, - dropout=getattr(args, "lora_dropout", 0.0), - lora_A_init_method=getattr(args, "lora_A_init_method", "xavier"), - lora_B_init_method=getattr(args, "lora_B_init_method", "zero"), - ) - - -def megatron_shard_name(tp_rank: int, pp_rank: int, ep_rank: int, ep_size: int) -> str: - """Adapter shard name for one (tp, pp, ep) coordinate; EP ranks hold different local - experts. The ep suffix is omitted at ep_size == 1 so legacy checkpoints stay loadable.""" - name = f"adapter_megatron_tp{tp_rank}_pp{pp_rank}" - if ep_size > 1: - name += f"_ep{ep_rank}" - return name + ".pt" - - -def adapter_shard_topology() -> tuple[bool, tuple[tuple[int, int, int], ...]]: - """Return ``(this_rank_writes_its_shard, realized (tp, pp, ep) coords)`` via one cached gloo all-gather.""" - global _shard_topology - if _shard_topology is not None: - return _shard_topology - parallel_state = get_parallel_state() - coords = (parallel_state.tp.rank, parallel_state.pp.rank, parallel_state.ep.rank) - if not dist.is_initialized(): - _shard_topology = (True, (coords,)) - return _shard_topology - - current_rank = dist.get_rank() - group = get_gloo_group() - gathered: list[object] = [None] * dist.get_world_size(group=group) - dist.all_gather_object(gathered, (coords, current_rank), group=group) - is_writer = current_rank == min(rank for entry_coords, rank in gathered if entry_coords == coords) - _shard_topology = (is_writer, tuple(sorted({entry_coords for entry_coords, _ in gathered}))) - return _shard_topology - - -def all_megatron_checkpoints_exist(step_dir: Path, shard_names) -> bool: - return all((step_dir / name).exists() for name in shard_names) - - -def find_latest_checkpoint(ckpt_dir: Path) -> tuple[Path | None, int]: - _, coords = adapter_shard_topology() - if not ckpt_dir.exists(): - return None, 0 - - parallel_state = get_parallel_state() - ep_size = parallel_state.ep.size - my_coords = (parallel_state.tp.rank, parallel_state.pp.rank, parallel_state.ep.rank) - - expected = {megatron_shard_name(*coord, ep_size) for coord in coords} - my_shard = megatron_shard_name(*my_coords, ep_size) - # Legacy pre-expert-adapter layout: no ep suffix; safe for all EP ranks to read (EP-replicated). - legacy = {megatron_shard_name(tp, pp, 0, 1) for tp, pp, _ in coords} - my_legacy = megatron_shard_name(my_coords[0], my_coords[1], 0, 1) - - def get_step(d): - return int(d.name.split("_")[1]) - - step_dirs = sorted( - [d for d in ckpt_dir.iterdir() if d.is_dir() and d.name.startswith("step_")], - key=get_step, - reverse=True, - ) - for step_dir in step_dirs: - step = get_step(step_dir) - if all_megatron_checkpoints_exist(step_dir, expected): - return step_dir / my_shard, step - if ep_size > 1 and all_megatron_checkpoints_exist(step_dir, legacy): - logger.info(f"[multilora] resuming from pre-expert-parallel shard layout in {step_dir}") - return step_dir / my_legacy, step - - return None, 0 - - -def zero_optimizer_state_for_adapter(optimizer, model, idx: int) -> None: - from megatron.bridge.peft.multi_lora_layers import MultiLoRALinear, _iter_multi_lora_modules - - target_main_params = set() - for module in _iter_multi_lora_modules(model): - if not isinstance(module, MultiLoRALinear): - continue - adapter = module.adapters[idx] - for param in adapter.parameters(): - main = getattr(param, "main_param", None) - target_main_params.add(id(main if main is not None else param)) - - chained = getattr(optimizer, "chained_optimizers", [optimizer]) - for chained_optimizer in chained: - inner = getattr(chained_optimizer, "optimizer", chained_optimizer) - if inner is None: - continue - # TE/apex FusedAdam tracks the Adam step per param GROUP, not per param; - # reset the retired slot's groups so the next tenant restarts bias correction. - for group in inner.param_groups: - if group.get("miles_multi_lora_slot") == idx and "step" in group: - if isinstance(group["step"], torch.Tensor): - group["step"].zero_() - else: - group["step"] = 0 - for param, state in inner.state.items(): - if id(param) not in target_main_params: - continue - if "exp_avg" in state: - state["exp_avg"].zero_() - if "exp_avg_sq" in state: - state["exp_avg_sq"].zero_() - # Bias correction restarts for the slot's next tenant. - if "step" in state: - if isinstance(state["step"], torch.Tensor): - state["step"].zero_() - else: - state["step"] = 0 - - -def slice_lora_to_rank(hf_name: str, tensor: torch.Tensor, adapter_rank: int) -> torch.Tensor: - """Trim a max-rank-padded LoRA tensor to ``adapter_rank`` on the rank axis, addressed - from the end so packed grouped-expert exports are not sliced on the expert axis.""" - if "lora_A" in hf_name: - rank_dim = tensor.ndim - 2 - if adapter_rank < tensor.shape[rank_dim]: - remainder = tensor.narrow(rank_dim, adapter_rank, tensor.shape[rank_dim] - adapter_rank) - assert remainder.abs().max() == 0, ( - f"lora_A padded dims are non-zero: {hf_name}, " - f"max={remainder.abs().max().item():.6e}, shape={tensor.shape}, rank={adapter_rank}" - ) - return tensor.narrow(rank_dim, 0, adapter_rank) - return tensor - if "lora_B" in hf_name: - rank_dim = tensor.ndim - 1 - if adapter_rank < tensor.shape[rank_dim]: - remainder = tensor.narrow(rank_dim, adapter_rank, tensor.shape[rank_dim] - adapter_rank) - assert remainder.abs().max() == 0, ( - f"lora_B padded dims are non-zero: {hf_name}, " - f"max={remainder.abs().max().item():.6e}, shape={tensor.shape}, rank={adapter_rank}" - ) - return tensor.narrow(rank_dim, 0, adapter_rank) - return tensor - return tensor - - -def save_multi_lora_checkpoints( - args, - model, - adapter_steps: Mapping[str, int], - adapters: Mapping[str, AdapterRun], -): - """Save per-adapter checkpoints in two formats per adapter. - - Layout (per adapter):: - - {adapter.save}/checkpoints/step_{iteration}/ - ├── adapter_megatron_tp{tp}_pp{pp}[_ep{ep}].pt ← per-rank shard, fast resume - ├── adapter_model.safetensors ← gathered HF, inference / external - └── adapter_config.json ← HF PEFT metadata (r, alpha, ...) - """ - from megatron.bridge import AutoBridge - from megatron.bridge.peft.multi_lora_layers import expose_adapter_slot - from safetensors.torch import save_file as save_safetensors - - from miles.backends.megatron_utils.lora_utils import convert_target_modules_to_hf - from miles.utils import megatron_bridge_utils - - parallel_state = get_parallel_state() - tp_rank = parallel_state.tp.rank - pp_rank = parallel_state.pp.rank - ep_rank = parallel_state.ep.rank - ep_size = parallel_state.ep.size - # Exactly one writer per (tp, pp, ep) shard; see adapter_shard_topology. - is_shard_writer, _ = adapter_shard_topology() - is_global_writer = is_shard_writer and tp_rank == 0 and pp_rank == 0 and ep_rank == 0 - - target_modules_hf = ( - convert_target_modules_to_hf(list(args.target_modules)) - if args.target_modules - else ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] - ) - - bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True) - - for adapter_name, adapter in adapters.items(): - config = adapter.config - log_prefix = f"[multilora] ({adapter_name})" - iteration = adapter_steps[adapter_name] - - if config.save is None: - logger.info(f"{log_prefix} skipping checkpoint (no save dir configured)") - continue - - final_dir = config.save / "checkpoints" / f"step_{iteration}" - tmp_dir = config.save / "checkpoints" / f"_tmp_step_{iteration}" - if is_shard_writer: - tmp_dir.mkdir(parents=True, exist_ok=True) - if dist.is_initialized(): - dist.barrier() - - with expose_adapter_slot(model, adapter.slot): - # Megatron checkpoints - if is_shard_writer: - shard: dict[str, torch.Tensor] = { - name: param.data.cpu() - for batch in model - for name, param in batch.named_parameters() - if ".adapter." in name - } - native_path = tmp_dir / megatron_shard_name(tp_rank, pp_rank, ep_rank, ep_size) - torch.save(shard, native_path) - logger.info(f"{log_prefix} saved Megatron shard " f"({len(shard)} tensors) to {native_path}") - - hf_state: dict[str, torch.Tensor] = {} - with megatron_bridge_utils.patch_megatron_model(model): - for hf_name, weight, _megatron_name in bridge.export_adapter_weights( - model, - cpu=True, - show_progress=False, - ): - # Slice from the shared --lora-rank down to this adapter's real rank to - # match adapter_config's r; clone() since safetensors rejects aliased views. - hf_state[hf_name] = slice_lora_to_rank(hf_name, weight, config.rank).clone() - - if is_global_writer: - save_safetensors( - hf_state, - str(tmp_dir / "adapter_model.safetensors"), - metadata={"format": "pt"}, - ) - adapter_config_json = { - "peft_type": "LORA", - "r": config.rank, - "lora_alpha": config.alpha, - "target_modules": target_modules_hf, - "lora_dropout": getattr(args, "lora_dropout", 0.0), - "bias": "none", - "task_type": "CAUSAL_LM", - } - with open(tmp_dir / "adapter_config.json", "w") as f: - json.dump(adapter_config_json, f, indent=2) - os.sync() - logger.info(f"{log_prefix} saved HF PEFT to {tmp_dir} " f"({len(hf_state)} tensors)") - - if dist.is_initialized(): - dist.barrier() - - # Write to a temp dir and move into place so readers never see a - # partially written checkpoint. - if is_global_writer: - if final_dir.exists(): - import shutil - - shutil.rmtree(final_dir) - os.replace(tmp_dir, final_dir) - logger.info(f"{log_prefix} promoted checkpoint to {final_dir}") - if dist.is_initialized(): - dist.barrier() - - -def _register_adapter(adapter: AdapterRun, model) -> int: - """Install one adapter on this rank's local model shard. Returns the step - of the checkpoint it resumed from (0 for a fresh adapter).""" - from megatron.bridge.peft.multi_lora_layers import init_adapter_slot, load_adapter - - name = adapter.name - config = adapter.config - slot = adapter.slot - log_prefix = f"[multilora] ({name})" - - step = 0 - if config.save is not None: - ckpt_root = config.save / "checkpoints" - ckpt, step = find_latest_checkpoint(ckpt_root) - else: - ckpt = None - - if ckpt is None: - logger.info(f"{log_prefix} no checkpoint, starting from random init") - step = 0 - else: - state_dict = torch.load(ckpt, map_location="cpu", weights_only=True) - loaded = load_adapter(model, slot, state_dict) - assert loaded > 0, ( - f"{log_prefix} loaded 0 tensors from {ckpt} " - f"(state_dict has {len(state_dict)} entries) — name mismatch?" - ) - logger.info(f"{log_prefix} loaded from {ckpt} ({loaded} tensors)") - - init_adapter_slot(model, slot, rank=config.rank, alpha=config.alpha) - logger.info(f"{log_prefix} installed at slot {slot}") - return step - - -def _deregister_adapter(adapter: AdapterRun, args, model, optimizer) -> None: - """Model-side cleanup for one adapter.""" - from megatron.bridge.peft.multi_lora_layers import clear_adapter_slot - - name = adapter.name - slot = adapter.slot - log_prefix = f"[multilora] ({name})" - - if args.save_interval is not None: - # The controller still holds the step count until free_slot runs. - step = ray.get(get_multi_lora_controller().adapter_step.remote(name)) - save_multi_lora_checkpoints(args, model, {name: step}, {name: adapter}) - logger.info(f"{log_prefix} saved final checkpoint at step {step}") - else: - logger.info(f"{log_prefix} save_interval unset; skipping final checkpoint") - - clear_adapter_slot(model, slot) - logger.info(f"{log_prefix} cleared adapter slot {slot}") - - # Prevent future slot tenants from inheriting optimizer momentum or the - # previous tenant's partially accumulated gradients. - from miles.backends.megatron_utils.multi_lora_optimizer import zero_adapter_slot_grads - - from miles.backends.megatron_utils.multi_lora_scheduler import drop_slot_scheduler - - zero_optimizer_state_for_adapter(optimizer, model, slot) - zero_adapter_slot_grads(model, slot) - drop_slot_scheduler(optimizer, slot) - optimizer.reload_model_params() - logger.info(f"{log_prefix} cleared optimizer state and retained grads for slot {slot}") - - -def load_adapters(args, model, optimizer, adapters) -> int: - """Load adapters into Megatron slots; resumes step counts from checkpoints.""" - from miles.backends.megatron_utils.initialize import is_first_replica_megatron_main_rank - from miles.utils.distributed_utils import get_gloo_group - - if dist.is_initialized(): - dist.barrier(group=get_gloo_group()) - if not adapters: - return 0 - from miles.backends.megatron_utils.multi_lora_scheduler import install_slot_scheduler - - resume_steps: dict[str, int] = {} - for adapter in adapters: - resume_steps[adapter.name] = _register_adapter(adapter, model) - # Per-adapter LR/WD schedule, positioned at the resumed step count. - install_slot_scheduler(args, optimizer, adapter, resume_steps[adapter.name]) - if dist.is_initialized(): - dist.barrier(group=get_gloo_group()) - optimizer.reload_model_params() - if is_first_replica_megatron_main_rank(): - for name, step in resume_steps.items(): - if step > 0: - ray.get(get_multi_lora_controller().set_adapter_step.remote(name, step)) - return len(adapters) - - -def cleanup_adapters(args, model, optimizer, adapters) -> int: - """Save final ckpt + clear Megatron slot, then free_slot on the controller.""" - from miles.backends.megatron_utils.initialize import is_first_replica_megatron_main_rank - from miles.utils.distributed_utils import get_gloo_group - - if dist.is_initialized(): - dist.barrier(group=get_gloo_group()) - if not adapters: - return 0 - for adapter in adapters: - _deregister_adapter(adapter, args, model, optimizer) - if dist.is_initialized(): - dist.barrier(group=get_gloo_group()) - if is_first_replica_megatron_main_rank(): - for adapter in adapters: - ray.get(get_multi_lora_controller().free_slot.remote(adapter.name)) - return len(adapters) - - -def step_stepped_adapter_slots(args, model, optimizer, rollout_data, rollout_id: int, step_id: int) -> float: - """Optimizer-step the slots whose adapter batch completes with this train batch and advance - their per-adapter LR/WD schedules. Returns the max grad norm across stepped slots (0.0 if none).""" - from miles.backends.megatron_utils.multi_lora_optimizer import step_adapter_slots - from miles.backends.megatron_utils.multi_lora_scheduler import step_slot_schedulers - from miles.utils.tracking_utils.structured_log import log_structured - - # slot -> adapter_global_batch_size for adapter batches completing now. - step_batch_sizes = dict(rollout_data.get("step_adapter_batch_sizes", {})) - grad_norms_by_slot = step_adapter_slots( - optimizer, - model, - step_batch_sizes, - clip_grad=args.clip_grad, - ) - - if lr_by_slot := step_slot_schedulers(optimizer, step_batch_sizes): - log_structured( - logger.info, - op="adapter_lr", - rollout=rollout_id, - step=step_id, - **{f"slot_{slot}": lr for slot, lr in lr_by_slot.items()}, - ) - return max(grad_norms_by_slot.values(), default=0.0) - - -def commit_trained_batch(rollout_data, rollout_id: int, pending_push: set) -> None: - """A train call landed: schedule the stepped adapters' engine push and - commit the batch on the controller (main rank only). The stepped set ships - with the train data, identical on all ranks.""" - from miles.backends.megatron_utils.initialize import is_first_replica_megatron_main_rank - - pending_push.update(rollout_data.get("step_adapter_names", [])) - if is_first_replica_megatron_main_rank(): - ray.get(get_multi_lora_controller().mark_batch_trained.remote(rollout_id)) - - -def save_due_adapter_checkpoints(args, model) -> bool: - """Save per-adapter checkpoints for adapters at a save-interval multiple - without a checkpoint on disk. Rank 0 picks and broadcasts, so the - collective export lines up. Returns False when nothing is due.""" - from miles.backends.megatron_utils.initialize import is_first_replica_megatron_main_rank - from miles.utils.distributed_utils import get_gloo_group - - due_buffer = [None] - if is_first_replica_megatron_main_rank() and args.save_interval is not None: - snapshot = ray.get(get_multi_lora_controller().snapshot.remote()) - adapters = {**snapshot["active"], **snapshot["retiring"]} - due_buffer[0] = { - name: adapter - for name, adapter in adapters.items() - if adapter.step > 0 - and adapter.step % args.save_interval == 0 - and adapter.config.save is not None - and not (Path(adapter.config.save) / "checkpoints" / f"step_{adapter.step}").exists() - } - if dist.is_initialized(): - dist.broadcast_object_list(due_buffer, src=0, group=get_gloo_group()) - due_adapters = due_buffer[0] - if not due_adapters: - return False - adapter_steps = {name: adapter.step for name, adapter in due_adapters.items()} - save_multi_lora_checkpoints(args, model, adapter_steps, due_adapters) - return True - - -def select_adapters_to_push(loaded_adapters: dict, pending_push: set, has_new_engines: bool) -> tuple[dict, list]: - """Pick the stale adapters to push (all loaded adapters when engines are new). Returns - (adapters to push keyed by name, names to version-bump — only those whose weights changed).""" - pending = pending_push & set(loaded_adapters) - push_names = set(loaded_adapters) if has_new_engines else pending - return {name: loaded_adapters[name] for name in sorted(push_names)}, sorted(pending) - - -def commit_weight_push(version_update_names: list, is_main_rank: bool) -> None: - """A weight push landed: bump the pushed adapters' slot versions on the - controller (promotes PENDING adapters to ACTIVE).""" - if version_update_names and is_main_rank: - ray.get(get_multi_lora_controller().record_weight_update.remote(version_update_names)) diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py index bbeea97f815..4a0be671eb2 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py @@ -279,7 +279,7 @@ def _send_one_multi_lora_adapter(self, adapter) -> None: from miles.utils.multi_lora import slot_lora_name - from ...multi_lora_utils import slice_lora_to_rank + from ...api_backends.multi_lora.model import slice_lora_to_rank adapter_rank = adapter.config.rank lora_config = build_lora_sync_config(self.args) | {"r": adapter_rank, "lora_alpha": adapter.config.alpha} @@ -302,12 +302,12 @@ def _send_one_multi_lora_adapter(self, adapter) -> None: self._update_multi_lora_weight_implementation( accumulated_named_tensors, - lora_name=slot_lora_name(adapter.slot), + lora_name=getattr(adapter, "serving_name", None) or slot_lora_name(adapter.slot), lora_config=lora_config, ) def _pause_and_prepare_engines(self) -> None: - """Pause rollout engines, flush cache, and open the weight-update session.""" + """Pause rollout engines and prepare base weights when they will be updated.""" self._weight_update_selector = weight_update_selector(self.args) if dist.get_rank() == 0: mode = self.args.pause_generation_mode @@ -315,10 +315,11 @@ def _pause_and_prepare_engines(self) -> None: if mode != "in_place": ray.get([engine.flush_cache.remote() for engine in self.rollout_engines]) - begin_weight_update(self.rollout_engines, self._weight_update_selector) + if not self.is_lora: + begin_weight_update(self.rollout_engines, self._weight_update_selector) def _finalize_and_resume_engines(self) -> None: - """Close the weight-update session and resume rollout engines.""" + """Finalize base weights when updated and resume rollout engines.""" if dist.get_rank() == 0: # unify update weight version here to cover both full param and lora update ray.get( @@ -327,7 +328,8 @@ def _finalize_and_resume_engines(self) -> None: for engine in self.rollout_engines ] ) - end_weight_update(self.rollout_engines) + if not self.is_lora: + end_weight_update(self.rollout_engines) ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) def pop_metrics(self) -> dict[str, float]: diff --git a/miles/backends/training_utils/data.py b/miles/backends/training_utils/data.py index 1bf6c743679..2eb5a6fd749 100644 --- a/miles/backends/training_utils/data.py +++ b/miles/backends/training_utils/data.py @@ -90,7 +90,7 @@ def get_rollout_data( rollout_data["max_seq_lens"] = [max_seq_len] * len(rollout_data["tokens"]) # Full-response SGLang OPD fields share rollout CP slicing but retain float32 precision. - for key in ("rollout_log_probs", "teacher_log_probs", "opd_reverse_kl"): + for key in ("rollout_log_probs", "teacher_log_probs", "opd_reverse_kl", "loss_weights", "advantages"): if key in rollout_data: dtype = _rollout_logprob_dtype(args) if key == "rollout_log_probs" else torch.float32 rollout_data[key] = [ @@ -155,13 +155,18 @@ def get_batch( assert "tokens" in keys # get_batch consumes adapter_slots itself (per-adapter token counts below); # fetch it here so callers don't have to know. None for non-multi-LoRA runs. - if "adapter_slots" not in keys: - keys = [*keys, "adapter_slots"] + for auto_key in ("adapter_slots", "tinker_operation_lanes"): + if auto_key not in keys: + keys = [*keys, auto_key] batch = data_iterator.get_next(keys) if "dynamic_global_batch_size" in data_iterator.rollout_data: batch["dynamic_global_batch_size"] = data_iterator.rollout_data["dynamic_global_batch_size"] + for key in ("tinker_loss_by_lane", "tinker_forward_only", "tinker_logprob_collector"): + if key in data_iterator.rollout_data: + batch[key] = data_iterator.rollout_data[key] + # No-op safety net if batches reach get_batch without rollout-level preprocessing. expand_multimodal_rollout_data_in_place(batch, qkv_format=qkv_format) diff --git a/miles/backends/training_utils/log_utils.py b/miles/backends/training_utils/log_utils.py index c9b2b8f0f68..87602ae82a6 100644 --- a/miles/backends/training_utils/log_utils.py +++ b/miles/backends/training_utils/log_utils.py @@ -205,9 +205,14 @@ def log_rollout_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatc "num_rollouts", "n_adapters", "adapter_slots", - "step_slots", - "step_adapter_names", - "step_adapter_batch_sizes", + "tinker_operation_lanes", + "tinker_loss_by_lane", + "operation_by_lane", + "registration_by_lane", + "batch_execution_lease", + "batch_kind", + "tinker_forward_only", + "tinker_logprob_collector", "prompt_group_sizes", ]: continue diff --git a/miles/backends/training_utils/loss.py b/miles/backends/training_utils/loss.py index b15311a377b..53f91e865c7 100644 --- a/miles/backends/training_utils/loss.py +++ b/miles/backends/training_utils/loss.py @@ -5,7 +5,7 @@ from miles.backends.training_utils.cp_utils import get_local_response_loss_masks, get_sum_of_sample_mean from miles.backends.training_utils.loss_hub.advantages import compute_advantages, normalize_advantages from miles.backends.training_utils.loss_hub.logit_processors import get_log_probs_and_entropy, get_values # noqa: F401 -from miles.backends.training_utils.loss_hub.losses import get_loss_function +from miles.backends.training_utils.loss_hub.losses import get_loss_function, tinker_loss_function from miles.backends.training_utils.loss_hub.math_utils import compute_approx_kl from miles.backends.training_utils.loss_hub.opd import apply_opd_kl_to_advantages from miles.backends.training_utils.parallel import get_parallel_state @@ -170,7 +170,10 @@ def loss_function( denominators=batch.get("rollout_mask_sums", None), ) - func = get_loss_function(args) + if batch.get("tinker_loss_by_lane"): + func = tinker_loss_function + else: + func = get_loss_function(args) if args.recompute_loss_function: loss, log = checkpoint( diff --git a/miles/backends/training_utils/loss_hub/losses.py b/miles/backends/training_utils/loss_hub/losses.py index 7296bde550e..54d103726d0 100644 --- a/miles/backends/training_utils/loss_hub/losses.py +++ b/miles/backends/training_utils/loss_hub/losses.py @@ -505,6 +505,80 @@ def sft_loss_function( ) +def tinker_loss_function( + args: Namespace, + batch: RolloutBatch, + logits: torch.Tensor, + sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor], +) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + specs_by_lane = batch["tinker_loss_by_lane"] + operation_lanes = batch["tinker_operation_lanes"] + response_lengths = batch["response_lengths"] + total_lengths = batch["total_lengths"] + max_seq_lens = batch.get("max_seq_lens", None) + + log_probs = get_log_probs_and_entropy( + logits, + args=args, + unconcat_tokens=batch["unconcat_tokens"], + total_lengths=total_lengths, + response_lengths=response_lengths, + with_entropy=False, + max_seq_lens=max_seq_lens, + )["log_probs"] + local_masks = get_local_response_loss_masks( + total_lengths, response_lengths, batch["loss_masks"], args.qkv_format, max_seq_lens + ) + + def channel(key: str, i: int, loss_fn: str) -> torch.Tensor: + values = batch.get(key) + if values is None or values[i] is None: + raise ValueError(f"tinker loss '{loss_fn}' needs per-token '{key}'") + return values[i] + + collector = batch.get("tinker_logprob_collector") + if collector is not None: + sample_indices = batch["sample_indices"] + for i, logp in enumerate(log_probs): + full = logp + if get_parallel_state().cp.size > 1: + full = all_gather_with_cp(logp, total_lengths[i], response_lengths[i]) + collector[(operation_lanes[i], sample_indices[i])] = full.detach().float().cpu().tolist() + + if batch.get("tinker_forward_only"): + loss = 0 * logits.sum() + return loss, {"loss": loss.clone().detach()} + + loss = None + for i, logp in enumerate(log_probs): + spec = specs_by_lane.get(operation_lanes[i]) + if spec is None: + raise ValueError(f"tinker backward batch has no loss spec for lane {operation_lanes[i]}") + loss_fn = spec.get("loss_fn", "cross_entropy") + config = spec.get("loss_fn_config") or {} + mask = local_masks[i].to(device=logp.device, dtype=logp.dtype) + if loss_fn == "cross_entropy": + sample_loss = -(logp * channel("loss_weights", i, loss_fn) * mask).sum() + elif loss_fn in ("importance_sampling", "ppo"): + ratio = torch.exp(logp - channel("rollout_log_probs", i, loss_fn)) + advantages = channel("advantages", i, loss_fn) + surrogate = ratio * advantages + if loss_fn == "ppo": + low = config.get("clip_low_threshold", 0.8) + high = config.get("clip_high_threshold", 1.2) + surrogate = torch.minimum(surrogate, ratio.clamp(low, high) * advantages) + sample_loss = -(surrogate * mask).sum() + else: + raise ValueError(f"tinker operation in lane {operation_lanes[i]} requests unknown loss_fn '{loss_fn}'") + loss = sample_loss if loss is None else loss + sample_loss + + if loss is None: + raise ValueError("tinker backward batch produced no loss terms; selections must be homogeneous") + loss = loss + 0 * logits.sum() + + return loss, {"loss": loss.clone().detach()} + + def get_loss_function(args: Namespace) -> LossFunction: match args.loss_type: case "policy_loss": diff --git a/miles/backends/training_utils/operation_execution.py b/miles/backends/training_utils/operation_execution.py new file mode 100644 index 00000000000..7b66862f3a0 --- /dev/null +++ b/miles/backends/training_utils/operation_execution.py @@ -0,0 +1,80 @@ +from dataclasses import dataclass +from typing import Protocol + +from miles.utils.operation_contract import BatchExecutionLease, BindingT + +# Adam defaults currently matching the Tinker protocol adapter's AdamParams. +ADAM_PARAM_DEFAULTS = dict(learning_rate=1e-4, beta1=0.9, beta2=0.95, eps=1e-12, weight_decay=0.0, grad_clip_norm=0.0) + + +def resolve_adam_params(adam_params: dict | None) -> dict: + return {**ADAM_PARAM_DEFAULTS, **{k: v for k, v in (adam_params or {}).items() if v is not None}} + + +@dataclass(frozen=True) +class StepRequest: + operation_id: str + adam_params: dict + + +class ParameterExecutor(Protocol[BindingT]): + def discard_many(self, lease: BatchExecutionLease[BindingT], operation_ids: list[str]) -> dict[str, dict]: ... + + def step_many(self, lease: BatchExecutionLease[BindingT], requests: list[StepRequest]) -> dict[str, dict]: ... + + +def run_optim_controls( + operations: list[dict], + lease: BatchExecutionLease[BindingT], + executor: ParameterExecutor[BindingT], +) -> dict[str, dict]: + all_optim = [op for op in operations if op["kind"] == "optim_step"] + results: dict[str, dict] = {} + + poisoned = [op for op in all_optim if op.get("poison")] + if poisoned: + discard_outcomes = executor.discard_many(lease, [op["operation_id"] for op in poisoned]) + for op in poisoned: + outcome = discard_outcomes.get(op["operation_id"]) + if outcome is None: + results[op["operation_id"]] = dict( + ok=False, + error=f"executor returned no discard outcome for operation '{op['operation_id']}'", + category="server", + ) + continue + results[op["operation_id"]] = ( + dict(ok=False, error=op["poison"], category="user", gradient_window_consumed=True) + if outcome.get("ok") + else outcome + ) + + clean = [op for op in all_optim if not op.get("poison")] + if clean: + requests = [ + StepRequest( + operation_id=op["operation_id"], + adam_params=resolve_adam_params((op.get("payload") or {}).get("adam_params")), + ) + for op in clean + ] + step_outcomes = executor.step_many(lease, requests) + for op in clean: + outcome = step_outcomes.get(op["operation_id"]) + if outcome is None: + outcome = dict( + ok=False, + error=f"executor returned no step outcome for operation '{op['operation_id']}'", + category="server", + ) + results[op["operation_id"]] = outcome + return results + + +def reset_grad_metadata_keep_grads(model_chunks) -> None: + for model_chunk in model_chunks: + if getattr(model_chunk.config, "cuda_graph_impl", "none") != "transformer_engine": + for param in model_chunk.params_with_grad: + param.grad_added_to_main_grad = False + for bucket_group in model_chunk.bucket_groups + model_chunk.expert_parallel_bucket_groups: + bucket_group.reset() diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index fb40d125098..85e5e680997 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -129,10 +129,13 @@ async def update_weights(self, rollout_id: int | None = None): await self._broadcast("update_weights", info=info) - async def reconcile_adapters(self) -> None: - """Multi-LoRA: reconcile loaded adapters with the controller's active set - (load new, cleanup gone). Called by the trainer before generate.""" - await self._broadcast("reconcile_adapters") + async def reconcile_tinker_adapters(self) -> None: + """Converge trainer residency to the tinker controller's registry.""" + await self._broadcast("reconcile_tinker_adapters") + + async def execute_tinker_controls(self, operations: list[dict], lease_metadata: dict) -> dict: + results = await self._broadcast("execute_tinker_controls", operations, lease_metadata) + return results[0] async def onload(self): await self._broadcast("wake_up") diff --git a/miles/ray/multi_lora/__init__.py b/miles/ray/multi_lora/__init__.py index e69de29bb2d..9d3953f4ab5 100644 --- a/miles/ray/multi_lora/__init__.py +++ b/miles/ray/multi_lora/__init__.py @@ -0,0 +1 @@ +"""Multi-LoRA operation control plane and fixed-slot residency.""" diff --git a/miles/ray/multi_lora/backend.py b/miles/ray/multi_lora/backend.py index bee97443c16..9648ad62ea1 100644 --- a/miles/ray/multi_lora/backend.py +++ b/miles/ray/multi_lora/backend.py @@ -1,189 +1,505 @@ -"""Multi-LoRA backend: the registry plus engine-facing aborts, shared by the -controller Ray actor and the HTTP server. Subclass via -``--multi-lora-backend-path``.""" - -import asyncio import logging +import math +import re from dataclasses import replace from pathlib import Path from typing import Any -import httpx - +from miles.ray.multi_lora.config import AdapterRunConfig +from miles.ray.multi_lora.gradient_windows import GradientWindowTracker +from miles.ray.multi_lora.identity import rid_prefix, serving_lora_name +from miles.ray.multi_lora.inference_admin import RouterInferenceAdmin +from miles.ray.multi_lora.operations import OperationLedger from miles.ray.multi_lora.registry import AdapterRegistry, AdapterState -from miles.utils.adapter_config import AdapterRunConfig -from miles.utils.http_utils import router_worker_base_urls -from miles.utils.multi_lora import RID_SEPARATOR, min_groups_per_dp_split +from miles.ray.multi_lora.residency import FixedSlotResidency, ResidentBinding, lease_from_metadata, lease_to_metadata +from miles.utils.operation_contract import BatchExecutionLease logger = logging.getLogger(__name__) +SUPPORTED_LOSS_FNS = ("cross_entropy", "importance_sampling", "ppo") +_ADAM_FIELDS = ("learning_rate", "beta1", "beta2", "eps", "weight_decay", "grad_clip_norm") +_SAMPLE_TENSOR_FIELDS = ("loss_mask", "loss_weights", "advantages", "rollout_log_probs") +_LOSS_REQUIRED_CHANNELS = { + "cross_entropy": ("loss_weights",), + "importance_sampling": ("rollout_log_probs", "advantages"), + "ppo": ("rollout_log_probs", "advantages"), +} + -class MultiLoRABackend: - """Registry + engine-facing aborts, shared by the Ray actor and HTTP server. - Subclass via --multi-lora-backend-path.""" +class MultiLoraOperationBackend: + """Multi-LoRA implementation selected by ``--multi-lora-backend-path``.""" def __init__(self, args: Any, router_url: str) -> None: self.args = args self.registry = AdapterRegistry(args.multi_lora_n_adapters) + self.operations = OperationLedger( + gap_timeout=getattr(args, "tinker_operation_gap_timeout", 600.0), + claimed_ttl=getattr(args, "tinker_operation_claimed_ttl", 1800.0), + ) + # Ledger lifetime rides the registry's completed ring: ring eviction purges the tenant's ledger state. + self.registry.on_completed_evicted = self.operations.drop_tenant + self.gradient_windows = GradientWindowTracker() + self.residency = FixedSlotResidency(self.registry) self.router_url = router_url.rstrip("/") - self.client: httpx.AsyncClient | None = None + self.inference_admin = RouterInferenceAdmin(self.router_url) + self.trainer_ready = False + + def mark_trainer_ready(self) -> None: + self.trainer_ready = True async def init(self) -> None: - self.client = httpx.AsyncClient(timeout=httpx.Timeout(30.0)) + await self.inference_admin.init() async def close(self) -> None: - if self.client is not None: - await self.client.aclose() - self.client = None + await self.inference_admin.close() + + # ---------------- registration ---------------- async def validate_adapter(self, name: str, config: Any) -> None: - """Override to reject adapter registrations (raise ValueError).""" + """Override to reject registrations (raise ValueError).""" def resolve_adapter_config(self, name: str, config: Any) -> Any: - """Resolve optional adapter-local values against process-wide defaults - and validate the batch shape against the trainer's DP layout. - - All batch-shape constraints are enforced here, at registration, so a - bad config fails immediately instead of crashing an arbitrary later - train batch. - """ + """Resolve client fields against deployment defaults. The public + surface takes rank/save/num_step/metadata only; alpha is server-set.""" if config is None or not isinstance(config, AdapterRunConfig): return config - rank = config.rank if config.rank is not None else getattr(self.args, "lora_rank", 1) - alpha = config.alpha if config.alpha is not None else getattr(self.args, "lora_alpha", rank) - rollout_batch_size = ( - config.rollout_batch_size - if config.rollout_batch_size is not None - else getattr(self.args, "rollout_batch_size", None) - ) - n_samples_per_prompt = ( - config.n_samples_per_prompt - if config.n_samples_per_prompt is not None - else getattr(self.args, "n_samples_per_prompt", 1) - ) - if type(rank) is not int or rank <= 0: raise ValueError(f"Adapter '{name}' rank must be a positive integer") if rank > getattr(self.args, "lora_rank", rank): - raise ValueError(f"Adapter '{name}' rank {rank} exceeds the allocated maximum rank {self.args.lora_rank}") - if alpha is None or alpha <= 0: - raise ValueError(f"Adapter '{name}' must have a positive alpha") - if type(rollout_batch_size) is not int or rollout_batch_size <= 0: - raise ValueError(f"Adapter '{name}' rollout_batch_size must be a positive integer (prompt groups)") - if type(n_samples_per_prompt) is not int or n_samples_per_prompt <= 0: - raise ValueError(f"Adapter '{name}' n_samples_per_prompt must be a positive integer") + raise ValueError(f"Adapter '{name}' rank {rank} exceeds the deployment maximum {self.args.lora_rank}") + if config.alpha is not None: + raise ValueError(f"Adapter '{name}' must not set alpha; it is deployment-configured (--lora-alpha)") + alpha = getattr(self.args, "lora_alpha", None) or rank if config.num_step is not None and (type(config.num_step) is not int or config.num_step <= 0): raise ValueError(f"Adapter '{name}' num_step must be a positive integer") - if config.num_epoch is not None and (type(config.num_epoch) is not int or config.num_epoch <= 0): - raise ValueError(f"Adapter '{name}' num_epoch must be a positive integer") - if config.num_step is not None and config.num_epoch is not None: - logger.warning(f"Adapter '{name}' sets both num_step and num_epoch; num_step takes precedence") - - # A bad data path or unresolvable reward config does not fail at this - # API otherwise: the data path kills the shared rollout producer thread - # and an empty reward config burns every generated sample, either way - # stalling ALL adapters behind a misleading empty-batch timeout. - if not Path(config.data).expanduser().exists(): - raise ValueError( - f"Adapter '{name}' data path '{config.data}' does not exist " - "(checked from the controller process, which runs on the head node with the rollout data source)" - ) - if ( - config.custom_rm_path is None - and not (config.rm_type or "").strip() - and getattr(self.args, "custom_rm_path", None) is None - and not (getattr(self.args, "rm_type", None) or "").strip() - ): - raise ValueError( - f"Adapter '{name}' has no reward config: set rm_type or custom_rm_path in the adapter " - "config, or launch with --rm-type / --custom-rm-path" - ) - - adapter_global_batch_size = rollout_batch_size * n_samples_per_prompt - if (max_batch := getattr(self.args, "multi_lora_max_adapter_global_batch_size", None)) is not None: - if adapter_global_batch_size > max_batch: - raise ValueError( - f"Adapter '{name}' consumes {adapter_global_batch_size} samples per step " - f"(rollout_batch_size {rollout_batch_size} x n_samples_per_prompt {n_samples_per_prompt}), " - f"exceeding --multi-lora-max-adapter-global-batch-size {max_batch}" - ) - if (dp_size := getattr(self.args, "multi_lora_dp_size", None)) is not None: - try: - group_multiple = min_groups_per_dp_split(n_samples_per_prompt, dp_size) - except ValueError as e: - raise ValueError(f"Adapter '{name}': {e}") from None - if rollout_batch_size % group_multiple != 0: - raise ValueError( - f"Adapter '{name}' rollout_batch_size {rollout_batch_size} must be a multiple of " - f"its min_groups_per_dp_split ({group_multiple} at dp_size={dp_size}), so the " - f"adapter batch can complete from evenly-splitting takes" - ) - save = Path(config.save) if config.save is not None else None if save is None: if getattr(self.args, "save", None) is None: - raise ValueError(f"Adapter '{name}' has no save dir: set 'save' in the adapter config or pass --save") + raise ValueError(f"Adapter '{name}' has no save dir: set 'save' in the config or pass --save") save = Path(self.args.save) / "adapters" / name - - return replace( - config, - rank=rank, - alpha=alpha, - rollout_batch_size=rollout_batch_size, - n_samples_per_prompt=n_samples_per_prompt, - save=save, - ) + return replace(config, rank=rank, alpha=alpha, save=save) async def register(self, name: str, config: Any) -> dict: config = self.resolve_adapter_config(name, config) await self.validate_adapter(name, config) result = self.registry.register(name, config) - resolved = getattr(config, "save", None) - if resolved is not None: - logger.info(f"Adapter '{name}' registered (slot {result['slot']}), checkpoints -> {resolved}") + self.gradient_windows.open(self.registry.records[name].tenant) + logger.info(f"[tinker] adapter '{name}' registered (slot {result['slot']})") return result - async def deregister(self, name: str) -> None: + async def deregister(self, name: str, expected_registration_id: str | None = None) -> None: + if expected_registration_id is not None: + record = self.registry.find(name) + if record is None or record.registration_id != expected_registration_id: + return # the handle's registration is already gone; never touch a successor self.registry.deregister(name) async def retire_adapters(self) -> list[str]: names = self.registry.retire_adapters() for name in names: - await self.abort_adapter_requests(name) + record = self.registry.records.get(name) + if record is not None: + # Fence before the engine abort: no operation of the dead + # registration may be claimed once retirement is underway. + self.operations.fence(name, record.registration_id) + await self.abort_adapter_requests(name, record.registration_id) return names async def free_slot(self, name: str) -> int: - """Free the adapter's slot after one final abort round: requests can survive the - ``retire_adapters`` abort (e.g. multi-turn groups), and must not leak to the slot's next tenant.""" + """Free the adapter's slot after one final abort round: requests can + survive the retire abort and must not leak to the slot's next tenant.""" record = self.registry.records.get(name) if record is not None and record.state is AdapterState.CLEANUP: - await self.abort_adapter_requests(name) - return self.registry.free_slot(name) - - async def worker_urls(self) -> list[str]: - assert self.client is not None - for endpoint, extract in ( - ("/list_workers", lambda body: body["urls"]), - ("/workers", lambda body: [worker["url"] for worker in body["workers"]]), - ): - try: - resp = await self.client.get(f"{self.router_url}{endpoint}") - if resp.status_code == 200: - return router_worker_base_urls(extract(resp.json())) - except Exception: - continue - return [] + await self.abort_adapter_requests(name, record.registration_id) + slot = self.registry.free_slot(name) + if record is not None and slot != -1: + self.gradient_windows.close(record.tenant) + return slot + + # ---------------- training-stream clocks ---------------- - async def abort_adapter_requests(self, adapter_name: str) -> None: - prefix = f"{adapter_name}{RID_SEPARATOR}" - urls = await self.worker_urls() - if not urls: - logger.warning(f"Abort for adapter '{adapter_name}': no workers discovered at {self.router_url}") + def set_adapter_step(self, name: str, step: int) -> None: + record = self.registry.find(name) + if record is None: return - results = await asyncio.gather( - *(self.client.post(f"{url}/abort_request", json={"rid": prefix, "prefix": True}) for url in urls), - return_exceptions=True, + self.gradient_windows.restore_step(record.tenant, step) + self.registry.set_step(name, step) + + def adapter_step(self, name: str) -> int: + record = self.registry.find(name) + return self.gradient_windows.step_of(record.tenant) if record is not None else 0 + + # ---------------- operation preflight (compatibility matrix) ---------------- + + def enqueue_operation( + self, + name: str, + operation_id: str, + ordinal: int, + kind: str, + payload: dict | None = None, + expected_registration_id: str | None = None, + ) -> dict: + record = self.registry.find(name) + if record is None or record.state not in (AdapterState.PENDING, AdapterState.READY): + raise ValueError(f"Adapter '{name}' is not accepting operations (not registered or retiring)") + self._check_expected_registration(name, record, expected_registration_id) + payload = payload or {} + self._preflight(name, kind, payload) + return self.operations.enqueue(operation_id, name, record.registration_id, ordinal, kind, payload) + + @staticmethod + def _check_expected_registration(name: str, record: Any, expected_registration_id: str | None) -> None: + if expected_registration_id is not None and record.registration_id != expected_registration_id: + raise ValueError( + f"Adapter '{name}' registration {expected_registration_id[:8]} was retired and the name " + f"re-registered ({record.registration_id[:8]}); operations from the stale handle are fenced" + ) + + def reject_operation( + self, + name: str, + operation_id: str, + ordinal: int, + kind: str, + payload: dict | None, + error: str, + expected_registration_id: str | None = None, + ) -> dict: + """Record a boundary-rejected submission as terminal FAILED(user) at + its ordinal (see ``OperationLedger.record_rejected``): a frontend that + refuses a request AFTER the client spent the ordinal must still keep + the registration's arrival sequence gap-free. Like ``enqueue_operation``, + a pinned ``expected_registration_id`` fences stale handles — a rejection + must never consume an ordinal slot of a same-name successor.""" + record = self.registry.find(name) + if record is None or record.state not in (AdapterState.PENDING, AdapterState.READY): + raise ValueError(f"Adapter '{name}' is not accepting operations (not registered or retiring)") + self._check_expected_registration(name, record, expected_registration_id) + return self.operations.record_rejected( + operation_id, name, record.registration_id, ordinal, kind, payload or {}, error + ) + + def _preflight(self, name: str, kind: str, payload: dict) -> None: + if kind in ("forward_backward", "forward"): + samples = payload.get("samples") + if not isinstance(samples, list) or not samples: + raise ValueError(f"{kind} payload needs a non-empty 'samples' list") + required_channels: tuple[str, ...] = () + if kind == "forward_backward": + loss = payload.get("loss") or {} + loss_fn = loss.get("loss_fn", "cross_entropy") + if loss_fn not in SUPPORTED_LOSS_FNS: + raise ValueError( + f"loss_fn '{loss_fn}' is not supported in v1; supported: {', '.join(SUPPORTED_LOSS_FNS)}" + ) + required_channels = _LOSS_REQUIRED_CHANNELS[loss_fn] + for i, sample in enumerate(samples): + self._preflight_sample(name, kind, i, sample, required_channels) + elif kind == "optim_step": + self._preflight_adam_params(payload.get("adam_params") or {}) + elif kind == "save_state": + tag = payload.get("tag") + if tag is not None: + if not isinstance(tag, str): + raise ValueError("save_state 'tag' must be a string") + # Containment: the tag is a single directory name under the + # adapter's states/ dir — '.'/'..' would escape it. + if not re.fullmatch(r"[A-Za-z0-9._-]{1,128}", tag) or tag in (".", ".."): + raise ValueError( + f"save_state tag '{tag}' is invalid: 1-128 chars of [A-Za-z0-9._-], not '.' or '..'" + ) + elif kind == "load_state": + if not isinstance(payload.get("path"), str) or not payload["path"]: + raise ValueError("load_state needs a 'path'") + elif kind == "save_weights_for_sampler": + pass + else: + raise ValueError(f"unknown operation kind '{kind}'") + + def _preflight_sample( + self, name: str, kind: str, index: int, sample: Any, required_channels: tuple[str, ...] = () + ) -> None: + where = f"{kind} sample[{index}]" + if not isinstance(sample, dict): + raise ValueError(f"{where} must be an object") + for banned in ("multimodal_inputs", "multimodal_train_inputs"): + if sample.get(banned): + raise ValueError(f"{where}: multimodal inputs are not supported in v1 (text-only)") + tokens = sample.get("tokens") + response_length = sample.get("response_length") + if not isinstance(tokens, list) or not tokens or not all(isinstance(t, int) for t in tokens): + raise ValueError(f"{where}: 'tokens' must be a non-empty list of ints (1-D; no top-K targets in v1)") + # Strictly below len(tokens): targets are shifted, so the first response + # token's logprob conditions on at least one preceding token. + if not isinstance(response_length, int) or not (0 < response_length < len(tokens)): + raise ValueError(f"{where}: 'response_length' must be an int in (0, len(tokens)) — shifted targets") + for field_name in required_channels: + if sample.get(field_name) is None: + raise ValueError(f"{where}: per-token '{field_name}' is required by this operation's loss_fn") + for field_name in _SAMPLE_TENSOR_FIELDS: + value = sample.get(field_name) + if value is None: + continue + if not isinstance(value, list) or len(value) != response_length: + raise ValueError(f"{where}: '{field_name}' must be a flat list of length response_length (1-D only)") + if any(isinstance(v, (list, dict)) for v in value): + raise ValueError(f"{where}: '{field_name}' must be 1-D; nested targets are not supported in v1") + + def _preflight_adam_params(self, adam: dict) -> None: + for field_name, value in adam.items(): + if field_name not in _ADAM_FIELDS: + raise ValueError(f"unknown adam_params field '{field_name}'") + if value is None: + continue + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + raise ValueError(f"adam_params.{field_name} must be a finite number") + for field_name in ("learning_rate", "weight_decay", "grad_clip_norm"): + if (value := adam.get(field_name)) is not None and value < 0: + raise ValueError(f"adam_params.{field_name} must be >= 0") + for field_name in ("beta1", "beta2"): + if (value := adam.get(field_name)) is not None and not (0 <= value < 1): + raise ValueError(f"adam_params.{field_name} must be in [0, 1)") + if (value := adam.get("eps")) is not None and value <= 0: + raise ValueError("adam_params.eps must be > 0") + + # ---------------- data-operation claims ---------------- + + def claim_data_operation(self, name: str, registration_id: str) -> dict | None: + binding = self.residency.binding_for((name, registration_id)) + if binding is None: + return None + operation = self.operations.claim_data_operation(name, registration_id) + if operation is None: + return None + operation["binding"] = binding + return operation + + def acquire_batch_lease(self, bindings_by_operation: list) -> BatchExecutionLease[ResidentBinding]: + return self.residency.acquire_batch( + tuple((operation_id, binding) for operation_id, binding in bindings_by_operation) ) - if failures := sum(isinstance(r, Exception) for r in results): - logger.warning(f"Abort for adapter '{adapter_name}': {failures}/{len(results)} posts failed") + + def release_batch_lease(self, lease_metadata: dict) -> None: + """Completion-boundary lifecycle hook; no-op under fixed residency.""" + self.residency.release_batch(lease_from_metadata(lease_metadata)) + + # ---------------- control-operation claims ---------------- + + EXECUTABLE_CONTROL_KINDS = ("optim_step", "save_weights_for_sampler", "save_state", "load_state") + DIRTY_GATED_KINDS = ("save_state", "load_state") + + def sweep_operation_timeouts(self) -> None: + # Both liveness backstops ride the same heartbeat: QUEUED gap holes and orphaned CLAIMED heads. + self.operations.sweep_gap_timeouts() + for view in self.operations.claimed_timeouts(): + error = ( + f"claimed-operation timeout: {view['kind']} '{view['operation_id']}' held CLAIMED for " + f"{view['claimed_age']:.0f}s (TTL {self.operations.claimed_ttl:.0f}s) without a terminal " + "outcome; its executor dispatch is presumed lost — resubmit the operation" + ) + logger.warning(f"[tinker] {error}") + self.fail_tinker_batch([view["operation_id"]], error) + + def claim_ready_control_operations(self) -> dict: + self.sweep_operation_timeouts() + ready: list[dict] = [] + bindings: list[tuple[str, ResidentBinding]] = [] + for name, registration_id in self.operations.claimable_control_tenants(): + record = self.registry.find(name) + if record is None or record.registration_id != registration_id: + continue + binding = self.residency.binding_for((name, registration_id)) + if binding is None: + continue + operation = self.operations.claim_control_operation( + name, registration_id, kinds=self.EXECUTABLE_CONTROL_KINDS + ) + if operation is None: + continue + if operation["kind"] == "optim_step": + blocker = self.operations.poisoned_window_blocker(name, registration_id, operation["ordinal"]) + if blocker is not None: + operation["poison"] = ( + f"a forward_backward in this gradient window failed ({blocker}); the window's " + "accumulated gradients were discarded — resubmit the batch and optim_step again" + ) + if operation["kind"] in self.DIRTY_GATED_KINDS and self.gradient_windows.is_dirty(record.tenant): + self.operations.fail( + operation["operation_id"], + f"adapter '{name}' holds unstepped gradients; optim_step (or deregister) before " + f"{operation['kind']}", + "user", + ) + continue + operation["step"] = self.gradient_windows.step_of(record.tenant) + operation["serving_version"] = record.serving_version + ready.append(operation) + bindings.append((operation["operation_id"], binding)) + if not ready: + return {"operations": [], "lease": None} + lease = self.residency.acquire_batch(tuple(bindings)) + return {"operations": ready, "lease": lease_to_metadata(lease)} + + def complete_control_operations(self, results: dict[str, dict]) -> None: + for operation_id, outcome in results.items(): + operation = self.operations.get(operation_id) + # Only still-CLAIMED operations complete: a swept/fenced (already terminal) one keeps its outcome. + if operation is None or operation["state"] != "CLAIMED": + continue + if outcome.get("ok"): + result = outcome.get("result") + if operation["kind"] == "save_weights_for_sampler": + # Completing after the push landed (the publish barrier): + # stamp the authoritative post-push serving identity. + record = self.registry.find(operation["name"]) + result = { + **(result or {}), + "serving_version": record.serving_version if record else None, + "serving_name": serving_lora_name(operation["name"], operation["registration_id"]), + } + self.operations.complete(operation_id, result) + key = (operation["name"], operation["registration_id"]) + if operation["kind"] == "optim_step": + self.operations.mark_window_consumed(operation_id) + step = self.gradient_windows.commit_step(key) + self.registry.on_step_committed(operation["name"], operation["registration_id"], step) + elif operation["kind"] == "load_state": + step = int((outcome.get("result") or {}).get("step", 0)) + self.gradient_windows.restore_step(key, step) + self.registry.set_step(operation["name"], step) + else: + self.operations.fail( + operation_id, outcome.get("error", "control operation failed"), outcome.get("category", "server") + ) + if operation["kind"] == "optim_step" and outcome.get("gradient_window_consumed"): + self.operations.mark_window_consumed(operation_id) + self.gradient_windows.clear_after_executed_optim((operation["name"], operation["registration_id"])) + self.registry.clear_dirty(operation["name"]) + + def commit_tinker_batch( + self, + accumulated: list[tuple[str, str]], + operation_ids: list[str], + logprobs_by_op: dict[str, list] | None = None, + ) -> None: + for name, registration_id in accumulated: + record = self.registry.find(name) + if record is None or record.registration_id != registration_id: + continue + self.gradient_windows.mark_forward_backward_succeeded(record.tenant) + # Multi-LoRA mirror: pin the accumulating slot's state immovable. + self.registry.mark_accumulated([name]) + logprobs_by_op = logprobs_by_op or {} + for operation_id in operation_ids: + operation = self.operations.get(operation_id) + if operation is not None and operation["state"] == "CLAIMED": + logprobs = logprobs_by_op.get(operation_id) + result = {"logprobs": logprobs} + if operation["kind"] == "forward_backward" and logprobs is not None: + result["metrics"] = operation_result_metrics(self.operations.payload(operation_id), logprobs) + self.operations.complete(operation_id, result) + + def fail_tinker_batch(self, operation_ids: list[str], error: str, lease_metadata: dict | None = None) -> None: + try: + for operation_id in operation_ids: + operation = self.operations.get(operation_id) + if operation is not None and operation["state"] == "CLAIMED": + self.operations.fail(operation_id, error, "server") + finally: + if lease_metadata is not None: + self.residency.release_batch(lease_from_metadata(lease_metadata)) + + # ---------------- engine-facing ---------------- + + async def abort_adapter_requests(self, adapter_name: str, registration_id: str) -> None: + await self.inference_admin.abort_registration(rid_prefix(adapter_name, registration_id)) + + # ---------------- frontend facade ---------------- + # The HTTP frontend sees projections and verbs only — never the registry, + # the ledger, or the router URL (codex-rollout-fullparameter-design-0810 + # §4.2; §3.7 dependency rule: frontend -> backend facade + sampling + # transport). A future lifecycle strategy replaces what sits behind these + # without forking the frontend. + + def registration_view(self, name: str) -> dict | None: + """Projection of the name's CURRENT registration: identity, lifecycle + state, resolved rank, bound-ness, and serving version.""" + record = self.registry.find(name) + if record is None: + return None + return dict( + name=record.name, + registration_id=record.registration_id, + state=record.state.value, + rank=getattr(record.config, "rank", None), + bound=record.slot is not None, + serving_version=record.serving_version, + ) + + def operation_view(self, operation_id: str) -> dict | None: + self.sweep_operation_timeouts() + view = self.operations.get(operation_id) + if view is not None and view["state"] == "QUEUED": + for stall in self.operations.gap_stalls(): + if (stall["name"], stall["registration_id"]) == (view["name"], view["registration_id"]): + view["waiting_on_ordinal"] = stall["missing_ordinal"] + view["gap_stalled_for"] = stall["stalled_for"] + return view + + def ack_operation(self, operation_id: str) -> None: + self.operations.ack(operation_id) + + def sampling_endpoint(self) -> str: + """Base URL sampling requests go to: the SGLang router today, the + InferenceController-provided endpoint after PR #1842.""" + return self.router_url + + # ---------------- info ---------------- + + def service_info(self) -> dict: + self.sweep_operation_timeouts() + args = self.args + return dict( + base_model=getattr(args, "hf_checkpoint", None), + lora_rank_max=getattr(args, "lora_rank", None), + n_adapters=getattr(args, "multi_lora_n_adapters", None), + occupied_slots=self.registry.slot_pool.occupied_slot_ids(), + ready_adapters=sorted(self.registry.in_state(AdapterState.READY)), + supported_loss_fns=list(SUPPORTED_LOSS_FNS), + operation_gap_timeout=self.operations.gap_timeout, + operation_claimed_ttl=self.operations.claimed_ttl, + gap_stalls=self.operations.gap_stalls(), + ) + + +def operation_result_metrics(payload: dict, logprobs: list[list[float]]) -> dict[str, float]: + spec = payload.get("loss") or {} + loss_fn = spec.get("loss_fn", "cross_entropy") + config = spec.get("loss_fn_config") or {} + total = 0.0 + weighted_tokens = 0.0 + loss_weight_sum = 0.0 + for sample, sample_logprobs in zip(payload.get("samples") or [], logprobs, strict=False): + mask = sample.get("loss_mask") or [1.0] * len(sample_logprobs) + weighted_tokens += sum(1.0 for m in mask if m) + if loss_fn == "cross_entropy": + weights = sample.get("loss_weights") or [] + total += sum(-lp * w * m for lp, w, m in zip(sample_logprobs, weights, mask, strict=False)) + loss_weight_sum += sum(w * m for w, m in zip(weights, mask, strict=False)) + else: + old = sample.get("rollout_log_probs") or [] + advantages = sample.get("advantages") or [] + for lp, old_lp, advantage, m in zip(sample_logprobs, old, advantages, mask, strict=False): + ratio = math.exp(min(lp - old_lp, 80.0)) + surrogate = ratio * advantage + if loss_fn == "ppo": + low = config.get("clip_low_threshold", 0.8) + high = config.get("clip_high_threshold", 1.2) + surrogate = min(surrogate, min(max(ratio, low), high) * advantage) + total += -surrogate * m + metrics = {"loss:sum": total, "unmasked_tokens:sum": weighted_tokens} + if loss_fn == "cross_entropy": + metrics["loss_weight:sum"] = loss_weight_sum + return metrics diff --git a/miles/ray/multi_lora/cache.py b/miles/ray/multi_lora/cache.py new file mode 100644 index 00000000000..21a6fda703e --- /dev/null +++ b/miles/ray/multi_lora/cache.py @@ -0,0 +1,33 @@ +"""Cached resident-adapter projection for rollout request routing.""" + +import time + +from miles.utils.misc import SingletonMeta + + +class AdaptersCache(metaclass=SingletonMeta): + """TTL-cache the controller's ready and retiring adapter registrations.""" + + def __init__(self, ttl_s: float = 1.0) -> None: + self.ttl_s = ttl_s + self.snapshot: dict = {"pending": {}, "ready": {}, "retiring": {}, "cleanup": []} + self.last_refresh: float | None = None + + async def get_snapshot(self) -> dict: + from miles.ray.multi_lora.controller import get_multi_lora_controller + + now = time.monotonic() + if self.last_refresh is None or now - self.last_refresh >= self.ttl_s: + try: + self.snapshot = await get_multi_lora_controller().snapshot.remote() + self.last_refresh = now + except Exception: + pass + return self.snapshot + + async def get_all(self) -> dict: + snapshot = await self.get_snapshot() + return {**snapshot.get("ready", {}), **snapshot.get("retiring", {})} + + async def get(self, adapter_name: str): + return (await self.get_all()).get(adapter_name) diff --git a/miles/ray/multi_lora/config.py b/miles/ray/multi_lora/config.py new file mode 100644 index 00000000000..99739963f8c --- /dev/null +++ b/miles/ray/multi_lora/config.py @@ -0,0 +1,52 @@ +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + + +@dataclass(frozen=True) +class AdapterRunConfig: + # LoRA rank; resolved against --lora-rank (ceiling) on register. + rank: int | None = None + # Server-internal: resolved from --lora-alpha; the public API never takes it. + alpha: int | None = None + # Checkpoint root; defaults to {--save}/adapters/{name}. + save: str | Path | None = None + # Optional client-set bound: auto-deregister after N optimizer steps. + num_step: int | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class AdapterRun: + """Read-only join view of a run's static config and current clocks.""" + + name: str + config: AdapterRunConfig + slot: int | None + version: int = 0 + step: int = 0 + # Unique per registration: a re-registered name is a new tenant, and any + # state stamped by the previous tenant must not carry over. + registration_id: str = "" + + @property + def serving_name(self) -> str: + from miles.ray.multi_lora.identity import serving_lora_name + + return serving_lora_name(self.name, self.registration_id) + + +def parse_adapter_run_yaml(path: Path) -> AdapterRunConfig: + import yaml + + with open(path) as f: + raw = yaml.safe_load(f) or {} + known = {"rank", "save", "num_step", "metadata"} + if unknown := set(raw) - known: + raise ValueError(f"adapter yaml {path} has unsupported fields: {sorted(unknown)} (allowed: {sorted(known)})") + return AdapterRunConfig( + rank=raw.get("rank"), + save=Path(raw["save"]) if raw.get("save") else None, + num_step=raw.get("num_step"), + metadata=raw.get("metadata") or {}, + ) diff --git a/miles/ray/multi_lora/controller.py b/miles/ray/multi_lora/controller.py index 7cbff2b5b90..05ec76a7aad 100644 --- a/miles/ray/multi_lora/controller.py +++ b/miles/ray/multi_lora/controller.py @@ -1,18 +1,16 @@ -"""Named Ray actor wrapping the multi-LoRA backend + HTTP server.""" +"""Ray actor for the Multi-LoRA operation control surface.""" -import time from functools import cache from typing import Any import ray -from miles.ray.multi_lora.backend import MultiLoRABackend -from miles.ray.multi_lora.http_server import MultiLoRAHTTPServer -from miles.utils.adapter_config import AdapterRun -from miles.utils.misc import SingletonMeta, get_current_node_ip, load_function +from miles.ray.multi_lora.backend import MultiLoraOperationBackend +from miles.ray.multi_lora.http_server import AdapterRunControlServer +from miles.utils.misc import load_function from miles.utils.ray_utils import compute_ray_pin_head_options -CONTROLLER_NAME = "miles_multi_lora_controller" +CONTROLLER_NAME = "miles_tinker_controller" CONTROLLER_NAMESPACE = "miles" @@ -21,33 +19,6 @@ def get_multi_lora_controller(): return ray.get_actor(CONTROLLER_NAME, namespace=CONTROLLER_NAMESPACE) -class AdaptersCache(metaclass=SingletonMeta): - """TTL-cached controller snapshot; get/get_all expose the sampleable - projection (active + retiring).""" - - def __init__(self, ttl_s: float = 1.0) -> None: - self.ttl_s = ttl_s - self.snapshot: dict = {"pending": {}, "active": {}, "retiring": {}, "cleanup": []} - self.last_refresh: float | None = None - - async def get_snapshot(self) -> dict: - now = time.monotonic() - if self.last_refresh is None or now - self.last_refresh >= self.ttl_s: - try: - self.snapshot = await get_multi_lora_controller().snapshot.remote() - self.last_refresh = now - except Exception: - pass - return self.snapshot - - async def get_all(self) -> dict[str, "AdapterRun"]: - snapshot = await self.get_snapshot() - return {**snapshot["active"], **snapshot["retiring"]} - - async def get(self, adapter_name: str) -> "AdapterRun | None": - return (await self.get_all()).get(adapter_name) - - def _load_subclass(path: str | None, base_cls): if not path: return base_cls @@ -57,10 +28,12 @@ def _load_subclass(path: str | None, base_cls): @ray.remote(num_cpus=0) -class MultiLoRAController: - def __init__(self, args, router_url: str, host: str = "0.0.0.0") -> None: - backend_cls = _load_subclass(getattr(args, "multi_lora_backend_path", None), MultiLoRABackend) - server_cls = _load_subclass(getattr(args, "multi_lora_http_server_path", None), MultiLoRAHTTPServer) +class MultiLoraOperationController: + # Loopback by default: the control plane executes client-referenced work + # and must be fronted by the (future) authenticated tinker frontend. + def __init__(self, args, router_url: str, host: str = "127.0.0.1") -> None: + backend_cls = _load_subclass(getattr(args, "multi_lora_backend_path", None), MultiLoraOperationBackend) + server_cls = _load_subclass(getattr(args, "multi_lora_http_server_path", None), AdapterRunControlServer) self.backend = backend_cls(args, router_url) self.server = server_cls(self.backend, host, api_port=getattr(args, "multi_lora_api_port", 0)) @@ -73,11 +46,13 @@ async def stop(self) -> None: await self.server.stop() await self.backend.close() + # ---------------- registration lifecycle ---------------- + async def register_adapter(self, name: str, config: Any) -> dict: return await self.backend.register(name, config) - async def deregister_adapter(self, name: str) -> None: - await self.backend.deregister(name) + async def deregister_adapter(self, name: str, expected_registration_id: str | None = None) -> None: + await self.backend.deregister(name, expected_registration_id) async def retire_adapters(self) -> list[str]: return await self.backend.retire_adapters() @@ -85,37 +60,108 @@ async def retire_adapters(self) -> list[str]: async def free_slot(self, name: str) -> int: return await self.backend.free_slot(name) - def record_weight_update(self, names: list[str]) -> None: - self.backend.registry.record_weight_update(names) + def bootstrap_pending(self) -> list[str]: + return self.backend.registry.bootstrap_pending() - def record_batch_adapters(self, rollout_id: int, groups: dict[str, int], step_names: list[str]) -> None: - self.backend.registry.record_batch_adapters(rollout_id, groups, step_names) + def mark_ready(self, names: list[str]) -> None: + self.backend.registry.mark_ready(names) - def mark_batch_trained(self, rollout_id: int) -> list[str]: - return self.backend.registry.mark_batch_trained(rollout_id) + def record_weight_update(self, names: list[str]) -> None: + self.backend.registry.record_weight_update(names) - def resolve_num_step(self, name: str, dataset_rows: int) -> None: - self.backend.registry.resolve_num_step(name, dataset_rows) + def set_trainer_ready(self) -> None: + self.backend.mark_trainer_ready() def set_adapter_step(self, name: str, step: int) -> None: - self.backend.registry.set_step(name, step) + self.backend.set_adapter_step(name, step) def adapter_step(self, name: str) -> int: - return self.backend.registry.step_count(name) + return self.backend.adapter_step(name) def snapshot(self) -> dict: return self.backend.registry.snapshot() + # ---------------- operations ---------------- + + def enqueue_operation( + self, + name: str, + operation_id: str, + ordinal: int, + kind: str, + payload: dict | None = None, + expected_registration_id: str | None = None, + ) -> dict: + return self.backend.enqueue_operation(name, operation_id, ordinal, kind, payload, expected_registration_id) + + def reject_operation( + self, + name: str, + operation_id: str, + ordinal: int, + kind: str, + payload: dict | None = None, + error: str = "", + expected_registration_id: str | None = None, + ) -> dict: + return self.backend.reject_operation( + name, operation_id, ordinal, kind, payload, error, expected_registration_id + ) + + def claim_data_operation(self, name: str, registration_id: str) -> dict | None: + # Claim-and-bind in this single actor call: no binding, no CLAIMED. + return self.backend.claim_data_operation(name, registration_id) + + def acquire_batch_lease(self, bindings_by_operation: list): + return self.backend.acquire_batch_lease(bindings_by_operation) + + def release_batch_lease(self, lease_metadata: dict) -> None: + self.backend.release_batch_lease(lease_metadata) + + def claim_ready_control_operations(self) -> list[dict]: + return self.backend.claim_ready_control_operations() + + def complete_control_operations(self, results: dict) -> None: + self.backend.complete_control_operations(results) + + def commit_tinker_batch(self, accumulated: list, operation_ids: list, logprobs_by_op: dict | None = None) -> None: + # ``accumulated`` is a list of exact (name, registration_id) keys; + # normalize sequence types that crossed the Ray boundary. + self.backend.commit_tinker_batch([tuple(key) for key in accumulated], list(operation_ids), logprobs_by_op) + + def fail_tinker_batch(self, operation_ids: list, error: str, lease_metadata: dict | None = None) -> None: + # The abnormal-outcome finalizer for a dispatched data batch that did + # not commit: still-CLAIMED operations terminal-fail typed server. + self.backend.fail_tinker_batch(list(operation_ids), error, lease_metadata) + + def complete_operation(self, operation_id: str, result: dict | None = None) -> None: + self.backend.operations.complete(operation_id, result) + + def fail_operation(self, operation_id: str, error: str, category: str = "server") -> None: + self.backend.operations.fail(operation_id, error, category) + + def cancel_operation(self, operation_id: str) -> dict: + return self.backend.operations.cancel(operation_id) + + def get_operation(self, operation_id: str) -> dict | None: + return self.backend.operation_view(operation_id) + + def ack_operation(self, operation_id: str) -> None: + self.backend.operations.ack(operation_id) + + def service_info(self) -> dict: + return self.backend.service_info() + def http_host(self) -> str: - return get_current_node_ip() + return self.server.advertised_host def api_port(self) -> int: return self.server.actual_api_port -def create_multilora_controller(args, router_url: str, host: str = "0.0.0.0"): +def create_multi_lora_controller(args, router_url: str, host: str = "127.0.0.1"): # Pinned to the head node so the API sits at a port-forwardable address. - return MultiLoRAController.options( + return MultiLoraOperationController.options( name=CONTROLLER_NAME, namespace=CONTROLLER_NAMESPACE, **compute_ray_pin_head_options(), diff --git a/miles/ray/multi_lora/gradient_windows.py b/miles/ray/multi_lora/gradient_windows.py new file mode 100644 index 00000000000..1c3800433f3 --- /dev/null +++ b/miles/ray/multi_lora/gradient_windows.py @@ -0,0 +1,58 @@ +from dataclasses import dataclass + +from miles.utils.operation_contract import RegistrationKey + + +@dataclass +class TrainingStreamState: + step: int = 0 + # True while the stream holds unstepped accumulated gradients. + dirty: bool = False + + +class GradientWindowTracker: + """Step/dirty authority for every live training stream.""" + + def __init__(self) -> None: + self._streams: dict[RegistrationKey, TrainingStreamState] = {} + + def _stream(self, key: RegistrationKey) -> TrainingStreamState: + return self._streams.setdefault(key, TrainingStreamState()) + + # ------------------------------ lifecycle ------------------------------ + + def open(self, key: RegistrationKey) -> None: + """Start tracking a registration's stream (idempotent).""" + self._stream(key) + + def close(self, key: RegistrationKey) -> None: + """Drop a retired registration's stream state.""" + self._streams.pop(key, None) + + # ------------------------------ queries ------------------------------ + + def step_of(self, key: RegistrationKey) -> int: + stream = self._streams.get(key) + return stream.step if stream is not None else 0 + + def is_dirty(self, key: RegistrationKey) -> bool: + stream = self._streams.get(key) + return stream is not None and stream.dirty + + # ------------------------------ transitions ------------------------------ + + def mark_forward_backward_succeeded(self, key: RegistrationKey) -> None: + self._stream(key).dirty = True + + def clear_after_executed_optim(self, key: RegistrationKey) -> None: + self._stream(key).dirty = False + + def commit_step(self, key: RegistrationKey) -> int: + stream = self._stream(key) + stream.step += 1 + stream.dirty = False + return stream.step + + def restore_step(self, key: RegistrationKey, step: int) -> None: + stream = self._stream(key) + stream.step = step diff --git a/miles/ray/multi_lora/http_server.py b/miles/ray/multi_lora/http_server.py index b209142e1ed..f3a4dd2e02e 100644 --- a/miles/ray/multi_lora/http_server.py +++ b/miles/ray/multi_lora/http_server.py @@ -1,35 +1,39 @@ -"""Multi-LoRA control-plane HTTP API over a MultiLoRABackend. - -Subclass via ``--multi-lora-http-server-path`` (override add_routes / -create_app).""" - import asyncio from dataclasses import asdict from pathlib import Path +from typing import Any import uvicorn from fastapi import FastAPI, HTTPException, Query, Request from fastapi.responses import JSONResponse from pydantic import BaseModel +from miles.ray.multi_lora.config import AdapterRunConfig, parse_adapter_run_yaml from miles.ray.multi_lora.registry import AdapterState -from miles.utils.adapter_config import AdapterRunConfig, parse_adapter_run_yaml + +_NAMES_QUERY = Query(default_factory=list) + + +class PublicRunConfig(BaseModel): + rank: int | None = None + save: str | None = None + num_step: int | None = None + metadata: dict[str, Any] = {} + + def to_config(self) -> AdapterRunConfig: + return AdapterRunConfig(rank=self.rank, save=self.save, num_step=self.num_step, metadata=self.metadata) class RegisterAdapterRequest(BaseModel): """Exactly one of ``config`` (inline) or ``yaml_path`` must be set.""" name: str - config: AdapterRunConfig | None = None + config: PublicRunConfig | None = None yaml_path: str | None = None -_NAMES_QUERY = Query(default_factory=list) - - -class MultiLoRAHTTPServer: - """Control-plane API over a MultiLoRABackend. Subclass via - --multi-lora-http-server-path (add_routes / create_app).""" +class AdapterRunControlServer: + """Subclass via --multi-lora-http-server-path (add_routes / create_app).""" def __init__(self, backend, host="127.0.0.1", api_port=0): self.backend = backend @@ -44,22 +48,26 @@ def actual_api_port(self) -> int: return self.api_server.servers[0].sockets[0].getsockname()[1] return self.api_port + @property + def advertised_host(self) -> str: + if self.host in ("0.0.0.0", "::", ""): + from miles.utils.misc import get_current_node_ip + + return get_current_node_ip() + return self.host + def create_app(self) -> FastAPI: - app = FastAPI(title="Miles Multi-LoRA Controller") + app = FastAPI(title="Miles Multi-LoRA operation backend") @app.exception_handler(ValueError) async def value_error_handler(request: Request, exc: ValueError): return JSONResponse({"detail": str(exc)}, status_code=400) - @app.exception_handler(RuntimeError) - async def runtime_error_handler(request: Request, exc: RuntimeError): - status = 409 if "No free adapter slots" in str(exc) else 500 - return JSONResponse({"detail": str(exc)}, status_code=status) - return app def add_routes(self, app: FastAPI) -> None: app.get("/health")(self.health) + app.get("/info")(self.service_info) app.get("/adapter_runs")(self.list_adapters) app.get("/adapter_runs/state")(self.adapter_states) # before /adapter_runs/{name} app.get("/adapter_runs/{name}")(self.get_adapter) @@ -112,13 +120,16 @@ async def get_adapter(self, name: str) -> dict: return status raise HTTPException(status_code=404, detail=f"Adapter '{name}' not registered") + async def service_info(self) -> dict: + return self.backend.service_info() + async def register_adapter(self, request: RegisterAdapterRequest) -> dict: if (request.config is None) == (request.yaml_path is None): raise HTTPException(status_code=400, detail="Exactly one of 'config' or 'yaml_path' must be set") if request.yaml_path is not None: config = parse_adapter_run_yaml(Path(request.yaml_path)) else: - config = request.config + config = request.config.to_config() return await self.backend.register(request.name, config) async def deregister_adapter(self, name: str) -> dict: diff --git a/miles/ray/multi_lora/identity.py b/miles/ray/multi_lora/identity.py new file mode 100644 index 00000000000..12e94967944 --- /dev/null +++ b/miles/ray/multi_lora/identity.py @@ -0,0 +1,25 @@ +"""Registration-scoped identities for the Multi-LoRA operation backend.""" + +import uuid + +RID_SEPARATOR = "::" + + +def make_rid(adapter_name: str, registration_id: str) -> str: + """Mint a request ID inside one exact adapter registration.""" + return f"{adapter_name}{RID_SEPARATOR}{registration_id}{RID_SEPARATOR}{uuid.uuid4().hex}" + + +def rid_prefix(adapter_name: str, registration_id: str) -> str: + """Return the abort namespace for one exact adapter registration.""" + return f"{adapter_name}{RID_SEPARATOR}{registration_id}{RID_SEPARATOR}" + + +def serving_lora_name(adapter_name: str, registration_id: str) -> str: + """Return the engine-side name for one exact adapter registration.""" + return f"__miles_adapter_{adapter_name}_{registration_id}" + + +def cache_extra_key(adapter_name: str, registration_id: str, serving_version: int) -> str: + """Return the registration- and version-scoped KV-cache namespace.""" + return f"{adapter_name}:{registration_id}:v{serving_version}" diff --git a/miles/ray/multi_lora/inference_admin.py b/miles/ray/multi_lora/inference_admin.py new file mode 100644 index 00000000000..7a5d567104c --- /dev/null +++ b/miles/ray/multi_lora/inference_admin.py @@ -0,0 +1,66 @@ +import asyncio +import logging +from typing import Protocol + +import httpx + +from miles.utils.http_utils import router_worker_base_urls + +logger = logging.getLogger(__name__) + + +class InferenceAdminPort(Protocol): + async def init(self) -> None: + """Open the transport; declared in the contract so a fake implementing the port never raises AttributeError.""" + ... + + async def close(self) -> None: + """Release the transport (idempotent).""" + ... + + async def abort_registration(self, rid_prefix: str) -> None: + """Abort in-flight engine requests carrying this registration's rid prefix (anti-ABA: never a successor's).""" + ... + + +class RouterInferenceAdmin: + """Current adapter: worker discovery via the router's + ``/list_workers``|``/workers`` and per-worker ``/abort_request`` posts.""" + + def __init__(self, router_url: str) -> None: + self.router_url = router_url.rstrip("/") + self.client: httpx.AsyncClient | None = None + + async def init(self) -> None: + self.client = httpx.AsyncClient(timeout=httpx.Timeout(30.0)) + + async def close(self) -> None: + if self.client is not None: + await self.client.aclose() + self.client = None + + async def worker_urls(self) -> list[str]: + assert self.client is not None + for endpoint, extract in ( + ("/list_workers", lambda body: body["urls"]), + ("/workers", lambda body: [worker["url"] for worker in body["workers"]]), + ): + try: + resp = await self.client.get(f"{self.router_url}{endpoint}") + if resp.status_code == 200: + return router_worker_base_urls(extract(resp.json())) + except Exception: + continue + return [] + + async def abort_registration(self, rid_prefix: str) -> None: + urls = await self.worker_urls() + if not urls: + logger.warning(f"[tinker] abort for '{rid_prefix}': no workers discovered at {self.router_url}") + return + results = await asyncio.gather( + *(self.client.post(f"{url}/abort_request", json={"rid": rid_prefix, "prefix": True}) for url in urls), + return_exceptions=True, + ) + if failures := sum(isinstance(r, Exception) for r in results): + logger.warning(f"[tinker] abort for '{rid_prefix}': {failures}/{len(results)} posts failed") diff --git a/miles/ray/multi_lora/operations.py b/miles/ray/multi_lora/operations.py new file mode 100644 index 00000000000..cf39af98688 --- /dev/null +++ b/miles/ray/multi_lora/operations.py @@ -0,0 +1,554 @@ +"""Per-registration ledger for the Multi-LoRA operation backend. + +Clients push protocol-neutral operations; data-bearing kinds ride the rollout +selection path through the queue child rollout fn, data-less kinds execute in +the driver's control phase. One registration is strictly serialized: an +operation is claimable only when every earlier operation reached a terminal +state, which carries the client's per-model ordering end to end. + +Arrival may be OUT OF ORDER (the tinker SDK deliberately posts the first +chunk of a large forward_backward last): operations buffer by ordinal and a +gap below the head blocks claims until it fills. Ordinals are consecutive +integers starting at 1 per registration. +NOTE(frontend): the tinker HTTP frontend forwards the SDK's per-model +seq_id verbatim as the ordinal (the counters are the same contract), so +this gap buffer IS the frontend's reorder point — out-of-order chunk +arrival lands here by design. A submission the frontend rejects still +consumes its ordinal via record_rejected (terminal on arrival), keeping +the sequence gap-free. + +Retries are fingerprinted: re-enqueueing a known operation_id with an +identical (kind, payload) returns the original operation; a different +fingerprint is a conflict error, never silently swallowed. + +All mutations run inside the controller actor between awaits, so ledger +methods are synchronous and atomic by construction. +""" + +import hashlib +import json +import logging +import time +from bisect import insort +from dataclasses import dataclass, field +from enum import Enum + +logger = logging.getLogger(__name__) + +Tenant = tuple[str, str] + + +class OperationKind(str, Enum): + FORWARD_BACKWARD = "forward_backward" + FORWARD = "forward" + OPTIM_STEP = "optim_step" + SAVE_WEIGHTS_FOR_SAMPLER = "save_weights_for_sampler" + SAVE_STATE = "save_state" + LOAD_STATE = "load_state" + + +# Ride the rollout/BatchPlan path (they carry Datums). +DATA_KINDS = frozenset({OperationKind.FORWARD_BACKWARD, OperationKind.FORWARD}) +# Execute in the driver's control phase (no Datums). +CONTROL_KINDS = frozenset(OperationKind) - DATA_KINDS + + +class OperationState(str, Enum): + QUEUED = "QUEUED" + CLAIMED = "CLAIMED" + SUCCEEDED = "SUCCEEDED" + FAILED = "FAILED" + CANCELLED = "CANCELLED" + + +TERMINAL_STATES = frozenset({OperationState.SUCCEEDED, OperationState.FAILED, OperationState.CANCELLED}) + + +class OperationBackpressure(RuntimeError): + """Queue or unacked-result capacity reached; the caller must retry later + (the HTTP layer maps this to 429 + Retry-After — 4xx families the tinker + SDK treats as fatal must never carry backpressure).""" + + +def payload_fingerprint(kind: str, payload: dict | None) -> str: + """Canonical digest of an operation's identity-relevant content.""" + canonical = json.dumps({"kind": kind, "payload": payload or {}}, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(canonical.encode()).hexdigest() + + +@dataclass +class SealedGap: + """Gap-timeout filler for a never-arrived ordinal: liveness, fence kept; never executes, poison-neutral.""" + + operation_id: str + ordinal: int + + +@dataclass +class Operation: + operation_id: str + name: str + registration_id: str + # Consecutive from 1 per registration; arrival may be out of order. + ordinal: int + kind: OperationKind + payload: dict = field(default_factory=dict) + fingerprint: str = "" + state: OperationState = OperationState.QUEUED + result: dict | None = None + error: str | None = None + # "user" (bad request / cancelled by lifecycle) or "server" (execution failure). + error_category: str | None = None + was_claimed: bool = False + window_consumed: bool = False + # Monotonic stamp of the QUEUED->CLAIMED transition; the claimed-TTL sweep ages against it. + claimed_at: float | None = None + + @property + def tenant(self) -> Tenant: + return (self.name, self.registration_id) + + @property + def terminal(self) -> bool: + return self.state in TERMINAL_STATES + + def view(self) -> dict: + return dict( + operation_id=self.operation_id, + name=self.name, + registration_id=self.registration_id, + ordinal=self.ordinal, + kind=self.kind.value, + state=self.state.value, + result=self.result, + error=self.error, + error_category=self.error_category, + ) + + def claimed_view(self) -> dict: + """Executor-facing view: the request payload rides only on claims + (forward_backward samples, adam_params, save/load targets) so poll + results stay lean.""" + return {**self.view(), "payload": self.payload} + + +@dataclass +class _RegistrationQueue: + """Ordinal-sorted operations of one registration, pending and terminal.""" + + operations: list[Operation] = field(default_factory=list) + by_ordinal: dict[int, "Operation | SealedGap"] = field(default_factory=dict) + fenced: bool = False + # Cached contiguity frontier; ordinals are never removed, so it only advances. + _contiguous: int = 0 + # Gap-stall clock: missing ordinal blocking the queue and when first observed; a new hole restarts it. + _stall_missing: int | None = None + _stall_since: float | None = None + + def insert(self, op: Operation) -> None: + insort(self.operations, op, key=lambda o: o.ordinal) + self.by_ordinal[op.ordinal] = op + + def contiguous_arrived(self) -> int: + """Largest K such that ordinals 1..K have all arrived.""" + k = self._contiguous + while (k + 1) in self.by_ordinal: + k += 1 + self._contiguous = k + return k + + def fills_blocking_gap(self, ordinal: int) -> bool: + if not self.operations or ordinal >= self.operations[-1].ordinal: + return False + return ordinal == self.contiguous_arrived() + 1 + + def first_open(self) -> Operation | None: + for op in self.operations: + if not op.terminal: + return op if op.ordinal <= self.contiguous_arrived() else None + return None + + def open_count(self) -> int: + return sum(1 for op in self.operations if not op.terminal) + + def unacked_terminal_count(self) -> int: + return sum(1 for op in self.operations if op.terminal) + + def gap_stall(self, now: float) -> tuple[int, float] | None: + if self.fenced or self.first_open() is not None or self.open_count() == 0: + self._stall_missing = self._stall_since = None + return None + missing = self.contiguous_arrived() + 1 + if self._stall_missing != missing: + self._stall_missing, self._stall_since = missing, now + return missing, now - self._stall_since + + +class OperationLedger: + """All registrations' queues plus the operation_id index.""" + + def __init__( + self, + max_pending: int = 256, + max_unacked_results: int = 4096, + gap_timeout: float | None = 600.0, + claimed_ttl: float | None = 1800.0, + time_fn=time.monotonic, + ) -> None: + self.max_pending = max_pending + self.max_unacked_results = max_unacked_results + self.gap_timeout = gap_timeout + self.claimed_ttl = claimed_ttl + self._time = time_fn + self.queues: dict[Tenant, _RegistrationQueue] = {} + self.by_id: dict[str, Operation] = {} + + # ------------------------------ enqueue ------------------------------ + + def enqueue( + self, + operation_id: str, + name: str, + registration_id: str, + ordinal: int, + kind: str, + payload: dict | None = None, + ) -> dict: + """Buffer one operation; idempotent on (operation_id, fingerprint).""" + fingerprint = payload_fingerprint(kind, payload) + if (existing := self.by_id.get(operation_id)) is not None: + if ( + existing.fingerprint != fingerprint + or existing.tenant != (name, registration_id) + or existing.ordinal != ordinal + ): + raise ValueError( + f"operation '{operation_id}' already exists with different content; " + "retries must resend the identical request" + ) + return existing.view() + + queue = self.queues.setdefault((name, registration_id), _RegistrationQueue()) + if queue.fenced: + raise ValueError(f"registration '{name}' ({registration_id[:8]}) is retired; operations are fenced") + if ordinal < 1: + raise ValueError(f"operation '{operation_id}' ordinal must be >= 1, got {ordinal}") + if (holder := queue.by_ordinal.get(ordinal)) is not None: + raise ValueError( + f"ordinal {ordinal} already taken by operation '{holder.operation_id}'; " + "per-registration ordinals are unique and consecutive" + ) + if queue.open_count() >= self.max_pending and not queue.fills_blocking_gap(ordinal): + raise OperationBackpressure(f"registration '{name}' has {self.max_pending} operations pending") + if queue.unacked_terminal_count() >= self.max_unacked_results: + raise OperationBackpressure( + f"registration '{name}' holds {self.max_unacked_results} unacknowledged results; ack or deregister" + ) + + op = Operation( + operation_id=operation_id, + name=name, + registration_id=registration_id, + ordinal=ordinal, + kind=OperationKind(kind), + payload=payload or {}, + fingerprint=fingerprint, + ) + queue.insert(op) + self.by_id[operation_id] = op + return op.view() + + def record_rejected( + self, + operation_id: str, + name: str, + registration_id: str, + ordinal: int, + kind: str, + payload: dict | None, + error: str, + ) -> dict: + """Consume an ordinal with an operation born terminal FAILED(user). + + A submission rejected at the boundary (bad payload, unsupported + feature) must still fill its slot in the arrival sequence: the client + has already spent the ordinal and moved on, so refusing to record it + would leave a gap no retry ever fills — every later operation of the + registration would buffer forever. Identity rules match ``enqueue`` + (idempotent on an identical retry, conflict on anything else). The + record bypasses the PENDING cap (terminal on arrival, it never occupies + execution capacity) but still answers to the unacked-results budget — + born-terminal records hold result memory, and an invalid-request flood + must backpressure like any other unretrieved pile-up. The one exception + is a true hole-filler, whose refusal could never clear (the buffered + tail above it stays unclaimable, so no capacity would ever free).""" + fingerprint = payload_fingerprint(kind, payload) + if (existing := self.by_id.get(operation_id)) is not None: + if ( + existing.fingerprint != fingerprint + or existing.tenant != (name, registration_id) + or existing.ordinal != ordinal + ): + raise ValueError( + f"operation '{operation_id}' already exists with different content; " + "retries must resend the identical request" + ) + return existing.view() + + queue = self.queues.setdefault((name, registration_id), _RegistrationQueue()) + if queue.fenced: + raise ValueError(f"registration '{name}' ({registration_id[:8]}) is retired; operations are fenced") + if ordinal < 1: + raise ValueError(f"operation '{operation_id}' ordinal must be >= 1, got {ordinal}") + if (holder := queue.by_ordinal.get(ordinal)) is not None: + raise ValueError( + f"ordinal {ordinal} already taken by operation '{holder.operation_id}'; " + "per-registration ordinals are unique and consecutive" + ) + if queue.unacked_terminal_count() >= self.max_unacked_results and not queue.fills_blocking_gap(ordinal): + raise OperationBackpressure( + f"registration '{name}' holds {self.max_unacked_results} unacknowledged results; ack or deregister" + ) + op = Operation( + operation_id=operation_id, + name=name, + registration_id=registration_id, + ordinal=ordinal, + kind=OperationKind(kind), + # The payload was rejected — only its fingerprint matters (retry identity). + payload={}, + fingerprint=fingerprint, + state=OperationState.FAILED, + error=error, + error_category="user", + ) + queue.insert(op) + self.by_id[operation_id] = op + return op.view() + + # ------------------------------ claims ------------------------------ + + def claim_data_operation(self, name: str, registration_id: str) -> dict | None: + queue = self.queues.get((name, registration_id)) + if queue is None: + return None + op = queue.first_open() + if op is None or op.state is not OperationState.QUEUED or op.kind not in DATA_KINDS: + return None + op.state = OperationState.CLAIMED + op.was_claimed = True + op.claimed_at = self._time() + return op.claimed_view() + + def claimable_control_tenants(self) -> list[Tenant]: + tenants = [] + for tenant, queue in self.queues.items(): + op = queue.first_open() + if op is not None and op.state is OperationState.QUEUED and op.kind in CONTROL_KINDS: + tenants.append(tenant) + return tenants + + def claim_control_operation( + self, name: str, registration_id: str, kinds: tuple[str, ...] | None = None + ) -> dict | None: + queue = self.queues.get((name, registration_id)) + if queue is None: + return None + op = queue.first_open() + if op is None or op.state is not OperationState.QUEUED or op.kind not in CONTROL_KINDS: + return None + if kinds is not None and op.kind.value not in kinds: + return None + op.state = OperationState.CLAIMED + op.was_claimed = True + op.claimed_at = self._time() + return op.claimed_view() + + def poisoned_window_blocker(self, name: str, registration_id: str, ordinal: int) -> str | None: + queue = self.queues.get((name, registration_id)) + if queue is None: + return None + for o in range(ordinal - 1, 0, -1): + op = queue.by_ordinal.get(o) + if op is None or isinstance(op, SealedGap): + continue + if op.kind is OperationKind.OPTIM_STEP and op.was_claimed and op.terminal and op.window_consumed: + return None + if op.kind is OperationKind.FORWARD_BACKWARD and op.terminal and op.state is not OperationState.SUCCEEDED: + return f"forward_backward ordinal {o} {op.state.value}: {op.error or 'failed'}" + return None + + # ------------------------------ gap stalls ------------------------------ + # A consumed ordinal can fail client-side before HTTP; no retry fills it — blocked ops terminal-fail, hole sealed. + + def gap_stalls(self, now: float | None = None) -> list[dict]: + """Current stalls (observability): registrations whose open operations + are all buffered above an arrival hole, with the hole's ordinal, its + age, and the number of operations blocked behind it.""" + now = self._time() if now is None else now + stalls = [] + for (name, registration_id), queue in self.queues.items(): + stall = queue.gap_stall(now) + if stall is not None: + missing, stalled_for = stall + stalls.append( + dict( + name=name, + registration_id=registration_id, + missing_ordinal=missing, + stalled_for=stalled_for, + blocked_operations=queue.open_count(), + ) + ) + return stalls + + def sweep_gap_timeouts(self, now: float | None = None) -> list[dict]: + now = self._time() if now is None else now + if self.gap_timeout is None or self.gap_timeout <= 0: + for queue in self.queues.values(): # keep stall clocks observable + queue.gap_stall(now) + return [] + events = [] + for stall in self.gap_stalls(now): + if stall["stalled_for"] >= self.gap_timeout: + events.append(self._expire_stall(stall)) + return events + + def _expire_stall(self, stall: dict) -> dict: + queue = self.queues[(stall["name"], stall["registration_id"])] + missing, stalled_for = stall["missing_ordinal"], stall["stalled_for"] + last_arrived = max(queue.by_ordinal) + sealed = [] + for ordinal in range(missing, last_arrived): + if ordinal not in queue.by_ordinal: + queue.by_ordinal[ordinal] = SealedGap( + operation_id=f"{stall['name']}:gap-sealed:{ordinal}", ordinal=ordinal + ) + sealed.append(ordinal) + failed = [] + for op in queue.operations: + if not op.terminal: # all QUEUED: nothing is claimable while the queue stalls + op.state = OperationState.FAILED + op.error = ( + f"gap timeout: stalled {stalled_for:.0f}s behind missing ordinal {missing}; resubmit as new ops" + ) + op.error_category = "user" + failed.append(op.operation_id) + queue._stall_missing = queue._stall_since = None + event = {**stall, "sealed_ordinals": sealed, "failed_operations": failed} + logger.warning( + f"[tinker] gap timeout on '{stall['name']}' ({stall['registration_id'][:8]}): ordinal {missing} " + f"never arrived in {stalled_for:.0f}s; sealed {sealed}, failed {failed}" + ) + return event + + # ------------------------------ claimed TTL ------------------------------ + + def claimed_timeouts(self, now: float | None = None) -> list[dict]: + """Over-age CLAIMED operations for the backend to terminal-fail (an orphaned claim blocks its queue forever).""" + now = self._time() if now is None else now + if self.claimed_ttl is None or self.claimed_ttl <= 0: + return [] + return [ + {**op.view(), "claimed_age": now - op.claimed_at} + for op in self.by_id.values() + if op.state is OperationState.CLAIMED + and op.claimed_at is not None + and now - op.claimed_at >= self.claimed_ttl + ] + + # ------------------------------ terminals ------------------------------ + + def complete(self, operation_id: str, result: dict | None = None) -> None: + op = self._open_op(operation_id) + op.state = OperationState.SUCCEEDED + op.result = result + + def fail(self, operation_id: str, error: str, category: str = "server") -> None: + op = self._open_op(operation_id) + op.state = OperationState.FAILED + op.error = error + op.error_category = category + + def mark_window_consumed(self, operation_id: str) -> None: + op = self.by_id.get(operation_id) + if op is not None: + op.window_consumed = True + + def cancel(self, operation_id: str) -> dict: + op = self.by_id.get(operation_id) + if op is None: + raise KeyError(f"unknown operation '{operation_id}'") + if op.state is not OperationState.QUEUED: + raise ValueError(f"operation '{operation_id}' is {op.state.value}; only QUEUED operations cancel") + op.state = OperationState.CANCELLED + op.error = "cancelled by client" + op.error_category = "user" + return op.view() + + def _open_op(self, operation_id: str) -> Operation: + op = self.by_id.get(operation_id) + if op is None: + raise KeyError(f"unknown operation '{operation_id}'") + if op.terminal: + raise ValueError(f"operation '{operation_id}' already terminal ({op.state.value})") + return op + + # ------------------------------ results ------------------------------ + + def get(self, operation_id: str) -> dict | None: + op = self.by_id.get(operation_id) + return op.view() if op is not None else None + + def payload(self, operation_id: str) -> dict | None: + """The stored request payload (metrics recomputation at completion).""" + op = self.by_id.get(operation_id) + return op.payload if op is not None else None + + def ack(self, operation_id: str) -> None: + op = self.by_id.get(operation_id) + if op is None: + return + if not op.terminal: + raise ValueError(f"operation '{operation_id}' is {op.state.value}; ack applies to terminal operations") + self.by_id.pop(operation_id, None) + # The ordinal slot stays reserved (contiguity/uniqueness), but an acked + # record's payload and result are released — they can be large. + op.payload = {} + op.result = None + queue = self.queues.get(op.tenant) + if queue is not None: + queue.operations = [o for o in queue.operations if o.operation_id != operation_id] + # by_ordinal keeps the slot so contiguity and ordinal uniqueness survive the ack. + if not queue.operations and queue.fenced: + self.queues.pop(op.tenant, None) + + # ------------------------------ fencing ------------------------------ + + def fence(self, name: str, registration_id: str) -> list[str]: + queue = self.queues.get((name, registration_id)) + if queue is None or queue.fenced: + return [] + queue.fenced = True + failed = [] + for op in queue.operations: + if not op.terminal: + op.state = OperationState.FAILED + op.error = "registration retired before the operation ran" + op.error_category = "user" + failed.append(op.operation_id) + # Fenced ops are never claimed: release the payload now (the fingerprint alone carries retry identity). + op.payload = {} + return failed + + def drop_tenant(self, name: str, registration_id: str) -> None: + """Purge a dead registration the registry evicted from its completed ring; its results stop being pollable.""" + queue = self.queues.pop((name, registration_id), None) + if queue is None: + return + for op in queue.operations: + self.by_id.pop(op.operation_id, None) + + def queue_view(self, name: str, registration_id: str) -> list[dict]: + queue = self.queues.get((name, registration_id)) + return [op.view() for op in queue.operations] if queue is not None else [] diff --git a/miles/ray/multi_lora/registry.py b/miles/ray/multi_lora/registry.py index 4c8723c29d7..77b84189691 100644 --- a/miles/ray/multi_lora/registry.py +++ b/miles/ray/multi_lora/registry.py @@ -1,28 +1,28 @@ -"""Multi-LoRA adapter registry: the controller-owned lifecycle state machine. - -One record per adapter name, walking PENDING -> ACTIVE -> RETIRING -> CLEANUP --> COMPLETED. Slots are reused across registrations but ``slot_versions`` -never reset, so a (slot, version) pair never recurs. -""" +"""Controller-owned Multi-LoRA run lifecycle under fixed slot residency. +Serving identity includes the registration ID to prevent same-name aliasing.""" import logging import re import uuid -from dataclasses import dataclass, field, replace +from collections.abc import Callable +from dataclasses import dataclass, field from enum import Enum from pathlib import Path from typing import Any -from miles.utils.adapter_config import AdapterRun, AdapterRunConfig +from miles.ray.multi_lora.config import AdapterRun +from miles.ray.multi_lora.slot_pool import SlotPool logger = logging.getLogger(__name__) VALID_ADAPTER_NAME = re.compile(r"^[A-Za-z0-9._-]+$") +DIRTY_PIN = "dirty-grads" + class AdapterState(str, Enum): PENDING = "PENDING" - ACTIVE = "ACTIVE" + READY = "READY" RETIRING = "RETIRING" CLEANUP = "CLEANUP" COMPLETED = "COMPLETED" @@ -31,43 +31,41 @@ class AdapterState(str, Enum): # States that hold a slot. LIVE_STATES = ( AdapterState.PENDING, - AdapterState.ACTIVE, + AdapterState.READY, AdapterState.RETIRING, AdapterState.CLEANUP, ) +MAX_COMPLETED_RECORDS = 1024 + @dataclass class AdapterRecord: name: str - slot: int - config: Any + config: Any = None + # Bound trainer slot; None while queued behind a full pool. + slot: int | None = None step: int = 0 - # Baseline step for relative num_step stopping (supports checkpoint resume). + # Baseline step for the relative num_step bound (supports state resume). start_step: int = 0 - # Committed prompt groups accumulated toward the current optimizer step. - # Only advanced by mark_batch_trained (after a successful train call). - accumulated_groups: int = 0 + serving_version: int = 0 state: AdapterState = AdapterState.PENDING - # Unique per registration: a re-registered name is a new tenant, and - # rollout-side state stamped by the previous tenant must not carry over. registration_id: str = field(default_factory=lambda: uuid.uuid4().hex) - -MAX_BATCH_RECORDS = 16 -MAX_COMPLETED_RECORDS = 1024 + @property + def tenant(self) -> tuple[str, str]: + return (self.name, self.registration_id) class AdapterRegistry: - """One record per name; ``slot_versions`` never reset, so (slot, version) - never recurs across slot reuse.""" + """One record per name; slot tenancy delegated to the SlotPool.""" def __init__(self, max_adapters: int) -> None: self.max_adapters = max_adapters - self.free_slots: set[int] = set(range(max_adapters)) - self.slot_versions: list[int] = [0] * max_adapters + self.slot_pool = SlotPool(max_adapters) self.records: dict[str, AdapterRecord] = {} - self.batch_records: dict[int, dict] = {} + # Fires (name, registration_id) when a COMPLETED record leaves the ring; the backend wires ledger purging. + self.on_completed_evicted: Callable[[str, str], None] | None = None def in_state(self, *states: AdapterState) -> dict[str, AdapterRecord]: return {name: r for name, r in self.records.items() if r.state in states} @@ -76,15 +74,13 @@ def find(self, name: str) -> AdapterRecord | None: record = self.records.get(name) return record if record is not None and record.state in LIVE_STATES else None - def is_active(self, name: str) -> bool: - record = self.records.get(name) - return record is not None and record.state in (AdapterState.ACTIVE, AdapterState.RETIRING) + # ---------------------- registration lifecycle ---------------------- def register(self, name: str, config: Any) -> dict: if not VALID_ADAPTER_NAME.match(name) or name in (".", ".."): raise ValueError(f"Adapter name '{name}' is invalid: use only letters, digits, '.', '_' and '-'") if (existing := self.records.get(name)) is not None: - if existing.state in (AdapterState.PENDING, AdapterState.ACTIVE): + if existing.state in (AdapterState.PENDING, AdapterState.READY): raise ValueError(f"Adapter '{name}' already registered") if existing.state in (AdapterState.RETIRING, AdapterState.CLEANUP): raise ValueError(f"Adapter '{name}' is still cleaning up; retry shortly") @@ -95,17 +91,39 @@ def register(self, name: str, config: Any) -> dict: raise ValueError( f"Adapter '{name}' save dir '{save_dir}' is already used by adapter '{record.name}'" ) - if not self.free_slots: - raise RuntimeError(f"No free adapter slots (max {self.max_adapters})") - slot = min(self.free_slots) - self.free_slots.remove(slot) - self.records.pop(name, None) - self.records[name] = AdapterRecord(name=name, slot=slot, config=config) - return {"name": name, "slot": slot} + record = AdapterRecord(name=name, config=config) + # Fixed residency: a full pool queues the registration unbound; + # bootstrap_pending binds it when a slot frees at retirement. + record.slot = self.slot_pool.bind_immediately(record.tenant) + if name in self.records: + self._evict_completed(name) + self.records[name] = record + if record.slot is None: + logger.info(f"[tinker] adapter '{name}' queued unbound: all {self.max_adapters} slots busy") + return {"name": name, "slot": record.slot} + + def bootstrap_pending(self) -> list[str]: + bound = [] + for name, record in self.in_state(AdapterState.PENDING).items(): + if record.slot is not None: + continue + slot = self.slot_pool.bind_immediately(record.tenant) + if slot is None: + break + record.slot = slot + bound.append(name) + logger.info(f"[tinker] adapter '{name}' bound to freed slot {slot}") + return bound + + def mark_ready(self, names: list[str]) -> None: + for name in names: + record = self.find(name) + if record is not None and record.state is AdapterState.PENDING and record.slot is not None: + record.state = AdapterState.READY def deregister(self, name: str) -> None: record = self.records.get(name) - if record is not None and record.state in (AdapterState.PENDING, AdapterState.ACTIVE): + if record is not None and record.state in (AdapterState.PENDING, AdapterState.READY): record.state = AdapterState.RETIRING def retire_adapters(self) -> list[str]: @@ -118,14 +136,19 @@ def free_slot(self, name: str) -> int: record = self.records.get(name) if record is None or record.state is not AdapterState.CLEANUP: return -1 - self.free_slots.add(record.slot) + self.slot_pool.release(record.tenant) record.state = AdapterState.COMPLETED self.records[name] = self.records.pop(name) completed = self.in_state(AdapterState.COMPLETED) - for oldest in list(completed)[: len(completed) - MAX_COMPLETED_RECORDS]: - self.records.pop(oldest) + for oldest in list(completed)[: max(0, len(completed) - MAX_COMPLETED_RECORDS)]: + self._evict_completed(oldest) return record.slot + def _evict_completed(self, name: str) -> None: + evicted = self.records.pop(name) + if self.on_completed_evicted is not None: + self.on_completed_evicted(evicted.name, evicted.registration_id) + def adapter_state(self, name: str) -> AdapterState | None: record = self.records.get(name) if record is None: @@ -134,109 +157,71 @@ def adapter_state(self, name: str) -> AdapterState | None: self.records[name] = self.records.pop(name) return record.state + # ---------------------- clocks and serving ---------------------- + def record_weight_update(self, names: list[str]) -> None: - """A weight push landed: bump slot versions, promote PENDING to ACTIVE.""" + """A weight push landed on the engines: bump the serving version. + Publication is orthogonal to readiness (no state promotion here).""" for name in names: record = self.find(name) - if record is None: - continue - self.slot_versions[record.slot] += 1 - if record.state is AdapterState.PENDING: - record.state = AdapterState.ACTIVE - - def record_batch_adapters(self, rollout_id: int, groups: dict[str, int], step_names: list[str]) -> None: - """Register what a train batch contains before it trains. - - ``groups`` maps adapter name -> prompt groups riding in this batch; - ``step_names`` lists adapters whose adapter batch completes with - this batch (decided by the collection loop, which caps per-adapter - contributions at the adapter's remaining groups). - """ - unknown = set(step_names) - set(groups) - assert not unknown, f"step adapters {sorted(unknown)} not present in batch groups" - self.batch_records[rollout_id] = {"groups": dict(groups), "step_names": list(step_names)} - while len(self.batch_records) > MAX_BATCH_RECORDS: - self.batch_records.pop(next(iter(self.batch_records))) - - def mark_batch_trained(self, rollout_id: int) -> list[str]: - """Bank the batch's trained groups and fire steps; returns adapters that stepped. Only place - accumulation/step state advances, so a failed/retried train call leaves the registry untouched.""" - record_entry = self.batch_records.pop(rollout_id, None) - if record_entry is None: - return [] - stepped = [] - reached_num_step = [] - for name, n_groups in record_entry["groups"].items(): - record = self.records.get(name) - if record is None or record.state not in ( - AdapterState.ACTIVE, - AdapterState.RETIRING, - AdapterState.CLEANUP, - ): - continue - record.accumulated_groups += n_groups - if name in record_entry["step_names"]: - target = record.config.rollout_batch_size - if record.accumulated_groups != target: - logger.warning( - f"Adapter '{name}' stepped with accumulated_groups={record.accumulated_groups} " - f"!= rollout_batch_size={target}; adapter batch accounting drifted" - ) - record.step += 1 - record.accumulated_groups = 0 - stepped.append(name) - if ( - getattr(record.config, "num_step", None) is not None - and record.state is AdapterState.ACTIVE - and (record.step - record.start_step) >= record.config.num_step - ): - reached_num_step.append(name) - for name in reached_num_step: - logger.info( - f"Adapter '{name}' reached num_step={self.records[name].config.num_step} " - f"(start_step={self.records[name].start_step}, step={self.records[name].step}), deregistering" - ) - self.deregister(name) - return stepped + if record is not None: + record.serving_version += 1 - def resolve_num_step(self, name: str, dataset_rows: int) -> None: - """Derive num_step from num_epoch once the data source knows the - post-filter dataset length. No-op when num_step was set explicitly.""" + def on_step_committed(self, name: str, registration_id: str, step: int) -> None: record = self.find(name) - if record is None or not isinstance(record.config, AdapterRunConfig): + if record is None or record.registration_id != registration_id: return - if record.config.num_step is not None: - return - num_epoch = record.config.num_epoch or 1 - num_step = max(1, num_epoch * dataset_rows // record.config.rollout_batch_size) - record.config = replace(record.config, num_step=num_step) - logger.info(f"Adapter '{name}': num_epoch={num_epoch} x {dataset_rows} rows -> num_step={num_step}") + record.step = step + self.slot_pool.unpin(record.tenant, DIRTY_PIN) + if ( + getattr(record.config, "num_step", None) is not None + and record.state is AdapterState.READY + and (record.step - record.start_step) >= record.config.num_step + ): + logger.info(f"[tinker] adapter '{name}' reached num_step={record.config.num_step}, deregistering") + self.deregister(name) def set_step(self, name: str, step: int) -> None: + """Mirror hook: a restore (load_state / sidecar resume) repositioned + the stream's clock and its num_step baseline.""" if (record := self.find(name)) is not None: record.step = step record.start_step = step - def step_count(self, name: str) -> int: + # ---------------------- gradient-state pins ---------------------- + + def mark_accumulated(self, names: list[str]) -> None: + for name in names: + record = self.find(name) + if record is not None: + self.slot_pool.pin(record.tenant, DIRTY_PIN) + + def clear_dirty(self, name: str) -> None: record = self.find(name) - return record.step if record is not None else 0 + if record is not None: + self.slot_pool.unpin(record.tenant, DIRTY_PIN) + + def is_dirty(self, name: str) -> bool: + record = self.find(name) + return record is not None and self.slot_pool.is_pinned(record.tenant, DIRTY_PIN) + + # ---------------------- views ---------------------- def view(self, record: AdapterRecord) -> AdapterRun: return AdapterRun( name=record.name, config=record.config, slot=record.slot, - version=self.slot_versions[record.slot], + version=record.serving_version, step=record.step, - accumulated_groups=record.accumulated_groups, registration_id=record.registration_id, ) - def active_adapters(self) -> dict[str, AdapterRun]: - """Sampleable view: RETIRING keeps serving until retired.""" + def ready_adapters(self) -> dict[str, AdapterRun]: + """Operation-executable view: RETIRING keeps draining until retired.""" return { name: self.view(record) - for name, record in self.in_state(AdapterState.ACTIVE, AdapterState.RETIRING).items() + for name, record in self.in_state(AdapterState.READY, AdapterState.RETIRING).items() } def snapshot(self) -> dict: @@ -245,7 +230,7 @@ def views(state: AdapterState) -> dict[str, AdapterRun]: return { "pending": views(AdapterState.PENDING), - "active": views(AdapterState.ACTIVE), + "ready": views(AdapterState.READY), "retiring": views(AdapterState.RETIRING), "cleanup": list(self.in_state(AdapterState.CLEANUP)), "completed": list(self.in_state(AdapterState.COMPLETED)), diff --git a/miles/ray/multi_lora/residency.py b/miles/ray/multi_lora/residency.py new file mode 100644 index 00000000000..76913bb1ab6 --- /dev/null +++ b/miles/ray/multi_lora/residency.py @@ -0,0 +1,86 @@ +import logging +import uuid +from dataclasses import dataclass + +from miles.ray.multi_lora.registry import AdapterRegistry, AdapterState +from miles.utils.operation_contract import BatchExecutionLease, RegistrationKey + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class ResidentBinding: + """Multi-LoRA execution binding (registration -> fixed trainer slot); opaque above the residency port.""" + + registration_key: RegistrationKey + training_slot: int + + +class FixedSlotResidency: + """TrainerResidencyPort[ResidentBinding] over the adapter registry.""" + + def __init__(self, registry: AdapterRegistry) -> None: + self.registry = registry + + def binding_for(self, key: RegistrationKey) -> ResidentBinding | None: + name, registration_id = key + record = self.registry.find(name) + if ( + record is None + or record.registration_id != registration_id + or record.state is not AdapterState.READY + or record.slot is None + ): + return None + return ResidentBinding(registration_key=key, training_slot=record.slot) + + def acquire_batch( + self, bindings_by_operation: tuple[tuple[str, ResidentBinding], ...] + ) -> BatchExecutionLease[ResidentBinding]: + for operation_id, binding in bindings_by_operation: + if not self._owns_slot(binding): + raise ValueError( + f"operation '{operation_id}': registration " + f"{binding.registration_key} no longer owns trainer slot {binding.training_slot}" + ) + return BatchExecutionLease( + dispatch_id=uuid.uuid4().hex, + bindings_by_operation=tuple(bindings_by_operation), + ) + + def release_batch(self, lease: BatchExecutionLease[ResidentBinding]) -> None: + """No-op lifecycle hook (nothing to free under fixed residency).""" + + def _owns_slot(self, binding: ResidentBinding) -> bool: + name, registration_id = binding.registration_key + record = self.registry.records.get(name) + return ( + record is not None + and record.registration_id == registration_id + and record.slot == binding.training_slot + and record.state in (AdapterState.READY, AdapterState.RETIRING) + ) + + +# ---------------- data-plane encoding ---------------- +# The lease crosses rollout -> store -> trainer as plain data; typed leases stay at controller/adapter boundaries. + + +def lease_to_metadata(lease: BatchExecutionLease[ResidentBinding]) -> dict: + return { + "dispatch_id": lease.dispatch_id, + "bindings_by_operation": [ + [op_id, [binding.registration_key[0], binding.registration_key[1], binding.training_slot]] + for op_id, binding in lease.bindings_by_operation + ], + } + + +def lease_from_metadata(data: dict) -> BatchExecutionLease[ResidentBinding]: + return BatchExecutionLease( + dispatch_id=data["dispatch_id"], + bindings_by_operation=tuple( + (op_id, ResidentBinding(registration_key=(name, registration_id), training_slot=slot)) + for op_id, (name, registration_id, slot) in data["bindings_by_operation"] + ), + ) diff --git a/miles/ray/multi_lora/slot_pool.py b/miles/ray/multi_lora/slot_pool.py new file mode 100644 index 00000000000..43da6ed118f --- /dev/null +++ b/miles/ray/multi_lora/slot_pool.py @@ -0,0 +1,67 @@ +from dataclasses import dataclass, field + +# (adapter name, registration id): a re-registered name is a different tenant. +Tenant = tuple[str, str] + + +@dataclass +class SlotEntry: + slot: int + tenant: Tenant | None = None + # Non-empty pins mark the slot's state as immovable (e.g. "dirty-grads": + # accumulated gradients that no checkpoint carries). + pins: set = field(default_factory=set) + + +class SlotPool: + def __init__(self, n_slots: int) -> None: + self.entries = [SlotEntry(slot=i) for i in range(n_slots)] + + # -------------------------- queries -------------------------- + + def entry_of(self, tenant: Tenant) -> SlotEntry | None: + for entry in self.entries: + if entry.tenant == tenant: + return entry + return None + + def free_slot_ids(self) -> set[int]: + return {e.slot for e in self.entries if e.tenant is None} + + def occupied_slot_ids(self) -> list[int]: + return sorted(e.slot for e in self.entries if e.tenant is not None) + + def is_pinned(self, tenant: Tenant, reason: str) -> bool: + entry = self.entry_of(tenant) + return entry is not None and reason in entry.pins + + # ---------------------- tenancy ---------------------- + + def bind_immediately(self, tenant: Tenant) -> int | None: + """Bind to the lowest free slot; None when the pool is full (the + registration queues until another tenant releases).""" + free = [e for e in self.entries if e.tenant is None] + if not free: + return None + entry = free[0] + entry.tenant = tenant + return entry.slot + + def release(self, tenant: Tenant) -> int | None: + """Return the tenant's slot to the free pool (retirement path).""" + entry = self.entry_of(tenant) + if entry is None: + return None + entry.tenant = None + entry.pins.clear() + return entry.slot + + # -------------------------- pins -------------------------- + + def pin(self, tenant: Tenant, reason: str) -> None: + if (entry := self.entry_of(tenant)) is not None: + entry.pins.add(reason) + + def unpin(self, tenant: Tenant, reason: str) -> None: + if (entry := self.entry_of(tenant)) is not None: + entry.pins.discard(reason) diff --git a/miles/ray/rollout/components.py b/miles/ray/rollout/components.py new file mode 100644 index 00000000000..e2a47d2588a --- /dev/null +++ b/miles/ray/rollout/components.py @@ -0,0 +1,95 @@ +"""Role-separated rollout construction: PR #1842 role names now, Legacy adapters over one combined RolloutManager.""" + +from dataclasses import dataclass +from typing import Protocol + + +@dataclass(frozen=True) +class InferenceEndpoint: + """Where sampling requests go (the SGLang router).""" + + host: str + port: int + + @property + def base_url(self) -> str: + return f"http://{self.host}:{self.port}" + + +class InferenceControllerPort(Protocol): + async def get_inference_endpoint(self) -> InferenceEndpoint: ... + + async def prepare_rollout(self, rollout_id: int) -> None: + """Called before every generate; no-op in the legacy adapter (PR #1842 moves engine preparation here).""" + ... + + +class RolloutExecutorPort(Protocol): + async def generate(self, rollout_id: int): ... + + +class RolloutLifecyclePort(Protocol): + async def dispose_once(self) -> None: ... + + +class LegacyInferenceControllerAdapter: + """Inference-owner role view over the combined RolloutManager; the raw handle rides only weight_update_owner.""" + + def __init__(self, manager) -> None: + self._manager = manager + + async def get_inference_endpoint(self) -> InferenceEndpoint: + host, port = await self._manager.get_router_address.remote() + return InferenceEndpoint(host=host, port=port) + + async def prepare_rollout(self, rollout_id: int) -> None: + """No-op: the combined manager prepares inside generate(); PR #1842 moves that preparation here.""" + + +class LegacyRolloutExecutorAdapter: + """Execution role view over the same combined RolloutManager.""" + + def __init__(self, manager) -> None: + self._manager = manager + + async def generate(self, rollout_id: int): + return await self._manager.generate.remote(rollout_id) + + +class LegacyRolloutLifecycle: + """Exactly-once disposal of the SHARED underlying actor: two role views must never each dispose it.""" + + def __init__(self, manager) -> None: + self._manager = manager + self._disposed = False + + async def dispose_once(self) -> None: + if self._disposed: + return + self._disposed = True + await self._manager.dispose.remote() + + +@dataclass +class RolloutComponents: + inference_controller: InferenceControllerPort + rollout_executor: RolloutExecutorPort + lifecycle: RolloutLifecyclePort + # Opaque weight-update owner/target (today the combined manager handle); passed verbatim, never introspected. + weight_update_owner: object + + async def dispose(self) -> None: + await self.lifecycle.dispose_once() + + +def create_rollout_components(args, pg) -> RolloutComponents: + """One construction seam: legacy manager + role views today, PR #1842's pair later; call sites never change.""" + from miles.ray.placement_group import create_rollout_manager + + rollout_manager, _num_rollout_per_epoch = create_rollout_manager(args, pg) + return RolloutComponents( + inference_controller=LegacyInferenceControllerAdapter(rollout_manager), + rollout_executor=LegacyRolloutExecutorAdapter(rollout_manager), + lifecycle=LegacyRolloutLifecycle(rollout_manager), + weight_update_owner=rollout_manager, + ) diff --git a/miles/ray/rollout/metrics.py b/miles/ray/rollout/metrics.py index 14261be89fc..cc9fe8acaaa 100644 --- a/miles/ray/rollout/metrics.py +++ b/miles/ray/rollout/metrics.py @@ -173,6 +173,11 @@ def _compute_zero_std_metrics(args, all_samples: list[Sample]): if args.advantage_estimator == "ppo": return {} + # Reward-less batches (e.g. tinker client operations) have no reward + # plane: zero-std over missing rewards is meaningless, not zero. + if any(sample.get_reward_value(args) is None for sample in all_samples): + return {} + def _is_zero_std(samples: list[Sample]): rewards = [sample.get_reward_value(args) for sample in samples] return len(rewards) == 0 or all(rewards[0] == r for r in rewards) diff --git a/miles/ray/rollout/rollout_data_conversion.py b/miles/ray/rollout/rollout_data_conversion.py index b948856dd3a..2857c9e852c 100644 --- a/miles/ray/rollout/rollout_data_conversion.py +++ b/miles/ray/rollout/rollout_data_conversion.py @@ -1,3 +1,4 @@ +import copy import itertools import logging @@ -7,24 +8,23 @@ logger = logging.getLogger(__name__) -def postprocess_rollout_data(args, data, train_parallel_config): +def postprocess_rollout_data(args, data, train_parallel_config, pad_to_dp: bool = False): metadata = {} validate_compact_rollout_ids(data) - # Multi-LoRA: record group boundaries (heterogeneous per-adapter group sizes) - # and lift the collection loop's batch-level step decision out of sample metadata, - # both before flattening. + # Multi-LoRA: record group boundaries (heterogeneous per-adapter group + # sizes) before flattening. if is_multi_lora_enabled(args) and isinstance(data[0], list): metadata["prompt_group_sizes"] = [_nested_sample_count(group) for group in data] - head = _first_sample(data[0]) - metadata["step_slots"] = list(head.metadata.pop("step_slots", [])) - metadata["step_adapter_names"] = list(head.metadata.pop("step_adapter_names", [])) # flatten the data if it is a list of lists while isinstance(data[0], list): data = list(itertools.chain.from_iterable(data)) + if pad_to_dp and (dp_size := (train_parallel_config or {}).get("dp_size")): + data = _pad_samples_to_dp(data, dp_size) + # Compact rollouts must not be trimmed by sample count; the schedule drops # whole trailing rollouts instead. is_compact = any(s.rollout_id is not None for s in data) @@ -72,16 +72,31 @@ def validate_compact_rollout_ids(node, depth=0): validate_compact_rollout_ids(item, depth + 1) -def _first_sample(group): - return _first_sample(group[0]) if isinstance(group[0], list) else group[0] - - def _nested_sample_count(group) -> int: if not isinstance(group, list): return 1 return sum(_nested_sample_count(item) for item in group) +def _pad_samples_to_dp(data: list[Sample], dp_size: int) -> list[Sample]: + deficit = -len(data) % dp_size + if deficit == 0: + return data + donor = data[-1] + padded = list(data) + for _ in range(deficit): + pad = copy.deepcopy(donor) + pad.index = -1 # sentinel: the result plane filters row < 0 + pad.rollout_id = None + pad.loss_mask = [0] * pad.response_length + for channel in ("loss_weights", "advantages"): + if getattr(pad, channel) is not None: + setattr(pad, channel, [0.0] * len(getattr(pad, channel))) + padded.append(pad) + logger.info(f"[tinker] padded batch from {len(data)} to {len(padded)} samples for DP alignment") + return padded + + def _compute_dynamic_global_batch_size(args, train_parallel_config, num_samples: int) -> int: """Calculate dynamic global_batch_size to ensure only one training step. @@ -92,13 +107,13 @@ def _compute_dynamic_global_batch_size(args, train_parallel_config, num_samples: original_gbs = args.global_batch_size if is_multi_lora_enabled(args): - # Batches take groups in multiples of each adapter's - # min_groups_per_dp_split, so this holds by construction; a violation - # means a generate fn's group shape broke the invariant. + # Multi-LoRA batches are built from whole prompt groups sized to split + # evenly across DP ranks; a violation means a generate fn's group + # shape broke that invariant. if num_samples % dp_size != 0: raise ValueError( f"Multi-LoRA batch of {num_samples} samples is not divisible by dp_size={dp_size}; " - "the min_groups_per_dp_split invariant was violated (variable-size generate fn output?)" + "whole prompt groups must split evenly across ranks (variable-size generate fn output?)" ) return num_samples diff --git a/miles/ray/rollout/rollout_manager.py b/miles/ray/rollout/rollout_manager.py index 2f74e4096fa..a4ac0056e2f 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -19,12 +19,14 @@ ROLLOUT_DATA_VALUE_SPEC, convert_samples_to_train_data, split_train_data_by_dp, + tinker_dispatch_summary, ) from miles.ray.utils import Lock from miles.rollout.base_types import ( RolloutFnConstructorInput, RolloutFnEvalInput, RolloutFnTrainInput, + RolloutPostprocessOptions, call_rollout_fn, ) from miles.rollout.checkpoint_eval import CheckpointEvalFn, EvalSkip @@ -162,11 +164,15 @@ async def generate(self, rollout_id): custom_reward_post_process_func=self.custom_reward_post_process_func, ) sample_indices = data.get("sample_indices") + dispatch = tinker_dispatch_summary(data) if self.args.delay_split_train_data_by_dp: data_ref = object_store.get_instance().put(value=data, value_spec=ROLLOUT_DATA_VALUE_SPEC) else: data_ref = split_train_data_by_dp(self.args, data, self.train_parallel_config) - return dict(sample_indices=sample_indices, data_ref=data_ref) + if dispatch is not None: + return dict(sample_indices=sample_indices, data_ref=data_ref, tinker_dispatch=dispatch) + else: + return dict(sample_indices=sample_indices, data_ref=data_ref) async def eval( self, @@ -256,10 +262,16 @@ async def _get_rollout_data(self, rollout_id): call_rollout_fn, self.generate_rollout, self.args, rollout_id, self.data_source, evaluation=False ) metrics = data.metrics + conversion_metadata = getattr(data, "conversion_metadata", None) or {} + postprocess = getattr(data, "postprocess", None) or RolloutPostprocessOptions() data = data.samples data, metadata = postprocess_rollout_data( - self.args, data, train_parallel_config=self.train_parallel_config + self.args, + data, + train_parallel_config=self.train_parallel_config, + pad_to_dp=postprocess.pad_to_dp, ) + metadata.update(conversion_metadata) if RolloutDataInjectionUtil.should_inject(self.args, rollout_id): generated_data = data data, metadata = RolloutDataInjectionUtil.load(self.args, rollout_id=rollout_id) diff --git a/miles/ray/rollout/train_data_conversion.py b/miles/ray/rollout/train_data_conversion.py index e3c95a92999..f3be9cb4b8c 100644 --- a/miles/ray/rollout/train_data_conversion.py +++ b/miles/ray/rollout/train_data_conversion.py @@ -19,6 +19,10 @@ "rollout_log_probs": "float32", "teacher_log_probs": "float32", "opd_reverse_kl": "float32", + # Client-supplied per-token channels (tinker adapters); the binary + # loss_masks stay int32, these carry the float semantics. + "loss_weights": "float32", + "advantages": "float32", "rollout_routed_experts": "int32", "rollout_indexer_topk": "int32", } @@ -60,12 +64,18 @@ def convert_samples_to_train_data( if (f := custom_convert_samples_to_train_data_func) is not None: return f(args, samples) - raw_rewards, rewards = _post_process_rewards( - args, - samples, - custom_reward_post_process_func=custom_reward_post_process_func, - prompt_group_sizes=metadata.get("prompt_group_sizes"), - ) + tinker = metadata.get("batch_kind") == "tinker" + if tinker: + # Tinker batches carry no rewards: losses come from client-supplied + # per-token channels, never from reward post-processing. + raw_rewards = rewards = [0.0] * len(samples) + else: + raw_rewards, rewards = _post_process_rewards( + args, + samples, + custom_reward_post_process_func=custom_reward_post_process_func, + prompt_group_sizes=metadata.get("prompt_group_sizes"), + ) assert len(raw_rewards) == len(samples) assert len(rewards) == len(samples) @@ -109,7 +119,12 @@ def convert_samples_to_train_data( train_data["round_number"] = [sample.metadata["round_number"] for sample in samples] # Add rollout log probabilities for off-policy correction - if samples[0].rollout_log_probs is not None: + if tinker and any(sample.rollout_log_probs is not None for sample in samples): + train_data["rollout_log_probs"] = [ + sample.rollout_log_probs if sample.rollout_log_probs is not None else [0.0] * sample.response_length + for sample in samples + ] + elif samples[0].rollout_log_probs is not None: train_data["rollout_log_probs"] = [sample.rollout_log_probs for sample in samples] if samples[0].rollout_routed_experts is not None: @@ -130,20 +145,37 @@ def convert_samples_to_train_data( if samples[0].teacher_log_probs is not None: train_data["teacher_log_probs"] = [sample.teacher_log_probs for sample in samples] + # Client-supplied per-token channels (tinker); absent tensors default to zeros so CE and IS/PPO adapters can mix. + if any(sample.loss_weights is not None for sample in samples): + train_data["loss_weights"] = [ + sample.loss_weights if sample.loss_weights is not None else [0.0] * sample.response_length + for sample in samples + ] + if any(sample.advantages is not None for sample in samples): + train_data["advantages"] = [ + sample.advantages if sample.advantages is not None else [0.0] * sample.response_length + for sample in samples + ] + + if tinker: + train_data["batch_kind"] = "tinker" + train_data["tinker_operation_lanes"] = _tinker_sample_lanes(metadata["tinker_operation_lanes"], len(samples)) + train_data["tinker_loss_by_lane"] = metadata["tinker_loss_by_lane"] + train_data["operation_by_lane"] = metadata["operation_by_lane"] + train_data["registration_by_lane"] = metadata["registration_by_lane"] + if (lease := metadata.get("batch_execution_lease")) is not None: + train_data["batch_execution_lease"] = lease + if metadata.get("tinker_forward_only"): + train_data["tinker_forward_only"] = True + if any(sample.adapter is not None for sample in samples): assert all(sample.adapter is not None for sample in samples), "Cannot mix adapter and adapter-less samples" - train_data["adapter_slots"] = [sample.adapter.slot for sample in samples] - # Slots whose adapter batch completes with this batch: the trainer scales their - # accumulated gradients by 1/adapter-batch-size and advances the LR schedule. - step_slots = sorted(metadata.get("step_slots", [])) - train_data["step_slots"] = step_slots - train_data["step_adapter_names"] = sorted(metadata.get("step_adapter_names", [])) - step_slot_set = set(step_slots) - train_data["step_adapter_batch_sizes"] = { - sample.adapter.slot: sample.metadata["adapter_global_batch_size"] - for sample in samples - if sample.adapter.slot in step_slot_set - } + # Adapter batches only come from the tinker rollout fn, whose lease is mandatory; stamped-slot fallback removed. + if not tinker or metadata.get("batch_execution_lease") is None: + raise ValueError("adapter-stamped batch without a tinker batch lease; BatchPlan slot routing is required") + train_data["adapter_slots"] = _adapter_slots_from_lease( + metadata, train_data["tinker_operation_lanes"], samples + ) if (prompt_group_sizes := metadata.get("prompt_group_sizes")) is not None: train_data["prompt_group_sizes"] = prompt_group_sizes @@ -159,6 +191,44 @@ def convert_samples_to_train_data( return train_data +def tinker_dispatch_summary(train_data: dict[str, Any]) -> dict[str, Any] | None: + if train_data.get("batch_kind") != "tinker": + return None + return { + "operation_ids": [op_id for op_id in train_data.get("operation_by_lane", {}).values() if op_id], + "lease": train_data.get("batch_execution_lease"), + } + + +def _adapter_slots_from_lease(metadata: dict, sample_lanes: list[int], samples: list[Sample]) -> list[int]: + lease = metadata["batch_execution_lease"] + binding_by_op = {op_id: tuple(binding) for op_id, binding in lease["bindings_by_operation"]} + operation_by_lane = metadata["operation_by_lane"] + lane_ops = list(operation_by_lane.values()) + if len(set(lane_ops)) != len(lane_ops) or set(lane_ops) != set(binding_by_op): + raise ValueError( + f"batch lease and lane plan disagree: lanes carry {sorted(lane_ops)}, " + f"lease carries {sorted(binding_by_op)}" + ) + slots = [] + for sample, lane in zip(samples, sample_lanes, strict=True): + name, registration_id, slot = binding_by_op[operation_by_lane[lane]] + if sample.adapter.name != name or sample.adapter.registration_id != registration_id: + raise ValueError( + f"sample stamped for adapter '{sample.adapter.name}' " + f"(registration '{sample.adapter.registration_id}') rides lane {lane}, " + f"which the batch lease binds to '{name}' (registration '{registration_id}')" + ) + slots.append(slot) + return slots + + +def _tinker_sample_lanes(lanes: list[int], num_samples: int) -> list[int]: + if not lanes or len(lanes) > num_samples: + raise ValueError(f"tinker selection has {len(lanes)} planned rows but {num_samples} samples") + return lanes + [lanes[-1]] * (num_samples - len(lanes)) + + def _compute_rollout_mask_sums(rollout_ids: list[int], loss_masks: list[list[int]]) -> list[int]: """Whole-rollout loss-mask total per sample: every sibling of one rollout carries the sum over all of that rollout's samples, so the loss reducer reconstructs one @@ -371,9 +441,14 @@ def _package_shards(args, data: dict[str, Any], partitions) -> list[dict[str, An "prompt", "teacher_log_probs", "opd_reverse_kl", + # Client-supplied per-token channels (tinker adapters). + "loss_weights", + "advantages", "seq_witness_ids", "weight_versions", "adapter_slots", + # Per-sample batch-local operation lane (tinker correlation plane). + "tinker_operation_lanes", ]: if key not in data: continue @@ -384,9 +459,12 @@ def _package_shards(args, data: dict[str, Any], partitions) -> list[dict[str, An "raw_reward", "total_lengths", "dynamic_global_batch_size", - "step_slots", - "step_adapter_names", - "step_adapter_batch_sizes", + "tinker_loss_by_lane", + "operation_by_lane", + "registration_by_lane", + "batch_execution_lease", + "tinker_forward_only", + "batch_kind", "prompt_group_sizes", ]: if key not in data: diff --git a/miles/ray/tinker_frontend/__init__.py b/miles/ray/tinker_frontend/__init__.py new file mode 100644 index 00000000000..aa62283394d --- /dev/null +++ b/miles/ray/tinker_frontend/__init__.py @@ -0,0 +1,8 @@ +"""HTTP frontend speaking the official tinker SDK's REST protocol (/api/v1). + +Verified against ``tinker==0.24.1`` wheel source and live captured traffic: +an UNMODIFIED SDK pointed at this server (``base_url`` + ``api_key``) drives +training and sampling. The frontend is a thin protocol gateway — every +training verb becomes one operation on the backend ledger, sampling proxies +to the sglang router, and no training semantics live here. +""" diff --git a/miles/ray/tinker_frontend/http_server.py b/miles/ray/tinker_frontend/http_server.py new file mode 100644 index 00000000000..b12e526b0a9 --- /dev/null +++ b/miles/ray/tinker_frontend/http_server.py @@ -0,0 +1,229 @@ +"""HTTP surface for the tinker frontend: /api/v1 as ``tinker==0.24.1`` speaks it. + +Extends the controller's registration server (selected via +``--tinker-frontend`` / ``--multi-lora-http-server-path``), so the SDK +protocol and the operator plane share one uvicorn on the head node — but not +one trust domain: the operator routes (/adapter_runs*, /info) accept +loopback peers only, whatever the bind. The SDK ``X-API-Key`` authenticates +/api/v1/* and never grants the operator plane (which reads server-local +yaml_path files, chooses save paths, and deregisters tenants) to a remote +caller. When a key is configured, every route except the health probes +additionally requires it; a non-loopback bind without a key refuses to +start (fail closed). + +Error mapping (what the 0.24.1 SDK does with each status, observed): +- 429 + Retry-After <- backend backpressure (SDK retries with backoff) +- 422 <- same-identity/different-payload conflicts (fatal to + the SDK; 409 must never be used — the SDK retries it) +- 400/404/401 <- malformed/unknown/unauthenticated (fatal) +- 410 <- expired/unknown future. The SDK does NOT re-run the + original training request: it raises a retryable + "promise expired/broken" toward the caller. Delivered + results answer 410 from a fingerprint tombstone, so + an identical late retry is typed instead of silently + re-executing. +- payload rejections on a spent seq_id are NOT HTTP errors: they become + terminal FAILED(user) futures so the ordinal stays consumed. +""" + +import hmac +import os +from typing import Any + +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse + +from miles.ray.multi_lora.http_server import AdapterRunControlServer +from miles.ray.multi_lora.operations import OperationBackpressure +from miles.ray.tinker_frontend import wire +from miles.ray.tinker_frontend.service import ApiError, TinkerFrontend + +AUTH_EXEMPT_PATHS = ("/health", "/api/v1/healthz") +API_KEY_ENV = "MILES_TINKER_API_KEY" +# The operator plane stays node-local even on a public bind; the SDK key is +# a client credential, not an operator one. +LOOPBACK_PEERS = ("127.0.0.1", "::1", "localhost") + + +def is_sdk_path(path: str) -> bool: + """/api/v1/* plus the base liveness probe; everything else is operator.""" + return path.startswith("/api/v1/") or path == "/health" + + +def resolve_api_key(args: Any) -> str | None: + return getattr(args, "tinker_api_key", None) or os.environ.get(API_KEY_ENV) or None + + +def resolve_sampling_max_context(args: Any) -> int | None: + """Static engine context limit for the sampling preflight: the explicit + tinker flag wins, else the context length this deployment itself launched + its engines with (--sglang-context-length). None defers to lazy discovery + from the router's /get_server_info on the first sample.""" + return getattr(args, "tinker_sampling_max_context", None) or getattr(args, "sglang_context_length", None) or None + + +class TinkerFrontendHTTPServer(AdapterRunControlServer): + """The registration server + the official tinker SDK protocol.""" + + def __init__(self, backend, host="127.0.0.1", api_port=0): + super().__init__(backend, host, api_port) + args = backend.args + self.frontend = TinkerFrontend( + backend, + # Aggregate sampling cap across ALL SDK clients, in sub-generation + # units (the per-client SDK limit of 64 never bounded the sum). + sampling_max_active_subgenerations=getattr(args, "tinker_sampling_max_active_subgenerations", 64), + sampling_max_context=resolve_sampling_max_context(args), + session_idle_ttl_s=getattr(args, "tinker_session_idle_ttl", 3600.0), + future_unpolled_ttl_s=getattr(args, "tinker_future_unpolled_ttl", 900.0), + future_undelivered_ttl_s=getattr(args, "tinker_future_undelivered_ttl", 3600.0), + ) + self.api_key = resolve_api_key(backend.args) + + async def start(self) -> None: + if self.host not in ("127.0.0.1", "localhost", "::1") and not self.api_key: + raise RuntimeError( + f"refusing to bind the tinker frontend to '{self.host}' without an API key: " + f"pass --tinker-api-key or set {API_KEY_ENV}" + ) + await super().start() + # The reaper + metrics-summary loop lives with the serving surface: + # started only once the server accepts traffic, torn down by stop() + # through frontend.close(). + self.frontend.start_maintenance() + + async def stop(self) -> None: + # Order matters: stop ACCEPTING first (uvicorn), then drain the + # frontend (cancel + await in-flight samples, close the transport). + # Closing the frontend first would let a late request lazily reopen + # the transport it just closed. Idempotent: a second stop is a no-op. + if getattr(self, "_stopped", False): + return + self._stopped = True + await super().stop() + await self.frontend.close() + + def create_app(self) -> FastAPI: + app = super().create_app() + + @app.exception_handler(ApiError) + async def api_error_handler(request: Request, exc: ApiError): + return JSONResponse({"detail": exc.detail}, status_code=exc.status_code) + + @app.exception_handler(OperationBackpressure) + async def backpressure_handler(request: Request, exc: OperationBackpressure): + # Retryable by contract: the SDK backs off and resends the same + # request, which the deterministic request ids dedupe. + return JSONResponse({"detail": str(exc)}, status_code=429, headers={"Retry-After": "1"}) + + key = self.api_key.encode() if self.api_key is not None else None + + @app.middleware("http") + async def guard(request: Request, call_next): + path = request.url.path + if key is not None and path not in AUTH_EXEMPT_PATHS: + supplied = request.headers.get("x-api-key", "").encode() + if not hmac.compare_digest(supplied, key): + return JSONResponse({"detail": "invalid or missing X-API-Key"}, status_code=401) + if not is_sdk_path(path): + # Operator plane: node-local only, key or no key. A missing + # peer identity fails closed. + client = request.client + if client is None or client.host not in LOOPBACK_PEERS: + return JSONResponse( + {"detail": "operator routes are loopback-only; the SDK surface is /api/v1/*"}, + status_code=403, + ) + return await call_next(request) + + return app + + def add_routes(self, app: FastAPI) -> None: + super().add_routes(app) + frontend = self.frontend + + # -------- bootstrap / session -------- + @app.get("/api/v1/healthz") + async def healthz() -> dict: + return frontend.health() + + @app.get("/api/v1/get_server_capabilities") + async def get_server_capabilities() -> dict: + return frontend.capabilities() + + @app.post("/api/v1/client/config") + async def client_config(request: wire.ClientConfigRequest) -> dict: + return frontend.client_config(request) + + @app.post("/api/v1/create_session") + async def create_session(request: wire.CreateSessionRequest) -> dict: + return frontend.create_session(request) + + @app.post("/api/v1/session_heartbeat") + async def session_heartbeat(request: wire.SessionHeartbeatRequest) -> dict: + return frontend.session_heartbeat(request) + + @app.post("/api/v1/telemetry") + async def telemetry(request: Request) -> dict: + return frontend.telemetry(await request.body()) + + # -------- models -------- + @app.post("/api/v1/create_model") + async def create_model(request: wire.CreateModelRequest) -> dict: + return await frontend.create_model(request) + + @app.post("/api/v1/get_info") + async def get_info(request: wire.GetInfoRequest) -> dict: + return frontend.get_info(request) + + @app.post("/api/v1/unload_model") + async def unload_model(request: wire.UnloadModelRequest) -> dict: + return await frontend.unload_model(request) + + # -------- training -------- + @app.post("/api/v1/forward_backward") + async def forward_backward(request: wire.ForwardBackwardRequest) -> dict: + return frontend.forward_backward(request) + + @app.post("/api/v1/forward") + async def forward(request: wire.ForwardRequest) -> dict: + return frontend.forward(request) + + @app.post("/api/v1/optim_step") + async def optim_step(request: wire.OptimStepRequest) -> dict: + return frontend.optim_step(request) + + # -------- checkpoints -------- + @app.post("/api/v1/save_weights") + async def save_weights(request: wire.SaveWeightsRequest) -> dict: + return frontend.save_weights(request) + + @app.post("/api/v1/load_weights") + async def load_weights(request: wire.LoadWeightsRequest) -> dict: + return frontend.load_weights(request) + + @app.post("/api/v1/weights_info") + async def weights_info(request: wire.WeightsInfoRequest) -> dict: + return frontend.weights_info(request) + + # -------- sampling -------- + @app.post("/api/v1/save_weights_for_sampler") + async def save_weights_for_sampler(request: wire.SaveWeightsForSamplerRequest) -> dict: + return frontend.save_weights_for_sampler(request) + + @app.post("/api/v1/create_sampling_session") + async def create_sampling_session(request: wire.CreateSamplingSessionRequest) -> dict: + return frontend.create_sampling_session(request) + + @app.get("/api/v1/samplers/{sampler_id}") + async def get_sampler(sampler_id: str) -> dict: + return frontend.get_sampler(sampler_id) + + @app.post("/api/v1/asample") + async def asample(request: wire.SampleRequest) -> dict: + return frontend.sample(request) + + # -------- futures -------- + @app.post("/api/v1/retrieve_future") + async def retrieve_future(request: wire.FutureRetrieveRequest) -> dict: + return await frontend.retrieve_future(request) diff --git a/miles/ray/tinker_frontend/sampling.py b/miles/ray/tinker_frontend/sampling.py new file mode 100644 index 00000000000..320ae835b40 --- /dev/null +++ b/miles/ray/tinker_frontend/sampling.py @@ -0,0 +1,97 @@ +"""Sampling transport for the tinker frontend +(codex-rollout-fullparameter-design-0810 §4.6). + +The sampling hot path stays frontend -> router: /asample answers with a +future immediately and a background task posts the generation itself. This +port isolates WHERE that post goes — the SGLang router today, whatever +endpoint the InferenceController advertises after PR #1842 — without ever +proxying per-sample traffic through a rollout component. Serving identity, +versions, and session invalidation stay in the tinker backend/frontend: +only the HTTP hop lives here.""" + +import asyncio +from typing import Protocol + +import httpx + + +class SamplingTransport(Protocol): + async def generate(self, payload: dict) -> dict: ... + + async def server_info(self) -> dict: ... + + async def close(self) -> None: ... + + +class SGLangRouterSamplingTransport: + """Direct router transport with an explicit hard bound on in-flight + generations (lazy client creation on the first request, like before). + + The previous default-configured client carried an implicit + ``max_connections=100`` pool with a 10-second pool timeout: above 100 + concurrent generations (2 SDK clients x 64, before ``num_samples`` + fan-out) request #101 died waiting for a connection — an empty-message + ``PoolTimeout`` the frontend turned into a terminal server failure the + SDK never retries (the Tau 100/28 sampling cliff). The bound here is the + transport-level invariant behind the frontend's weighted admission: even + a caller that bypasses admission cannot stampede the router.""" + + def __init__(self, base_url: str, max_inflight: int = 64) -> None: + self.base_url = base_url.rstrip("/") + self.max_inflight = max_inflight + # Acquired INSIDE each per-sample generation task (not at submit): + # `async with` guarantees a sibling-cancelled or shutdown-cancelled + # generation releases its permit on the way out. + self._gate = asyncio.Semaphore(max_inflight) + # The pool matches the gate, and pool=None removes the 10s pool + # deadline. That is safe ONLY because the semaphore keeps in-flight + # requests <= max_connections, so a request never actually queues on + # the pool: legal, bounded waiting happens on the gate instead of + # being misclassified as a terminal PoolTimeout. Read stays at 600s + # (the value this frontend always used) — deriving it from router + # config is deliberately out of scope here. + self.limits = httpx.Limits(max_connections=max_inflight, max_keepalive_connections=max_inflight) + self.timeout = httpx.Timeout(connect=10.0, read=600.0, write=60.0, pool=None) + self._http: httpx.AsyncClient | None = None + + async def generate(self, payload: dict) -> dict: + async with self._gate: + if self._http is None: + self._http = httpx.AsyncClient(limits=self.limits, timeout=self.timeout) + response = await self._http.post(f"{self.base_url}/generate", json=payload) + response.raise_for_status() + return response.json() + + async def server_info(self) -> dict: + """Engine server info for the frontend's context-limit discovery, + via a dedicated short-timeout client — an info probe must neither + take a generation permit nor wait behind a saturated pool. + + Two shapes exist behind one URL (verified live on H200): a bare + SGLang engine answers /get_server_info with its ServerArgs + + scheduler_info (context_length / max_req_input_len present), while + sglang-router >= 0.3 answers with router metadata + ({"router_manager": true, ...}) and keeps the engines one hop away + behind /workers. When the first answer carries no engine fields, + hop to the first healthy worker — miles deployments run homogeneous + engines, so any worker's limit is the deployment's limit.""" + async with httpx.AsyncClient(timeout=10.0) as client: + response = await client.get(f"{self.base_url}/get_server_info") + response.raise_for_status() + info = response.json() + if isinstance(info, dict) and ("context_length" in info or "max_req_input_len" in info): + return info + workers_response = await client.get(f"{self.base_url}/workers") + workers_response.raise_for_status() + workers = (workers_response.json() or {}).get("workers") or [] + urls = [row.get("url") for row in workers if row.get("url") and row.get("is_healthy", True)] + if not urls: + return info if isinstance(info, dict) else {} + response = await client.get(f"{urls[0].rstrip('/')}/get_server_info") + response.raise_for_status() + return response.json() + + async def close(self) -> None: + if self._http is not None: + await self._http.aclose() + self._http = None diff --git a/miles/ray/tinker_frontend/service.py b/miles/ray/tinker_frontend/service.py new file mode 100644 index 00000000000..dab2923e7ab --- /dev/null +++ b/miles/ray/tinker_frontend/service.py @@ -0,0 +1,1237 @@ +"""The tinker frontend service: official SDK verbs -> backend operations. + +Request -> ordinal mapping (the D5 note in operations.py): the 0.24.1 SDK +holds one per-model counter — every training verb (each forward_backward +chunk, forward chunk, optim_step, save/load, sampler publish) consumes one +``seq_id``, consecutive from 1 — which is exactly the backend ledger's +per-registration ordinal contract. The frontend therefore forwards +``ordinal = seq_id`` verbatim; chunks the SDK posts out of order (first +chunk last, by design) arrive out of order and the ledger gap-buffers them. +A submission this layer rejects still consumes its ordinal as a terminal +FAILED(user) ledger record, so one bad chunk can never leave a gap that +starves the registration. + +Future protocol: every heavy verb returns ``{"request_id"}`` and the SDK +polls /api/v1/retrieve_future. Request ids are deterministic in the SDK's +own coordinates ((session, model_seq_id) / (model, seq_id)), so a resent +submission lands on its original record: identical -> replay, different -> +422 (the SDK treats 409 as retryable, so a real conflict must never be 409). +Terminal bodies are stored for replay BEFORE the backend record is acked — +a response lost on the wire is re-polled and must find the same bytes. + +This layer is deliberately thin: datum/loss validation and translation live +in translation.py, execution semantics live behind the controller surface +(register/deregister/enqueue/reject/get/ack + registry state), and sampling +proxies to the sglang router under the registration-scoped serving name. + +Sampling is additionally guarded by a context preflight (prompt + max_tokens +against the engine context limit — configured or discovered, typed 400 +before identity consumption) and observed through SamplingAdmission/ +SamplingStats counters; a background maintenance loop reaps orphaned +sessions and futures without ever freeing an identity (code-0815 §6/§7). +""" + +import asyncio +import logging +import time +from collections.abc import Callable +from typing import Any + +from miles.ray.multi_lora.config import AdapterRunConfig +from miles.ray.multi_lora.identity import cache_extra_key, make_rid, serving_lora_name +from miles.ray.multi_lora.operations import OperationBackpressure +from miles.ray.tinker_frontend import translation, wire +from miles.ray.tinker_frontend.sampling import SamplingTransport, SGLangRouterSamplingTransport +from miles.ray.tinker_frontend.state import ( + CheckpointCatalog, + CheckpointRecord, + ConflictError, + ExpiredError, + FutureRecord, + FutureStore, + ModelRecord, + ModelStore, + SamplingSessionRecord, + SamplingSessionStore, + SessionStore, + fingerprint_of, +) +from miles.ray.tinker_frontend.translation import UserInputError + +logger = logging.getLogger(__name__) + +_LEDGER_CONFLICT_MARKS = ("different content", "already taken") + + +class ApiError(Exception): + """Maps to an HTTP error response (submit-time failures the SDK should + see as a status code, not a terminal future).""" + + def __init__(self, status_code: int, detail: str) -> None: + super().__init__(detail) + self.status_code = status_code + self.detail = detail + + +class SamplingAdmission: + """Global fail-fast sampling admission, counted in sub-generations: a + logical request weighs ``num_samples`` because each sample fans out into + its own router call. + + The SDK's per-client ``sample_max_concurrent_requests=64`` bounds ONE + client; the aggregate across clients was unbounded, and >100 concurrent + generations hit the shared router client's implicit 100-connection/10s + pool deadline as empty terminal failures the SDK never retries (the Tau + sampling cliff). Rejecting here — BEFORE the request consumes its seq + identity or mints a FutureRecord — maps to HTTP 429 + Retry-After, which + the SDK retries with backoff using the SAME seq id, so an admitted + request still executes exactly once. Single event loop, no awaits + between check and acquire: admission is atomic with submission.""" + + def __init__(self, capacity: int) -> None: + self.capacity = capacity + self.in_use = 0 + self.rejected = 0 # total backpressured submissions (429s) + self.admitted = 0 # admitted logical requests + self.admitted_weight = 0 # admitted sub-generations (sum of weights) + self.peak_in_use = 0 # high-water of concurrently active sub-generations + + def try_acquire(self, weight: int) -> bool: + if self.in_use + weight > self.capacity: + self.rejected += 1 + return False + self.in_use += weight + self.admitted += 1 + self.admitted_weight += weight + if self.in_use > self.peak_in_use: + self.peak_in_use = self.in_use + return True + + def release(self, weight: int) -> None: + self.in_use -= weight + + +class SamplingStats: + """Aggregate sampling terminal counters (the code-0815 §6.1 minimal set; + admission-side counts live on SamplingAdmission). Latencies are per + logical request: submit -> first completed sub-generation (the closest + observable to time-to-first-token over a non-streaming router hop) and + submit -> terminal.""" + + def __init__(self) -> None: + self.completed = 0 + self.failed = 0 + self.failures_by_class: dict[str, int] = {} + self.first_result_s_sum = 0.0 + self.first_result_s_max = 0.0 + self.first_result_count = 0 + self.total_s_sum = 0.0 + self.total_s_max = 0.0 + + def record_latency(self, first_result_s: float | None, total_s: float) -> None: + self.total_s_sum += total_s + self.total_s_max = max(self.total_s_max, total_s) + if first_result_s is not None: + self.first_result_s_sum += first_result_s + self.first_result_s_max = max(self.first_result_s_max, first_result_s) + self.first_result_count += 1 + + def record_failure(self, failure_class: str) -> None: + self.failed += 1 + self.failures_by_class[failure_class] = self.failures_by_class.get(failure_class, 0) + 1 + + +# The engine context limit out of sglang's /get_server_info. The response is +# ``{**asdict(ServerArgs), **scheduler_info, ...}``: ``context_length`` echoes +# an explicitly configured limit (null when derived from the model config), +# and the scheduler always reports ``max_req_input_len``, which it computes as +# ``min(context_len - 1, kv_pool_tokens - 1) - 5`` — so ``+ 6`` reconstructs +# the effective context (folding in the KV-pool bound when that is tighter). +def _context_limit_from_server_info(info: Any) -> int | None: + if not isinstance(info, dict): + return None + limits = [] + context_length = info.get("context_length") + if isinstance(context_length, int) and not isinstance(context_length, bool) and context_length > 0: + limits.append(context_length) + max_req_input_len = info.get("max_req_input_len") + if isinstance(max_req_input_len, int) and not isinstance(max_req_input_len, bool) and max_req_input_len > 0: + limits.append(max_req_input_len + 6) + return min(limits, default=None) + + +def _note_first_result(task: asyncio.Task, record: "FutureRecord") -> None: + """Done-callback on each sub-generation: stamps when the request's FIRST + sub-generation finished (queue-to-first-result latency). Cancellations + are not results.""" + if not task.cancelled() and record.first_result_at is None: + record.first_result_at = time.time() + + +class TinkerFrontend: + """One instance per controller; single event loop, no cross-await state + mutation inside a submit or resolve step.""" + + def __init__( + self, + backend: Any, + poll_window_s: float = 15.0, + poll_interval_s: float = 0.1, + sampling_transport: SamplingTransport | None = None, + sampling_max_active_subgenerations: int = 64, + sampling_max_context: int | None = None, + session_idle_ttl_s: float = 3600.0, + future_unpolled_ttl_s: float = 900.0, + future_undelivered_ttl_s: float = 3600.0, + maintenance_interval_s: float = 15.0, + ) -> None: + self.backend = backend + self.poll_window_s = poll_window_s + self.poll_interval_s = poll_interval_s + # One capacity, two layers: fail-fast admission at submit (429 before + # identity consumption) and the transport's hard in-flight bound + # (last-resort invariant). 64 is the GPU-validated safe default, not + # a universal optimum — deployments tune it via + # --tinker-sampling-max-active-subgenerations. + self.sampling_admission = SamplingAdmission(sampling_max_active_subgenerations) + self.sampling_stats = SamplingStats() + # Engine context limit for the sampling preflight (prompt + max_tokens + # must fit): statically configured here, or discovered lazily from the + # transport's server_info on the first sample. None = not yet known; + # the preflight only ever rejects against a KNOWN limit. + self._context_limit = sampling_max_context + self._context_limit_source = "configured" if sampling_max_context is not None else None + self._context_discovery_task: asyncio.Task | None = None + self._context_discovery_attempts = 0 + # Orphan reaping TTLs (<= 0 disables that class of reaping). + self.session_idle_ttl_s = session_idle_ttl_s + self.future_unpolled_ttl_s = future_unpolled_ttl_s + self.future_undelivered_ttl_s = future_undelivered_ttl_s + self.maintenance_interval_s = maintenance_interval_s + self._maintenance_task: asyncio.Task | None = None + self._stats_logged: tuple | None = None + # Injected sampling hop (frontend -> router); the default preserves + # the direct-router transport this frontend always used. + self.sampling_transport = ( + sampling_transport + if sampling_transport is not None + else SGLangRouterSamplingTransport( + backend.sampling_endpoint(), max_inflight=sampling_max_active_subgenerations + ) + ) + self.sessions = SessionStore() + self.models = ModelStore() + self.futures = FutureStore() + self.checkpoints = CheckpointCatalog() + self.samplers = SamplingSessionStore() + self._sample_tasks: set[asyncio.Task] = set() + # request_id -> task, so the reaper can cancel one orphaned sample. + self._sample_task_by_request: dict[str, asyncio.Task] = {} + self._closing = False + + async def close(self) -> None: + """Idempotent shutdown barrier: gate new samples, stop the background + maintenance/discovery tasks, cancel AND await every in-flight sample + task (so the transport observes cancellation before it is closed + under it), then close the transport.""" + self._closing = True + background = [task for task in (self._maintenance_task, self._context_discovery_task) if task is not None] + self._maintenance_task = None + self._context_discovery_task = None + for task in background: + task.cancel() + if background: + await asyncio.gather(*background, return_exceptions=True) + tasks = list(self._sample_tasks) + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + # The done-callbacks discard too, but only on a later loop tick; + # close() must return with the set verifiably drained. + self._sample_tasks.difference_update(tasks) + self._sample_task_by_request.clear() + await self.sampling_transport.close() + + # ---------------- maintenance: orphan reaping + metrics summary ---------------- + + def start_maintenance(self) -> None: + """Start the background maintenance loop (idempotent). Owned by the + HTTP server's start — a frontend embedded in tests drives reap_once + directly with an injected clock instead.""" + if self._maintenance_task is None and not self._closing: + self._maintenance_task = asyncio.get_running_loop().create_task(self._maintenance_loop()) + + async def _maintenance_loop(self) -> None: + while True: + await asyncio.sleep(self.maintenance_interval_s) + try: + self.reap_once() + self._log_sampling_summary() + except Exception: # noqa: BLE001 — maintenance must never die silently mid-run + logger.exception("[tinker] frontend maintenance tick failed") + + def reap_once(self, now: float | None = None) -> dict[str, int]: + """One reaping pass (code-0815 §7), replay-idempotency preserved by + construction — reaping frees bytes and capacity without permitting + re-execution: + + - idle sessions (no heartbeat past the TTL): the session record and + its sampling sessions go together; old sampler ids fail closed; + - orphaned sample futures (client stopped polling past the TTL): the + server-side generation is cancelled (releasing admission permits + and transport slots via the existing done-callbacks) and the future + resolves typed with the reap reason — the seq was spent at submit + and stays spent; + - unpolled operation-family futures past the same TTL: polled once on + the client's behalf, which stores the terminal bytes BEFORE acking + the ledger record (the existing ack-based retention order), so the + unacked-results budget drains for vanished clients; + - terminal futures never retrieved past the undelivered TTL: evicted + to a fingerprint tombstone — a late retry gets a typed 410, never a + re-execution. + """ + now = time.time() if now is None else now + counts = {"sessions": 0, "cancelled_samples": 0, "undelivered": 0} + if self.session_idle_ttl_s > 0: + idle_sessions = self.sessions.reap_idle(self.session_idle_ttl_s, now) + self.samplers.remove_for_sessions({session.session_id for session in idle_sessions}) + for session in idle_sessions: + counts["sessions"] += 1 + logger.info( + f"[tinker] reaped idle session '{session.session_id}' (no heartbeat for " + f"{now - session.last_heartbeat:.0f}s; its sampling sessions were retired)" + ) + for record in list(self.futures.records.values()): + if record.terminal is None: + if self.future_unpolled_ttl_s <= 0: + continue + idle_s = now - max(record.created_at, record.last_polled_at) + if idle_s <= self.future_unpolled_ttl_s: + continue + if record.kind == "sample": + task = self._sample_task_by_request.get(record.request_id) + if task is not None and not task.done(): + record.cancel_reason = ( + f"sampling request '{record.request_id}' was orphaned (not polled for " + f"{idle_s:.0f}s) and its generation was cancelled by the reaper; the seq " + "identity stays spent — resubmitting it will not re-run the generation" + ) + task.cancel() + counts["cancelled_samples"] += 1 + logger.warning( + f"[tinker] reaped orphaned sample '{record.request_id}': cancelled its " + f"generation after {idle_s:.0f}s without a poll" + ) + else: + # The training/lifecycle ledger owns execution — never + # cancel it. Resolving on the vanished client's behalf + # moves the terminal bytes here and acks the ledger. + self._poll(record) + elif self.future_undelivered_ttl_s > 0 and not self.futures.is_delivered(record.request_id): + age_s = now - max(record.resolved_at or record.created_at, record.last_polled_at) + if age_s > self.future_undelivered_ttl_s: + self.futures.reap_undelivered(record) + counts["undelivered"] += 1 + logger.info( + f"[tinker] reaped undelivered terminal future '{record.request_id}' " + f"({record.kind}, unretrieved for {age_s:.0f}s); a tombstone keeps its identity" + ) + return counts + + def _log_sampling_summary(self) -> None: + """Periodic aggregate line (INFO), only when something changed since + the last tick — the per-request lines are DEBUG/WARNING.""" + admission, stats = self.sampling_admission, self.sampling_stats + snapshot = (admission.admitted, admission.rejected, stats.completed, stats.failed) + if snapshot == self._stats_logged: + return + self._stats_logged = snapshot + finished = stats.completed + stats.failed + mean_total = stats.total_s_sum / finished if finished else 0.0 + mean_first = stats.first_result_s_sum / stats.first_result_count if stats.first_result_count else 0.0 + logger.info( + f"[tinker] sampling summary: admitted={admission.admitted} ({admission.admitted_weight} " + f"sub-generations) rejected_429={admission.rejected} active={admission.in_use}" + f"/{admission.capacity} peak={admission.peak_in_use} completed={stats.completed} " + f"failed={stats.failed} failures_by_class={stats.failures_by_class} " + f"queue_to_first_result_s(mean/max)={mean_first:.3f}/{stats.first_result_s_max:.3f} " + f"total_s(mean/max)={mean_total:.3f}/{stats.total_s_max:.3f} " + f"context_limit={self._context_limit}" + ) + + # ---------------- bootstrap ---------------- + + def health(self) -> dict: + # Readiness, not liveness (/health): the socket accepting connections + # says nothing about the trainer, which starts later and can fail. + if not getattr(self.backend, "trainer_ready", True): + raise ApiError(503, "trainer is initializing; the service is not ready for SDK traffic yet") + return {"status": "ok"} + + def _check_sdk_version(self, sdk_version: str) -> None: + # Exact pin: this frontend mirrors the request shapes tinker==0.24.1 + # actually POSTs. A different patch of 0.24.x is untested wire surface + # (and 0.25+ switches forward_backward to protobuf mid-run) — reject + # at bootstrap, where the version travels with the request. + if sdk_version != wire.TINKER_SDK_VERSION_PIN: + raise ApiError( + 400, + f"unsupported tinker SDK version '{sdk_version}': this deployment serves exactly " + f"tinker=={wire.TINKER_SDK_VERSION_PIN}. Pin tinker=={wire.TINKER_SDK_VERSION_PIN}.", + ) + + def client_config(self, request: wire.ClientConfigRequest) -> dict: + self._check_sdk_version(request.sdk_version) + return dict(wire.CLIENT_CONFIG_FLAGS) + + def capabilities(self) -> dict: + info = self.backend.service_info() + # None until the engine context limit is configured or discovered. + model = {"model_name": info.get("base_model"), "max_context_length": self._context_limit} + return {"supported_models": [model]} + + def create_session(self, request: wire.CreateSessionRequest) -> dict: + self._check_sdk_version(request.sdk_version) + record = self.sessions.create(request.sdk_version, request.tags, request.user_metadata) + return {"type": "create_session", "session_id": record.session_id} + + def session_heartbeat(self, request: wire.SessionHeartbeatRequest) -> dict: + if not self.sessions.heartbeat(request.session_id): + raise ApiError(404, f"unknown session '{request.session_id}'") + return {"type": "session_heartbeat"} + + def telemetry(self, _body: Any) -> dict: + return {"status": "accepted"} + + # ---------------- models ---------------- + + def _base_model(self) -> str: + return self.backend.service_info().get("base_model") or "" + + def _model_for(self, model_id: str | None) -> ModelRecord: + model = self.models.get(model_id) if model_id else None + if model is None: + raise ApiError(404, f"unknown model_id '{model_id}'") + return model + + async def create_model(self, request: wire.CreateModelRequest) -> dict: + session = self.sessions.get(request.session_id) + if session is None: + raise ApiError(404, f"unknown session '{request.session_id}'") + fingerprint = fingerprint_of(request.model_dump(mode="json")) + name = f"t{session.short}-m{request.model_seq_id}" + request_id = f"{name}:create" + if (existing := self._existing(request_id, fingerprint)) is not None: + return wire.untyped_future(request_id, existing.model.model_id if existing.model else None) + + lora = request.lora_config + if lora is None: + raise ApiError(400, "lora_config is required: this deployment serves LoRA training runs only") + if lora.seed is not None: + raise ApiError(400, "lora_config.seed cannot be honored by this deployment; omit it") + if not (lora.train_unembed and lora.train_mlp and lora.train_attn): + raise ApiError( + 400, + "per-module train flags cannot be honored: trained modules are deployment-wide " + "(--target-modules); leave train_unembed/train_mlp/train_attn at their defaults", + ) + base_model = self._base_model() + if request.base_model != base_model: + raise ApiError( + 400, f"base_model '{request.base_model}' is not served; this deployment serves '{base_model}'" + ) + + metadata = {"session_id": request.session_id, "model_seq_id": request.model_seq_id} + if request.user_metadata: + metadata["user_metadata"] = request.user_metadata + try: + await self.backend.register(name, AdapterRunConfig(rank=lora.rank, metadata=metadata)) + except ValueError as exc: + # A concurrent identical create may have raced this one. + if (existing := self._existing(request_id, fingerprint)) is not None: + return wire.untyped_future(request_id, existing.model.model_id if existing.model else None) + raise ApiError(400, str(exc)) from exc + registered = self.backend.registration_view(name) + model = ModelRecord( + model_id=f"{request.session_id}:train:{request.model_seq_id}", + session_id=request.session_id, + model_seq_id=request.model_seq_id, + name=name, + registration_id=registered["registration_id"], + base_model=base_model, + rank=registered["rank"], + fingerprint=fingerprint, + ) + self.models.add(model) + self.futures.put( + FutureRecord(request_id=request_id, kind="create_model", fingerprint=fingerprint, model=model) + ) + return wire.untyped_future(request_id, model.model_id) + + def get_info(self, request: wire.GetInfoRequest) -> dict: + model = self._model_for(request.model_id) + return { + "type": "get_info", + "model_id": model.model_id, + "model_data": {"arch": None, "model_name": model.base_model, "tokenizer_id": model.base_model}, + "is_lora": True, + "lora_rank": model.rank, + "model_name": model.base_model, + } + + async def unload_model(self, request: wire.UnloadModelRequest) -> dict: + model = self._model_for(request.model_id) + fingerprint = fingerprint_of(request.model_dump(mode="json")) + request_id = f"{model.name}.{model.rid8}:unload" + if self._existing(request_id, fingerprint) is not None: + return wire.untyped_future(request_id, model.model_id) + # Registration-pinned: a same-name successor must never be retired + # by a stale handle (the backend re-checks under the same pin). + await self.backend.deregister(model.name, model.registration_id) + self.futures.put( + FutureRecord(request_id=request_id, kind="unload_model", fingerprint=fingerprint, model=model) + ) + return wire.untyped_future(request_id, model.model_id) + + # ---------------- training operations ---------------- + + def forward_backward(self, request: wire.ForwardBackwardRequest) -> dict: + return self._submit_operation( + request, + request.model_id, + request.seq_id, + "forward_backward", + lambda: translation.fb_input_to_payload(request.forward_backward_input), + ) + + def forward(self, request: wire.ForwardRequest) -> dict: + def prepare(record: FutureRecord, payload: dict) -> None: + # The backend attaches loss metrics to forward_backward results + # only; keep the request payload for the forward recompute. + record.forward_payload = payload + + return self._submit_operation( + request, + request.model_id, + request.seq_id, + "forward", + lambda: translation.fb_input_to_payload(request.forward_input), + prepare=prepare, + ) + + def optim_step(self, request: wire.OptimStepRequest) -> dict: + return self._submit_operation( + request, + request.model_id, + request.seq_id, + "optim_step", + lambda: translation.adam_params_to_payload(request.adam_params), + ) + + def save_weights(self, request: wire.SaveWeightsRequest) -> dict: + def build() -> dict: + if request.overwrite: + raise UserInputError("overwrite=true is not supported: named states are immutable") + if request.ttl_seconds is not None: + # No reaper runs in v1: accepting a TTL would promise an expiry + # that never happens. Same typed rejection as sampler publishes. + raise UserInputError("ttl_seconds is not supported in v1 (checkpoints never expire); omit it") + payload: dict = {} + if request.path is not None: + payload["tag"] = request.path + return payload + + return self._submit_operation(request, request.model_id, request.seq_id, "save_state", build) + + def load_weights(self, request: wire.LoadWeightsRequest) -> dict: + if request.model_id is None: + # create_model_via_load_weights is advertised off; the SDK only + # sends session addressing when the server enables that flag. + raise ApiError(400, "load_weights requires model_id (session-addressed creation is not supported)") + + def build() -> dict: + if not request.optimizer: + raise UserInputError( + "weights-only restore is not supported in v1 (the backend restores the full training " + "state); use load_state_with_optimizer / create_training_client_from_state_with_optimizer" + ) + checkpoint = self.checkpoints.get(request.path) + if checkpoint is None: + raise UserInputError( + f"unknown checkpoint '{request.path}'; v1 resolves paths minted during this service lifetime" + ) + return {"path": checkpoint.backend_path} + + def prepare(record: FutureRecord, payload: dict) -> None: + record.tinker_path = request.path + # Redaction: failures echo the trainer-side path; swap it back for + # the public URI before the error body reaches the client. + record.backend_target = {"path": payload["path"]} + + return self._submit_operation(request, request.model_id, request.seq_id, "load_state", build, prepare=prepare) + + def save_weights_for_sampler(self, request: wire.SaveWeightsForSamplerRequest) -> dict: + model = self._model_for(request.model_id) + + def build() -> dict: + if self.sessions.get(model.session_id) is None: + raise UserInputError("the parent session expired; create a new session before publishing a sampler") + if request.path is not None: + raise UserInputError( + "named sampler checkpoints are not supported in v1 (latest-only serving); use " + "save_weights_and_get_sampling_client for ephemeral sampling" + ) + if request.sampling_session_seq_id is None: + raise UserInputError("save_weights_for_sampler without a path needs sampling_session_seq_id") + if request.ttl_seconds is not None: + raise UserInputError("ttl_seconds is not supported for sampler publishes in v1") + return {} + + def prepare(record: FutureRecord, payload: dict) -> None: + session = self.sessions.get(model.session_id) + short = session.short if session is not None else model.session_id[:12] + record.sampling_session_id = f"samp-{short}-ss{request.sampling_session_seq_id}" + + # The official 0.24.1 client increments its sampling counter INSIDE + # the HTTP retry closure (training_client.py: _send_request mints a + # fresh sampling_session_seq_id per attempt) while the operation + # seq_id stays fixed. A response lost on the wire therefore retries + # the SAME operation identity with a different sampling sequence — + # fingerprinting that field would turn the retry into a fatal 422. + # The operation seq_id remains authoritative; replay returns the + # originally minted sampler id. + fingerprint_dump = request.model_dump(mode="json") + fingerprint_dump.pop("sampling_session_seq_id", None) + return self._submit_operation( + request, + request.model_id, + request.seq_id, + "save_weights_for_sampler", + build, + prepare=prepare, + fingerprint_dump=fingerprint_dump, + ) + + def _existing(self, request_id: str, fingerprint: str) -> FutureRecord | None: + try: + return self.futures.existing(request_id, fingerprint) + except ExpiredError as exc: + raise ApiError(410, str(exc)) from exc + except ConflictError as exc: + raise ApiError(422, str(exc)) from exc + + def _submit_operation( + self, + request: wire.WireModel, + model_id: str | None, + seq_id: int | None, + kind: str, + build_payload: Callable[[], dict], + prepare: Callable[[FutureRecord, dict], None] | None = None, + fingerprint_dump: dict | None = None, + ) -> dict: + model = self._model_for(model_id) + if seq_id is None or seq_id < 1: + raise ApiError(400, f"{kind} needs a seq_id >= 1") + request_dump = request.model_dump(mode="json") + # ``fingerprint_dump`` lets a verb exclude fields the official SDK + # regenerates per retry attempt (save_weights_for_sampler's + # sampling_session_seq_id) from the retry-identity fingerprint. + fingerprint = fingerprint_of(fingerprint_dump if fingerprint_dump is not None else request_dump) + request_id = f"{model.name}.{model.rid8}:op{seq_id}" + if self._existing(request_id, fingerprint) is not None: + return wire.untyped_future(request_id, model.model_id) + + record = FutureRecord( + request_id=request_id, + kind="operation", + fingerprint=fingerprint, + model=model, + operation_id=request_id, + operation_kind=kind, + ) + try: + payload = build_payload() + if prepare is not None: + prepare(record, payload) + # Registration-pinned (anti-ABA): a stale model handle must fence, + # never bind to a same-name successor registration. + self.backend.enqueue_operation(model.name, request_id, seq_id, kind, payload, model.registration_id) + except UserInputError as exc: + # The client spent this ordinal: consume it as terminal + # FAILED(user) so later operations never wait behind a gap. + self._reject_into_ledger(record, model, seq_id, kind, request_dump, str(exc)) + except ValueError as exc: + message = str(exc) + if any(mark in message for mark in _LEDGER_CONFLICT_MARKS): + raise ApiError(422, message) from exc + if "not accepting operations" in message or "fenced" in message: + record.resolve(wire.terminal_failure(message, "user")) + else: + self._reject_into_ledger(record, model, seq_id, kind, request_dump, message) + self.futures.put(record) + return wire.untyped_future(request_id, model.model_id) + + def _reject_into_ledger( + self, record: FutureRecord, model: ModelRecord, seq_id: int, kind: str, request_dump: dict, error: str + ) -> None: + # The wire dump is the reject payload: deterministic across retries, + # so a resend after a frontend restart matches the ledger fingerprint. + try: + self.backend.reject_operation( + model.name, record.operation_id, seq_id, kind, {"wire": request_dump}, error, model.registration_id + ) + except ValueError: + record.resolve(wire.terminal_failure(error, "user")) + + # ---------------- checkpoints ---------------- + + def weights_info(self, request: wire.WeightsInfoRequest) -> dict: + checkpoint = self.checkpoints.get(request.tinker_path) + if checkpoint is None: + raise ApiError( + 404, + f"unknown checkpoint '{request.tinker_path}'; v1 resolves paths minted during this service lifetime", + ) + return { + "base_model": checkpoint.base_model, + "is_lora": True, + "lora_rank": checkpoint.rank, + "train_unembed": None, + "train_mlp": None, + "train_attn": None, + } + + # ---------------- sampling ---------------- + + def create_sampling_session(self, request: wire.CreateSamplingSessionRequest) -> dict: + session = self.sessions.get(request.session_id) + if session is None: + raise ApiError(404, f"unknown session '{request.session_id}'") + fingerprint = fingerprint_of(request.model_dump(mode="json")) + sampling_session_id = f"samp-{session.short}-ss{request.sampling_session_seq_id}" + try: + existing = self.samplers.existing(sampling_session_id, fingerprint) + except ConflictError as exc: + raise ApiError(422, str(exc)) from exc + if existing is not None: + return {"type": "create_sampling_session", "sampling_session_id": sampling_session_id} + if request.model_path is not None: + raise ApiError( + 400, + "sampling from saved checkpoints is not supported in v1 (latest-only serving); use " + "save_weights_and_get_sampling_client on the training client, or a base_model session", + ) + base_model = self._base_model() + if request.base_model != base_model: + raise ApiError( + 400, f"base_model '{request.base_model}' is not served; this deployment serves '{base_model}'" + ) + self.samplers.add( + SamplingSessionRecord( + sampling_session_id=sampling_session_id, + session_id=request.session_id, + fingerprint=fingerprint, + base_model=base_model, + ) + ) + return {"type": "create_sampling_session", "sampling_session_id": sampling_session_id} + + def get_sampler(self, sampler_id: str) -> dict: + sampler = self.samplers.get(sampler_id) + if sampler is None: + raise ApiError(404, f"unknown sampler '{sampler_id}'") + return {"sampler_id": sampler.sampling_session_id, "base_model": sampler.base_model, "model_path": None} + + def sample(self, request: wire.SampleRequest) -> dict: + if self._closing: + raise ApiError(503, "the service is shutting down; no new samples are accepted") + if request.sampling_session_id is None: + raise ApiError(400, "asample requires sampling_session_id (create a sampling session first)") + sampler = self.samplers.get(request.sampling_session_id) + if sampler is None: + raise ApiError(404, f"unknown sampler '{request.sampling_session_id}'") + if request.seq_id is None or request.seq_id < 0: + raise ApiError(400, "asample needs a seq_id >= 0") + fingerprint = fingerprint_of(request.model_dump(mode="json")) + request_id = f"{sampler.sampling_session_id}:s{request.seq_id}" + if self._existing(request_id, fingerprint) is not None: + return wire.untyped_future(request_id) + if sampler.is_spent(request.seq_id): + # The replay bytes AND the fingerprint tombstone are gone (bounded + # retention rolled over), but the per-session spent-sequence fence + # still knows this identity executed: answer a typed terminal + # failure instead of silently re-running the generation. + record = self.futures.put(FutureRecord(request_id=request_id, kind="sample", fingerprint=fingerprint)) + record.resolve( + wire.terminal_failure( + f"sample seq {request.seq_id} of '{sampler.sampling_session_id}' was already executed " + "and its result expired from the replay window; it cannot be re-run", + "user", + ) + ) + return wire.untyped_future(request_id) + + record = FutureRecord(request_id=request_id, kind="sample", fingerprint=fingerprint) + try: + if request.topk_prompt_logprobs: + raise UserInputError("topk_prompt_logprobs is not supported in v1") + if request.num_samples < 1: + raise UserInputError("num_samples must be >= 1") + prompt_tokens = translation._input_tokens("prompt", request.prompt) + sglang_params = translation.sampling_params_to_sglang(request.sampling_params) + seed = request.sampling_params.seed + if seed is not None and seed + request.num_samples - 1 >= 2**63: + raise UserInputError("sampling_params.seed + num_samples must fit in a signed 64-bit integer") + except UserInputError as exc: + # Invalid payloads still consume the seq as a typed terminal (the + # http_server contract) — but never a permit: nothing will run. + sampler.mark_spent(request.seq_id) + self.futures.put(record) + record.resolve(wire.terminal_failure(str(exc), "user")) + return wire.untyped_future(request_id) + + admission = self.sampling_admission + if request.num_samples > admission.capacity: + # Would 429 forever — fail typed and non-retryable, without + # consuming the seq, so the client can split into waves. + raise ApiError( + 400, + f"num_samples={request.num_samples} exceeds this deployment's sampling capacity of " + f"{admission.capacity} concurrent sub-generations; split the request into smaller waves", + ) + # Context preflight (code-0815 §6.2): a prompt that leaves no decode + # budget must fail HERE, typed and non-retryable — the engine itself + # silently truncates max_new_tokens to whatever fits (near zero for + # an oversized accumulated context) and returns garbage. Like the + # num_samples cap above: a deterministic 400 before the seq identity + # is consumed, so nothing executes and nothing gaps. + limit = self._context_limit + if limit is None: + self._ensure_context_limit_discovery() + else: + max_new_tokens = sglang_params["max_new_tokens"] + if len(prompt_tokens) + max_new_tokens > limit: + raise ApiError( + 400, + f"prompt ({len(prompt_tokens)} tokens) + max_tokens ({max_new_tokens}) exceeds this " + f"deployment's engine context limit of {limit} tokens ({self._context_limit_source}); " + "shorten the prompt or lower max_tokens — the engine would silently truncate the " + "decode budget instead of honoring the request", + ) + if not admission.try_acquire(request.num_samples): + # BEFORE mark_spent/FutureRecord: the identity stays unconsumed, + # so the SDK's backoff retry of the SAME seq id is safe. The HTTP + # layer maps this to 429 + Retry-After. + raise OperationBackpressure( + f"sampling capacity reached ({admission.in_use}/{admission.capacity} sub-generations " + "active); retry the identical request" + ) + # No await from try_acquire to create_task: admission, identity + # consumption, and FutureRecord creation are one atomic submission + # step (two identical racing requests cannot both execute). + sampler.mark_spent(request.seq_id) + self.futures.put(record) + task = asyncio.get_running_loop().create_task( + self._run_sample( + record, + sampler, + prompt_tokens, + sglang_params, + request.num_samples, + request.sampling_params.seed, + prompt_logprobs=bool(request.prompt_logprobs), + ) + ) + self._sample_tasks.add(task) + self._sample_task_by_request[request_id] = task + task.add_done_callback( + lambda done: self._terminalize_prestart_cancelled_sample( + done, + record, + request.num_samples, + len(prompt_tokens), + sglang_params.get("max_new_tokens"), + ) + ) + task.add_done_callback(self._sample_tasks.discard) + task.add_done_callback(lambda _task, rid=request_id: self._sample_task_by_request.pop(rid, None)) + # Release via done-callback, not inside the coroutine: a task + # cancelled before its first step never enters the coroutine body, so + # a `finally` there could leak the permit on shutdown. + task.add_done_callback(lambda _task, weight=request.num_samples: admission.release(weight)) + return wire.untyped_future(request_id) + + # ---------------- engine context discovery ---------------- + + _CONTEXT_DISCOVERY_MAX_ATTEMPTS = 3 + + def _ensure_context_limit_discovery(self) -> None: + """Single-flight, non-blocking: sample submission stays synchronous + (admission atomicity), so discovery runs as a background task kicked + off by the first sample. Until it lands the preflight admits + everything (a permissive window, never a false reject).""" + if ( + self._context_limit is not None + or self._closing + or self._context_discovery_task is not None + or self._context_discovery_attempts >= self._CONTEXT_DISCOVERY_MAX_ATTEMPTS + ): + return + server_info = getattr(self.sampling_transport, "server_info", None) + if server_info is None: + self._context_discovery_attempts = self._CONTEXT_DISCOVERY_MAX_ATTEMPTS + logger.warning( + "[tinker] sampling context preflight disabled: the sampling transport exposes no " + "server_info; pass --tinker-sampling-max-context to enforce a limit" + ) + return + self._context_discovery_task = asyncio.get_running_loop().create_task( + self._discover_context_limit(server_info) + ) + + async def _discover_context_limit(self, server_info: Callable) -> None: + self._context_discovery_attempts += 1 + attempt = f"attempt {self._context_discovery_attempts}/{self._CONTEXT_DISCOVERY_MAX_ATTEMPTS}" + try: + info = await server_info() + except asyncio.CancelledError: + raise + except Exception as exc: # noqa: BLE001 — discovery must never take sampling down + if self._context_discovery_attempts >= self._CONTEXT_DISCOVERY_MAX_ATTEMPTS: + logger.warning( + f"[tinker] sampling context preflight disabled: engine context discovery failed " + f"({attempt}: {type(exc).__name__}: {exc}); pass --tinker-sampling-max-context " + "to enforce a limit" + ) + else: + logger.info(f"[tinker] engine context discovery failed ({attempt}), will retry: {exc}") + return + finally: + # Cleared AFTER the outcome is recorded: the next sample may + # re-trigger discovery only while attempts remain. + self._context_discovery_task = None + limit = _context_limit_from_server_info(info) + if limit is None: + self._context_discovery_attempts = self._CONTEXT_DISCOVERY_MAX_ATTEMPTS + logger.warning( + "[tinker] sampling context preflight disabled: /get_server_info carried neither " + "context_length nor max_req_input_len; pass --tinker-sampling-max-context to enforce a limit" + ) + return + self._context_limit = limit + self._context_limit_source = "discovered from the engine" + logger.info(f"[tinker] sampling context preflight active: engine context limit {limit} tokens (discovered)") + + def _terminalize_prestart_cancelled_sample( + self, + task: asyncio.Task, + record: FutureRecord, + num_samples: int, + prompt_tokens: int, + max_new_tokens: int | None, + ) -> None: + """Resolve a task cancelled before its coroutine body ever ran.""" + if not task.cancelled() or record.terminal is not None: + return + record.failure_class = "Cancelled" + record.resolve( + wire.terminal_failure(record.cancel_reason or "sampling cancelled: the service is shutting down", "server") + ) + self._account_sample_terminal(record, num_samples, prompt_tokens, max_new_tokens) + + async def _run_sample( + self, + record: FutureRecord, + sampler: SamplingSessionRecord, + tokens: list[int], + params: dict, + num_samples: int, + seed: int | None = None, + prompt_logprobs: bool = False, + ) -> None: + try: + await self._execute_sample(record, sampler, tokens, params, num_samples, seed, prompt_logprobs) + except asyncio.CancelledError: + # Reaper cancellation carries its reason on the record; anything + # else is the shutdown barrier. Either way the future resolves so + # a client polling it sees a typed terminal, never an identity + # that silently stops progressing. + record.failure_class = "Cancelled" + record.resolve( + wire.terminal_failure( + record.cancel_reason or "sampling cancelled: the service is shutting down", "server" + ) + ) + raise + except Exception as exc: # noqa: BLE001 — every failure must resolve the future + # Always name the exception class: str(httpx.PoolTimeout()) is + # empty, and a bare "sampling failed: " is undiagnosable. + record.failure_class = type(exc).__name__ + record.resolve(wire.terminal_failure(f"sampling failed ({type(exc).__name__}): {exc}", "server")) + finally: + self._account_sample_terminal(record, num_samples, len(tokens), params.get("max_new_tokens")) + + async def _execute_sample( + self, + record: FutureRecord, + sampler: SamplingSessionRecord, + tokens: list[int], + params: dict, + num_samples: int, + seed: int | None = None, + prompt_logprobs: bool = False, + ) -> None: + payload: dict = {"input_ids": tokens, "sampling_params": params, "return_logprob": True} + if prompt_logprobs: + # sglang natively scores the prompt: input_token_logprobs from position 0. + payload["logprob_start_len"] = 0 + if sampler.name is not None: + live = self.backend.registration_view(sampler.name) + if live is None or live["registration_id"] != sampler.registration_id: + record.resolve( + wire.terminal_failure("sampler weights are no longer live (registration retired)", "user") + ) + return + if live["serving_version"] != sampler.serving_version: + record.resolve( + wire.terminal_failure( + "stale ephemeral sampler: the model was republished and this backend serves the " + "latest weights only — create a new sampling client after each publish", + "user", + ) + ) + return + payload["lora_path"] = sampler.serving_name + payload["extra_key"] = cache_extra_key(sampler.name, sampler.registration_id, sampler.serving_version) + + def per_sample_payload(index: int) -> dict: + one = dict(payload) + if seed is not None: + # Deterministic per request, still diverse across samples. + one["sampling_params"] = {**params, "sampling_seed": seed + index} + if sampler.name is not None: + one["rid"] = make_rid(sampler.name, sampler.registration_id) + return one + + # Not a bare gather: the first exception must not leave siblings + # running untracked — cancel them and AWAIT their cancellation + # before this future turns terminal, so no generation outlives + # its request's resolution. + generation_tasks = [ + asyncio.get_running_loop().create_task(self.sampling_transport.generate(per_sample_payload(index))) + for index in range(num_samples) + ] + for generation_task in generation_tasks: + generation_task.add_done_callback(lambda task, r=record: _note_first_result(task, r)) + try: + generations = await asyncio.gather(*generation_tasks) + except BaseException: + for task in generation_tasks: + task.cancel() + await asyncio.gather(*generation_tasks, return_exceptions=True) + raise + if sampler.name is not None and not self._sampler_still_live(sampler): + # Re-checked AFTER generation: a republish that landed while + # the request was in flight swapped the engine-side weights + # under the same serving name (latest-only serving), so the + # output cannot be attributed to the pinned version. Fail loud + # rather than return cross-version samples. (A publish + # committing between this check and delivery remains possible + # — the serving identity is versioned, not leased; see README.) + record.resolve( + wire.terminal_failure( + "the model was republished while this sample was in flight; create a new sampling " + "client after each publish and resample", + "user", + ) + ) + return + sequences = [translation.generation_to_sequence(generation) for generation in generations] + # The prompt is shared across the fan-out, so any generation's scores serve. + scored = translation.prompt_logprobs_from_generation(generations[0], len(tokens)) if prompt_logprobs else None + record.resolve(translation.sequences_to_sample_response(sequences, scored)) + + def _account_sample_terminal( + self, record: FutureRecord, num_samples: int, prompt_tokens: int, max_new_tokens: int | None + ) -> None: + """Single terminal choke point for every task-executed sample: the + §6.1 counters plus one per-request line carrying the latencies. Per + request at DEBUG (high-volume), failures at WARNING with their class.""" + body = record.terminal or {} + stats = self.sampling_stats + admission = self.sampling_admission + terminal_at = record.resolved_at or time.time() + total_s = terminal_at - record.created_at + first_result_s = (record.first_result_at - record.created_at) if record.first_result_at is not None else None + first_result = f"{first_result_s:.3f}" if first_result_s is not None else "n/a" + detail = ( + f"request='{record.request_id}' num_samples={num_samples} prompt_tokens={prompt_tokens} " + f"max_tokens={max_new_tokens} queue_to_first_result_s={first_result} total_s={total_s:.3f} " + f"active={admission.in_use}/{admission.capacity} peak={admission.peak_in_use} " + f"admitted={admission.admitted} rejected_429={admission.rejected}" + ) + stats.record_latency(first_result_s, total_s) + if "error" in body: + failure_class = record.failure_class or body.get("category") or "unknown" + stats.record_failure(failure_class) + logger.warning( + f"[tinker] sample terminal failure class={failure_class} category={body.get('category')} " + f"{detail} error={body.get('error')!r}" + ) + else: + stats.completed += 1 + logger.debug(f"[tinker] sample terminal ok {detail}") + + def _sampler_still_live(self, sampler: SamplingSessionRecord) -> bool: + live = self.backend.registration_view(sampler.name) + return ( + live is not None + and live["registration_id"] == sampler.registration_id + and live["serving_version"] == sampler.serving_version + ) + + # ---------------- future retrieval ---------------- + + async def retrieve_future(self, request: wire.FutureRetrieveRequest) -> dict: + """Long-poll: resolve inside the window when possible, else try_again.""" + deadline = time.monotonic() + self.poll_window_s + while True: + record = self.futures.get(request.request_id) + if record is None: + if self.futures.expired_fingerprint(request.request_id) is not None: + raise ApiError( + 410, + f"request '{request.request_id}' was already delivered and its replay window expired", + ) + if self.futures.reaped_fingerprint(request.request_id) is not None: + raise ApiError( + 410, + f"request '{request.request_id}' completed but was never retrieved within its " + "retention TTL and was reaped", + ) + raise ApiError( + 410, f"unknown request '{request.request_id}' (expired or from a previous service lifetime)" + ) + # Liveness for the orphan reaper: an actively polled future is + # never an orphan, whatever its age. + record.last_polled_at = time.time() + if record.terminal is None: + self._poll(record) + if record.terminal is not None: + body = record.terminal + self.futures.mark_delivered(record) + return body + if time.monotonic() >= deadline: + return wire.try_again(self._queue_state(record)) + await asyncio.sleep(self.poll_interval_s) + + def _queue_state(self, record: FutureRecord) -> str: + if record.kind == "create_model" and record.model is not None: + live = self.backend.registration_view(record.model.name) + if live is not None and not live["bound"]: + return "paused_capacity" + return "active" + + def _poll(self, record: FutureRecord) -> None: + if record.kind == "operation": + self._poll_operation(record) + elif record.kind == "create_model": + self._poll_create_model(record) + elif record.kind == "unload_model": + self._poll_unload_model(record) + # "sample" resolves from its own task. + + def _poll_operation(self, record: FutureRecord) -> None: + view = self.backend.operation_view(record.operation_id) + if view is None: + record.resolve(wire.terminal_failure("operation record lost before retrieval", "server")) + return + state = view["state"] + if state in ("QUEUED", "CLAIMED"): + return + if state == "SUCCEEDED": + record.resolve(self._success_body(record, view.get("result") or {})) + else: # FAILED | CANCELLED + error = view.get("error") or "operation failed" + if record.backend_target and record.tinker_path: + # Clients know the tinker:// URI, not the trainer's filesystem. + error = error.replace(record.backend_target["path"], record.tinker_path) + record.resolve(wire.terminal_failure(error, view.get("error_category") or "server")) + # Ack only after the terminal body is stored: a lost response replays + # from the future store, never from a record the ack released. + self.backend.ack_operation(record.operation_id) + + def _success_body(self, record: FutureRecord, result: dict) -> dict: + kind, model = record.operation_kind, record.model + if kind in ("forward_backward", "forward"): + return translation.fb_result_to_response(result, record.forward_payload) + if kind == "optim_step": + return translation.optim_result_to_response(result) + if kind == "save_state": + backend_path = str(result.get("path")) + tag = backend_path.rstrip("/").rsplit("/", 1)[-1] + tinker_path = f"tinker://{model.name}.{model.rid8}/weights/{tag}" + self.checkpoints.add( + CheckpointRecord( + tinker_path=tinker_path, + backend_path=backend_path, + name=model.name, + registration_id=model.registration_id, + base_model=model.base_model, + rank=model.rank, + step=int(result.get("step") or 0), + ) + ) + return translation.save_weights_result_to_response(tinker_path) + if kind == "load_state": + return translation.load_weights_result_to_response(record.tinker_path, model.model_id) + if kind == "save_weights_for_sampler": + if self.sessions.get(model.session_id) is None: + return wire.terminal_failure( + "the parent session expired before sampler publication completed; create a new session", "user" + ) + existing = self.samplers.get(record.sampling_session_id) + if existing is not None and existing.fingerprint != record.fingerprint: + # Never overwrite a live sampler identity: a base sampler (or + # another publish) already owns this namespace, and silently + # rebinding it would swap the weights under an existing + # client. The weights are live (the publish itself landed); + # only the sampler minting fails, typed. + return wire.terminal_failure( + f"sampling session '{record.sampling_session_id}' already exists; publish with a fresh " + "sampling_session_seq_id to mint a new sampler", + "user", + ) + self.samplers.add( + SamplingSessionRecord( + sampling_session_id=record.sampling_session_id, + session_id=model.session_id, + fingerprint=record.fingerprint, + base_model=model.base_model, + name=model.name, + registration_id=model.registration_id, + serving_name=result.get("serving_name") or serving_lora_name(model.name, model.registration_id), + serving_version=result.get("serving_version"), + ) + ) + return translation.sampler_publish_result_to_response(record.sampling_session_id) + return wire.terminal_failure(f"no translator for operation kind '{kind}'", "server") + + def _poll_create_model(self, record: FutureRecord) -> None: + model = record.model + live = self.backend.registration_view(model.name) + if live is None or live["registration_id"] != model.registration_id: + record.resolve(wire.terminal_failure("registration retired before the model became ready", "user")) + return + if live["state"] == "READY": + record.resolve({"type": "create_model", "model_id": model.model_id}) + elif live["state"] != "PENDING": + record.resolve(wire.terminal_failure(f"registration is {live['state']}; model creation failed", "user")) + + def _poll_unload_model(self, record: FutureRecord) -> None: + model = record.model + live = self.backend.registration_view(model.name) + if live is None or live["registration_id"] != model.registration_id: + record.resolve({"type": "unload_model", "model_id": model.model_id}) diff --git a/miles/ray/tinker_frontend/state.py b/miles/ray/tinker_frontend/state.py new file mode 100644 index 00000000000..690657105c9 --- /dev/null +++ b/miles/ray/tinker_frontend/state.py @@ -0,0 +1,332 @@ +"""Frontend-owned protocol state: sessions, models, futures, checkpoints, +sampling sessions. + +Identity is deterministic wherever the SDK retries: the SDK addresses work +by (session_id, model_seq_id) and (model, seq_id), so request ids derive +from those coordinates and a resent submission finds its original record. +Every record carries the fingerprint of the request that minted it — an +identical retry replays, a different payload at the same coordinates is a +conflict (422; the SDK treats 409 as retryable, so a true conflict must +never be a 409). + +All state is in-memory and single-writer: the frontend runs on the +controller actor's event loop, and store mutations never straddle an await. +Terminal future bodies are kept for replay (a response lost on the wire is +re-polled) inside a bounded LRU of delivered results; eviction keeps a +compact fingerprint tombstone so an expired identity answers a typed 410 +instead of silently re-executing. +""" + +import hashlib +import json +import time +import uuid +from collections import OrderedDict +from dataclasses import dataclass, field +from typing import Any + + +def fingerprint_of(payload: Any) -> str: + return hashlib.sha256(json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()).hexdigest() + + +class ConflictError(ValueError): + """Same identity, different content: the client must not silently retry.""" + + +class ExpiredError(ValueError): + """The result was delivered and its replay window has expired. The exact + terminal bytes are gone, so neither replay nor re-execution is possible — + re-running would break idempotency (a fresh sample for a spent seq, an + ordinal the ledger already consumed). Maps to a typed 410.""" + + +def _check_fingerprint(kind: str, key: str, existing: str, incoming: str) -> None: + if existing != incoming: + raise ConflictError(f"{kind} '{key}' already exists with a different request; retries must be identical") + + +@dataclass +class SessionRecord: + session_id: str + sdk_version: str = "" + tags: list[str] = field(default_factory=list) + user_metadata: dict | None = None + created_at: float = field(default_factory=time.time) + last_heartbeat: float = field(default_factory=time.time) + + @property + def short(self) -> str: + return self.session_id.removeprefix("sess-")[:12] + + +class SessionStore: + def __init__(self) -> None: + self.records: dict[str, SessionRecord] = {} + + def create(self, sdk_version: str, tags: list[str], user_metadata: dict | None) -> SessionRecord: + record = SessionRecord( + session_id=f"sess-{uuid.uuid4().hex[:16]}", + sdk_version=sdk_version, + tags=tags, + user_metadata=user_metadata, + ) + self.records[record.session_id] = record + return record + + def get(self, session_id: str) -> SessionRecord | None: + return self.records.get(session_id) + + def heartbeat(self, session_id: str) -> bool: + record = self.records.get(session_id) + if record is None: + return False + record.last_heartbeat = time.time() + return True + + def reap_idle(self, ttl_s: float, now: float) -> list[SessionRecord]: + """Remove sessions whose client stopped heartbeating for ``ttl_s``. + Child sampling sessions are retired separately by their lifecycle + owner; their old ids then fail closed instead of becoming reusable.""" + idle = [record for record in self.records.values() if now - record.last_heartbeat > ttl_s] + for record in idle: + del self.records[record.session_id] + return idle + + +@dataclass +class ModelRecord: + """One SDK training client == one backend registration.""" + + model_id: str # public: "{session_id}:train:{model_seq_id}" (official shape) + session_id: str + model_seq_id: int + name: str # backend adapter name + registration_id: str + base_model: str + rank: int + fingerprint: str + + @property + def rid8(self) -> str: + return self.registration_id[:8] + + +class ModelStore: + def __init__(self) -> None: + self.by_model_id: dict[str, ModelRecord] = {} + + def add(self, record: ModelRecord) -> None: + self.by_model_id[record.model_id] = record + + def get(self, model_id: str) -> ModelRecord | None: + return self.by_model_id.get(model_id) + + +@dataclass +class FutureRecord: + """One retrievable request_id. ``terminal`` holds the exact JSON body to + replay once resolved; until then ``kind`` picks the resolver.""" + + request_id: str + kind: str # "operation" | "create_model" | "unload_model" | "sample" + fingerprint: str + model: ModelRecord | None = None + operation_id: str | None = None + operation_kind: str | None = None + # forward results need a metrics recompute from the request payload (the + # backend attaches metrics to forward_backward only); dropped when terminal. + forward_payload: dict | None = None + # save/load bookkeeping minted at submit time. + tinker_path: str | None = None + backend_target: dict | None = None + # ephemeral publish: the sampling session to mint at completion. + sampling_session_id: str | None = None + terminal: dict | None = None + created_at: float = field(default_factory=time.time) + # Lifecycle observability (metrics + the orphan reaper): when the client + # last long-polled this record, when the first sub-generation finished, + # and when the record turned terminal. + last_polled_at: float = field(default_factory=time.time) + first_result_at: float | None = None + resolved_at: float | None = None + # The reaper writes WHY it cancelled here before task.cancel(); the + # sample task's CancelledError handler resolves with this message so a + # late poll sees the true reason, not a generic shutdown notice. + cancel_reason: str | None = None + # Exception class of a task failure (terminal-failures-by-class metric). + failure_class: str | None = None + + def resolve(self, body: dict) -> dict: + self.terminal = body + self.forward_payload = None + self.resolved_at = time.time() + return body + + +class FutureStore: + """request_id -> FutureRecord with bounded retention of delivered + terminal results (replay window for lost responses). Eviction leaves a + compact tombstone (request_id -> fingerprint): the record's identity + outlives its bytes, so a late identical retry gets a truthful typed 410 + instead of silently re-executing (samples would re-generate, training + ordinals would collide with the ledger) or a misleading conflict.""" + + def __init__(self, max_delivered: int = 4096, max_expired: int = 65536) -> None: + self.records: dict[str, FutureRecord] = {} + self.max_delivered = max_delivered + self.max_expired = max_expired + self._delivered: OrderedDict[str, None] = OrderedDict() + self._expired: OrderedDict[str, str] = OrderedDict() + # Reaped-before-delivery tombstones (terminal results whose client + # never retrieved them within the retention TTL): same identity + # preservation as ``_expired``, but a late retry must hear the truth + # — the result was reaped unclaimed, not delivered. + self._reaped: OrderedDict[str, str] = OrderedDict() + + def put(self, record: FutureRecord) -> FutureRecord: + self.records[record.request_id] = record + return record + + def get(self, request_id: str) -> FutureRecord | None: + return self.records.get(request_id) + + def expired_fingerprint(self, request_id: str) -> str | None: + return self._expired.get(request_id) + + def reaped_fingerprint(self, request_id: str) -> str | None: + return self._reaped.get(request_id) + + def is_delivered(self, request_id: str) -> bool: + return request_id in self._delivered + + def existing(self, request_id: str, fingerprint: str) -> FutureRecord | None: + """The idempotent-retry lookup: same id + same fingerprint replays, + same id + different content conflicts, delivered-then-evicted (or + reaped-unclaimed) expires.""" + record = self.records.get(request_id) + if record is None: + expired = self._expired.get(request_id) + if expired is not None: + _check_fingerprint("request", request_id, expired, fingerprint) + raise ExpiredError( + f"request '{request_id}' was already delivered and its replay window expired; " + "the original result cannot be reproduced" + ) + reaped = self._reaped.get(request_id) + if reaped is not None: + _check_fingerprint("request", request_id, reaped, fingerprint) + raise ExpiredError( + f"request '{request_id}' completed but was never retrieved within its retention " + "TTL and was reaped; the original result cannot be reproduced" + ) + return None + _check_fingerprint("request", request_id, record.fingerprint, fingerprint) + return record + + def reap_undelivered(self, record: FutureRecord) -> None: + """Evict a terminal-but-never-delivered record, keeping its identity + as a compact tombstone: the reaper frees the (potentially large) + terminal bytes without ever freeing the identity — a late identical + retry answers a typed 410 instead of silently re-executing.""" + self.records.pop(record.request_id, None) + self._reaped[record.request_id] = record.fingerprint + self._reaped.move_to_end(record.request_id) + while len(self._reaped) > self.max_expired: + self._reaped.popitem(last=False) + + def mark_delivered(self, record: FutureRecord) -> None: + if record.terminal is None: + return + self._delivered[record.request_id] = None + self._delivered.move_to_end(record.request_id) + while len(self._delivered) > self.max_delivered: + evicted, _ = self._delivered.popitem(last=False) + dropped = self.records.pop(evicted, None) + if dropped is not None: + self._expired[evicted] = dropped.fingerprint + while len(self._expired) > self.max_expired: + self._expired.popitem(last=False) + + +@dataclass +class CheckpointRecord: + tinker_path: str # public "tinker://{run}/weights/{tag}" + backend_path: str # trainer-side state directory + name: str + registration_id: str + base_model: str + rank: int + step: int + + +class CheckpointCatalog: + """tinker:// URI -> backend state path. In-memory: paths minted by this + controller lifetime resolve; the artifacts themselves persist on disk.""" + + def __init__(self) -> None: + self.records: dict[str, CheckpointRecord] = {} + + def add(self, record: CheckpointRecord) -> None: + self.records[record.tinker_path] = record + + def get(self, tinker_path: str) -> CheckpointRecord | None: + return self.records.get(tinker_path) + + +@dataclass +class SamplingSessionRecord: + sampling_session_id: str + session_id: str + fingerprint: str + base_model: str + # None for base-model sessions; set for ephemeral LoRA publishes. + name: str | None = None + registration_id: str | None = None + serving_name: str | None = None + serving_version: int | None = None + # Compact spent-sequence fence: sample identities outlive the bounded + # future/tombstone retention. Every seq <= spent_fence has executed; + # spent_sparse holds executed seqs above the fence (out-of-order arrival + # gaps only, so it stays tiny for the SDK's monotonic counters). A retry + # of a spent seq whose bytes and tombstone are both gone gets a typed + # terminal failure instead of silently re-running the generation. + spent_fence: int = -1 + spent_sparse: set = field(default_factory=set) + + def is_spent(self, seq_id: int) -> bool: + return seq_id <= self.spent_fence or seq_id in self.spent_sparse + + def mark_spent(self, seq_id: int) -> None: + if self.is_spent(seq_id): + return + self.spent_sparse.add(seq_id) + while self.spent_fence + 1 in self.spent_sparse: + self.spent_fence += 1 + self.spent_sparse.discard(self.spent_fence) + + +class SamplingSessionStore: + def __init__(self) -> None: + self.records: dict[str, SamplingSessionRecord] = {} + + def add(self, record: SamplingSessionRecord) -> SamplingSessionRecord: + self.records[record.sampling_session_id] = record + return record + + def get(self, sampling_session_id: str) -> SamplingSessionRecord | None: + return self.records.get(sampling_session_id) + + def existing(self, sampling_session_id: str, fingerprint: str) -> SamplingSessionRecord | None: + record = self.records.get(sampling_session_id) + if record is None: + return None + _check_fingerprint("sampling session", sampling_session_id, record.fingerprint, fingerprint) + return record + + def remove_for_sessions(self, session_ids: set[str]) -> None: + """Retire child sampler namespaces for multiple parents in one pass.""" + if not session_ids: + return + for sampling_session_id in [key for key, record in self.records.items() if record.session_id in session_ids]: + del self.records[sampling_session_id] diff --git a/miles/ray/tinker_frontend/translation.py b/miles/ray/tinker_frontend/translation.py new file mode 100644 index 00000000000..6819b401e62 --- /dev/null +++ b/miles/ray/tinker_frontend/translation.py @@ -0,0 +1,280 @@ +"""Official tinker payloads <-> backend operation payloads. + +The SDK's Datum is (model_input tokens, per-token ``loss_fn_inputs`` of +length N, next-token targets); the backend's sample is (tokens, trailing +``response_length`` span, per-token channels on that span). The bridge: + + input_tokens = concat(encoded_text chunks) # length N + target_tokens = loss_fn_inputs["target_tokens"] # length N + tokens = input_tokens + [target_tokens[-1]] # length N + 1 + response_length = N + +so the trainer's shifted logprob for response position i is exactly the +logprob of ``tokens[i+1]`` given the first i+1 tokens — the official +"logprob of target i given the input prefix" for every position where +``target_tokens[i] == input_tokens[i+1]``. Positions with a non-zero loss +contribution MUST satisfy that next-token alignment (rejected otherwise); +zero-weighted positions (canonical RL pads prompt targets with 0) are +normalized to the next input token, and their returned logprob refers to +that normalized target. + +Every rejection raises ``UserInputError`` — the caller records it as a +terminal FAILED(user) operation so the client's ordinal is still consumed. +""" + +import math + +from miles.ray.tinker_frontend import wire + +SUPPORTED_LOSS_FNS = ("cross_entropy", "importance_sampling", "ppo") + +# Official loss_fn_inputs channel -> backend per-token channel. +_CHANNEL_TO_BACKEND = { + "weights": "loss_weights", + "advantages": "advantages", + "logprobs": "rollout_log_probs", +} +_REQUIRED_CHANNELS = { + "cross_entropy": ("weights",), + "importance_sampling": ("logprobs", "advantages"), + "ppo": ("logprobs", "advantages"), +} +# Which channel decides whether a position contributes loss (and therefore +# must be a true next-token target). +_ACTIVE_CHANNEL = {"cross_entropy": "weights", "importance_sampling": "advantages", "ppo": "advantages"} + + +class UserInputError(ValueError): + """Typed client-payload rejection (never a server fault).""" + + +def _decode_1d(name: str, where: str, tensor: wire.TensorData, expect_len: int, integer: bool) -> list: + if tensor.sparse_crow_indices is not None or tensor.sparse_col_indices is not None: + raise UserInputError(f"{where}: sparse (CSR) '{name}' is not supported in v1 — send dense 1-D tensors") + if tensor.shape is not None and (len(tensor.shape) != 1 or tensor.shape[0] != len(tensor.data)): + raise UserInputError( + f"{where}: '{name}' must be 1-D with shape matching its data " + f"(got shape={tensor.shape}, len={len(tensor.data)}) — nested/top-K targets are not supported in v1" + ) + if len(tensor.data) != expect_len: + raise UserInputError( + f"{where}: '{name}' must have one value per input token (got {len(tensor.data)}, want {expect_len})" + ) + values = [] + for value in tensor.data: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise UserInputError(f"{where}: '{name}' must contain only finite numbers") + if integer: + if isinstance(value, float) and not value.is_integer(): + raise UserInputError(f"{where}: '{name}' must contain integer token ids") + value = int(value) + if value < 0: + # No tokenizer has negative ids; vocab UPPER bounds are the + # engine's to enforce (the frontend never loads the tokenizer). + raise UserInputError(f"{where}: '{name}' token ids must be non-negative (got {value})") + values.append(value) + else: + value = float(value) + if not math.isfinite(value): + raise UserInputError(f"{where}: '{name}' must contain only finite numbers") + values.append(value) + return values + + +def _input_tokens(where: str, model_input: wire.ModelInput) -> list[int]: + tokens: list[int] = [] + for chunk in model_input.chunks: + if chunk.type != "encoded_text": + raise UserInputError(f"{where}: model_input chunk type '{chunk.type}' is not supported in v1 (text-only)") + tokens.extend(chunk.tokens) + if not tokens or not all(isinstance(t, int) and not isinstance(t, bool) for t in tokens): + raise UserInputError(f"{where}: model_input must carry at least one encoded-text token") + if any(t < 0 for t in tokens): + raise UserInputError(f"{where}: model_input token ids must be non-negative") + return tokens + + +def datum_to_sample(index: int, datum: wire.Datum, loss_fn: str) -> dict: + where = f"data[{index}]" + input_tokens = _input_tokens(where, datum.model_input) + n = len(input_tokens) + + known = {"target_tokens", *_CHANNEL_TO_BACKEND} + if unknown := sorted(set(datum.loss_fn_inputs) - known): + raise UserInputError(f"{where}: unsupported loss_fn_inputs {unknown}; v1 accepts {sorted(known)}") + if "target_tokens" not in datum.loss_fn_inputs: + raise UserInputError(f"{where}: loss_fn_inputs must include 'target_tokens'") + for required in _REQUIRED_CHANNELS[loss_fn]: + if required not in datum.loss_fn_inputs: + raise UserInputError(f"{where}: loss_fn '{loss_fn}' requires loss_fn_inputs['{required}']") + + targets = _decode_1d("target_tokens", where, datum.loss_fn_inputs["target_tokens"], n, integer=True) + channels = { + official: _decode_1d(official, where, tensor, n, integer=False) + for official, tensor in datum.loss_fn_inputs.items() + if official != "target_tokens" + } + + # Positions that contribute loss must be true next-token targets; the + # rest (canonical RL zero-weights its prompt span) are normalized to the + # next input token, which is what their returned logprob refers to. + active = channels[_ACTIVE_CHANNEL[loss_fn]] + for i in range(n - 1): + if active[i] != 0.0 and targets[i] != input_tokens[i + 1]: + raise UserInputError( + f"{where}: target_tokens[{i}]={targets[i]} has non-zero loss weight but is not the next input " + f"token ({input_tokens[i + 1]}); v1 serves next-token targets only" + ) + + sample = { + "tokens": input_tokens + [targets[-1]], + "response_length": n, + "loss_mask": [1] * n, + } + for official, backend_channel in _CHANNEL_TO_BACKEND.items(): + if official in channels: + sample[backend_channel] = channels[official] + return sample + + +def fb_input_to_payload(fb_input: wire.ForwardBackwardInput) -> dict: + """Backend payload for one forward_backward/forward operation. The loss + spec rides along for forward too: the trainer ignores it structurally + (no gradients), and result translation recomputes the loss metrics from + it (the backend attaches metrics only to forward_backward results).""" + if fb_input.loss_fn not in SUPPORTED_LOSS_FNS: + raise UserInputError( + f"loss_fn '{fb_input.loss_fn}' is not supported in v1; supported: {', '.join(SUPPORTED_LOSS_FNS)}" + ) + if not fb_input.data: + raise UserInputError("forward_backward needs at least one datum") + loss: dict = {"loss_fn": fb_input.loss_fn} + if fb_input.loss_fn_config is not None: + loss["loss_fn_config"] = dict(fb_input.loss_fn_config) + return { + "samples": [datum_to_sample(i, datum, fb_input.loss_fn) for i, datum in enumerate(fb_input.data)], + "loss": loss, + } + + +def adam_params_to_payload(adam: wire.AdamParams) -> dict: + return {"adam_params": adam.model_dump()} + + +# ---------------- results: backend operation -> SDK terminal JSON ---------------- + + +def fb_result_to_response(result: dict, payload: dict | None = None) -> dict: + """ForwardBackwardOutput JSON. logprobs arrive in the operation's datum + order, one row per datum, one value per input token. ``payload`` (the + operation's request) triggers a metrics recompute for forward results, + which the backend completes without metrics.""" + logprobs = result.get("logprobs") or [] + metrics = result.get("metrics") + if metrics is None and payload is not None: + from miles.ray.multi_lora.backend import operation_result_metrics + + metrics = operation_result_metrics(payload, logprobs) + return { + "type": "forward_backward", + "loss_fn_output_type": "ArrayRecord", + "loss_fn_outputs": [{"logprobs": {"data": row, "dtype": "float32", "shape": [len(row)]}} for row in logprobs], + "metrics": metrics or {}, + } + + +def optim_result_to_response(result: dict) -> dict: + metrics = { + key: float(value) + for key, value in (result or {}).items() + if isinstance(value, (int, float)) and not isinstance(value, bool) + } + return {"type": "optim_step", "metrics": metrics} + + +def save_weights_result_to_response(tinker_path: str) -> dict: + return {"type": "save_weights", "path": tinker_path} + + +def load_weights_result_to_response(tinker_path: str, model_id: str) -> dict: + return {"type": "load_weights", "path": tinker_path, "model_id": model_id} + + +def sampler_publish_result_to_response(sampling_session_id: str) -> dict: + return {"type": "save_weights_for_sampler", "path": None, "sampling_session_id": sampling_session_id} + + +# ---------------- sampling: SDK request <-> sglang router ---------------- + +_FINISH_TO_STOP_REASON = {"stop": "stop", "length": "length"} + + +def sampling_params_to_sglang(params: wire.SamplingParams) -> dict: + """Per-request sglang sampling_params. ``seed`` is handled by the caller + (each fanned-out sample i gets ``sampling_seed = seed + i``: deterministic + per request, still diverse across num_samples).""" + if params.max_tokens is None or params.max_tokens < 1: + raise UserInputError("sampling_params.max_tokens is required (>= 1) in v1") + if not math.isfinite(params.temperature) or params.temperature < 0: + raise UserInputError("sampling_params.temperature must be a non-negative finite number") + if not math.isfinite(params.top_p) or not 0 < params.top_p <= 1: + raise UserInputError("sampling_params.top_p must be a finite number in (0, 1]") + if params.top_k != -1 and params.top_k < 1: + raise UserInputError("sampling_params.top_k must be -1 or at least 1") + if params.seed is not None and not -(2**63) <= params.seed < 2**63: + raise UserInputError("sampling_params.seed must fit in a signed 64-bit integer") + sglang_params: dict = { + "max_new_tokens": params.max_tokens, + "temperature": params.temperature, + "top_p": params.top_p, + "top_k": params.top_k, + } + stop = params.stop + if stop is not None: + if isinstance(stop, str): + sglang_params["stop"] = [stop] + elif all(isinstance(s, str) for s in stop): + sglang_params["stop"] = list(stop) + elif all(isinstance(s, int) and not isinstance(s, bool) for s in stop): + if any(s < 0 for s in stop): + raise UserInputError("sampling_params.stop token ids must be non-negative") + sglang_params["stop_token_ids"] = list(stop) + else: + raise UserInputError("sampling_params.stop must be a string, list of strings, or list of token ids") + return sglang_params + + +def generation_to_sequence(generation: dict) -> dict: + """One sglang /generate response -> one SampledSequence JSON.""" + meta = generation.get("meta_info") or {} + finish = (meta.get("finish_reason") or {}).get("type") + stop_reason = _FINISH_TO_STOP_REASON.get(finish) + if stop_reason is None: + raise RuntimeError(f"generation finished with '{finish}'") + token_logprobs = meta.get("output_token_logprobs") or [] + return { + "stop_reason": stop_reason, + "tokens": [int(entry[1]) for entry in token_logprobs], + "logprobs": [float(entry[0]) for entry in token_logprobs], + } + + +def prompt_logprobs_from_generation(generation: dict, prompt_len: int) -> list[float | None]: + """meta_info.input_token_logprobs (logprob_start_len=0) -> one float-or-None per prompt token.""" + entries = (generation.get("meta_info") or {}).get("input_token_logprobs") + if not entries: + raise RuntimeError("the engine returned no input_token_logprobs for a prompt_logprobs request") + if len(entries) != prompt_len: + raise RuntimeError(f"the engine returned {len(entries)} prompt logprobs for {prompt_len} prompt tokens") + # The first entry has no context, so sglang reports None there; keep it. + return [None if entry[0] is None else float(entry[0]) for entry in entries] + + +def sequences_to_sample_response(sequences: list[dict], prompt_logprobs: list[float | None] | None = None) -> dict: + return { + "type": "sample", + "sequences": sequences, + "prompt_logprobs": prompt_logprobs, + "topk_prompt_logprobs": None, + "prompt_cache_hit_tokens": 0, + } diff --git a/miles/ray/tinker_frontend/wire.py b/miles/ray/tinker_frontend/wire.py new file mode 100644 index 00000000000..b0e3b18ab35 --- /dev/null +++ b/miles/ray/tinker_frontend/wire.py @@ -0,0 +1,227 @@ +"""Wire models of the tinker SDK's REST protocol (server side). + +Mirrors the request shapes ``tinker==0.24.1`` actually POSTs (verified from +the wheel source and captured traffic, not from documentation). Requests are +parsed permissively (``extra="ignore"``) so additive SDK fields never break +the server; everything the backend relies on is validated explicitly in the +translation layer. Responses are plain dicts built by the service — the SDK +deserializes JSON terminal results against its own pydantic models, so the +literal ``type`` discriminators below must match its expectations exactly. +""" + +from typing import Any + +from pydantic import BaseModel, ConfigDict + +TINKER_SDK_VERSION_PIN = "0.24.1" + +# Flags returned from /api/v1/client/config. They steer the 0.24.1 SDK onto +# the pure-JSON protocol this frontend implements: +# - proto_write_fwdbwd=False keeps forward_backward on JSON (the wheel's own +# default) and forward on the legacy JSON /api/v1/forward route; +# - fwd_via_fwdbwd must then also be False (the SDK asserts forward_only +# requires the proto path); +# - parallel_fwdbwd_chunks=True lets the SDK post fwdbwd chunks concurrently, +# first chunk last — exactly the out-of-order arrival the backend ledger +# gap-buffers by design; +# - use_pyqwest_transport=False keeps the SDK on the plain httpx transport. +CLIENT_CONFIG_FLAGS = { + "pjwt_auth_enabled": False, + "credential_default_source": "api_key", + "parallel_fwdbwd_chunks": True, + "proto_write_fwdbwd": False, + "proto_compress_fwdbwd": False, + "fwd_via_fwdbwd": False, + "use_pyqwest_transport": False, + "create_model_via_load_weights": False, + "sample_no_retries": False, + "sample_max_concurrent_requests": 64, +} + + +class WireModel(BaseModel): + # protected_namespaces cleared: the protocol is full of model_* fields. + model_config = ConfigDict(extra="ignore", protected_namespaces=()) + + +class CreateSessionRequest(WireModel): + tags: list[str] = [] + user_metadata: dict[str, Any] | None = None + sdk_version: str = "" + project_id: str | None = None + + +class SessionHeartbeatRequest(WireModel): + session_id: str + + +class ClientConfigRequest(WireModel): + sdk_version: str = "" + + +class LoraConfig(WireModel): + rank: int + seed: int | None = None + train_unembed: bool = True + train_mlp: bool = True + train_attn: bool = True + + +class CreateModelRequest(WireModel): + session_id: str + model_seq_id: int + base_model: str + user_metadata: dict[str, Any] | None = None + lora_config: LoraConfig | None = None + + +class GetInfoRequest(WireModel): + model_id: str + + +class UnloadModelRequest(WireModel): + model_id: str + + +class TensorData(WireModel): + data: list[int | float] + dtype: str = "float32" + shape: list[int] | None = None + sparse_crow_indices: list[int] | None = None + sparse_col_indices: list[int] | None = None + + +class ModelInputChunk(WireModel): + # Non-text chunk types (image, dmel, ...) carry other fields; the type + # tag alone is enough to reject them at the boundary. + type: str = "encoded_text" + tokens: list[int] = [] + + +class ModelInput(WireModel): + chunks: list[ModelInputChunk] + + +class Datum(WireModel): + model_input: ModelInput + loss_fn_inputs: dict[str, TensorData] + + +class ForwardBackwardInput(WireModel): + data: list[Datum] + loss_fn: str + loss_fn_config: dict[str, float] | None = None + + +class ForwardBackwardRequest(WireModel): + forward_backward_input: ForwardBackwardInput + model_id: str + seq_id: int | None = None + + +class ForwardRequest(WireModel): + forward_input: ForwardBackwardInput + model_id: str + seq_id: int | None = None + + +class AdamParams(WireModel): + learning_rate: float = 1e-4 + beta1: float = 0.9 + beta2: float = 0.95 + eps: float = 1e-12 + weight_decay: float = 0.0 + grad_clip_norm: float = 0.0 + + +class OptimStepRequest(WireModel): + adam_params: AdamParams + model_id: str + seq_id: int | None = None + + +class SaveWeightsRequest(WireModel): + model_id: str + path: str | None = None + seq_id: int | None = None + ttl_seconds: int | None = None + overwrite: bool = False + + +class LoadWeightsRequest(WireModel): + model_id: str | None = None + seq_id: int | None = None + session_id: str | None = None + model_seq_id: int | None = None + base_model: str | None = None + user_metadata: dict[str, Any] | None = None + path: str + optimizer: bool + weights_access_token: str | None = None + + +class SaveWeightsForSamplerRequest(WireModel): + model_id: str + path: str | None = None + sampling_session_seq_id: int | None = None + seq_id: int | None = None + ttl_seconds: int | None = None + + +class WeightsInfoRequest(WireModel): + tinker_path: str + + +class CreateSamplingSessionRequest(WireModel): + session_id: str + sampling_session_seq_id: int + base_model: str | None = None + model_path: str | None = None + + +class SamplingParams(WireModel): + max_tokens: int | None = None + seed: int | None = None + stop: str | list[str] | list[int] | None = None + temperature: float = 1.0 + top_k: int = -1 + top_p: float = 1.0 + + +class SampleRequest(WireModel): + num_samples: int = 1 + prompt: ModelInput + sampling_params: SamplingParams + base_model: str | None = None + model_path: str | None = None + sampling_session_id: str | None = None + seq_id: int | None = None + prompt_logprobs: bool | None = None + topk_prompt_logprobs: int = 0 + + +class FutureRetrieveRequest(WireModel): + request_id: str + allow_metadata_only: bool = False + model_id: str | None = None + + +def untyped_future(request_id: str, model_id: str | None = None) -> dict: + body: dict = {"request_id": request_id} + if model_id is not None: + body["model_id"] = model_id + return body + + +def try_again(queue_state: str = "active", reason: str | None = None) -> dict: + body: dict = {"type": "try_again", "queue_state": queue_state} + if reason is not None: + body["queue_state_reason"] = reason + return body + + +def terminal_failure(error: str, category: str = "user") -> dict: + # RequestErrorCategory on the SDK side accepts exactly unknown|server|user. + if category not in ("unknown", "server", "user"): + category = "unknown" + return {"error": error, "category": category} diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index 373533f679e..9ea69139e10 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -1,7 +1,7 @@ from __future__ import annotations from argparse import Namespace -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any from miles.rollout.data_source import DataSource @@ -49,11 +49,21 @@ def evaluation(self): return True +@dataclass(frozen=True) +class RolloutPostprocessOptions: + """Postprocess policy declared by the rollout fn; pad_to_dp zero-weight pads to the DP grid instead of trimming.""" + + pad_to_dp: bool = False + + # TODO make it frozen @dataclass class RolloutFnTrainOutput: samples: list[list[Sample]] metrics: dict[str, Any] = None + metadata: dict[str, Any] | None = None + conversion_metadata: dict[str, Any] | None = None + postprocess: RolloutPostprocessOptions = field(default_factory=RolloutPostprocessOptions) # TODO make it frozen diff --git a/miles/rollout/generate_utils/sample_utils.py b/miles/rollout/generate_utils/sample_utils.py index d9d3cb2fd91..e4ea9e357c3 100644 --- a/miles/rollout/generate_utils/sample_utils.py +++ b/miles/rollout/generate_utils/sample_utils.py @@ -161,6 +161,10 @@ def _merge_metadata(): rollout_log_probs=a.rollout_log_probs + [0.0] * obs_len + b.rollout_log_probs, teacher_log_probs=_merge_optional_per_token("teacher_log_probs"), opd_reverse_kl=_merge_optional_per_token("opd_reverse_kl"), + # Tinker per-token channels: response-aligned like the OPD lists; + # zero weight/advantage over the injected observation span. + loss_weights=_merge_optional_per_token("loss_weights"), + advantages=_merge_optional_per_token("advantages"), rollout_routed_experts=b.rollout_routed_experts, rollout_indexer_topk=b.rollout_indexer_topk, remove_sample=_merge_equal_value("remove_sample"), diff --git a/miles/rollout/multi_lora/async_rollout.py b/miles/rollout/multi_lora/async_rollout.py deleted file mode 100644 index 21bb36179c6..00000000000 --- a/miles/rollout/multi_lora/async_rollout.py +++ /dev/null @@ -1,584 +0,0 @@ -"""Fully-async multi-LoRA rollout: a background producer fills per-adapter buffers; batches are collected -round-robin in ``min_groups_per_dp_split`` multiples without overshooting any adapter's remaining batch.""" - -import asyncio -import itertools -import logging -import threading -import time -from collections import defaultdict, deque -from collections.abc import Callable -from dataclasses import dataclass -from typing import Any - -from miles.ray.multi_lora.controller import AdaptersCache, get_multi_lora_controller -from miles.rollout.base_types import RolloutFnTrainOutput -from miles.rollout.filter_hub.base_types import call_dynamic_filter -from miles.rollout.generate_utils.prefill_logprobs import recompute_samples_rollout_logprobs_via_prefill -from miles.rollout.sglang_rollout import GenerateState, generate_and_rm_group, get_model_url -from miles.utils.async_utils import run -from miles.utils.metric_utils import compute_statistics, dict_add_prefix -from miles.utils.misc import load_function -from miles.utils.multi_lora import EmptyBatchTimeoutError, min_groups_per_dp_split -from miles.utils.tracking_utils import tracking -from miles.utils.types import Sample - -logger = logging.getLogger(__name__) - -GenerateFn = Callable[..., Any] - -# Generate fns may return several samples per rollout; the manager flattens later. -Group = list[Sample | list[Sample]] - - -def iter_group_samples(group: Group): - return itertools.chain.from_iterable(item if isinstance(item, list) else (item,) for item in group) - - -def first_sample(group: Group) -> Sample: - return group[0][0] if isinstance(group[0], list) else group[0] - - -def group_adapter_name(group: Group) -> str | None: - head = first_sample(group) if group else None - return head.adapter.name if head is not None and head.adapter else None - - -def group_sample_count(group: Group) -> int: - return sum(1 for _ in iter_group_samples(group)) - - -# Safety valve, same convention as fully_async's queue.Queue(maxsize=1000): -# never hit in practice, just bounds memory if training stalls entirely. -MAX_BUFFERED_GROUPS = 1000 -EMPTY_BATCH_TIMEOUT_S = 30.0 - - -class GroupBuffer: - """One adapter's FIFO of completed prompt groups; bounded — the oldest group is dropped when full.""" - - def __init__(self) -> None: - self._groups: deque[Group] = deque(maxlen=MAX_BUFFERED_GROUPS) - - def __len__(self) -> int: - return len(self._groups) - - def put(self, group: Group) -> None: - self._groups.append(group) - - def get(self, n_groups: int) -> list[Group]: - """Remove and return the n oldest groups (queue.Queue-style API).""" - return [self._groups.popleft() for _ in range(n_groups)] - - def drop_foreign(self, registration_id: str) -> int: - """Drop groups stamped by a different registration of this adapter - name: an in-flight generation of a retired tenant can land after the - buffer was reset for a same-name re-registration. Unstamped groups - (no adapter view at submission time) are kept. Returns the drop count.""" - if not self._groups: - return 0 - kept: deque[Group] = deque(maxlen=MAX_BUFFERED_GROUPS) - dropped = 0 - for group in self._groups: - stamped = first_sample(group).metadata.get("registration_id") - if stamped is not None and stamped != registration_id: - dropped += 1 - else: - kept.append(group) - self._groups = kept - return dropped - - def drop_stale(self, current_version: int, max_staleness: int | None) -> list[int]: - """Drop groups generated too many weight versions ago; returns the - staleness of each dropped group (for metrics).""" - if max_staleness is None or not self._groups: - return [] - kept: deque[Group] = deque(maxlen=MAX_BUFFERED_GROUPS) - dropped: list[int] = [] - for group in self._groups: - stamped = first_sample(group).metadata.get("slot_version") - staleness = current_version - stamped if stamped is not None else 0 - if stamped is not None and staleness > max_staleness: - for sample in iter_group_samples(group): - sample.reset_for_retry() - dropped.append(staleness) - else: - kept.append(group) - self._groups = kept - return dropped - - -@dataclass -class TrainBatch: - """One train batch: the groups for one train call, with its per-adapter bookkeeping.""" - - groups: list[Group] - group_counts: dict[str, int] # prompt groups per adapter in this batch - step_names: list[str] # adapters whose adapter batch completes -> they step - step_slots: list[int] - - -def remaining_groups(adapter) -> int: - """Groups still needed to complete the adapter's batch.""" - remaining = adapter.config.rollout_batch_size - adapter.accumulated_groups - assert remaining > 0, ( - f"adapter '{adapter.name}' accumulated_groups={adapter.accumulated_groups} >= " - f"rollout_batch_size={adapter.config.rollout_batch_size}; batch accounting drifted" - ) - return remaining - - -async def process_group( - args, group: list[Sample], sampling_params: dict, generate_fn: GenerateFn, data_source -) -> Group | None: - """Generate a group; returns None for aborted groups. The slot version is - stamped at submission time (what the staleness filter compares against).""" - adapter_name = group[0].adapter.name if group and group[0].adapter else None - submission_version: int | None = None - submission_registration: str | None = None - if adapter_name is not None: - adapter = await AdaptersCache().get(adapter_name) - submission_version = adapter.version if adapter is not None else None - submission_registration = adapter.registration_id if adapter is not None else None - - if submission_version is not None: - for s in group: - s.metadata["slot_version"] = submission_version - s.metadata["registration_id"] = submission_registration - - result = await generate_fn(args, group, sampling_params) - - if submission_version is not None: - for s in iter_group_samples(result): - s.metadata["slot_version"] = submission_version - s.metadata["registration_id"] = submission_registration - - if any(s.status == Sample.Status.ABORTED for s in iter_group_samples(result)): - for s in iter_group_samples(result): - s.reset_for_retry() - # Re-queuing is not wired up (the per-adapter source is read-only). - return None - return result - - -class MultiLoRAWorkerMetrics: - """Cross-batch metric state; locked because the producer thread records while the trainer thread drains.""" - - def __init__(self) -> None: - self.lock = threading.Lock() - self.dynamic_filter_drop_counts: dict[str, int] = defaultdict(int) - # Staleness of dropped groups per adapter, drained every batch. - self.staleness_values: dict[str, list[int]] = defaultdict(list) - # Per-adapter shipped-sample values, flushed as step statistics when the adapter steps. - self.step_rewards: dict[str, list[float]] = defaultdict(list) - self.step_response_lens: dict[str, list[float]] = defaultdict(list) - # Per-sample mean engine log prob (rough per-adapter entropy trend). - self.step_log_prob_means: dict[str, list[float]] = defaultdict(list) - # Group outcomes for zero-std rates: shipped group counts and each uniform-reward group's reward. - self.step_group_counts: dict[str, int] = defaultdict(int) - self.step_zero_std_rewards: dict[str, list[float]] = defaultdict(list) - - def record_dynamic_filter_drop(self, reason: str) -> None: - with self.lock: - self.dynamic_filter_drop_counts[reason] += 1 - - def record_stale_drops(self, name: str, staleness_values: list[int]) -> None: - with self.lock: - self.staleness_values[name] += staleness_values - - def pop_stale_drops(self) -> dict[str, list[int]]: - """Drain the staleness values of groups dropped since the last batch.""" - with self.lock: - drained = dict(self.staleness_values) - self.staleness_values.clear() - return drained - - def record_shipped_samples( - self, args, data: list[Group], step_names: list[str], adapters: dict - ) -> dict[str, dict[str, float]]: - """Accumulate shipped rewards/response lengths per adapter; flush whole-adapter-batch statistics - for adapters stepping with this batch. Returns {adapter name: flushed metrics}.""" - with self.lock: - for group in data: - name = group_adapter_name(group) - if name is None: - continue - group_rewards = [] - for sample in iter_group_samples(group): - reward = sample.get_reward_value(args) - group_rewards.append(reward) - self.step_rewards[name].append(reward) - self.step_response_lens[name].append(sample.effective_response_length) - if sample.rollout_log_probs: - self.step_log_prob_means[name].append( - sum(sample.rollout_log_probs) / len(sample.rollout_log_probs) - ) - self.step_group_counts[name] += 1 - if len(group_rewards) > 1 and all(reward == group_rewards[0] for reward in group_rewards): - self.step_zero_std_rewards[name].append(round(group_rewards[0], 1)) - - flushed: dict[str, dict[str, float]] = {} - for name in step_names: - rewards = self.step_rewards.pop(name, []) - response_lens = self.step_response_lens.pop(name, []) - log_prob_means = self.step_log_prob_means.pop(name, []) - total_groups = self.step_group_counts.pop(name, 0) - zero_std_rewards = self.step_zero_std_rewards.pop(name, []) - if not rewards: - continue - expected = adapters[name].config.adapter_global_batch_size - if len(rewards) != expected: - logger.warning( - f"Adapter '{name}' stepped with {len(rewards)} shipped samples, expected " - f"adapter_global_batch_size={expected}; batch accounting drifted" - ) - # Single-segment keys so "{name}/" matches the "{name}/*" glob (server globs one segment). - flushed[name] = { - **dict_add_prefix(compute_statistics(rewards), "raw_reward_"), - **dict_add_prefix(compute_statistics(response_lens), "response_len_"), - } - if log_prob_means: - flushed[name]["log_probs"] = sum(log_prob_means) / len(log_prob_means) - if total_groups: - zero = sum(1 for reward in zero_std_rewards if reward == 0.0) - one = sum(1 for reward in zero_std_rewards if reward == 1.0) - flushed[name]["zero_std_all_zero_percentage"] = zero / total_groups - flushed[name]["zero_std_all_one_percentage"] = one / total_groups - return flushed - - def discard_adapter(self, name: str) -> None: - """Drop a retired adapter's partial step accumulation.""" - with self.lock: - self.step_rewards.pop(name, None) - self.step_response_lens.pop(name, None) - self.step_log_prob_means.pop(name, None) - self.step_group_counts.pop(name, None) - self.step_zero_std_rewards.pop(name, None) - self.staleness_values.pop(name, None) - - def pop_metrics(self) -> dict[str, float]: - with self.lock: - metrics = { - f"rollout/dynamic_filter/drop_{reason}": count - for reason, count in self.dynamic_filter_drop_counts.items() - } - self.dynamic_filter_drop_counts.clear() - return metrics - - -class AsyncMultiLoRAWorker: - """Background producer filling bounded per-adapter completed-group buffers; - the collection loop pops from them via ``get_groups``.""" - - global_worker = None - worker_lock = threading.Lock() - - def __init__(self, args, data_source, generate_fn: GenerateFn, concurrency: int = None) -> None: - self.args = args - self.data_source = data_source - self.generate_fn = generate_fn - self.concurrency = concurrency or args.rollout_batch_size - self.running = True - self.worker_thread: threading.Thread | None = None - self.state = GenerateState(args) - self.dynamic_filter = ( - load_function(args.dynamic_sampling_filter_path) if args.dynamic_sampling_filter_path else None - ) - # Guards the buffers: the producer thread puts while get_groups (trainer side) pops. - self.buffer_lock = threading.Lock() - self.buffers: dict[str, GroupBuffer] = defaultdict(GroupBuffer) - # Round-robin cursor over adapters, persisting across get_groups calls and batches. - self.rotation: deque[str] = deque() - self.metrics = MultiLoRAWorkerMetrics() - # Last seen registration id per adapter name; a change means re-registration -> drop inherited state. - self.registrations: dict[str, str] = {} - # Set when run_loop dies; collect_batch surfaces it instead of a misleading empty-batch timeout. - self.failure: Exception | None = None - - @classmethod - def get_or_create(cls, args, data_source, generate_fn: GenerateFn, concurrency: int = None): - with cls.worker_lock: - if cls.global_worker is None or not cls.global_worker.worker_thread.is_alive(): - cls.global_worker = cls(args, data_source, generate_fn, concurrency) - cls.global_worker.start() - return cls.global_worker - - def start(self) -> None: - self.worker_thread = threading.Thread(target=self.thread_main, daemon=True) - self.worker_thread.start() - - def stop(self) -> None: - self.running = False - if self.worker_thread and self.worker_thread.is_alive(): - self.worker_thread.join(timeout=5) - - @classmethod - def stop_global(cls) -> None: - with cls.worker_lock: - if cls.global_worker is None: - return - cls.global_worker.stop() - cls.global_worker = None - - def thread_main(self) -> None: - asyncio.run(self.run_loop()) - - async def run_loop(self) -> None: - active: set[asyncio.Task] = set() - max_concurrent = self.concurrency - try: - while self.running: - done = {t for t in active if t.done()} - for t in done: - try: - t.result() - except Exception as e: - logger.warning(f"generate task failed: {e}") - active.discard(t) - - while len(active) < max_concurrent and self.running: - samples = self.data_source.get_samples(1) - if not samples: - break - active.add(asyncio.create_task(self.process_and_enqueue(samples[0]))) - - await asyncio.sleep(0.01) - except Exception as e: - # Typically the data source: this stops production for EVERY - # adapter, so record the cause for collect_batch to surface. - self.failure = e - logger.exception("multi-LoRA producer failed; generation is stopped") - finally: - for task in active: - task.cancel() - if active: - await asyncio.gather(*active, return_exceptions=True) - - async def process_and_enqueue(self, group: list[Sample]) -> None: - result = await process_group(self.args, group, self.state.sampling_params, self.generate_fn, self.data_source) - if result is None: - return - - filter_result = call_dynamic_filter(self.dynamic_filter, self.args, result) - if not filter_result.keep: - if filter_result.reason: - self.metrics.record_dynamic_filter_drop(filter_result.reason) - return - - adapter_name = group_adapter_name(result) - if adapter_name is None: - return - with self.buffer_lock: - self.buffers[adapter_name].put(result) - - def queue_size(self) -> int: - with self.buffer_lock: - return sum(len(buffer) for buffer in self.buffers.values()) - - def queue_sizes(self) -> dict[str, int]: - """Buffered (completed, not yet shipped) prompt groups per adapter.""" - with self.buffer_lock: - return {name: len(buffer) for name, buffer in self.buffers.items()} - - def get_groups( - self, snapshot: dict, num_samples: int, group_counts: dict[str, int] - ) -> tuple[list[Group], dict[str, int]]: - """Pop groups round-robin in ``min_groups_per_dp_split`` multiples until ``num_samples`` is covered or - nothing is poppable; returns them with an updated ``group_counts`` copy (prevents adapter overshoot).""" - adapters = {**snapshot["active"], **snapshot["retiring"]} - dp_size = self.args.multi_lora_dp_size - max_staleness = getattr(self.args, "max_weight_staleness", None) - group_counts = dict(group_counts) # updated copy; the argument is not modified - popped: list[Group] = [] - popped_samples = 0 - - with self.buffer_lock: - # Retired adapters: discard their buffered tail and partial reward stats. - for name in list(self.buffers): - if name not in adapters: - self.buffers.pop(name) - self.metrics.discard_adapter(name) - self.registrations.pop(name, None) - - # A re-registered name is a new tenant: drop buffered groups and - # partial stats inherited from the old tenant. - for name, adapter in adapters.items(): - previous = self.registrations.get(name) - if previous is not None and previous != adapter.registration_id: - self.buffers.pop(name, None) - self.metrics.discard_adapter(name) - logger.warning(f"Adapter '{name}' was re-registered; dropped the previous tenant's buffered state") - self.registrations[name] = adapter.registration_id - - # Keep the rotation in sync with live adapters. - self.rotation = deque(name for name in self.rotation if name in adapters) - for name in sorted(set(adapters) - set(self.rotation)): - self.rotation.append(name) - - while popped_samples < num_samples: - made_progress = False - for _ in range(len(self.rotation)): - name = self.rotation[0] - self.rotation.rotate(-1) - adapter = adapters[name] - buffer = self.buffers[name] - if dropped := buffer.drop_stale(adapter.version, max_staleness): - self.metrics.record_stale_drops(name, dropped) - # In-flight stragglers of a retired same-name tenant that - # landed after the re-registration sweep reset the buffer. - if foreign := buffer.drop_foreign(adapter.registration_id): - logger.warning(f"Dropped {foreign} buffered groups from a previous registration of '{name}'") - min_groups_per_pop = min_groups_per_dp_split(adapter.config.n_samples_per_prompt, dp_size) - trainable_groups = len(buffer) // min_groups_per_pop * min_groups_per_pop - remaining_allowed_groups = max(0, remaining_groups(adapter) - group_counts.get(name, 0)) - groups_to_pop = min(min_groups_per_pop, trainable_groups, remaining_allowed_groups) - if groups_to_pop <= 0: - continue - popped.extend(buffer.get(groups_to_pop)) - popped_samples += groups_to_pop * adapter.config.n_samples_per_prompt - group_counts[name] = group_counts.get(name, 0) + groups_to_pop - made_progress = True - break - if not made_progress: - break # a full pass over rotation yielded nothing - return popped, group_counts - - -async def collect_batch(args, worker: AsyncMultiLoRAWorker, snapshot: dict) -> TrainBatch: - """Pop group multiples until the batch reaches ``--global-batch-size`` samples, or it is non-empty and - stalls for ``--multi-lora-max-coalesce-wait-s`` (the target can be unreachable; ship what there is).""" - adapters = {**snapshot["active"], **snapshot["retiring"]} - target_samples = args.global_batch_size - wait_s = getattr(args, "multi_lora_max_coalesce_wait_s", 0.5) - empty_wait_s = getattr(args, "multi_lora_max_empty_wait_s", EMPTY_BATCH_TIMEOUT_S) - - collected: list[Group] = [] - group_counts: dict[str, int] = {} - total_samples = 0 - last_progress = time.time() - last_warning = time.time() - - while total_samples < target_samples: - if worker.failure is not None: - raise RuntimeError( - "multi-LoRA producer thread died; generation is stalled for every adapter" - ) from worker.failure - groups, group_counts = worker.get_groups(snapshot, target_samples - total_samples, group_counts) - if groups: - collected.extend(groups) - total_samples += sum(adapters[group_adapter_name(g)].config.n_samples_per_prompt for g in groups) - last_progress = time.time() - continue - stalled_s = time.time() - last_progress - if collected and stalled_s > wait_s: - break - if not collected and stalled_s > empty_wait_s: - raise EmptyBatchTimeoutError( - "No poppable groups collected before empty timeout; this likely means every live adapter is " - "below min_groups_per_dp_split (or sources are exhausted). " - f"queue={worker.queue_size()} active={sorted(snapshot['active'])} retiring={sorted(snapshot['retiring'])}" - ) - if not collected and time.time() - last_warning > 30: - logger.warning( - "No completed groups for 30s. " - f"queue={worker.queue_size()} active={sorted(snapshot['active'])} " - f"retiring={sorted(snapshot['retiring'])}" - ) - last_warning = time.time() - await asyncio.sleep(0.01) - - step_names = sorted(name for name, count in group_counts.items() if count == remaining_groups(adapters[name])) - return TrainBatch( - groups=collected, - group_counts=group_counts, - step_names=step_names, - step_slots=sorted(adapters[name].slot for name in step_names), - ) - - -async def generate_rollout_multi_lora_async( - args, rollout_id: int, data_source, generate_fn: GenerateFn = generate_and_rm_group -) -> RolloutFnTrainOutput: - """Collect one train batch and record its contents on the controller.""" - assert args.rollout_global_dataset - - state = GenerateState(args) - worker = AsyncMultiLoRAWorker.get_or_create(args, data_source, generate_fn) - start_time = time.time() - queue_sizes = worker.queue_sizes() - - # Driver contract: adapter state only changes between generate calls, so one snapshot serves the collection. - snapshot = await get_multi_lora_controller().snapshot.remote() - assert snapshot["active"] or snapshot["retiring"], "generate called with no live adapters" - - batch = await collect_batch(args, worker, snapshot) - - data = sorted( - batch.groups, - key=lambda group: ( - first_sample(group).adapter.slot if first_sample(group).adapter is not None else -1, - first_sample(group).index, - ), - ) - - # Per-sample adapter batch size (drives loss normalization) and batch-level step - # decision (drives selective optimizer stepping), shipped via sample metadata. - adapters = {**snapshot["active"], **snapshot["retiring"]} - for group in data: - adapter = adapters[group_adapter_name(group)] - for sample in iter_group_samples(group): - sample.metadata["adapter_global_batch_size"] = adapter.config.adapter_global_batch_size - if data: - head = first_sample(data[0]) - head.metadata["step_slots"] = list(batch.step_slots) - head.metadata["step_adapter_names"] = list(batch.step_names) - - await get_multi_lora_controller().record_batch_adapters.remote(rollout_id, batch.group_counts, batch.step_names) - - if (x := args.rollout_sample_filter_path) is not None: - load_function(x)(args, data) - - await recompute_samples_rollout_logprobs_via_prefill( - args, - [s for g in data for s in iter_group_samples(g)], - url=get_model_url(args, "default"), - sampling_params=state.sampling_params, - ) - - # Adapter metrics ride the adapter's own optimizer-step axis ({name}/step); this batch completes step + 1. - for name, step_metrics in worker.metrics.record_shipped_samples(args, data, batch.step_names, adapters).items(): - step_key = f"{name}/step" - log_dict = {step_key: adapters[name].step + 1} - log_dict |= {f"{name}/{key}": value for key, value in step_metrics.items()} - tracking.log(args, log_dict, step_key=step_key) - - stale_drops = worker.metrics.pop_stale_drops() - all_staleness = [staleness for values in stale_drops.values() for staleness in values] - metrics = { - **worker.metrics.pop_metrics(), - "perf/fully_async/queue_length": sum(queue_sizes.values()), - "perf/fully_async/stale_dropped": len(all_staleness), - # {name}/perf/* rides rollout/step; two segments under {name}/ keep these off the step axis. - **{f"{name}/perf/queue_length": size for name, size in queue_sizes.items()}, - **{f"{name}/perf/stale_dropped": len(stale_drops.get(name, [])) for name in adapters}, - "perf/fully_async/batch_wait_time": time.time() - start_time, - "perf/fully_async/batch_n_adapters": len(batch.group_counts), - "perf/fully_async/batch_n_groups": len(data), - "perf/fully_async/batch_n_samples": sum(group_sample_count(group) for group in data), - "perf/fully_async/batch_n_adapters_to_step": len(batch.step_names), - } - if all_staleness: - metrics["perf/fully_async/stale_dropped_avg_staleness"] = sum(all_staleness) / len(all_staleness) - metrics["perf/fully_async/stale_dropped_max_staleness"] = max(all_staleness) - for name, values in stale_drops.items(): - if values: - metrics[f"{name}/perf/stale_dropped_avg_staleness"] = sum(values) / len(values) - metrics[f"{name}/perf/stale_dropped_max_staleness"] = max(values) - - return RolloutFnTrainOutput(samples=data, metrics=metrics) - - -def generate_rollout_multi_lora(args, rollout_id: int, data_source, evaluation: bool = False): - if evaluation: - raise ValueError("Evaluation not supported in multi-LoRA async rollout") - return run(generate_rollout_multi_lora_async(args, rollout_id, data_source)) diff --git a/miles/rollout/multi_lora/data_source.py b/miles/rollout/multi_lora/data_source.py deleted file mode 100644 index 426b7af2452..00000000000 --- a/miles/rollout/multi_lora/data_source.py +++ /dev/null @@ -1,137 +0,0 @@ -"""Round-robin per-adapter data source. Deregistration is step-based and -lives in the controller (``mark_batch_trained``); every adapter gets a -``num_step`` at registration, explicit or derived from ``num_epoch``.""" - -import copy -import logging -from argparse import Namespace -from collections import deque -from concurrent.futures import ThreadPoolExecutor - -import ray - -from miles.ray.multi_lora.controller import get_multi_lora_controller -from miles.rollout.data_source import DataSource, RolloutDataSource -from miles.utils.adapter_config import AdapterRun -from miles.utils.types import AdapterRef, RewardSpec, Sample - -logger = logging.getLogger(__name__) - -MAX_RECONCILE_WORKERS = 16 - - -def fetch_snapshot() -> dict: - return ray.get(get_multi_lora_controller().snapshot.remote()) - - -def sampleable(snapshot: dict) -> dict[str, AdapterRun]: - return {**snapshot["active"], **snapshot["retiring"]} - - -class MultiLoRAAsyncDataSource(DataSource): - def __init__(self, args: Namespace): - self.args = args - self.sources: dict[str, RolloutDataSource] = {} - self.source_queue: deque = deque() - - def reconcile(self, adapters: dict[str, AdapterRun]) -> None: - for name in list(self.sources): - if name not in adapters: - del self.sources[name] - logger.info(f"Removed data source for adapter '{name}'") - pending = [(name, a) for name, a in adapters.items() if name not in self.sources] - if pending: - workers = min(MAX_RECONCILE_WORKERS, len(pending)) - if workers > 1: - with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="mlora-ds") as ex: - built = list(ex.map(lambda na: (na[0], self.create_source(na[1])), pending)) - else: - built = [(name, self.create_source(a)) for name, a in pending] - for name, source in built: - self.sources[name] = source - logger.info(f"Created data source for adapter '{name}'") - # Post-filter dataset length; the controller derives num_step - # from num_epoch for adapters that didn't set it. - ray.get(get_multi_lora_controller().resolve_num_step.remote(name, len(source.dataset))) - self.update_queue(set(adapters)) - - def create_source(self, adapter: AdapterRun) -> RolloutDataSource: - config = adapter.config - adapter_args = copy.copy(self.args) - adapter_args.prompt_data = config.data - adapter_args.input_key = config.input_key or self.args.input_key - adapter_args.label_key = config.label_key or self.args.label_key - adapter_args.metadata_key = config.metadata_key or self.args.metadata_key - adapter_args.save = config.save or self.args.save - adapter_args.load = config.save or self.args.load - adapter_args.n_samples_per_prompt = config.n_samples_per_prompt or self.args.n_samples_per_prompt - adapter_args.start_rollout_id = 0 - return RolloutDataSource(adapter_args) - - def update_queue(self, active_names: set[str]) -> None: - new_queue: deque = deque() - in_queue: set[str] = set() - while self.source_queue: - if (name := self.source_queue.popleft()) in active_names: - new_queue.append(name) - in_queue.add(name) - for name in active_names: - if name not in in_queue: - new_queue.append(name) - self.source_queue = new_queue - - def get_samples(self, num_samples: int = 1) -> list[list[Sample]]: - """Return the next prompt group, round-robined across adapters. - - One rotation of the queue: pull one group from the first adapter that - yields, stamp it, and return. Empty list when no adapter can produce. - """ - assert num_samples == 1, "the async producer dispatches one prompt group at a time" - snapshot = fetch_snapshot() - adapters = sampleable(snapshot) - self.reconcile(adapters) - self.update_queue(set(self.sources)) - - for _ in range(len(self.source_queue)): - name = self.source_queue.popleft() - self.source_queue.append(name) - source = self.sources[name] - groups = source.get_samples(1) - if not groups: - continue - - adapter = adapters[name] - config = adapter.config - ref = AdapterRef(name=name, slot=adapter.slot) - reward_spec = RewardSpec(rm_type=config.rm_type, custom_rm_path=config.custom_rm_path) - for sample in groups[0]: - sample.adapter = ref - sample.reward_spec = reward_spec - sample.metadata = {**config.metadata, **sample.metadata} - - return groups - - return [] - - def add_samples(self, samples: list[list[Sample]]) -> None: - """Recycle retried/aborted groups; drop groups for deregistered adapters.""" - adapters = sampleable(fetch_snapshot()) - self.reconcile(adapters) - for group in samples: - name = group[0].adapter.name if group and group[0].adapter else None - if not name or name not in self.sources or name not in adapters: - continue - self.sources[name].add_samples([group]) - - def save(self, rollout_id): - for source in self.sources.values(): - source.save(rollout_id) - - def load(self, rollout_id=None): - for source in self.sources.values(): - source.load(rollout_id) - - def close(self) -> None: - from miles.rollout.multi_lora.async_rollout import AsyncMultiLoRAWorker - - AsyncMultiLoRAWorker.stop_global() diff --git a/miles/rollout/multi_lora/operation_port.py b/miles/rollout/multi_lora/operation_port.py new file mode 100644 index 00000000000..98db6eb944f --- /dev/null +++ b/miles/rollout/multi_lora/operation_port.py @@ -0,0 +1,59 @@ +import asyncio +from typing import Protocol + +import ray + +from miles.utils.operation_contract import BindingT, RegistrationKey + + +class OperationQueuePort(Protocol[BindingT]): + """Ledger claims: ready_streams lists READY streams (head kind unknown); claim_data claim-and-binds in one call.""" + + async def ready_streams(self) -> dict: ... + + async def claim_data(self, key: RegistrationKey) -> dict | None: ... + + async def fail(self, operation_id: str, error: str, category: str) -> None: ... + + +class BatchResidencyPort(Protocol[BindingT]): + """Async transport face of controller-side TrainerResidencyPort: one immutable dispatch receipt per selection.""" + + async def acquire_batch(self, bindings_by_operation: list) -> object: ... + + +class RayMultiLoraOperationQueue: + """Only this class (and its residency sibling) knows the Ray controller, + .remote(), and ray.get.""" + + async def ready_streams(self) -> dict: + from miles.ray.multi_lora.controller import get_multi_lora_controller + + snapshot = await asyncio.to_thread(ray.get, get_multi_lora_controller().snapshot.remote()) + return snapshot["ready"] + + async def claim_data(self, key: RegistrationKey) -> dict | None: + from miles.ray.multi_lora.controller import get_multi_lora_controller + + name, registration_id = key + return await asyncio.to_thread( + ray.get, get_multi_lora_controller().claim_data_operation.remote(name, registration_id) + ) + + async def fail(self, operation_id: str, error: str, category: str) -> None: + from miles.ray.multi_lora.controller import get_multi_lora_controller + + await asyncio.to_thread( + ray.get, get_multi_lora_controller().fail_operation.remote(operation_id, error, category) + ) + + +class RayTrainerResidencyPort: + """Thin async proxy to the backend-owned FixedSlotResidency.""" + + async def acquire_batch(self, bindings_by_operation: list) -> object: + from miles.ray.multi_lora.controller import get_multi_lora_controller + + return await asyncio.to_thread( + ray.get, get_multi_lora_controller().acquire_batch_lease.remote(list(bindings_by_operation)) + ) diff --git a/miles/rollout/multi_lora/rollout_fn.py b/miles/rollout/multi_lora/rollout_fn.py new file mode 100644 index 00000000000..13e5572905c --- /dev/null +++ b/miles/rollout/multi_lora/rollout_fn.py @@ -0,0 +1,369 @@ +import asyncio +import logging +import time +from collections import deque +from dataclasses import dataclass +from typing import Any + +from miles.ray.multi_lora.config import AdapterRun +from miles.ray.multi_lora.residency import lease_to_metadata +from miles.rollout.base_types import ( + RolloutFnConstructorInput, + RolloutFnInput, + RolloutFnTrainOutput, + RolloutPostprocessOptions, +) +from miles.rollout.multi_lora.operation_port import ( + BatchResidencyPort, + OperationQueuePort, + RayMultiLoraOperationQueue, + RayTrainerResidencyPort, +) +from miles.utils.operation_contract import EmptyBatchTimeoutError +from miles.utils.types import AdapterRef, Sample + +logger = logging.getLogger(__name__) + + +def batch_plan_to_metadata(batch_plan: list[dict], lease) -> dict[str, Any]: + kinds = {entry["operation_kind"] for entry in batch_plan} + if len(kinds) != 1 or not kinds <= {"forward_backward", "forward"}: + raise ValueError(f"tinker selection must be one homogeneous data kind, got {sorted(kinds)}") + metadata: dict[str, Any] = { + "batch_kind": "tinker", + # Per-sample lanes in selection order (each entry's rows are contiguous). + "tinker_operation_lanes": [ + lane for lane, entry in enumerate(batch_plan) for _ in range(entry["sample_count"]) + ], + "tinker_loss_by_lane": {lane: entry.get("loss_spec") or {} for lane, entry in enumerate(batch_plan)}, + # The trainer completes these operations after the batch lands. + "operation_by_lane": {lane: entry["operation_id"] for lane, entry in enumerate(batch_plan)}, + # Exact registration per lane: the batch commit dirties these streams, + # never a trainer-reported name list. + "registration_by_lane": { + lane: (entry["name"], entry["registration_id"]) for lane, entry in enumerate(batch_plan) + }, + # The lease is mandatory: a batch without its dispatch receipt is one + # the trainer must reject, so the optional path may not exist here. + "batch_execution_lease": lease_to_metadata(lease), + } + if kinds == {"forward"}: + metadata["tinker_forward_only"] = True + return metadata + + +_CLAIM_POLL_S = 0.5 +# A FAILED child runtime returns to IDLE after this cooldown instead of starving its adapter until deregister. +_FAILED_RELAUNCH_COOLDOWN_S = 5.0 + +Tenant = tuple[str, str] + +DATA_OPERATION_KINDS = ("forward_backward", "forward") + + +@dataclass(frozen=True) +class ClaimedOperationBatch: + """One claimed operation as a complete batch; its binding (claim-and-bind) is the one dispatch truth.""" + + operation_id: str + kind: str + loss_spec: dict | None + binding: Any # duck-typed port binding; production ships ResidentBinding + samples: list[list[Sample]] + + +def decode_operation(operation: dict, run: AdapterRun) -> ClaimedOperationBatch: + if operation["kind"] not in DATA_OPERATION_KINDS: + raise ValueError(f"operation kind '{operation['kind']}' is not a data operation") + payload = operation.get("payload") or {} + raw_samples = payload.get("samples") + if not raw_samples: + raise ValueError(f"{operation['kind']} payload carries no samples") + ref = AdapterRef( + name=run.name, + registration_id=run.registration_id, + serving_version=run.version, + slot=run.slot, + ) + groups: list[list[Sample]] = [] + for i, raw in enumerate(raw_samples): + raw = dict(raw) + raw.setdefault("status", Sample.Status.COMPLETED.value) + raw["index"] = i + sample = Sample.from_dict(raw) + sample.adapter = ref + sample.metadata = {**run.config.metadata, **sample.metadata} + groups.append([sample]) + return ClaimedOperationBatch( + operation_id=operation["operation_id"], + kind=operation["kind"], + loss_spec=payload.get("loss"), + binding=operation["binding"], + samples=groups, + ) + + +class TinkerNullDataSource: + """Dataset-less data source for tinker runs; only satisfies the manager's save/load/close surface.""" + + dataset = () + + def __init__(self, args): + self.args = args + + def get_samples(self, num_samples: int): + raise RuntimeError("tinker runs have no dataset; data arrives as client operations") + + def add_samples(self, samples) -> None: + pass + + def save(self, rollout_id) -> None: + pass + + def load(self, rollout_id=None) -> None: + pass + + +class AdapterRolloutRuntime: + """One per registration: at most one in-flight child claim task and one + ready output.""" + + IDLE = "IDLE" + IN_FLIGHT = "IN_FLIGHT" + READY = "READY" + SELECTED = "SELECTED" + FAILED = "FAILED" + + def __init__(self, run: AdapterRun): + self.run = run + self.state = self.IDLE + self.ready_output: ClaimedOperationBatch | None = None + self.task: asyncio.Task | None = None + self.last_failure: float | None = None + + @property + def ready_kind(self) -> str | None: + if self.ready_output is None: + return None + return self.ready_output.kind + + def refresh(self, run: AdapterRun) -> None: + """Serving version advances between batches; identity stays fixed.""" + self.run = run + + async def aclose(self) -> None: + if self.task is not None and not self.task.done(): + self.task.cancel() + try: + await self.task + except (asyncio.CancelledError, Exception): # noqa: BLE001 - teardown must not raise + pass + self.task = None + + +class MultiLoraOperationBatchFn: + def __init__( + self, + input: RolloutFnConstructorInput, + operations: OperationQueuePort | None = None, + residency: BatchResidencyPort | None = None, + ): + self.args = input.args + self.operations = operations if operations is not None else RayMultiLoraOperationQueue() + self.residency = residency if residency is not None else RayTrainerResidencyPort() + self.runtimes: dict[Tenant, AdapterRolloutRuntime] = {} + self.rotation: deque[Tenant] = deque() + self._ready = asyncio.Event() + + # ------------------------------ lifecycle ------------------------------ + + async def __call__(self, input: RolloutFnInput) -> RolloutFnTrainOutput: + if input.evaluation: + raise ValueError( + "MultiLoraOperationBatchFn does not serve eval; tinker runs have no server-side eval loop" + ) + # READY streams only: a retiring registration's queued ops are fenced terminal, so a claim never returns. + adapters = await self.operations.ready_streams() + await self._reconcile(adapters) + self._launch_idle_children() + selected = await self._select() + return await self._merge(selected) + + async def aclose(self) -> None: + for runtime in list(self.runtimes.values()): + await runtime.aclose() + self.runtimes.clear() + self.rotation.clear() + + # ------------------------------ runtimes ------------------------------ + + async def _reconcile(self, adapters: dict[str, AdapterRun]) -> None: + live = {(name, run.registration_id) for name, run in adapters.items()} + for tenant in [t for t in self.runtimes if t not in live]: + # Deregistered or re-registered: close the old tenant's runtime; + # its late results are dropped with it (registration fencing). + await self.runtimes.pop(tenant).aclose() + logger.info(f"[tinker] closed child runtime for '{tenant[0]}' ({tenant[1][:8]})") + for name, run in adapters.items(): + tenant = (name, run.registration_id) + if tenant in self.runtimes: + self.runtimes[tenant].refresh(run) + continue + self.runtimes[tenant] = AdapterRolloutRuntime(run) + logger.info(f"[tinker] created child runtime for '{name}' ({run.registration_id[:8]})") + self._sync_rotation() + + def _sync_rotation(self) -> None: + in_queue = set() + kept: deque[Tenant] = deque() + while self.rotation: + if (tenant := self.rotation.popleft()) in self.runtimes and tenant not in in_queue: + kept.append(tenant) + in_queue.add(tenant) + for tenant in self.runtimes: + if tenant not in in_queue: + kept.append(tenant) + self.rotation = kept + + def _launch_idle_children(self) -> None: + now = time.monotonic() + for runtime in self.runtimes.values(): + if runtime.state == AdapterRolloutRuntime.FAILED and ( + runtime.last_failure is None or now - runtime.last_failure >= _FAILED_RELAUNCH_COOLDOWN_S + ): + # FAILED is transient: after the cooldown the child relaunches instead of starving the adapter. + runtime.state = AdapterRolloutRuntime.IDLE + if runtime.state == AdapterRolloutRuntime.IDLE: + runtime.state = AdapterRolloutRuntime.IN_FLIGHT + runtime.task = asyncio.create_task(self._run_child(runtime)) + + async def _claim_batch(self, runtime: AdapterRolloutRuntime) -> ClaimedOperationBatch: + key = (runtime.run.name, runtime.run.registration_id) + while True: + operation = await self.operations.claim_data(key) + if operation is None: + await asyncio.sleep(_CLAIM_POLL_S) + continue + try: + return decode_operation(operation, runtime.run) + except asyncio.CancelledError: + raise + except Exception as e: # noqa: BLE001 - a bad payload fails its op, not the adapter + logger.exception(f"[tinker] ({key[0]}) operation '{operation['operation_id']}' rejected: {e}") + await self.operations.fail(operation["operation_id"], f"invalid operation payload: {e}", "user") + + async def _run_child(self, runtime: AdapterRolloutRuntime) -> None: + try: + output = await self._claim_batch(runtime) + runtime.ready_output = output + runtime.state = AdapterRolloutRuntime.READY + except asyncio.CancelledError: + runtime.state = AdapterRolloutRuntime.IDLE + raise + except Exception as e: + # Child failure isolates to this adapter; other adapters keep going. + logger.exception(f"[tinker] child for '{runtime.run.name}' failed: {e}") + runtime.last_failure = time.monotonic() + runtime.state = AdapterRolloutRuntime.FAILED + finally: + self._ready.set() + + # ------------------------------ selection ------------------------------ + + async def _select(self) -> list[AdapterRolloutRuntime]: + soft_target = self.args.rollout_batch_size * self.args.n_samples_per_prompt + coalesce_wait = self.args.tinker_max_coalesce_wait_s + empty_deadline = time.monotonic() + self.args.tinker_max_empty_wait_s + selected: list[AdapterRolloutRuntime] = [] + kind_lock: str | None = None + collected = 0 + coalesce_deadline: float | None = None + + while True: + runtime = self._pop_next_ready(kind_lock) + if runtime is not None: + selected.append(runtime) + # Leave READY immediately or the round-robin would re-select + # the same batch until the target is met (duplicated samples). + runtime.state = AdapterRolloutRuntime.SELECTED + kind_lock = runtime.ready_kind + collected += sum(len(group) for group in runtime.ready_output.samples) + if coalesce_deadline is None: + coalesce_deadline = time.monotonic() + coalesce_wait + # Whole batches only: overshoot past the soft target is allowed, + # trimming is not. + if collected >= soft_target or len(selected) >= len(self.runtimes): + break + continue + + now = time.monotonic() + if selected: + if now >= coalesce_deadline: + break + timeout = coalesce_deadline - now + else: + if now >= empty_deadline: + raise EmptyBatchTimeoutError( + "no adapter produced a batch within " + f"--tinker-max-empty-wait-s ({self.args.tinker_max_empty_wait_s}s)" + ) + timeout = empty_deadline - now + self._ready.clear() + try: + await asyncio.wait_for(self._ready.wait(), timeout=timeout) + except TimeoutError: + continue + return selected + + def _pop_next_ready(self, kind_lock: str | None) -> AdapterRolloutRuntime | None: + for _ in range(len(self.rotation)): + tenant = self.rotation.popleft() + self.rotation.append(tenant) + runtime = self.runtimes.get(tenant) + if runtime is None or runtime.state != AdapterRolloutRuntime.READY: + continue + if kind_lock is not None and runtime.ready_kind != kind_lock: + continue + return runtime + return None + + # ------------------------------ merge ------------------------------ + + async def _merge(self, selected: list[AdapterRolloutRuntime]) -> RolloutFnTrainOutput: + data: list[list[Sample]] = [] + batch_plan: list[dict] = [] + metrics: dict = {} + try: + for runtime in selected: + claim = runtime.ready_output + data.extend(claim.samples) + name, registration_id = claim.binding.registration_key + batch_plan.append( + dict( + name=name, + registration_id=registration_id, + operation_id=claim.operation_id, + operation_kind=claim.kind, + loss_spec=claim.loss_spec, + sample_count=sum(len(group) for group in claim.samples), + binding=claim.binding, + ) + ) + metrics[f"{runtime.run.name}/operation_samples"] = sum(len(group) for group in claim.samples) + lease = await self.residency.acquire_batch( + [(entry["operation_id"], entry["binding"]) for entry in batch_plan] + ) + except BaseException: + for runtime in selected: + runtime.state = AdapterRolloutRuntime.READY + raise + # Acquisition succeeded: NOW consume the outputs. + for runtime in selected: + runtime.ready_output = None + runtime.state = AdapterRolloutRuntime.IDLE # relaunches at the NEXT generate call + return RolloutFnTrainOutput( + samples=data, + metrics=metrics, + conversion_metadata=batch_plan_to_metadata(batch_plan, lease), + postprocess=RolloutPostprocessOptions(pad_to_dp=True), + ) diff --git a/miles/rollout/sglang_rollout.py b/miles/rollout/sglang_rollout.py index ee5ad83642d..091ee25aa43 100644 --- a/miles/rollout/sglang_rollout.py +++ b/miles/rollout/sglang_rollout.py @@ -182,7 +182,7 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A payload["top_logprobs_num"] = opd_top_k if sample.adapter is not None: - from miles.ray.multi_lora.controller import AdaptersCache + from miles.ray.multi_lora.cache import AdaptersCache if (adapter := await AdaptersCache().get(sample.adapter.name)) is None: # Adapter deregistered: don't POST, or an orphan the abort round can't see diff --git a/miles/utils/adapter_config.py b/miles/utils/adapter_config.py deleted file mode 100644 index 6c4d224a86d..00000000000 --- a/miles/utils/adapter_config.py +++ /dev/null @@ -1,93 +0,0 @@ -"""Adapter config parsing for multi-LoRA training. - -``AdapterRunConfig`` carries only static, YAML-sourced configuration; the -mutable slot is owned by the controller and exposed through ``AdapterRun`` -views. -""" - -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any - -import yaml - - -@dataclass(frozen=True) -class AdapterRunConfig: - - data: str - - # resolves them to CLI defaults if None (--lora-rank / --lora-alpha) on register. - rank: int | None = None - alpha: int | None = None - - # Prompt groups consumed per optimizer step for this adapter (group units, - # like --rollout-batch-size, which it defaults to). The samples-per-step - # analog of --global-batch-size is derived: adapter_global_batch_size = - # rollout_batch_size * n_samples_per_prompt. - rollout_batch_size: int | None = None - n_samples_per_prompt: int | None = None - - save: str | Path | None = None - - input_key: str = "text" - label_key: str | None = None - metadata_key: str | None = None - - rm_type: str | None = None - custom_rm_path: str | None = None - - # Stop after N optimizer steps; derived from num_epoch (default 1) when absent. - num_step: int | None = None - num_epoch: int | None = None - - metadata: dict[str, Any] = field(default_factory=dict) - - @property - def adapter_global_batch_size(self) -> int: - """Samples per optimizer step (per-adapter analog of --global-batch-size).""" - assert self.rollout_batch_size is not None and self.n_samples_per_prompt is not None - return self.rollout_batch_size * self.n_samples_per_prompt - - -@dataclass(frozen=True) -class AdapterRun: - """Read-only join view of a run's static config and current slot.""" - - name: str - config: AdapterRunConfig - slot: int - version: int = 0 - step: int = 0 - # Committed prompt groups accumulated toward the current optimizer step. - accumulated_groups: int = 0 - # Unique per registration (see AdapterRecord.registration_id): lets the - # rollout worker tell a re-registered name apart from the previous tenant. - registration_id: str = "" - - -def parse_adapter_run_yaml(path: Path) -> AdapterRunConfig: - """Parse a single adapter.yaml file. - - ``rank``, ``alpha`` and ``save`` are optional in the YAML; when absent the - caller (e.g. the multi-LoRA controller) is responsible for resolving them. - """ - with open(path) as f: - raw = yaml.safe_load(f) - - return AdapterRunConfig( - rank=raw.get("rank"), - alpha=raw.get("alpha"), - data=raw["data"], - rollout_batch_size=raw.get("rollout_batch_size"), - n_samples_per_prompt=raw.get("n_samples_per_prompt"), - save=Path(raw["save"]) if raw.get("save", None) else None, - input_key=raw.get("input_key", "text"), - label_key=raw.get("label_key"), - metadata_key=raw.get("metadata_key"), - rm_type=raw.get("rm_type"), - custom_rm_path=raw.get("custom_rm_path"), - num_step=raw.get("num_step"), - num_epoch=raw.get("num_epoch"), - metadata=raw.get("metadata") or {}, - ) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 62941feb353..1c06d1fb2bd 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1798,12 +1798,40 @@ def add_lora_arguments(parser): help="Maximum number of concurrent adapter slots for multi-LoRA. Set to 0 to disable multi-LoRA (default: 0)", ) parser.add_argument( - "--multi-lora-adapter", - nargs=2, - action="append", - type=str, - dest="multi_lora_adapters", - default=[], + "--tinker-backend", + action="store_true", + default=False, + help="Enable the Tinker protocol adapter for Multi-LoRA (requires --multi-lora-n-adapters > 0)", + ) + parser.add_argument( + "--tinker-max-coalesce-wait-s", + type=float, + default=2.0, + help="Keep coalescing ready batches into the same train call this long after the first (default: 2.0)", + ) + parser.add_argument( + "--tinker-max-empty-wait-s", + type=float, + default=5.0, + help="Idle window before EmptyBatchTimeoutError; short so control ops never wait (default: 5.0)", + ) + parser.add_argument( + "--tinker-operation-gap-timeout", + type=float, + default=600.0, + help="Gap-stall seconds before blocked ops fail and the hole seals; <= 0 disables (default: 600)", + ) + parser.add_argument( + "--tinker-operation-claimed-ttl", + type=float, + default=1800.0, + help="Liveness backstop: fail orphaned CLAIMED ops after this long; <= 0 disables (default: 1800)", + ) + parser.add_argument( + "--multi-lora-max-consecutive-generate-failures", + type=int, + default=10, + help="Consecutive generate failures skipped before re-raising; 0 fails fast (default: 10)", ) parser.add_argument( "--multi-lora-idle-poll-s", @@ -1816,18 +1844,15 @@ def add_lora_arguments(parser): type=str, default=None, help=( - "Dotted path to a MultiLoRAHTTPServer subclass to use for the multi-LoRA " - "controller's HTTP server (default: MultiLoRAHTTPServer)" + "Dotted path to an AdapterRunControlServer subclass to use for the multi-LoRA " + "controller's HTTP server (default: AdapterRunControlServer)" ), ) parser.add_argument( "--multi-lora-backend-path", type=str, default=None, - help=( - "Dotted path to a MultiLoRABackend subclass for the multi-LoRA controller, " - "e.g. to add custom adapter validation via validate_adapter (default: MultiLoRABackend)" - ), + help="Dotted path to a MultiLoraOperationBackend subclass (e.g. custom validate_adapter)", ) parser.add_argument( "--multi-lora-api-port", @@ -1836,38 +1861,68 @@ def add_lora_arguments(parser): help="Port for the multi-LoRA controller's control-plane API, served from the head node (default: 8068)", ) parser.add_argument( - "--multi-lora-disable-service-mode", - action="store_false", - dest="multi_lora_service_mode", - help="Disable service mode. By default, the trainer waits indefinitely for new adapters. With this flag, it exits after all adapters have been processed.", + "--tinker-frontend", + action="store_true", + default=False, + help="Serve the official tinker SDK REST protocol (/api/v1) on the tinker " + "controller's HTTP server: an unmodified `tinker` client pointed at it " + "(base_url + api_key) drives training and sampling", ) parser.add_argument( - "--multi-lora-max-adapter-global-batch-size", + "--tinker-api-key", + type=str, + default=None, + help="API key the tinker frontend requires in X-API-Key (single-tenant; the SDK " + "needs a 'tml-' prefix). Falls back to $MILES_TINKER_API_KEY. Required for a " + "non-loopback bind (fail closed)", + ) + parser.add_argument( + "--tinker-sampling-max-active-subgenerations", + type=int, + default=64, + help="Global cap on concurrently executing sampling sub-generations across ALL " + "SDK clients (one request counts num_samples). Submissions over the cap get a " + "retryable 429 before consuming their identity, and the router transport holds " + "the same hard bound; the SDK's per-client limit of 64 never bounded the " + "aggregate (default: 64, validated on H200)", + ) + parser.add_argument( + "--tinker-sampling-max-context", type=int, default=None, - help=( - "Registration-time upper bound on an adapter's samples per optimizer " - "step (rollout_batch_size x n_samples_per_prompt). Defaults to 4x " - "--global-batch-size." - ), + help="Engine context limit (tokens) the tinker frontend preflights sample " + "requests against: prompt + max_tokens over the limit is a typed 400 before " + "the seq identity is consumed (the engine would otherwise silently truncate " + "the decode budget and return garbage). Default: --sglang-context-length when " + "set, else discovered from the router's /get_server_info on the first sample", ) parser.add_argument( - "--multi-lora-max-coalesce-wait-s", + "--tinker-session-idle-ttl", type=float, - default=0.5, - help=( - "Maximum time ready groups wait for the batch to fill toward " - "--global-batch-size before training starts on what is ready (default: 0.5)." - ), + default=3600.0, + help="Seconds without a session heartbeat before the tinker frontend reaps the " + "session and its sampling sessions (the SDK heartbeats continuously while the " + "client lives). Old sampler ids then fail closed, so nothing a vanished client " + "executed can re-execute. <= 0 disables (default: 3600)", ) parser.add_argument( - "--multi-lora-max-empty-wait-s", + "--tinker-future-unpolled-ttl", type=float, - default=30.0, - help=( - "How long a generate call waits for the first poppable group before " - "failing with an empty-batch timeout (default: 30)." - ), + default=900.0, + help="Seconds without a retrieve_future poll before the tinker frontend treats " + "a pending future as orphaned: an orphaned sample's server-side generation is " + "cancelled (SDK future cancellation never reaches the engine on its own) and " + "the future resolves typed; orphaned training futures are polled on the " + "client's behalf so the ledger's unacked-results budget drains. <= 0 disables " + "(default: 900)", + ) + parser.add_argument( + "--tinker-future-undelivered-ttl", + type=float, + default=3600.0, + help="Seconds a terminal-but-never-retrieved future result is retained before " + "the reaper evicts it to a fingerprint tombstone (a late retry then gets a " + "typed 410, never a silent re-execution). <= 0 disables (default: 3600)", ) return parser @@ -3133,6 +3188,10 @@ def miles_validate_args(args): validate_multi_lora_args(args) + from miles.utils.tinker import validate_tinker_args + + validate_tinker_args(args) + assert not (args.kl_coef != 0 and args.kl_loss_coef != 0), "Only one of kl_coef and kl_loss_coef can be set" if args.advantage_estimator in ["reinforce_plus_plus", "reinforce_plus_plus_baseline"]: diff --git a/miles/utils/multi_lora.py b/miles/utils/multi_lora.py index 8a0a2d42f57..92f6b98c747 100644 --- a/miles/utils/multi_lora.py +++ b/miles/utils/multi_lora.py @@ -1,8 +1,4 @@ -"""Small multi-LoRA helpers shared across the rollout, trainer, and controller. - -The controller-side machinery (AdapterRegistry, MultiLoRABackend, -MultiLoRAHTTPServer) lives in ``miles/ray/multi_lora/``. -""" +"""Small Multi-LoRA helpers shared across rollout, trainer, and controller.""" import logging import uuid @@ -11,14 +7,12 @@ logger = logging.getLogger(__name__) __all__ = [ - "EmptyBatchTimeoutError", "RID_SEPARATOR", - "define_new_adapter_metrics", "is_multi_lora_enabled", "make_rid", - "min_groups_per_dp_split", - "parse_adapter", "slot_lora_name", + "targets_expert_leaves", + "uses_multi_lora_operation_executor", "validate_multi_lora_args", ] @@ -27,23 +21,15 @@ RID_SEPARATOR = "::" -class EmptyBatchTimeoutError(RuntimeError): - """No trainable groups arrived before empty-wait timeout.""" - - def is_multi_lora_enabled(args: Any) -> bool: return getattr(args, "multi_lora", False) -def define_new_adapter_metrics(snapshot: dict) -> None: - """Declare metric axes for new adapters ({name}/* -> {name}/step, {name}/perf/* -> rollout/step); must run - in the primary tracking writer. Already-declared adapters are skipped, so calling every snapshot is free.""" - # lazy import tracking deps - from miles.utils.tracking_utils.tracking import define_step_key_metric_group +def uses_multi_lora_operation_executor(args: Any) -> bool: + """Whether explicit operations execute on fixed Multi-LoRA slots.""" + from miles.utils.tinker import uses_explicit_training_operations - for name in {**snapshot["pending"], **snapshot["active"], **snapshot["retiring"]}: - define_step_key_metric_group(prefix=name, step_key=f"{name}/step") - define_step_key_metric_group(prefix=f"{name}/perf", step_key="rollout/step") + return uses_explicit_training_operations(args) and getattr(args, "multi_lora_n_adapters", 0) > 0 # Leaf module names that can live inside MoE experts (they also name the dense MLP @@ -64,17 +50,13 @@ def targets_expert_leaves(target_modules: Any) -> bool: def validate_multi_lora_args(args: Any) -> None: - """Set ``args.multi_lora``, then validate and default the multi-LoRA arg - surface. Called from ``miles_validate_args``; a no-op for normal runs.""" args.multi_lora = getattr(args, "multi_lora_n_adapters", 0) > 0 if not args.multi_lora: return - # Swap in the multi-LoRA rollout fn and data source unless the user pointed these flags elsewhere. - if args.rollout_function_path is None: - args.rollout_function_path = "miles.rollout.multi_lora.async_rollout.generate_rollout_multi_lora" - if args.data_source_path == "miles.rollout.data_source.RolloutDataSourceWithBuffer": - args.data_source_path = "miles.rollout.multi_lora.data_source.MultiLoRAAsyncDataSource" + assert getattr( + args, "tinker_backend", False + ), "multi-LoRA now requires --tinker-backend: the dataset-driven adapter-sample-level path was removed" # The per-adapter data source is inherently global (the controller owns # what is sampleable); rollout workers must not shard it. args.rollout_global_dataset = True @@ -129,47 +111,17 @@ def validate_multi_lora_args(args: Any) -> None: "(sample-mean); per-token loss normalization would make adapter batch weights " "depend on batch contents. Drop --calculate-per-token-loss." ) - assert args.multi_lora_max_coalesce_wait_s >= 0, "--multi-lora-max-coalesce-wait-s must be non-negative" assert (getattr(args, "optimizer", "adam") or "adam").lower() == "adam", ( "Multi-LoRA requires --optimizer adam: the per-slot optimizer isolation " - "(build_multi_lora_optimizer, slot retirement state cleanup) only implements " + "(slot optimizer construction, slot retirement state cleanup) only implements " f"Adam semantics; got --optimizer {args.optimizer}" ) from miles.utils.environ import enable_experimental_ft_trainer assert not enable_experimental_ft_trainer(), ( "Multi-LoRA is not supported with MILES_EXPERIMENTAL_FT_TRAINER=1: the v2 " - "train group has no reconcile_adapters and does not return train outcomes" + "train group has no adapter reconcile verbs and does not return train outcomes" ) - # --global-batch-size may legitimately be unset (Megatron derives it later); - # leave the adapter cap unset too rather than multiplying None. - if args.multi_lora_max_adapter_global_batch_size is None and getattr(args, "global_batch_size", None) is not None: - args.multi_lora_max_adapter_global_batch_size = 4 * args.global_batch_size - if args.multi_lora_max_adapter_global_batch_size is not None: - assert ( - args.multi_lora_max_adapter_global_batch_size > 0 - ), "--multi-lora-max-adapter-global-batch-size must be positive" - - # Trainer DP size, used to validate adapter batch shapes; guarded for harnesses without megatron args set. - if all( - hasattr(args, name) - for name in ( - "world_size", - "tensor_model_parallel_size", - "pipeline_model_parallel_size", - "context_parallel_size", - ) - ): - from miles.utils.megatron_args_utils import compute_megatron_world_size_except_dp - - model_parallel = compute_megatron_world_size_except_dp(args) - assert ( - args.world_size % model_parallel == 0 - ), f"actor world size {args.world_size} is not divisible by tp*pp*cp {model_parallel}" - args.multi_lora_dp_size = args.world_size // model_parallel - else: - args.multi_lora_dp_size = None - # Batches are variable-sized; carry the exact sample # count through rollout conversion instead of trimming to --global-batch-size. assert not args.disable_rollout_trim_samples, ( @@ -184,31 +136,7 @@ def make_rid(adapter_name: str) -> str: return f"{adapter_name}{RID_SEPARATOR}{uuid.uuid4().hex}" -def parse_adapter(rid: str) -> str: - return rid.rsplit(RID_SEPARATOR, 1)[0] - - def slot_lora_name(slot: int) -> str: """Engine-side LoRA adapter name for a controller slot. Weight pushes and every inference request (rollout and prefill scoring) must agree on this.""" return f"__miles_slot_{slot}" - - -def min_groups_per_dp_split(n_samples_per_prompt: int, dp_size: int) -> int: - """Minimum prompt-group count that splits cleanly across data-parallel - ranks. - - Train batches only pop groups in multiples of this value, so each popped - slice has a sample count divisible by ``dp_size`` with no trimming. - - Requires ``n_samples_per_prompt`` and ``dp_size`` to divide each other - (one must be a multiple of the other). - """ - larger = max(dp_size, n_samples_per_prompt) - smaller = min(dp_size, n_samples_per_prompt) - if larger % smaller == 0: - return larger // n_samples_per_prompt - raise ValueError( - f"n_samples_per_prompt={n_samples_per_prompt} must be a divisor or a multiple of " - f"the data-parallel size {dp_size} so whole prompt groups can split evenly across ranks" - ) diff --git a/miles/utils/operation_contract.py b/miles/utils/operation_contract.py new file mode 100644 index 00000000000..c35b8458af6 --- /dev/null +++ b/miles/utils/operation_contract.py @@ -0,0 +1,44 @@ +from dataclasses import dataclass +from typing import Generic, Protocol, TypeVar + +RegistrationKey = tuple[str, str] + +BindingT = TypeVar("BindingT") + + +@dataclass(frozen=True) +class BatchExecutionLease(Generic[BindingT]): + """Immutable logical-operation to physical-binding receipt for one batch.""" + + dispatch_id: str + bindings_by_operation: tuple[tuple[str, BindingT], ...] + + def binding_of(self, operation_id: str) -> BindingT | None: + for op_id, binding in self.bindings_by_operation: + if op_id == operation_id: + return binding + return None + + +class TrainerResidencyPort(Protocol[BindingT]): + """Resolve, snapshot, validate, and release opaque trainer bindings.""" + + def binding_for(self, key: RegistrationKey) -> BindingT | None: + """Return the key's dispatchable binding, or ``None``.""" + ... + + def acquire_batch(self, bindings_by_operation: tuple[tuple[str, BindingT], ...]) -> BatchExecutionLease[BindingT]: + """Snapshot validated bindings into one immutable dispatch receipt.""" + ... + + def validate(self, lease: BatchExecutionLease[BindingT]) -> bool: + """Re-check a receipt before physical mutation.""" + ... + + def release_batch(self, lease: BatchExecutionLease[BindingT]) -> None: + """Release any physical reservation represented by ``lease``.""" + ... + + +class EmptyBatchTimeoutError(RuntimeError): + """No registration produced a claimable data operation within the wait.""" diff --git a/miles/utils/tinker.py b/miles/utils/tinker.py new file mode 100644 index 00000000000..8377150c683 --- /dev/null +++ b/miles/utils/tinker.py @@ -0,0 +1,36 @@ +def uses_explicit_training_operations(args) -> bool: + """Whether the Tinker protocol drives explicit training operations.""" + return bool(getattr(args, "tinker_backend", False)) + + +def is_tinker_enabled(args) -> bool: + """Whether the current Tinker-to-Multi-LoRA composition is enabled.""" + from miles.utils.multi_lora import uses_multi_lora_operation_executor + + return uses_multi_lora_operation_executor(args) + + +def validate_tinker_args(args) -> None: + """Validate the Tinker adapter and select its queue-backed rollout path.""" + if not getattr(args, "tinker_backend", False): + assert not getattr(args, "tinker_frontend", False), "--tinker-frontend requires --tinker-backend" + assert not getattr(args, "tinker_api_key", None), "--tinker-api-key requires --tinker-frontend" + return + + assert not ( + getattr(args, "tinker_api_key", None) and not getattr(args, "tinker_frontend", False) + ), "--tinker-api-key requires --tinker-frontend (only the SDK frontend authenticates requests)" + + from miles.utils.environ import use_legacy_rollout_v1 + + assert getattr(args, "multi_lora_n_adapters", 0) > 0, "--tinker-backend requires --multi-lora-n-adapters > 0" + assert ( + not use_legacy_rollout_v1() + ), "--tinker-backend needs the class-based rollout API (the default); unset MILES_USE_LEGACY_ROLLOUT_V1" + if getattr(args, "tinker_frontend", False) and not getattr(args, "multi_lora_http_server_path", None): + args.multi_lora_http_server_path = "miles.ray.tinker_frontend.http_server.TinkerFrontendHTTPServer" + if args.rollout_function_path is None: + args.rollout_function_path = "miles.rollout.multi_lora.rollout_fn.MultiLoraOperationBatchFn" + if args.data_source_path == "miles.rollout.data_source.RolloutDataSourceWithBuffer": + args.data_source_path = "miles.rollout.multi_lora.rollout_fn.TinkerNullDataSource" + args.use_dynamic_global_batch_size = True diff --git a/miles/utils/types.py b/miles/utils/types.py index 78d24e49f10..94777e7c265 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -12,6 +12,10 @@ class AdapterRef: name: str slot: int + # Registration-scoped serving identity (tinker): a re-registered name is a + # new tenant (anti-ABA), and the serving version keys the KV cache. + registration_id: str = "" + serving_version: int = 0 @dataclass(frozen=True) @@ -53,6 +57,11 @@ class Sample: remove_sample: bool = False teacher_log_probs: list[float] | None = None # Log probabilities from teacher model for OPD opd_reverse_kl: list[float] | None = None # Precomputed per-token OPD reverse-KL estimate + # Client-supplied per-token channels (tinker adapters): linear-CE + # coefficients and precomputed advantages, response-aligned like loss_mask. + # Distinct from the binary loss_mask — weights may be fractional or negative. + loss_weights: list[float] | None = None + advantages: list[float] | None = None class Status(Enum): PENDING = "pending" @@ -175,7 +184,8 @@ def from_dict(data: dict): return sample def get_reward_value(self, args) -> float: - return self.reward if not args.reward_key else self.reward[args.reward_key] + reward_key = getattr(args, "reward_key", None) + return self.reward if not reward_key else self.reward[reward_key] @property def effective_response_length(self): diff --git a/tests/ci/requirements-ci-cpu.txt b/tests/ci/requirements-ci-cpu.txt index e812c695af8..d5d728d95ae 100644 --- a/tests/ci/requirements-ci-cpu.txt +++ b/tests/ci/requirements-ci-cpu.txt @@ -9,6 +9,10 @@ partial_json_parser==0.2.1.1.post7 pyzmq==27.1.0 sentencepiece==0.2.1 tiktoken==0.13.0 +# The official tinker SDK wheel: the frontend contract tests +# (tests/fast/ray/tinker_frontend/test_sdk_contract.py) and the backend +# metrics-combiner contract test importorskip without the pin. +tinker==0.24.1 torch==2.11.0 torchvision==0.26.0 xgrammar==0.2.1 diff --git a/tests/ci/run_suite.py b/tests/ci/run_suite.py index a984a7b677d..0dbcd77d86b 100644 --- a/tests/ci/run_suite.py +++ b/tests/ci/run_suite.py @@ -174,12 +174,22 @@ def pretty_print_tests( def build_cpu_pytest_cmd(filenames: list[str], continue_on_error: bool) -> list[str]: """Build the single pytest invocation for a CPU suite. + Files are passed in sorted order so every package directory's arguments + stay contiguous. pytest 9.1 binds conftest fixtures to the package + collector *instance*, yet re-collects a package's children -- overwriting + the collection cache -- whenever an argument is a file directly inside it. + An interleaving like ``pkg/test_a.py ancestor/test_b.py pkg/test_c.py`` + therefore rebuilds pkg's collector chain after its conftest fixtures were + bound to the old instance, and pkg/test_c.py dies at setup with "fixture + not found". Sorted paths keep each package's block contiguous, so a + package is never re-entered after an ancestor-level file re-collect. + `-x` (stop at first failure) is the default regular-run behavior. With continue_on_error -- e.g. a PR carrying the `bypass-fastfail` label -- drop `-x` so every file runs; pytest still exits non-zero if any failed, so the stage stays red. """ - cmd = ["pytest", *filenames, "-v"] + cmd = ["pytest", *sorted(filenames), "-v"] if not continue_on_error: cmd.append("-x") return cmd diff --git a/tests/ci/test/test_run_suite.py b/tests/ci/test/test_run_suite.py index 03dbff18358..ec7a44aa80d 100644 --- a/tests/ci/test/test_run_suite.py +++ b/tests/ci/test/test_run_suite.py @@ -87,6 +87,28 @@ def test_x_dropped_on_continue_on_error(self): assert cmd[0] == "pytest" assert "tests/fast/a.py" in cmd and "tests/fast/b.py" in cmd + def test_files_sorted_so_package_args_stay_contiguous(self): + # pytest 9.1 re-collects a package's children (clobbering the + # collection cache) when an argument is a file directly inside it, + # while conftest fixtures stay bound to the original collector + # instance. The order "pkg file, ancestor-level file, pkg file" then + # errors at setup with "fixture not found". Sorting keeps each + # package's arguments contiguous, so that interleave cannot occur. + cmd = build_cpu_pytest_cmd( + [ + "tests/fast/ray/rollout/test_z.py", + "tests/fast/test_mid.py", + "tests/fast/ray/rollout/test_a.py", + ], + continue_on_error=False, + ) + files = [part for part in cmd if part.endswith(".py")] + assert files == [ + "tests/fast/ray/rollout/test_a.py", + "tests/fast/ray/rollout/test_z.py", + "tests/fast/test_mid.py", + ] + # --- CI_SUITES locked to the stage taxonomy --------------------------------- diff --git a/tests/ci/verify_source_resolution.py b/tests/ci/verify_source_resolution.py index 8ab42287e89..5c7d8aff33e 100644 --- a/tests/ci/verify_source_resolution.py +++ b/tests/ci/verify_source_resolution.py @@ -1,4 +1,7 @@ +import ast +import importlib import os +import pkgutil from importlib.util import find_spec from pathlib import Path @@ -11,6 +14,99 @@ "megatron.training": "MEGATRON_SOURCE_ROOT", } +# Namespaces that only exist once the multi-lora stack lands; walked in full when present. +OPTIONAL_PACKAGES = ( + "miles.backends.megatron_utils.api_backends", + "miles.ray.multi_lora", + "miles.rollout.multi_lora", + "miles.ray.tinker_frontend", +) + + +def _miles_roots() -> tuple[Path, Path]: + spec = find_spec("miles") + if spec is None or spec.origin is None: + raise RuntimeError("cannot resolve miles") + miles_root = Path(spec.origin).resolve().parent + return miles_root, miles_root.parent + + +def _module_file_exists(repo_root: Path, dotted: str) -> bool: + path = repo_root.joinpath(*dotted.split(".")) + return path.with_suffix(".py").is_file() or (path / "__init__.py").is_file() + + +def _resolve_relative(miles_root: Path, repo_root: Path, py_file: Path, node: ast.ImportFrom) -> str | None: + if not py_file.is_relative_to(miles_root): + return None + parts = list(py_file.relative_to(repo_root).parts) + package = parts[:-1] + if node.level > 1: + package = package[: -(node.level - 1)] + return ".".join(package + node.module.split(".")) if node.module else ".".join(package) + + +def _iter_miles_import_targets(miles_root: Path, repo_root: Path, py_file: Path): + tree = ast.parse(py_file.read_text(), filename=str(py_file)) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name.partition(".")[0] == "miles": + yield node.lineno, alias.name + elif isinstance(node, ast.ImportFrom): + target = _resolve_relative(miles_root, repo_root, py_file, node) if node.level else node.module + if target and target.partition(".")[0] == "miles": + yield node.lineno, target + + +def _python_files(root: Path): + return (p for p in sorted(root.rglob("*.py")) if "__pycache__" not in p.parts) + + +def verify_import_sites_resolve() -> None: + miles_root, repo_root = _miles_roots() + stale = [] + roots = [miles_root] + ([repo_root / "examples"] if (repo_root / "examples").is_dir() else []) + for root in roots: + for py_file in _python_files(root): + for lineno, target in _iter_miles_import_targets(miles_root, repo_root, py_file): + if not _module_file_exists(repo_root, target): + stale.append(f"{py_file.relative_to(repo_root)}:{lineno}: {target}") + if stale: + raise RuntimeError("stale miles-internal imports:\n" + "\n".join(stale)) + print("import-integrity: every miles import site resolves") + + +def verify_optional_namespaces_import() -> None: + for package_name in OPTIONAL_PACKAGES: + if find_spec(package_name) is None: + continue + package = importlib.import_module(package_name) + for info in pkgutil.walk_packages(package.__path__, prefix=package_name + "."): + importlib.import_module(info.name) + print(f"import-integrity: {package_name} imports in full") + + +def verify_update_weight_lazy_imports() -> None: + miles_root, repo_root = _miles_roots() + update_weight_dir = miles_root / "backends" / "megatron_utils" / "update_weight" + targets = sorted( + { + target + for py_file in _python_files(update_weight_dir) + for _, target in _iter_miles_import_targets(miles_root, repo_root, py_file) + } + ) + if not targets: + raise RuntimeError("expected function-local miles imports under update_weight/") + for target in targets: + try: + importlib.import_module(target) + except ModuleNotFoundError as exc: + if (exc.name or "").partition(".")[0] == "miles": + raise RuntimeError(f"update_weight lazy import target does not import: {target}") from exc + print("import-integrity: update_weight lazy imports resolve") + def main() -> None: for module_name, root_env in MODULE_ROOT_ENV.items(): @@ -24,6 +120,9 @@ def main() -> None: except ValueError as exc: raise RuntimeError(f"{module_name} resolved to {origin}, expected {expected_root}") from exc print(f"{module_name}: {origin}") + verify_import_sites_resolve() + verify_optional_namespaces_import() + verify_update_weight_lazy_imports() if __name__ == "__main__": diff --git a/tests/e2e/multi_lora_operations/multi_lora_e2e_client.py b/tests/e2e/multi_lora_operations/multi_lora_e2e_client.py new file mode 100644 index 00000000000..5793c3a3859 --- /dev/null +++ b/tests/e2e/multi_lora_operations/multi_lora_e2e_client.py @@ -0,0 +1,658 @@ +#!/usr/bin/env python3 +import argparse +import json +import math +import os +import sys +import time +import urllib.error +import urllib.request +import uuid + +import ray + +API = "http://127.0.0.1:8068" +ROUTER = "http://127.0.0.1:20080" # rebound to the head-node IP at startup +NAME = "e2e_a" +SAVE_ROOT = "/personal/tinker_e2e/save" # rebound from --save-root + +PASS: list[str] = [] +FAIL: list[str] = [] + + +def report(phase: str, ok: bool, detail: str) -> None: + tag = "PASS" if ok else "FAIL" + (PASS if ok else FAIL).append(phase) + print(f"[{tag}] {phase}: {detail}", flush=True) + if not ok: + print("--- aborting on first failure ---", flush=True) + sys.exit(1) + + +def http(method: str, path: str, body: dict | None = None, base: str = API) -> dict: + req = urllib.request.Request( + base + path, + method=method, + data=json.dumps(body).encode() if body is not None else None, + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(req, timeout=120) as resp: + return json.loads(resp.read()) + + +def wait_state(name: str, want: str, timeout_s: float = 300) -> str: + deadline = time.monotonic() + timeout_s + state = None + while time.monotonic() < deadline: + state = http("GET", f"/adapter_runs/state?names={name}")["states"].get(name) + if state == want: + return state + time.sleep(2) + raise TimeoutError(f"adapter '{name}' state {state!r}, wanted {want!r} within {timeout_s}s") + + +class Ops: + """Operation plane over the controller Ray actor.""" + + def __init__(self): + self.controller = ray.get_actor("miles_tinker_controller", namespace="miles") + # Ordinals are consecutive from 1 PER REGISTRATION; a re-registered + # name is a new tenant and restarts at 1 (reset_ordinals). + self.ordinals: dict[str, int] = {} + + def enqueue(self, kind: str, payload: dict | None = None, name: str = NAME) -> str: + ordinal = self.ordinals.get(name, 0) + 1 + self.ordinals[name] = ordinal + op_id = f"op-{name}-{ordinal}-{kind}-{uuid.uuid4().hex[:8]}" + view = ray.get(self.controller.enqueue_operation.remote(name, op_id, ordinal, kind, payload)) + assert view["state"] == "QUEUED", view + return op_id + + def wait(self, op_id: str, timeout_s: float = 600) -> dict: + deadline = time.monotonic() + timeout_s + view = None + while time.monotonic() < deadline: + view = ray.get(self.controller.get_operation.remote(op_id)) + if view is not None and view["state"] in ("SUCCEEDED", "FAILED", "CANCELLED"): + return view + time.sleep(1) + raise TimeoutError(f"operation {op_id} not terminal within {timeout_s}s: {view}") + + def ack(self, op_id: str) -> None: + ray.get(self.controller.ack_operation.remote(op_id)) + + def run(self, kind: str, payload: dict | None = None, name: str = NAME, timeout_s: float = 600) -> dict: + op_id = self.enqueue(kind, payload, name=name) + view = self.wait(op_id, timeout_s) + self.ack(op_id) + return view + + def snapshot(self) -> dict: + return ray.get(self.controller.snapshot.remote()) + + def step_of(self, name: str) -> int: + return ray.get(self.controller.adapter_step.remote(name)) + + def reset_ordinals(self, name: str) -> None: + self.ordinals.pop(name, None) + + +def fb_payload(sample_lens: list[tuple[int, int]], base_token: int = 2000) -> dict: + """CE forward_backward payload: (total_len, response_len) per sample.""" + samples = [] + for i, (total, resp) in enumerate(sample_lens): + tokens = [base_token + i * 100 + j for j in range(total)] + samples.append( + dict( + tokens=tokens, + response_length=resp, + loss_mask=[1] * resp, + loss_weights=[1.0 / resp] * resp, + ) + ) + return dict(samples=samples, loss=dict(loss_fn="cross_entropy")) + + +def check_fb_result(view: dict, sample_lens: list[tuple[int, int]], phase: str) -> list[list[float]]: + ok = view["state"] == "SUCCEEDED" + detail = f"state={view['state']}" + logprobs, loss = None, None + if ok: + result = view["result"] or {} + logprobs = result.get("logprobs") + metrics = result.get("metrics") or {} + loss = metrics.get("loss:sum") + shapes_ok = ( + isinstance(logprobs, list) + and len(logprobs) == len(sample_lens) + and all(len(lp) == resp for lp, (_, resp) in zip(logprobs, sample_lens, strict=True)) + ) + loss_ok = isinstance(loss, float) and math.isfinite(loss) + ok = shapes_ok and loss_ok + detail = ( + f"state=SUCCEEDED shapes={[len(lp) for lp in logprobs] if isinstance(logprobs, list) else None} " + f"want={[r for _, r in sample_lens]} loss:sum={loss} " + f"unmasked_tokens:sum={metrics.get('unmasked_tokens:sum')}" + ) + else: + detail += f" error={view.get('error')}" + report(phase, ok, detail) + return logprobs + + +def check_optim(view: dict, phase: str, expect_zero: bool = False) -> float | None: + result = view.get("result") or {} + grad_norm = result.get("grad_norm") + finite = isinstance(grad_norm, float) and math.isfinite(grad_norm) + ok = view["state"] == "SUCCEEDED" and finite and (grad_norm == 0.0 if expect_zero else grad_norm > 0) + report( + phase, + ok, + f"state={view['state']} grad_norm={grad_norm} lr={result.get('learning_rate')} error={view.get('error')}", + ) + return grad_norm + + +def max_logprob_delta(a: list[list[float]], b: list[list[float]]) -> float: + return max(abs(x - y) for row_a, row_b in zip(a, b, strict=True) for x, y in zip(row_a, row_b, strict=True)) + + +def register(ops: Ops, name: str, rank: int = 8) -> dict: + """Register and wait READY; returns {slot, registration_id}. A fresh + registration is a new tenant: its operation ordinals restart at 1.""" + ops.reset_ordinals(name) + reg = http("POST", "/adapter_runs", {"name": name, "config": {"rank": rank}}) + wait_state(name, "READY", timeout_s=600) + info = http("GET", f"/adapter_runs/{name}") + return {"slot": reg.get("slot"), "registration_id": info["registration_id"]} + + +def deregister(ops: Ops, name: str, timeout_s: float = 300) -> str: + http("DELETE", f"/adapter_runs/{name}") + deadline = time.monotonic() + timeout_s + state = None + while time.monotonic() < deadline: + state = http("GET", f"/adapter_runs/state?names={name}")["states"].get(name) + if state == "COMPLETED": + return state + time.sleep(2) + raise TimeoutError(f"adapter '{name}' not COMPLETED within {timeout_s}s (state={state})") + + +def sidecar_manifest(name: str) -> str: + return f"{SAVE_ROOT}/adapters/{name}/slot_state/manifest.pt" + + +# Adapter lifecycle at DP=2. + + +def phase_a(ops: Ops) -> None: + from miles.ray.multi_lora.identity import serving_lora_name # noqa: PLC0415 + + reg = http("POST", "/adapter_runs", {"name": NAME, "config": {"rank": 8}}) + slot_bound = reg.get("slot") is not None + state = wait_state(NAME, "READY", timeout_s=600) + info = http("GET", f"/adapter_runs/{NAME}") + registration_id = info["registration_id"] + report( + "phase1-register", + slot_bound and state == "READY", + f"slot={reg.get('slot')} state={state} rid={registration_id[:8]}", + ) + + fb_shapes = [ + [(24, 16), (20, 12), (28, 16)], # 3 samples: count not divisible by DP=2 + [(16, 8), (32, 24)], + [(24, 16), (24, 16), (20, 10), (30, 20)], + ] + fb3_payload = fb_payload(fb_shapes[2], base_token=5000) + payloads = [fb_payload(fb_shapes[0]), fb_payload(fb_shapes[1], base_token=3500), fb3_payload] + fb3_logprobs = None + for i, (shapes, payload) in enumerate(zip(fb_shapes, payloads, strict=True), start=1): + view = ops.run("forward_backward", payload) + lp = check_fb_result(view, shapes, f"phase2-fb{i}") + if i == 3: + fb3_logprobs = lp + + # One-sample fb: at DP=2 one whole rank runs only the zero-weight padding + # row; the result plane must carry exactly the client's single row. + view = ops.run("forward_backward", fb_payload([(22, 14)], base_token=7000)) + check_fb_result(view, [(22, 14)], "phase2-fb-odd1") + + view = ops.run("optim_step", dict(adam_params=dict(learning_rate=1e-4))) + check_optim(view, "phase3-optim_step") + + view = ops.run("save_weights_for_sampler", {}) + result = view.get("result") or {} + serving_version = result.get("serving_version") + serving_name = result.get("serving_name") + expected_name = serving_lora_name(NAME, registration_id) + ok = view["state"] == "SUCCEEDED" and serving_version == 1 and serving_name == expected_name + report( + "phase4-save_weights_for_sampler", + ok, + f"state={view['state']} serving_version={serving_version} serving_name={serving_name} error={view.get('error')}", + ) + + sample_body = dict( + text="The capital of France is", + sampling_params=dict(max_new_tokens=8, temperature=0.0), + lora_path=serving_name, + ) + try: + gen = http("POST", "/generate", sample_body, base=ROUTER) + text = gen.get("text") + report("phase4-sample", isinstance(text, str) and len(text) > 0, f"text={text!r}") + except urllib.error.HTTPError as e: + report("phase4-sample", False, f"HTTP {e.code}: {e.read().decode()[:500]}") + + view = ops.run("save_state", dict(tag="e2e-t0")) + result = view.get("result") or {} + state_path = result.get("path") + manifest_ok = bool(state_path) and os.path.exists(os.path.join(state_path, "manifest.pt")) + ok = view["state"] == "SUCCEEDED" and manifest_ok and result.get("step") == 1 + report( + "phase5-save_state", + ok, + f"state={view['state']} path={state_path} manifest={manifest_ok} step={result.get('step')} error={view.get('error')}", + ) + + view = ops.run("load_state", dict(path=state_path)) + result = view.get("result") or {} + ok = view["state"] == "SUCCEEDED" and result.get("step") == 1 + report("phase6-load_state", ok, f"state={view['state']} step={result.get('step')} error={view.get('error')}") + + view = ops.run("forward_backward", fb3_payload) + fb4_logprobs = check_fb_result(view, fb_shapes[2], "phase6-fb-post-restore") + + # weights actually moved: identical payload, logprobs differ pre/post optim + max_delta = max_logprob_delta(fb3_logprobs, fb4_logprobs) + report("phase6-weights-moved", max_delta > 1e-9, f"max |dlogprob| fb3 vs post-optim fb = {max_delta:.6g}") + + view = ops.run("optim_step", dict(adam_params=dict(learning_rate=1e-4))) + check_optim(view, "phase6-optim-post-restore") + + http("DELETE", f"/adapter_runs/{NAME}") + deadline = time.monotonic() + 300 + final_state, snapshot = None, None + while time.monotonic() < deadline: + snapshot = ops.snapshot() + final_state = http("GET", f"/adapter_runs/state?names={NAME}")["states"].get(NAME) + if final_state == "COMPLETED": + break + time.sleep(2) + slot_free = ( + NAME not in {**snapshot["pending"], **snapshot["ready"], **snapshot["retiring"]} + and NAME not in snapshot["cleanup"] + ) + + sidecar_ok = os.path.exists(sidecar_manifest(NAME)) + + rejected = False + reject_detail = "enqueue unexpectedly accepted" + try: + ops.enqueue("forward_backward", fb_payload([(16, 8)])) + except Exception as e: # noqa: BLE001 + rejected = "not accepting operations" in str(e) or "fenced" in str(e) + reject_detail = str(e).splitlines()[-1][:200] + ok = final_state == "COMPLETED" and slot_free and sidecar_ok and rejected + report( + "phase7-deregister", + ok, + f"final_state={final_state} slot_free={slot_free} sidecar={sidecar_ok} post-dereg-enqueue-rejected={rejected} ({reject_detail})", + ) + + reg_b = http("POST", "/adapter_runs", {"name": "e2e_b", "config": {"rank": 8}}) + state_b = wait_state("e2e_b", "READY", timeout_s=600) + http("DELETE", "/adapter_runs/e2e_b") + report( + "phase7-second-adapter", + reg_b.get("slot") is not None and state_b == "READY", + f"slot={reg_b.get('slot')} state={state_b}", + ) + wait_state("e2e_b", "COMPLETED", timeout_s=300) + + +# Forward-only operations and empty optimizer steps. + + +def phase_b(ops: Ops) -> None: + name = "e2e_f" + reg = register(ops, name) + report("phaseB-register", reg["slot"] is not None, f"slot={reg['slot']} rid={reg['registration_id'][:8]}") + + shapes = [(24, 16), (20, 12)] + payload = fb_payload(shapes, base_token=9000) + + view = ops.run("forward", dict(samples=payload["samples"]), name=name) + result = view.get("result") or {} + fwd_logprobs = result.get("logprobs") + shapes_ok = ( + view["state"] == "SUCCEEDED" + and isinstance(fwd_logprobs, list) + and len(fwd_logprobs) == len(shapes) + and all(len(lp) == resp for lp, (_, resp) in zip(fwd_logprobs, shapes, strict=True)) + ) + report( + "phaseB-forward", + shapes_ok, + f"state={view['state']} shapes={[len(lp) for lp in fwd_logprobs] if isinstance(fwd_logprobs, list) else None} " + f"metrics={result.get('metrics')} error={view.get('error')}", + ) + + # no dirty pin: a save_state right after the forward must not be rejected + # by the unstepped-gradients gate + view = ops.run("save_state", dict(tag="b-nodirty"), name=name) + dirty_gated = "unstepped gradients" in (view.get("error") or "") + report( + "phaseB-no-dirty-pin", + view["state"] == "SUCCEEDED" and not dirty_gated, + f"state={view['state']} error={view.get('error')}", + ) + + # optim_step with nothing accumulated: the backend contract is an empty + # step — SUCCEEDED with grad_norm == 0.0 (fresh Adam moments: weights + # cannot move), never a user-side rejection + view = ops.run("optim_step", dict(adam_params=dict(learning_rate=1e-4)), name=name) + check_optim(view, "phaseB-optim-after-forward-only", expect_zero=True) + step = ops.step_of(name) + report("phaseB-empty-step-clock", step == 1, f"step={step} (empty optim_step advances the clock)") + + # identical payload through forward_backward: same weights (the empty step + # moved nothing), so the logprob planes must agree + view = ops.run("forward_backward", payload, name=name) + fb_logprobs = check_fb_result(view, shapes, "phaseB-fb-same-payload") + delta = max_logprob_delta(fwd_logprobs, fb_logprobs) + report("phaseB-forward-vs-fb-logprobs", delta <= 1e-4, f"max |dlogprob| forward vs fb = {delta:.6g}") + + # the fb DID pin dirty (contrast with the forward): its optim_step has real gradients + view = ops.run("optim_step", dict(adam_params=dict(learning_rate=1e-4)), name=name) + check_optim(view, "phaseB-optim-after-fb") + + deregister(ops, name) + report("phaseB-deregister", True, "COMPLETED") + + +# Slot-state ownership at DP=2. + + +def _rank_swapped_copy(state_path: str, dest: str) -> str: + """A byte-identical copy of a two-rank state with the rank shards swapped: + same save generation, same shapes, but each rank now reads a shard whose + recorded per-rank ownership signature is the OTHER rank's — exactly the + 'sharded with a different per-rank parameter ownership' condition.""" + import shutil # noqa: PLC0415 + + if os.path.isdir(dest): + shutil.rmtree(dest) + shutil.copytree(state_path, dest) + r0, r1 = os.path.join(dest, "shard_rank00000.pt"), os.path.join(dest, "shard_rank00001.pt") + tmp = os.path.join(dest, "shard_rank_tmp.pt") + os.rename(r0, tmp) + os.rename(r1, r0) + os.rename(tmp, r1) + return dest + + +def phase_c(ops: Ops) -> None: + import torch # noqa: PLC0415 + + reg = register(ops, "e2e_c") + report("phaseC-register-slot0", reg["slot"] == 0, f"slot={reg['slot']}") + ops.run("forward_backward", fb_payload([(24, 16), (20, 12)], base_token=11000), name="e2e_c") + view = ops.run("optim_step", dict(adam_params=dict(learning_rate=1e-4)), name="e2e_c") + check_optim(view, "phaseC-seed-optim") + view = ops.run("save_state", dict(tag="c0"), name="e2e_c") + state_path = (view.get("result") or {}).get("path") + report( + "phaseC-save-slot0-state", + view["state"] == "SUCCEEDED" and (view.get("result") or {}).get("step") == 1, + f"state={view['state']} path={state_path} step={(view.get('result') or {}).get('step')}", + ) + + # LayerWise DP sharding is real: the two rank shards carry disjoint, + # non-trivial ownership signatures + sig = [ + torch.load(os.path.join(state_path, f"shard_rank{r:05d}.pt"), map_location="cpu", weights_only=True)[ + "optimizer_param_names" + ] + for r in (0, 1) + ] + flat0 = {name for child in sig[0] for name in child} + flat1 = {name for child in sig[1] for name in child} + report( + "phaseC-dp-sharding-real", + sig[0] != sig[1] and flat0 and flat1 and not (flat0 & flat1), + f"rank0 owns {len(flat0)} params, rank1 owns {len(flat1)}, overlap {len(flat0 & flat1)}", + ) + deregister(ops, "e2e_c") + + reg1 = register(ops, "e2e_c1") + reg2 = register(ops, "e2e_c2") + report("phaseC-slot-arrangement", reg1["slot"] == 0 and reg2["slot"] == 1, f"c1={reg1['slot']} c2={reg2['slot']}") + + # Cross-slot restore under MATCHING signatures: on this deployment every + # numel-class block is a multiple of 4, so slot 0 and slot 1 get identical + # per-rank ownership in LayerWise's DP-2 ping-pong — the fence must allow + # the restore (the contract is signature equality, not same-slot). + view = ops.run("load_state", dict(path=state_path), name="e2e_c2") + restored = view["state"] == "SUCCEEDED" and (view.get("result") or {}).get("step") == 1 + report( + "phaseC-cross-slot-matching-sig-restore", + restored and ops.step_of("e2e_c2") == 1, + f"state={view['state']} step={(view.get('result') or {}).get('step')} error={view.get('error')}", + ) + + # ... and bitwise-correctly: a state saved back out of slot 1 carries the + # same weights and optimizer tensors (only the slot tag differs) + view = ops.run("save_state", dict(tag="c2snap"), name="e2e_c2") + snap_path = (view.get("result") or {}).get("path") + mismatch = None + for r in (0, 1): + shard = f"shard_rank{r:05d}.pt" + before = torch.load(os.path.join(state_path, shard), map_location="cpu", weights_only=True) + after = torch.load(os.path.join(snap_path, shard), map_location="cpu", weights_only=True) + for key in ("weights", "optimizer_state", "optimizer_param_names"): + mismatch = mismatch or _payload_tensors_equal(before[key], after[key], f"{shard}:{key}") + report( + "phaseC-cross-slot-restore-correct", + mismatch is None, + "slot0 state == slot1 re-save (bitwise)" if mismatch is None else mismatch, + ) + + # A state with a genuinely DIFFERENT per-rank ownership (the same save + # with its rank shards swapped) must be refused by the ownership fence: + # clean user-category failure, unanimous across ranks, nothing mutated. + swapped = _rank_swapped_copy(state_path, os.path.join(os.path.dirname(state_path), "c0-rankswap")) + view = ops.run("load_state", dict(path=swapped), name="e2e_c2") + fence_msg = view.get("error") or "" + refused = ( + view["state"] == "FAILED" + and view.get("error_category") == "user" + and "different per-rank parameter ownership" in fence_msg + ) + report( + "phaseC-ownership-fence-refused", + refused, + f"state={view['state']} category={view.get('error_category')} error={fence_msg[:220]}", + ) + + # trainer stayed healthy: the refused tenant keeps training + ops.run("forward_backward", fb_payload([(18, 10)], base_token=12000), name="e2e_c2") + view = ops.run("optim_step", dict(adam_params=dict(learning_rate=1e-4)), name="e2e_c2") + check_optim(view, "phaseC-post-refusal-train") + + view = ops.run("load_state", dict(path=state_path), name="e2e_c1") + restored = view["state"] == "SUCCEEDED" and (view.get("result") or {}).get("step") == 1 + step = ops.step_of("e2e_c1") + report( + "phaseC-same-slot-restore", + restored and step == 1, + f"state={view['state']} result_step={(view.get('result') or {}).get('step')} registry_step={step} " + f"error={view.get('error')}", + ) + ops.run("forward_backward", fb_payload([(24, 16)], base_token=13000), name="e2e_c1") + view = ops.run("optim_step", dict(adam_params=dict(learning_rate=1e-4)), name="e2e_c1") + check_optim(view, "phaseC-post-restore-train") + + deregister(ops, "e2e_c1") + deregister(ops, "e2e_c2") + + # sidecar variant of the fence: swap the retired tenant's sidecar shards + # so its recorded ownership is foreign on every rank; re-registration must + # fall back to a fresh init (no crash, step 0) instead of resuming it + sidecar_base = os.path.dirname(sidecar_manifest("e2e_c2")) + r0, r1 = os.path.join(sidecar_base, "shard_rank00000.pt"), os.path.join(sidecar_base, "shard_rank00001.pt") + tmp = os.path.join(sidecar_base, "shard_rank_tmp.pt") + os.rename(r0, tmp) + os.rename(r1, r0) + os.rename(tmp, r1) + reg2b = register(ops, "e2e_c2") + step = ops.step_of("e2e_c2") + report( + "phaseC-foreign-sidecar-fresh-init", + reg2b["slot"] == 0 and step == 0, + f"slot={reg2b['slot']} step={step} (rank-swapped sidecar refused by the fence; reconcile fresh-inits)", + ) + ops.run("forward_backward", fb_payload([(18, 10)], base_token=14000), name="e2e_c2") + view = ops.run("optim_step", dict(adam_params=dict(learning_rate=1e-4)), name="e2e_c2") + check_optim(view, "phaseC-fresh-init-train") + deregister(ops, "e2e_c2") + + +# Sidecar resume preserves step, weights, and FP32 masters. + + +def _payload_tensors_equal(a, b, where: str = "") -> str | None: + """First mismatch path between two saved payload subtrees, or None.""" + import torch # noqa: PLC0415 + + if isinstance(a, torch.Tensor) or isinstance(b, torch.Tensor): + if not (isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor)): + return f"{where}: tensor vs {type(b).__name__}" + return None if torch.equal(a, b) else f"{where}: tensors differ (max|d|={(a - b).abs().max().item():.3g})" + if isinstance(a, dict) and isinstance(b, dict): + if a.keys() != b.keys(): + return f"{where}: keys {sorted(a)} != {sorted(b)}" + for key in a: + if key == "miles_multi_lora_slot": # the destination slot's tag: differs across slots by design + continue + if (m := _payload_tensors_equal(a[key], b[key], f"{where}.{key}")) is not None: + return m + return None + if isinstance(a, (list, tuple)) and isinstance(b, (list, tuple)): + if len(a) != len(b): + return f"{where}: length {len(a)} != {len(b)}" + for i, (x, y) in enumerate(zip(a, b, strict=True)): + if (m := _payload_tensors_equal(x, y, f"{where}[{i}]")) is not None: + return m + return None + return None if a == b else f"{where}: {a!r} != {b!r}" + + +def phase_d(ops: Ops) -> None: + import torch # noqa: PLC0415 + + name = "e2e_d" + probe = dict(samples=fb_payload([(26, 18), (20, 12)], base_token=15000)["samples"]) + + reg = register(ops, name) + report("phaseD-register", reg["slot"] is not None and ops.step_of(name) == 0, f"slot={reg['slot']} step=0") + + # two real steps so the resume has a non-trivial clock and Adam state + for i, base in enumerate((16000, 17000), start=1): + ops.run("forward_backward", fb_payload([(24, 16), (28, 18)], base_token=base), name=name) + view = ops.run("optim_step", dict(adam_params=dict(learning_rate=1e-4)), name=name) + check_optim(view, f"phaseD-optim{i}") + + view = ops.run("save_weights_for_sampler", {}, name=name) + serving_version = (view.get("result") or {}).get("serving_version") + report( + "phaseD-publish", + view["state"] == "SUCCEEDED" and serving_version == 1, + f"state={view['state']} serving_version={serving_version}", + ) + + view = ops.run("forward", probe, name=name) + probe_before = (view.get("result") or {}).get("logprobs") + report("phaseD-probe-before", view["state"] == "SUCCEEDED" and probe_before is not None, "captured L1") + + deregister(ops, name) + sidecar_ok = os.path.exists(sidecar_manifest(name)) + report("phaseD-final-sidecar", sidecar_ok, sidecar_manifest(name)) + + # re-register the SAME name: reconcile must auto-resume from the sidecar + reg2 = register(ops, name) + step = ops.step_of(name) + report("phaseD-resume-step", step == 2, f"slot={reg2['slot']} restored step={step} (want 2)") + + view = ops.run("forward", probe, name=name) + probe_after = (view.get("result") or {}).get("logprobs") + delta = max_logprob_delta(probe_before, probe_after) + report("phaseD-resume-logprobs", delta <= 1e-6, f"max |dlogprob| pre-dereg vs post-resume = {delta:.6g}") + + # the resumed masters are the checkpoint's fp32 masters, NOT re-quantized + # through bf16: a state saved now must carry bitwise-identical weights and + # optimizer state (fp32 masters + Adam moments) to the retirement sidecar + view = ops.run("save_state", dict(tag="d-resumed"), name=name) + resumed_path = (view.get("result") or {}).get("path") + report("phaseD-save-resumed", view["state"] == "SUCCEEDED" and resumed_path is not None, f"path={resumed_path}") + + sidecar_base = os.path.dirname(sidecar_manifest(name)) + shards = sorted(f for f in os.listdir(sidecar_base) if f.startswith("shard_rank")) + mismatch, compared = None, 0 + for shard in shards: + before = torch.load(os.path.join(sidecar_base, shard), map_location="cpu", weights_only=True) + after = torch.load(os.path.join(resumed_path, shard), map_location="cpu", weights_only=True) + for key in ("weights", "optimizer_state", "optimizer_param_names"): + mismatch = mismatch or _payload_tensors_equal(before[key], after[key], f"{shard}:{key}") + compared += 1 + report( + "phaseD-fp32-masters-preserved", + compared > 0 and mismatch is None, + f"{compared} rank shards bitwise-compared (weights + optimizer fp32 masters/moments): " + + ("identical" if mismatch is None else mismatch), + ) + + # and training continues from the restored state + ops.run("forward_backward", fb_payload([(24, 16), (20, 12)], base_token=18000), name=name) + view = ops.run("optim_step", dict(adam_params=dict(learning_rate=1e-4)), name=name) + check_optim(view, "phaseD-post-resume-train") + step = ops.step_of(name) + report("phaseD-post-resume-step", step == 3, f"step={step} (want 3)") + + deregister(ops, name) + report("phaseD-deregister", True, "COMPLETED") + + +PHASES = {"a": phase_a, "b": phase_b, "c": phase_c, "d": phase_d} + + +def main() -> None: + global SAVE_ROOT, ROUTER + parser = argparse.ArgumentParser() + parser.add_argument("--ray-address", default="auto") + parser.add_argument("--phases", default="a,b,c,d", help="comma-separated subset of a,b,c,d") + parser.add_argument("--save-root", default=SAVE_ROOT, help="the service's --save dir (sidecar/state paths)") + args = parser.parse_args() + ray.init(address=args.ray_address, namespace="miles", ignore_reinit_error=True, log_to_driver=False) + + SAVE_ROOT = args.save_root.rstrip("/") + + ops = Ops() + # The sglang router binds the node IP (the control API advertises its own + # loopback bind host, which never reaches the router's socket). + from miles.utils.misc import get_current_node_ip # noqa: PLC0415 + + ROUTER = f"http://{get_current_node_ip()}:20080" + print(f"router: {ROUTER}", flush=True) + + for phase in args.phases.split(","): + print(f"\n=== phase {phase.upper()} ===", flush=True) + PHASES[phase.strip().lower()](ops) + + print(f"\n=== E2E SUMMARY: {len(PASS)} passed, {len(FAIL)} failed ===", flush=True) + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/multi_lora_operations/multi_lora_rl_quality.py b/tests/e2e/multi_lora_operations/multi_lora_rl_quality.py new file mode 100644 index 00000000000..5205d5445a4 --- /dev/null +++ b/tests/e2e/multi_lora_operations/multi_lora_rl_quality.py @@ -0,0 +1,475 @@ +#!/usr/bin/env python3 +"""Run four concurrent client-driven GRPO loops against the Multi-LoRA backend.""" + +import argparse +import csv +import json +import os +import statistics +import sys +import threading +import time +import urllib.error +import urllib.request +import uuid +from dataclasses import dataclass, field + +import ray + +API = "http://127.0.0.1:8068" + +DEFAULT_SPECS = [ + # Each adapter uses a disjoint quarter of the GSM8K training split. + dict(name="rl_a", rank=8, lr=1e-5, shard=0), + dict(name="rl_b", rank=16, lr=2e-5, shard=1), + dict(name="rl_c", rank=16, lr=4e-5, shard=2), + dict(name="rl_d", rank=32, lr=1e-5, shard=3), +] + + +def http(method: str, path: str, body: dict | None = None, base: str = API, timeout: float = 900) -> dict: + req = urllib.request.Request( + base + path, + method=method, + data=json.dumps(body).encode() if body is not None else None, + headers={"Content-Type": "application/json"}, + ) + try: + with urllib.request.urlopen(req, timeout=timeout) as resp: + return json.loads(resp.read()) + except urllib.error.HTTPError as e: + detail = "" + try: + detail = e.read().decode()[:500] + except Exception: # noqa: BLE001,S110 - the status code alone still identifies the failure + pass + raise RuntimeError(f"HTTP {e.code} on {method} {path}: {detail}") from e + + +def discover_router(explicit: str | None, candidates=(20080, 30080)) -> str: + """The sglang router binds the node IP; the port may differ from the + requested one when something (e.g. a stray nginx) squats it. Probe the + worker-listing endpoints to find the live router.""" + if explicit: + return explicit.rstrip("/") + from miles.utils.misc import get_current_node_ip # noqa: PLC0415 + + ip = get_current_node_ip() + for port in candidates: + base = f"http://{ip}:{port}" + for endpoint in ("/list_workers", "/workers"): + try: + body = http("GET", endpoint, base=base, timeout=5) + if isinstance(body, dict) and ("urls" in body or "workers" in body): + return base + except Exception: # noqa: BLE001,S112 - probe failures just move to the next candidate + continue + raise RuntimeError(f"no router found on ports {candidates} at {ip}") + + +class Ops: + """Operation plane over the controller Ray actor (one shared handle; each + adapter thread only touches its own name's ordinal counter).""" + + def __init__(self): + self.controller = ray.get_actor("miles_tinker_controller", namespace="miles") + self.ordinals: dict[str, int] = {} + + def enqueue(self, name: str, kind: str, payload: dict | None = None) -> str: + ordinal = self.ordinals.get(name, 0) + 1 + self.ordinals[name] = ordinal + op_id = f"op-{name}-{ordinal}-{kind}-{uuid.uuid4().hex[:8]}" + view = ray.get(self.controller.enqueue_operation.remote(name, op_id, ordinal, kind, payload)) + assert view["state"] == "QUEUED", view + return op_id + + def wait(self, op_id: str, timeout_s: float = 1800) -> dict: + deadline = time.monotonic() + timeout_s + view = None + while time.monotonic() < deadline: + view = ray.get(self.controller.get_operation.remote(op_id)) + if view is not None and view["state"] in ("SUCCEEDED", "FAILED", "CANCELLED"): + return view + time.sleep(1) + raise TimeoutError(f"operation {op_id} not terminal within {timeout_s}s: {view}") + + def run(self, name: str, kind: str, payload: dict | None = None, timeout_s: float = 1800) -> dict: + op_id = self.enqueue(name, kind, payload) + view = self.wait(op_id, timeout_s) + ray.get(self.controller.ack_operation.remote(op_id)) + return view + + def step_of(self, name: str) -> int: + return ray.get(self.controller.adapter_step.remote(name)) + + +@dataclass +class StepRecord: + step: int + t_start: float + dt_s: float + n_prompts: int + n_samples: int + reward_mean: float + reward_std: float + mean_resp_len: float + frac_stop: float + frac_zero_adv: float + loss_sum: float | None + grad_norm: float | None + logprob_absdiff_mean: float | None + serving_version: int | None + note: str = "" + + +@dataclass +class AdapterRun: + spec: dict + registration_id: str = "" + serving_name: str = "" + records: list[StepRecord] = field(default_factory=list) + error: str | None = None + final_step_clock: int | None = None + final_serving_version: int | None = None + + +def group_advantages(rewards: list[float], group_size: int) -> list[float]: + """GRPO-style per-prompt advantages: mean baseline, std-normalized.""" + advantages = [] + for start in range(0, len(rewards), group_size): + group = rewards[start : start + group_size] + mean = sum(group) / len(group) + std = statistics.pstdev(group) + advantages.extend([(r - mean) / (std + 1e-6) if std > 0 else 0.0 for r in group]) + return advantages + + +def sample_batch(router: str, serving_name: str, prompts: list[list[int]], args) -> list[dict]: + """One batched /generate: each prompt replicated n times, temperature 1.0 + so the returned logprobs are the sampling distribution's.""" + input_ids = [ids for ids in prompts for _ in range(args.samples_per_prompt)] + body = dict( + input_ids=input_ids, + sampling_params=dict( + temperature=1.0, + top_p=1.0, + top_k=-1, + max_new_tokens=args.max_new_tokens, + ), + lora_path=serving_name, + return_logprob=True, + ) + outputs = http("POST", "/generate", body, base=router, timeout=args.sample_timeout_s) + assert isinstance(outputs, list) and len(outputs) == len(input_ids), f"batch size mismatch: {len(outputs)}" + return outputs + + +def adapter_loop(run: AdapterRun, ops: Ops, router: str, dataset: list[dict], grade, args, log) -> None: + spec = run.spec + name = spec["name"] + + reg = http("POST", "/adapter_runs", {"name": name, "config": {"rank": spec["rank"]}}) + deadline = time.monotonic() + 900 + while time.monotonic() < deadline: + if http("GET", f"/adapter_runs/state?names={name}")["states"].get(name) == "READY": + break + time.sleep(2) + else: + raise TimeoutError(f"adapter '{name}' never became READY") + info = http("GET", f"/adapter_runs/{name}") + run.registration_id = info["registration_id"] + from miles.ray.multi_lora.identity import serving_lora_name # noqa: PLC0415 + + run.serving_name = serving_lora_name(name, run.registration_id) + log( + f"({name}) registered: slot={reg.get('slot')} rank={spec['rank']} lr={spec['lr']} rid={run.registration_id[:8]}" + ) + + # Publish the fresh (identity) adapter before the first sampling round so + # the serving name exists on the engines. + view = ops.run(name, "save_weights_for_sampler", {}) + assert view["state"] == "SUCCEEDED", f"({name}) initial publish failed: {view.get('error')}" + + cursor = 0 + step = 0 + while step < args.steps: + t0 = time.time() + note = "" + + prompts, labels = [], [] + while len(prompts) < args.prompts_per_step: + row = dataset[cursor % len(dataset)] + cursor += 1 + ids = row["input_ids"] + if 0 < len(ids) <= args.max_prompt_tokens: + prompts.append(ids) + labels.append(row["label"]) + + try: + outputs = sample_batch(router, run.serving_name, prompts, args) + except (urllib.error.URLError, RuntimeError, TimeoutError, AssertionError) as e: + log(f"({name}) step {step + 1}: sampling failed ({e}); retrying next round") + time.sleep(5) + continue + + samples, rewards, resp_lens, stops = [], [], [], 0 + for i, out in enumerate(outputs): + prompt_ids = prompts[i // args.samples_per_prompt] + label = labels[i // args.samples_per_prompt] + token_logprobs = (out.get("meta_info") or {}).get("output_token_logprobs") or [] + resp_tokens = [int(t[1]) for t in token_logprobs] + resp_logprobs = [float(t[0]) for t in token_logprobs] + reward = 1.0 if resp_tokens and grade(out.get("text") or "", label) else 0.0 + finish = ((out.get("meta_info") or {}).get("finish_reason") or {}).get("type") + stops += finish == "stop" + rewards.append(reward) + resp_lens.append(len(resp_tokens)) + samples.append( + dict( + tokens=prompt_ids + resp_tokens, + response_length=len(resp_tokens), + loss_mask=[1] * len(resp_tokens), + rollout_log_probs=resp_logprobs, + ) + ) + + # Grouped advantages; sample-mean scaling folds GRPO's normalization + # into the per-token channel (the backend's loss is a plain token sum). + advantages = group_advantages(rewards, args.samples_per_prompt) + usable = [i for i, s in enumerate(samples) if s["response_length"] > 0] + n_usable = len(usable) + for i in usable: + r_len = samples[i]["response_length"] + per_token = advantages[i] / (r_len * n_usable) + samples[i]["advantages"] = [per_token] * r_len + + reward_mean = sum(rewards) / len(rewards) + reward_std = statistics.pstdev(rewards) + frac_zero_adv = sum(1 for i in usable if advantages[i] == 0.0) / max(n_usable, 1) + + loss_sum = grad_norm = absdiff = version = None + optim_ok = False + try: + fb = ops.run( + name, + "forward_backward", + dict(samples=[samples[i] for i in usable], loss=dict(loss_fn="importance_sampling")), + timeout_s=args.op_timeout_s, + ) + if fb["state"] != "SUCCEEDED": + raise RuntimeError(f"forward_backward FAILED: {fb.get('error')}") + metrics = (fb.get("result") or {}).get("metrics") or {} + loss_sum = metrics.get("loss:sum") + train_logprobs = (fb.get("result") or {}).get("logprobs") or [] + diffs = [ + abs(tr - ro) + for lp_row, i in zip(train_logprobs, usable, strict=True) + for tr, ro in zip(lp_row, samples[i]["rollout_log_probs"], strict=True) + ] + absdiff = sum(diffs) / len(diffs) if diffs else None + + optim = ops.run( + name, + "optim_step", + dict(adam_params=dict(learning_rate=spec["lr"], grad_clip_norm=1.0)), + timeout_s=args.op_timeout_s, + ) + if optim["state"] != "SUCCEEDED": + raise RuntimeError(f"optim_step FAILED: {optim.get('error')}") + optim_ok = True + grad_norm = (optim.get("result") or {}).get("grad_norm") + + publish = ops.run(name, "save_weights_for_sampler", {}, timeout_s=args.op_timeout_s) + if publish["state"] != "SUCCEEDED": + note = f"publish FAILED (sampling stays on previous version): {publish.get('error')}" + else: + version = (publish.get("result") or {}).get("serving_version") + except Exception as e: # noqa: BLE001 - an op failure is a per-step finding; the loop continues + note = f"{type(e).__name__}: {str(e)[:300]}" + log(f"({name}) step {step + 1}: {note}") + if not optim_ok: + # No optimizer step landed: this round is not a step. + time.sleep(2) + continue + + step += 1 + rec = StepRecord( + step=step, + t_start=t0, + dt_s=time.time() - t0, + n_prompts=len(prompts), + n_samples=n_usable, + reward_mean=reward_mean, + reward_std=reward_std, + mean_resp_len=sum(resp_lens) / max(len(resp_lens), 1), + frac_stop=stops / len(outputs), + frac_zero_adv=frac_zero_adv, + loss_sum=loss_sum, + grad_norm=grad_norm, + logprob_absdiff_mean=absdiff, + serving_version=version, + note=note, + ) + run.records.append(rec) + log( + f"({name}) step {step}/{args.steps}: reward={reward_mean:.3f} grad_norm={grad_norm} " + f"absdiff={absdiff if absdiff is None else round(absdiff, 4)} version={version} dt={rec.dt_s:.1f}s" + ) + + run.final_step_clock = ops.step_of(name) + run.final_serving_version = http("GET", f"/adapter_runs/{name}").get("version") + if args.deregister: + http("DELETE", f"/adapter_runs/{name}") + + +def least_squares_slope(ys: list[float]) -> float: + n = len(ys) + if n < 2: + return 0.0 + xs = range(1, n + 1) + mean_x, mean_y = (n + 1) / 2, sum(ys) / n + num = sum((x - mean_x) * (y - mean_y) for x, y in zip(xs, ys, strict=True)) + den = sum((x - mean_x) ** 2 for x in xs) + return num / den + + +def write_csv(run: AdapterRun, out_dir: str) -> str: + path = os.path.join(out_dir, f"{run.spec['name']}.csv") + fields = [f for f in StepRecord.__dataclass_fields__] + with open(path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fields) + writer.writeheader() + for rec in run.records: + writer.writerow({k: getattr(rec, k) for k in fields}) + return path + + +def main() -> None: + global API + parser = argparse.ArgumentParser() + parser.add_argument("--ray-address", default="auto") + parser.add_argument("--api", default=API) + parser.add_argument("--router", default=None, help="router base URL; discovered from 20080/30080 when omitted") + parser.add_argument("--data", default="/root/gsm8k/train.parquet") + parser.add_argument("--tokenizer", default="/root/models/Qwen3-4B") + parser.add_argument("--out-dir", required=True) + parser.add_argument("--steps", type=int, default=50) + parser.add_argument("--prompts-per-step", type=int, default=8) + parser.add_argument("--samples-per-prompt", type=int, default=4) + parser.add_argument("--max-new-tokens", type=int, default=512) + parser.add_argument("--max-prompt-tokens", type=int, default=1024) + parser.add_argument("--sample-timeout-s", type=float, default=900) + parser.add_argument("--op-timeout-s", type=float, default=1800) + parser.add_argument("--deregister", action="store_true", help="deregister adapters after the run") + parser.add_argument( + "--enable-thinking", + action="store_true", + help="Qwen3 thinking mode. With a tight max_new_tokens budget the base policy mostly truncates " + "(low initial reward), which is exactly the headroom the reward-growth check needs; non-thinking " + "GSM8K starts near 0.9 and has almost no group variance left to learn from.", + ) + args = parser.parse_args() + + API = args.api + os.makedirs(args.out_dir, exist_ok=True) + + ray.init(address=args.ray_address, namespace="miles", ignore_reinit_error=True, log_to_driver=False) + router = discover_router(args.router) + print(f"router: {router}", flush=True) + + import pandas as pd # noqa: PLC0415 + from transformers import AutoTokenizer # noqa: PLC0415 + + from miles.rollout.rm_hub.math_utils import grade_answer_verl # noqa: PLC0415 + + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer) + df = pd.read_parquet(args.data) + + specs = DEFAULT_SPECS + shards: dict[int, list[dict]] = {} + for spec in specs: + rows = df.iloc[spec["shard"] :: len(specs)] + shard = [] + for _, row in rows.iterrows(): + messages = [dict(m) for m in row["messages"]] + encoded = tokenizer.apply_chat_template( + messages, tokenize=True, add_generation_prompt=True, enable_thinking=args.enable_thinking + ) + # transformers >= 5 returns a BatchEncoding; earlier versions a flat list. + input_ids = encoded["input_ids"] if not isinstance(encoded, list) else encoded + if input_ids and isinstance(input_ids[0], list): + input_ids = input_ids[0] + shard.append(dict(input_ids=[int(t) for t in input_ids], label=str(row["label"]))) + shards[spec["shard"]] = shard + print(f"shard {spec['shard']}: {len(shard)} prompts", flush=True) + + ops = Ops() + log_lock = threading.Lock() + + def log(msg: str) -> None: + with log_lock: + print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True) + + runs = [AdapterRun(spec=spec) for spec in specs] + threads = [] + for run in runs: + thread = threading.Thread( + target=_thread_main, + args=(run, ops, router, shards[run.spec["shard"]], grade_answer_verl, args, log), + name=run.spec["name"], + daemon=True, + ) + thread.start() + threads.append(thread) + for thread in threads: + thread.join() + + summary = {} + for run in runs: + rewards = [rec.reward_mean for rec in run.records] + first10 = rewards[:10] + last10 = rewards[-10:] + summary[run.spec["name"]] = dict( + spec={k: v for k, v in run.spec.items()}, + steps_recorded=len(run.records), + step_clock=run.final_step_clock, + serving_version=run.final_serving_version, + reward_first10_mean=sum(first10) / len(first10) if first10 else None, + reward_last10_mean=sum(last10) / len(last10) if last10 else None, + reward_slope_per_step=least_squares_slope(rewards), + logprob_absdiff_mean=( + sum(r.logprob_absdiff_mean for r in run.records if r.logprob_absdiff_mean is not None) + / max(sum(1 for r in run.records if r.logprob_absdiff_mean is not None), 1) + ), + mean_step_dt_s=sum(r.dt_s for r in run.records) / max(len(run.records), 1), + failures=[f"step {r.step}: {r.note}" for r in run.records if r.note], + error=run.error, + csv=write_csv(run, args.out_dir), + ) + with open(os.path.join(args.out_dir, "summary.json"), "w") as f: + json.dump(summary, f, indent=2) + print(json.dumps(summary, indent=2), flush=True) + + grew = sum( + 1 + for s in summary.values() + if s["reward_first10_mean"] is not None and s["reward_last10_mean"] > s["reward_first10_mean"] + ) + print(f"\n=== RL QUALITY: reward grew (last10 > first10) on {grew}/{len(runs)} adapters ===", flush=True) + + # An aborted loop is recorded in run.error by its thread; it must fail the process, not just the summary. + aborted = [run.spec["name"] for run in runs if run.error] + if aborted: + sys.exit(f"RL quality FAILED: adapter loop(s) aborted: {', '.join(aborted)}") + + +def _thread_main(run: AdapterRun, ops: Ops, router: str, dataset, grade, args, log) -> None: + try: + adapter_loop(run, ops, router, dataset, grade, args, log) + except Exception as e: # noqa: BLE001 - a dead loop is a finding, not a crash of the harness + run.error = f"{type(e).__name__}: {e}" + log(f"({run.spec['name']}) LOOP ABORTED: {run.error}") + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/tinker_frontend/tinker_sdk_mini_loop.py b/tests/e2e/tinker_frontend/tinker_sdk_mini_loop.py new file mode 100644 index 00000000000..636663a2943 --- /dev/null +++ b/tests/e2e/tinker_frontend/tinker_sdk_mini_loop.py @@ -0,0 +1,242 @@ +#!/usr/bin/env python3 +"""Golden-acceptance mini-loop: the UNMODIFIED official ``tinker==0.24.1`` SDK +drives the miles tinker frontend end to end, cookbook style. + + ServiceClient(base_url, api_key) + -> get_server_capabilities (the deployment's one base model) + -> create_lora_training_client(rank=16) + -> ~10x [forward_backward(cross_entropy, teacher-forced prompt-masked + SFT datums: prompt weight 0, completion weight 1) + + optim_step(AdamParams(lr=1e-4))] loss:sum must decrease + -> save_weights_and_get_sampling_client -> sample coherent continuation + -> save_state -> load_state_with_optimizer -> one more fb/optim + -> out-of-order large fb (>MAX_CHUNK_LEN datums: the SDK splits chunks + and posts the first one LAST; the backend ledger reorders) + -> a deliberate channel-mismatch datum surfacing as a typed SDK error; + it poisons its gradient window (#2258 §5) so the window's optim_step + fails as a discard, and the next round steps normally + +The SFT per-token loss divides by ``loss_weight:sum`` (Σ weight·mask), NOT by +``unmasked_tokens:sum`` — the latter counts the weight-0 prompt positions too +(codex-0817-sft-fix §7). The prompt masking here keeps the two metrics +distinct, so this loop regression-tests the denominator on real GPUs: with +the old all-ones weights they were equal and the bug was invisible. + +Run on the head node from a venv with ``tinker==0.24.1`` installed: + python tests/e2e/tinker_frontend/tinker_sdk_mini_loop.py --out-dir +""" + +import argparse +import json +import os +import time + +import tinker +from tinker import types + +CORPUS = [ + "The old lighthouse keeper climbed the spiral stairs every evening at dusk.", + "He lit the great lamp so that ships could find their way home through the fog.", + "One autumn night a fierce storm rolled in from the north and shook the tower.", + "The keeper held his lantern steady and watched the waves crash on the rocks.", + "By morning the sea was calm again and a small fishing boat waved its thanks.", + "The keeper smiled, poured his tea, and wrote the night's story in his logbook.", + "Years later his granddaughter found the logbook and read every page aloud.", + "She decided then that she too would keep the light burning for the ships.", +] + +SAMPLE_PROMPT = "The old lighthouse keeper climbed" + + +def ce_datum(tokens: list[int]) -> types.Datum: + """Plain LM datum: model_input = tokens[:-1], next-token targets, weight 1.""" + inputs, targets = tokens[:-1], tokens[1:] + return types.Datum( + model_input=types.ModelInput.from_ints(inputs), + loss_fn_inputs={"target_tokens": targets, "weights": [1.0] * len(targets)}, + ) + + +def sft_datum(prompt_tokens: list[int], completion_tokens: list[int]) -> tuple[types.Datum, float, int]: + """Teacher-forced SFT datum (the correct shape, codex-0817-sft-fix §2): + position i predicts tokens[i+1], so the prompt-internal next-token + positions get weight 0 and the completion positions weight 1. Returns the + datum plus its CE weight sum and its total target-position count.""" + tokens = prompt_tokens + completion_tokens + weights = [0.0] * (len(prompt_tokens) - 1) + [1.0] * len(completion_tokens) + datum = types.Datum( + model_input=types.ModelInput.from_ints(tokens[:-1]), + loss_fn_inputs={"target_tokens": tokens[1:], "weights": weights}, + ) + return datum, sum(weights), len(weights) + + +def split_prompt_completion(text: str) -> tuple[str, str]: + """First half of the words is the prompt (weight 0), the rest completion.""" + words = text.split() + split = max(1, len(words) // 2) + return " ".join(words[:split]), " " + " ".join(words[split:]) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--base-url", default="http://127.0.0.1:8068") + parser.add_argument("--api-key", default=os.environ.get("MILES_TINKER_API_KEY", "tml-miles-gpu-acceptance")) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--iterations", type=int, default=10) + parser.add_argument("--lr", type=float, default=1e-4) + parser.add_argument("--large-fb-datums", type=int, default=1030, help=">1024 forces multi-chunk posting") + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + summary: dict = {} + + def log(msg: str) -> None: + print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True) + + service = tinker.ServiceClient(base_url=args.base_url, api_key=args.api_key) + + # ---- capabilities: the deployment serves exactly one base model ---- + capabilities = service.get_server_capabilities() + base_models = [m.model_name for m in capabilities.supported_models] + log(f"server capabilities: supported_models={base_models}") + assert len(base_models) == 1 and base_models[0], base_models + base_model = base_models[0] + summary["base_model"] = base_model + + client = service.create_lora_training_client(base_model=base_model, rank=16) + info = client.get_info() + assert info.lora_rank == 16, info + log(f"training client ready: model_id={client.model_id} rank={info.lora_rank}") + + tokenizer = client.get_tokenizer() + pairs = [split_prompt_completion(text) for text in CORPUS] + built = [sft_datum(tokenizer.encode(prompt), tokenizer.encode(completion)) for prompt, completion in pairs] + data = [datum for datum, _, _ in built] + expected_weight_sum = sum(weight_sum for _, weight_sum, _ in built) + expected_positions = sum(positions for _, _, positions in built) + assert expected_positions > expected_weight_sum > 0, (expected_positions, expected_weight_sum) + n_tokens = sum(len(d.model_input.to_ints()) for d in data) + log( + f"corpus: {len(data)} prompt-masked SFT datums, {n_tokens} input tokens, " + f"{expected_weight_sum:.0f} completion positions of {expected_positions} targets" + ) + + # ---- supervised mini-loop: loss must decrease ---- + losses: list[float] = [] + t0 = time.time() + for iteration in range(1, args.iterations + 1): + fb_future = client.forward_backward(data, "cross_entropy") + optim_future = client.optim_step(types.AdamParams(learning_rate=args.lr)) + fb = fb_future.result() + optim = optim_future.result() + loss_sum = fb.metrics["loss:sum"] + # The SFT denominator is the CE weight sum (completion positions), + # not unmasked_tokens:sum, which also counts the weight-0 prompt + # (codex-0817-sft-fix §7). Guarded: weights are arbitrary floats. + weight_sum = fb.metrics["loss_weight:sum"] + unmasked = fb.metrics["unmasked_tokens:sum"] + assert abs(weight_sum - expected_weight_sum) < 1e-6, (weight_sum, expected_weight_sum) + assert abs(unmasked - expected_positions) < 1e-6, (unmasked, expected_positions) + assert unmasked > weight_sum, "prompt masking must keep the two denominators distinct" + per_token = loss_sum / weight_sum if weight_sum > 0 else None + losses.append(loss_sum) + log( + f"iter {iteration:2d}/{args.iterations}: loss:sum={loss_sum:.3f} " + f"per_token={per_token:.4f} grad_norm={optim.metrics.get('grad_norm')}" + ) + train_dt = time.time() - t0 + summary["losses"] = losses + summary["loss_weight_sum"] = expected_weight_sum + summary["unmasked_tokens"] = expected_positions + summary["train_seconds"] = round(train_dt, 1) + assert losses[-1] < losses[0], f"loss did not decrease: {losses}" + assert all(b <= a * 1.02 for a, b in zip(losses, losses[1:], strict=False)), f"loss not (near-)monotone: {losses}" + log(f"loss decreased {losses[0]:.3f} -> {losses[-1]:.3f} over {args.iterations} iterations ({train_dt:.0f}s)") + + # ---- publish + sample: the tuned adapter must speak ---- + sampling = client.save_weights_and_get_sampling_client() + assert sampling.get_base_model() == base_model + prompt_ids = tokenizer.encode(SAMPLE_PROMPT) + response = sampling.sample( + prompt=types.ModelInput.from_ints(prompt_ids), + num_samples=2, + sampling_params=types.SamplingParams(max_tokens=24, temperature=0.0), + ).result() + continuations = [tokenizer.decode(seq.tokens) for seq in response.sequences] + for i, (seq, text) in enumerate(zip(response.sequences, continuations, strict=True)): + log(f"sample[{i}] stop={seq.stop_reason} logprobs[:3]={[round(p, 3) for p in (seq.logprobs or [])[:3]]}") + log(f"sample[{i}] text: {SAMPLE_PROMPT}{text!s}") + assert seq.tokens and seq.logprobs and len(seq.logprobs) == len(seq.tokens) + summary["sample_prompt"] = SAMPLE_PROMPT + summary["sample_continuations"] = continuations + + # ---- save_state -> load_state_with_optimizer -> training continues ---- + path = client.save_state("mini-loop-golden").result().path + log(f"save_state -> {path}") + assert path.startswith("tinker://") + client.load_state_with_optimizer(path).result() + fb = client.forward_backward(data, "cross_entropy").result() + client.optim_step(types.AdamParams(learning_rate=args.lr)).result() + resumed_loss = fb.metrics["loss:sum"] + summary["checkpoint_path"] = path + summary["loss_after_restore"] = resumed_loss + # The restored state is the post-loop state: its loss must match the + # trained trajectory, not the untrained start. + assert resumed_loss < losses[0], (resumed_loss, losses[0]) + log(f"restored from checkpoint; fb/optim after load works (loss:sum={resumed_loss:.3f})") + + # ---- large out-of-order fb: SDK chunks >MAX_CHUNK_LEN and posts the ---- + # ---- first chunk last; the backend gap-buffers and reassembles. ---- + short = tokenizer.encode("The sea was calm.") + big = [ce_datum(short) for _ in range(args.large_fb_datums)] + t1 = time.time() + result = client.forward_backward(big, "cross_entropy").result() + client.optim_step(types.AdamParams(learning_rate=0.0)).result() # release the dirty-grad pin + assert len(result.loss_fn_outputs) == args.large_fb_datums, len(result.loss_fn_outputs) + row = result.loss_fn_outputs[0]["logprobs"].tolist() + assert len(row) == len(short) - 1 + summary["large_fb"] = {"datums": args.large_fb_datums, "seconds": round(time.time() - t1, 1)} + log( + f"large fb: {args.large_fb_datums} datums (multi-chunk, out-of-order) -> " + f"{len(result.loss_fn_outputs)} rows in {summary['large_fb']['seconds']}s" + ) + + # ---- deliberate user error: channel mismatch -> typed SDK error, no hang ---- + bad = types.Datum( + model_input=types.ModelInput.from_ints(short[:-1]), + loss_fn_inputs={"target_tokens": short[1:], "advantages": [1.0] * (len(short) - 1)}, + # importance_sampling requires 'logprobs'; it is deliberately missing. + ) + t2 = time.time() + try: + client.forward_backward([bad], "importance_sampling").result() + raise AssertionError("channel-mismatch datum was accepted") + except tinker.RequestFailedError as exc: + err_dt = time.time() - t2 + summary["typed_user_error"] = {"error": str(exc)[:200], "seconds": round(err_dt, 1)} + log(f"typed user error in {err_dt:.1f}s (no hang): {str(exc)[:120]}") + # The rejected submission consumed its ordinal AND poisoned its gradient + # window (#2258 §5): the window's optim_step must discard, not step. + good = client.forward_backward(data[:2], "cross_entropy").result() + assert len(good.loss_fn_outputs) == 2 + try: + client.optim_step(types.AdamParams(learning_rate=args.lr)).result() + raise AssertionError("optim_step on a poisoned window succeeded") + except tinker.RequestFailedError as exc: + assert "gradient window" in str(exc), exc + summary["poisoned_optim_error"] = str(exc)[:200] + log(f"poisoned-window optim_step failed typed: {str(exc)[:120]}") + # The discard reset the window: the next round steps normally. + good = client.forward_backward(data[:2], "cross_entropy").result() + client.optim_step(types.AdamParams(learning_rate=0.0)).result() + assert len(good.loss_fn_outputs) == 2 + log("post-error round stepped: the discard left no residue and no gap") + + summary["ok"] = True + with open(os.path.join(args.out_dir, "mini_loop_summary.json"), "w") as f: + json.dump(summary, f, indent=2) + log("=== MINI-LOOP GOLDEN ACCEPTANCE: PASS ===") + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/tinker_frontend/tinker_sdk_poison_window.py b/tests/e2e/tinker_frontend/tinker_sdk_poison_window.py new file mode 100644 index 00000000000..91ac06b5407 --- /dev/null +++ b/tests/e2e/tinker_frontend/tinker_sdk_poison_window.py @@ -0,0 +1,430 @@ +#!/usr/bin/env python3 +"""Poison-window GPU acceptance: a FAILED forward_backward chunk poisons the +registration's gradient window (#2258 §5) — the window's ``optim_step`` is +rejected ("gradient window ... discarded") and the trainer executes the +discard (``zero_adapter_slot_grads``) on EVERY rank instead of stepping. + +The CPU contract tests prove the control flow; this client proves the +collective semantics on a live DP>1 deployment through the UNMODIFIED +``tinker==0.24.1`` SDK: + + 1. baseline rank-8 client, 3x good fb+optim: finite losses/grad_norms, + step clock exactly 3, publish bumps serving_version to 1 + 2. poison capture probe logprobs L0 on a fixed payload (forward: no + gradients, no dirty pin) and a clean-window reference + grad_norm for the SAME batch; then good fb (EXECUTES into + the window) + channel-mismatch fb (typed reject) + optim. + The fb error is typed; the optim FAILS with the poison + message; step clock and serving version hold; the probe + re-reads EXACTLY L0 — the good chunk's gradients were + discarded on both ranks, no half-applied update + 3. recovery the same batch again: optim SUCCEEDS and its grad_norm + matches the clean reference exactly — the discard left no + residue on any rank (residue would double the norm). A + real step then MOVES the probe (sensitivity control) + 4. isolation a second adapter runs the poison sequence CONCURRENTLY + while the first trains normally: the victim's poison never + perturbs the neighbor's losses or step clock + 5. late chunk a 1030-datum fb whose LATE chunk carries the bad datum: the + SDK splits at 1024 and posts the first chunk last; the + 1024-datum chunk lands (real gradients on both ranks) + before the poison is seen — same discard assertions + +Step/serving clocks come from the operator plane (``GET /adapter_runs`` on +the same uvicorn; loopback-only), so run this on the head node from a venv +with ``tinker==0.24.1``: + python tests/e2e/tinker_frontend/tinker_sdk_poison_window.py --out-dir + +Numeric tolerances: dense deployments are bitwise-deterministic per forward, +so every probe comparison defaults to EXACT (0.0) and the grad-norm reference +comparison to rel_tol=1e-6. MoE deployments (e.g. GPT-OSS grouped-GEMM/Triton +kernels) have inherent run-to-run forward nondeterminism at the BASE model +(measured 0.09-0.21 max |dlogprob| on 4xH200 GPT-OSS 20B, pre-existing before +any multi-LoRA change), which fails the probe-stability precondition before +any mechanism is tested. For those deployments pass ``--probe-tolerance`` (and +``--grad-norm-rtol``) calibrated to the measured noise; the client then also +REQUIRES the real-step sensitivity to clear that tolerance (reporting the +margin), so a discard check can never hide a real update inside the noise +band — measured on 4xH200 GPT-OSS 20B at LR=1e-4: noise 0.130, real-step +movement 0.406 (3.1x the noise, 1.6x a 2x-noise tolerance). Every MECHANISM +assertion — typed fb/optim failures, step/serving clocks held, discard +executed, neighbor isolation, no-hang — stays exact regardless of tolerance. +""" + +import argparse +import json +import math +import os +import threading +import time +import urllib.request + +import tinker +from tinker import types + +CORPUS = [ + "The old lighthouse keeper climbed the spiral stairs every evening at dusk.", + "He lit the great lamp so that ships could find their way home through the fog.", + "One autumn night a fierce storm rolled in from the north and shook the tower.", + "The keeper held his lantern steady and watched the waves crash on the rocks.", + "By morning the sea was calm again and a small fishing boat waved its thanks.", + "The keeper smiled, poured his tea, and wrote the night's story in his logbook.", + "Years later his granddaughter found the logbook and read every page aloud.", + "She decided then that she too would keep the light burning for the ships.", +] + +LR = 1e-4 + +# Deployment noise tolerances; overridden from --probe-tolerance / +# --grad-norm-rtol in main(). 0.0 / 1e-6 = the exact dense-deployment contract. +PROBE_TOLERANCE = 0.0 +GRAD_NORM_RTOL = 1e-6 + + +def assert_probe_still(delta: float, what: str) -> None: + """A probe that must NOT have moved (discard/isolation checks): exact on + dense deployments, within the deployment's measured forward-noise band on + nondeterministic (MoE) ones.""" + assert delta <= PROBE_TOLERANCE, f"{what}: max|dlogprob|={delta} > tolerance {PROBE_TOLERANCE}" + + +def log(msg: str) -> None: + print(f"[{time.strftime('%H:%M:%S')}] [{threading.current_thread().name}] {msg}", flush=True) + + +def ce_datum(tokens: list[int]) -> types.Datum: + inputs, targets = tokens[:-1], tokens[1:] + return types.Datum( + model_input=types.ModelInput.from_ints(inputs), + loss_fn_inputs={"target_tokens": targets, "weights": [1.0] * len(targets)}, + ) + + +def channel_mismatch_datum(tokens: list[int]) -> types.Datum: + """importance_sampling requires 'logprobs'; deliberately missing -> the + frontend rejects the chunk typed, consuming (and poisoning) its ordinal.""" + return types.Datum( + model_input=types.ModelInput.from_ints(tokens[:-1]), + loss_fn_inputs={"target_tokens": tokens[1:], "advantages": [1.0] * (len(tokens) - 1)}, + ) + + +def bad_target_datum(tokens: list[int]) -> types.Datum: + """cross_entropy datum whose active target is not the next input token.""" + inputs, targets = tokens[:-1], list(tokens[1:]) + targets[0] += 7 # non-next-token target with non-zero weight -> typed reject + return types.Datum( + model_input=types.ModelInput.from_ints(inputs), + loss_fn_inputs={"target_tokens": targets, "weights": [1.0] * len(targets)}, + ) + + +# ---------------- operator plane (loopback, same uvicorn) ---------------- + + +def adapter_record(base_url: str, api_key: str, session_id: str) -> dict: + req = urllib.request.Request(f"{base_url}/adapter_runs", headers={"X-API-Key": api_key}) + with urllib.request.urlopen(req, timeout=30) as resp: + adapters = json.load(resp)["adapters"] + for status in adapters: + if (status.get("metadata") or {}).get("session_id") == session_id: + return status + raise AssertionError(f"no adapter registered for session {session_id}") + + +def session_of(client) -> str: + return client.model_id.split(":")[0] + + +def clocks(args, client) -> tuple[int, int, int]: + record = adapter_record(args.base_url, args.api_key, session_of(client)) + return record["step"], record["version"], record["slot"] + + +def wait_version(args, client, version: int, timeout: float = 180.0) -> None: + deadline = time.time() + timeout + while time.time() < deadline: + if clocks(args, client)[1] == version: + return + time.sleep(2) + raise AssertionError(f"serving version never reached {version}") + + +# ---------------- probes and typed-failure helpers ---------------- + + +def probe_rows(client, probe_data) -> list[list[float]]: + forward = client.forward(probe_data, "cross_entropy").result() + return [out["logprobs"].tolist() for out in forward.loss_fn_outputs] + + +def max_abs_delta(a: list[list[float]], b: list[list[float]]) -> float: + return max(abs(x - y) for ra, rb in zip(a, b, strict=True) for x, y in zip(ra, rb, strict=True)) + + +def expect_typed_failure(future, needle: str, what: str) -> str: + try: + future.result() + except tinker.RequestFailedError as exc: + message = str(exc) + assert needle in message, f"{what}: expected {needle!r} in: {message}" + return message + raise AssertionError(f"{what}: expected a typed RequestFailedError, got success") + + +def poison_round(args, client, data, bad_datum, bad_loss_fn, bad_needle) -> tuple[str, str]: + """Submit good fb + bad fb + optim in one window (cookbook style: all + posted before any await). Returns (fb_error, optim_error); asserts the + step clock and serving version held still.""" + step_pre, version_pre, _ = clocks(args, client) + good_future = client.forward_backward(data, "cross_entropy") + bad_future = client.forward_backward([bad_datum], bad_loss_fn) + optim_future = client.optim_step(types.AdamParams(learning_rate=LR)) + good = good_future.result() # the good chunk EXECUTED: gradients are live on every rank + assert len(good.loss_fn_outputs) == len(data) + t0 = time.time() + fb_error = expect_typed_failure(bad_future, bad_needle, "bad fb chunk") + optim_error = expect_typed_failure(optim_future, "gradient window", "poisoned optim_step") + assert "discarded" in optim_error, optim_error + log(f"typed fb reject + poisoned optim discard in {time.time() - t0:.1f}s (no hang)") + step_post, version_post, _ = clocks(args, client) + assert (step_post, version_post) == (step_pre, version_pre), ( + f"clocks moved across a poisoned window: step {step_pre}->{step_post}, " + f"version {version_pre}->{version_post}" + ) + return fb_error, optim_error + + +def train_round(client, data, lr: float = LR) -> tuple[float, float]: + fb_future = client.forward_backward(data, "cross_entropy") + optim_future = client.optim_step(types.AdamParams(learning_rate=lr)) + fb = fb_future.result() + optim = optim_future.result() + loss = fb.metrics["loss:sum"] + grad_norm = optim.metrics["grad_norm"] + assert math.isfinite(loss) and math.isfinite(grad_norm) and grad_norm > 0, (loss, grad_norm) + return loss, grad_norm + + +def assert_close(observed: float, reference: float, what: str) -> None: + assert math.isclose( + observed, reference, rel_tol=GRAD_NORM_RTOL + ), f"{what}: {observed} != {reference} (rel_tol {GRAD_NORM_RTOL})" + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--base-url", default="http://127.0.0.1:8068") + parser.add_argument("--api-key", default=os.environ.get("MILES_TINKER_API_KEY", "tml-miles-gpu-acceptance")) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--large-fb-datums", type=int, default=1030, help=">1024 forces multi-chunk posting") + parser.add_argument( + "--probe-tolerance", + type=float, + default=0.0, + help="allowed max |dlogprob| for probe comparisons; 0.0 (exact) for dense deployments, " + "the measured base-model forward-noise band for nondeterministic MoE kernels", + ) + parser.add_argument( + "--grad-norm-rtol", + type=float, + default=1e-6, + help="rel_tol for grad-norm reference comparisons (recovery/residue checks)", + ) + args = parser.parse_args() + global PROBE_TOLERANCE, GRAD_NORM_RTOL + PROBE_TOLERANCE = args.probe_tolerance + GRAD_NORM_RTOL = args.grad_norm_rtol + os.makedirs(args.out_dir, exist_ok=True) + summary: dict = {} + + service_a = tinker.ServiceClient(base_url=args.base_url, api_key=args.api_key) + capabilities = service_a.get_server_capabilities() + [base_model] = [m.model_name for m in capabilities.supported_models] + summary["base_model"] = base_model + + # ================= phase 1: baseline sanity ================= + client_a = service_a.create_lora_training_client(base_model=base_model, rank=8) + assert client_a.get_info().lora_rank == 8 + tokenizer = client_a.get_tokenizer() + data = [ce_datum(tokenizer.encode(text)) for text in CORPUS] + probe_data = [ce_datum(tokenizer.encode(text)) for text in CORPUS[:4]] + log(f"adapter A ready: model_id={client_a.model_id} rank=8") + + baseline = [train_round(client_a, data) for _ in range(3)] + step, version, slot_a = clocks(args, client_a) + assert step == 3, f"baseline step clock: {step} != 3" + sampling = client_a.save_weights_and_get_sampling_client() + assert sampling.get_base_model() == base_model + wait_version(args, client_a, 1) + summary["phase1_baseline"] = {"rounds": baseline, "step": step, "serving_version": 1, "slot": slot_a} + log(f"baseline: 3 rounds, losses {[round(loss, 3) for loss, _ in baseline]}, step=3, published version=1") + + # ================= phase 2: poison the window ================= + l0 = probe_rows(client_a, probe_data) + _, grad_norm_ref = train_round(client_a, data, lr=0.0) # clean-window reference, weights unchanged + l0_control = probe_rows(client_a, probe_data) + control_delta = max_abs_delta(l0_control, l0) + # Precondition: the deployment's inherent forward noise must sit inside + # the configured tolerance, or every later stillness check is meaningless. + assert_probe_still(control_delta, "probe not stable across an lr=0 round") + step_pre, version_pre, _ = clocks(args, client_a) + assert (step_pre, version_pre) == (4, 1) + + fb_error, optim_error = poison_round( + args, client_a, data, channel_mismatch_datum(tokenizer.encode(CORPUS[0])), "importance_sampling", "logprobs" + ) + l1 = probe_rows(client_a, probe_data) + poison_delta = max_abs_delta(l1, l0) + assert_probe_still(poison_delta, "weights moved across a poisoned window") + summary["phase2_poison"] = { + "grad_norm_ref": grad_norm_ref, + "control_probe_delta": control_delta, + "fb_error": fb_error[:200], + "optim_error": optim_error[:200], + "step_held": step_pre, + "version_held": version_pre, + "probe_delta_after_discard": poison_delta, + } + log(f"poison: optim rejected, step/version held at {step_pre}/{version_pre}, probe delta {poison_delta}") + + # ================= phase 3: recovery, no residue ================= + loss_rec, grad_norm_rec = train_round(client_a, data) # same batch, same weights + assert_close(grad_norm_rec, grad_norm_ref, "recovery grad_norm vs clean reference (residue would double it)") + step, version, _ = clocks(args, client_a) + assert step == step_pre + 1, f"recovery step clock: {step} != {step_pre + 1}" + l2 = probe_rows(client_a, probe_data) + sensitivity = max_abs_delta(l2, l0) + # The minimum meaningful bar: a real update must be distinguishable from + # the configured noise band, or the stillness checks above prove nothing. + assert sensitivity > PROBE_TOLERANCE, ( + f"probe blind: a real optim step moved the logprobs by {sensitivity}, " + f"inside the noise tolerance {PROBE_TOLERANCE}" + ) + if PROBE_TOLERANCE > 0.0: + log(f"sensitivity margin: real step moved {sensitivity:.4f} = {sensitivity / PROBE_TOLERANCE:.2f}x tolerance") + summary["phase3_recovery"] = { + "loss": loss_rec, + "grad_norm": grad_norm_rec, + "grad_norm_ref": grad_norm_ref, + "step": step, + "probe_moved_by_real_step": sensitivity, + } + log(f"recovery: grad_norm {grad_norm_rec} == ref {grad_norm_ref}, step->{step}, probe moved {sensitivity:.4f}") + + # ================= phase 4: concurrent isolation ================= + service_b = tinker.ServiceClient(base_url=args.base_url, api_key=args.api_key) + client_b = service_b.create_lora_training_client(base_model=base_model, rank=8) + _, _, slot_b = clocks(args, client_b) + assert slot_b != slot_a, (slot_a, slot_b) + lb0 = probe_rows(client_b, probe_data) + _, grad_norm_ref_b = train_round(client_b, data, lr=0.0) # quiet reference for B + assert_probe_still(max_abs_delta(probe_rows(client_b, probe_data), lb0), "B probe not stable") + step_a_pre = clocks(args, client_a)[0] + + barrier = threading.Barrier(2) + neighbor_rounds: list[tuple[float, float]] = [] + victim_errors: list[str] = [] + failures: list[BaseException] = [] + + def neighbor() -> None: + try: + barrier.wait(timeout=60) + for _ in range(4): + neighbor_rounds.append(train_round(client_a, data)) + except BaseException as exc: # noqa: BLE001 - surfaced after join + failures.append(exc) + + def victim() -> None: + try: + barrier.wait(timeout=60) + errors = poison_round( + args, + client_b, + data, + channel_mismatch_datum(tokenizer.encode(CORPUS[1])), + "importance_sampling", + "logprobs", + ) + victim_errors.extend(errors) + except BaseException as exc: # noqa: BLE001 - surfaced after join + failures.append(exc) + + threads = [ + threading.Thread(target=neighbor, name="neighbor-A"), + threading.Thread(target=victim, name="victim-B"), + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=600) + assert not thread.is_alive(), f"{thread.name} hung" + assert not failures, failures + + step_a_post = clocks(args, client_a)[0] + assert step_a_post == step_a_pre + 4, f"neighbor step clock: {step_a_pre}->{step_a_post}, expected +4" + neighbor_losses = [loss for loss, _ in neighbor_rounds] + assert neighbor_losses[-1] < neighbor_losses[0], f"neighbor loss did not decrease: {neighbor_losses}" + assert all(b <= a * 1.02 for a, b in zip(neighbor_losses, neighbor_losses[1:], strict=False)), neighbor_losses + lb1 = probe_rows(client_b, probe_data) + victim_delta = max_abs_delta(lb1, lb0) + assert_probe_still(victim_delta, "victim weights moved") + _, grad_norm_rec_b = train_round(client_b, data) # victim recovery (quiet) + assert_close(grad_norm_rec_b, grad_norm_ref_b, "victim recovery grad_norm vs quiet reference") + step_b, version_b, _ = clocks(args, client_b) + assert (step_b, version_b) == (2, 0), (step_b, version_b) + summary["phase4_isolation"] = { + "neighbor_losses": neighbor_losses, + "neighbor_grad_norms": [grad_norm for _, grad_norm in neighbor_rounds], + "neighbor_steps": [step_a_pre, step_a_post], + "victim_errors": [error[:200] for error in victim_errors], + "victim_probe_delta": victim_delta, + "victim_recovery_grad_norm": grad_norm_rec_b, + "victim_grad_norm_ref": grad_norm_ref_b, + } + log( + f"isolation: neighbor stepped {step_a_pre}->{step_a_post} losses {[round(loss, 2) for loss in neighbor_losses]}; " + f"victim poisoned+discarded (delta {victim_delta}), recovered grad_norm {grad_norm_rec_b}" + ) + + # ================= phase 5: late chunk fails after early chunk landed ================= + lb2 = probe_rows(client_b, probe_data) + _, grad_norm_ref_late = train_round(client_b, data, lr=0.0) + step_pre_b, version_pre_b, _ = clocks(args, client_b) + short = tokenizer.encode("The sea was calm.") + big = [ce_datum(short) for _ in range(args.large_fb_datums - 1)] + [bad_target_datum(short)] + fb_future = client_b.forward_backward(big, "cross_entropy") # 2 chunks; the bad datum rides the late one + optim_future = client_b.optim_step(types.AdamParams(learning_rate=LR)) + t0 = time.time() + fb_error = expect_typed_failure(fb_future, "next input", "late bad chunk") + optim_error = expect_typed_failure(optim_future, "gradient window", "poisoned optim after landed chunk") + log(f"late-chunk poison surfaced in {time.time() - t0:.1f}s") + step_post_b, version_post_b, _ = clocks(args, client_b) + assert (step_post_b, version_post_b) == (step_pre_b, version_pre_b) + lb3 = probe_rows(client_b, probe_data) + late_delta = max_abs_delta(lb3, lb2) + assert_probe_still(late_delta, "1024 landed datums leaked into the weights") + loss_late, grad_norm_late = train_round(client_b, data) # residue of 1024 datums would explode this + assert_close(grad_norm_late, grad_norm_ref_late, "post-late-chunk recovery grad_norm vs quiet reference") + summary["phase5_late_chunk"] = { + "datums": args.large_fb_datums, + "fb_error": fb_error[:200], + "optim_error": optim_error[:200], + "step_held": step_pre_b, + "probe_delta": late_delta, + "recovery_grad_norm": grad_norm_late, + "grad_norm_ref": grad_norm_ref_late, + "recovery_loss": loss_late, + } + log( + f"late chunk: discard held (delta {late_delta}), recovery grad_norm {grad_norm_late} == ref {grad_norm_ref_late}" + ) + + summary["ok"] = True + with open(os.path.join(args.out_dir, "poison_window_summary.json"), "w") as f: + json.dump(summary, f, indent=2) + log("=== POISON-WINDOW ACCEPTANCE (DP=2): PASS ===") + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/tinker_frontend/tinker_sdk_rl_quality.py b/tests/e2e/tinker_frontend/tinker_sdk_rl_quality.py new file mode 100644 index 00000000000..345ee2d15a8 --- /dev/null +++ b/tests/e2e/tinker_frontend/tinker_sdk_rl_quality.py @@ -0,0 +1,415 @@ +#!/usr/bin/env python3 +"""4-adapter RL training-quality acceptance, driven END TO END by the +UNMODIFIED official ``tinker==0.24.1`` SDK against the miles tinker frontend. + +The SDK port of tests/e2e/multi_lora_operations/multi_lora_rl_quality.py: four adapters +run concurrent, fully independent GRPO loops on disjoint GSM8K shards +(different ranks/learning rates), 50 optimizer steps each, Qwen3 thinking +mode with a tight max_tokens budget (the learnable regime). Per step and per +adapter, everything goes over /api/v1: + + SamplingClient.sample (num_samples per prompt, temp 1.0, logprobs back) + -> client-side math grading (reward 1/0) + -> grouped advantages (per-prompt mean baseline, std-normalized, + sample-mean token scaling) + -> TrainingClient.forward_backward(loss_fn="importance_sampling", + per-token advantages + the sampler's logprobs) + -> TrainingClient.optim_step(AdamParams(lr, grad_clip_norm=1.0)) + -> save_weights_and_get_sampling_client (publish barrier: the loop stays + on-policy, and the frontend fails stale samplers loudly by design) + +Serving version / step clock come from the operator /adapter_runs routes +(same uvicorn, X-API-Key). One CSV per adapter + summary.json, the same +schema as the raw-op acceptance run. + +Run on the head node from a venv with ``tinker==0.24.1`` installed +(PYTHONPATH must include the miles tree for the math grader): + python tests/e2e/tinker_frontend/tinker_sdk_rl_quality.py --out-dir +""" + +import argparse +import csv +import json +import os +import statistics +import threading +import time +import urllib.error +import urllib.request +from dataclasses import dataclass, field + +import tinker +from tinker import types + +DEFAULT_SPECS = [ + # name, lora rank, learning rate, gsm8k shard (disjoint quarter of train) + dict(name="rl_a", rank=8, lr=1e-5, shard=0), + dict(name="rl_b", rank=16, lr=2e-5, shard=1), + dict(name="rl_c", rank=16, lr=4e-5, shard=2), + dict(name="rl_d", rank=32, lr=1e-5, shard=3), +] + + +@dataclass +class StepRecord: + step: int + t_start: float + dt_s: float + n_prompts: int + n_samples: int + reward_mean: float + reward_std: float + mean_resp_len: float + frac_stop: float + frac_zero_adv: float + loss_sum: float | None + grad_norm: float | None + logprob_absdiff_mean: float | None + serving_version: int | None + note: str = "" + + +@dataclass +class AdapterRun: + spec: dict + model_id: str = "" + adapter_name: str = "" + registration_id: str = "" + records: list[StepRecord] = field(default_factory=list) + error: str | None = None + final_step_clock: int | None = None + final_serving_version: int | None = None + + +class OperatorApi: + """The registration control plane (same uvicorn as /api/v1); used only to + READ acceptance evidence: adapter name, step clock, serving version.""" + + def __init__(self, base: str, api_key: str) -> None: + self.base = base.rstrip("/") + self.api_key = api_key + + def get(self, path: str) -> dict: + req = urllib.request.Request(self.base + path, headers={"X-API-Key": self.api_key}) + with urllib.request.urlopen(req, timeout=60) as resp: + return json.loads(resp.read()) + + def find_adapter(self, model_id: str) -> dict: + session_id, seq = model_id.rsplit(":train:", 1) + for status in self.get("/adapter_runs")["adapters"]: + metadata = status.get("metadata") or {} + if metadata.get("session_id") == session_id and str(metadata.get("model_seq_id")) == seq: + return status + raise RuntimeError(f"no registration found for model '{model_id}'") + + def status_of(self, name: str) -> dict: + return self.get(f"/adapter_runs/{name}") + + +def group_advantages(rewards: list[float], group_size: int) -> list[float]: + """GRPO-style per-prompt advantages: mean baseline, std-normalized.""" + advantages = [] + for start in range(0, len(rewards), group_size): + group = rewards[start : start + group_size] + mean = sum(group) / len(group) + std = statistics.pstdev(group) + advantages.extend([(r - mean) / (std + 1e-6) if std > 0 else 0.0 for r in group]) + return advantages + + +def rl_datum(prompt_ids: list[int], resp_tokens: list[int], resp_logprobs: list[float], per_token_adv: float): + """Importance-sampling datum over the full sequence: zero advantage (and + zero rollout logprob) on the prompt span, the sampler's logprobs and the + scaled advantage on the response span. Next-token alignment holds by + construction, which is exactly what the frontend validates.""" + full = prompt_ids + resp_tokens + n_prompt = len(prompt_ids) + return types.Datum( + model_input=types.ModelInput.from_ints(full[:-1]), + loss_fn_inputs={ + "target_tokens": full[1:], + "logprobs": [0.0] * (n_prompt - 1) + resp_logprobs, + "advantages": [0.0] * (n_prompt - 1) + [per_token_adv] * len(resp_tokens), + }, + ) + + +def adapter_loop(run: AdapterRun, base_url, api_key, operator, dataset, tokenizer, grade, args, log): + spec = run.spec + name = spec["name"] + + # One ServiceClient (= one SDK session) per adapter: fully independent. + service = tinker.ServiceClient(base_url=base_url, api_key=api_key) + base_model = service.get_server_capabilities().supported_models[0].model_name + client = service.create_lora_training_client(base_model=base_model, rank=spec["rank"]) + run.model_id = str(client.model_id) + status = operator.find_adapter(run.model_id) + run.adapter_name = status["name"] + run.registration_id = status["registration_id"] + log( + f"({name}) model {run.model_id} -> registration '{run.adapter_name}' " + f"slot={status.get('slot')} rank={spec['rank']} lr={spec['lr']} rid={run.registration_id[:8]}" + ) + + # Publish the fresh (identity) adapter before the first sampling round. + sampling = client.save_weights_and_get_sampling_client() + + params = types.SamplingParams(max_tokens=args.max_new_tokens, temperature=1.0, top_p=1.0, top_k=-1) + cursor = 0 + step = 0 + while step < args.steps: + t0 = time.time() + note = "" + + prompts, labels = [], [] + while len(prompts) < args.prompts_per_step: + row = dataset[cursor % len(dataset)] + cursor += 1 + ids = row["input_ids"] + if 0 < len(ids) <= args.max_prompt_tokens: + prompts.append(ids) + labels.append(row["label"]) + + try: + futures = [ + sampling.sample( + prompt=types.ModelInput.from_ints(ids), + num_samples=args.samples_per_prompt, + sampling_params=params, + ) + for ids in prompts + ] + responses = [future.result() for future in futures] + except Exception as e: # noqa: BLE001 - a failed round is retried, not a crash + log(f"({name}) step {step + 1}: sampling failed ({type(e).__name__}: {str(e)[:200]}); retrying") + time.sleep(5) + continue + + datums, rewards, resp_lens, stops, n_seqs = [], [], [], 0, 0 + sample_rows = [] # (prompt_index, resp_tokens, resp_logprobs) + for prompt_index, response in enumerate(responses): + label = labels[prompt_index] + for seq in response.sequences: + n_seqs += 1 + resp_tokens = list(seq.tokens) + resp_logprobs = list(seq.logprobs or []) + reward = 1.0 if resp_tokens and grade(tokenizer.decode(resp_tokens), label) else 0.0 + stops += seq.stop_reason == "stop" + rewards.append(reward) + resp_lens.append(len(resp_tokens)) + sample_rows.append((prompt_index, resp_tokens, resp_logprobs)) + + advantages = group_advantages(rewards, args.samples_per_prompt) + usable = [i for i, (_, toks, _) in enumerate(sample_rows) if len(toks) > 0] + n_usable = len(usable) + for i in usable: + prompt_index, resp_tokens, resp_logprobs = sample_rows[i] + per_token = advantages[i] / (len(resp_tokens) * n_usable) + datums.append(rl_datum(prompts[prompt_index], resp_tokens, resp_logprobs, per_token)) + + reward_mean = sum(rewards) / len(rewards) + reward_std = statistics.pstdev(rewards) + frac_zero_adv = sum(1 for i in usable if advantages[i] == 0.0) / max(n_usable, 1) + + loss_sum = grad_norm = absdiff = version = None + optim_ok = False + try: + fb_future = client.forward_backward(datums, "importance_sampling") + optim_future = client.optim_step(types.AdamParams(learning_rate=spec["lr"], grad_clip_norm=1.0)) + fb = fb_future.result() + optim = optim_future.result() + optim_ok = True + loss_sum = fb.metrics.get("loss:sum") + grad_norm = optim.metrics.get("grad_norm") + + diffs = [] + for row_index, i in enumerate(usable): + prompt_index, resp_tokens, resp_logprobs = sample_rows[i] + train_row = fb.loss_fn_outputs[row_index]["logprobs"].tolist() + train_tail = train_row[len(prompts[prompt_index]) - 1 :] + diffs.extend(abs(tr - ro) for tr, ro in zip(train_tail, resp_logprobs, strict=True)) + absdiff = sum(diffs) / len(diffs) if diffs else None + + sampling = client.save_weights_and_get_sampling_client() + version = operator.status_of(run.adapter_name).get("version") + except Exception as e: # noqa: BLE001 - an op failure is a per-step finding + note = f"{type(e).__name__}: {str(e)[:300]}" + log(f"({name}) step {step + 1}: {note}") + if not optim_ok: + time.sleep(2) + continue + + step += 1 + rec = StepRecord( + step=step, + t_start=t0, + dt_s=time.time() - t0, + n_prompts=len(prompts), + n_samples=n_usable, + reward_mean=reward_mean, + reward_std=reward_std, + mean_resp_len=sum(resp_lens) / max(len(resp_lens), 1), + frac_stop=stops / max(n_seqs, 1), + frac_zero_adv=frac_zero_adv, + loss_sum=loss_sum, + grad_norm=grad_norm, + logprob_absdiff_mean=absdiff, + serving_version=version, + note=note, + ) + run.records.append(rec) + log( + f"({name}) step {step}/{args.steps}: reward={reward_mean:.3f} grad_norm={grad_norm} " + f"absdiff={absdiff if absdiff is None else round(absdiff, 4)} version={version} dt={rec.dt_s:.1f}s" + ) + + final = operator.status_of(run.adapter_name) + run.final_step_clock = final.get("step") + run.final_serving_version = final.get("version") + + +def least_squares_slope(ys: list[float]) -> float: + n = len(ys) + if n < 2: + return 0.0 + xs = range(1, n + 1) + mean_x, mean_y = (n + 1) / 2, sum(ys) / n + num = sum((x - mean_x) * (y - mean_y) for x, y in zip(xs, ys, strict=True)) + den = sum((x - mean_x) ** 2 for x in xs) + return num / den + + +def write_csv(run: AdapterRun, out_dir: str) -> str: + path = os.path.join(out_dir, f"{run.spec['name']}.csv") + fields = [f for f in StepRecord.__dataclass_fields__] + with open(path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fields) + writer.writeheader() + for rec in run.records: + writer.writerow({k: getattr(rec, k) for k in fields}) + return path + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--base-url", default="http://127.0.0.1:8068") + parser.add_argument("--api-key", default=os.environ.get("MILES_TINKER_API_KEY", "tml-miles-gpu-acceptance")) + parser.add_argument("--data", default="/root/datasets/gsm8k/train.parquet") + parser.add_argument("--tokenizer", default="/root/models/Qwen3-4B") + parser.add_argument("--out-dir", required=True) + parser.add_argument("--steps", type=int, default=50) + parser.add_argument("--prompts-per-step", type=int, default=8) + parser.add_argument("--samples-per-prompt", type=int, default=4) + parser.add_argument("--max-new-tokens", type=int, default=512) + parser.add_argument("--max-prompt-tokens", type=int, default=1024) + parser.add_argument( + "--enable-thinking", + action="store_true", + help="Qwen3 thinking mode: with a tight max_tokens budget the base policy mostly truncates " + "(low initial reward), which is the headroom the reward-growth check needs.", + ) + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + import pandas as pd # noqa: PLC0415 + from transformers import AutoTokenizer # noqa: PLC0415 + + from miles.rollout.rm_hub.math_utils import grade_answer_verl # noqa: PLC0415 + + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer) + df = pd.read_parquet(args.data) + + specs = DEFAULT_SPECS + shards: dict[int, list[dict]] = {} + for spec in specs: + rows = df.iloc[spec["shard"] :: len(specs)] + shard = [] + for _, row in rows.iterrows(): + messages = [dict(m) for m in row["messages"]] + encoded = tokenizer.apply_chat_template( + messages, tokenize=True, add_generation_prompt=True, enable_thinking=args.enable_thinking + ) + input_ids = encoded["input_ids"] if not isinstance(encoded, list) else encoded + if input_ids and isinstance(input_ids[0], list): + input_ids = input_ids[0] + shard.append(dict(input_ids=[int(t) for t in input_ids], label=str(row["label"]))) + shards[spec["shard"]] = shard + print(f"shard {spec['shard']}: {len(shard)} prompts", flush=True) + + operator = OperatorApi(args.base_url, args.api_key) + log_lock = threading.Lock() + + def log(msg: str) -> None: + with log_lock: + print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True) + + runs = [AdapterRun(spec=spec) for spec in specs] + threads = [] + for run in runs: + thread = threading.Thread( + target=_thread_main, + args=( + run, + args.base_url, + args.api_key, + operator, + shards[run.spec["shard"]], + tokenizer, + grade_answer_verl, + args, + log, + ), + name=run.spec["name"], + daemon=True, + ) + thread.start() + threads.append(thread) + for thread in threads: + thread.join() + + summary = {} + for run in runs: + rewards = [rec.reward_mean for rec in run.records] + first10 = rewards[:10] + last10 = rewards[-10:] + summary[run.spec["name"]] = dict( + spec={k: v for k, v in run.spec.items()}, + model_id=run.model_id, + registration=run.adapter_name, + steps_recorded=len(run.records), + step_clock=run.final_step_clock, + serving_version=run.final_serving_version, + reward_first10_mean=sum(first10) / len(first10) if first10 else None, + reward_last10_mean=sum(last10) / len(last10) if last10 else None, + reward_slope_per_step=least_squares_slope(rewards), + logprob_absdiff_mean=( + sum(r.logprob_absdiff_mean for r in run.records if r.logprob_absdiff_mean is not None) + / max(sum(1 for r in run.records if r.logprob_absdiff_mean is not None), 1) + ), + mean_step_dt_s=sum(r.dt_s for r in run.records) / max(len(run.records), 1), + failures=[f"step {r.step}: {r.note}" for r in run.records if r.note], + error=run.error, + csv=write_csv(run, args.out_dir), + ) + with open(os.path.join(args.out_dir, "summary.json"), "w") as f: + json.dump(summary, f, indent=2) + print(json.dumps(summary, indent=2), flush=True) + + grew = sum( + 1 + for s in summary.values() + if s["reward_first10_mean"] is not None and s["reward_last10_mean"] > s["reward_first10_mean"] + ) + print(f"\n=== RL QUALITY (SDK): reward grew (last10 > first10) on {grew}/{len(runs)} adapters ===", flush=True) + + +def _thread_main(run, base_url, api_key, operator, dataset, tokenizer, grade, args, log) -> None: + try: + adapter_loop(run, base_url, api_key, operator, dataset, tokenizer, grade, args, log) + except Exception as e: # noqa: BLE001 - a dead loop is a finding, not a harness crash + run.error = f"{type(e).__name__}: {e}" + log(f"({run.spec['name']}) LOOP ABORTED: {run.error}") + + +if __name__ == "__main__": + main() diff --git a/tests/fast-gpu/_layerwise_expert_dependency_worker.py b/tests/fast-gpu/_layerwise_expert_dependency_worker.py new file mode 100644 index 00000000000..4fed28fbe49 --- /dev/null +++ b/tests/fast-gpu/_layerwise_expert_dependency_worker.py @@ -0,0 +1,109 @@ +import os + +import pytest +import torch +import torch.distributed as dist + +from megatron.bridge.peft.utils import GroupedExpertLinearAdapter +from megatron.core import parallel_state +from megatron.core.model_parallel_config import ModelParallelConfig +from megatron.core.optimizer.layer_wise_optimizer import LayerWiseDistributedOptimizer +from megatron.core.optimizer.optimizer import FP32Optimizer +from megatron.core.optimizer.optimizer_config import OptimizerConfig +from megatron.core.process_groups_config import ProcessGroupCollection + + +def main() -> None: + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="nccl") + torch.manual_seed(1234) + torch.cuda.manual_seed(1234) + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=2, + expert_model_parallel_size=1, + expert_tensor_parallel_size=1, + ) + try: + config = ModelParallelConfig( + tensor_model_parallel_size=2, + expert_tensor_parallel_size=1, + params_dtype=torch.float32, + ) + adapter = GroupedExpertLinearAdapter( + in_features=4, + out_features=4, + dim=2, + num_local_experts=2, + base_linear_name="decoder.layers.0.mlp.experts.linear_fc2", + activation="identity", + input_is_parallel=True, + model_parallel_config=config, + params_device=torch.device("cuda", local_rank), + params_dtype=torch.float32, + ) + params = [adapter.linear_in.weight, adapter.linear_out.weight] + with torch.no_grad(): + for index, param in enumerate(params, start=1): + param.fill_(float(index)) + assert all(param.allreduce is False for param in params) + assert all(param.tensor_model_parallel is True for param in params) + + optimizer_config = OptimizerConfig( + optimizer="sgd", + lr=0.1, + min_lr=0.0, + weight_decay=0.0, + sgd_momentum=0.0, + clip_grad=1.0, + bf16=False, + use_distributed_optimizer=False, + params_dtype=torch.float32, + ) + base_optimizer = torch.optim.SGD( + [{"params": params, "is_expert_parallel": True}], + lr=optimizer_config.lr, + ) + optimizer = LayerWiseDistributedOptimizer( + [FP32Optimizer(base_optimizer, optimizer_config, None)], + optimizer_config, + ProcessGroupCollection.use_mpu_process_groups(["tp", "expt_tp", "dp_cp", "expt_dp"]), + ) + + assert optimizer.dp_cp_params_list is None + assert optimizer.expt_dp_params_list is not None + local_owners = torch.tensor( + len(optimizer.chained_optimizers[0].get_parameters()), + device="cuda", + dtype=torch.int64, + ) + dist.all_reduce(local_owners) + assert local_owners.item() == len(params) + + for param in params: + param.main_grad = torch.full_like(param, 3.0) + true_norm = (sum(param.numel() * 3.0**2 for param in params)) ** 0.5 + before = [param.detach().clone() for param in params] + + update_successful, grad_norm, _ = optimizer.step() + + assert update_successful + assert grad_norm == pytest.approx(true_norm, rel=1e-6, abs=1e-6) + clip_coefficient = 1.0 / (true_norm + 1.0e-6) + for previous, param in zip(before, params, strict=True): + torch.testing.assert_close( + param, + previous - optimizer_config.lr * 3.0 * clip_coefficient, + rtol=1e-6, + atol=1e-6, + ) + replicas = [torch.empty_like(param) for _ in range(dist.get_world_size())] + dist.all_gather(replicas, param) + torch.testing.assert_close(replicas[0], replicas[1], rtol=0, atol=0) + finally: + parallel_state.destroy_model_parallel() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/fast-gpu/test_layerwise_expert_dependencies.py b/tests/fast-gpu/test_layerwise_expert_dependencies.py new file mode 100644 index 00000000000..bbcb9996275 --- /dev/null +++ b/tests/fast-gpu/test_layerwise_expert_dependencies.py @@ -0,0 +1,31 @@ +from tests.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=90, suite="stage-b-2-gpu-h200", labels=["lora"]) + +import os +import subprocess +import sys +from pathlib import Path + + +def test_grouped_expert_lora_layerwise_norm_and_clip() -> None: + worker = Path(__file__).with_name("_layerwise_expert_dependency_worker.py") + repo_root = Path(__file__).parents[2] + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join(filter(None, [str(repo_root), env.get("PYTHONPATH")])) + result = subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nproc-per-node=2", + str(worker), + ], + env=env, + capture_output=True, + text=True, + timeout=180, + ) + + assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/fast/backends/megatron_utils/api_backends/__init__.py b/tests/fast/backends/megatron_utils/api_backends/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/fast/backends/megatron_utils/api_backends/full_parameter/test_executor.py b/tests/fast/backends/megatron_utils/api_backends/full_parameter/test_executor.py new file mode 100644 index 00000000000..6d7210a7291 --- /dev/null +++ b/tests/fast/backends/megatron_utils/api_backends/full_parameter/test_executor.py @@ -0,0 +1,315 @@ +from dataclasses import FrozenInstanceError +from types import SimpleNamespace + +import pytest +import torch +from miles.backends.megatron_utils.api_backends.full_parameter.executor import ( + FullParameterBinding, + FullParameterExecutor, +) +from miles.backends.training_utils.operation_execution import StepRequest, run_optim_controls +from miles.utils.operation_contract import BatchExecutionLease + + +class FakeModelChunk: + def __init__(self, *, zero_error: Exception | None = None, gradient: torch.Tensor | None = None): + self.zero_calls = 0 + self.zero_error = zero_error + self.parameter = SimpleNamespace(main_grad=gradient, grad=None, decoupled_grad=None) + + def zero_grad_buffer(self): + self.zero_calls += 1 + if self.zero_error is not None: + raise self.zero_error + + def parameters(self): + return [self.parameter] + + +class FakeOptimizer: + def __init__( + self, + *, + step_result=(True, 3.5, 0), + step_error: Exception | None = None, + optimizer_name: str = "adam", + ): + self.param_groups = [dict(lr=9.0, params=[]), dict(lr=8.0, params=[])] + self.config = SimpleNamespace(clip_grad=17.0, optimizer=optimizer_name) + self.step_result = step_result + self.step_error = step_error + self.step_calls = 0 + self.zero_calls = 0 + self.seen_groups = None + self.seen_clip = None + + def step(self): + self.step_calls += 1 + self.seen_groups = [dict(group) for group in self.param_groups] + self.seen_clip = self.config.clip_grad + if self.step_error is not None: + raise self.step_error + return self.step_result(self) if callable(self.step_result) else self.step_result + + def zero_grad(self): + self.zero_calls += 1 + + +TARGET = FullParameterBinding(target_id="actor") + + +def make_lease(operation_id="op", binding=TARGET, *, extras=()): + return BatchExecutionLease( + dispatch_id="dispatch", + bindings_by_operation=((operation_id, binding), *extras), + ) + + +def make_request(operation_id="op", **overrides): + adam = dict( + learning_rate=0.25, + beta1=0.7, + beta2=0.8, + eps=1e-7, + weight_decay=0.03, + grad_clip_norm=2.5, + ) + adam.update(overrides) + return StepRequest(operation_id=operation_id, adam_params=adam) + + +def make_executor(*, gradient: torch.Tensor | None = None, **optimizer_kwargs): + model = [FakeModelChunk(gradient=gradient), FakeModelChunk()] + optimizer = FakeOptimizer(**optimizer_kwargs) + return FullParameterExecutor(model_chunks=model, optimizer=optimizer, binding=TARGET), model, optimizer + + +def test_binding_and_executor_configuration_are_immutable(): + binding = FullParameterBinding(target_id="actor") + assert binding == TARGET + with pytest.raises(FrozenInstanceError): + binding.target_id = "slot-0" + + +def test_discard_clears_model_buffers_and_optimizer_gradients(): + executor, model, optimizer = make_executor() + + assert executor.discard_many(make_lease(), ["op"]) == {"op": {"ok": True, "gradient_window_consumed": True}} + assert [chunk.zero_calls for chunk in model] == [1, 1] + assert optimizer.zero_calls == 1 + assert optimizer.step_calls == 0 + + +@pytest.mark.parametrize( + ("lease", "operation_ids"), + [ + (make_lease("leased"), ["requested"]), + (make_lease(binding=FullParameterBinding(target_id="other")), ["op"]), + (make_lease(extras=(("other", TARGET),)), ["op"]), + (make_lease(), ["op", "other"]), + (make_lease(), ["op", "op"]), + ], +) +def test_invalid_or_non_singleton_discard_is_refused_before_mutation(lease, operation_ids): + executor, model, optimizer = make_executor() + + outcomes = executor.discard_many(lease, operation_ids) + + assert set(outcomes) == set(operation_ids) + assert all(outcome["ok"] is False for outcome in outcomes.values()) + assert all("gradient_window_consumed" not in outcome for outcome in outcomes.values()) + assert [chunk.zero_calls for chunk in model] == [0, 0] + assert optimizer.zero_calls == 0 + + +def test_step_applies_per_call_adam_uses_temporary_clip_and_clears_window(): + executor, model, optimizer = make_executor() + + outcome = executor.step_many(make_lease(), [make_request()])["op"] + + assert outcome == { + "ok": True, + "gradient_window_consumed": True, + "result": {"grad_norm": 3.5, "learning_rate": 0.25}, + } + assert optimizer.step_calls == 1 + assert optimizer.seen_clip == 2.5 + assert optimizer.config.clip_grad == 17.0 + for group in optimizer.seen_groups: + assert group["lr"] == 0.25 + assert group["betas"] == (0.7, 0.8) + assert group["eps"] == 1e-7 + assert group["weight_decay"] == 0.03 + assert [chunk.zero_calls for chunk in model] == [1, 1] + assert optimizer.zero_calls == 1 + + +def test_zero_clip_uses_infinite_stock_clip_to_measure_norm_without_scaling(): + def direct_optimizer_result(optimizer): + return (True, 4.25, 0) if optimizer.config.clip_grad == float("inf") else (True, None, 0) + + executor, _, optimizer = make_executor(step_result=direct_optimizer_result) + + outcome = executor.step_many(make_lease(), [make_request(grad_clip_norm=0.0)])["op"] + + assert outcome["ok"] is True + assert outcome["result"]["grad_norm"] == 4.25 + assert optimizer.seen_clip == float("inf") + assert optimizer.config.clip_grad == 17.0 + + +def test_success_without_stock_grad_norm_is_fail_stop(): + executor, model, optimizer = make_executor(step_result=(True, None, 0)) + + with pytest.raises(RuntimeError, match="did not report a gradient norm"): + executor.step_many(make_lease(), [make_request()]) + + assert optimizer.config.clip_grad == 17.0 + assert optimizer.zero_calls == 1 + assert [chunk.zero_calls for chunk in model] == [1, 1] + + +def test_generic_coordinator_refuses_poisoned_and_clean_shared_whole_lease_without_mutation(): + executor, model, optimizer = make_executor() + operations = [ + dict(kind="optim_step", operation_id="poisoned", poison="bad gradient window"), + dict(kind="optim_step", operation_id="clean", payload=dict(adam_params=dict(learning_rate=0.2))), + ] + lease = BatchExecutionLease( + dispatch_id="mixed", + bindings_by_operation=(("poisoned", TARGET), ("clean", TARGET)), + ) + + outcomes = run_optim_controls(operations, lease, executor) + + assert set(outcomes) == {"poisoned", "clean"} + assert all(outcome["ok"] is False for outcome in outcomes.values()) + assert all(outcome["category"] == "server" for outcome in outcomes.values()) + assert all("singleton whole-model lease" in outcome["error"] for outcome in outcomes.values()) + assert all("gradient_window_consumed" not in outcome for outcome in outcomes.values()) + assert optimizer.step_calls == 0 + assert optimizer.zero_calls == 0 + assert [chunk.zero_calls for chunk in model] == [0, 0] + + +def test_optimizer_veto_fails_closed_and_consumes_the_window(): + executor, model, optimizer = make_executor(step_result=(False, None, 0)) + + outcome = executor.step_many(make_lease(), [make_request()])["op"] + + assert outcome["ok"] is False + assert outcome["category"] == "server" + assert outcome["gradient_window_consumed"] is True + assert optimizer.config.clip_grad == 17.0 + assert [chunk.zero_calls for chunk in model] == [1, 1] + assert optimizer.zero_calls == 1 + + +@pytest.mark.parametrize("step_kwargs", [dict(step_error=RuntimeError("boom")), dict(step_result=(True, 1.0))]) +def test_step_fault_is_fail_stop_after_restoring_clip_and_clearing_window(step_kwargs): + executor, model, optimizer = make_executor(**step_kwargs) + + with pytest.raises(RuntimeError): + executor.step_many(make_lease(), [make_request()]) + + assert optimizer.config.clip_grad == 17.0 + assert [chunk.zero_calls for chunk in model] == [1, 1] + assert optimizer.zero_calls == 1 + + +def test_cleanup_failure_is_fail_stop(): + model = [FakeModelChunk(zero_error=RuntimeError("cannot clear")), FakeModelChunk()] + optimizer = FakeOptimizer() + executor = FullParameterExecutor(model_chunks=model, optimizer=optimizer, binding=TARGET) + + with pytest.raises(RuntimeError, match="cannot clear"): + executor.step_many(make_lease(), [make_request()]) + assert [chunk.zero_calls for chunk in model] == [1, 1] + assert optimizer.zero_calls == 1 + + +def test_discard_cleanup_failure_is_fail_stop(): + model = [FakeModelChunk(zero_error=RuntimeError("cannot discard"))] + optimizer = FakeOptimizer() + executor = FullParameterExecutor(model_chunks=model, optimizer=optimizer, binding=TARGET) + + with pytest.raises(RuntimeError, match="cannot discard"): + executor.discard_many(make_lease(), ["op"]) + assert optimizer.zero_calls == 1 + + +def test_nonfinite_gradient_vetoes_before_physical_step_and_clears_window(): + executor, model, optimizer = make_executor(gradient=torch.tensor([float("nan")])) + + outcome = executor.step_many(make_lease(), [make_request()])["op"] + + assert outcome == { + "ok": False, + "error": "non-finite gradient norm; step vetoed and gradients cleared", + "category": "server", + "gradient_window_consumed": True, + } + assert optimizer.step_calls == 0 + assert optimizer.config.clip_grad == 17.0 + assert optimizer.zero_calls == 1 + assert [chunk.zero_calls for chunk in model] == [1, 1] + + +def test_non_adam_optimizer_is_refused_before_mutation(): + executor, model, optimizer = make_executor(optimizer_name="sgd") + + outcome = executor.step_many(make_lease(), [make_request()])["op"] + + assert outcome["ok"] is False + assert "require an Adam optimizer" in outcome["error"] + assert optimizer.step_calls == 0 + assert optimizer.zero_calls == 0 + assert [chunk.zero_calls for chunk in model] == [0, 0] + + +def test_empty_model_is_refused_before_mutation(): + optimizer = FakeOptimizer() + executor = FullParameterExecutor(model_chunks=[], optimizer=optimizer, binding=TARGET) + + outcome = executor.step_many(make_lease(), [make_request()])["op"] + + assert outcome["ok"] is False + assert "at least one model chunk" in outcome["error"] + assert optimizer.step_calls == 0 + assert optimizer.zero_calls == 0 + + +def test_malformed_adam_is_refused_before_mutation(): + executor, model, optimizer = make_executor() + request = StepRequest(operation_id="op", adam_params=[("learning_rate", 0.1)]) + + outcome = executor.step_many(make_lease(), [request])["op"] + + assert outcome["ok"] is False + assert "invalid Adam parameters" in outcome["error"] + assert optimizer.step_calls == 0 + assert optimizer.zero_calls == 0 + assert [chunk.zero_calls for chunk in model] == [0, 0] + + +def test_non_singleton_step_refuses_every_operation_without_mutation(): + executor, model, optimizer = make_executor() + lease = make_lease(extras=(("other", TARGET),)) + + outcomes = executor.step_many(lease, [make_request(), make_request("other")]) + + assert set(outcomes) == {"op", "other"} + assert all(outcome["ok"] is False for outcome in outcomes.values()) + assert optimizer.step_calls == 0 + assert optimizer.zero_calls == 0 + assert [chunk.zero_calls for chunk in model] == [0, 0] + + +def test_empty_control_batch_is_a_noop(): + executor, model, optimizer = make_executor() + + assert executor.discard_many(make_lease(), []) == {} + assert executor.step_many(make_lease(), []) == {} + assert optimizer.step_calls == 0 + assert optimizer.zero_calls == 0 + assert [chunk.zero_calls for chunk in model] == [0, 0] diff --git a/tests/fast/backends/megatron_utils/api_backends/multi_lora/__init__.py b/tests/fast/backends/megatron_utils/api_backends/multi_lora/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/fast/backends/megatron_utils/api_backends/multi_lora/test_checkpoint.py b/tests/fast/backends/megatron_utils/api_backends/multi_lora/test_checkpoint.py new file mode 100644 index 00000000000..bc92140c4aa --- /dev/null +++ b/tests/fast/backends/megatron_utils/api_backends/multi_lora/test_checkpoint.py @@ -0,0 +1,166 @@ +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import torch + +import miles.backends.megatron_utils.api_backends.multi_lora.checkpoint as tc +from miles.backends.megatron_utils.api_backends.multi_lora.checkpoint import ( + FORMAT, + find_slot_state, + stable_slot_param_name, +) + + +class TestStableName: + def test_strips_exactly_the_target_slot(self): + name = "decoder.layers.0.self_attention.linear_qkv.adapters.3.linear_in.weight" + assert stable_slot_param_name(name, 3) == "decoder.layers.0.self_attention.linear_qkv.adapter.linear_in.weight" + assert stable_slot_param_name(name, 2) == name + assert ".adapter." in stable_slot_param_name("m.adapters.0.linear_out.weight", 0) + assert stable_slot_param_name("m.adapters.12.linear_in.weight", 12) == "m.adapter.linear_in.weight" + assert stable_slot_param_name("m.adapters.12.linear_in.weight", 1) == "m.adapters.12.linear_in.weight" + + +def make_adapter(tmp_path, name="a", rank=8, alpha=16): + config = SimpleNamespace(save=tmp_path, rank=rank, alpha=alpha) + return SimpleNamespace(name=name, registration_id="r1", slot=0, step=3, version=2, config=config) + + +def write_manifest(base, **overrides): + manifest = {"format": FORMAT, "name": "a", "rank_lora": 8, "alpha": 16, "optimizer_step": 3, "world_size": 1} + manifest.update(overrides) + base.mkdir(parents=True, exist_ok=True) + torch.save(manifest, base / "manifest.pt") + + +class TestManifestGating: + def test_missing_dir_or_manifest_means_no_state(self, tmp_path): + assert find_slot_state(SimpleNamespace(config=SimpleNamespace(save=None))) is None + adapter = make_adapter(tmp_path) + (tmp_path / "slot_state").mkdir() + assert find_slot_state(adapter) is None + + def test_foreign_name_is_loadable_but_foreign_shape_is_not(self, tmp_path): + adapter = make_adapter(tmp_path) + base = tmp_path / "slot_state" + write_manifest(base, name="someone-else") + assert find_slot_state(adapter) == base + + write_manifest(base, rank_lora=4) + assert find_slot_state(adapter) is None + + write_manifest(base, world_size=8) + assert find_slot_state(adapter) is None + + write_manifest(base, format="something-old") + assert find_slot_state(adapter) is None + + +class TestSlotStateRoundTrip: + class _FakeChild: + def __init__(self, slot: int, moment: float): + self.param_groups = [{"params": [0], "miles_multi_lora_slot": slot, "step": 0}] + self.moment = torch.full((2,), moment) + + def state_dict(self): + return { + "optimizer": { + "state": {0: {"exp_avg": self.moment.clone()}}, + "param_groups": [dict(group) for group in self.param_groups], + } + } + + def load_state_dict(self, state): + self.moment.copy_(state["optimizer"]["state"][0]["exp_avg"]) + for group, saved in zip(self.param_groups, state["optimizer"]["param_groups"], strict=True): + group.update({key: value for key, value in saved.items() if key != "params"}) + + def _round_trip(self, tmp_path, monkeypatch, target_children, ttl_seconds=None, after_save=None): + adapter = make_adapter(tmp_path) + adapter.step = 7 + + source = [self._FakeChild(slot=0, moment=1.5)] + source[0].param_groups[0]["step"] = 7 + children_by_slot = {0: source, 1: target_children} + monkeypatch.setattr(tc, "_slot_children", lambda optimizer, slot: children_by_slot[slot]) + monkeypatch.setattr( + tc, + "named_adapter_slot_parameters", + lambda model, slot: iter([("m.adapter.linear_in.weight", torch.ones(2))]), + ) + bridge = ModuleType("megatron.bridge.peft.multi_lora_layers") + loads: dict = {} + bridge.load_adapter = lambda model, slot, weights: loads.update(weights=weights) or len(weights) + bridge.init_adapter_slot = lambda model, slot, rank, alpha: loads.update(rank=rank, alpha=alpha) + monkeypatch.setitem(sys.modules, "megatron.bridge.peft.multi_lora_layers", bridge) + + tc.save_slot_state( + args=SimpleNamespace(), model=[], optimizer=None, adapter=adapter, reason="state", ttl_seconds=ttl_seconds + ) + if after_save is not None: + after_save() + adapter.slot = 1 + step = tc.load_slot_state(args=SimpleNamespace(), model=[], optimizer=None, adapter=adapter) + return step, loads, adapter + + def test_optimizer_state_restores_into_another_slot(self, tmp_path, monkeypatch): + target = [self._FakeChild(slot=1, moment=0.0)] + step, loads, _ = self._round_trip(tmp_path, monkeypatch, target) + assert step == 7 + assert loads["rank"] == 8 and loads["alpha"] == 16 + assert torch.equal(loads["weights"]["m.adapter.linear_in.weight"], torch.ones(2)) + assert torch.equal(target[0].moment, torch.full((2,), 1.5)) + group = target[0].param_groups[0] + assert group["step"] == 7 + assert group["miles_multi_lora_slot"] == 1 + + def test_child_count_mismatch_is_refused(self, tmp_path, monkeypatch): + two_children = [self._FakeChild(slot=1, moment=0.0), self._FakeChild(slot=1, moment=0.0)] + with pytest.raises(ValueError, match="refusing partial restore"): + self._round_trip(tmp_path, monkeypatch, two_children) + + def test_torn_save_is_refused(self, tmp_path, monkeypatch): + def cross_generation_manifest(): + manifest_path = tmp_path / "slot_state" / "manifest.pt" + manifest = torch.load(manifest_path, weights_only=True) + manifest["save_id"] = "another-generation" + torch.save(manifest, manifest_path) + + target = [self._FakeChild(slot=1, moment=0.0)] + with pytest.raises(ValueError, match="torn"): + self._round_trip(tmp_path, monkeypatch, target, after_save=cross_generation_manifest) + + def test_ownership_signature_mismatch_is_refused_before_mutation(self, tmp_path, monkeypatch): + adapter = make_adapter(tmp_path) + param_a, param_b = torch.zeros(1), torch.zeros(1) + + def child_with(param, slot): + child = self._FakeChild(slot=slot, moment=0.0) + child.param_groups[0]["params"] = [param] + return child + + children_by_slot = {0: [child_with(param_a, 0)], 1: [child_with(param_b, 1)]} + names_by_slot = { + 0: [("m.adapter.linear_in.weight", param_a)], + 1: [("m.adapter.linear_out.weight", param_b)], + } + monkeypatch.setattr(tc, "_slot_children", lambda optimizer, slot: children_by_slot[slot]) + monkeypatch.setattr(tc, "named_adapter_slot_parameters", lambda model, slot: iter(names_by_slot[slot])) + bridge = ModuleType("megatron.bridge.peft.multi_lora_layers") + loads: dict = {} + bridge.load_adapter = lambda model, slot, weights: loads.update(weights=weights) or len(weights) + bridge.init_adapter_slot = lambda model, slot, rank, alpha: loads.update(rank=rank, alpha=alpha) + monkeypatch.setitem(sys.modules, "megatron.bridge.peft.multi_lora_layers", bridge) + + tc.save_slot_state(args=SimpleNamespace(), model=[], optimizer=None, adapter=adapter, reason="state") + adapter.slot = 1 + with pytest.raises(ValueError, match="ownership"): + tc.load_slot_state(args=SimpleNamespace(), model=[], optimizer=None, adapter=adapter) + assert loads == {} + + def test_ttl_is_recorded_in_the_manifest(self, tmp_path, monkeypatch): + target = [self._FakeChild(slot=1, moment=0.0)] + self._round_trip(tmp_path, monkeypatch, target, ttl_seconds=3600) + manifest = torch.load(tmp_path / "slot_state" / "manifest.pt", weights_only=True) + assert manifest["ttl_seconds"] == 3600 diff --git a/tests/fast/backends/megatron_utils/api_backends/multi_lora/test_executor.py b/tests/fast/backends/megatron_utils/api_backends/multi_lora/test_executor.py new file mode 100644 index 00000000000..0d5b0fedd5f --- /dev/null +++ b/tests/fast/backends/megatron_utils/api_backends/multi_lora/test_executor.py @@ -0,0 +1,83 @@ +from types import SimpleNamespace + +import miles.backends.megatron_utils.api_backends.multi_lora.executor as executor_module +from miles.backends.megatron_utils.api_backends.multi_lora.executor import MultiLoraParameterExecutor +from miles.backends.training_utils.operation_execution import StepRequest +from miles.ray.multi_lora.residency import ResidentBinding +from miles.utils.operation_contract import BatchExecutionLease + + +def loaded(name="A", registration_id="r-A", slot=0): + return {name: SimpleNamespace(registration_id=registration_id, slot=slot)} + + +def make_executor(loaded_adapters=None): + return MultiLoraParameterExecutor(model=object(), optimizer=object(), loaded_adapters=loaded_adapters or loaded()) + + +def lease_of(*bindings): + return BatchExecutionLease(dispatch_id="d", bindings_by_operation=tuple(bindings)) + + +def binding(name="A", registration_id="r-A", slot=0): + return ResidentBinding(registration_key=(name, registration_id), training_slot=slot) + + +def step(op_id, lr=1e-4): + return StepRequest(operation_id=op_id, adam_params={"learning_rate": lr}) + + +class TestStepMany: + def test_step_and_veto_both_report_the_window_consumed(self, monkeypatch): + monkeypatch.setattr( + executor_module, "step_adapter_slots", lambda optimizer, model, adam: ({0: 1.5}, {1}, set()) + ) + executor = make_executor({**loaded("A", "r-A", 0), **loaded("B", "r-B", 1)}) + lease = lease_of(("op-A", binding("A", "r-A", 0)), ("op-B", binding("B", "r-B", 1))) + outcomes = executor.step_many(lease, [step("op-A"), step("op-B")]) + + assert outcomes["op-A"]["ok"] is True + assert outcomes["op-A"]["gradient_window_consumed"] is True + assert outcomes["op-A"]["result"]["grad_norm"] == 1.5 + assert outcomes["op-B"]["ok"] is False + assert outcomes["op-B"]["gradient_window_consumed"] is True + + def test_stale_binding_refusal_does_not_claim_consumption(self): + executor = make_executor() + lease = lease_of(("op-A", binding("A", "stale-registration", 0))) + outcomes = executor.step_many(lease, [step("op-A")]) + assert outcomes["op-A"]["ok"] is False and outcomes["op-A"]["category"] == "server" + assert not outcomes["op-A"].get("gradient_window_consumed") + + def test_duplicate_physical_step_targets_never_silently_drop_an_operation(self, monkeypatch): + stepped = [] + monkeypatch.setattr( + executor_module, + "step_adapter_slots", + lambda optimizer, model, adam: (stepped.append(dict(adam)) or ({s: 1.0 for s in adam}, set(), set())), + ) + executor = make_executor() + lease = lease_of(("op-1", binding("A", "r-A", 0)), ("op-2", binding("A", "r-A", 0))) + outcomes = executor.step_many(lease, [step("op-1", 1e-4), step("op-2", 2e-4)]) + + assert set(outcomes) == {"op-1", "op-2"} + for op_id in ("op-1", "op-2"): + assert outcomes[op_id]["ok"] is False and outcomes[op_id]["category"] == "server" + assert not outcomes[op_id].get("gradient_window_consumed") + assert stepped == [] + + +class TestDiscardMany: + def test_successful_discard_reports_the_window_consumed(self, monkeypatch): + cleared = [] + monkeypatch.setattr(executor_module, "zero_adapter_slot_grads", lambda model, slot: cleared.append(slot)) + executor = make_executor() + outcomes = executor.discard_many(lease_of(("op-A", binding())), ["op-A"]) + assert outcomes["op-A"] == dict(ok=True, gradient_window_consumed=True) + assert cleared == [0] + + def test_refused_discard_does_not_claim_consumption(self): + executor = make_executor() + outcomes = executor.discard_many(lease_of(("op-A", binding(slot=5))), ["op-A"]) + assert outcomes["op-A"]["ok"] is False + assert not outcomes["op-A"].get("gradient_window_consumed") diff --git a/tests/fast/backends/megatron_utils/api_backends/multi_lora/test_optimizer.py b/tests/fast/backends/megatron_utils/api_backends/multi_lora/test_optimizer.py new file mode 100644 index 00000000000..3ba5f0d6595 --- /dev/null +++ b/tests/fast/backends/megatron_utils/api_backends/multi_lora/test_optimizer.py @@ -0,0 +1,173 @@ +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import torch + +import miles.backends.megatron_utils.api_backends.multi_lora.optimizer as multi_lora_optimizer +from miles.backends.megatron_utils.api_backends.multi_lora.optimizer import ( + _found_inf_anywhere, + apply_adam_params_to_slot, + build_multi_lora_operation_optimizer, + step_adapter_slots, +) + + +class FakeChild: + def __init__(self, grads, found_inf=False): + self.params = [torch.nn.Parameter(torch.zeros(len(g))) for g in grads] + for param, grad in zip(self.params, grads, strict=True): + param.grad = torch.tensor(grad, dtype=torch.float32) + self.found_inf = found_inf + self.stepped = 0 + self.param_groups = [{"params": self.params, "lr": 0.0}] + + def prepare_grads(self): + return self.found_inf + + def get_parameters(self): + return self.params + + def get_main_grads_for_grad_norm(self): + return [p.grad for p in self.params] + + def step_with_ready_grads(self): + self.stepped += 1 + + +class FakeChained: + def __init__(self, children_by_slot): + self.chained_optimizers = [child for children in children_by_slot.values() for child in children] + self.miles_slot_child_indices, i = {}, 0 + for slot, children in children_by_slot.items(): + self.miles_slot_child_indices[slot] = list(range(i, i + len(children))) + i += len(children) + self.allgathered = 0 + + def allgather_params(self): + self.allgathered += 1 + + +@pytest.fixture() +def torch_clip_grads(monkeypatch): + fake = ModuleType("megatron.core.optimizer.clip_grads") + + def get_grad_norm_fp32(grads, grad_stats_parallel_group=None): + return torch.norm(torch.stack([torch.norm(g) for g in grads])).item() if grads else 0.0 + + def clip_grad_by_total_norm_fp32(params, max_norm, total_norm, _): + coeff = max_norm / (total_norm + 1e-6) + if coeff < 1.0: + for p in params: + p.grad.mul_(coeff) + + fake.get_grad_norm_fp32 = get_grad_norm_fp32 + fake.clip_grad_by_total_norm_fp32 = clip_grad_by_total_norm_fp32 + monkeypatch.setitem(sys.modules, "megatron.core.optimizer.clip_grads", fake) + return fake + + +@pytest.fixture() +def no_slot_traversal(monkeypatch): + monkeypatch.setattr(multi_lora_optimizer, "named_adapter_slot_parameters", lambda model, slot: iter(())) + + +class TestAdamParams: + def test_lands_on_every_group_of_the_slot_only(self): + mine, other = FakeChild([[1.0]]), FakeChild([[1.0]]) + chained = FakeChained({0: [mine], 1: [other]}) + apply_adam_params_to_slot(chained, 0, {"learning_rate": 5e-5, "beta1": 0.8, "weight_decay": 0.01}) + group = mine.param_groups[0] + assert group["lr"] == 5e-5 and group["betas"] == (0.8, 0.95) and group["weight_decay"] == 0.01 + assert other.param_groups[0]["lr"] == 0.0 + + +class TestStep: + def test_gradient_sum_is_never_count_normalized(self, torch_clip_grads, no_slot_traversal): + child = FakeChild([[3.0, 4.0]]) + chained = FakeChained({0: [child]}) + norms, vetoed, norm_blind = step_adapter_slots(chained, model=None, adam_params_by_slot={0: {}}) + assert vetoed == set() + assert norms[0] == pytest.approx(5.0) + assert child.stepped == 1 and chained.allgathered == 1 + + def test_per_call_clip_scales_the_update(self, torch_clip_grads, no_slot_traversal): + child = FakeChild([[3.0, 4.0]]) + chained = FakeChained({0: [child]}) + norms, _, _ = step_adapter_slots(chained, None, {0: {"grad_clip_norm": 1.0}}) + assert norms[0] == pytest.approx(5.0) + assert torch.allclose(child.params[0].grad, torch.tensor([0.6, 0.8]), atol=1e-4) + + def test_zero_clip_means_no_clip(self, torch_clip_grads, no_slot_traversal): + child = FakeChild([[30.0, 40.0]]) + chained = FakeChained({0: [child]}) + step_adapter_slots(chained, None, {0: {"grad_clip_norm": 0.0}}) + assert torch.allclose(child.params[0].grad, torch.tensor([30.0, 40.0])) + + def test_nonfinite_slot_is_vetoed_neighbours_step(self, torch_clip_grads, no_slot_traversal): + bad = FakeChild([[float("nan"), 1.0]]) + good = FakeChild([[1.0, 0.0]]) + chained = FakeChained({0: [bad], 1: [good]}) + norms, vetoed, _ = step_adapter_slots(chained, None, {0: {}, 1: {}}) + assert vetoed == {0} and bad.stepped == 0 + assert list(norms) == [1] and good.stepped == 1 + assert chained.allgathered == 1 + + def test_found_inf_from_prepare_grads_vetoes(self, torch_clip_grads, no_slot_traversal): + child = FakeChild([[1.0]], found_inf=True) + chained = FakeChained({0: [child]}) + norms, vetoed, _ = step_adapter_slots(chained, None, {0: {}}) + assert vetoed == {0} and norms == {} and child.stepped == 0 + assert chained.allgathered == 0 + + def test_untouched_slots_retain_grads(self, torch_clip_grads, no_slot_traversal): + stepped, retained = FakeChild([[1.0]]), FakeChild([[7.0]]) + chained = FakeChained({0: [stepped], 1: [retained]}) + step_adapter_slots(chained, None, {0: {}}) + assert retained.stepped == 0 + assert torch.allclose(retained.params[0].grad, torch.tensor([7.0])) + + def test_norm_blind_slot_is_refused_not_silently_stepped(self, torch_clip_grads, no_slot_traversal): + class NormBlindChild(FakeChild): + def get_main_grads_for_grad_norm(self): + return [] + + child = NormBlindChild([[3.0, 4.0]]) + chained = FakeChained({0: [child]}) + norms, vetoed, norm_blind = step_adapter_slots(chained, None, {0: {"grad_clip_norm": 1.0}}) + assert norm_blind == {0} and vetoed == set() and norms == {} + assert child.stepped == 0 + + def test_truly_zero_gradients_step_with_a_truthful_zero_norm(self, torch_clip_grads, no_slot_traversal): + class NormBlindChild(FakeChild): + def get_main_grads_for_grad_norm(self): + return [] + + child = NormBlindChild([[0.0, 0.0]]) + chained = FakeChained({0: [child]}) + norms, vetoed, norm_blind = step_adapter_slots(chained, None, {0: {}}) + assert norms == {0: 0.0} and vetoed == set() and norm_blind == set() + assert child.stepped == 1 + + +def test_found_inf_passthrough_without_dist(): + assert _found_inf_anywhere(True) is True + assert _found_inf_anywhere(False) is False + + +class TestBuildGuards: + def make(self, **overrides): + config = SimpleNamespace(use_distributed_optimizer=False, fp16=False, bf16=True, optimizer="adam") + config.__dict__.update(overrides) + args = SimpleNamespace(multi_lora_n_adapters=2, use_gloo_process_groups=False) + return args, config + + def test_rejects_distributed_optimizer_fp16_and_non_adam(self): + for overrides, message in [ + (dict(use_distributed_optimizer=True), "use_distributed_optimizer=False"), + (dict(fp16=True), "bf16"), + (dict(optimizer="sgd"), "Adam semantics"), + ]: + args, config = self.make(**overrides) + with pytest.raises(AssertionError, match=message): + build_multi_lora_operation_optimizer(args, config, model_chunks=[]) diff --git a/tests/fast/backends/megatron_utils/api_backends/multi_lora/test_trainer.py b/tests/fast/backends/megatron_utils/api_backends/multi_lora/test_trainer.py new file mode 100644 index 00000000000..825a9eaa2bb --- /dev/null +++ b/tests/fast/backends/megatron_utils/api_backends/multi_lora/test_trainer.py @@ -0,0 +1,261 @@ +from pathlib import Path +from types import SimpleNamespace + +import pytest + +import miles.backends.megatron_utils.api_backends.multi_lora.executor as executor_module +import miles.backends.megatron_utils.api_backends.multi_lora.trainer as trainer +from miles.ray.multi_lora.config import AdapterRun, AdapterRunConfig + + +def make_run(name="X", slot=0, step=3, save="/tmp/tinker-trainer-test"): + config = AdapterRunConfig(rank=8, alpha=16, save=Path(save) / name if save else None) + return AdapterRun(name=name, config=config, slot=slot, step=step, registration_id="reg1") + + +def control_op(kind, name="X", slot=0, op_id="op1", payload=None, step=3, serving_version=1): + return dict( + operation_id=op_id, + name=name, + kind=kind, + payload=payload, + step=step, + serving_version=serving_version, + _lease_slot=slot, + ) + + +@pytest.fixture() +def harness(monkeypatch): + calls = SimpleNamespace(step_args=None, saved=[], loaded=[], backups=0) + + def fake_step(optimizer, model, adam_params_by_slot): + calls.step_args = adam_params_by_slot + vetoed = {slot for slot, adam in adam_params_by_slot.items() if (adam or {}).get("veto")} + return {slot: 1.25 for slot in adam_params_by_slot if slot not in vetoed}, vetoed, set() + + monkeypatch.setattr(executor_module, "step_adapter_slots", fake_step) + monkeypatch.setattr(trainer, "save_slot_state", lambda *a, **k: calls.saved.append(k) or Path("/saved")) + monkeypatch.setattr(trainer, "load_slot_state", lambda *a, base=None, **k: 42 if "good" in str(base) else None) + + loaded = {"X": make_run(), "Y": make_run("Y", slot=1)} + pending: set = set() + backuper = SimpleNamespace(backup=lambda tag: setattr(calls, "backups", calls.backups + 1)) + + def run(operations): + lease = { + "dispatch_id": "lease-t", + "bindings_by_operation": [ + [op["operation_id"], [op["name"], "reg1", op.pop("_lease_slot", 0)]] for op in operations + ], + } + return trainer.execute_controls(SimpleNamespace(), None, None, loaded, pending, backuper, operations, lease) + + return SimpleNamespace(run=run, calls=calls, loaded=loaded, pending=pending) + + +class TestExecuteControls: + def test_optim_steps_apply_per_call_adam_and_report_norms(self, harness): + results = harness.run([control_op("optim_step", payload={"adam_params": {"learning_rate": 3e-4}})]) + assert harness.calls.step_args[0]["learning_rate"] == 3e-4 + assert harness.calls.step_args[0]["beta1"] == 0.9 + assert results["op1"] == dict( + ok=True, gradient_window_consumed=True, result=dict(grad_norm=1.25, learning_rate=3e-4) + ) + + def test_poisoned_optim_discards_the_window_and_never_steps(self, harness, monkeypatch): + zeroed = [] + monkeypatch.setattr(executor_module, "zero_adapter_slot_grads", lambda model, slot: zeroed.append(slot)) + poison = "a forward_backward in this gradient window failed; the window's gradients were discarded" + results = harness.run( + [ + {**control_op("optim_step", op_id="bad", payload={"adam_params": {}}), "poison": poison}, + control_op( + "optim_step", name="Y", op_id="good", slot=1, payload={"adam_params": {"learning_rate": 2e-4}} + ), + ] + ) + assert zeroed == [0] + assert set(harness.calls.step_args) == {1} + assert harness.calls.step_args[1]["learning_rate"] == 2e-4 + assert results["bad"] == dict(ok=False, error=poison, category="user", gradient_window_consumed=True) + assert results["good"]["ok"] is True + + def test_vetoed_slot_fails_as_server_error(self, harness): + results = harness.run([control_op("optim_step", payload={"adam_params": {"veto": True}})]) + assert results["op1"]["ok"] is False and results["op1"]["category"] == "server" + assert "vetoed" in results["op1"]["error"] + + def test_publish_stages_the_push_and_defers(self, harness): + results = harness.run([control_op("save_weights_for_sampler")]) + assert results["op1"] == dict(ok=True, deferred="publish") + assert harness.pending == {"X"} + + def test_non_resident_adapter_is_a_server_error(self, harness): + results = harness.run([control_op("save_state", name="ghost", slot=2)]) + assert results["op1"]["ok"] is False and "not resident" in results["op1"]["error"] + + def test_lease_binding_must_match_the_loaded_registration_and_slot(self, harness): + wrong_slot = harness.run([control_op("optim_step", slot=1)]) + assert wrong_slot["op1"]["ok"] is False and "not resident" in wrong_slot["op1"]["error"] + assert harness.calls.step_args is None + + def test_state_operation_validates_the_binding_name_before_mutation(self): + from miles.ray.multi_lora.residency import ResidentBinding + from miles.utils.operation_contract import BatchExecutionLease + + lease = BatchExecutionLease( + dispatch_id="lease-t", + bindings_by_operation=(("op1", ResidentBinding(("B", "reg1"), 0)),), + ) + pending: set = set() + result = trainer._execute_state_op( + dict(operation_id="op1", name="A", kind="save_weights_for_sampler"), + lease, + None, + None, + None, + {"A": make_run("A")}, + pending, + ) + assert result["ok"] is False and result["category"] == "server" + assert pending == set() + + def test_operation_missing_from_the_lease_is_refused(self, harness): + op = control_op("optim_step") + op.pop("_lease_slot") + lease = {"dispatch_id": "lease-t", "bindings_by_operation": []} + results = trainer.execute_controls( + SimpleNamespace(), + None, + None, + harness.loaded, + harness.pending, + SimpleNamespace(backup=lambda t: None), + [op], + lease, + ) + assert results["op1"]["ok"] is False and "no binding in the batch lease" in results["op1"]["error"] + + def test_save_state_validates_tag_and_immutability(self, harness, tmp_path, monkeypatch): + results = harness.run([control_op("save_state", payload={"tag": "../evil"})]) + assert "invalid state tag" in results["op1"]["error"] and results["op1"]["category"] == "user" + + harness.loaded["X"] = make_run(save=None) + results = harness.run([control_op("save_state", payload={"tag": "t0"})]) + assert "no save dir" in results["op1"]["error"] + + harness.loaded["X"] = make_run(save=tmp_path) + existing = tmp_path / "X" / "states" / "t0" + existing.mkdir(parents=True) + (existing / "manifest.pt").touch() + results = harness.run([control_op("save_state", payload={"tag": "t0"})]) + assert "immutable" in results["op1"]["error"] + + results = harness.run([control_op("save_state", payload={"tag": "t1"})]) + assert results["op1"] == dict(ok=True, result=dict(path=str(tmp_path / "X" / "states" / "t1"), step=3)) + assert harness.calls.saved[0]["reason"] == "state:t1" + + def test_load_state_restores_step_and_stages_republish(self, harness): + results = harness.run([control_op("load_state", payload={"path": "/good/state"})]) + assert results["op1"] == dict(ok=True, deferred="publish", result=dict(step=42, path="/good/state")) + assert harness.pending == {"X"} + assert harness.calls.backups == 1 + + results = harness.run([control_op("load_state", op_id="op2", payload={"path": "/missing"})]) + assert results["op2"]["ok"] is False and results["op2"]["category"] == "user" + + def test_unknown_kind_fails_every_leftover(self, harness): + results = harness.run([control_op("compile_model")]) + assert results["op1"]["ok"] is False and "no executor" in results["op1"]["error"] + + +class TestLoadAdapters: + def test_master_reload_skips_restored_slots(self, monkeypatch): + import sys + from types import ModuleType + + restored = {"fresh": None, "resumed": 9, "resumed-at-zero": 0} + inits: list = [] + reloaded: list = [] + bridge = ModuleType("megatron.bridge.peft.multi_lora_layers") + bridge.init_adapter_slot = lambda model, slot, rank, alpha: inits.append(slot) + monkeypatch.setitem(sys.modules, "megatron.bridge.peft.multi_lora_layers", bridge) + monkeypatch.setattr(trainer, "load_slot_state", lambda args, model, optimizer, adapter: restored[adapter.name]) + monkeypatch.setattr(trainer, "reload_adapter_slot_model_params", lambda optimizer, slot: reloaded.append(slot)) + import miles.backends.megatron_utils.initialize as megatron_initialize + + monkeypatch.setattr(megatron_initialize, "is_first_replica_megatron_main_rank", lambda: False) + + adapters = [make_run("fresh", slot=0), make_run("resumed", slot=1), make_run("resumed-at-zero", slot=2)] + assert trainer.load_adapters(SimpleNamespace(), None, None, adapters) == 3 + assert inits == [0] + assert reloaded == [0] + + +class TestGatherAndCommit: + def test_gather_groups_rows_per_operation_in_order(self): + rollout_data = { + "tinker_logprob_collector": {(0, 1): [-2.0], (0, 0): [-1.0], (1, 0): [-9.0], (0, -1): [-7.0]}, + "operation_by_lane": {0: "fb1", 1: "fb2", 2: None}, + } + assert trainer._gather_logprobs(rollout_data) == {"fb1": [[-1.0], [-2.0]], "fb2": [[-9.0]]} + + def test_commit_pins_accumulators_and_completes_ops(self, monkeypatch): + committed = {} + + class FakeController: + class commit_tinker_batch: # noqa: N801 - mimics the .remote handle + @staticmethod + def remote(accumulated, operation_ids, logprobs_by_op): + committed.update( + accumulated=accumulated, operation_ids=operation_ids, logprobs_by_op=logprobs_by_op + ) + + monkeypatch.setattr(trainer, "get_multi_lora_controller", lambda: FakeController) + monkeypatch.setattr(trainer.ray, "get", lambda ref: ref) + import miles.backends.megatron_utils.initialize as megatron_initialize + + monkeypatch.setattr(megatron_initialize, "is_first_replica_megatron_main_rank", lambda: True) + + rollout_data = { + "registration_by_lane": {0: ("A", "r-A"), 1: ("B", "r-B")}, + "operation_by_lane": {0: "fb1", 1: None}, + "tinker_logprob_collector": {(0, 0): [-1.0]}, + } + trainer.commit_batch(rollout_data, pending_push=set()) + assert committed["accumulated"] == [("A", "r-A"), ("B", "r-B")] + assert committed["operation_ids"] == ["fb1"] + assert committed["logprobs_by_op"] == {"fb1": [[-1.0]]} + + committed.clear() + trainer.commit_batch({**rollout_data, "tinker_forward_only": True}, pending_push=set()) + assert committed["accumulated"] == [] + + +class TestPushPlumbing: + def test_select_pushes_only_staged_unless_new_engines(self): + loaded = {"A": make_run("A"), "B": make_run("B", slot=1)} + pushes, bumps = trainer.select_adapters_to_push(loaded, {"B", "gone"}, has_new_engines=False) + assert list(pushes) == ["B"] and bumps == ["B"] + + pushes, bumps = trainer.select_adapters_to_push(loaded, {"B"}, has_new_engines=True) + assert list(pushes) == ["A", "B"] + assert bumps == ["B"] + + def test_commit_weight_push_only_on_main_rank(self, monkeypatch): + recorded = [] + + class FakeController: + class record_weight_update: # noqa: N801 + @staticmethod + def remote(names): + recorded.append(names) + + monkeypatch.setattr(trainer, "get_multi_lora_controller", lambda: FakeController) + monkeypatch.setattr(trainer.ray, "get", lambda ref: ref) + trainer.commit_weight_push(["A"], is_main_rank=False) + trainer.commit_weight_push([], is_main_rank=True) + assert recorded == [] + trainer.commit_weight_push(["A"], is_main_rank=True) + assert recorded == [["A"]] diff --git a/tests/fast/backends/megatron_utils/test_lora_model_branches.py b/tests/fast/backends/megatron_utils/test_lora_model_branches.py index 6453e39225e..5559a406d3c 100644 --- a/tests/fast/backends/megatron_utils/test_lora_model_branches.py +++ b/tests/fast/backends/megatron_utils/test_lora_model_branches.py @@ -163,6 +163,34 @@ def test_lora_raw_mode_skips_bridge(self, mock_lora_setup, mock_get_model, mock_ mock_lora_setup.assert_not_called() mock_get_model.assert_called_once() + @patch(f"{_MODEL_MODULE}.get_optimizer_param_scheduler") + @patch("miles.backends.megatron_utils.api_backends.multi_lora.optimizer.build_multi_lora_operation_optimizer") + @patch(f"{_MODEL_MODULE}.get_megatron_optimizer") + @patch(f"{_MODEL_MODULE}._setup_lora_model_via_bridge") + def test_multi_lora_operations_route_to_canonical_optimizer_builder( + self, mock_lora_setup, mock_megatron_opt, mock_operation_opt, mock_sched + ): + from miles.backends.megatron_utils.model import setup_model_and_optimizer + + model = [MagicMock()] + optimizer = MagicMock() + mock_lora_setup.return_value = model + mock_operation_opt.return_value = optimizer + mock_sched.return_value = MagicMock() + + args = self._make_args(lora_rank=32, role="actor", mode="bridge") + args.multi_lora = True + args.multi_lora_n_adapters = 2 + args.tinker_backend = True + + _, actual_optimizer, _ = setup_model_and_optimizer(args, role="actor") + + mock_operation_opt.assert_called_once() + assert mock_operation_opt.call_args.args[0] is args + assert mock_operation_opt.call_args.args[2] is model + assert actual_optimizer is optimizer + mock_megatron_opt.assert_not_called() + # --------------------------------------------------------------------------- # save — LoRA vs regular branch diff --git a/tests/fast/backends/megatron_utils/test_lora_weight_sync_validation.py b/tests/fast/backends/megatron_utils/test_lora_weight_sync_validation.py index 1c59f200917..0dfe8224a0d 100644 --- a/tests/fast/backends/megatron_utils/test_lora_weight_sync_validation.py +++ b/tests/fast/backends/megatron_utils/test_lora_weight_sync_validation.py @@ -455,6 +455,39 @@ def test_lora_loaded_stays_false_when_implementation_raises(self): assert fake_self._lora_loaded is False +@pytest.mark.parametrize("is_lora", [False, True]) +def test_distributed_lora_skips_base_weight_update_session(is_lora): + engine = MagicMock() + updater = SimpleNamespace( + args=SimpleNamespace(pause_generation_mode="retract"), + rollout_engines=[engine], + weight_version=7, + is_lora=is_lora, + ) + + with ( + patch(f"{_MIXIN_MODULE}.dist") as dist_mock, + patch(f"{_MIXIN_MODULE}.ray") as ray_mock, + patch(f"{_MIXIN_MODULE}.begin_weight_update") as begin_mock, + patch(f"{_MIXIN_MODULE}.end_weight_update") as end_mock, + ): + dist_mock.get_rank.return_value = 0 + ray_mock.get.side_effect = lambda refs: refs + DistBucketedWeightUpdateMixin._pause_and_prepare_engines(updater) + DistBucketedWeightUpdateMixin._finalize_and_resume_engines(updater) + + if is_lora: + begin_mock.assert_not_called() + end_mock.assert_not_called() + else: + begin_mock.assert_called_once_with([engine], "all") + end_mock.assert_called_once_with([engine]) + engine.pause_generation.remote.assert_called_once_with(mode="retract") + engine.flush_cache.remote.assert_called_once_with() + engine.update_weight_version.remote.assert_called_once_with(weight_version="7") + engine.continue_generation.remote.assert_called_once_with() + + class TestBroadcastLoraImplementation: """Broadcast transport ``UpdateWeightFromDistributed._update_lora_weight_implementation``: send metadata over Ray, then ``dist.broadcast`` each adapter tensor over the diff --git a/tests/fast/backends/megatron_utils/test_model_initialize.py b/tests/fast/backends/megatron_utils/test_model_initialize.py index d3e4ee9f113..b461cc6d2e7 100644 --- a/tests/fast/backends/megatron_utils/test_model_initialize.py +++ b/tests/fast/backends/megatron_utils/test_model_initialize.py @@ -133,7 +133,9 @@ def _mock_megatron_environment(): _stub_module("miles.backends.megatron_utils.model_provider", {"get_model_provider_func": MagicMock()}) yield finally: - sys.modules.clear() + for name in [n for n in sys.modules if n not in original_modules]: + if name.split(".")[0] in ("miles", "megatron", "sglang"): + del sys.modules[name] sys.modules.update(original_modules) diff --git a/tests/fast/backends/megatron_utils/test_multi_lora_checkpoint_naming.py b/tests/fast/backends/megatron_utils/test_multi_lora_checkpoint_naming.py deleted file mode 100644 index 97ab9f3c5ec..00000000000 --- a/tests/fast/backends/megatron_utils/test_multi_lora_checkpoint_naming.py +++ /dev/null @@ -1,46 +0,0 @@ -"""Adapter shards are keyed by (tp, pp, ep): EP ranks hold different local experts, and -the realized coordinates are not the tp x pp x ep cross product when ETP < TP.""" - -from miles.backends.megatron_utils.multi_lora_utils import all_megatron_checkpoints_exist, megatron_shard_name - - -def _names(coords, ep_size): - return {megatron_shard_name(*coord, ep_size) for coord in coords} - - -def test_shard_name_omits_ep_suffix_without_expert_parallelism(): - # Checkpoints written before expert adapters existed must stay loadable. - assert megatron_shard_name(0, 0, 0, ep_size=1) == "adapter_megatron_tp0_pp0.pt" - assert megatron_shard_name(1, 2, 0, ep_size=1) == "adapter_megatron_tp1_pp2.pt" - - -def test_shard_name_is_unique_per_expert_parallel_rank(): - names = {megatron_shard_name(0, 0, ep, ep_size=4) for ep in range(4)} - assert len(names) == 4 - assert megatron_shard_name(0, 0, 2, ep_size=4) == "adapter_megatron_tp0_pp0_ep2.pt" - - -def test_completeness_check_requires_every_realized_shard(tmp_path): - coords = [(0, 0, 0), (0, 0, 1), (0, 0, 2)] - for coord in coords[:2]: - (tmp_path / megatron_shard_name(*coord, 3)).touch() - - assert not all_megatron_checkpoints_exist(tmp_path, _names(coords, 3)) - - (tmp_path / megatron_shard_name(*coords[2], 3)).touch() - assert all_megatron_checkpoints_exist(tmp_path, _names(coords, 3)) - - -def test_completeness_ignores_unrealized_coordinates(tmp_path): - # TP=2, EP=2, ETP=1: only (0,0,0) and (1,0,1) exist; a cross-product check - # would demand four shards and never resume. - coords = [(0, 0, 0), (1, 0, 1)] - for coord in coords: - (tmp_path / megatron_shard_name(*coord, 2)).touch() - - assert all_megatron_checkpoints_exist(tmp_path, _names(coords, 2)) - - -def test_completeness_check_with_single_shard(tmp_path): - (tmp_path / "adapter_megatron_tp0_pp0.pt").touch() - assert all_megatron_checkpoints_exist(tmp_path, _names([(0, 0, 0)], 1)) diff --git a/tests/fast/backends/megatron_utils/test_multi_lora_scheduler.py b/tests/fast/backends/megatron_utils/test_multi_lora_scheduler.py deleted file mode 100644 index be48d7118e8..00000000000 --- a/tests/fast/backends/megatron_utils/test_multi_lora_scheduler.py +++ /dev/null @@ -1,114 +0,0 @@ -"""Per-adapter LR schedules: parameters come from the global args, position is per adapter. -Pins two fixes: late loads don't inherit the decayed position; resume rebuilds position from committed steps.""" - -from types import SimpleNamespace - -import pytest - -from miles.backends.megatron_utils.multi_lora_scheduler import install_slot_scheduler, step_slot_schedulers - -LR = 2e-5 - - -def make_args(**overrides) -> SimpleNamespace: - args = SimpleNamespace( - lr=LR, - min_lr=0.0, - lr_warmup_init=0.0, - lr_warmup_fraction=None, - lr_warmup_iters=0, - lr_decay_style="cosine", - start_weight_decay=0.1, - end_weight_decay=0.1, - weight_decay_incr_style="constant", - lr_wsd_decay_iters=None, - lr_wsd_decay_style=None, - ) - for key, value in overrides.items(): - setattr(args, key, value) - return args - - -def make_optimizer(n_slots: int = 2) -> SimpleNamespace: - children = [SimpleNamespace(param_groups=[{"lr": 0.0, "weight_decay": 0.0}]) for _ in range(n_slots)] - return SimpleNamespace( - chained_optimizers=children, - miles_slot_child_indices={slot: [slot] for slot in range(n_slots)}, - ) - - -def make_adapter(slot: int, num_step: int | None, samples_per_step: int = 64) -> SimpleNamespace: - config = SimpleNamespace(num_step=num_step, adapter_global_batch_size=samples_per_step) - return SimpleNamespace(slot=slot, name=f"a{slot}", config=config) - - -def slot_lr(optimizer, slot: int) -> float: - return optimizer.chained_optimizers[slot].param_groups[0]["lr"] - - -def test_decaying_adapter_walks_its_own_cosine_schedule(): - optimizer = make_optimizer() - adapter = make_adapter(slot=0, num_step=10) - install_slot_scheduler(make_args(), optimizer, adapter, resume_step=0) - - assert slot_lr(optimizer, 0) == pytest.approx(LR) # fresh: top of the schedule - - step_slot_schedulers(optimizer, {0: 5 * 64}) # half the horizon - assert slot_lr(optimizer, 0) == pytest.approx(LR / 2) - - step_slot_schedulers(optimizer, {0: 100 * 64}) # far past the horizon - assert slot_lr(optimizer, 0) == pytest.approx(0.0) # clamped at min_lr - - -def test_adapter_without_num_step_holds_constant(): - optimizer = make_optimizer() - install_slot_scheduler(make_args(), optimizer, make_adapter(slot=0, num_step=None), resume_step=0) - - step_slot_schedulers(optimizer, {0: 12345 * 64}) - assert slot_lr(optimizer, 0) == pytest.approx(LR) # no horizon: never decays - - -def test_resume_position_is_deterministic_from_committed_steps(): - stepped = make_optimizer() - install_slot_scheduler(make_args(), stepped, make_adapter(slot=0, num_step=10), resume_step=0) - step_slot_schedulers(stepped, {0: 5 * 64}) - - resumed = make_optimizer() - install_slot_scheduler(make_args(), resumed, make_adapter(slot=0, num_step=10), resume_step=5) - - assert slot_lr(resumed, 0) == pytest.approx(slot_lr(stepped, 0)) - - -def test_only_stepped_slots_advance(): - optimizer = make_optimizer() - args = make_args() - install_slot_scheduler(args, optimizer, make_adapter(slot=0, num_step=10), resume_step=0) - install_slot_scheduler(args, optimizer, make_adapter(slot=1, num_step=10), resume_step=0) - - lr_by_slot = step_slot_schedulers(optimizer, {0: 5 * 64}) - - assert set(lr_by_slot) == {0} - assert slot_lr(optimizer, 0) == pytest.approx(LR / 2) - assert slot_lr(optimizer, 1) == pytest.approx(LR) # co-tenant untouched - - -def test_slot_reuse_installs_a_fresh_schedule(): - optimizer = make_optimizer() - args = make_args() - install_slot_scheduler(args, optimizer, make_adapter(slot=0, num_step=10), resume_step=0) - step_slot_schedulers(optimizer, {0: 5 * 64}) - - install_slot_scheduler(args, optimizer, make_adapter(slot=0, num_step=20), resume_step=0) - assert slot_lr(optimizer, 0) == pytest.approx(LR) # next tenant starts at the top - - -def test_warmup_ramps_from_init_lr(): - optimizer = make_optimizer() - args = make_args(lr_warmup_iters=2) # 2 adapter steps of warmup - install_slot_scheduler(args, optimizer, make_adapter(slot=0, num_step=10), resume_step=0) - - assert slot_lr(optimizer, 0) == pytest.approx(0.0) # init_lr - step_slot_schedulers(optimizer, {0: 64}) - assert slot_lr(optimizer, 0) == pytest.approx(LR / 2) # mid-warmup - step_slot_schedulers(optimizer, {0: 64}) - assert slot_lr(optimizer, 0) == pytest.approx(LR) # warmed up diff --git a/tests/fast/backends/megatron_utils/test_multi_lora_slot_cleanup.py b/tests/fast/backends/megatron_utils/test_multi_lora_slot_cleanup.py deleted file mode 100644 index 93fcd3bd933..00000000000 --- a/tests/fast/backends/megatron_utils/test_multi_lora_slot_cleanup.py +++ /dev/null @@ -1,91 +0,0 @@ -"""zero_optimizer_state_for_adapter must reset a retired slot's Adam moments and step clock -(group-level FusedAdam or per-param torch AdamW) while leaving co-tenant slots untouched.""" - -import sys -import types -from types import SimpleNamespace - -import pytest -import torch - -from miles.backends.megatron_utils.multi_lora_utils import zero_optimizer_state_for_adapter - -MLL_MODULE = "megatron.bridge.peft.multi_lora_layers" - - -class FakeAdapter: - def __init__(self, params): - self._params = list(params) - - def parameters(self): - return self._params - - -class FakeMultiLoRALinear: - def __init__(self, adapters): - self.adapters = adapters - - -@pytest.fixture() -def rig(monkeypatch): - # Stub the lazily imported bridge module so the test needs no bridge build that ships multi-LoRA. - p0 = torch.nn.Parameter(torch.ones(4)) - p1 = torch.nn.Parameter(torch.ones(4)) - module = FakeMultiLoRALinear({0: FakeAdapter([p0]), 1: FakeAdapter([p1])}) - stub = types.ModuleType(MLL_MODULE) - stub.MultiLoRALinear = FakeMultiLoRALinear - stub._iter_multi_lora_modules = lambda model: [module] - monkeypatch.setitem(sys.modules, MLL_MODULE, stub) - return SimpleNamespace(p0=p0, p1=p1, model=object()) - - -def make_optimizer(groups, state): - inner = SimpleNamespace(param_groups=groups, state=state) - return inner, SimpleNamespace(chained_optimizers=[SimpleNamespace(optimizer=inner)]) - - -def test_group_level_fused_adam_clock_resets_only_for_the_retired_slot(rig): - inner, optimizer = make_optimizer( - groups=[ - {"params": [rig.p0], "miles_multi_lora_slot": 0, "step": 50}, - {"params": [rig.p1], "miles_multi_lora_slot": 1, "step": 50}, - ], - state={ - rig.p0: {"exp_avg": torch.ones(4), "exp_avg_sq": torch.ones(4)}, - rig.p1: {"exp_avg": torch.ones(4), "exp_avg_sq": torch.ones(4)}, - }, - ) - - zero_optimizer_state_for_adapter(optimizer, rig.model, 0) - - assert inner.param_groups[0]["step"] == 0 - assert inner.param_groups[1]["step"] == 50 # co-tenant slot untouched - assert float(inner.state[rig.p0]["exp_avg"].abs().sum()) == 0.0 - assert float(inner.state[rig.p0]["exp_avg_sq"].abs().sum()) == 0.0 - assert float(inner.state[rig.p1]["exp_avg"].abs().sum()) == 4.0 - - -def test_tensor_valued_group_clock_resets_in_place(rig): - step = torch.tensor(50) - inner, optimizer = make_optimizer( - groups=[{"params": [rig.p0], "miles_multi_lora_slot": 0, "step": step}], - state={rig.p0: {"exp_avg": torch.ones(4), "exp_avg_sq": torch.ones(4)}}, - ) - - zero_optimizer_state_for_adapter(optimizer, rig.model, 0) - - assert int(step) == 0 # zeroed in place, no rebinding needed - - -def test_per_param_adamw_fallback_clock_resets(rig): - # torch.optim.AdamW keeps the clock per param; groups carry no "step". - inner, optimizer = make_optimizer( - groups=[{"params": [rig.p0], "miles_multi_lora_slot": 0}], - state={ - rig.p0: {"exp_avg": torch.ones(4), "exp_avg_sq": torch.ones(4), "step": torch.tensor(50.0)}, - }, - ) - - zero_optimizer_state_for_adapter(optimizer, rig.model, 0) - - assert float(inner.state[rig.p0]["step"]) == 0.0 diff --git a/tests/fast/backends/megatron_utils/test_shared_ppo_lifecycle.py b/tests/fast/backends/megatron_utils/test_shared_ppo_lifecycle.py index 6823755e5de..9402ac62fa0 100644 --- a/tests/fast/backends/megatron_utils/test_shared_ppo_lifecycle.py +++ b/tests/fast/backends/megatron_utils/test_shared_ppo_lifecycle.py @@ -146,7 +146,6 @@ def test_save_model_does_not_manage_lifecycle(actor_module, monkeypatch): reload_groups = Mock() destroy_groups = Mock() monkeypatch.setattr(actor_module, "save", save) - monkeypatch.setattr(actor_module, "is_multi_lora_enabled", lambda _args: False) monkeypatch.setattr(actor_module, "reload_process_groups", reload_groups) monkeypatch.setattr(actor_module, "destroy_process_groups", destroy_groups) @@ -329,6 +328,7 @@ def test_actor_logprob_forward_is_explicit_single_step_opt_in( "witness_info": None, "attempt": 0, "ft_test_action_executor": None, + "forward_only": False, } diff --git a/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py b/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py index 2df14f54b1f..ce0864e5889 100644 --- a/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py +++ b/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py @@ -4,7 +4,7 @@ import pytest import torch -from miles.backends.megatron_utils.multi_lora_utils import slice_lora_to_rank +from miles.backends.megatron_utils.api_backends.multi_lora.model import slice_lora_to_rank def _padded(shape, live_rows=None, live_cols=None): diff --git a/tests/fast/backends/sglang_utils/test_sglang_engine.py b/tests/fast/backends/sglang_utils/test_sglang_engine.py index a5b6c138e90..177d19f9cc9 100644 --- a/tests/fast/backends/sglang_utils/test_sglang_engine.py +++ b/tests/fast/backends/sglang_utils/test_sglang_engine.py @@ -1,4 +1,5 @@ import time +from types import SimpleNamespace import pytest import requests @@ -30,3 +31,26 @@ def test_flush_cache_sleeps_between_pending_request_retries(monkeypatch): f"expected the loop to back off on every one of its 60 attempts, got {len(sleep_calls)} sleeps " "-- a 400 response (pending requests) must not skip the retry delay" ) + + +def test_update_weight_version_does_not_abort_in_flight_requests(monkeypatch): + pytest.importorskip("sglang") + from miles.backends.sglang_utils.sglang_engine import SGLangEngine + + engine = SGLangEngine.__new__(SGLangEngine) + engine.node_rank = 0 + engine.server_host = "fake-host" + engine.server_port = 1234 + posts = [] + + def fake_post(url, json=None): + posts.append((url, json)) + return SimpleNamespace(raise_for_status=lambda: None, json=lambda: {}) + + monkeypatch.setattr(requests, "post", fake_post) + + engine.update_weight_version("3") + + assert posts == [ + ("http://fake-host:1234/update_weight_version", {"new_version": "3", "abort_all_requests": False}) + ] diff --git a/tests/fast/backends/training_utils/loss/test_tinker_loss.py b/tests/fast/backends/training_utils/loss/test_tinker_loss.py new file mode 100644 index 00000000000..a23fcbcb9a7 --- /dev/null +++ b/tests/fast/backends/training_utils/loss/test_tinker_loss.py @@ -0,0 +1,190 @@ +import pytest +import torch + +from miles.backends.training_utils.loss_hub.logit_processors import get_log_probs_and_entropy +from miles.backends.training_utils.loss_hub.losses import tinker_loss_function + +from .loss_test_utils import make_args, make_inputs, make_parallel_state + +VOCAB = 32 + + +def make_batch(seed=7, prompt_lens=(4, 6), response_lens=(3, 5)): + make_parallel_state() + args = make_args(loss_type="custom_loss") + inputs = make_inputs( + seed=seed, + batch_size=len(prompt_lens), + prompt_lens=list(prompt_lens), + response_lens=list(response_lens), + vocab_size=VOCAB, + args=args, + ) + batch = dict( + unconcat_tokens=inputs["unconcat_tokens"], + total_lengths=inputs["total_lens"], + response_lengths=list(response_lens), + loss_masks=[torch.ones(rl, dtype=torch.int32) for rl in response_lens], + rollout_log_probs=inputs["rollout_log_probs"], + tinker_operation_lanes=[0] * len(prompt_lens), + tinker_loss_by_lane={0: {"loss_fn": "cross_entropy"}}, + ) + return args, batch, inputs["policy_logits"].requires_grad_(True) + + +def reference_log_probs(args, batch, logits): + return get_log_probs_and_entropy( + logits, + args=args, + unconcat_tokens=batch["unconcat_tokens"], + total_lengths=batch["total_lengths"], + response_lengths=batch["response_lengths"], + with_entropy=False, + max_seq_lens=batch.get("max_seq_lens", None), + )["log_probs"] + + +def run(args, batch, logits): + loss, metrics = tinker_loss_function(args, batch, logits, sum_of_sample_mean=None) + return loss, metrics + + +def test_linear_cross_entropy_is_a_plain_weighted_sum(): + args, batch, logits = make_batch() + weights = [torch.tensor([0.5, 0.0, 2.0]), torch.tensor([1.0, 1.0, 0.0, -1.0, 0.25])] + batch["loss_weights"] = weights + + loss, metrics = run(args, batch, logits) + expected = sum(-(lp * w).sum() for lp, w in zip(reference_log_probs(args, batch, logits), weights, strict=True)) + assert torch.allclose(loss, expected) + assert torch.allclose(metrics["loss"], expected) + assert loss.requires_grad + + +def test_binary_mask_still_gates_tokens(): + args, batch, logits = make_batch() + batch["loss_weights"] = [torch.ones(3), torch.ones(5)] + batch["loss_masks"] = [torch.tensor([1, 0, 1], dtype=torch.int32), torch.zeros(5, dtype=torch.int32)] + + loss, _ = run(args, batch, logits) + lp = reference_log_probs(args, batch, logits) + expected = -(lp[0] * torch.tensor([1.0, 0.0, 1.0])).sum() + assert torch.allclose(loss, expected) + + +def test_importance_sampling_and_ppo_clip(): + args, batch, logits = make_batch() + advantages = [torch.tensor([1.0, -1.0, 2.0]), torch.tensor([0.5, 0.5, -0.5, 1.0, 0.0])] + batch["advantages"] = advantages + batch["tinker_loss_by_lane"] = {0: {"loss_fn": "importance_sampling"}} + + loss, _ = run(args, batch, logits) + lp = reference_log_probs(args, batch, logits) + ratios = [torch.exp(new - old) for new, old in zip(lp, batch["rollout_log_probs"], strict=True)] + expected = sum(-(r * a).sum() for r, a in zip(ratios, advantages, strict=True)) + assert torch.allclose(loss, expected) + + batch["tinker_loss_by_lane"] = { + 0: {"loss_fn": "ppo", "loss_fn_config": {"clip_low_threshold": 0.9, "clip_high_threshold": 1.1}} + } + loss_ppo, _ = run(args, batch, logits) + expected_ppo = sum( + -torch.minimum(r * a, r.clamp(0.9, 1.1) * a).sum() for r, a in zip(ratios, advantages, strict=True) + ) + assert torch.allclose(loss_ppo, expected_ppo) + assert not torch.allclose(loss_ppo, loss) + + +def test_mixed_lanes_dispatch_independently(): + args, batch, logits = make_batch() + batch["tinker_operation_lanes"] = [0, 1] + batch["loss_weights"] = [torch.ones(3), torch.zeros(5)] + batch["advantages"] = [torch.zeros(3), torch.ones(5)] + batch["tinker_loss_by_lane"] = { + 0: {"loss_fn": "cross_entropy"}, + 1: {"loss_fn": "importance_sampling"}, + } + + loss, _ = run(args, batch, logits) + lp = reference_log_probs(args, batch, logits) + ratio = torch.exp(lp[1] - batch["rollout_log_probs"][1]) + expected = -(lp[0].sum()) + -(ratio.sum()) + assert torch.allclose(loss, expected) + + +def test_sum_reduction_is_chunk_additive(): + args, batch, logits = make_batch() + batch["loss_weights"] = [torch.ones(3) * 0.5, torch.ones(5) * 1.5] + full_loss, _ = run(args, batch, logits) + + total = 0.0 + offset = 0 + for i, total_len in enumerate(batch["total_lengths"]): + sub_logits = logits[:, offset : offset + total_len] + sub = dict( + unconcat_tokens=[batch["unconcat_tokens"][i]], + total_lengths=[total_len], + response_lengths=[batch["response_lengths"][i]], + loss_masks=[batch["loss_masks"][i]], + loss_weights=[batch["loss_weights"][i]], + tinker_operation_lanes=[0], + tinker_loss_by_lane=batch["tinker_loss_by_lane"], + ) + sub_loss, _ = run(args, sub, sub_logits) + total += sub_loss + offset += total_len + assert torch.allclose(full_loss, total) + + +def test_zero_weight_padding_contributes_nothing(): + args, batch, logits = make_batch() + batch["loss_weights"] = [torch.ones(3), torch.zeros(5)] + loss, _ = run(args, batch, logits) + lp = reference_log_probs(args, batch, logits) + assert torch.allclose(loss, -(lp[0].sum())) + + +def test_missing_channel_missing_spec_and_unknown_loss_fail_loudly(): + args, batch, logits = make_batch() + with pytest.raises(ValueError, match="needs per-token 'loss_weights'"): + run(args, batch, logits) + + batch["loss_weights"] = [torch.ones(3), torch.ones(5)] + batch["tinker_operation_lanes"] = [0, 3] + with pytest.raises(ValueError, match="no loss spec for lane 3"): + run(args, batch, logits) + + batch["tinker_operation_lanes"] = [0, 0] + batch["tinker_loss_by_lane"] = {0: {"loss_fn": "dro"}} + with pytest.raises(ValueError, match="unknown loss_fn 'dro'"): + run(args, batch, logits) + + +def test_collector_captures_per_datum_logprobs_in_row_order(): + args, batch, logits = make_batch() + batch["loss_weights"] = [torch.ones(3), torch.ones(5)] + batch["sample_indices"] = [0, 1] + collector: dict = {} + batch["tinker_logprob_collector"] = collector + + run(args, batch, logits) + lp = reference_log_probs(args, batch, logits) + assert set(collector) == {(0, 0), (0, 1)} + assert collector[(0, 0)] == pytest.approx(lp[0].tolist()) + assert collector[(0, 1)] == pytest.approx(lp[1].tolist()) + + +def test_forward_only_batch_collects_logprobs_without_client_loss_terms(): + args, batch, logits = make_batch() + batch["tinker_operation_lanes"] = [0, 1] + batch["tinker_loss_by_lane"] = {} + batch["tinker_forward_only"] = True + batch["sample_indices"] = [0, 0] + collector: dict = {} + batch["tinker_logprob_collector"] = collector + + loss, metrics = run(args, batch, logits) + lp = reference_log_probs(args, batch, logits) + assert loss.item() == 0.0 and metrics["loss"].item() == 0.0 + assert collector[(0, 0)] == pytest.approx(lp[0].tolist()) + assert collector[(1, 0)] == pytest.approx(lp[1].tolist()) diff --git a/tests/fast/backends/training_utils/test_get_batch_multi_lora_cp.py b/tests/fast/backends/training_utils/test_get_batch_multi_lora_cp.py index cdc4e07ca17..f86db947c49 100644 --- a/tests/fast/backends/training_utils/test_get_batch_multi_lora_cp.py +++ b/tests/fast/backends/training_utils/test_get_batch_multi_lora_cp.py @@ -24,7 +24,7 @@ def __init__(self, batch: dict, n_adapters: int): self.rollout_data = {"n_adapters": n_adapters} def get_next(self, keys): - return {key: self._batch[key] for key in keys} + return {key: self._batch.get(key) for key in keys} KEYS = ["tokens", "loss_masks", "total_lengths", "response_lengths", "adapter_slots"] diff --git a/tests/fast/backends/training_utils/test_log_rollout_data_tinker_keys.py b/tests/fast/backends/training_utils/test_log_rollout_data_tinker_keys.py new file mode 100644 index 00000000000..157d841485f --- /dev/null +++ b/tests/fast/backends/training_utils/test_log_rollout_data_tinker_keys.py @@ -0,0 +1,61 @@ +from argparse import Namespace +from types import SimpleNamespace + +import torch + +from miles.backends.training_utils import cp_utils, log_utils + + +def test_every_tinker_conversion_key_is_handled(monkeypatch): + parallel_state = SimpleNamespace( + tp=SimpleNamespace(rank=0), + cp=SimpleNamespace(size=1), + intra_dp=SimpleNamespace(size=1), + is_pp_last_stage=True, + ) + monkeypatch.setattr(log_utils, "get_parallel_state", lambda: parallel_state) + monkeypatch.setattr(cp_utils, "get_parallel_state", lambda: parallel_state) + monkeypatch.setattr(log_utils, "gather_log_data", lambda *a, **k: None) + + rollout_data = { + "tokens": [torch.tensor([1, 2, 3])], + "total_lengths": [3], + "response_lengths": [2], + "rewards": [0.0], + "raw_reward": [0.0], + "truncated": [0], + "loss_masks": [torch.tensor([1, 1], dtype=torch.int32)], + "sample_indices": [0], + "rollout_ids": [0], + "rollout_mask_sums": torch.tensor([2]), + "loss_weights": [torch.tensor([1.0, 1.0])], + "advantages": [torch.tensor([0.0, 0.0])], + "adapter_slots": [0], + "batch_kind": "tinker", + "tinker_operation_lanes": [0], + "tinker_loss_by_lane": {0: {"loss_fn": "cross_entropy"}}, + "operation_by_lane": {0: "op-A"}, + "registration_by_lane": {0: ("A", "r-A")}, + "batch_execution_lease": { + "dispatch_id": "d", + "bindings_by_operation": [["op-A", ["A", "r-A", 0]]], + }, + "tinker_forward_only": True, + "tinker_logprob_collector": {}, + "dynamic_global_batch_size": 1, + "n_adapters": 2, + } + + log_utils.log_rollout_data( + 0, + Namespace( + ci_test=False, + ci_disable_logprobs_checker=True, + true_on_policy_mode=False, + qkv_format="thd", + log_multi_turn=False, + log_passrate=False, + log_correct_samples=False, + ), + rollout_data, + ) diff --git a/tests/fast/backends/training_utils/test_operation_execution.py b/tests/fast/backends/training_utils/test_operation_execution.py new file mode 100644 index 00000000000..6962e3cae82 --- /dev/null +++ b/tests/fast/backends/training_utils/test_operation_execution.py @@ -0,0 +1,97 @@ +from miles.backends.training_utils.operation_execution import ( + ADAM_PARAM_DEFAULTS, + StepRequest, + resolve_adam_params, + run_optim_controls, +) +from miles.utils.operation_contract import BatchExecutionLease + + +class FakeExecutor: + def __init__(self, step_outcomes=None, discard_outcomes=None): + self.discarded: list[str] = [] + self.stepped: list[StepRequest] = [] + self._step_outcomes = step_outcomes or {} + self._discard_outcomes = discard_outcomes + + def discard_many(self, lease, operation_ids): + self.discarded.extend(operation_ids) + if self._discard_outcomes is not None: + return self._discard_outcomes + return {op_id: dict(ok=True) for op_id in operation_ids} + + def step_many(self, lease, requests): + self.stepped.extend(requests) + return { + request.operation_id: self._step_outcomes.get(request.operation_id, dict(ok=True, result={})) + for request in requests + } + + +LEASE = BatchExecutionLease(dispatch_id="d", bindings_by_operation=(("opt1", "opaque-1"), ("opt2", "opaque-2"))) + + +def optim(op_id, adam=None, poison=None): + op = dict(operation_id=op_id, kind="optim_step", payload={"adam_params": adam} if adam else {}) + if poison: + op["poison"] = poison + return op + + +class TestResolveAdamParams: + def test_defaults_fill_and_none_is_absent(self): + resolved = resolve_adam_params({"learning_rate": 3e-4, "grad_clip_norm": None}) + assert resolved["learning_rate"] == 3e-4 + assert resolved["grad_clip_norm"] == ADAM_PARAM_DEFAULTS["grad_clip_norm"] + assert resolve_adam_params(None) == ADAM_PARAM_DEFAULTS + + +class TestRunOptimControls: + def test_poisoned_steps_discard_and_fail_as_user_errors(self): + executor = FakeExecutor() + results = run_optim_controls( + [optim("opt1", poison="window poisoned"), optim("opt2", adam={"learning_rate": 2e-4})], + LEASE, + executor, + ) + assert executor.discarded == ["opt1"] + assert results["opt1"] == dict( + ok=False, error="window poisoned", category="user", gradient_window_consumed=True + ) + [request] = executor.stepped + assert request.operation_id == "opt2" and request.adam_params["learning_rate"] == 2e-4 + assert results["opt2"]["ok"] is True + + def test_executor_refusal_wins_over_the_poison_policy(self): + executor = FakeExecutor(discard_outcomes={"opt1": dict(ok=False, error="stale binding", category="server")}) + results = run_optim_controls([optim("opt1", poison="poisoned")], LEASE, executor) + assert results["opt1"] == dict(ok=False, error="stale binding", category="server") + assert not results["opt1"].get("gradient_window_consumed") + + def test_missing_discard_outcome_fails_closed_as_a_server_error(self): + executor = FakeExecutor(discard_outcomes={}) + results = run_optim_controls([optim("opt1", poison="poisoned")], LEASE, executor) + outcome = results["opt1"] + assert outcome["ok"] is False and outcome["category"] == "server" + assert "discard" in outcome["error"] + assert not outcome.get("gradient_window_consumed") + + def test_missing_step_outcome_fails_closed_as_a_server_error(self): + class SilentExecutor(FakeExecutor): + def step_many(self, lease, requests): + return {} + + results = run_optim_controls([optim("opt1")], LEASE, SilentExecutor()) + outcome = results["opt1"] + assert outcome["ok"] is False and outcome["category"] == "server" + assert not outcome.get("gradient_window_consumed") + + def test_clean_step_needs_no_prior_fb(self): + executor = FakeExecutor() + results = run_optim_controls([optim("opt1")], LEASE, executor) + assert results["opt1"]["ok"] is True + + def test_non_optim_operations_are_not_the_coordinators_business(self): + executor = FakeExecutor() + results = run_optim_controls([dict(operation_id="save1", kind="save_state")], LEASE, executor) + assert results == {} and executor.stepped == [] and executor.discarded == [] diff --git a/tests/fast/ray/multi_lora/test_backend.py b/tests/fast/ray/multi_lora/test_backend.py new file mode 100644 index 00000000000..1e424482d39 --- /dev/null +++ b/tests/fast/ray/multi_lora/test_backend.py @@ -0,0 +1,523 @@ +import asyncio +from types import SimpleNamespace + +import pytest + +from miles.ray.multi_lora.backend import MultiLoraOperationBackend +from miles.ray.multi_lora.config import AdapterRunConfig +from miles.ray.multi_lora.registry import AdapterState + + +def make_backend(max_adapters: int = 4) -> MultiLoraOperationBackend: + args = SimpleNamespace( + multi_lora_n_adapters=max_adapters, + save="/tmp/tinker-test-save", + lora_rank=32, + lora_alpha=64, + hf_checkpoint="Qwen/Qwen3-0.6B", + ) + return MultiLoraOperationBackend(args, "http://unused") + + +def register(backend, name="X", **overrides) -> dict: + return asyncio.run(backend.register(name, AdapterRunConfig(**overrides))) + + +def ready_backend(num_step=None): + backend = make_backend() + register(backend, num_step=num_step) + backend.registry.mark_ready(["X"]) + return backend + + +def reg_key(backend, name="X"): + return (name, backend.registry.find(name).registration_id) + + +def fb_payload(n=1, loss_fn="cross_entropy"): + return { + "samples": [ + {"tokens": [1, 2, 3, 4], "response_length": 2, "loss_mask": [1, 1], "loss_weights": [1.0, 1.0]} + for _ in range(n) + ], + "loss": {"loss_fn": loss_fn}, + } + + +class TestRegistration: + def test_resolves_rank_alpha_and_save(self): + backend = make_backend() + result = register(backend, rank=8) + assert result == {"name": "X", "slot": 0} + config = backend.registry.find("X").config + assert config.rank == 8 and config.alpha == 64 + assert str(config.save).endswith("adapters/X") + + def test_rank_ceiling_and_client_alpha_rejected(self): + backend = make_backend() + with pytest.raises(ValueError, match="exceeds the deployment maximum"): + register(backend, rank=64) + with pytest.raises(ValueError, match="must not set alpha"): + register(backend, alpha=16) + + +class TestPreflight: + def test_unsupported_loss_is_a_boundary_error(self): + backend = ready_backend() + with pytest.raises(ValueError, match="not supported in v1"): + backend.enqueue_operation("X", "op1", 1, "forward_backward", fb_payload(loss_fn="cispo")) + + def test_multimodal_and_nested_targets_rejected(self): + backend = ready_backend() + bad = fb_payload() + bad["samples"][0]["multimodal_inputs"] = {"image": "..."} + with pytest.raises(ValueError, match="text-only"): + backend.enqueue_operation("X", "op1", 1, "forward_backward", bad) + nested = fb_payload() + nested["samples"][0]["loss_weights"] = [[1.0, 2.0], [3.0, 4.0]] + with pytest.raises(ValueError, match="1-D"): + backend.enqueue_operation("X", "op1", 1, "forward_backward", nested) + + def test_channel_length_must_match_response(self): + backend = ready_backend() + bad = fb_payload() + bad["samples"][0]["advantages"] = [1.0] + with pytest.raises(ValueError, match="length response_length"): + backend.enqueue_operation("X", "op1", 1, "forward_backward", bad) + + def test_adam_params_validated(self): + backend = ready_backend() + with pytest.raises(ValueError, match="unknown adam_params field"): + backend.enqueue_operation("X", "op1", 1, "optim_step", {"adam_params": {"lr": 1e-4}}) + with pytest.raises(ValueError, match="finite number"): + backend.enqueue_operation("X", "op1", 1, "optim_step", {"adam_params": {"learning_rate": "fast"}}) + + def test_adam_params_domain_checked_at_the_boundary(self): + backend = ready_backend() + rejected = [ + {"learning_rate": float("nan")}, + {"learning_rate": float("inf")}, + {"learning_rate": -1e-4}, + {"beta1": 2.0}, + {"beta2": -0.1}, + {"beta1": 1.0}, + {"eps": 0.0}, + {"eps": -1e-8}, + {"weight_decay": float("nan")}, + {"weight_decay": -0.1}, + {"grad_clip_norm": -1.0}, + {"learning_rate": True}, + ] + for adam in rejected: + with pytest.raises(ValueError, match="adam_params"): + backend.enqueue_operation("X", "op1", 1, "optim_step", {"adam_params": adam}) + ok = {"learning_rate": 3e-4, "beta1": 0.9, "beta2": 0.95, "eps": 1e-12, "weight_decay": 0.0} + assert backend.enqueue_operation("X", "op1", 1, "optim_step", {"adam_params": ok})["state"] == "QUEUED" + + def test_loss_required_channels_preflighted(self): + backend = ready_backend() + ce = fb_payload() + del ce["samples"][0]["loss_weights"] + with pytest.raises(ValueError, match="loss_weights"): + backend.enqueue_operation("X", "op1", 1, "forward_backward", ce) + for missing in ("rollout_log_probs", "advantages"): + for loss_fn in ("importance_sampling", "ppo"): + bad = fb_payload(loss_fn=loss_fn) + del bad["samples"][0]["loss_weights"] + bad["samples"][0]["rollout_log_probs"] = [-1.0, -1.0] + bad["samples"][0]["advantages"] = [0.5, 0.5] + del bad["samples"][0][missing] + with pytest.raises(ValueError, match=missing): + backend.enqueue_operation("X", "op1", 1, "forward_backward", bad) + bare = {"samples": [{"tokens": [1, 2, 3, 4], "response_length": 2}]} + assert backend.enqueue_operation("X", "op2", 1, "forward", bare)["state"] == "QUEUED" + + def test_response_must_leave_a_context_token(self): + backend = ready_backend() + bad = fb_payload() + bad["samples"][0].update(response_length=4, loss_mask=[1] * 4, loss_weights=[1.0] * 4) + with pytest.raises(ValueError, match="response_length"): + backend.enqueue_operation("X", "op1", 1, "forward_backward", bad) + + def test_unknown_kind_and_missing_path(self): + backend = ready_backend() + with pytest.raises(ValueError, match="unknown operation kind"): + backend.enqueue_operation("X", "op1", 1, "publish_snapshot") + with pytest.raises(ValueError, match="needs a 'path'"): + backend.enqueue_operation("X", "op1", 1, "load_state", {}) + + def test_save_state_tag_must_stay_inside_states(self): + backend = ready_backend() + for bad in ("..", ".", "a/b", "a" * 129, ""): + with pytest.raises(ValueError, match="tag"): + backend.enqueue_operation("X", f"save-{len(bad)}", 1, "save_state", {"tag": bad}) + assert backend.enqueue_operation("X", "save-ok", 1, "save_state", {"tag": "step_5.final"}) + + +class TestControlClaims: + def test_claim_requires_ready_and_serialization(self): + backend = make_backend() + register(backend) + backend.enqueue_operation("X", "opt1", 1, "optim_step") + assert backend.claim_ready_control_operations() == {"operations": [], "lease": None} + backend.registry.mark_ready(["X"]) + claimed = backend.claim_ready_control_operations() + [op] = claimed["operations"] + assert op["operation_id"] == "opt1" + assert "slot" not in op + rid = backend.registry.find("X").registration_id + assert claimed["lease"]["bindings_by_operation"] == [["opt1", ["X", rid, 0]]] + + def test_claim_carries_authoritative_clocks(self): + backend = ready_backend() + backend.set_adapter_step("X", 7) + backend.registry.record_weight_update(["X"]) + backend.enqueue_operation("X", "pub1", 1, "save_weights_for_sampler") + [op] = backend.claim_ready_control_operations()["operations"] + assert op["step"] == 7 and op["serving_version"] == 1 + + def test_dirty_slot_fails_state_moves_but_allows_publish(self): + backend = ready_backend() + backend.commit_tinker_batch([reg_key(backend)], []) + backend.enqueue_operation("X", "save1", 1, "save_state", {"tag": "t0"}) + assert backend.claim_ready_control_operations() == {"operations": [], "lease": None} + view = backend.operations.get("save1") + assert view["state"] == "FAILED" and "unstepped gradients" in view["error"] + + backend.enqueue_operation("X", "pub1", 2, "save_weights_for_sampler") + [op] = backend.claim_ready_control_operations()["operations"] + assert op["operation_id"] == "pub1" + + def test_success_advances_step_and_releases_pin(self): + backend = ready_backend(num_step=2) + backend.commit_tinker_batch([reg_key(backend)], []) + backend.enqueue_operation("X", "opt1", 1, "optim_step") + [op] = backend.claim_ready_control_operations()["operations"] + backend.complete_control_operations({op["operation_id"]: dict(ok=True, result={"grad_norm": 0.5})}) + record = backend.registry.find("X") + assert record.step == 1 and not backend.registry.is_dirty("X") + + def test_veto_fails_without_advancing(self): + backend = ready_backend() + backend.commit_tinker_batch([reg_key(backend)], []) + backend.enqueue_operation("X", "opt1", 1, "optim_step") + [op] = backend.claim_ready_control_operations()["operations"] + backend.complete_control_operations( + {op["operation_id"]: dict(ok=False, error="veto", category="server", gradient_window_consumed=True)} + ) + assert backend.registry.find("X").step == 0 + assert not backend.registry.is_dirty("X") + + def test_failed_chunk_poisons_the_pending_optim(self): + backend = ready_backend() + rid = backend.registry.find("X").registration_id + backend.enqueue_operation("X", "fb1", 1, "forward_backward", fb_payload()) + backend.operations.claim_data_operation("X", rid) + backend.operations.fail("fb1", "bad chunk", "user") + backend.enqueue_operation("X", "opt2", 2, "optim_step") + [op] = backend.claim_ready_control_operations()["operations"] + assert "gradient window" in op["poison"] and "discarded" in op["poison"] + backend.complete_control_operations( + {"opt2": dict(ok=False, error=op["poison"], category="user", gradient_window_consumed=True)} + ) + assert backend.registry.find("X").step == 0 + + backend.enqueue_operation("X", "fb3", 3, "forward_backward", fb_payload()) + backend.operations.claim_data_operation("X", rid) + backend.commit_tinker_batch([reg_key(backend)], ["fb3"], {"fb3": [[-0.1, -0.2]]}) + backend.enqueue_operation("X", "opt4", 4, "optim_step") + [clean] = backend.claim_ready_control_operations()["operations"] + assert clean["operation_id"] == "opt4" and "poison" not in clean + + def test_pre_mutation_refusal_keeps_dirty_and_poison(self): + backend = ready_backend() + rid = backend.registry.find("X").registration_id + backend.enqueue_operation("X", "fb1", 1, "forward_backward", fb_payload()) + backend.claim_data_operation("X", rid) + backend.commit_tinker_batch([reg_key(backend)], ["fb1"], {"fb1": [[-0.1, -0.2]]}) + backend.enqueue_operation("X", "fb2", 2, "forward_backward", fb_payload()) + backend.claim_data_operation("X", rid) + backend.operations.fail("fb2", "partial backward failed", "server") + + backend.enqueue_operation("X", "opt3", 3, "optim_step") + [poisoned] = backend.claim_ready_control_operations()["operations"] + assert poisoned.get("poison") + backend.complete_control_operations( + {"opt3": dict(ok=False, error="stale binding: no gradients were cleared", category="server")} + ) + + backend.enqueue_operation("X", "opt4", 4, "optim_step") + [next_optim] = backend.claim_ready_control_operations()["operations"] + assert backend.gradient_windows.is_dirty(("X", rid)) + assert next_optim.get("poison"), "a refused optimizer dispatch is not a window delimiter" + + def test_stale_registration_handle_is_fenced(self): + backend = ready_backend() + rid1 = backend.registry.find("X").registration_id + assert backend.enqueue_operation("X", "op1", 1, "optim_step", None, expected_registration_id=rid1) + backend.registry.deregister("X") + backend.registry.retire_adapters() + backend.registry.free_slot("X") + register(backend, "X") + rid2 = backend.registry.records["X"].registration_id + assert rid2 != rid1 + with pytest.raises(ValueError, match="fenced"): + backend.enqueue_operation("X", "op9", 1, "optim_step", None, expected_registration_id=rid1) + assert backend.operations.queue_view("X", rid2) == [] + asyncio.run(backend.deregister("X", rid1)) + assert backend.registry.records["X"].state is AdapterState.PENDING + + def test_publish_completion_stamps_post_push_serving_identity(self): + backend = ready_backend() + backend.registry.record_weight_update(["X"]) + backend.enqueue_operation("X", "pub1", 1, "save_weights_for_sampler") + [op] = backend.claim_ready_control_operations()["operations"] + backend.complete_control_operations({op["operation_id"]: dict(ok=True, result={})}) + result = backend.operations.get("pub1")["result"] + assert result["serving_version"] == 1 + reg_id = backend.registry.find("X").registration_id + assert result["serving_name"] == f"__miles_adapter_X_{reg_id}" + + def test_load_state_repositions_the_clock(self): + backend = ready_backend() + backend.enqueue_operation("X", "load1", 1, "load_state", {"path": "/tmp/state"}) + [op] = backend.claim_ready_control_operations()["operations"] + backend.complete_control_operations({op["operation_id"]: dict(ok=True, result={"step": 42})}) + record = backend.registry.find("X") + assert record.step == 42 and record.start_step == 42 + + +class TestCommitAndFence: + def test_commit_completes_data_ops_with_row_ordered_logprobs(self): + backend = ready_backend() + reg_id = backend.registry.find("X").registration_id + backend.enqueue_operation("X", "fb1", 1, "forward_backward", fb_payload()) + backend.operations.claim_data_operation("X", reg_id) + backend.commit_tinker_batch([reg_key(backend)], ["fb1"], {"fb1": [[-0.1, -0.2]]}) + result = backend.operations.get("fb1")["result"] + assert result["logprobs"] == [[-0.1, -0.2]] + assert result["metrics"]["loss:sum"] == pytest.approx(0.1 + 0.2) + assert backend.registry.is_dirty("X") + + def test_retirement_fences_open_operations(self, monkeypatch): + backend = ready_backend() + backend.enqueue_operation("X", "op1", 1, "forward_backward", fb_payload()) + + async def no_abort(name, registration_id): + pass + + monkeypatch.setattr(backend, "abort_adapter_requests", no_abort) + asyncio.run(backend.deregister("X")) + asyncio.run(backend.retire_adapters()) + view = backend.operations.get("op1") + assert view["state"] == "FAILED" and view["error_category"] == "user" + with pytest.raises(ValueError, match="not accepting operations"): + backend.enqueue_operation("X", "op2", 2, "forward_backward", fb_payload()) + assert backend.registry.records["X"].state is AdapterState.CLEANUP + + +class TestFailTinkerBatch: + def _claimed_batch(self, backend): + rid = backend.registry.find("X").registration_id + backend.enqueue_operation("X", "fb1", 1, "forward_backward", fb_payload()) + claim = backend.claim_data_operation("X", rid) + lease = backend.acquire_batch_lease([("fb1", claim["binding"])]) + from miles.ray.multi_lora.residency import lease_to_metadata + + return lease_to_metadata(lease) + + def test_uncommitted_batch_terminal_fails_claimed_operations_typed_server(self): + backend = ready_backend() + lease_metadata = self._claimed_batch(backend) + backend.fail_tinker_batch(["fb1"], "train step finished without committing", lease_metadata) + view = backend.operations.get("fb1") + assert view["state"] == "FAILED" and view["error_category"] == "server" + assert "without committing" in view["error"] + + def test_finalized_forward_backward_is_poison_evidence_for_the_next_optim(self): + backend = ready_backend() + lease_metadata = self._claimed_batch(backend) + backend.fail_tinker_batch(["fb1"], "abnormal train outcome", lease_metadata) + backend.enqueue_operation("X", "opt2", 2, "optim_step") + [op] = backend.claim_ready_control_operations()["operations"] + assert "forward_backward ordinal 1" in op["poison"] + + def test_already_terminal_operations_are_left_untouched(self): + backend = ready_backend() + lease_metadata = self._claimed_batch(backend) + backend.commit_tinker_batch([reg_key(backend)], ["fb1"], {"fb1": [[-0.1, -0.2]]}) + backend.fail_tinker_batch(["fb1"], "late failure", lease_metadata) + view = backend.operations.get("fb1") + assert view["state"] == "SUCCEEDED" and view["result"]["logprobs"] == [[-0.1, -0.2]] + + def test_lease_releases_even_when_the_ledger_walk_raises(self): + backend = ready_backend() + lease_metadata = self._claimed_batch(backend) + released = [] + backend.residency.release_batch = lambda lease: released.append(lease.dispatch_id) + + def boom(operation_id, error, category="server"): + raise RuntimeError("ledger unavailable") + + backend.operations.fail = boom + with pytest.raises(RuntimeError, match="ledger unavailable"): + backend.fail_tinker_batch(["fb1"], "abnormal train outcome", lease_metadata) + assert released == [lease_metadata["dispatch_id"]] + + def test_unknown_operation_ids_and_missing_lease_are_tolerated(self): + backend = ready_backend() + backend.fail_tinker_batch(["ghost"], "abnormal train outcome", None) + + +def test_service_info_reports_the_v1_matrix(): + backend = ready_backend() + info = backend.service_info() + assert info["base_model"] == "Qwen/Qwen3-0.6B" + assert info["lora_rank_max"] == 32 and info["n_adapters"] == 4 + assert info["occupied_slots"] == [0] and info["ready_adapters"] == ["X"] + assert info["supported_loss_fns"] == ["cross_entropy", "importance_sampling", "ppo"] + + +def test_engine_aborts_go_through_the_inference_admin_port(): + backend = make_backend() + aborted = [] + + class FakeAdmin: + async def abort_registration(self, rid_prefix): + aborted.append(rid_prefix) + + backend.inference_admin = FakeAdmin() + asyncio.run(backend.abort_adapter_requests("X", "reg-1")) + assert aborted == ["X::reg-1::"] + + +def test_trainer_readiness_flag_flips_once_marked(): + backend = make_backend() + assert backend.trainer_ready is False + backend.mark_trainer_ready() + assert backend.trainer_ready is True + + +def test_advertised_host_is_the_bind_host(): + from miles.ray.multi_lora.http_server import AdapterRunControlServer + + assert AdapterRunControlServer(None, host="127.0.0.1").advertised_host == "127.0.0.1" + + +class TestRejectOperation: + def test_rejects_into_the_ledger_only_for_live_registrations(self): + backend = ready_backend() + view = backend.reject_operation("X", "op1", 1, "optim_step", {"adam_params": {}}, "unsupported") + assert view["state"] == "FAILED" and view["error_category"] == "user" + backend.enqueue_operation("X", "op2", 2, "forward_backward", fb_payload()) + assert backend.operations.claim_data_operation("X", view["registration_id"])["operation_id"] == "op2" + backend.registry.deregister("X") + with pytest.raises(ValueError, match="not accepting operations"): + backend.reject_operation("X", "op3", 3, "optim_step", {}, "late") + + def test_reject_from_a_stale_handle_never_lands_on_a_successor(self): + backend = ready_backend() + rid1 = backend.registry.find("X").registration_id + backend.registry.deregister("X") + backend.registry.retire_adapters() + backend.registry.free_slot("X") + register(backend, "X") + rid2 = backend.registry.records["X"].registration_id + with pytest.raises(ValueError, match="fenced"): + backend.reject_operation("X", "op1", 1, "optim_step", {}, "bad", expected_registration_id=rid1) + assert backend.operations.queue_view("X", rid2) == [] + + +class TestGapTimeoutSurface: + def stalled_backend(self, timeout=30.0): + backend = ready_backend() + backend.operations.gap_timeout = timeout + clock = {"now": 1000.0} + backend.operations._time = lambda: clock["now"] + backend.enqueue_operation("X", "fb1", 1, "forward_backward", fb_payload()) + backend.claim_data_operation(*reg_key(backend)) + backend.operations.complete("fb1", {}) + backend.enqueue_operation("X", "opt3", 3, "optim_step", {"adam_params": {"learning_rate": 1e-4}}) + assert backend.claim_ready_control_operations()["operations"] == [] + return backend, clock + + def test_flag_reaches_the_ledger_with_a_default(self): + assert make_backend().operations.gap_timeout == 600.0 + args = SimpleNamespace(multi_lora_n_adapters=4, tinker_operation_gap_timeout=5.0) + assert MultiLoraOperationBackend(args, "http://unused").operations.gap_timeout == 5.0 + + def test_stall_is_typed_and_observable_before_expiry(self): + backend, clock = self.stalled_backend() + clock["now"] += 10 + info = backend.service_info() + assert info["operation_gap_timeout"] == 30.0 + [stall] = info["gap_stalls"] + assert stall["missing_ordinal"] == 2 and stall["blocked_operations"] == 1 + view = backend.operation_view("opt3") + assert view["state"] == "QUEUED" + assert view["waiting_on_ordinal"] == 2 and view["gap_stalled_for"] == pytest.approx(10.0) + + def test_control_claim_heartbeat_expires_the_stall(self): + backend, clock = self.stalled_backend() + clock["now"] += 31 + assert backend.claim_ready_control_operations()["operations"] == [] + view = backend.operation_view("opt3") + assert view["state"] == "FAILED" and view["error_category"] == "user" + assert "missing ordinal 2" in view["error"] + assert backend.service_info()["gap_stalls"] == [] + backend.enqueue_operation("X", "opt4", 4, "optim_step", {"adam_params": {"learning_rate": 1e-4}}) + [operation] = backend.claim_ready_control_operations()["operations"] + assert operation["operation_id"] == "opt4" and "poison" not in operation + + +class TestClaimedTtlSurface: + def orphaned_backend(self, ttl=60.0): + backend = ready_backend() + backend.operations.claimed_ttl = ttl + clock = {"now": 1000.0} + backend.operations._time = lambda: clock["now"] + backend.enqueue_operation("X", "fb1", 1, "forward_backward", fb_payload()) + assert backend.claim_data_operation(*reg_key(backend)) is not None + return backend, clock + + def test_flag_reaches_the_ledger_with_a_default(self): + assert make_backend().operations.claimed_ttl == 1800.0 + args = SimpleNamespace(multi_lora_n_adapters=4, tinker_operation_claimed_ttl=5.0) + assert MultiLoraOperationBackend(args, "http://unused").operations.claimed_ttl == 5.0 + + def test_heartbeat_fails_the_orphan_typed_server_and_unblocks_the_queue(self): + backend, clock = self.orphaned_backend() + clock["now"] += 61 + backend.enqueue_operation("X", "opt2", 2, "optim_step") + [op] = backend.claim_ready_control_operations()["operations"] + assert op["operation_id"] == "opt2" + view = backend.operations.get("fb1") + assert view["state"] == "FAILED" and view["error_category"] == "server" + assert "'fb1'" in view["error"] and "61s" in view["error"] and "forward_backward" in view["error"] + + def test_sweep_routes_through_the_lease_releasing_batch_finalizer(self): + backend, clock = self.orphaned_backend() + calls = [] + original = backend.fail_tinker_batch + + def spy(operation_ids, error, lease_metadata=None): + calls.append((operation_ids, lease_metadata)) + original(operation_ids, error, lease_metadata) + + backend.fail_tinker_batch = spy + clock["now"] += 61 + assert backend.service_info()["operation_claimed_ttl"] == 60.0 + assert calls == [(["fb1"], None)] + + def test_younger_claim_survives_the_sweep(self): + backend, clock = self.orphaned_backend() + clock["now"] += 59 + backend.service_info() + assert backend.operations.get("fb1")["state"] == "CLAIMED" + + def test_late_completion_of_a_swept_operation_is_ignored_not_a_crash(self): + backend, clock = self.orphaned_backend() + clock["now"] += 61 + backend.service_info() + backend.complete_control_operations({"fb1": dict(ok=True, result={})}) + assert backend.operations.get("fb1")["state"] == "FAILED" diff --git a/tests/fast/ray/multi_lora/test_controller_backend.py b/tests/fast/ray/multi_lora/test_controller_backend.py deleted file mode 100644 index fe75f985eff..00000000000 --- a/tests/fast/ray/multi_lora/test_controller_backend.py +++ /dev/null @@ -1,392 +0,0 @@ -"""Fast tests for AdapterRegistry + MultiLoRABackend validation -(no Ray, no HTTP I/O, no SGLang, no torch).""" - -from types import SimpleNamespace - -from tests.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=60, suite="stage-a-cpu") - -import pytest - -from miles.ray.multi_lora.backend import MultiLoRABackend -from miles.ray.multi_lora.registry import AdapterRegistry, AdapterState -from miles.utils.adapter_config import AdapterRunConfig -from miles.utils.multi_lora import make_rid, min_groups_per_dp_split, parse_adapter - - -# Registration validates that the data path exists; the test file itself is a -# convenient always-present stand-in. -DATA_FILE = __file__ - - -def make_args(max_adapters: int = 4, save: str | None = None, dp_size: int = 2) -> SimpleNamespace: - return SimpleNamespace( - multi_lora_n_adapters=max_adapters, - save=save, - lora_rank=32, - lora_alpha=32, - rollout_batch_size=16, - n_samples_per_prompt=4, - multi_lora_dp_size=dp_size, - multi_lora_max_adapter_global_batch_size=256, - ) - - -def make_backend(max_adapters: int = 4, save: str | None = None, dp_size: int = 2) -> MultiLoRABackend: - return MultiLoRABackend(make_args(max_adapters, save, dp_size), "http://unused") - - -def make_config(save: str | None = None, **overrides) -> AdapterRunConfig: - kwargs = dict( - rank=8, - alpha=16, - data=DATA_FILE, - rollout_batch_size=4, - n_samples_per_prompt=4, - save=save, - input_key="text", - label_key="label", - rm_type="math", - ) - kwargs.update(overrides) - return AdapterRunConfig(**kwargs) - - -def register_and_promote(registry: AdapterRegistry, name: str, config=None) -> None: - registry.register(name, config) - registry.record_weight_update([name]) - - -def test_rid_roundtrip_preserves_names_with_underscores(): - for name in ["a", "adapter_a", "weird__name", "x_y_z"]: - assert parse_adapter(make_rid(name)) == name - - -def test_register_starts_pending_and_push_promotes(): - registry = AdapterRegistry(max_adapters=4) - result = registry.register("A", config={"rm_type": "x"}) - assert result == {"name": "A", "slot": 0} - assert registry.active_adapters() == {} # pending: not sampleable - - registry.record_weight_update(["A"]) - assert registry.active_adapters()["A"].slot == 0 - view = registry.active_adapters()["A"] - assert view.slot == 0 - assert view.config == {"rm_type": "x"} - assert view.version == 1 - - -def test_snapshot_reports_sets_in_registry_vocabulary(): - registry = AdapterRegistry(max_adapters=4) - register_and_promote(registry, "A") - registry.register("B", None) - snapshot = registry.snapshot() - assert set(snapshot["active"]) == {"A"} - assert set(snapshot["pending"]) == {"B"} - assert snapshot["retiring"] == {} - assert snapshot["cleanup"] == [] - assert set(registry.active_adapters()) == {"A"} # only active adapters are sampleable - - -def test_slot_version_is_monotonic_across_slot_reuse(): - registry = AdapterRegistry(max_adapters=2) - register_and_promote(registry, "A") # slot 0, version 1 - registry.record_weight_update(["A"]) # version 2 - registry.deregister("A") - registry.retire_adapters() - registry.free_slot("A") - - registry.register("A2", None) # reuses slot 0 - assert registry.snapshot()["pending"]["A2"].version == 2 # inherits, not reset - registry.record_weight_update(["A2"]) - assert registry.active_adapters()["A2"].version == 3 - - -def test_record_weight_update_only_touches_reported_names(): - registry = AdapterRegistry(max_adapters=4) - register_and_promote(registry, "A") - register_and_promote(registry, "B") - registry.record_weight_update(["A"]) - assert registry.active_adapters()["A"].version == 2 - assert registry.active_adapters()["B"].version == 1 - - -def test_register_name_rejected_until_cleanup_done(): - registry = AdapterRegistry(max_adapters=4) - register_and_promote(registry, "A") - registry.deregister("A") - with pytest.raises(ValueError, match="cleaning up"): - registry.register("A", None) # retiring - registry.retire_adapters() - with pytest.raises(ValueError, match="cleaning up"): - registry.register("A", None) # cleanup - registry.free_slot("A") - assert registry.register("A", None) == {"name": "A", "slot": 0} - - -def test_deregister_retires_but_keeps_serving_until_demoted(): - registry = AdapterRegistry(max_adapters=4) - register_and_promote(registry, "A") - registry.deregister("A") - assert registry.adapter_state("A") == AdapterState.RETIRING - assert "A" in registry.active_adapters() # still sampleable this iteration - assert "A" in registry.snapshot()["retiring"] - assert registry.retire_adapters() == ["A"] - assert registry.active_adapters() == {} - assert registry.adapter_state("A") == AdapterState.CLEANUP - assert registry.retire_adapters() == [] # idempotent - - -# make_config(): rollout_batch_size=4 groups/step, n_samples_per_prompt=4. - - -def test_mark_batch_trained_accumulates_and_steps_on_completion(): - registry = AdapterRegistry(max_adapters=4) - register_and_promote(registry, "A", make_config()) - register_and_promote(registry, "B", make_config()) - - # Two partial batches accumulate; the third completes the adapter batch. - registry.record_batch_adapters(1, {"A": 1, "B": 2}, step_names=[]) - assert registry.mark_batch_trained(1) == [] - assert registry.records["A"].accumulated_groups == 1 - assert registry.records["B"].accumulated_groups == 2 - - registry.record_batch_adapters(2, {"A": 1}, step_names=[]) - assert registry.mark_batch_trained(2) == [] - assert registry.records["A"].accumulated_groups == 2 - - registry.record_batch_adapters(3, {"A": 2, "B": 2}, step_names=["A", "B"]) - assert registry.mark_batch_trained(3) == ["A", "B"] - assert registry.step_count("A") == 1 - assert registry.step_count("B") == 1 - assert registry.records["A"].accumulated_groups == 0 - assert registry.records["B"].accumulated_groups == 0 - - assert registry.mark_batch_trained(3) == [] # record consumed - - -def test_batch_trained_counts_deregistered_adapter_until_freed(): - registry = AdapterRegistry(max_adapters=4) - register_and_promote(registry, "A", make_config()) - registry.record_batch_adapters(3, {"A": 4}, step_names=["A"]) - registry.deregister("A") # deregistered while its batch is training - assert registry.mark_batch_trained(3) == ["A"] - assert registry.step_count("A") == 1 # final ckpt reads this - registry.retire_adapters() - assert registry.step_count("A") == 1 # cleanup record still holds it - registry.free_slot("A") - assert registry.step_count("A") == 0 - - -def test_set_step_on_resume(): - registry = AdapterRegistry(max_adapters=2) - registry.register("A", make_config()) - registry.set_step("A", 40) - registry.record_batch_adapters(1, {"A": 4}, step_names=["A"]) - registry.record_weight_update(["A"]) - registry.mark_batch_trained(1) - assert registry.step_count("A") == 41 - - -def test_num_step_deregisters_on_committed_steps(): - registry = AdapterRegistry(max_adapters=2) - register_and_promote(registry, "A", make_config(num_step=2)) - registry.record_batch_adapters(1, {"A": 4}, step_names=["A"]) - assert registry.mark_batch_trained(1) == ["A"] - assert registry.adapter_state("A") == AdapterState.ACTIVE - - registry.record_batch_adapters(2, {"A": 4}, step_names=["A"]) - assert registry.mark_batch_trained(2) == ["A"] - assert registry.step_count("A") == 2 - assert registry.adapter_state("A") == AdapterState.RETIRING - - -def test_num_step_is_relative_to_resume_step(): - registry = AdapterRegistry(max_adapters=2) - register_and_promote(registry, "A", make_config(num_step=2)) - registry.set_step("A", 40) - - registry.record_batch_adapters(1, {"A": 4}, step_names=["A"]) - registry.mark_batch_trained(1) - assert registry.step_count("A") == 41 - assert registry.adapter_state("A") == AdapterState.ACTIVE - - registry.record_batch_adapters(2, {"A": 4}, step_names=["A"]) - registry.mark_batch_trained(2) - assert registry.step_count("A") == 42 - assert registry.adapter_state("A") == AdapterState.RETIRING - - -def test_min_groups_per_dp_split(): - assert min_groups_per_dp_split(n_samples_per_prompt=4, dp_size=8) == 2 # divisor - assert min_groups_per_dp_split(n_samples_per_prompt=8, dp_size=8) == 1 # equal - assert min_groups_per_dp_split(n_samples_per_prompt=16, dp_size=8) == 1 # multiple - with pytest.raises(ValueError, match="divisor or a multiple"): - min_groups_per_dp_split(n_samples_per_prompt=6, dp_size=8) - - -@pytest.mark.asyncio -async def test_register_resolves_batch_shape_defaults(tmp_path): - backend = make_backend(save=str(tmp_path)) - await backend.register("A", AdapterRunConfig(data=DATA_FILE, rm_type="math")) - config = backend.registry.records["A"].config - assert config.rollout_batch_size == 16 # <- args.rollout_batch_size - assert config.n_samples_per_prompt == 4 # <- args.n_samples_per_prompt - assert config.rank == 32 and config.alpha == 32 - assert config.adapter_global_batch_size == 64 - - -@pytest.mark.asyncio -async def test_register_rejects_bad_batch_shapes(tmp_path): - backend = make_backend(save=str(tmp_path), dp_size=8) - with pytest.raises(ValueError, match="divisor or a multiple"): - await backend.register("B", make_config(n_samples_per_prompt=6, rollout_batch_size=4)) - with pytest.raises(ValueError, match="min_groups_per_dp_split"): - # dp=8, n_samples=4 -> multiple of 2 groups; 3 groups is not - await backend.register("C", make_config(rollout_batch_size=3)) - with pytest.raises(ValueError, match="exceeding"): - await backend.register("D", make_config(rollout_batch_size=128)) # 512 samples > cap 256 - with pytest.raises(ValueError, match="exceeds the allocated maximum rank"): - await backend.register("E", make_config(rank=64)) - with pytest.raises(ValueError, match="positive integer"): - await backend.register("F", make_config(rollout_batch_size=0)) - with pytest.raises(ValueError, match="num_step must be a positive integer"): - await backend.register("G", make_config(num_step=0)) - with pytest.raises(ValueError, match="num_epoch must be a positive integer"): - await backend.register("H", make_config(num_epoch=0)) - # A valid shape registers fine. - await backend.register("OK", make_config(rollout_batch_size=8)) - - -def test_deregister_holds_slot_until_free_slot(): - registry = AdapterRegistry(max_adapters=2) - register_and_promote(registry, "A") # slot 0 - register_and_promote(registry, "B") # slot 1 - registry.deregister("A") - registry.retire_adapters() - assert not registry.free_slots # slot 0 held until cleanup - with pytest.raises(RuntimeError, match="No free adapter slots"): - registry.register("C", None) - registry.free_slot("A") - assert registry.register("C", None) == {"name": "C", "slot": 0} - - -@pytest.mark.asyncio -async def test_free_slot_reaborts_before_releasing_slot(): - """Requests can survive the single retire-time abort (multi-turn groups - submitting between turns, engine tokenizer-adapter batch misses); free_slot must - fire one more abort round before the slot becomes reusable.""" - backend = make_backend() - aborted: list[str] = [] - - async def record_abort(name: str) -> None: - aborted.append(name) - - backend.abort_adapter_requests = record_abort - - register_and_promote(backend.registry, "A") - await backend.deregister("A") - await backend.retire_adapters() - assert aborted == ["A"] - - assert await backend.free_slot("A") == 0 - assert aborted == ["A", "A"] - assert backend.registry.free_slots == {0, 1, 2, 3} - - -@pytest.mark.asyncio -async def test_free_slot_skips_abort_when_not_in_cleanup(): - backend = make_backend() - aborted: list[str] = [] - - async def record_abort(name: str) -> None: - aborted.append(name) - - backend.abort_adapter_requests = record_abort - - register_and_promote(backend.registry, "A") # ACTIVE, not CLEANUP - assert await backend.free_slot("A") == -1 - assert await backend.free_slot("never-registered") == -1 - assert aborted == [] - - -@pytest.mark.asyncio -async def test_custom_backend_validation_rejects(): - class StrictBackend(MultiLoRABackend): - async def validate_adapter(self, name, config): - if not config: - raise ValueError("adapter config is required") - - backend = StrictBackend(make_args(), "http://unused") - with pytest.raises(ValueError, match="config is required"): - await backend.register("A", None) - assert backend.registry.active_adapters() == {} - - result = await backend.register("A", {"rm_type": "x"}) - assert result == {"name": "A", "slot": 0} - - -def test_register_rejects_unsafe_names(): - registry = AdapterRegistry(max_adapters=4) - for bad in ["a/b", "..", "a::b", "a b", ""]: - with pytest.raises(ValueError, match="invalid"): - registry.register(bad, None) - registry.register("ok-name_1.2", None) - - -def test_register_rejects_duplicate_save_dir(tmp_path): - registry = AdapterRegistry(max_adapters=4) - registry.register("A", make_config(save=tmp_path / "x")) - with pytest.raises(ValueError, match="already used by adapter 'A'"): - registry.register("B", make_config(save=tmp_path / "x")) - registry.register("C", make_config(save=tmp_path / "y")) - - -@pytest.mark.asyncio -async def test_save_dir_defaults_under_save_root(tmp_path): - backend = make_backend(save=str(tmp_path)) - await backend.register("A", make_config()) - saved = backend.registry.records["A"].config.save - assert saved == tmp_path / "adapters" / "A" - - -@pytest.mark.asyncio -async def test_explicit_save_dir_wins_over_root(tmp_path): - backend = make_backend(save=str(tmp_path)) - await backend.register("A", make_config(save=tmp_path / "custom")) - assert backend.registry.records["A"].config.save == tmp_path / "custom" - - -@pytest.mark.asyncio -async def test_register_fails_without_any_save_dir(): - backend = make_backend(save=None) - with pytest.raises(ValueError, match="no save dir"): - await backend.register("A", make_config()) - - -@pytest.mark.asyncio -async def test_register_rejects_missing_data_path(tmp_path): - # A nonexistent data path would otherwise kill the shared rollout producer - # thread at the first get_samples, stalling every adapter. - backend = make_backend(save=str(tmp_path)) - with pytest.raises(ValueError, match="data path"): - await backend.register("A", make_config(data=str(tmp_path / "missing.jsonl"))) - - -@pytest.mark.asyncio -async def test_register_rejects_unresolvable_reward_config(tmp_path): - # No adapter rm_type/custom_rm_path and no process-wide --rm-type: every - # sample would fail reward computation and be dropped. - backend = make_backend(save=str(tmp_path)) - with pytest.raises(ValueError, match="reward config"): - await backend.register("A", make_config(rm_type=None)) - - -@pytest.mark.asyncio -async def test_register_accepts_reward_config_from_global_args(tmp_path): - args = make_args(save=str(tmp_path)) - args.rm_type = "math" - backend = MultiLoRABackend(args, "http://unused") - await backend.register("A", make_config(rm_type=None)) - assert backend.registry.records["A"].config.rm_type is None # resolved at reward time via args diff --git a/tests/fast/ray/multi_lora/test_controller_http.py b/tests/fast/ray/multi_lora/test_controller_http.py deleted file mode 100644 index b2700175970..00000000000 --- a/tests/fast/ray/multi_lora/test_controller_http.py +++ /dev/null @@ -1,239 +0,0 @@ -"""HTTP tests for the MultiLoRAHTTPServer control plane with a mock router -(no Ray, no SGLang).""" - -import json -from contextlib import asynccontextmanager -from pathlib import Path -from types import SimpleNamespace - -import aiohttp -import pytest -from aiohttp import web - -from tests.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=60, suite="stage-a-cpu") - -from miles.ray.multi_lora.backend import MultiLoRABackend -from miles.ray.multi_lora.http_server import MultiLoRAHTTPServer -from miles.utils.adapter_config import AdapterRunConfig -from miles.utils.multi_lora import RID_SEPARATOR - - -# Registration validates that the data path exists; the test file itself is a -# convenient always-present stand-in. -DATA_FILE = __file__ - - -def minimal_config(name: str) -> dict: - return {"data": DATA_FILE, "rm_type": "math", "save": f"/tmp/adapters/{name}"} - - -class ControllerHarness: - """Running control plane (backend + API listener) against a mock router - that serves /list_workers and records /abort_request posts.""" - - def __init__(self, session: aiohttp.ClientSession, backend: MultiLoRABackend, srv: MultiLoRAHTTPServer): - self.session = session - self.backend = backend - self.srv = srv - self.aborts: list[dict] = [] - - @property - def api_base(self) -> str: - return f"http://127.0.0.1:{self.srv.actual_api_port}" - - async def api_post(self, path: str, payload: dict) -> tuple[int, dict]: - async with self.session.post(f"{self.api_base}{path}", json=payload) as resp: - return resp.status, await resp.json() - - async def api_get(self, path: str) -> tuple[int, dict, dict]: - async with self.session.get(f"{self.api_base}{path}") as resp: - headers = {k.lower(): v for k, v in resp.headers.items()} - return resp.status, await resp.json(), headers - - async def api_delete(self, path: str) -> tuple[int, dict]: - async with self.session.delete(f"{self.api_base}{path}") as resp: - return resp.status, await resp.json() - - async def register(self, name: str) -> tuple[int, dict]: - status, body = await self.api_post("/adapter_runs", {"name": name, "config": minimal_config(name)}) - # Registered adapters start pending; a weight push promotes them. - self.backend.registry.record_weight_update([name]) - return status, body - - async def deregister(self, name: str) -> tuple[int, dict]: - return await self.api_delete(f"/adapter_runs/{name}") - - async def active(self) -> dict: - _, body, _ = await self.api_get("/adapter_runs") - return { - s["name"]: {"slot": s["slot"], "version": s["version"], "step": s["step"]} - for s in body["adapters"] - if s["state"] == "ACTIVE" - } - - -@asynccontextmanager -async def running_controller(server_cls=MultiLoRAHTTPServer): - router_url = "" - harness: ControllerHarness | None = None - - async def router_handler(request): - if request.path == "/list_workers": - return web.json_response({"urls": [router_url]}) - if request.path == "/abort_request": - harness.aborts.append(json.loads(await request.read())) - return web.json_response({}) - return web.json_response({}, status=404) - - app = web.Application() - app.router.add_resource("/{tail:.*}").add_route("*", router_handler) - runner = web.AppRunner(app) - await runner.setup() - site = web.TCPSite(runner, "127.0.0.1", 0) - await site.start() - router_url = f"http://127.0.0.1:{site._server.sockets[0].getsockname()[1]}" - - backend = MultiLoRABackend( - SimpleNamespace( - multi_lora_n_adapters=4, - save=None, - lora_rank=32, - lora_alpha=32, - rollout_batch_size=16, - n_samples_per_prompt=4, - multi_lora_dp_size=2, - multi_lora_max_adapter_global_batch_size=256, - ), - router_url, - ) - srv = server_cls(backend) - await backend.init() - await srv.start() - try: - async with aiohttp.ClientSession() as session: - harness = ControllerHarness(session, backend, srv) - yield harness - finally: - await srv.stop() - await backend.close() - await runner.cleanup() - - -@pytest.mark.asyncio -async def test_register_and_active_view(): - async with running_controller() as ctl: - status, body = await ctl.register("A") - assert status == 200 - assert body["slot"] == 0 - assert await ctl.active() == {"A": {"slot": 0, "version": 1, "step": 0}} - - -@pytest.mark.asyncio -async def test_deregister_marks_and_retire_adapters_aborts(): - """Deregistration only marks; the driver-synced apply performs the - demotion and fans out one prefix abort per worker.""" - async with running_controller() as ctl: - await ctl.register("A") - status, _ = await ctl.deregister("A") - assert status == 200 - assert ctl.aborts == [] # still serving until the sync point - assert "A" in ctl.backend.registry.active_adapters() - - applied = await ctl.backend.retire_adapters() - assert applied == ["A"] - assert ctl.aborts == [{"rid": f"A{RID_SEPARATOR}", "prefix": True}] - assert ctl.backend.registry.active_adapters() == {} - - -@pytest.mark.asyncio -async def test_register_json_config_validates_to_adapter_config(): - """FastAPI validates the JSON body straight into AdapterRunConfig (422 on bad - payloads).""" - async with running_controller() as ctl: - config = { - "rank": 8, - "data": DATA_FILE, - "save": "/tmp/adapters/A", - "rm_type": "math", - } - status, _ = await ctl.api_post("/adapter_runs", {"name": "A", "config": config}) - assert status == 200 - record = ctl.backend.registry.find("A") - assert isinstance(record.config, AdapterRunConfig) - assert record.config.data == DATA_FILE - assert Path(record.config.save) == Path("/tmp/adapters/A") - assert record.config.input_key == "text" # dataclass default - - status, _ = await ctl.api_post("/adapter_runs", {"name": "B", "config": {"rank": 8}}) - assert status == 422 # data is required - - status, _ = await ctl.api_post("/adapter_runs", {"name": "C"}) - assert status == 400 # exactly one of config/yaml_path - - -@pytest.mark.asyncio -async def test_state_endpoint_reports_lifecycle_and_completed(): - """States walk PENDING -> ACTIVE -> RETIRING -> CLEANUP -> COMPLETED; - unknown names report null; COMPLETED is retained after free_slot.""" - async with running_controller() as ctl: - await ctl.api_post("/adapter_runs", {"name": "A", "config": minimal_config("A")}) - - async def state_of(name): - _, body, _ = await ctl.api_get(f"/adapter_runs/state?names={name}") - return body["states"][name] - - assert await state_of("A") == "PENDING" - ctl.backend.registry.record_weight_update(["A"]) - assert await state_of("A") == "ACTIVE" - - await ctl.deregister("A") - assert await state_of("A") == "RETIRING" - await ctl.backend.retire_adapters() - assert await state_of("A") == "CLEANUP" - - ctl.backend.registry.free_slot("A") - assert await state_of("A") == "COMPLETED" - assert await state_of("nope") is None - - # GET by name serves the completed record; DELETE of unknown 404s. - status, body, _ = await ctl.api_get("/adapter_runs/A") - assert status == 200 and body["state"] == "COMPLETED" - status, _ = await ctl.api_delete("/adapter_runs/nope") - assert status == 404 - - # Re-registration reclaims the name; the completed record is dropped. - status, _ = await ctl.api_post( - "/adapter_runs", - {"name": "A", "config": {"data": DATA_FILE, "rm_type": "math", "save": "/tmp/adapters/A2"}}, - ) - assert status == 200 - assert await state_of("A") == "PENDING" - - -@pytest.mark.asyncio -async def test_custom_server_subclass_adds_routes(): - class CustomServer(MultiLoRAHTTPServer): - def create_app(self): - app = super().create_app() - - @app.middleware("http") - async def tag_response(request, call_next): - response = await call_next(request) - response.headers["X-Custom-Server"] = "1" - return response - - return app - - def add_routes(self, app): - super().add_routes(app) - app.get("/custom_status")(self.custom_status) - - async def custom_status(self): - return {"custom": True, "active": sorted(self.backend.registry.active_adapters())} - - async with running_controller(server_cls=CustomServer) as ctl: - _, body, headers = await ctl.api_get("/custom_status") - assert headers.get("x-custom-server") == "1" - assert body == {"custom": True, "active": []} diff --git a/tests/fast/ray/multi_lora/test_gradient_windows.py b/tests/fast/ray/multi_lora/test_gradient_windows.py new file mode 100644 index 00000000000..f856c8f2c3e --- /dev/null +++ b/tests/fast/ray/multi_lora/test_gradient_windows.py @@ -0,0 +1,65 @@ +from miles.ray.multi_lora.gradient_windows import GradientWindowTracker + +KEY_A = ("A", "reg-1") +KEY_A2 = ("A", "reg-2") +KEY_B = ("B", "reg-1") + + +class TestDirtyFlag: + def test_successful_fb_sets_dirty_and_forward_never_calls_in(self): + tracker = GradientWindowTracker() + tracker.open(KEY_A) + assert not tracker.is_dirty(KEY_A) + tracker.mark_forward_backward_succeeded(KEY_A) + assert tracker.is_dirty(KEY_A) + + def test_committed_step_consumes_the_window(self): + tracker = GradientWindowTracker() + tracker.mark_forward_backward_succeeded(KEY_A) + assert tracker.commit_step(KEY_A) == 1 + assert not tracker.is_dirty(KEY_A) + assert tracker.step_of(KEY_A) == 1 + + def test_executed_optim_without_commit_clears_without_advancing(self): + tracker = GradientWindowTracker() + tracker.mark_forward_backward_succeeded(KEY_A) + tracker.clear_after_executed_optim(KEY_A) + assert not tracker.is_dirty(KEY_A) + assert tracker.step_of(KEY_A) == 0 + + def test_clean_commit_needs_no_prior_fb(self): + tracker = GradientWindowTracker() + assert tracker.commit_step(KEY_A) == 1 + + +class TestStreamIdentity: + def test_registrations_of_the_same_name_are_different_streams(self): + tracker = GradientWindowTracker() + tracker.mark_forward_backward_succeeded(KEY_A) + assert not tracker.is_dirty(KEY_A2) + assert tracker.commit_step(KEY_A2) == 1 + assert tracker.is_dirty(KEY_A) + + def test_streams_are_independent_across_names(self): + tracker = GradientWindowTracker() + tracker.commit_step(KEY_A) + tracker.commit_step(KEY_A) + tracker.mark_forward_backward_succeeded(KEY_B) + assert tracker.step_of(KEY_A) == 2 and not tracker.is_dirty(KEY_A) + assert tracker.step_of(KEY_B) == 0 and tracker.is_dirty(KEY_B) + + def test_close_drops_the_stream_and_queries_go_inert(self): + tracker = GradientWindowTracker() + tracker.commit_step(KEY_A) + tracker.mark_forward_backward_succeeded(KEY_A) + tracker.close(KEY_A) + assert tracker.step_of(KEY_A) == 0 + assert not tracker.is_dirty(KEY_A) + + +class TestRestore: + def test_restore_moves_the_clock(self): + tracker = GradientWindowTracker() + tracker.restore_step(KEY_A, 42) + assert tracker.step_of(KEY_A) == 42 + assert tracker.commit_step(KEY_A) == 43 diff --git a/tests/fast/ray/multi_lora/test_metrics_contract.py b/tests/fast/ray/multi_lora/test_metrics_contract.py new file mode 100644 index 00000000000..1644cede5f7 --- /dev/null +++ b/tests/fast/ray/multi_lora/test_metrics_contract.py @@ -0,0 +1,127 @@ +import math + +import pytest + +from miles.ray.multi_lora.backend import operation_result_metrics + + +def ce_payload(weights_by_sample, masks=None): + samples = [] + for i, weights in enumerate(weights_by_sample): + sample = {"tokens": [1] * (len(weights) + 2), "response_length": len(weights), "loss_weights": weights} + if masks is not None: + sample["loss_mask"] = masks[i] + samples.append(sample) + return {"samples": samples, "loss": {"loss_fn": "cross_entropy"}} + + +class TestMetricsValues: + def test_cross_entropy_matches_hand_sum(self): + payload = ce_payload([[0.5, 2.0], [1.0]]) + logprobs = [[-1.0, -2.0], [-3.0]] + metrics = operation_result_metrics(payload, logprobs) + assert metrics["loss:sum"] == pytest.approx(0.5 * 1.0 + 2.0 * 2.0 + 1.0 * 3.0) + assert metrics["unmasked_tokens:sum"] == 3.0 + + def test_mask_gates_tokens(self): + payload = ce_payload([[1.0, 1.0]], masks=[[1, 0]]) + metrics = operation_result_metrics(payload, [[-1.0, -9.0]]) + assert metrics["loss:sum"] == pytest.approx(1.0) + assert metrics["unmasked_tokens:sum"] == 1.0 + assert metrics["loss_weight:sum"] == pytest.approx(1.0) + + def test_importance_sampling_and_ppo_clip(self): + base = { + "tokens": [1, 1, 1], + "response_length": 2, + "rollout_log_probs": [-1.0, -1.0], + "advantages": [1.0, -2.0], + } + logprobs = [[-0.5, -1.5]] + ratios = [math.exp(0.5), math.exp(-0.5)] + + metrics = operation_result_metrics({"samples": [base], "loss": {"loss_fn": "importance_sampling"}}, logprobs) + assert metrics["loss:sum"] == pytest.approx(-(ratios[0] * 1.0) - (ratios[1] * -2.0)) + + spec = {"loss_fn": "ppo", "loss_fn_config": {"clip_low_threshold": 0.9, "clip_high_threshold": 1.1}} + metrics_ppo = operation_result_metrics({"samples": [base], "loss": spec}, logprobs) + expected = -min(ratios[0] * 1.0, 1.1 * 1.0) - min(ratios[1] * -2.0, 0.9 * -2.0) + assert metrics_ppo["loss:sum"] == pytest.approx(expected) + assert metrics_ppo["loss:sum"] != pytest.approx(metrics["loss:sum"]) + + def test_degenerate_ratio_cannot_overflow_the_recompute(self): + sample = { + "tokens": [1, 1, 1], + "response_length": 2, + "rollout_log_probs": [-1000.0, -1.0], + "advantages": [1.0, 1.0], + } + payload = {"samples": [sample], "loss": {"loss_fn": "importance_sampling"}} + metrics = operation_result_metrics(payload, [[0.0, -1.0]]) + assert math.isfinite(metrics["loss:sum"]) + + def test_sum_metrics_are_chunk_additive(self): + whole = ce_payload([[0.5, 2.0], [1.0, 1.0, 1.0]]) + whole_logprobs = [[-1.0, -2.0], [-3.0, -4.0, -5.0]] + chunks = [ + (ce_payload([[0.5, 2.0]]), [whole_logprobs[0]]), + (ce_payload([[1.0, 1.0, 1.0]]), [whole_logprobs[1]]), + ] + whole_metrics = operation_result_metrics(whole, whole_logprobs) + for key, value in whole_metrics.items(): + assert value == pytest.approx(sum(operation_result_metrics(p, lp)[key] for p, lp in chunks)), key + + +def test_sdk_combiner_merges_our_chunked_metrics(): + helpers = pytest.importorskip("tinker.lib.chunked_fwdbwd_helpers") + types = pytest.importorskip("tinker.types") + + whole = ce_payload([[0.5, 2.0], [1.0, 1.0, 1.0], [3.0]]) + whole_logprobs = [[-1.0, -2.0], [-3.0, -4.0, -5.0], [-0.25]] + chunk_rows = [(0, 2), (2, 3)] + + def chunk_output(start, stop): + payload = {"samples": whole["samples"][start:stop], "loss": whole["loss"]} + metrics = operation_result_metrics(payload, whole_logprobs[start:stop]) + for key in metrics: + assert key.split(":")[1] in helpers.REDUCE_MAP, f"SDK cannot reduce '{key}'" + return types.ForwardBackwardOutput( + loss_fn_output_type="scalar", + metrics=metrics, + loss_fn_outputs=[{} for _ in range(stop - start)], + ) + + combined = helpers.combine_fwd_bwd_output_results([chunk_output(*rows) for rows in chunk_rows]) + whole_metrics = operation_result_metrics(whole, whole_logprobs) + assert combined.metrics["loss:sum"] == pytest.approx(whole_metrics["loss:sum"]) + assert combined.metrics["unmasked_tokens:sum"] == pytest.approx(whole_metrics["unmasked_tokens:sum"]) + assert combined.metrics["loss_weight:sum"] == pytest.approx(whole_metrics["loss_weight:sum"]) + assert len(combined.loss_fn_outputs) == 3 + + +class TestLossWeightSum: + def test_prompt_masked_sft_gets_the_completion_denominator(self): + payload = ce_payload([[0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0]]) + metrics = operation_result_metrics(payload, [[-0.5] * 7]) + assert metrics["unmasked_tokens:sum"] == 7.0 + assert metrics["loss_weight:sum"] == pytest.approx(4.0) + assert metrics["loss:sum"] / metrics["loss_weight:sum"] == pytest.approx(0.5) + + def test_fractional_weights_get_a_weighted_mean_denominator(self): + metrics = operation_result_metrics(ce_payload([[0.0, 0.5, 0.0, 2.0]]), [[-0.5] * 4]) + assert metrics["loss:sum"] == pytest.approx(1.25) + assert metrics["loss_weight:sum"] == pytest.approx(2.5) + + def test_all_zero_weight_chunk_still_reports_the_key(self): + metrics = operation_result_metrics(ce_payload([[0.0, 0.0]]), [[-1.0, -1.0]]) + assert metrics["loss_weight:sum"] == 0.0 + + def test_non_ce_losses_do_not_report_it(self): + sample = { + "tokens": [1, 1, 1], + "response_length": 2, + "rollout_log_probs": [-1.0, -1.0], + "advantages": [1.0, 1.0], + } + payload = {"samples": [sample], "loss": {"loss_fn": "importance_sampling"}} + assert "loss_weight:sum" not in operation_result_metrics(payload, [[-0.5, -1.5]]) diff --git a/tests/fast/ray/multi_lora/test_operations.py b/tests/fast/ray/multi_lora/test_operations.py new file mode 100644 index 00000000000..459bc5d4790 --- /dev/null +++ b/tests/fast/ray/multi_lora/test_operations.py @@ -0,0 +1,536 @@ +import pytest + +from miles.ray.multi_lora.operations import OperationBackpressure, OperationLedger + + +def enqueue(ledger, op_id, ordinal, kind="forward_backward", name="A", reg="ra", payload=None): + return ledger.enqueue(op_id, name, reg, ordinal, kind, payload) + + +class TestArrivalBuffering: + def test_out_of_order_arrival_executes_in_ordinal_order(self): + ledger = OperationLedger() + enqueue(ledger, "op2", 2) + enqueue(ledger, "op3", 3) + assert ledger.claim_data_operation("A", "ra") is None + enqueue(ledger, "op1", 1) + assert ledger.claim_data_operation("A", "ra")["operation_id"] == "op1" + ledger.complete("op1", {}) + assert ledger.claim_data_operation("A", "ra")["operation_id"] == "op2" + + def test_gap_blocks_control_claims_too(self): + ledger = OperationLedger() + enqueue(ledger, "opt2", 2, "optim_step") + assert ledger.claimable_control_tenants() == [] + enqueue(ledger, "fb1", 1) + assert ledger.claim_data_operation("A", "ra")["operation_id"] == "fb1" + + def test_duplicate_ordinal_is_a_conflict(self): + ledger = OperationLedger() + enqueue(ledger, "op1", 1) + with pytest.raises(ValueError, match="already taken"): + enqueue(ledger, "op1b", 1) + + def test_ordinals_start_at_one(self): + ledger = OperationLedger() + with pytest.raises(ValueError, match=">= 1"): + enqueue(ledger, "op0", 0) + + +class TestFingerprintedIdempotency: + def test_identical_retry_returns_the_original(self): + ledger = OperationLedger() + first = enqueue(ledger, "op1", 1, payload={"samples": [1]}) + retry = enqueue(ledger, "op1", 1, payload={"samples": [1]}) + assert retry == first + + def test_same_id_different_payload_is_a_conflict(self): + ledger = OperationLedger() + enqueue(ledger, "op1", 1, payload={"samples": [1]}) + with pytest.raises(ValueError, match="different content"): + enqueue(ledger, "op1", 1, payload={"samples": [2]}) + + def test_same_id_different_kind_is_a_conflict(self): + ledger = OperationLedger() + enqueue(ledger, "op1", 1, "forward_backward") + with pytest.raises(ValueError, match="different content"): + enqueue(ledger, "op1", 1, "optim_step") + + def test_same_id_different_ordinal_is_a_conflict(self): + ledger = OperationLedger() + enqueue(ledger, "op1", 1, payload={"samples": [1]}) + with pytest.raises(ValueError, match="different content"): + enqueue(ledger, "op1", 2, payload={"samples": [1]}) + + +class TestClaimViews: + def test_claims_carry_the_request_payload(self): + ledger = OperationLedger() + enqueue(ledger, "fb", 1, payload={"samples": [{"tokens": [1, 2]}]}) + enqueue(ledger, "optim", 2, "optim_step", payload={"adam_params": {"learning_rate": 2e-4}}) + assert ledger.claim_data_operation("A", "ra")["payload"] == {"samples": [{"tokens": [1, 2]}]} + ledger.complete("fb", {}) + assert ledger.claim_control_operation("A", "ra")["payload"] == {"adam_params": {"learning_rate": 2e-4}} + assert "payload" not in ledger.get("optim") + + +class TestSerialization: + def test_nothing_overtakes_an_open_operation(self): + ledger = OperationLedger() + enqueue(ledger, "fb", 1, "forward_backward") + enqueue(ledger, "optim", 2, "optim_step") + assert ledger.claim_control_operation("A", "ra") is None + claimed = ledger.claim_data_operation("A", "ra") + assert claimed["operation_id"] == "fb" + assert ledger.claim_control_operation("A", "ra") is None + assert ledger.claim_data_operation("A", "ra") is None + ledger.complete("fb", {}) + assert ledger.claim_control_operation("A", "ra")["operation_id"] == "optim" + + def test_control_head_blocks_data_claims(self): + ledger = OperationLedger() + enqueue(ledger, "optim", 1, "optim_step") + enqueue(ledger, "fb", 2, "forward_backward") + assert ledger.claim_data_operation("A", "ra") is None + assert ("A", "ra") in ledger.claimable_control_tenants() + ledger.claim_control_operation("A", "ra") + ledger.complete("optim", {}) + assert ledger.claim_data_operation("A", "ra")["operation_id"] == "fb" + + def test_control_claim_kind_filter(self): + ledger = OperationLedger() + enqueue(ledger, "save", 1, "save_state") + assert ledger.claim_control_operation("A", "ra", kinds=("optim_step",)) is None + assert ledger.claim_control_operation("A", "ra", kinds=("save_state",))["operation_id"] == "save" + + def test_registrations_are_independent(self): + ledger = OperationLedger() + enqueue(ledger, "a1", 1, name="A", reg="ra") + enqueue(ledger, "b1", 1, name="B", reg="rb") + assert ledger.claim_data_operation("A", "ra")["operation_id"] == "a1" + assert ledger.claim_data_operation("B", "rb")["operation_id"] == "b1" + + +class TestPoisonedWindow: + def fail_fb(self, ledger, op_id, ordinal, category="user"): + enqueue(ledger, op_id, ordinal, "forward_backward") + claimed = ledger.claim_data_operation("A", "ra") + assert claimed["operation_id"] == op_id + ledger.fail(op_id, "bad chunk", category) + + def complete_fb(self, ledger, op_id, ordinal): + enqueue(ledger, op_id, ordinal, "forward_backward") + ledger.claim_data_operation("A", "ra") + ledger.complete(op_id, {}) + + def test_failed_chunk_poisons_and_success_does_not(self): + ledger = OperationLedger() + self.complete_fb(ledger, "fb1", 1) + assert ledger.poisoned_window_blocker("A", "ra", 2) is None + self.fail_fb(ledger, "fb2", 2) + blocker = ledger.poisoned_window_blocker("A", "ra", 3) + assert blocker is not None and "ordinal 2" in blocker + + def test_executed_optim_delimits_the_window(self): + ledger = OperationLedger() + self.fail_fb(ledger, "fb1", 1) + enqueue(ledger, "opt2", 2, "optim_step") + ledger.claim_control_operation("A", "ra") + ledger.fail("opt2", "window poisoned", "user") + assert ledger.poisoned_window_blocker("A", "ra", 4) is not None + ledger.mark_window_consumed("opt2") + self.complete_fb(ledger, "fb3", 3) + assert ledger.poisoned_window_blocker("A", "ra", 4) is None + + def test_cancelled_optim_is_no_delimiter_and_cancelled_fb_poisons(self): + ledger = OperationLedger() + self.fail_fb(ledger, "fb1", 1) + enqueue(ledger, "opt2", 2, "optim_step") + ledger.cancel("opt2") + assert ledger.poisoned_window_blocker("A", "ra", 3) is not None + + enqueue(ledger, "fb3", 3, "forward_backward") + ledger.cancel("fb3") + blocker = ledger.poisoned_window_blocker("A", "ra", 4) + assert blocker is not None and "ordinal 3" in blocker + + def test_failed_forward_does_not_poison(self): + ledger = OperationLedger() + enqueue(ledger, "fw1", 1, "forward") + ledger.claim_data_operation("A", "ra") + ledger.fail("fw1", "bad forward", "user") + assert ledger.poisoned_window_blocker("A", "ra", 2) is None + + +class TestTerminals: + def test_cancel_applies_only_to_queued_and_keeps_contiguity(self): + ledger = OperationLedger() + enqueue(ledger, "op1", 1) + enqueue(ledger, "op2", 2) + assert ledger.cancel("op2")["state"] == "CANCELLED" + ledger.claim_data_operation("A", "ra") + with pytest.raises(ValueError, match="only QUEUED"): + ledger.cancel("op1") + ledger.complete("op1", {}) + enqueue(ledger, "op3", 3) + assert ledger.claim_data_operation("A", "ra")["operation_id"] == "op3" + + def test_fail_records_error_and_category(self): + ledger = OperationLedger() + enqueue(ledger, "op1", 1) + ledger.claim_data_operation("A", "ra") + ledger.fail("op1", "bad payload", "user") + view = ledger.get("op1") + assert view["state"] == "FAILED" and view["error_category"] == "user" + + def test_double_terminal_is_rejected(self): + ledger = OperationLedger() + enqueue(ledger, "op1", 1) + ledger.claim_data_operation("A", "ra") + ledger.complete("op1", {}) + with pytest.raises(ValueError, match="already terminal"): + ledger.fail("op1", "late failure") + + +class TestBackpressureAndRetention: + def test_pending_depth_backpressure(self): + ledger = OperationLedger(max_pending=2) + enqueue(ledger, "op1", 1) + enqueue(ledger, "op2", 2) + with pytest.raises(OperationBackpressure): + enqueue(ledger, "op3", 3) + + def test_gap_filler_bypasses_the_pending_cap(self): + ledger = OperationLedger(max_pending=2) + enqueue(ledger, "op2", 2) + enqueue(ledger, "op3", 3) + assert ledger.claim_data_operation("A", "ra") is None + enqueue(ledger, "op1", 1) + assert ledger.claim_data_operation("A", "ra")["operation_id"] == "op1" + with pytest.raises(OperationBackpressure): + enqueue(ledger, "op4", 4) + + def test_ack_releases_the_payload_and_result(self): + ledger = OperationLedger() + enqueue(ledger, "op1", 1, payload={"samples": ["x" * 64]}) + ledger.claim_data_operation("A", "ra") + ledger.complete("op1", {"logprobs": [[0.0] * 64]}) + ledger.ack("op1") + residue = ledger.queues[("A", "ra")].by_ordinal[1] + assert residue.payload == {} and residue.result is None + + def test_unacked_results_backpressure_and_ack_release(self): + ledger = OperationLedger(max_unacked_results=1) + enqueue(ledger, "op1", 1) + ledger.claim_data_operation("A", "ra") + ledger.complete("op1", {"ok": True}) + with pytest.raises(OperationBackpressure, match="unacknowledged"): + enqueue(ledger, "op2", 2) + ledger.ack("op1") + enqueue(ledger, "op2", 2) + assert ledger.claim_data_operation("A", "ra")["operation_id"] == "op2" + + def test_ack_drops_only_terminal_records(self): + ledger = OperationLedger() + enqueue(ledger, "op1", 1) + with pytest.raises(ValueError, match="ack applies to terminal"): + ledger.ack("op1") + ledger.claim_data_operation("A", "ra") + ledger.complete("op1", {}) + ledger.ack("op1") + assert ledger.get("op1") is None + ledger.ack("op1") + + +class TestFencing: + def test_fence_fails_open_ops_and_refuses_new_ones(self): + ledger = OperationLedger() + enqueue(ledger, "done", 1) + ledger.claim_data_operation("A", "ra") + ledger.complete("done", {"kept": True}) + enqueue(ledger, "pending", 2) + assert ledger.fence("A", "ra") == ["pending"] + assert ledger.get("pending")["state"] == "FAILED" + assert ledger.get("pending")["error_category"] == "user" + assert ledger.get("done")["result"] == {"kept": True} + with pytest.raises(ValueError, match="fenced"): + enqueue(ledger, "late", 3) + + def test_a_new_registration_of_the_same_name_starts_fresh(self): + ledger = OperationLedger() + enqueue(ledger, "old", 1, name="A", reg="ra") + ledger.fence("A", "ra") + fresh = enqueue(ledger, "new", 1, name="A", reg="rb") + assert fresh["state"] == "QUEUED" + + +class TestRecordRejected: + def test_rejected_ordinal_keeps_the_sequence_gap_free(self): + ledger = OperationLedger() + enqueue(ledger, "op1", 1) + rejected = ledger.record_rejected("op2", "A", "ra", 2, "optim_step", {"adam_params": {}}, "bad params") + assert rejected["state"] == "FAILED" + assert rejected["error_category"] == "user" + enqueue(ledger, "op3", 3) + assert ledger.claim_data_operation("A", "ra")["operation_id"] == "op1" + ledger.complete("op1", {}) + assert ledger.claim_data_operation("A", "ra")["operation_id"] == "op3" + + def test_identical_retry_replays_the_terminal_record(self): + ledger = OperationLedger() + first = ledger.record_rejected("op1", "A", "ra", 1, "forward", {"samples": []}, "empty") + again = ledger.record_rejected("op1", "A", "ra", 1, "forward", {"samples": []}, "empty") + assert again == first + + def test_different_payload_at_the_same_id_is_a_conflict(self): + ledger = OperationLedger() + ledger.record_rejected("op1", "A", "ra", 1, "forward", {"samples": []}, "empty") + with pytest.raises(ValueError, match="different content"): + ledger.record_rejected("op1", "A", "ra", 1, "forward", {"samples": [1]}, "empty") + + def test_taken_ordinal_and_fence_still_refuse(self): + ledger = OperationLedger() + enqueue(ledger, "op1", 1) + with pytest.raises(ValueError, match="already taken"): + ledger.record_rejected("op1b", "A", "ra", 1, "forward", {}, "x") + ledger.fence("A", "ra") + with pytest.raises(ValueError, match="fenced"): + ledger.record_rejected("op2", "A", "ra", 2, "forward", {}, "x") + + def test_rejected_flood_hits_the_unacked_results_budget(self): + ledger = OperationLedger(max_unacked_results=8) + accepted = 0 + for i in range(1, 1001): + try: + ledger.record_rejected(f"op{i}", "A", "ra", i, "forward_backward", {"i": i}, "bad") + accepted += 1 + except OperationBackpressure: + break + assert accepted == 8 + assert ledger.queues[("A", "ra")].unacked_terminal_count() == 8 + ledger.ack("op1") + assert ledger.record_rejected("op9", "A", "ra", 9, "forward_backward", {"i": 9}, "bad")["state"] == "FAILED" + + def test_rejected_hole_filler_bypasses_the_unacked_budget(self): + ledger = OperationLedger(max_unacked_results=1) + enqueue(ledger, "fb1", 1) + enqueue(ledger, "fb3", 3) + ledger.claim_data_operation("A", "ra") + ledger.fail("fb1", "boom", "user") + with pytest.raises(OperationBackpressure): + ledger.record_rejected("tail", "A", "ra", 4, "forward_backward", {}, "bad") + ledger.record_rejected("hole", "A", "ra", 2, "forward_backward", {}, "bad") + assert ledger.claim_data_operation("A", "ra")["operation_id"] == "fb3" + + def test_born_terminal_optim_is_no_window_delimiter(self): + ledger = OperationLedger() + enqueue(ledger, "fb1", 1) + ledger.claim_data_operation("A", "ra") + ledger.fail("fb1", "bad chunk", "user") + ledger.record_rejected("opt2", "A", "ra", 2, "optim_step", {"adam_params": {"beta1": 9}}, "bad params") + assert ledger.poisoned_window_blocker("A", "ra", 3) is not None + + def test_rejection_bypasses_backpressure_like_a_hole_filler(self): + ledger = OperationLedger(max_pending=1) + enqueue(ledger, "op1", 1) + with pytest.raises(OperationBackpressure): + enqueue(ledger, "op3", 3) + assert ledger.record_rejected("op2", "A", "ra", 2, "forward", {}, "x")["state"] == "FAILED" + + def test_rejected_record_is_ackable(self): + ledger = OperationLedger() + ledger.record_rejected("op1", "A", "ra", 1, "forward", {}, "x") + ledger.ack("op1") + assert ledger.get("op1") is None + enqueue(ledger, "op2", 2) + assert ledger.claim_data_operation("A", "ra")["operation_id"] == "op2" + + +class Clock: + def __init__(self, now: float = 1000.0) -> None: + self.now = now + + def __call__(self) -> float: + return self.now + + +class TestGapTimeout: + def gapped(self, timeout=10.0): + clock = Clock() + ledger = OperationLedger(gap_timeout=timeout, time_fn=clock) + enqueue(ledger, "fb1", 1) + ledger.claim_data_operation("A", "ra") + ledger.complete("fb1", {}) + enqueue(ledger, "opt3", 3, "optim_step") + ledger.sweep_gap_timeouts() + return ledger, clock + + def test_stall_is_observable_before_expiry(self): + ledger, clock = self.gapped() + clock.now += 4 + [stall] = ledger.gap_stalls() + assert stall["missing_ordinal"] == 2 and stall["blocked_operations"] == 1 + assert stall["stalled_for"] == pytest.approx(4.0) + assert ledger.sweep_gap_timeouts() == [] + assert ledger.get("opt3")["state"] == "QUEUED" + + def test_legit_out_of_order_fill_beats_the_timeout(self): + ledger, clock = self.gapped() + clock.now += 9 + enqueue(ledger, "fb2", 2) + assert ledger.sweep_gap_timeouts() == [] + assert ledger.gap_stalls() == [] + assert ledger.claim_data_operation("A", "ra")["operation_id"] == "fb2" + + def test_expiry_fails_blocked_ops_typed_and_seals_the_hole(self): + ledger, clock = self.gapped() + clock.now += 11 + [event] = ledger.sweep_gap_timeouts() + assert event["missing_ordinal"] == 2 + assert event["sealed_ordinals"] == [2] and event["failed_operations"] == ["opt3"] + view = ledger.get("opt3") + assert view["state"] == "FAILED" and view["error_category"] == "user" + assert "missing ordinal 2" in view["error"] and "resubmit" in view["error"] + with pytest.raises(ValueError, match="already taken"): + enqueue(ledger, "late2", 2, "optim_step") + enqueue(ledger, "opt4", 4, "optim_step") + assert ledger.claimable_control_tenants() == [("A", "ra")] + assert ledger.claim_control_operation("A", "ra")["operation_id"] == "opt4" + + def test_expiry_seals_every_hole_below_the_arrived_tail(self): + clock = Clock() + ledger = OperationLedger(gap_timeout=10.0, time_fn=clock) + enqueue(ledger, "fb1", 1) + ledger.claim_data_operation("A", "ra") + ledger.complete("fb1", {}) + enqueue(ledger, "fb3", 3) + enqueue(ledger, "fb5", 5) + ledger.sweep_gap_timeouts() + clock.now += 11 + [event] = ledger.sweep_gap_timeouts() + assert event["sealed_ordinals"] == [2, 4] + assert sorted(event["failed_operations"]) == ["fb3", "fb5"] + enqueue(ledger, "fb6", 6) + assert ledger.claim_data_operation("A", "ra")["operation_id"] == "fb6" + + def test_sealed_hole_is_poison_neutral_and_no_delimiter(self): + ledger, clock = self.gapped() + clock.now += 11 + ledger.sweep_gap_timeouts() + enqueue(ledger, "opt4", 4, "optim_step") + assert ledger.poisoned_window_blocker("A", "ra", 4) is None + + def test_gap_failed_forward_backward_still_poisons_its_window(self): + clock = Clock() + ledger = OperationLedger(gap_timeout=10.0, time_fn=clock) + enqueue(ledger, "fb1", 1) + ledger.claim_data_operation("A", "ra") + ledger.complete("fb1", {}) + enqueue(ledger, "fb3", 3) + ledger.sweep_gap_timeouts() + clock.now += 11 + [event] = ledger.sweep_gap_timeouts() + assert event["failed_operations"] == ["fb3"] + enqueue(ledger, "opt4", 4, "optim_step") + blocker = ledger.poisoned_window_blocker("A", "ra", 4) + assert blocker is not None and "ordinal 3" in blocker + + def test_disabled_timeout_reports_but_never_expires(self): + ledger, clock = self.gapped(timeout=0) + clock.now += 10_000 + assert ledger.sweep_gap_timeouts() == [] + [stall] = ledger.gap_stalls() + assert stall["missing_ordinal"] == 2 and stall["stalled_for"] == pytest.approx(10_000.0) + assert ledger.get("opt3")["state"] == "QUEUED" + + def test_a_new_hole_restarts_the_stall_clock(self): + ledger, clock = self.gapped() + clock.now += 9 + enqueue(ledger, "fb2", 2) + for op_id in ("fb2", "opt3"): + if op_id == "fb2": + ledger.claim_data_operation("A", "ra") + else: + ledger.claim_control_operation("A", "ra") + ledger.complete(op_id, {}) + enqueue(ledger, "fb5", 5) + assert ledger.sweep_gap_timeouts() == [] + clock.now += 9 + assert ledger.sweep_gap_timeouts() == [] + clock.now += 2 + [event] = ledger.sweep_gap_timeouts() + assert event["missing_ordinal"] == 4 + + def test_fenced_queue_never_stalls(self): + ledger, clock = self.gapped() + ledger.fence("A", "ra") + clock.now += 100 + assert ledger.gap_stalls() == [] and ledger.sweep_gap_timeouts() == [] + + +class TestClaimedTimeout: + def claimed(self, ttl=100.0): + clock = Clock() + ledger = OperationLedger(gap_timeout=10.0, claimed_ttl=ttl, time_fn=clock) + enqueue(ledger, "fb1", 1) + ledger.claim_data_operation("A", "ra") + return ledger, clock + + def test_over_age_claimed_is_reported_with_its_age(self): + ledger, clock = self.claimed() + clock.now += 101 + [view] = ledger.claimed_timeouts() + assert view["operation_id"] == "fb1" and view["state"] == "CLAIMED" + assert view["claimed_age"] == pytest.approx(101.0) + + def test_younger_claimed_is_untouched(self): + ledger, clock = self.claimed() + clock.now += 99 + assert ledger.claimed_timeouts() == [] + assert ledger.get("fb1")["state"] == "CLAIMED" + + def test_control_claims_age_too(self): + ledger, clock = self.claimed() + ledger.complete("fb1", {}) + enqueue(ledger, "opt2", 2, "optim_step") + ledger.claim_control_operation("A", "ra") + clock.now += 101 + [view] = ledger.claimed_timeouts() + assert view["operation_id"] == "opt2" + + def test_disabled_ttl_never_reports(self): + ledger, clock = self.claimed(ttl=0) + clock.now += 1_000_000 + assert ledger.claimed_timeouts() == [] + assert ledger.get("fb1")["state"] == "CLAIMED" + + def test_queued_operations_age_by_gap_rules_only(self): + clock = Clock() + ledger = OperationLedger(claimed_ttl=100.0, time_fn=clock) + enqueue(ledger, "fb1", 1) + clock.now += 1000 + assert ledger.claimed_timeouts() == [] + assert ledger.get("fb1")["state"] == "QUEUED" + + def test_a_claimed_head_is_not_a_gap_stall(self): + ledger, clock = self.claimed() + clock.now += 1000 + assert ledger.gap_stalls() == [] and ledger.sweep_gap_timeouts() == [] + + +class TestTenantEviction: + def test_drop_tenant_purges_the_dead_registration_only(self): + ledger = OperationLedger() + enqueue(ledger, "old1", 1, payload={"samples": ["x" * 64]}, name="A", reg="ra") + ledger.complete("old1", {"kept": True}) + ledger.fence("A", "ra") + assert ledger.by_id["old1"].payload == {} + assert ledger.get("old1")["result"] == {"kept": True} + enqueue(ledger, "young1", 1, name="B", reg="rb") + ledger.complete("young1", {}) + ledger.fence("B", "rb") + ledger.drop_tenant("A", "ra") + assert ledger.get("old1") is None and ("A", "ra") not in ledger.queues + assert not any(op.tenant == ("A", "ra") for op in ledger.by_id.values()) + assert ledger.get("young1")["state"] == "SUCCEEDED" + ledger.drop_tenant("A", "ra") diff --git a/tests/fast/ray/multi_lora/test_registry.py b/tests/fast/ray/multi_lora/test_registry.py new file mode 100644 index 00000000000..897be5e4a10 --- /dev/null +++ b/tests/fast/ray/multi_lora/test_registry.py @@ -0,0 +1,156 @@ +import pytest + +from miles.ray.multi_lora.config import AdapterRunConfig +from miles.ray.multi_lora.registry import AdapterRegistry, AdapterState +from miles.ray.multi_lora.slot_pool import SlotPool + + +class TestSlotPool: + def test_binds_lowest_free_and_queues_when_full(self): + pool = SlotPool(2) + assert pool.bind_immediately(("a", "r1")) == 0 + assert pool.bind_immediately(("b", "r1")) == 1 + assert pool.bind_immediately(("c", "r1")) is None + assert pool.release(("a", "r1")) == 0 + assert pool.bind_immediately(("c", "r1")) == 0 + + def test_release_clears_pins(self): + pool = SlotPool(1) + pool.bind_immediately(("a", "r1")) + pool.pin(("a", "r1"), "dirty-grads") + assert pool.is_pinned(("a", "r1"), "dirty-grads") + pool.release(("a", "r1")) + pool.bind_immediately(("b", "r1")) + assert not pool.is_pinned(("b", "r1"), "dirty-grads") + + def test_occupied_ids(self): + pool = SlotPool(3) + pool.bind_immediately(("a", "r1")) + pool.bind_immediately(("b", "r1")) + assert pool.occupied_slot_ids() == [0, 1] + assert pool.free_slot_ids() == {2} + + +def config(**overrides) -> AdapterRunConfig: + return AdapterRunConfig(**overrides) + + +def register_ready(registry, name): + registry.register(name, config()) + registry.mark_ready([name]) + return registry.find(name) + + +class TestLifecycle: + def test_ready_comes_from_trainer_load_not_from_a_publish(self): + registry = AdapterRegistry(2) + registry.register("A", config()) + assert registry.find("A").state is AdapterState.PENDING + registry.record_weight_update(["A"]) + assert registry.find("A").state is AdapterState.PENDING + assert registry.find("A").serving_version == 1 + registry.mark_ready(["A"]) + assert registry.find("A").state is AdapterState.READY + + def test_unbound_pending_cannot_become_ready(self): + registry = AdapterRegistry(1) + registry.register("A", config()) + registry.register("B", config()) + assert registry.find("B").slot is None + registry.mark_ready(["B"]) + assert registry.find("B").state is AdapterState.PENDING + + def test_queue_drains_at_retirement(self): + registry = AdapterRegistry(1) + registry.register("A", config()) + registry.register("B", config()) + registry.deregister("A") + assert registry.retire_adapters() == ["A"] + assert registry.free_slot("A") == 0 + assert registry.bootstrap_pending() == ["B"] + assert registry.find("B").slot == 0 + + def test_queue_drains_in_arrival_order_not_name_order(self): + registry = AdapterRegistry(1) + registry.register("A", config()) + registry.register("Z", config()) + registry.register("B", config()) + registry.deregister("A") + registry.retire_adapters() + registry.free_slot("A") + assert registry.bootstrap_pending() == ["Z"] + registry.deregister("Z") + registry.retire_adapters() + registry.free_slot("Z") + assert registry.bootstrap_pending() == ["B"] + + def test_duplicate_and_invalid_names_rejected(self): + registry = AdapterRegistry(2) + registry.register("A", config()) + with pytest.raises(ValueError, match="already registered"): + registry.register("A", config()) + with pytest.raises(ValueError, match="invalid"): + registry.register("bad name", config()) + + def test_save_dir_conflict_rejected(self): + registry = AdapterRegistry(2) + registry.register("A", config(save="/tmp/x")) + with pytest.raises(ValueError, match="already used"): + registry.register("B", config(save="/tmp/x")) + + +class TestClocksAndPins: + def test_committed_step_mirrors_clock_and_releases_the_pin(self): + registry = AdapterRegistry(1) + record = register_ready(registry, "A") + registry.mark_accumulated(["A"]) + assert registry.is_dirty("A") + registry.on_step_committed("A", record.registration_id, 1) + assert not registry.is_dirty("A") + assert record.step == 1 + + def test_hook_ignores_a_stale_registration(self): + registry = AdapterRegistry(1) + record = register_ready(registry, "A") + registry.on_step_committed("A", "not-the-registration", 7) + assert record.step == 0 + + def test_veto_path_clears_dirty_without_advancing(self): + registry = AdapterRegistry(1) + record = register_ready(registry, "A") + registry.mark_accumulated(["A"]) + registry.clear_dirty("A") + assert not registry.is_dirty("A") + assert record.step == 0 + + def test_num_step_bound_deregisters(self): + registry = AdapterRegistry(1) + registry.register("A", config(num_step=2)) + registry.mark_ready(["A"]) + rid = registry.find("A").registration_id + registry.on_step_committed("A", rid, 1) + assert registry.find("A").state is AdapterState.READY + registry.on_step_committed("A", rid, 2) + assert registry.records["A"].state is AdapterState.RETIRING + + def test_set_step_repositions_baseline(self): + registry = AdapterRegistry(1) + registry.register("A", config(num_step=2)) + registry.mark_ready(["A"]) + rid = registry.find("A").registration_id + registry.set_step("A", 10) + registry.on_step_committed("A", rid, 11) + assert registry.records["A"].state is AdapterState.READY + registry.on_step_committed("A", rid, 12) + assert registry.records["A"].state is AdapterState.RETIRING + + +class TestViews: + def test_snapshot_vocabulary(self): + registry = AdapterRegistry(2) + register_ready(registry, "A") + registry.register("B", config()) + snap = registry.snapshot() + assert list(snap["ready"]) == ["A"] and list(snap["pending"]) == ["B"] + assert snap["ready"]["A"].registration_id + assert registry.ready_adapters()["A"].slot == 0 diff --git a/tests/fast/ray/multi_lora/test_residency.py b/tests/fast/ray/multi_lora/test_residency.py new file mode 100644 index 00000000000..8121858b182 --- /dev/null +++ b/tests/fast/ray/multi_lora/test_residency.py @@ -0,0 +1,180 @@ +import asyncio +import copy +from types import SimpleNamespace + +import pytest + +from miles.ray.multi_lora.backend import MultiLoraOperationBackend +from miles.ray.multi_lora.config import AdapterRunConfig +from miles.ray.multi_lora.registry import AdapterRegistry +from miles.ray.multi_lora.residency import FixedSlotResidency, ResidentBinding, lease_from_metadata, lease_to_metadata + + +def make_registry(n=1) -> AdapterRegistry: + return AdapterRegistry(n) + + +def register_ready(registry, name) -> tuple[str, str]: + registry.register(name, AdapterRunConfig()) + registry.mark_ready([name]) + return (name, registry.find(name).registration_id) + + +def make_backend(max_adapters=1) -> MultiLoraOperationBackend: + args = SimpleNamespace( + multi_lora_n_adapters=max_adapters, + save="/tmp/tinker-test-save", + lora_rank=32, + lora_alpha=64, + hf_checkpoint="Qwen/Qwen3-0.6B", + ) + return MultiLoraOperationBackend(args, "http://unused") + + +def fb_payload(): + return { + "samples": [{"tokens": [1, 2, 3, 4], "response_length": 2, "loss_mask": [1, 1], "loss_weights": [1.0, 1.0]}], + "loss": {"loss_fn": "cross_entropy"}, + } + + +class TestBindingFor: + def test_exact_ready_with_slot_only(self): + registry = make_registry(2) + key = register_ready(registry, "A") + residency = FixedSlotResidency(registry) + assert residency.binding_for(key) == ResidentBinding(registration_key=key, training_slot=0) + + def test_every_other_state_is_rejected_without_mutation(self): + registry = make_registry(1) + residency = FixedSlotResidency(registry) + + registry.register("A", AdapterRunConfig()) + key_a = ("A", registry.find("A").registration_id) + assert residency.binding_for(key_a) is None + + registry.register("B", AdapterRunConfig()) + key_b = ("B", registry.find("B").registration_id) + assert residency.binding_for(key_b) is None + + registry.mark_ready(["A"]) + assert residency.binding_for(("A", "not-the-registration")) is None + + registry.deregister("A") + assert residency.binding_for(key_a) is None + + registry.retire_adapters() + assert residency.binding_for(key_a) is None + + assert registry.records["A"].slot == 0 + assert registry.records["B"].slot is None + before = copy.deepcopy(registry.snapshot()) + residency.binding_for(key_a) + assert registry.snapshot() == before + + +class TestClaimAndBind: + def test_data_claim_carries_the_binding(self): + backend = make_backend() + asyncio.run(backend.register("A", AdapterRunConfig())) + backend.registry.mark_ready(["A"]) + rid = backend.registry.find("A").registration_id + backend.enqueue_operation("A", "fb1", 1, "forward_backward", fb_payload()) + claim = backend.claim_data_operation("A", rid) + assert claim["operation_id"] == "fb1" + assert claim["binding"] == ResidentBinding(registration_key=("A", rid), training_slot=0) + + def test_unbound_pending_is_never_claimed_and_head_stays_queued(self): + backend = make_backend(max_adapters=1) + asyncio.run(backend.register("A", AdapterRunConfig())) + backend.registry.mark_ready(["A"]) + asyncio.run(backend.register("B", AdapterRunConfig())) + rid_b = backend.registry.find("B").registration_id + backend.enqueue_operation("B", "b-fb1", 1, "forward_backward", fb_payload()) + + assert backend.claim_data_operation("B", rid_b) is None + assert backend.operations.get("b-fb1")["state"] == "QUEUED" + + backend.registry.deregister("A") + backend.registry.retire_adapters() + backend.registry.free_slot("A") + assert backend.registry.bootstrap_pending() == ["B"] + backend.registry.mark_ready(["B"]) + claim = backend.claim_data_operation("B", rid_b) + assert claim["operation_id"] == "b-fb1" + assert claim["binding"].training_slot == 0 + + def test_control_claims_still_require_ready_and_slot(self): + backend = make_backend(max_adapters=1) + asyncio.run(backend.register("A", AdapterRunConfig())) + backend.registry.mark_ready(["A"]) + asyncio.run(backend.register("B", AdapterRunConfig())) + backend.enqueue_operation("B", "b-opt1", 1, "optim_step") + assert backend.claim_ready_control_operations() == {"operations": [], "lease": None} + assert backend.operations.get("b-opt1")["state"] == "QUEUED" + + +class TestBatchLease: + def test_acquire_release_roundtrip(self): + registry = make_registry(2) + key_a = register_ready(registry, "A") + key_b = register_ready(registry, "B") + residency = FixedSlotResidency(registry) + lease = residency.acquire_batch( + ( + ("op-A", residency.binding_for(key_a)), + ("op-B", residency.binding_for(key_b)), + ) + ) + assert lease.binding_of("op-A").training_slot == 0 + assert lease.binding_of("op-B").training_slot == 1 + assert lease.binding_of("op-unknown") is None + before = copy.deepcopy(registry.snapshot()) + residency.release_batch(lease) + assert registry.snapshot() == before + assert lease_from_metadata(lease_to_metadata(lease)) == lease + + def test_retiring_after_claim_keeps_the_receipt_valid(self): + registry = make_registry(1) + key = register_ready(registry, "A") + residency = FixedSlotResidency(registry) + binding = residency.binding_for(key) + + registry.deregister("A") + lease = residency.acquire_batch((("op-A", binding),)) + assert lease.binding_of("op-A") is binding + + registry.retire_adapters() + registry.free_slot("A") + with pytest.raises(ValueError, match="no longer owns trainer slot"): + residency.acquire_batch((("op-A", binding),)) + + def test_wrong_slot_or_foreign_registration_is_refused(self): + registry = make_registry(2) + key = register_ready(registry, "A") + residency = FixedSlotResidency(registry) + with pytest.raises(ValueError, match="no longer owns"): + residency.acquire_batch((("op-A", ResidentBinding(registration_key=key, training_slot=1)),)) + with pytest.raises(ValueError, match="no longer owns"): + residency.acquire_batch( + (("op-A", ResidentBinding(registration_key=("A", "stale-registration"), training_slot=0)),) + ) + + +class TestTrainerLocalValidation: + def test_lease_must_match_locally_loaded_adapters(self): + from miles.backends.megatron_utils.api_backends.multi_lora.trainer import validate_batch_lease + + loaded = {"A": SimpleNamespace(registration_id="r-A", slot=0)} + good = {"batch_execution_lease": {"dispatch_id": "d", "bindings_by_operation": [["op-A", ["A", "r-A", 0]]]}} + validate_batch_lease(good, loaded) + + for name, rid, slot in [("A", "r-A", 1), ("A", "r-OLD", 0), ("Z", "r-Z", 0)]: + bad = { + "batch_execution_lease": {"dispatch_id": "d", "bindings_by_operation": [["op-A", [name, rid, slot]]]} + } + with pytest.raises(RuntimeError, match="does not match"): + validate_batch_lease(bad, loaded) + + with pytest.raises(RuntimeError, match="no execution lease"): + validate_batch_lease({}, loaded) diff --git a/tests/fast/ray/rollout/conftest.py b/tests/fast/ray/rollout/conftest.py index 7da77908280..dc83c1ed81a 100644 --- a/tests/fast/ray/rollout/conftest.py +++ b/tests/fast/ray/rollout/conftest.py @@ -296,12 +296,3 @@ def _alloc(start_port: int = 15000, consecutive: int = 1): e._get_current_node_ip_and_free_port.remote.side_effect = lambda **kw: _alloc(**kw) return e - - -@pytest.fixture -def patch_ray_get(monkeypatch): - """Make ``ray.get(remote_call(...))`` return the MagicMock's value directly, - so allocator tests don't need a real Ray cluster.""" - import miles.ray.rollout.addr_allocator as mod - - monkeypatch.setattr(mod.ray, "get", lambda x: x) diff --git a/tests/fast/ray/rollout/real_ray/test_rollout_manager.py b/tests/fast/ray/rollout/real_ray/test_rollout_manager.py index 9584bc21250..3c32994701e 100644 --- a/tests/fast/ray/rollout/real_ray/test_rollout_manager.py +++ b/tests/fast/ray/rollout/real_ray/test_rollout_manager.py @@ -20,7 +20,13 @@ class behind ``@ray.remote``) — that keeps the manager in the test process so from tests.fast.ray.rollout.conftest import make_args, make_samples_grouped from miles.ray.rollout.rollout_manager import RolloutManager -from miles.rollout.base_types import RolloutFnEvalInput, RolloutFnEvalOutput, RolloutFnTrainInput, RolloutFnTrainOutput +from miles.rollout.base_types import ( + RolloutFnEvalInput, + RolloutFnEvalOutput, + RolloutFnTrainInput, + RolloutFnTrainOutput, + RolloutPostprocessOptions, +) @pytest.fixture @@ -626,6 +632,34 @@ def fake_rollout_fn(input): # 8 samples / 2 dp = 4 per rank assert len(partition["tokens"]) == 4 + async def test_typed_postprocess_options_drive_dp_padding( + self, + ray_local_mode, + placement_group_factory, + tmp_path, + patch_low_level, + ): + args = _make_test_args(tmp_path, models=[("actor", True)]) + args.global_batch_size = 8 + pg = placement_group_factory(2) + + manager = _make_manager(args, pg) + manager.train_parallel_config = {"dp_size": 2} + + def fake_rollout_fn(input): + return RolloutFnTrainOutput( + samples=[make_samples_grouped(n_groups=7, group_size=1)], + postprocess=RolloutPostprocessOptions(pad_to_dp=True), + ) + + manager.generate_rollout = fake_rollout_fn + + result = await manager.generate(rollout_id=7) + + assert result["sample_indices"] == [0, 1, 2, 3, 4, 5, 6, -1] + partitions = ray.get([box.inner for box in result["data_ref"]]) + assert [len(p["tokens"]) for p in partitions] == [4, 4] + @pytest.mark.asyncio class TestEval: diff --git a/tests/fast/ray/rollout/test_addr_allocator.py b/tests/fast/ray/rollout/test_addr_allocator.py index 088b677982d..f2eaecdc7f9 100644 --- a/tests/fast/ray/rollout/test_addr_allocator.py +++ b/tests/fast/ray/rollout/test_addr_allocator.py @@ -2,6 +2,8 @@ from unittest.mock import MagicMock +import pytest + from tests.fast.ray.rollout.conftest import fake_engine, make_args from miles.ray.rollout.addr_allocator import ( @@ -11,6 +13,13 @@ ) +@pytest.fixture +def patch_ray_get(monkeypatch): + import miles.ray.rollout.addr_allocator as mod + + monkeypatch.setattr(mod.ray, "get", lambda x: x) + + class TestPortCursors: def test_empty_has_no_values(self): c = PortCursors.empty() diff --git a/tests/fast/ray/rollout/test_components.py b/tests/fast/ray/rollout/test_components.py new file mode 100644 index 00000000000..5612b6ad6e0 --- /dev/null +++ b/tests/fast/ray/rollout/test_components.py @@ -0,0 +1,55 @@ +import asyncio +from types import SimpleNamespace + +from miles.ray.rollout.components import InferenceEndpoint, create_rollout_components + + +class Remote: + def __init__(self, log, name, value=None): + self._log, self._name, self._value = log, name, value + + async def remote(self, *args): + self._log.append((self._name, args)) + return self._value + + +def make_fake_manager(log): + return SimpleNamespace( + get_router_address=Remote(log, "get_router_address", ("10.0.0.7", 30001)), + generate=Remote(log, "generate", {"batch": 1}), + dispose=Remote(log, "dispose"), + ) + + +def build(monkeypatch, log): + manager = make_fake_manager(log) + monkeypatch.setattr( + "miles.ray.placement_group.create_rollout_manager", lambda args, pg: (manager, 7), raising=True + ) + components = create_rollout_components(SimpleNamespace(), pg=None) + return components, manager + + +def test_factory_builds_two_role_views_over_one_legacy_handle(monkeypatch): + log: list = [] + components, manager = build(monkeypatch, log) + + assert components.inference_controller is not components.rollout_executor + assert components.weight_update_owner is manager + assert not hasattr(components.inference_controller, "manager") + + endpoint = asyncio.run(components.inference_controller.get_inference_endpoint()) + assert endpoint == InferenceEndpoint(host="10.0.0.7", port=30001) + assert endpoint.base_url == "http://10.0.0.7:30001" + + asyncio.run(components.inference_controller.prepare_rollout(3)) + assert asyncio.run(components.rollout_executor.generate(3)) == {"batch": 1} + assert ("generate", (3,)) in log + + +def test_bundle_disposes_the_shared_actor_exactly_once(monkeypatch): + log: list = [] + components, _ = build(monkeypatch, log) + asyncio.run(components.dispose()) + asyncio.run(components.dispose()) + assert [name for name, _ in log].count("dispose") == 1 diff --git a/tests/fast/ray/rollout/test_multi_lora_batch_collection.py b/tests/fast/ray/rollout/test_multi_lora_batch_collection.py deleted file mode 100644 index 35cb66b88ee..00000000000 --- a/tests/fast/ray/rollout/test_multi_lora_batch_collection.py +++ /dev/null @@ -1,290 +0,0 @@ -"""Unit tests for multi-LoRA batch collection (get_groups + collect_batch): -group-multiple math, adapter batch capping, step stamping, coalesce timeout, -round-robin fairness, retirement, and staleness filtering. No Ray, no engines: -the worker is built bare.""" - -import asyncio -import threading -import time -from collections import defaultdict, deque -from types import SimpleNamespace - -import pytest - -from miles.rollout.multi_lora.async_rollout import ( - AsyncMultiLoRAWorker, - GroupBuffer, - MultiLoRAWorkerMetrics, - collect_batch, - group_adapter_name, -) -from miles.utils.adapter_config import AdapterRun, AdapterRunConfig -from miles.utils.types import AdapterRef, Sample - - -def make_args(**overrides) -> SimpleNamespace: - args = SimpleNamespace( - global_batch_size=16, - multi_lora_dp_size=4, - multi_lora_max_coalesce_wait_s=0.05, - max_weight_staleness=None, - ) - for key, value in overrides.items(): - setattr(args, key, value) - return args - - -def make_worker(args=None) -> AsyncMultiLoRAWorker: - worker = AsyncMultiLoRAWorker.__new__(AsyncMultiLoRAWorker) - worker.args = args or make_args() - worker.buffer_lock = threading.Lock() - worker.buffers = defaultdict(GroupBuffer) - worker.rotation = deque() - worker.dynamic_filter = None - worker.metrics = MultiLoRAWorkerMetrics() - worker.registrations = {} - worker.failure = None - return worker - - -def adapter_run( - name: str, - slot: int, - rollout_batch_size: int = 4, - n_samples_per_prompt: int = 4, - accumulated_groups: int = 0, - version: int = 1, - registration_id: str = "", -) -> AdapterRun: - config = AdapterRunConfig( - data="/d", - rank=8, - alpha=16, - rollout_batch_size=rollout_batch_size, - n_samples_per_prompt=n_samples_per_prompt, - ) - return AdapterRun( - name=name, - config=config, - slot=slot, - version=version, - step=0, - accumulated_groups=accumulated_groups, - registration_id=registration_id, - ) - - -def make_group( - adapter: AdapterRun, slot_version: int | None = None, registration_id: str | None = None -) -> list[Sample]: - samples = [] - for _ in range(adapter.config.n_samples_per_prompt): - sample = Sample(prompt="p", adapter=AdapterRef(adapter.name, adapter.slot)) - if slot_version is not None: - sample.metadata["slot_version"] = slot_version - if registration_id is not None: - sample.metadata["registration_id"] = registration_id - samples.append(sample) - return samples - - -def buffer_groups( - worker, adapter: AdapterRun, count: int, slot_version: int | None = None, registration_id: str | None = None -): - for _ in range(count): - worker.buffers[adapter.name].put(make_group(adapter, slot_version, registration_id)) - - -def snapshot_of(*adapters: AdapterRun, retiring: tuple[AdapterRun, ...] = ()) -> dict: - return { - "active": {a.name: a for a in adapters}, - "retiring": {a.name: a for a in retiring}, - "cleanup": [], - } - - -def collect(worker, snapshot): - return asyncio.run(collect_batch(worker.args, worker, snapshot)) - - -def test_no_pop_until_a_whole_group_multiple_is_buffered(): - # dp=8 with n_samples=4 -> multiple = 2 groups; one buffered group is below the multiple. - worker = make_worker(make_args(multi_lora_dp_size=8)) - a = adapter_run("A", 0, rollout_batch_size=4, n_samples_per_prompt=4) - buffer_groups(worker, a, count=1) - groups, counts = worker.get_groups(snapshot_of(a), 16, {}) - assert (groups, counts) == ([], {}) - - buffer_groups(worker, a, count=1) - groups, counts = worker.get_groups(snapshot_of(a), 16, {}) - assert len(groups) == 2 - assert counts == {"A": 2} - - -def test_reaching_target_stops_collecting(): - worker = make_worker() - a = adapter_run("A", 0, rollout_batch_size=8) # adapter batch: 8 groups - buffer_groups(worker, a, count=5) # 20 samples > 16 target - start = time.monotonic() - batch = collect(worker, snapshot_of(a)) - assert time.monotonic() - start < worker.args.multi_lora_max_coalesce_wait_s # no timeout waited - assert batch.group_counts == {"A": 4} # stops once 16 samples are reached - assert batch.step_names == [] # adapter batch (8 groups) not complete - assert len(worker.buffers["A"]) == 1 - - -def test_below_target_ships_after_no_progress_timeout(): - worker = make_worker() - a = adapter_run("A", 0, rollout_batch_size=8) - buffer_groups(worker, a, count=1) # 4 samples < 16 target - start = time.monotonic() - batch = collect(worker, snapshot_of(a)) - assert time.monotonic() - start >= worker.args.multi_lora_max_coalesce_wait_s - assert batch.group_counts == {"A": 1} - - -def test_collection_capped_at_remaining_groups_and_step_stamped(): - worker = make_worker() - # Adapter batch = 4 groups; 3 already banked -> 1 remaining, despite 4 buffered. - a = adapter_run("A", 0, rollout_batch_size=4, accumulated_groups=3) - buffer_groups(worker, a, count=4) - batch = collect(worker, snapshot_of(a)) - assert batch.group_counts == {"A": 1} - assert batch.step_names == ["A"] - assert batch.step_slots == [0] - assert len(worker.buffers["A"]) == 3 # surplus stays buffered - - -def test_batch_never_overshoots_adapter_batch_across_fetches(): - """Groups arriving after an adapter's remaining groups are already in the - batch must not be popped into the same batch.""" - worker = make_worker() - a = adapter_run("A", 0, rollout_batch_size=2) - buffer_groups(worker, a, count=2) - groups, counts = worker.get_groups(snapshot_of(a), 16, {}) - assert len(groups) == 2 # whole remaining batch - - buffer_groups(worker, a, count=2) # fresh arrivals mid-collection - groups, counts = worker.get_groups(snapshot_of(a), 16, counts) - assert groups == [] - - groups, _counts = worker.get_groups(snapshot_of(a), 16, {}) # next batch may pop them - assert len(groups) == 2 - - -def test_pops_interleave_adapters_round_robin(): - worker = make_worker() - a = adapter_run("A", 0, rollout_batch_size=16) - b = adapter_run("B", 1, rollout_batch_size=16) - buffer_groups(worker, a, count=2) - buffer_groups(worker, b, count=2) - groups, counts = worker.get_groups(snapshot_of(a, b), 16, {}) - assert [group_adapter_name(g) for g in groups] == ["A", "B", "A", "B"] - assert counts == {"A": 2, "B": 2} - groups, counts = worker.get_groups(snapshot_of(a, b), 16, counts) - assert groups == [] # buffers drained - - -def test_cursor_persists_across_batches(): - worker = make_worker(make_args(global_batch_size=8)) - a = adapter_run("A", 0, rollout_batch_size=16) - b = adapter_run("B", 1, rollout_batch_size=16) - buffer_groups(worker, a, count=4) - buffer_groups(worker, b, count=4) - - # 8-sample target = 2 groups per batch; collection interleaves A and B. - batch = collect(worker, snapshot_of(a, b)) - assert batch.group_counts == {"A": 1, "B": 1} - - # The next batch continues from the cursor, not from A again. - batch = collect(worker, snapshot_of(a, b)) - assert batch.group_counts == {"A": 1, "B": 1} - assert len(worker.buffers["A"]) == 2 - assert len(worker.buffers["B"]) == 2 - - -def test_retiring_adapter_remains_selectable_until_retired(): - """RETIRING adapters keep serving until the reconcile sync point (base - deregistration semantics): buffered groups stay poppable.""" - worker = make_worker() - a = adapter_run("A", 0, rollout_batch_size=4) - buffer_groups(worker, a, count=4) - batch = collect(worker, snapshot_of(retiring=(a,))) - assert batch.group_counts == {"A": 4} - assert batch.step_names == ["A"] - - -def test_retired_adapter_buffers_are_discarded(): - """Once an adapter leaves the snapshot (retired at reconcile), its buffered - tail is dropped.""" - worker = make_worker() - a = adapter_run("A", 0, rollout_batch_size=4) - b = adapter_run("B", 1, rollout_batch_size=4) - buffer_groups(worker, a, count=3) - groups, _counts = worker.get_groups(snapshot_of(b), 16, {}) # A gone from snapshot - assert groups == [] - assert "A" not in worker.buffers # tail discarded with the adapter - - -def test_stale_buffered_groups_are_dropped(): - worker = make_worker(make_args(max_weight_staleness=1)) - a = adapter_run("A", 0, rollout_batch_size=4, version=5) - buffer_groups(worker, a, count=2, slot_version=3) # staleness 2 > 1 - buffer_groups(worker, a, count=1, slot_version=5) # fresh - batch = collect(worker, snapshot_of(a)) - assert batch.group_counts == {"A": 1} # only the fresh group ships - - -def test_empty_collection_times_out_instead_of_spinning_forever(): - worker = make_worker(make_args(multi_lora_max_empty_wait_s=0.02)) - a = adapter_run("A", 0, rollout_batch_size=4) - with pytest.raises(RuntimeError, match="No poppable groups collected before empty timeout"): - collect(worker, snapshot_of(a)) - - -def test_re_registered_name_drops_previous_tenant_buffer_and_metrics(): - # A retires while its buffer still holds groups; the driver idles (no - # generate), then the operator re-registers the same name. The new - # tenant's first get_groups must not ship the old tenant's groups nor - # inherit its partial step statistics. - worker = make_worker() - old = adapter_run("A", 0, registration_id="reg-old") - buffer_groups(worker, old, count=2, registration_id="reg-old") - worker.get_groups(snapshot_of(old), 0, {}) # worker has seen the old tenant - worker.metrics.step_rewards["A"].append(1.0) # old tenant's partial step stats - - new = adapter_run("A", 0, registration_id="reg-new") - groups, counts = worker.get_groups(snapshot_of(new), 16, {}) - - assert (groups, counts) == ([], {}) - assert len(worker.buffers["A"]) == 0 - assert "A" not in worker.metrics.step_rewards - - -def test_straggler_group_of_previous_registration_is_dropped(): - # An in-flight generation of the old tenant lands in the buffer after the - # re-registration sweep already reset it; only the new tenant's groups ship. - worker = make_worker() - new = adapter_run("A", 0, registration_id="reg-new") - worker.get_groups(snapshot_of(new), 0, {}) # sweep records the new registration - buffer_groups(worker, new, count=1, registration_id="reg-old") # straggler - buffer_groups(worker, new, count=1, registration_id="reg-new") - - groups, counts = worker.get_groups(snapshot_of(new), 16, {}) - - assert counts == {"A": 1} - assert [s.metadata["registration_id"] for g in groups for s in g] == ["reg-new"] * 4 - - -def test_dead_producer_surfaces_its_cause_instead_of_timing_out(): - # A producer-thread failure (e.g. an adapter whose dataset vanished) stops - # generation for every adapter; collect_batch must raise the recorded cause - # immediately, not wait out the empty-batch timeout. - worker = make_worker(make_args(multi_lora_max_empty_wait_s=30.0)) - worker.failure = RuntimeError("dataset gone") - a = adapter_run("A", 0) - start = time.monotonic() - with pytest.raises(RuntimeError, match="producer thread died") as excinfo: - collect(worker, snapshot_of(a)) - assert time.monotonic() - start < 1.0 # no timeout wait - assert "dataset gone" in repr(excinfo.value.__cause__) diff --git a/tests/fast/ray/rollout/test_multi_lora_operation_train_data.py b/tests/fast/ray/rollout/test_multi_lora_operation_train_data.py new file mode 100644 index 00000000000..d685ded0d2b --- /dev/null +++ b/tests/fast/ray/rollout/test_multi_lora_operation_train_data.py @@ -0,0 +1,261 @@ +from types import SimpleNamespace + +import pytest + +from miles.ray.multi_lora.residency import ResidentBinding +from miles.ray.rollout.rollout_data_conversion import postprocess_rollout_data +from miles.ray.rollout.train_data_conversion import convert_samples_to_train_data +from miles.rollout.multi_lora.rollout_fn import batch_plan_to_metadata +from miles.utils.operation_contract import BatchExecutionLease +from miles.utils.types import AdapterRef, Sample + + +def plan_lease(batch_plan) -> BatchExecutionLease: + return BatchExecutionLease( + dispatch_id="lease-test", + bindings_by_operation=tuple( + ( + entry["operation_id"], + ResidentBinding((entry["name"], entry["registration_id"]), entry["bound_slot"]), + ) + for entry in batch_plan + ), + ) + + +def plan_metadata(batch_plan) -> dict: + return batch_plan_to_metadata(batch_plan, plan_lease(batch_plan)) + + +def plan_entry(name="A", slot=0, kind="forward_backward", op_id="op-A", loss=None, sample_count=1): + return dict( + name=name, + registration_id=f"r-{name}", + bound_slot=slot, + operation_id=op_id, + operation_kind=kind, + loss_spec=loss, + sample_count=sample_count, + ) + + +class TestBatchPlanToMetadata: + def test_forward_backward_plan(self): + plan = [plan_entry("A", 0, loss={"loss_fn": "ppo"}), plan_entry("B", 3, op_id="op-B")] + metadata = batch_plan_to_metadata(plan, plan_lease(plan)) + assert metadata["batch_kind"] == "tinker" + assert metadata["tinker_operation_lanes"] == [0, 1] + assert metadata["tinker_loss_by_lane"] == {0: {"loss_fn": "ppo"}, 1: {}} + assert metadata["operation_by_lane"] == {0: "op-A", 1: "op-B"} + assert metadata["registration_by_lane"] == {0: ("A", "r-A"), 1: ("B", "r-B")} + assert metadata["batch_execution_lease"]["bindings_by_operation"] == [ + ["op-A", ["A", "r-A", 0]], + ["op-B", ["B", "r-B", 3]], + ] + assert "tinker_forward_only" not in metadata + + def test_lanes_expand_per_sample_counts(self): + plan = [plan_entry("A", 0, sample_count=2), plan_entry("B", 3, op_id="op-B", sample_count=3)] + metadata = batch_plan_to_metadata(plan, plan_lease(plan)) + assert metadata["tinker_operation_lanes"] == [0, 0, 1, 1, 1] + + def test_all_forward_sets_the_flag(self): + plan = [plan_entry(kind="forward")] + metadata = batch_plan_to_metadata(plan, plan_lease(plan)) + assert metadata["tinker_forward_only"] is True + + def test_mixed_kinds_are_structurally_rejected(self): + plan = [plan_entry("A", 0), plan_entry("B", 1, kind="forward")] + with pytest.raises(ValueError, match="homogeneous"): + batch_plan_to_metadata(plan, plan_lease(plan)) + plan = [plan_entry(kind="optim_step")] + with pytest.raises(ValueError, match="homogeneous"): + batch_plan_to_metadata(plan, plan_lease(plan)) + + +def make_sample(name="A", index=0, stale_slot=9, loss_weights=None, advantages=None): + sample = Sample( + tokens=[1, 2, 3, 4], + response_length=2, + loss_mask=[1, 1], + index=index, + status=Sample.Status.COMPLETED, + loss_weights=loss_weights, + advantages=advantages, + ) + sample.adapter = AdapterRef(name=name, registration_id=f"r-{name}", serving_version=1, slot=stale_slot) + return sample + + +def convert(samples, metadata): + args = SimpleNamespace(use_dynamic_global_batch_size=False) + return convert_samples_to_train_data( + args, + samples, + metadata=metadata, + custom_convert_samples_to_train_data_func=None, + custom_reward_post_process_func=None, + ) + + +class TestConvert: + def test_tinker_batch_skips_rewards_and_routes_by_plan_slot(self): + metadata = plan_metadata([plan_entry("A", 5, sample_count=2)]) + samples = [make_sample("A", i, stale_slot=9, loss_weights=[0.5, 1.5]) for i in range(2)] + data = convert(samples, metadata) + assert data["rewards"] == [0.0, 0.0] + assert data["adapter_slots"] == [5, 5] + assert data["loss_weights"] == [[0.5, 1.5], [0.5, 1.5]] + assert data["sample_indices"] == [0, 1] + assert data["batch_kind"] == "tinker" + assert data["tinker_operation_lanes"] == [0, 0] + assert data["tinker_loss_by_lane"] == {0: {}} + assert data["operation_by_lane"] == {0: "op-A"} + assert data["registration_by_lane"] == {0: ("A", "r-A")} + assert data["batch_execution_lease"]["bindings_by_operation"] == [["op-A", ["A", "r-A", 5]]] + assert "step_slots" not in data + + def test_two_operations_may_share_one_physical_slot(self): + plan = [ + plan_entry("A", 5, op_id="op-A1"), + plan_entry("A", 5, op_id="op-A2", loss={"loss_fn": "ppo"}), + ] + metadata = plan_metadata(plan) + assert metadata["operation_by_lane"] == {0: "op-A1", 1: "op-A2"} + assert metadata["tinker_loss_by_lane"] == {0: {}, 1: {"loss_fn": "ppo"}} + samples = [make_sample("A", 0, loss_weights=[1.0, 1.0]), make_sample("A", 0, loss_weights=[2.0, 2.0])] + data = convert(samples, metadata) + assert data["adapter_slots"] == [5, 5] + assert data["tinker_operation_lanes"] == [0, 1] + + def test_unplanned_adapter_fails_loudly(self): + metadata = plan_metadata([plan_entry("A", 5)]) + with pytest.raises(ValueError, match="batch lease binds"): + convert([make_sample("ghost")], metadata) + + def test_stale_same_name_registration_is_rejected_before_slot_routing(self): + metadata = plan_metadata([plan_entry("A", 5)]) + stale = make_sample("A") + stale.adapter = AdapterRef(name="A", registration_id="r-old", serving_version=1, slot=9) + with pytest.raises(ValueError, match="registration"): + convert([stale], metadata) + + def test_lease_binding_no_lane_references_is_a_plan_mismatch(self): + metadata = plan_metadata([plan_entry("A", 5)]) + metadata["batch_execution_lease"]["bindings_by_operation"].append(["op-ghost", ["G", "r-G", 7]]) + with pytest.raises(ValueError, match="disagree"): + convert([make_sample("A")], metadata) + + def test_mixed_channels_default_to_zeros(self): + plan = [ + plan_entry("A", 0, loss={"loss_fn": "cross_entropy"}), + plan_entry("B", 1, op_id="op-B", loss={"loss_fn": "importance_sampling"}), + ] + samples = [ + make_sample("A", 0, loss_weights=[1.0, 1.0]), + make_sample("B", 0, advantages=[0.5, -0.5]), + ] + pure_ce_data = convert(samples[:1], plan_metadata(plan[:1])) + assert "rollout_log_probs" not in pure_ce_data + + samples[1].rollout_log_probs = [-0.1, -0.2] + data = convert(samples, plan_metadata(plan)) + assert data["loss_weights"] == [[1.0, 1.0], [0.0, 0.0]] + assert data["advantages"] == [[0.0, 0.0], [0.5, -0.5]] + assert data["rollout_log_probs"] == [[0.0, 0.0], [-0.1, -0.2]] + + reversed_data = convert(samples[::-1], plan_metadata(plan[::-1])) + assert reversed_data["loss_weights"] == [[0.0, 0.0], [1.0, 1.0]] + assert reversed_data["advantages"] == [[0.5, -0.5], [0.0, 0.0]] + assert reversed_data["rollout_log_probs"] == [[-0.1, -0.2], [0.0, 0.0]] + + def test_legacy_batch_keeps_first_sample_optional_channel_semantics(self): + samples = [make_sample("A"), make_sample("B")] + for sample in samples: + sample.adapter = None + samples[1].rollout_log_probs = [-0.1, -0.2] + + data = convert_samples_to_train_data( + SimpleNamespace( + advantage_estimator="grpo", rewards_normalization=False, use_dynamic_global_batch_size=False + ), + samples, + metadata={}, + custom_convert_samples_to_train_data_func=None, + custom_reward_post_process_func=None, + ) + + assert "rollout_log_probs" not in data + + def test_client_channels_survive_the_dp_shard_split(self): + from miles.ray.rollout.train_data_conversion import split_train_data_by_dp_raw + + metadata = plan_metadata([plan_entry("A", 0, sample_count=2)]) + samples = [make_sample("A", i, loss_weights=[0.5, 1.5], advantages=[1.0, -1.0]) for i in range(2)] + data = convert(samples, metadata) + args = SimpleNamespace(balance_data=False, multi_lora_n_adapters=2) + shards = split_train_data_by_dp_raw(args, data, dp_size=2) + for shard in shards: + assert shard["loss_weights"] == [[0.5, 1.5]] + assert shard["advantages"] == [[1.0, -1.0]] + assert shard["tinker_operation_lanes"] == [0] + assert shard["tinker_loss_by_lane"] == {0: {}} + assert shard["operation_by_lane"] == {0: "op-A"} + + +class TestPadding: + def tinker_args(self): + return SimpleNamespace( + multi_lora=True, + use_dynamic_global_batch_size=True, + disable_rollout_trim_samples=False, + global_batch_size=8, + ) + + def samples(self, n): + return [make_sample("A", i, loss_weights=[0.5, 1.5]) for i in range(n)] + + def postprocess(self, n, pad_to_dp=True, args=None): + return postprocess_rollout_data( + args or self.tinker_args(), + self.samples(n), + train_parallel_config={"dp_size": 4}, + pad_to_dp=pad_to_dp, + ) + + def test_pads_to_dp_size_with_inert_rows(self): + data, metadata = self.postprocess(n=2) + assert metadata["dynamic_global_batch_size"] == len(data) == 4 + assert [s.index for s in data] == [0, 1, -1, -1] + assert data[2].loss_mask == [0, 0] and data[3].loss_weights == [0.0, 0.0] + assert data[2].rollout_id is None + assert data[0].loss_mask == [1, 1] and data[1].loss_weights == [0.5, 1.5] + assert all(s.adapter.name == "A" for s in data) + + def test_pads_to_the_next_multiple_not_just_dp_size(self): + data, _ = self.postprocess(n=5) + assert len(data) == 8 + assert [s.index for s in data] == [0, 1, 2, 3, 4, -1, -1, -1] + + def test_noop_when_batch_is_an_exact_multiple(self): + data, metadata = self.postprocess(n=4) + assert [s.index for s in data] == [0, 1, 2, 3] + assert metadata["dynamic_global_batch_size"] == 4 + + +class TestTinkerDispatchSummary: + def test_summary_carries_operation_ids_and_lease(self): + from miles.ray.rollout.train_data_conversion import tinker_dispatch_summary + + lease = {"dispatch_id": "d1", "bindings_by_operation": [["op-A", ["A", "r-A", 0]]]} + train_data = { + "batch_kind": "tinker", + "operation_by_lane": {0: "op-A", 1: "op-B"}, + "batch_execution_lease": lease, + } + assert tinker_dispatch_summary(train_data) == {"operation_ids": ["op-A", "op-B"], "lease": lease} + + def test_non_tinker_batches_have_no_summary(self): + from miles.ray.rollout.train_data_conversion import tinker_dispatch_summary + + assert tinker_dispatch_summary({"tokens": [[1]]}) is None diff --git a/tests/fast/ray/rollout/test_multi_lora_process_group.py b/tests/fast/ray/rollout/test_multi_lora_process_group.py deleted file mode 100644 index b97ff2238a0..00000000000 --- a/tests/fast/ray/rollout/test_multi_lora_process_group.py +++ /dev/null @@ -1,56 +0,0 @@ -"""Pins process_group's submission-time slot-version stamping: the staleness -filter compares against the version live when the group was submitted, not -when it completed.""" - -import pytest - -import miles.rollout.multi_lora.async_rollout as mod -from miles.rollout.multi_lora.async_rollout import process_group -from miles.utils.types import AdapterRef, Sample - - -class FakeDataSource: - def __init__(self) -> None: - self.added: list = [] - - def add_samples(self, groups) -> None: - self.added.extend(groups) - - -class FakeAdapterView: - def __init__(self, version: int, registration_id: str = "reg-1") -> None: - self.version = version - self.registration_id = registration_id - - -class FakeAdaptersCache: - def __init__(self, versions: dict[str, int]) -> None: - self.versions = versions - - def bump(self, name: str, to: int) -> None: - self.versions[name] = to - - async def get(self, adapter_name: str) -> FakeAdapterView | None: - version = self.versions.get(adapter_name) - return FakeAdapterView(version) if version is not None else None - - -@pytest.mark.asyncio -async def test_process_group_stamps_submission_version(monkeypatch): - """The stamp is the version live at submission (5), not completion (7).""" - cache = FakeAdaptersCache({"A": 5}) - - async def gen(args, group, sampling_params): - cache.bump("A", 7) # update lands mid-generation - for s in group: - s.status = Sample.Status.COMPLETED - return group - - monkeypatch.setattr(mod, "AdaptersCache", lambda: cache) - - g = [Sample(prompt="p", adapter=AdapterRef("A", 0))] - result = await process_group(None, g, {}, gen, FakeDataSource()) - - assert result is g - assert g[0].metadata["slot_version"] == 5 - assert g[0].metadata["registration_id"] == "reg-1" diff --git a/tests/fast/ray/rollout/test_multi_lora_train_data.py b/tests/fast/ray/rollout/test_multi_lora_train_data.py index 86f259b3b8a..9fed0d0f546 100644 --- a/tests/fast/ray/rollout/test_multi_lora_train_data.py +++ b/tests/fast/ray/rollout/test_multi_lora_train_data.py @@ -1,7 +1,3 @@ -"""Multi-LoRA train-data pipeline: batch metadata extraction, exact dynamic -batch size, per-adapter batch loss scales, step stamping, and per-group reward -normalization with heterogeneous group sizes.""" - import pytest from tests.ci.ci_register import register_cpu_ci @@ -29,7 +25,6 @@ def adapter_group( name: str, slot: int, n_samples: int, - adapter_global_batch_size: int, rewards: list[float], start_index: int = 0, ): @@ -38,44 +33,24 @@ def adapter_group( for k in range(n_samples): sample = make_sample(index=start_index + k, reward=rewards[k]) sample.adapter = AdapterRef(name, slot) - sample.metadata = {"adapter_global_batch_size": adapter_global_batch_size} group.append(sample) return group def make_batch(): - """Two adapters, heterogeneous group sizes: A steps this batch, B doesn't.""" - groups = [ - adapter_group("A", 0, 4, 16, [1.0, 0.0, 1.0, 0.0], start_index=0), - adapter_group("A", 0, 4, 16, [1.0, 1.0, 1.0, 1.0], start_index=4), - adapter_group("B", 1, 2, 32, [3.0, 1.0], start_index=8), + return [ + adapter_group("A", 0, 4, [1.0, 0.0, 1.0, 0.0], start_index=0), + adapter_group("A", 0, 4, [1.0, 1.0, 1.0, 1.0], start_index=4), + adapter_group("B", 1, 2, [3.0, 1.0], start_index=8), ] - groups[0][0].metadata["step_slots"] = [0] - groups[0][0].metadata["step_adapter_names"] = ["A"] - return groups - - -def run_pipeline(dp_size: int = 2): - args = multi_lora_args() - data, metadata = postprocess_rollout_data(args, make_batch(), train_parallel_config={"dp_size": dp_size}) - train_data = convert_samples_to_train_data( - args, - data, - metadata=metadata, - custom_convert_samples_to_train_data_func=None, - custom_reward_post_process_func=None, - ) - return data, metadata, train_data def test_postprocess_extracts_batch_metadata_and_exact_batch_size(): - data, metadata, _ = run_pipeline() + args = multi_lora_args() + data, metadata = postprocess_rollout_data(args, make_batch(), train_parallel_config={"dp_size": 2}) assert metadata["prompt_group_sizes"] == [4, 4, 2] - assert metadata["step_slots"] == [0] - assert metadata["step_adapter_names"] == ["A"] assert metadata["dynamic_global_batch_size"] == 10 # exact batch size, no trim assert len(data) == 10 # flattened - assert "step_slots" not in data[0].metadata # lifted out def test_multi_lora_rejects_dp_indivisible_batch(): @@ -84,23 +59,14 @@ def test_multi_lora_rejects_dp_indivisible_batch(): postprocess_rollout_data(args, make_batch(), train_parallel_config={"dp_size": 4}) -def test_step_fields(): - _, _, train_data = run_pipeline() - assert train_data["adapter_slots"] == [0] * 8 + [1] * 2 - assert train_data["step_slots"] == [0] - assert train_data["step_adapter_names"] == ["A"] - # Only A steps: the trainer scales slot 0's accumulated gradient by 1/16. - assert train_data["step_adapter_batch_sizes"] == {0: 16} - assert train_data["prompt_group_sizes"] == [4, 4, 2] - - -def test_rewards_normalize_within_heterogeneous_groups(): - _, _, train_data = run_pipeline() - rewards = train_data["rewards"] - # Group boundaries: [0:4], [4:8], [8:10] — each zero-mean. - for start, end in [(0, 4), (4, 8), (8, 10)]: - assert sum(rewards[start:end]) == pytest.approx(0.0, abs=1e-6) - # Constant group (all 1.0) normalizes to zeros, not NaN. - assert rewards[4:8] == pytest.approx([0.0] * 4) - # Singleton-free std normalization applied to group 1 (n=4, mixed). - assert max(abs(r) for r in rewards[0:4]) > 0.5 +def test_adapter_batch_without_tinker_lease_is_rejected(): + args = multi_lora_args() + data, metadata = postprocess_rollout_data(args, make_batch(), train_parallel_config={"dp_size": 2}) + with pytest.raises(ValueError, match="batch lease"): + convert_samples_to_train_data( + args, + data, + metadata=metadata, + custom_convert_samples_to_train_data_func=None, + custom_reward_post_process_func=None, + ) diff --git a/tests/fast/ray/rollout/test_train_data_conversion.py b/tests/fast/ray/rollout/test_train_data_conversion.py index 6e85c66e88a..8d5cb03f77a 100644 --- a/tests/fast/ray/rollout/test_train_data_conversion.py +++ b/tests/fast/ray/rollout/test_train_data_conversion.py @@ -606,8 +606,9 @@ def test_ppo_path_is_identity(self, n, seed): class TestSplitTrainDataByDp: @pytest.fixture(autouse=True) - def _init_object_store(self): + def _init_object_store(self, monkeypatch): """split_train_data_by_dp puts through the object store singleton.""" + monkeypatch.setattr(object_store, "_INSTANCE", None) object_store.init_instance(make_args()) def test_strided_partition_when_balance_data_off(self): diff --git a/tests/fast/ray/tinker_frontend/__init__.py b/tests/fast/ray/tinker_frontend/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/fast/ray/tinker_frontend/fake_stack.py b/tests/fast/ray/tinker_frontend/fake_stack.py new file mode 100644 index 00000000000..bcb1b22eabd --- /dev/null +++ b/tests/fast/ray/tinker_frontend/fake_stack.py @@ -0,0 +1,140 @@ +"""Test stack for the tinker frontend: a real MultiLoraOperationBackend (registry + +ledger + validation) driven by a fake trainer loop. Only the Ray/trainer/GPU +boundary is faked — the fake driver speaks exactly the documented controller +verbs the Megatron driver uses (claim/commit/complete/retire/bootstrap/ +mark_ready/record_weight_update), so ordering, dirty pins, fencing, and the +publish barrier behave like production.""" + +import asyncio +from types import SimpleNamespace + +from miles.ray.multi_lora.backend import MultiLoraOperationBackend +from miles.ray.multi_lora.registry import AdapterState + + +def make_backend(router_url: str = "http://127.0.0.1:9", save_root: str = "/tmp/tinker-frontend-test", **overrides): + args = SimpleNamespace( + multi_lora_n_adapters=4, + save=save_root, + lora_rank=32, + lora_alpha=64, + hf_checkpoint="Qwen/Qwen3-0.6B", + tinker_api_key=None, + ) + for key, value in overrides.items(): + setattr(args, key, value) + return MultiLoraOperationBackend(args, router_url) + + +class FakeDriver: + def __init__(self, backend: MultiLoraOperationBackend, base_logprob: float = -0.5) -> None: + self.backend = backend + self.base_logprob = base_logprob + self.saved_states: dict[str, int] = {} + self.paused = False + backend.mark_trainer_ready() + + async def run(self, interval: float = 0.005) -> None: + while True: + if not self.paused: + await self.tick() + await asyncio.sleep(interval) + + async def tick(self) -> None: + registry = self.backend.registry + await self.backend.retire_adapters() + for name in sorted(registry.in_state(AdapterState.CLEANUP)): + await self.backend.free_slot(name) + registry.bootstrap_pending() + registry.mark_ready( + [name for name, r in registry.in_state(AdapterState.PENDING).items() if r.slot is not None] + ) + self._run_data_operations() + self._run_control_operations() + + def _row(self, name: str, length: int) -> list[float]: + step = self.backend.adapter_step(name) + return [self.base_logprob - 0.01 * step] * length + + def _run_data_operations(self) -> None: + for name, run in list(self.backend.registry.ready_adapters().items()): + while (op := self.backend.claim_data_operation(name, run.registration_id)) is not None: + rows = [self._row(name, sample["response_length"]) for sample in op["payload"]["samples"]] + accumulated = [(name, run.registration_id)] if op["kind"] == "forward_backward" else [] + self.backend.commit_tinker_batch(accumulated, [op["operation_id"]], {op["operation_id"]: rows}) + + def _run_control_operations(self) -> None: + claimed = self.backend.claim_ready_control_operations() + for op in claimed["operations"]: + kind, name, payload = op["kind"], op["name"], op.get("payload") or {} + if kind == "optim_step": + if op.get("poison"): + result = dict(ok=False, error=op["poison"], category="user", gradient_window_consumed=True) + else: + adam = payload.get("adam_params") or {} + result = dict(ok=True, result=dict(grad_norm=0.125, learning_rate=adam.get("learning_rate", 1e-4))) + elif kind == "save_state": + tag = str(payload.get("tag") or f"step_{op['step']}") + save_dir = self.backend.registry.find(name).config.save + path = f"{save_dir}/{tag}" + if path in self.saved_states: + result = dict( + ok=False, error=f"state '{tag}' already exists; states are immutable", category="user" + ) + else: + self.saved_states[path] = op["step"] + result = dict(ok=True, result=dict(path=path, step=op["step"])) + elif kind == "load_state": + path = payload.get("path") + if path not in self.saved_states: + result = dict(ok=False, error=f"no state at '{path}'", category="user") + else: + result = dict(ok=True, result=dict(step=self.saved_states[path], path=path)) + elif kind == "save_weights_for_sampler": + self.backend.registry.record_weight_update([name]) + result = dict(ok=True) + else: + result = dict(ok=False, error=f"fake driver cannot run '{kind}'", category="server") + self.backend.complete_control_operations({op["operation_id"]: result}) + if claimed["lease"] is not None: + self.backend.release_batch_lease(claimed["lease"]) + + +class FakeRouter: + def __init__(self, max_req_input_len: int = 4090) -> None: + self.requests: list[dict] = [] + self.max_req_input_len = max_req_input_len + self.server_info_calls = 0 + + def app(self): + from fastapi import FastAPI, Request + + app = FastAPI() + + @app.post("/generate") + async def generate(request: Request) -> dict: + payload = await request.json() + self.requests.append(payload) + return self.response_for(payload) + + @app.get("/get_server_info") + async def get_server_info() -> dict: + self.server_info_calls += 1 + return {"context_length": None, "max_req_input_len": self.max_req_input_len, "status": "ready"} + + return app + + def response_for(self, payload: dict) -> dict: + max_new = int((payload.get("sampling_params") or {}).get("max_new_tokens") or 4) + n = min(max_new, 3) + input_ids = payload.get("input_ids") or [] + meta_info = { + "finish_reason": {"type": "length" if n == max_new else "stop"}, + "output_token_logprobs": [[-0.25 * (i + 1), 1000 + i, None] for i in range(n)], + "prompt_tokens": len(input_ids), + } + if payload.get("logprob_start_len") == 0: + meta_info["input_token_logprobs"] = [ + [None if i == 0 else -0.125 * i, token, None] for i, token in enumerate(input_ids) + ] + return {"text": "ok", "meta_info": meta_info} diff --git a/tests/fast/ray/tinker_frontend/test_http_server.py b/tests/fast/ray/tinker_frontend/test_http_server.py new file mode 100644 index 00000000000..91b187b4f58 --- /dev/null +++ b/tests/fast/ray/tinker_frontend/test_http_server.py @@ -0,0 +1,87 @@ +"""The frontend HTTP guard: SDK-key auth on /api/v1, the loopback-only +operator plane, readiness vs liveness probes, and the CLI flag contract — +all over ASGI so peer addresses can be faked.""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=30, suite="stage-a-cpu") + +import asyncio +from types import SimpleNamespace + +import httpx +import pytest +from tests.fast.ray.tinker_frontend.fake_stack import FakeDriver, make_backend + +from miles.ray.multi_lora.http_server import AdapterRunControlServer +from miles.ray.tinker_frontend.http_server import TinkerFrontendHTTPServer +from miles.utils.tinker import validate_tinker_args + +API_KEY = "tml-test-key" + + +def test_frontend_extends_the_canonical_operation_control_server(): + assert issubclass(TinkerFrontendHTTPServer, AdapterRunControlServer) + + +def make_app(api_key=API_KEY, ready=True): + backend = make_backend(tinker_api_key=api_key) + if ready: + FakeDriver(backend) + server = TinkerFrontendHTTPServer(backend, host="127.0.0.1", api_port=0) + app = server.create_app() + server.add_routes(app) + return app + + +def get(app, path, peer="127.0.0.1", **headers): + async def go(): + transport = httpx.ASGITransport(app=app, client=(peer, 40000)) + async with httpx.AsyncClient(transport=transport, base_url="http://frontend") as client: + return await client.get(path, headers=headers) + + return asyncio.run(go()) + + +class TestGuard: + def test_sdk_routes_require_the_key_from_any_peer(self): + app = make_app() + assert get(app, "/api/v1/get_server_capabilities").status_code == 401 + for peer in ("127.0.0.1", "203.0.113.9"): + response = get(app, "/api/v1/get_server_capabilities", peer=peer, **{"x-api-key": API_KEY}) + assert response.status_code == 200, peer + + def test_operator_plane_is_loopback_only_even_with_the_sdk_key(self): + app = make_app() + for path in ("/adapter_runs", "/info"): + assert get(app, path, peer="203.0.113.9", **{"x-api-key": API_KEY}).status_code == 403, path + assert get(app, path, peer="127.0.0.1", **{"x-api-key": API_KEY}).status_code == 200, path + assert get(app, "/adapter_runs").status_code == 401 + + def test_health_probes_are_exempt_from_auth(self): + app = make_app() + assert get(app, "/health").status_code == 200 + assert get(app, "/api/v1/healthz").status_code == 200 + + def test_healthz_is_503_until_the_trainer_is_ready(self): + app = make_app(ready=False) + assert get(app, "/health").status_code == 200 + assert get(app, "/api/v1/healthz").status_code == 503 + + +class TestLaunchFlags: + def args(self, **overrides): + values = dict(tinker_backend=False, tinker_frontend=False, tinker_api_key=None) + values.update(overrides) + return SimpleNamespace(**values) + + def test_frontend_alone_fails_loud_instead_of_a_silent_noop(self): + with pytest.raises(AssertionError, match="requires --tinker-backend"): + validate_tinker_args(self.args(tinker_frontend=True)) + + def test_api_key_requires_the_frontend(self): + with pytest.raises(AssertionError, match="requires --tinker-frontend"): + validate_tinker_args(self.args(tinker_api_key="tml-x")) + + def test_plain_run_still_validates(self): + validate_tinker_args(self.args()) diff --git a/tests/fast/ray/tinker_frontend/test_sampling_admission.py b/tests/fast/ray/tinker_frontend/test_sampling_admission.py new file mode 100644 index 00000000000..09c31b4c2c1 --- /dev/null +++ b/tests/fast/ray/tinker_frontend/test_sampling_admission.py @@ -0,0 +1,336 @@ +"""Sampling admission + transport bound (the Tau sampling-stall P0 fix): +global weighted fail-fast admission (429 BEFORE identity consumption), the +transport's hard in-flight invariant, permit release on every exit path +(success, failure, stale, sibling cancellation, shutdown), and typed error +classification for empty-message exceptions like httpx.PoolTimeout.""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=60, suite="stage-a-cpu") + +import asyncio + +import httpx +import pytest +from tests.fast.ray.tinker_frontend.fake_stack import make_backend + +from miles.ray.multi_lora.operations import OperationBackpressure +from miles.ray.tinker_frontend import wire +from miles.ray.tinker_frontend.sampling import SGLangRouterSamplingTransport +from miles.ray.tinker_frontend.service import ApiError, TinkerFrontend + +BASE = "Qwen/Qwen3-0.6B" + + +class GatedTransport: + def __init__(self) -> None: + self.calls = 0 + self.started = asyncio.Event() + self.release = asyncio.Event() + + async def generate(self, payload: dict) -> dict: + self.calls += 1 + self.started.set() + await self.release.wait() + return { + "meta_info": { + "finish_reason": {"type": "stop"}, + "output_token_logprobs": [[-0.25, 1000, None]], + } + } + + async def close(self) -> None: + pass + + +class FailingTransport: + def __init__(self, exc: BaseException) -> None: + self.calls = 0 + self.exc = exc + + async def generate(self, payload: dict) -> dict: + self.calls += 1 + raise self.exc + + async def close(self) -> None: + pass + + +async def make_frontend(transport, cap: int): + backend = make_backend() + await backend.init() + frontend = TinkerFrontend( + backend, + poll_window_s=0.2, + poll_interval_s=0.001, + sampling_transport=transport, + sampling_max_active_subgenerations=cap, + ) + session_id = frontend.create_session(wire.CreateSessionRequest(sdk_version="0.24.1"))["session_id"] + sampler_id = frontend.create_sampling_session( + wire.CreateSamplingSessionRequest(session_id=session_id, sampling_session_seq_id=0, base_model=BASE) + )["sampling_session_id"] + return backend, frontend, sampler_id + + +def sample_request(sampler_id, seq=0, num_samples=1): + return wire.SampleRequest.model_validate( + { + "sampling_session_id": sampler_id, + "seq_id": seq, + "num_samples": num_samples, + "prompt": {"chunks": [{"type": "encoded_text", "tokens": [5, 6]}]}, + "sampling_params": {"max_tokens": 1}, + } + ) + + +async def retrieve(frontend, request_id): + return await frontend.retrieve_future(wire.FutureRetrieveRequest(request_id=request_id)) + + +async def drain_callbacks(): + await asyncio.sleep(0) + await asyncio.sleep(0) + + +class TestWeightedAdmission: + def test_num_samples_weighs_the_quota(self): + async def main(): + transport = GatedTransport() + backend, frontend, sampler_id = await make_frontend(transport, cap=4) + try: + first = frontend.sample(sample_request(sampler_id, seq=0, num_samples=3)) + assert frontend.sampling_admission.in_use == 3 + with pytest.raises(OperationBackpressure): + frontend.sample(sample_request(sampler_id, seq=1, num_samples=2)) + second = frontend.sample(sample_request(sampler_id, seq=2, num_samples=1)) + assert frontend.sampling_admission.in_use == 4 + transport.release.set() + assert (await retrieve(frontend, first["request_id"]))["type"] == "sample" + assert (await retrieve(frontend, second["request_id"]))["type"] == "sample" + await drain_callbacks() + assert frontend.sampling_admission.in_use == 0 + finally: + transport.release.set() + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_backpressure_precedes_identity_consumption_and_the_retry_runs_once(self): + async def main(): + transport = GatedTransport() + backend, frontend, sampler_id = await make_frontend(transport, cap=1) + try: + first = frontend.sample(sample_request(sampler_id, seq=0)) + with pytest.raises(OperationBackpressure): + frontend.sample(sample_request(sampler_id, seq=1)) + assert frontend.futures.get(f"{sampler_id}:s1") is None + assert not frontend.samplers.get(sampler_id).is_spent(1) + assert frontend.sampling_admission.rejected == 1 + + transport.release.set() + assert (await retrieve(frontend, first["request_id"]))["type"] == "sample" + await drain_callbacks() + retried = frontend.sample(sample_request(sampler_id, seq=1)) + assert (await retrieve(frontend, retried["request_id"]))["type"] == "sample" + assert transport.calls == 2 + finally: + transport.release.set() + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_exact_replay_bypasses_a_full_quota(self): + async def main(): + transport = GatedTransport() + transport.release.set() + backend, frontend, sampler_id = await make_frontend(transport, cap=1) + try: + done = frontend.sample(sample_request(sampler_id, seq=0)) + body = await retrieve(frontend, done["request_id"]) + assert body["type"] == "sample" + await drain_callbacks() + + transport.release.clear() + transport.started.clear() + frontend.sample(sample_request(sampler_id, seq=1)) + await transport.started.wait() + assert frontend.sampling_admission.in_use == 1 + replay = frontend.sample(sample_request(sampler_id, seq=0)) + assert replay["request_id"] == done["request_id"] + assert await retrieve(frontend, replay["request_id"]) == body + assert frontend.sampling_admission.in_use == 1 + assert transport.calls == 2 + finally: + transport.release.set() + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_spent_but_evicted_seq_answers_typed_terminal_without_a_permit(self): + async def main(): + transport = GatedTransport() + transport.release.set() + backend, frontend, sampler_id = await make_frontend(transport, cap=1) + frontend.futures.max_delivered = 1 + frontend.futures.max_expired = 1 + try: + for seq in range(3): + done = frontend.sample(sample_request(sampler_id, seq=seq)) + await retrieve(frontend, done["request_id"]) + await drain_callbacks() + calls = transport.calls + + transport.release.clear() + transport.started.clear() + frontend.sample(sample_request(sampler_id, seq=3)) + await transport.started.wait() + resent = frontend.sample(sample_request(sampler_id, seq=0)) + body = await retrieve(frontend, resent["request_id"]) + assert body["category"] == "user" and "already executed" in body["error"] + assert transport.calls == calls + 1 + assert frontend.sampling_admission.in_use == 1 + finally: + transport.release.set() + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_num_samples_over_capacity_is_a_nonretryable_400(self): + async def main(): + transport = GatedTransport() + transport.release.set() + backend, frontend, sampler_id = await make_frontend(transport, cap=4) + try: + with pytest.raises(ApiError) as excinfo: + frontend.sample(sample_request(sampler_id, seq=0, num_samples=5)) + assert excinfo.value.status_code == 400 and "exceeds" in excinfo.value.detail + assert not frontend.samplers.get(sampler_id).is_spent(0) + assert frontend.sampling_admission.rejected == 0 + + fits = frontend.sample(sample_request(sampler_id, seq=0, num_samples=4)) + assert (await retrieve(frontend, fits["request_id"]))["type"] == "sample" + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + +class TestPermitLifecycle: + def test_transport_failure_releases_permits_and_names_the_exception_class(self): + async def main(): + transport = FailingTransport(httpx.PoolTimeout("")) + backend, frontend, sampler_id = await make_frontend(transport, cap=4) + try: + failed = frontend.sample(sample_request(sampler_id, seq=0, num_samples=2)) + body = await retrieve(frontend, failed["request_id"]) + assert body["category"] == "server" + assert "sampling failed (PoolTimeout):" in body["error"] + await drain_callbacks() + assert frontend.sampling_admission.in_use == 0 + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_ambiguous_midbody_failure_is_terminal_and_never_reissued(self): + async def main(): + transport = FailingTransport(httpx.RemoteProtocolError("peer closed connection mid-body")) + backend, frontend, sampler_id = await make_frontend(transport, cap=4) + try: + failed = frontend.sample(sample_request(sampler_id, seq=0)) + body = await retrieve(frontend, failed["request_id"]) + assert body["category"] == "server" and "(RemoteProtocolError)" in body["error"] + await drain_callbacks() + assert transport.calls == 1 + assert frontend.sampling_admission.in_use == 0 + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_stale_registration_releases_the_permit_before_any_router_call(self): + async def main(): + transport = GatedTransport() + backend, frontend, sampler_id = await make_frontend(transport, cap=4) + record = frontend.samplers.get(sampler_id) + record.name, record.registration_id = "ghost", "r-gone" + try: + stale = frontend.sample(sample_request(sampler_id, seq=0)) + body = await retrieve(frontend, stale["request_id"]) + assert body["category"] == "user" and "no longer live" in body["error"] + await drain_callbacks() + assert transport.calls == 0 + assert frontend.sampling_admission.in_use == 0 + finally: + transport.release.set() + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_shutdown_cancellation_drains_permits_deterministically(self): + async def main(): + transport = GatedTransport() + backend, frontend, sampler_id = await make_frontend(transport, cap=4) + future = frontend.sample(sample_request(sampler_id, seq=0, num_samples=3)) + await transport.started.wait() + assert frontend.sampling_admission.in_use == 3 + try: + await frontend.close() + assert frontend.sampling_admission.in_use == 0 + body = await retrieve(frontend, future["request_id"]) + assert body["category"] == "server" and "shutting down" in body["error"] + finally: + await backend.close() + + asyncio.run(main()) + + +class TestTransportBound: + def test_limits_and_timeouts_match_the_configured_bound(self): + transport = SGLangRouterSamplingTransport("http://router:9/", max_inflight=7) + assert transport.base_url == "http://router:9" + assert transport.limits.max_connections == 7 + assert transport.limits.max_keepalive_connections == 7 + assert transport.timeout.pool is None + assert transport._gate._value == 7 + assert transport.timeout.connect == 10.0 + assert transport.timeout.read == 600.0 + assert transport.timeout.write == 60.0 + + def test_sibling_cancellation_releases_the_gate(self): + async def main(): + transport = SGLangRouterSamplingTransport("http://unused:9", max_inflight=1) + started = asyncio.Event() + + class HangingClient: + async def post(self, url, json): + started.set() + await asyncio.Event().wait() + + async def aclose(self): + pass + + transport._http = HangingClient() + holder = asyncio.create_task(transport.generate({})) + await started.wait() + waiter = asyncio.create_task(transport.generate({})) + await asyncio.sleep(0) + assert transport._gate.locked() + waiter.cancel() + holder.cancel() + await asyncio.gather(holder, waiter, return_exceptions=True) + assert not transport._gate.locked() + assert transport._gate._value == 1 + await transport.close() + + asyncio.run(main()) diff --git a/tests/fast/ray/tinker_frontend/test_sampling_context_preflight.py b/tests/fast/ray/tinker_frontend/test_sampling_context_preflight.py new file mode 100644 index 00000000000..8d5ec58a14b --- /dev/null +++ b/tests/fast/ray/tinker_frontend/test_sampling_context_preflight.py @@ -0,0 +1,336 @@ +"""Sampling context preflight (code-0815 §6.2): prompt + max_tokens must fit +the engine context limit, enforced as a typed 400 BEFORE the seq identity is +consumed — the engine itself silently truncates the decode budget of an +oversized request (near zero for an accumulated Tau context) and returns +garbage instead of failing. The limit is statically configured +(--tinker-sampling-max-context / --sglang-context-length) or discovered +lazily from the router's /get_server_info; while unknown, the preflight +admits everything (permissive, never a false reject).""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=60, suite="stage-a-cpu") + +import asyncio +from types import SimpleNamespace + +import pytest +from tests.fast.ray.tinker_frontend.fake_stack import make_backend + +from miles.ray.tinker_frontend import wire +from miles.ray.tinker_frontend.http_server import resolve_sampling_max_context +from miles.ray.tinker_frontend.service import ApiError, TinkerFrontend, _context_limit_from_server_info + +BASE = "Qwen/Qwen3-0.6B" + + +class InfoTransport: + def __init__(self, info: dict | None = None, info_exc: Exception | None = None) -> None: + self.info = info + self.info_exc = info_exc + self.generate_calls = 0 + self.info_calls = 0 + + async def generate(self, payload: dict) -> dict: + self.generate_calls += 1 + return { + "meta_info": { + "finish_reason": {"type": "stop"}, + "output_token_logprobs": [[-0.25, 1000, None]], + } + } + + async def server_info(self) -> dict: + self.info_calls += 1 + if self.info_exc is not None: + raise self.info_exc + return self.info + + async def close(self) -> None: + pass + + +class NoInfoTransport(InfoTransport): + server_info = None + + +async def make_frontend(transport, max_context=None, cap=8): + backend = make_backend() + await backend.init() + frontend = TinkerFrontend( + backend, + poll_window_s=0.2, + poll_interval_s=0.001, + sampling_transport=transport, + sampling_max_active_subgenerations=cap, + sampling_max_context=max_context, + ) + session_id = frontend.create_session(wire.CreateSessionRequest(sdk_version="0.24.1"))["session_id"] + sampler_id = frontend.create_sampling_session( + wire.CreateSamplingSessionRequest(session_id=session_id, sampling_session_seq_id=0, base_model=BASE) + )["sampling_session_id"] + return backend, frontend, sampler_id + + +def sample_request(sampler_id, seq=0, prompt_len=2, max_tokens=1, num_samples=1): + return wire.SampleRequest.model_validate( + { + "sampling_session_id": sampler_id, + "seq_id": seq, + "num_samples": num_samples, + "prompt": {"chunks": [{"type": "encoded_text", "tokens": list(range(5, 5 + prompt_len))}]}, + "sampling_params": {"max_tokens": max_tokens}, + } + ) + + +async def retrieve(frontend, request_id): + return await frontend.retrieve_future(wire.FutureRetrieveRequest(request_id=request_id)) + + +async def wait_discovery(frontend, timeout_s=2.0): + deadline = asyncio.get_running_loop().time() + timeout_s + while frontend._context_limit is None and frontend._context_discovery_task is not None: + if asyncio.get_running_loop().time() > deadline: + raise TimeoutError("context discovery never settled") + await asyncio.sleep(0.001) + + +class TestConfiguredLimit: + def test_oversized_is_a_typed_400_before_identity_and_the_boundary_is_inclusive(self): + async def main(): + transport = InfoTransport() + backend, frontend, sampler_id = await make_frontend(transport, max_context=64) + try: + with pytest.raises(ApiError) as excinfo: + frontend.sample(sample_request(sampler_id, seq=0, prompt_len=60, max_tokens=8)) + assert excinfo.value.status_code == 400 + assert "context limit of 64" in excinfo.value.detail + assert frontend.futures.get(f"{sampler_id}:s0") is None + assert not frontend.samplers.get(sampler_id).is_spent(0) + assert frontend.sampling_admission.rejected == 0 + assert transport.generate_calls == 0 + + fits = frontend.sample(sample_request(sampler_id, seq=0, prompt_len=56, max_tokens=8)) + assert (await retrieve(frontend, fits["request_id"]))["type"] == "sample" + assert transport.generate_calls == 1 + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_a_configured_limit_never_queries_the_transport(self): + async def main(): + transport = InfoTransport(info={"context_length": 999}) + backend, frontend, sampler_id = await make_frontend(transport, max_context=64) + try: + done = frontend.sample(sample_request(sampler_id, seq=0)) + await retrieve(frontend, done["request_id"]) + assert transport.info_calls == 0 + assert frontend._context_limit == 64 + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_capabilities_advertise_the_known_limit(self): + async def main(): + backend, frontend, _ = await make_frontend(InfoTransport(), max_context=64) + try: + [model] = frontend.capabilities()["supported_models"] + assert model["max_context_length"] == 64 + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + +class TestDiscovery: + def test_tighter_discovered_limit_wins(self): + async def main(): + transport = InfoTransport(info={"context_length": 128, "max_req_input_len": 100}) + backend, frontend, sampler_id = await make_frontend(transport) + try: + [model] = frontend.capabilities()["supported_models"] + assert model["max_context_length"] is None + done = frontend.sample(sample_request(sampler_id, seq=0)) + await wait_discovery(frontend) + assert frontend._context_limit == 106 + await retrieve(frontend, done["request_id"]) + + with pytest.raises(ApiError, match="context limit of 106"): + frontend.sample(sample_request(sampler_id, seq=1, prompt_len=120, max_tokens=16)) + assert transport.info_calls == 1 + [model] = frontend.capabilities()["supported_models"] + assert model["max_context_length"] == 106 + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_null_context_length_reconstructs_from_max_req_input_len(self): + assert _context_limit_from_server_info({"context_length": None, "max_req_input_len": 122}) == 128 + assert _context_limit_from_server_info({"context_length": 256, "max_req_input_len": 122}) == 128 + assert _context_limit_from_server_info({"context_length": True, "max_req_input_len": True}) is None + assert _context_limit_from_server_info({"status": "ready"}) is None + assert _context_limit_from_server_info(["not", "a", "dict"]) is None + + def test_preflight_is_permissive_until_discovery_lands(self): + async def main(): + release = asyncio.Event() + + class SlowInfoTransport(InfoTransport): + async def server_info(self): + self.info_calls += 1 + await release.wait() + return {"context_length": 8} + + transport = SlowInfoTransport() + backend, frontend, sampler_id = await make_frontend(transport) + try: + admitted = frontend.sample(sample_request(sampler_id, seq=0, prompt_len=100, max_tokens=50)) + assert (await retrieve(frontend, admitted["request_id"]))["type"] == "sample" + release.set() + await wait_discovery(frontend) + with pytest.raises(ApiError, match="context limit of 8"): + frontend.sample(sample_request(sampler_id, seq=1, prompt_len=100, max_tokens=50)) + finally: + release.set() + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_discovery_failure_disables_the_preflight_after_bounded_attempts(self): + async def main(): + transport = InfoTransport(info_exc=RuntimeError("router not ready")) + backend, frontend, sampler_id = await make_frontend(transport) + try: + for seq in range(TinkerFrontend._CONTEXT_DISCOVERY_MAX_ATTEMPTS + 2): + done = frontend.sample(sample_request(sampler_id, seq=seq, prompt_len=100, max_tokens=100)) + await wait_discovery(frontend) + assert (await retrieve(frontend, done["request_id"]))["type"] == "sample" + assert transport.info_calls == TinkerFrontend._CONTEXT_DISCOVERY_MAX_ATTEMPTS + assert frontend._context_limit is None + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_a_transport_without_server_info_disables_the_preflight(self): + async def main(): + transport = NoInfoTransport() + backend, frontend, sampler_id = await make_frontend(transport) + try: + done = frontend.sample(sample_request(sampler_id, seq=0, prompt_len=100, max_tokens=100)) + assert (await retrieve(frontend, done["request_id"]))["type"] == "sample" + assert frontend._context_discovery_task is None + assert frontend._context_limit is None + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + +class TestLaunchResolution: + def test_the_tinker_flag_wins_then_the_sglang_context_then_discovery(self): + flagged = SimpleNamespace(tinker_sampling_max_context=32768, sglang_context_length=65536) + assert resolve_sampling_max_context(flagged) == 32768 + deployed = SimpleNamespace(tinker_sampling_max_context=None, sglang_context_length=65536) + assert resolve_sampling_max_context(deployed) == 65536 + bare = SimpleNamespace(tinker_sampling_max_context=None) + assert resolve_sampling_max_context(bare) is None + + +class TestTransportDiscoveryHop: + @staticmethod + async def _serve(app): + import uvicorn + + server = uvicorn.Server(uvicorn.Config(app, host="127.0.0.1", port=0, log_level="critical", access_log=False)) + task = asyncio.get_running_loop().create_task(server.serve()) + while not server.started: + if task.done(): + task.result() + await asyncio.sleep(0.005) + return server, task, server.servers[0].sockets[0].getsockname()[1] + + def test_router_metadata_hops_to_the_first_healthy_worker(self): + from fastapi import FastAPI + + from miles.ray.tinker_frontend.sampling import SGLangRouterSamplingTransport + + async def main(): + worker = FastAPI() + + @worker.get("/get_server_info") + async def worker_info() -> dict: + return {"context_length": 8192, "max_req_input_len": 8186} + + worker_server, worker_task, worker_port = await self._serve(worker) + + router = FastAPI() + + @router.get("/get_server_info") + async def router_info() -> dict: + return {"router_manager": True, "routers_count": 1, "workers_count": 1} + + @router.get("/workers") + async def workers() -> dict: + return { + "workers": [ + {"url": "http://127.0.0.1:1", "is_healthy": False}, + {"url": f"http://127.0.0.1:{worker_port}", "is_healthy": True}, + ] + } + + router_server, router_task, router_port = await self._serve(router) + transport = SGLangRouterSamplingTransport(f"http://127.0.0.1:{router_port}") + try: + info = await transport.server_info() + assert info["context_length"] == 8192 + finally: + await transport.close() + router_server.should_exit = True + worker_server.should_exit = True + await asyncio.gather(router_task, worker_task, return_exceptions=True) + + asyncio.run(main()) + + def test_engine_shape_answers_without_a_hop(self): + from fastapi import FastAPI + + from miles.ray.tinker_frontend.sampling import SGLangRouterSamplingTransport + + async def main(): + engine = FastAPI() + workers_calls = 0 + + @engine.get("/get_server_info") + async def engine_info() -> dict: + return {"context_length": None, "max_req_input_len": 40954} + + @engine.get("/workers") + async def workers() -> dict: + nonlocal workers_calls + workers_calls += 1 + return {"workers": []} + + engine_server, engine_task, engine_port = await self._serve(engine) + transport = SGLangRouterSamplingTransport(f"http://127.0.0.1:{engine_port}") + try: + info = await transport.server_info() + assert info["max_req_input_len"] == 40954 + assert workers_calls == 0 + finally: + await transport.close() + engine_server.should_exit = True + await asyncio.gather(engine_task, return_exceptions=True) + + asyncio.run(main()) diff --git a/tests/fast/ray/tinker_frontend/test_sampling_reaper.py b/tests/fast/ray/tinker_frontend/test_sampling_reaper.py new file mode 100644 index 00000000000..04d31c7cbaf --- /dev/null +++ b/tests/fast/ray/tinker_frontend/test_sampling_reaper.py @@ -0,0 +1,503 @@ +"""Orphan reaper + sampling observability (code-0815 §7 / §6.1). + +The reaper frees bytes and capacity without permitting re-execution: a reaped +sample's seq stays spent while its parent session is live, a reaped result +leaves a fingerprint tombstone, and reaped parent sessions retire their whole +sampler namespace fail-closed. Unpolled operation futures are polled on the vanished +client's behalf, which stores the terminal bytes BEFORE acking the ledger — +the existing retention order, so the unacked-results budget drains without +ever acking an undelivered result away.""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=60, suite="stage-a-cpu") + +import asyncio +import logging +import time + +import httpx +import pytest +from tests.fast.ray.tinker_frontend.fake_stack import FakeDriver, make_backend + +from miles.ray.tinker_frontend import wire +from miles.ray.tinker_frontend.service import ApiError, TinkerFrontend + +BASE = "Qwen/Qwen3-0.6B" +SERVICE_LOGGER = "miles.ray.tinker_frontend.service" + + +class GatedTransport: + def __init__(self) -> None: + self.calls = 0 + self.started = asyncio.Event() + self.release = asyncio.Event() + + async def generate(self, payload: dict) -> dict: + self.calls += 1 + self.started.set() + await self.release.wait() + return { + "meta_info": { + "finish_reason": {"type": "stop"}, + "output_token_logprobs": [[-0.25, 1000, None]], + } + } + + async def close(self) -> None: + pass + + +class FailingTransport: + def __init__(self, exc: BaseException) -> None: + self.exc = exc + + async def generate(self, payload: dict) -> dict: + raise self.exc + + async def close(self) -> None: + pass + + +async def make_frontend(transport, cap=4, **ttl_overrides): + backend = make_backend() + await backend.init() + frontend = TinkerFrontend( + backend, + poll_window_s=0.2, + poll_interval_s=0.001, + sampling_transport=transport, + sampling_max_active_subgenerations=cap, + **ttl_overrides, + ) + session_id = frontend.create_session(wire.CreateSessionRequest(sdk_version="0.24.1"))["session_id"] + sampler_id = frontend.create_sampling_session( + wire.CreateSamplingSessionRequest(session_id=session_id, sampling_session_seq_id=0, base_model=BASE) + )["sampling_session_id"] + return backend, frontend, sampler_id + + +def sample_request(sampler_id, seq=0, num_samples=1): + return wire.SampleRequest.model_validate( + { + "sampling_session_id": sampler_id, + "seq_id": seq, + "num_samples": num_samples, + "prompt": {"chunks": [{"type": "encoded_text", "tokens": [5, 6]}]}, + "sampling_params": {"max_tokens": 1}, + } + ) + + +async def retrieve(frontend, request_id): + return await frontend.retrieve_future(wire.FutureRetrieveRequest(request_id=request_id)) + + +async def drain_callbacks(): + await asyncio.sleep(0) + await asyncio.sleep(0) + + +class TestOrphanedSamples: + def test_unpolled_sample_is_cancelled_typed_and_its_identity_stays_spent(self): + async def main(): + transport = GatedTransport() + backend, frontend, sampler_id = await make_frontend(transport, cap=4) + try: + submitted = frontend.sample(sample_request(sampler_id, seq=0, num_samples=3)) + request_id = submitted["request_id"] + await transport.started.wait() + assert frontend.sampling_admission.in_use == 3 + task = frontend._sample_task_by_request[request_id] + + counts = frontend.reap_once(now=time.time() + frontend.future_unpolled_ttl_s + 1) + assert counts["cancelled_samples"] == 1 + await asyncio.gather(task, return_exceptions=True) + await drain_callbacks() + + assert frontend.sampling_admission.in_use == 0 + assert request_id not in frontend._sample_task_by_request + body = await retrieve(frontend, request_id) + assert body["category"] == "server" and "orphaned" in body["error"] + assert frontend.samplers.get(sampler_id).is_spent(0) + + calls = transport.calls + replay = frontend.sample(sample_request(sampler_id, seq=0, num_samples=3)) + assert replay["request_id"] == request_id + assert (await retrieve(frontend, request_id)) == body + assert transport.calls == calls + finally: + transport.release.set() + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_prestart_orphan_cancellation_still_terminalizes_the_future(self): + async def main(): + transport = GatedTransport() + backend, frontend, sampler_id = await make_frontend(transport, cap=4) + try: + submitted = frontend.sample(sample_request(sampler_id, seq=0)) + request_id = submitted["request_id"] + task = frontend._sample_task_by_request[request_id] + counts = frontend.reap_once(now=time.time() + frontend.future_unpolled_ttl_s + 1) + assert counts["cancelled_samples"] == 1 + await asyncio.gather(task, return_exceptions=True) + await drain_callbacks() + + assert not transport.started.is_set() + assert frontend.sampling_admission.in_use == 0 + assert frontend.sampling_stats.failures_by_class == {"Cancelled": 1} + body = await retrieve(frontend, request_id) + assert body["category"] == "server" and "orphaned" in body["error"] + finally: + transport.release.set() + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_an_actively_polled_sample_is_never_an_orphan(self): + async def main(): + transport = GatedTransport() + backend, frontend, sampler_id = await make_frontend(transport, cap=4) + try: + submitted = frontend.sample(sample_request(sampler_id, seq=0)) + await transport.started.wait() + record = frontend.futures.get(submitted["request_id"]) + record.created_at -= frontend.future_unpolled_ttl_s * 10 + await retrieve(frontend, submitted["request_id"]) + + counts = frontend.reap_once() + assert counts["cancelled_samples"] == 0 + assert not frontend._sample_task_by_request[submitted["request_id"]].done() + + transport.release.set() + assert (await retrieve(frontend, submitted["request_id"]))["type"] == "sample" + finally: + transport.release.set() + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_ttl_zero_disables_reaping(self): + async def main(): + transport = GatedTransport() + backend, frontend, sampler_id = await make_frontend( + transport, cap=4, session_idle_ttl_s=0.0, future_unpolled_ttl_s=0.0, future_undelivered_ttl_s=0.0 + ) + try: + frontend.sample(sample_request(sampler_id, seq=0)) + await transport.started.wait() + counts = frontend.reap_once(now=time.time() + 10_000_000) + assert counts == {"sessions": 0, "cancelled_samples": 0, "undelivered": 0} + assert len(frontend.sessions.records) == 1 + finally: + transport.release.set() + await frontend.close() + await backend.close() + + asyncio.run(main()) + + +class TestUndeliveredResults: + def test_reaped_result_leaves_a_typed_tombstone_and_never_reexecutes(self): + async def main(): + transport = GatedTransport() + transport.release.set() + backend, frontend, sampler_id = await make_frontend(transport, cap=4, session_idle_ttl_s=0) + frontend.futures.max_expired = 1 + try: + submitted = frontend.sample(sample_request(sampler_id, seq=0)) + request_id = submitted["request_id"] + for _ in range(200): + record = frontend.futures.get(request_id) + if record.terminal is not None: + break + await asyncio.sleep(0.001) + assert record.terminal is not None + calls = transport.calls + + counts = frontend.reap_once(now=time.time() + frontend.future_undelivered_ttl_s + 1) + assert counts["undelivered"] == 1 + assert frontend.futures.get(request_id) is None + + with pytest.raises(ApiError) as repoll: + await retrieve(frontend, request_id) + assert repoll.value.status_code == 410 and "reaped" in repoll.value.detail + with pytest.raises(ApiError) as resent: + frontend.sample(sample_request(sampler_id, seq=0)) + assert resent.value.status_code == 410 + assert transport.calls == calls + + done = frontend.sample(sample_request(sampler_id, seq=1)) + await retrieve(frontend, done["request_id"]) + await drain_callbacks() + second = frontend.futures.get(done["request_id"]) + frontend.futures.reap_undelivered(second) + assert frontend.futures.reaped_fingerprint(request_id) is None + calls = transport.calls + fenced = frontend.sample(sample_request(sampler_id, seq=0)) + body = await retrieve(frontend, fenced["request_id"]) + assert body["category"] == "user" and "already executed" in body["error"] + assert transport.calls == calls + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_delivered_results_stay_in_the_replay_window(self): + async def main(): + transport = GatedTransport() + transport.release.set() + backend, frontend, sampler_id = await make_frontend(transport, cap=4) + try: + submitted = frontend.sample(sample_request(sampler_id, seq=0)) + body = await retrieve(frontend, submitted["request_id"]) + assert body["type"] == "sample" + counts = frontend.reap_once(now=time.time() + 10_000_000) + assert counts["undelivered"] == 0 + assert (await retrieve(frontend, submitted["request_id"])) == body + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + +class TestIdleSessions: + def test_idle_session_retires_all_child_samplers_fail_closed(self): + async def main(): + transport = GatedTransport() + transport.release.set() + backend, frontend, sampler_id = await make_frontend(transport, cap=4) + try: + session_id = frontend.samplers.get(sampler_id).session_id + sampler_ids = [sampler_id] + for seq in range(1, 257): + sampler_ids.append( + frontend.create_sampling_session( + wire.CreateSamplingSessionRequest( + session_id=session_id, + sampling_session_seq_id=seq, + base_model=BASE, + ) + )["sampling_session_id"] + ) + done = frontend.sample(sample_request(sampler_id, seq=0)) + body = await retrieve(frontend, done["request_id"]) + await drain_callbacks() + + counts = frontend.reap_once(now=time.time() + frontend.session_idle_ttl_s + 1) + assert counts["sessions"] == 1 + with pytest.raises(ApiError) as heartbeat: + frontend.session_heartbeat(wire.SessionHeartbeatRequest(session_id=session_id)) + assert heartbeat.value.status_code == 404 + assert all(frontend.samplers.get(sampler) is None for sampler in sampler_ids) + calls = transport.calls + assert (await retrieve(frontend, done["request_id"])) == body + with pytest.raises(ApiError) as get_sampler: + frontend.get_sampler(sampler_id) + assert get_sampler.value.status_code == 404 + with pytest.raises(ApiError) as resubmit: + frontend.sample(sample_request(sampler_id, seq=0)) + assert resubmit.value.status_code == 404 + assert transport.calls == calls + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_a_heartbeating_session_is_not_reaped(self): + async def main(): + backend, frontend, sampler_id = await make_frontend(GatedTransport(), cap=4) + try: + session_id = frontend.samplers.get(sampler_id).session_id + frontend.sessions.get(session_id).last_heartbeat = time.time() + assert frontend.reap_once()["sessions"] == 0 + assert frontend.sessions.get(session_id) is not None + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + +class TestVanishedClientOperations: + def test_unpolled_operation_future_is_resolved_and_acked_then_tombstoned(self): + async def main(): + backend = make_backend() + await backend.init() + driver = FakeDriver(backend) + frontend = TinkerFrontend(backend, poll_window_s=0.5, poll_interval_s=0.002) + session_id = frontend.create_session(wire.CreateSessionRequest(sdk_version="0.24.1"))["session_id"] + driver_task = asyncio.create_task(driver.run(interval=0.002)) + try: + create = await frontend.create_model( + wire.CreateModelRequest( + session_id=session_id, model_seq_id=0, base_model=BASE, lora_config=wire.LoraConfig(rank=8) + ) + ) + model_body = await retrieve(frontend, create["request_id"]) + model_id = model_body["model_id"] + fb = frontend.forward_backward( + wire.ForwardBackwardRequest.model_validate( + { + "forward_backward_input": { + "data": [ + { + "model_input": {"chunks": [{"type": "encoded_text", "tokens": [1, 2, 3]}]}, + "loss_fn_inputs": { + "target_tokens": {"data": [2, 3, 99], "dtype": "int64", "shape": [3]}, + "weights": {"data": [1.0, 1.0, 1.0], "dtype": "float32", "shape": [3]}, + }, + } + ], + "loss_fn": "cross_entropy", + }, + "model_id": model_id, + "seq_id": 1, + } + ) + ) + operation_id = fb["request_id"] + for _ in range(500): + view = backend.operation_view(operation_id) + if view is not None and view["state"] == "SUCCEEDED": + break + await asyncio.sleep(0.002) + assert backend.operation_view(operation_id)["state"] == "SUCCEEDED" + + frontend.reap_once(now=time.time() + frontend.future_unpolled_ttl_s + 1) + record = frontend.futures.get(operation_id) + assert record.terminal is not None and record.terminal["type"] == "forward_backward" + assert backend.operation_view(operation_id) is None + + assert (await retrieve(frontend, operation_id))["type"] == "forward_backward" + finally: + driver_task.cancel() + await asyncio.gather(driver_task, return_exceptions=True) + await frontend.close() + await backend.close() + + asyncio.run(main()) + + +class TestMaintenanceLoop: + def test_start_is_idempotent_and_close_tears_it_down(self): + async def main(): + backend, frontend, _ = await make_frontend(GatedTransport(), cap=4) + try: + frontend.start_maintenance() + task = frontend._maintenance_task + assert task is not None + frontend.start_maintenance() + assert frontend._maintenance_task is task + finally: + await frontend.close() + await backend.close() + assert frontend._maintenance_task is None + assert task.cancelled() + + asyncio.run(main()) + + +class TestSamplingMetrics: + def test_admission_counters_and_high_water(self): + async def main(): + transport = GatedTransport() + backend, frontend, sampler_id = await make_frontend(transport, cap=4) + try: + first = frontend.sample(sample_request(sampler_id, seq=0, num_samples=3)) + admission = frontend.sampling_admission + assert (admission.admitted, admission.admitted_weight, admission.peak_in_use) == (1, 3, 3) + second = frontend.sample(sample_request(sampler_id, seq=1, num_samples=1)) + assert (admission.admitted, admission.admitted_weight, admission.peak_in_use) == (2, 4, 4) + transport.release.set() + await retrieve(frontend, first["request_id"]) + await retrieve(frontend, second["request_id"]) + await drain_callbacks() + assert admission.in_use == 0 and admission.peak_in_use == 4 + assert frontend.sampling_stats.completed == 2 + assert frontend.sampling_stats.failed == 0 + finally: + transport.release.set() + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_terminal_failures_are_counted_by_exception_class(self): + async def main(): + backend, frontend, sampler_id = await make_frontend(FailingTransport(httpx.PoolTimeout("")), cap=4) + try: + failed = frontend.sample(sample_request(sampler_id, seq=0, num_samples=2)) + body = await retrieve(frontend, failed["request_id"]) + assert body["category"] == "server" + await drain_callbacks() + assert frontend.sampling_stats.failed == 1 + assert frontend.sampling_stats.failures_by_class == {"PoolTimeout": 1} + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_per_request_latencies_are_stamped(self): + async def main(): + transport = GatedTransport() + transport.release.set() + backend, frontend, sampler_id = await make_frontend(transport, cap=4) + try: + done = frontend.sample(sample_request(sampler_id, seq=0)) + await retrieve(frontend, done["request_id"]) + await drain_callbacks() + record = frontend.futures.get(done["request_id"]) + assert record.first_result_at is not None and record.resolved_at is not None + assert record.created_at <= record.first_result_at <= record.resolved_at + stats = frontend.sampling_stats + assert stats.first_result_count == 1 + assert stats.total_s_max >= stats.first_result_s_max >= 0.0 + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_summary_logs_only_when_something_changed(self): + async def main(): + transport = GatedTransport() + transport.release.set() + backend, frontend, sampler_id = await make_frontend(transport, cap=4) + logger = logging.getLogger(SERVICE_LOGGER) + captured: list[str] = [] + + class Capture(logging.Handler): + def emit(self, record): + captured.append(record.getMessage()) + + handler = Capture(level=logging.INFO) + logger.addHandler(handler) + previous_level = logger.level + logger.setLevel(logging.INFO) + try: + done = frontend.sample(sample_request(sampler_id, seq=0)) + await retrieve(frontend, done["request_id"]) + await drain_callbacks() + frontend._log_sampling_summary() + summaries = [line for line in captured if "sampling summary" in line] + assert len(summaries) == 1 + assert "admitted=1" in summaries[0] and "completed=1" in summaries[0] + frontend._log_sampling_summary() + assert len([line for line in captured if "sampling summary" in line]) == 1 + finally: + logger.removeHandler(handler) + logger.setLevel(previous_level) + await frontend.close() + await backend.close() + + asyncio.run(main()) diff --git a/tests/fast/ray/tinker_frontend/test_sdk_contract.py b/tests/fast/ray/tinker_frontend/test_sdk_contract.py new file mode 100644 index 00000000000..7b11e4e4ed3 --- /dev/null +++ b/tests/fast/ray/tinker_frontend/test_sdk_contract.py @@ -0,0 +1,384 @@ +"""Contract tests: the REAL, unmodified ``tinker`` SDK (pinned wire behavior +of 0.24.1) drives the frontend over a live localhost HTTP server. + +The stack is the production one minus GPUs and Ray: TinkerFrontendHTTPServer +-> TinkerFrontend -> real MultiLoraOperationBackend (registry + ledger + validation), +executed by the FakeDriver (the documented trainer verbs), sampling proxied +to a stub sglang router. The SDK is never mocked, monkeypatched, or called +below its public surface (the one exception: models.unload is a low-level +``AsyncTinker`` resource because no high-level client exposes it). + +Skipped when the ``tinker`` wheel is not installed (hosted CPU CI); install +``tinker==0.24.1`` to run. +""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=180, suite="stage-a-cpu") + +import asyncio +import threading +from types import SimpleNamespace + +import pytest + +tinker = pytest.importorskip("tinker") + +import uvicorn # noqa: E402 +from tests.fast.ray.tinker_frontend.fake_stack import FakeDriver, FakeRouter, make_backend # noqa: E402 +from tinker import types # noqa: E402 + +from miles.ray.tinker_frontend.http_server import TinkerFrontendHTTPServer # noqa: E402 + +API_KEY = "tml-test-key" +BASE = "Qwen/Qwen3-0.6B" + + +@pytest.fixture(scope="module") +def stack(tmp_path_factory): + loop = asyncio.new_event_loop() + threading.Thread(target=loop.run_forever, daemon=True).start() + + def run(coro, timeout=60): + return asyncio.run_coroutine_threadsafe(coro, loop).result(timeout) + + router = FakeRouter() + router_server = uvicorn.Server( + uvicorn.Config(router.app(), host="127.0.0.1", port=0, log_level="warning", access_log=False) + ) + + router_task: dict = {} + + async def start_router(): + task = asyncio.get_running_loop().create_task(router_server.serve()) + router_task["serve"] = task + while not router_server.started: + if task.done(): + task.result() + await asyncio.sleep(0.01) + return router_server.servers[0].sockets[0].getsockname()[1] + + router_port = run(start_router()) + backend = make_backend( + router_url=f"http://127.0.0.1:{router_port}", + save_root=str(tmp_path_factory.mktemp("tinker-save")), + multi_lora_n_adapters=16, + tinker_api_key=API_KEY, + ) + run(backend.init()) + driver = FakeDriver(backend) + + async def spawn_driver(): + return asyncio.get_running_loop().create_task(driver.run(interval=0.002)) + + driver_task = run(spawn_driver()) + server = TinkerFrontendHTTPServer(backend, host="127.0.0.1", api_port=0) + run(server.start()) + yield SimpleNamespace( + base_url=f"http://127.0.0.1:{server.actual_api_port}", + backend=backend, + driver=driver, + router=router, + frontend=server.frontend, + run=run, + ) + + async def stop_background_tasks(): + driver_task.cancel() + await asyncio.gather(driver_task, return_exceptions=True) + router_server.should_exit = True + await asyncio.gather(router_task["serve"], return_exceptions=True) + + run(stop_background_tasks()) + run(server.stop()) + run(backend.close()) + loop.call_soon_threadsafe(loop.stop) + + +@pytest.fixture() +def service_client(stack): + return tinker.ServiceClient(base_url=stack.base_url, api_key=API_KEY) + + +def make_datum(tokens, weights=None, targets=None): + targets = targets if targets is not None else tokens[1:] + [99] + weights = weights if weights is not None else [1.0] * len(tokens) + return types.Datum( + model_input=types.ModelInput.from_ints(tokens), + loss_fn_inputs={"target_tokens": targets, "weights": weights}, + ) + + +class TestBootstrap: + def test_capabilities_list_the_deployment_base_model(self, service_client): + capabilities = service_client.get_server_capabilities() + assert [m.model_name for m in capabilities.supported_models] == [BASE] + + def test_a_wrong_api_key_is_a_clean_auth_failure(self, stack): + bad = tinker.ServiceClient(base_url=stack.base_url, api_key="tml-wrong-key") + with pytest.raises(Exception, match="401|X-API-Key"): + bad.get_server_capabilities() + + +class TestTrainingChain: + def test_fb_optim_forward_chain(self, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=8) + assert client.get_info().lora_rank == 8 + + data = [make_datum([1, 2, 3]), make_datum([4, 5, 6, 7])] + fb_future = client.forward_backward(data, "cross_entropy") + optim_future = client.optim_step(types.AdamParams(learning_rate=1e-4)) + fb = fb_future.result() + optim = optim_future.result() + + rows = [output["logprobs"].tolist() for output in fb.loss_fn_outputs] + assert rows == [[-0.5] * 3, [-0.5] * 4] + assert fb.metrics["loss:sum"] == pytest.approx(3.5) + assert fb.metrics["unmasked_tokens:sum"] == pytest.approx(7.0) + assert optim.metrics["grad_norm"] == pytest.approx(0.125) + + forward = client.forward([make_datum([1, 2, 3])], "cross_entropy").result() + assert forward.loss_fn_outputs[0]["logprobs"].tolist() == pytest.approx([-0.51] * 3) + assert forward.metrics["loss:sum"] == pytest.approx(1.53) + + def test_multi_chunk_forward_backward_posts_out_of_order(self, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=4) + count = 1030 + data = [make_datum([10, 11]) for _ in range(count)] + result = client.forward_backward(data, "cross_entropy").result() + assert len(result.loss_fn_outputs) == count + assert result.metrics["unmasked_tokens:sum"] == pytest.approx(2.0 * count) + + def test_importance_sampling_and_ppo(self, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=4) + datum = types.Datum( + model_input=types.ModelInput.from_ints([1, 2, 3]), + loss_fn_inputs={ + "target_tokens": [2, 3, 99], + "logprobs": [-0.4, -0.4, -0.4], + "advantages": [0.0, 1.0, 1.0], + }, + ) + is_result = client.forward_backward([datum], "importance_sampling").result() + assert "loss:sum" in is_result.metrics + ppo_result = client.forward_backward( + [datum], "ppo", loss_fn_config={"clip_low_threshold": 0.8, "clip_high_threshold": 1.2} + ).result() + assert "loss:sum" in ppo_result.metrics + + def test_user_error_is_typed_and_leaves_no_gap(self, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=4) + bad = make_datum([1, 2, 3], targets=[9, 3, 99]) + with pytest.raises(tinker.RequestFailedError, match="next input"): + client.forward_backward([bad], "cross_entropy").result() + good = client.forward_backward([make_datum([1, 2, 3])], "cross_entropy").result() + assert len(good.loss_fn_outputs) == 1 + with pytest.raises(tinker.RequestFailedError, match="gradient window"): + client.optim_step(types.AdamParams()).result() + client.forward_backward([make_datum([1, 2, 3])], "cross_entropy").result() + assert client.optim_step(types.AdamParams()).result().metrics["grad_norm"] == pytest.approx(0.125) + + def test_failed_chunk_never_partial_steps_the_window(self, stack, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=4) + data = [make_datum([10, 11]) for _ in range(1024)] + data.append(make_datum([1, 2, 3], targets=[9, 3, 99])) + fb_future = client.forward_backward(data, "cross_entropy") + optim_future = client.optim_step(types.AdamParams(learning_rate=1e-4)) + with pytest.raises(tinker.RequestFailedError, match="next input"): + fb_future.result() + with pytest.raises(tinker.RequestFailedError, match="gradient window"): + optim_future.result() + name = client.model_id.split(":")[0] + [record] = [ + r for n, r in stack.backend.registry.records.items() if r.config.metadata.get("session_id") == name + ] + assert record.step == 0 + + def test_backpressure_429_retries_to_success(self, stack, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=4) + + async def throttle(): + stack.driver.paused = True + stack.backend.operations.max_pending = 1 + + async def release(): + stack.driver.paused = False + stack.backend.operations.max_pending = 256 + + stack.run(throttle()) + try: + fb_future = client.forward_backward([make_datum([1, 2, 3])], "cross_entropy") + optim_future = client.optim_step(types.AdamParams()) + stack.run(asyncio.sleep(0.2)) + finally: + stack.run(release()) + assert len(fb_future.result().loss_fn_outputs) == 1 + assert optim_future.result().metrics["grad_norm"] == pytest.approx(0.125) + + +class TestCheckpoints: + def test_save_then_resume_with_optimizer(self, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=8) + client.forward_backward([make_datum([1, 2, 3])], "cross_entropy").result() + client.optim_step(types.AdamParams()).result() + path = client.save_state("resume-me").result().path + assert path.startswith("tinker://") and path.endswith("/weights/resume-me") + + resumed = service_client.create_training_client_from_state_with_optimizer(path) + assert resumed.get_info().lora_rank == 8 + result = resumed.forward_backward([make_datum([1, 2, 3])], "cross_entropy").result() + assert result.loss_fn_outputs[0]["logprobs"].tolist() == pytest.approx([-0.51] * 3) + + def test_weights_only_resume_is_a_typed_rejection(self, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=8) + path = client.save_state("no-optim").result().path + with pytest.raises(tinker.RequestFailedError, match="weights-only"): + service_client.create_training_client_from_state(path) + + def test_immutable_states_and_load_after_unload(self, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=8) + client.save_state("once").result() + with pytest.raises(tinker.RequestFailedError, match="immutable"): + client.save_state("once").result() + + +class TestSampling: + def test_publish_then_sample(self, stack, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=8) + client.forward_backward([make_datum([1, 2, 3])], "cross_entropy").result() + client.optim_step(types.AdamParams()).result() + sampling = client.save_weights_and_get_sampling_client() + response = sampling.sample( + prompt=types.ModelInput.from_ints([5, 6, 7]), + num_samples=2, + sampling_params=types.SamplingParams(max_tokens=3, temperature=0.5, top_p=0.9), + ).result() + assert len(response.sequences) == 2 + for sequence in response.sequences: + assert sequence.tokens == [1000, 1001, 1002] + assert sequence.logprobs == [-0.25, -0.5, -0.75] + assert sequence.stop_reason == "length" + generated = stack.router.requests[-1] + assert generated["lora_path"].startswith("__miles_adapter_") + assert generated["extra_key"].endswith(":v1") + assert generated["sampling_params"] == { + "max_new_tokens": 3, + "temperature": 0.5, + "top_p": 0.9, + "top_k": -1, + } + assert sampling.get_base_model() == BASE + + def test_base_model_sampling_session(self, stack, service_client): + sampling = service_client.create_sampling_client(base_model=BASE) + response = sampling.sample( + prompt=types.ModelInput.from_ints([8]), + num_samples=1, + sampling_params=types.SamplingParams(max_tokens=2), + ).result() + assert response.sequences[0].tokens == [1000, 1001] + assert "lora_path" not in stack.router.requests[-1] + + def test_compute_logprobs_scores_every_prompt_token(self, stack, service_client): + sampling = service_client.create_sampling_client(base_model=BASE) + prompt = [5, 6, 7, 8] + logprobs = sampling.compute_logprobs(types.ModelInput.from_ints(prompt)).result() + assert logprobs == [None, -0.125, -0.25, -0.375] + assert len(logprobs) == len(prompt) + assert all(isinstance(lp, float) for lp in logprobs[1:]) + sent = stack.router.requests[-1] + assert sent["input_ids"] == prompt + assert sent["logprob_start_len"] == 0 and sent["return_logprob"] is True + assert sent["sampling_params"]["max_new_tokens"] == 1 + + def test_sample_with_prompt_logprobs_returns_both(self, service_client): + sampling = service_client.create_sampling_client(base_model=BASE) + response = sampling.sample( + prompt=types.ModelInput.from_ints([5, 6, 7]), + num_samples=2, + sampling_params=types.SamplingParams(max_tokens=3), + include_prompt_logprobs=True, + ).result() + assert len(response.sequences) == 2 + assert response.sequences[0].tokens == [1000, 1001, 1002] + assert response.prompt_logprobs == [None, -0.125, -0.25] + + def test_topk_prompt_logprobs_is_a_typed_rejection(self, service_client): + sampling = service_client.create_sampling_client(base_model=BASE) + future = sampling.sample( + prompt=types.ModelInput.from_ints([5, 6]), + num_samples=1, + sampling_params=types.SamplingParams(max_tokens=2), + topk_prompt_logprobs=2, + ) + with pytest.raises(tinker.RequestFailedError, match="topk_prompt_logprobs"): + future.result() + + def test_stale_ephemeral_sampler_fails_loud_after_republish(self, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=8) + old = client.save_weights_and_get_sampling_client() + client.save_weights_and_get_sampling_client() + future = old.sample( + prompt=types.ModelInput.from_ints([5]), + num_samples=1, + sampling_params=types.SamplingParams(max_tokens=2), + ) + with pytest.raises(tinker.RequestFailedError, match="republished"): + future.result() + + def test_oversized_context_is_a_typed_rejection_not_silent_truncation(self, stack, service_client): + sampling = service_client.create_sampling_client(base_model=BASE) + small = sampling.sample( + prompt=types.ModelInput.from_ints([9]), + num_samples=1, + sampling_params=types.SamplingParams(max_tokens=2), + ).result() + assert small.sequences[0].tokens == [1000, 1001] + + async def discovered(): + for _ in range(200): + if stack.frontend._context_limit is not None: + return stack.frontend._context_limit + await asyncio.sleep(0.01) + raise TimeoutError("context discovery never landed") + + assert stack.run(discovered()) == stack.router.max_req_input_len + 6 == 4096 + + with pytest.raises(Exception, match="context limit of 4096"): + sampling.sample( + prompt=types.ModelInput.from_ints(list(range(1, 4001))), + num_samples=1, + sampling_params=types.SamplingParams(max_tokens=2048), + ).result() + again = sampling.sample( + prompt=types.ModelInput.from_ints([11]), + num_samples=1, + sampling_params=types.SamplingParams(max_tokens=2), + ).result() + assert again.sequences[0].stop_reason in ("length", "stop") + + +class TestUnload: + def test_low_level_unload_retires_the_registration(self, stack, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=4) + model_id = client.model_id + + async def unload_and_poll(): + from tinker._client import AsyncTinker + + low_level = AsyncTinker(base_url=stack.base_url, api_key=API_KEY) + future = await low_level.models.unload(request=types.UnloadModelRequest(model_id=model_id)) + for _ in range(200): + raw = await low_level.futures.with_raw_response.retrieve( + request=types.FutureRetrieveRequest(request_id=future.request_id) + ) + body = await raw.json() + if body.get("type") != "try_again": + return body + await asyncio.sleep(0.02) + raise TimeoutError("unload future never resolved") + + body = stack.run(unload_and_poll()) + assert body == {"type": "unload_model", "model_id": model_id} + with pytest.raises(tinker.RequestFailedError): + client.forward_backward([make_datum([1, 2, 3])], "cross_entropy").result() diff --git a/tests/fast/ray/tinker_frontend/test_sdk_sampling_saturation.py b/tests/fast/ray/tinker_frontend/test_sdk_sampling_saturation.py new file mode 100644 index 00000000000..862072c0350 --- /dev/null +++ b/tests/fast/ray/tinker_frontend/test_sdk_sampling_saturation.py @@ -0,0 +1,156 @@ +"""Aggregate-saturation regression with the REAL tinker SDK over live HTTP: +two SamplingClients each admit 64 concurrent requests (the SDK's per-client +ceiling — it never bounded the aggregate), against the PRODUCTION router +transport and a shared slow router that holds every generation longer than +the old implicit 10-second pool deadline. + +Before the admission/transport fix this exact load was the Tau sampling +cliff: the shared httpx client's default 100-connection pool timed request +#101+ out before it ever reached the router — exactly 100/128 succeeded and +28 died as terminal, empty-message failures ("sampling failed: ") the SDK +never retries. Now the frontend 429s the overflow BEFORE the request +consumes its seq identity, the SDK retries on the same seq ids with backoff, +and all 128 complete exactly once inside the configured bound. + +Skipped when the ``tinker`` wheel is not installed; install tinker==0.24.1 +(pinned in tests/ci/requirements-ci-cpu.txt) to run.""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=180, suite="stage-a-cpu") + +import asyncio +import logging + +import pytest + +tinker = pytest.importorskip("tinker") + +import uvicorn # noqa: E402 +from fastapi import FastAPI, Request # noqa: E402 +from tests.fast.ray.tinker_frontend.fake_stack import FakeDriver, make_backend # noqa: E402 +from tinker import types # noqa: E402 + +from miles.ray.tinker_frontend.http_server import TinkerFrontendHTTPServer # noqa: E402 + +API_KEY = "tml-test-key" +BASE = "Qwen/Qwen3-0.6B" +CLIENTS = 2 +PER_CLIENT = 64 +ROUTER_DELAY_S = 11.0 +CAP = 64 + + +class SlowRouter: + def __init__(self, delay_s: float) -> None: + self.delay_s = delay_s + self.calls = 0 + self.active = 0 + self.max_active = 0 + + def app(self) -> FastAPI: + app = FastAPI() + + @app.post("/generate") + async def generate(request: Request) -> dict: + payload = await request.json() + self.calls += 1 + self.active += 1 + self.max_active = max(self.max_active, self.active) + try: + await asyncio.sleep(self.delay_s) + return { + "text": "ok", + "meta_info": { + "finish_reason": {"type": "length"}, + "output_token_logprobs": [[-0.25, int(payload["input_ids"][0]) + 10_000, None]], + "prompt_tokens": len(payload["input_ids"]), + }, + } + finally: + self.active -= 1 + + return app + + +def test_aggregate_sdk_load_completes_within_the_bound_instead_of_the_pool_cliff(tmp_path): + logging.getLogger("tinker.lib.retry_handler").setLevel(logging.CRITICAL) + logging.getLogger("tinker.lib.api_future_impl").setLevel(logging.ERROR) + + async def main(): + router = SlowRouter(ROUTER_DELAY_S) + router_server = uvicorn.Server( + uvicorn.Config(router.app(), host="127.0.0.1", port=0, log_level="critical", access_log=False) + ) + serve_task = asyncio.get_running_loop().create_task(router_server.serve()) + while not router_server.started: + if serve_task.done(): + serve_task.result() + await asyncio.sleep(0.005) + router_port = router_server.servers[0].sockets[0].getsockname()[1] + + backend = make_backend( + router_url=f"http://127.0.0.1:{router_port}", + save_root=str(tmp_path), + multi_lora_n_adapters=16, + tinker_api_key=API_KEY, + ) + await backend.init() + FakeDriver(backend) + server = TinkerFrontendHTTPServer(backend, host="127.0.0.1", api_port=0) + await server.start() + frontend = server.frontend + assert frontend.sampling_admission.capacity == CAP + try: + base_url = f"http://127.0.0.1:{server.actual_api_port}" + service = await asyncio.to_thread(tinker.ServiceClient, base_url=base_url, api_key=API_KEY) + clients = [ + await asyncio.to_thread(service.create_sampling_client, base_model=BASE) for _ in range(CLIENTS) + ] + holder = service._session_holder + session = frontend.sessions.get(holder._session_id) + heartbeat_before = session.last_heartbeat + + params = types.SamplingParams(max_tokens=1, seed=7) + tasks = [ + asyncio.create_task( + client.sample_async( + prompt=types.ModelInput.from_ints([1_000 + index * PER_CLIENT + i]), + num_samples=1, + sampling_params=params, + ) + ) + for index, client in enumerate(clients) + for i in range(PER_CLIENT) + ] + outcomes = await asyncio.gather(*tasks, return_exceptions=True) + + failures = [item for item in outcomes if isinstance(item, BaseException)] + assert not failures, [f"{type(item).__name__}: {item}" for item in failures[:3]] + assert sum(isinstance(item, types.SampleResponse) for item in outcomes) == CLIENTS * PER_CLIENT + + assert frontend.sampling_admission.rejected > 0 + assert router.max_active <= CAP + assert router.calls == CLIENTS * PER_CLIENT + + for client in clients: + record = frontend.samplers.get(client._sampling_session_id) + assert record.spent_fence == PER_CLIENT - 1 and not record.spent_sparse + + for _ in range(200): + if frontend.sampling_admission.in_use == 0 and not frontend._sample_tasks: + break + await asyncio.sleep(0.01) + assert frontend.sampling_admission.in_use == 0 + assert not frontend._sample_tasks + + assert session.last_heartbeat > heartbeat_before + holder.close() + await asyncio.sleep(0.05) + finally: + await server.stop() + await backend.close() + router_server.should_exit = True + await asyncio.gather(serve_task, return_exceptions=True) + + asyncio.run(main()) diff --git a/tests/fast/ray/tinker_frontend/test_sdk_sft_contract.py b/tests/fast/ray/tinker_frontend/test_sdk_sft_contract.py new file mode 100644 index 00000000000..cafd062d8e1 --- /dev/null +++ b/tests/fast/ray/tinker_frontend/test_sdk_sft_contract.py @@ -0,0 +1,190 @@ +"""SFT-only contract probes: the REAL, unmodified ``tinker==0.24.1`` SDK +drives the live HTTP stack through the teacher-forced cross-entropy path — +accumulation windows, prompt masking, checkpoint gating, rejected-Adam +recovery — plus the two verified pre-HTTP SDK failure modes and what the +server does (gap timeout) and cannot do (immediate cancel) about them. + +Permanent adaptation of the codex-0817-sft-fix §3.2 adversarial suite; the +stack fixture (frontend -> real backend -> FakeDriver) comes from +test_sdk_contract. Skipped when the ``tinker`` wheel is not installed.""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=120, suite="stage-a-cpu") + +from concurrent.futures import TimeoutError as FutureTimeoutError # noqa: E402 + +import pytest # noqa: E402 + +tinker = pytest.importorskip("tinker") + +from tests.fast.ray.tinker_frontend import test_sdk_contract as sdk_contract # noqa: E402 +from tinker import types # noqa: E402 + +from miles.ray.multi_lora.operations import SealedGap # noqa: E402 + +BASE = sdk_contract.BASE +make_datum = sdk_contract.make_datum +stack = sdk_contract.stack +service_client = sdk_contract.service_client + + +def _record_for(stack, client): + session = client.model_id.split(":", 1)[0] + [record] = [ + record + for record in stack.backend.registry.records.values() + if record.config.metadata.get("session_id") == session + ] + return record + + +async def _set_driver_paused(stack, paused): + stack.driver.paused = paused + + +def sft_datum(prompt_tokens, completion_tokens): + tokens = prompt_tokens + completion_tokens + return types.Datum( + model_input=types.ModelInput.from_ints(tokens[:-1]), + loss_fn_inputs={ + "target_tokens": tokens[1:], + "weights": [0.0] * (len(prompt_tokens) - 1) + [1.0] * len(completion_tokens), + }, + ) + + +class TestSftTrainingContract: + def test_three_fb_accumulate_then_one_optim(self, stack, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=8) + fbs = [client.forward_backward([make_datum([10 + i, 20 + i, 30 + i])], "cross_entropy") for i in range(3)] + optim = client.optim_step(types.AdamParams(learning_rate=2e-4)) + + results = [future.result() for future in fbs] + step = optim.result() + + assert [result.metrics["loss:sum"] for result in results] == pytest.approx([1.5, 1.5, 1.5]) + assert step.metrics["learning_rate"] == pytest.approx(2e-4) + assert _record_for(stack, client).step == 1 + forward = client.forward([make_datum([1, 2, 3])], "cross_entropy").result() + assert forward.loss_fn_outputs[0]["logprobs"].tolist() == pytest.approx([-0.51] * 3) + + def test_prompt_masked_sft_datum_separates_the_two_denominators(self, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=4) + result = client.forward_backward([sft_datum([11, 12, 13, 14], [15, 16, 17])], "cross_entropy").result() + assert result.metrics["unmasked_tokens:sum"] == pytest.approx(6.0) + assert result.metrics["loss_weight:sum"] == pytest.approx(3.0) + assert result.metrics["loss:sum"] / result.metrics["loss_weight:sum"] == pytest.approx(0.5) + client.optim_step(types.AdamParams(learning_rate=0.0)).result() + + def test_zero_weight_prefix_and_fractional_ce_weights(self, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=4) + datum = make_datum( + [10, 11, 12, 13], + targets=[999, 12, 888, 77], + weights=[0.0, 0.5, 0.0, 2.0], + ) + result = client.forward_backward([datum], "cross_entropy").result() + + assert result.loss_fn_outputs[0]["logprobs"].tolist() == pytest.approx([-0.5] * 4) + assert result.metrics["loss:sum"] == pytest.approx(1.25) + assert result.metrics["unmasked_tokens:sum"] == pytest.approx(4.0) + assert result.metrics["loss_weight:sum"] == pytest.approx(2.5) + client.optim_step(types.AdamParams(learning_rate=0.0)).result() + + def test_dirty_save_rejection_preserves_gradients_for_later_step(self, stack, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=4) + client.forward_backward([make_datum([1, 2, 3])], "cross_entropy").result() + + with pytest.raises(tinker.RequestFailedError, match="unstepped gradients"): + client.save_state("must-not-save-dirty").result() + + result = client.optim_step(types.AdamParams(learning_rate=3e-4)).result() + assert result.metrics["grad_norm"] == pytest.approx(0.125) + assert _record_for(stack, client).step == 1 + assert client.save_state("clean-after-step").result().path.endswith("/weights/clean-after-step") + + def test_rejected_adam_params_do_not_drop_prior_gradients(self, stack, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=4) + client.forward_backward([make_datum([1, 2, 3])], "cross_entropy").result() + + with pytest.raises(tinker.RequestFailedError, match="learning_rate.*>= 0"): + client.optim_step(types.AdamParams(learning_rate=-1.0)).result() + + assert _record_for(stack, client).step == 0 + result = client.optim_step(types.AdamParams(learning_rate=1e-4)).result() + assert result.metrics["grad_norm"] == pytest.approx(0.125) + assert _record_for(stack, client).step == 1 + + def test_forward_is_no_grad_and_checkpointable(self, stack, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=4) + result = client.forward([make_datum([1, 2, 3])], "cross_entropy").result() + assert result.metrics["loss:sum"] == pytest.approx(1.5) + assert _record_for(stack, client).step == 0 + assert client.save_state("after-forward").result().path.endswith("/weights/after-forward") + + def test_result_timeout_is_non_destructive(self, stack, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=4) + stack.run(_set_driver_paused(stack, True)) + future = client.forward_backward([make_datum([1, 2, 3])], "cross_entropy") + try: + with pytest.raises(FutureTimeoutError): + future.result(timeout=0.02) + finally: + stack.run(_set_driver_paused(stack, False)) + assert future.result(timeout=5).metrics["loss:sum"] == pytest.approx(1.5) + client.optim_step(types.AdamParams(learning_rate=0.0)).result() + + +class TestPreHttpSdkFailureModes: + def test_pre_http_serialization_hole_gap_times_out_typed_then_the_same_client_resubmits( + self, stack, service_client + ): + client = service_client.create_lora_training_client(base_model=BASE, rank=4) + client.forward_backward([make_datum([1, 2, 3])], "cross_entropy").result() + + async def set_gap_timeout(value): + previous = stack.backend.operations.gap_timeout + stack.backend.operations.gap_timeout = value + return previous + + original = stack.run(set_gap_timeout(0.3)) + try: + bad = client.optim_step(types.AdamParams(learning_rate=float("nan"))) + with pytest.raises(ValueError, match="Out of range float values|JSON compliant"): + bad.result(timeout=2) + + later = client.optim_step(types.AdamParams(learning_rate=1e-4)) + with pytest.raises(tinker.RequestFailedError, match="missing ordinal 2"): + later.result(timeout=30) + finally: + stack.run(set_gap_timeout(original)) + + record = _record_for(stack, client) + assert record.step == 0 + + result = client.optim_step(types.AdamParams(learning_rate=1e-4)).result(timeout=30) + assert result.metrics["grad_norm"] == pytest.approx(0.125) + assert record.step == 1 + + async def sealed_ordinals(): + queue = stack.backend.operations.queues[(record.name, record.registration_id)] + return [ordinal for ordinal, holder in queue.by_ordinal.items() if isinstance(holder, SealedGap)] + + assert stack.run(sealed_ordinals()) == [2] + + def test_immediate_sdk_future_cancel_spends_turn_and_wedges_later_work(self, stack, service_client): + client = service_client.create_lora_training_client(base_model=BASE, rank=4) + cancelled = [] + for _ in range(32): + future = client.forward_backward([make_datum([1, 2, 3])], "cross_entropy") + cancelled.append(future.future().cancel()) + assert any(cancelled) + + later = client.optim_step(types.AdamParams(learning_rate=0.0)) + with pytest.raises(FutureTimeoutError): + later.result(timeout=0.5) + later.future().cancel() + + record = _record_for(stack, client) + assert stack.backend.operations.queue_view(record.name, record.registration_id) == [] diff --git a/tests/fast/ray/tinker_frontend/test_service.py b/tests/fast/ray/tinker_frontend/test_service.py new file mode 100644 index 00000000000..f799f9804f8 --- /dev/null +++ b/tests/fast/ray/tinker_frontend/test_service.py @@ -0,0 +1,762 @@ +"""TinkerFrontend against a real backend + fake driver: the future protocol, +seq->ordinal mapping (incl. out-of-order chunk arrival and rejected-seq +contiguity), idempotent retries, checkpoints, publish->sample, and fences.""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=120, suite="stage-a-cpu") + +import asyncio + +import pytest +from tests.fast.ray.tinker_frontend.fake_stack import FakeDriver, FakeRouter, make_backend + +from miles.ray.multi_lora.config import AdapterRunConfig +from miles.ray.multi_lora.operations import OperationBackpressure +from miles.ray.tinker_frontend import wire +from miles.ray.tinker_frontend.service import ApiError, TinkerFrontend + +BASE = "Qwen/Qwen3-0.6B" + + +class Stack: + def __init__(self, frontend, driver, router): + self.frontend = frontend + self.driver = driver + self.router = router + self.session_id = frontend.create_session(wire.CreateSessionRequest(sdk_version="0.24.1"))["session_id"] + + async def create_model(self, model_seq_id=0, rank=8, **lora_overrides): + request = wire.CreateModelRequest( + session_id=self.session_id, + model_seq_id=model_seq_id, + base_model=BASE, + lora_config=wire.LoraConfig(rank=rank, **lora_overrides), + ) + future = await self.frontend.create_model(request) + body = await self.retrieve(future["request_id"]) + assert body == {"type": "create_model", "model_id": f"{self.session_id}:train:{model_seq_id}"} + return body["model_id"] + + async def retrieve(self, request_id): + return await self.frontend.retrieve_future(wire.FutureRetrieveRequest(request_id=request_id)) + + def fb_request(self, model_id, seq_id, tokens=(1, 2, 3), weights=(0.0, 1.0, 1.0), targets=None): + targets = targets if targets is not None else list(tokens[1:]) + [99] + return wire.ForwardBackwardRequest.model_validate( + { + "forward_backward_input": { + "data": [ + { + "model_input": {"chunks": [{"type": "encoded_text", "tokens": list(tokens)}]}, + "loss_fn_inputs": { + "target_tokens": {"data": targets, "dtype": "int64", "shape": [len(targets)]}, + "weights": {"data": list(weights), "dtype": "float32", "shape": [len(weights)]}, + }, + } + ], + "loss_fn": "cross_entropy", + }, + "model_id": model_id, + "seq_id": seq_id, + } + ) + + def optim_request(self, model_id, seq_id, lr=1e-4): + return wire.OptimStepRequest.model_validate( + {"adam_params": {"learning_rate": lr}, "model_id": model_id, "seq_id": seq_id} + ) + + +class RouterSamplingTransport: + def __init__(self, router): + self.router = router + self.closed = False + + async def generate(self, payload: dict) -> dict: + self.router.requests.append(payload) + return self.router.response_for(payload) + + async def close(self) -> None: + self.closed = True + + +def run(scenario, poll_window_s=5.0, **backend_overrides): + async def main(): + router = FakeRouter() + backend = make_backend(**backend_overrides) + await backend.init() + driver = FakeDriver(backend) + frontend = TinkerFrontend( + backend, + poll_window_s=poll_window_s, + poll_interval_s=0.002, + sampling_transport=RouterSamplingTransport(router), + ) + stack = Stack(frontend, driver, router) + driver_task = asyncio.create_task(driver.run(interval=0.002)) + try: + await asyncio.wait_for(scenario(stack), timeout=30) + finally: + driver_task.cancel() + await asyncio.gather(driver_task, return_exceptions=True) + await frontend.close() + await backend.close() + + asyncio.run(main()) + + +class TestTrainingChain: + def test_out_of_order_chunks_then_optim(self): + async def scenario(stack): + model_id = await stack.create_model() + fb2 = stack.frontend.forward_backward(stack.fb_request(model_id, 2, tokens=(5, 6, 7))) + fb1 = stack.frontend.forward_backward(stack.fb_request(model_id, 1)) + optim = stack.frontend.optim_step(stack.optim_request(model_id, 3)) + body1 = await stack.retrieve(fb1["request_id"]) + body2 = await stack.retrieve(fb2["request_id"]) + body3 = await stack.retrieve(optim["request_id"]) + for body in (body1, body2): + (row,) = [output["logprobs"]["data"] for output in body["loss_fn_outputs"]] + assert row == [-0.5, -0.5, -0.5] + assert body["metrics"]["loss:sum"] == pytest.approx(1.0) + assert body["metrics"]["unmasked_tokens:sum"] == pytest.approx(3.0) + assert body3 == {"type": "optim_step", "metrics": {"grad_norm": 0.125, "learning_rate": 1e-4}} + fb4 = stack.frontend.forward_backward(stack.fb_request(model_id, 4)) + body4 = await stack.retrieve(fb4["request_id"]) + assert body4["loss_fn_outputs"][0]["logprobs"]["data"] == [-0.51, -0.51, -0.51] + + run(scenario) + + def test_forward_recomputes_metrics_and_takes_no_dirty_pin(self): + async def scenario(stack): + model_id = await stack.create_model() + forward = stack.frontend.forward( + wire.ForwardRequest.model_validate( + { + **stack.fb_request(model_id, 1).model_dump(exclude={"forward_backward_input"}), + "forward_input": stack.fb_request(model_id, 1).forward_backward_input.model_dump(), + "seq_id": 1, + } + ) + ) + body = await stack.retrieve(forward["request_id"]) + assert body["metrics"]["loss:sum"] == pytest.approx(1.0) + save = stack.frontend.save_weights( + wire.SaveWeightsRequest(model_id=model_id, path="after-forward", seq_id=2) + ) + saved = await stack.retrieve(save["request_id"]) + assert saved["type"] == "save_weights" + + run(scenario) + + def test_idempotent_retry_and_conflict(self): + async def scenario(stack): + model_id = await stack.create_model() + request = stack.fb_request(model_id, 1) + first = stack.frontend.forward_backward(request) + again = stack.frontend.forward_backward(request) + assert again == first + with pytest.raises(ApiError) as excinfo: + stack.frontend.forward_backward(stack.fb_request(model_id, 1, tokens=(7, 8, 9))) + assert excinfo.value.status_code == 422 + body = await stack.retrieve(first["request_id"]) + replay = await stack.retrieve(first["request_id"]) + assert replay == body + + run(scenario) + + def test_rejected_seq_still_consumes_its_ordinal(self): + async def scenario(stack): + model_id = await stack.create_model() + fb1 = stack.frontend.forward_backward(stack.fb_request(model_id, 1)) + bad = stack.frontend.forward_backward( + stack.fb_request(model_id, 2, tokens=(1, 2, 3), weights=(1.0, 1.0, 1.0), targets=[9, 3, 99]) + ) + fb3 = stack.frontend.forward_backward(stack.fb_request(model_id, 3)) + failed = await stack.retrieve(bad["request_id"]) + assert failed["category"] == "user" and "next input" in failed["error"] + assert (await stack.retrieve(fb3["request_id"]))["type"] == "forward_backward" + assert (await stack.retrieve(fb1["request_id"]))["type"] == "forward_backward" + + run(scenario) + + def test_failed_chunk_poisons_the_gradient_window(self): + async def scenario(stack): + model_id = await stack.create_model() + good = stack.frontend.forward_backward(stack.fb_request(model_id, 1)) + bad = stack.frontend.forward_backward( + stack.fb_request(model_id, 2, weights=(1.0, 1.0, 1.0), targets=[9, 3, 99]) + ) + optim = stack.frontend.optim_step(stack.optim_request(model_id, 3)) + assert (await stack.retrieve(good["request_id"]))["type"] == "forward_backward" + assert (await stack.retrieve(bad["request_id"]))["category"] == "user" + poisoned = await stack.retrieve(optim["request_id"]) + assert poisoned["category"] == "user" and "gradient window" in poisoned["error"] + record = stack.frontend.backend.registry.find(stack.frontend.models.get(model_id).name) + assert record.step == 0 + + fb4 = stack.frontend.forward_backward(stack.fb_request(model_id, 4)) + optim5 = stack.frontend.optim_step(stack.optim_request(model_id, 5)) + assert (await stack.retrieve(fb4["request_id"]))["type"] == "forward_backward" + assert (await stack.retrieve(optim5["request_id"]))["type"] == "optim_step" + assert record.step == 1 + + run(scenario) + + def test_backpressure_is_retryable_not_terminal(self): + async def scenario(stack): + model_id = await stack.create_model() + stack.driver.paused = True + stack.frontend.backend.operations.max_pending = 1 + stack.frontend.forward_backward(stack.fb_request(model_id, 1)) + with pytest.raises(OperationBackpressure): + stack.frontend.optim_step(stack.optim_request(model_id, 2)) + stack.driver.paused = False + for _ in range(500): + try: + retried = stack.frontend.optim_step(stack.optim_request(model_id, 2)) + break + except OperationBackpressure: + await asyncio.sleep(0.005) + assert (await stack.retrieve(retried["request_id"]))["type"] == "optim_step" + + run(scenario) + + +class TestCheckpoints: + def test_save_load_roundtrip_mints_and_resolves_tinker_paths(self): + async def scenario(stack): + model_id = await stack.create_model() + save = stack.frontend.save_weights(wire.SaveWeightsRequest(model_id=model_id, path="ckpt-0", seq_id=1)) + saved = await stack.retrieve(save["request_id"]) + path = saved["path"] + assert path.startswith("tinker://") and path.endswith("/weights/ckpt-0") + info = stack.frontend.weights_info(wire.WeightsInfoRequest(tinker_path=path)) + assert info == { + "base_model": BASE, + "is_lora": True, + "lora_rank": 8, + "train_unembed": None, + "train_mlp": None, + "train_attn": None, + } + load = stack.frontend.load_weights( + wire.LoadWeightsRequest(model_id=model_id, path=path, optimizer=True, seq_id=2) + ) + loaded = await stack.retrieve(load["request_id"]) + assert loaded == {"type": "load_weights", "path": path, "model_id": model_id} + + run(scenario) + + def test_weights_only_load_is_a_typed_user_failure_without_a_gap(self): + async def scenario(stack): + model_id = await stack.create_model() + save = stack.frontend.save_weights(wire.SaveWeightsRequest(model_id=model_id, path="s0", seq_id=1)) + path = (await stack.retrieve(save["request_id"]))["path"] + load = stack.frontend.load_weights( + wire.LoadWeightsRequest(model_id=model_id, path=path, optimizer=False, seq_id=2) + ) + failed = await stack.retrieve(load["request_id"]) + assert failed["category"] == "user" and "weights-only" in failed["error"] + fb = stack.frontend.forward_backward(stack.fb_request(model_id, 3)) + assert (await stack.retrieve(fb["request_id"]))["type"] == "forward_backward" + + run(scenario) + + def test_ttl_is_a_typed_rejection_no_reaper_runs(self): + async def scenario(stack): + model_id = await stack.create_model() + save = stack.frontend.save_weights( + wire.SaveWeightsRequest(model_id=model_id, path="t0", seq_id=1, ttl_seconds=3600) + ) + failed = await stack.retrieve(save["request_id"]) + assert failed["category"] == "user" and "ttl_seconds" in failed["error"] + + run(scenario) + + def test_load_failure_redacts_the_backend_path(self): + async def scenario(stack): + model_id = await stack.create_model() + save = stack.frontend.save_weights(wire.SaveWeightsRequest(model_id=model_id, path="lost", seq_id=1)) + tinker_path = (await stack.retrieve(save["request_id"]))["path"] + backend_path = stack.frontend.checkpoints.get(tinker_path).backend_path + del stack.driver.saved_states[backend_path] + load = stack.frontend.load_weights( + wire.LoadWeightsRequest(model_id=model_id, path=tinker_path, optimizer=True, seq_id=2) + ) + failed = await stack.retrieve(load["request_id"]) + assert failed["category"] == "user" + assert tinker_path in failed["error"] and backend_path not in failed["error"] + + run(scenario) + + def test_overwrite_and_unknown_paths_are_typed_rejections(self): + async def scenario(stack): + model_id = await stack.create_model() + save = stack.frontend.save_weights( + wire.SaveWeightsRequest(model_id=model_id, path="x", seq_id=1, overwrite=True) + ) + assert "immutable" in (await stack.retrieve(save["request_id"]))["error"] + load = stack.frontend.load_weights( + wire.LoadWeightsRequest(model_id=model_id, path="tinker://nope/weights/x", optimizer=True, seq_id=2) + ) + assert "unknown checkpoint" in (await stack.retrieve(load["request_id"]))["error"] + with pytest.raises(ApiError) as excinfo: + stack.frontend.weights_info(wire.WeightsInfoRequest(tinker_path="tinker://nope/weights/x")) + assert excinfo.value.status_code == 404 + + run(scenario) + + +class TestSampling: + async def publish(self, stack, model_id, seq_id, sampling_session_seq_id): + publish = stack.frontend.save_weights_for_sampler( + wire.SaveWeightsForSamplerRequest( + model_id=model_id, seq_id=seq_id, sampling_session_seq_id=sampling_session_seq_id + ) + ) + body = await stack.retrieve(publish["request_id"]) + assert body["type"] == "save_weights_for_sampler" and body["path"] is None + return body["sampling_session_id"] + + def sample_request(self, sampler_id, seq_id=0, num_samples=1, **params): + return wire.SampleRequest.model_validate( + { + "sampling_session_id": sampler_id, + "seq_id": seq_id, + "num_samples": num_samples, + "prompt": {"chunks": [{"type": "encoded_text", "tokens": [5, 6]}]}, + "sampling_params": {"max_tokens": 3, **params}, + } + ) + + def test_publish_then_sample_carries_serving_identity(self): + async def scenario(stack): + model_id = await stack.create_model() + sampler_id = await self.publish(stack, model_id, seq_id=1, sampling_session_seq_id=0) + future = stack.frontend.sample(self.sample_request(sampler_id, num_samples=2)) + body = await stack.retrieve(future["request_id"]) + assert body["type"] == "sample" and len(body["sequences"]) == 2 + assert body["sequences"][0]["tokens"] == [1000, 1001, 1002] + request = stack.router.requests[0] + assert request["lora_path"].startswith("__miles_adapter_") + assert request["extra_key"].endswith(":v1") + assert request["return_logprob"] is True + info = stack.frontend.get_sampler(sampler_id) + assert info["base_model"] == BASE + + run(scenario) + + def test_client_supplied_routing_identity_never_reaches_the_router(self): + async def scenario(stack): + model_id = await stack.create_model() + sampler_id = await self.publish(stack, model_id, seq_id=1, sampling_session_seq_id=0) + request = wire.SampleRequest.model_validate( + { + "sampling_session_id": sampler_id, + "seq_id": 0, + "num_samples": 1, + "prompt": {"chunks": [{"type": "encoded_text", "tokens": [5, 6]}]}, + "sampling_params": {"max_tokens": 3, "lora_path": "../../pwn", "extra_key": "x", "rid": "x"}, + "lora_path": "../../pwn", + "extra_key": "hijacked", + "rid": "chosen-rid", + } + ) + future = stack.frontend.sample(request) + body = await stack.retrieve(future["request_id"]) + assert body["type"] == "sample" + sent = stack.router.requests[0] + assert sent["lora_path"].startswith("__miles_adapter_") + assert sent["extra_key"] != "hijacked" and sent["rid"] != "chosen-rid" + assert set(sent["sampling_params"]) == {"max_new_tokens", "temperature", "top_p", "top_k"} + + run(scenario) + + def test_republish_makes_the_old_session_fail_loud(self): + async def scenario(stack): + model_id = await stack.create_model() + old = await self.publish(stack, model_id, seq_id=1, sampling_session_seq_id=0) + await self.publish(stack, model_id, seq_id=2, sampling_session_seq_id=1) + future = stack.frontend.sample(self.sample_request(old)) + body = await stack.retrieve(future["request_id"]) + assert body["category"] == "user" and "republished" in body["error"] + + run(scenario) + + def test_republish_mid_generation_fails_the_inflight_sample(self): + async def scenario(stack): + model_id = await stack.create_model() + sampler_id = await self.publish(stack, model_id, seq_id=1, sampling_session_seq_id=0) + name = stack.frontend.samplers.get(sampler_id).name + + gate = asyncio.Event() + transport = stack.frontend.sampling_transport + original = transport.generate + + async def delayed(payload): + await gate.wait() + return await original(payload) + + transport.generate = delayed + future = stack.frontend.sample(self.sample_request(sampler_id)) + await asyncio.sleep(0.02) + stack.frontend.backend.registry.record_weight_update([name]) + gate.set() + body = await stack.retrieve(future["request_id"]) + assert body["category"] == "user" and "republished while this sample was in flight" in body["error"] + + run(scenario) + + def test_named_sampler_path_is_a_typed_rejection(self): + async def scenario(stack): + model_id = await stack.create_model() + publish = stack.frontend.save_weights_for_sampler( + wire.SaveWeightsForSamplerRequest(model_id=model_id, seq_id=1, path="final") + ) + body = await stack.retrieve(publish["request_id"]) + assert body["category"] == "user" and "latest-only" in body["error"] + + run(scenario) + + def test_base_model_session_and_unsupported_probes(self): + async def scenario(stack): + request = wire.CreateSamplingSessionRequest( + session_id=stack.session_id, sampling_session_seq_id=0, base_model=BASE + ) + sampler_id = stack.frontend.create_sampling_session(request)["sampling_session_id"] + assert stack.frontend.create_sampling_session(request)["sampling_session_id"] == sampler_id + future = stack.frontend.sample(self.sample_request(sampler_id, num_samples=2, seed=40)) + body = await stack.retrieve(future["request_id"]) + assert body["type"] == "sample" + assert "lora_path" not in stack.router.requests[-1] + seeds = sorted(r["sampling_params"]["sampling_seed"] for r in stack.router.requests[-2:]) + assert seeds == [40, 41] + calls = len(stack.router.requests) + overflow = self.sample_request(sampler_id, seq_id=1, num_samples=2, seed=2**63 - 1) + failed = await stack.retrieve(stack.frontend.sample(overflow)["request_id"]) + assert failed["category"] == "user" and "signed 64-bit" in failed["error"] + assert len(stack.router.requests) == calls + + probe = self.sample_request(sampler_id, seq_id=2) + probe.prompt_logprobs = True + body = await stack.retrieve(stack.frontend.sample(probe)["request_id"]) + assert body["type"] == "sample" and body["prompt_logprobs"] == [None, -0.125] + assert stack.router.requests[-1]["logprob_start_len"] == 0 + assert all("logprob_start_len" not in r for r in stack.router.requests[:-1]) + + topk_probe = self.sample_request(sampler_id, seq_id=3) + topk_probe.topk_prompt_logprobs = 2 + failed = await stack.retrieve(stack.frontend.sample(topk_probe)["request_id"]) + assert failed["category"] == "user" and "topk_prompt_logprobs" in failed["error"] + + run(scenario) + + +class TestReplayExpiry: + def test_training_resubmit_after_eviction_is_410_not_conflict(self): + async def scenario(stack): + stack.frontend.futures.max_delivered = 1 + model_id = await stack.create_model() + first = stack.frontend.forward_backward(stack.fb_request(model_id, 1)) + await stack.retrieve(first["request_id"]) + second = stack.frontend.forward_backward(stack.fb_request(model_id, 2)) + await stack.retrieve(second["request_id"]) + + with pytest.raises(ApiError) as repoll: + await stack.retrieve(first["request_id"]) + assert repoll.value.status_code == 410 and "already delivered" in repoll.value.detail + with pytest.raises(ApiError) as resubmit: + stack.frontend.forward_backward(stack.fb_request(model_id, 1)) + assert resubmit.value.status_code == 410 + with pytest.raises(ApiError) as conflict: + stack.frontend.forward_backward(stack.fb_request(model_id, 1, tokens=(7, 8, 9))) + assert conflict.value.status_code == 422 + + run(scenario) + + def test_sample_resubmit_after_eviction_never_regenerates(self): + async def scenario(stack): + stack.frontend.futures.max_delivered = 1 + model_id = await stack.create_model() + publish = stack.frontend.save_weights_for_sampler( + wire.SaveWeightsForSamplerRequest(model_id=model_id, seq_id=1, sampling_session_seq_id=0) + ) + sampler_id = (await stack.retrieve(publish["request_id"]))["sampling_session_id"] + request = wire.SampleRequest.model_validate( + { + "sampling_session_id": sampler_id, + "seq_id": 0, + "num_samples": 1, + "prompt": {"chunks": [{"type": "encoded_text", "tokens": [5, 6]}]}, + "sampling_params": {"max_tokens": 3}, + } + ) + future = stack.frontend.sample(request) + await stack.retrieve(future["request_id"]) + generated = len(stack.router.requests) + fb = stack.frontend.forward_backward(stack.fb_request(model_id, 2)) + await stack.retrieve(fb["request_id"]) + with pytest.raises(ApiError) as excinfo: + stack.frontend.sample(request) + assert excinfo.value.status_code == 410 + await asyncio.sleep(0.05) + assert len(stack.router.requests) == generated + + run(scenario) + + +class TestLifecycle: + def test_unload_fences_and_resolves(self): + async def scenario(stack): + model_id = await stack.create_model() + unload = await stack.frontend.unload_model(wire.UnloadModelRequest(model_id=model_id)) + body = await stack.retrieve(unload["request_id"]) + assert body == {"type": "unload_model", "model_id": model_id} + follow_up = stack.frontend.forward_backward(stack.fb_request(model_id, 1)) + failed = await stack.retrieve(follow_up["request_id"]) + assert failed["category"] == "user" + + run(scenario) + + def test_create_model_rejections(self): + async def scenario(stack): + base = wire.CreateModelRequest( + session_id=stack.session_id, model_seq_id=0, base_model=BASE, lora_config=wire.LoraConfig(rank=8) + ) + for broken, match in ( + (base.model_copy(update={"base_model": "other/model"}), 400), + (base.model_copy(update={"lora_config": wire.LoraConfig(rank=8, seed=7)}), 400), + (base.model_copy(update={"lora_config": wire.LoraConfig(rank=8, train_mlp=False)}), 400), + (base.model_copy(update={"lora_config": None}), 400), + (base.model_copy(update={"session_id": "sess-unknown"}), 404), + ): + with pytest.raises(ApiError) as excinfo: + await stack.frontend.create_model(broken) + assert excinfo.value.status_code == match + + run(scenario) + + def test_unknown_future_is_410(self): + async def scenario(stack): + with pytest.raises(ApiError) as excinfo: + await stack.retrieve("nope") + assert excinfo.value.status_code == 410 + + run(scenario) + + def test_stale_model_handle_never_binds_to_a_same_name_successor(self): + async def scenario(stack): + model_id = await stack.create_model() + record = stack.frontend.models.get(model_id) + name, rid1 = record.name, record.registration_id + unload = await stack.frontend.unload_model(wire.UnloadModelRequest(model_id=model_id)) + await stack.retrieve(unload["request_id"]) + for _ in range(500): + if stack.frontend.backend.registry.find(name) is None: + break + await asyncio.sleep(0.005) + await stack.frontend.backend.register(name, AdapterRunConfig(rank=8)) + rid2 = stack.frontend.backend.registry.find(name).registration_id + assert rid2 != rid1 + + submitted = stack.frontend.optim_step(stack.optim_request(model_id, 1)) + body = await stack.retrieve(submitted["request_id"]) + assert body["category"] == "user" and "fenced" in body["error"] + assert stack.frontend.backend.operations.queue_view(name, rid2) == [] + + run(scenario) + + def test_unsupported_sdk_version_is_rejected_at_bootstrap(self): + async def scenario(stack): + for request in ( + lambda: stack.frontend.client_config(wire.ClientConfigRequest(sdk_version="0.25.0")), + lambda: stack.frontend.create_session(wire.CreateSessionRequest(sdk_version="0.25.0")), + lambda: stack.frontend.create_session(wire.CreateSessionRequest()), + ): + with pytest.raises(ApiError) as excinfo: + request() + assert excinfo.value.status_code == 400 and "tinker==0.24.1" in excinfo.value.detail + + run(scenario) + + def test_healthz_reports_readiness_not_liveness(self): + async def scenario(stack): + assert stack.frontend.health() == {"status": "ok"} + stack.frontend.backend.trainer_ready = False + with pytest.raises(ApiError) as excinfo: + stack.frontend.health() + assert excinfo.value.status_code == 503 + stack.frontend.backend.mark_trainer_ready() + assert stack.frontend.health() == {"status": "ok"} + + run(scenario) + + def test_rejected_flood_backpressures_instead_of_growing_without_bound(self): + async def scenario(stack): + model_id = await stack.create_model() + stack.driver.paused = True + stack.frontend.backend.operations.max_unacked_results = 8 + accepted = throttled = 0 + for seq in range(1, 101): + bad = stack.fb_request(model_id, seq, weights=(1.0, 1.0, 1.0), targets=[9, 3, 99]) + try: + stack.frontend.forward_backward(bad) + accepted += 1 + except OperationBackpressure: + throttled += 1 + assert accepted == 8 and throttled == 92 + assert len(stack.frontend.futures.records) <= 8 + 1 + + run(scenario) + + def test_bootstrap_surfaces(self): + async def scenario(stack): + assert stack.frontend.health() == {"status": "ok"} + capabilities = stack.frontend.capabilities() + assert capabilities["supported_models"][0]["model_name"] == BASE + config = stack.frontend.client_config(wire.ClientConfigRequest(sdk_version="0.24.1")) + assert config["proto_write_fwdbwd"] is False and config["pjwt_auth_enabled"] is False + assert stack.frontend.session_heartbeat(wire.SessionHeartbeatRequest(session_id=stack.session_id)) == { + "type": "session_heartbeat" + } + + run(scenario) + + def test_get_info(self): + async def scenario(stack): + model_id = await stack.create_model() + info = stack.frontend.get_info(wire.GetInfoRequest(model_id=model_id)) + assert info["model_id"] == model_id and info["lora_rank"] == 8 + assert info["model_data"]["model_name"] == BASE + + run(scenario) + + +async def until_terminal(stack, request_id): + while (body := await stack.retrieve(request_id)).get("type") == "try_again": + pass + return body + + +class TestCapacityQueue: + def test_unbound_create_reports_paused_capacity_until_the_slot_frees(self): + async def scenario(stack): + model_a = await stack.create_model(model_seq_id=0) + future_b = await stack.frontend.create_model( + wire.CreateModelRequest( + session_id=stack.session_id, + model_seq_id=1, + base_model=BASE, + lora_config=wire.LoraConfig(rank=8), + ) + ) + model_b = f"{stack.session_id}:train:1" + paused = {"type": "try_again", "queue_state": "paused_capacity"} + assert await stack.retrieve(future_b["request_id"]) == paused + + fb_b = stack.frontend.forward_backward(stack.fb_request(model_b, 1)) + assert (await stack.retrieve(fb_b["request_id"]))["type"] == "try_again" + assert await stack.retrieve(future_b["request_id"]) == paused + + unload = await stack.frontend.unload_model(wire.UnloadModelRequest(model_id=model_a)) + assert await until_terminal(stack, unload["request_id"]) == { + "type": "unload_model", + "model_id": model_a, + } + assert await until_terminal(stack, future_b["request_id"]) == { + "type": "create_model", + "model_id": model_b, + } + body = await until_terminal(stack, fb_b["request_id"]) + (row,) = [output["logprobs"]["data"] for output in body["loss_fn_outputs"]] + assert row == [-0.5, -0.5, -0.5] + + run(scenario, poll_window_s=0.2, multi_lora_n_adapters=1) + + +def test_seq_to_ordinal_documented_mapping(): + from miles.ray.tinker_frontend import service + + assert "ordinal = seq_id" in service.__doc__ + + +def test_frontend_reads_the_backend_facade_only(): + import inspect + + from miles.ray.tinker_frontend import service + + source = inspect.getsource(service) + for internal in ("self.backend.registry", "self.backend.operations", "self.backend.router_url"): + assert internal not in source, f"frontend must not read {internal}" + + +def test_injected_sampling_transport_receives_the_exact_router_payload(): + import asyncio + + from tests.fast.ray.tinker_frontend.fake_stack import FakeDriver, FakeRouter, make_backend + + from miles.ray.tinker_frontend.service import TinkerFrontend + + class FakeTransport: + def __init__(self, router): + self.router = router + self.payloads = [] + self.release = asyncio.Event() + + async def generate(self, payload): + self.payloads.append(payload) + await self.release.wait() + return self.router.response_for(payload) + + async def close(self): + pass + + async def main(): + router = FakeRouter() + backend = make_backend() + await backend.init() + driver = FakeDriver(backend) + transport = FakeTransport(router) + frontend = TinkerFrontend(backend, poll_window_s=0.3, poll_interval_s=0.002, sampling_transport=transport) + stack = Stack(frontend, driver, router) + driver_task = asyncio.create_task(driver.run(interval=0.002)) + try: + model_id = await stack.create_model() + publish = frontend.save_weights_for_sampler( + wire.SaveWeightsForSamplerRequest(model_id=model_id, seq_id=1, sampling_session_seq_id=0) + ) + publish_body = await stack.retrieve(publish["request_id"]) + sampler_id = publish_body["sampling_session_id"] + + request = wire.SampleRequest.model_validate( + { + "sampling_session_id": sampler_id, + "seq_id": 0, + "prompt": {"chunks": [{"type": "encoded_text", "tokens": [1, 2, 3]}]}, + "sampling_params": {"max_tokens": 4, "temperature": 0.0}, + "num_samples": 1, + } + ) + future = frontend.sample(request) + assert future["request_id"] + for _ in range(200): + if transport.payloads: + break + await asyncio.sleep(0.002) + [payload] = transport.payloads + assert payload["input_ids"] == [1, 2, 3] + assert payload["return_logprob"] is True + assert payload["sampling_params"]["max_new_tokens"] == 4 + assert payload["lora_path"].startswith("__miles_adapter_") + assert payload["rid"].count("::") == 2 + transport.release.set() + body = await stack.retrieve(future["request_id"]) + assert body["sequences"] + finally: + driver_task.cancel() + await frontend.close() + await backend.close() + + asyncio.run(main()) diff --git a/tests/fast/ray/tinker_frontend/test_service_failure_paths.py b/tests/fast/ray/tinker_frontend/test_service_failure_paths.py new file mode 100644 index 00000000000..7b502a8217c --- /dev/null +++ b/tests/fast/ray/tinker_frontend/test_service_failure_paths.py @@ -0,0 +1,309 @@ +"""Frontend failure-path contracts (external adversarial review): lost-response +publish retries, sampler-identity retention, sibling cancellation, shutdown +barriers, bounded-idempotency fences, and the exact SDK patch pin — the +behaviors happy-path/equivalence tests cannot see.""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=90, suite="stage-a-cpu") + +import asyncio +import time + +import pytest +from tests.fast.ray.tinker_frontend.fake_stack import make_backend + +from miles.ray.tinker_frontend import wire +from miles.ray.tinker_frontend.service import ApiError, TinkerFrontend + +BASE = "Qwen/Qwen3-0.6B" + + +class StaticTransport: + def __init__(self) -> None: + self.calls = 0 + self.closed = False + + async def generate(self, payload: dict) -> dict: + self.calls += 1 + return { + "meta_info": { + "finish_reason": {"type": "stop"}, + "output_token_logprobs": [[-0.25, 1000, None]], + } + } + + async def close(self) -> None: + self.closed = True + + +async def make_frontend(transport): + backend = make_backend() + await backend.init() + frontend = TinkerFrontend(backend, poll_window_s=0.2, poll_interval_s=0.001, sampling_transport=transport) + session_id = frontend.create_session(wire.CreateSessionRequest(sdk_version="0.24.1"))["session_id"] + return backend, frontend, session_id + + +async def create_ready_model(backend, frontend, session_id): + submitted = await frontend.create_model( + wire.CreateModelRequest( + session_id=session_id, + model_seq_id=0, + base_model=BASE, + lora_config=wire.LoraConfig(rank=8), + ) + ) + model_id = f"{session_id}:train:0" + model = frontend.models.get(model_id) + backend.registry.mark_ready([model.name]) + await frontend.retrieve_future(wire.FutureRetrieveRequest(request_id=submitted["request_id"])) + return model_id, model + + +def base_sampler(frontend, session_id, seq=0): + return frontend.create_sampling_session( + wire.CreateSamplingSessionRequest( + session_id=session_id, + sampling_session_seq_id=seq, + base_model=BASE, + ) + )["sampling_session_id"] + + +def sample_request(sampler_id, seq=0, num_samples=1): + return wire.SampleRequest.model_validate( + { + "sampling_session_id": sampler_id, + "seq_id": seq, + "num_samples": num_samples, + "prompt": {"chunks": [{"type": "encoded_text", "tokens": [5, 6]}]}, + "sampling_params": {"max_tokens": 1}, + } + ) + + +class TestExactSdkPin: + def test_only_the_pinned_patch_version_is_accepted(self): + async def main(): + backend = make_backend() + frontend = TinkerFrontend(backend, sampling_transport=StaticTransport()) + try: + for version in ("0.24.0", "0.24.2"): + with pytest.raises(ApiError, match="0.24.1"): + frontend.create_session(wire.CreateSessionRequest(sdk_version=version)) + assert frontend.create_session(wire.CreateSessionRequest(sdk_version="0.24.1"))["session_id"] + finally: + await frontend.close() + + asyncio.run(main()) + + +class TestPublishRetryIdempotency: + def test_lost_response_retry_replays_the_original_future(self): + async def main(): + backend, frontend, session_id = await make_frontend(StaticTransport()) + try: + model_id, _ = await create_ready_model(backend, frontend, session_id) + first = frontend.save_weights_for_sampler( + wire.SaveWeightsForSamplerRequest(model_id=model_id, seq_id=1, sampling_session_seq_id=0) + ) + retry = frontend.save_weights_for_sampler( + wire.SaveWeightsForSamplerRequest(model_id=model_id, seq_id=1, sampling_session_seq_id=1) + ) + assert retry == first + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_a_different_operation_at_the_same_seq_still_conflicts(self): + async def main(): + backend, frontend, session_id = await make_frontend(StaticTransport()) + try: + model_id, _ = await create_ready_model(backend, frontend, session_id) + frontend.save_weights_for_sampler( + wire.SaveWeightsForSamplerRequest(model_id=model_id, seq_id=1, sampling_session_seq_id=0) + ) + with pytest.raises(ApiError) as excinfo: + frontend.save_weights_for_sampler( + wire.SaveWeightsForSamplerRequest(model_id=model_id, seq_id=1, path="named") + ) + assert excinfo.value.status_code == 422 + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + +class TestSamplerIdentityRetention: + def test_publish_cannot_overwrite_an_existing_base_sampler(self): + async def main(): + backend, frontend, session_id = await make_frontend(StaticTransport()) + try: + sampler_id = base_sampler(frontend, session_id, seq=0) + model_id, model = await create_ready_model(backend, frontend, session_id) + publish = frontend.save_weights_for_sampler( + wire.SaveWeightsForSamplerRequest(model_id=model_id, seq_id=1, sampling_session_seq_id=0) + ) + claimed = backend.claim_ready_control_operations()["operations"] + backend.registry.record_weight_update([model.name]) + backend.complete_control_operations({claimed[0]["operation_id"]: {"ok": True}}) + body = await frontend.retrieve_future(wire.FutureRetrieveRequest(request_id=publish["request_id"])) + assert body["category"] == "user" and "already exists" in body["error"] + assert frontend.samplers.get(sampler_id).name is None + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_publish_completion_after_parent_reap_cannot_recreate_a_sampler(self): + async def main(): + backend, frontend, session_id = await make_frontend(StaticTransport()) + try: + model_id, model = await create_ready_model(backend, frontend, session_id) + publish = frontend.save_weights_for_sampler( + wire.SaveWeightsForSamplerRequest(model_id=model_id, seq_id=1, sampling_session_seq_id=0) + ) + frontend.reap_once(now=time.time() + frontend.session_idle_ttl_s + 1) + assert frontend.sessions.get(session_id) is None + + claimed = backend.claim_ready_control_operations()["operations"] + backend.registry.record_weight_update([model.name]) + backend.complete_control_operations({claimed[0]["operation_id"]: {"ok": True}}) + body = await frontend.retrieve_future(wire.FutureRetrieveRequest(request_id=publish["request_id"])) + assert body["category"] == "user" and "parent session expired" in body["error"] + assert not frontend.samplers.records + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_sample_identity_does_not_reexecute_after_tombstone_rollover(self): + async def main(): + transport = StaticTransport() + backend, frontend, session_id = await make_frontend(transport) + frontend.futures.max_delivered = 1 + frontend.futures.max_expired = 1 + try: + sampler_id = base_sampler(frontend, session_id) + for seq in range(3): + future = frontend.sample(sample_request(sampler_id, seq=seq)) + await frontend.retrieve_future(wire.FutureRetrieveRequest(request_id=future["request_id"])) + assert transport.calls == 3 + + retried = frontend.sample(sample_request(sampler_id, seq=0)) + body = await frontend.retrieve_future(wire.FutureRetrieveRequest(request_id=retried["request_id"])) + assert transport.calls == 3 + assert body["category"] == "user" and "already executed" in body["error"] + finally: + await frontend.close() + await backend.close() + + asyncio.run(main()) + + +class PartialFailureTransport: + def __init__(self) -> None: + self.calls = 0 + self.second_started = asyncio.Event() + self.second_cancelled = asyncio.Event() + self.release = asyncio.Event() + + async def generate(self, payload: dict) -> dict: + index = self.calls + self.calls += 1 + if index == 0: + await self.second_started.wait() + raise RuntimeError("first generation failed") + self.second_started.set() + try: + await self.release.wait() + except asyncio.CancelledError: + self.second_cancelled.set() + raise + return await StaticTransport().generate(payload) + + async def close(self) -> None: + pass + + +class BlockingTransport: + def __init__(self) -> None: + self.started = asyncio.Event() + self.cancelled = asyncio.Event() + + async def generate(self, payload: dict) -> dict: + self.started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.cancelled.set() + raise + raise AssertionError("unreachable") + + async def close(self) -> None: + pass + + +class TestAsyncLifecycle: + def test_partial_multisample_failure_cancels_sibling_generation(self): + async def main(): + transport = PartialFailureTransport() + backend, frontend, session_id = await make_frontend(transport) + try: + sampler_id = base_sampler(frontend, session_id) + future = frontend.sample(sample_request(sampler_id, num_samples=2)) + body = await frontend.retrieve_future(wire.FutureRetrieveRequest(request_id=future["request_id"])) + assert body["category"] == "server" + assert transport.second_cancelled.is_set() + finally: + transport.release.set() + await frontend.close() + await backend.close() + + asyncio.run(main()) + + def test_close_awaits_inflight_sample_cancellation_and_gates_new_ones(self): + async def main(): + transport = BlockingTransport() + backend, frontend, session_id = await make_frontend(transport) + sampler_id = base_sampler(frontend, session_id) + future = frontend.sample(sample_request(sampler_id)) + await transport.started.wait() + try: + await frontend.close() + assert transport.cancelled.is_set() + assert not frontend._sample_tasks + body = await frontend.retrieve_future(wire.FutureRetrieveRequest(request_id=future["request_id"])) + assert body["category"] == "server" and "shutting down" in body["error"] + with pytest.raises(ApiError) as excinfo: + frontend.sample(sample_request(sampler_id, seq=1)) + assert excinfo.value.status_code == 503 + await frontend.close() + finally: + await backend.close() + + asyncio.run(main()) + + def test_close_terminalizes_sample_cancelled_before_its_first_step(self): + async def main(): + transport = BlockingTransport() + backend, frontend, session_id = await make_frontend(transport) + sampler_id = base_sampler(frontend, session_id) + future = frontend.sample(sample_request(sampler_id)) + try: + await frontend.close() + assert not transport.started.is_set() + assert frontend.sampling_admission.in_use == 0 + assert frontend.sampling_stats.failures_by_class == {"Cancelled": 1} + body = await frontend.retrieve_future(wire.FutureRetrieveRequest(request_id=future["request_id"])) + assert body["category"] == "server" and "shutting down" in body["error"] + finally: + await backend.close() + + asyncio.run(main()) diff --git a/tests/fast/ray/tinker_frontend/test_state.py b/tests/fast/ray/tinker_frontend/test_state.py new file mode 100644 index 00000000000..9bca3164391 --- /dev/null +++ b/tests/fast/ray/tinker_frontend/test_state.py @@ -0,0 +1,106 @@ +"""Frontend state stores: fingerprint identity, conflicts, replay retention.""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=30, suite="stage-a-cpu") + +import pytest + +from miles.ray.tinker_frontend.state import ( + ConflictError, + ExpiredError, + FutureRecord, + FutureStore, + SamplingSessionRecord, + SamplingSessionStore, + SessionStore, + fingerprint_of, +) + + +def record(request_id="r1", fingerprint="f1", terminal=None): + rec = FutureRecord(request_id=request_id, kind="operation", fingerprint=fingerprint) + if terminal is not None: + rec.resolve(terminal) + return rec + + +class TestFutureStore: + def test_existing_replays_identical_and_conflicts_on_divergence(self): + store = FutureStore() + store.put(record()) + assert store.existing("r1", "f1") is not None + assert store.existing("r2", "f1") is None + with pytest.raises(ConflictError, match="identical"): + store.existing("r1", "OTHER") + + def test_delivered_terminal_records_are_evicted_lru(self): + store = FutureStore(max_delivered=2) + for i in range(3): + rec = store.put(record(f"r{i}", terminal={"n": i})) + store.mark_delivered(rec) + assert store.get("r0") is None + assert store.get("r1").terminal == {"n": 1} + assert store.get("r2").terminal == {"n": 2} + + def test_pending_records_are_never_evicted(self): + store = FutureStore(max_delivered=1) + pending = store.put(record("pending")) + store.mark_delivered(pending) + for i in range(3): + store.mark_delivered(store.put(record(f"r{i}", terminal={}))) + assert store.get("pending") is pending + + def test_eviction_leaves_a_typed_tombstone(self): + store = FutureStore(max_delivered=1) + store.mark_delivered(store.put(record("r1", "f1", terminal={"n": 1}))) + store.mark_delivered(store.put(record("r2", "f2", terminal={"n": 2}))) + assert store.get("r1") is None + assert store.expired_fingerprint("r1") == "f1" + with pytest.raises(ExpiredError, match="already delivered"): + store.existing("r1", "f1") + with pytest.raises(ConflictError, match="identical"): + store.existing("r1", "OTHER") + + def test_tombstones_are_bounded(self): + store = FutureStore(max_delivered=1, max_expired=2) + for i in range(4): + store.mark_delivered(store.put(record(f"r{i}", f"f{i}", terminal={}))) + assert store.expired_fingerprint("r0") is None + assert store.expired_fingerprint("r2") == "f2" + assert store.existing("r0", "f0") is None + + def test_resolve_drops_the_forward_payload(self): + rec = record() + rec.forward_payload = {"samples": []} + rec.resolve({"ok": True}) + assert rec.forward_payload is None + + +class TestSessions: + def test_heartbeat_only_touches_known_sessions(self): + store = SessionStore() + session = store.create("0.24.1", [], None) + assert store.heartbeat(session.session_id) + assert not store.heartbeat("sess-nope") + + def test_fingerprints_are_canonical(self): + assert fingerprint_of({"a": 1, "b": 2}) == fingerprint_of({"b": 2, "a": 1}) + assert fingerprint_of({"a": 1}) != fingerprint_of({"a": 2}) + + def test_child_sampler_namespaces_are_retired_in_one_bulk_pass(self): + store = SamplingSessionStore() + for session_id in ("sess-a", "sess-b", "sess-live"): + for suffix in range(2): + store.add( + SamplingSessionRecord( + sampling_session_id=f"{session_id}:sample:{suffix}", + session_id=session_id, + fingerprint=f"fp-{session_id}-{suffix}", + base_model="test-model", + ) + ) + + store.remove_for_sessions({"sess-a", "sess-b"}) + + assert set(store.records) == {"sess-live:sample:0", "sess-live:sample:1"} diff --git a/tests/fast/ray/tinker_frontend/test_translation.py b/tests/fast/ray/tinker_frontend/test_translation.py new file mode 100644 index 00000000000..316020e7db8 --- /dev/null +++ b/tests/fast/ray/tinker_frontend/test_translation.py @@ -0,0 +1,196 @@ +"""Datum/result/sampling translation: official wire shapes <-> backend +payloads, with every v1 boundary rejection typed as UserInputError.""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=60, suite="stage-a-cpu") + +import pytest + +from miles.ray.tinker_frontend import translation, wire +from miles.ray.tinker_frontend.translation import UserInputError + + +def tensor(data, dtype="float32", **kwargs): + return {"data": data, "dtype": dtype, "shape": [len(data)], **kwargs} + + +def datum(tokens, targets, **channels): + loss_fn_inputs = {"target_tokens": tensor(targets, "int64")} + for name, values in channels.items(): + loss_fn_inputs[name] = tensor(values) + return wire.Datum.model_validate( + {"model_input": {"chunks": [{"type": "encoded_text", "tokens": tokens}]}, "loss_fn_inputs": loss_fn_inputs} + ) + + +def fb_input(data, loss_fn="cross_entropy", config=None): + return wire.ForwardBackwardInput.model_validate( + {"data": [d.model_dump() for d in data], "loss_fn": loss_fn, "loss_fn_config": config} + ) + + +class TestDatumToSample: + def test_shifted_targets_extend_the_token_sequence(self): + sample = translation.datum_to_sample(0, datum([1, 2, 3], [2, 3, 4], weights=[0.0, 1.0, 1.0]), "cross_entropy") + assert sample == { + "tokens": [1, 2, 3, 4], + "response_length": 3, + "loss_mask": [1, 1, 1], + "loss_weights": [0.0, 1.0, 1.0], + } + + def test_active_position_must_be_next_token(self): + with pytest.raises(UserInputError, match="next input"): + translation.datum_to_sample(0, datum([1, 2, 3], [9, 3, 4], weights=[1.0, 1.0, 1.0]), "cross_entropy") + + def test_negative_token_ids_are_rejected(self): + with pytest.raises(UserInputError, match="non-negative"): + translation.datum_to_sample(0, datum([-1, 2, 3], [2, 3, 4], weights=[0.0, 1.0, 1.0]), "cross_entropy") + with pytest.raises(UserInputError, match="non-negative"): + translation.datum_to_sample(0, datum([1, 2, 3], [2, 3, -4], weights=[0.0, 1.0, 1.0]), "cross_entropy") + + def test_zero_weighted_mismatch_is_normalized_not_rejected(self): + sample = translation.datum_to_sample(0, datum([1, 2, 3], [0, 3, 4], weights=[0.0, 1.0, 1.0]), "cross_entropy") + assert sample["tokens"] == [1, 2, 3, 4] + + def test_importance_sampling_channels_map_to_backend_names(self): + d = datum([1, 2], [2, 5], logprobs=[-0.5, -0.5], advantages=[0.0, 1.0]) + sample = translation.datum_to_sample(0, d, "importance_sampling") + assert sample["rollout_log_probs"] == [-0.5, -0.5] + assert sample["advantages"] == [0.0, 1.0] + assert "loss_weights" not in sample + + def test_missing_required_channel_is_rejected(self): + with pytest.raises(UserInputError, match="requires loss_fn_inputs\\['weights'\\]"): + translation.datum_to_sample(0, datum([1, 2], [2, 3]), "cross_entropy") + + def test_sparse_csr_is_rejected(self): + d = datum([1, 2], [2, 3], weights=[1.0, 1.0]) + d.loss_fn_inputs["target_tokens"].sparse_crow_indices = [0, 1, 2] + with pytest.raises(UserInputError, match="sparse"): + translation.datum_to_sample(0, d, "cross_entropy") + + def test_top_k_shaped_targets_are_rejected(self): + d = datum([1, 2], [2, 3], weights=[1.0, 1.0]) + d.loss_fn_inputs["target_tokens"].shape = [2, 1] + with pytest.raises(UserInputError, match="1-D"): + translation.datum_to_sample(0, d, "cross_entropy") + + def test_non_text_chunks_are_rejected(self): + d = datum([1, 2], [2, 3], weights=[1.0, 1.0]) + d.model_input.chunks[0].type = "image" + with pytest.raises(UserInputError, match="text-only"): + translation.datum_to_sample(0, d, "cross_entropy") + + def test_unknown_channels_and_length_mismatches_are_rejected(self): + d = datum([1, 2], [2, 3], weights=[1.0, 1.0]) + d.loss_fn_inputs["mystery"] = d.loss_fn_inputs["weights"] + with pytest.raises(UserInputError, match="unsupported loss_fn_inputs"): + translation.datum_to_sample(0, d, "cross_entropy") + with pytest.raises(UserInputError, match="one value per input token"): + translation.datum_to_sample(0, datum([1, 2, 3], [2, 3], weights=[1.0, 1.0]), "cross_entropy") + + +class TestFbPayload: + def test_payload_carries_samples_and_loss_spec(self): + payload = translation.fb_input_to_payload( + fb_input([datum([1, 2], [2, 3], weights=[1.0, 1.0])], config={"clip_low_threshold": 0.8}) + ) + assert payload["loss"] == {"loss_fn": "cross_entropy", "loss_fn_config": {"clip_low_threshold": 0.8}} + assert len(payload["samples"]) == 1 + + def test_unsupported_loss_fns_are_rejected(self): + for loss_fn in ("cispo", "dro", "nope"): + with pytest.raises(UserInputError, match="not supported"): + translation.fb_input_to_payload(fb_input([datum([1, 2], [2, 3], weights=[1.0, 1.0])], loss_fn)) + + def test_empty_data_is_rejected(self): + with pytest.raises(UserInputError, match="at least one datum"): + translation.fb_input_to_payload(fb_input([])) + + +class TestResults: + def test_fb_result_uses_backend_metrics(self): + body = translation.fb_result_to_response({"logprobs": [[-0.5, -0.25]], "metrics": {"loss:sum": 0.75}}) + assert body["metrics"] == {"loss:sum": 0.75} + assert body["loss_fn_outputs"] == [{"logprobs": {"data": [-0.5, -0.25], "dtype": "float32", "shape": [2]}}] + + def test_forward_result_recomputes_metrics_from_the_request(self): + payload = translation.fb_input_to_payload(fb_input([datum([1, 2, 3], [2, 3, 4], weights=[0.0, 1.0, 1.0])])) + body = translation.fb_result_to_response({"logprobs": [[-0.5, -0.5, -0.5]]}, payload) + assert body["metrics"]["loss:sum"] == pytest.approx(1.0) + assert body["metrics"]["unmasked_tokens:sum"] == pytest.approx(3.0) + + def test_optim_result_projects_numeric_metrics(self): + assert translation.optim_result_to_response({"grad_norm": 0.5, "learning_rate": 1e-4}) == { + "type": "optim_step", + "metrics": {"grad_norm": 0.5, "learning_rate": 1e-4}, + } + + +class TestSampling: + def params(self, **kwargs): + return wire.SamplingParams.model_validate({"max_tokens": 8, **kwargs}) + + def test_params_map_to_sglang(self): + params = translation.sampling_params_to_sglang(self.params(temperature=0.5, top_p=0.9, stop="\n")) + assert params == {"max_new_tokens": 8, "temperature": 0.5, "top_p": 0.9, "top_k": -1, "stop": ["\n"]} + + def test_stop_token_ids(self): + assert translation.sampling_params_to_sglang(self.params(stop=[7, 8]))["stop_token_ids"] == [7, 8] + with pytest.raises(UserInputError, match="non-negative"): + translation.sampling_params_to_sglang(self.params(stop=[7, -8])) + + def test_missing_max_tokens_is_rejected_and_seed_stays_out_of_base_params(self): + with pytest.raises(UserInputError, match="max_tokens"): + translation.sampling_params_to_sglang(wire.SamplingParams()) + assert "sampling_seed" not in translation.sampling_params_to_sglang(self.params(seed=1)) + + @pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"temperature": -1.0}, "temperature"), + ({"temperature": float("nan")}, "temperature"), + ({"top_p": 0.0}, "top_p"), + ({"top_p": float("inf")}, "top_p"), + ({"top_k": 0}, "top_k"), + ({"seed": -(2**63) - 1}, "seed"), + ({"seed": 2**63}, "seed"), + ], + ) + def test_invalid_sampling_ranges_are_rejected_locally(self, overrides, message): + with pytest.raises(UserInputError, match=message): + translation.sampling_params_to_sglang(self.params(**overrides)) + + def test_generation_maps_tokens_logprobs_and_stop_reason(self): + sequence = translation.generation_to_sequence( + { + "meta_info": { + "finish_reason": {"type": "length"}, + "output_token_logprobs": [[-0.1, 11, None], [-0.2, 12, None]], + } + } + ) + assert sequence == {"stop_reason": "length", "tokens": [11, 12], "logprobs": [-0.1, -0.2]} + + def test_aborted_generation_raises(self): + with pytest.raises(RuntimeError, match="abort"): + translation.generation_to_sequence({"meta_info": {"finish_reason": {"type": "abort"}}}) + + def test_prompt_logprobs_map_per_token_with_leading_none(self): + generation = {"meta_info": {"input_token_logprobs": [[None, 5, None], [-0.5, 6, None], [-1.25, 7, None]]}} + assert translation.prompt_logprobs_from_generation(generation, 3) == [None, -0.5, -1.25] + + def test_prompt_logprobs_missing_from_the_engine_is_a_server_fault(self): + with pytest.raises(RuntimeError, match="no input_token_logprobs"): + translation.prompt_logprobs_from_generation({"meta_info": {}}, 2) + + def test_prompt_logprobs_length_mismatch_is_a_server_fault(self): + generation = {"meta_info": {"input_token_logprobs": [[None, 5, None]]}} + with pytest.raises(RuntimeError, match="1 prompt logprobs for 2 prompt tokens"): + translation.prompt_logprobs_from_generation(generation, 2) + + def test_sample_response_carries_prompt_logprobs_only_when_scored(self): + assert translation.sequences_to_sample_response([])["prompt_logprobs"] is None + assert translation.sequences_to_sample_response([], [None, -0.5])["prompt_logprobs"] == [None, -0.5] diff --git a/tests/fast/rollout/multi_lora/__init__.py b/tests/fast/rollout/multi_lora/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/fast/rollout/multi_lora/test_rollout_fn.py b/tests/fast/rollout/multi_lora/test_rollout_fn.py new file mode 100644 index 00000000000..75f8e66e9fb --- /dev/null +++ b/tests/fast/rollout/multi_lora/test_rollout_fn.py @@ -0,0 +1,278 @@ +import asyncio +from types import SimpleNamespace + +import pytest + +from miles.ray.multi_lora.config import AdapterRun, AdapterRunConfig +from miles.ray.multi_lora.residency import ResidentBinding +from miles.rollout.base_types import RolloutFnConstructorInput, RolloutFnTrainOutput +from miles.rollout.multi_lora.rollout_fn import AdapterRolloutRuntime, ClaimedOperationBatch, MultiLoraOperationBatchFn +from miles.utils.operation_contract import BatchExecutionLease, EmptyBatchTimeoutError + + +def make_run(name="X", reg="rx", slot=3, version=2) -> AdapterRun: + config = AdapterRunConfig(rank=8, alpha=16, metadata={"team": "t1"}) + return AdapterRun(name=name, config=config, slot=slot, version=version, registration_id=reg) + + +def claim_batch(run: AdapterRun, operations) -> ClaimedOperationBatch: + fn = MultiLoraOperationBatchFn( + RolloutFnConstructorInput(args=SimpleNamespace(), data_source=None), + operations=operations, + residency=FakeResidency(), + ) + return asyncio.run(fn._claim_batch(AdapterRolloutRuntime(run))) + + +def sample_payload(n=2) -> dict: + return { + "batch_id": "batch-7", + "samples": [ + {"prompt": "p", "tokens": [1, 2, 3, 4], "response_length": 2, "loss_mask": [1, 1]} for _ in range(n) + ], + "loss": {"loss_fn": "cross_entropy"}, + } + + +class FakeOperationQueue: + def __init__(self, claims=(), ready=None): + self._claims = list(claims) + self._ready = ready or {} + self.failed: list[tuple] = [] + + async def ready_streams(self) -> dict: + return self._ready + + async def claim_data(self, key): + return self._claims.pop(0) if self._claims else None + + async def fail(self, operation_id, error, category): + self.failed.append((operation_id, error, category)) + + +class FakeResidency: + def __init__(self): + self.leases: list[tuple] = [] + + async def acquire_batch(self, bindings_by_operation): + self.leases.append(tuple(bindings_by_operation)) + return BatchExecutionLease(dispatch_id="lease-1", bindings_by_operation=tuple(bindings_by_operation)) + + +@pytest.fixture() +def fast_poll(monkeypatch): + import miles.rollout.multi_lora.rollout_fn as rollout_module + + monkeypatch.setattr(rollout_module, "_CLAIM_POLL_S", 0.01) + + +def op(op_id="op1", kind="forward_backward", payload=None, slot=3): + return dict( + operation_id=op_id, + name="X", + registration_id="rx", + kind=kind, + payload=sample_payload() if payload is None else payload, + state="CLAIMED", + binding=ResidentBinding(registration_key=("X", "rx"), training_slot=slot), + ) + + +class TestClaimBatch: + def test_one_operation_becomes_one_stamped_batch(self): + output = claim_batch(make_run(), FakeOperationQueue([op()])) + + assert len(output.samples) == 2 and all(len(group) == 1 for group in output.samples) + stamped = output.samples[0][0] + assert (stamped.adapter.name, stamped.adapter.registration_id) == ("X", "rx") + assert stamped.adapter.serving_version == 2 and stamped.adapter.slot == 3 + assert stamped.metadata["team"] == "t1" + assert stamped.status == stamped.Status.COMPLETED + assert [group[0].index for group in output.samples] == [0, 1] + assert isinstance(output, ClaimedOperationBatch) + assert output.operation_id == "op1" + assert output.kind == "forward_backward" + assert output.loss_spec == {"loss_fn": "cross_entropy"} + assert output.binding == ResidentBinding(registration_key=("X", "rx"), training_slot=3) + + def test_client_supplied_row_index_is_overwritten(self): + payload = sample_payload() + payload["samples"][0]["index"] = -1 + payload["samples"][1]["index"] = 0 + queue = FakeOperationQueue([op(payload=payload)]) + output = claim_batch(make_run(), queue) + assert [group[0].index for group in output.samples] == [0, 1] + + def test_child_waits_for_a_claim(self, fast_poll): + queue = FakeOperationQueue([None, None, op()]) + output = claim_batch(make_run(), queue) + assert output.operation_id == "op1" + + def test_bad_payload_fails_its_operation_and_the_child_continues(self): + queue = FakeOperationQueue([op("bad", payload={"samples": []}), op("good")]) + output = claim_batch(make_run(), queue) + + assert output.operation_id == "good" + [(failed_id, error, category)] = queue.failed + assert failed_id == "bad" and category == "user" and "no samples" in error + + def test_forward_operations_build_batches_too(self): + payload = {"samples": [{"prompt": "p", "tokens": [1, 2], "response_length": 1, "loss_mask": [1]}]} + queue = FakeOperationQueue([op("fwd", kind="forward", payload=payload)]) + output = claim_batch(make_run(), queue) + assert output.kind == "forward" + assert output.loss_spec is None + assert queue.failed == [] + + +def ready_runtime(fn: MultiLoraOperationBatchFn, name: str, slot: int, kind: str) -> AdapterRolloutRuntime: + run = make_run(name=name, reg=f"r-{name}", slot=9) + runtime = AdapterRolloutRuntime(run) + runtime.state = AdapterRolloutRuntime.READY + runtime.ready_output = ClaimedOperationBatch( + operation_id=f"op-{name}", + kind=kind, + loss_spec=None, + binding=ResidentBinding(registration_key=(name, f"r-{name}"), training_slot=slot), + samples=[[SimpleNamespace(adapter=None, metadata={})]], + ) + fn.runtimes[(run.name, run.registration_id)] = runtime + fn._sync_rotation() + return runtime + + +def merge(fn: MultiLoraOperationBatchFn, selected) -> RolloutFnTrainOutput: + return asyncio.run(fn._merge(selected)) + + +def make_fn(soft_target=100) -> MultiLoraOperationBatchFn: + args = SimpleNamespace( + rollout_batch_size=soft_target, + n_samples_per_prompt=1, + tinker_max_coalesce_wait_s=0.05, + tinker_max_empty_wait_s=0.05, + ) + return MultiLoraOperationBatchFn( + RolloutFnConstructorInput(args=args, data_source=None), + operations=FakeOperationQueue(), + residency=FakeResidency(), + ) + + +class TestSelectionKindLock: + def test_first_ready_locks_the_kind(self): + fn = make_fn() + ready_runtime(fn, "A", 0, "forward_backward") + other = ready_runtime(fn, "B", 1, "forward") + ready_runtime(fn, "C", 2, "forward_backward") + + selected = asyncio.run(fn._select()) + assert sorted(r.run.name for r in selected) == ["A", "C"] + assert other.state == AdapterRolloutRuntime.READY + + def test_soft_target_stops_collection_but_never_trims(self): + fn = make_fn(soft_target=1) + ready_runtime(fn, "A", 0, "forward_backward") + ready_runtime(fn, "B", 1, "forward_backward") + selected = asyncio.run(fn._select()) + assert len(selected) == 1 + + def test_empty_selection_times_out(self): + fn = make_fn() + with pytest.raises(EmptyBatchTimeoutError): + asyncio.run(fn._select()) + + def test_merge_ships_the_converted_plan_and_pad_policy(self): + fn = make_fn() + first = ready_runtime(fn, "A", 0, "forward_backward") + selected = asyncio.run(fn._select()) + output = merge(fn, selected) + assert output.conversion_metadata == { + "batch_kind": "tinker", + "tinker_operation_lanes": [0], + "tinker_loss_by_lane": {0: {}}, + "operation_by_lane": {0: "op-A"}, + "registration_by_lane": {0: ("A", "r-A")}, + "batch_execution_lease": { + "dispatch_id": "lease-1", + "bindings_by_operation": [["op-A", ["A", "r-A", 0]]], + }, + } + assert output.postprocess.pad_to_dp is True + assert first.state == AdapterRolloutRuntime.IDLE and first.ready_output is None + + def test_failed_lease_acquisition_keeps_claimed_output_retryable(self): + class RefusingOnceResidency(FakeResidency): + def __init__(self): + super().__init__() + self.refusals_left = 1 + + async def acquire_batch(self, bindings_by_operation): + if self.refusals_left: + self.refusals_left -= 1 + raise ValueError("stale binding") + return await super().acquire_batch(bindings_by_operation) + + fn = make_fn() + fn.residency = RefusingOnceResidency() + runtime = ready_runtime(fn, "A", 0, "forward_backward") + selected = asyncio.run(fn._select()) + + with pytest.raises(ValueError, match="stale binding"): + merge(fn, selected) + assert runtime.state == AdapterRolloutRuntime.READY + assert runtime.ready_output is not None + + selected = asyncio.run(fn._select()) + output = merge(fn, selected) + assert output.conversion_metadata["operation_by_lane"] == {0: "op-A"} + assert runtime.state == AdapterRolloutRuntime.IDLE and runtime.ready_output is None + + def test_merge_of_a_forward_selection_marks_forward_only(self): + fn = make_fn() + ready_runtime(fn, "A", 0, "forward") + ready_runtime(fn, "B", 1, "forward") + selected = asyncio.run(fn._select()) + output = merge(fn, selected) + assert output.conversion_metadata["tinker_forward_only"] is True + assert output.conversion_metadata["operation_by_lane"] == {0: "op-A", 1: "op-B"} + assert output.conversion_metadata["tinker_operation_lanes"] == [0, 1] + assert output.postprocess.pad_to_dp is True + + +class TestFailedRuntimeSelfHeal: + def test_child_failure_stamps_the_cooldown_clock(self): + class BoomQueue(FakeOperationQueue): + async def claim_data(self, key): + raise RuntimeError("transient engine failure") + + fn = make_fn() + fn.operations = BoomQueue() + runtime = AdapterRolloutRuntime(make_run(name="A", reg="r-A")) + asyncio.run(fn._run_child(runtime)) + assert runtime.state == AdapterRolloutRuntime.FAILED + assert runtime.last_failure is not None + + def test_failed_runtime_relaunches_after_the_cooldown_not_before(self, monkeypatch): + import time + + import miles.rollout.multi_lora.rollout_fn as rollout_module + + fn = make_fn() + runtime = AdapterRolloutRuntime(make_run(name="A", reg="r-A")) + runtime.state = AdapterRolloutRuntime.FAILED + runtime.last_failure = time.monotonic() + fn.runtimes[("A", "r-A")] = runtime + fn._sync_rotation() + + async def scenario(): + monkeypatch.setattr(rollout_module, "_FAILED_RELAUNCH_COOLDOWN_S", 3600.0) + fn._launch_idle_children() + assert runtime.state == AdapterRolloutRuntime.FAILED and runtime.task is None + + monkeypatch.setattr(rollout_module, "_FAILED_RELAUNCH_COOLDOWN_S", 0.0) + fn._launch_idle_children() + assert runtime.state == AdapterRolloutRuntime.IN_FLIGHT and runtime.task is not None + await fn.aclose() + + asyncio.run(scenario()) diff --git a/tests/fast/test_multi_lora_operation_driver.py b/tests/fast/test_multi_lora_operation_driver.py new file mode 100644 index 00000000000..49472992816 --- /dev/null +++ b/tests/fast/test_multi_lora_operation_driver.py @@ -0,0 +1,262 @@ +import asyncio +from types import SimpleNamespace + +import pytest +import ray +from train_multi_lora_operations import ActorGroupWeightUpdater, generate_with_failure_cap, run_control_phase + +from miles.utils.operation_contract import EmptyBatchTimeoutError + + +class Remote: + def __init__(self, log, name, value=None): + self._log, self._name, self._value = log, name, value + + async def remote(self, *args, **kwargs): + self._log.append((self._name, args)) + return self._value + + +def test_control_phase_completes_deferred_publishes_only_after_the_push(): + log: list = [] + + operations = [ + dict(operation_id="opt1", name="A", kind="optim_step"), + dict(operation_id="pub1", name="A", kind="save_weights_for_sampler"), + dict(operation_id="load1", name="A", kind="load_state"), + ] + lease = { + "dispatch_id": "lease-7", + "bindings_by_operation": [["opt1", ["A", "r-A", 0]], ["pub1", ["A", "r-A", 0]], ["load1", ["A", "r-A", 0]]], + } + controller = SimpleNamespace( + claim_ready_control_operations=Remote(log, "claim", {"operations": operations, "lease": lease}), + complete_control_operations=Remote(log, "complete"), + release_batch_lease=Remote(log, "release"), + ) + + async def execute(ops, lease_metadata): + log.append(("execute", tuple(op["operation_id"] for op in ops))) + assert lease_metadata == lease + return { + "opt1": dict(ok=True, result=dict(grad_norm=1.0, learning_rate=1e-4)), + "pub1": dict(ok=True, deferred="publish"), + "load1": dict(ok=True, deferred="publish", result=dict(step=4, path="/s")), + } + + async def update_weights(): + log.append(("update_weights", ())) + + actor_model = SimpleNamespace(execute_tinker_controls=execute, update_weights=update_weights) + asyncio.run(run_control_phase(actor_model, controller, ActorGroupWeightUpdater(actor_model))) + + order = [name for name, _ in log] + assert order == ["claim", "execute", "complete", "update_weights", "complete", "release"] + first_complete = log[2][1][0] + assert set(first_complete) == {"opt1"} + deferred_complete = log[4][1][0] + assert deferred_complete == { + "pub1": dict(ok=True), + "load1": dict(ok=True, result=dict(step=4, path="/s")), + } + assert log[5][1] == (lease,) + + +def test_immediate_only_batch_releases_at_its_completion_boundary(): + log: list = [] + operations = [dict(operation_id="opt1", name="A", kind="optim_step")] + lease = {"dispatch_id": "lease-8", "bindings_by_operation": [["opt1", ["A", "r-A", 0]]]} + controller = SimpleNamespace( + claim_ready_control_operations=Remote(log, "claim", {"operations": operations, "lease": lease}), + complete_control_operations=Remote(log, "complete"), + release_batch_lease=Remote(log, "release"), + ) + + async def execute(ops, lease_metadata): + log.append(("execute", ())) + return {"opt1": dict(ok=True, result=dict(grad_norm=1.0, learning_rate=1e-4))} + + async def update_weights(): + log.append(("update_weights", ())) + + actor_model = SimpleNamespace(execute_tinker_controls=execute, update_weights=update_weights) + asyncio.run(run_control_phase(actor_model, controller, ActorGroupWeightUpdater(actor_model))) + assert [name for name, _ in log] == ["claim", "execute", "complete", "release", "update_weights"] + + +def test_control_phase_still_pushes_with_no_operations(): + log: list = [] + controller = SimpleNamespace( + claim_ready_control_operations=Remote(log, "claim", {"operations": [], "lease": None}), + complete_control_operations=Remote(log, "complete"), + release_batch_lease=Remote(log, "release"), + ) + + async def update_weights(): + log.append(("update_weights", ())) + + actor_model = SimpleNamespace(execute_tinker_controls=None, update_weights=update_weights) + asyncio.run(run_control_phase(actor_model, controller, ActorGroupWeightUpdater(actor_model))) + assert [name for name, _ in log] == ["claim", "update_weights"] + + +def test_validate_tinker_args_defaults_the_rollout_plane(): + from miles.rollout.multi_lora.rollout_fn import MultiLoraOperationBatchFn, TinkerNullDataSource + from miles.utils.misc import load_function + from miles.utils.tinker import validate_tinker_args + + args = SimpleNamespace( + tinker_backend=True, + multi_lora_n_adapters=4, + rollout_function_path=None, + data_source_path="miles.rollout.data_source.RolloutDataSourceWithBuffer", + use_dynamic_global_batch_size=False, + ) + validate_tinker_args(args) + assert args.rollout_function_path == "miles.rollout.multi_lora.rollout_fn.MultiLoraOperationBatchFn" + assert args.data_source_path == "miles.rollout.multi_lora.rollout_fn.TinkerNullDataSource" + assert args.use_dynamic_global_batch_size is True + assert load_function(args.rollout_function_path) is MultiLoraOperationBatchFn + assert load_function(args.data_source_path) is TinkerNullDataSource + + args.rollout_function_path = "my.custom.Fn" + args.data_source_path = "my.custom.Source" + validate_tinker_args(args) + assert args.rollout_function_path == "my.custom.Fn" + assert args.data_source_path == "my.custom.Source" + + off = SimpleNamespace(tinker_backend=False) + validate_tinker_args(off) + + +class TestDataBatchFinalizer: + def _pack(self): + lease = { + "dispatch_id": "lease-9", + "bindings_by_operation": [["fb1", ["A", "r-A", 0]], ["fb2", ["B", "r-B", 1]]], + } + pack = {"data_ref": None, "tinker_dispatch": {"operation_ids": ["fb1", "fb2"], "lease": lease}} + return pack, lease + + def test_normal_outcome_never_calls_the_finalizer(self): + from train_multi_lora_operations import train_data_batch + + from miles.backends.megatron_utils.ft.types import TrainStepOutcome + + log: list = [] + controller = SimpleNamespace(fail_tinker_batch=Remote(log, "fail")) + + async def train(rollout_id, rollout_data): + return [TrainStepOutcome.NORMAL, TrainStepOutcome.NORMAL] + + pack, _ = self._pack() + asyncio.run(train_data_batch(SimpleNamespace(train=train), controller, 0, pack)) + assert log == [] + + def test_abnormal_outcome_fails_the_batch_operations_and_releases_the_lease(self): + from train_multi_lora_operations import train_data_batch + + from miles.backends.megatron_utils.ft.types import TrainStepOutcome + + log: list = [] + controller = SimpleNamespace(fail_tinker_batch=Remote(log, "fail")) + + async def train(rollout_id, rollout_data): + return [TrainStepOutcome.NORMAL, TrainStepOutcome.DISCARDED_SHOULD_RETRY] + + pack, lease = self._pack() + asyncio.run(train_data_batch(SimpleNamespace(train=train), controller, 3, pack)) + [(name, (operation_ids, error, lease_arg))] = log + assert name == "fail" and operation_ids == ["fb1", "fb2"] and lease_arg == lease + assert "discarded_should_retry" in error and "resubmit" in error + + def test_train_exception_finalizes_then_reraises(self): + import pytest + from train_multi_lora_operations import train_data_batch + + log: list = [] + controller = SimpleNamespace(fail_tinker_batch=Remote(log, "fail")) + + async def train(rollout_id, rollout_data): + raise RuntimeError("trainer rank died") + + pack, lease = self._pack() + with pytest.raises(RuntimeError, match="trainer rank died"): + asyncio.run(train_data_batch(SimpleNamespace(train=train), controller, 3, pack)) + [(name, (operation_ids, error, lease_arg))] = log + assert name == "fail" and operation_ids == ["fb1", "fb2"] and lease_arg == lease + assert "trainer rank died" in error and "poisoned" in error + + def test_missing_dispatch_summary_still_finalizes_with_empty_ids(self): + from train_multi_lora_operations import train_data_batch + + from miles.backends.megatron_utils.ft.types import TrainStepOutcome + + log: list = [] + controller = SimpleNamespace(fail_tinker_batch=Remote(log, "fail")) + + async def train(rollout_id, rollout_data): + return [TrainStepOutcome.DISCARDED_SHOULD_RETRY] + + asyncio.run(train_data_batch(SimpleNamespace(train=train), controller, 0, {"data_ref": None})) + [(name, (operation_ids, error, lease_arg))] = log + assert operation_ids == [] and lease_arg is None + + +class FakeRayTaskError(ray.exceptions.RayTaskError): + def __init__(self, cause): + Exception.__init__(self, str(cause)) + self.cause = cause + self.function_name = "generate" + self.traceback_str = f"fake traceback: {cause}" + + def as_instanceof_cause(self): + return self.cause + + +class TestGenerateFailureCap: + class Executor: + def __init__(self, outcomes): + self.outcomes = list(outcomes) + + async def generate(self, rollout_id): + outcome = self.outcomes.pop(0) + if isinstance(outcome, Exception): + raise outcome + return outcome + + def attempt(self, executor, streak, cap=3): + return asyncio.run(generate_with_failure_cap(executor, 0, streak, cap)) + + def test_a_failure_below_the_cap_skips_the_round(self): + executor = self.Executor([FakeRayTaskError(RuntimeError("engine died"))]) + assert self.attempt(executor, streak=0) == (None, 1) + + def test_a_success_resets_the_streak(self): + executor = self.Executor([{"batch": 1}]) + assert self.attempt(executor, streak=2) == ({"batch": 1}, 0) + + def test_the_cap_reraises(self): + executor = self.Executor([FakeRayTaskError(RuntimeError("engine died"))]) + with pytest.raises(ray.exceptions.RayTaskError): + self.attempt(executor, streak=2, cap=3) + + def test_zero_cap_fails_fast(self): + executor = self.Executor([FakeRayTaskError(RuntimeError("engine died"))]) + with pytest.raises(ray.exceptions.RayTaskError): + self.attempt(executor, streak=0, cap=0) + + def test_empty_batch_timeout_neither_counts_nor_resets(self): + executor = self.Executor([FakeRayTaskError(EmptyBatchTimeoutError("idle"))]) + assert self.attempt(executor, streak=2) == (None, 2) + + def test_interleaved_successes_keep_the_loop_alive(self): + executor = self.Executor( + [FakeRayTaskError(RuntimeError("a")), {"batch": 1}, FakeRayTaskError(RuntimeError("b"))] + ) + data, streak = self.attempt(executor, streak=0, cap=2) + assert data is None and streak == 1 + data, streak = self.attempt(executor, streak=streak, cap=2) + assert data == {"batch": 1} and streak == 0 + data, streak = self.attempt(executor, streak=streak, cap=2) + assert data is None and streak == 1 diff --git a/tests/fast/utils/test_arguments.py b/tests/fast/utils/test_arguments.py index c26945b4653..b33f3229d24 100644 --- a/tests/fast/utils/test_arguments.py +++ b/tests/fast/utils/test_arguments.py @@ -574,6 +574,7 @@ def _parse(self, extra): [ "--multi-lora-n-adapters", "2", + "--tinker-backend", "--lora-rank", "8", "--target-modules", @@ -585,6 +586,26 @@ def _parse(self, extra): + REQUIRED_ARGS ) + def test_rejects_multi_lora_without_tinker_backend(self): + parser = argparse.ArgumentParser() + get_miles_extra_args_provider()(parser) + args = parser.parse_args( + [ + "--multi-lora-n-adapters", + "2", + "--lora-rank", + "8", + "--target-modules", + "linear_qkv", + "--num-rollout", + "1", + ] + + REQUIRED_ARGS + ) + + with pytest.raises(AssertionError, match="requires --tinker-backend"): + miles_validate_args(args) + def test_rejects_multiple_tokenizer_workers(self): # Each sglang tokenizer worker holds its own LoRA registry, so per-step # upserts fail non-deterministically; fail at launch, not first push. @@ -600,13 +621,13 @@ def test_accepts_default_single_tokenizer_worker(self): assert args.multi_lora is True - def test_defaults_rollout_fn_and_data_source_to_multi_lora(self): + def test_defaults_rollout_fn_and_data_source_to_tinker(self): args = self._parse([]) miles_validate_args(args) - assert args.rollout_function_path == "miles.rollout.multi_lora.async_rollout.generate_rollout_multi_lora" - assert args.data_source_path == "miles.rollout.multi_lora.data_source.MultiLoRAAsyncDataSource" + assert args.rollout_function_path == "miles.rollout.multi_lora.rollout_fn.MultiLoraOperationBatchFn" + assert args.data_source_path == "miles.rollout.multi_lora.rollout_fn.TinkerNullDataSource" assert args.rollout_global_dataset is True def test_keeps_user_supplied_rollout_fn_and_data_source(self): @@ -620,8 +641,8 @@ def test_keeps_user_supplied_rollout_fn_and_data_source(self): assert args.data_source_path == "my.custom.DataSource" def test_empty_wait_is_a_registered_argument(self): - assert self._parse([]).multi_lora_max_empty_wait_s == 30.0 - assert self._parse(["--multi-lora-max-empty-wait-s", "5"]).multi_lora_max_empty_wait_s == 5.0 + assert self._parse([]).tinker_max_empty_wait_s == 5.0 + assert self._parse(["--tinker-max-empty-wait-s", "9"]).tinker_max_empty_wait_s == 9.0 def test_rejects_non_adam_optimizer(self): # Per-slot optimizer isolation (state init, retirement cleanup, step diff --git a/tests/fast/utils/test_tinker_predicates.py b/tests/fast/utils/test_tinker_predicates.py new file mode 100644 index 00000000000..266b34efaed --- /dev/null +++ b/tests/fast/utils/test_tinker_predicates.py @@ -0,0 +1,37 @@ +from types import SimpleNamespace + +import pytest + +from miles.utils.multi_lora import uses_multi_lora_operation_executor, validate_multi_lora_args +from miles.utils.tinker import uses_explicit_training_operations, validate_tinker_args + + +def _args(tinker_backend: bool, n_adapters: int) -> SimpleNamespace: + return SimpleNamespace( + tinker_backend=tinker_backend, + multi_lora_n_adapters=n_adapters, + multi_lora=n_adapters > 0, + ) + + +class TestPredicateRoles: + def test_operation_semantics_is_the_protocol_flag_alone(self): + assert uses_explicit_training_operations(_args(True, 0)) + assert uses_explicit_training_operations(_args(True, 4)) + assert not uses_explicit_training_operations(_args(False, 4)) + assert not uses_explicit_training_operations(_args(False, 0)) + + def test_executor_requires_protocol_and_slots(self): + assert uses_multi_lora_operation_executor(_args(True, 4)) + assert not uses_multi_lora_operation_executor(_args(True, 0)) + assert not uses_multi_lora_operation_executor(_args(False, 4)) + + +class TestValidationClosesTheGap: + def _validate(self, args) -> None: + validate_multi_lora_args(args) + validate_tinker_args(args) + + def test_tinker_without_slots_is_rejected(self): + with pytest.raises(AssertionError, match="--multi-lora-n-adapters"): + self._validate(_args(True, 0)) diff --git a/tests/fast/utils/test_tinker_sample_channels.py b/tests/fast/utils/test_tinker_sample_channels.py new file mode 100644 index 00000000000..748b87d9889 --- /dev/null +++ b/tests/fast/utils/test_tinker_sample_channels.py @@ -0,0 +1,59 @@ +from miles.ray.rollout.train_data_conversion import ROLLOUT_DATA_TENSOR_DTYPES +from miles.utils.types import Sample + + +def test_wire_dtypes_keep_binary_mask_and_add_float_channels(): + assert ROLLOUT_DATA_TENSOR_DTYPES["loss_masks"] == "int32" + assert ROLLOUT_DATA_TENSOR_DTYPES["loss_weights"] == "float32" + assert ROLLOUT_DATA_TENSOR_DTYPES["advantages"] == "float32" + + +def test_sample_round_trips_the_channels(): + sample = Sample.from_dict( + { + "prompt": "p", + "tokens": [1, 2, 3], + "response_length": 2, + "loss_mask": [1, 1], + "loss_weights": [0.5, -1.0], + "advantages": [2.0, 0.0], + "status": "completed", + } + ) + assert sample.loss_weights == [0.5, -1.0] + assert sample.advantages == [2.0, 0.0] + assert Sample.from_dict(sample.to_dict()).loss_weights == [0.5, -1.0] + + +def test_merge_pads_the_channels_over_the_observation_span(monkeypatch): + from miles.rollout.generate_utils.sample_utils import merge_samples + + class _Tok: + def decode(self, tokens): + return "obs" + + a = Sample( + prompt="p", + status=Sample.Status.COMPLETED, + tokens=[1, 2, 3], + response="x", + response_length=1, + loss_mask=[1], + loss_weights=[0.5], + advantages=[1.0], + rollout_log_probs=[-0.1], + ) + b = Sample( + prompt="p", + status=Sample.Status.COMPLETED, + tokens=[1, 2, 3, 9, 4, 5], + response="y", + response_length=2, + loss_mask=[1, 1], + loss_weights=[1.5, 2.5], + advantages=[0.0, -1.0], + rollout_log_probs=[-0.2, -0.3], + ) + merged = merge_samples([a, b], tokenizer=_Tok()) + assert merged.loss_weights == [0.5, 0.0, 1.5, 2.5] + assert merged.advantages == [1.0, 0.0, 0.0, -1.0] diff --git a/train_multi_lora_async.py b/train_multi_lora_async.py deleted file mode 100644 index 0f8cdf54c8f..00000000000 --- a/train_multi_lora_async.py +++ /dev/null @@ -1,106 +0,0 @@ -"""Fully-async multi-LoRA trainer driver.""" - -import asyncio -import logging -from pathlib import Path - -import ray - -from miles.ray.multi_lora.controller import create_multilora_controller, get_multi_lora_controller -from miles.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models -from miles.utils import object_store -from miles.utils.adapter_config import parse_adapter_run_yaml -from miles.utils.arguments import parse_args -from miles.utils.audit_utils.process_identity import MainProcessIdentity -from miles.utils.data import remove_rollout_data_refs -from miles.utils.logging_utils import configure_logger -from miles.utils.multi_lora import EmptyBatchTimeoutError, define_new_adapter_metrics -from miles.utils.tracking_utils.tracking import init_tracking - -logger = logging.getLogger(__name__) - - -def _is_empty_batch_timeout(task_error: ray.exceptions.RayTaskError) -> bool: - cause = getattr(task_error, "cause", None) - if isinstance(cause, EmptyBatchTimeoutError): - return True - return isinstance(task_error.as_instanceof_cause(), EmptyBatchTimeoutError) - - -async def main(args): - assert ( - not args.colocate - ), "Colocation is not supported for fully-async training (generation needs continuous GPU; colocate time-shares)." - configure_logger(args, source=MainProcessIdentity()) - - # The multi-LoRA rollout fn / data source / global dataset flags are - # defaulted by miles_validate_args when --multi-lora-n-adapters > 0. - pgs = create_placement_groups(args) - object_store.init_instance(args, contribute_segment=False) - init_tracking(args) - rollout_manager, _num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"]) - - # Create a controller nclusing MultiLoRAController and MultiLoRAHTTPServer to manage lora - router_ip, router_port = await rollout_manager.get_router_address.remote() - args.sglang_router_ip, args.sglang_router_port = router_ip, router_port - controller = create_multilora_controller(args, f"http://{router_ip}:{router_port}") - await controller.start.remote() - host = await controller.http_host.remote() - api_port = await controller.api_port.remote() - logger.info(f"Multi-LoRA control API listening on http://{host}:{api_port} (head node)") - - actor_model, _ = await create_training_models(args, pgs, rollout_manager) - - # CLI-registered adapters are loaded and pushed by the loop's first - # reconcile + update_weights. - for name, path in args.multi_lora_adapters: - config = parse_adapter_run_yaml(Path(path)) - await controller.register_adapter.remote(name, config) - - rollout_id = 0 - while True: - snapshot = await get_multi_lora_controller().snapshot.remote() - - # handle dynamic metrics in tracking backend - define_new_adapter_metrics(snapshot) - if not (snapshot["pending"] or snapshot["active"] or snapshot["retiring"] or snapshot["cleanup"]): - if not args.multi_lora_service_mode: - logger.info("No adapters; exiting.") - break - logger.info(f"No adapters; sleeping for {args.multi_lora_idle_poll_s}s...") - await asyncio.sleep(args.multi_lora_idle_poll_s) - continue - - # Reconcile + push before generate: the push promotes pending adapters, - # and only then does the data source sample them. The actor pushes only - # stale adapter weights (newly loaded, or stepped by the last batch). - await actor_model.reconcile_adapters() - await actor_model.update_weights() - - # With nothing active, generate would wait forever. - post_update = await get_multi_lora_controller().snapshot.remote() - if not (post_update["active"] or post_update["retiring"]): - continue - - try: - rollout_data = await rollout_manager.generate.remote(rollout_id) - except ray.exceptions.RayTaskError as e: - if _is_empty_batch_timeout(e): - logger.warning(f"Generate timed out with no trainable groups; retrying reconcile/update. {e}") - continue - raise - await actor_model.train(rollout_id, rollout_data) - remove_rollout_data_refs(args, rollout_data) - - # Per-adapter save cadence decided inside save_model. - await actor_model.save_model(rollout_id) - - rollout_id += 1 - - await rollout_manager.dispose.remote() - await controller.stop.remote() - - -if __name__ == "__main__": - args = parse_args() - asyncio.run(main(args)) diff --git a/train_multi_lora_operations.py b/train_multi_lora_operations.py new file mode 100644 index 00000000000..fd2645e793c --- /dev/null +++ b/train_multi_lora_operations.py @@ -0,0 +1,177 @@ +import asyncio +import logging + +import ray + +from miles.ray.multi_lora.controller import create_multi_lora_controller +from miles.ray.placement_group import create_placement_groups, create_training_models +from miles.ray.rollout.components import create_rollout_components +from miles.utils import object_store +from miles.utils.arguments import parse_args +from miles.utils.audit_utils.process_identity import MainProcessIdentity +from miles.utils.data import remove_rollout_data_refs +from miles.utils.logging_utils import configure_logger +from miles.utils.operation_contract import EmptyBatchTimeoutError +from miles.utils.tracking_utils.tracking import init_tracking + +logger = logging.getLogger(__name__) + + +def _is_empty_batch_timeout(task_error: ray.exceptions.RayTaskError) -> bool: + cause = getattr(task_error, "cause", None) + if isinstance(cause, EmptyBatchTimeoutError): + return True + return isinstance(task_error.as_instanceof_cause(), EmptyBatchTimeoutError) + + +class ActorGroupWeightUpdater: + def __init__(self, actor_model) -> None: + self._actor_model = actor_model + + async def update_weights(self) -> None: + await self._actor_model.update_weights() + + +async def train_data_batch(actor_model, controller, rollout_id: int, rollout_data) -> None: + from miles.backends.megatron_utils.ft.types import TrainStepOutcome + + dispatch = rollout_data.get("tinker_dispatch") or {} + operation_ids = list(dispatch.get("operation_ids") or []) + lease = dispatch.get("lease") + + try: + outcomes = await actor_model.train(rollout_id, rollout_data) + except Exception as e: + await controller.fail_tinker_batch.remote( + operation_ids, + f"train dispatch raised on the trainer: {e}; the batch did not commit and its " + "gradient window is poisoned — resubmit the batch and optim_step again", + lease, + ) + raise + outcomes = outcomes if isinstance(outcomes, list) else [outcomes] + abnormal = sorted({str(outcome) for outcome in outcomes if outcome != TrainStepOutcome.NORMAL}) + if abnormal: + await controller.fail_tinker_batch.remote( + operation_ids, + f"train step finished without committing (outcome {', '.join(abnormal)}); the batch's " + "gradient window is poisoned — resubmit the batch and optim_step again", + lease, + ) + + +async def run_control_phase(actor_model, controller, weight_updater) -> None: + claimed = await controller.claim_ready_control_operations.remote() + operations, lease = claimed["operations"], claimed["lease"] + released = lease is None + try: + deferred: list[str] = [] + if operations: + results = await actor_model.execute_tinker_controls(operations, lease) + deferred = [op_id for op_id, outcome in results.items() if outcome.get("deferred") == "publish"] + immediate = {op_id: outcome for op_id, outcome in results.items() if op_id not in deferred} + if immediate: + await controller.complete_control_operations.remote(immediate) + if not deferred and not released: + released = True + await controller.release_batch_lease.remote(lease) + + # Push staged weights (publishes and load_state re-publishes); a no-op + # when nothing is staged. Serving versions bump as the push commits. + await weight_updater.update_weights() + + if deferred: + await controller.complete_control_operations.remote( + { + op_id: {key: value for key, value in results[op_id].items() if key != "deferred"} + for op_id in deferred + } + ) + released = True + await controller.release_batch_lease.remote(lease) + finally: + if not released: + await controller.release_batch_lease.remote(lease) + + +async def generate_with_failure_cap(rollout_executor, rollout_id: int, failure_streak: int, cap: int): + """One tolerated generate attempt; returns (rollout_data or None, updated consecutive-failure streak).""" + try: + return await rollout_executor.generate(rollout_id), 0 + except ray.exceptions.RayTaskError as e: + if _is_empty_batch_timeout(e): + # The data queue is idle; yield to the control phase so queued optim/save/load never wait behind it. + return None, failure_streak + failure_streak += 1 + if failure_streak >= cap: + raise + # Skipping the round self-heals: failure paths restore unconsumed claims to READY for re-dispatch. + logger.exception( + f"[tinker] generate failed ({failure_streak} consecutive, cap {cap}); " + f"keeping the multi-tenant service alive: {e}" + ) + return None, failure_streak + + +async def main(args): + assert ( + not args.colocate + ), "Colocation is not supported for Multi-LoRA operations (generation needs continuous GPU; colocate time-shares)." + configure_logger(args, source=MainProcessIdentity()) + + pgs = create_placement_groups(args) + object_store.init_instance(args, contribute_segment=False) + init_tracking(args) + rollout_components = create_rollout_components(args, pgs["rollout"]) + inference_controller = rollout_components.inference_controller + rollout_executor = rollout_components.rollout_executor + + inference_endpoint = await inference_controller.get_inference_endpoint() + args.sglang_router_ip, args.sglang_router_port = inference_endpoint.host, inference_endpoint.port + multi_lora_controller = create_multi_lora_controller(args, inference_endpoint.base_url) + await multi_lora_controller.start.remote() + host = await multi_lora_controller.http_host.remote() + api_port = await multi_lora_controller.api_port.remote() + logger.info(f"Tinker control API listening on http://{host}:{api_port} (head node)") + + # As in train_async.py, actor_model is the actor RayTrainGroup, with the weight-update owner wired in. + actor_model, _ = await create_training_models(args, pgs, rollout_components.weight_update_owner) + weight_updater = ActorGroupWeightUpdater(actor_model) + + # The trainer is up: flip readiness so /api/v1/healthz stops answering 503. + await multi_lora_controller.set_trainer_ready.remote() + + rollout_id = 0 + generate_failures = 0 + while True: + # This handle is the controller's only owning reference; rebinding it would let Ray reap the actor. + snapshot = await multi_lora_controller.snapshot.remote() + if not (snapshot["pending"] or snapshot["ready"] or snapshot["retiring"] or snapshot["cleanup"]): + logger.info(f"No adapters; sleeping for {args.multi_lora_idle_poll_s}s...") + await asyncio.sleep(args.multi_lora_idle_poll_s) + continue + + # Residency first: retire deregistered adapters, then load bound registrations. + await actor_model.reconcile_tinker_adapters() + + await run_control_phase(actor_model, multi_lora_controller, weight_updater) + + post_control = await multi_lora_controller.snapshot.remote() + if not post_control["ready"]: + continue + + # Per-rollout engine preparation; a no-op behind today's combined manager. + await inference_controller.prepare_rollout(rollout_id) + rollout_data, generate_failures = await generate_with_failure_cap( + rollout_executor, rollout_id, generate_failures, args.multi_lora_max_consecutive_generate_failures + ) + if rollout_data is None: + continue + await train_data_batch(actor_model, multi_lora_controller, rollout_id, rollout_data) + remove_rollout_data_refs(args, rollout_data) + rollout_id += 1 + + +if __name__ == "__main__": + args = parse_args() + asyncio.run(main(args))