From fe5a088260d70759a3b63b6ee9808569763122b2 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Wed, 14 Oct 2026 14:53:27 +0800 Subject: [PATCH 01/46] feat: add on-demand training with isolated lifecycle management Preserve legacy expiration semantics, isolate user initialization failures, and keep recovery nonblocking. Bound retry receipts and reconcile lost completion notifications. --- .../client/server/on_demand_training.yaml | 40 ++ docs/design/on-demand-training.md | 169 +++++ .../validation/on-demand-training/README.md | 35 + .../Usage Guide/Server and Client/Server.md | 6 + .../\346\234\215\345\212\241\347\253\257.md" | 6 + src/twinkle/infra/_ray/ray_helper.py | 20 +- src/twinkle/infra/_ray/resource_manager.py | 31 + src/twinkle/model/megatron/megatron.py | 3 + .../model/transformers/transformers.py | 9 + src/twinkle/protocol/types/model.py | 47 +- src/twinkle/server/config/application_spec.py | 56 +- src/twinkle/server/config/hybrid.py | 21 + src/twinkle/server/config/server_config.py | 20 + src/twinkle/server/config/training.py | 127 ++++ .../server/config/training_normalization.py | 87 +++ src/twinkle/server/gateway/app.py | 6 +- src/twinkle/server/gateway/tinker_handlers.py | 8 +- .../server/gateway/twinkle_handlers.py | 9 +- .../server/launcher/env_propagation.py | 1 + .../server/launcher/server_launcher.py | 71 +- src/twinkle/server/lifecycle/envelope.py | 5 + src/twinkle/server/lifecycle/submit.py | 19 +- src/twinkle/server/model/app.py | 83 +-- src/twinkle/server/model/tinker_handlers.py | 68 +- src/twinkle/server/model/twinkle_handlers.py | 80 +-- src/twinkle/server/state/backend/factory.py | 5 +- .../server/state/backend/memory_backend.py | 13 +- src/twinkle/server/state/model_manager.py | 4 + src/twinkle/server/state/models.py | 8 + src/twinkle/server/state/sequence.py | 27 + src/twinkle/server/state/server_state.py | 37 +- src/twinkle/server/task_queue/mixin.py | 49 +- src/twinkle/server/task_queue/types.py | 2 + src/twinkle/server/task_queue/worker.py | 27 +- src/twinkle/server/training/__init__.py | 2 + src/twinkle/server/training/actor.py | 156 ++++ src/twinkle/server/training/allocation.py | 93 +++ src/twinkle/server/training/catalog.py | 44 ++ src/twinkle/server/training/controller.py | 668 +++++++++++++++++ src/twinkle/server/training/driver.py | 172 +++++ src/twinkle/server/training/errors.py | 14 + src/twinkle/server/training/records.py | 58 ++ src/twinkle/server/training/router.py | 263 +++++++ src/twinkle/server/training/runtime.py | 214 ++++++ src/twinkle/server/training/store.py | 101 +++ src/twinkle_client/auto/agent/server_tools.py | 79 +- .../model/multi_lora_transformers.py | 55 +- tests/server/cli/test_cli.py | 2 +- tests/server/config/test_server_config.py | 29 +- tests/server/config/test_training_config.py | 187 +++++ tests/server/contract/test_error_wire.py | 3 + .../server/integration/test_full_param_e2e.py | 6 +- .../integration/test_mock_mode_startup.py | 100 +-- .../lifecycle/test_retrieve_endpoint.py | 3 + .../server/lifecycle/test_run_submit_dedup.py | 6 +- .../test_tinker_retrieve_regression.py | 3 + tests/server/model/test_replica_lifecycle.py | 13 +- tests/server/start_e2e_server.py | 2 + tests/server/test_gateway_services.py | 5 +- tests/server/training/test_allocation.py | 23 + tests/server/training/test_controller.py | 677 ++++++++++++++++++ .../server/training/test_http_integration.py | 16 + tests/server/training/test_launcher.py | 39 + tests/server/training/test_router.py | 87 +++ tests/server/training/test_runtime_queue.py | 189 +++++ tests/server/training/test_sequence.py | 70 ++ .../server/training/test_storage_recovery.py | 64 ++ .../test_resident_training_reset.py | 39 + tests/twinkle_client/test_async_components.py | 57 +- .../on_demand_training/allocation_probe.py | 88 +++ .../on_demand_training/device_validation.py | 440 ++++++++++++ .../on_demand_training/fragmentation_probe.py | 84 +++ .../on_demand_training/http_probe.py | 246 +++++++ .../on_demand_training/on_demand_sdk.py | 18 + 74 files changed, 5128 insertions(+), 486 deletions(-) create mode 100644 cookbook/client/server/on_demand_training.yaml create mode 100644 docs/design/on-demand-training.md create mode 100644 docs/design/validation/on-demand-training/README.md create mode 100644 src/twinkle/server/config/hybrid.py create mode 100644 src/twinkle/server/config/training.py create mode 100644 src/twinkle/server/config/training_normalization.py create mode 100644 src/twinkle/server/state/sequence.py create mode 100644 src/twinkle/server/training/__init__.py create mode 100644 src/twinkle/server/training/actor.py create mode 100644 src/twinkle/server/training/allocation.py create mode 100644 src/twinkle/server/training/catalog.py create mode 100644 src/twinkle/server/training/controller.py create mode 100644 src/twinkle/server/training/driver.py create mode 100644 src/twinkle/server/training/errors.py create mode 100644 src/twinkle/server/training/records.py create mode 100644 src/twinkle/server/training/router.py create mode 100644 src/twinkle/server/training/runtime.py create mode 100644 src/twinkle/server/training/store.py create mode 100644 tests/server/config/test_training_config.py create mode 100644 tests/server/training/test_allocation.py create mode 100644 tests/server/training/test_controller.py create mode 100644 tests/server/training/test_http_integration.py create mode 100644 tests/server/training/test_launcher.py create mode 100644 tests/server/training/test_router.py create mode 100644 tests/server/training/test_runtime_queue.py create mode 100644 tests/server/training/test_sequence.py create mode 100644 tests/server/training/test_storage_recovery.py create mode 100644 tests/transformers/test_resident_training_reset.py create mode 100644 tools/validation/on_demand_training/allocation_probe.py create mode 100644 tools/validation/on_demand_training/device_validation.py create mode 100644 tools/validation/on_demand_training/fragmentation_probe.py create mode 100644 tools/validation/on_demand_training/http_probe.py create mode 100644 tools/validation/on_demand_training/on_demand_sdk.py diff --git a/cookbook/client/server/on_demand_training.yaml b/cookbook/client/server/on_demand_training.yaml new file mode 100644 index 000000000..2646d3aa2 --- /dev/null +++ b/cookbook/client/server/on_demand_training.yaml @@ -0,0 +1,40 @@ +# On-demand training: applications stays the public configuration entry. +training: + allocation_timeout_s: 3 + startup_timeout_s: 1800 + shutdown_timeout_s: 60 + sweep_interval_s: 5 +applications: + - name: model-qwen + route_prefix: /api/v1/model/Qwen/Qwen3-1.7B + import_path: model + args: + model_id: ms://Qwen/Qwen3-1.7B + backend: transformers + device_type: cuda + modes: + lora: + num_devices: 2 + device_mesh: {device_type: cuda, fsdp_size: 2, dp_size: 1} + min_replicas: 1 + max_loras: 5 + max_r: 32 + idle_timeout_s: 1800 + target_modules: all-linear + full: + num_devices: 4 + device_mesh: {device_type: cuda, fsdp_size: 4, dp_size: 1} + min_replicas: 0 + idle_timeout_s: 1800 + - name: model-small + route_prefix: /api/v1/model/Qwen/Qwen3-0.6B + import_path: model + args: + model_id: ms://Qwen/Qwen3-0.6B + backend: megatron + device_type: cuda + modes: + lora: + num_devices: 2 + device_mesh: {device_type: cuda, tp_size: 2, dp_size: 1} + max_loras: 5 diff --git a/docs/design/on-demand-training.md b/docs/design/on-demand-training.md new file mode 100644 index 000000000..892767829 --- /dev/null +++ b/docs/design/on-demand-training.md @@ -0,0 +1,169 @@ +# On-demand training and device reuse + +The catalog, controller, allocation, and SDK routing are implemented. Review fixes +are under local validation. Real-device reruns are deferred at the user's request; +see the [validation record](validation/on-demand-training/README.md). + +## Scope and requirements + +A fixed device pool serves a catalog larger than its resident capacity. The demand +unit is a training handle; the allocation unit is a model instance on N devices. +Sampling retains its existing deployment lifecycle. This change does not introduce +resource waiting queues, preemption, training batching, or automatic checkpointing. + +| Requirement | Behavior | +| --- | --- | +| R01 | Support the existing Transformers and Megatron training backends. | +| R02 | Administrators declare model sources, modes, parallelism, and N devices. | +| R03 | New mode specifications default to lazy startup; `min_replicas` keeps resident capacity. | +| R04 | Initialize residents before accepting dynamic training; all instances use Ray resources. | +| R05 | An instance reserves N devices on one node; independent bases do not share that reservation. | +| R06 | A model may offer both LoRA and full training through separate instances. | +| R07 | Fill compatible LoRA instances first; expand when slots and devices permit. | +| R08 | Each full handle exclusively owns an instance, subject to user quota. | +| R09 | Capacity failures are terminal; callers explicitly issue a new create to retry. | +| R10 | New specifications use business inactivity, defaulting to 1800 seconds. | +| R11 | Heartbeats, polling, health, and information queries do not renew business inactivity. | +| R12 | Remove an empty dynamic instance immediately; retain clean resident capacity. | +| R13 | Expired handles stay invalid; users explicitly recreate and restore saved state. | +| R14 | Worker failure or controller restart invalidates affected handles without replaying training. | +| R15 | Retain the existing instance queue and backend execution gate. | +| R16 | Keep `applications`, old model arguments, and SDK methods; use one controller. | +| R17 | Every create gets a distinct training ID; identical adapter labels cannot reattach old state. | +| R18 | Validate both backends and modes on real devices before marking the PR ready. | + +## Configuration and compatibility + +The public entry point remains `applications` with `import_path: model`. +The [configuration example](../../cookbook/client/server/on_demand_training.yaml) +shows `args.modes.lora/full`. Each mode sets `num_devices`, `device_mesh`, +`min_replicas`, `max_replicas`, and `idle_timeout_s`. LoRA additionally sets +`max_loras`, `max_r`, and its existing backend options. + +- New modes default to `min_replicas=0`, `idle_timeout_s=1800`, five LoRA slots, + and maximum LoRA rank 32. Full capacity is one handle. +- Mesh world size must equal N; N and capacities are positive integers. Timeouts + must be finite and positive. Minimum replicas cannot exceed the maximum. +- `modes` cannot be mixed with legacy single-mode settings or deployment replicas. +- Legacy `train_mode`, device counts, mesh, strategy, LoRA parameters, and environment + variables normalize into the same runtime catalog. +- Existing deployment replica bounds keep their resident-minimum/instance-limit + meaning. Explicit `args.min_replicas` and `args.max_replicas` take precedence. +- **Legacy `adapter_timeout` remains a session heartbeat timeout.** It is not converted + to business inactivity. `adapter_max_lifetime` remains an absolute handle lifetime. + Their existing defaults are 1800 and 36000 seconds respectively. +- A legacy configuration enables business inactivity only by explicitly setting + `args.idle_timeout_s`. Its heartbeat and absolute lifetime bounds still apply. +- The Native SDK and server must upgrade together for the internal training-ID + handshake. Native and Tinker retain their existing initialization differences. + +Controller allocation/startup/shutdown/sweep bounds default to 3/1800/60/5 seconds +under `training`. Base startup and user initialization each receive a bounded +startup window; waiting for a base does not consume the user's initialization window. +The create future covers both phases without allowing late success. + +## Ownership and lifecycle + +`TrainingCatalog` performs pure configuration lookup. `TrainingController` is the +single lifecycle writer. `RayInstanceDriver` performs Ray/Serve effects, and +`ModelManagement` executes backend operations on the existing queue. + +The controller is a named actor owned by the launcher, not a detached service that +another launcher adopts. Initialization failure cleans and terminates that actor. +While its owner is alive, a controller restart starts a new epoch and invalidates +old handles. Runtime constructors register an incarnation once; Serve cannot +silently replace a process while preserving its old training state. + +The state lock covers decisions and commits. Model startup, user initialization, +cleanup, and resident rebuilding run outside it. A resident rebuilding for one +model/mode does not block another specification or the sweep loop. + +Every instance owns a named Ray Placement Group with one accelerator/CPU bundle +and `STRICT_PACK`. The allocator bounds `pg.ready()`, obtains actual device IDs +through a scheduled probe, and injects that allocation into workers without a +second reservation. Worker, application, registry, and PG names identify their +server namespace, controller epoch, and instance ID. + +A model has a stable CPU router. Creating/releasing handles enters the controller; +training requests route to their fixed instance. Internal model lifecycle HTTP +routes explicitly reject direct creation or unloading. Backend workers do not +restart or replay stateful calls automatically. + +| Record | States | +| --- | --- | +| Training | PENDING → ACTIVE → CLOSING → RELEASED / EXPIRED / FAILED / CANCELLED | +| Instance | ALLOCATING → STARTING → READY; RESETTING for resident full reset; DRAINING → STOPPED | +| Unconfirmed shutdown | QUARANTINED, with devices/quota still owned until death is confirmed | + +## Creation, isolation, and expiry + +Create identity is `(client type, token, session, create key)`: Native uses its UUID +and Tinker uses `model_seq_id`. Identical retries return the original future; +changed parameters are rejected. Pending users count toward quota and LoRA slots. +Compatible pending creates share a base startup. Full handles never share an instance. + +A failed user initialization closes only that user. Its slot remains reserved while +queued or late initialization may still execute. Cleanup follows initialization in +the serial lifecycle queue, and a late result cannot publish an ACTIVE handle. +Worker death, base startup failure, or unsafe backend execution can fail an entire +instance; a create RPC timeout alone is not evidence of that failure. + +Controller tickets protect accepted operations from expiry. The runtime exposes +queued/executing membership through the compute worker's public interface. Sweep +reconciles both uncommitted admission and committed operations whose completion +notification was lost. A terminal future alone does not establish that a background +call stopped. Reconciled activity uses the original completion time, not sweep time. + +Native and Tinker keep bounded sequence retry receipts per training stream. An +expired-sequence fence survives while the handle is live, so expiry cannot cause an +old optimizer operation to run again. Retained receipts tolerate out-of-order +requests; claims are released only when submission never created a future. Retiring +the handle removes its stream ledgers. Redis enumeration uses the existing SCAN API. + +Expiry closes inactive users only after their operations drain. Business inactivity +uses executed operations; heartbeat-only clients expire under the new policy. +Legacy heartbeat and maximum-lifetime checks retain their original meanings. + +The last user leaving closes a dynamic instance. Clean residents keep their base; +full residents reset weights, optimizer state, gradients, and template hooks before +another user is admitted. A resident modified by an instance-level patch is rebuilt +when empty. A cleanup failure cannot release a slot or kill another active LoRA user +merely because its cleanup RPC timed out. + +Shutdown kills owned workers, deletes the owned Serve app, removes the PG, confirms +worker/resource termination, and then retires quota. Failed confirmation preserves +ownership and retries in the background. Recovery never clears a shared registry or +shuts down another component's Ray cluster. Launcher shutdown retains its normal +Serve teardown behavior. + +## Acceptance + +| Case | Required observation | +| --- | --- | +| A01 | Startup/query/health on a lazy catalog loads no base or training PG. | +| A02 | Concurrent first creates coalesce into the required LoRA slot count. | +| A03 | Full LoRA slots expand with devices available; exhausted capacity fails terminally. | +| A04 | Both backends and SDKs isolate LoRA tenants and exclusive full handles. | +| A05 | Disjoint allocations are released and reused by another model. | +| A06 | Fragmented resources across nodes cannot satisfy one instance. | +| A07 | Heartbeats do not renew new-policy business inactivity; legacy semantics remain unchanged. | +| A08 | Idle LoRA cleanup preserves an active tenant's optimizer continuation. | +| A09 | Queued/running/save operations drain before expiry or release. | +| A10 | Unsafe background execution retains ownership until workers terminate. | +| A11 | Partial startup, cancellation, and allocation failure leave correct quota and ownership. | +| A12 | Worker/controller failure fences old handles and preserves unrelated instances. | +| A13 | Same-label recreation, old requests, and sequence retries cannot mutate another handle. | +| A14 | Explicit recreation restores saved weights and optimizer state. | +| A15 | Resident full reset leaves a clean base for the next user. | +| A16 | Ten cross-model reload cycles leave no accumulated workers, PGs, or device memory. | +| A17 | Invalid mode/config, expired IDs, and cross-user access fail without allocation. | +| A18 | Related SDK, YAML, checkpoint, and repository regressions pass together. | +| A19 | Late startup/user initialization cannot publish success or prematurely reuse a slot. | +| A20 | Memory and Redis recovery invalidate old state and rebuild quota correctly. | +| A21 | Checkpoint upload needs no base; patched residents rebuild before reuse. | + +Review regressions also cover nonblocking resident recovery, launcher initialization +rollback, lost completion notification, and sequence-receipt expiry. CPU/logical Ray +resources are not device training evidence. The revised code still requires both +backend device reruns, including an SP=true check for the Megatron LoRA forward path. +Raw device evidence is retained outside the public source tree. diff --git a/docs/design/validation/on-demand-training/README.md b/docs/design/validation/on-demand-training/README.md new file mode 100644 index 000000000..99a8bead2 --- /dev/null +++ b/docs/design/validation/on-demand-training/README.md @@ -0,0 +1,35 @@ +# On-demand training validation + +The review fixes passed local validation with CPU models, real Ray scheduling, +the Native/Tinker SDKs over HTTP, and isolated memory/Redis state: + +- Server, model, Transformers, and client CPU regression: 807 passed, no failures, + and 21 hardware-specific tests skipped. +- HTTP mock acceptance: passed for both SDKs in shared LoRA and exclusive full + modes, including release/recreation, runtime lifecycle rejection, controller + restart, state loss, and ten reuse cycles. +- Existing YAML startup: passed with memory state and with Redis state. +- The separate backend prerequisite (#315): 17 focused CPU regressions passed + without failures or skips in an isolated checkout. + +Repository pre-commit checks passed. Raw local outputs are retained outside the +repository. + +Real-device validation of the revised code is deferred at the user's request +because the PPU server is running other tests. PR #312 remains a draft. + +The previous runtime (`1bcf657e1534f6023b5de605b5a204c94705729b`) passed nine +acceptance cases on each of Transformers and Megatron, with no failures or skips. +Those results do not validate the revised lifecycle code. Raw logs, JUnit files, +server configurations, and numerical results are retained outside the repository. + +Standalone validation programs live in `tools/validation/on_demand_training/`. +The collected CPU tests under `tests/server/training/` run their allocation and +HTTP programs in isolated processes. Device validation is explicitly invoked: + +```bash +python -m pytest tools/validation/on_demand_training/device_validation.py -vv +``` + +It requires the backend, model paths, a separate evidence directory, and free +devices configured by `TWINKLE_DEVICE_*`. It must not run against occupied devices. diff --git a/docs/source_en/Usage Guide/Server and Client/Server.md b/docs/source_en/Usage Guide/Server and Client/Server.md index dd0064aba..4701f24f6 100644 --- a/docs/source_en/Usage Guide/Server and Client/Server.md +++ b/docs/source_en/Usage Guide/Server and Client/Server.md @@ -479,3 +479,9 @@ raises. During that window newly arriving tasks fail fast with a `server`/503 er Each persisted future stores its immutable `absolute_deadline` when it is created. Cleanup therefore reaches the same decision regardless of which deployment process holds the cleanup lease. + +## On-demand training and device reclamation + +Existing `applications` entries with `import_path: model` remain valid. Set `args.min_replicas: 0` for lazy startup and explicitly set `args.idle_timeout_s` for business inactivity. Legacy `adapter_config.adapter_timeout` remains a heartbeat timeout and `adapter_max_lifetime` retains its absolute lifetime. New `modes` default to 1800 seconds of business inactivity; heartbeats and polling do not renew that deadline; queued and running operations protect training state. Empty instances above the resident minimum are destroyed so other models can reuse their devices. Existing deployment min/max replicas become training instance minimums/limits. + +To offer both LoRA and full training for one model, configure `num_devices`, `device_mesh`, `min_replicas`, and `idle_timeout_s` separately under `args.modes.lora/full`. LoRA users share compatible instances; each full-training handle owns an instance. Sampling configuration stays the same. See the [specification](../../../design/on-demand-training.md) and `cookbook/client/server/on_demand_training.yaml`. diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" index 81fcacfe5..8125400ca 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" @@ -466,3 +466,9 @@ twinkle-server check-config -c server_config.yaml 任务记录进入终态后，per-replica 准入闸门额外保持关闭的最长时长为 `max(0, Collect_Width × R − T)`：此时记录已是终态，但泄漏的 executor 线程可能仍持有闸门，直到其 `ray.get` 返回或抛出。在该窗口内新到达的任务会以 `server`/503 错误快速失败。 每条持久化 future 在创建时写入不可变的 `absolute_deadline`，因此无论哪个 deployment 进程持有 cleanup lease，清理结果都由任务自身契约决定。 + +## 按需训练与设备回收 + +原有 `applications` 和 `import_path: model` 配置继续有效。单模式配置增加 `args.min_replicas: 0` 即可懒启动，`args.idle_timeout_s` 显式设置训练句柄的业务闲置时间。旧 `adapter_config.adapter_timeout` 仍按心跳判断会话存活，`adapter_max_lifetime` 保留绝对寿命；不再把旧字段转换成业务闲置时间。新增 `modes` 的业务闲置默认 1800 秒，心跳和轮询不延长该期限，运行中或已排队的任务保护训练状态。最后一个用户释放或过期后，超出保底数量的实例立即销毁，释放的卡可用于其他模型。已有 deployment 的 min/max_replicas 转成训练实例的保底/上限。 + +同一模型需要同时提供 LoRA 和全参数训练时，在 `args.modes.lora/full` 内分别设置 `num_devices`、`device_mesh`、`min_replicas` 和 `idle_timeout_s`；不用改成另一套顶层模型目录。LoRA 用户共享兼容实例，全参数用户各占一个实例。采样服务继续使用原配置。详见 [完整规范](../../../design/on-demand-training.md) 与仓库 `cookbook/client/server/on_demand_training.yaml`。 diff --git a/src/twinkle/infra/_ray/ray_helper.py b/src/twinkle/infra/_ray/ray_helper.py index 4cc5f6a66..d676f4661 100644 --- a/src/twinkle/infra/_ray/ray_helper.py +++ b/src/twinkle/infra/_ray/ray_helper.py @@ -71,13 +71,19 @@ def get_config(self, key: str): def clear(self): self.config.clear() + registry_name = os.environ.get('TWINKLE_RAY_REGISTRY_NAME', 'config_registry') + registry_namespace = os.environ.get('TWINKLE_RAY_NAMESPACE') try: - RayHelper._registry = ray.get_actor('config_registry') + RayHelper._registry = ray.get_actor(registry_name, namespace=registry_namespace) except ValueError: try: - RayHelper._registry = WorkerRegistry.options(name='config_registry', lifetime='detached').remote() + RayHelper._registry = WorkerRegistry.options( + name=registry_name, + namespace=registry_namespace, + lifetime='detached', + num_cpus=0.01 if registry_namespace else 1).remote() except ValueError: - RayHelper._registry = ray.get_actor('config_registry') + RayHelper._registry = ray.get_actor(registry_name, namespace=registry_namespace) assert RayHelper._registry is not None @staticmethod @@ -359,7 +365,7 @@ def create_workers(worker_cls: Type[T], 'scheduling_strategy': PlacementGroupSchedulingStrategy(placement_group=deploy_pg['placement_group']), 'name': worker_name, - 'namespace': 'default', + 'namespace': os.environ.get('TWINKLE_RAY_NAMESPACE', 'default'), 'runtime_env': runtime_env, 'num_cpus': 0.01, } @@ -383,6 +389,8 @@ def create_workers(worker_cls: Type[T], if max_concurrency is not None: worker_options['max_concurrency'] = max_concurrency + if os.environ.get('TWINKLE_RAY_NAMESPACE'): + worker_options.update(max_restarts=0, max_task_retries=0) worker = worker_cls.options(**worker_options).remote(*args, **kwargs) workers.append(worker) else: @@ -417,7 +425,7 @@ def create_workers(worker_cls: Type[T], 'scheduling_strategy': PlacementGroupSchedulingStrategy(placement_group=deploy_pg['placement_group']), 'name': worker_name, - 'namespace': 'default', + 'namespace': os.environ.get('TWINKLE_RAY_NAMESPACE', 'default'), 'runtime_env': runtime_env, 'num_cpus': 0.01, } @@ -425,6 +433,8 @@ def create_workers(worker_cls: Type[T], if max_concurrency is not None: worker_options['max_concurrency'] = max_concurrency + if os.environ.get('TWINKLE_RAY_NAMESPACE'): + worker_options.update(max_restarts=0, max_task_retries=0) worker = worker_cls.options(**worker_options).remote(*args, **kwargs) workers.append(worker) return workers diff --git a/src/twinkle/infra/_ray/resource_manager.py b/src/twinkle/infra/_ray/resource_manager.py index 9b8ba0cef..c9013c972 100644 --- a/src/twinkle/infra/_ray/resource_manager.py +++ b/src/twinkle/infra/_ray/resource_manager.py @@ -20,6 +20,35 @@ def _gpu_placement_group_cpus(node_cpu: int, nproc_per_node: int) -> int: class ResourceManager: + @classmethod + def from_allocation(cls, group: DeviceGroup, placement_group, devices: list[str]): + """Use an allocation owned by the training controller, without reserving again. + + Ranks are local to this allocation; devices are the actual Ray-selected IDs. + The controller owns removal, so this view's destroy method is a no-op. + """ + import ray + manager = cls.__new__(cls) + ranks = list(range(len(devices))) + group.ranks = ranks + manager.group_configs = [group] + manager.placement_groups = [placement_group] + manager.cpu_placement_groups = [] + manager._external_allocation = True + manager.device_groups = { + group.name: [{ + 'gpu_rank': [device], + 'placement_group': placement_group, + 'ray_address': str(ray.get_runtime_context().gcs_address) + } for device in devices] + } + if (group.device_type or '').upper() == 'CPU': + manager.device_groups[group.name] = [{ + 'placement_group': placement_group, + 'ray_address': str(ray.get_runtime_context().gcs_address) + } for _ in devices] + return manager + def __init__(self, nproc_per_node: int, ncpu_proc_per_node: int, groups: List[DeviceGroup]): # CPU placement group default strategy: # - Old approach: use node_cpu//4 as CPU bundle per node, even if only 1~2 CPU processes are needed, @@ -296,6 +325,8 @@ def get_group(self, group: str): return self.device_groups[group] def destroy_placement_group(self): + if getattr(self, '_external_allocation', False): + return import ray for pg in self.placement_groups: ray.util.remove_placement_group(pg) diff --git a/src/twinkle/model/megatron/megatron.py b/src/twinkle/model/megatron/megatron.py index e48243b15..2587e27b6 100644 --- a/src/twinkle/model/megatron/megatron.py +++ b/src/twinkle/model/megatron/megatron.py @@ -1007,6 +1007,9 @@ def reload_initial_weights(self, **kwargs): for _model in self.strategy.unwrap_model(self.model): for param in _model.parameters(): param.grad = None + if kwargs.get('reset_training_state', False): + self.optimizer_group = {_default_adapter_name: self._construct_default_optimizer_group()} + self.optimizer_group[_default_adapter_name].adapter_name = _default_adapter_name if dist.is_initialized(): dist.barrier() diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index e86b90745..2dba85f5e 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -1495,6 +1495,15 @@ def reload_initial_weights(self, **kwargs): for param in self.model.parameters(): param.grad = None + if kwargs.get('reset_training_state', False): + for optimizer_group in self.optimizer_group.values(): + self.unregister_mm_forward_hook(optimizer_group) + self.optimizer_group = {_default_adapter_name: self._construct_default_optimizer_group()} + self.optimizer_group[_default_adapter_name].adapter_name = _default_adapter_name + self.active_group = _default_adapter_name + if self._model_wrapped: + self.register_mm_forward_hook(self.optimizer_group[_default_adapter_name]) + def _load_optimizer(self, checkpoint_dir, **kwargs): adapter_name = kwargs.pop('adapter_name', _default_adapter_name) strict = kwargs.pop('strict', False) diff --git a/src/twinkle/protocol/types/model.py b/src/twinkle/protocol/types/model.py index f2e6bd976..1e267d066 100644 --- a/src/twinkle/protocol/types/model.py +++ b/src/twinkle/protocol/types/model.py @@ -28,7 +28,9 @@ class CreateRequest(StrictRequest): - """Body of ``POST /twinkle/create``: a session-establishing no-op.""" + """Body of ``POST /twinkle/create``: CPU-only capability handshake.""" + + training_handles: bool = False # --------------------------------------------------------------------------- # @@ -36,7 +38,13 @@ class CreateRequest(StrictRequest): # --------------------------------------------------------------------------- # -class AdapterRequest(StrictRequest): +class TrainingRequest(StrictRequest): + """A training handle is control metadata, never a backend kwarg.""" + + training_id: str | None = None + + +class AdapterRequest(TrainingRequest): """The shared shape of an adapter-scoped operation. ``seq_id`` is an idempotency key, not a backend parameter: the submit shell @@ -83,7 +91,7 @@ class ClipGradAndStepRequest(ClipGradNormRequest): optim_params: dict[str, JsonValue] | None = backend_kwarg(default=None) -class CalculateMetricRequest(StrictRequest): +class CalculateMetricRequest(TrainingRequest): """Body of ``POST /twinkle/calculate_metric``.""" adapter_name: str @@ -105,7 +113,7 @@ class CalculateMetricRequest(StrictRequest): # --------------------------------------------------------------------------- # -class _InlineForwardBase(StrictRequest): +class _InlineForwardBase(TrainingRequest): """Fields common to the three inline forward endpoints.""" inputs: WireInputBatch @@ -161,7 +169,7 @@ class ForwardBackwardTaskRequest(_InlineForwardBase): # --------------------------------------------------------------------------- # -class DataPlaneForwardRequest(StrictRequest): +class DataPlaneForwardRequest(TrainingRequest): """Body of the ``*_from_data_plane`` forward endpoints.""" input_refs: list[DataRef] = Field(min_length=1) @@ -180,7 +188,7 @@ class DataPlaneForwardRequest(StrictRequest): loss_kwargs: dict[str, JsonValue] = passthrough() -class DataPlaneForwardOnlyRequest(StrictRequest): +class DataPlaneForwardOnlyRequest(TrainingRequest): """Body of ``POST /twinkle/forward_only_from_data_plane``. This endpoint is read-only, so it deliberately has no ``seq_id`` idempotency @@ -222,25 +230,25 @@ def validate_output(self) -> DataPlaneForwardOnlyRequest: # --------------------------------------------------------------------------- # -class SetLossRequest(StrictRequest): +class SetLossRequest(TrainingRequest): loss_cls: str adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() -class SetOptimizerRequest(StrictRequest): +class SetOptimizerRequest(TrainingRequest): optimizer_cls: str adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() -class SetLrSchedulerRequest(StrictRequest): +class SetLrSchedulerRequest(TrainingRequest): scheduler_cls: str adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() -class SetTemplateRequest(StrictRequest): +class SetTemplateRequest(TrainingRequest): """Body of ``POST /twinkle/set_template``. No top-level ``model_id``: the backend always overrides it with its own @@ -254,20 +262,20 @@ class SetTemplateRequest(StrictRequest): init_kwargs: dict[str, JsonValue] = passthrough() -class SetProcessorRequest(StrictRequest): +class SetProcessorRequest(TrainingRequest): processor_cls: str adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() -class AddMetricRequest(StrictRequest): +class AddMetricRequest(TrainingRequest): metric_cls: str adapter_name: str is_training: bool | None = None init_kwargs: dict[str, JsonValue] = passthrough() -class ApplyPatchRequest(StrictRequest): +class ApplyPatchRequest(TrainingRequest): patch_cls: str adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() @@ -278,7 +286,7 @@ class ApplyPatchRequest(StrictRequest): # --------------------------------------------------------------------------- # -class SaveRequest(StrictRequest): +class SaveRequest(TrainingRequest): adapter_name: str name: str | None = None save_optimizer: bool = False @@ -287,7 +295,7 @@ class SaveRequest(StrictRequest): merge_lora: bool | None = backend_only('megatron', default=None) -class LoadRequest(StrictRequest): +class LoadRequest(TrainingRequest): adapter_name: str name: str load_optimizer: bool = False @@ -296,7 +304,7 @@ class LoadRequest(StrictRequest): strict: bool | None = backend_only('transformers', default=None) -class ResumeFromCheckpointRequest(StrictRequest): +class ResumeFromCheckpointRequest(TrainingRequest): """Body of ``POST /twinkle/resume_from_checkpoint``.""" name: str @@ -305,6 +313,7 @@ class ResumeFromCheckpointRequest(StrictRequest): class AddAdapterRequest(StrictRequest): + client_create_id: str | None = None adapter_name: str # ``config`` is None for full-parameter training (no LoRA adapter) and a # serialized LoraConfig string for LoRA training. @@ -418,6 +427,12 @@ class TrainingProgressResponse(ResponseModel): # --- Other responses --- +class TrainingHandleResponse(ResponseModel): + training_id: str + adapter_name: str + + class CreateResponse(ResponseModel): """Response for /create endpoint.""" status: str = 'ok' + training_handles: bool = False diff --git a/src/twinkle/server/config/application_spec.py b/src/twinkle/server/config/application_spec.py index de5c8bc3a..ea2094f62 100644 --- a/src/twinkle/server/config/application_spec.py +++ b/src/twinkle/server/config/application_spec.py @@ -11,12 +11,13 @@ """ from __future__ import annotations -import math import os -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic import BaseModel, ConfigDict, Field, model_validator from typing import Any, Literal from twinkle.server.task_queue.config import TaskQueueConfig +from .hybrid import SpectralHybridArgs +from .training import TrainingModes # Env var keys the launcher sets from the gateway ``server_config`` so that any # Ray worker (model / sampler / processor), not just the gateway, applies the @@ -55,21 +56,6 @@ class HttpOptions(BaseModel): # ---------- per-deployment args schemas ------------------------------------ # -class SpectralHybridArgs(_ArgsBase): - """Strict server-owned Spectral Hybrid allocation and optimizer defaults.""" - - allocation_path: str = Field(min_length=1) - default_lr_lora: float = Field(default=2.5e-5, gt=0) - default_lr_fft: float = Field(default=1.0e-6, gt=0) - - @field_validator('default_lr_lora', 'default_lr_fft') - @classmethod - def _finite_learning_rate(cls, value: float) -> float: - if not math.isfinite(value): - raise ValueError('learning rate must be finite') - return value - - class ModelArgs(_ArgsBase): """Args for the ``model`` deployment. @@ -83,6 +69,12 @@ class ModelArgs(_ArgsBase): device_group: dict[str, Any] device_mesh: dict[str, Any] backend: Literal['mock', 'transformers', 'megatron'] + device_type: Literal['cuda', 'npu', 'mps', 'cpu'] | None = None + env_vars: dict[str, str] = Field(default_factory=dict) + modes: TrainingModes | None = None + min_replicas: int | None = Field(default=None, ge=0, strict=True) + max_replicas: int | None = Field(default=None, gt=0, strict=True) + idle_timeout_s: float | None = Field(default=None, gt=0, allow_inf_nan=False) train_mode: Literal['lora', 'hybrid', 'full'] = 'lora' strategy: Literal['accelerate', 'native_fsdp'] | None = None adapter_config: dict[str, Any] | None = None @@ -94,8 +86,18 @@ class ModelArgs(_ArgsBase): hybrid: SpectralHybridArgs | None = None data_plane_url: str | None = None + @model_validator(mode='before') + @classmethod + def _mode_resources(cls, data): + if isinstance(data, dict) and data.get('modes') is not None: + return {'device_group': {}, 'device_mesh': {}, **data} + return data + @model_validator(mode='after') def _validate_spectral_backend(self): + if self.modes is not None and any( + mode.idle_timeout_s is None for mode in (self.modes.lora, self.modes.full) if mode is not None): + raise ValueError('modes require a finite positive idle_timeout_s') if self.strategy is not None and self.backend != 'transformers': raise ValueError('strategy is only supported by the transformers backend') if self.hybrid is not None: @@ -269,23 +271,3 @@ def _coerce_args_to_schema(cls, data: Any) -> Any: # ``schema.model_validate`` rejects a non-dict itself with a clean # error, so no separate non-dict guard is needed here. return {**data, 'args': schema.model_validate(raw_args)} - - @model_validator(mode='after') - def _validate_full_mode_single_replica(self) -> ApplicationSpec: - """``train_mode: full`` requires exactly one replica. - - The exclusive-tenant lock lives in per-replica memory - (``ModelManagement._resource_records``), so more than one replica would - silently allow one full-parameter tenant per replica, each rewriting - its own copy of the base weights. Reject that at config-load time. - """ - if not (isinstance(self.args, ModelArgs) and self.args.train_mode == 'full'): - return self - for dep in self.deployments: - num_replicas = dep.get('num_replicas') - max_replicas = (dep.get('autoscaling_config') or {}).get('max_replicas') - if (num_replicas or 1) > 1 or (max_replicas or 1) > 1: - raise ValueError(f"Application '{self.name}': train_mode='full' is an exclusive single-tenant " - 'mode and requires a single replica; set num_replicas/autoscaling_config.' - 'max_replicas to 1.') - return self diff --git a/src/twinkle/server/config/hybrid.py b/src/twinkle/server/config/hybrid.py new file mode 100644 index 000000000..1456a8913 --- /dev/null +++ b/src/twinkle/server/config/hybrid.py @@ -0,0 +1,21 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shared Spectral Hybrid configuration.""" +import math +from pydantic import BaseModel, ConfigDict, Field, field_validator + + +class SpectralHybridArgs(BaseModel): + """Strict server-owned Spectral Hybrid allocation and optimizer defaults.""" + + model_config = ConfigDict(extra='forbid') + + allocation_path: str = Field(min_length=1) + default_lr_lora: float = Field(default=2.5e-5, gt=0) + default_lr_fft: float = Field(default=1.0e-6, gt=0) + + @field_validator('default_lr_lora', 'default_lr_fft') + @classmethod + def _finite_learning_rate(cls, value: float) -> float: + if not math.isfinite(value): + raise ValueError('learning rate must be finite') + return value diff --git a/src/twinkle/server/config/server_config.py b/src/twinkle/server/config/server_config.py index 82bac4d44..5ac9aad36 100644 --- a/src/twinkle/server/config/server_config.py +++ b/src/twinkle/server/config/server_config.py @@ -14,6 +14,7 @@ """ from __future__ import annotations +from functools import cached_property from pathlib import Path from pydantic import BaseModel, ConfigDict, Field, model_validator from typing import Any @@ -23,6 +24,8 @@ from .application_spec import ApplicationSpec, HttpOptions from .persistence import PersistenceConfig from .telemetry import TelemetryConfig +from .training import TrainingConfig, TrainingOptions +from .training_normalization import normalize_model_application class ServerConfig(BaseModel): @@ -36,6 +39,7 @@ class ServerConfig(BaseModel): telemetry: TelemetryConfig = Field(default_factory=TelemetryConfig) persistence: PersistenceConfig = Field(default_factory=PersistenceConfig) task_queue: TaskQueueConfig = Field(default_factory=TaskQueueConfig) + training: TrainingOptions = Field(default_factory=TrainingOptions) applications: list[ApplicationSpec] = Field(default_factory=list) # ---- loading ---------------------------------------------------------- # @@ -71,8 +75,24 @@ def from_yaml(cls, path: str | Path) -> ServerConfig: def _validate_cross_field(self) -> ServerConfig: if self.persistence.mode == 'redis' and not self.persistence.redis_url: raise ValueError("persistence.redis_url is required when persistence.mode == 'redis'", ) + if self.training_config.models: + static = [app for app in self.applications if app.import_path != 'model'] + names = {app.name for app in static} + for model in self.training_config.models: + if model.name in names: + raise ValueError(f'training model name conflicts with application: {model.name}') + path = model.route_prefix + if any(app.route_prefix.rstrip('/') == path for app in static): + raise ValueError(f'training route conflicts with application: {path}') return self + @cached_property + def training_config(self) -> TrainingConfig: + """Normalize old and extended model args into one runtime catalog.""" + return TrainingConfig( + **self.training.model_dump(), + models=[normalize_model_application(app) for app in self.applications if app.import_path == 'model']) + # ---- round-trip / serialization -------------------------------------- # def to_yaml_dict(self) -> dict[str, Any]: diff --git a/src/twinkle/server/config/training.py b/src/twinkle/server/config/training.py new file mode 100644 index 000000000..e68929c62 --- /dev/null +++ b/src/twinkle/server/config/training.py @@ -0,0 +1,127 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Administrator-owned training catalog. Parsing never allocates devices.""" +from __future__ import annotations + +from pydantic import BaseModel, ConfigDict, Field, model_validator +from typing import Literal + +from twinkle.server.task_queue.config import TaskQueueConfig +from .hybrid import SpectralHybridArgs + + +class _TrainingConfig(BaseModel): + model_config = ConfigDict(extra='forbid', frozen=True) + + +class TrainingMesh(_TrainingConfig): + """Sizes accepted by DeviceMesh.from_sizes, including overlapping EP axes.""" + + device_type: str | None = None + dp_size: int = Field(default=1, gt=0, strict=True) + fsdp_size: int | None = Field(default=None, gt=0, strict=True) + tp_size: int | None = Field(default=None, gt=0, strict=True) + pp_size: int | None = Field(default=None, gt=0, strict=True) + cp_size: int | None = Field(default=None, gt=0, strict=True) + ep_size: int | None = Field(default=None, gt=0, strict=True) + etp_size: int = Field(default=1, gt=0, strict=True) + ep_fsdp_size: int | None = Field(default=None, gt=0, strict=True) + vpp_size: int | None = Field(default=None, gt=0, strict=True) + ulysses_size: int | None = Field(default=None, gt=0, strict=True) + sequence_parallel: bool = False + + def world_size(self) -> int: + from twinkle.utils.device_mesh import DeviceMesh + return DeviceMesh.from_sizes(**self.model_dump(exclude_none=True)).world_size + + +class TrainingModeConfig(_TrainingConfig): + num_devices: int = Field(gt=0, strict=True) + device_mesh: TrainingMesh + min_replicas: int = Field(default=0, ge=0, strict=True) + max_replicas: int | None = Field(default=None, gt=0, strict=True) + idle_timeout_s: float | None = Field(default=1800, gt=0, allow_inf_nan=False) + strategy: Literal['accelerate', 'native_fsdp'] | None = None + queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) + max_length: int | None = Field(default=None, gt=0, strict=True) + + @model_validator(mode='after') + def _validate_resources(self): + if self.max_replicas is not None and self.min_replicas > self.max_replicas: + raise ValueError('min_replicas must not exceed max_replicas') + if self.device_mesh.world_size() != self.num_devices: + raise ValueError('device_mesh world_size must equal num_devices') + return self + + +class LoraTrainingConfig(TrainingModeConfig): + max_loras: int = Field(default=5, gt=0, strict=True) + max_r: int = Field(default=32, gt=0, strict=True) + target_modules: str | list[str] = 'all-linear' + hybrid: SpectralHybridArgs | None = None + + +class TrainingModes(_TrainingConfig): + lora: LoraTrainingConfig | None = None + full: TrainingModeConfig | None = None + + @model_validator(mode='after') + def _not_empty(self): + if self.lora is None and self.full is None: + raise ValueError('at least one training mode is required') + return self + + +class LegacyExpiration(_TrainingConfig): + """Internal compatibility policy derived from existing adapter_config fields.""" + session_timeout_s: float | None = Field(default=1800, gt=0, allow_inf_nan=False) + max_lifetime_s: float | None = Field(default=36000, gt=0, allow_inf_nan=False) + + +class TrainingModelConfig(_TrainingConfig): + name: str = Field(min_length=1) + model_id: str = Field(min_length=1) + backend: Literal['mock', 'transformers', 'megatron'] + device_type: Literal['cuda', 'npu', 'mps', 'cpu'] + modes: TrainingModes + legacy_expiration: LegacyExpiration | None = None + route_prefix: str | None = None + data_plane_url: str | None = None + env_vars: dict[str, str] = Field(default_factory=dict) + + @model_validator(mode='after') + def _validate_backend(self): + if (self.name.startswith('/') or any(part in ('', '.', '..') for part in self.name.split('/')) + or any(char.isspace() or char in '?#%\\' for char in self.name)): + raise ValueError('name must be a nonempty relative model path without dot segments') + if self.device_type == 'cpu' and self.backend != 'mock': + raise ValueError('cpu is supported only by the mock training backend') + for mode in (self.modes.lora, self.modes.full): + if mode is None: + continue + if mode.device_mesh.device_type not in (None, self.device_type): + raise ValueError('device_mesh.device_type must match device_type') + if mode.strategy is not None and self.backend != 'transformers': + raise ValueError('strategy requires the transformers backend') + if self.modes.lora and self.modes.lora.hybrid and self.backend != 'transformers': + raise ValueError('hybrid requires the transformers backend') + return self + + +class TrainingOptions(_TrainingConfig): + allocation_timeout_s: float = Field(default=3, gt=0, allow_inf_nan=False) + startup_timeout_s: float = Field(default=1800, gt=0, allow_inf_nan=False) + shutdown_timeout_s: float = Field(default=60, gt=0, allow_inf_nan=False) + sweep_interval_s: float = Field(default=5, gt=0, allow_inf_nan=False) + + +class TrainingConfig(TrainingOptions): + """Internal normalized catalog; public model entries stay in applications.""" + + models: tuple[TrainingModelConfig, ...] = () + + @model_validator(mode='after') + def _unique_names(self): + names = [model.name for model in self.models] + if len(names) != len(set(names)): + raise ValueError('training model names must be unique') + return self diff --git a/src/twinkle/server/config/training_normalization.py b/src/twinkle/server/config/training_normalization.py new file mode 100644 index 000000000..da7487cab --- /dev/null +++ b/src/twinkle/server/config/training_normalization.py @@ -0,0 +1,87 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Translate model application configuration without creating a second runtime.""" +from twinkle.hub.model_alias import extract_model_alias +from .application_spec import ApplicationSpec, ModelArgs +from .training import LegacyExpiration, TrainingModelConfig, TrainingModes + + +def normalize_model_application(app: ApplicationSpec) -> TrainingModelConfig: + args = app.args + assert isinstance(args, ModelArgs) + if args.modes is not None: + baseline = ModelArgs(model_id=args.model_id, backend=args.backend, device_group={}, device_mesh={}) + common = { + 'model_id', 'backend', 'modes', 'data_plane_url', 'device_group', 'device_mesh', 'device_type', 'env_vars' + } + legacy_fields = { + key + for key in type(args).model_fields if key not in common and getattr(args, key) != getattr(baseline, key) + } + if legacy_fields or args.device_group or args.device_mesh: + raise ValueError(f'{app.name}: modes cannot be mixed with single-mode training parameters') + if app.deployments: + raise ValueError(f'{app.name}: configure instance replicas inside modes, not deployments') + modes = args.modes + legacy_expiration = None + device_types = {mode.device_mesh.device_type for mode in (modes.lora, modes.full) if mode is not None} + device_types.discard(None) + if len(device_types) > 1: + raise ValueError(f'{app.name}: modes must use the same device type') + device_type = args.device_type or next(iter(device_types), 'cpu' if args.backend == 'mock' else 'cuda').lower() + env_vars = args.env_vars + else: + if len(app.deployments) > 1: + raise ValueError(f'{app.name}: a model application must have one deployment configuration') + deployment = app.deployments[0] if app.deployments else {} + scaling = deployment.get('autoscaling_config') or {} + ranks = args.device_group.get('ranks', args.nproc_per_node) + if isinstance(ranks, list): + raise ValueError(f'{app.name}: use a device count instead of fixed ranks for pooled training') + device_type = args.device_type or str( + args.device_group.get('device_type', + args.device_mesh.get('device_type', + 'cpu' if args.backend == 'mock' else 'cuda'))).lower() + mesh = {**args.device_mesh, 'device_type': device_type} + mode = { + 'num_devices': + ranks, + 'device_mesh': + mesh, + 'min_replicas': + args.min_replicas if args.min_replicas is not None else scaling.get('min_replicas', + deployment.get('num_replicas', 0)), + 'max_replicas': + args.max_replicas if args.max_replicas is not None else scaling.get('max_replicas', + deployment.get('num_replicas')), + 'idle_timeout_s': + args.idle_timeout_s, + 'strategy': + args.strategy, + 'queue_config': + args.queue_config, + 'max_length': + args.max_length, + } + mode_name = 'full' if args.train_mode == 'full' else 'lora' + if mode_name == 'lora': + mode.update( + max_loras=args.max_loras, max_r=args.max_r, target_modules=args.target_modules, hybrid=args.hybrid) + modes = TrainingModes.model_validate({mode_name: mode}) + legacy_expiration = LegacyExpiration( + session_timeout_s=(args.adapter_config or {}).get('adapter_timeout', 1800), + max_lifetime_s=(args.adapter_config or {}).get('adapter_max_lifetime', 36000)) + env_vars = { + **(deployment.get('ray_actor_options') or {}).get('runtime_env', {}).get('env_vars', {}), + **args.env_vars + } + name = extract_model_alias(app.route_prefix) or app.name + return TrainingModelConfig( + name=name, + model_id=args.model_id, + backend=args.backend, + device_type=device_type, + modes=modes, + legacy_expiration=legacy_expiration, + route_prefix=app.route_prefix.rstrip('/') or '/', + data_plane_url=args.data_plane_url, + env_vars=env_vars) diff --git a/src/twinkle/server/gateway/app.py b/src/twinkle/server/gateway/app.py index 0c54f506b..0f630055c 100644 --- a/src/twinkle/server/gateway/app.py +++ b/src/twinkle/server/gateway/app.py @@ -41,6 +41,8 @@ def __init__(self, self.http_options = http_options or {} self.proxy = ServiceProxy(http_options=http_options, route_prefix=self.route_prefix) self.supported_models = self._normalize_models(supported_models) + self.training_models = self._normalize_models(kwargs.get('training_models', supported_models)) + self._training_model_names = frozenset(m.model_name for m in self.training_models) self._supported_model_names = frozenset(m.model_name for m in self.supported_models) self._modelscope_config_lock = asyncio.Lock() self._state_cleanup_started = False @@ -89,11 +91,11 @@ def _normalize_models(self, supported_models): return normalized def _validate_base_model(self, base_model: str) -> None: - if base_model not in self._supported_model_names: + if base_model not in self._training_model_names: raise HTTPException( status_code=400, detail=f"Base model '{base_model}' is not supported. " - f"Supported models: {', '.join(sorted(self._supported_model_names))}") + f"Supported models: {', '.join(sorted(self._training_model_names))}") async def _get_base_model(self, model_id: str) -> str: metadata = await self.state.get_model_metadata(model_id) diff --git a/src/twinkle/server/gateway/tinker_handlers.py b/src/twinkle/server/gateway/tinker_handlers.py index 93809d780..74f81f32c 100644 --- a/src/twinkle/server/gateway/tinker_handlers.py +++ b/src/twinkle/server/gateway/tinker_handlers.py @@ -29,6 +29,11 @@ # Keys must equal ``state.models.FAILURE_REASON_CODES`` (guarded by test_envelope). _TINKER_FAILURE_WIRE: dict[str, tuple[int, str]] = { + 'capacity_unavailable': (409, 'user'), + 'training_expired': (404, 'user'), + 'startup_failed': (500, 'server'), + 'startup_timeout': (504, 'server'), + 'protocol_mismatch': (400, 'user'), 'invalid_request': (400, 'user'), 'request_rejected': (400, 'user'), 'resource_not_found': (404, 'user'), @@ -83,7 +88,7 @@ async def get_server_capabilities( self: GatewayServer = Depends(self_fn), ) -> types.GetServerCapabilitiesResponse: # Convert twinkle.protocol.types.SupportedModel to tinker.types.SupportedModel - tinker_supported_models = [types.SupportedModel(model_name=m.model_name) for m in self.supported_models] + tinker_supported_models = [types.SupportedModel(model_name=m.model_name) for m in self.training_models] return types.GetServerCapabilitiesResponse(supported_models=tinker_supported_models) @app.post('/telemetry') @@ -121,6 +126,7 @@ async def retrieve_future(request: Request, self: GatewayServer = Depends(self_fn)) -> Any: """Retrieve the result of an async task with long polling.""" request_id = body.request_id + await self.state.assert_future_owner(request_id, get_token_from_request(request)) outcome = await poll_future(self.state, request_id) record = outcome.record if outcome.timed_out: diff --git a/src/twinkle/server/gateway/twinkle_handlers.py b/src/twinkle/server/gateway/twinkle_handlers.py index 3df44ffc8..91962c8d6 100644 --- a/src/twinkle/server/gateway/twinkle_handlers.py +++ b/src/twinkle/server/gateway/twinkle_handlers.py @@ -57,11 +57,14 @@ async def healthz_deep( results = {} all_healthy = True - for model in self.supported_models: + for model in self.training_models: model_name = model.model_name try: resp = await self.proxy.proxy_request(request, 'healthz', model_name, 'model') healthy = (resp.status_code == 200) + if healthy: + import json + healthy = json.loads(resp.body).get('healthy', True) if not healthy: all_healthy = False results[model_name] = { @@ -86,7 +89,7 @@ async def get_server_capabilities( self: GatewayServer = Depends(self_fn), ) -> types.GetServerCapabilitiesResponse: return types.GetServerCapabilitiesResponse( - supported_models=self.supported_models, + supported_models=self.training_models, protocol_version=1, features=types.ClientFeatures( task_envelope=True, @@ -139,6 +142,7 @@ async def retrieve_future( finish. See ``poll_config`` for the numbers. """ request_id = body.request_id + await self.state.assert_future_owner(request_id, get_token_from_request(request)) outcome = await poll_future(self.state, request_id) if outcome.record is None: raise ResourceNotFoundError(f'request_id {request_id} not found or expired') @@ -156,6 +160,7 @@ async def cancel_future( running or already-terminal task is reported but never interrupted, so cancel can never corrupt in-flight GPU/optimizer state. """ + await self.state.assert_future_owner(body.request_id, get_token_from_request(request)) result = await self.state.cancel_future(body.request_id) return types.CancelResponse(**result) diff --git a/src/twinkle/server/launcher/env_propagation.py b/src/twinkle/server/launcher/env_propagation.py index c478444d9..61941669c 100644 --- a/src/twinkle/server/launcher/env_propagation.py +++ b/src/twinkle/server/launcher/env_propagation.py @@ -12,6 +12,7 @@ # Telemetry env var keys that need to be propagated to Ray worker processes. TELEMETRY_ENV_KEYS: tuple[str, ...] = ( + 'TWINKLE_STATE_RAY_NAMESPACE', 'TWINKLE_TELEMETRY_ENABLED', 'TWINKLE_TELEMETRY_DEBUG', 'TWINKLE_TELEMETRY_SERVICE', diff --git a/src/twinkle/server/launcher/server_launcher.py b/src/twinkle/server/launcher/server_launcher.py index 75fd0bd52..47921d70b 100644 --- a/src/twinkle/server/launcher/server_launcher.py +++ b/src/twinkle/server/launcher/server_launcher.py @@ -17,7 +17,7 @@ from typing import Any from twinkle import get_logger -from twinkle.hub.model_alias import MODEL_ID_ALIASES_ENV, build_model_alias_map +from twinkle.hub.model_alias import MODEL_ID_ALIASES_ENV, build_model_alias_map, extract_model_alias from twinkle.patch.ray_serve import apply_ray_serve_patches, get_runtime_env_for_patches from twinkle.server.config import ServerConfig from twinkle.server.config.application_spec import ApplicationSpec @@ -146,6 +146,10 @@ def _deploy_application(self, app_spec: ApplicationSpec) -> None: name = app_spec.name route_prefix = app_spec.route_prefix import_path = app_spec.import_path + if import_path == 'model': + # Model applications are specifications. _deploy_training creates + # their CPU routes and the controller owns all device instances. + return # Shallow-dump the typed args to a kwargs dict WITHOUT recursing into # nested models: ``dict(model)`` yields top-level (field, value) pairs # with nested models left as instances, so ``queue_config`` stays a @@ -187,6 +191,12 @@ def _deploy_application(self, app_spec: ApplicationSpec) -> None: # Pass http_options to server apps for internal proxy routing if import_path == 'server': args.setdefault('http_options', self.config.http_options.model_dump()) + if self.config.training_config.models: + args['training_models'] = [model.name for model in self.config.training_config.models] + sampler_apps = [app for app in self.config.applications if app.import_path == 'sampler'] + args['supported_models'] = list( + dict.fromkeys( + alias for app in sampler_apps if (alias := extract_model_alias(app.route_prefix, 'sampler')))) app = builder(deploy_options=deploy_options, **args) @@ -202,6 +212,9 @@ def launch(self) -> None: """ # Apply Ray Serve patches before initializing Ray apply_ray_serve_patches() + if self.config.training_config.models: + os.environ['TWINKLE_STATE_RAY_NAMESPACE'] = ( + self.ray_namespace or self.config.ray_namespace or 'twinkle_cluster') # Initialize telemetry if configured telemetry = self.config.telemetry @@ -241,6 +254,7 @@ def launch(self) -> None: logger.info(f'ServerState policy exported to worker env: {server_state_env}') model_alias_map = build_model_alias_map(self.config.applications) + model_alias_map.update({model.name: model.model_id for model in self.config.training_config.models}) if model_alias_map: os.environ[MODEL_ID_ALIASES_ENV] = json.dumps(model_alias_map, ensure_ascii=False) else: @@ -250,12 +264,16 @@ def launch(self) -> None: self._start_serve() applications = self.config.applications - if not applications: + if not applications and not self.config.training_config.models: logger.warning('No applications configured') return for app_spec in applications: - self._deploy_application(app_spec) + if app_spec.import_path != 'model': + self._deploy_application(app_spec) + + if self.config.training_config.models: + self._deploy_training() host = self.config.http_options.host port = self.config.http_options.port @@ -287,6 +305,10 @@ def _shutdown(self) -> None: telemetry-shutdown failure cannot mask the user-facing shutdown path. """ from ray import serve + if getattr(self, '_training_controller', None) is not None: + import ray + ray.get(self._training_controller.shutdown.remote()) + ray.kill(self._training_controller, no_restart=True) try: serve.shutdown() logger.info('Ray Serve shut down successfully') @@ -298,6 +320,49 @@ def _shutdown(self) -> None: from twinkle.server.telemetry import flush_telemetry_safely flush_telemetry_safely() + def _deploy_training(self) -> None: + import ray + from ray import serve + + from twinkle.server.training.actor import TrainingControllerActor, controller_name + from twinkle.server.training.router import build_training_router + + namespace = self.ray_namespace or self.config.ray_namespace or 'twinkle_cluster' + env = build_propagated_env_vars() + for key in ('PYTHONPATH', 'TWINKLE_DEFAULT_SAVE_DIR', 'TMPDIR', 'OMP_NUM_THREADS'): + if key in os.environ: + env.setdefault(key, os.environ[key]) + controller = TrainingControllerActor.options( + name=controller_name(namespace), namespace=namespace, runtime_env={ + 'env_vars': env + }).remote(self.config.training_config.model_dump(), namespace) + try: + ray.get(controller.initialize.remote()) + except Exception: + try: + ray.get(controller.shutdown.remote()) + except Exception: + logger.exception('Training controller cleanup failed during initialization rollback') + finally: + ray.kill(controller, no_restart=True) + raise + self._training_controller = controller + for model in self.config.training_config.models: + app = build_training_router(self.config.training_config, model, namespace, + self.config.http_options.model_dump(), { + 'num_replicas': 1, + 'ray_actor_options': { + 'num_cpus': 0.1, + 'runtime_env': { + 'env_vars': env + } + } + }) + serve.run( + app, + name=f'training-router-{namespace}-{model.name}', + route_prefix=model.route_prefix or f'/api/v1/model/{model.name}') + def launch_server( config: ServerConfig | None = None, diff --git a/src/twinkle/server/lifecycle/envelope.py b/src/twinkle/server/lifecycle/envelope.py index 25343c9b0..d88fff666 100644 --- a/src/twinkle/server/lifecycle/envelope.py +++ b/src/twinkle/server/lifecycle/envelope.py @@ -17,6 +17,11 @@ # Keys must equal ``state.models.FAILURE_REASON_CODES`` (guarded by test_envelope). _FAILURE_WIRE: dict[str, tuple[int, ErrorCategory]] = { + 'capacity_unavailable': (409, ErrorCategory.User), + 'training_expired': (404, ErrorCategory.User), + 'startup_failed': (500, ErrorCategory.Server), + 'startup_timeout': (504, ErrorCategory.Server), + 'protocol_mismatch': (400, ErrorCategory.User), 'invalid_request': (400, ErrorCategory.User), 'request_rejected': (400, ErrorCategory.User), 'resource_not_found': (404, ErrorCategory.User), diff --git a/src/twinkle/server/lifecycle/submit.py b/src/twinkle/server/lifecycle/submit.py index b1b409a47..b4da7ecb5 100644 --- a/src/twinkle/server/lifecycle/submit.py +++ b/src/twinkle/server/lifecycle/submit.py @@ -168,7 +168,7 @@ async def run_submit( not implemented by every backend. """ token = await self._on_request_start(request) - adapter_name = resolve_twinkle_adapter_name(request, body.adapter_name) + adapter_name = getattr(body, 'training_id', None) or resolve_twinkle_adapter_name(request, body.adapter_name) # ---- Preflight: decidable from the body plus this deployment's backend, so it # runs before any state write and before the enqueue. ---- @@ -197,9 +197,15 @@ async def _task(): dedup_key = None if seq_id is not None: session_id = get_session_id_from_request(request) or request.state.request_id - dedup_key = f'seq::{session_id}::{adapter_name or "-"}::{seq_id}' - ttl = int(self.task_queue_config.effective_execution_timeout) + 60 - prior_request_id = await self.state.claim_seq(dedup_key, request_id, ttl) + dedup_key = f'seq::{session_id}::{adapter_name or "-"}::{task_type}::{seq_id}' + training_id = getattr(body, 'training_id', None) + if training_id is not None: + stream = f'native::{session_id}::{task_type}' + retention = self.sequence_retention_s + prior_request_id = await self.state.claim_training_seq(training_id, stream, seq_id, request_id, retention) + else: + ttl = int(self.task_queue_config.effective_execution_timeout) + 60 + prior_request_id = await self.state.claim_seq(dedup_key, request_id, ttl) if prior_request_id is not None: return await self._peek_terminal(prior_request_id, fallback_status='pending') @@ -216,5 +222,8 @@ async def _task(): # there would let a retry enqueue a duplicate -> the exact double-apply this # dedup prevents. When unsure (record exists), keep the claim. if dedup_key is not None and await self.state.get_future(request_id) is None: - await self.state.release_seq(dedup_key) + if training_id is not None: + await self.state.release_training_seq(training_id, stream, seq_id, request_id) + else: + await self.state.release_seq(dedup_key) raise diff --git a/src/twinkle/server/model/app.py b/src/twinkle/server/model/app.py index 1b71add8e..d0049fee4 100644 --- a/src/twinkle/server/model/app.py +++ b/src/twinkle/server/model/app.py @@ -8,21 +8,18 @@ from __future__ import annotations import asyncio -from fastapi import FastAPI, HTTPException, Request +from fastapi import FastAPI, Request from ray import serve -from ray.serve.config import RequestRouterConfig from typing import Any from twinkle import DeviceGroup from twinkle.server.config.backend_dispatch import BackendSelector -from twinkle.server.deployment import LazyCleanupMixin, bind_deployment, build_deployment_app -from twinkle.server.exceptions import FullModeBusyError -from twinkle.server.middleware.auth import get_token_from_request -from twinkle.server.model.routing import StickyLoraRequestRouter +from twinkle.server.deployment import bind_deployment, build_deployment_app from twinkle.server.runtime import init_twinkle_runtime from twinkle.server.session_resource import AdapterManagerMixin from twinkle.server.state import ServerState, get_server_state from twinkle.server.task_queue import TaskQueueConfig, TaskQueueMixin +from twinkle.server.training.runtime import ManagedTrainingRuntime from twinkle.server.utils import wrap_builder_with_device_group_env from twinkle.utils.logger import get_logger from .tinker_handlers import _register_model_tinker_routes @@ -82,7 +79,7 @@ def _make_megatron_model(kw: dict[str, Any]) -> Any: ) -class ModelManagement(LazyCleanupMixin, TaskQueueMixin, AdapterManagerMixin): +class ModelManagement(ManagedTrainingRuntime, TaskQueueMixin, AdapterManagerMixin): """Unified model management service. Handles: @@ -105,13 +102,16 @@ async def __init__(self, **kwargs): self.backend = backend self.device_group = DeviceGroup(**device_group) + self.training_context = kwargs.pop('training_context') + await self.initialize_managed_runtime(self.training_context, kwargs.pop('placement_group'), + kwargs.pop('allocation_devices')) self.device_mesh = init_twinkle_runtime( is_mock=(backend == 'mock'), nproc_per_node=nproc_per_node, device_group=self.device_group, device_mesh_dict=device_mesh, ) - self.replica_id = serve.get_replica_context().replica_id.unique_id + self.replica_id = self.training_context['instance_id'] self.max_loras = kwargs.get('max_loras', 5) self.base_model = model_id # ``train_mode`` selects LoRA multi-tenant (default) or exclusive @@ -130,6 +130,7 @@ async def __init__(self, if self.device_mesh is not None: ctor_kwargs['device_mesh'] = self.device_mesh self.model = MODEL_SELECTOR.construct(backend, ctor_kwargs) + await self.wait_for_training_workers() self.state: ServerState = get_server_state() from twinkle.server.data_plane import DataPlaneProxy @@ -149,7 +150,6 @@ async def __init__(self, self.model._ray_get_timeout = self.task_queue_config.effective_execution_timeout self._init_adapter_manager(**(adapter_config or {})) await self._register_replica_on_startup() - # Note: countdown task is started lazily in _ensure_sticky() @property def data_world_size(self) -> int: @@ -164,34 +164,15 @@ def data_world_size(self) -> int: async def _register_replica_on_startup(self) -> None: """Register this replica's capacity before Ray Serve marks it ready.""" if not self._replica_registered: - await self.state.register_replica(self.replica_id, self.max_loras) + await self.state.register_replica(self.replica_id, 0 if self.is_full_mode else self.max_loras) await self.state.touch_replica_last_seen(self.replica_id) self._replica_registered = True - @serve.multiplexed(max_num_models_per_replica=5) - async def _sticky_entry(self, sticky_key: str): - return sticky_key - - async def _ensure_sticky(self): - sticky_key = serve.get_multiplexed_model_id() - await self._sticky_entry(sticky_key) - # Lazy-start countdown task on first request (requires running event loop) - self._ensure_countdown_started() - async def _on_request_start(self, request: Request) -> str: - await self._ensure_sticky() - await self.state.touch_replica_last_seen(self.replica_id) - await self._ensure_state_cleanup_started() - if self._model_unhealthy: - raise HTTPException(status_code=503, detail='Model actors are unavailable') - return get_token_from_request(request) + return await self.managed_request_start(request) async def shutdown(self) -> None: """Explicit async cleanup — called via FastAPI shutdown event.""" - try: - await self.state.unregister_replica(self.replica_id) - except Exception: - pass await self.shutdown_task_queue() await self.data_plane.close() @@ -199,8 +180,7 @@ async def _run_model_health_probe(self) -> dict: try: result = await self.call_backend(self.model.ping, admit=False) if result is True: - self._model_unhealthy = False - return {'healthy': True, 'detail': 'model actors alive'} + return {'healthy': not self._model_unhealthy, 'detail': 'model actors alive'} self._model_unhealthy = True return {'healthy': False, 'detail': f'unexpected ping result: {result}'} except Exception as e: @@ -219,24 +199,7 @@ def mark_unhealthy(self) -> None: self._model_unhealthy = True async def _probe_after_timeout(self) -> None: - """Fired by ComputeWorker on a backend timeout: probe and log liveness.""" - result = await self.check_model_health() - logger.warning('[Model] post-timeout liveness probe: %s', result) - - async def _cleanup_adapter(self, adapter_name: str) -> None: - if self.get_resource_info(adapter_name): - if self.train_mode == 'full': - # No PEFT adapter to remove; restore clean base weights so the - # next tenant does not inherit this tenant's trained weights. - # Takes the Admission_Gate: this path is driven by the background - # countdown and never enters Task_Queue, so the gate is what keeps - # it from colliding with an in-flight training call. - await self.call_backend(self.model.reload_initial_weights) - else: - await self.call_backend(self.model.remove_adapter, adapter_name) - self.clear_resource_state(adapter_name) - self.unregister_resource(adapter_name) - await self.state.unload_model(adapter_name) + self._model_unhealthy = True @property def is_full_mode(self) -> bool: @@ -253,25 +216,8 @@ def resolve_model_adapter_name(self, adapter_name: str | None) -> str | None: return '' return adapter_name - def assert_full_mode_available(self, adapter_name: str | None = None) -> None: - """In full mode, ensure no other tenant currently holds this deployment. - - Call before registering any state so a rejected request leaves nothing - behind. ``adapter_name`` (when given) is excluded from the check so a - tenant re-issuing against its own resource is not blocked. - - Raises: - FullModeBusyError: another (non-expiring) tenant already holds it. - """ - if not self.is_full_mode: - return - holder = self.find_active_resource(exclude=adapter_name) - if holder is not None: - raise FullModeBusyError(holder) - async def _on_adapter_expired(self, adapter_name: str) -> None: - self.fail_pending_tasks_for_model(adapter_name, reason='Adapter expired') - await self._cleanup_adapter(adapter_name) + raise RuntimeError('Managed training expiry must be handled by the training controller') def build_model_app(model_id: str, @@ -326,7 +272,6 @@ async def _on_shutdown(servable: Any) -> None: ModelManagement, deploy_options, deployment_name='ModelManagement', - request_router_config=RequestRouterConfig(request_router_class=StickyLoraRequestRouter), bind_args=(model_id, nproc_per_node, device_group, device_mesh, backend, adapter_config, queue_config, data_plane_url), bind_kwargs=kwargs, diff --git a/src/twinkle/server/model/tinker_handlers.py b/src/twinkle/server/model/tinker_handlers.py index b455e63d2..4566aafa5 100644 --- a/src/twinkle/server/model/tinker_handlers.py +++ b/src/twinkle/server/model/tinker_handlers.py @@ -20,9 +20,7 @@ from .app import ModelManagement from twinkle.server.checkpoint import create_checkpoint_manager, create_training_run_manager -from twinkle.server.exceptions import FullModeBusyError -from twinkle.server.task_queue.types import UserTaskError -from twinkle.server.utils import get_template_for_model +from twinkle.server.exceptions import RequestRejectedError from twinkle.server.validation.tinker_loss import validate_tinker_loss from twinkle.utils.logger import get_logger @@ -42,60 +40,7 @@ async def create_model( body: types.CreateModelRequest, self: ModelManagement = Depends(self_fn), ) -> types.UntypedAPIFuture: - token = await self._on_request_start(request) - - async def _create_adapter(): - _model_id = None - try: - # Validate lora_config against the deployment's train_mode up front. - if self.is_full_mode and body.lora_config: - raise UserTaskError('This deployment runs in full-parameter (exclusive) mode; do not pass ' - 'lora_config. Use create_full_training_client (or omit lora_config).') - if (not self.is_full_mode) and (not body.lora_config): - raise UserTaskError('This deployment runs in LoRA mode; lora_config is required. ' - 'Use create_lora_training_client.') - # Exclusive full-parameter training: reject early (before touching - # state) if another tenant already holds the deployment. - if self.is_full_mode: - self.assert_full_mode_available() - - _model_id = await self.state.register_model( - body.model_dump(), token=token, replica_id=self.replica_id, session_id=body.session_id) - adapter_name = self.get_adapter_name(adapter_name=_model_id) - model_adapter = self.resolve_model_adapter_name(adapter_name) - # Select template based on model type - template = get_template_for_model(self.base_model) - if self.is_full_mode: - self.register_resource(adapter_name, token, session_id=body.session_id) - await self.call_backend( - self.model.set_template, template, adapter_name=model_adapter, model_id=self.base_model) - await self.call_backend(self.model.set_processor, 'InputProcessor', adapter_name=model_adapter) - await self.call_backend(self.model.set_optimizer, 'Adam', adapter_name=model_adapter) - self.set_resource_state(adapter_name, 'grad_ready', False) - else: - from peft import LoraConfig - lora_cfg = LoraConfig(r=body.lora_config.rank, target_modules='all-linear') - self.register_resource(adapter_name, token, session_id=body.session_id) - await self.call_backend( - self.model.add_adapter_to_model, adapter_name=adapter_name, config_or_dir=lora_cfg) - await self.call_backend( - self.model.set_template, template, adapter_name=adapter_name, model_id=self.base_model) - await self.call_backend(self.model.set_processor, 'InputProcessor', adapter_name=adapter_name) - await self.call_backend(self.model.set_optimizer, 'Adam', adapter_name=adapter_name) - self.set_resource_state(adapter_name, 'grad_ready', False) - training_run_manager = create_training_run_manager(token, client_type='tinker') - training_run_manager.save(_model_id, body) - return types.CreateModelResponse(model_id=_model_id) - except FullModeBusyError as e: - raise UserTaskError(str(e)) from e - except Exception: - if _model_id: - adapter_name = self.get_adapter_name(adapter_name=_model_id) - await self._cleanup_adapter(adapter_name) - logger.error(traceback.format_exc()) - raise - - return await self.schedule_task(_create_adapter, token=token, task_type='create_model') + raise RequestRejectedError('Training lifecycle operations must use the catalog endpoint') @app.post('/tinker/get_info') async def get_info(request: Request, body: types.GetInfoRequest, @@ -123,14 +68,7 @@ async def unload_model( body: types.UnloadModelRequest, self: ModelManagement = Depends(self_fn), ) -> types.UntypedAPIFuture: - token = await self._on_request_start(request) - - async def _do_unload(): - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - await self._cleanup_adapter(adapter_name) - return types.UnloadModelResponse(model_id=body.model_id) - - return await self.schedule_task(_do_unload, model_id=body.model_id, token=token, task_type='unload_model') + raise RequestRejectedError('Training lifecycle operations must use the catalog endpoint') @app.post('/tinker/forward') async def forward(request: Request, body: types.ForwardRequest, diff --git a/src/twinkle/server/model/twinkle_handlers.py b/src/twinkle/server/model/twinkle_handlers.py index 926c09b06..2e1a98d49 100644 --- a/src/twinkle/server/model/twinkle_handlers.py +++ b/src/twinkle/server/model/twinkle_handlers.py @@ -22,12 +22,9 @@ import twinkle.protocol.types as types from twinkle.protocol.serialize import deserialize_object -from twinkle.server.checkpoint import (_resolve_client_save_dir, create_checkpoint_manager, create_training_run_manager, - validate_user_path) -from twinkle.server.exceptions import RequestRejectedError, TrainModeMismatchError -from twinkle.server.lifecycle.submit import (backend_kwargs, input_metrics, resolve_twinkle_adapter_name, run_submit, - to_backend_inputs) -from twinkle.server.middleware.auth import get_session_id_from_request +from twinkle.server.checkpoint import create_checkpoint_manager +from twinkle.server.exceptions import RequestRejectedError +from twinkle.server.lifecycle.submit import backend_kwargs, input_metrics, run_submit, to_backend_inputs from twinkle.server.model.data_plane_inputs import (data_plane_request_shape, merge_forward_kwargs, resolve_data_plane_model_inputs, select_output_rows) from twinkle.server.validation import BackendCapability @@ -74,8 +71,7 @@ async def model_healthz( @app.post('/twinkle/create', response_model=types.CreateResponse) async def create(request: Request, body: types.CreateRequest, self: ModelManagement = Depends(self_fn)) -> types.CreateResponse: - await self._on_request_start(request) - return types.CreateResponse() + raise RequestRejectedError('Training lifecycle operations must use the catalog endpoint') # ------------------------------------------------------------------ # # Inline data / forward family @@ -514,14 +510,7 @@ async def _call(self, body, adapter_name, token): @app.post('/twinkle/remove_adapter', response_model=types.TaskEnvelope) async def remove_adapter( request: Request, body: types.AdapterRequest, self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - """Release a drained tenant's in-memory training adapter.""" - - async def _call(self, body, adapter_name, token): - await self._cleanup_adapter(adapter_name) - return {'status': 'ok'} - - return await run_submit( - self, request, body, task_type='remove_adapter', backend_call=_call, assert_resource=False) + raise RequestRejectedError('Training lifecycle operations must use the catalog endpoint') @app.post('/twinkle/add_adapter_to_model', response_model=types.TaskEnvelope) async def add_adapter_to_model( @@ -538,64 +527,7 @@ async def add_adapter_to_model( # alone, so it owes the caller a real 400. A bare `assert` would surface as a 500 # ('the server broke') and would vanish entirely under `python -O`, letting an # empty adapter_name through to the backend. - if not body.adapter_name: - raise RequestRejectedError('`adapter_name` is required and must be non-empty.') - token = await self._on_request_start(request) - if not validate_user_path(token, body.adapter_name): - raise RequestRejectedError(f'Invalid adapter_name: {body.adapter_name}') - adapter_name = resolve_twinkle_adapter_name(request, body.adapter_name) - session_id = get_session_id_from_request(request) - try: - resolved_save_dir = _resolve_client_save_dir(body.save_dir).as_posix() if body.save_dir else None - except ValueError as exc: - raise RequestRejectedError(str(exc)) from exc - - config = deserialize_object(body.config) - - # ---- Decision_Boundary left: validate against the deployment's train_mode ---- - if self.is_full_mode and config is not None: - raise TrainModeMismatchError('This deployment runs in full-parameter (exclusive) mode; pass ' - 'config=None (do not send a LoraConfig).') - if (not self.is_full_mode) and config is None: - raise TrainModeMismatchError('This deployment runs in LoRA mode; a LoraConfig is required.') - if self.is_full_mode: - # Raises FullModeBusyError (409) if another tenant holds the exclusive deployment. - self.assert_full_mode_available(adapter_name) - - async def _task(): - from peft import LoraConfig - extra_kwargs = backend_kwargs(body) - training_run_manager = create_training_run_manager(token, client_type='twinkle') - lora_config = None - if isinstance(config, LoraConfig): - lora_config = types.LoraConfig(rank=config.r, train_unembed=False, train_mlp=True, train_attn=True) - run_config = types.CreateModelRequest( - base_model=self.base_model, - lora_config=lora_config, - save_dir=resolved_save_dir, - user_metadata={'adapter_name': body.adapter_name}) - await self.state.register_model( - run_config.model_dump(), - token=token, - model_id=adapter_name, - replica_id=self.replica_id, - session_id=session_id, - ) - try: - self.register_resource(adapter_name, token, session_id) - if self.is_full_mode: - # No PEFT adapter to add; the default optimizer group is used. - self.set_resource_state(adapter_name, 'grad_ready', False) - else: - await self.call_backend(self.model.add_adapter_to_model, adapter_name, config, **extra_kwargs) - except Exception: - self.unregister_resource(adapter_name) - await self.state.unload_model(adapter_name) - raise - training_run_manager.save(adapter_name, run_config) - return {'status': 'ok', 'adapter_name': adapter_name} - - return await self.submit_and_peek(_task, model_id=adapter_name, token=token, task_type='add_adapter_to_model') + raise RequestRejectedError('Training lifecycle operations must use the catalog endpoint') # ------------------------------------------------------------------ # # Hub upload (pure I/O -> background task; state-tracked via Retrieve_Endpoint) diff --git a/src/twinkle/server/state/backend/factory.py b/src/twinkle/server/state/backend/factory.py index 842a67dfe..d91a29039 100644 --- a/src/twinkle/server/state/backend/factory.py +++ b/src/twinkle/server/state/backend/factory.py @@ -1,6 +1,8 @@ """Backend factory for creating StateBackend instances based on configuration.""" from __future__ import annotations +import os + from twinkle.server.config.persistence import PersistenceConfig from twinkle.utils import get_logger from .base import StateBackend @@ -29,7 +31,8 @@ def create_backend(config: PersistenceConfig | None = None) -> StateBackend: # Deferred import keeps the module-level import graph light; the Twinkle # server always runs on Ray Serve, so ``ray`` is available here. from .memory_backend import RayActorBackend - return RayActorBackend(key_prefix=config.key_prefix) + return RayActorBackend( + key_prefix=config.key_prefix, namespace=os.environ.get('TWINKLE_STATE_RAY_NAMESPACE')) case 'redis': if not config.redis_url: raise ValueError('redis_url is required for redis persistence mode') diff --git a/src/twinkle/server/state/backend/memory_backend.py b/src/twinkle/server/state/backend/memory_backend.py index 433f8baa4..dde709339 100644 --- a/src/twinkle/server/state/backend/memory_backend.py +++ b/src/twinkle/server/state/backend/memory_backend.py @@ -121,21 +121,26 @@ class RayActorBackend(StateBackend): persistence mode literal for this backend is still ``"memory"``.) """ - def __init__(self, key_prefix: str = '') -> None: + def __init__(self, key_prefix: str = '', namespace: str | None = None) -> None: if not ray.is_initialized(): raise RuntimeError('RayActorBackend requires an initialized Ray runtime — call ' 'ray.init() first, switch persistence to "redis", or rely on ' 'the deployment launcher to start Ray.') name = _actor_name(key_prefix) try: - self._actor = ray.get_actor(name) + self._actor = ray.get_actor(name, namespace=namespace) except ValueError: try: - self._actor = _StateActor.options(name=name, lifetime='detached').remote() + self._actor = _StateActor.options( + name=name, + namespace=namespace, + lifetime='detached', + max_restarts=-1 if namespace else 0, + max_task_retries=0).remote() except ValueError: # Lost the create race against another worker — the actor is # now resolvable. - self._actor = ray.get_actor(name) + self._actor = ray.get_actor(name, namespace=namespace) async def set(self, key: str, value: Any, ttl: int | None = None) -> None: await self._actor.set.remote(key, value, ttl) diff --git a/src/twinkle/server/state/model_manager.py b/src/twinkle/server/state/model_manager.py index aaf520fca..2b01af9f0 100644 --- a/src/twinkle/server/state/model_manager.py +++ b/src/twinkle/server/state/model_manager.py @@ -94,6 +94,7 @@ async def rebuild_indexes(self) -> None: async def get_capacity_info(self) -> dict[str, int]: """Return global LoRA capacity across all registered replicas.""" replicas = await self._replicas.get_all() + replicas = {identity: capacity for identity, capacity in replicas.items() if capacity > 0} loaded_per_replica = await self._loaded_per_replica() total_max = sum(replicas.values()) total_used = sum(loaded_per_replica.get(rid, 0) for rid in replicas) @@ -220,6 +221,9 @@ async def cleanup_expired(self, cutoff_time: float, expired_session_ids: list[st all_records = await self.get_all() expired_ids: list[str] = [] for model_id, record in all_records.items(): + if record.controller_epoch is not None: + # Managed training has one lifecycle writer: TrainingController. + continue if record.session_id and record.session_id in session_set: expired_ids.append(model_id) continue diff --git a/src/twinkle/server/state/models.py b/src/twinkle/server/state/models.py index 0364aa6ce..427dee87e 100644 --- a/src/twinkle/server/state/models.py +++ b/src/twinkle/server/state/models.py @@ -35,6 +35,9 @@ class ModelRecord(BaseModel): lora_config: Any = None replica_id: str | None = None created_at: str = Field(default_factory=_now_iso) + controller_epoch: str | None = None + training_namespace: str | None = None + instance_id: str | None = None class SamplingSessionRecord(BaseModel): @@ -61,6 +64,11 @@ class FutureFailureRecord(BaseModel): # gateways each map this same key set to their own wire vocabularies; keeping the set # in one place lets a consistency test catch a map that drifts out of coverage. FAILURE_REASON_CODES: frozenset[str] = frozenset({ + 'capacity_unavailable', + 'training_expired', + 'startup_failed', + 'startup_timeout', + 'protocol_mismatch', 'invalid_request', 'request_rejected', 'resource_not_found', diff --git a/src/twinkle/server/state/sequence.py b/src/twinkle/server/state/sequence.py new file mode 100644 index 000000000..78efab55b --- /dev/null +++ b/src/twinkle/server/state/sequence.py @@ -0,0 +1,27 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Bounded training retry receipts with a persistent expired-sequence fence.""" + + +def claim_sequence(current, *, sequence, request_id, now, retention): + floor = current.get('floor') if current else None + receipts = {} + for key, receipt in (current or {}).get('receipts', {}).items(): + if receipt['expires_at'] <= now: + floor = int(key) if floor is None else max(floor, int(key)) + else: + receipts[key] = receipt + key = str(sequence) + # Unexpired receipts remain retryable even if another sequence expired first. + if key not in receipts and (floor is None or sequence > floor): + receipts[key] = {'request_id': request_id, 'expires_at': now + retention} + return {'floor': floor, 'receipts': receipts} + + +def release_sequence(current, *, sequence, request_id): + if current is None: + return None + receipts = dict(current['receipts']) + key = str(sequence) + if receipts.get(key, {}).get('request_id') == request_id: + receipts.pop(key) + return {'floor': current['floor'], 'receipts': receipts} diff --git a/src/twinkle/server/state/server_state.py b/src/twinkle/server/state/server_state.py index dd926aa57..461d9a7a4 100644 --- a/src/twinkle/server/state/server_state.py +++ b/src/twinkle/server/state/server_state.py @@ -10,7 +10,7 @@ from typing import Any from twinkle.server.config.persistence import PersistenceConfig -from twinkle.server.exceptions import ResourceQuotaExceededError +from twinkle.server.exceptions import RequestRejectedError, ResourceNotFoundError, ResourceQuotaExceededError from twinkle.server.telemetry.correlation import (BASE_MODEL, MODEL_ID, REPLICA_ID, SAMPLING_SESSION_ID, SESSION_ID, TOKEN_ID) from twinkle.server.telemetry.tracing import traced_operation @@ -173,6 +173,17 @@ async def _on_lose_leader(self) -> None: async def get_capacity_info(self) -> dict[str, int]: return await self._model_mgr.get_capacity_info() + @property + def backend(self) -> StateBackend: + return self._backend + + async def register_model_record(self, model_id: str, record: ModelRecord) -> None: + """Reserve quota for a complete model record, including managed training state.""" + await self._model_mgr.add(model_id, record) + + async def rebuild_model_indexes(self) -> None: + await self._model_mgr.rebuild_indexes() + # ----- Session Management ----- async def create_session(self, payload: dict[str, Any]) -> str: @@ -429,7 +440,12 @@ async def get_future(self, request_id: str) -> dict[str, Any] | None: record = await self._future_mgr.get(request_id) return record.model_dump() if record is not None else None - async def claim_seq(self, dedup_key: str, request_id: str, ttl: int) -> str | None: + async def assert_future_owner(self, request_id: str, token: str) -> None: + owner = await self._backend.get(f'future_owner::{request_id}') + if owner is not None and owner != token: + raise ResourceNotFoundError('Unknown future') + + async def claim_seq(self, dedup_key: str, request_id: str, ttl: int | None) -> str | None: """Idempotency claim for a client seq_id. Atomically records ``dedup_key -> request_id`` if unseen and returns ``None`` @@ -446,6 +462,23 @@ async def release_seq(self, dedup_key: str) -> None: preflight rejected it) so a retry can be admitted rather than see a phantom.""" await self._backend.delete(dedup_key) + async def claim_training_seq(self, training_id: str, stream: str, sequence: int, request_id: str, + retention: float) -> str | None: + from .sequence import claim_sequence + ledger = await self._backend.update_atomic( + f'seq_ledger::{training_id}::{stream}', + functools.partial( + claim_sequence, sequence=sequence, request_id=request_id, now=time.time(), retention=retention)) + receipt = ledger['receipts'].get(str(sequence)) + if receipt is None: + raise RequestRejectedError('Sequence retry window expired; the operation will not be executed again') + return None if receipt['request_id'] == request_id else receipt['request_id'] + + async def release_training_seq(self, training_id: str, stream: str, sequence: int, request_id: str) -> None: + from .sequence import release_sequence + await self._backend.update_atomic(f'seq_ledger::{training_id}::{stream}', + functools.partial(release_sequence, sequence=sequence, request_id=request_id)) + async def cancel_future(self, request_id: str) -> dict[str, Any]: """Best-effort cancel: drop the task iff it has not started running. diff --git a/src/twinkle/server/task_queue/mixin.py b/src/twinkle/server/task_queue/mixin.py index f26ec9c5c..be24ac9d9 100644 --- a/src/twinkle/server/task_queue/mixin.py +++ b/src/twinkle/server/task_queue/mixin.py @@ -139,6 +139,10 @@ def task_queue_config(self) -> TaskQueueConfig: """ return self._task_queue_config + @property + def sequence_retention_s(self) -> float: + return self.state.expiration_timeout + self._future_absolute_ttl + async def call_backend(self, fn: Callable[..., Any], /, *args: Any, admit: bool = True, **kwargs: Any) -> Any: """Run one backend call outside the event loop. @@ -227,9 +231,32 @@ async def _schedule_task( batch_size_multiple=batch_size_multiple, ) + task = QueuedTask( + request_id=request_id, + coro_factory=coro_factory, + model_id=model_id, + token=token, + input_tokens=input_tokens, + task_type=task_type, + created_at=time.monotonic()) + await self._prepare_queued_task(task) + if self._event_loop is None: self._event_loop = asyncio.get_running_loop() + try: + await self._enqueue_prepared_task(task) + except BaseException: + if task.on_finished is not None and await self.state.get_future(request_id) is None: + await task.on_finished(False, False) + raise + return {'request_id': request_id, 'model_id': model_id} + + async def _prepare_queued_task(self, task: QueuedTask) -> None: + """Optional owner admission, after preflight and before any future write.""" + + async def _enqueue_prepared_task(self, task: QueuedTask) -> None: + request_id, model_id = task.request_id, task.model_id await self.state.store_future_status( request_id, TaskStatus.PENDING.value, @@ -239,35 +266,27 @@ async def _schedule_task( absolute_deadline=self._future_deadline(), ) - queue_key = self._queue_key(model_id=model_id, token=token) + queue_key = self._queue_key(model_id=model_id, token=task.token) self._compute_worker.ensure_queue_registered(queue_key) await self._compute_worker.ensure_started() q = self._compute_worker.get_queue(queue_key) - await q.put( - QueuedTask( - request_id=request_id, - coro_factory=coro_factory, - model_id=model_id, - token=token, - input_tokens=input_tokens, - task_type=task_type, - created_at=time.monotonic(), - )) + self._compute_worker.track_task(request_id) + await q.put(task) + if task.on_enqueued is not None: + await task.on_enqueued() await self.state.store_future_status( request_id, TaskStatus.QUEUED.value, model_id, queue_state=QueueState.ACTIVE.value, ) - logger.info(f'[TaskQueue] Task {request_id} queued, type={task_type or "unknown"}, ' + logger.info(f'[TaskQueue] Task {request_id} queued, type={task.task_type or "unknown"}, ' f'model_id={model_id}, queue_key={queue_key}, ' - f'queue_depth={q.qsize()}, input_tokens={input_tokens}') + f'queue_depth={q.qsize()}, input_tokens={task.input_tokens}') self._compute_worker.notify_new_task() - return {'request_id': request_id, 'model_id': model_id} - async def schedule_task( self, coro_factory: Callable[[], Coroutine], diff --git a/src/twinkle/server/task_queue/types.py b/src/twinkle/server/task_queue/types.py index 936a433a3..6e522785e 100644 --- a/src/twinkle/server/task_queue/types.py +++ b/src/twinkle/server/task_queue/types.py @@ -68,3 +68,5 @@ class QueuedTask: input_tokens: int task_type: str | None created_at: float + on_finished: Callable[[bool, bool], Coroutine] | None = None + on_enqueued: Callable[[], Coroutine] | None = None diff --git a/src/twinkle/server/task_queue/worker.py b/src/twinkle/server/task_queue/worker.py index 93d4e165f..fd9704feb 100644 --- a/src/twinkle/server/task_queue/worker.py +++ b/src/twinkle/server/task_queue/worker.py @@ -85,6 +85,7 @@ def __init__( self._on_backend_timeout = on_backend_timeout self.task_queues: dict[str, asyncio.Queue] = {} + self._request_ids: set[str] = set() self.queue_order: Deque[str] = deque() self.new_task_event: asyncio.Event = asyncio.Event() @@ -149,6 +150,13 @@ def get_queue(self, queue_key: str) -> asyncio.Queue: """ return self.task_queues[queue_key] + def track_task(self, request_id: str) -> None: + self._request_ids.add(request_id) + + def has_task(self, request_id: str) -> bool: + """True while a task is queued or executing, including a terminal-future race.""" + return request_id in self._request_ids + def total_queued(self) -> int: """Total number of pending tasks across all per-key queues.""" return sum(q.qsize() for q in self.task_queues.values()) @@ -246,6 +254,9 @@ async def fail_queue_tasks(self, queue_key: str, reason: str) -> None: for task in drained: await self._store_task_failed(task, reason, QueueState.UNKNOWN.value, queue_state_reason=reason) q.task_done() + self._request_ids.discard(task.request_id) + if task.on_finished is not None: + await task.on_finished(False, False) self.task_queues.pop(queue_key, None) try: @@ -272,6 +283,9 @@ async def _fail_timed_out_task(self, task: QueuedTask, waited: float, q: asyncio await self._store_task_failed(task, error, QueueState.PAUSED_CAPACITY.value, queue_state_reason=error) self._record_tasks_total(task.task_type or 'unknown', 'timeout') q.task_done() + self._request_ids.discard(task.request_id) + if task.on_finished is not None: + await task.on_finished(False, False) async def _execute_task(self, task: QueuedTask, queue_key: str, q: asyncio.Queue) -> None: """Execute a dequeued task, then release its queue slot.""" @@ -369,6 +383,7 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: exec_start = time.monotonic() task_status = 'completed' exec_time = 0.0 + executed = False # One span per queued task execution, tagged `.` # (e.g. `model.forward`, `sampler.sample`). queue_attrs = { @@ -399,6 +414,7 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: logger.debug(f'[ComputeWorker] Task {task.request_id} executing, ' f'type={task_type}, queue_key={queue_key}') with traced_operation(handler_span_name, attrs=handler_attrs): + executed = True coro = task.coro_factory() # effective_execution_timeout is always positive (0 -> finite fallback), # so wait_for is always in effect. @@ -492,6 +508,9 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: traceback_text=traceback.format_exc(), ) finally: + self._request_ids.discard(task.request_id) + if task.on_finished is not None: + await task.on_finished(executed, executed and task_status in ('timeout', 'cancelled')) self._record_execution_time(task_type, exec_time) self._record_tasks_total(task_type, task_status) @@ -529,10 +548,16 @@ async def _try_run_one(self) -> bool: if await self._is_record_terminal(task.request_id): logger.info(f'[ComputeWorker] Task {task.request_id} already terminal on dequeue; skipping.') q.task_done() + self._request_ids.discard(task.request_id) + if task.on_finished is not None: + await task.on_finished(False, False) continue # Execute the task (serial: stops after the first execution) - await self._execute_task(task, queue_key, q) + try: + await self._execute_task(task, queue_key, q) + finally: + self._request_ids.discard(task.request_id) return True return False diff --git a/src/twinkle/server/training/__init__.py b/src/twinkle/server/training/__init__.py new file mode 100644 index 000000000..33dc7caf2 --- /dev/null +++ b/src/twinkle/server/training/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""On-demand training: catalog, lifecycle controller and runtime ownership.""" diff --git a/src/twinkle/server/training/actor.py b/src/twinkle/server/training/actor.py new file mode 100644 index 000000000..4ceb1e392 --- /dev/null +++ b/src/twinkle/server/training/actor.py @@ -0,0 +1,156 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""One named controller actor per server namespace.""" +from __future__ import annotations + +import asyncio +import ray +import time + +from twinkle.server.config.training import TrainingConfig +from twinkle.server.state import get_server_state +from twinkle.server.state.models import FutureFailureRecord +from twinkle.utils.logger import get_logger +from .controller import TrainingController +from .driver import RayInstanceDriver +from .records import CreateTraining +from .store import TrainingStore + +logger = get_logger() + + +def controller_name(namespace): + return f'training-controller-{namespace}' + + +@ray.remote(num_cpus=0.1, max_restarts=-1, max_task_retries=0) +class TrainingControllerActor: + + def __init__(self, config, namespace): + self.config = TrainingConfig.model_validate(config) + self.namespace = namespace + self.store = TrainingStore(get_server_state(), namespace) + self.driver = RayInstanceDriver(self.config, self.store, namespace, controller_name(namespace)) + self.core = TrainingController(self.config, self.store, self.driver) + self._monitors = [] + self._unreachable_since = {} + self._initialization = asyncio.get_event_loop().create_task(self._initialize()) + + async def initialize(self): + await asyncio.shield(self._initialization) + + async def _initialize(self): + await self.core.initialize() + self._monitors = [asyncio.create_task(self._heartbeat()), asyncio.create_task(self._sweep())] + + async def _heartbeat(self): + while True: + try: + active = await self.store.backend.get(self.store.key('controller', 'active')) + if active is None: + await self._recover_lost_state() + return + await self.store.state.touch_replica_last_seen(self.core.epoch) + except Exception: + logger.exception('Training state heartbeat failed; admission remains fenced by the state backend') + await asyncio.sleep(5) + continue + for instance in list(self.core.instances.values()): + if instance.state not in ('READY', 'RESETTING'): + continue + try: + identity = await self.driver.health(instance) + if identity['unhealthy']: + await self.core.fail_instance( + instance.instance_id, + FutureFailureRecord( + reason_code='orphaned_replica', + attribution='server', + message='Training runtime reported a failed worker')) + continue + await self.store.state.touch_replica_last_seen(instance.instance_id) + self._unreachable_since.pop(instance.instance_id, None) + except Exception: + since = self._unreachable_since.setdefault(instance.instance_id, time.monotonic()) + if time.monotonic() - since >= 30: + await self.core.fail_instance( + instance.instance_id, + FutureFailureRecord( + reason_code='orphaned_replica', + attribution='server', + message='Training runtime became unavailable')) + await asyncio.sleep(5) + + async def _sweep(self): + while True: + try: + await self.core.sweep() + except Exception: + logger.exception('Training lifecycle sweep failed; resources remain owned') + await asyncio.sleep(self.config.sweep_interval_s) + + async def _recover_lost_state(self): + # Reconstruct enough ownership for the next epoch to fail pending + # futures and terminate compute. Sessions were lost, so old handles + # cannot be reused, even while these rows are being restored. + failure = FutureFailureRecord( + reason_code='orphaned_replica', + attribution='server', + message='Training state was lost; create a new session and training handle.') + for instance in self.core.instances.values(): + if instance.state != 'STOPPED': + await self.store.save_instance(instance) + for training in self.core.trainings.values(): + if training.state in ('PENDING', 'ACTIVE', 'CLOSING'): + await self.store.save_training(training) + for request_id in (training.create_request_id, *training.pending_operation_ids): + await self.store.own_future(request_id, training.token) + await self.store.future(request_id, 'failed', training.training_id, failure=failure) + ray.kill(ray.get_actor(controller_name(self.namespace), namespace=self.namespace), no_restart=False) + + async def register_runtime(self, instance_id, epoch, incarnation): + await self.core.register_runtime(instance_id, epoch, incarnation) + + async def create(self, request): + await self.initialize() + return await self.core.create(CreateTraining.model_validate(request)) + + async def route(self, *args): + await self.initialize() + return await self.core.route(*args) + + async def admit(self, *args, **kwargs): + await self.initialize() + return await self.core.admit(*args, **kwargs) + + async def commit_operation(self, *args): + return await self.core.commit_operation(*args) + + async def finish_operation(self, *args, **kwargs): + return await self.core.finish_operation(*args, **kwargs) + + async def release(self, *args): + return await self.core.release(*args) + + async def fail_instance(self, *args): + return await self.core.fail_instance(*args) + + async def mark_dirty(self, instance_id, epoch): + await self.core.mark_dirty(instance_id, epoch) + + async def diagnostics(self, base_model=None): + await self.initialize() + instances = [ + instance.model_dump() for instance in self.core.instances.values() + if base_model is None or instance.base_model == base_model + ] + return { + 'ready': self.core._ready and self.core.residents_ready, + 'epoch': self.core.epoch, + 'instances': instances + } + + async def shutdown(self): + for task in self._monitors: + task.cancel() + await asyncio.gather(*self._monitors, return_exceptions=True) + await self.core.shutdown() diff --git a/src/twinkle/server/training/allocation.py b/src/twinkle/server/training/allocation.py new file mode 100644 index 000000000..08bdf8573 --- /dev/null +++ b/src/twinkle/server/training/allocation.py @@ -0,0 +1,93 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Single-node device allocation. Only the owner removes placement groups.""" +from __future__ import annotations + +import asyncio +from dataclasses import dataclass + +from twinkle import Platform +from twinkle.infra._ray.resource_manager import _gpu_placement_group_cpus +from .catalog import TrainingSpec +from .errors import CapacityUnavailableError + + +@dataclass(frozen=True) +class DeviceAllocation: + name: str + placement_group_id: str + node_id: str + devices: tuple[str, ...] + + +def _probe_allocation(resource: str, device_type: str): + import os + import ray + context = ray.get_runtime_context() + if resource == 'CPU': + devices = [] + else: + env = os.environ.get(Platform.get_platform(device_type).visible_device_env()) + devices = env.split(',') if env else context.get_accelerator_ids().get(resource, []) + return context.get_node_id(), [str(device) for device in devices] + + +class DevicePool: + + def __init__(self, namespace: str, *, allocation_timeout: float, shutdown_timeout: float): + self.namespace = namespace + self.allocation_timeout = allocation_timeout + self.shutdown_timeout = shutdown_timeout + + def name(self, epoch: str, instance_id: str) -> str: + return f'training-{self.namespace}-{epoch}-{instance_id}' + + async def allocate(self, spec: TrainingSpec, epoch: str, instance_id: str) -> DeviceAllocation: + import ray + from ray.util.placement_group import placement_group + from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + + resource = 'CPU' if spec.model.device_type == 'cpu' else Platform.get_platform(spec.model.device_type).__name__ + count = spec.config.num_devices + candidates = [node for node in ray.nodes() if node.get('Alive') and node['Resources'].get(resource, 0) >= count] + available = ray.available_resources() + if not candidates or available.get(resource, 0) < count: + raise CapacityUnavailableError(f'No single-node allocation available for {count} {resource} devices') + cpus = min(_gpu_placement_group_cpus(int(node['Resources'].get('CPU', 0)), count) for node in candidates) + bundle = {resource: count, 'CPU': count if resource == 'CPU' else cpus} + name = self.name(epoch, instance_id) + pg = placement_group([bundle], strategy='STRICT_PACK', name=name) + try: + await asyncio.wait_for(pg.ready(), self.allocation_timeout) + options = {'num_cpus': 0.01, 'scheduling_strategy': PlacementGroupSchedulingStrategy(placement_group=pg)} + if resource == 'GPU': + options['num_gpus'] = count + elif resource != 'CPU': + options['resources'] = {resource: count} + node, devices = await asyncio.wait_for( + ray.remote(_probe_allocation).options(**options).remote(resource, spec.model.device_type), + self.shutdown_timeout) + if resource == 'CPU': + devices = [str(index) for index in range(count)] + if len(devices) != count or len(set(devices)) != count: + raise RuntimeError('Ray allocation did not return the requested unique device IDs') + return DeviceAllocation(name, pg.id.hex(), node, tuple(devices)) + except BaseException as exc: + await self.remove(name) + if isinstance(exc, TimeoutError): + raise CapacityUnavailableError( + 'Single-node allocation could not be confirmed; retry a new create later.') from exc + raise + + async def remove(self, name: str) -> None: + from ray.util.placement_group import get_placement_group, placement_group_table, remove_placement_group + try: + pg = get_placement_group(name) + except ValueError: + return + remove_placement_group(pg) + deadline = asyncio.get_running_loop().time() + self.shutdown_timeout + while asyncio.get_running_loop().time() < deadline: + if placement_group_table(pg).get('state') == 'REMOVED': + return + await asyncio.sleep(0.05) + raise TimeoutError(f'Placement group removal was not confirmed: {name}') diff --git a/src/twinkle/server/training/catalog.py b/src/twinkle/server/training/catalog.py new file mode 100644 index 000000000..1be0548e6 --- /dev/null +++ b/src/twinkle/server/training/catalog.py @@ -0,0 +1,44 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Pure lookup and request validation, separate from runtime allocation.""" +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +from twinkle.server.config.training import LoraTrainingConfig, TrainingConfig, TrainingModeConfig, TrainingModelConfig +from twinkle.server.exceptions import RequestRejectedError + + +@dataclass(frozen=True) +class TrainingSpec: + model: TrainingModelConfig + mode: Literal['lora', 'full'] + config: TrainingModeConfig + + @property + def key(self) -> tuple[str, str]: + return self.model.name, self.mode + + @property + def capacity(self) -> int: + return self.config.max_loras if isinstance(self.config, LoraTrainingConfig) else 1 + + +class TrainingCatalog: + + def __init__(self, config: TrainingConfig): + self._models = {model.name: model for model in config.models} + + def get(self, name: str, mode: Literal['lora', 'full']) -> TrainingSpec: + model = self._models.get(name) + if model is None: + raise RequestRejectedError(f'Unknown training model: {name}') + config = getattr(model.modes, mode) + if config is None: + raise RequestRejectedError(f'Training model {name} does not provide {mode} mode') + return TrainingSpec(model, mode, config) + + def specs(self) -> tuple[TrainingSpec, ...]: + return tuple( + TrainingSpec(model, mode, config) for model in self._models.values() for mode in ('lora', 'full') + if (config := getattr(model.modes, mode)) is not None) diff --git a/src/twinkle/server/training/controller.py b/src/twinkle/server/training/controller.py new file mode 100644 index 000000000..2b8dee8a8 --- /dev/null +++ b/src/twinkle/server/training/controller.py @@ -0,0 +1,668 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Single-writer training lifecycle; slow runtime work stays outside the lock.""" +from __future__ import annotations + +import asyncio +import logging +import time +import traceback +import uuid +from collections.abc import Callable +from datetime import datetime +from typing import Any, Protocol + +from twinkle.server.config.training import TrainingConfig +from twinkle.server.exceptions import RequestRejectedError, ResourceNotFoundError, ResourceQuotaExceededError +from twinkle.server.state.models import FutureFailureRecord +from .catalog import TrainingCatalog, TrainingSpec +from .errors import CapacityUnavailableError, TrainingExpiredError +from .records import LIVE_INSTANCES, TERMINAL_TRAINING, CreateTraining, InstanceRecord, TrainingRecord +from .store import TrainingStore + +logger = logging.getLogger(__name__) + + +class InstanceDriver(Protocol): + + async def start(self, instance: InstanceRecord, spec: TrainingSpec) -> dict[str, Any]: + ... + + async def create_training(self, instance: InstanceRecord, training: TrainingRecord, + request: CreateTraining) -> None: + ... + + async def release_training(self, instance: InstanceRecord, training: TrainingRecord) -> None: + ... + + async def shutdown(self, instance: InstanceRecord) -> None: + ... + + async def verify_operation(self, instance: InstanceRecord, request_id: str) -> dict[str, Any]: + ... + + async def recover_orphans(self) -> None: + ... + + +def failure_for(exc: Exception) -> FutureFailureRecord: + if isinstance(exc, CapacityUnavailableError): + reason, attribution = 'capacity_unavailable', 'user' + elif isinstance(exc, ResourceQuotaExceededError): + reason, attribution = 'resource_quota_exceeded', 'user' + elif isinstance(exc, TimeoutError): + reason, attribution = 'startup_timeout', 'server' + elif isinstance(exc, RequestRejectedError): + reason, attribution = 'invalid_request', 'user' + else: + reason, attribution = 'startup_failed', 'server' + return FutureFailureRecord( + reason_code=reason, + message=str(exc)[:1024], + attribution=attribution, + diagnostic=traceback.format_exc() if attribution == 'server' else None) + + +class TrainingController: + + def __init__(self, + config: TrainingConfig, + store: TrainingStore, + driver: InstanceDriver, + *, + clock: Callable[[], float] = time.monotonic): + self.config = config + self.catalog = TrainingCatalog(config) + self.store = store + self.driver = driver + self.clock = clock + self.epoch = 'controller_' + uuid.uuid4().hex + self.instances: dict[str, InstanceRecord] = {} + self.trainings: dict[str, TrainingRecord] = {} + self._startup: dict[str, asyncio.Task] = {} + self._jobs: set[asyncio.Task] = set() + self._cleaning: set[str] = set() + self._closing: set[str] = set() + self._reconciling: set[str] = set() + self._lock = asyncio.Lock() + self._ready = False + + def _spawn(self, coroutine) -> asyncio.Task: + task = asyncio.create_task(coroutine) + self._jobs.add(task) + + def done(completed): + self._jobs.discard(completed) + if not completed.cancelled(): + completed.exception() + + task.add_done_callback(done) + return task + + async def initialize(self) -> None: + """Discard old state only after owned compute has been confirmed dead.""" + old_instances = await self.store.instances() + old_trainings = await self.store.trainings() + old_controller = await self.store.backend.get(self.store.key('controller', 'active')) + await self.store.begin_epoch(self.epoch) + failure = FutureFailureRecord( + reason_code='orphaned_replica', + message='Training controller restarted; create a new training handle.', + attribution='server') + for training in old_trainings: + for request_id in (training.create_request_id, *training.pending_operation_ids): + await self.store.future(request_id, 'failed', training.training_id, failure=failure) + await self.driver.recover_orphans() + for instance in old_instances: + if instance.state != 'STOPPED': + await asyncio.wait_for(self.driver.shutdown(instance), self.config.shutdown_timeout_s) + instance.state = 'STOPPED' + instance.stopped_at = time.time() + await self.store.save_instance(instance) + for training in old_trainings: + training.state = 'FAILED' + training.terminal_reason = 'orphaned_replica' + training.terminal_at = time.time() + await self.store.retire(training) + await self.store.state.rebuild_model_indexes() + if old_controller: + await self.store.state.unregister_replica(old_controller['epoch']) + await self.store.state.register_replica(self.epoch, 0) + await self.store.state.touch_replica_last_seen(self.epoch) + try: + for spec in self.catalog.specs(): + for _ in range(spec.config.min_replicas): + async with self._lock: + instance = await self._new_instance(spec, resident=True) + await asyncio.shield(self._startup[instance.instance_id]) + except Exception: + await self.shutdown() + raise + self._ready = True + + @property + def residents_ready(self) -> bool: + return all(self.residents_ready_for(spec) for spec in self.catalog.specs()) + + def residents_ready_for(self, spec: TrainingSpec) -> bool: + return sum( + instance.resident and (instance.base_model, instance.mode) == spec.key and instance.state in ('READY', + 'RESETTING') + for instance in self.instances.values()) >= spec.config.min_replicas + + async def register_runtime(self, instance_id: str, epoch: str, incarnation: str) -> None: + async with self._lock: + await self._assert_epoch() + instance = self.instances.get(instance_id) + if (epoch != self.epoch or instance is None or instance.incarnation is not None + or instance.state not in ('ALLOCATING', 'STARTING')): + raise ResourceNotFoundError('Runtime incarnation is no longer allowed') + instance.incarnation = incarnation + instance.state = 'STARTING' + await self.store.save_instance(instance) + + async def mark_dirty(self, instance_id: str, epoch: str) -> None: + async with self._lock: + await self._assert_epoch() + instance = self.instances.get(instance_id) + if epoch != self.epoch or instance is None: + raise ResourceNotFoundError('Unknown runtime') + instance.dirty = True + await self.store.save_instance(instance) + + async def _assert_epoch(self) -> None: + active = await self.store.backend.get(self.store.key('controller', 'active')) + if not active or active['epoch'] != self.epoch: + raise ResourceNotFoundError('Training controller epoch is no longer active') + + async def _new_instance(self, spec: TrainingSpec, *, resident: bool = False) -> InstanceRecord: + instance = InstanceRecord( + instance_id='instance_' + uuid.uuid4().hex, + controller_epoch=self.epoch, + base_model=spec.model.name, + mode=spec.mode, + resident=resident, + startup_deadline=self.clock() + self.config.startup_timeout_s) + await self.store.save_instance(instance) + self.instances[instance.instance_id] = instance + self._startup[instance.instance_id] = self._spawn(self._start_instance(instance)) + return instance + + async def _start_instance(self, instance: InstanceRecord) -> None: + spec = self.catalog.get(instance.base_model, instance.mode) + try: + details = await asyncio.wait_for(self.driver.start(instance, spec), self.config.startup_timeout_s) + async with self._lock: + await self._assert_epoch() + if self.clock() >= instance.startup_deadline: + raise TimeoutError('Training runtime became ready after its startup deadline') + if instance.state not in ('ALLOCATING', 'STARTING'): + raise RuntimeError('Instance became unavailable during startup') + for name, value in details.items(): + if name not in ('application_name', 'placement_group_id', 'node_id', 'devices'): + raise ValueError(f'Unknown runtime allocation field: {name}') + setattr(instance, name, value) + instance.state = 'READY' + await self.store.save_instance(instance) + except Exception as exc: + await self.fail_instance(instance.instance_id, failure_for(exc)) + raise + + def _select(self, spec: TrainingSpec) -> InstanceRecord | None: + candidates = [ + instance for instance in self.instances.values() + if (instance.base_model, instance.mode) == spec.key and instance.state in ( + 'READY', 'ALLOCATING', 'STARTING') and len(instance.training_ids) < spec.capacity + ] + return min( + candidates, + key=lambda instance: (instance.state != 'READY', -len(instance.training_ids), instance.instance_id), + default=None) + + async def create(self, request: CreateTraining) -> dict[str, str]: + # Catalog/mode errors are rejected before a future or any resource exists. + spec = self.catalog.get(request.base_model, request.mode) + identity, fingerprint = self.store.create_identity(request) + async with self._lock: + await self._assert_epoch() + if not self._ready: + raise RequestRejectedError('Training service is not ready') + previous = await self.store.get_create(identity) + if previous is not None: + if previous['fingerprint'] != fingerprint: + raise RequestRejectedError('Create identity was reused with different parameters') + return {'request_id': previous['request_id']} + request_id = 'req_' + uuid.uuid4().hex + training_id = 'training_' + uuid.uuid4().hex + await self.store.future( + request_id, + 'pending', + training_id, + epoch=self.epoch, + deadline=time.time() + 2 * self.config.startup_timeout_s) + await self.store.own_future(request_id, request.token) + await self.store.save_create(identity, fingerprint, request_id, request.session_id) + training = None + try: + # Reserve quota before starting an instance, but leave no usable route. + training = TrainingRecord( + training_id=training_id, + token=request.token, + session_id=request.session_id, + base_model=request.base_model, + mode=request.mode, + controller_epoch=self.epoch, + training_namespace=self.store.prefix, + instance_id='', + adapter_label=request.adapter_label, + create_request_id=request_id, + creation_deadline=self.clock() + self.config.startup_timeout_s, + last_activity_at=self.clock(), + lora_config=request.payload.get('lora_config'), + user_metadata=request.payload.get('user_metadata') or {}) + await self.store.reserve(training) + instance = self._select(spec) + if instance is None: + if not self.residents_ready_for(spec): + raise CapacityUnavailableError( + 'Resident training capacity is recovering; retry a new create later.') + count = sum((item.base_model, item.mode) == spec.key and item.state != 'STOPPED' + for item in self.instances.values()) + if spec.config.max_replicas is not None and count >= spec.config.max_replicas: + raise CapacityUnavailableError('Training replica limit reached; retry a new create later.') + instance = await self._new_instance(spec) + training.instance_id = instance.instance_id + training.replica_id = instance.instance_id + instance.training_ids.append(training_id) + self.trainings[training_id] = training + await self.store.save_training(training) + await self.store.save_instance(instance) + self._spawn(self._create_user(training, request)) + except Exception as exc: + if training is not None: + self.trainings.pop(training_id, None) + if training.instance_id in self.instances: + instance = self.instances[training.instance_id] + if training_id in instance.training_ids: + instance.training_ids.remove(training_id) + await self.store.save_instance(instance) + if not instance.training_ids and not instance.resident: + self._spawn(self.fail_instance(instance.instance_id, failure_for(exc))) + training.state = 'FAILED' + training.terminal_reason = failure_for(exc).reason_code + await self.store.retire(training) + await self.store.future(request_id, 'failed', training_id, failure=failure_for(exc)) + return {'request_id': request_id} + + async def _create_user(self, training: TrainingRecord, request: CreateTraining) -> None: + instance = self.instances[training.instance_id] + try: + await asyncio.shield(self._startup[instance.instance_id]) + async with self._lock: + future = await self.store.state.get_future(training.create_request_id) + if future and future['status'] == 'cancelled': + training.state = 'CLOSING' + training.terminal_reason = 'cancelled' + if training.state != 'PENDING' or instance.state != 'READY': + raise RequestRejectedError('Training creation is no longer pending') + training.creation_deadline = self.clock() + self.config.startup_timeout_s + await self.store.save_training(training) + await self.store.future(training.create_request_id, 'running', training.training_id) + remaining = training.creation_deadline - self.clock() + await asyncio.wait_for(self.driver.create_training(instance, training, request), max(0, remaining)) + # Driver includes training-run metadata commit, before public success. + async with self._lock: + await self._assert_epoch() + if self.clock() >= training.creation_deadline: + raise TimeoutError('Training initialization completed after its creation deadline') + if training.state != 'PENDING' or instance.state != 'READY': + raise RuntimeError('Training became unavailable during initialization') + training.state = 'ACTIVE' + training.last_activity_at = self.clock() + await self.store.save_training(training) + result = ({ + 'model_id': training.training_id + } if request.client_type == 'tinker' else { + 'training_id': training.training_id, + 'adapter_name': training.adapter_label + }) + await self.store.future(training.create_request_id, 'completed', training.training_id, result=result) + except Exception as exc: + if training.state not in TERMINAL_TRAINING: + await self.store.future( + training.create_request_id, 'failed', training.training_id, failure=failure_for(exc)) + async with self._lock: + training.state = 'CLOSING' + training.terminal_reason = ('cancelled' if training.terminal_reason == 'cancelled' else + failure_for(exc).reason_code) + await self.store.save_training(training) + await self._cleanup_user(training) + + async def _expired(self, training: TrainingRecord) -> bool: + spec = self.catalog.get(training.base_model, training.mode) + config, legacy = spec.config, spec.model.legacy_expiration + heartbeat = await self.store.state.get_session_last_heartbeat(training.session_id) + if heartbeat is None: + return True + if legacy is not None: + now = time.time() + if legacy.session_timeout_s is not None and now - heartbeat >= legacy.session_timeout_s: + return True + if legacy.max_lifetime_s is not None: + created_at = datetime.fromisoformat(training.created_at).timestamp() + if now - created_at >= legacy.max_lifetime_s: + return True + return config.idle_timeout_s is not None and self.clock() - training.last_activity_at >= config.idle_timeout_s + + async def _authorized(self, training_id: str, token: str, session_id: str | None) -> TrainingRecord: + await self._assert_epoch() + training = self.trainings.get(training_id) + if training is None or training.token != token or (session_id is not None + and session_id != training.session_id): + raise ResourceNotFoundError('Unknown training handle') + if training.state == 'EXPIRED': + raise TrainingExpiredError('Training handle expired; create a new handle and explicitly load a checkpoint.') + if training.state != 'ACTIVE': + raise ResourceNotFoundError('Training handle is no longer active') + if await self.store.state.get_session_last_heartbeat(training.session_id) is None: + raise TrainingExpiredError('Training session expired; create a new training handle') + return training + + async def admit(self, + training_id: str, + token: str, + session_id: str | None, + request_id: str, + *, + renew: bool = True) -> dict[str, str]: + async with self._lock: + training = await self._authorized(training_id, token, session_id) + if not training.pending_operation_ids and await self._expired(training): + training.state = 'CLOSING' + training.terminal_reason = 'expired' + await self.store.save_training(training) + self._spawn(self._cleanup_user(training)) + raise TrainingExpiredError( + 'Training handle expired; create a new handle and explicitly load a checkpoint.') + instance = self.instances[training.instance_id] + if instance.state != 'READY': + raise ResourceNotFoundError('Training runtime is unavailable') + if request_id not in training.pending_operation_ids: + training.pending_operation_ids[request_id] = { + 'accepted_at': self.clock(), + 'committed': False, + 'renew': renew + } + await self.store.save_training(training) + await self.store.own_future(request_id, token) + return {'instance_id': instance.instance_id, 'epoch': self.epoch} + + async def commit_operation(self, training_id: str, request_id: str, epoch: str) -> None: + async with self._lock: + await self._assert_epoch() + training = self.trainings.get(training_id) + if epoch != self.epoch or training is None: + raise ResourceNotFoundError('Operation admission no longer exists') + if request_id not in training.pending_operation_ids: + # A very short task can finish while the enqueue acknowledgment + # is in flight. Its original terminal future is conclusive. + future = await self.store.state.get_future(request_id) + if future and future['status'] in ('completed', 'failed', 'cancelled'): + return + raise ResourceNotFoundError('Operation admission no longer exists') + training.pending_operation_ids[request_id]['committed'] = True + await self.store.save_training(training) + + async def finish_operation(self, + training_id: str, + request_id: str, + epoch: str, + *, + executed: bool, + finished_at: float | None = None) -> None: + async with self._lock: + try: + await self._assert_epoch() + except ResourceNotFoundError: + return + training = self.trainings.get(training_id) + if epoch != self.epoch or training is None: + return + ticket = training.pending_operation_ids.pop(request_id, None) + if ticket is None: + return + if executed and ticket['renew']: + age = max(0, time.time() - finished_at) if finished_at is not None else 0 + training.last_activity_at = max(training.last_activity_at, self.clock() - age) + await self.store.save_training(training) + cleanup = training.state == 'CLOSING' and not training.pending_operation_ids + if cleanup: + self._spawn(self._cleanup_user(training)) + + async def route(self, training_id: str, token: str, session_id: str | None = None) -> dict[str, str]: + async with self._lock: + training = await self._authorized(training_id, token, session_id) + return {'instance_id': training.instance_id, 'base_model': training.base_model, 'epoch': self.epoch} + + async def release(self, training_id: str, token: str, session_id: str | None = None) -> None: + async with self._lock: + training = await self._authorized(training_id, token, session_id) + training.state = 'CLOSING' + training.terminal_reason = 'released' + await self.store.save_training(training) + cleanup = not training.pending_operation_ids + if cleanup: + await self._cleanup_user(training) + config = self.catalog.get(training.base_model, training.mode).config.queue_config + deadline = asyncio.get_running_loop().time() + config.absolute_future_ttl(1) + self.config.shutdown_timeout_s + while training.state == 'CLOSING' or self.instances[training.instance_id].state == 'DRAINING': + if self.instances[training.instance_id].state == 'QUARANTINED': + raise RuntimeError('Runtime cleanup is quarantined; device release is not confirmed') + if asyncio.get_running_loop().time() >= deadline: + raise TimeoutError('Training cleanup is still pending; its slot remains reserved') + await asyncio.sleep(config.inline_fast_path_timeout) + + async def _cleanup_user(self, training: TrainingRecord) -> None: + instance = self.instances[training.instance_id] + # Each training must have exactly one cleanup, including concurrent finish callbacks. + async with self._lock: + if (instance.state not in ('READY', 'RESETTING') or training.state != 'CLOSING' + or training.pending_operation_ids or training.training_id in self._cleaning): + return + self._cleaning.add(training.training_id) + resetting = instance.resident and instance.mode == 'full' and not instance.dirty + if resetting: + instance.state = 'RESETTING' + await self.store.save_instance(instance) + closing_full = instance.mode == 'full' and (not instance.resident or instance.dirty) + if closing_full: + instance.state = 'DRAINING' + await self.store.save_instance(instance) + if closing_full: + await self._close_instance(instance) + self._cleaning.discard(training.training_id) + return + try: + spec = self.catalog.get(training.base_model, training.mode) + await asyncio.wait_for( + self.driver.release_training(instance, training), + spec.config.queue_config.absolute_future_ttl(spec.config.num_devices)) + except Exception as exc: + # A timed-out create/cleanup RPC does not prove that shared workers failed. + # Keep this slot reserved until serial cleanup can be confirmed. + if instance.mode == 'lora' and any( + self.trainings[identity].state in ('ACTIVE', + 'PENDING') or self.trainings[identity].pending_operation_ids + for identity in instance.training_ids if identity != training.training_id): + logger.warning('Training cleanup remains pending for %s: %s', training.training_id, exc) + else: + await self.fail_instance(instance.instance_id, failure_for(exc)) + self._cleaning.discard(training.training_id) + return + async with self._lock: + training.state = { + 'expired': 'EXPIRED', + 'released': 'RELEASED', + 'cancelled': 'CANCELLED' + }.get(training.terminal_reason, 'FAILED') + training.terminal_at = time.time() + await self.store.retire(training) + self._cleaning.discard(training.training_id) + if training.training_id in instance.training_ids: + instance.training_ids.remove(training.training_id) + close = not instance.training_ids and (not instance.resident or instance.dirty) + if close: + instance.state = 'DRAINING' + elif resetting: + instance.state = 'READY' + await self.store.save_instance(instance) + if close: + await self._close_instance(instance) + + async def _close_instance(self, instance: InstanceRecord) -> None: + if instance.instance_id in self._closing: + return + self._closing.add(instance.instance_id) + try: + try: + await asyncio.wait_for(self.driver.shutdown(instance), self.config.shutdown_timeout_s) + except Exception as exc: + async with self._lock: + instance.state = 'QUARANTINED' + instance.failure = str(exc) + await self.store.save_instance(instance) + return + async with self._lock: + instance.state = 'STOPPED' + instance.stopped_at = time.time() + await self.store.save_instance(instance) + for training_id in list(instance.training_ids): + training = self.trainings[training_id] + training.pending_operation_ids.clear() + training.state = { + 'expired': 'EXPIRED', + 'released': 'RELEASED', + 'cancelled': 'CANCELLED' + }.get(training.terminal_reason, 'FAILED') + training.terminal_at = time.time() + await self.store.retire(training) + instance.training_ids.clear() + await self.store.save_instance(instance) + finally: + self._closing.discard(instance.instance_id) + + async def fail_instance(self, instance_id: str, failure: FutureFailureRecord) -> None: + async with self._lock: + instance = self.instances[instance_id] + if instance.state in ('DRAINING', 'STOPPED', 'QUARANTINED'): + return + instance.state = 'DRAINING' + instance.failure = failure.message + await self.store.save_instance(instance) + for training_id in instance.training_ids: + training = self.trainings[training_id] + training.state = 'CLOSING' + training.terminal_reason = failure.reason_code + for request_id in (training.create_request_id, *training.pending_operation_ids): + await self.store.future(request_id, 'failed', training_id, failure=failure) + await self.store.save_training(training) + startup = self._startup.get(instance_id) + if startup is not None and startup is not asyncio.current_task() and not startup.done(): + startup.cancel() + await asyncio.gather(startup, return_exceptions=True) + await self._close_instance(instance) + + async def sweep(self) -> None: + async with self._lock: + await self._assert_epoch() + for training in self.trainings.values(): + if training.state != 'ACTIVE' or training.pending_operation_ids: + continue + if await self._expired(training): + training.state = 'CLOSING' + training.terminal_reason = 'expired' + await self.store.save_training(training) + uncertain = [(training, request_id) for training in self.trainings.values() + for request_id, ticket in training.pending_operation_ids.items() + if request_id not in self._reconciling and self.clock() + - ticket['accepted_at'] >= self.config.sweep_interval_s] + self._reconciling.update(request_id for _, request_id in uncertain) + closing = [ + training for training in self.trainings.values() if training.state == 'CLOSING' + and not training.pending_operation_ids and training.training_id not in self._cleaning + ] + for training in closing: + self._spawn(self._cleanup_user(training)) + for training, request_id in uncertain: + self._spawn(self._reconcile_operation(training, request_id)) + for instance in list(self.instances.values()): + if instance.state == 'QUARANTINED': + self._spawn(self._close_instance(instance)) + # Rebuilding one specification must not stop expiry or other models. + for spec in self.catalog.specs(): + residents = sum( + instance.resident and (instance.base_model, + instance.mode) == spec.key and instance.state in LIVE_INSTANCES + for instance in self.instances.values()) + for _ in range(spec.config.min_replicas - residents): + async with self._lock: + instance = await self._new_instance(spec, resident=True) + await self.store.cleanup_receipts() + # Terminal diagnostics expire like their backend tombstones. Repeated + # model creation must not grow controller memory for the server lifetime. + async with self._lock: + cutoff = time.time() - self.store.state.expiration_timeout + for identity, training in list(self.trainings.items()): + if training.state in TERMINAL_TRAINING and training.terminal_at is not None: + if training.terminal_at <= cutoff: + self.trainings.pop(identity) + for identity, instance in list(self.instances.items()): + if (instance.state == 'STOPPED' and instance.stopped_at is not None and instance.stopped_at <= cutoff + and not instance.training_ids): + self.instances.pop(identity) + self._startup.pop(identity, None) + + async def _reconcile_operation(self, training: TrainingRecord, request_id: str) -> None: + try: + instance = self.instances[training.instance_id] + if instance.state != 'READY': + return + result = await asyncio.wait_for( + self.driver.verify_operation(instance, request_id), self.config.shutdown_timeout_s) + if result['state'] == 'pending': + await self.commit_operation(training.training_id, request_id, self.epoch) + elif result['state'] == 'finished': + await self.finish_operation( + training.training_id, + request_id, + self.epoch, + executed=result['executed'], + finished_at=result.get('finished_at')) + elif result['state'] == 'missing': + await self.store.future( + request_id, + 'failed', + training.training_id, + failure=FutureFailureRecord( + reason_code='internal_error', + attribution='server', + message='Operation was not committed to the runtime queue')) + await self.finish_operation(training.training_id, request_id, self.epoch, executed=False) + # An unsafe/unreachable runtime is handled by the worker health monitor. + except Exception: + logger.exception('Operation reconciliation deferred for %s', request_id) + finally: + self._reconciling.discard(request_id) + + async def wait_for_jobs(self) -> None: + while self._jobs: + jobs = tuple(self._jobs) + await asyncio.gather(*jobs, return_exceptions=True) + self._jobs.difference_update(task for task in jobs if task.done()) + + async def shutdown(self) -> None: + self._ready = False + for instance in list(self.instances.values()): + await self.fail_instance( + instance.instance_id, + FutureFailureRecord( + reason_code='orphaned_replica', message='Training service is shutting down', attribution='server')) diff --git a/src/twinkle/server/training/driver.py b/src/twinkle/server/training/driver.py new file mode 100644 index 000000000..cab2139c6 --- /dev/null +++ b/src/twinkle/server/training/driver.py @@ -0,0 +1,172 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Ray/Serve effects behind the controller's small runtime interface.""" +from __future__ import annotations + +import asyncio +import os +import re + +from twinkle.server.launcher.env_propagation import build_propagated_env_vars +from .allocation import DevicePool +from .records import InstanceRecord + + +class RayInstanceDriver: + + def __init__(self, config, store, namespace, controller_name): + self.config, self.store = config, store + self.namespace, self.controller_name = namespace, controller_name + self.pool = DevicePool( + namespace, allocation_timeout=config.allocation_timeout_s, shutdown_timeout=config.shutdown_timeout_s) + + def application_name(self, instance): + return self.pool.name(instance.controller_epoch, instance.instance_id) + + async def start(self, instance, spec): + from ray import serve + from ray.util.placement_group import get_placement_group + + from twinkle.server.model.app import build_model_app + + allocation = await self.pool.allocate(spec, instance.controller_epoch, instance.instance_id) + name = self.application_name(instance) + # Persist ownership before the first worker can be constructed. + instance.application_name = name + instance.placement_group_id = allocation.placement_group_id + instance.node_id, instance.devices = allocation.node_id, list(allocation.devices) + await self.store.save_instance(instance) + env = build_propagated_env_vars() + env.update(spec.model.env_vars) + for key in ('PYTHONPATH', 'TWINKLE_DEFAULT_SAVE_DIR', 'TMPDIR', 'OMP_NUM_THREADS'): + if key in os.environ: + env.setdefault(key, os.environ[key]) + env.update({'TWINKLE_RAY_REGISTRY_NAME': name + '-registry', 'TWINKLE_RAY_NAMESPACE': self.namespace}) + context = { + 'instance_id': instance.instance_id, + 'epoch': instance.controller_epoch, + 'namespace': self.namespace, + 'controller_name': self.controller_name, + 'registry_name': name + '-registry' + } + kwargs = spec.config.model_dump( + exclude={'num_devices', 'device_mesh', 'min_replicas', 'max_replicas', 'idle_timeout_s', 'queue_config'}, + exclude_none=True) + if 'hybrid' in kwargs: + kwargs['hybrid'] = spec.config.hybrid + kwargs['train_mode'] = 'hybrid' if kwargs.get('hybrid') else spec.mode + app = build_model_app( + model_id=spec.model.model_id, + backend=spec.model.backend, + nproc_per_node=spec.config.num_devices, + device_group={ + 'name': 'model', + 'ranks': spec.config.num_devices, + 'device_type': spec.model.device_type + }, + device_mesh={ + **spec.config.device_mesh.model_dump(exclude_none=True), 'device_type': spec.model.device_type + }, + deploy_options={ + 'num_replicas': 1, + 'ray_actor_options': { + 'num_cpus': 0.1, + 'runtime_env': { + 'env_vars': env + } + } + }, + queue_config=spec.config.queue_config, + data_plane_url=spec.model.data_plane_url, + training_context=context, + placement_group=get_placement_group(allocation.name), + allocation_devices=list(allocation.devices), + **kwargs) + # Shield the actual thread so its late return can still be cleaned. + deployment = asyncio.create_task( + asyncio.to_thread( + serve.run, app, name=name, route_prefix=f'/api/v1/training-runtime/{instance.instance_id}')) + try: + await asyncio.shield(deployment) + except asyncio.CancelledError: + + async def clean_late(): + try: + await deployment + except Exception: + pass + await self.shutdown(instance) + + asyncio.create_task(clean_late()) + raise + return { + 'application_name': name, + 'placement_group_id': allocation.placement_group_id, + 'node_id': allocation.node_id, + 'devices': list(allocation.devices) + } + + def handle(self, instance): + from ray import serve + return serve.get_app_handle(instance.application_name or self.application_name(instance)) + + async def create_training(self, instance, training, request): + await self.handle(instance).initialize_training.remote(training.model_dump(), request.model_dump()) + + async def release_training(self, instance, training): + await self.handle(instance).release_training.remote(training.model_dump(), instance.resident + and not instance.dirty) + + async def verify_operation(self, instance, request_id): + return await self.handle(instance).verify_operation.remote(request_id) + + async def shutdown(self, instance: InstanceRecord): + import ray + from ray import serve + from ray.util import list_named_actors + name = self.application_name(instance) + owned = [ + actor for actor in list_named_actors(all_namespaces=True) + if actor['namespace'] == self.namespace and instance.instance_id in actor['name'] + ] + actors = [] + for actor in owned: + try: + handle = ray.get_actor(actor['name'], namespace=self.namespace) + except ValueError: + continue + actors.append(handle) + ray.kill(handle, no_restart=True) + await asyncio.to_thread(serve.delete, instance.application_name or name) + await self.pool.remove(name) + # The PG kills even unnamed remaining children. Named workers must also + # have failed readiness RPCs before the controller retires quota/tickets. + for handle in actors: + try: + await asyncio.wait_for(handle.__ray_ready__.remote(), 5) + except ray.exceptions.RayActorError: + continue + raise RuntimeError('Owned worker death could not be confirmed') + await self.store.state.unregister_replica(instance.instance_id) + + async def health(self, instance): + return await asyncio.wait_for(self.handle(instance).runtime_identity.remote(), 5) + + async def recover_orphans(self): + """Ownership names survive loss of the memory state actor.""" + from ray import serve + from ray.util.placement_group import placement_group_table + pattern = re.compile(r'training-' + re.escape(self.namespace) + + r'-(controller_[0-9a-f]{32})-(instance_[0-9a-f]{32})$') + names = set(serve.status().applications) + names.update(row.get('name', '') for row in placement_group_table().values() if row.get('state') != 'REMOVED') + for name in sorted(names): + match = pattern.fullmatch(name) + if match: + orphan = InstanceRecord( + controller_epoch=match[1], + instance_id=match[2], + application_name=name, + base_model='', + mode='full', + startup_deadline=0) + await asyncio.wait_for(self.shutdown(orphan), self.config.shutdown_timeout_s) diff --git a/src/twinkle/server/training/errors.py b/src/twinkle/server/training/errors.py new file mode 100644 index 000000000..4d7cfc128 --- /dev/null +++ b/src/twinkle/server/training/errors.py @@ -0,0 +1,14 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from twinkle.server.exceptions import RequestRejectedError, ResourceNotFoundError + + +class CapacityUnavailableError(RequestRejectedError): + error_code = 409 + + +class TrainingExpiredError(ResourceNotFoundError): + pass + + +class ProtocolMismatchError(RequestRejectedError): + pass diff --git a/src/twinkle/server/training/records.py b/src/twinkle/server/training/records.py new file mode 100644 index 000000000..4b7130ed4 --- /dev/null +++ b/src/twinkle/server/training/records.py @@ -0,0 +1,58 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Serializable identities and states. Ray handles never enter these records.""" +from __future__ import annotations + +from pydantic import BaseModel, Field +from typing import Any, Literal + +from twinkle.server.state.models import ModelRecord + +TrainingState = Literal['PENDING', 'ACTIVE', 'CLOSING', 'RELEASED', 'EXPIRED', 'FAILED', 'CANCELLED'] +InstanceState = Literal['ALLOCATING', 'STARTING', 'READY', 'RESETTING', 'DRAINING', 'STOPPED', 'QUARANTINED'] +TERMINAL_TRAINING = frozenset({'RELEASED', 'EXPIRED', 'FAILED', 'CANCELLED'}) +LIVE_INSTANCES = frozenset({'ALLOCATING', 'STARTING', 'READY', 'RESETTING'}) + + +class TrainingRecord(ModelRecord): + training_id: str + controller_epoch: str + instance_id: str + mode: Literal['lora', 'full'] + adapter_label: str | None = None + state: TrainingState = 'PENDING' + create_request_id: str + creation_deadline: float + last_activity_at: float + pending_operation_ids: dict[str, dict[str, Any]] = Field(default_factory=dict) + terminal_reason: str | None = None + terminal_at: float | None = None + + +class InstanceRecord(BaseModel): + instance_id: str + controller_epoch: str + base_model: str + mode: Literal['lora', 'full'] + resident: bool = False + state: InstanceState = 'ALLOCATING' + training_ids: list[str] = Field(default_factory=list) + startup_deadline: float + application_name: str | None = None + placement_group_id: str | None = None + node_id: str | None = None + devices: list[str] = Field(default_factory=list) + incarnation: str | None = None + dirty: bool = False + failure: str | None = None + stopped_at: float | None = None + + +class CreateTraining(BaseModel): + base_model: str + mode: Literal['lora', 'full'] + token: str + session_id: str + create_key: str + client_type: Literal['twinkle', 'tinker'] + payload: dict[str, Any] + adapter_label: str | None = None diff --git a/src/twinkle/server/training/router.py b/src/twinkle/server/training/router.py new file mode 100644 index 000000000..fd20c8cc6 --- /dev/null +++ b/src/twinkle/server/training/router.py @@ -0,0 +1,263 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Stable CPU endpoint for one catalog model; requests route by training ID.""" +from __future__ import annotations + +import asyncio +import json +import ray +import uuid +from fastapi import Depends, FastAPI, Request +from fastapi.exceptions import RequestValidationError +from pydantic import ValidationError +from tinker import types as tinker +from types import SimpleNamespace + +from twinkle.hub import HubOperation +from twinkle.protocol.types import model as native +from twinkle.server.checkpoint import _resolve_client_save_dir, create_checkpoint_manager, validate_user_path +from twinkle.server.deployment import LazyCleanupMixin, bind_deployment, build_deployment_app +from twinkle.server.exceptions import RequestRejectedError, ResourceNotFoundError +from twinkle.server.gateway.proxy import ServiceProxy +from twinkle.server.lifecycle.envelope import envelope_from_record +from twinkle.server.middleware.auth import get_session_id_from_request, get_token_from_request +from twinkle.server.state import get_server_state +from twinkle.server.task_queue import TaskQueueMixin +from twinkle.server.validation import assert_request_supported +from .actor import controller_name +from .catalog import TrainingCatalog +from .errors import ProtocolMismatchError +from .records import CreateTraining +from .store import TrainingStore + +NATIVE_REQUESTS = { + 'forward': native.ForwardRequest, + 'forward_only': native.ForwardOnlyRequest, + 'forward_backward': native.ForwardBackwardTaskRequest, + 'calculate_loss': native.AdapterRequest, + 'backward': native.AdapterRequest, + 'forward_from_data_plane': native.DataPlaneForwardRequest, + 'forward_backward_from_data_plane': native.DataPlaneForwardRequest, + 'forward_only_from_data_plane': native.DataPlaneForwardOnlyRequest, + 'step': native.StepRequest, + 'zero_grad': native.AdapterRequest, + 'lr_step': native.LrStepRequest, + 'clip_grad_norm': native.ClipGradNormRequest, + 'clip_grad_and_step': native.ClipGradAndStepRequest, + 'get_train_configs': native.AdapterRequest, + 'set_loss': native.SetLossRequest, + 'set_optimizer': native.SetOptimizerRequest, + 'set_lr_scheduler': native.SetLrSchedulerRequest, + 'set_template': native.SetTemplateRequest, + 'set_processor': native.SetProcessorRequest, + 'add_metric': native.AddMetricRequest, + 'apply_patch': native.ApplyPatchRequest, + 'calculate_metric': native.CalculateMetricRequest, + 'save': native.SaveRequest, + 'load': native.LoadRequest, + 'resume_from_checkpoint': native.ResumeFromCheckpointRequest, + 'remove_adapter': native.AdapterRequest, + 'add_adapter_to_model': native.AddAdapterRequest, + 'upload_to_hub': native.UploadToHubRequest, + 'create': native.CreateRequest, +} +TINKER_REQUESTS = { + 'create_model': tinker.CreateModelRequest, + 'get_info': tinker.GetInfoRequest, + 'unload_model': tinker.UnloadModelRequest, + 'forward': tinker.ForwardRequest, + 'forward_backward': tinker.ForwardBackwardRequest, + 'optim_step': tinker.OptimStepRequest, + 'save_weights': tinker.SaveWeightsRequest, + 'load_weights': tinker.LoadWeightsRequest, + 'save_weights_for_sampler': tinker.SaveWeightsForSamplerRequest, +} + + +class TrainingRouter(TaskQueueMixin, LazyCleanupMixin): + + async def __init__(self, config, model, namespace, http_options): + self.catalog = TrainingCatalog(config) + self.model = model + self.backend = model.backend + self.state = get_server_state() + self.store = TrainingStore(self.state, namespace) + self.controller = ray.get_actor(controller_name(namespace), namespace=namespace) + self.proxy = ServiceProxy(http_options) + self.replica_id = 'training-router_' + uuid.uuid4().hex + self._init_task_queue(None, deployment_name='TrainingRouter') + await self.state.register_replica(self.replica_id, 0) + self._heartbeat = asyncio.create_task(self._keep_alive()) + + async def _keep_alive(self): + while True: + await self.state.touch_replica_last_seen(self.replica_id) + await asyncio.sleep(5) + + async def shutdown(self): + self._heartbeat.cancel() + await asyncio.gather(self._heartbeat, return_exceptions=True) + await self.shutdown_task_queue() + await self.proxy.close() + await self.state.unregister_replica(self.replica_id) + + async def _check_session(self, session): + if not session or await self.state.get_session_last_heartbeat(session) is None: + raise ResourceNotFoundError('Unknown or expired session') + + async def native_request(self, request, endpoint): + schema = NATIVE_REQUESTS.get(endpoint) + if schema is None: + raise ResourceNotFoundError('Unknown training endpoint') + try: + body = schema.model_validate(await request.json()) + except ValidationError as exc: + raise RequestValidationError(exc.errors()) from exc + token = get_token_from_request(request) + if endpoint == 'create': + if not body.training_handles: + raise ProtocolMismatchError('Upgrade Twinkle SDK: training_handles handshake is required') + return native.CreateResponse(training_handles=True) + if endpoint == 'add_adapter_to_model': + if not body.client_create_id: + raise ProtocolMismatchError('Upgrade Twinkle SDK: client_create_id is required') + if not body.adapter_name or not validate_user_path(token, body.adapter_name): + raise RequestRejectedError('Invalid adapter_name') + try: + config = json.loads(body.config) if body.config is not None else None + except (ValueError, TypeError) as exc: + raise RequestRejectedError('Expected a serialized LoRA config') from exc + mode = 'full' if config is None else 'lora' + spec = self.catalog.get(self.model.name, mode) + if config is not None: + rank = config.get('r') if isinstance(config, dict) else None + if (not isinstance(config, dict) or config.get('_TWINKLE_TYPE_') != 'LoraConfig' + or type(rank) is not int or not 0 < rank <= spec.config.max_r): + raise RequestRejectedError('LoRA config/rank is not supported by this training specification') + payload = body.model_dump(mode='json') + if body.save_dir: + try: + payload['save_dir'] = _resolve_client_save_dir(body.save_dir).as_posix() + except ValueError as exc: + raise RequestRejectedError(str(exc)) from exc + session = get_session_id_from_request(request) + await self._check_session(session) + ref = await self.controller.create.remote( + CreateTraining( + base_model=self.model.name, + mode=mode, + token=token, + session_id=session, + create_key=body.client_create_id, + client_type='twinkle', + adapter_label=body.adapter_name, + payload=payload).model_dump()) + return envelope_from_record(ref['request_id'], await self.state.get_future(ref['request_id'])) + if endpoint == 'upload_to_hub': + + async def upload(): + checkpoint = body.checkpoint_dir + if checkpoint.startswith('twinkle://'): + checkpoint = create_checkpoint_manager( + token, client_type='twinkle').resolve_checkpoint_dir(checkpoint) + await self.call_backend( + HubOperation.push_to_hub, + repo_id=body.hub_model_id, + folder_path=checkpoint, + token=body.hub_token or token, + private=True) + + ref = await self.schedule_background_task(upload, task_type=endpoint) + await self.store.own_future(ref['request_id'], token) + return envelope_from_record(ref['request_id'], await self.state.get_future(ref['request_id'])) + if not body.training_id: + raise ProtocolMismatchError('Upgrade Twinkle SDK: training_id is required') + session = get_session_id_from_request(request) + await self._check_session(session) + if endpoint == 'remove_adapter': + return await self._release(body.training_id, token, session, 'twinkle') + capability = endpoint if endpoint in ('forward', 'backward', 'calculate_loss') else None + assert_request_supported(SimpleNamespace(backend=self.backend), body, capability=capability) + route = await self.controller.route.remote(body.training_id, token, session) + if route['base_model'] != self.model.name: + raise ResourceNotFoundError('Training handle belongs to another model endpoint') + return await self.proxy.proxy_request(request, f'twinkle/{endpoint}', route['instance_id'], 'training-runtime') + + async def tinker_request(self, request, endpoint): + schema = TINKER_REQUESTS.get(endpoint) + if schema is None: + raise ResourceNotFoundError('Unknown training endpoint') + try: + body = schema.model_validate(await request.json()) + except ValidationError as exc: + raise RequestValidationError(exc.errors()) from exc + token = get_token_from_request(request) + if endpoint == 'create_model': + if body.base_model != self.model.name: + raise RequestRejectedError('Model does not match the catalog endpoint') + mode = 'lora' if body.lora_config is not None else 'full' + spec = self.catalog.get(self.model.name, mode) + if body.lora_config is not None and not 0 < body.lora_config.rank <= spec.config.max_r: + raise RequestRejectedError('LoRA rank exceeds the training specification') + await self._check_session(body.session_id) + return await self.controller.create.remote( + CreateTraining( + base_model=self.model.name, + mode=mode, + token=token, + session_id=body.session_id, + create_key=str(body.model_seq_id), + client_type='tinker', + payload=body.model_dump(mode='json')).model_dump()) + if endpoint == 'unload_model': + return await self._release(body.model_id, token, None, 'tinker') + route = await self.controller.route.remote(body.model_id, token) + if route['base_model'] != self.model.name: + raise ResourceNotFoundError('Training handle belongs to another model endpoint') + return await self.proxy.proxy_request(request, f'tinker/{endpoint}', route['instance_id'], 'training-runtime') + + async def _release(self, training_id, token, session, client_type): + await self.controller.route.remote(training_id, token, session) + + async def release(): + await self.controller.release.remote(training_id, token, session) + return {'model_id': training_id} if client_type == 'tinker' else None + + ref = await self.schedule_background_task(release, task_type='release_training') + await self.store.own_future(ref['request_id'], token) + if client_type == 'tinker': + return {'request_id': ref['request_id']} + return envelope_from_record(ref['request_id'], await self.state.get_future(ref['request_id'])) + + +def build_training_router(config, model, namespace, http_options, deploy_options): + + def routes(app: FastAPI, get_self): + + @app.get('/healthz') + async def health(self=Depends(get_self)): + diagnostic = await self.controller.diagnostics.remote(self.model.name) + active = [item for item in diagnostic['instances'] if item['state'] != 'STOPPED'] + return { + 'healthy': diagnostic['ready'] and not any(item['state'] == 'QUARANTINED' for item in active), + 'status': 'unloaded' if not active else 'loaded', + 'instances': active + } + + @app.post('/twinkle/{endpoint}') + async def native_route(request: Request, endpoint: str, self=Depends(get_self)): + return await self.native_request(request, endpoint) + + @app.post('/tinker/{endpoint}') + async def tinker_route(request: Request, endpoint: str, self=Depends(get_self)): + return await self.tinker_request(request, endpoint) + + async def shutdown(servable): + await servable.shutdown() + + app = build_deployment_app('TrainingRouter', routes, on_shutdown=shutdown) + return bind_deployment( + app, + TrainingRouter, + deploy_options, + deployment_name='TrainingRouter', + bind_args=(config, model, namespace, http_options)) diff --git a/src/twinkle/server/training/runtime.py b/src/twinkle/server/training/runtime.py new file mode 100644 index 000000000..455ccf040 --- /dev/null +++ b/src/twinkle/server/training/runtime.py @@ -0,0 +1,214 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Managed training operations on the existing model queue and backend gate.""" +from __future__ import annotations + +import asyncio +import os +import uuid +from contextvars import ContextVar +from datetime import datetime + +from twinkle.protocol.serialize import deserialize_object +from twinkle.protocol.types import model as native +from twinkle.protocol.types.training import CreateModelRequest +from twinkle.protocol.types.training import LoraConfig as RunLoraConfig +from twinkle.server.checkpoint import create_training_run_manager +from twinkle.server.exceptions import ResourceNotFoundError +from twinkle.server.lifecycle.submit import backend_kwargs +from twinkle.server.middleware.auth import get_session_id_from_request, get_token_from_request +from twinkle.server.state.models import FutureFailureRecord +from twinkle.server.utils import get_template_for_model +from twinkle.utils.logger import get_logger +from .records import CreateTraining, TrainingRecord + +_tinker_operation = ContextVar('managed_tinker_operation', default=None) +logger = get_logger() + + +class ManagedTrainingRuntime: + + async def wait_for_training_workers(self): + # Constructing a remote_class returns actor handles before its worker + # constructors (including model loading) finish. Full creation otherwise + # performs no backend call to establish this readiness barrier. + actors = getattr(self.model, '_actors', None) or [] + await asyncio.gather(*(actor.__ray_ready__.remote() for actor in actors)) + + async def initialize_managed_runtime(self, context, placement_group, devices): + import ray + + from twinkle.infra._ray.ray_helper import RayHelper + from twinkle.infra._ray.resource_manager import ResourceManager + self._training_controller = ray.get_actor(context['controller_name'], namespace=context['namespace']) + await self._training_controller.register_runtime.remote(context['instance_id'], context['epoch'], + uuid.uuid4().hex) + os.environ['TWINKLE_RAY_REGISTRY_NAME'] = context['registry_name'] + os.environ['TWINKLE_RAY_NAMESPACE'] = context['namespace'] + RayHelper.resource_manager = ResourceManager.from_allocation(self.device_group, placement_group, devices) + + async def managed_request_start(self, request): + token = get_token_from_request(request) + if self._model_unhealthy: + raise ResourceNotFoundError('Training runtime is unavailable') + body = await request.json() + training_id = body.get('training_id') or body.get('model_id') + session = get_session_id_from_request(request) if '/twinkle/' in request.url.path else None + if '/twinkle/' in request.url.path and not session: + raise ResourceNotFoundError('Training requests require the creating session') + route = await self._training_controller.route.remote(training_id, token, session) + if route['instance_id'] != self.training_context['instance_id'] or route['epoch'] != self.training_context[ + 'epoch']: + raise ResourceNotFoundError('Runtime does not own this training handle') + _tinker_operation.set((training_id, body.get('seq_id')) if '/tinker/' in request.url.path else None) + return token + + async def schedule_task(self, coro_factory, **kwargs): + operation = _tinker_operation.get() + if operation is None or operation[1] is None: + return await super().schedule_task(coro_factory, **kwargs) + training_id, seq_id = operation + task_type = kwargs.get('task_type') + stream = f'tinker::{task_type}' + request_id = 'req_' + uuid.uuid4().hex + retention = self.sequence_retention_s + previous = await self.state.claim_training_seq(training_id, stream, seq_id, request_id, retention) + if previous is not None: + return {'request_id': previous, 'model_id': training_id} + try: + return await super().schedule_task(coro_factory, request_id=request_id, **kwargs) + except Exception: + if await self.state.get_future(request_id) is None: + await self.state.release_training_seq(training_id, stream, seq_id, request_id) + raise + + async def _prepare_queued_task(self, task): + if task.model_id is None: + return + context = self.training_context + await self._training_controller.admit.remote( + task.model_id, + task.token, + None, + task.request_id, + renew=task.task_type not in ('get_info', 'get_train_configs')) + + async def enqueued(): + await self._training_controller.commit_operation.remote(task.model_id, task.request_id, context['epoch']) + + async def finished(executed, unsafe): + if unsafe: + self._model_unhealthy = True + self._training_controller.fail_instance.remote( + context['instance_id'], + FutureFailureRecord( + reason_code='execution_timeout', + attribution='server', + message='Backend execution timed out; runtime is being terminated.')) + else: + try: + await self._training_controller.finish_operation.remote( + task.model_id, task.request_id, context['epoch'], executed=executed) + except Exception: + # The terminal future and queue membership allow the controller + # to reconcile a lost notification without stopping compute. + logger.exception('Training completion notification failed for %s', task.request_id) + + task.on_enqueued, task.on_finished = enqueued, finished + if task.task_type == 'apply_patch': + await self._training_controller.mark_dirty.remote(context['instance_id'], context['epoch']) + + async def _managed_serial(self, factory, task_type): + ref = await self.schedule_task(factory, task_type=task_type) + deadline = asyncio.get_running_loop().time() + self.task_queue_config.absolute_future_ttl(1) + while True: + future = await self.state.get_future(ref['request_id']) + if future and future['status'] in ('failed', 'cancelled'): + raise RuntimeError(future['failure']['message']) + if future and future['status'] == 'completed': + return future['result'] + if asyncio.get_running_loop().time() >= deadline: + raise TimeoutError('Training lifecycle operation has not finished') + await asyncio.sleep(self.task_queue_config.inline_fast_path_timeout) + + async def initialize_training(self, training_data, request_data): + training = TrainingRecord.model_validate(training_data) + request = CreateTraining.model_validate(request_data) + + async def initialize(): + adapter = training.training_id + backend_adapter = self.resolve_model_adapter_name(adapter) + self.register_resource(adapter, training.token, training.session_id) + self.set_resource_state(adapter, 'grad_ready', False) + if request.client_type == 'twinkle': + body = native.AddAdapterRequest.model_validate(request.payload) + config = deserialize_object(body.config) + if not self.is_full_mode: + await self.call_backend(self.model.add_adapter_to_model, adapter, config, **backend_kwargs(body)) + lora = None if config is None else RunLoraConfig( + rank=config.r, train_unembed=False, train_mlp=True, train_attn=True) + run_config = CreateModelRequest( + base_model=training.base_model, + lora_config=lora, + save_dir=body.save_dir, + user_metadata={'adapter_name': training.adapter_label}) + else: + from peft import LoraConfig + from tinker import types + run_config = types.CreateModelRequest.model_validate(request.payload) + if not self.is_full_mode: + await self.call_backend( + self.model.add_adapter_to_model, + adapter_name=adapter, + config_or_dir=LoraConfig(r=run_config.lora_config.rank, target_modules='all-linear')) + await self.call_backend( + self.model.set_template, + get_template_for_model(self.base_model), + adapter_name=backend_adapter, + model_id=self.base_model) + await self.call_backend(self.model.set_processor, 'InputProcessor', adapter_name=backend_adapter) + await self.call_backend(self.model.set_optimizer, 'Adam', adapter_name=backend_adapter) + create_training_run_manager(training.token, client_type=request.client_type).save(adapter, run_config) + + await self._managed_serial(initialize, 'initialize_training') + + async def release_training(self, training_data, reset_full): + training = TrainingRecord.model_validate(training_data) + + async def release(): + adapter = training.training_id + if self.get_resource_info(adapter) is None: + return + if self.is_full_mode: + if reset_full: + await self.call_backend(self.model.reload_initial_weights, reset_training_state=True) + else: + await self.call_backend(self.model.remove_adapter, adapter) + self.clear_resource_state(adapter) + self.unregister_resource(adapter) + + await self._managed_serial(release, 'release_training') + + async def runtime_identity(self): + import ray + actors = getattr(self.model, '_actors', None) or [] + states = await asyncio.to_thread( + lambda: [ray._private.state.actors(actor._actor_id.hex()).get('State') for actor in actors]) + return { + 'pid': os.getpid(), + 'instance_id': self.training_context['instance_id'], + 'unhealthy': self._model_unhealthy or any(state != 'ALIVE' for state in states) + } + + async def verify_operation(self, request_id): + if self._model_unhealthy: + return {'state': 'unsafe'} + if self._compute_worker.has_task(request_id): + return {'state': 'pending'} + future = await self.state.get_future(request_id) + if future and future['status'] in ('completed', 'failed', 'cancelled'): + return { + 'state': 'finished', + 'executed': future['status'] != 'cancelled', + 'finished_at': datetime.fromisoformat(future['updated_at']).timestamp() + } + return {'state': 'missing'} diff --git a/src/twinkle/server/training/store.py b/src/twinkle/server/training/store.py new file mode 100644 index 000000000..f44f1a7f4 --- /dev/null +++ b/src/twinkle/server/training/store.py @@ -0,0 +1,101 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Controller-owned records on the existing shared state backend.""" +from __future__ import annotations + +import hashlib +import json +import time +from typing import Any + +from twinkle.server.state import ServerState +from twinkle.server.state.models import FutureFailureRecord +from .records import CreateTraining, InstanceRecord, TrainingRecord + + +class TrainingStore: + + def __init__(self, state: ServerState, namespace: str): + self.state = state + self.backend = state.backend + self.prefix = f'training::{namespace}::' + + def key(self, kind: str, identity: str) -> str: + return f'{self.prefix}{kind}::{identity}' + + @staticmethod + def create_identity(request: CreateTraining) -> tuple[str, str]: + identity = json.dumps([request.client_type, request.token, request.session_id, request.create_key]) + fingerprint = json.dumps(request.model_dump(mode='json'), sort_keys=True, separators=(',', ':')) + return hashlib.sha256(identity.encode()).hexdigest(), hashlib.sha256(fingerprint.encode()).hexdigest() + + async def reserve(self, record: TrainingRecord) -> None: + # ModelManager applies the existing global quota atomically, including PENDING. + await self.state.register_model_record(record.training_id, record) + + async def save_training(self, record: TrainingRecord) -> None: + await self.backend.set(f'model::{record.training_id}', record.model_dump()) + + async def retire(self, record: TrainingRecord) -> None: + # Keep a bounded diagnostic tombstone after freeing the live quota row. + await self.backend.set( + self.key('terminal', record.training_id), + record.model_dump(), + ttl=max(1, int(self.state.expiration_timeout))) + await self.state.unload_model(record.training_id) + for pattern in (f'seq::*::{record.training_id}::*', f'tinker_seq::{record.training_id}::*', + f'seq_ledger::{record.training_id}::*'): + for key in await self.backend.keys(pattern): + await self.backend.delete(key) + + async def save_instance(self, record: InstanceRecord) -> None: + ttl = max(1, int(self.state.expiration_timeout)) if record.state == 'STOPPED' else None + await self.backend.set(self.key('instance', record.instance_id), record.model_dump(), ttl=ttl) + + async def instances(self) -> list[InstanceRecord]: + keys = await self.backend.keys(self.key('instance', '*')) + return [InstanceRecord.model_validate(value) for value in await self.backend.mget(keys) if value is not None] + + async def trainings(self) -> list[TrainingRecord]: + keys = await self.backend.keys('model::*') + records = await self.backend.mget(keys) + return [ + TrainingRecord.model_validate(value) for value in records if isinstance(value, dict) + and value.get('controller_epoch') and value.get('training_namespace') == self.prefix + ] + + async def get_create(self, identity: str) -> dict[str, str] | None: + return await self.backend.get(self.key('create', identity)) + + async def save_create(self, identity: str, fingerprint: str, request_id: str, session_id: str) -> None: + # Heartbeats can keep a session alive longer than expiration_timeout. + # Keep its creation receipts until session cleanup, never recreate on retry. + await self.backend.set( + self.key('create', identity), { + 'fingerprint': fingerprint, + 'request_id': request_id, + 'session_id': session_id + }) + + async def cleanup_receipts(self) -> None: + for key in await self.backend.keys(self.key('create', '*')): + receipt = await self.backend.get(key) + if receipt and await self.state.get_session_last_heartbeat(receipt['session_id']) is None: + await self.backend.delete(key) + + async def future(self, + request_id: str, + status: str, + model_id: str | None, + *, + epoch: str | None = None, + result: Any = None, + failure: FutureFailureRecord | None = None, + deadline: float | None = None) -> None: + await self.state.store_future_status( + request_id, status, model_id, result=result, failure=failure, replica_id=epoch, absolute_deadline=deadline) + + async def own_future(self, request_id: str, token: str) -> None: + await self.backend.set(f'future_owner::{request_id}', token, ttl=max(1, int(self.state.expiration_timeout))) + + async def begin_epoch(self, epoch: str) -> None: + await self.backend.set(self.key('controller', 'active'), {'epoch': epoch, 'started_at': time.time()}) diff --git a/src/twinkle_client/auto/agent/server_tools.py b/src/twinkle_client/auto/agent/server_tools.py index 09786c274..b8f11333f 100644 --- a/src/twinkle_client/auto/agent/server_tools.py +++ b/src/twinkle_client/auto/agent/server_tools.py @@ -373,8 +373,6 @@ def _generate_server_config( def _short(mid: str) -> str: return mid.split('/')[-1] if '/' in mid else mid - model_short = _short(model_id) - # Collect all model IDs for supported_models all_model_ids = [model_id] + [s['model_id'] for s in sampler_list] @@ -409,62 +407,41 @@ def _short(mid: str) -> str: }], }) - # 2. Build GPU-requiring applications (model + samplers), + # 2. Build static GPU-requiring sampler applications, # then sort by GPU count DESCENDING before appending. # Largest PG deploys first → it has the fewest node choices → # avoids GPU scheduling deadlock on single-machine multi-node. gpu_apps: list[tuple[int, dict]] = [] # (gpu_count, app_config) - # 2a. Training model worker (student) - gpu_apps.append(( - train_gpus, - { - 'name': - f'models-{model_short}', - 'route_prefix': - f'/api/v1/model/{model_id}', - 'import_path': - 'model', - 'args': { - 'backend': backend, - 'model_id': f'ms://{model_id}', - 'max_length': 500000, # total tokens per forward pass (must match max_input_tokens) - 'nproc_per_node': train_gpus, - 'device_group': { - 'name': 'model', - 'ranks': train_gpus, - 'device_type': 'cuda', - }, - 'device_mesh': { - 'device_type': 'cuda', - 'dp_size': train_gpus, - }, - 'queue_config': { - 'rps_limit': 100, - 'tps_limit': 100000, - 'max_input_tokens': 500000, - }, - 'adapter_config': { - 'adapter_timeout': 600, - }, + applications.append({ + 'name': f'model-{_short(model_id)}', + 'route_prefix': f'/api/v1/model/{model_id}', + 'import_path': 'model', + 'args': { + 'model_id': f'ms://{model_id}', + 'backend': backend, + 'device_type': 'cuda', + 'env_vars': { + 'TWINKLE_TRUST_REMOTE_CODE': '1' }, - 'deployments': [{ - 'name': 'ModelManagement', - 'autoscaling_config': { - 'min_replicas': 1, - 'max_replicas': 1, - 'target_ongoing_requests': 16, - }, - 'ray_actor_options': { - 'num_cpus': 0.1, - 'runtime_env': { - 'env_vars': { - 'TWINKLE_TRUST_REMOTE_CODE': '1' - }, + 'modes': { + 'lora': { + 'num_devices': train_gpus, + 'device_mesh': { + 'device_type': 'cuda', + 'dp_size': train_gpus }, - }, - }], - })) + 'max_length': 500000, + 'idle_timeout_s': 600, + 'queue_config': { + 'rps_limit': 100, + 'tps_limit': 100000, + 'max_input_tokens': 500000 + }, + } + }, + }, + }) # 2b. Sampler/teacher models sampler_name_count: dict[str, int] = {} diff --git a/src/twinkle_client/model/multi_lora_transformers.py b/src/twinkle_client/model/multi_lora_transformers.py index a9fbfd8ed..b445aca5c 100644 --- a/src/twinkle_client/model/multi_lora_transformers.py +++ b/src/twinkle_client/model/multi_lora_transformers.py @@ -3,6 +3,7 @@ import itertools import logging import threading +import uuid from collections.abc import Mapping from pathlib import Path from typing import TYPE_CHECKING, Any, Dict, Optional @@ -58,6 +59,9 @@ def __init__( self.model_id = model_id self.server_url = f'{self._transport.context.base_url}/model/{model_id}/twinkle' self.adapter_name = None + self._training_ids: dict[str, str] = {} + self._creating: set[str] = set() + self._training_lock = threading.Lock() # Per-client monotonic sequence for idempotent dedup of stateful training ops: # the server dedups on (session_id, seq_id) so a retried grad/step call is # applied at most once. Reserved once per call and reused on retry. @@ -84,12 +88,21 @@ def _ensure_created(self) -> None: with self._create_lock: if self._created: return - self._transport.post(f'{self.server_url}/create') + response = self._transport.post(f'{self.server_url}/create', json_data={'training_handles': True}) + if response.json().get('training_handles') is not True: + raise RuntimeError('Server does not support training handles; upgrade Twinkle server and SDK together.') self._created = True def _submit(self, endpoint: str, model_cls, response_cls, **values): """Build, send, and resolve one twinkle-native request.""" self._ensure_created() + if issubclass(model_cls, model_types.TrainingRequest): + name = values.get('adapter_name') + with self._training_lock: + training_id = values.get('training_id') or self._training_ids.get(name) + if training_id is None: + raise RuntimeError('Create this training with add_adapter_to_model on this object before using it.') + values['training_id'] = training_id body = build_request(model_cls, **values) response = self._transport.post_model(f'{self.server_url}/{endpoint}', body) return self._await_task(response, response_cls) @@ -132,22 +145,38 @@ def add_adapter_to_model( save_dir = kwargs.pop('save_dir', None) if save_dir: save_dir = Path(save_dir).expanduser().resolve().as_posix() - self._submit( - 'add_adapter_to_model', - model_types.AddAdapterRequest, - None, - adapter_name=adapter_name, - config=config, - save_dir=save_dir, - **kwargs) - self.adapter_name = adapter_name + with self._training_lock: + if adapter_name in self._creating: + raise RuntimeError('A training with this adapter_name is already being created on this object.') + self._creating.add(adapter_name) + try: + result = self._submit( + 'add_adapter_to_model', + model_types.AddAdapterRequest, + model_types.TrainingHandleResponse, + client_create_id=uuid.uuid4().hex, + adapter_name=adapter_name, + config=config, + save_dir=save_dir, + **kwargs) + with self._training_lock: + self._training_ids[adapter_name] = result.training_id + self.adapter_name = adapter_name + finally: + with self._training_lock: + self._creating.discard(adapter_name) def remove_adapter(self, adapter_name: str | None = None) -> None: """Release one client-owned adapter from the training component.""" name = adapter_name or self.adapter_name - self._submit('remove_adapter', model_types.AdapterRequest, None, adapter_name=name) - if name == self.adapter_name: - self.adapter_name = None + with self._training_lock: + training_id = self._training_ids.get(name) + self._submit('remove_adapter', model_types.AdapterRequest, None, adapter_name=name, training_id=training_id) + with self._training_lock: + if self._training_ids.get(name) == training_id: + self._training_ids.pop(name, None) + if name == self.adapter_name: + self.adapter_name = None # ------------------------------------------------------------------ # # Inline forward family diff --git a/tests/server/cli/test_cli.py b/tests/server/cli/test_cli.py index a49ae372e..ef859abca 100644 --- a/tests/server/cli/test_cli.py +++ b/tests/server/cli/test_cli.py @@ -89,5 +89,5 @@ def test_env_var_overrides_when_flag_omitted(monkeypatch) -> None: def test_example_config_loads_via_server_config() -> None: cfg = ServerConfig.from_yaml(EXAMPLE) assert isinstance(cfg, ServerConfig) - assert any(a.import_path == 'model' for a in cfg.applications) + assert cfg.training_config.models[0].backend == 'mock' assert any(a.import_path == 'sampler' for a in cfg.applications) diff --git a/tests/server/config/test_server_config.py b/tests/server/config/test_server_config.py index 927089ff2..8e47461b0 100644 --- a/tests/server/config/test_server_config.py +++ b/tests/server/config/test_server_config.py @@ -33,38 +33,25 @@ }), ) -_MODEL_APP = st.fixed_dictionaries({ - 'name': - st.just('m'), - 'route_prefix': - st.just('/api/v1/m'), - 'import_path': - st.just('model'), - 'args': - st.fixed_dictionaries({ +_TRAINING_MODEL = st.fixed_dictionaries({ + 'name': st.just('m'), 'import_path': st.just('model'), 'route_prefix': st.just('/api/v1/model/m'), + 'args': st.fixed_dictionaries({ 'model_id': st.just('model-id'), - 'device_group': st.just({ - 'name': 'g', - 'ranks': 1, - 'device_type': 'CPU' - }), - 'device_mesh': st.just({ - 'device_type': 'CPU', - 'dp_size': 1 - }), 'backend': st.sampled_from(['mock', 'transformers', 'megatron']), + 'device_type': st.just('cuda'), + 'modes': st.just({'lora': {'num_devices': 1, 'device_mesh': {'dp_size': 1}}}), }), }) _VALID_CONFIG = st.fixed_dictionaries({ 'persistence': _PERSISTENCE_VARIANTS, - 'applications': st.lists(_MODEL_APP, min_size=0, max_size=3), + 'applications': st.lists(_TRAINING_MODEL, min_size=0, max_size=1), }) # ---------- valid → fully validated ---------------------------------------- # -@settings(max_examples=100) +@settings(max_examples=100, deadline=None) @given(payload=_VALID_CONFIG) def test_valid_payload_yields_full_instance(payload: dict) -> None: cfg = ServerConfig.model_validate(payload) @@ -141,7 +128,7 @@ def test_sampler_docstring_values_match_selector() -> None: # ---------- round-trip fidelity ----------------------------------------- # -@settings(max_examples=100) +@settings(max_examples=100, deadline=None) @given(payload=_VALID_CONFIG) def test_round_trip_fidelity(payload: dict) -> None: cfg = ServerConfig.model_validate(payload) diff --git a/tests/server/config/test_training_config.py b/tests/server/config/test_training_config.py new file mode 100644 index 000000000..0a313079c --- /dev/null +++ b/tests/server/config/test_training_config.py @@ -0,0 +1,187 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +import copy +import math + +import pytest +from pydantic import ValidationError + +from twinkle.server.config import ServerConfig +from twinkle.server.config.training import TrainingConfig +from twinkle.server.exceptions import RequestRejectedError +from twinkle.server.training.catalog import TrainingCatalog + + +def model_config(): + return { + 'name': 'Qwen/small', 'model_id': 'ms://Qwen/small', 'backend': 'transformers', 'device_type': 'cuda', + 'modes': { + 'lora': {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}}, + 'full': {'num_devices': 4, 'device_mesh': {'fsdp_size': 4}}, + }, + } + + +def model_application(model=None): + model = dict(model or model_config()) + name = model.pop('name') + return {'name': 'model-' + name, 'route_prefix': '/api/v1/model/' + name, + 'import_path': 'model', 'args': model} + + +def test_catalog_defaults_and_mode_selection(): + config = ServerConfig(applications=[model_application()]) + catalog = TrainingCatalog(config.training_config) + shared = catalog.get('Qwen/small', 'lora') + exclusive = catalog.get('Qwen/small', 'full') + assert shared.capacity == 5 + assert exclusive.capacity == 1 + assert shared.config.min_replicas == exclusive.config.min_replicas == 0 + assert shared.config.idle_timeout_s == 1800 + assert config.training.allocation_timeout_s == 3 + assert config.to_yaml_dict()['applications'][0]['args']['model_id'] == 'ms://Qwen/small' + with pytest.raises(RequestRejectedError, match='Unknown training model'): + catalog.get('not-in-catalog', 'lora') + + +def test_existing_yaml_retains_heartbeat_timeout_and_max_lifetime(): + from pathlib import Path + paths = sorted(Path('cookbook/client/server').glob('**/server_config*.yaml')) + assert paths + for path in paths: + config = ServerConfig.from_yaml(path) + for app in config.applications: + if app.import_path != 'model': + continue + normalized = config.training_config.models[ + [item.model_id for item in config.training_config.models].index(app.args.model_id)] + mode = normalized.modes.full if app.args.train_mode == 'full' else normalized.modes.lora + legacy = app.args.adapter_config or {} + assert mode.idle_timeout_s == app.args.idle_timeout_s + assert normalized.legacy_expiration.session_timeout_s == legacy.get('adapter_timeout', 1800) + assert normalized.legacy_expiration.max_lifetime_s == legacy.get('adapter_max_lifetime', 36000) + + +@pytest.mark.parametrize('field,value', [ + ('num_devices', 0), ('num_devices', True), ('num_devices', 3), + ('min_replicas', -1), ('idle_timeout_s', math.inf), ('idle_timeout_s', math.nan), + ('idle_timeout_s', 0), ('ranks', [0, 1]), ('max_loras', 0), +]) +def test_invalid_mode_cannot_reach_allocation(field, value): + model = model_config() + model['modes']['lora'][field] = value + with pytest.raises(ValidationError): + TrainingConfig(models=[model]) + + +def test_ep_and_vpp_do_not_multiply_physical_device_count(): + model = model_config() + model['backend'] = 'megatron' + model['modes']['lora'] = { + 'num_devices': 8, + 'device_mesh': {'tp_size': 2, 'dp_size': 2, 'pp_size': 2, 'ep_size': 2, 'vpp_size': 2}, + } + assert TrainingConfig(models=[model]).models[0].modes.lora.num_devices == 8 + + +@pytest.mark.parametrize('mutation', ['duplicate', 'no_modes', 'full_lora_field', 'hybrid_megatron', 'mesh_type', 'min_max']) +def test_cross_field_constraints(mutation): + model = model_config() + models = [model] + if mutation == 'duplicate': + models.append(copy.deepcopy(model)) + elif mutation == 'no_modes': + model['modes'] = {} + elif mutation == 'full_lora_field': + model['modes']['full']['max_loras'] = 2 + elif mutation == 'hybrid_megatron': + model['backend'] = 'megatron' + model['modes']['lora']['hybrid'] = {'allocation_path': '/tmp/allocation.json'} + elif mutation == 'mesh_type': + model['modes']['lora']['device_mesh']['device_type'] = 'npu' + elif mutation == 'min_max': + model['modes']['lora'].update(min_replicas=2, max_replicas=1) + with pytest.raises(ValidationError): + TrainingConfig(models=models) + + +def test_full_only_is_explicit_and_lora_is_unavailable(): + model = model_config() + del model['modes']['lora'] + catalog = TrainingCatalog(TrainingConfig(models=[model])) + assert catalog.get('Qwen/small', 'full').capacity == 1 + with pytest.raises(RequestRejectedError, match='does not provide lora'): + catalog.get('Qwen/small', 'lora') + + +def test_training_namespace_conflicts_are_rejected(): + with pytest.raises(ValidationError, match='unique'): + ServerConfig(applications=[model_application(), model_application()]) + with pytest.raises(ValidationError, match='route conflicts'): + ServerConfig(applications=[model_application(), { + 'name': 'processor', 'import_path': 'processor', 'route_prefix': '/api/v1/model/Qwen/small', + }]) + with pytest.raises(ValidationError, match='models'): + ServerConfig(training={'models': [model_config()]}) + + +def test_existing_yaml_examples_keep_their_model_entries(monkeypatch): + from pathlib import Path + monkeypatch.setenv('TWINKLE_LOCAL_MODEL_PATH', '/tmp/model-cache') + paths = list(Path('cookbook/client/server').rglob('server_config*.yaml')) + paths += list(Path('cookbook/client/async_rl').glob('server_config.yaml')) + paths += list(Path('tests/server/config').glob('server_config*e2e*.yaml')) + paths += list(Path('tests/server/fixtures').glob('server_config_mock*.yaml')) + assert len(paths) == 9 + for path in paths: + config = ServerConfig.from_yaml(path) + models = [app for app in config.applications if app.import_path == 'model'] + assert len(models) == len(config.training_config.models) == 1 + model = config.training_config.models[0] + assert model.model_id == models[0].args.model_id + assert model.route_prefix == models[0].route_prefix + mode = model.modes.full or model.modes.lora + assert mode.min_replicas == 1 + assert ServerConfig.model_validate(config.to_yaml_dict()).training_config == config.training_config + + +def test_legacy_lifecycle_overrides_and_full_instances(): + app = {'name': 'model', 'route_prefix': '/api/v1/model/small', 'import_path': 'model', 'args': { + 'model_id': 'weights', 'backend': 'mock', 'train_mode': 'full', + 'device_group': {'ranks': 2, 'device_type': 'CPU'}, 'device_mesh': {'dp_size': 2}, + 'min_replicas': 0, 'max_replicas': 3, 'idle_timeout_s': 17, + 'adapter_config': {'adapter_timeout': 30}, + }, 'deployments': [{'num_replicas': 1}]} + config = ServerConfig(applications=[app]) + mode = config.training_config.models[0].modes.full + assert (mode.num_devices, mode.min_replicas, mode.max_replicas, mode.idle_timeout_s) == (2, 0, 3, 17) + assert config.training_config.models[0].modes.lora is None + + +def test_extended_modes_do_not_accept_ambiguous_single_mode_settings(): + app = model_application() + app['args']['train_mode'] = 'full' + with pytest.raises(ValidationError, match='cannot be mixed'): + ServerConfig(applications=[app]) + + +def test_launcher_keeps_training_and_sampling_capabilities_separate(monkeypatch): + from ray import serve + from twinkle.server.launcher.server_launcher import ServerLauncher + config = ServerConfig(applications=[model_application(), { + 'name': 'gateway', 'import_path': 'server', 'route_prefix': '/api/v1', + }, { + 'name': 'sampler', 'import_path': 'sampler', 'route_prefix': '/api/v1/sampler/teacher', + 'args': {'model_id': 'teacher', 'device_group': {}, 'device_mesh': {}, 'sampler_type': 'mock'}, + }]) + captured = {} + def builder(**kwargs): + captured.update(kwargs) + return 'app' + launcher = ServerLauncher(config) + monkeypatch.setattr(launcher, '_resolve_builder', lambda path: builder) + monkeypatch.setattr(serve, 'run', lambda *args, **kwargs: None) + launcher._deploy_application(config.applications[1]) + assert captured['training_models'] == ['Qwen/small'] + assert captured['supported_models'] == ['teacher'] diff --git a/tests/server/contract/test_error_wire.py b/tests/server/contract/test_error_wire.py index 41b394e42..d2ae84862 100644 --- a/tests/server/contract/test_error_wire.py +++ b/tests/server/contract/test_error_wire.py @@ -9,6 +9,9 @@ class _State: + async def assert_future_owner(self, request_id, token): + return None + async def get_future(self, request_id: str): return { 'status': 'failed', diff --git a/tests/server/integration/test_full_param_e2e.py b/tests/server/integration/test_full_param_e2e.py index 857fbc711..32d605191 100644 --- a/tests/server/integration/test_full_param_e2e.py +++ b/tests/server/integration/test_full_param_e2e.py @@ -117,7 +117,7 @@ def main() -> int: raise except Exception as e: # noqa: BLE001 msg = str(e) - assert 'full-parameter' in msg or 'lora_config' in msg, f'Phase B FAIL: unexpected error: {msg[:500]}' + assert 'lora' in msg.lower(), f'Phase B FAIL: unexpected error: {msg[:500]}' logger.info('Phase B OK: rejected with mode-mismatch error') # ── Phase C: second full tenant must be rejected (exclusive) ── @@ -130,8 +130,8 @@ def main() -> int: raise except Exception as e: # noqa: BLE001 msg = str(e) - assert 'exclusive' in msg or 'already' in msg, f'Phase C FAIL: unexpected error: {msg[:500]}' - logger.info('Phase C OK: rejected with exclusive-busy error') + assert 'replica limit' in msg.lower() or 'allocation' in msg.lower(), f'Phase C FAIL: unexpected error: {msg[:500]}' + logger.info('Phase C OK: rejected at configured instance capacity') # ── Phase D: save full weights for sampler + sample ── logger.info('Phase D: save_weights_and_get_sampling_client + sample') diff --git a/tests/server/integration/test_mock_mode_startup.py b/tests/server/integration/test_mock_mode_startup.py index d5f706061..e90af75d5 100644 --- a/tests/server/integration/test_mock_mode_startup.py +++ b/tests/server/integration/test_mock_mode_startup.py @@ -3,11 +3,11 @@ Launches the all-mock cookbook config inside the test process via Ray Serve, then asserts: -- Both Model and Sampler deployments report RUNNING within 30 seconds. +- Managed training and sampler routes become ready within the 120-second startup budget. - Repeated calls to the mock model and mock sampler over HTTP return byte-identical responses for identical input. -- The launch path imports cleanly even when ``transformers`` / ``vllm`` / - ``megatron`` would not be available — the mock branches don't pull them. +- The actual launcher normalizes the original applications YAML and deploys + CPU routes and managed mock instances without loading real model weights. This test is heavier than the property suite (boots a full Ray Serve cluster) and is gated behind ``TWINKLE_TEST_INTEGRATION=1`` so plain @@ -18,8 +18,6 @@ import httpx import os import pytest -import subprocess -import sys import time from tests.server.fixtures import MOCK_SERVER_CONFIG, MOCK_SERVER_CONFIG_REDIS @@ -37,23 +35,10 @@ SELECTED_CONFIG = ( MOCK_SERVER_CONFIG_REDIS if os.environ.get('TWINKLE_TEST_REDIS_PERSISTENCE', '0') == '1' else MOCK_SERVER_CONFIG) -READY_BUDGET_SECONDS = 30.0 +READY_BUDGET_SECONDS = 120.0 RAY_NODE_CPUS = 8 -def _run_ray_command(*args: str) -> None: - ray_bin = os.path.join(os.path.dirname(sys.executable), 'ray') - result = subprocess.run( - [ray_bin, *args], - check=False, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - ) - if result.returncode != 0: - raise RuntimeError(f'ray {" ".join(args)} failed with code {result.returncode}:\n{result.stdout}') - - @pytest.fixture(scope='module') def ray_cluster(): """Start a fresh local Ray head node for the duration of the module. @@ -78,18 +63,8 @@ def ray_cluster(): if ray.is_initialized(): ray.shutdown() - _run_ray_command('stop', '--force') - _run_ray_command( - 'start', - '--head', - '--port=0', - f'--num-cpus={RAY_NODE_CPUS}', - '--num-gpus=0', - '--include-dashboard=false', - '--disable-usage-stats', - ) ray.init( - address='auto', + address='local', num_cpus=RAY_NODE_CPUS, num_gpus=0, include_dashboard=False, runtime_env={'env_vars': persistence_env} if persistence_env else None, ) yield @@ -101,10 +76,6 @@ def ray_cluster(): ray.shutdown() except Exception: pass - try: - _run_ray_command('stop', '--force') - except Exception: - pass def _wait_until_healthy(serve_module, timeout: float) -> dict: @@ -124,12 +95,10 @@ def _http(url: str, method: str = 'GET', json: dict | None = None) -> httpx.Resp return httpx.request(method, url, json=json, timeout=10.0) -def test_mock_mode_reaches_ready_under_30s_and_is_deterministic(ray_cluster) -> None: +def test_mock_mode_reaches_ready_and_is_deterministic(ray_cluster, monkeypatch) -> None: + import ray from ray import serve - - from twinkle.server.gateway import build_gateway_app - from twinkle.server.model import build_model_app - from twinkle.server.sampler import build_sampler_app + from twinkle.server.launcher.server_launcher import ServerLauncher cfg = ServerConfig.from_yaml(SELECTED_CONFIG) @@ -139,44 +108,14 @@ def test_mock_mode_reaches_ready_under_30s_and_is_deterministic(ray_cluster) -> serve.start(http_options={'host': host, 'port': port}) started = time.monotonic() - deploys: list[tuple[str, str]] = [] - builders = { - 'server': build_gateway_app, - 'model': build_model_app, - 'sampler': build_sampler_app, - } + namespace = ray.get_runtime_context().namespace + monkeypatch.setenv('TWINKLE_STATE_RAY_NAMESPACE', namespace) + cfg = cfg.model_copy(update={'ray_namespace': namespace, + 'http_options': cfg.http_options.model_copy(update={'host': host, 'port': port})}) + launcher = ServerLauncher(cfg) for app_spec in cfg.applications: - builder = builders[app_spec.import_path] - # Shallow-dump so nested typed models (notably ``queue_config``) stay as - # ``TaskQueueConfig`` instances instead of being serialized back to a - # dict — this mirrors what the production launcher does at - # ``launcher/server_launcher.py:161`` and is required since the - # builders + deployment ``__init__`` accept ``TaskQueueConfig`` directly - # (the ``from_dict`` revival path was removed). - args = {k: v for k, v in dict(app_spec.args).items() if v is not None} - if app_spec.import_path == 'server': - # Gateway's ServiceProxy reads http_options.port to build internal - # proxy targets; the fixture hard-codes 8000 but the test runs on - # a randomized port, so override before passing in. - http_opts = cfg.http_options.model_dump() - http_opts['host'] = host - http_opts['port'] = port - args.setdefault('http_options', http_opts) - # Strip ray_actor_options runtime_env to keep the test light, BUT keep - # ``num_cpus`` low so the three deployments leave room for the mock - # distributed runtime workers in the dedicated local Ray node. - deploy_options: dict = {'ray_actor_options': {'num_cpus': 0.1}} - for raw in app_spec.deployments: - if isinstance(raw, dict): - deploy_options = { - k: v - for k, v in raw.items() if k not in ('name', 'ray_actor_options', 'autoscaling_config') - } - deploy_options['ray_actor_options'] = {'num_cpus': 0.1} - break - bound = builder(deploy_options=deploy_options, **args) - serve.run(bound, name=app_spec.name, route_prefix=app_spec.route_prefix) - deploys.append((app_spec.name, app_spec.route_prefix)) + launcher._deploy_application(app_spec) + launcher._deploy_training() statuses = _wait_until_healthy(serve, READY_BUDGET_SECONDS) elapsed = time.monotonic() - started @@ -261,13 +200,8 @@ def _exercise_twinkle_clients(base: str) -> None: model.resume_from_checkpoint(name=save_resp.twinkle_path) model.apply_patch('NoopPatch') - model.upload_to_hub( - checkpoint_dir=save_resp.twinkle_path, - hub_model_id='mock/dummy', - hub_token='EMPTY_TOKEN', - async_upload=False, - poll_interval=0.5, - ) + # Hub upload is checked with a local mocked HubOperation; this smoke + # must not publish files to an external account. # --- sampler service --- sampler = vLLMSampler(model_id='mock-model') diff --git a/tests/server/lifecycle/test_retrieve_endpoint.py b/tests/server/lifecycle/test_retrieve_endpoint.py index 84b8767e4..cde575825 100644 --- a/tests/server/lifecycle/test_retrieve_endpoint.py +++ b/tests/server/lifecycle/test_retrieve_endpoint.py @@ -17,6 +17,9 @@ class _State: + + async def assert_future_owner(self, request_id, token): + return None """A fake ServerState whose get_future returns a fixed record (or None).""" def __init__(self, record): diff --git a/tests/server/lifecycle/test_run_submit_dedup.py b/tests/server/lifecycle/test_run_submit_dedup.py index f9e91d76c..4d3b3291e 100644 --- a/tests/server/lifecycle/test_run_submit_dedup.py +++ b/tests/server/lifecycle/test_run_submit_dedup.py @@ -75,7 +75,7 @@ async def test_release_kept_when_task_already_enqueued(): with pytest.raises(RuntimeError): await run_submit(mgmt, _request(), _body(), task_type='forward_backward', backend_call=_call) assert mgmt.state.released == [], 'claim wrongly released for an already-enqueued task' - assert 'seq::sess-1::sess-1-ad::7' in mgmt.state.claimed + assert 'seq::sess-1::sess-1-ad::forward_backward::7' in mgmt.state.claimed @pytest.mark.asyncio @@ -84,7 +84,7 @@ async def test_release_when_never_enqueued(): mgmt = _FakeManagement(record_after_claim=None) with pytest.raises(RuntimeError): await run_submit(mgmt, _request(), _body(), task_type='forward_backward', backend_call=_call) - assert mgmt.state.released == ['seq::sess-1::sess-1-ad::7'], 'claim should be released when nothing enqueued' + assert mgmt.state.released == ['seq::sess-1::sess-1-ad::forward_backward::7'], 'claim should be released when nothing enqueued' @pytest.mark.asyncio @@ -107,5 +107,5 @@ async def test_dedup_key_is_scoped_per_adapter(): backend_call=_call) claimed = set(mgmt.state.claimed) - assert claimed == {'seq::sess-1::sess-1-lora-A::1', 'seq::sess-1::sess-1-lora-B::1'}, ( + assert claimed == {'seq::sess-1::sess-1-lora-A::forward_backward::1', 'seq::sess-1::sess-1-lora-B::forward_backward::1'}, ( f'adapters collided on one dedup key: {claimed}') diff --git a/tests/server/lifecycle/test_tinker_retrieve_regression.py b/tests/server/lifecycle/test_tinker_retrieve_regression.py index 4af4611f7..504ec294c 100644 --- a/tests/server/lifecycle/test_tinker_retrieve_regression.py +++ b/tests/server/lifecycle/test_tinker_retrieve_regression.py @@ -15,6 +15,9 @@ class _State: + async def assert_future_owner(self, request_id, token): + return None + def __init__(self, record): self._record = record diff --git a/tests/server/model/test_replica_lifecycle.py b/tests/server/model/test_replica_lifecycle.py index 84b961391..0b1160085 100644 --- a/tests/server/model/test_replica_lifecycle.py +++ b/tests/server/model/test_replica_lifecycle.py @@ -34,6 +34,7 @@ def _make_lifecycle_manager(state: _CapacityState, replica_id: str, max_loras: i manager.state = state manager.replica_id = replica_id manager.max_loras = max_loras + manager.train_mode = 'lora' manager._replica_registered = False manager.data_plane = SimpleNamespace(close=AsyncMock()) manager.shutdown_task_queue = AsyncMock() @@ -54,24 +55,27 @@ async def test_replica_lifecycle_updates_shared_capacity() -> None: assert await state.get_capacity_info() == {'max_loras': 6, 'used_loras': 0, 'free_loras': 6} await second.shutdown() + # Only the owner retires capacity after confirming the GPU workers stopped. + assert await state.get_capacity_info() == {'max_loras': 6, 'used_loras': 0, 'free_loras': 6} + await state.unregister_replica(second.replica_id) assert await state.get_capacity_info() == {'max_loras': 3, 'used_loras': 0, 'free_loras': 3} @pytest.mark.asyncio async def test_async_constructor_registers_replica_before_ready() -> None: state = SimpleNamespace(register_replica=AsyncMock(), touch_replica_last_seen=AsyncMock()) - replica_context = SimpleNamespace(replica_id=SimpleNamespace(unique_id='replica-1')) manager = ModelManagement.__new__(ModelManagement) manager._task_queue_config = SimpleNamespace(effective_execution_timeout=1800.0) with patch('twinkle.server.model.app.DeviceGroup', return_value=SimpleNamespace(name='group')), \ patch('twinkle.server.model.app.init_twinkle_runtime', return_value=None), \ - patch('twinkle.server.model.app.serve.get_replica_context', return_value=replica_context), \ patch.object(MODEL_SELECTOR, 'construct', return_value=MagicMock()), \ patch('twinkle.server.model.app.get_server_state', return_value=state), \ patch('twinkle.server.data_plane.DataPlaneProxy', return_value=MagicMock()), \ patch.object(ModelManagement, '_init_task_queue'), \ - patch.object(ModelManagement, '_init_adapter_manager'): + patch.object(ModelManagement, '_init_adapter_manager'), \ + patch.object(ModelManagement, 'initialize_managed_runtime', new_callable=AsyncMock), \ + patch.object(ModelManagement, 'wait_for_training_workers', new_callable=AsyncMock): await ModelManagement.__init__( manager, model_id='model', @@ -80,6 +84,9 @@ async def test_async_constructor_registers_replica_before_ready() -> None: device_mesh={}, backend='mock', max_loras=3, + training_context={'instance_id': 'replica-1'}, + placement_group=None, + allocation_devices=[], ) state.register_replica.assert_awaited_once_with('replica-1', 3) diff --git a/tests/server/start_e2e_server.py b/tests/server/start_e2e_server.py index 4af2e2b85..1e98c914a 100644 --- a/tests/server/start_e2e_server.py +++ b/tests/server/start_e2e_server.py @@ -38,6 +38,8 @@ MEGATRON_BASE = "tests/server/config/server_config_4b_e2e_megatron.yaml" _FULL = { "train_mode": "full", + "idle_timeout_s": 3600, + "max_replicas": 1, # Full-parameter training holds the deployment exclusively; use long # timeouts so the tenant is not expired mid-test. "adapter_config.adapter_timeout": 3600, diff --git a/tests/server/test_gateway_services.py b/tests/server/test_gateway_services.py index 5a5d032d8..e61dfd631 100644 --- a/tests/server/test_gateway_services.py +++ b/tests/server/test_gateway_services.py @@ -10,6 +10,9 @@ class _State: + async def assert_future_owner(self, request_id, token): + return None + def __init__(self, records): self.records = list(records) @@ -43,7 +46,7 @@ def test_gateway_services_do_not_import_protocol_models(): class _FakeGateway: """Minimal ``GatewayServer`` stand-in; ``poll_future`` is patched so state is unused.""" - state = None + state = _State([]) def _parity_client(monkeypatch, canonical_record): diff --git a/tests/server/training/test_allocation.py b/tests/server/training/test_allocation.py new file mode 100644 index 000000000..22d4f6860 --- /dev/null +++ b/tests/server/training/test_allocation.py @@ -0,0 +1,23 @@ +import json +import subprocess +import sys +from pathlib import Path + + +def test_production_allocator_uses_disjoint_devices_and_reuses_released_group(): + result = subprocess.run([sys.executable, str(Path(__file__).resolve().parents[3] / 'tools/validation/on_demand_training/allocation_probe.py')], + capture_output=True, text=True, timeout=120) + assert result.returncode == 0, result.stdout + result.stderr + evidence = json.loads(result.stdout.strip().splitlines()[-1]) + assert evidence['passed'] and len(set(evidence['devices'])) == 8 + assert len(evidence['reused']) == 2 + + +def test_fragmented_devices_cannot_be_combined_across_nodes(): + result = subprocess.run([sys.executable, str(Path(__file__).resolve().parents[3] / 'tools/validation/on_demand_training/fragmentation_probe.py')], + capture_output=True, text=True, timeout=120) + assert result.returncode == 0, result.stdout + result.stderr + evidence = json.loads(result.stdout.strip().splitlines()[-1]) + assert evidence['passed'] and evidence['nodes'] == 2 + assert evidence['fragmented_rejected'] and evidence['pending_removed'] + assert len(evidence['devices']) == 2 diff --git a/tests/server/training/test_controller.py b/tests/server/training/test_controller.py new file mode 100644 index 000000000..11987a10b --- /dev/null +++ b/tests/server/training/test_controller.py @@ -0,0 +1,677 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Exercise real controller/storage decisions against a controllable runtime.""" +from __future__ import annotations + +import asyncio + +import pytest + +from twinkle.server.config.training import TrainingConfig +from twinkle.server.exceptions import RequestRejectedError, ResourceNotFoundError +from twinkle.server.state import ServerState +from twinkle.server.state.backend.memory_backend import RayActorBackend +from twinkle.server.training.controller import TrainingController +from twinkle.server.training.errors import CapacityUnavailableError, TrainingExpiredError +from twinkle.server.training.records import CreateTraining +from twinkle.server.training.store import TrainingStore + + +class Clock: + now = 0.0 + + def __call__(self): + return self.now + + +class Runtime: + """Allocates a finite pool and stores actual per-user state for assertions.""" + + def __init__(self, devices=16): + self.devices = devices + self.allocations = {} + self.users = {} + self.starts = 0 + self.stopped = [] + self.start_gate = asyncio.Event() + self.start_gate.set() + self.fail_start = False + self.fail_shutdown = False + + async def start(self, instance, spec): + self.starts += 1 + if self.fail_start: + raise RuntimeError('injected startup failure') + used = sum(self.allocations.values()) + if used + spec.config.num_devices > self.devices: + raise CapacityUnavailableError('No devices available') + self.allocations[instance.instance_id] = spec.config.num_devices + await self.start_gate.wait() + return {} + + async def create_training(self, instance, training, request): + assert instance.instance_id in self.allocations + assert training.training_id not in self.users + self.users[training.training_id] = {'instance': instance.instance_id, 'label': training.adapter_label, 'steps': 0} + + async def release_training(self, instance, training): + self.users.pop(training.training_id, None) + + async def shutdown(self, instance): + if self.fail_shutdown: + raise RuntimeError('cannot confirm process death') + self.allocations.pop(instance.instance_id, None) + self.users = {key: value for key, value in self.users.items() if value['instance'] != instance.instance_id} + self.stopped.append(instance.instance_id) + + async def verify_operation(self, instance, request_id): + return {'state': 'pending'} + + async def recover_orphans(self): + pass + + +@pytest.mark.asyncio +async def test_initialization_timeout_only_retires_the_failing_tenant(): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, _ = await setup(backend=FakeBackend()) + active = await create(controller, state, 'active') + runtime.users[active]['steps'] = 7 + instance = controller.instances[controller.trainings[active].instance_id] + async def timeout(instance, training, request): + raise TimeoutError('User initialization timed out') + runtime.create_training = timeout + ref = await controller.create(request('failing')) + await controller.wait_for_jobs() + assert (await state.get_future(ref['request_id']))['failure']['reason_code'] == 'startup_timeout' + assert controller.trainings[active].state == 'ACTIVE' + assert runtime.users[active]['steps'] == 7 + assert instance.state == 'READY' and instance.training_ids == [active] + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_late_initializer_keeps_its_slot_reserved_until_serial_cleanup(): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, _ = await setup(backend=FakeBackend(), max_replicas=1) + active = await create(controller, state, 'active') + raw = controller.config.model_dump() + raw['startup_timeout_s'] = 0.03 + raw['models'][0]['modes']['lora']['max_loras'] = 2 + controller.config = TrainingConfig.model_validate(raw) + from twinkle.server.training.catalog import TrainingCatalog + controller.catalog = TrainingCatalog(controller.config) + entered, unblock = asyncio.Event(), asyncio.Event() + original_create, original_release = runtime.create_training, runtime.release_training + initializer = None + async def initialize(instance, training, req): + nonlocal initializer + async def late(): + entered.set() + await unblock.wait() + await original_create(instance, training, req) + initializer = asyncio.create_task(late()) + await asyncio.shield(initializer) + async def release(instance, training): + await initializer + await original_release(instance, training) + runtime.create_training, runtime.release_training = initialize, release + ref = await controller.create(request('late')) + await entered.wait() + for _ in range(100): + future = await state.get_future(ref['request_id']) + if future['status'] == 'failed': + break + await asyncio.sleep(0.005) + assert future['failure']['reason_code'] == 'startup_timeout' + instance = controller.instances[controller.trainings[active].instance_id] + assert len(instance.training_ids) == 2 and controller.trainings[active].state == 'ACTIVE' + refused = await controller.create(request('no-free-slot')) + assert (await state.get_future(refused['request_id']))['failure']['reason_code'] == 'capacity_unavailable' + unblock.set() + await controller.wait_for_jobs() + assert instance.training_ids == [active] and set(runtime.users) == {active} + runtime.create_training, runtime.release_training = original_create, original_release + replacement = await create(controller, state, 'replacement') + assert controller.trainings[replacement].instance_id == instance.instance_id + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_user_gets_full_initialization_window_after_base_model_startup(): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, clock = await setup(backend=FakeBackend()) + runtime.start_gate.clear() + ref = await controller.create(request('near-startup-deadline')) + instance = next(iter(controller.instances.values())) + clock.now = instance.startup_deadline - 1 + original = runtime.create_training + async def initialize(instance, training, req): + clock.now += 2 + await original(instance, training, req) + runtime.create_training = initialize + runtime.start_gate.set() + await controller.wait_for_jobs() + assert (await state.get_future(ref['request_id']))['status'] == 'completed' + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_resident_rebuild_does_not_block_expiry_or_another_mode(): + from tests.server.state.fake_backend import FakeBackend + from twinkle.server.training.catalog import TrainingCatalog + controller, state, runtime, clock = await setup(backend=FakeBackend()) + expired = await create(controller, state, 'expired-full', mode='full') + raw = controller.config.model_dump() + raw['models'][0]['modes']['lora']['min_replicas'] = 1 + controller.config = TrainingConfig.model_validate(raw) + controller.catalog = TrainingCatalog(controller.config) + runtime.start_gate.clear() + await asyncio.wait_for(controller.sweep(), 0.1) + for _ in range(20): + await asyncio.sleep(0) + if runtime.starts == 2: + break + clock.now = 11 + await asyncio.wait_for(controller.sweep(), 0.1) + for _ in range(20): + await asyncio.sleep(0) + if controller.trainings[expired].state == 'EXPIRED': + break + assert controller.trainings[expired].state == 'EXPIRED' + ref = await controller.create(request('unrelated-full', mode='full')) + assert (await state.get_future(ref['request_id']))['status'] == 'pending' + resident_ids = [item.instance_id for item in controller.instances.values() if item.resident] + assert len(resident_ids) == 1 + await controller.sweep() + assert [item.instance_id for item in controller.instances.values() if item.resident] == resident_ids + runtime.start_gate.set() + await controller.wait_for_jobs() + assert (await state.get_future(ref['request_id']))['status'] == 'completed' + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_reconciliation_of_committed_completion_does_not_extend_idle_time(): + import time + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, clock = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'active') + await controller.admit(identity, 'owner', 'session', 'lost-finish') + await controller.commit_operation(identity, 'lost-finish', controller.epoch) + async def finished(instance, request_id): + return {'state': 'finished', 'executed': True, 'finished_at': time.time() - 31} + runtime.verify_operation = finished + clock.now = 31 + await controller.sweep() + await controller.wait_for_jobs() + assert not controller.trainings[identity].pending_operation_ids + assert controller.trainings[identity].last_activity_at < 0.01 + await controller.sweep() + await controller.wait_for_jobs() + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_legacy_session_and_max_lifetime_bounds_keep_their_original_meaning(): + import time + from datetime import datetime, timezone + from tests.server.state.fake_backend import FakeBackend + from twinkle.server.training.catalog import TrainingCatalog + controller, state, runtime, clock = await setup(backend=FakeBackend()) + raw = controller.config.model_dump() + raw['models'][0]['modes']['lora']['idle_timeout_s'] = None + raw['models'][0]['legacy_expiration'] = {'session_timeout_s': 30, 'max_lifetime_s': 36000} + controller.config = TrainingConfig.model_validate(raw) + controller.catalog = TrainingCatalog(controller.config) + identity = await create(controller, state, 'legacy') + clock.now = 1801 + await state.touch_session('session') + await controller.sweep() + await controller.wait_for_jobs() + assert controller.trainings[identity].state == 'ACTIVE' + controller.trainings[identity].created_at = datetime.fromtimestamp(time.time() - 36001, timezone.utc).isoformat() + await controller.sweep() + await controller.wait_for_jobs() + assert controller.trainings[identity].state == 'EXPIRED' + second = await create(controller, state, 'stale-heartbeat') + session = await state._backend.get('session::session') + session['last_heartbeat'] = time.time() - 31 + await state._backend.set('session::session', session) + await controller.sweep() + await controller.wait_for_jobs() + assert controller.trainings[second].state == 'EXPIRED' + + +async def setup(*, devices=16, min_replicas=0, max_replicas=None, quota=100, backend=None): + mode = {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}, 'max_loras': 5, + 'idle_timeout_s': 10, 'min_replicas': min_replicas, 'max_replicas': max_replicas} + config = TrainingConfig(models=[{ + 'name': 'small', 'model_id': 'ms://small', 'backend': 'mock', 'device_type': 'cuda', + 'modes': {'lora': mode, 'full': {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}, 'idle_timeout_s': 10}}, + }]) + state = ServerState(backend=backend or RayActorBackend(), per_token_model_limit=quota) + await state.create_session({'session_id': 'session'}) + runtime, clock = Runtime(devices), Clock() + controller = TrainingController(config, TrainingStore(state, 'unit'), runtime, clock=clock) + await controller.initialize() + return controller, state, runtime, clock + + +def request(key, *, mode='lora', label='same', token='owner'): + return CreateTraining(base_model='small', mode=mode, token=token, session_id='session', create_key=str(key), + client_type='twinkle', payload={}, adapter_label=label) + + +async def create(controller, state, key, **kwargs): + ref = await controller.create(request(key, **kwargs)) + await controller.wait_for_jobs() + future = await state.get_future(ref['request_id']) + assert future['status'] == 'completed', future + return future['result']['training_id'] + + +@pytest.mark.asyncio +async def test_lazy_queries_allocate_nothing(): + controller, state, runtime, _ = await setup() + assert runtime.starts == 0 + assert controller.catalog.get('small', 'full').capacity == 1 + with pytest.raises(ResourceNotFoundError): + await controller.route('unknown', 'owner') + assert not runtime.allocations + assert await state.get_capacity_info() == {'max_loras': 0, 'used_loras': 0, 'free_loras': 0} + + +@pytest.mark.asyncio +async def test_32_concurrent_creates_coalesce_even_before_allocation_ready(): + controller, state, runtime, _ = await setup() + runtime.start_gate.clear() + refs = await asyncio.gather(*(controller.create(request(index)) for index in range(32))) + assert len(controller.instances) == 7 + assert sorted(len(instance.training_ids) for instance in controller.instances.values()) == [2, 5, 5, 5, 5, 5, 5] + runtime.start_gate.set() + await controller.wait_for_jobs() + records = [await state.get_future(ref['request_id']) for ref in refs] + assert all(record['status'] == 'completed' for record in records) + assert len({record['result']['training_id'] for record in records}) == 32 + assert runtime.starts == 7 + assert len(runtime.users) == 32 + + +@pytest.mark.asyncio +async def test_full_and_lora_isolation_and_device_reuse(): + controller, state, runtime, _ = await setup(devices=6) + shared = await create(controller, state, 1) + full1 = await create(controller, state, 2, mode='full') + full2 = await create(controller, state, 3, mode='full') + identities = {runtime.users[key]['instance'] for key in (shared, full1, full2)} + assert len(identities) == 3 + old_full = runtime.users[full1]['instance'] + await controller.release(full1, 'owner') + replacement = await create(controller, state, 4, mode='full') + assert runtime.users[replacement]['instance'] != old_full + assert shared in runtime.users and full2 in runtime.users + assert len(runtime.allocations) == 3 + + +@pytest.mark.asyncio +async def test_capacity_failure_is_terminal_and_does_not_hold_quota(): + controller, state, runtime, _ = await setup(devices=2) + await create(controller, state, 1, mode='full') + ref = await controller.create(request(2, mode='full')) + await controller.wait_for_jobs() + record = await state.get_future(ref['request_id']) + assert record['status'] == 'failed' + assert record['failure']['reason_code'] == 'capacity_unavailable' + assert await state._backend.get('token_count::owner') == 1 + assert len(runtime.allocations) == 1 + assert await controller.create(request(2, mode='full')) == ref + + +@pytest.mark.asyncio +async def test_create_idempotency_and_changed_payload_conflict(): + controller, state, runtime, _ = await setup() + refs = await asyncio.gather(*(controller.create(request('one')) for _ in range(16))) + assert len({ref['request_id'] for ref in refs}) == 1 + await controller.wait_for_jobs() + assert len(runtime.users) == 1 + changed = request('one') + changed.payload = {'different': True} + with pytest.raises(RequestRejectedError, match='different parameters'): + await controller.create(changed) + + +@pytest.mark.asyncio +async def test_pending_creates_count_toward_quota(): + controller, state, runtime, _ = await setup(quota=1) + runtime.start_gate.clear() + first = await controller.create(request(1)) + second = await controller.create(request(2)) + failure = await state.get_future(second['request_id']) + assert failure['failure']['reason_code'] == 'resource_quota_exceeded' + runtime.start_gate.set() + await controller.wait_for_jobs() + assert (await state.get_future(first['request_id']))['status'] == 'completed' + assert len(runtime.users) == 1 + + +@pytest.mark.asyncio +async def test_heartbeat_does_not_prevent_idle_expiry_or_affect_busy_user(): + controller, state, runtime, clock = await setup() + idle = await create(controller, state, 1) + busy = await create(controller, state, 2) + await controller.admit(busy, 'owner', 'session', 'op') + await controller.commit_operation(busy, 'op', controller.epoch) + clock.now = 11 + await state.touch_session('session') + await controller.sweep() + await controller.wait_for_jobs() + assert idle not in runtime.users and busy in runtime.users + with pytest.raises(TrainingExpiredError): + await controller.route(idle, 'owner') + await controller.finish_operation(busy, 'op', controller.epoch, executed=True) + assert controller.trainings[busy].last_activity_at == 11 + clock.now = 20 + await controller.sweep() + await controller.wait_for_jobs() + assert busy in runtime.users + clock.now = 22 + await controller.sweep() + await controller.wait_for_jobs() + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_release_drains_accepted_operation_before_cleanup(): + controller, state, runtime, _ = await setup() + identity = await create(controller, state, 1) + await controller.admit(identity, 'owner', 'session', 'op') + await controller.commit_operation(identity, 'op', controller.epoch) + releasing = asyncio.create_task(controller.release(identity, 'owner')) + while controller.trainings[identity].state == 'ACTIVE': + await asyncio.sleep(0) + assert identity in runtime.users + with pytest.raises(ResourceNotFoundError): + await controller.admit(identity, 'owner', 'session', 'next') + await controller.finish_operation(identity, 'op', controller.epoch, executed=True) + await controller.finish_operation(identity, 'op', controller.epoch, executed=True) + await controller.wait_for_jobs() + await releasing + assert not runtime.users and not runtime.allocations + assert len(runtime.stopped) == 1 + + +@pytest.mark.asyncio +async def test_same_label_recreation_and_ownership_do_not_reattach_old_handle(): + controller, state, runtime, _ = await setup() + old = await create(controller, state, 1) + with pytest.raises(ResourceNotFoundError): + await controller.route(old, 'other') + with pytest.raises(ResourceNotFoundError): + await controller.route(old, 'owner', 'other-session') + await controller.release(old, 'owner') + new = await create(controller, state, 2) + assert new != old and runtime.users[new]['label'] == 'same' + with pytest.raises(ResourceNotFoundError): + await controller.route(old, 'owner') + + +@pytest.mark.asyncio +async def test_resident_survives_empty_but_extra_replica_does_not(): + controller, state, runtime, _ = await setup(min_replicas=1) + refs = [await controller.create(request(index)) for index in range(6)] + await controller.wait_for_jobs() + identities = [(await state.get_future(ref['request_id']))['result']['training_id'] for ref in refs] + for identity in identities: + await controller.release(identity, 'owner') + assert len(runtime.allocations) == 1 + resident = next(instance for instance in controller.instances.values() if instance.resident) + assert resident.state == 'READY' + + +@pytest.mark.asyncio +async def test_startup_failure_cleans_all_waiters_and_quota(): + controller, state, runtime, _ = await setup() + runtime.fail_start = True + refs = await asyncio.gather(*(controller.create(request(index)) for index in range(5))) + await controller.wait_for_jobs() + records = [await state.get_future(ref['request_id']) for ref in refs] + assert all(record['status'] == 'failed' for record in records) + assert (await state._backend.get('token_count::owner') or 0) == 0 + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_managed_records_are_not_deleted_by_generic_training_ttl(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 1) + await state._model_mgr.cleanup_expired(10**12, ['session']) + assert await state.get_model_metadata(identity) is not None + assert identity in runtime.users + + +@pytest.mark.asyncio +async def test_restart_invalidates_handles_and_cleans_old_allocation(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 1) + replacement = TrainingController(controller.config, TrainingStore(state, 'unit'), runtime, clock=clock) + await replacement.initialize() + assert replacement.epoch != controller.epoch + assert not runtime.allocations and not runtime.users + assert await state.get_model_metadata(identity) is None + assert (await state._backend.get('token_count::owner') or 0) == 0 + with pytest.raises(ResourceNotFoundError): + await replacement.route(identity, 'owner') + + +@pytest.mark.asyncio +async def test_cancelled_create_never_initializes_user_and_keeps_receipt(): + controller, state, runtime, _ = await setup() + runtime.start_gate.clear() + ref = await controller.create(request('cancel')) + assert (await state.cancel_future(ref['request_id']))['cancelled'] + runtime.start_gate.set() + await controller.wait_for_jobs() + assert (await state.get_future(ref['request_id']))['status'] == 'cancelled' + assert not runtime.users and not runtime.allocations + assert await controller.create(request('cancel')) == ref + + +@pytest.mark.asyncio +async def test_unconfirmed_worker_shutdown_holds_devices_and_quota_until_retry(): + from twinkle.server.state.models import FutureFailureRecord + controller, state, runtime, _ = await setup(devices=2) + identity = await create(controller, state, 1, mode='full') + instance = controller.instances[controller.trainings[identity].instance_id] + await controller.admit(identity, 'owner', 'session', 'running') + await controller.commit_operation(identity, 'running', controller.epoch) + runtime.fail_shutdown = True + await controller.fail_instance(instance.instance_id, FutureFailureRecord( + reason_code='execution_timeout', message='RPC timed out', attribution='server')) + assert instance.state == 'QUARANTINED' + assert runtime.allocations and controller.trainings[identity].pending_operation_ids + assert await state.get_model_metadata(identity) is not None + failed = await controller.create(request(2, mode='full')) + await controller.wait_for_jobs() + assert (await state.get_future(failed['request_id']))['failure']['reason_code'] == 'capacity_unavailable' + runtime.fail_shutdown = False + await controller.sweep() + await controller.wait_for_jobs() + assert not runtime.allocations + assert instance.state == 'STOPPED' + assert (await state._backend.get('token_count::owner') or 0) == 0 + + +@pytest.mark.asyncio +async def test_uncertain_enqueue_is_failed_before_ticket_is_released(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 1) + await controller.admit(identity, 'owner', 'session', 'lost-submit') + await controller.store.future('lost-submit', 'pending', identity) + async def missing(instance, request_id): + return {'state': 'missing'} + runtime.verify_operation = missing + clock.now = 31 + await controller.sweep() + await controller.wait_for_jobs() + assert (await state.get_future('lost-submit'))['status'] == 'failed' + assert not controller.trainings[identity].pending_operation_ids + assert controller.trainings[identity].last_activity_at == 0 + await controller.sweep() + await controller.wait_for_jobs() + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_restart_fails_pending_operations_and_fences_old_controller(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 1) + await controller.admit(identity, 'owner', 'session', 'queued') + await controller.store.future('queued', 'queued', identity) + replacement = TrainingController(controller.config, TrainingStore(state, 'unit'), runtime, clock=clock) + await replacement.initialize() + assert (await state.get_future('queued'))['failure']['reason_code'] == 'orphaned_replica' + with pytest.raises(ResourceNotFoundError, match='epoch'): + await controller.create(request(2)) + with pytest.raises(ResourceNotFoundError, match='epoch'): + await controller.route(identity, 'owner') + assert await replacement.create(request(1)) == {'request_id': controller.trainings[identity].create_request_id} + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_info_does_not_renew_and_expiry_wins_before_next_sweep(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 1) + clock.now = 9 + await controller.admit(identity, 'owner', 'session', 'info', renew=False) + await controller.finish_operation(identity, 'info', controller.epoch, executed=True) + assert controller.trainings[identity].last_activity_at == 0 + clock.now = 10 + with pytest.raises(TrainingExpiredError): + await controller.admit(identity, 'owner', 'session', 'next') + await controller.wait_for_jobs() + assert not runtime.users + + +@pytest.mark.asyncio +async def test_full_resident_cannot_accept_user_while_resetting(): + controller, state, runtime, _ = await setup() + original = controller.config.model_dump() + original['models'][0]['modes']['full']['min_replicas'] = 1 + original['models'][0]['modes']['full']['max_replicas'] = 1 + controller.config = TrainingConfig.model_validate(original) + from twinkle.server.training.catalog import TrainingCatalog + controller.catalog = TrainingCatalog(controller.config) + await controller.sweep() + await controller.wait_for_jobs() + identity = await create(controller, state, 1, mode='full') + entered, resume = asyncio.Event(), asyncio.Event() + previous = runtime.release_training + async def resetting(instance, training): + entered.set() + await resume.wait() + await previous(instance, training) + runtime.release_training = resetting + releasing = asyncio.create_task(controller.release(identity, 'owner')) + await entered.wait() + instance = controller.instances[controller.trainings[identity].instance_id] + assert instance.state == 'RESETTING' + ref = await controller.create(request(2, mode='full')) + assert (await state.get_future(ref['request_id']))['failure']['reason_code'] == 'capacity_unavailable' + resume.set() + await releasing + assert instance.state == 'READY' + next_id = await create(controller, state, 3, mode='full') + assert controller.trainings[next_id].instance_id == instance.instance_id + + +@pytest.mark.asyncio +async def test_dirty_resident_is_rebuilt_after_last_user_leaves(): + controller, state, runtime, _ = await setup(min_replicas=1) + identity = await create(controller, state, 1) + instance = controller.instances[controller.trainings[identity].instance_id] + instance.dirty = True + await controller.release(identity, 'owner') + assert instance.state == 'STOPPED' + await controller.sweep() + await controller.wait_for_jobs() + ready = [item for item in controller.instances.values() if item.state == 'READY'] + assert len(ready) == 1 and ready[0].instance_id != instance.instance_id + + +@pytest.mark.asyncio +@pytest.mark.parametrize('late_phase', ['runtime', 'user']) +async def test_late_startup_result_cannot_publish_success(late_phase): + controller, state, runtime, clock = await setup() + if late_phase == 'runtime': + original = runtime.start + async def late(instance, spec): + result = await original(instance, spec) + clock.now = instance.startup_deadline + 1 + return result + runtime.start = late + else: + original = runtime.create_training + async def late(instance, training, request): + await original(instance, training, request) + clock.now = training.creation_deadline + 1 + runtime.create_training = late + ref = await controller.create(request('late')) + await controller.wait_for_jobs() + future = await state.get_future(ref['request_id']) + assert future['status'] == 'failed' and future['failure']['reason_code'] == 'startup_timeout' + assert not runtime.allocations and not runtime.users + assert all(instance.state == 'STOPPED' for instance in controller.instances.values()) + + +@pytest.mark.asyncio +async def test_failed_resident_startup_rolls_back_previously_loaded_residents(): + controller, state, runtime, clock = await setup(devices=2) + raw = controller.config.model_dump() + raw['models'][0]['modes']['lora']['min_replicas'] = 2 + replacement = TrainingController(TrainingConfig.model_validate(raw), TrainingStore(state, 'unit'), runtime, clock=clock) + with pytest.raises(CapacityUnavailableError): + await replacement.initialize() + assert not replacement._ready and not runtime.allocations + assert all(item.state == 'STOPPED' for item in replacement.instances.values()) + await replacement.wait_for_jobs() + + +@pytest.mark.asyncio +async def test_retired_diagnostics_are_bounded_without_discarding_live_training(): + controller, state, runtime, clock = await setup() + old = await create(controller, state, 'old') + await controller.release(old, 'owner', 'session') + retired_instance = controller.trainings[old].instance_id + live = await create(controller, state, 'live') + controller.trainings[old].terminal_at -= state.expiration_timeout + 1 + controller.instances[retired_instance].stopped_at -= state.expiration_timeout + 1 + await controller.sweep() + await controller.wait_for_jobs() + assert old not in controller.trainings and retired_instance not in controller.instances + assert retired_instance not in controller._startup + assert (await controller.route(live, 'owner', 'session'))['instance_id'] in runtime.allocations + await controller.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('startup_blocked', [False, True]) +async def test_creation_record_write_failure_closes_new_empty_instance(monkeypatch, startup_blocked): + controller, state, runtime, clock = await setup() + if startup_blocked: + runtime.start_gate.clear() + original = state._backend.set + failed = False + async def failing_set(key, value, *args, **kwargs): + nonlocal failed + if (not failed and key.startswith('model::') and isinstance(value, dict) + and value.get('instance_id') and value.get('state') == 'PENDING'): + failed = True + raise RuntimeError('injected training record write failure') + return await original(key, value, *args, **kwargs) + monkeypatch.setattr(state._backend, 'set', failing_set) + ref = await controller.create(request('write-failure')) + await controller.wait_for_jobs() + assert (await state.get_future(ref['request_id']))['status'] == 'failed' + assert not runtime.allocations + assert await state._backend.get('token_count::owner') == 0 diff --git a/tests/server/training/test_http_integration.py b/tests/server/training/test_http_integration.py new file mode 100644 index 000000000..367555749 --- /dev/null +++ b/tests/server/training/test_http_integration.py @@ -0,0 +1,16 @@ +import json +import subprocess +import sys +from pathlib import Path + + +def test_installed_native_and_tinker_sdks_use_managed_http_routes(): + result = subprocess.run([sys.executable, str(Path(__file__).resolve().parents[3] / 'tools/validation/on_demand_training/http_probe.py')], + capture_output=True, text=True, timeout=420) + assert result.returncode == 0, result.stdout + result.stderr + evidence = next(json.loads(line) for line in result.stdout.splitlines() if line.startswith('{"passed":')) + assert evidence['passed'] and evidence['native_recreation'] and evidence['tinker_shared'] + assert evidence['native_lora'] + assert evidence['runtime_lifecycle_rejected'] + assert evidence['tinker_full'] + assert evidence['restart'] and evidence['state_loss'] and evidence['cycles'] == 10 diff --git a/tests/server/training/test_launcher.py b/tests/server/training/test_launcher.py new file mode 100644 index 000000000..d212e43e0 --- /dev/null +++ b/tests/server/training/test_launcher.py @@ -0,0 +1,39 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import socket +import uuid + +import pytest +import ray + +from twinkle.server.config import ServerConfig +from twinkle.server.launcher.server_launcher import ServerLauncher +from twinkle.server.training.actor import controller_name + + +def test_failed_initialization_does_not_leave_an_actor_that_blocks_a_retry(): + namespace = 'launcher-retry-' + uuid.uuid4().hex + with socket.socket() as listener: + listener.bind(('127.0.0.1', 0)) + port = listener.getsockname()[1] + devices = int(ray.cluster_resources()['CPU']) + 1 + payload = { + 'ray_namespace': namespace, 'http_options': {'host': '127.0.0.1', 'port': port}, + 'applications': [{'name': 'retry', 'route_prefix': '/api/v1/model/retry', 'import_path': 'model', + 'args': {'model_id': 'mock', 'backend': 'mock', 'device_type': 'cpu', + 'modes': {'full': {'num_devices': devices, 'device_mesh': {'dp_size': devices}, + 'min_replicas': 1}}}}], + } + launcher = ServerLauncher(ServerConfig.model_validate(payload)) + launcher._start_serve() + try: + with pytest.raises(ray.exceptions.RayTaskError, match='No single-node allocation'): + launcher._deploy_training() + with pytest.raises(ValueError): + ray.get_actor(controller_name(namespace), namespace=namespace) + payload['applications'][0]['args']['modes']['full']['min_replicas'] = 0 + launcher.config = ServerConfig.model_validate(payload) + launcher._deploy_training() + diagnostic = ray.get(launcher._training_controller.diagnostics.remote()) + assert diagnostic['ready'] and diagnostic['instances'] == [] + finally: + launcher._shutdown() diff --git a/tests/server/training/test_router.py b/tests/server/training/test_router.py new file mode 100644 index 000000000..08e541ad0 --- /dev/null +++ b/tests/server/training/test_router.py @@ -0,0 +1,87 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +from unittest.mock import Mock +from types import SimpleNamespace + +import pytest +from starlette.requests import Request + +from tests.server.state.fake_backend import FakeBackend +from tests.server.training.test_runtime_queue import settle +from twinkle.hub import HubOperation +from twinkle.server.state import ServerState +from twinkle.server.training.router import TrainingRouter +from twinkle.server.training.store import TrainingStore +from twinkle.server.training.catalog import TrainingCatalog +from twinkle.server.config.training import TrainingConfig + + +def http_request(body): + async def receive(): + return {'type': 'http.request', 'body': json.dumps(body).encode(), 'more_body': False} + request = Request({'type': 'http', 'method': 'POST', 'path': '/upload_to_hub', 'headers': []}, receive) + request.state.token = 'owner' + request.state.session_id = '' + return request + + +@pytest.mark.asyncio +async def test_upload_after_recycle_is_cpu_only_and_preserves_future_ownership(monkeypatch, tmp_path): + router = TrainingRouter.__new__(TrainingRouter) + router.state = ServerState(backend=FakeBackend()) + router.store = TrainingStore(router.state, 'upload-test') + router.replica_id = 'router' + await router.state.register_replica(router.replica_id, 0) + await router.state.touch_replica_last_seen(router.replica_id) + router._init_task_queue() + router.controller = Mock(side_effect=AssertionError('Upload must not create or route a training')) + upload = Mock(return_value=None) + monkeypatch.setattr(HubOperation, 'push_to_hub', upload) + try: + response = await router.native_request(http_request({ + 'checkpoint_dir': str(tmp_path), 'hub_model_id': 'owner/saved', 'hub_token': 'hub-token', + }), 'upload_to_hub') + future = await settle(router.state, response.request_id) + assert future['status'] == 'completed' + upload.assert_called_once_with(repo_id='owner/saved', folder_path=str(tmp_path), token='hub-token', private=True) + assert await router.state._backend.get('future_owner::' + response.request_id) == 'owner' + assert not await router.state._backend.keys('model::*') + finally: + await router.shutdown_task_queue() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('rank', [8, 0, 33, True]) +async def test_native_lora_admission_does_not_import_peft(monkeypatch, rank): + import builtins + from unittest.mock import AsyncMock + from twinkle.server.exceptions import RequestRejectedError + router = TrainingRouter.__new__(TrainingRouter) + router.model = SimpleNamespace(name='small') + router.state = ServerState(backend=FakeBackend()) + await router.state.create_session({'session_id': 'session'}) + router.catalog = TrainingCatalog(TrainingConfig(models=[{ + 'name': 'small', 'model_id': 'weights', 'backend': 'transformers', 'device_type': 'cuda', + 'modes': {'lora': {'num_devices': 1, 'device_mesh': {'dp_size': 1}}}, + }])) + async def create(payload): + await router.state.store_future_status('created', 'pending', None) + return {'request_id': 'created'} + router.controller = SimpleNamespace(create=SimpleNamespace(remote=AsyncMock(side_effect=create))) + original_import = builtins.__import__ + def cpu_import(name, *args, **kwargs): + if name == 'peft' or name.startswith('peft.'): + raise AssertionError('CPU admission must not import PEFT') + return original_import(name, *args, **kwargs) + monkeypatch.setattr(builtins, '__import__', cpu_import) + request = http_request({'client_create_id': 'create', 'adapter_name': 'same', + 'config': json.dumps({'_TWINKLE_TYPE_': 'LoraConfig', 'r': rank})}) + request.state.session_id = 'session' + if type(rank) is int and rank == 8: + result = await router.native_request(request, 'add_adapter_to_model') + assert result.request_id == 'created' + router.controller.create.remote.assert_awaited_once() + else: + with pytest.raises(RequestRejectedError, match='rank'): + await router.native_request(request, 'add_adapter_to_model') + router.controller.create.remote.assert_not_awaited() diff --git a/tests/server/training/test_runtime_queue.py b/tests/server/training/test_runtime_queue.py new file mode 100644 index 000000000..0994e5ffe --- /dev/null +++ b/tests/server/training/test_runtime_queue.py @@ -0,0 +1,189 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Lifecycle protection through the actual compute worker and backend thread.""" +import asyncio +import threading +from types import SimpleNamespace + +import pytest + +from twinkle.server.task_queue import TaskQueueConfig, TaskQueueMixin +from twinkle.server.training.runtime import ManagedTrainingRuntime, _tinker_operation +from tests.server.training.test_controller import create, setup + + +class RPC: + def __init__(self, method): + self.method = method + + def remote(self, *args, **kwargs): + return asyncio.create_task(self.method(*args, **kwargs)) + + +class RuntimeQueue(ManagedTrainingRuntime, TaskQueueMixin): + def __init__(self, controller, state, instance): + self.state, self.replica_id = state, instance.instance_id + self.training_context = {'instance_id': instance.instance_id, 'epoch': controller.epoch} + self._model_unhealthy = False + self._training_controller = SimpleNamespace(**{ + name: RPC(getattr(controller, name)) + for name in ('admit', 'commit_operation', 'finish_operation', 'fail_instance') + }) + self._init_task_queue(TaskQueueConfig(execution_timeout=0.1, inline_fast_path_timeout=0.01), + enable_admission_gate=True) + + +async def settle(state, request_id): + for _ in range(300): + future = await state.get_future(request_id) + if future and future['status'] in ('completed', 'failed', 'cancelled'): + return future + await asyncio.sleep(0.01) + raise AssertionError('Future did not settle') + + +@pytest.mark.asyncio +async def test_execution_timeout_keeps_live_thread_and_allocation_quarantined(): + controller, state, runtime, _ = await setup(devices=2) + identity = await create(controller, state, 1) + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + entered, unblock = threading.Event(), threading.Event() + runtime.fail_shutdown = True + def blocking(): + entered.set() + unblock.wait(3) + async def operation(): + await queue.call_backend(blocking) + try: + ref = await queue.schedule_task(operation, model_id=identity, token='owner', task_type='forward_backward') + future = await settle(state, ref['request_id']) + assert entered.is_set() and not unblock.is_set() + assert future['failure']['reason_code'] == 'execution_timeout' + for _ in range(100): + if instance.state == 'QUARANTINED':break + await asyncio.sleep(0.01) + assert instance.state == 'QUARANTINED' and runtime.allocations + assert controller.trainings[identity].pending_operation_ids + assert queue._model_unhealthy + finally: + unblock.set() + await queue.shutdown_task_queue() + runtime.fail_shutdown = False + await controller.sweep() + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_tinker_seq_retry_runs_once_and_different_operations_do_not_collide(): + controller, state, runtime, _ = await setup() + identity = await create(controller, state, 1) + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + counter = [] + async def operation(): + counter.append('executed') + return {'ok': True} + context = _tinker_operation.set((identity, 1)) + try: + first, retry = await asyncio.gather(*( + queue.schedule_task(operation, model_id=identity, token='owner', task_type='forward_backward') + for _ in range(2))) + assert first == retry + assert (await settle(state, first['request_id']))['status'] == 'completed' + step = await queue.schedule_task(operation, model_id=identity, token='owner', task_type='optim_step') + assert step['request_id'] != first['request_id'] + await settle(state, step['request_id']) + assert len(counter) == 2 + finally: + _tinker_operation.reset(context) + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_cancelled_queued_operation_releases_ticket_without_business_renewal(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 1) + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + queue._task_queue_config = TaskQueueConfig(execution_timeout=2) + queue._compute_worker._config = queue._task_queue_config + entered, unblock = asyncio.Event(), asyncio.Event() + executions = [] + async def busy(): + entered.set() + await unblock.wait() + async def cancelled(): + executions.append('must not run') + try: + first = await queue.schedule_task(busy, model_id=identity, token='owner', task_type='forward_backward') + await entered.wait() + second = await queue.schedule_task(cancelled, model_id=identity, token='owner', task_type='step') + assert (await state.cancel_future(second['request_id']))['cancelled'] + clock.now = 9 + unblock.set() + await settle(state, first['request_id']) + for _ in range(100): + if not controller.trainings[identity].pending_operation_ids:break + await asyncio.sleep(0.01) + assert not executions and not controller.trainings[identity].pending_operation_ids + assert controller.trainings[identity].last_activity_at == 9 + assert (await state.get_future(second['request_id']))['status'] == 'cancelled' + finally: + unblock.set() + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_readiness_waits_for_real_ray_worker_constructor(): + import ray + import time + @ray.remote(num_cpus=0.1) + class Worker: + def __init__(self): + time.sleep(0.3) + self.loaded = True + def is_loaded(self): + return self.loaded + worker = Worker.remote() + runtime = ManagedTrainingRuntime() + runtime.model = SimpleNamespace(_actors=[worker]) + try: + await runtime.wait_for_training_workers() + assert await worker.is_loaded.remote() + assert ray._private.state.actors(worker._actor_id.hex())['State'] == 'ALIVE' + finally: + ray.kill(worker, no_restart=True) + + +@pytest.mark.asyncio +async def test_lost_finish_notification_is_reconciled_without_stopping_compute(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 'lost-notification') + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + async def lost(*args, **kwargs): + raise RuntimeError('Completion RPC lost') + queue._training_controller.finish_operation = RPC(lost) + async def status(instance, request_id): + return await queue.verify_operation(request_id) + runtime.verify_operation = status + async def compute(): + return {'updated': True} + try: + ref = await queue.schedule_task(compute, model_id=identity, token='owner', task_type='step') + assert (await settle(state, ref['request_id']))['status'] == 'completed' + assert controller.trainings[identity].pending_operation_ids + clock.now = 6 + await controller.sweep() + await controller.wait_for_jobs() + assert not controller.trainings[identity].pending_operation_ids + assert controller.trainings[identity].state == 'ACTIVE' and not queue._model_unhealthy + clock.now += 11 + await controller.sweep() + await controller.wait_for_jobs() + assert not runtime.allocations + finally: + await queue.shutdown_task_queue() + await controller.shutdown() diff --git a/tests/server/training/test_sequence.py b/tests/server/training/test_sequence.py new file mode 100644 index 000000000..904bc3df0 --- /dev/null +++ b/tests/server/training/test_sequence.py @@ -0,0 +1,70 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +import time +import uuid + +import pytest +import pytest_asyncio + +from twinkle.server.exceptions import RequestRejectedError +from twinkle.server.state import ServerState +from twinkle.server.state.backend.memory_backend import RayActorBackend +from twinkle.server.state.backend.redis_backend import RedisBackend + + +@pytest_asyncio.fixture(params=['memory', 'redis']) +async def state(request): + if request.param == 'redis': + url = os.environ.get('TWINKLE_TEST_REDIS_URL') + if not url: + pytest.skip('TWINKLE_TEST_REDIS_URL is required for the Redis regression') + backend = RedisBackend(url, key_prefix='sequence-test-' + uuid.uuid4().hex + ':') + else: + backend = RayActorBackend() + yield ServerState(backend=backend) + if request.param == 'redis': + for key in await backend.keys('*'): + await backend.delete(key) + await backend.close() + + +@pytest.mark.asyncio +async def test_expired_retry_receipts_cannot_execute_old_sequences_again(state, monkeypatch): + backend = state.backend + now = time.time() + monkeypatch.setattr('twinkle.server.state.server_state.time.time', lambda: now) + assert await state.claim_training_seq('training', 'native::session::step', 1, 'first', 10) is None + assert await state.claim_training_seq('training', 'native::session::step', 1, 'retry', 10) == 'first' + now += 11 + assert await state.claim_training_seq('training', 'native::session::step', 2, 'second', 10) is None + with pytest.raises(RequestRejectedError, match='retry window expired'): + await state.claim_training_seq('training', 'native::session::step', 1, 'late', 10) + ledger = await backend.get('seq_ledger::training::native::session::step') + assert list(ledger['receipts']) == ['2'] + assert len(await backend.keys('seq_ledger::*')) == 1 + + +@pytest.mark.asyncio +async def test_out_of_order_sequences_and_preflight_release_remain_retryable(state): + assert await state.claim_training_seq('training', 'tinker::step', 2, 'second', 10) is None + assert await state.claim_training_seq('training', 'tinker::step', 1, 'first', 10) is None + await state.release_training_seq('training', 'tinker::step', 1, 'another-request') + assert await state.claim_training_seq('training', 'tinker::step', 1, 'retry', 10) == 'first' + await state.release_training_seq('training', 'tinker::step', 1, 'first') + assert await state.claim_training_seq('training', 'tinker::step', 1, 'retry', 10) is None + assert await state.claim_training_seq('other-training', 'tinker::step', 1, 'other', 10) is None + + +@pytest.mark.asyncio +async def test_retiring_a_handle_removes_native_and_tinker_stream_ledgers(state): + from tests.server.training.test_controller import create, setup + controller, state, _, _ = await setup(backend=state.backend) + try: + identity = await create(controller, state, 'dedup-retirement') + for stream in ('native::session::step', 'tinker::optim_step'): + await state.claim_training_seq(identity, stream, 1, stream, 10) + assert len(await state.backend.keys(f'seq_ledger::{identity}::*')) == 2 + await controller.release(identity, 'owner', 'session') + assert await state.backend.keys(f'seq_ledger::{identity}::*') == [] + finally: + await controller.shutdown() diff --git a/tests/server/training/test_storage_recovery.py b/tests/server/training/test_storage_recovery.py new file mode 100644 index 000000000..6740348de --- /dev/null +++ b/tests/server/training/test_storage_recovery.py @@ -0,0 +1,64 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""The lifecycle uses the same recovery semantics on real memory and Redis backends.""" +import os +import uuid + +import pytest + +from twinkle.server.state.backend.memory_backend import RayActorBackend +from twinkle.server.state.backend.redis_backend import RedisBackend +from twinkle.server.training.controller import TrainingController +from twinkle.server.training.store import TrainingStore +from tests.server.training.test_controller import create, request, setup + + +@pytest.mark.asyncio +@pytest.mark.parametrize('mode', ['memory', 'redis']) +async def test_pending_operation_recovery_preserves_creation_identity_and_quota(mode): + if mode == 'redis' and not os.environ.get('TWINKLE_TEST_REDIS_URL'): + pytest.skip('Set TWINKLE_TEST_REDIS_URL to a dedicated local Redis') + backend = (RayActorBackend() if mode == 'memory' else + RedisBackend(os.environ['TWINKLE_TEST_REDIS_URL'], key_prefix='training-test-' + uuid.uuid4().hex + ':')) + try: + controller, state, runtime, clock = await setup(backend=backend) + identity = await create(controller, state, 'retry') + original = await controller.create(request('retry')) + await controller.admit(identity, 'owner', 'session', 'queued') + await controller.store.future('queued', 'queued', identity) + replacement = TrainingController(controller.config, TrainingStore(state, 'unit'), runtime, clock=clock) + await replacement.initialize() + assert (await state.get_future('queued'))['failure']['reason_code'] == 'orphaned_replica' + assert await replacement.create(request('retry')) == original + assert not runtime.allocations + assert await state.get_model_metadata(identity) is None + assert (await backend.get('token_count::owner') or 0) == 0 + new = await create(replacement, state, 'new') + assert new != identity + assert await backend.get('token_count::owner') == 1 + await replacement.shutdown() + finally: + if mode == 'redis': + for key in await backend.keys('*'): + await backend.delete(key) + await backend.close() + + +@pytest.mark.asyncio +async def test_managed_memory_actor_restart_reconnects_existing_handles(): + import asyncio + import ray + backend = RayActorBackend(namespace='memory-recovery-' + uuid.uuid4().hex) + try: + await backend.set('before', 'old') + ray.kill(backend._actor, no_restart=False) + for _ in range(100): + try: + if await backend.health_check() and await backend.get('before') is None:break + except ray.exceptions.RayError:pass + await asyncio.sleep(0.1) + else:raise AssertionError('State actor did not restart') + assert await backend.get('before') is None + await backend.set('after', 'new') + assert await backend.get('after') == 'new' + finally: + ray.kill(backend._actor, no_restart=True) diff --git a/tests/transformers/test_resident_training_reset.py b/tests/transformers/test_resident_training_reset.py new file mode 100644 index 000000000..c35570071 --- /dev/null +++ b/tests/transformers/test_resident_training_reset.py @@ -0,0 +1,39 @@ +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from twinkle.model.transformers.transformers import TransformersModel + + +def test_resident_reset_removes_previous_templates_forward_hook(): + class ScalingTemplate: + def __init__(self, scale): + self.scale = scale + + def pre_forward_hook(self, model, args, kwargs): + return (args[0] * self.scale,), kwargs + + def group(scale): + return SimpleNamespace(template=ScalingTemplate(scale), _handler=None) + + model = torch.nn.Linear(1, 1, bias=False) + model.weight.data.fill_(1) + wrapper = object.__new__(TransformersModel) + wrapper.__dict__['model'] = model + wrapper.model_id = 'local-reset-test' + wrapper._model_wrapped = True + wrapper.strategy = SimpleNamespace(unwrap_model=lambda value: value, + load_full_state_dict=lambda value, state: value.load_state_dict(state)) + wrapper.optimizer_group = {'': group(3)} + object.__setattr__(wrapper, '_construct_default_optimizer_group', lambda: group(1)) + wrapper.register_mm_forward_hook(wrapper.optimizer_group['']) + assert model(torch.ones(1, 1)).item() == 3 + initial = {'weight': torch.ones(1, 1)} + with patch('twinkle.model.transformers.transformers._read_hf_state_dict', return_value=initial): + for _ in range(2): + model.weight.grad = torch.ones_like(model.weight) + wrapper.reload_initial_weights(reset_training_state=True) + assert model(torch.ones(1, 1)).item() == 1 + assert model.weight.grad is None + assert len(model._forward_pre_hooks) == 1 diff --git a/tests/twinkle_client/test_async_components.py b/tests/twinkle_client/test_async_components.py index 0a092627c..313986191 100644 --- a/tests/twinkle_client/test_async_components.py +++ b/tests/twinkle_client/test_async_components.py @@ -41,6 +41,10 @@ def _recorder(calls, result_factory): def post(url, headers=None, data=None, timeout=None, **kwargs): body = json.loads(data) if data else kwargs.get('json') or {} calls.append((url, body)) + if '/model/' in url and url.endswith('/create'): + return _Response({'training_handles': True}) + if url.endswith('/add_adapter_to_model'): + return _Response(_completed({'training_id': 'training-test', 'adapter_name': body['adapter_name']})) return _Response(result_factory(url)) return post @@ -65,7 +69,7 @@ def test_model_forward_backward_sends_multiple_data_refs(monkeypatch) -> None: }})) model = module.MultiLoraTransformersModel('ms://base') - model.adapter_name = 'adapter' + model.add_adapter_to_model('adapter', None) refs = [ DataRef(ref_id='data-1', size=2, fields=['train_input']), DataRef(ref_id='data-2', size=2, fields=['train_input']), @@ -91,7 +95,7 @@ def test_model_inline_forward_methods_keep_the_original_endpoints(monkeypatch) - _patch_transport(monkeypatch, calls, lambda url: {} if url.endswith('/create') else _completed({'result': {}})) model = module.MultiLoraTransformersModel('ms://base') - model.adapter_name = 'adapter' + model.add_adapter_to_model('adapter', None) inputs = [{'input_ids': [1, 2]}] model.forward(inputs, return_logits=True) model.forward_only(inputs, disable_lora=True) @@ -140,7 +144,7 @@ def test_undeclared_forward_arguments_are_routed_to_loss_kwargs(monkeypatch) -> _patch_transport(monkeypatch, calls, lambda url: {} if url.endswith('/create') else _completed({'result': {}})) model = module.MultiLoraTransformersModel('ms://base') - model.adapter_name = 'adapter' + model.add_adapter_to_model('adapter', None) model.forward_backward([{'input_ids': [1, 2]}], advantages=[0.5], old_logps=[[-1.0, -2.0]]) _, body = calls[-1] @@ -158,7 +162,7 @@ def test_model_data_plane_forward_uses_a_separate_api(monkeypatch) -> None: }})) model = module.MultiLoraTransformersModel('ms://base') - model.adapter_name = 'adapter' + model.add_adapter_to_model('adapter', None) ref = DataRef(ref_id='data-1', size=2, fields=['train_input']) model.forward_from_data_plane(ref, input_field='train_input') @@ -180,7 +184,7 @@ def test_model_data_plane_forward_only_can_append_selected_outputs(monkeypatch) if url.endswith('/create') else _completed({'result': updated_ref.model_dump()})) model = module.MultiLoraTransformersModel('ms://base') - model.adapter_name = 'adapter' + model.add_adapter_to_model('adapter', None) result = model.forward_only_from_data_plane( ref, output_ref=ref, @@ -223,3 +227,46 @@ def test_sampler_async_data_plane_path_returns_reference_without_materializing(m assert url.endswith('/sampler/base/twinkle/sample_to_data_plane') assert body['num_samples'] == 4 assert body['group_ids'] == ['group-1'] + + +def test_old_remove_completion_does_not_clear_new_same_name_binding(): + import threading + from twinkle_client.http import ClientContext, ClientTransport + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + entered, resume = threading.Event(), threading.Event() + class Transport(ClientTransport): + def __init__(self): + super().__init__(ClientContext(base_url='http://server', api_key='owner')) + self.created = 0 + self.operations = [] + def post(self, url, json_data): + return _Response({'training_handles': True}) + def post_model(self, url, body): + self.operations.append((url.rsplit('/', 1)[-1], body.model_dump())) + if url.endswith('/add_adapter_to_model'): + self.created += 1 + return _Response(_completed({'training_id': f'training-{self.created}', 'adapter_name': 'same'})) + if url.endswith('/remove_adapter'): + entered.set() + assert resume.wait(3) + return _Response(_completed(None)) + transport = Transport() + model = MultiLoraTransformersModel('small', transport=transport) + model.add_adapter_to_model('same', None) + errors = [] + def remove(): + try:model.remove_adapter() + except Exception as exc:errors.append(exc) + releasing = threading.Thread(target=remove) + releasing.start() + try: + assert entered.wait(3) + model.add_adapter_to_model('same', None) + finally: + resume.set() + releasing.join(3) + assert not releasing.is_alive() and not errors + model.zero_grad() + assert transport.operations[-1][1]['training_id'] == 'training-2' + removal = next(body for endpoint, body in transport.operations if endpoint == 'remove_adapter') + assert removal['training_id'] == 'training-1' diff --git a/tools/validation/on_demand_training/allocation_probe.py b/tools/validation/on_demand_training/allocation_probe.py new file mode 100644 index 000000000..633785610 --- /dev/null +++ b/tools/validation/on_demand_training/allocation_probe.py @@ -0,0 +1,88 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Run in its own process: test production allocation against logical Ray GPUs.""" +import asyncio +import json +import ray +import uuid + +from twinkle import DeviceGroup +from twinkle.infra._ray.resource_manager import ResourceManager +from twinkle.server.config.training import TrainingConfig +from twinkle.server.training.allocation import DevicePool +from twinkle.server.training.catalog import TrainingCatalog +from twinkle.server.training.errors import CapacityUnavailableError + + +async def main(): + namespace = 'alloc_' + uuid.uuid4().hex + ray.init( + address='local', namespace=namespace, num_cpus=16, num_gpus=8, include_dashboard=False, log_to_driver=False) + pool = DevicePool(namespace, allocation_timeout=3, shutdown_timeout=10) + allocations = [] + try: + for index, count in enumerate((2, 2, 4)): + config = TrainingConfig(models=[{ + 'name': str(index), + 'model_id': str(index), + 'backend': 'mock', + 'device_type': 'cuda', + 'modes': { + 'full': { + 'num_devices': count, + 'device_mesh': { + 'fsdp_size': count + } + } + }, + }]) + spec = TrainingCatalog(config).get(str(index), 'full') + allocation = await pool.allocate(spec, 'epoch', str(index)) + allocations.append(allocation) + devices = [device for allocation in allocations for device in allocation.devices] + assert len(set(devices)) == 8, devices + assert len({allocation.node_id for allocation in allocations}) == 1 + assert ray.available_resources().get('GPU', 0) == 0 + try: + await pool.allocate(spec, 'epoch', 'overflow') + raise AssertionError('over-allocation was admitted') + except CapacityUnavailableError: + pass + from ray.util.placement_group import get_placement_group + group = DeviceGroup(name='owned', ranks=2, device_type='cuda') + view = ResourceManager.from_allocation(group, get_placement_group(allocations[0].name), + list(allocations[0].devices)) + assert view.get_config('owned').ranks == [0, 1] + assert [entry['gpu_rank'][0] for entry in view.get_group('owned')] == list(allocations[0].devices) + view.destroy_placement_group() + assert ray.available_resources().get('GPU', 0) == 0 + freed = allocations[0] + await pool.remove(freed.name) + two = TrainingCatalog( + TrainingConfig(models=[{ + 'name': 'replacement', + 'model_id': 'replacement', + 'backend': 'mock', + 'device_type': 'cuda', + 'modes': { + 'full': { + 'num_devices': 2, + 'device_mesh': { + 'fsdp_size': 2 + } + } + }, + }])).get('replacement', 'full') + replacement = await pool.allocate(two, 'epoch', 'replacement') + allocations.append(replacement) + assert set(replacement.devices) == set(freed.devices) + print(json.dumps({'passed': True, 'devices': devices, 'reused': list(replacement.devices)})) + finally: + try: + for allocation in allocations: + await pool.remove(allocation.name) + finally: + ray.shutdown() + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/tools/validation/on_demand_training/device_validation.py b/tools/validation/on_demand_training/device_validation.py new file mode 100644 index 000000000..b496c816f --- /dev/null +++ b/tools/validation/on_demand_training/device_validation.py @@ -0,0 +1,440 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Explicit real-device acceptance; run this file only after local mock checks. + +Required: TWINKLE_DEVICE_BACKEND, TWINKLE_DEVICE_MODEL, TWINKLE_DEVICE_OTHER_MODEL, +TWINKLE_DEVICE_RUN_DIR and four exclusively available visible devices. +This is intentionally outside pytest's default test filename pattern. +""" +import json +import numpy as np +import os +import pytest +import ray +import requests +import socket +import time +import uuid +from on_demand_sdk import unload_tinker_model +from pathlib import Path +from ray import serve + +from twinkle.server.config import ServerConfig +from twinkle.server.launcher.server_launcher import ServerLauncher +from twinkle_client import init_tinker_client +from twinkle_client.http import ClientTransport +from twinkle_client.http.context import ClientContext +from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + + +@pytest.fixture(scope='module') +def server(): + backend = os.environ['TWINKLE_DEVICE_BACKEND'] + assert backend in ('transformers', 'megatron') + run_dir = Path(os.environ['TWINKLE_DEVICE_RUN_DIR']) + run_dir.mkdir(parents=True, exist_ok=True) + namespace = 'device_' + uuid.uuid4().hex[:8] + with socket.socket() as listener: + listener.bind(('127.0.0.1', 0)) + port = listener.getsockname()[1] + os.environ['TWINKLE_STATE_RAY_NAMESPACE'] = namespace + os.environ['TWINKLE_DEFAULT_SAVE_DIR'] = str(run_dir / 'checkpoints') + os.environ['TWINKLE_PERSISTENCE_MODE'] = 'memory' + os.environ['TWINKLE_LONG_POLL_TIMEOUT'] = '2' + mesh = {'fsdp_size': 2} if backend == 'transformers' else {'tp_size': 2} + apps = [] + for alias, variable in [('small', 'TWINKLE_DEVICE_MODEL'), ('other', 'TWINKLE_DEVICE_OTHER_MODEL')]: + source = os.environ[variable] + assert Path(source, 'config.json').is_file() + mode = {'num_devices': 2, 'device_mesh': mesh, 'idle_timeout_s': 600 if alias == 'small' else 30} + if backend == 'transformers': + mode['strategy'] = 'native_fsdp' + apps.append({ + 'name': alias, + 'route_prefix': '/api/v1/model/' + alias, + 'import_path': 'model', + 'args': { + 'model_id': source, + 'backend': backend, + 'device_type': 'cuda', + 'modes': { + 'lora': { + **mode, 'max_loras': 2, + 'max_r': 8 + }, + 'full': mode + }, + } + }) + apps.append({ + 'name': 'gateway', + 'route_prefix': '/api/v1', + 'import_path': 'server', + 'args': { + 'server_config': { + 'per_token_model_limit': 10 + } + }, + 'deployments': [{ + 'ray_actor_options': { + 'num_cpus': 0.1 + } + }] + }) + os.environ['TWINKLE_MODEL_ID_ALIASES'] = json.dumps( + {app['name']: app['args']['model_id'] + for app in apps if app['import_path'] == 'model'}) + config = ServerConfig(ray_namespace=namespace, http_options={'host': '127.0.0.1', 'port': port}, applications=apps) + import yaml + (run_dir / 'server.yaml').write_text(yaml.safe_dump(config.to_yaml_dict(), sort_keys=False)) + ray.init( + address='local', + namespace=namespace, + num_cpus=16, + num_gpus=4, + _temp_dir='/tmp/twd-' + namespace[-8:], + include_dashboard=False, + log_to_driver=True) + launcher = ServerLauncher(config) + launcher._start_serve() + evidence = { + 'backend': backend, + 'namespace': namespace, + 'visible_devices': os.environ.get('CUDA_VISIBLE_DEVICES'), + 'cases': [] + } + try: + for app in config.applications: + launcher._deploy_application(app) + launcher._deploy_training() + controller = launcher._training_controller + base = f'http://127.0.0.1:{port}' + assert ray.get(controller.diagnostics.remote())['instances'] == [] + evidence['ray_address'] = ray.get_runtime_context().gcs_address + yield base, controller, evidence, launcher + finally: + if getattr(launcher, '_training_controller', None) is not None: + ray.get(launcher._training_controller.shutdown.remote(), timeout=180) + ray.kill(launcher._training_controller, no_restart=True) + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + serve.shutdown() + ray.shutdown() + + +class Handle: + + def __init__(self, base, protocol, mode, alias='small'): + self.base, self.protocol = base, protocol + self.token = 'device-validation-owner' + session = requests.post(base + '/api/v1/twinkle/create_session', json={}, timeout=20).json()['session_id'] + self.transport = ClientTransport(ClientContext(base_url=base, api_key=self.token, session_id=session)) + self.rate = 1e-3 if mode == 'lora' else 5e-5 + if protocol == 'native': + from peft import LoraConfig + self.client = MultiLoraTransformersModel(alias, transport=self.transport) + self.client.add_adapter_to_model('same', + LoraConfig(r=8, target_modules='all-linear') if mode == 'lora' else None) + self.identity = self.client._training_ids['same'] + self.client.set_loss('CrossEntropyLoss') + self.client.set_optimizer('Adam', lr=self.rate) + self.client.set_processor('InputProcessor', padding_side='right') + else: + init_tinker_client() + from tinker import ServiceClient + self.service = ServiceClient(base_url=base, api_key=self.token) + self.client = ( + self.service.create_lora_training_client(base_model=alias, rank=8) + if mode == 'lora' else self.service.create_full_training_client(base_model=alias)) + self.identity = self.client.model_id + # Fixed token IDs need no external dataset, and every position has a target. + self.inputs = [{ + 'input_ids': [100, 200, 300, 400, 500, 600, 700, 800], + 'labels': [200, 300, 400, 500, 600, 700, 800, 900], + 'position_ids': list(range(8)), + 'attention_mask': [1] * 8 + }] * 2 + + def datum(self): + from tinker import types + feature = self.inputs[0] + return [ + types.Datum( + model_input=types.ModelInput.from_ints(feature['input_ids']), + loss_fn_inputs={ + 'target_tokens': types.TensorData(data=feature['labels'], dtype='int64', shape=[8]), + 'weights': types.TensorData(data=[1.] * 8, dtype='float32', shape=[8]), + }) for _ in self.inputs + ] + + def evaluate(self): + if self.protocol == 'native': + result = self.client.forward_only(self.inputs).result + values = result['logps'] + else: + result = self.client.forward(self.datum(), 'cross_entropy').result() + values = result.loss_fn_outputs[0]['logprobs'].tolist() + values = np.asarray(values, dtype=float).reshape(-1) + assert values.size and np.isfinite(values).all(), values + return values + + def step(self, *, contrast=False): + labels = self.inputs[0]['labels'] + if contrast: + self.inputs[0]['labels'] = list(reversed(labels)) + try: + if self.protocol == 'native': + self.client.forward_backward(self.inputs) + self.client.clip_grad_and_step() + else: + from tinker import types + self.client.forward_backward(self.datum(), 'cross_entropy').result() + self.client.optim_step(types.AdamParams(learning_rate=self.rate)).result() + finally: + self.inputs[0]['labels'] = labels + + def save(self): + name = 'device-' + uuid.uuid4().hex[:8] + return (self.client.save(name, save_optimizer=True).twinkle_path + if self.protocol == 'native' else self.client.save_state(name).result().path) + + def restore(self, path): + if self.protocol == 'native': + self.client.resume_from_checkpoint(path) + else: + self.client.load_state_with_optimizer(path).result() + + def restore_weights(self, path): + if self.protocol == 'native': + self.client.resume_from_checkpoint(path, resume_only_model=True) + else: + self.client.load_state(path).result() + + def close(self): + if self.protocol == 'native': + self.client.remove_adapter() + else: + assert unload_tinker_model(self.service, self.identity).model_id == self.identity + self.transport.close() + + +def active_instances(controller): + return [row for row in ray.get(controller.diagnostics.remote())['instances'] if row['state'] != 'STOPPED'] + + +@pytest.mark.parametrize('protocol', ['native', 'tinker']) +@pytest.mark.parametrize('mode', ['lora', 'full']) +def test_training_isolation_checkpoint_and_device_reclamation(server, protocol, mode): + base, controller, evidence, _ = server + first, second = Handle(base, protocol, mode), Handle(base, protocol, mode) + assert first.identity != second.identity + instances = active_instances(controller) + assert len(instances) == (1 if mode == 'lora' else 2) + assert all(len(row['devices']) == 2 for row in instances) + devices = [device for row in instances for device in row['devices']] + assert len(set(devices)) == len(devices) + original, untouched = first.evaluate(), second.evaluate() + untouched_checkpoint = second.save() if mode == 'lora' else None + for _ in range(4): + first.step() + trained = first.evaluate() + assert float(trained.mean()) > float(original.mean()) + 1e-4, (original, trained) + np.testing.assert_allclose(second.evaluate(), untouched, rtol=0, atol=2e-3) + if mode == 'lora': + third, fourth = Handle(base, protocol, mode), Handle(base, protocol, mode) + assert len(active_instances(controller)) == 2 + assert ray.available_resources().get('GPU', 0) == 0 + with pytest.raises(Exception, match='(?i)allocation'): + Handle(base, protocol, mode) + fourth.close() + third.restore(untouched_checkpoint) + second.step() + third.step() + np.testing.assert_allclose(second.evaluate(), third.evaluate(), rtol=0, atol=2e-2) + np.testing.assert_allclose(first.evaluate(), trained, rtol=0, atol=2e-3) + third.close() + assert len(active_instances(controller)) == 1 + checkpoint = first.save() + first.step(contrast=True) + continuation = first.evaluate() + first.restore(checkpoint) + np.testing.assert_allclose(first.evaluate(), trained, rtol=0, atol=2e-3) + first.step(contrast=True) + np.testing.assert_allclose(first.evaluate(), continuation, rtol=0, atol=2e-2) + old_id = first.identity + first.close() + assert len(active_instances(controller)) == 1 + replacement = Handle(base, protocol, mode) + assert replacement.identity != old_id + replacement.restore_weights(checkpoint) + np.testing.assert_allclose(replacement.evaluate(), trained, rtol=0, atol=2e-3) + replacement.step(contrast=True) + optimizer_reset = replacement.evaluate() + assert float(np.max(np.abs(optimizer_reset - continuation))) > 2e-2, 'Optimizer continuation oracle is not decisive' + replacement.restore(checkpoint) + np.testing.assert_allclose(replacement.evaluate(), trained, rtol=0, atol=2e-3) + replacement.step(contrast=True) + np.testing.assert_allclose(replacement.evaluate(), continuation, rtol=0, atol=2e-2) + replacement.close() + second.close() + assert not active_instances(controller) + assert ray.available_resources().get('GPU', 0) == 4 + evidence['cases'].append({ + 'protocol': protocol, + 'mode': mode, + 'initial_logps': original.tolist(), + 'trained_logps': trained.tolist(), + 'continuation_logps': continuation.tolist(), + 'optimizer_reset_logps': optimizer_reset.tolist(), + 'checkpoint': checkpoint, + 'devices': devices, + 'passed': True + }) + + +def test_cross_model_reuse_and_heartbeat_does_not_prevent_expiry(server): + base, controller, evidence, _ = server + before = ray.available_resources().get('GPU', 0) + for index in range(10): + handle = Handle(base, 'native', 'full', alias='small' if index % 2 == 0 else 'other') + values = handle.evaluate() + handle.close() + assert not active_instances(controller) + assert ray.available_resources().get('GPU', 0) == before + print(json.dumps({'cycle': index + 1, 'logps': values.tolist()}), flush=True) + expired = Handle(base, 'native', 'lora', alias='other') + expired.evaluate() + deadline = time.monotonic() + 60 + while active_instances(controller) and time.monotonic() < deadline: + response = expired.transport.post( + base + '/api/v1/twinkle/session_heartbeat', json_data={'session_id': expired.transport.context.session_id}) + response.raise_for_status() + time.sleep(2) + assert not active_instances(controller), 'Heartbeats incorrectly kept idle training alive' + with pytest.raises(Exception): + expired.client.zero_grad() + assert ray.available_resources().get('GPU', 0) == before + expired.transport.close() + evidence['cases'].append({'cycles': 10, 'idle_reclaimed': True, 'passed': True}) + + +@pytest.mark.parametrize('protocol', ['native', 'tinker']) +def test_idle_lora_user_expires_without_resetting_active_user(server, protocol): + base, controller, evidence, _ = server + idle = Handle(base, protocol, 'lora', alias='other') + active = Handle(base, protocol, 'lora', alias='other') + instance = active_instances(controller)[0] + assert instance['training_ids'] == [idle.identity, active.identity] + original = idle.evaluate() + for _ in range(4): + active.step() + trained = active.evaluate() + assert trained.mean() > original.mean() + 1e-4 + np.testing.assert_allclose(idle.evaluate(), original, rtol=0, atol=2e-3) + checkpoint = active.save() + active.step(contrast=True) + expected = active.evaluate() + active.restore(checkpoint) + deadline = time.monotonic() + 60 + while True: + np.testing.assert_allclose(active.evaluate(), trained, rtol=0, atol=2e-3) + response = idle.transport.post( + base + '/api/v1/twinkle/session_heartbeat', json_data={'session_id': idle.transport.context.session_id}) + response.raise_for_status() + instances = active_instances(controller) + assert len(instances) == 1 and instances[0]['instance_id'] == instance['instance_id'] + if idle.identity not in instances[0]['training_ids']: + break + assert time.monotonic() < deadline, 'Idle LoRA user did not expire while the other user stayed active' + time.sleep(2) + with pytest.raises(Exception): + idle.evaluate() + active.step(contrast=True) + np.testing.assert_allclose(active.evaluate(), expected, rtol=0, atol=2e-2) + replacement = Handle(base, protocol, 'lora', alias='other') + assert replacement.identity != idle.identity + assert active_instances(controller)[0]['instance_id'] == instance['instance_id'] + np.testing.assert_allclose(replacement.evaluate(), original, rtol=0, atol=2e-3) + replacement.close() + np.testing.assert_allclose(active.evaluate(), expected, rtol=0, atol=2e-3) + active.close() + idle.transport.close() + assert not active_instances(controller) + assert ray.available_resources().get('GPU', 0) == 4 + evidence['cases'].append({ + 'protocol': protocol, + 'idle_lora_expired': True, + 'active_optimizer_preserved': True, + 'slot_reused': True, + 'passed': True + }) + + +def test_resident_full_reset_preserves_other_instances(server): + base, controller, evidence, launcher = server + raw = launcher.config.to_yaml_dict() + raw['applications'][0]['args']['modes']['full']['min_replicas'] = 1 + ray.get(controller.shutdown.remote(), timeout=180) + ray.kill(controller, no_restart=True) + from twinkle.server.training.actor import controller_name + deadline = time.monotonic() + 30 + while True: + try: + ray.get_actor(controller_name(launcher.config.ray_namespace), namespace=launcher.config.ray_namespace) + except ValueError: + break + assert time.monotonic() < deadline, 'Previous controller name was not released' + time.sleep(0.1) + replacement = ServerLauncher(ServerConfig.model_validate(raw)) + replacement._deploy_training() + controller = replacement._training_controller + launcher._training_controller = controller + evidence['resident_controller_epoch'] = ray.get(controller.diagnostics.remote())['epoch'] + resident, extra = Handle(base, 'native', 'full'), Handle(base, 'native', 'full') + instances = active_instances(controller) + resident_instance = next(row for row in instances if row['resident']) + original, untouched = resident.evaluate(), extra.evaluate() + for _ in range(2): + resident.step() + assert resident.evaluate().mean() > original.mean() + 1e-4 + resident.close() + restored = Handle(base, 'native', 'full') + instances = active_instances(controller) + assert next(row for row in instances if row['resident'])['instance_id'] == resident_instance['instance_id'] + np.testing.assert_allclose(restored.evaluate(), original, rtol=0, atol=2e-3) + np.testing.assert_allclose(extra.evaluate(), untouched, rtol=0, atol=2e-3) + restored.close() + extra.close() + instances = active_instances(controller) + assert len(instances) == 1 and instances[0]['resident'] + assert ray.available_resources().get('GPU', 0) == 2 + evidence['cases'].append({'resident_reset': True, 'extra_preserved': True, 'passed': True}) + + +def test_worker_failure_reclaims_devices_and_preserves_resident(server): + from ray.util import list_named_actors + base, _, evidence, launcher = server + controller = launcher._training_controller + resident = Handle(base, 'native', 'full') + original = resident.evaluate() + failing = Handle(base, 'native', 'lora') + failing.evaluate() + instance = next(row for row in active_instances(controller) if row['mode'] == 'lora') + actors = [ + row for row in list_named_actors(all_namespaces=True) if row['namespace'] == evidence['namespace'] + and instance['instance_id'] in row['name'] and not row['name'].endswith('-registry') + ] + assert len(actors) == 2, actors + ray.kill(ray.get_actor(actors[0]['name'], namespace=evidence['namespace']), no_restart=True) + deadline = time.monotonic() + 90 + while any(row['instance_id'] == instance['instance_id'] for row in active_instances(controller)): + assert time.monotonic() < deadline, 'Failed runtime devices were not reclaimed' + time.sleep(1) + with pytest.raises(Exception): + failing.client.zero_grad() + np.testing.assert_allclose(resident.evaluate(), original, rtol=0, atol=2e-3) + recovered = Handle(base, 'native', 'lora') + recovered.evaluate() + recovered.close() + resident.close() + failing.transport.close() + assert ray.available_resources().get('GPU', 0) == 2 + evidence['cases'].append({'worker_failure': True, 'resident_preserved': True, 'recreated': True, 'passed': True}) diff --git a/tools/validation/on_demand_training/fragmentation_probe.py b/tools/validation/on_demand_training/fragmentation_probe.py new file mode 100644 index 000000000..02dd22ab8 --- /dev/null +++ b/tools/validation/on_demand_training/fragmentation_probe.py @@ -0,0 +1,84 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Exercise single-node admission with real Ray nodes and logical devices.""" +import asyncio +import json +import ray +import uuid +from ray.cluster_utils import Cluster +from ray.util.placement_group import placement_group, placement_group_table, remove_placement_group + +from twinkle.server.config.training import TrainingConfig +from twinkle.server.training.allocation import DevicePool +from twinkle.server.training.catalog import TrainingCatalog +from twinkle.server.training.errors import CapacityUnavailableError + + +async def main(): + cluster = Cluster() + occupied = [] + allocated = None + pool = None + try: + cluster.add_node(num_cpus=4, num_gpus=4, resources={'left': 1}, include_dashboard=False) + cluster.add_node(num_cpus=4, num_gpus=4, resources={'right': 1}) + ray.init(address=cluster.address, namespace='fragmentation_' + uuid.uuid4().hex, log_to_driver=False) + pool = DevicePool(ray.get_runtime_context().namespace, allocation_timeout=1, shutdown_timeout=10) + spec = TrainingCatalog( + TrainingConfig(models=[{ + 'name': 'two', + 'model_id': 'two', + 'backend': 'mock', + 'device_type': 'cuda', + 'modes': { + 'full': { + 'num_devices': 2, + 'device_mesh': { + 'fsdp_size': 2 + } + } + }, + }])).get('two', 'full') + for marker in ('left', 'right'): + pg = placement_group([{'GPU': 3, marker: 0.01}], strategy='STRICT_PACK') + occupied.append(pg) + await asyncio.wait_for(pg.ready(), 20) + assert ray.available_resources()['GPU'] == 2 + try: + await pool.allocate(spec, 'epoch', 'fragmented') + except CapacityUnavailableError: + pass + else: + raise AssertionError('Two nodes with one free GPU each admitted a two-GPU instance') + groups = placement_group_table().values() + assert not any(group['name'] == pool.name('epoch', 'fragmented') and group['state'] != 'REMOVED' + for group in groups) + remove_placement_group(occupied[0]) + for _ in range(200): + if ray.available_resources().get('GPU', 0) == 5: + break + await asyncio.sleep(0.05) + else: + raise AssertionError('Owned GPU reservation was not released') + allocated = await pool.allocate(spec, 'epoch', 'available') + assert len(set(allocated.devices)) == 2 + print( + json.dumps({ + 'passed': True, + 'nodes': len(ray.nodes()), + 'fragmented_rejected': True, + 'pending_removed': True, + 'devices': allocated.devices + })) + finally: + try: + if allocated is not None: + await pool.remove(allocated.name) + for pg in occupied: + remove_placement_group(pg) + finally: + ray.shutdown() + cluster.shutdown() + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/tools/validation/on_demand_training/http_probe.py b/tools/validation/on_demand_training/http_probe.py new file mode 100644 index 000000000..5f4e765c0 --- /dev/null +++ b/tools/validation/on_demand_training/http_probe.py @@ -0,0 +1,246 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""HTTP integration with both installed SDKs and a CPU mock model.""" +import json +import os +import ray +import requests +import socket +import tempfile +import time +import uuid +from on_demand_sdk import unload_tinker_model +from ray import serve + +from twinkle.server.config import ServerConfig +from twinkle.server.launcher.server_launcher import ServerLauncher +from twinkle_client import init_tinker_client +from twinkle_client.http import ClientTransport +from twinkle_client.http.context import ClientContext +from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + + +def main(): + namespace = 'http_' + uuid.uuid4().hex[:8] + with socket.socket() as listener: + listener.bind(('127.0.0.1', 0)) + port = listener.getsockname()[1] + os.environ['TWINKLE_STATE_RAY_NAMESPACE'] = namespace + os.environ['TWINKLE_DEFAULT_SAVE_DIR'] = tempfile.mkdtemp(prefix='twinkle-training-http-') + os.environ['TWINKLE_LONG_POLL_TIMEOUT'] = '2' + os.environ['TWINKLE_PERSISTENCE_MODE'] = 'memory' + os.environ['PYTHONPATH'] = str(__import__('pathlib').Path.cwd() / 'src') + ray.init(address='local', namespace=namespace, num_cpus=16, include_dashboard=False, log_to_driver=False) + config = ServerConfig.model_validate({ + 'ray_namespace': + namespace, + 'http_options': { + 'host': '127.0.0.1', + 'port': port + }, + 'training': { + 'startup_timeout_s': 120, + 'shutdown_timeout_s': 30 + }, + 'applications': [{ + 'name': 'small', + 'import_path': 'model', + 'route_prefix': '/api/v1/model/small', + 'args': { + 'model_id': 'small', + 'backend': 'mock', + 'device_type': 'cpu', + 'modes': { + 'lora': { + 'num_devices': 1, + 'device_mesh': { + 'dp_size': 1 + }, + 'max_loras': 2, + 'max_replicas': 1 + }, + 'full': { + 'num_devices': 1, + 'device_mesh': { + 'dp_size': 1 + } + } + } + }, + }, { + 'name': 'other', + 'import_path': 'model', + 'route_prefix': '/api/v1/model/other', + 'args': { + 'model_id': 'other', + 'backend': 'mock', + 'device_type': 'cpu', + 'modes': { + 'full': { + 'num_devices': 1, + 'device_mesh': { + 'dp_size': 1 + } + } + } + }, + }, { + 'name': 'gateway', + 'import_path': 'server', + 'route_prefix': '/api/v1', + 'deployments': [{ + 'ray_actor_options': { + 'num_cpus': 0.1 + } + }] + }], + }) + launcher = ServerLauncher(config) + launcher._start_serve() + base = f'http://127.0.0.1:{port}' + + def wait_epoch(controller, old_epoch): + deadline = time.monotonic() + 90 + last_error = None + while time.monotonic() < deadline: + try: + diagnostic = ray.get(controller.diagnostics.remote(), timeout=2) + if diagnostic['epoch'] != old_epoch: + return diagnostic + except ray.exceptions.RayError as exc: + last_error = exc + time.sleep(0.1) + raise AssertionError(f'Controller did not recover: {last_error}') + + try: + for spec in config.applications: + if spec.import_path != 'model': + launcher._deploy_application(spec) + launcher._deploy_training() + controller = launcher._training_controller + assert ray.get(controller.diagnostics.remote())['instances'] == [] + session = requests.post(base + '/api/v1/twinkle/create_session', json={}, timeout=10).json()['session_id'] + transport = ClientTransport(ClientContext(base_url=base, api_key='probe-owner', session_id=session)) + first = MultiLoraTransformersModel('small', transport=transport) + second = MultiLoraTransformersModel('small', transport=transport) + first.add_adapter_to_model('same', None) + instance = next( + item for item in ray.get(controller.diagnostics.remote())['instances'] + if first._training_ids['same'] in item['training_ids']) + bypass = requests.post( + base + f"/api/v1/training-runtime/{instance['instance_id']}/twinkle/add_adapter_to_model", + json={ + 'adapter_name': 'bypass', + 'config': None + }, + headers={'Authorization': 'Bearer probe-owner'}, + timeout=10) + assert bypass.status_code == 400 and 'catalog endpoint' in bypass.text, bypass.text + second.add_adapter_to_model('same', None) + old_id = first._training_ids['same'] + assert old_id != second._training_ids['same'] + assert json.loads(first.get_train_configs().result) == {} + first.zero_grad() + first.remove_adapter() + first.add_adapter_to_model('same', None) + assert old_id != first._training_ids['same'] + first.remove_adapter() + second.remove_adapter() + from peft import LoraConfig + native_lora = MultiLoraTransformersModel('small', transport=transport) + native_lora.add_adapter_to_model('native-lora', LoraConfig(r=8, target_modules='all-linear')) + native_lora.zero_grad() + native_lora.remove_adapter() + init_tinker_client() + from tinker import ServiceClient + service = ServiceClient(base_url=base, api_key='probe-tinker') + shared1 = service.create_lora_training_client(base_model='small', rank=8) + shared2 = service.create_lora_training_client(base_model='small', rank=8) + assert shared1.model_id != shared2.model_id + assert shared1.get_info().is_lora + assert shared2.get_info().model_name == 'small' + from tinker import types + datum = types.Datum( + model_input=types.ModelInput.from_ints([1, 2, 3]), + loss_fn_inputs={ + 'target_tokens': types.TensorData(data=[2, 3, 4], dtype='int64', shape=[3]), + 'weights': types.TensorData(data=[1.0, 1.0, 1.0], dtype='float32', shape=[3]), + }) + assert shared1.forward_backward([datum], 'cross_entropy').result().metrics + shared1.optim_step(types.AdamParams(learning_rate=1e-4)).result() + exclusive = service.create_full_training_client(base_model='small') + assert not exclusive.get_info().is_lora + assert exclusive.forward_backward([datum], 'cross_entropy').result().metrics + exclusive.optim_step(types.AdamParams(learning_rate=1e-4)).result() + assert unload_tinker_model(service, exclusive.model_id).model_id == exclusive.model_id + assert shared1.get_info().is_lora and shared2.get_info().is_lora + try: + service.create_lora_training_client(base_model='small', rank=8) + raise AssertionError('A third LoRA user bypassed configured capacity') + except Exception as exc: + assert 'replica limit' in str(exc).lower(), exc + assert unload_tinker_model(service, shared2.model_id).model_id == shared2.model_id + replacement = service.create_lora_training_client(base_model='small', rank=8) + assert replacement.model_id != shared2.model_id + assert shared1.get_info().is_lora + old_epoch = ray.get(controller.diagnostics.remote())['epoch'] + ray.kill(controller, no_restart=False) + diagnostic = wait_epoch(controller, old_epoch) + assert not [item for item in diagnostic['instances'] if item['state'] != 'STOPPED'] + stale = requests.post( + base + '/api/v1/get_info', + json={ + 'model_id': shared1.model_id, + 'type': 'get_info' + }, + headers={ + 'Authorization': 'Bearer probe-tinker', + 'x-request-id': uuid.uuid4().hex + }, + timeout=10) + assert stale.status_code == 404, stale.text + before = ray.available_resources().get('CPU', 0) + for index in range(10): + model = MultiLoraTransformersModel('small' if index % 2 == 0 else 'other', transport=transport) + model.add_adapter_to_model('cycle', None) + model.zero_grad() + model.remove_adapter() + print(json.dumps({'phase': 'cycles', 'completed': index + 1}), flush=True) + assert ray.available_resources().get('CPU', 0) >= before - 0.02 + final = ray.get(controller.diagnostics.remote()) + assert all(item['state'] == 'STOPPED' for item in final['instances']) + live = MultiLoraTransformersModel('other', transport=transport) + live.add_adapter_to_model('state-loss', None) + ray.kill(ray.get_actor('twinkle_state_actor', namespace=namespace), no_restart=False) + after_loss = wait_epoch(controller, final['epoch']) + assert not [item for item in after_loss['instances'] if item['state'] != 'STOPPED'] + session = requests.post(base + '/api/v1/twinkle/create_session', json={}, timeout=10).json()['session_id'] + fresh_transport = ClientTransport(ClientContext(base_url=base, api_key='probe-owner', session_id=session)) + fresh = MultiLoraTransformersModel('other', transport=fresh_transport) + fresh.add_adapter_to_model('state-loss', None) + assert fresh._training_ids['state-loss'] != live._training_ids['state-loss'] + fresh.remove_adapter() + fresh_transport.close() + print( + json.dumps({ + 'passed': True, + 'native_recreation': True, + 'native_lora': True, + 'tinker_shared': True, + 'tinker_full': True, + 'restart': True, + 'state_loss': True, + 'runtime_lifecycle_rejected': True, + 'cycles': 10, + 'instances': len(final['instances']) + })) + transport.close() + finally: + if getattr(launcher, '_training_controller', None) is not None: + ray.get(launcher._training_controller.shutdown.remote(), timeout=60) + ray.kill(launcher._training_controller, no_restart=True) + serve.shutdown() + ray.shutdown() + + +if __name__ == '__main__': + main() diff --git a/tools/validation/on_demand_training/on_demand_sdk.py b/tools/validation/on_demand_training/on_demand_sdk.py new file mode 100644 index 000000000..d8e30f8aa --- /dev/null +++ b/tools/validation/on_demand_training/on_demand_sdk.py @@ -0,0 +1,18 @@ +"""SDK operations shared by the mock and real-device acceptance harnesses.""" +import time + + +def unload_tinker_model(service, model_id): + from tinker import types + from tinker.lib.api_future_impl import _APIFuture + from tinker.lib.client_connection_pool_type import ClientConnectionPoolType + + async def unload(): + started = time.time() + with service.holder.aclient(ClientConnectionPoolType.TRAIN) as client: + future = await client.models.unload(request=types.UnloadModelRequest(model_id=model_id)) + return await _APIFuture( + types.UnloadModelResponse, service.holder, future, request_start_time=started, + request_type='UnloadModel').result_async() + + return service.holder.run_coroutine_threadsafe(unload()).result(timeout=180) From dd96d85464b63e37d195820393785bbf98cf93b6 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Thu, 15 Oct 2026 09:43:46 +0800 Subject: [PATCH 02/46] fix: keep on-demand training independent of checkpoint repairs --- .../client/server/on_demand_training.yaml | 6 +- docs/design/on-demand-training.md | 10 +- .../validation/on-demand-training/README.md | 45 ++- src/twinkle/model/megatron/megatron.py | 4 +- .../model/megatron/strategy/megatron.py | 6 +- src/twinkle/server/model/app.py | 2 +- src/twinkle/server/task_queue/worker.py | 5 +- .../server/integration/test_actor_recovery.py | 15 +- .../integration/test_mock_mode_startup.py | 7 +- .../on_demand_training/device_validation.py | 296 +++++------------- 10 files changed, 128 insertions(+), 268 deletions(-) diff --git a/cookbook/client/server/on_demand_training.yaml b/cookbook/client/server/on_demand_training.yaml index 2646d3aa2..12018f1f7 100644 --- a/cookbook/client/server/on_demand_training.yaml +++ b/cookbook/client/server/on_demand_training.yaml @@ -15,7 +15,7 @@ applications: modes: lora: num_devices: 2 - device_mesh: {device_type: cuda, fsdp_size: 2, dp_size: 1} + device_mesh: {device_type: cuda, dp_size: 2} min_replicas: 1 max_loras: 5 max_r: 32 @@ -23,7 +23,7 @@ applications: target_modules: all-linear full: num_devices: 4 - device_mesh: {device_type: cuda, fsdp_size: 4, dp_size: 1} + device_mesh: {device_type: cuda, dp_size: 4} min_replicas: 0 idle_timeout_s: 1800 - name: model-small @@ -36,5 +36,5 @@ applications: modes: lora: num_devices: 2 - device_mesh: {device_type: cuda, tp_size: 2, dp_size: 1} + device_mesh: {device_type: cuda, dp_size: 2} max_loras: 5 diff --git a/docs/design/on-demand-training.md b/docs/design/on-demand-training.md index 892767829..b5605c723 100644 --- a/docs/design/on-demand-training.md +++ b/docs/design/on-demand-training.md @@ -25,7 +25,7 @@ resource waiting queues, preemption, training batching, or automatic checkpointi | R10 | New specifications use business inactivity, defaulting to 1800 seconds. | | R11 | Heartbeats, polling, health, and information queries do not renew business inactivity. | | R12 | Remove an empty dynamic instance immediately; retain clean resident capacity. | -| R13 | Expired handles stay invalid; users explicitly recreate and restore saved state. | +| R13 | Expired handles stay invalid; users explicitly recreate. No checkpoint is resumed implicitly. | | R14 | Worker failure or controller restart invalidates affected handles without replaying training. | | R15 | Retain the existing instance queue and backend execution gate. | | R16 | Keep `applications`, old model arguments, and SDK methods; use one controller. | @@ -153,17 +153,17 @@ Serve teardown behavior. | A11 | Partial startup, cancellation, and allocation failure leave correct quota and ownership. | | A12 | Worker/controller failure fences old handles and preserves unrelated instances. | | A13 | Same-label recreation, old requests, and sequence retries cannot mutate another handle. | -| A14 | Explicit recreation restores saved weights and optimizer state. | +| A14 | Explicit recreation starts with clean weights and optimizer state; checkpoint restoration is separate. | | A15 | Resident full reset leaves a clean base for the next user. | | A16 | Ten cross-model reload cycles leave no accumulated workers, PGs, or device memory. | | A17 | Invalid mode/config, expired IDs, and cross-user access fail without allocation. | -| A18 | Related SDK, YAML, checkpoint, and repository regressions pass together. | +| A18 | Related SDK, YAML, and repository CPU regressions pass together. | | A19 | Late startup/user initialization cannot publish success or prematurely reuse a slot. | | A20 | Memory and Redis recovery invalidate old state and rebuild quota correctly. | | A21 | Checkpoint upload needs no base; patched residents rebuild before reuse. | Review regressions also cover nonblocking resident recovery, launcher initialization rollback, lost completion notification, and sequence-receipt expiry. CPU/logical Ray -resources are not device training evidence. The revised code still requires both -backend device reruns, including an SP=true check for the Megatron LoRA forward path. +resources are not device training evidence. The standalone resource lifecycle is accepted with DP on both backends. +Checkpoint restoration and FSDP/TP/SP backend repairs are tracked separately in #315. Raw device evidence is retained outside the public source tree. diff --git a/docs/design/validation/on-demand-training/README.md b/docs/design/validation/on-demand-training/README.md index 99a8bead2..ff7d2b4a8 100644 --- a/docs/design/validation/on-demand-training/README.md +++ b/docs/design/validation/on-demand-training/README.md @@ -1,35 +1,30 @@ # On-demand training validation -The review fixes passed local validation with CPU models, real Ray scheduling, -the Native/Tinker SDKs over HTTP, and isolated memory/Redis state: +PR #312 is based directly on main. It does not require the checkpoint, +FSDP optimizer, or Megatron TP/SP repairs in #315. Its backend changes are +limited to resetting resident full-training state and using the supported +Transformers configuration import needed to initialize Megatron. -- Server, model, Transformers, and client CPU regression: 807 passed, no failures, - and 21 hardware-specific tests skipped. -- HTTP mock acceptance: passed for both SDKs in shared LoRA and exclusive full - modes, including release/recreation, runtime lifecycle rejection, controller - restart, state loss, and ten reuse cycles. -- Existing YAML startup: passed with memory state and with Redis state. -- The separate backend prerequisite (#315): 17 focused CPU regressions passed - without failures or skips in an isolated checkout. - -Repository pre-commit checks passed. Raw local outputs are retained outside the -repository. - -Real-device validation of the revised code is deferred at the user's request -because the PPU server is running other tests. PR #312 remains a draft. - -The previous runtime (`1bcf657e1534f6023b5de605b5a204c94705729b`) passed nine -acceptance cases on each of Transformers and Megatron, with no failures or skips. -Those results do not validate the revised lifecycle code. Raw logs, JUnit files, -server configurations, and numerical results are retained outside the repository. +Collected CPU tests under `tests/server/training/` exercise the controller, +allocation, recovery, and both installed SDKs over HTTP using a mock backend. +The HTTP and allocation probes run in isolated subprocesses. Existing +applications YAML startup can also be checked with memory and Redis state. Standalone validation programs live in `tools/validation/on_demand_training/`. -The collected CPU tests under `tests/server/training/` run their allocation and -HTTP programs in isolated processes. Device validation is explicitly invoked: +The DP resource-lifecycle device suite is explicitly invoked: ```bash python -m pytest tools/validation/on_demand_training/device_validation.py -vv ``` -It requires the backend, model paths, a separate evidence directory, and free -devices configured by `TWINKLE_DEVICE_*`. It must not run against occupied devices. +It runs DP=2, TP=1, and no FSDP for both backends and both SDKs in shared +LoRA/exclusive full modes. Assertions cover actual training and tenant isolation, +release, cross-model reuse, business-idle reclamation, expired handles, and clean +recreation. Checkpoint I/O and FSDP/TP repairs belong to #315, not this acceptance. +Results from older revisions containing #315 do not validate standalone #312. + +Device validation requires explicit `TWINKLE_DEVICE_*` configuration, a separate +evidence directory, and exclusively available devices. Record the exact source +SHA, import origins, raw JUnit/logs, numerical assertions, and resource cleanup. +Current results and remaining limitations are recorded in the English PR body; +results from an older revision are not acceptance evidence for the new head. diff --git a/src/twinkle/model/megatron/megatron.py b/src/twinkle/model/megatron/megatron.py index 2587e27b6..152e7350e 100644 --- a/src/twinkle/model/megatron/megatron.py +++ b/src/twinkle/model/megatron/megatron.py @@ -17,7 +17,7 @@ from peft.tuners.lora import Linear as LoraLinear from torch.optim import Optimizer from torch.optim.lr_scheduler import LRScheduler -from transformers import PreTrainedConfig +from transformers import PretrainedConfig from typing import Any, Callable, Dict, Generator, List, Literal, Optional, Tuple, Type, Union import twinkle @@ -101,7 +101,7 @@ class MegatronModel(TwinkleModel, nn.Module, CheckpointEngineMixin): def __init__( self, model_id: str, - config: Optional[PreTrainedConfig] = None, + config: Optional[PretrainedConfig] = None, ddp_config: Optional[Dict[str, Any]] = None, device_mesh: Optional[DeviceMesh] = None, mixed_precision: Literal['no', 'fp16', 'bf16'] = 'bf16', diff --git a/src/twinkle/model/megatron/strategy/megatron.py b/src/twinkle/model/megatron/strategy/megatron.py index 59f545a2c..7b8c822bd 100644 --- a/src/twinkle/model/megatron/strategy/megatron.py +++ b/src/twinkle/model/megatron/strategy/megatron.py @@ -1,7 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import torch import torch.nn as nn -from transformers import PreTrainedConfig +from transformers import PretrainedConfig from typing import Any, Dict, List, Literal, Optional from twinkle import DeviceMesh, Platform, torch_util @@ -44,7 +44,7 @@ def __init__( mixed_precision: Literal['no', 'fp16', 'bf16'] = 'bf16', seed: int = 42, variable_seq_lengths: bool = True, - config: PreTrainedConfig = None, + config: PretrainedConfig = None, ddp_config: Dict[str, Any] = None, **kwargs, ): @@ -263,7 +263,7 @@ def reduce_loss(self, local_loss, local_count, logits, logps): def get_model_config( self, - hf_config: PreTrainedConfig, + hf_config: PretrainedConfig, parallel_kwargs: Dict[str, Any], **kwargs, ): diff --git a/src/twinkle/server/model/app.py b/src/twinkle/server/model/app.py index d0049fee4..34f9c3a5c 100644 --- a/src/twinkle/server/model/app.py +++ b/src/twinkle/server/model/app.py @@ -195,7 +195,7 @@ async def check_model_health(self) -> dict: return await asyncio.shield(self._health_probe_task) def mark_unhealthy(self) -> None: - """Flag the deployment unhealthy; /healthz returns 503 until a probe recovers it.""" + """Flag this training runtime unhealthy until the controller replaces it.""" self._model_unhealthy = True async def _probe_after_timeout(self) -> None: diff --git a/src/twinkle/server/task_queue/worker.py b/src/twinkle/server/task_queue/worker.py index fd9704feb..91edc50c8 100644 --- a/src/twinkle/server/task_queue/worker.py +++ b/src/twinkle/server/task_queue/worker.py @@ -554,10 +554,7 @@ async def _try_run_one(self) -> bool: continue # Execute the task (serial: stops after the first execution) - try: - await self._execute_task(task, queue_key, q) - finally: - self._request_ids.discard(task.request_id) + await self._execute_task(task, queue_key, q) return True return False diff --git a/tests/server/integration/test_actor_recovery.py b/tests/server/integration/test_actor_recovery.py index 16527bbae..dfc03a74b 100644 --- a/tests/server/integration/test_actor_recovery.py +++ b/tests/server/integration/test_actor_recovery.py @@ -40,7 +40,7 @@ async def call_backend(self, fn, /, *args, admit: bool = True, **kwargs): @pytest.mark.asyncio -async def test_timeout_probe_marks_unhealthy_then_recovers(): +async def test_timeout_probe_keeps_runtime_unhealthy_even_if_actor_recovers(): model = _MockModel() h = _HealthHarness(model) @@ -54,11 +54,12 @@ async def test_timeout_probe_marks_unhealthy_then_recovers(): await h._probe_after_timeout() assert h._model_unhealthy is True # /healthz would return 503 - # Actor recovers; one successful probe clears the bit (no restart needed). + # Actor liveness cannot prove that timed-out training stopped mutating state. + # The controller must replace this runtime before accepting more work. model.alive = True result = await h.check_model_health() - assert result['healthy'] is True - assert h._model_unhealthy is False + assert result['healthy'] is False + assert h._model_unhealthy is True @pytest.mark.asyncio @@ -76,11 +77,11 @@ async def test_health_route_returns_503_when_probe_fails(): @pytest.mark.asyncio -async def test_mark_unhealthy_is_cleared_by_successful_probe(): +async def test_mark_unhealthy_is_not_cleared_by_successful_probe(): h = _HealthHarness(_MockModel()) h.mark_unhealthy() assert h._model_unhealthy is True result = await h.check_model_health() - assert result['healthy'] is True - assert h._model_unhealthy is False + assert result['healthy'] is False + assert h._model_unhealthy is True diff --git a/tests/server/integration/test_mock_mode_startup.py b/tests/server/integration/test_mock_mode_startup.py index e90af75d5..b1c38259a 100644 --- a/tests/server/integration/test_mock_mode_startup.py +++ b/tests/server/integration/test_mock_mode_startup.py @@ -196,6 +196,10 @@ def _exercise_twinkle_clients(base: str) -> None: save_resp = model.save(name='step-1') assert save_resp.twinkle_path and save_resp.twinkle_path.startswith('twinkle://') + # Mock save creates an empty directory; identify it as the LoRA checkpoint + # expected by the sampler's checkpoint classification introduced on main. + assert save_resp.checkpoint_dir + lora_cfg.save_pretrained(save_resp.checkpoint_dir) model.load(name=save_resp.twinkle_path) model.resume_from_checkpoint(name=save_resp.twinkle_path) model.apply_patch('NoopPatch') @@ -221,8 +225,7 @@ def _exercise_twinkle_clients(base: str) -> None: ) assert samples and samples[0].sequences and len(samples[0].sequences[0].tokens) == 4 - # adapter_uri triggers reset_prefix_cache; reuse the training-weights - # path because mock ``save()`` materialized that directory on disk. + # Exercise adapter_uri resolution using the mock LoRA checkpoint directory. sampler.sample( inputs=[{ 'input_ids': [1, 2, 3] diff --git a/tools/validation/on_demand_training/device_validation.py b/tools/validation/on_demand_training/device_validation.py index b496c816f..2fe2c17c6 100644 --- a/tools/validation/on_demand_training/device_validation.py +++ b/tools/validation/on_demand_training/device_validation.py @@ -3,6 +3,7 @@ Required: TWINKLE_DEVICE_BACKEND, TWINKLE_DEVICE_MODEL, TWINKLE_DEVICE_OTHER_MODEL, TWINKLE_DEVICE_RUN_DIR and four exclusively available visible devices. +Tests cover DP resource lifecycle without checkpoint I/O. This is intentionally outside pytest's default test filename pattern. """ import json @@ -40,14 +41,14 @@ def server(): os.environ['TWINKLE_DEFAULT_SAVE_DIR'] = str(run_dir / 'checkpoints') os.environ['TWINKLE_PERSISTENCE_MODE'] = 'memory' os.environ['TWINKLE_LONG_POLL_TIMEOUT'] = '2' - mesh = {'fsdp_size': 2} if backend == 'transformers' else {'tp_size': 2} + mesh = {'dp_size': 2} apps = [] for alias, variable in [('small', 'TWINKLE_DEVICE_MODEL'), ('other', 'TWINKLE_DEVICE_OTHER_MODEL')]: source = os.environ[variable] assert Path(source, 'config.json').is_file() mode = {'num_devices': 2, 'device_mesh': mesh, 'idle_timeout_s': 600 if alias == 'small' else 30} if backend == 'transformers': - mode['strategy'] = 'native_fsdp' + mode['strategy'] = 'accelerate' apps.append({ 'name': alias, 'route_prefix': '/api/v1/model/' + alias, @@ -176,37 +177,14 @@ def evaluate(self): assert values.size and np.isfinite(values).all(), values return values - def step(self, *, contrast=False): - labels = self.inputs[0]['labels'] - if contrast: - self.inputs[0]['labels'] = list(reversed(labels)) - try: - if self.protocol == 'native': - self.client.forward_backward(self.inputs) - self.client.clip_grad_and_step() - else: - from tinker import types - self.client.forward_backward(self.datum(), 'cross_entropy').result() - self.client.optim_step(types.AdamParams(learning_rate=self.rate)).result() - finally: - self.inputs[0]['labels'] = labels - - def save(self): - name = 'device-' + uuid.uuid4().hex[:8] - return (self.client.save(name, save_optimizer=True).twinkle_path - if self.protocol == 'native' else self.client.save_state(name).result().path) - - def restore(self, path): - if self.protocol == 'native': - self.client.resume_from_checkpoint(path) - else: - self.client.load_state_with_optimizer(path).result() - - def restore_weights(self, path): + def step(self): if self.protocol == 'native': - self.client.resume_from_checkpoint(path, resume_only_model=True) + self.client.forward_backward(self.inputs) + self.client.clip_grad_and_step() else: - self.client.load_state(path).result() + from tinker import types + self.client.forward_backward(self.datum(), 'cross_entropy').result() + self.client.optim_step(types.AdamParams(learning_rate=self.rate)).result() def close(self): if self.protocol == 'native': @@ -222,219 +200,105 @@ def active_instances(controller): @pytest.mark.parametrize('protocol', ['native', 'tinker']) @pytest.mark.parametrize('mode', ['lora', 'full']) -def test_training_isolation_checkpoint_and_device_reclamation(server, protocol, mode): +def test_dp_resource_lifecycle(server, protocol, mode): base, controller, evidence, _ = server + assert not active_instances(controller) + assert ray.available_resources().get('GPU', 0) == 4 first, second = Handle(base, protocol, mode), Handle(base, protocol, mode) - assert first.identity != second.identity instances = active_instances(controller) assert len(instances) == (1 if mode == 'lora' else 2) assert all(len(row['devices']) == 2 for row in instances) - devices = [device for row in instances for device in row['devices']] - assert len(set(devices)) == len(devices) - original, untouched = first.evaluate(), second.evaluate() - untouched_checkpoint = second.save() if mode == 'lora' else None + first_instance_ids = {row['instance_id'] for row in instances} + baseline, untouched = first.evaluate(), second.evaluate() for _ in range(4): first.step() trained = first.evaluate() - assert float(trained.mean()) > float(original.mean()) + 1e-4, (original, trained) + assert float(trained.mean()) > float(baseline.mean()) + 1e-4 np.testing.assert_allclose(second.evaluate(), untouched, rtol=0, atol=2e-3) + second.step() + peer_trained = second.evaluate() + assert float(peer_trained.mean()) > float(untouched.mean()) + 1e-4 + np.testing.assert_allclose(first.evaluate(), trained, rtol=0, atol=2e-3) if mode == 'lora': third, fourth = Handle(base, protocol, mode), Handle(base, protocol, mode) - assert len(active_instances(controller)) == 2 + third.evaluate() + instances = active_instances(controller) + assert len(instances) == 2 assert ray.available_resources().get('GPU', 0) == 0 with pytest.raises(Exception, match='(?i)allocation'): Handle(base, protocol, mode) + allocated = [device for row in instances for device in row['devices']] + assert len(allocated) == 4 and len(set(allocated)) == 4 + if mode == 'lora': fourth.close() - third.restore(untouched_checkpoint) - second.step() - third.step() - np.testing.assert_allclose(second.evaluate(), third.evaluate(), rtol=0, atol=2e-2) - np.testing.assert_allclose(first.evaluate(), trained, rtol=0, atol=2e-3) third.close() - assert len(active_instances(controller)) == 1 - checkpoint = first.save() - first.step(contrast=True) - continuation = first.evaluate() - first.restore(checkpoint) - np.testing.assert_allclose(first.evaluate(), trained, rtol=0, atol=2e-3) - first.step(contrast=True) - np.testing.assert_allclose(first.evaluate(), continuation, rtol=0, atol=2e-2) - old_id = first.identity first.close() assert len(active_instances(controller)) == 1 - replacement = Handle(base, protocol, mode) - assert replacement.identity != old_id - replacement.restore_weights(checkpoint) - np.testing.assert_allclose(replacement.evaluate(), trained, rtol=0, atol=2e-3) - replacement.step(contrast=True) - optimizer_reset = replacement.evaluate() - assert float(np.max(np.abs(optimizer_reset - continuation))) > 2e-2, 'Optimizer continuation oracle is not decisive' - replacement.restore(checkpoint) - np.testing.assert_allclose(replacement.evaluate(), trained, rtol=0, atol=2e-3) - replacement.step(contrast=True) - np.testing.assert_allclose(replacement.evaluate(), continuation, rtol=0, atol=2e-2) - replacement.close() + assert ray.available_resources().get('GPU', 0) == 2 + np.testing.assert_allclose(second.evaluate(), peer_trained, rtol=0, atol=2e-3) + second.step() + peer_continued = second.evaluate() + assert float(peer_continued.mean()) > float(peer_trained.mean()) + 1e-4 second.close() assert not active_instances(controller) assert ray.available_resources().get('GPU', 0) == 4 - evidence['cases'].append({ - 'protocol': protocol, - 'mode': mode, - 'initial_logps': original.tolist(), - 'trained_logps': trained.tolist(), - 'continuation_logps': continuation.tolist(), - 'optimizer_reset_logps': optimizer_reset.tolist(), - 'checkpoint': checkpoint, - 'devices': devices, - 'passed': True - }) - -def test_cross_model_reuse_and_heartbeat_does_not_prevent_expiry(server): - base, controller, evidence, _ = server - before = ray.available_resources().get('GPU', 0) - for index in range(10): - handle = Handle(base, 'native', 'full', alias='small' if index % 2 == 0 else 'other') - values = handle.evaluate() - handle.close() - assert not active_instances(controller) - assert ray.available_resources().get('GPU', 0) == before - print(json.dumps({'cycle': index + 1, 'logps': values.tolist()}), flush=True) - expired = Handle(base, 'native', 'lora', alias='other') - expired.evaluate() - deadline = time.monotonic() + 60 - while active_instances(controller) and time.monotonic() < deadline: - response = expired.transport.post( - base + '/api/v1/twinkle/session_heartbeat', json_data={'session_id': expired.transport.context.session_id}) - response.raise_for_status() - time.sleep(2) - assert not active_instances(controller), 'Heartbeats incorrectly kept idle training alive' - with pytest.raises(Exception): - expired.client.zero_grad() - assert ray.available_resources().get('GPU', 0) == before - expired.transport.close() - evidence['cases'].append({'cycles': 10, 'idle_reclaimed': True, 'passed': True}) - - -@pytest.mark.parametrize('protocol', ['native', 'tinker']) -def test_idle_lora_user_expires_without_resetting_active_user(server, protocol): - base, controller, evidence, _ = server - idle = Handle(base, protocol, 'lora', alias='other') - active = Handle(base, protocol, 'lora', alias='other') - instance = active_instances(controller)[0] - assert instance['training_ids'] == [idle.identity, active.identity] - original = idle.evaluate() + alternate = Handle(base, protocol, mode, alias='other') + alternate_instances = active_instances(controller) + assert len(alternate_instances) == 1 + reused = alternate_instances[0]['devices'] + assert len(reused) == 2 and set(reused).issubset(allocated) + other_initial = alternate.evaluate() for _ in range(4): - active.step() - trained = active.evaluate() - assert trained.mean() > original.mean() + 1e-4 - np.testing.assert_allclose(idle.evaluate(), original, rtol=0, atol=2e-3) - checkpoint = active.save() - active.step(contrast=True) - expected = active.evaluate() - active.restore(checkpoint) - deadline = time.monotonic() + 60 - while True: - np.testing.assert_allclose(active.evaluate(), trained, rtol=0, atol=2e-3) - response = idle.transport.post( - base + '/api/v1/twinkle/session_heartbeat', json_data={'session_id': idle.transport.context.session_id}) + alternate.step() + other_trained = alternate.evaluate() + assert float(other_trained.mean()) > float(other_initial.mean()) + 1e-4 + health = requests.get(base + '/api/v1/twinkle/healthz/deep', timeout=30) + assert health.status_code == 200 and health.json()['healthy'] + deadline = time.monotonic() + 90 + while active_instances(controller): + assert time.monotonic() < deadline, 'Business-idle instance did not release DP devices' + response = alternate.transport.post( + base + '/api/v1/twinkle/session_heartbeat', + json_data={'session_id': alternate.transport.context.session_id}) response.raise_for_status() - instances = active_instances(controller) - assert len(instances) == 1 and instances[0]['instance_id'] == instance['instance_id'] - if idle.identity not in instances[0]['training_ids']: - break - assert time.monotonic() < deadline, 'Idle LoRA user did not expire while the other user stayed active' time.sleep(2) - with pytest.raises(Exception): - idle.evaluate() - active.step(contrast=True) - np.testing.assert_allclose(active.evaluate(), expected, rtol=0, atol=2e-2) - replacement = Handle(base, protocol, 'lora', alias='other') - assert replacement.identity != idle.identity - assert active_instances(controller)[0]['instance_id'] == instance['instance_id'] - np.testing.assert_allclose(replacement.evaluate(), original, rtol=0, atol=2e-3) - replacement.close() - np.testing.assert_allclose(active.evaluate(), expected, rtol=0, atol=2e-3) - active.close() - idle.transport.close() + assert ray.available_resources().get('GPU', 0) == 4 + with pytest.raises(Exception) as expired: + alternate.evaluate() + alternate.transport.close() + + fresh = Handle(base, protocol, mode) + fresh_instances = active_instances(controller) + assert len(fresh_instances) == 1 + assert fresh_instances[0]['instance_id'] not in first_instance_ids + np.testing.assert_allclose(fresh.evaluate(), baseline, rtol=0, atol=2e-3) + fresh.close() assert not active_instances(controller) assert ray.available_resources().get('GPU', 0) == 4 evidence['cases'].append({ 'protocol': protocol, - 'idle_lora_expired': True, - 'active_optimizer_preserved': True, - 'slot_reused': True, + 'mode': mode, + 'parallelism': { + 'dp_size': 2, + 'tp_size': 1, + 'fsdp_size': 1 + }, + 'initial_logps': baseline.tolist(), + 'trained_logps': trained.tolist(), + 'peer_initial_logps': untouched.tolist(), + 'peer_trained_logps': peer_trained.tolist(), + 'peer_continued_logps': peer_continued.tolist(), + 'other_initial_logps': other_initial.tolist(), + 'other_trained_logps': other_trained.tolist(), + 'original_devices': allocated, + 'cross_model_reused_devices': reused, + 'health': health.json(), + 'idle_expired_error': str(expired.value)[:400], + 'lazy_start': True, + 'shared_or_exclusive_isolation': True, + 'idle_reclaimed': True, + 'fresh_recreated': True, 'passed': True }) - - -def test_resident_full_reset_preserves_other_instances(server): - base, controller, evidence, launcher = server - raw = launcher.config.to_yaml_dict() - raw['applications'][0]['args']['modes']['full']['min_replicas'] = 1 - ray.get(controller.shutdown.remote(), timeout=180) - ray.kill(controller, no_restart=True) - from twinkle.server.training.actor import controller_name - deadline = time.monotonic() + 30 - while True: - try: - ray.get_actor(controller_name(launcher.config.ray_namespace), namespace=launcher.config.ray_namespace) - except ValueError: - break - assert time.monotonic() < deadline, 'Previous controller name was not released' - time.sleep(0.1) - replacement = ServerLauncher(ServerConfig.model_validate(raw)) - replacement._deploy_training() - controller = replacement._training_controller - launcher._training_controller = controller - evidence['resident_controller_epoch'] = ray.get(controller.diagnostics.remote())['epoch'] - resident, extra = Handle(base, 'native', 'full'), Handle(base, 'native', 'full') - instances = active_instances(controller) - resident_instance = next(row for row in instances if row['resident']) - original, untouched = resident.evaluate(), extra.evaluate() - for _ in range(2): - resident.step() - assert resident.evaluate().mean() > original.mean() + 1e-4 - resident.close() - restored = Handle(base, 'native', 'full') - instances = active_instances(controller) - assert next(row for row in instances if row['resident'])['instance_id'] == resident_instance['instance_id'] - np.testing.assert_allclose(restored.evaluate(), original, rtol=0, atol=2e-3) - np.testing.assert_allclose(extra.evaluate(), untouched, rtol=0, atol=2e-3) - restored.close() - extra.close() - instances = active_instances(controller) - assert len(instances) == 1 and instances[0]['resident'] - assert ray.available_resources().get('GPU', 0) == 2 - evidence['cases'].append({'resident_reset': True, 'extra_preserved': True, 'passed': True}) - - -def test_worker_failure_reclaims_devices_and_preserves_resident(server): - from ray.util import list_named_actors - base, _, evidence, launcher = server - controller = launcher._training_controller - resident = Handle(base, 'native', 'full') - original = resident.evaluate() - failing = Handle(base, 'native', 'lora') - failing.evaluate() - instance = next(row for row in active_instances(controller) if row['mode'] == 'lora') - actors = [ - row for row in list_named_actors(all_namespaces=True) if row['namespace'] == evidence['namespace'] - and instance['instance_id'] in row['name'] and not row['name'].endswith('-registry') - ] - assert len(actors) == 2, actors - ray.kill(ray.get_actor(actors[0]['name'], namespace=evidence['namespace']), no_restart=True) - deadline = time.monotonic() + 90 - while any(row['instance_id'] == instance['instance_id'] for row in active_instances(controller)): - assert time.monotonic() < deadline, 'Failed runtime devices were not reclaimed' - time.sleep(1) - with pytest.raises(Exception): - failing.client.zero_grad() - np.testing.assert_allclose(resident.evaluate(), original, rtol=0, atol=2e-3) - recovered = Handle(base, 'native', 'lora') - recovered.evaluate() - recovered.close() - resident.close() - failing.transport.close() - assert ray.available_resources().get('GPU', 0) == 2 - evidence['cases'].append({'worker_failure': True, 'resident_preserved': True, 'recreated': True, 'passed': True}) + (Path(os.environ['TWINKLE_DEVICE_RUN_DIR']) / 'progress.json').write_text(json.dumps(evidence, indent=2)) From fd959e30a7bd36bd4a8694452bd13f0f6c012c80 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Mon, 19 Oct 2026 04:09:52 +0000 Subject: [PATCH 03/46] fix(server): harden on-demand training lifecycle and resource boundaries - Split runtime allocation from deployment so the controller persists ownership before any worker starts; the driver no longer writes records. - Make cleanup, close and failure transitions persist-before-mutate, track failed instances for sweep retries, fence admission, and exclude failed residents when rebuilding minimum replicas. - Never lose a worker wake-up when enqueue callbacks or queued-status writes fail; log background job failures with tracebacks. - Return 503/Server while the controller is not ready; drop unreleased reason codes; reset resident training state before model-id branching. - Move generic device allocation to server/resources, managed runtime to server/model, training records/store to server/state, TaskQueueConfig to server/config, and shared proxy/env propagation to the server root, with no compatibility re-exports. - Encapsulate ServerState backend access and batch heartbeat reads outside the controller lock. - Guard all server component dependencies with a strict grimp layer test. - Require ray[serve]>=2.55 for async deployment constructors. - Stop tracking local-only docs/design notes. --- .dev_scripts/ci_container_test.sh | 2 +- .gitignore | 3 + docs/design/on-demand-training.md | 169 ---------- .../validation/on-demand-training/README.md | 30 -- pyproject.toml | 4 +- .../model/transformers/transformers.py | 18 +- src/twinkle/server/config/__init__.py | 2 + src/twinkle/server/config/application_spec.py | 2 +- src/twinkle/server/config/server_config.py | 2 +- .../config.py => config/task_queue.py} | 20 +- src/twinkle/server/config/training.py | 2 +- src/twinkle/server/env_propagation.py | 37 +++ src/twinkle/server/gateway/app.py | 2 +- src/twinkle/server/gateway/proxy.py | 230 ------------- src/twinkle/server/gateway/tinker_handlers.py | 2 - .../server/launcher/env_propagation.py | 53 --- .../server/launcher/server_launcher.py | 2 +- src/twinkle/server/lifecycle/envelope.py | 2 - src/twinkle/server/lifecycle/protocols.py | 2 +- src/twinkle/server/model/app.py | 16 +- .../runtime.py => model/managed_runtime.py} | 33 +- src/twinkle/server/resources/__init__.py | 2 + .../{training => resources}/allocation.py | 21 +- src/twinkle/server/sampler/app.py | 2 +- src/twinkle/server/service_proxy.py | 105 ++++++ src/twinkle/server/state/model_manager.py | 16 + src/twinkle/server/state/models.py | 2 - src/twinkle/server/state/server_state.py | 25 +- src/twinkle/server/state/session_manager.py | 7 + .../records.py => state/training_records.py} | 4 +- .../store.py => state/training_store.py} | 38 +-- src/twinkle/server/task_queue/__init__.py | 3 - src/twinkle/server/task_queue/mixin.py | 31 +- src/twinkle/server/task_queue/worker.py | 2 +- src/twinkle/server/training/actor.py | 21 +- src/twinkle/server/training/controller.py | 309 ++++++++++++------ src/twinkle/server/training/driver.py | 54 +-- src/twinkle/server/training/errors.py | 4 - src/twinkle/server/training/router.py | 7 +- tests/server/gateway/test_proxy.py | 2 +- .../lifecycle/test_preflight_rejection.py | 2 +- .../server/lifecycle/test_submit_peek_e2e.py | 2 +- tests/server/lifecycle/test_timing_bounds.py | 2 +- tests/server/model/test_tinker_handlers.py | 2 +- .../sampler/test_sampling_concurrency.py | 2 +- tests/server/test_dependencies.py | 59 ++++ .../training/test_actor_initialization.py | 53 +++ tests/server/training/test_controller.py | 24 +- .../server/training/test_failure_recovery.py | 292 +++++++++++++++++ tests/server/training/test_router.py | 18 +- .../training/test_runtime_initialization.py | 49 +++ tests/server/training/test_runtime_queue.py | 5 +- tests/server/training/test_sequence.py | 8 +- .../server/training/test_storage_recovery.py | 3 +- tests/server/utils/task_queue/test_config.py | 2 +- tests/server/utils/test_task_queue_mixin.py | 2 +- .../test_resident_training_reset.py | 13 + .../on_demand_training/allocation_probe.py | 9 +- .../on_demand_training/device_validation.py | 35 +- .../on_demand_training/fragmentation_probe.py | 7 +- .../on_demand_training/matrix_validation.py | 117 +++++++ 61 files changed, 1209 insertions(+), 785 deletions(-) delete mode 100644 docs/design/on-demand-training.md delete mode 100644 docs/design/validation/on-demand-training/README.md rename src/twinkle/server/{task_queue/config.py => config/task_queue.py} (78%) create mode 100644 src/twinkle/server/env_propagation.py delete mode 100644 src/twinkle/server/gateway/proxy.py delete mode 100644 src/twinkle/server/launcher/env_propagation.py rename src/twinkle/server/{training/runtime.py => model/managed_runtime.py} (88%) create mode 100644 src/twinkle/server/resources/__init__.py rename src/twinkle/server/{training => resources}/allocation.py (84%) create mode 100644 src/twinkle/server/service_proxy.py rename src/twinkle/server/{training/records.py => state/training_records.py} (93%) rename src/twinkle/server/{training/store.py => state/training_store.py} (69%) create mode 100644 tests/server/test_dependencies.py create mode 100644 tests/server/training/test_actor_initialization.py create mode 100644 tests/server/training/test_failure_recovery.py create mode 100644 tests/server/training/test_runtime_initialization.py create mode 100644 tools/validation/on_demand_training/matrix_validation.py diff --git a/.dev_scripts/ci_container_test.sh b/.dev_scripts/ci_container_test.sh index 722e977a4..04cd9d4ac 100644 --- a/.dev_scripts/ci_container_test.sh +++ b/.dev_scripts/ci_container_test.sh @@ -26,7 +26,7 @@ if [ "$MODELSCOPE_SDK_DEBUG" == "True" ]; then pip uninstall autoawq -y pip uninstall lmdeploy -y pip uninstall tensorflow -y - pip install ray==2.48 + pip install ray==2.55.1 pip install optimum # test with install diff --git a/.gitignore b/.gitignore index 6ba7b5381..dac8c4337 100644 --- a/.gitignore +++ b/.gitignore @@ -77,6 +77,9 @@ instance/ # Sphinx documentation docs/_build/ +# Internal design notes are local-only and must not be committed. +/docs/design/ + # PyBuilder target/ diff --git a/docs/design/on-demand-training.md b/docs/design/on-demand-training.md deleted file mode 100644 index b5605c723..000000000 --- a/docs/design/on-demand-training.md +++ /dev/null @@ -1,169 +0,0 @@ -# On-demand training and device reuse - -The catalog, controller, allocation, and SDK routing are implemented. Review fixes -are under local validation. Real-device reruns are deferred at the user's request; -see the [validation record](validation/on-demand-training/README.md). - -## Scope and requirements - -A fixed device pool serves a catalog larger than its resident capacity. The demand -unit is a training handle; the allocation unit is a model instance on N devices. -Sampling retains its existing deployment lifecycle. This change does not introduce -resource waiting queues, preemption, training batching, or automatic checkpointing. - -| Requirement | Behavior | -| --- | --- | -| R01 | Support the existing Transformers and Megatron training backends. | -| R02 | Administrators declare model sources, modes, parallelism, and N devices. | -| R03 | New mode specifications default to lazy startup; `min_replicas` keeps resident capacity. | -| R04 | Initialize residents before accepting dynamic training; all instances use Ray resources. | -| R05 | An instance reserves N devices on one node; independent bases do not share that reservation. | -| R06 | A model may offer both LoRA and full training through separate instances. | -| R07 | Fill compatible LoRA instances first; expand when slots and devices permit. | -| R08 | Each full handle exclusively owns an instance, subject to user quota. | -| R09 | Capacity failures are terminal; callers explicitly issue a new create to retry. | -| R10 | New specifications use business inactivity, defaulting to 1800 seconds. | -| R11 | Heartbeats, polling, health, and information queries do not renew business inactivity. | -| R12 | Remove an empty dynamic instance immediately; retain clean resident capacity. | -| R13 | Expired handles stay invalid; users explicitly recreate. No checkpoint is resumed implicitly. | -| R14 | Worker failure or controller restart invalidates affected handles without replaying training. | -| R15 | Retain the existing instance queue and backend execution gate. | -| R16 | Keep `applications`, old model arguments, and SDK methods; use one controller. | -| R17 | Every create gets a distinct training ID; identical adapter labels cannot reattach old state. | -| R18 | Validate both backends and modes on real devices before marking the PR ready. | - -## Configuration and compatibility - -The public entry point remains `applications` with `import_path: model`. -The [configuration example](../../cookbook/client/server/on_demand_training.yaml) -shows `args.modes.lora/full`. Each mode sets `num_devices`, `device_mesh`, -`min_replicas`, `max_replicas`, and `idle_timeout_s`. LoRA additionally sets -`max_loras`, `max_r`, and its existing backend options. - -- New modes default to `min_replicas=0`, `idle_timeout_s=1800`, five LoRA slots, - and maximum LoRA rank 32. Full capacity is one handle. -- Mesh world size must equal N; N and capacities are positive integers. Timeouts - must be finite and positive. Minimum replicas cannot exceed the maximum. -- `modes` cannot be mixed with legacy single-mode settings or deployment replicas. -- Legacy `train_mode`, device counts, mesh, strategy, LoRA parameters, and environment - variables normalize into the same runtime catalog. -- Existing deployment replica bounds keep their resident-minimum/instance-limit - meaning. Explicit `args.min_replicas` and `args.max_replicas` take precedence. -- **Legacy `adapter_timeout` remains a session heartbeat timeout.** It is not converted - to business inactivity. `adapter_max_lifetime` remains an absolute handle lifetime. - Their existing defaults are 1800 and 36000 seconds respectively. -- A legacy configuration enables business inactivity only by explicitly setting - `args.idle_timeout_s`. Its heartbeat and absolute lifetime bounds still apply. -- The Native SDK and server must upgrade together for the internal training-ID - handshake. Native and Tinker retain their existing initialization differences. - -Controller allocation/startup/shutdown/sweep bounds default to 3/1800/60/5 seconds -under `training`. Base startup and user initialization each receive a bounded -startup window; waiting for a base does not consume the user's initialization window. -The create future covers both phases without allowing late success. - -## Ownership and lifecycle - -`TrainingCatalog` performs pure configuration lookup. `TrainingController` is the -single lifecycle writer. `RayInstanceDriver` performs Ray/Serve effects, and -`ModelManagement` executes backend operations on the existing queue. - -The controller is a named actor owned by the launcher, not a detached service that -another launcher adopts. Initialization failure cleans and terminates that actor. -While its owner is alive, a controller restart starts a new epoch and invalidates -old handles. Runtime constructors register an incarnation once; Serve cannot -silently replace a process while preserving its old training state. - -The state lock covers decisions and commits. Model startup, user initialization, -cleanup, and resident rebuilding run outside it. A resident rebuilding for one -model/mode does not block another specification or the sweep loop. - -Every instance owns a named Ray Placement Group with one accelerator/CPU bundle -and `STRICT_PACK`. The allocator bounds `pg.ready()`, obtains actual device IDs -through a scheduled probe, and injects that allocation into workers without a -second reservation. Worker, application, registry, and PG names identify their -server namespace, controller epoch, and instance ID. - -A model has a stable CPU router. Creating/releasing handles enters the controller; -training requests route to their fixed instance. Internal model lifecycle HTTP -routes explicitly reject direct creation or unloading. Backend workers do not -restart or replay stateful calls automatically. - -| Record | States | -| --- | --- | -| Training | PENDING → ACTIVE → CLOSING → RELEASED / EXPIRED / FAILED / CANCELLED | -| Instance | ALLOCATING → STARTING → READY; RESETTING for resident full reset; DRAINING → STOPPED | -| Unconfirmed shutdown | QUARANTINED, with devices/quota still owned until death is confirmed | - -## Creation, isolation, and expiry - -Create identity is `(client type, token, session, create key)`: Native uses its UUID -and Tinker uses `model_seq_id`. Identical retries return the original future; -changed parameters are rejected. Pending users count toward quota and LoRA slots. -Compatible pending creates share a base startup. Full handles never share an instance. - -A failed user initialization closes only that user. Its slot remains reserved while -queued or late initialization may still execute. Cleanup follows initialization in -the serial lifecycle queue, and a late result cannot publish an ACTIVE handle. -Worker death, base startup failure, or unsafe backend execution can fail an entire -instance; a create RPC timeout alone is not evidence of that failure. - -Controller tickets protect accepted operations from expiry. The runtime exposes -queued/executing membership through the compute worker's public interface. Sweep -reconciles both uncommitted admission and committed operations whose completion -notification was lost. A terminal future alone does not establish that a background -call stopped. Reconciled activity uses the original completion time, not sweep time. - -Native and Tinker keep bounded sequence retry receipts per training stream. An -expired-sequence fence survives while the handle is live, so expiry cannot cause an -old optimizer operation to run again. Retained receipts tolerate out-of-order -requests; claims are released only when submission never created a future. Retiring -the handle removes its stream ledgers. Redis enumeration uses the existing SCAN API. - -Expiry closes inactive users only after their operations drain. Business inactivity -uses executed operations; heartbeat-only clients expire under the new policy. -Legacy heartbeat and maximum-lifetime checks retain their original meanings. - -The last user leaving closes a dynamic instance. Clean residents keep their base; -full residents reset weights, optimizer state, gradients, and template hooks before -another user is admitted. A resident modified by an instance-level patch is rebuilt -when empty. A cleanup failure cannot release a slot or kill another active LoRA user -merely because its cleanup RPC timed out. - -Shutdown kills owned workers, deletes the owned Serve app, removes the PG, confirms -worker/resource termination, and then retires quota. Failed confirmation preserves -ownership and retries in the background. Recovery never clears a shared registry or -shuts down another component's Ray cluster. Launcher shutdown retains its normal -Serve teardown behavior. - -## Acceptance - -| Case | Required observation | -| --- | --- | -| A01 | Startup/query/health on a lazy catalog loads no base or training PG. | -| A02 | Concurrent first creates coalesce into the required LoRA slot count. | -| A03 | Full LoRA slots expand with devices available; exhausted capacity fails terminally. | -| A04 | Both backends and SDKs isolate LoRA tenants and exclusive full handles. | -| A05 | Disjoint allocations are released and reused by another model. | -| A06 | Fragmented resources across nodes cannot satisfy one instance. | -| A07 | Heartbeats do not renew new-policy business inactivity; legacy semantics remain unchanged. | -| A08 | Idle LoRA cleanup preserves an active tenant's optimizer continuation. | -| A09 | Queued/running/save operations drain before expiry or release. | -| A10 | Unsafe background execution retains ownership until workers terminate. | -| A11 | Partial startup, cancellation, and allocation failure leave correct quota and ownership. | -| A12 | Worker/controller failure fences old handles and preserves unrelated instances. | -| A13 | Same-label recreation, old requests, and sequence retries cannot mutate another handle. | -| A14 | Explicit recreation starts with clean weights and optimizer state; checkpoint restoration is separate. | -| A15 | Resident full reset leaves a clean base for the next user. | -| A16 | Ten cross-model reload cycles leave no accumulated workers, PGs, or device memory. | -| A17 | Invalid mode/config, expired IDs, and cross-user access fail without allocation. | -| A18 | Related SDK, YAML, and repository CPU regressions pass together. | -| A19 | Late startup/user initialization cannot publish success or prematurely reuse a slot. | -| A20 | Memory and Redis recovery invalidate old state and rebuild quota correctly. | -| A21 | Checkpoint upload needs no base; patched residents rebuild before reuse. | - -Review regressions also cover nonblocking resident recovery, launcher initialization -rollback, lost completion notification, and sequence-receipt expiry. CPU/logical Ray -resources are not device training evidence. The standalone resource lifecycle is accepted with DP on both backends. -Checkpoint restoration and FSDP/TP/SP backend repairs are tracked separately in #315. -Raw device evidence is retained outside the public source tree. diff --git a/docs/design/validation/on-demand-training/README.md b/docs/design/validation/on-demand-training/README.md deleted file mode 100644 index ff7d2b4a8..000000000 --- a/docs/design/validation/on-demand-training/README.md +++ /dev/null @@ -1,30 +0,0 @@ -# On-demand training validation - -PR #312 is based directly on main. It does not require the checkpoint, -FSDP optimizer, or Megatron TP/SP repairs in #315. Its backend changes are -limited to resetting resident full-training state and using the supported -Transformers configuration import needed to initialize Megatron. - -Collected CPU tests under `tests/server/training/` exercise the controller, -allocation, recovery, and both installed SDKs over HTTP using a mock backend. -The HTTP and allocation probes run in isolated subprocesses. Existing -applications YAML startup can also be checked with memory and Redis state. - -Standalone validation programs live in `tools/validation/on_demand_training/`. -The DP resource-lifecycle device suite is explicitly invoked: - -```bash -python -m pytest tools/validation/on_demand_training/device_validation.py -vv -``` - -It runs DP=2, TP=1, and no FSDP for both backends and both SDKs in shared -LoRA/exclusive full modes. Assertions cover actual training and tenant isolation, -release, cross-model reuse, business-idle reclamation, expired handles, and clean -recreation. Checkpoint I/O and FSDP/TP repairs belong to #315, not this acceptance. -Results from older revisions containing #315 do not validate standalone #312. - -Device validation requires explicit `TWINKLE_DEVICE_*` configuration, a separate -evidence directory, and exclusively available devices. Record the exact source -SHA, import origins, raw JUnit/logs, numerical assertions, and resource cleanup. -Current results and remaining limitations are recorded in the English PR body; -results from an older revision are not acceptance evidence for the new head. diff --git a/pyproject.toml b/pyproject.toml index 32ea8e643..a5dafee76 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,11 +30,11 @@ megatron = ["megatron-core>=0.12.0", "transformer-engine[pytorch]", "mcore_bridg data = ["py-data-juicer"] rl = [ "vllm>=0.11", - "ray[serve]" + "ray[serve]>=2.55" ] async-rl = [ "vllm>=0.11", - "ray[serve]", + "ray[serve]>=2.55", "TransferQueue>=0.1.9.dev0", "math-verify==0.8.0", "swanlab>=0.8", diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 5f4c84f99..f255b5325 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -1501,6 +1501,15 @@ def reload_initial_weights(self, **kwargs): (exclusive) model, so the next tenant starts from clean pretrained weights instead of the previous tenant's trained weights. """ + if kwargs.get('reset_training_state', False): + for optimizer_group in self.optimizer_group.values(): + self.unregister_mm_forward_hook(optimizer_group) + self.optimizer_group = {_default_adapter_name: self._construct_default_optimizer_group()} + self.optimizer_group[_default_adapter_name].adapter_name = _default_adapter_name + self.active_group = _default_adapter_name + if self._model_wrapped: + self.register_mm_forward_hook(self.optimizer_group[_default_adapter_name]) + if not self.model_id: logger.warning('reload_initial_weights skipped: model_id is not set (blank model).') return @@ -1511,15 +1520,6 @@ def reload_initial_weights(self, **kwargs): for param in self.model.parameters(): param.grad = None - if kwargs.get('reset_training_state', False): - for optimizer_group in self.optimizer_group.values(): - self.unregister_mm_forward_hook(optimizer_group) - self.optimizer_group = {_default_adapter_name: self._construct_default_optimizer_group()} - self.optimizer_group[_default_adapter_name].adapter_name = _default_adapter_name - self.active_group = _default_adapter_name - if self._model_wrapped: - self.register_mm_forward_hook(self.optimizer_group[_default_adapter_name]) - def _load_optimizer(self, checkpoint_dir, **kwargs): adapter_name = kwargs.pop('adapter_name', _default_adapter_name) strict = kwargs.pop('strict', False) diff --git a/src/twinkle/server/config/__init__.py b/src/twinkle/server/config/__init__.py index 3e3fe2f71..9eb6aa09e 100644 --- a/src/twinkle/server/config/__init__.py +++ b/src/twinkle/server/config/__init__.py @@ -5,6 +5,7 @@ ServerArgs, SpectralHybridArgs) from .persistence import PersistenceConfig from .server_config import ServerConfig +from .task_queue import TaskQueueConfig from .telemetry import TelemetryConfig __all__ = [ @@ -18,5 +19,6 @@ 'ServerArgs', 'SpectralHybridArgs', 'ServerConfig', + 'TaskQueueConfig', 'TelemetryConfig', ] diff --git a/src/twinkle/server/config/application_spec.py b/src/twinkle/server/config/application_spec.py index ea2094f62..924beece0 100644 --- a/src/twinkle/server/config/application_spec.py +++ b/src/twinkle/server/config/application_spec.py @@ -15,8 +15,8 @@ from pydantic import BaseModel, ConfigDict, Field, model_validator from typing import Any, Literal -from twinkle.server.task_queue.config import TaskQueueConfig from .hybrid import SpectralHybridArgs +from .task_queue import TaskQueueConfig from .training import TrainingModes # Env var keys the launcher sets from the gateway ``server_config`` so that any diff --git a/src/twinkle/server/config/server_config.py b/src/twinkle/server/config/server_config.py index 5ac9aad36..a87e7f037 100644 --- a/src/twinkle/server/config/server_config.py +++ b/src/twinkle/server/config/server_config.py @@ -20,9 +20,9 @@ from typing import Any from twinkle.server.exceptions import ConfigParseError -from twinkle.server.task_queue.config import TaskQueueConfig from .application_spec import ApplicationSpec, HttpOptions from .persistence import PersistenceConfig +from .task_queue import TaskQueueConfig from .telemetry import TelemetryConfig from .training import TrainingConfig, TrainingOptions from .training_normalization import normalize_model_application diff --git a/src/twinkle/server/task_queue/config.py b/src/twinkle/server/config/task_queue.py similarity index 78% rename from src/twinkle/server/task_queue/config.py rename to src/twinkle/server/config/task_queue.py index 76f1813e8..4dede86b6 100644 --- a/src/twinkle/server/task_queue/config.py +++ b/src/twinkle/server/config/task_queue.py @@ -1,12 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -""" -Task queue configuration. - -Provides TaskQueueConfig (Pydantic) for controlling rate limits, timeouts, -and queue behavior. Constraints are validated at construction time so an -invalid YAML/dict value is rejected before the deployment reaches a ready -state. -""" +"""Task queue configuration.""" from __future__ import annotations from pydantic import BaseModel, ConfigDict, Field @@ -55,20 +48,13 @@ class TaskQueueConfig(BaseModel): @property def effective_execution_timeout(self) -> float: - """The single source of the execution time bound. - - ``0`` is not rejected (that would fail existing deployments); it is read - as "no configured limit" and replaced by a finite fallback so the bound - is always positive. This value feeds both ``_ray_get_timeout`` and the - ComputeWorker's ``asyncio.wait_for`` -- there is no second, independently - configurable timeout. - """ + """Return the single execution time bound used by all queue paths.""" if self.execution_timeout > 0: return self.execution_timeout return _ZERO_EXECUTION_TIMEOUT_FALLBACK def absolute_future_ttl(self, collect_width: int) -> float: - """Conservative lifetime for a non-terminal future record.""" + """Return a conservative lifetime for a non-terminal future record.""" ray_timeout = max(self.effective_execution_timeout, _MAX_DECLARED_BACKEND_TIMEOUT) resource_bound = max(1, collect_width) * ray_timeout return _ABSOLUTE_TTL_MULTIPLIER * (self.queue_timeout + resource_bound) diff --git a/src/twinkle/server/config/training.py b/src/twinkle/server/config/training.py index e68929c62..7c7a10bf8 100644 --- a/src/twinkle/server/config/training.py +++ b/src/twinkle/server/config/training.py @@ -5,8 +5,8 @@ from pydantic import BaseModel, ConfigDict, Field, model_validator from typing import Literal -from twinkle.server.task_queue.config import TaskQueueConfig from .hybrid import SpectralHybridArgs +from .task_queue import TaskQueueConfig class _TrainingConfig(BaseModel): diff --git a/src/twinkle/server/env_propagation.py b/src/twinkle/server/env_propagation.py new file mode 100644 index 000000000..9482acc90 --- /dev/null +++ b/src/twinkle/server/env_propagation.py @@ -0,0 +1,37 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Collect environment variables propagated to Ray worker processes.""" +from __future__ import annotations + +import os + +TELEMETRY_ENV_KEYS: tuple[str, ...] = ( + 'TWINKLE_STATE_RAY_NAMESPACE', + 'TWINKLE_TELEMETRY_ENABLED', + 'TWINKLE_TELEMETRY_DEBUG', + 'TWINKLE_TELEMETRY_SERVICE', + 'TWINKLE_TELEMETRY_ENDPOINT', + 'TWINKLE_TELEMETRY_INTERVAL', + 'TWINKLE_MODEL_ID_ALIASES', +) + + +def build_telemetry_env_vars() -> dict[str, str]: + return {key: os.environ[key] for key in TELEMETRY_ENV_KEYS if key in os.environ} + + +def build_persistence_env_vars() -> dict[str, str]: + from twinkle.server.config.persistence import PERSISTENCE_ENV_KEYS + return {key: os.environ[key] for key in PERSISTENCE_ENV_KEYS if key in os.environ} + + +def build_server_state_env_vars() -> dict[str, str]: + from twinkle.server.config.application_spec import SERVER_STATE_ENV_KEYS + return {key: os.environ[key] for key in SERVER_STATE_ENV_KEYS if key in os.environ} + + +def build_propagated_env_vars() -> dict[str, str]: + merged: dict[str, str] = {} + merged.update(build_telemetry_env_vars()) + merged.update(build_persistence_env_vars()) + merged.update(build_server_state_env_vars()) + return merged diff --git a/src/twinkle/server/gateway/app.py b/src/twinkle/server/gateway/app.py index 0f630055c..d284f83dd 100644 --- a/src/twinkle/server/gateway/app.py +++ b/src/twinkle/server/gateway/app.py @@ -13,10 +13,10 @@ import twinkle.protocol.types as types from twinkle.server.deployment import LazyCleanupMixin, bind_deployment, build_deployment_app +from twinkle.server.service_proxy import ServiceProxy from twinkle.server.state import get_server_state from twinkle.utils.logger import get_logger from .openai_handlers import _register_openai_routes -from .proxy import ServiceProxy from .tinker_handlers import _register_gateway_tinker_routes from .twinkle_handlers import _register_gateway_twinkle_routes diff --git a/src/twinkle/server/gateway/proxy.py b/src/twinkle/server/gateway/proxy.py deleted file mode 100644 index 664d68d23..000000000 --- a/src/twinkle/server/gateway/proxy.py +++ /dev/null @@ -1,230 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -""" -Proxy utilities for forwarding requests to internal model/sampler services. - -Moved from tinker/proxy.py. Updated proxy_to_model and proxy_to_sampler -to prepend the 'tinker/' prefix to endpoints so they route to /tinker/* paths -on the unified model/sampler deployments. -""" -from __future__ import annotations - -import httpx -from fastapi import Request, Response -from fastapi.responses import JSONResponse -from typing import Any - -from twinkle.protocol.headers import H_MULTIPLEX, H_MULTIPLEX_LEGACY, H_REQUEST_ID, H_REQUEST_ID_LEGACY -from twinkle.protocol.types.errors import ErrorCategory, ErrorPayload -from twinkle.server.telemetry.tracing import inject_context -from twinkle.utils.logger import get_logger -from . import routes - -logger = get_logger() - - -class ServiceProxy: - """HTTP proxy for routing requests to internal model and sampler services. - - Handles: - 1. URL construction using localhost to avoid external routing loops - 2. Header forwarding with appropriate cleanup - 3. Debug logging for troubleshooting - 4. Error handling and response forwarding - - Tinker endpoints are routed to /tinker/ on the unified deployments. - """ - - def __init__( - self, - http_options: dict[str, Any] | None = None, - route_prefix: str = '/api/v1', - ): - self.http_options = http_options or {} - self.route_prefix = route_prefix - # Disable proxy env vars to avoid external routing - self.client = httpx.AsyncClient(timeout=None, trust_env=False) - - async def close(self) -> None: - """Close the underlying httpx.AsyncClient to release connections.""" - await self.client.aclose() - - def _build_target_url(self, service_type: str, base_model: str, endpoint: str) -> str: - """Build the target URL for internal service routing. - - Args: - service_type: Either 'model' or 'sampler' - base_model: The base model name for routing - endpoint: The target endpoint name (already includes tinker/ or twinkle/ prefix) - - Returns: - Complete target URL for the internal service - """ - host = self.http_options.get('host', 'localhost') - port = self.http_options.get('port', 8000) - - if host == '0.0.0.0': - host = 'localhost' - - base_url = f'http://{host}:{port}' - return f'{base_url}{routes.target_url(self.route_prefix, service_type, base_model, endpoint)}' - - def _prepare_headers(self, request_headers) -> dict[str, str]: - """Prepare headers for proxying by removing problematic headers.""" - logger.debug('prepare_headers request_headers=%s', request_headers) - headers = dict(request_headers) - headers.pop('host', None) - headers.pop('content-length', None) - request_id = request_headers.get(H_REQUEST_ID) or request_headers.get(H_REQUEST_ID_LEGACY) - if request_id is not None and not request_headers.get(H_MULTIPLEX_LEGACY): - headers[H_MULTIPLEX_LEGACY] = request_id - if request_id is not None: - headers.setdefault(H_REQUEST_ID, request_id) - headers.setdefault(H_MULTIPLEX, request_id) - return headers - - async def proxy_request( - self, - request: Request, - endpoint: str, - base_model: str, - service_type: str, - *, - body_override: bytes | None = None, - extra_headers: dict[str, str] | None = None, - ) -> Response: - """Generic proxy method to forward requests to model or sampler services. - - Args: - request: The incoming FastAPI request - endpoint: The target endpoint path (e.g., 'tinker/create_model') - base_model: The base model name for routing - service_type: Either 'model' or 'sampler' - body_override: If set, use this as the request body instead of reading from request - extra_headers: Additional headers to merge into the forwarded request - - Returns: - Proxied response from the target service - """ - body_bytes = body_override if body_override is not None else await request.body() - target_url = self._build_target_url(service_type, base_model, endpoint) - headers = self._prepare_headers(request.headers) - if extra_headers: - headers.update(extra_headers) - - # Inject current trace context into outgoing headers for distributed tracing. - # When telemetry is not initialized, this is a noop. - inject_context(headers) - - try: - logger.debug( - 'proxy_request service=%s endpoint=%s target_url=%s request_id=%s', - service_type, - endpoint, - target_url, - headers.get(H_MULTIPLEX_LEGACY), - ) - - response = await self.client.request( - method=request.method, - url=target_url, - content=body_bytes, - headers=headers, - params=request.query_params, - ) - - logger.debug( - 'proxy_response status=%s body_preview=%s', - response.status_code, - response.text[:200], - ) - - return Response( - content=response.content, - status_code=response.status_code, - headers=dict(response.headers), - media_type=response.headers.get('content-type'), - ) - except Exception as e: - logger.error('Proxy error: %s', str(e), exc_info=True) - # The gateway could not reach the upstream deployment. Return the - # unified ``ErrorPayload`` (502/Server) instead of a plain-text body - # so every gateway failure has the same wire shape. Upstream error - # responses are passed through unchanged above, preserving their own - # ErrorPayload body. - request_id = request.headers.get(H_REQUEST_ID) or request.headers.get(H_REQUEST_ID_LEGACY) or '' - payload = ErrorPayload( - error=f'Proxy Error: {str(e)}', - category=ErrorCategory.Server, - error_code=502, - request_id=request_id, - ) - return JSONResponse(status_code=502, content=payload.model_dump(mode='json', exclude_none=True)) - - async def proxy_request_stream( - self, - request: Request, - endpoint: str, - base_model: str, - service_type: str, - *, - body_override: bytes | None = None, - extra_headers: dict[str, str] | None = None, - ): - """Proxy with httpx streaming — returns an async generator of response lines. - - Yields: - str lines from the upstream response (newline-delimited JSON). - - Raises: - httpx.HTTPStatusError: If the upstream returns a non-2xx status. - """ - body_bytes = body_override if body_override is not None else await request.body() - target_url = self._build_target_url(service_type, base_model, endpoint) - headers = self._prepare_headers(request.headers) - if extra_headers: - headers.update(extra_headers) - inject_context(headers) - - logger.debug( - 'proxy_request_stream service=%s endpoint=%s target_url=%s', - service_type, - endpoint, - target_url, - ) - - async with self.client.stream( - method='POST', - url=target_url, - content=body_bytes, - headers=headers, - ) as response: - if response.status_code >= 400: - body = await response.aread() - raise httpx.HTTPStatusError( - f'Upstream returned {response.status_code}', - request=response.request, - response=httpx.Response(status_code=response.status_code, content=body), - ) - async for line in response.aiter_lines(): - if line: - yield line - - async def proxy_to_model(self, request: Request, endpoint: str, base_model: str) -> Response: - """Proxy request to model's tinker endpoint (/tinker/). - - Args: - request: The incoming FastAPI request - endpoint: The tinker endpoint name (e.g., 'create_model', 'forward') - base_model: The base model name for routing - """ - return await self.proxy_request(request, routes.tinker_endpoint(endpoint), base_model, 'model') - - async def proxy_to_sampler(self, request: Request, endpoint: str, base_model: str) -> Response: - """Proxy request to sampler's tinker endpoint (/tinker/). - - Args: - request: The incoming FastAPI request - endpoint: The tinker endpoint name (e.g., 'asample') - base_model: The base model name for routing - """ - return await self.proxy_request(request, routes.tinker_endpoint(endpoint), base_model, 'sampler') diff --git a/src/twinkle/server/gateway/tinker_handlers.py b/src/twinkle/server/gateway/tinker_handlers.py index 74f81f32c..1a16e45aa 100644 --- a/src/twinkle/server/gateway/tinker_handlers.py +++ b/src/twinkle/server/gateway/tinker_handlers.py @@ -30,10 +30,8 @@ # Keys must equal ``state.models.FAILURE_REASON_CODES`` (guarded by test_envelope). _TINKER_FAILURE_WIRE: dict[str, tuple[int, str]] = { 'capacity_unavailable': (409, 'user'), - 'training_expired': (404, 'user'), 'startup_failed': (500, 'server'), 'startup_timeout': (504, 'server'), - 'protocol_mismatch': (400, 'user'), 'invalid_request': (400, 'user'), 'request_rejected': (400, 'user'), 'resource_not_found': (404, 'user'), diff --git a/src/twinkle/server/launcher/env_propagation.py b/src/twinkle/server/launcher/env_propagation.py deleted file mode 100644 index 61941669c..000000000 --- a/src/twinkle/server/launcher/env_propagation.py +++ /dev/null @@ -1,53 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Collection of telemetry / persistence env vars for propagation to Ray workers. - -These variables are read inside each Ray Serve worker process — telemetry by -``ensure_telemetry_initialized()`` and persistence -by ``PersistenceConfig.from_env()`` — so the chosen backend / telemetry config -is independent of deployment startup order. -""" -from __future__ import annotations - -import os - -# Telemetry env var keys that need to be propagated to Ray worker processes. -TELEMETRY_ENV_KEYS: tuple[str, ...] = ( - 'TWINKLE_STATE_RAY_NAMESPACE', - 'TWINKLE_TELEMETRY_ENABLED', - 'TWINKLE_TELEMETRY_DEBUG', - 'TWINKLE_TELEMETRY_SERVICE', - 'TWINKLE_TELEMETRY_ENDPOINT', - 'TWINKLE_TELEMETRY_INTERVAL', - 'TWINKLE_MODEL_ID_ALIASES', -) - - -def build_telemetry_env_vars() -> dict[str, str]: - """Collect telemetry env vars from ``os.environ`` for worker propagation.""" - return {k: os.environ[k] for k in TELEMETRY_ENV_KEYS if k in os.environ} - - -def build_persistence_env_vars() -> dict[str, str]: - """Collect persistence env vars from ``os.environ`` for worker propagation.""" - from twinkle.server.config.persistence import PERSISTENCE_ENV_KEYS - return {k: os.environ[k] for k in PERSISTENCE_ENV_KEYS if k in os.environ} - - -def build_server_state_env_vars() -> dict[str, str]: - """Collect ServerState-policy env vars from ``os.environ`` for worker propagation. - - Read inside each worker by ``ServerStateArgs.from_env()`` (via - ``get_server_state``) so the configured quota / expiry / metrics interval is - applied everywhere, not only in the gateway that first built the state. - """ - from twinkle.server.config.application_spec import SERVER_STATE_ENV_KEYS - return {k: os.environ[k] for k in SERVER_STATE_ENV_KEYS if k in os.environ} - - -def build_propagated_env_vars() -> dict[str, str]: - """Aggregate all env vars that must reach Ray worker processes.""" - merged: dict[str, str] = {} - merged.update(build_telemetry_env_vars()) - merged.update(build_persistence_env_vars()) - merged.update(build_server_state_env_vars()) - return merged diff --git a/src/twinkle/server/launcher/server_launcher.py b/src/twinkle/server/launcher/server_launcher.py index 47921d70b..342fe469b 100644 --- a/src/twinkle/server/launcher/server_launcher.py +++ b/src/twinkle/server/launcher/server_launcher.py @@ -21,8 +21,8 @@ from twinkle.patch.ray_serve import apply_ray_serve_patches, get_runtime_env_for_patches from twinkle.server.config import ServerConfig from twinkle.server.config.application_spec import ApplicationSpec +from twinkle.server.env_propagation import build_propagated_env_vars from .builder_registry import get_builders, resolve_builder -from .env_propagation import build_propagated_env_vars logger = get_logger() diff --git a/src/twinkle/server/lifecycle/envelope.py b/src/twinkle/server/lifecycle/envelope.py index d88fff666..f089e29ec 100644 --- a/src/twinkle/server/lifecycle/envelope.py +++ b/src/twinkle/server/lifecycle/envelope.py @@ -18,10 +18,8 @@ # Keys must equal ``state.models.FAILURE_REASON_CODES`` (guarded by test_envelope). _FAILURE_WIRE: dict[str, tuple[int, ErrorCategory]] = { 'capacity_unavailable': (409, ErrorCategory.User), - 'training_expired': (404, ErrorCategory.User), 'startup_failed': (500, ErrorCategory.Server), 'startup_timeout': (504, ErrorCategory.Server), - 'protocol_mismatch': (400, ErrorCategory.User), 'invalid_request': (400, ErrorCategory.User), 'request_rejected': (400, ErrorCategory.User), 'resource_not_found': (404, ErrorCategory.User), diff --git a/src/twinkle/server/lifecycle/protocols.py b/src/twinkle/server/lifecycle/protocols.py index c05ea12bb..3bee91ade 100644 --- a/src/twinkle/server/lifecycle/protocols.py +++ b/src/twinkle/server/lifecycle/protocols.py @@ -19,9 +19,9 @@ from fastapi import Request from typing import Any, Protocol, runtime_checkable +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.lifecycle.envelope import TaskEnvelope from twinkle.server.state import ServerState -from twinkle.server.task_queue.config import TaskQueueConfig @runtime_checkable diff --git a/src/twinkle/server/model/app.py b/src/twinkle/server/model/app.py index 34f9c3a5c..dfdb7e5b7 100644 --- a/src/twinkle/server/model/app.py +++ b/src/twinkle/server/model/app.py @@ -9,19 +9,19 @@ import asyncio from fastapi import FastAPI, Request -from ray import serve from typing import Any from twinkle import DeviceGroup from twinkle.server.config.backend_dispatch import BackendSelector +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.deployment import bind_deployment, build_deployment_app from twinkle.server.runtime import init_twinkle_runtime from twinkle.server.session_resource import AdapterManagerMixin from twinkle.server.state import ServerState, get_server_state -from twinkle.server.task_queue import TaskQueueConfig, TaskQueueMixin -from twinkle.server.training.runtime import ManagedTrainingRuntime +from twinkle.server.task_queue import TaskQueueMixin from twinkle.server.utils import wrap_builder_with_device_group_env from twinkle.utils.logger import get_logger +from .managed_runtime import ManagedTrainingRuntime from .tinker_handlers import _register_model_tinker_routes from .twinkle_handlers import _register_model_twinkle_routes @@ -86,7 +86,7 @@ class ModelManagement(ManagedTrainingRuntime, TaskQueueMixin, AdapterManagerMixi - Base model and multiple LoRA adapters (multi-user) - Tinker training operations via /tinker/* endpoints (async/polling) - Twinkle training operations via /twinkle/* endpoints (synchronous) - - Adapter lifecycle via AdapterManagerMixin + - AdapterManagerMixin provides local registration only; expiry is handled by the training controller. - Per-user rate limiting via TaskQueueMixin """ @@ -96,7 +96,6 @@ async def __init__(self, device_group: dict[str, Any], device_mesh: dict[str, Any], backend: str, - adapter_config: dict[str, Any] | None = None, queue_config: TaskQueueConfig | None = None, data_plane_url: str | None = None, **kwargs): @@ -148,7 +147,7 @@ async def __init__(self, collect_width=len(actors) if actors else 1, ) self.model._ray_get_timeout = self.task_queue_config.effective_execution_timeout - self._init_adapter_manager(**(adapter_config or {})) + self._init_adapter_manager() await self._register_replica_on_startup() @property @@ -226,7 +225,6 @@ def build_model_app(model_id: str, device_mesh: dict[str, Any], deploy_options: dict[str, Any], backend: str, - adapter_config: dict[str, Any] | None = None, queue_config: TaskQueueConfig | None = None, data_plane_url: str | None = None, **kwargs): @@ -243,7 +241,6 @@ def build_model_app(model_id: str, backend: Model backend selector — ``mock`` | ``transformers`` | ``megatron``. Validated up front; bad values raise :class:`ConfigError` before any side effect. - adapter_config: Adapter lifecycle config (timeout, per-token limits) queue_config: Task queue configuration (rate limiting, etc.) **kwargs: Additional model initialization arguments @@ -272,8 +269,7 @@ async def _on_shutdown(servable: Any) -> None: ModelManagement, deploy_options, deployment_name='ModelManagement', - bind_args=(model_id, nproc_per_node, device_group, device_mesh, backend, adapter_config, queue_config, - data_plane_url), + bind_args=(model_id, nproc_per_node, device_group, device_mesh, backend, queue_config, data_plane_url), bind_kwargs=kwargs, ) diff --git a/src/twinkle/server/training/runtime.py b/src/twinkle/server/model/managed_runtime.py similarity index 88% rename from src/twinkle/server/training/runtime.py rename to src/twinkle/server/model/managed_runtime.py index 455ccf040..3c4c54c04 100644 --- a/src/twinkle/server/training/runtime.py +++ b/src/twinkle/server/model/managed_runtime.py @@ -17,9 +17,9 @@ from twinkle.server.lifecycle.submit import backend_kwargs from twinkle.server.middleware.auth import get_session_id_from_request, get_token_from_request from twinkle.server.state.models import FutureFailureRecord +from twinkle.server.state.training_records import CreateTraining, TrainingRecord from twinkle.server.utils import get_template_for_model from twinkle.utils.logger import get_logger -from .records import CreateTraining, TrainingRecord _tinker_operation = ContextVar('managed_tinker_operation', default=None) logger = get_logger() @@ -28,9 +28,6 @@ class ManagedTrainingRuntime: async def wait_for_training_workers(self): - # Constructing a remote_class returns actor handles before its worker - # constructors (including model loading) finish. Full creation otherwise - # performs no backend call to establish this readiness barrier. actors = getattr(self.model, '_actors', None) or [] await asyncio.gather(*(actor.__ray_ready__.remote() for actor in actors)) @@ -69,7 +66,7 @@ async def schedule_task(self, coro_factory, **kwargs): training_id, seq_id = operation task_type = kwargs.get('task_type') stream = f'tinker::{task_type}' - request_id = 'req_' + uuid.uuid4().hex + request_id = kwargs.pop('request_id', None) or 'req_' + uuid.uuid4().hex retention = self.sequence_retention_s previous = await self.state.claim_training_seq(training_id, stream, seq_id, request_id, retention) if previous is not None: @@ -93,7 +90,11 @@ async def _prepare_queued_task(self, task): renew=task.task_type not in ('get_info', 'get_train_configs')) async def enqueued(): - await self._training_controller.commit_operation.remote(task.model_id, task.request_id, context['epoch']) + try: + await self._training_controller.commit_operation.remote(task.model_id, task.request_id, + context['epoch']) + except Exception: + logger.exception('Training enqueue notification failed for %s', task.request_id) async def finished(executed, unsafe): if unsafe: @@ -109,8 +110,6 @@ async def finished(executed, unsafe): await self._training_controller.finish_operation.remote( task.model_id, task.request_id, context['epoch'], executed=executed) except Exception: - # The terminal future and queue membership allow the controller - # to reconcile a lost notification without stopping compute. logger.exception('Training completion notification failed for %s', task.request_id) task.on_enqueued, task.on_finished = enqueued, finished @@ -141,7 +140,8 @@ async def initialize(): self.set_resource_state(adapter, 'grad_ready', False) if request.client_type == 'twinkle': body = native.AddAdapterRequest.model_validate(request.payload) - config = deserialize_object(body.config) + # Dynamic imports (e.g. PEFT) and deserialization must not block the health-probe event loop. + config = await asyncio.to_thread(deserialize_object, body.config) if not self.is_full_mode: await self.call_backend(self.model.add_adapter_to_model, adapter, config, **backend_kwargs(body)) lora = None if config is None else RunLoraConfig( @@ -152,14 +152,16 @@ async def initialize(): save_dir=body.save_dir, user_metadata={'adapter_name': training.adapter_label}) else: - from peft import LoraConfig from tinker import types run_config = types.CreateModelRequest.model_validate(request.payload) if not self.is_full_mode: + def build_lora_config(): + from peft import LoraConfig + return LoraConfig(r=run_config.lora_config.rank, target_modules='all-linear') + + config = await asyncio.to_thread(build_lora_config) await self.call_backend( - self.model.add_adapter_to_model, - adapter_name=adapter, - config_or_dir=LoraConfig(r=run_config.lora_config.rank, target_modules='all-linear')) + self.model.add_adapter_to_model, adapter_name=adapter, config_or_dir=config) await self.call_backend( self.model.set_template, get_template_for_model(self.base_model), @@ -167,7 +169,10 @@ async def initialize(): model_id=self.base_model) await self.call_backend(self.model.set_processor, 'InputProcessor', adapter_name=backend_adapter) await self.call_backend(self.model.set_optimizer, 'Adam', adapter_name=backend_adapter) - create_training_run_manager(training.token, client_type=request.client_type).save(adapter, run_config) + def save_training_metadata(): + create_training_run_manager(training.token, client_type=request.client_type).save(adapter, run_config) + + await asyncio.to_thread(save_training_metadata) await self._managed_serial(initialize, 'initialize_training') diff --git a/src/twinkle/server/resources/__init__.py b/src/twinkle/server/resources/__init__.py new file mode 100644 index 000000000..5fe876191 --- /dev/null +++ b/src/twinkle/server/resources/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shared device-resource allocation for training, sampling, and future workloads.""" diff --git a/src/twinkle/server/training/allocation.py b/src/twinkle/server/resources/allocation.py similarity index 84% rename from src/twinkle/server/training/allocation.py rename to src/twinkle/server/resources/allocation.py index 08bdf8573..bfd887378 100644 --- a/src/twinkle/server/training/allocation.py +++ b/src/twinkle/server/resources/allocation.py @@ -1,5 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Single-node device allocation. Only the owner removes placement groups.""" +"""Single-node device allocation; only the resource owner removes placement groups.""" from __future__ import annotations import asyncio @@ -7,8 +7,11 @@ from twinkle import Platform from twinkle.infra._ray.resource_manager import _gpu_placement_group_cpus -from .catalog import TrainingSpec -from .errors import CapacityUnavailableError +from twinkle.server.exceptions import RequestRejectedError + + +class CapacityUnavailableError(RequestRejectedError): + error_code = 409 @dataclass(frozen=True) @@ -39,15 +42,17 @@ def __init__(self, namespace: str, *, allocation_timeout: float, shutdown_timeou self.shutdown_timeout = shutdown_timeout def name(self, epoch: str, instance_id: str) -> str: - return f'training-{self.namespace}-{epoch}-{instance_id}' + return f'resources-{self.namespace}-{epoch}-{instance_id}' - async def allocate(self, spec: TrainingSpec, epoch: str, instance_id: str) -> DeviceAllocation: + async def allocate(self, num_devices: int, device_type: str, epoch: str, instance_id: str) -> DeviceAllocation: import ray from ray.util.placement_group import placement_group from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy - resource = 'CPU' if spec.model.device_type == 'cpu' else Platform.get_platform(spec.model.device_type).__name__ - count = spec.config.num_devices + if type(num_devices) is not int or num_devices <= 0: + raise ValueError('num_devices must be a positive integer') + resource = 'CPU' if device_type == 'cpu' else Platform.get_platform(device_type).__name__ + count = num_devices candidates = [node for node in ray.nodes() if node.get('Alive') and node['Resources'].get(resource, 0) >= count] available = ray.available_resources() if not candidates or available.get(resource, 0) < count: @@ -64,7 +69,7 @@ async def allocate(self, spec: TrainingSpec, epoch: str, instance_id: str) -> De elif resource != 'CPU': options['resources'] = {resource: count} node, devices = await asyncio.wait_for( - ray.remote(_probe_allocation).options(**options).remote(resource, spec.model.device_type), + ray.remote(_probe_allocation).options(**options).remote(resource, device_type), self.shutdown_timeout) if resource == 'CPU': devices = [str(index) for index in range(count)] diff --git a/src/twinkle/server/sampler/app.py b/src/twinkle/server/sampler/app.py index f0ee357df..42e3be7a3 100644 --- a/src/twinkle/server/sampler/app.py +++ b/src/twinkle/server/sampler/app.py @@ -13,11 +13,11 @@ from twinkle import DeviceGroup from twinkle.server.config.backend_dispatch import BackendSelector +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.deployment import LazyCleanupMixin, bind_deployment, build_deployment_app from twinkle.server.middleware.auth import get_token_from_request from twinkle.server.runtime import init_twinkle_runtime from twinkle.server.state import ServerState, get_server_state -from twinkle.server.task_queue.config import TaskQueueConfig from twinkle.server.task_queue.mixin import TaskQueueMixin from twinkle.server.utils import get_template_for_model, wrap_builder_with_device_group_env from twinkle.utils.logger import get_logger diff --git a/src/twinkle/server/service_proxy.py b/src/twinkle/server/service_proxy.py new file mode 100644 index 000000000..071d2253f --- /dev/null +++ b/src/twinkle/server/service_proxy.py @@ -0,0 +1,105 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shared HTTP forwarding between services; independent of gateway and training controller.""" +from __future__ import annotations + +import httpx +from fastapi import Request, Response +from fastapi.responses import JSONResponse +from typing import Any + +from twinkle.protocol.headers import H_MULTIPLEX, H_MULTIPLEX_LEGACY, H_REQUEST_ID, H_REQUEST_ID_LEGACY +from twinkle.protocol.types.errors import ErrorCategory, ErrorPayload +from twinkle.server.telemetry.tracing import inject_context +from twinkle.utils.logger import get_logger + +logger = get_logger() + + +class ServiceProxy: + + def __init__(self, http_options: dict[str, Any] | None = None, route_prefix: str = '/api/v1'): + self.http_options = http_options or {} + self.route_prefix = route_prefix + self.client = httpx.AsyncClient(timeout=None, trust_env=False) + + async def close(self) -> None: + await self.client.aclose() + + def _build_target_url(self, service_type: str, base_model: str, endpoint: str) -> str: + host = self.http_options.get('host', 'localhost') + port = self.http_options.get('port', 8000) + if host == '0.0.0.0': + host = 'localhost' + return f'http://{host}:{port}{self.route_prefix.rstrip("/")}/{service_type}/{base_model}/{endpoint}' + + def _prepare_headers(self, request_headers) -> dict[str, str]: + headers = dict(request_headers) + headers.pop('host', None) + headers.pop('content-length', None) + request_id = request_headers.get(H_REQUEST_ID) or request_headers.get(H_REQUEST_ID_LEGACY) + if request_id is not None and not request_headers.get(H_MULTIPLEX_LEGACY): + headers[H_MULTIPLEX_LEGACY] = request_id + if request_id is not None: + headers.setdefault(H_REQUEST_ID, request_id) + headers.setdefault(H_MULTIPLEX, request_id) + return headers + + async def proxy_request(self, + request: Request, + endpoint: str, + base_model: str, + service_type: str, + *, + body_override: bytes | None = None, + extra_headers: dict[str, str] | None = None) -> Response: + body_bytes = body_override if body_override is not None else await request.body() + target_url = self._build_target_url(service_type, base_model, endpoint) + headers = self._prepare_headers(request.headers) + if extra_headers: + headers.update(extra_headers) + inject_context(headers) + try: + response = await self.client.request( + method=request.method, url=target_url, content=body_bytes, headers=headers, params=request.query_params) + return Response( + content=response.content, + status_code=response.status_code, + headers=dict(response.headers), + media_type=response.headers.get('content-type')) + except Exception as exc: + logger.error('Proxy error: %s', str(exc), exc_info=True) + request_id = request.headers.get(H_REQUEST_ID) or request.headers.get(H_REQUEST_ID_LEGACY) or '' + payload = ErrorPayload( + error=f'Proxy Error: {exc}', category=ErrorCategory.Server, error_code=502, request_id=request_id) + return JSONResponse(status_code=502, content=payload.model_dump(mode='json', exclude_none=True)) + + async def proxy_request_stream(self, + request: Request, + endpoint: str, + base_model: str, + service_type: str, + *, + body_override: bytes | None = None, + extra_headers: dict[str, str] | None = None): + body_bytes = body_override if body_override is not None else await request.body() + target_url = self._build_target_url(service_type, base_model, endpoint) + headers = self._prepare_headers(request.headers) + if extra_headers: + headers.update(extra_headers) + inject_context(headers) + async with self.client.stream(method='POST', url=target_url, content=body_bytes, headers=headers) as response: + if response.status_code >= 400: + body = await response.aread() + raise httpx.HTTPStatusError( + f'Upstream returned {response.status_code}', + request=response.request, + response=httpx.Response(status_code=response.status_code, content=body)) + async for line in response.aiter_lines(): + if line: + yield line + + async def proxy_to_model(self, request: Request, endpoint: str, base_model: str) -> Response: + return await self.proxy_request(request, f'tinker/{endpoint}', base_model, 'model') + + async def proxy_to_sampler(self, request: Request, endpoint: str, base_model: str) -> Response: + return await self.proxy_request(request, f'tinker/{endpoint}', base_model, 'sampler') diff --git a/src/twinkle/server/state/model_manager.py b/src/twinkle/server/state/model_manager.py index 2b01af9f0..de8f872c7 100644 --- a/src/twinkle/server/state/model_manager.py +++ b/src/twinkle/server/state/model_manager.py @@ -117,6 +117,22 @@ async def unregister_replica(self, replica_id: str) -> None: await self.remove(model_id) await self._replicas.unregister(replica_id) + async def unregister_replica_capacity(self, replica_id: str) -> None: + """Remove capacity and heartbeat only; model records are retired by the owning lifecycle controller.""" + await self._replicas.unregister(replica_id) + + async def save_record(self, model_id: str, record: ModelRecord) -> None: + # Update without re-applying quota; cannot reuse add(). + await self._backend.set(self._make_key(model_id), record.model_dump()) + + async def get_training_records(self, namespace: str): + from .training_records import TrainingRecord + keys = await self._backend.keys('model::*') + return [ + TrainingRecord.model_validate(value) for value in await self._backend.mget(keys) if isinstance(value, dict) + and value.get('controller_epoch') and value.get('training_namespace') == namespace + ] + async def touch_replica_last_seen(self, replica_id: str) -> None: """Refresh a replica's liveness timestamp.""" await self._replicas.touch_last_seen(replica_id) diff --git a/src/twinkle/server/state/models.py b/src/twinkle/server/state/models.py index 427dee87e..ced4475e1 100644 --- a/src/twinkle/server/state/models.py +++ b/src/twinkle/server/state/models.py @@ -65,10 +65,8 @@ class FutureFailureRecord(BaseModel): # in one place lets a consistency test catch a map that drifts out of coverage. FAILURE_REASON_CODES: frozenset[str] = frozenset({ 'capacity_unavailable', - 'training_expired', 'startup_failed', 'startup_timeout', - 'protocol_mismatch', 'invalid_request', 'request_rejected', 'resource_not_found', diff --git a/src/twinkle/server/state/server_state.py b/src/twinkle/server/state/server_state.py index 461d9a7a4..150d95f9f 100644 --- a/src/twinkle/server/state/server_state.py +++ b/src/twinkle/server/state/server_state.py @@ -173,9 +173,25 @@ async def _on_lose_leader(self) -> None: async def get_capacity_info(self) -> dict[str, int]: return await self._model_mgr.get_capacity_info() - @property - def backend(self) -> StateBackend: - return self._backend + def create_training_store(self, namespace: str): + from .training_store import TrainingStore + return TrainingStore(self, namespace, self._backend) + + async def save_model_record(self, model_id: str, record: ModelRecord) -> None: + await self._model_mgr.save_record(model_id, record) + + async def get_training_records(self, namespace: str): + return await self._model_mgr.get_training_records(namespace) + + async def set_future_owner(self, request_id: str, token: str) -> None: + await self._backend.set(f'future_owner::{request_id}', token, ttl=max(1, int(self.expiration_timeout))) + + async def delete_training_sequences(self, training_id: str) -> None: + for key in await self._backend.keys(f'seq_ledger::{training_id}::*'): + await self._backend.delete(key) + + async def unregister_replica_capacity(self, replica_id: str) -> None: + await self._model_mgr.unregister_replica_capacity(replica_id) async def register_model_record(self, model_id: str, record: ModelRecord) -> None: """Reserve quota for a complete model record, including managed training state.""" @@ -224,6 +240,9 @@ async def get_session_last_heartbeat(self, session_id: str) -> float | None: """ return await self._session_mgr.get_last_heartbeat(session_id) + async def get_session_last_heartbeats(self, session_ids: list[str]) -> dict[str, float | None]: + return await self._session_mgr.get_last_heartbeats(session_ids) + # ----- Model Registration ----- async def register_model(self, diff --git a/src/twinkle/server/state/session_manager.py b/src/twinkle/server/state/session_manager.py index 0e4beeb37..6249ffc9c 100644 --- a/src/twinkle/server/state/session_manager.py +++ b/src/twinkle/server/state/session_manager.py @@ -67,6 +67,13 @@ async def get_last_heartbeat(self, session_id: str) -> float | None: return None return record.last_heartbeat + async def get_last_heartbeats(self, session_ids: list[str]) -> dict[str, float | None]: + records = await self._backend.mget([self._make_key(identity) for identity in session_ids]) + return { + identity: SessionRecord.model_validate(record).last_heartbeat if record is not None else None + for identity, record in zip(session_ids, records, strict=True) + } + # ----- Cleanup ----- def _is_expired(self, record: SessionRecord, cutoff_time: float) -> bool: diff --git a/src/twinkle/server/training/records.py b/src/twinkle/server/state/training_records.py similarity index 93% rename from src/twinkle/server/training/records.py rename to src/twinkle/server/state/training_records.py index 4b7130ed4..e89eb3d3d 100644 --- a/src/twinkle/server/training/records.py +++ b/src/twinkle/server/state/training_records.py @@ -1,11 +1,11 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Serializable identities and states. Ray handles never enter these records.""" +"""Serializable training identities and states shared by controller and model layers.""" from __future__ import annotations from pydantic import BaseModel, Field from typing import Any, Literal -from twinkle.server.state.models import ModelRecord +from .models import ModelRecord TrainingState = Literal['PENDING', 'ACTIVE', 'CLOSING', 'RELEASED', 'EXPIRED', 'FAILED', 'CANCELLED'] InstanceState = Literal['ALLOCATING', 'STARTING', 'READY', 'RESETTING', 'DRAINING', 'STOPPED', 'QUARANTINED'] diff --git a/src/twinkle/server/training/store.py b/src/twinkle/server/state/training_store.py similarity index 69% rename from src/twinkle/server/training/store.py rename to src/twinkle/server/state/training_store.py index f44f1a7f4..c814e6a73 100644 --- a/src/twinkle/server/training/store.py +++ b/src/twinkle/server/state/training_store.py @@ -1,22 +1,25 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Controller-owned records on the existing shared state backend.""" +"""Controller-owned persistence; other key-spaces are accessed through ServerState.""" from __future__ import annotations import hashlib import json import time -from typing import Any +from typing import TYPE_CHECKING, Any -from twinkle.server.state import ServerState -from twinkle.server.state.models import FutureFailureRecord -from .records import CreateTraining, InstanceRecord, TrainingRecord +from .backend.base import StateBackend +from .models import FutureFailureRecord +from .training_records import CreateTraining, InstanceRecord, TrainingRecord + +if TYPE_CHECKING: + from .server_state import ServerState class TrainingStore: - def __init__(self, state: ServerState, namespace: str): + def __init__(self, state: ServerState, namespace: str, backend: StateBackend): self.state = state - self.backend = state.backend + self.backend = backend self.prefix = f'training::{namespace}::' def key(self, kind: str, identity: str) -> str: @@ -29,23 +32,19 @@ def create_identity(request: CreateTraining) -> tuple[str, str]: return hashlib.sha256(identity.encode()).hexdigest(), hashlib.sha256(fingerprint.encode()).hexdigest() async def reserve(self, record: TrainingRecord) -> None: - # ModelManager applies the existing global quota atomically, including PENDING. await self.state.register_model_record(record.training_id, record) async def save_training(self, record: TrainingRecord) -> None: - await self.backend.set(f'model::{record.training_id}', record.model_dump()) + await self.state.save_model_record(record.training_id, record) async def retire(self, record: TrainingRecord) -> None: - # Keep a bounded diagnostic tombstone after freeing the live quota row. await self.backend.set( self.key('terminal', record.training_id), record.model_dump(), ttl=max(1, int(self.state.expiration_timeout))) + await self.state.delete_training_sequences(record.training_id) + # Release the quota row last so a failed earlier step can be idempotently retried. await self.state.unload_model(record.training_id) - for pattern in (f'seq::*::{record.training_id}::*', f'tinker_seq::{record.training_id}::*', - f'seq_ledger::{record.training_id}::*'): - for key in await self.backend.keys(pattern): - await self.backend.delete(key) async def save_instance(self, record: InstanceRecord) -> None: ttl = max(1, int(self.state.expiration_timeout)) if record.state == 'STOPPED' else None @@ -56,19 +55,12 @@ async def instances(self) -> list[InstanceRecord]: return [InstanceRecord.model_validate(value) for value in await self.backend.mget(keys) if value is not None] async def trainings(self) -> list[TrainingRecord]: - keys = await self.backend.keys('model::*') - records = await self.backend.mget(keys) - return [ - TrainingRecord.model_validate(value) for value in records if isinstance(value, dict) - and value.get('controller_epoch') and value.get('training_namespace') == self.prefix - ] + return await self.state.get_training_records(self.prefix) async def get_create(self, identity: str) -> dict[str, str] | None: return await self.backend.get(self.key('create', identity)) async def save_create(self, identity: str, fingerprint: str, request_id: str, session_id: str) -> None: - # Heartbeats can keep a session alive longer than expiration_timeout. - # Keep its creation receipts until session cleanup, never recreate on retry. await self.backend.set( self.key('create', identity), { 'fingerprint': fingerprint, @@ -95,7 +87,7 @@ async def future(self, request_id, status, model_id, result=result, failure=failure, replica_id=epoch, absolute_deadline=deadline) async def own_future(self, request_id: str, token: str) -> None: - await self.backend.set(f'future_owner::{request_id}', token, ttl=max(1, int(self.state.expiration_timeout))) + await self.state.set_future_owner(request_id, token) async def begin_epoch(self, epoch: str) -> None: await self.backend.set(self.key('controller', 'active'), {'epoch': epoch, 'started_at': time.time()}) diff --git a/src/twinkle/server/task_queue/__init__.py b/src/twinkle/server/task_queue/__init__.py index bf5391611..6d7dd81be 100644 --- a/src/twinkle/server/task_queue/__init__.py +++ b/src/twinkle/server/task_queue/__init__.py @@ -5,11 +5,9 @@ Public exports: - TaskStatus - task lifecycle enum - QueueState - queue state enum for tinker client compatibility -- TaskQueueConfig - queue and rate-limit configuration dataclass - TaskQueueMixin - mixin with schedule_task / schedule_background_task - RateLimiter - sliding-window rate limiter """ -from .config import TaskQueueConfig from .mixin import TaskQueueMixin from .rate_limiter import RateLimiter from .types import QueuedTask, QueueState, TaskStatus, UserTaskError @@ -20,7 +18,6 @@ 'QueueState', 'UserTaskError', 'QueuedTask', - 'TaskQueueConfig', 'TaskQueueMixin', 'ComputeWorker', 'RateLimiter', diff --git a/src/twinkle/server/task_queue/mixin.py b/src/twinkle/server/task_queue/mixin.py index be24ac9d9..ff3da6e79 100644 --- a/src/twinkle/server/task_queue/mixin.py +++ b/src/twinkle/server/task_queue/mixin.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Any from twinkle.protocol.types.lifecycle import TERMINAL_STATUSES, TaskEnvelope +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.exceptions import BatchSizeError, ConfigError, InputTokensExceededError, RateLimitExceededError from twinkle.server.lifecycle.envelope import envelope_from_record from twinkle.server.lifecycle.poll_config import long_poll_window @@ -22,7 +23,6 @@ from twinkle.server.telemetry.metrics import get_task_metrics from twinkle.utils.logger import get_logger from .backend_gate import BackendGate -from .config import TaskQueueConfig from .rate_limiter import RateLimiter from .types import QueuedTask, QueueState, TaskStatus from .worker import ComputeWorker @@ -273,19 +273,22 @@ async def _enqueue_prepared_task(self, task: QueuedTask) -> None: q = self._compute_worker.get_queue(queue_key) self._compute_worker.track_task(request_id) await q.put(task) - if task.on_enqueued is not None: - await task.on_enqueued() - await self.state.store_future_status( - request_id, - TaskStatus.QUEUED.value, - model_id, - queue_state=QueueState.ACTIVE.value, - ) - logger.info(f'[TaskQueue] Task {request_id} queued, type={task.task_type or "unknown"}, ' - f'model_id={model_id}, queue_key={queue_key}, ' - f'queue_depth={q.qsize()}, input_tokens={task.input_tokens}') - - self._compute_worker.notify_new_task() + try: + if task.on_enqueued is not None: + await task.on_enqueued() + await self.state.store_future_status( + request_id, + TaskStatus.QUEUED.value, + model_id, + queue_state=QueueState.ACTIVE.value, + ) + logger.info(f'[TaskQueue] Task {request_id} queued, type={task.task_type or "unknown"}, ' + f'model_id={model_id}, queue_key={queue_key}, ' + f'queue_depth={q.qsize()}, input_tokens={task.input_tokens}') + finally: + # Once q.put succeeds the sleeping worker must always be woken, even + # if an acknowledgement or status write fails afterwards. + self._compute_worker.notify_new_task() async def schedule_task( self, diff --git a/src/twinkle/server/task_queue/worker.py b/src/twinkle/server/task_queue/worker.py index 91edc50c8..6dbc4ee89 100644 --- a/src/twinkle/server/task_queue/worker.py +++ b/src/twinkle/server/task_queue/worker.py @@ -16,6 +16,7 @@ from typing import TYPE_CHECKING, Any, Callable, Deque from twinkle.protocol.types.errors import ErrorCategory +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.exceptions import (BatchSizeError, EndpointUnavailableError, FullModeBusyError, InputTokensExceededError, RateLimitExceededError, RequestRejectedError, ResourceExhaustedError, ResourceNotFoundError, ResourceQuotaExceededError, @@ -25,7 +26,6 @@ from twinkle.server.telemetry.correlation import MODEL_ID, TOKEN_ID from twinkle.server.telemetry.tracing import traced_operation from twinkle.utils.logger import get_logger -from .config import TaskQueueConfig from .types import BACKEND_TIMEOUT_EXCEPTIONS, BackendBusyError, QueuedTask, QueueState, TaskStatus, UserTaskError if TYPE_CHECKING: diff --git a/src/twinkle/server/training/actor.py b/src/twinkle/server/training/actor.py index 4ceb1e392..b26d5a88e 100644 --- a/src/twinkle/server/training/actor.py +++ b/src/twinkle/server/training/actor.py @@ -9,11 +9,10 @@ from twinkle.server.config.training import TrainingConfig from twinkle.server.state import get_server_state from twinkle.server.state.models import FutureFailureRecord +from twinkle.server.state.training_records import CreateTraining from twinkle.utils.logger import get_logger from .controller import TrainingController from .driver import RayInstanceDriver -from .records import CreateTraining -from .store import TrainingStore logger = get_logger() @@ -28,7 +27,7 @@ class TrainingControllerActor: def __init__(self, config, namespace): self.config = TrainingConfig.model_validate(config) self.namespace = namespace - self.store = TrainingStore(get_server_state(), namespace) + self.store = get_server_state().create_training_store(namespace) self.driver = RayInstanceDriver(self.config, self.store, namespace, controller_name(namespace)) self.core = TrainingController(self.config, self.store, self.driver) self._monitors = [] @@ -39,9 +38,19 @@ async def initialize(self): await asyncio.shield(self._initialization) async def _initialize(self): - await self.core.initialize() + try: + await self.core.initialize() + except Exception: + logger.exception('Training controller initialization failed; restart is deferred') + # First-time startup is rolled back by the launcher; auto-restart waits before retrying. + self._restart_task = asyncio.create_task(self._restart_after_failure()) + raise self._monitors = [asyncio.create_task(self._heartbeat()), asyncio.create_task(self._sweep())] + async def _restart_after_failure(self): + await asyncio.sleep(max(5.0, self.config.sweep_interval_s)) + ray.kill(ray.get_actor(controller_name(self.namespace), namespace=self.namespace), no_restart=False) + async def _heartbeat(self): while True: try: @@ -150,6 +159,10 @@ async def diagnostics(self, base_model=None): } async def shutdown(self): + restart = getattr(self, '_restart_task', None) + if restart is not None: + restart.cancel() + await asyncio.gather(restart, return_exceptions=True) for task in self._monitors: task.cancel() await asyncio.gather(*self._monitors, return_exceptions=True) diff --git a/src/twinkle/server/training/controller.py b/src/twinkle/server/training/controller.py index 2b8dee8a8..15cecd414 100644 --- a/src/twinkle/server/training/controller.py +++ b/src/twinkle/server/training/controller.py @@ -12,19 +12,25 @@ from typing import Any, Protocol from twinkle.server.config.training import TrainingConfig -from twinkle.server.exceptions import RequestRejectedError, ResourceNotFoundError, ResourceQuotaExceededError +from twinkle.server.exceptions import (RequestRejectedError, ResourceExhaustedError, ResourceNotFoundError, + ResourceQuotaExceededError) +from twinkle.server.resources.allocation import CapacityUnavailableError from twinkle.server.state.models import FutureFailureRecord +from twinkle.server.state.training_records import (LIVE_INSTANCES, TERMINAL_TRAINING, CreateTraining, InstanceRecord, + TrainingRecord) +from twinkle.server.state.training_store import TrainingStore from .catalog import TrainingCatalog, TrainingSpec -from .errors import CapacityUnavailableError, TrainingExpiredError -from .records import LIVE_INSTANCES, TERMINAL_TRAINING, CreateTraining, InstanceRecord, TrainingRecord -from .store import TrainingStore +from .errors import TrainingExpiredError logger = logging.getLogger(__name__) class InstanceDriver(Protocol): - async def start(self, instance: InstanceRecord, spec: TrainingSpec) -> dict[str, Any]: + async def allocate(self, instance: InstanceRecord, spec: TrainingSpec) -> dict[str, Any]: + ... + + async def deploy(self, instance: InstanceRecord, spec: TrainingSpec) -> None: ... async def create_training(self, instance: InstanceRecord, training: TrainingRecord, @@ -83,6 +89,7 @@ def __init__(self, self._cleaning: set[str] = set() self._closing: set[str] = set() self._reconciling: set[str] = set() + self._failed_instances: dict[str, FutureFailureRecord] = {} self._lock = asyncio.Lock() self._ready = False @@ -92,8 +99,11 @@ def _spawn(self, coroutine) -> asyncio.Task: def done(completed): self._jobs.discard(completed) - if not completed.cancelled(): - completed.exception() + if completed.cancelled(): + return + exc = completed.exception() + if exc is not None: + logger.error('Training background job failed', exc_info=(type(exc), exc, exc.__traceback__)) task.add_done_callback(done) return task @@ -145,8 +155,8 @@ def residents_ready(self) -> bool: def residents_ready_for(self, spec: TrainingSpec) -> bool: return sum( - instance.resident and (instance.base_model, instance.mode) == spec.key and instance.state in ('READY', - 'RESETTING') + instance.resident and instance.instance_id not in self._failed_instances and ( + instance.base_model, instance.mode) == spec.key and instance.state in ('READY', 'RESETTING') for instance in self.instances.values()) >= spec.config.min_replicas async def register_runtime(self, instance_id: str, epoch: str, incarnation: str) -> None: @@ -190,19 +200,33 @@ async def _new_instance(self, spec: TrainingSpec, *, resident: bool = False) -> async def _start_instance(self, instance: InstanceRecord) -> None: spec = self.catalog.get(instance.base_model, instance.mode) try: - details = await asyncio.wait_for(self.driver.start(instance, spec), self.config.startup_timeout_s) + + async def startup(): + details = await self.driver.allocate(instance, spec) + async with self._lock: + await self._assert_epoch() + if instance.state != 'ALLOCATING': + raise RuntimeError('Instance became unavailable during allocation') + for name in details: + if name not in ('application_name', 'placement_group_id', 'node_id', 'devices'): + raise ValueError(f'Unknown runtime allocation field: {name}') + allocated = instance.model_copy(deep=True, update=details) + await self.store.save_instance(allocated) + for name, value in details.items(): + setattr(instance, name, value) + # Ownership is persisted before any worker is constructed; the driver does not write records. + await self.driver.deploy(instance.model_copy(deep=True), spec) + + await asyncio.wait_for(startup(), self.config.startup_timeout_s) async with self._lock: await self._assert_epoch() if self.clock() >= instance.startup_deadline: raise TimeoutError('Training runtime became ready after its startup deadline') if instance.state not in ('ALLOCATING', 'STARTING'): raise RuntimeError('Instance became unavailable during startup') - for name, value in details.items(): - if name not in ('application_name', 'placement_group_id', 'node_id', 'devices'): - raise ValueError(f'Unknown runtime allocation field: {name}') - setattr(instance, name, value) + ready = instance.model_copy(update={'state': 'READY'}) + await self.store.save_instance(ready) instance.state = 'READY' - await self.store.save_instance(instance) except Exception as exc: await self.fail_instance(instance.instance_id, failure_for(exc)) raise @@ -210,8 +234,8 @@ async def _start_instance(self, instance: InstanceRecord) -> None: def _select(self, spec: TrainingSpec) -> InstanceRecord | None: candidates = [ instance for instance in self.instances.values() - if (instance.base_model, instance.mode) == spec.key and instance.state in ( - 'READY', 'ALLOCATING', 'STARTING') and len(instance.training_ids) < spec.capacity + if instance.instance_id not in self._failed_instances and (instance.base_model, instance.mode) == spec.key + and instance.state in ('READY', 'ALLOCATING', 'STARTING') and len(instance.training_ids) < spec.capacity ] return min( candidates, @@ -225,7 +249,7 @@ async def create(self, request: CreateTraining) -> dict[str, str]: async with self._lock: await self._assert_epoch() if not self._ready: - raise RequestRejectedError('Training service is not ready') + raise ResourceExhaustedError('Training service is not ready') previous = await self.store.get_create(identity) if previous is not None: if previous['fingerprint'] != fingerprint: @@ -338,9 +362,12 @@ async def _create_user(self, training: TrainingRecord, request: CreateTraining) await self._cleanup_user(training) async def _expired(self, training: TrainingRecord) -> bool: + heartbeat = await self.store.state.get_session_last_heartbeat(training.session_id) + return self._expired_with_heartbeat(training, heartbeat) + + def _expired_with_heartbeat(self, training: TrainingRecord, heartbeat: float | None) -> bool: spec = self.catalog.get(training.base_model, training.mode) config, legacy = spec.config, spec.model.legacy_expiration - heartbeat = await self.store.state.get_session_last_heartbeat(training.session_id) if heartbeat is None: return True if legacy is not None: @@ -361,7 +388,7 @@ async def _authorized(self, training_id: str, token: str, session_id: str | None raise ResourceNotFoundError('Unknown training handle') if training.state == 'EXPIRED': raise TrainingExpiredError('Training handle expired; create a new handle and explicitly load a checkpoint.') - if training.state != 'ACTIVE': + if training.state != 'ACTIVE' or training.instance_id in self._failed_instances: raise ResourceNotFoundError('Training handle is no longer active') if await self.store.state.get_session_last_heartbeat(training.session_id) is None: raise TrainingExpiredError('Training session expired; create a new training handle') @@ -463,60 +490,82 @@ async def release(self, training_id: str, token: str, session_id: str | None = N async def _cleanup_user(self, training: TrainingRecord) -> None: instance = self.instances[training.instance_id] - # Each training must have exactly one cleanup, including concurrent finish callbacks. - async with self._lock: - if (instance.state not in ('READY', 'RESETTING') or training.state != 'CLOSING' - or training.pending_operation_ids or training.training_id in self._cleaning): - return - self._cleaning.add(training.training_id) - resetting = instance.resident and instance.mode == 'full' and not instance.dirty - if resetting: - instance.state = 'RESETTING' - await self.store.save_instance(instance) - closing_full = instance.mode == 'full' and (not instance.resident or instance.dirty) - if closing_full: - instance.state = 'DRAINING' - await self.store.save_instance(instance) - if closing_full: - await self._close_instance(instance) - self._cleaning.discard(training.training_id) - return + claimed = False + resetting = False + closing_full = False try: - spec = self.catalog.get(training.base_model, training.mode) - await asyncio.wait_for( - self.driver.release_training(instance, training), - spec.config.queue_config.absolute_future_ttl(spec.config.num_devices)) - except Exception as exc: - # A timed-out create/cleanup RPC does not prove that shared workers failed. - # Keep this slot reserved until serial cleanup can be confirmed. - if instance.mode == 'lora' and any( - self.trainings[identity].state in ('ACTIVE', - 'PENDING') or self.trainings[identity].pending_operation_ids - for identity in instance.training_ids if identity != training.training_id): - logger.warning('Training cleanup remains pending for %s: %s', training.training_id, exc) - else: - await self.fail_instance(instance.instance_id, failure_for(exc)) - self._cleaning.discard(training.training_id) - return - async with self._lock: - training.state = { - 'expired': 'EXPIRED', - 'released': 'RELEASED', - 'cancelled': 'CANCELLED' - }.get(training.terminal_reason, 'FAILED') - training.terminal_at = time.time() - await self.store.retire(training) - self._cleaning.discard(training.training_id) - if training.training_id in instance.training_ids: - instance.training_ids.remove(training.training_id) - close = not instance.training_ids and (not instance.resident or instance.dirty) + # Each training must have exactly one cleanup, including concurrent finish callbacks. + async with self._lock: + if (instance.state not in ('READY', 'RESETTING') or training.state != 'CLOSING' + or training.pending_operation_ids or training.training_id in self._cleaning): + return + self._cleaning.add(training.training_id) + claimed = True + resetting = instance.resident and instance.mode == 'full' and not instance.dirty + closing_full = instance.mode == 'full' and (not instance.resident or instance.dirty) + if resetting: + updated_instance = instance.model_copy(update={'state': 'RESETTING'}) + await self.store.save_instance(updated_instance) + instance.state = 'RESETTING' + if closing_full: + updated_instance = instance.model_copy(update={'state': 'DRAINING'}) + await self.store.save_instance(updated_instance) + instance.state = 'DRAINING' + if closing_full: + await self._close_instance(instance) + return + try: + spec = self.catalog.get(training.base_model, training.mode) + await asyncio.wait_for( + self.driver.release_training(instance, training), + spec.config.queue_config.absolute_future_ttl(spec.config.num_devices)) + except Exception as exc: + # A timed-out create/cleanup RPC does not prove that shared workers failed. + # Keep this slot reserved until serial cleanup can be confirmed. + if instance.mode == 'lora' and any( + self.trainings[identity].state in ('ACTIVE', + 'PENDING') or self.trainings[identity].pending_operation_ids + for identity in instance.training_ids if identity != training.training_id): + logger.warning('Training cleanup remains pending for %s: %s', training.training_id, exc) + else: + await self.fail_instance( + instance.instance_id, + FutureFailureRecord( + reason_code='internal_error', + attribution='server', + message=f'Training cleanup failed: {exc}'[:1024], + diagnostic=traceback.format_exc())) + return + async with self._lock: + terminal = training.model_copy( + deep=True, + update={ + 'state': { + 'expired': 'EXPIRED', + 'released': 'RELEASED', + 'cancelled': 'CANCELLED' + }.get(training.terminal_reason, 'FAILED'), + 'terminal_at': time.time() + }) + await self.store.retire(terminal) + remaining = [identity for identity in instance.training_ids if identity != training.training_id] + close = not remaining and (not instance.resident or instance.dirty) + updated_instance = instance.model_copy( + deep=True, + update={ + 'training_ids': remaining, + 'state': 'DRAINING' if close else ('READY' if resetting else instance.state) + }) + await self.store.save_instance(updated_instance) + training.state = terminal.state + training.terminal_at = terminal.terminal_at + instance.training_ids = remaining + instance.state = updated_instance.state if close: - instance.state = 'DRAINING' - elif resetting: - instance.state = 'READY' - await self.store.save_instance(instance) - if close: - await self._close_instance(instance) + await self._close_instance(instance) + finally: + if claimed: + self._cleaning.discard(training.training_id) async def _close_instance(self, instance: InstanceRecord) -> None: if instance.instance_id in self._closing: @@ -527,60 +576,105 @@ async def _close_instance(self, instance: InstanceRecord) -> None: await asyncio.wait_for(self.driver.shutdown(instance), self.config.shutdown_timeout_s) except Exception as exc: async with self._lock: + quarantined = instance.model_copy(update={'state': 'QUARANTINED', 'failure': str(exc)}) + await self.store.save_instance(quarantined) instance.state = 'QUARANTINED' instance.failure = str(exc) - await self.store.save_instance(instance) return async with self._lock: - instance.state = 'STOPPED' - instance.stopped_at = time.time() - await self.store.save_instance(instance) + terminal_records = [] + terminal_at = time.time() + failure = self._failed_instances.get(instance.instance_id) for training_id in list(instance.training_ids): training = self.trainings[training_id] + if failure is not None: + for request_id in (training.create_request_id, *training.pending_operation_ids): + await self.store.future(request_id, 'failed', training_id, failure=failure) + terminal = training.model_copy( + deep=True, + update={ + 'pending_operation_ids': {}, + 'state': { + 'expired': 'EXPIRED', + 'released': 'RELEASED', + 'cancelled': 'CANCELLED' + }.get(training.terminal_reason, 'FAILED'), + 'terminal_at': terminal_at + }) + await self.store.retire(terminal) + terminal_records.append((training, terminal)) + stopped = instance.model_copy( + deep=True, update={ + 'state': 'STOPPED', + 'stopped_at': terminal_at, + 'training_ids': [] + }) + await self.store.save_instance(stopped) + for training, terminal in terminal_records: training.pending_operation_ids.clear() - training.state = { - 'expired': 'EXPIRED', - 'released': 'RELEASED', - 'cancelled': 'CANCELLED' - }.get(training.terminal_reason, 'FAILED') - training.terminal_at = time.time() - await self.store.retire(training) + training.state = terminal.state + training.terminal_at = terminal.terminal_at + instance.state = 'STOPPED' + instance.stopped_at = terminal_at instance.training_ids.clear() - await self.store.save_instance(instance) + self._failed_instances.pop(instance.instance_id, None) finally: self._closing.discard(instance.instance_id) async def fail_instance(self, instance_id: str, failure: FutureFailureRecord) -> None: async with self._lock: instance = self.instances[instance_id] - if instance.state in ('DRAINING', 'STOPPED', 'QUARANTINED'): + if instance.state == 'STOPPED': + self._failed_instances.pop(instance_id, None) return - instance.state = 'DRAINING' - instance.failure = failure.message - await self.store.save_instance(instance) - for training_id in instance.training_ids: - training = self.trainings[training_id] - training.state = 'CLOSING' - training.terminal_reason = failure.reason_code - for request_id in (training.create_request_id, *training.pending_operation_ids): - await self.store.future(request_id, 'failed', training_id, failure=failure) - await self.store.save_training(training) - startup = self._startup.get(instance_id) - if startup is not None and startup is not asyncio.current_task() and not startup.done(): - startup.cancel() - await asyncio.gather(startup, return_exceptions=True) - await self._close_instance(instance) + # Record the retry intent without promoting unpersisted state; block new routing. + self._failed_instances[instance_id] = failure + if instance.state not in ('DRAINING', 'QUARANTINED'): + draining = instance.model_copy(update={'state': 'DRAINING', 'failure': failure.message}) + await self.store.save_instance(draining) + instance.state = 'DRAINING' + instance.failure = failure.message + try: + async with self._lock: + for training_id in instance.training_ids: + training = self.trainings[training_id] + closing = training.model_copy( + deep=True, update={ + 'state': 'CLOSING', + 'terminal_reason': failure.reason_code + }) + await self.store.save_training(closing) + training.state = 'CLOSING' + training.terminal_reason = failure.reason_code + for request_id in (training.create_request_id, *training.pending_operation_ids): + await self.store.future(request_id, 'failed', training_id, failure=failure) + startup = self._startup.get(instance_id) + if startup is not None and startup is not asyncio.current_task() and not startup.done(): + startup.cancel() + await asyncio.gather(startup, return_exceptions=True) + finally: + # Once DRAINING is durable, shutdown must be attempted even if a + # later training/future write fails. Sweep will retry if this fails. + await self._close_instance(instance) async def sweep(self) -> None: + # Batch I/O outside the decision lock; re-verify state after re-acquiring, ignore handles added since snapshot. + session_ids = list({ + training.session_id + for training in self.trainings.values() if training.state == 'ACTIVE' and not training.pending_operation_ids + }) + heartbeats = await self.store.state.get_session_last_heartbeats(session_ids) async with self._lock: await self._assert_epoch() for training in self.trainings.values(): - if training.state != 'ACTIVE' or training.pending_operation_ids: + if (training.state != 'ACTIVE' or training.pending_operation_ids + or training.session_id not in heartbeats): continue - if await self._expired(training): + if self._expired_with_heartbeat(training, heartbeats[training.session_id]): + closing_record = training.model_copy(update={'state': 'CLOSING', 'terminal_reason': 'expired'}) + await self.store.save_training(closing_record) training.state = 'CLOSING' training.terminal_reason = 'expired' - await self.store.save_training(training) uncertain = [(training, request_id) for training in self.trainings.values() for request_id, ticket in training.pending_operation_ids.items() if request_id not in self._reconciling and self.clock() @@ -595,13 +689,16 @@ async def sweep(self) -> None: for training, request_id in uncertain: self._spawn(self._reconcile_operation(training, request_id)) for instance in list(self.instances.values()): - if instance.state == 'QUARANTINED': + failure = self._failed_instances.get(instance.instance_id) + if failure is not None: + self._spawn(self.fail_instance(instance.instance_id, failure)) + elif instance.state in ('DRAINING', 'QUARANTINED'): self._spawn(self._close_instance(instance)) # Rebuilding one specification must not stop expiry or other models. for spec in self.catalog.specs(): residents = sum( - instance.resident and (instance.base_model, - instance.mode) == spec.key and instance.state in LIVE_INSTANCES + instance.resident and instance.instance_id not in self._failed_instances and ( + instance.base_model, instance.mode) == spec.key and instance.state in LIVE_INSTANCES for instance in self.instances.values()) for _ in range(spec.config.min_replicas - residents): async with self._lock: diff --git a/src/twinkle/server/training/driver.py b/src/twinkle/server/training/driver.py index cab2139c6..8a0a9b752 100644 --- a/src/twinkle/server/training/driver.py +++ b/src/twinkle/server/training/driver.py @@ -6,9 +6,12 @@ import os import re -from twinkle.server.launcher.env_propagation import build_propagated_env_vars -from .allocation import DevicePool -from .records import InstanceRecord +from twinkle.server.env_propagation import build_propagated_env_vars +from twinkle.server.resources.allocation import DevicePool +from twinkle.server.state.training_records import InstanceRecord +from twinkle.utils.logger import get_logger + +logger = get_logger() class RayInstanceDriver: @@ -16,25 +19,30 @@ class RayInstanceDriver: def __init__(self, config, store, namespace, controller_name): self.config, self.store = config, store self.namespace, self.controller_name = namespace, controller_name + self._cleanup_tasks: set[asyncio.Task] = set() self.pool = DevicePool( namespace, allocation_timeout=config.allocation_timeout_s, shutdown_timeout=config.shutdown_timeout_s) def application_name(self, instance): return self.pool.name(instance.controller_epoch, instance.instance_id) - async def start(self, instance, spec): + async def allocate(self, instance, spec): + allocation = await self.pool.allocate(spec.config.num_devices, spec.model.device_type, + instance.controller_epoch, instance.instance_id) + return { + 'application_name': allocation.name, + 'placement_group_id': allocation.placement_group_id, + 'node_id': allocation.node_id, + 'devices': list(allocation.devices) + } + + async def deploy(self, instance, spec): from ray import serve from ray.util.placement_group import get_placement_group from twinkle.server.model.app import build_model_app - allocation = await self.pool.allocate(spec, instance.controller_epoch, instance.instance_id) name = self.application_name(instance) - # Persist ownership before the first worker can be constructed. - instance.application_name = name - instance.placement_group_id = allocation.placement_group_id - instance.node_id, instance.devices = allocation.node_id, list(allocation.devices) - await self.store.save_instance(instance) env = build_propagated_env_vars() env.update(spec.model.env_vars) for key in ('PYTHONPATH', 'TWINKLE_DEFAULT_SAVE_DIR', 'TMPDIR', 'OMP_NUM_THREADS'): @@ -78,8 +86,8 @@ async def start(self, instance, spec): queue_config=spec.config.queue_config, data_plane_url=spec.model.data_plane_url, training_context=context, - placement_group=get_placement_group(allocation.name), - allocation_devices=list(allocation.devices), + placement_group=get_placement_group(name), + allocation_devices=list(instance.devices), **kwargs) # Shield the actual thread so its late return can still be cleaned. deployment = asyncio.create_task( @@ -96,14 +104,18 @@ async def clean_late(): pass await self.shutdown(instance) - asyncio.create_task(clean_late()) + task = asyncio.create_task(clean_late()) + self._cleanup_tasks.add(task) + + def done(completed): + self._cleanup_tasks.discard(completed) + if not completed.cancelled(): + exc = completed.exception() + if exc is not None: + logger.error('Late deployment cleanup failed', exc_info=(type(exc), exc, exc.__traceback__)) + + task.add_done_callback(done) raise - return { - 'application_name': name, - 'placement_group_id': allocation.placement_group_id, - 'node_id': allocation.node_id, - 'devices': list(allocation.devices) - } def handle(self, instance): from ray import serve @@ -146,7 +158,7 @@ async def shutdown(self, instance: InstanceRecord): except ray.exceptions.RayActorError: continue raise RuntimeError('Owned worker death could not be confirmed') - await self.store.state.unregister_replica(instance.instance_id) + await self.store.state.unregister_replica_capacity(instance.instance_id) async def health(self, instance): return await asyncio.wait_for(self.handle(instance).runtime_identity.remote(), 5) @@ -155,7 +167,7 @@ async def recover_orphans(self): """Ownership names survive loss of the memory state actor.""" from ray import serve from ray.util.placement_group import placement_group_table - pattern = re.compile(r'training-' + re.escape(self.namespace) + pattern = re.compile(r'resources-' + re.escape(self.namespace) + r'-(controller_[0-9a-f]{32})-(instance_[0-9a-f]{32})$') names = set(serve.status().applications) names.update(row.get('name', '') for row in placement_group_table().values() if row.get('state') != 'REMOVED') diff --git a/src/twinkle/server/training/errors.py b/src/twinkle/server/training/errors.py index 4d7cfc128..d0d508bc6 100644 --- a/src/twinkle/server/training/errors.py +++ b/src/twinkle/server/training/errors.py @@ -2,10 +2,6 @@ from twinkle.server.exceptions import RequestRejectedError, ResourceNotFoundError -class CapacityUnavailableError(RequestRejectedError): - error_code = 409 - - class TrainingExpiredError(ResourceNotFoundError): pass diff --git a/src/twinkle/server/training/router.py b/src/twinkle/server/training/router.py index fd20c8cc6..125336da5 100644 --- a/src/twinkle/server/training/router.py +++ b/src/twinkle/server/training/router.py @@ -17,17 +17,16 @@ from twinkle.server.checkpoint import _resolve_client_save_dir, create_checkpoint_manager, validate_user_path from twinkle.server.deployment import LazyCleanupMixin, bind_deployment, build_deployment_app from twinkle.server.exceptions import RequestRejectedError, ResourceNotFoundError -from twinkle.server.gateway.proxy import ServiceProxy from twinkle.server.lifecycle.envelope import envelope_from_record from twinkle.server.middleware.auth import get_session_id_from_request, get_token_from_request +from twinkle.server.service_proxy import ServiceProxy from twinkle.server.state import get_server_state +from twinkle.server.state.training_records import CreateTraining from twinkle.server.task_queue import TaskQueueMixin from twinkle.server.validation import assert_request_supported from .actor import controller_name from .catalog import TrainingCatalog from .errors import ProtocolMismatchError -from .records import CreateTraining -from .store import TrainingStore NATIVE_REQUESTS = { 'forward': native.ForwardRequest, @@ -80,7 +79,7 @@ async def __init__(self, config, model, namespace, http_options): self.model = model self.backend = model.backend self.state = get_server_state() - self.store = TrainingStore(self.state, namespace) + self.store = self.state.create_training_store(namespace) self.controller = ray.get_actor(controller_name(namespace), namespace=namespace) self.proxy = ServiceProxy(http_options) self.replica_id = 'training-router_' + uuid.uuid4().hex diff --git a/tests/server/gateway/test_proxy.py b/tests/server/gateway/test_proxy.py index 9b03e4a98..a111b5096 100644 --- a/tests/server/gateway/test_proxy.py +++ b/tests/server/gateway/test_proxy.py @@ -13,7 +13,7 @@ from starlette.requests import Request from unittest.mock import AsyncMock -from twinkle.server.gateway.proxy import ServiceProxy +from twinkle.server.service_proxy import ServiceProxy from twinkle.protocol.headers import H_MULTIPLEX, H_MULTIPLEX_LEGACY, H_REQUEST_ID diff --git a/tests/server/lifecycle/test_preflight_rejection.py b/tests/server/lifecycle/test_preflight_rejection.py index 74639c137..eb997cd8a 100644 --- a/tests/server/lifecycle/test_preflight_rejection.py +++ b/tests/server/lifecycle/test_preflight_rejection.py @@ -14,7 +14,7 @@ from twinkle.server.deployment import twinkle_server_error_handler from twinkle.server.exceptions import (BatchSizeError, InputTokensExceededError, RateLimitExceededError, RequestRejectedError, TwinkleServerError) -from twinkle.server.task_queue.config import TaskQueueConfig +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.task_queue.mixin import TaskQueueMixin diff --git a/tests/server/lifecycle/test_submit_peek_e2e.py b/tests/server/lifecycle/test_submit_peek_e2e.py index c87aba599..b3de7aa6c 100644 --- a/tests/server/lifecycle/test_submit_peek_e2e.py +++ b/tests/server/lifecycle/test_submit_peek_e2e.py @@ -13,7 +13,7 @@ ray = pytest.importorskip('ray') from twinkle.server.state import ServerState # noqa: E402 -from twinkle.server.task_queue.config import TaskQueueConfig # noqa: E402 +from twinkle.server.config.task_queue import TaskQueueConfig # noqa: E402 from twinkle.server.task_queue.mixin import TaskQueueMixin # noqa: E402 from twinkle_client import _future # noqa: E402 from twinkle_client.exceptions import TaskFailedError # noqa: E402 diff --git a/tests/server/lifecycle/test_timing_bounds.py b/tests/server/lifecycle/test_timing_bounds.py index b2f5f5464..7664b5bb4 100644 --- a/tests/server/lifecycle/test_timing_bounds.py +++ b/tests/server/lifecycle/test_timing_bounds.py @@ -21,7 +21,7 @@ from twinkle.server.lifecycle.poll_config import long_poll_window, retrieve_poll_interval # noqa: E402 from twinkle.server.state import ServerState # noqa: E402 -from twinkle.server.task_queue.config import TaskQueueConfig # noqa: E402 +from twinkle.server.config.task_queue import TaskQueueConfig # noqa: E402 from twinkle.server.task_queue.mixin import TaskQueueMixin # noqa: E402 _WINDOW = 0.05 diff --git a/tests/server/model/test_tinker_handlers.py b/tests/server/model/test_tinker_handlers.py index 6daa6d566..9d116ed51 100644 --- a/tests/server/model/test_tinker_handlers.py +++ b/tests/server/model/test_tinker_handlers.py @@ -9,7 +9,7 @@ from twinkle.server.model.tinker_handlers import _register_model_tinker_routes from twinkle.server.deployment import twinkle_server_error_handler from twinkle.server.exceptions import BatchSizeError, TwinkleServerError -from twinkle.server.task_queue.config import TaskQueueConfig +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.task_queue.mixin import TaskQueueMixin diff --git a/tests/server/sampler/test_sampling_concurrency.py b/tests/server/sampler/test_sampling_concurrency.py index f1697ec9c..cf57479b0 100644 --- a/tests/server/sampler/test_sampling_concurrency.py +++ b/tests/server/sampler/test_sampling_concurrency.py @@ -21,7 +21,7 @@ from twinkle.server.sampler.tinker_handlers import _register_tinker_sampler_routes from twinkle.server.sampler.twinkle_handlers import _register_twinkle_sampler_routes from twinkle.server.task_queue.backend_gate import BackendGate -from twinkle.server.task_queue.config import TaskQueueConfig +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.task_queue.types import BackendBusyError from twinkle.server.task_queue.worker import ComputeWorker diff --git a/tests/server/test_dependencies.py b/tests/server/test_dependencies.py new file mode 100644 index 000000000..3100a71b4 --- /dev/null +++ b/tests/server/test_dependencies.py @@ -0,0 +1,59 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Executable dependency architecture for every ``twinkle.server`` component.""" +import grimp + +_ROOT = 'twinkle.server' +_ROOT_COMPONENT = '__root__' + +# Dependencies must point from a later layer to an earlier layer. Components in +# the same layer are independent and therefore may not import one another. +_LAYERS = ( + frozenset({'exceptions'}), + frozenset({'config', 'middleware', 'resources', 'runtime', 'task_errors', 'utils', 'validation'}), + frozenset({'telemetry'}), + frozenset({'env_propagation', 'service_proxy', 'state'}), + frozenset({'deployment', 'lifecycle', 'session_resource'}), + frozenset({'task_queue'}), + frozenset({'checkpoint', 'data_plane'}), + frozenset({'gateway', 'model', 'processor'}), + frozenset({'sampler', 'training'}), + frozenset({'launcher'}), + frozenset({'cli'}), + frozenset({'__main__', _ROOT_COMPONENT}), +) + + +def _component(module: str) -> str: + if module == _ROOT: + return _ROOT_COMPONENT + return module.removeprefix(f'{_ROOT}.').split('.', 1)[0] + + +def test_server_components_form_declared_dag(): + """Reject cycles, upward imports, same-layer coupling, and unclassified components.""" + graph = grimp.build_graph('twinkle', include_external_packages=False) + server_modules = { + module for module in graph.modules if module == _ROOT or module.startswith(f'{_ROOT}.') + } + observed_components = {_component(module) for module in server_modules} + declared_components = set().union(*_LAYERS) + assert observed_components == declared_components, ( + f'Update the server DAG when adding or removing components: ' + f'undeclared={sorted(observed_components - declared_components)}, ' + f'missing={sorted(declared_components - observed_components)}') + + layer_by_component = { + component: layer_index for layer_index, layer in enumerate(_LAYERS) for component in layer + } + component_edges = { + (_component(importer), _component(imported)) + for importer in server_modules + for imported in graph.find_modules_directly_imported_by(importer) + if imported in server_modules and _component(importer) != _component(imported) + } + violations = sorted( + (importer, imported) for importer, imported in component_edges + if layer_by_component[importer] <= layer_by_component[imported]) + assert not violations, ( + 'Server dependencies must point strictly down the declared DAG; ' + f'invalid edges={violations}') diff --git a/tests/server/training/test_actor_initialization.py b/tests/server/training/test_actor_initialization.py new file mode 100644 index 000000000..211d925e1 --- /dev/null +++ b/tests/server/training/test_actor_initialization.py @@ -0,0 +1,53 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Initialization failures and backoff during automatic controller restarts.""" +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + +from twinkle.server.training.actor import TrainingControllerActor + + +@pytest.mark.asyncio +async def test_failed_initialization_restarts_only_after_backoff(monkeypatch): + cls = TrainingControllerActor.__ray_metadata__.modified_class + gate, sleeping = asyncio.Event(), asyncio.Event() + delays = [] + handle = object() + kill = Mock() + get_actor = Mock(return_value=handle) + original_sleep = asyncio.sleep + + async def sleep(delay): + delays.append(delay) + sleeping.set() + await gate.wait() + + monkeypatch.setattr('twinkle.server.training.actor.asyncio.sleep', sleep) + monkeypatch.setattr('twinkle.server.training.actor.ray.get_actor', get_actor) + monkeypatch.setattr('twinkle.server.training.actor.ray.kill', kill) + actor = SimpleNamespace(core=SimpleNamespace(initialize=AsyncMock(side_effect=RuntimeError('allocation failed'))), + config=SimpleNamespace(sweep_interval_s=7), namespace='initialization-test', _monitors=[]) + actor._restart_after_failure = lambda: cls._restart_after_failure(actor) + with pytest.raises(RuntimeError, match='allocation failed'): + await cls._initialize(actor) + await sleeping.wait() + assert delays == [7] and not actor._monitors + kill.assert_not_called() + gate.set() + await actor._restart_task + await original_sleep(0) + get_actor.assert_called_once_with('training-controller-initialization-test', namespace='initialization-test') + kill.assert_called_once_with(handle, no_restart=False) + + +@pytest.mark.asyncio +async def test_shutdown_cancels_pending_restart(): + cls = TrainingControllerActor.__ray_metadata__.modified_class + wait = asyncio.Event() + restart = asyncio.create_task(wait.wait()) + actor = SimpleNamespace(core=SimpleNamespace(shutdown=AsyncMock()), _monitors=[], _restart_task=restart) + await cls.shutdown(actor) + assert restart.cancelled() + actor.core.shutdown.assert_awaited_once() diff --git a/tests/server/training/test_controller.py b/tests/server/training/test_controller.py index 11987a10b..358c88e15 100644 --- a/tests/server/training/test_controller.py +++ b/tests/server/training/test_controller.py @@ -11,9 +11,9 @@ from twinkle.server.state import ServerState from twinkle.server.state.backend.memory_backend import RayActorBackend from twinkle.server.training.controller import TrainingController -from twinkle.server.training.errors import CapacityUnavailableError, TrainingExpiredError -from twinkle.server.training.records import CreateTraining -from twinkle.server.training.store import TrainingStore +from twinkle.server.training.errors import TrainingExpiredError +from twinkle.server.resources.allocation import CapacityUnavailableError +from twinkle.server.state.training_records import CreateTraining class Clock: @@ -37,7 +37,7 @@ def __init__(self, devices=16): self.fail_start = False self.fail_shutdown = False - async def start(self, instance, spec): + async def allocate(self, instance, spec): self.starts += 1 if self.fail_start: raise RuntimeError('injected startup failure') @@ -45,9 +45,11 @@ async def start(self, instance, spec): if used + spec.config.num_devices > self.devices: raise CapacityUnavailableError('No devices available') self.allocations[instance.instance_id] = spec.config.num_devices - await self.start_gate.wait() return {} + async def deploy(self, instance, spec): + await self.start_gate.wait() + async def create_training(self, instance, training, request): assert instance.instance_id in self.allocations assert training.training_id not in self.users @@ -252,7 +254,7 @@ async def setup(*, devices=16, min_replicas=0, max_replicas=None, quota=100, bac state = ServerState(backend=backend or RayActorBackend(), per_token_model_limit=quota) await state.create_session({'session_id': 'session'}) runtime, clock = Runtime(devices), Clock() - controller = TrainingController(config, TrainingStore(state, 'unit'), runtime, clock=clock) + controller = TrainingController(config, state.create_training_store('unit'), runtime, clock=clock) await controller.initialize() return controller, state, runtime, clock @@ -453,7 +455,7 @@ async def test_managed_records_are_not_deleted_by_generic_training_ttl(): async def test_restart_invalidates_handles_and_cleans_old_allocation(): controller, state, runtime, clock = await setup() identity = await create(controller, state, 1) - replacement = TrainingController(controller.config, TrainingStore(state, 'unit'), runtime, clock=clock) + replacement = TrainingController(controller.config, state.create_training_store('unit'), runtime, clock=clock) await replacement.initialize() assert replacement.epoch != controller.epoch assert not runtime.allocations and not runtime.users @@ -527,7 +529,7 @@ async def test_restart_fails_pending_operations_and_fences_old_controller(): identity = await create(controller, state, 1) await controller.admit(identity, 'owner', 'session', 'queued') await controller.store.future('queued', 'queued', identity) - replacement = TrainingController(controller.config, TrainingStore(state, 'unit'), runtime, clock=clock) + replacement = TrainingController(controller.config, state.create_training_store('unit'), runtime, clock=clock) await replacement.initialize() assert (await state.get_future('queued'))['failure']['reason_code'] == 'orphaned_replica' with pytest.raises(ResourceNotFoundError, match='epoch'): @@ -604,12 +606,12 @@ async def test_dirty_resident_is_rebuilt_after_last_user_leaves(): async def test_late_startup_result_cannot_publish_success(late_phase): controller, state, runtime, clock = await setup() if late_phase == 'runtime': - original = runtime.start + original = runtime.deploy async def late(instance, spec): result = await original(instance, spec) clock.now = instance.startup_deadline + 1 return result - runtime.start = late + runtime.deploy = late else: original = runtime.create_training async def late(instance, training, request): @@ -629,7 +631,7 @@ async def test_failed_resident_startup_rolls_back_previously_loaded_residents(): controller, state, runtime, clock = await setup(devices=2) raw = controller.config.model_dump() raw['models'][0]['modes']['lora']['min_replicas'] = 2 - replacement = TrainingController(TrainingConfig.model_validate(raw), TrainingStore(state, 'unit'), runtime, clock=clock) + replacement = TrainingController(TrainingConfig.model_validate(raw), state.create_training_store('unit'), runtime, clock=clock) with pytest.raises(CapacityUnavailableError): await replacement.initialize() assert not replacement._ready and not runtime.allocations diff --git a/tests/server/training/test_failure_recovery.py b/tests/server/training/test_failure_recovery.py new file mode 100644 index 000000000..d0b78ed9e --- /dev/null +++ b/tests/server/training/test_failure_recovery.py @@ -0,0 +1,292 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Fault injection ensures persistence and notification failures do not leak resources.""" +import asyncio +import logging +from unittest.mock import AsyncMock + +import pytest + +from tests.server.state.fake_backend import FakeBackend +from tests.server.training.test_controller import create, request, setup +from tests.server.training.test_runtime_queue import RPC, RuntimeQueue, settle +from twinkle.protocol.types.errors import ErrorCategory +from twinkle.server.exceptions import ResourceExhaustedError, ResourceNotFoundError +from twinkle.server.state.models import FAILURE_REASON_CODES, FutureFailureRecord + + +async def wait_idle(queue): + worker = queue._compute_worker + for _ in range(100): + if not worker._request_ids and not worker.new_task_event.is_set(): + await asyncio.sleep(0) + return + await asyncio.sleep(0.001) + raise AssertionError('Worker did not enter idle wait') + + +@pytest.mark.asyncio +@pytest.mark.parametrize('fault', ['commit', 'queued_status']) +async def test_warm_idle_worker_wakes_after_enqueue_failure(fault, monkeypatch): + controller, state, _, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'queue-failure') + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + executed = [] + + async def operation(): + executed.append(True) + return {'ok': True} + + try: + warmup = await queue.schedule_task(operation, model_id=identity, token='owner', task_type='step') + await settle(state, warmup['request_id']) + await wait_idle(queue) + executed.clear() + if fault == 'commit': + async def broken(*args): + raise RuntimeError('injected commit failure') + queue._training_controller.commit_operation = RPC(broken) + ref = await queue.schedule_task(operation, model_id=identity, token='owner', task_type='step') + else: + previous = state.store_future_status + + async def broken(request_id, status, *args, **kwargs): + if status == 'queued': + raise RuntimeError('injected queued write failure') + return await previous(request_id, status, *args, **kwargs) + + monkeypatch.setattr(state, 'store_future_status', broken) + with pytest.raises(RuntimeError, match='queued write'): + await queue.schedule_task(operation, model_id=identity, token='owner', task_type='step', + request_id='queued-failure') + ref = {'request_id': 'queued-failure'} + assert (await settle(state, ref['request_id']))['status'] == 'completed' + await wait_idle(queue) + assert executed == [True] + assert not controller.trainings[identity].pending_operation_ids + assert queue._compute_worker.total_queued() == 0 + finally: + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('fault,mode', [ + ('retire', 'lora'), ('save_draining', 'full'), ('save_stopped', 'full'), + ('close_retire', 'full'), ('post_retire_save', 'lora') +]) +async def test_cleanup_storage_failure_is_retryable(fault, mode, monkeypatch): + controller, state, runtime, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, fault, mode=mode) + training = controller.trainings[identity] + instance = controller.instances[training.instance_id] + previous_retire = controller.store.retire + previous_save = controller.store.save_instance + failed = False + + async def retire(record): + nonlocal failed + if fault in ('retire', 'close_retire') and not failed: + failed = True + raise RuntimeError('injected retirement failure') + return await previous_retire(record) + + async def save(record): + nonlocal failed + matches = ((fault == 'save_draining' and record.state == 'DRAINING') + or (fault == 'save_stopped' and record.state == 'STOPPED') + or (fault == 'post_retire_save' and not record.training_ids)) + if matches and not failed: + failed = True + raise RuntimeError('injected instance write failure') + return await previous_save(record) + + monkeypatch.setattr(controller.store, 'retire', retire) + monkeypatch.setattr(controller.store, 'save_instance', save) + with pytest.raises(RuntimeError, match='injected'): + await controller.release(identity, 'owner', 'session') + assert failed + assert identity not in controller._cleaning + assert instance.instance_id not in controller._closing + assert training.state == 'CLOSING' + assert instance.state != 'STOPPED' + for _ in range(3): + await controller.sweep() + await controller.wait_for_jobs() + if training.state == 'RELEASED': + break + assert training.state == 'RELEASED' + assert not instance.training_ids + assert instance.state == 'STOPPED' + assert not runtime.allocations + assert await state.get_model_metadata(identity) is None + assert (await state._backend.get('token_count::owner') or 0) == 0 + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_failure_transition_write_is_retried_and_admission_fenced(monkeypatch): + controller, state, runtime, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'failure-transition') + training = controller.trainings[identity] + instance = controller.instances[training.instance_id] + await controller.admit(identity, 'owner', 'session', 'pending-op') + await controller.store.future('pending-op', 'pending', identity) + original = controller.store.save_instance + failed = False + + async def save(record): + nonlocal failed + if record.state == 'DRAINING' and not failed: + failed = True + raise RuntimeError('injected draining write failure') + return await original(record) + + monkeypatch.setattr(controller.store, 'save_instance', save) + failure = FutureFailureRecord(reason_code='orphaned_replica', attribution='server', message='Worker died') + with pytest.raises(RuntimeError, match='draining write'): + await controller.fail_instance(instance.instance_id, failure) + assert instance.state == 'READY' + with pytest.raises(ResourceNotFoundError): + await controller.route(identity, 'owner') + await controller.sweep() + await controller.wait_for_jobs() + assert instance.state == 'STOPPED' and training.state == 'FAILED' + assert not runtime.allocations and not training.pending_operation_ids + assert (await state.get_future('pending-op'))['failure']['reason_code'] == 'orphaned_replica' + assert not controller._failed_instances + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_failed_resident_is_excluded_when_rebuilding_minimum_replicas(monkeypatch): + controller, _, runtime, _ = await setup(backend=FakeBackend(), min_replicas=1) + failed_instance = next(iter(controller.instances.values())) + original = controller.store.save_instance + failed = False + + async def save(record): + nonlocal failed + if record.state == 'DRAINING' and not failed: + failed = True + raise RuntimeError('injected draining write failure') + return await original(record) + + monkeypatch.setattr(controller.store, 'save_instance', save) + failure = FutureFailureRecord(reason_code='orphaned_replica', attribution='server', message='Worker died') + with pytest.raises(RuntimeError, match='draining write'): + await controller.fail_instance(failed_instance.instance_id, failure) + + assert failed_instance.state == 'READY' + assert failed_instance.instance_id in controller._failed_instances + await controller.sweep() + await controller.wait_for_jobs() + + ready_residents = [ + instance for instance in controller.instances.values() if instance.resident and instance.state == 'READY' + ] + assert failed_instance.state == 'STOPPED' + assert len(ready_residents) == 1 + assert ready_residents[0].instance_id != failed_instance.instance_id + assert set(runtime.allocations) == {ready_residents[0].instance_id} + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_cleanup_error_is_not_classified_as_startup_failure(): + controller, state, runtime, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'release-failure') + instance = controller.instances[controller.trainings[identity].instance_id] + runtime.release_training = AsyncMock(side_effect=TimeoutError('release timed out')) + await controller.release(identity, 'owner') + assert instance.failure == 'Training cleanup failed: release timed out' + assert controller.trainings[identity].terminal_reason == 'internal_error' + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_unready_service_is_503_server_and_reason_sets_have_no_dead_codes(): + controller, _, _, _ = await setup(backend=FakeBackend()) + controller._ready = False + with pytest.raises(ResourceExhaustedError) as error: + await controller.create(request('unready')) + assert error.value.error_code == 503 and error.value.category is ErrorCategory.Server + assert not {'protocol_mismatch', 'training_expired'} & FAILURE_REASON_CODES + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_spawn_logs_exception_with_traceback(caplog, monkeypatch): + # The ``twinkle`` logger disables propagation, so attach caplog to the controller logger directly. + target = logging.getLogger('twinkle.server.training.controller') + monkeypatch.setattr(target, 'handlers', [*target.handlers, caplog.handler]) + controller, _, _, _ = await setup(backend=FakeBackend()) + + async def broken(): + raise RuntimeError('background sentinel') + + with caplog.at_level(logging.ERROR): + controller._spawn(broken()) + await controller.wait_for_jobs() + records = [record for record in caplog.records if record.message == 'Training background job failed'] + assert records and records[0].exc_info[1].args == ('background sentinel',) + assert not controller._jobs + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_sweep_heartbeat_io_does_not_hold_route_lock(monkeypatch): + controller, state, _, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'slow-heartbeat') + entered, resume = asyncio.Event(), asyncio.Event() + original = state.get_session_last_heartbeats + + async def slow(identities): + entered.set() + await resume.wait() + return await original(identities) + + monkeypatch.setattr(state, 'get_session_last_heartbeats', slow) + sweeping = asyncio.create_task(controller.sweep()) + try: + await entered.wait() + assert await asyncio.wait_for(controller.route(identity, 'owner'), 0.1) + finally: + resume.set() + await sweeping + await controller.wait_for_jobs() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_explicit_tinker_request_id_is_preserved(): + from twinkle.server.model.managed_runtime import _tinker_operation + controller, state, _, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'explicit-id') + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + context = _tinker_operation.set((identity, 10)) + try: + ref = await queue.schedule_task(AsyncMock(return_value='ok'), model_id=identity, token='owner', + task_type='step', request_id='explicit-id') + assert ref['request_id'] == 'explicit-id' + assert (await settle(state, 'explicit-id'))['status'] == 'completed' + finally: + _tinker_operation.reset(context) + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_capacity_unregister_does_not_retire_managed_model(): + controller, state, _, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'capacity-only') + instance = controller.instances[controller.trainings[identity].instance_id] + await state.register_replica(instance.instance_id, 5) + await state.unregister_replica_capacity(instance.instance_id) + assert await state.get_model_metadata(identity) is not None + assert await state._backend.get('token_count::owner') == 1 + await controller.release(identity, 'owner') + assert await state.get_model_metadata(identity) is None + assert await state._backend.get('token_count::owner') == 0 + await controller.shutdown() diff --git a/tests/server/training/test_router.py b/tests/server/training/test_router.py index 08e541ad0..37741d0b7 100644 --- a/tests/server/training/test_router.py +++ b/tests/server/training/test_router.py @@ -11,7 +11,6 @@ from twinkle.hub import HubOperation from twinkle.server.state import ServerState from twinkle.server.training.router import TrainingRouter -from twinkle.server.training.store import TrainingStore from twinkle.server.training.catalog import TrainingCatalog from twinkle.server.config.training import TrainingConfig @@ -29,7 +28,7 @@ async def receive(): async def test_upload_after_recycle_is_cpu_only_and_preserves_future_ownership(monkeypatch, tmp_path): router = TrainingRouter.__new__(TrainingRouter) router.state = ServerState(backend=FakeBackend()) - router.store = TrainingStore(router.state, 'upload-test') + router.store = router.state.create_training_store('upload-test') router.replica_id = 'router' await router.state.register_replica(router.replica_id, 0) await router.state.touch_replica_last_seen(router.replica_id) @@ -85,3 +84,18 @@ def cpu_import(name, *args, **kwargs): with pytest.raises(RequestRejectedError, match='rank'): await router.native_request(request, 'add_adapter_to_model') router.controller.create.remote.assert_not_awaited() + + +def test_router_schemas_match_model_runtime_routes(): + from fastapi import FastAPI + from twinkle.server.model.tinker_handlers import _register_model_tinker_routes + from twinkle.server.model.twinkle_handlers import _register_model_twinkle_routes + from twinkle.server.training.router import NATIVE_REQUESTS, TINKER_REQUESTS + + app = FastAPI() + _register_model_tinker_routes(app, lambda: None) + _register_model_twinkle_routes(app, lambda: None) + for protocol, schemas in [('twinkle', NATIVE_REQUESTS), ('tinker', TINKER_REQUESTS)]: + endpoints = {route.path.rsplit('/', 1)[-1] for route in app.routes + if route.path.startswith(f'/{protocol}/') and 'POST' in route.methods} + assert endpoints == set(schemas) diff --git a/tests/server/training/test_runtime_initialization.py b/tests/server/training/test_runtime_initialization.py new file mode 100644 index 000000000..27c8ec219 --- /dev/null +++ b/tests/server/training/test_runtime_initialization.py @@ -0,0 +1,49 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""The runtime event loop remains responsive during the first dynamic config import.""" +import asyncio +import threading +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from tests.server.state.fake_backend import FakeBackend +from tests.server.training.test_controller import create, setup +from twinkle.server.model.managed_runtime import ManagedTrainingRuntime +from twinkle.server.state.training_records import CreateTraining + + +@pytest.mark.asyncio +async def test_configuration_deserialization_does_not_block_event_loop(monkeypatch): + controller, _, _, _ = await setup(backend=FakeBackend()) + identity = await create(controller, controller.store.state, 'configuration-import', mode='full') + training = controller.trainings[identity] + entered, resume, finished = threading.Event(), threading.Event(), threading.Event() + + def deserialize(config): + entered.set() + resume.wait(3) + finished.set() + return None + + monkeypatch.setattr('twinkle.server.model.managed_runtime.deserialize_object', deserialize) + manager = SimpleNamespace(save=Mock()) + monkeypatch.setattr('twinkle.server.model.managed_runtime.create_training_run_manager', lambda *args, **kw: manager) + + async def serial(factory, task_type): + return await factory() + + runtime = SimpleNamespace(_managed_serial=serial, resolve_model_adapter_name=lambda _: '', + register_resource=Mock(), set_resource_state=Mock(), is_full_mode=True) + request = CreateTraining(base_model='small', mode='full', token='owner', session_id='session', create_key='key', + client_type='twinkle', payload={'adapter_name': 'same', 'config': None}) + initializing = asyncio.create_task(ManagedTrainingRuntime.initialize_training( + runtime, training.model_dump(), request.model_dump())) + try: + assert await asyncio.to_thread(entered.wait, 2) + assert not finished.is_set(), 'Synchronous import blocked the event loop' + finally: + resume.set() + await initializing + await controller.shutdown() + manager.save.assert_called_once() diff --git a/tests/server/training/test_runtime_queue.py b/tests/server/training/test_runtime_queue.py index 0994e5ffe..e489f3782 100644 --- a/tests/server/training/test_runtime_queue.py +++ b/tests/server/training/test_runtime_queue.py @@ -6,8 +6,9 @@ import pytest -from twinkle.server.task_queue import TaskQueueConfig, TaskQueueMixin -from twinkle.server.training.runtime import ManagedTrainingRuntime, _tinker_operation +from twinkle.server.config.task_queue import TaskQueueConfig +from twinkle.server.task_queue import TaskQueueMixin +from twinkle.server.model.managed_runtime import ManagedTrainingRuntime, _tinker_operation from tests.server.training.test_controller import create, setup diff --git a/tests/server/training/test_sequence.py b/tests/server/training/test_sequence.py index 904bc3df0..924258bc6 100644 --- a/tests/server/training/test_sequence.py +++ b/tests/server/training/test_sequence.py @@ -30,7 +30,7 @@ async def state(request): @pytest.mark.asyncio async def test_expired_retry_receipts_cannot_execute_old_sequences_again(state, monkeypatch): - backend = state.backend + backend = state._backend now = time.time() monkeypatch.setattr('twinkle.server.state.server_state.time.time', lambda: now) assert await state.claim_training_seq('training', 'native::session::step', 1, 'first', 10) is None @@ -58,13 +58,13 @@ async def test_out_of_order_sequences_and_preflight_release_remain_retryable(sta @pytest.mark.asyncio async def test_retiring_a_handle_removes_native_and_tinker_stream_ledgers(state): from tests.server.training.test_controller import create, setup - controller, state, _, _ = await setup(backend=state.backend) + controller, state, _, _ = await setup(backend=state._backend) try: identity = await create(controller, state, 'dedup-retirement') for stream in ('native::session::step', 'tinker::optim_step'): await state.claim_training_seq(identity, stream, 1, stream, 10) - assert len(await state.backend.keys(f'seq_ledger::{identity}::*')) == 2 + assert len(await state._backend.keys(f'seq_ledger::{identity}::*')) == 2 await controller.release(identity, 'owner', 'session') - assert await state.backend.keys(f'seq_ledger::{identity}::*') == [] + assert await state._backend.keys(f'seq_ledger::{identity}::*') == [] finally: await controller.shutdown() diff --git a/tests/server/training/test_storage_recovery.py b/tests/server/training/test_storage_recovery.py index 6740348de..c2529a6ca 100644 --- a/tests/server/training/test_storage_recovery.py +++ b/tests/server/training/test_storage_recovery.py @@ -8,7 +8,6 @@ from twinkle.server.state.backend.memory_backend import RayActorBackend from twinkle.server.state.backend.redis_backend import RedisBackend from twinkle.server.training.controller import TrainingController -from twinkle.server.training.store import TrainingStore from tests.server.training.test_controller import create, request, setup @@ -25,7 +24,7 @@ async def test_pending_operation_recovery_preserves_creation_identity_and_quota( original = await controller.create(request('retry')) await controller.admit(identity, 'owner', 'session', 'queued') await controller.store.future('queued', 'queued', identity) - replacement = TrainingController(controller.config, TrainingStore(state, 'unit'), runtime, clock=clock) + replacement = TrainingController(controller.config, state.create_training_store('unit'), runtime, clock=clock) await replacement.initialize() assert (await state.get_future('queued'))['failure']['reason_code'] == 'orphaned_replica' assert await replacement.create(request('retry')) == original diff --git a/tests/server/utils/task_queue/test_config.py b/tests/server/utils/task_queue/test_config.py index 60787f567..511d6bde5 100644 --- a/tests/server/utils/task_queue/test_config.py +++ b/tests/server/utils/task_queue/test_config.py @@ -13,7 +13,7 @@ from hypothesis import strategies as st from pydantic import ValidationError -from twinkle.server.task_queue.config import TaskQueueConfig +from twinkle.server.config.task_queue import TaskQueueConfig # ---------- defaults snapshot used by the default-value test -------------- # diff --git a/tests/server/utils/test_task_queue_mixin.py b/tests/server/utils/test_task_queue_mixin.py index 033c372b4..a947065d1 100644 --- a/tests/server/utils/test_task_queue_mixin.py +++ b/tests/server/utils/test_task_queue_mixin.py @@ -2,7 +2,7 @@ import pytest import pytest_asyncio -from twinkle.server.task_queue.config import TaskQueueConfig +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.task_queue.mixin import TaskQueueMixin from twinkle.server.task_queue.types import UserTaskError from twinkle.server.task_queue.worker import ComputeWorker diff --git a/tests/transformers/test_resident_training_reset.py b/tests/transformers/test_resident_training_reset.py index c35570071..ed1846c5a 100644 --- a/tests/transformers/test_resident_training_reset.py +++ b/tests/transformers/test_resident_training_reset.py @@ -37,3 +37,16 @@ def group(scale): assert model(torch.ones(1, 1)).item() == 1 assert model.weight.grad is None assert len(model._forward_pre_hooks) == 1 + + +def test_blank_model_still_resets_training_state(): + wrapper = object.__new__(TransformersModel) + wrapper.model_id = '' + wrapper._model_wrapped = False + previous = SimpleNamespace(template=None, _handler=None, optimizer=object()) + fresh = SimpleNamespace(template=None, _handler=None, optimizer=None) + wrapper.optimizer_group = {'': previous} + object.__setattr__(wrapper, '_construct_default_optimizer_group', lambda: fresh) + wrapper.reload_initial_weights(reset_training_state=True) + assert wrapper.optimizer_group == {'': fresh} + assert wrapper.active_group == '' and fresh.adapter_name == '' diff --git a/tools/validation/on_demand_training/allocation_probe.py b/tools/validation/on_demand_training/allocation_probe.py index 633785610..3b21dc0bd 100644 --- a/tools/validation/on_demand_training/allocation_probe.py +++ b/tools/validation/on_demand_training/allocation_probe.py @@ -8,9 +8,8 @@ from twinkle import DeviceGroup from twinkle.infra._ray.resource_manager import ResourceManager from twinkle.server.config.training import TrainingConfig -from twinkle.server.training.allocation import DevicePool +from twinkle.server.resources.allocation import CapacityUnavailableError, DevicePool from twinkle.server.training.catalog import TrainingCatalog -from twinkle.server.training.errors import CapacityUnavailableError async def main(): @@ -36,14 +35,14 @@ async def main(): }, }]) spec = TrainingCatalog(config).get(str(index), 'full') - allocation = await pool.allocate(spec, 'epoch', str(index)) + allocation = await pool.allocate(spec.config.num_devices, spec.model.device_type, 'epoch', str(index)) allocations.append(allocation) devices = [device for allocation in allocations for device in allocation.devices] assert len(set(devices)) == 8, devices assert len({allocation.node_id for allocation in allocations}) == 1 assert ray.available_resources().get('GPU', 0) == 0 try: - await pool.allocate(spec, 'epoch', 'overflow') + await pool.allocate(spec.config.num_devices, spec.model.device_type, 'epoch', 'overflow') raise AssertionError('over-allocation was admitted') except CapacityUnavailableError: pass @@ -72,7 +71,7 @@ async def main(): } }, }])).get('replacement', 'full') - replacement = await pool.allocate(two, 'epoch', 'replacement') + replacement = await pool.allocate(two.config.num_devices, two.model.device_type, 'epoch', 'replacement') allocations.append(replacement) assert set(replacement.devices) == set(freed.devices) print(json.dumps({'passed': True, 'devices': devices, 'reused': list(replacement.devices)})) diff --git a/tools/validation/on_demand_training/device_validation.py b/tools/validation/on_demand_training/device_validation.py index 2fe2c17c6..01695893f 100644 --- a/tools/validation/on_demand_training/device_validation.py +++ b/tools/validation/on_demand_training/device_validation.py @@ -132,8 +132,10 @@ def __init__(self, base, protocol, mode, alias='small'): if protocol == 'native': from peft import LoraConfig self.client = MultiLoraTransformersModel(alias, transport=self.transport) - self.client.add_adapter_to_model('same', - LoraConfig(r=8, target_modules='all-linear') if mode == 'lora' else None) + self.client.add_adapter_to_model( + 'same', + LoraConfig(r=8, target_modules='all-linear') if mode == 'lora' else None, + gradient_accumulation_steps=2) self.identity = self.client._training_ids['same'] self.client.set_loss('CrossEntropyLoss') self.client.set_optimizer('Adam', lr=self.rate) @@ -179,8 +181,9 @@ def evaluate(self): def step(self): if self.protocol == 'native': - self.client.forward_backward(self.inputs) - self.client.clip_grad_and_step() + for _ in range(2): + self.client.forward_backward(self.inputs, gradient_accumulation_steps=2) + self.client.clip_grad_and_step(gradient_accumulation_steps=2) else: from tinker import types self.client.forward_backward(self.datum(), 'cross_entropy').result() @@ -198,6 +201,23 @@ def active_instances(controller): return [row for row in ray.get(controller.diagnostics.remote())['instances'] if row['state'] != 'STOPPED'] +@pytest.fixture(autouse=True) +def cleanup_failed_case(server): + from twinkle.server.state.models import FutureFailureRecord + _, controller, _, _ = server + yield + # Release this case's resources even when a numeric assertion fails, so later combinations start clean. + for instance in active_instances(controller): + ray.get( + controller.fail_instance.remote( + instance['instance_id'], + FutureFailureRecord( + reason_code='orphaned_replica', attribution='server', message='Device test teardown')), + timeout=180) + assert not active_instances(controller) + assert ray.available_resources().get('GPU', 0) == 4 + + @pytest.mark.parametrize('protocol', ['native', 'tinker']) @pytest.mark.parametrize('mode', ['lora', 'full']) def test_dp_resource_lifecycle(server, protocol, mode): @@ -215,7 +235,9 @@ def test_dp_resource_lifecycle(server, protocol, mode): trained = first.evaluate() assert float(trained.mean()) > float(baseline.mean()) + 1e-4 np.testing.assert_allclose(second.evaluate(), untouched, rtol=0, atol=2e-3) - second.step() + # A single LoRA update may be invisible in BF16 outputs, so both tenants use the same training window. + for _ in range(4): + second.step() peer_trained = second.evaluate() assert float(peer_trained.mean()) > float(untouched.mean()) + 1e-4 np.testing.assert_allclose(first.evaluate(), trained, rtol=0, atol=2e-3) @@ -236,7 +258,8 @@ def test_dp_resource_lifecycle(server, protocol, mode): assert len(active_instances(controller)) == 1 assert ray.available_resources().get('GPU', 0) == 2 np.testing.assert_allclose(second.evaluate(), peer_trained, rtol=0, atol=2e-3) - second.step() + for _ in range(4): + second.step() peer_continued = second.evaluate() assert float(peer_continued.mean()) > float(peer_trained.mean()) + 1e-4 second.close() diff --git a/tools/validation/on_demand_training/fragmentation_probe.py b/tools/validation/on_demand_training/fragmentation_probe.py index 02dd22ab8..267a539a2 100644 --- a/tools/validation/on_demand_training/fragmentation_probe.py +++ b/tools/validation/on_demand_training/fragmentation_probe.py @@ -8,9 +8,8 @@ from ray.util.placement_group import placement_group, placement_group_table, remove_placement_group from twinkle.server.config.training import TrainingConfig -from twinkle.server.training.allocation import DevicePool +from twinkle.server.resources.allocation import CapacityUnavailableError, DevicePool from twinkle.server.training.catalog import TrainingCatalog -from twinkle.server.training.errors import CapacityUnavailableError async def main(): @@ -44,7 +43,7 @@ async def main(): await asyncio.wait_for(pg.ready(), 20) assert ray.available_resources()['GPU'] == 2 try: - await pool.allocate(spec, 'epoch', 'fragmented') + await pool.allocate(spec.config.num_devices, spec.model.device_type, 'epoch', 'fragmented') except CapacityUnavailableError: pass else: @@ -59,7 +58,7 @@ async def main(): await asyncio.sleep(0.05) else: raise AssertionError('Owned GPU reservation was not released') - allocated = await pool.allocate(spec, 'epoch', 'available') + allocated = await pool.allocate(spec.config.num_devices, spec.model.device_type, 'epoch', 'available') assert len(set(allocated.devices)) == 2 print( json.dumps({ diff --git a/tools/validation/on_demand_training/matrix_validation.py b/tools/validation/on_demand_training/matrix_validation.py new file mode 100644 index 000000000..aa9921d1e --- /dev/null +++ b/tools/validation/on_demand_training/matrix_validation.py @@ -0,0 +1,117 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Run the full training matrix on an isolated server; runtimes start from zero on client demand.""" +import argparse +import json +import os +import socket +import subprocess +import sys +import time +import uuid +from pathlib import Path +from xml.etree import ElementTree + +import ray +import yaml +from ray import serve + +from twinkle.server.config import ServerConfig +from twinkle.server.launcher.server_launcher import ServerLauncher + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--backend', required=True, choices=['transformers', 'megatron']) + parser.add_argument('--model', required=True) + parser.add_argument('--run-dir', required=True) + args = parser.parse_args() + root = Path(__file__).resolve().parents[3] + run_dir = Path(args.run_dir) + run_dir.mkdir(parents=True, exist_ok=True) + namespace = 'matrix_' + uuid.uuid4().hex[:8] + with socket.socket() as listener: + listener.bind(('127.0.0.1', 0)) + port = listener.getsockname()[1] + source = root / 'tests/server/config' / ( + 'server_config_4b_e2e.yaml' if args.backend == 'transformers' else 'server_config_4b_e2e_megatron.yaml') + raw = yaml.safe_load(source.read_text()) + raw['ray_namespace'] = namespace + raw['http_options'] = {'host': '127.0.0.1', 'port': port} + for app in raw['applications']: + params = app['args'] + if app['import_path'] == 'model': + params['model_id'] = args.model + params['nproc_per_node'] = 2 + params['device_group']['ranks'] = 2 + params['device_mesh'] = {'device_type': 'cuda', 'dp_size': 2} + params['min_replicas'] = 0 + params['max_replicas'] = 2 + params['idle_timeout_s'] = 3600 + params['adapter_config'] = {'adapter_timeout': 3600, 'adapter_max_lifetime': 36000} + params['env_vars'] = {'TWINKLE_FAIL_FAST': '0', 'TWINKLE_TRUST_REMOTE_CODE': '1'} + app.pop('deployments', None) + elif app['import_path'] == 'sampler': + params['model_id'] = args.model + params['engine_args']['attention_config'] = {'flash_attn_version': 3} + config = ServerConfig.model_validate(raw) + (run_dir / 'server.yaml').write_text(yaml.safe_dump(config.to_yaml_dict(), sort_keys=False)) + os.environ.update({ + 'TWINKLE_STATE_RAY_NAMESPACE': namespace, + 'TWINKLE_PERSISTENCE_MODE': 'memory', + 'TWINKLE_DEFAULT_SAVE_DIR': str(run_dir / 'checkpoints'), + 'TWINKLE_MODEL_ID_ALIASES': json.dumps({'Qwen/Qwen3.5-4B': args.model}), + 'TWINKLE_TRUST_REMOTE_CODE': '1', + 'TWINKLE_FAIL_FAST': '0', + 'TWINKLE_TEST_GPU_E2E': '1', + 'TWINKLE_TEST_BACKEND': args.backend, + 'TWINKLE_SERVER_URL': f'http://127.0.0.1:{port}', + 'TWINKLE_TEST_OPERATION_TIMEOUT': '600', + 'PYTHONPATH': str(root / 'src'), + }) + evidence = {'backend': args.backend, 'namespace': namespace, 'cases': [], 'lazy_start': False} + ray.init(address='local', namespace=namespace, num_cpus=24, num_gpus=3, + _temp_dir='/tmp/twm-' + namespace[-8:], include_dashboard=False) + launcher = ServerLauncher(config) + try: + launcher._start_serve() + for app in config.applications: + launcher._deploy_application(app) + launcher._deploy_training() + assert ray.get(launcher._training_controller.diagnostics.remote())['instances'] == [] + evidence['lazy_start'] = True + for filename in ['test_sft_e2e.py', 'test_dpo_e2e.py', 'test_grpo_e2e.py', 'test_full_cycle_e2e.py']: + start = time.monotonic() + with (run_dir / (filename + '.log')).open('w') as output: + result = subprocess.run([ + sys.executable, '-m', 'pytest', str(root / 'tests/server/integration' / filename), + '-vv', '-s', '--junitxml=' + str(run_dir / (filename + '.xml')), + ], stdout=output, stderr=subprocess.STDOUT, env=os.environ.copy(), timeout=7200) + junit = run_dir / (filename + '.xml') + cases = ElementTree.parse(junit).getroot().findall('.//testcase') if junit.exists() else [] + skipped = sum(case.find('skipped') is not None for case in cases) + evidence['cases'].append({'file': filename, 'returncode': result.returncode, + 'tests': len(cases), 'skipped': skipped, + 'seconds': time.monotonic() - start}) + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + assert result.returncode == 0, f'{filename} failed; see {run_dir / (filename + ".log")}' + assert len(cases) == (1 if filename == 'test_full_cycle_e2e.py' else 2) and skipped == 0 + finally: + try: + controller = getattr(launcher, '_training_controller', None) + if controller is not None: + ray.get(controller.shutdown.remote(), timeout=180) + diagnostic = ray.get(controller.diagnostics.remote(), timeout=30) + assert all(row['state'] == 'STOPPED' for row in diagnostic['instances']) + evidence['training_resources_released'] = True + launcher._shutdown() + finally: + serve.shutdown() + ray.shutdown() + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + evidence['passed'] = True + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) From 709e51acb885c875fafd641ae1d191966a3ea7dc Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Mon, 19 Oct 2026 08:33:16 +0000 Subject: [PATCH 04/46] fix(server): wait for concurrent training instance shutdown --- src/twinkle/server/training/controller.py | 107 ++++++++++-------- .../server/training/test_failure_recovery.py | 29 +++++ 2 files changed, 86 insertions(+), 50 deletions(-) diff --git a/src/twinkle/server/training/controller.py b/src/twinkle/server/training/controller.py index 15cecd414..499cc0f21 100644 --- a/src/twinkle/server/training/controller.py +++ b/src/twinkle/server/training/controller.py @@ -87,7 +87,7 @@ def __init__(self, self._startup: dict[str, asyncio.Task] = {} self._jobs: set[asyncio.Task] = set() self._cleaning: set[str] = set() - self._closing: set[str] = set() + self._closing: dict[str, asyncio.Task] = {} self._reconciling: set[str] = set() self._failed_instances: dict[str, FutureFailureRecord] = {} self._lock = asyncio.Lock() @@ -568,58 +568,65 @@ async def _cleanup_user(self, training: TrainingRecord) -> None: self._cleaning.discard(training.training_id) async def _close_instance(self, instance: InstanceRecord) -> None: - if instance.instance_id in self._closing: - return - self._closing.add(instance.instance_id) + close_task = self._closing.get(instance.instance_id) + if close_task is None: + close_task = self._spawn(self._close_instance_once(instance)) + self._closing[instance.instance_id] = close_task + + def clear(completed: asyncio.Task) -> None: + if self._closing.get(instance.instance_id) is completed: + self._closing.pop(instance.instance_id, None) + + close_task.add_done_callback(clear) + await asyncio.shield(close_task) + + async def _close_instance_once(self, instance: InstanceRecord) -> None: try: - try: - await asyncio.wait_for(self.driver.shutdown(instance), self.config.shutdown_timeout_s) - except Exception as exc: - async with self._lock: - quarantined = instance.model_copy(update={'state': 'QUARANTINED', 'failure': str(exc)}) - await self.store.save_instance(quarantined) - instance.state = 'QUARANTINED' - instance.failure = str(exc) - return + await asyncio.wait_for(self.driver.shutdown(instance), self.config.shutdown_timeout_s) + except Exception as exc: async with self._lock: - terminal_records = [] - terminal_at = time.time() - failure = self._failed_instances.get(instance.instance_id) - for training_id in list(instance.training_ids): - training = self.trainings[training_id] - if failure is not None: - for request_id in (training.create_request_id, *training.pending_operation_ids): - await self.store.future(request_id, 'failed', training_id, failure=failure) - terminal = training.model_copy( - deep=True, - update={ - 'pending_operation_ids': {}, - 'state': { - 'expired': 'EXPIRED', - 'released': 'RELEASED', - 'cancelled': 'CANCELLED' - }.get(training.terminal_reason, 'FAILED'), - 'terminal_at': terminal_at - }) - await self.store.retire(terminal) - terminal_records.append((training, terminal)) - stopped = instance.model_copy( - deep=True, update={ - 'state': 'STOPPED', - 'stopped_at': terminal_at, - 'training_ids': [] + quarantined = instance.model_copy(update={'state': 'QUARANTINED', 'failure': str(exc)}) + await self.store.save_instance(quarantined) + instance.state = 'QUARANTINED' + instance.failure = str(exc) + return + async with self._lock: + terminal_records = [] + terminal_at = time.time() + failure = self._failed_instances.get(instance.instance_id) + for training_id in list(instance.training_ids): + training = self.trainings[training_id] + if failure is not None: + for request_id in (training.create_request_id, *training.pending_operation_ids): + await self.store.future(request_id, 'failed', training_id, failure=failure) + terminal = training.model_copy( + deep=True, + update={ + 'pending_operation_ids': {}, + 'state': { + 'expired': 'EXPIRED', + 'released': 'RELEASED', + 'cancelled': 'CANCELLED' + }.get(training.terminal_reason, 'FAILED'), + 'terminal_at': terminal_at }) - await self.store.save_instance(stopped) - for training, terminal in terminal_records: - training.pending_operation_ids.clear() - training.state = terminal.state - training.terminal_at = terminal.terminal_at - instance.state = 'STOPPED' - instance.stopped_at = terminal_at - instance.training_ids.clear() - self._failed_instances.pop(instance.instance_id, None) - finally: - self._closing.discard(instance.instance_id) + await self.store.retire(terminal) + terminal_records.append((training, terminal)) + stopped = instance.model_copy( + deep=True, update={ + 'state': 'STOPPED', + 'stopped_at': terminal_at, + 'training_ids': [] + }) + await self.store.save_instance(stopped) + for training, terminal in terminal_records: + training.pending_operation_ids.clear() + training.state = terminal.state + training.terminal_at = terminal.terminal_at + instance.state = 'STOPPED' + instance.stopped_at = terminal_at + instance.training_ids.clear() + self._failed_instances.pop(instance.instance_id, None) async def fail_instance(self, instance_id: str, failure: FutureFailureRecord) -> None: async with self._lock: diff --git a/tests/server/training/test_failure_recovery.py b/tests/server/training/test_failure_recovery.py index d0b78ed9e..bf93488c1 100644 --- a/tests/server/training/test_failure_recovery.py +++ b/tests/server/training/test_failure_recovery.py @@ -192,6 +192,35 @@ async def save(record): await controller.shutdown() +@pytest.mark.asyncio +async def test_shutdown_waits_for_an_in_flight_instance_close(): + controller, _, runtime, _ = await setup(backend=FakeBackend(), min_replicas=1) + instance = next(iter(controller.instances.values())) + original_shutdown = runtime.shutdown + entered, resume = asyncio.Event(), asyncio.Event() + + async def slow_shutdown(target): + entered.set() + await resume.wait() + await original_shutdown(target) + + runtime.shutdown = slow_shutdown + closing = asyncio.create_task(controller._close_instance(instance)) + await entered.wait() + shutdown = asyncio.create_task(controller.shutdown()) + for _ in range(100): + if instance.state == 'DRAINING' or shutdown.done(): + break + await asyncio.sleep(0) + + assert instance.state == 'DRAINING' + assert not shutdown.done() + resume.set() + await asyncio.gather(closing, shutdown) + assert instance.state == 'STOPPED' + assert instance.instance_id not in controller._closing + + @pytest.mark.asyncio async def test_cleanup_error_is_not_classified_as_startup_failure(): controller, state, runtime, _ = await setup(backend=FakeBackend()) From ef290ed118e2ff34eadd31f2bc59a6c12b7aafc6 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Tue, 20 Oct 2026 02:11:07 +0000 Subject: [PATCH 05/46] test(server): allow PPU training resources to stop cleanly --- tools/validation/on_demand_training/matrix_validation.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tools/validation/on_demand_training/matrix_validation.py b/tools/validation/on_demand_training/matrix_validation.py index d55f8d791..387da6371 100644 --- a/tools/validation/on_demand_training/matrix_validation.py +++ b/tools/validation/on_demand_training/matrix_validation.py @@ -34,6 +34,7 @@ def main(): source = root / 'tests/server/config' / ('server_config_4b_e2e.yaml' if args.backend == 'transformers' else 'server_config_4b_e2e_megatron.yaml') raw = yaml.safe_load(source.read_text()) + raw.setdefault('training', {})['shutdown_timeout_s'] = 180 raw['ray_namespace'] = namespace raw['http_options'] = {'host': '127.0.0.1', 'port': port} for app in raw['applications']: From e8f42aecb8b0a102e40dfccf9004e368ed36bf24 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Wed, 21 Oct 2026 00:17:39 +0000 Subject: [PATCH 06/46] fix(server): harden on-demand training recovery and validation Fix optimizer metrics, wrapped parameter groups and safe Redis reconnects. Add lifecycle metrics, isolated Redis/GPU acceptance and stress coverage. Move shared probes into tests and untrack local tools. --- .gitignore | 3 +- .../client/server/on_demand_training.yaml | 3 +- src/twinkle/model/optimizer_group.py | 16 +- .../model/transformers/transformers.py | 5 +- .../server/state/backend/redis_backend.py | 56 ++- src/twinkle/server/telemetry/metrics.py | 105 +++++ src/twinkle/server/training/actor.py | 9 + src/twinkle/server/training/controller.py | 75 +++ src/twinkle/server/training/driver.py | 43 +- .../on_demand_training/__init__.py | 1 + .../on_demand_training/acceptance_support.py | 163 +++++++ .../on_demand_training/allocation_probe.py | 20 +- .../on_demand_training/fragmentation_probe.py | 5 +- .../on_demand_training/http_probe.py | 14 +- .../on_demand_training/managed_redis.py | 65 +++ .../on_demand_training/matrix_validation.py | 441 ++++++++++++++++++ .../on_demand_training/on_demand_sdk.py | 0 .../on_demand_training/stress_validation.py | 102 ++++ .../server/integration/test_full_cycle_e2e.py | 7 +- .../test_on_demand_gpu_redis_e2e.py | 53 +++ .../server/telemetry/test_training_metrics.py | 152 ++++++ .../training/test_acceptance_harness.py | 84 ++++ tests/server/training/test_allocation.py | 4 +- .../server/training/test_http_integration.py | 20 +- .../training/test_redis_restart_reads.py | 114 +++++ tests/server/training/test_runtime_queue.py | 1 + .../test_optimizer_group_regression.py | 48 ++ .../on_demand_training/device_validation.py | 327 ------------- .../on_demand_training/matrix_validation.py | 135 ------ 29 files changed, 1565 insertions(+), 506 deletions(-) create mode 100644 tests/server/integration/on_demand_training/__init__.py create mode 100644 tests/server/integration/on_demand_training/acceptance_support.py rename {tools/validation => tests/server/integration}/on_demand_training/allocation_probe.py (83%) rename {tools/validation => tests/server/integration}/on_demand_training/fragmentation_probe.py (94%) rename {tools/validation => tests/server/integration}/on_demand_training/http_probe.py (97%) create mode 100644 tests/server/integration/on_demand_training/managed_redis.py create mode 100644 tests/server/integration/on_demand_training/matrix_validation.py rename {tools/validation => tests/server/integration}/on_demand_training/on_demand_sdk.py (100%) create mode 100644 tests/server/integration/on_demand_training/stress_validation.py create mode 100644 tests/server/integration/test_on_demand_gpu_redis_e2e.py create mode 100644 tests/server/telemetry/test_training_metrics.py create mode 100644 tests/server/training/test_acceptance_harness.py create mode 100644 tests/server/training/test_redis_restart_reads.py create mode 100644 tests/transformers/test_optimizer_group_regression.py delete mode 100644 tools/validation/on_demand_training/device_validation.py delete mode 100644 tools/validation/on_demand_training/matrix_validation.py diff --git a/.gitignore b/.gitignore index dac8c4337..cc52a2f6a 100644 --- a/.gitignore +++ b/.gitignore @@ -77,8 +77,9 @@ instance/ # Sphinx documentation docs/_build/ -# Internal design notes are local-only and must not be committed. +# Internal design notes and developer tools are local-only. /docs/design/ +/tools/ # PyBuilder target/ diff --git a/cookbook/client/server/on_demand_training.yaml b/cookbook/client/server/on_demand_training.yaml index 12018f1f7..60ff1caba 100644 --- a/cookbook/client/server/on_demand_training.yaml +++ b/cookbook/client/server/on_demand_training.yaml @@ -2,7 +2,8 @@ training: allocation_timeout_s: 3 startup_timeout_s: 1800 - shutdown_timeout_s: 60 + # PPU actor/PlacementGroup death confirmation can exceed the 60s framework default. + shutdown_timeout_s: 180 sweep_interval_s: 5 applications: - name: model-qwen diff --git a/src/twinkle/model/optimizer_group.py b/src/twinkle/model/optimizer_group.py index 150e694bc..59b4b331b 100644 --- a/src/twinkle/model/optimizer_group.py +++ b/src/twinkle/model/optimizer_group.py @@ -64,15 +64,15 @@ def accumulate_metrics(self, is_training): """Accumulate metrics for train/eval status. Override in subclass if needed.""" status = self.train_status if is_training else self.eval_status if len(status.metrics) > 0 and status.inputs is not None and status.outputs is not None: + metric_kwargs = { + **status.forward_kwargs, + 'lr': self._get_lr(), + 'step': self.cur_step - 1, + 'gradient_accumulation_steps': self.gradient_accumulation_steps, + 'grad_norm': self._last_grad_norm, + } for metric in status.metrics: - metric.accumulate( - status.inputs, - status.outputs, - lr=self._get_lr(), - step=self.cur_step - 1, - gradient_accumulation_steps=self.gradient_accumulation_steps, - grad_norm=self._last_grad_norm, - **status.forward_kwargs) + metric.accumulate(status.inputs, status.outputs, **metric_kwargs) def calculate_metrics(self, is_training): """Calculate and return metrics.""" diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index f255b5325..5a8a38561 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -1127,7 +1127,10 @@ def get_parameter_names(model, forbidden_layer_types, forbidden_layer_names=None return result forbidden_name_patterns = [r'bias', r'layernorm', r'rmsnorm', r'(?:^|\.)norm(?:$|\.)', r'_norm(?:$|\.)'] - decay_parameters = get_parameter_names(self.model, [torch.nn.LayerNorm], forbidden_name_patterns) + # Match names from the same unwrapped model used by _get_trainable_parameters. + # DDP's module prefix otherwise changes optimizer groups for later tenants. + decay_parameters = get_parameter_names( + self.strategy.unwrap_model(self.model), [torch.nn.LayerNorm], forbidden_name_patterns) params = self._get_trainable_parameters(adapter_name) decay_param_names = [n for n, p in params.items() if (n in decay_parameters and p.requires_grad)] no_decay_param_names = [n for n, p in params.items() if (n not in decay_parameters and p.requires_grad)] diff --git a/src/twinkle/server/state/backend/redis_backend.py b/src/twinkle/server/state/backend/redis_backend.py index e6769551a..62c6a7be1 100644 --- a/src/twinkle/server/state/backend/redis_backend.py +++ b/src/twinkle/server/state/backend/redis_backend.py @@ -6,10 +6,13 @@ from collections.abc import Callable from typing import Any +from twinkle.server.telemetry.metrics import get_state_backend_metrics from .base import ConcurrencyError, StateBackend try: import redis.asyncio as aioredis + from redis.exceptions import ConnectionError as RedisConnectionError + from redis.exceptions import TimeoutError as RedisTimeoutError from redis.exceptions import WatchError _REDIS_AVAILABLE = True @@ -38,6 +41,7 @@ def __init__(self, redis_url: str, key_prefix: str = '') -> None: raise ImportError('redis package required. Install with: pip install redis') self._client = aioredis.from_url(redis_url, decode_responses=True) self._prefix = key_prefix + self._metrics = get_state_backend_metrics() def _make_key(self, key: str) -> str: """Add namespace prefix to key.""" @@ -56,10 +60,27 @@ async def set(self, key: str, value: Any, ttl: int | None = None) -> None: else: await self._client.set(real_key, data) + async def _read_with_reconnect(self, operation, *args): + """Retry only read commands after a stale pooled connection is discarded. + + Writes/EXEC are deliberately excluded: a lost acknowledgement does not + prove that a mutation failed, so replay could apply it twice. + """ + for attempt in range(3): + try: + result = await operation(*args) + self._metrics.redis_health.set(1) + return result + except (RedisConnectionError, RedisTimeoutError): + self._metrics.redis_health.set(0) + if attempt == 2: + raise + await asyncio.sleep(0.05 * (attempt + 1)) + async def get(self, key: str) -> Any | None: """Retrieve value, return None if not found or expired.""" real_key = self._make_key(key) - raw = await self._client.get(real_key) + raw = await self._read_with_reconnect(self._client.get, real_key) if raw is None: return None return json.loads(raw) @@ -72,7 +93,7 @@ async def delete(self, key: str) -> None: async def exists(self, key: str) -> bool: """Check if key exists and is not expired.""" real_key = self._make_key(key) - return bool(await self._client.exists(real_key)) + return bool(await self._read_with_reconnect(self._client.exists, real_key)) async def keys(self, pattern: str) -> list[str]: """Return all key names matching the pattern using SCAN. @@ -116,8 +137,10 @@ async def update_atomic( / heartbeat) or reraise (lease renewal). """ real_key = self._make_key(key) + connection_failures = 0 for attempt in range(_UPDATE_ATOMIC_MAX_RETRIES): async with self._client.pipeline() as pipe: + executing = False try: await pipe.watch(real_key) raw = await pipe.get(real_key) @@ -125,21 +148,40 @@ async def update_atomic( new_value = transform(current) if new_value is None: await pipe.unwatch() + self._metrics.update_atomic_retries.observe(attempt) return current pipe.multi() if ttl is not None: pipe.set(real_key, json.dumps(new_value), ex=ttl) else: pipe.set(real_key, json.dumps(new_value)) + executing = True await pipe.execute() + self._metrics.update_atomic_retries.observe(attempt) return new_value - except WatchError: + except WatchError as exc: + # redis-py also wraps a lost watched connection as WatchError. + # Once EXEC was sent, its outcome may be unknown, not aborted. + transport_failure = isinstance(exc.__context__, + (RedisConnectionError, RedisTimeoutError)) or (str(exc).startswith( + ('A ConnectionError occurred', 'A TimeoutError occurred'))) + if executing and transport_failure: + raise RedisConnectionError('Atomic update outcome could not be confirmed') from exc + self._metrics.update_atomic_conflicts.inc() backoff = min( _UPDATE_ATOMIC_BASE_BACKOFF_SECONDS * (2**attempt), _UPDATE_ATOMIC_MAX_BACKOFF_SECONDS, ) await asyncio.sleep(backoff * random.random()) continue + except (RedisConnectionError, RedisTimeoutError): + self._metrics.redis_health.set(0) + connection_failures += 1 + if executing or connection_failures >= 3: + raise + # Before EXEC there was no mutation, so reconnecting is safe. + await asyncio.sleep(0.05 * connection_failures) + self._metrics.update_atomic_retries.observe(_UPDATE_ATOMIC_MAX_RETRIES) raise ConcurrencyError(f'update_atomic exhausted {_UPDATE_ATOMIC_MAX_RETRIES} retries on key {key!r}') async def mget(self, keys: list[str]) -> list[Any | None]: @@ -147,7 +189,7 @@ async def mget(self, keys: list[str]) -> list[Any | None]: if not keys: return [] real_keys = [self._make_key(k) for k in keys] - raw_values = await self._client.mget(real_keys) + raw_values = await self._read_with_reconnect(self._client.mget, real_keys) return [json.loads(v) if v is not None else None for v in raw_values] async def close(self) -> None: @@ -157,6 +199,8 @@ async def close(self) -> None: async def health_check(self) -> bool: """Check if Redis is healthy and available.""" try: - return bool(await self._client.ping()) + healthy = bool(await self._read_with_reconnect(self._client.ping)) except Exception: - return False + healthy = False + self._metrics.redis_health.set(1 if healthy else 0) + return healthy diff --git a/src/twinkle/server/telemetry/metrics.py b/src/twinkle/server/telemetry/metrics.py index 586277708..1c287ab9e 100644 --- a/src/twinkle/server/telemetry/metrics.py +++ b/src/twinkle/server/telemetry/metrics.py @@ -111,6 +111,33 @@ class TaskMetrics(BaseModel): rate_limiter_active_tokens: _Gauge +class TrainingMetrics(BaseModel): + """Low-cardinality metrics for the managed training lifecycle.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + instances_active: _Gauge + startup_seconds: _Histogram + shutdown_seconds: _Histogram + allocation_seconds: _Histogram + removal_seconds: _Histogram + quarantined_total: _Counter + sweep_total: _Counter + sweep_seconds: _Histogram + devices_allocated: _Gauge + placement_groups_active: _Gauge + + +class StateBackendMetrics(BaseModel): + """Metrics for persistent state backend availability and contention.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + redis_health: _Gauge + update_atomic_conflicts: _Counter + update_atomic_retries: _Histogram + + class _RequestMetrics(BaseModel): """HTTP request metrics container (internal).""" @@ -177,6 +204,45 @@ def __init__(self) -> None: description='Number of tokens currently tracked by the rate limiter', ) + # === Managed training === + self.training_instances_active = meter.create_up_down_counter( + 'twinkle.training.instances.active', description='Active managed training instances by lifecycle state') + self.training_startup_seconds = meter.create_histogram( + 'twinkle.training.instance.startup.duration_seconds', + description='Training runtime startup duration', + unit='s') + self.training_shutdown_seconds = meter.create_histogram( + 'twinkle.training.instance.shutdown.duration_seconds', + description='Training runtime shutdown duration', + unit='s') + self.training_allocation_seconds = meter.create_histogram( + 'twinkle.training.placement_group.allocation.duration_seconds', + description='Placement group allocation duration', + unit='s') + self.training_removal_seconds = meter.create_histogram( + 'twinkle.training.placement_group.removal.duration_seconds', + description='Placement group removal duration', + unit='s') + self.training_quarantined_total = meter.create_counter( + 'twinkle.training.quarantined.total', description='Training instances quarantined during cleanup') + self.training_sweep_total = meter.create_counter( + 'twinkle.training.sweep.total', description='Training lifecycle sweep runs by status') + self.training_sweep_seconds = meter.create_histogram( + 'twinkle.training.sweep.duration_seconds', description='Training lifecycle sweep duration', unit='s') + self.training_devices_allocated = meter.create_up_down_counter( + 'twinkle.training.devices.allocated', description='Logical devices owned by managed training instances') + self.training_placement_groups_active = meter.create_up_down_counter( + 'twinkle.training.placement_groups.active', + description='Placement groups owned by managed training instances') + + # === Redis state backend === + self.redis_health = meter.create_up_down_counter( + 'twinkle.state.redis.health', description='Last observed Redis health (1 healthy, 0 unhealthy)') + self.redis_update_atomic_conflicts = meter.create_counter( + 'twinkle.state.redis.update_atomic.conflicts.total', description='Redis WATCH conflicts') + self.redis_update_atomic_retries = meter.create_histogram( + 'twinkle.state.redis.update_atomic.retries', description='Retries used by Redis atomic updates') + # === Resources (ObservableGauge backed by the push cache) === # OTEL holds its own references to the gauges via the meter, so we # only need to keep the cache the callbacks read from. @@ -192,6 +258,8 @@ def __init__(self) -> None: # swaps this singleton to rebind instruments) invalidates them. self._task_metrics: dict[str, TaskMetrics] = {} self._request_metrics: dict[str, _RequestMetrics] = {} + self._training_metrics: TrainingMetrics | None = None + self._state_backend_metrics: StateBackendMetrics | None = None def _make_gauge_callback(self, name: str): """Build the sync OTEL callback that reads ``_resource_cache[name]``.""" @@ -227,6 +295,33 @@ def request_metrics(self, deployment: str) -> _RequestMetrics: ) return cached + def training_metrics(self) -> TrainingMetrics: + """Return process-local managed training metric adapters.""" + if self._training_metrics is None: + self._training_metrics = TrainingMetrics( + instances_active=_Gauge(self.training_instances_active), + startup_seconds=_Histogram(self.training_startup_seconds), + shutdown_seconds=_Histogram(self.training_shutdown_seconds), + allocation_seconds=_Histogram(self.training_allocation_seconds), + removal_seconds=_Histogram(self.training_removal_seconds), + quarantined_total=_Counter(self.training_quarantined_total), + sweep_total=_Counter(self.training_sweep_total), + sweep_seconds=_Histogram(self.training_sweep_seconds), + devices_allocated=_Gauge(self.training_devices_allocated), + placement_groups_active=_Gauge(self.training_placement_groups_active), + ) + return self._training_metrics + + def state_backend_metrics(self) -> StateBackendMetrics: + """Return process-local state backend metric adapters.""" + if self._state_backend_metrics is None: + self._state_backend_metrics = StateBackendMetrics( + redis_health=_Gauge(self.redis_health), + update_atomic_conflicts=_Counter(self.redis_update_atomic_conflicts), + update_atomic_retries=_Histogram(self.redis_update_atomic_retries), + ) + return self._state_backend_metrics + # ----- Push API for the cleanup leader ----- def set_resource_count(self, name: str, value: int) -> None: @@ -268,3 +363,13 @@ def get_task_metrics(deployment: str) -> TaskMetrics: def get_request_metrics(deployment: str) -> _RequestMetrics: """Return the per-deployment HTTP request metric adapters.""" return MetricsRegistry.get().request_metrics(deployment) + + +def get_training_metrics() -> TrainingMetrics: + """Return managed training metric adapters.""" + return MetricsRegistry.get().training_metrics() + + +def get_state_backend_metrics() -> StateBackendMetrics: + """Return persistent state backend metric adapters.""" + return MetricsRegistry.get().state_backend_metrics() diff --git a/src/twinkle/server/training/actor.py b/src/twinkle/server/training/actor.py index b26d5a88e..2e109d1cf 100644 --- a/src/twinkle/server/training/actor.py +++ b/src/twinkle/server/training/actor.py @@ -10,6 +10,8 @@ from twinkle.server.state import get_server_state from twinkle.server.state.models import FutureFailureRecord from twinkle.server.state.training_records import CreateTraining +from twinkle.server.telemetry.metrics import get_training_metrics +from twinkle.server.telemetry.worker_init import ensure_telemetry_initialized from twinkle.utils.logger import get_logger from .controller import TrainingController from .driver import RayInstanceDriver @@ -25,11 +27,13 @@ def controller_name(namespace): class TrainingControllerActor: def __init__(self, config, namespace): + ensure_telemetry_initialized() self.config = TrainingConfig.model_validate(config) self.namespace = namespace self.store = get_server_state().create_training_store(namespace) self.driver = RayInstanceDriver(self.config, self.store, namespace, controller_name(namespace)) self.core = TrainingController(self.config, self.store, self.driver) + self._metrics = get_training_metrics() self._monitors = [] self._unreachable_since = {} self._initialization = asyncio.get_event_loop().create_task(self._initialize()) @@ -91,10 +95,15 @@ async def _heartbeat(self): async def _sweep(self): while True: + started, status = time.monotonic(), 'success' try: await self.core.sweep() except Exception: + status = 'error' logger.exception('Training lifecycle sweep failed; resources remain owned') + finally: + self._metrics.sweep_total.inc(tags={'status': status}) + self._metrics.sweep_seconds.observe(time.monotonic() - started, tags={'status': status}) await asyncio.sleep(self.config.sweep_interval_s) async def _recover_lost_state(self): diff --git a/src/twinkle/server/training/controller.py b/src/twinkle/server/training/controller.py index 499cc0f21..455c04027 100644 --- a/src/twinkle/server/training/controller.py +++ b/src/twinkle/server/training/controller.py @@ -19,6 +19,7 @@ from twinkle.server.state.training_records import (LIVE_INSTANCES, TERMINAL_TRAINING, CreateTraining, InstanceRecord, TrainingRecord) from twinkle.server.state.training_store import TrainingStore +from twinkle.server.telemetry.metrics import get_training_metrics from .catalog import TrainingCatalog, TrainingSpec from .errors import TrainingExpiredError @@ -92,6 +93,53 @@ def __init__(self, self._failed_instances: dict[str, FutureFailureRecord] = {} self._lock = asyncio.Lock() self._ready = False + self._metrics = get_training_metrics() + self._reported_metric_keys: set[tuple[str, str, str, str]] = set() + self._reported_resource_metric_keys: set[tuple[str, str, str]] = set() + + def _instance_metric_tags(self, instance: InstanceRecord) -> dict[str, str]: + spec = self.catalog.get(instance.base_model, instance.mode) + return { + 'backend': spec.model.backend, + 'mode': instance.mode, + 'state': instance.state, + 'resident': str(instance.resident).lower(), + } + + def _report_instance_metrics(self) -> None: + counts: dict[tuple[str, str, str, str], int] = {} + resources: dict[tuple[str, str], tuple[int, int]] = {} + for instance in self.instances.values(): + if instance.state == 'STOPPED': + continue + tags = self._instance_metric_tags(instance) + key = tuple(tags[name] for name in ('backend', 'mode', 'state', 'resident')) + counts[key] = counts.get(key, 0) + 1 + resource_key = tags['backend'], tags['mode'], self.catalog.get(instance.base_model, + instance.mode).model.device_type + devices, groups = resources.get(resource_key, (0, 0)) + resources[resource_key] = (devices + len(instance.devices), groups + int(bool(instance.application_name))) + for key in self._reported_metric_keys | counts.keys(): + tags = dict(zip(('backend', 'mode', 'state', 'resident'), key)) + self._metrics.instances_active.set(counts.get(key, 0), tags=tags) + for key in self._reported_resource_metric_keys | resources.keys(): + tags = dict(zip(('backend', 'mode', 'device_type'), key)) + devices, groups = resources.get(key, (0, 0)) + self._metrics.devices_allocated.set(devices, tags=tags) + self._metrics.placement_groups_active.set(groups, tags=tags) + self._reported_metric_keys = set(counts) + self._reported_resource_metric_keys = set(resources) + + @staticmethod + def _quarantine_reason(exc: Exception) -> str: + message = str(exc).lower() + if isinstance(exc, TimeoutError): + return 'timeout' + if 'worker death' in message: + return 'worker_death' + if 'placement group' in message: + return 'pg_removal' + return 'other' def _spawn(self, coroutine) -> asyncio.Task: task = asyncio.create_task(coroutine) @@ -169,6 +217,7 @@ async def register_runtime(self, instance_id: str, epoch: str, incarnation: str) instance.incarnation = incarnation instance.state = 'STARTING' await self.store.save_instance(instance) + self._report_instance_metrics() async def mark_dirty(self, instance_id: str, epoch: str) -> None: async with self._lock: @@ -194,11 +243,14 @@ async def _new_instance(self, spec: TrainingSpec, *, resident: bool = False) -> startup_deadline=self.clock() + self.config.startup_timeout_s) await self.store.save_instance(instance) self.instances[instance.instance_id] = instance + self._report_instance_metrics() self._startup[instance.instance_id] = self._spawn(self._start_instance(instance)) return instance async def _start_instance(self, instance: InstanceRecord) -> None: spec = self.catalog.get(instance.base_model, instance.mode) + started, status = time.monotonic(), 'error' + tags = {'backend': spec.model.backend, 'mode': spec.mode} try: async def startup(): @@ -214,6 +266,7 @@ async def startup(): await self.store.save_instance(allocated) for name, value in details.items(): setattr(instance, name, value) + self._report_instance_metrics() # Ownership is persisted before any worker is constructed; the driver does not write records. await self.driver.deploy(instance.model_copy(deep=True), spec) @@ -227,9 +280,13 @@ async def startup(): ready = instance.model_copy(update={'state': 'READY'}) await self.store.save_instance(ready) instance.state = 'READY' + self._report_instance_metrics() + status = 'success' except Exception as exc: await self.fail_instance(instance.instance_id, failure_for(exc)) raise + finally: + self._metrics.startup_seconds.observe(time.monotonic() - started, tags={**tags, 'status': status}) def _select(self, spec: TrainingSpec) -> InstanceRecord | None: candidates = [ @@ -511,6 +568,8 @@ async def _cleanup_user(self, training: TrainingRecord) -> None: updated_instance = instance.model_copy(update={'state': 'DRAINING'}) await self.store.save_instance(updated_instance) instance.state = 'DRAINING' + if resetting or closing_full: + self._report_instance_metrics() if closing_full: await self._close_instance(instance) return @@ -561,6 +620,7 @@ async def _cleanup_user(self, training: TrainingRecord) -> None: training.terminal_at = terminal.terminal_at instance.training_ids = remaining instance.state = updated_instance.state + self._report_instance_metrics() if close: await self._close_instance(instance) finally: @@ -581,6 +641,9 @@ def clear(completed: asyncio.Task) -> None: await asyncio.shield(close_task) async def _close_instance_once(self, instance: InstanceRecord) -> None: + started = time.monotonic() + tags = self._instance_metric_tags(instance) + timing_tags = {'backend': tags['backend'], 'mode': tags['mode']} try: await asyncio.wait_for(self.driver.shutdown(instance), self.config.shutdown_timeout_s) except Exception as exc: @@ -589,7 +652,16 @@ async def _close_instance_once(self, instance: InstanceRecord) -> None: await self.store.save_instance(quarantined) instance.state = 'QUARANTINED' instance.failure = str(exc) + self._report_instance_metrics() + tags = self._instance_metric_tags(instance) + self._metrics.quarantined_total.inc(tags={ + 'backend': tags['backend'], + 'mode': tags['mode'], + 'reason': self._quarantine_reason(exc), + }) + self._metrics.shutdown_seconds.observe(time.monotonic() - started, tags={**timing_tags, 'status': 'error'}) return + self._metrics.shutdown_seconds.observe(time.monotonic() - started, tags={**timing_tags, 'status': 'success'}) async with self._lock: terminal_records = [] terminal_at = time.time() @@ -627,6 +699,7 @@ async def _close_instance_once(self, instance: InstanceRecord) -> None: instance.stopped_at = terminal_at instance.training_ids.clear() self._failed_instances.pop(instance.instance_id, None) + self._report_instance_metrics() async def fail_instance(self, instance_id: str, failure: FutureFailureRecord) -> None: async with self._lock: @@ -641,6 +714,7 @@ async def fail_instance(self, instance_id: str, failure: FutureFailureRecord) -> await self.store.save_instance(draining) instance.state = 'DRAINING' instance.failure = failure.message + self._report_instance_metrics() try: async with self._lock: for training_id in instance.training_ids: @@ -724,6 +798,7 @@ async def sweep(self) -> None: and not instance.training_ids): self.instances.pop(identity) self._startup.pop(identity, None) + self._report_instance_metrics() async def _reconcile_operation(self, training: TrainingRecord, request_id: str) -> None: try: diff --git a/src/twinkle/server/training/driver.py b/src/twinkle/server/training/driver.py index 8a0a9b752..0341c548a 100644 --- a/src/twinkle/server/training/driver.py +++ b/src/twinkle/server/training/driver.py @@ -5,11 +5,14 @@ import asyncio import os import re +import time from twinkle.server.env_propagation import build_propagated_env_vars from twinkle.server.resources.allocation import DevicePool from twinkle.server.state.training_records import InstanceRecord +from twinkle.server.telemetry.metrics import get_training_metrics from twinkle.utils.logger import get_logger +from .catalog import TrainingCatalog logger = get_logger() @@ -20,6 +23,8 @@ def __init__(self, config, store, namespace, controller_name): self.config, self.store = config, store self.namespace, self.controller_name = namespace, controller_name self._cleanup_tasks: set[asyncio.Task] = set() + self._metrics = get_training_metrics() + self._catalog = TrainingCatalog(config) self.pool = DevicePool( namespace, allocation_timeout=config.allocation_timeout_s, shutdown_timeout=config.shutdown_timeout_s) @@ -27,14 +32,20 @@ def application_name(self, instance): return self.pool.name(instance.controller_epoch, instance.instance_id) async def allocate(self, instance, spec): - allocation = await self.pool.allocate(spec.config.num_devices, spec.model.device_type, - instance.controller_epoch, instance.instance_id) - return { - 'application_name': allocation.name, - 'placement_group_id': allocation.placement_group_id, - 'node_id': allocation.node_id, - 'devices': list(allocation.devices) - } + started, status = time.monotonic(), 'error' + tags = {'backend': spec.model.backend, 'mode': spec.mode, 'device_type': spec.model.device_type} + try: + allocation = await self.pool.allocate(spec.config.num_devices, spec.model.device_type, + instance.controller_epoch, instance.instance_id) + status = 'success' + return { + 'application_name': allocation.name, + 'placement_group_id': allocation.placement_group_id, + 'node_id': allocation.node_id, + 'devices': list(allocation.devices) + } + finally: + self._metrics.allocation_seconds.observe(time.monotonic() - started, tags={**tags, 'status': status}) async def deploy(self, instance, spec): from ray import serve @@ -149,7 +160,21 @@ async def shutdown(self, instance: InstanceRecord): actors.append(handle) ray.kill(handle, no_restart=True) await asyncio.to_thread(serve.delete, instance.application_name or name) - await self.pool.remove(name) + removal_started, removal_status = time.monotonic(), 'error' + try: + await self.pool.remove(name) + removal_status = 'success' + finally: + try: + spec = self._catalog.get(instance.base_model, instance.mode) + except Exception: + tags = {'backend': 'unknown', 'mode': instance.mode, 'device_type': 'unknown'} + else: + tags = {'backend': spec.model.backend, 'mode': spec.mode, 'device_type': spec.model.device_type} + self._metrics.removal_seconds.observe( + time.monotonic() - removal_started, tags={ + **tags, 'status': removal_status + }) # The PG kills even unnamed remaining children. Named workers must also # have failed readiness RPCs before the controller retires quota/tickets. for handle in actors: diff --git a/tests/server/integration/on_demand_training/__init__.py b/tests/server/integration/on_demand_training/__init__.py new file mode 100644 index 000000000..0bb7387ac --- /dev/null +++ b/tests/server/integration/on_demand_training/__init__.py @@ -0,0 +1 @@ +"""On-demand training integration probes.""" diff --git a/tests/server/integration/on_demand_training/acceptance_support.py b/tests/server/integration/on_demand_training/acceptance_support.py new file mode 100644 index 000000000..b24b5d9f0 --- /dev/null +++ b/tests/server/integration/on_demand_training/acceptance_support.py @@ -0,0 +1,163 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Real SDK operations and bounded resource assertions for acceptance probes.""" +import json +import time + +import requests + +from tests.server.integration.on_demand_training.on_demand_sdk import unload_tinker_model + + +class TrainingHandle: + """A small fixed-token workload with optimizer checkpoint support.""" + + def __init__(self, base_url, model, protocol, label, mode='lora'): + from twinkle_client import init_tinker_client + from twinkle_client.http import ClientTransport + from twinkle_client.http.context import ClientContext + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + + self.base_url, self.model, self.protocol = base_url, model, protocol + self.label, self.mode = label, mode + self.token = 'acceptance-' + protocol + self.transport = None + self.closed = False + self.consumed_samples = 0 + if protocol == 'twinkle': + from peft import LoraConfig + response = requests.post(base_url + '/api/v1/twinkle/create_session', json={}, timeout=30) + response.raise_for_status() + self.transport = ClientTransport(ClientContext( + base_url=base_url, api_key=self.token, session_id=response.json()['session_id'])) + self.client = MultiLoraTransformersModel(model, transport=self.transport) + self.client.add_adapter_to_model( + label, LoraConfig(r=8, target_modules='all-linear') if mode == 'lora' else None, + gradient_accumulation_steps=2) + self.identity = self.client._training_ids[label] + self.configure() + else: + init_tinker_client() + from tinker import ServiceClient + self.service = ServiceClient(base_url=base_url, api_key=self.token) + self.client = (self.service.create_lora_training_client(base_model=model, rank=8) + if mode == 'lora' else self.service.create_full_training_client(base_model=model)) + self.identity = self.client.model_id + self.inputs = [dict(input_ids=[100, 200, 300, 400], labels=[200, 300, 400, 500], + position_ids=list(range(4)), attention_mask=[1] * 4) for _ in range(2)] + + def configure(self): + if self.protocol == 'twinkle': + # Qwen3.5 is multimodal: this hook connects the visual branch even + # for text-only batches, as in the established SFT helpers. + self.client.set_template('Qwen3_5Template', max_length=512) + self.client.set_loss('CrossEntropyLoss') + self.client.set_optimizer('Adam', lr=1e-4) + self.client.set_processor('InputProcessor', padding_side='right') + + def datums(self): + from twinkle.server.model.tinker_datum import input_feature_to_datum + return [input_feature_to_datum(feature) for feature in self.inputs] + + def step(self): + if self.protocol == 'twinkle': + for _ in range(3): + self.client.forward_backward(self.inputs, gradient_accumulation_steps=2) + self.client.clip_grad_and_step(gradient_accumulation_steps=2) + self.consumed_samples += len(self.inputs) + else: + from tinker import types + self.client.forward_backward(self.datums(), 'cross_entropy').result() + self.client.optim_step(types.AdamParams(learning_rate=1e-4)).result() + self.consumed_samples += len(self.inputs) + + def evaluate(self): + if self.protocol == 'twinkle': + result = self.client.forward_only(self.inputs).result + values = result['logps'] + else: + result = self.client.forward(self.datums(), 'cross_entropy').result() + values = [output['logprobs'].tolist() for output in result.loss_fn_outputs] + import numpy as np + values = np.asarray(values, dtype=float).reshape(-1) + assert values.size and np.isfinite(values).all(), 'Non-finite checkpoint probe' + return values + + def save(self, name): + if self.protocol == 'twinkle': + path = self.client.save( + name=name, save_optimizer=True, consumed_train_samples=self.consumed_samples).twinkle_path + else: + path = self.client.save_state(name).result().path + assert path, 'Checkpoint path is empty' + return path + + def resume(self, path): + if self.protocol == 'twinkle': + progress = self.client.resume_from_checkpoint(path) + self.consumed_samples = progress['consumed_train_samples'] + return progress + old = self.identity + self.client = self.service.create_training_client_from_state_with_optimizer(path=path) + self.identity = self.client.model_id + unload_tinker_model(self.service, old) + + def assert_stale(self): + try: + if self.protocol == 'twinkle': + self.client.get_train_configs().result + else: + self.client.get_info() + except Exception as exc: + assert any(term in str(exc).lower() for term in ('404', 'not found', 'handle', 'active', 'unavailable')), exc + else: + raise AssertionError('An old training handle remained usable') + + def close(self, stale=False): + if self.closed: + return + try: + if not stale: + if self.protocol == 'twinkle': + self.client.remove_adapter() + else: + unload_tinker_model(self.service, self.identity) + finally: + if self.transport is not None: + self.transport.close() + elif hasattr(self, 'service'): + self.service.holder.close() + self.closed = True + + +def wait_until(predicate, timeout, label): + deadline = time.monotonic() + timeout + while True: + result = predicate() + if result: + return result + if time.monotonic() >= deadline: + raise TimeoutError(label) + time.sleep(0.5) + + +def assert_released(controller, namespace, redis_url=None, prefix=None): + import ray + from ray import serve + from ray.util.placement_group import placement_group_table + diagnostic = ray.get(controller.diagnostics.remote(), timeout=30) + assert all(row['state'] == 'STOPPED' for row in diagnostic['instances']), diagnostic + owned_prefix = 'resources-' + namespace + '-' + assert not [name for name in serve.status().applications if name.startswith(owned_prefix)] + assert not [row for row in placement_group_table().values() + if row.get('name', '').startswith(owned_prefix) and row.get('state') != 'REMOVED'] + if redis_url: + import redis + client = redis.Redis.from_url(redis_url, decode_responses=True, socket_timeout=5) + try: + assert not list(client.scan_iter(match=prefix + 'model::training_*')) + assert not list(client.scan_iter(match=prefix + 'seq_ledger::training_*')) + for key in client.scan_iter(match=prefix + 'token_count::*'): + assert int(json.loads(client.get(key)) or 0) == 0, 'Quota drift' + finally: + client.close() + return True diff --git a/tools/validation/on_demand_training/allocation_probe.py b/tests/server/integration/on_demand_training/allocation_probe.py similarity index 83% rename from tools/validation/on_demand_training/allocation_probe.py rename to tests/server/integration/on_demand_training/allocation_probe.py index 3b21dc0bd..d3cc874ec 100644 --- a/tools/validation/on_demand_training/allocation_probe.py +++ b/tests/server/integration/on_demand_training/allocation_probe.py @@ -2,6 +2,7 @@ """Run in its own process: test production allocation against logical Ray GPUs.""" import asyncio import json +import os import ray import uuid @@ -13,9 +14,18 @@ async def main(): + # These are synthetic Ray GPU reservations, not CUDA workloads. + os.environ.pop('CUDA_VISIBLE_DEVICES', None) namespace = 'alloc_' + uuid.uuid4().hex ray.init( - address='local', namespace=namespace, num_cpus=16, num_gpus=8, include_dashboard=False, log_to_driver=False) + address='local', + namespace=namespace, + num_cpus=4, + num_gpus=8, + _temp_dir='/tmp/twa-' + namespace[-8:], + include_dashboard=False, + dashboard_port=0, + log_to_driver=False) pool = DevicePool(namespace, allocation_timeout=3, shutdown_timeout=10) allocations = [] try: @@ -40,7 +50,6 @@ async def main(): devices = [device for allocation in allocations for device in allocation.devices] assert len(set(devices)) == 8, devices assert len({allocation.node_id for allocation in allocations}) == 1 - assert ray.available_resources().get('GPU', 0) == 0 try: await pool.allocate(spec.config.num_devices, spec.model.device_type, 'epoch', 'overflow') raise AssertionError('over-allocation was admitted') @@ -48,12 +57,13 @@ async def main(): pass from ray.util.placement_group import get_placement_group group = DeviceGroup(name='owned', ranks=2, device_type='cuda') - view = ResourceManager.from_allocation(group, get_placement_group(allocations[0].name), - list(allocations[0].devices)) + owned_pg = get_placement_group(allocations[0].name) + view = ResourceManager.from_allocation(group, owned_pg, list(allocations[0].devices)) assert view.get_config('owned').ranks == [0, 1] assert [entry['gpu_rank'][0] for entry in view.get_group('owned')] == list(allocations[0].devices) + # A ResourceManager view cannot release a controller-owned allocation. view.destroy_placement_group() - assert ray.available_resources().get('GPU', 0) == 0 + assert get_placement_group(allocations[0].name).id == owned_pg.id freed = allocations[0] await pool.remove(freed.name) two = TrainingCatalog( diff --git a/tools/validation/on_demand_training/fragmentation_probe.py b/tests/server/integration/on_demand_training/fragmentation_probe.py similarity index 94% rename from tools/validation/on_demand_training/fragmentation_probe.py rename to tests/server/integration/on_demand_training/fragmentation_probe.py index 267a539a2..f88132ea3 100644 --- a/tools/validation/on_demand_training/fragmentation_probe.py +++ b/tests/server/integration/on_demand_training/fragmentation_probe.py @@ -2,6 +2,7 @@ """Exercise single-node admission with real Ray nodes and logical devices.""" import asyncio import json +import os import ray import uuid from ray.cluster_utils import Cluster @@ -13,12 +14,14 @@ async def main(): + # Synthetic reservations must not inherit a hidden CUDA-device mask. + os.environ.pop('CUDA_VISIBLE_DEVICES', None) cluster = Cluster() occupied = [] allocated = None pool = None try: - cluster.add_node(num_cpus=4, num_gpus=4, resources={'left': 1}, include_dashboard=False) + cluster.add_node(num_cpus=4, num_gpus=4, resources={'left': 1}, include_dashboard=False, dashboard_port=0) cluster.add_node(num_cpus=4, num_gpus=4, resources={'right': 1}) ray.init(address=cluster.address, namespace='fragmentation_' + uuid.uuid4().hex, log_to_driver=False) pool = DevicePool(ray.get_runtime_context().namespace, allocation_timeout=1, shutdown_timeout=10) diff --git a/tools/validation/on_demand_training/http_probe.py b/tests/server/integration/on_demand_training/http_probe.py similarity index 97% rename from tools/validation/on_demand_training/http_probe.py rename to tests/server/integration/on_demand_training/http_probe.py index 5f4e765c0..ebeda6cc8 100644 --- a/tools/validation/on_demand_training/http_probe.py +++ b/tests/server/integration/on_demand_training/http_probe.py @@ -29,7 +29,13 @@ def main(): os.environ['TWINKLE_LONG_POLL_TIMEOUT'] = '2' os.environ['TWINKLE_PERSISTENCE_MODE'] = 'memory' os.environ['PYTHONPATH'] = str(__import__('pathlib').Path.cwd() / 'src') - ray.init(address='local', namespace=namespace, num_cpus=16, include_dashboard=False, log_to_driver=False) + ray.init( + address='local', + namespace=namespace, + num_cpus=8, + _temp_dir='/tmp/twh-' + namespace[-8:], + include_dashboard=False, + log_to_driver=False) config = ServerConfig.model_validate({ 'ray_namespace': namespace, @@ -38,8 +44,8 @@ def main(): 'port': port }, 'training': { - 'startup_timeout_s': 120, - 'shutdown_timeout_s': 30 + 'startup_timeout_s': 600, + 'shutdown_timeout_s': 180 }, 'applications': [{ 'name': 'small', @@ -236,7 +242,7 @@ def wait_epoch(controller, old_epoch): transport.close() finally: if getattr(launcher, '_training_controller', None) is not None: - ray.get(launcher._training_controller.shutdown.remote(), timeout=60) + ray.get(launcher._training_controller.shutdown.remote(), timeout=180) ray.kill(launcher._training_controller, no_restart=True) serve.shutdown() ray.shutdown() diff --git a/tests/server/integration/on_demand_training/managed_redis.py b/tests/server/integration/on_demand_training/managed_redis.py new file mode 100644 index 000000000..c832bbaee --- /dev/null +++ b/tests/server/integration/on_demand_training/managed_redis.py @@ -0,0 +1,65 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""An owned Redis subprocess for safe, durable restart fault injection.""" +import socket +import subprocess +import time +from pathlib import Path + + +class ManagedRedis: + """Never stop an externally supplied Redis server.""" + + def __init__(self, directory): + self.directory = Path(directory).resolve() + self.directory.mkdir(parents=True, exist_ok=True) + with socket.socket() as listener: + listener.bind(('127.0.0.1', 0)) + self.port = listener.getsockname()[1] + self.url = f'redis://127.0.0.1:{self.port}/0' + self.process = None + self.log = None + + def start(self): + import redis + assert self.process is None + self.log = (self.directory / 'redis.log').open('a') + self.process = subprocess.Popen([ + 'redis-server', '--bind', '127.0.0.1', '--port', str(self.port), + '--dir', str(self.directory), '--appendonly', 'yes', '--appendfsync', 'always', + '--save', '', '--daemonize', 'no'], stdout=self.log, stderr=subprocess.STDOUT) + client = redis.Redis.from_url(self.url, socket_timeout=1, socket_connect_timeout=1) + try: + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + if self.process.poll() is not None: + raise RuntimeError('Owned test Redis exited during startup') + try: + if client.ping(): + return + except redis.exceptions.ConnectionError: + pass + time.sleep(0.1) + raise TimeoutError('Owned test Redis did not start') + finally: + client.close() + + def stop(self): + if self.process is not None: + self.process.terminate() + try: + self.process.wait(timeout=30) + except subprocess.TimeoutExpired: + self.process.kill() + self.process.wait(timeout=5) + raise + finally: + self.log.close() + self.process = None + self.log = None + + def fault(self, probe): + self.stop() + try: + probe() + finally: + self.start() diff --git a/tests/server/integration/on_demand_training/matrix_validation.py b/tests/server/integration/on_demand_training/matrix_validation.py new file mode 100644 index 000000000..35a30ee6d --- /dev/null +++ b/tests/server/integration/on_demand_training/matrix_validation.py @@ -0,0 +1,441 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Run the full training matrix on an isolated on-demand server.""" +import argparse +import atexit +import json +import math +import os +import ray +import socket +import subprocess +import sys +import time +import uuid +import yaml +from pathlib import Path +from ray import serve +from xml.etree import ElementTree + +_ROOT = Path(__file__).resolve().parents[4] +if str(_ROOT) not in sys.path: + sys.path.insert(0, str(_ROOT)) + +from twinkle.server.config import ServerConfig +from twinkle.server.launcher.server_launcher import ServerLauncher + + +def validate_controller_restart(base_url, controller, model_name): + """Train, restart, restore optimizer checkpoints and recover a failed worker.""" + import numpy as np + from ray.util import list_named_actors + from tests.server.integration.on_demand_training.acceptance_support import ( + TrainingHandle, assert_released, wait_until) + + namespace = ray.get_runtime_context().namespace + old = [TrainingHandle(base_url, model_name, protocol, 'restart-' + protocol) + for protocol in ('twinkle', 'tinker')] + snapshots = [] + try: + for handle in old: + handle.step() + snapshots.append((handle.save('restart-state'), handle.evaluate(), handle.consumed_samples)) + old_epoch = ray.get(controller.diagnostics.remote())['epoch'] + ray.kill(controller, no_restart=False) + + def restarted(): + try: + diagnostic = ray.get(controller.diagnostics.remote(), timeout=10) + return diagnostic if diagnostic['epoch'] != old_epoch else None + except ray.exceptions.RayError: + return None + + diagnostic = wait_until(restarted, 600, 'Controller restart timed out') + assert_released(controller, namespace) + for handle in old: + handle.assert_stale() + restored = [] + try: + # Keep restored peers on one runtime: checkpoint/optimizer isolation + # is checked without reloading the base model for each protocol. + for protocol, (path, values, consumed) in zip(('twinkle', 'tinker'), snapshots): + fresh = TrainingHandle(base_url, model_name, protocol, 'restored-' + protocol) + restored.append(fresh) + progress = fresh.resume(path) + if protocol == 'twinkle': + assert progress['consumed_train_samples'] == consumed + np.testing.assert_allclose(fresh.evaluate(), values, rtol=0, atol=2e-3) + fresh.step() + active = next(row for row in ray.get(controller.diagnostics.remote())['instances'] + if restored[0].identity in row['training_ids']) + assert set(active['training_ids']) == {handle.identity for handle in restored} + owned = [row for row in list_named_actors(all_namespaces=True) + if row['namespace'] == namespace and active['instance_id'] in row['name'] + and row['name'].startswith('model-') and not row['name'].endswith('-registry')] + assert owned, 'No model worker available for fault injection' + ray.kill(ray.get_actor(owned[0]['name'], namespace=namespace), no_restart=True) + wait_until(lambda: all(row['state'] == 'STOPPED' + for row in ray.get(controller.diagnostics.remote())['instances']), + 600, 'Worker failure cleanup timed out') + for handle in restored: + handle.assert_stale() + assert_released(controller, namespace) + finally: + for handle in restored: + handle.close(stale=True) + fresh = TrainingHandle(base_url, model_name, 'twinkle', 'after-worker-fault') + try: + fresh.step() + finally: + fresh.close() + finally: + for handle in old: + handle.close(stale=True) + return {'epoch': diagnostic['epoch'], 'worker_fault_recovered': True, 'checkpoint_restored': True} + + +def validate_idempotent_create(base_url, controller, model_name): + """Concurrent HTTP retries share one receipt, handle and allocation.""" + from concurrent.futures import ThreadPoolExecutor + import requests + from peft import LoraConfig + from twinkle.protocol.serialize import serialize_object + from twinkle.protocol.types.model import TrainingHandleResponse + from twinkle_client._future import resolve_response + from twinkle_client.http import ClientTransport + from twinkle_client.http.context import ClientContext + from tests.server.integration.on_demand_training.acceptance_support import assert_released + + session_reply = requests.post(base_url + '/api/v1/twinkle/create_session', json={}, timeout=30) + session_reply.raise_for_status() + session = session_reply.json()['session_id'] + transport = ClientTransport(ClientContext(base_url=base_url, api_key='idempotency', session_id=session)) + headers = {'Authorization': 'Bearer idempotency', 'X-Twinkle-Session-Id': session} + payload = {'client_create_id': uuid.uuid4().hex, 'adapter_name': 'duplicate', + 'config': serialize_object(LoraConfig(r=8, target_modules='all-linear')), + 'gradient_accumulation_steps': 2} + endpoint = base_url + '/api/v1/model/' + model_name + '/twinkle/add_adapter_to_model' + + def submit(_): + response = requests.post(endpoint, json=payload, + headers={**headers, 'x-request-id': uuid.uuid4().hex}, timeout=60) + response.raise_for_status() + return response + + try: + with ThreadPoolExecutor(max_workers=8) as executor: + replies = list(executor.map(submit, range(8))) + assert len({reply.json()['request_id'] for reply in replies}) == 1 + handle = resolve_response(replies[0], TrainingHandleResponse, transport=transport) + live = [row for row in ray.get(controller.diagnostics.remote())['instances'] if row['state'] != 'STOPPED'] + assert len(live) == 1 and live[0]['training_ids'] == [handle.training_id] + reply = requests.post(base_url + '/api/v1/model/' + model_name + '/twinkle/remove_adapter', + json={'training_id': handle.training_id, 'adapter_name': 'duplicate'}, + headers={**headers, 'x-request-id': uuid.uuid4().hex}, timeout=60) + reply.raise_for_status() + resolve_response(reply, None, transport=transport) + assert_released(controller, ray.get_runtime_context().namespace) + finally: + transport.close() + return True + + +def validate_capacity(base_url, controller, model_name, capacity): + """Assert shared LoRA capacity and full-training exclusivity on real devices.""" + from tests.server.integration.on_demand_training.acceptance_support import TrainingHandle, assert_released + handles = [] + try: + for index in range(capacity): + handles.append(TrainingHandle(base_url, model_name, ('twinkle', 'tinker')[index % 2], + 'capacity-' + str(index))) + live = [row for row in ray.get(controller.diagnostics.remote())['instances'] if row['state'] != 'STOPPED'] + assert len(live) == 1 and len(live[0]['training_ids']) == capacity + try: + overflow = TrainingHandle(base_url, model_name, 'tinker', 'overflow') + except Exception as exc: + assert 'replica limit' in str(exc).lower(), exc + else: + overflow.close() + raise AssertionError('LoRA capacity was exceeded') + handles[0].step() + finally: + for handle in reversed(handles): + handle.close() + assert_released(controller, ray.get_runtime_context().namespace) + exclusive = TrainingHandle(base_url, model_name, 'twinkle', 'exclusive', mode='full') + try: + exclusive.step() + live = [row for row in ray.get(controller.diagnostics.remote())['instances'] if row['state'] != 'STOPPED'] + assert len(live) == 1 and live[0]['mode'] == 'full' and len(live[0]['training_ids']) == 1 + try: + overflow = TrainingHandle(base_url, model_name, 'tinker', 'exclusive-overflow', mode='full') + except Exception as exc: + assert 'replica limit' in str(exc).lower(), exc + else: + overflow.close() + raise AssertionError('Full-training instance was shared') + finally: + exclusive.close() + assert_released(controller, ray.get_runtime_context().namespace) + return True + + +def validate_redis_restart(base_url, controller, model_name, managed): + """Only interrupt a Redis process started and owned by this harness.""" + import requests + from tests.server.integration.on_demand_training.acceptance_support import TrainingHandle + handle = TrainingHandle(base_url, model_name, 'twinkle', 'redis-fault') + try: + handle.step() + before = ray.get(controller.diagnostics.remote()) + + def unavailable(): + response = requests.post( + base_url + '/api/v1/model/' + model_name + '/twinkle/get_train_configs', + json={'training_id': handle.identity, 'adapter_name': handle.label}, + headers={'Authorization': 'Bearer ' + handle.token, 'x-request-id': uuid.uuid4().hex, + 'X-Twinkle-Session-Id': handle.transport.context.session_id}, timeout=15) + assert response.status_code >= 500, response.text + + managed.fault(unavailable) + after = ray.get(controller.diagnostics.remote(), timeout=30) + assert after['epoch'] == before['epoch'] + assert {row['instance_id'] for row in before['instances']} == { + row['instance_id'] for row in after['instances']} + handle.step() + finally: + handle.close() + return True + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--backend', required=True, choices=['transformers', 'megatron']) + parser.add_argument('--model', required=True) + parser.add_argument('--run-dir', required=True) + parser.add_argument('--persistence', choices=['memory', 'redis'], default='memory') + parser.add_argument('--managed-redis', action='store_true', help='Start an owned Redis for restart injection') + parser.add_argument('--redis-url', default=os.environ.get('TWINKLE_TEST_REDIS_URL')) + parser.add_argument('--startup-timeout-s', type=float, default=1800) + parser.add_argument('--shutdown-timeout-s', type=float, default=180) + parser.add_argument('--recovery-rounds', type=int, default=1) + parser.add_argument('--skip-full-cycle', action='store_true') + parser.add_argument('--recovery-only', action='store_true') + parser.add_argument('--full-cycle-only', action='store_true') + parser.add_argument('--stress-iterations', type=int, default=0) + parser.add_argument('--soak-seconds', type=float, default=0) + args = parser.parse_args() + if args.persistence == 'redis' and not args.redis_url and not args.managed_redis: + parser.error('--redis-url or --managed-redis is required for Redis persistence') + if args.managed_redis and args.persistence != 'redis': + parser.error('--managed-redis requires --persistence redis') + if any(not math.isfinite(value) or value <= 0 + for value in (args.startup_timeout_s, args.shutdown_timeout_s)): + parser.error('startup and shutdown timeouts must be finite and positive') + if args.recovery_rounds < 0 or args.stress_iterations < 0 or args.soak_seconds < 0: + parser.error('validation budgets must be non-negative') + root = Path(__file__).resolve().parents[4] + run_dir = Path(args.run_dir) + run_dir.mkdir(parents=True, exist_ok=True) + (run_dir / 'checkpoints').mkdir(parents=True, exist_ok=True) + managed = None + if args.managed_redis: + from tests.server.integration.on_demand_training.managed_redis import ManagedRedis + managed = ManagedRedis(run_dir / 'redis') + atexit.register(managed.stop) + managed.start() + args.redis_url = managed.url + namespace = 'matrix_' + uuid.uuid4().hex[:8] + with socket.socket() as listener: + listener.bind(('127.0.0.1', 0)) + port = listener.getsockname()[1] + source = root / 'tests/server/config' / ('server_config_4b_e2e.yaml' if args.backend == 'transformers' else + 'server_config_4b_e2e_megatron.yaml') + raw = yaml.safe_load(source.read_text()) + raw.setdefault('training', {}).update({ + 'startup_timeout_s': args.startup_timeout_s, + 'shutdown_timeout_s': args.shutdown_timeout_s, + }) + redis_prefix = f'on-demand-test::{namespace}::' + raw['persistence'] = ({ + 'mode': 'redis', + 'redis_url': args.redis_url, + 'key_prefix': redis_prefix, + } if args.persistence == 'redis' else {'mode': 'memory'}) + raw['ray_namespace'] = namespace + raw['http_options'] = {'host': '127.0.0.1', 'port': port} + for app in raw['applications']: + params = app['args'] + if app['import_path'] == 'model': + common = { + 'num_devices': 2, 'device_mesh': {'device_type': 'cuda', 'dp_size': 2}, + 'min_replicas': 0, 'max_replicas': 1, 'idle_timeout_s': 3600, + 'max_length': params.get('max_length'), + 'queue_config': {**params.get('queue_config', {}), 'execution_timeout': 600}, + } + app['args'] = { + 'model_id': args.model, 'backend': args.backend, 'device_type': 'cuda', + 'modes': {'lora': {**common, 'max_loras': 10, 'max_r': 32}, 'full': common}, + 'env_vars': {'TWINKLE_FAIL_FAST': '0', 'TWINKLE_TRUST_REMOTE_CODE': '1'}, + } + app.pop('deployments', None) + elif app['import_path'] == 'sampler': + params['model_id'] = args.model + params['engine_args']['attention_config'] = {'flash_attn_version': 3} + if args.recovery_only: + raw['applications'] = [app for app in raw['applications'] if app['import_path'] in ('server', 'model')] + config = ServerConfig.model_validate(raw) + evidence_config = config.to_yaml_dict() + if args.persistence == 'redis': + evidence_config['persistence']['redis_url'] = '' + (run_dir / 'server.yaml').write_text(yaml.safe_dump(evidence_config, sort_keys=False)) + os.environ.update({ + 'TWINKLE_STATE_RAY_NAMESPACE': namespace, + 'TWINKLE_PERSISTENCE_MODE': args.persistence, + 'TWINKLE_DEFAULT_SAVE_DIR': str(run_dir / 'checkpoints'), + 'TWINKLE_MODEL_ID_ALIASES': json.dumps({'Qwen/Qwen3.5-4B': args.model}), + 'TWINKLE_TRUST_REMOTE_CODE': '1', + 'TWINKLE_FAIL_FAST': '0', + 'TWINKLE_TEST_GPU_E2E': '1', + 'TWINKLE_TEST_BACKEND': args.backend, + 'TWINKLE_SERVER_URL': f'http://127.0.0.1:{port}', + 'TWINKLE_TEST_OPERATION_TIMEOUT': '600', + 'PYTHONPATH': os.pathsep.join((str(root / 'src'), str(root))), + }) + if args.persistence == 'redis': + os.environ.update({ + 'TWINKLE_PERSISTENCE_REDIS_URL': args.redis_url, + 'TWINKLE_PERSISTENCE_KEY_PREFIX': redis_prefix, + }) + else: + os.environ.pop('TWINKLE_PERSISTENCE_REDIS_URL', None) + os.environ.pop('TWINKLE_PERSISTENCE_KEY_PREFIX', None) + evidence = { + 'backend': args.backend, + 'persistence': args.persistence, + 'namespace': namespace, + 'cases': [], + 'lazy_start': False, + } + ray.init( + address='local', + namespace=namespace, + num_cpus=24, + num_gpus=2 if args.recovery_only else 3, + _temp_dir='/tmp/twm-' + namespace[-8:], + include_dashboard=False) + launcher = ServerLauncher(config) + try: + launcher._start_serve() + for app in config.applications: + launcher._deploy_application(app) + launcher._deploy_training() + assert ray.get(launcher._training_controller.diagnostics.remote())['instances'] == [] + evidence['lazy_start'] = True + filenames = ['test_sft_e2e.py', 'test_dpo_e2e.py', 'test_grpo_e2e.py'] + if not args.skip_full_cycle: + filenames.append('test_full_cycle_e2e.py') + if args.recovery_only: + filenames = [] + elif args.full_cycle_only: + filenames = ['test_full_cycle_e2e.py'] + for filename in filenames: + start = time.monotonic() + with (run_dir / (filename + '.log')).open('w') as output: + result = subprocess.run([ + sys.executable, + '-m', + 'pytest', + '--noconftest', # HTTP clients do not need a second local Ray fixture. + str(root / 'tests/server/integration' / filename), + '-vv', + '-s', + '--junitxml=' + str(run_dir / (filename + '.xml')), + ], + stdout=output, + stderr=subprocess.STDOUT, + env=os.environ.copy(), + timeout=7200) + junit = run_dir / (filename + '.xml') + cases = ElementTree.parse(junit).getroot().findall('.//testcase') if junit.exists() else [] + skipped = sum(case.find('skipped') is not None for case in cases) + evidence['cases'].append({ + 'file': filename, + 'returncode': result.returncode, + 'tests': len(cases), + 'skipped': skipped, + 'seconds': time.monotonic() - start + }) + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + assert result.returncode == 0, f'{filename} failed; see {run_dir / (filename + ".log")}' + assert len(cases) == (1 if filename == 'test_full_cycle_e2e.py' else 2) and skipped == 0 + base_url = f'http://127.0.0.1:{port}' + model_name = config.training_config.models[0].name + evidence['recovery'] = [] + for _ in range(args.recovery_rounds): + recovery = validate_controller_restart(base_url, launcher._training_controller, model_name) + evidence['recovery'].append(recovery) + evidence['restart_epoch'] = recovery['epoch'] + evidence['worker_fault_recovered'] = recovery['worker_fault_recovered'] + if managed is not None: + evidence['redis_fault_rounds'] = evidence.get('redis_fault_rounds', 0) + int( + validate_redis_restart(base_url, launcher._training_controller, model_name, managed)) + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + if args.recovery_rounds: + evidence['idempotency_checked'] = validate_idempotent_create( + base_url, launcher._training_controller, model_name) + evidence['capacity_checked'] = validate_capacity( + base_url, launcher._training_controller, model_name, + config.training_config.models[0].modes.lora.max_loras) + if args.stress_iterations or args.soak_seconds: + from tests.server.integration.on_demand_training.acceptance_support import assert_released + controller = launcher._training_controller + assert_released(controller, namespace, + args.redis_url if args.persistence == 'redis' else None, redis_prefix) + from tests.server.integration.on_demand_training.stress_validation import run_stress + evidence['stress'] = [] + budgets = ([('short-' + str(index), args.stress_iterations, 0) for index in range(3)] + if args.stress_iterations else []) + if args.soak_seconds: + budgets.append(('soak', 0, args.soak_seconds)) + for name, iterations, duration in budgets: + evidence['stress'].append(run_stress( + f'http://127.0.0.1:{port}', config.training_config.models[0].name, + concurrency=2, iterations=iterations, duration=duration, + controller=launcher._training_controller, namespace=namespace, + redis_url=args.redis_url if args.persistence == 'redis' else None, prefix=redis_prefix, + progress_path=run_dir / (name + '.json'))) + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + finally: + try: + controller = getattr(launcher, '_training_controller', None) + if controller is not None: + ray.get(controller.shutdown.remote(), timeout=args.shutdown_timeout_s) + diagnostic = ray.get(controller.diagnostics.remote(), timeout=30) + assert all(row['state'] == 'STOPPED' for row in diagnostic['instances']) + from tests.server.integration.on_demand_training.acceptance_support import assert_released + assert_released(controller, namespace, + args.redis_url if args.persistence == 'redis' else None, redis_prefix) + evidence['training_resources_released'] = True + launcher._shutdown() + finally: + serve.shutdown() + ray.shutdown() + if args.persistence == 'redis' and evidence.get('training_resources_released'): + import redis + client = redis.Redis.from_url(args.redis_url) + try: + keys = list(client.scan_iter(match=redis_prefix + '*', count=500)) + if keys: + client.delete(*keys) + finally: + client.close() + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + if managed is not None: + managed.stop() + evidence['passed'] = True + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/tools/validation/on_demand_training/on_demand_sdk.py b/tests/server/integration/on_demand_training/on_demand_sdk.py similarity index 100% rename from tools/validation/on_demand_training/on_demand_sdk.py rename to tests/server/integration/on_demand_training/on_demand_sdk.py diff --git a/tests/server/integration/on_demand_training/stress_validation.py b/tests/server/integration/on_demand_training/stress_validation.py new file mode 100644 index 000000000..3dc2a3156 --- /dev/null +++ b/tests/server/integration/on_demand_training/stress_validation.py @@ -0,0 +1,102 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Concurrent and soak validation against an already running single-node server.""" +import argparse +import concurrent.futures +import json +import requests +import sys +import time +from pathlib import Path + +_ROOT = Path(__file__).resolve().parents[4] +if str(_ROOT) not in sys.path: + sys.path.insert(0, str(_ROOT)) + +from tests.server.integration.on_demand_training.acceptance_support import TrainingHandle, assert_released + + +def run_stress(base_url, model, concurrency=2, iterations=100, duration=0, controller=None, + namespace=None, redis_url=None, prefix=None, progress_path=None, protocols=('twinkle', 'tinker')): + """Run actual training; duration mode is never truncated by an iteration limit.""" + started, completed = time.monotonic(), 0 + evidence = {'completed': 0, 'failures': [], 'passed': False, 'duration_target': duration} + + def cycle(protocol, index): + handle = TrainingHandle(base_url, model, protocol, 'stress-' + str(index)) + try: + handle.step() + if (index // len(protocols)) % 10 == 0: + path = handle.save('soak-' + str(index)) + handle.resume(path) + handle.step() + finally: + handle.close() + + # Keep a peer alive so cycles test slot reuse instead of spending hours + # repeatedly loading the same base model. Its final release must free the PG. + anchor = TrainingHandle(base_url, model, protocols[0], 'soak-anchor') + try: + anchor.step() + # Measure sustained traffic time, not model cold-start/JIT time. + started = time.monotonic() + with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as executor: + while (time.monotonic() - started < duration) if duration else (completed < iterations): + count = concurrency if duration else min(concurrency, iterations - completed) + jobs = [executor.submit(cycle, protocols[(completed + slot) % len(protocols)], completed + slot) + for slot in range(count)] + for job in jobs: + job.result() # Do not convert incomplete cycles into successful evidence. + completed += count + # Session heartbeats do not renew the business-idle lease. + # Keep the warm peer active during a multi-hour soak. + if completed % 10 == 0: + anchor.step() + if controller is not None: + import ray + live = [row for row in ray.get(controller.diagnostics.remote(), timeout=30)['instances'] + if row['state'] != 'STOPPED'] + assert len(live) == 1 and live[0]['state'] == 'READY' + assert live[0]['training_ids'] == [anchor.identity], live + assert len(live[0]['devices']) == 2 + evidence.update(completed=completed, seconds=time.monotonic() - started) + if progress_path: + Path(progress_path).write_text(json.dumps(evidence, indent=2)) + print(json.dumps({'phase': 'stress', **evidence}), flush=True) + finally: + anchor.close() + if controller is not None: + assert_released(controller, namespace, redis_url, prefix) + evidence['passed'] = True + if duration: + assert evidence['seconds'] >= duration + if progress_path: + Path(progress_path).write_text(json.dumps(evidence, indent=2)) + return evidence + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--backend', required=True, choices=['transformers', 'megatron']) + parser.add_argument('--client', choices=['twinkle', 'tinker', 'both'], default='both') + parser.add_argument('--server-url', required=True) + parser.add_argument('--model', required=True) + parser.add_argument('--concurrency', type=int, default=2) + parser.add_argument('--iterations', type=int, default=100) + parser.add_argument('--duration', type=float, default=0, help='Run for this many seconds instead of fixed iterations') + parser.add_argument('--results', required=True) + args = parser.parse_args() + if args.concurrency <= 0 or args.iterations <= 0 or args.duration < 0: + parser.error('concurrency and iterations must be positive; duration must be non-negative') + + protocols = ['twinkle', 'tinker'] if args.client == 'both' else [args.client] + evidence = run_stress(args.server_url, args.model, args.concurrency, args.iterations, args.duration, + progress_path=args.results, protocols=protocols) + response = requests.get(args.server_url + '/api/v1/model/' + args.model + '/healthz', timeout=30) + response.raise_for_status() + assert response.json()['healthy'] and not response.json()['instances'] + print(json.dumps(evidence), flush=True) + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/tests/server/integration/test_full_cycle_e2e.py b/tests/server/integration/test_full_cycle_e2e.py index 6fa31f006..12d17de9c 100644 --- a/tests/server/integration/test_full_cycle_e2e.py +++ b/tests/server/integration/test_full_cycle_e2e.py @@ -6,8 +6,7 @@ Backend selection via env var TWINKLE_TEST_BACKEND: - "transformers" (default): all 6 phases run strictly - - "megatron": Phase C/D skipped (known multi-LoRA strict-load bug), - Phase E/F best-effort (GPU OOM possible) + - "megatron": all 6 phases run strictly (requires a working optimizer-resume backend) Phase A — initial training (STEPS_PHASE_A steps) Phase B — keep training STEPS_PHASE_B more steps, save again @@ -58,9 +57,9 @@ assert BACKEND in ('transformers', 'megatron'), f'Invalid TWINKLE_TEST_BACKEND={BACKEND!r}' BASE_MODEL = 'Qwen/Qwen3.5-4B' -BASE_URL = 'http://localhost:9000' +BASE_URL = os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:9000') API_KEY = 'EMPTY_API_KEY' -SAVE_DIR = '/mnt/nas2/yunlin.myl/twinkle/output/twinkle_e2e_full_cycle' +SAVE_DIR = os.environ.get('TWINKLE_DEFAULT_SAVE_DIR', 'output/twinkle_e2e_full_cycle') STEPS_PHASE_A = 60 STEPS_PHASE_B = 4 STEPS_PHASE_D = 4 diff --git a/tests/server/integration/test_on_demand_gpu_redis_e2e.py b/tests/server/integration/test_on_demand_gpu_redis_e2e.py new file mode 100644 index 000000000..69fa201a6 --- /dev/null +++ b/tests/server/integration/test_on_demand_gpu_redis_e2e.py @@ -0,0 +1,53 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Opt-in single-node GPU matrix using the real Redis state backend.""" +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.skipif( + os.environ.get('TWINKLE_TEST_GPU_E2E') != '1' or not os.environ.get('TWINKLE_TEST_REDIS_URL'), + reason='Set TWINKLE_TEST_GPU_E2E=1 and TWINKLE_TEST_REDIS_URL for the GPU Redis matrix', +) + + +@pytest.mark.parametrize('backend', ['transformers', 'megatron']) +def test_on_demand_training_gpu_redis_matrix(tmp_path, backend): + model = os.environ.get('TWINKLE_TEST_MODEL') or os.environ.get('TWINKLE_DEVICE_MODEL') + if not model: + pytest.skip('Set TWINKLE_TEST_MODEL or TWINKLE_DEVICE_MODEL') + root = Path(__file__).resolve().parents[3] + run_dir = tmp_path / backend + script = Path(__file__).with_name('on_demand_training') / 'matrix_validation.py' + result = subprocess.run([ + sys.executable, + str(script), + '--backend', + backend, + '--model', + model, + '--run-dir', + str(run_dir), + '--persistence', + 'redis', + '--startup-timeout-s', + '1800', + '--shutdown-timeout-s', + '180', + ], + cwd=root, + capture_output=True, + text=True, + timeout=4 * 60 * 60) + assert result.returncode == 0, result.stdout + result.stderr + evidence = json.loads((run_dir / 'results.json').read_text()) + assert evidence['passed'] + assert evidence['persistence'] == 'redis' + assert evidence['lazy_start'] and evidence['training_resources_released'] + assert evidence['restart_epoch'] and evidence['worker_fault_recovered'] + assert all(row['checkpoint_restored'] for row in evidence['recovery']) + assert evidence['capacity_checked'] and evidence['idempotency_checked'] + assert all(case['returncode'] == 0 and not case['skipped'] for case in evidence['cases']) diff --git a/tests/server/telemetry/test_training_metrics.py b/tests/server/telemetry/test_training_metrics.py new file mode 100644 index 000000000..ad446cfeb --- /dev/null +++ b/tests/server/telemetry/test_training_metrics.py @@ -0,0 +1,152 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Managed training metrics stay low-cardinality and survive registry resets.""" +from collections import defaultdict +from types import SimpleNamespace + +import pytest + +from twinkle.server.telemetry import metrics +from twinkle.server.training.controller import TrainingController + +_ALLOWED_ATTRIBUTES = {'backend', 'mode', 'state', 'status', 'reason', 'device_type', 'resident'} + + +@pytest.fixture(autouse=True) +def reset_registry(): + yield + metrics.MetricsRegistry.reset() + + +class FakeInstrument: + + def __init__(self, name): + self.name = name + self.calls = [] + + def add(self, value, attributes=None): + self.calls.append((value, attributes or {})) + + def record(self, value, attributes=None): + self.calls.append((value, attributes or {})) + + +class FakeMeter: + + def __init__(self): + self.instruments = {} + + def _create(self, name, **kwargs): + instrument = FakeInstrument(name) + self.instruments[name] = instrument + return instrument + + create_counter = _create + create_histogram = _create + create_up_down_counter = _create + + def create_observable_gauge(self, name, **kwargs): + return self._create(name, **kwargs) + + +def build_registry(monkeypatch): + meter = FakeMeter() + monkeypatch.setattr(metrics, 'get_meter', lambda name: meter) + metrics.MetricsRegistry.reset() + return metrics.MetricsRegistry.get(), meter + + +def test_training_and_redis_instruments_are_registered(monkeypatch): + registry, meter = build_registry(monkeypatch) + training = registry.training_metrics() + state = registry.state_backend_metrics() + + training.instances_active.set(2, tags={'backend': 'transformers', 'mode': 'lora', 'state': 'READY', + 'resident': 'false'}) + training.startup_seconds.observe(1.5, tags={'backend': 'transformers', 'mode': 'lora', 'status': 'success'}) + training.quarantined_total.inc(tags={'backend': 'megatron', 'mode': 'full', 'reason': 'timeout'}) + state.redis_health.set(1) + state.update_atomic_conflicts.inc() + state.update_atomic_retries.observe(2) + + expected = { + 'twinkle.training.instances.active', + 'twinkle.training.instance.startup.duration_seconds', + 'twinkle.training.instance.shutdown.duration_seconds', + 'twinkle.training.placement_group.allocation.duration_seconds', + 'twinkle.training.placement_group.removal.duration_seconds', + 'twinkle.training.quarantined.total', + 'twinkle.training.sweep.total', + 'twinkle.training.sweep.duration_seconds', + 'twinkle.training.devices.allocated', + 'twinkle.training.placement_groups.active', + 'twinkle.state.redis.health', + 'twinkle.state.redis.update_atomic.conflicts.total', + 'twinkle.state.redis.update_atomic.retries', + } + assert expected <= meter.instruments.keys() + for instrument in meter.instruments.values(): + for _, attributes in instrument.calls: + assert set(attributes) <= _ALLOWED_ATTRIBUTES + assert not {'token', 'session_id', 'training_id', 'instance_id', 'request_id'} & set(attributes) + + +def test_training_metric_adapters_rebind_after_registry_reset(monkeypatch): + first, _ = build_registry(monkeypatch) + first_metrics = first.training_metrics() + metrics.MetricsRegistry.reset() + second_metrics = metrics.MetricsRegistry.get().training_metrics() + assert first_metrics is not second_metrics + assert first_metrics.startup_seconds._instrument is not second_metrics.startup_seconds._instrument + + +@pytest.mark.parametrize('backend', ['transformers', 'megatron']) +def test_controller_metrics_clear_old_states_and_released_resources(monkeypatch, backend): + registry, meter = build_registry(monkeypatch) + spec = SimpleNamespace(model=SimpleNamespace(backend=backend, device_type='cuda')) + controller = TrainingController.__new__(TrainingController) + controller.catalog = SimpleNamespace(get=lambda base_model, mode: spec) + controller._metrics = registry.training_metrics() + controller._reported_metric_keys = set() + controller._reported_resource_metric_keys = set() + controller.instances = { + str(index): SimpleNamespace( + instance_id=f'instance_{index}', base_model='private-model', mode='lora', + state='READY', resident=False, devices=[index * 2, index * 2 + 1], + application_name=f'private-app-{index}', training_ids=[f'training_{index}']) + for index in range(2) + } + + def totals(name): + values = defaultdict(float) + for delta, attributes in meter.instruments[name].calls: + assert set(attributes) <= _ALLOWED_ATTRIBUTES + assert not {'instance_id', 'base_model', 'application_name', 'training_id'} & set(attributes) + values[tuple(sorted(attributes.items()))] += delta + return values + + instances_name = 'twinkle.training.instances.active' + devices_name = 'twinkle.training.devices.allocated' + groups_name = 'twinkle.training.placement_groups.active' + controller._report_instance_metrics() + first_calls = {name: len(meter.instruments[name].calls) + for name in (instances_name, devices_name, groups_name)} + controller._report_instance_metrics() + assert first_calls == {name: len(meter.instruments[name].calls) for name in first_calls} + assert sum(totals(instances_name).values()) == 2 + assert sum(totals(devices_name).values()) == 4 + assert sum(totals(groups_name).values()) == 2 + + controller.instances['0'].state = 'DRAINING' + controller._report_instance_metrics() + state_totals = {dict(key)['state']: value for key, value in totals(instances_name).items()} + assert state_totals == {'READY': 1, 'DRAINING': 1} + assert sum(totals(devices_name).values()) == 4 + assert sum(totals(groups_name).values()) == 2 + + for instance in controller.instances.values(): + instance.state = 'STOPPED' + controller._report_instance_metrics() + for name in (instances_name, devices_name, groups_name): + assert all(value == 0 for value in totals(name).values()) + assert not controller._reported_metric_keys + assert not controller._reported_resource_metric_keys diff --git a/tests/server/training/test_acceptance_harness.py b/tests/server/training/test_acceptance_harness.py new file mode 100644 index 000000000..92b8b972c --- /dev/null +++ b/tests/server/training/test_acceptance_harness.py @@ -0,0 +1,84 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Acceptance harness budgets and Redis ownership are independently testable.""" +import shutil + +import pytest + +from tests.server.integration.on_demand_training import stress_validation +from tests.server.integration.on_demand_training.managed_redis import ManagedRedis + + +@pytest.fixture(scope='session', autouse=True) +def _ray_runtime(): + # These tests deliberately use no Ray runtime. + yield + + +class FakeHandle: + calls = [] + + def __init__(self, base, model, protocol, label): + self.calls.append(('create', protocol)) + + def step(self): + self.calls.append(('step', None)) + + def save(self, name): + return 'checkpoint' + + def resume(self, path): + assert path == 'checkpoint' + + def close(self): + self.calls.append(('close', None)) + + +def test_iteration_budget_does_not_round_up(monkeypatch): + FakeHandle.calls = [] + monkeypatch.setattr(stress_validation, 'TrainingHandle', FakeHandle) + evidence = stress_validation.run_stress('unused', 'unused', concurrency=2, iterations=3) + assert evidence['completed'] == 3 + assert len([call for call in FakeHandle.calls if call[0] == 'close']) == 4 # Includes the warm peer. + + +def test_duration_is_not_truncated_by_iterations(monkeypatch): + now = [0] + + def clock(): + now[0] += 1 + return now[0] + + monkeypatch.setattr(stress_validation.time, 'monotonic', clock) + monkeypatch.setattr(stress_validation, 'TrainingHandle', FakeHandle) + evidence = stress_validation.run_stress('unused', 'unused', concurrency=1, iterations=1, duration=10) + assert evidence['completed'] > 1 and evidence['seconds'] >= 10 + + +def test_failed_training_never_reports_success(monkeypatch): + class Broken(FakeHandle): + def step(self): + raise RuntimeError('injected failure') + + monkeypatch.setattr(stress_validation, 'TrainingHandle', Broken) + with pytest.raises(RuntimeError, match='injected failure'): + stress_validation.run_stress('unused', 'unused', iterations=1) + + +@pytest.mark.skipif(not shutil.which('redis-server'), reason='redis-server executable required') +def test_owned_redis_restart_preserves_state(tmp_path): + redis = pytest.importorskip('redis') + managed = ManagedRedis(tmp_path) + managed.start() + client = redis.Redis.from_url(managed.url, socket_timeout=1) + try: + client.set('checkpoint', 'durable') + + def unavailable(): + with pytest.raises(redis.exceptions.ConnectionError): + client.get('checkpoint') + + managed.fault(unavailable) + assert client.get('checkpoint') == b'durable' + finally: + client.close() + managed.stop() diff --git a/tests/server/training/test_allocation.py b/tests/server/training/test_allocation.py index 22d4f6860..519347b8c 100644 --- a/tests/server/training/test_allocation.py +++ b/tests/server/training/test_allocation.py @@ -5,7 +5,7 @@ def test_production_allocator_uses_disjoint_devices_and_reuses_released_group(): - result = subprocess.run([sys.executable, str(Path(__file__).resolve().parents[3] / 'tools/validation/on_demand_training/allocation_probe.py')], + result = subprocess.run([sys.executable, str(Path(__file__).resolve().parents[3] / 'tests/server/integration/on_demand_training/allocation_probe.py')], capture_output=True, text=True, timeout=120) assert result.returncode == 0, result.stdout + result.stderr evidence = json.loads(result.stdout.strip().splitlines()[-1]) @@ -14,7 +14,7 @@ def test_production_allocator_uses_disjoint_devices_and_reuses_released_group(): def test_fragmented_devices_cannot_be_combined_across_nodes(): - result = subprocess.run([sys.executable, str(Path(__file__).resolve().parents[3] / 'tools/validation/on_demand_training/fragmentation_probe.py')], + result = subprocess.run([sys.executable, str(Path(__file__).resolve().parents[3] / 'tests/server/integration/on_demand_training/fragmentation_probe.py')], capture_output=True, text=True, timeout=120) assert result.returncode == 0, result.stdout + result.stderr evidence = json.loads(result.stdout.strip().splitlines()[-1]) diff --git a/tests/server/training/test_http_integration.py b/tests/server/training/test_http_integration.py index 367555749..b20f14b29 100644 --- a/tests/server/training/test_http_integration.py +++ b/tests/server/training/test_http_integration.py @@ -1,12 +1,28 @@ import json import subprocess import sys +import uuid from pathlib import Path def test_installed_native_and_tinker_sdks_use_managed_http_routes(): - result = subprocess.run([sys.executable, str(Path(__file__).resolve().parents[3] / 'tools/validation/on_demand_training/http_probe.py')], - capture_output=True, text=True, timeout=420) + # This probe owns a separate Ray cluster. Stop the session fixture's cluster + # first so two local GCS/raylet pairs do not evict each other on a slow host. + import ray + ray.shutdown() + try: + result = subprocess.run([ + sys.executable, + str(Path(__file__).resolve().parents[3] + / 'tests/server/integration/on_demand_training/http_probe.py') + ], capture_output=True, text=True, timeout=1800) + finally: + ray.init( + ignore_reinit_error=True, + num_cpus=2, + log_to_driver=False, + namespace='twinkle_test', + _temp_dir='/tmp/twinkle-test-' + uuid.uuid4().hex[:8]) assert result.returncode == 0, result.stdout + result.stderr evidence = next(json.loads(line) for line in result.stdout.splitlines() if line.startswith('{"passed":')) assert evidence['passed'] and evidence['native_recreation'] and evidence['tinker_shared'] diff --git a/tests/server/training/test_redis_restart_reads.py b/tests/server/training/test_redis_restart_reads.py new file mode 100644 index 000000000..bfde03295 --- /dev/null +++ b/tests/server/training/test_redis_restart_reads.py @@ -0,0 +1,114 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Read-only Redis reconnects are bounded; uncertain writes are never replayed.""" +import shutil +from unittest.mock import AsyncMock, Mock + +import pytest + +redis = pytest.importorskip('redis') +ConnectionError = redis.exceptions.ConnectionError + +from twinkle.server.state.backend.redis_backend import RedisBackend +from tests.server.integration.on_demand_training.managed_redis import ManagedRedis + + +@pytest.fixture(scope='session', autouse=True) +def _ray_runtime(): + yield # These backend tests need no Ray cluster. + + +@pytest.mark.asyncio +async def test_stale_read_is_retried_but_write_is_not(): + backend = RedisBackend('redis://127.0.0.1:1/0') + try: + backend._client.get = AsyncMock(side_effect=[ConnectionError('stale socket'), '{"value": 1}']) + assert await backend.get('owned') == {'value': 1} + assert backend._client.get.await_count == 2 + backend._client.set = AsyncMock(side_effect=ConnectionError('lost acknowledgement')) + with pytest.raises(ConnectionError): + await backend.set('owned', {'value': 2}) + assert backend._client.set.await_count == 1 + finally: + await backend.close() + + +@pytest.mark.asyncio +async def test_persistent_outage_does_not_look_like_missing_state(): + backend = RedisBackend('redis://127.0.0.1:1/0') + try: + backend._client.get = AsyncMock(side_effect=ConnectionError('offline')) + with pytest.raises(ConnectionError): + await backend.get('controller-active') + assert backend._client.get.await_count == 3 + finally: + await backend.close() + + +@pytest.mark.asyncio +@pytest.mark.skipif(not shutil.which('redis-server'), reason='redis-server executable required') +async def test_real_restart_recovers_old_pooled_readers(tmp_path): + managed = ManagedRedis(tmp_path) + managed.start() + first, stale, atomic = (RedisBackend(managed.url) for _ in range(3)) + try: + await first.set('owner', {'epoch': 'original'}) + assert await stale.get('owner') == {'epoch': 'original'} + assert await atomic.get('owner') == {'epoch': 'original'} + managed.stop() + with pytest.raises(ConnectionError): + await first.get('owner') + managed.start() + assert await stale.get('owner') == {'epoch': 'original'} + assert await stale.exists('owner') + assert await first.mget(['owner', 'missing']) == [{'epoch': 'original'}, None] + assert await stale.health_check() + assert await atomic.update_atomic('counter', lambda value: int(value or 0) + 1) == 1 + assert await atomic.get('counter') == 1 + finally: + await first.close() + await stale.close() + await atomic.close() + managed.stop() + + +class PipelineStub: + def __init__(self): + self.watch = AsyncMock() + self.get = AsyncMock(return_value='0') + self.multi = Mock() + self.set = Mock() + self.execute = AsyncMock(return_value=[]) + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + +@pytest.mark.asyncio +async def test_transport_watch_error_after_exec_is_never_replayed(): + backend = RedisBackend('redis://127.0.0.1:1/0') + pipe = PipelineStub() + pipe.execute.side_effect = redis.exceptions.WatchError( + 'A ConnectionError occurred while watching one or more keys') + backend._client.pipeline = lambda: pipe + try: + with pytest.raises(ConnectionError, match='outcome could not be confirmed'): + await backend.update_atomic('counter', lambda value: value + 1) + assert pipe.execute.await_count == 1 + finally: + await backend.close() + + +@pytest.mark.asyncio +async def test_connection_loss_before_exec_is_safe_to_retry(): + backend = RedisBackend('redis://127.0.0.1:1/0') + pipe = PipelineStub() + pipe.watch.side_effect = [ConnectionError('stale'), None] + backend._client.pipeline = lambda: pipe + try: + assert await backend.update_atomic('counter', lambda value: value + 1) == 1 + assert pipe.watch.await_count == 2 and pipe.execute.await_count == 1 + finally: + await backend.close() diff --git a/tests/server/training/test_runtime_queue.py b/tests/server/training/test_runtime_queue.py index e489f3782..57f02083f 100644 --- a/tests/server/training/test_runtime_queue.py +++ b/tests/server/training/test_runtime_queue.py @@ -71,6 +71,7 @@ async def operation(): await queue.shutdown_task_queue() runtime.fail_shutdown = False await controller.sweep() + await controller.wait_for_jobs() assert not runtime.allocations diff --git a/tests/transformers/test_optimizer_group_regression.py b/tests/transformers/test_optimizer_group_regression.py new file mode 100644 index 000000000..23922c3a1 --- /dev/null +++ b/tests/transformers/test_optimizer_group_regression.py @@ -0,0 +1,48 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Optimizer groups must not depend on DDP wrapping or repeated GA kwargs.""" +from types import SimpleNamespace +from unittest.mock import Mock + +import torch + +from twinkle.model.optimizer_group import BaseOptimizerGroup, TrainStatus +from twinkle.model.transformers.transformers import TransformersModel + + +def test_forward_kwargs_do_not_duplicate_metric_control_parameters(): + metric = Mock() + kwargs = {'gradient_accumulation_steps': 2, 'lr': 99, 'step': 99, 'grad_norm': 99, 'loss_reduction': 'sum'} + status = TrainStatus(inputs=[{}], outputs={'loss': 1.0}, metrics=[metric], forward_kwargs=kwargs) + group = BaseOptimizerGroup(train_status=status, cur_step=3, gradient_accumulation_steps=2) + group.accumulate_metrics(True) + metric.accumulate.assert_called_once_with( + status.inputs, status.outputs, gradient_accumulation_steps=2, lr=[], step=2, grad_norm=0.0, + loss_reduction='sum') + assert kwargs['lr'] == 99 + + +def test_parameter_groups_and_optimizer_restore_are_stable_after_wrapping(): + model = torch.nn.Sequential(torch.nn.Linear(4, 4), torch.nn.LayerNorm(4)) + + class Wrapper(torch.nn.Module): + def __init__(self, module): + super().__init__() + self.module = module + + harness = SimpleNamespace( + model=model, + strategy=SimpleNamespace(unwrap_model=lambda wrapped: getattr(wrapped, 'module', wrapped)), + _get_trainable_parameters=lambda adapter: dict(model.named_parameters()), + ) + original_groups = TransformersModel._create_param_group(harness, '', weight_decay=0.1) + optimizer = torch.optim.Adam(original_groups, lr=1e-4) + model(torch.ones(2, 4)).sum().backward() + optimizer.step() + state = optimizer.state_dict() + harness.model = Wrapper(model) + late_groups = TransformersModel._create_param_group(harness, '', weight_decay=0.1) + assert [g['param_names'] for g in late_groups] == [g['param_names'] for g in original_groups] + assert original_groups[0]['param_names'] == ['0.weight'] + restored = torch.optim.Adam(late_groups, lr=1e-4) + restored.load_state_dict(state) + assert len(restored.state) == len(optimizer.state) diff --git a/tools/validation/on_demand_training/device_validation.py b/tools/validation/on_demand_training/device_validation.py deleted file mode 100644 index 01695893f..000000000 --- a/tools/validation/on_demand_training/device_validation.py +++ /dev/null @@ -1,327 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Explicit real-device acceptance; run this file only after local mock checks. - -Required: TWINKLE_DEVICE_BACKEND, TWINKLE_DEVICE_MODEL, TWINKLE_DEVICE_OTHER_MODEL, -TWINKLE_DEVICE_RUN_DIR and four exclusively available visible devices. -Tests cover DP resource lifecycle without checkpoint I/O. -This is intentionally outside pytest's default test filename pattern. -""" -import json -import numpy as np -import os -import pytest -import ray -import requests -import socket -import time -import uuid -from on_demand_sdk import unload_tinker_model -from pathlib import Path -from ray import serve - -from twinkle.server.config import ServerConfig -from twinkle.server.launcher.server_launcher import ServerLauncher -from twinkle_client import init_tinker_client -from twinkle_client.http import ClientTransport -from twinkle_client.http.context import ClientContext -from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel - - -@pytest.fixture(scope='module') -def server(): - backend = os.environ['TWINKLE_DEVICE_BACKEND'] - assert backend in ('transformers', 'megatron') - run_dir = Path(os.environ['TWINKLE_DEVICE_RUN_DIR']) - run_dir.mkdir(parents=True, exist_ok=True) - namespace = 'device_' + uuid.uuid4().hex[:8] - with socket.socket() as listener: - listener.bind(('127.0.0.1', 0)) - port = listener.getsockname()[1] - os.environ['TWINKLE_STATE_RAY_NAMESPACE'] = namespace - os.environ['TWINKLE_DEFAULT_SAVE_DIR'] = str(run_dir / 'checkpoints') - os.environ['TWINKLE_PERSISTENCE_MODE'] = 'memory' - os.environ['TWINKLE_LONG_POLL_TIMEOUT'] = '2' - mesh = {'dp_size': 2} - apps = [] - for alias, variable in [('small', 'TWINKLE_DEVICE_MODEL'), ('other', 'TWINKLE_DEVICE_OTHER_MODEL')]: - source = os.environ[variable] - assert Path(source, 'config.json').is_file() - mode = {'num_devices': 2, 'device_mesh': mesh, 'idle_timeout_s': 600 if alias == 'small' else 30} - if backend == 'transformers': - mode['strategy'] = 'accelerate' - apps.append({ - 'name': alias, - 'route_prefix': '/api/v1/model/' + alias, - 'import_path': 'model', - 'args': { - 'model_id': source, - 'backend': backend, - 'device_type': 'cuda', - 'modes': { - 'lora': { - **mode, 'max_loras': 2, - 'max_r': 8 - }, - 'full': mode - }, - } - }) - apps.append({ - 'name': 'gateway', - 'route_prefix': '/api/v1', - 'import_path': 'server', - 'args': { - 'server_config': { - 'per_token_model_limit': 10 - } - }, - 'deployments': [{ - 'ray_actor_options': { - 'num_cpus': 0.1 - } - }] - }) - os.environ['TWINKLE_MODEL_ID_ALIASES'] = json.dumps( - {app['name']: app['args']['model_id'] - for app in apps if app['import_path'] == 'model'}) - config = ServerConfig(ray_namespace=namespace, http_options={'host': '127.0.0.1', 'port': port}, applications=apps) - import yaml - (run_dir / 'server.yaml').write_text(yaml.safe_dump(config.to_yaml_dict(), sort_keys=False)) - ray.init( - address='local', - namespace=namespace, - num_cpus=16, - num_gpus=4, - _temp_dir='/tmp/twd-' + namespace[-8:], - include_dashboard=False, - log_to_driver=True) - launcher = ServerLauncher(config) - launcher._start_serve() - evidence = { - 'backend': backend, - 'namespace': namespace, - 'visible_devices': os.environ.get('CUDA_VISIBLE_DEVICES'), - 'cases': [] - } - try: - for app in config.applications: - launcher._deploy_application(app) - launcher._deploy_training() - controller = launcher._training_controller - base = f'http://127.0.0.1:{port}' - assert ray.get(controller.diagnostics.remote())['instances'] == [] - evidence['ray_address'] = ray.get_runtime_context().gcs_address - yield base, controller, evidence, launcher - finally: - if getattr(launcher, '_training_controller', None) is not None: - ray.get(launcher._training_controller.shutdown.remote(), timeout=180) - ray.kill(launcher._training_controller, no_restart=True) - (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) - serve.shutdown() - ray.shutdown() - - -class Handle: - - def __init__(self, base, protocol, mode, alias='small'): - self.base, self.protocol = base, protocol - self.token = 'device-validation-owner' - session = requests.post(base + '/api/v1/twinkle/create_session', json={}, timeout=20).json()['session_id'] - self.transport = ClientTransport(ClientContext(base_url=base, api_key=self.token, session_id=session)) - self.rate = 1e-3 if mode == 'lora' else 5e-5 - if protocol == 'native': - from peft import LoraConfig - self.client = MultiLoraTransformersModel(alias, transport=self.transport) - self.client.add_adapter_to_model( - 'same', - LoraConfig(r=8, target_modules='all-linear') if mode == 'lora' else None, - gradient_accumulation_steps=2) - self.identity = self.client._training_ids['same'] - self.client.set_loss('CrossEntropyLoss') - self.client.set_optimizer('Adam', lr=self.rate) - self.client.set_processor('InputProcessor', padding_side='right') - else: - init_tinker_client() - from tinker import ServiceClient - self.service = ServiceClient(base_url=base, api_key=self.token) - self.client = ( - self.service.create_lora_training_client(base_model=alias, rank=8) - if mode == 'lora' else self.service.create_full_training_client(base_model=alias)) - self.identity = self.client.model_id - # Fixed token IDs need no external dataset, and every position has a target. - self.inputs = [{ - 'input_ids': [100, 200, 300, 400, 500, 600, 700, 800], - 'labels': [200, 300, 400, 500, 600, 700, 800, 900], - 'position_ids': list(range(8)), - 'attention_mask': [1] * 8 - }] * 2 - - def datum(self): - from tinker import types - feature = self.inputs[0] - return [ - types.Datum( - model_input=types.ModelInput.from_ints(feature['input_ids']), - loss_fn_inputs={ - 'target_tokens': types.TensorData(data=feature['labels'], dtype='int64', shape=[8]), - 'weights': types.TensorData(data=[1.] * 8, dtype='float32', shape=[8]), - }) for _ in self.inputs - ] - - def evaluate(self): - if self.protocol == 'native': - result = self.client.forward_only(self.inputs).result - values = result['logps'] - else: - result = self.client.forward(self.datum(), 'cross_entropy').result() - values = result.loss_fn_outputs[0]['logprobs'].tolist() - values = np.asarray(values, dtype=float).reshape(-1) - assert values.size and np.isfinite(values).all(), values - return values - - def step(self): - if self.protocol == 'native': - for _ in range(2): - self.client.forward_backward(self.inputs, gradient_accumulation_steps=2) - self.client.clip_grad_and_step(gradient_accumulation_steps=2) - else: - from tinker import types - self.client.forward_backward(self.datum(), 'cross_entropy').result() - self.client.optim_step(types.AdamParams(learning_rate=self.rate)).result() - - def close(self): - if self.protocol == 'native': - self.client.remove_adapter() - else: - assert unload_tinker_model(self.service, self.identity).model_id == self.identity - self.transport.close() - - -def active_instances(controller): - return [row for row in ray.get(controller.diagnostics.remote())['instances'] if row['state'] != 'STOPPED'] - - -@pytest.fixture(autouse=True) -def cleanup_failed_case(server): - from twinkle.server.state.models import FutureFailureRecord - _, controller, _, _ = server - yield - # Release this case's resources even when a numeric assertion fails, so later combinations start clean. - for instance in active_instances(controller): - ray.get( - controller.fail_instance.remote( - instance['instance_id'], - FutureFailureRecord( - reason_code='orphaned_replica', attribution='server', message='Device test teardown')), - timeout=180) - assert not active_instances(controller) - assert ray.available_resources().get('GPU', 0) == 4 - - -@pytest.mark.parametrize('protocol', ['native', 'tinker']) -@pytest.mark.parametrize('mode', ['lora', 'full']) -def test_dp_resource_lifecycle(server, protocol, mode): - base, controller, evidence, _ = server - assert not active_instances(controller) - assert ray.available_resources().get('GPU', 0) == 4 - first, second = Handle(base, protocol, mode), Handle(base, protocol, mode) - instances = active_instances(controller) - assert len(instances) == (1 if mode == 'lora' else 2) - assert all(len(row['devices']) == 2 for row in instances) - first_instance_ids = {row['instance_id'] for row in instances} - baseline, untouched = first.evaluate(), second.evaluate() - for _ in range(4): - first.step() - trained = first.evaluate() - assert float(trained.mean()) > float(baseline.mean()) + 1e-4 - np.testing.assert_allclose(second.evaluate(), untouched, rtol=0, atol=2e-3) - # A single LoRA update may be invisible in BF16 outputs, so both tenants use the same training window. - for _ in range(4): - second.step() - peer_trained = second.evaluate() - assert float(peer_trained.mean()) > float(untouched.mean()) + 1e-4 - np.testing.assert_allclose(first.evaluate(), trained, rtol=0, atol=2e-3) - if mode == 'lora': - third, fourth = Handle(base, protocol, mode), Handle(base, protocol, mode) - third.evaluate() - instances = active_instances(controller) - assert len(instances) == 2 - assert ray.available_resources().get('GPU', 0) == 0 - with pytest.raises(Exception, match='(?i)allocation'): - Handle(base, protocol, mode) - allocated = [device for row in instances for device in row['devices']] - assert len(allocated) == 4 and len(set(allocated)) == 4 - if mode == 'lora': - fourth.close() - third.close() - first.close() - assert len(active_instances(controller)) == 1 - assert ray.available_resources().get('GPU', 0) == 2 - np.testing.assert_allclose(second.evaluate(), peer_trained, rtol=0, atol=2e-3) - for _ in range(4): - second.step() - peer_continued = second.evaluate() - assert float(peer_continued.mean()) > float(peer_trained.mean()) + 1e-4 - second.close() - assert not active_instances(controller) - assert ray.available_resources().get('GPU', 0) == 4 - - alternate = Handle(base, protocol, mode, alias='other') - alternate_instances = active_instances(controller) - assert len(alternate_instances) == 1 - reused = alternate_instances[0]['devices'] - assert len(reused) == 2 and set(reused).issubset(allocated) - other_initial = alternate.evaluate() - for _ in range(4): - alternate.step() - other_trained = alternate.evaluate() - assert float(other_trained.mean()) > float(other_initial.mean()) + 1e-4 - health = requests.get(base + '/api/v1/twinkle/healthz/deep', timeout=30) - assert health.status_code == 200 and health.json()['healthy'] - deadline = time.monotonic() + 90 - while active_instances(controller): - assert time.monotonic() < deadline, 'Business-idle instance did not release DP devices' - response = alternate.transport.post( - base + '/api/v1/twinkle/session_heartbeat', - json_data={'session_id': alternate.transport.context.session_id}) - response.raise_for_status() - time.sleep(2) - assert ray.available_resources().get('GPU', 0) == 4 - with pytest.raises(Exception) as expired: - alternate.evaluate() - alternate.transport.close() - - fresh = Handle(base, protocol, mode) - fresh_instances = active_instances(controller) - assert len(fresh_instances) == 1 - assert fresh_instances[0]['instance_id'] not in first_instance_ids - np.testing.assert_allclose(fresh.evaluate(), baseline, rtol=0, atol=2e-3) - fresh.close() - assert not active_instances(controller) - assert ray.available_resources().get('GPU', 0) == 4 - evidence['cases'].append({ - 'protocol': protocol, - 'mode': mode, - 'parallelism': { - 'dp_size': 2, - 'tp_size': 1, - 'fsdp_size': 1 - }, - 'initial_logps': baseline.tolist(), - 'trained_logps': trained.tolist(), - 'peer_initial_logps': untouched.tolist(), - 'peer_trained_logps': peer_trained.tolist(), - 'peer_continued_logps': peer_continued.tolist(), - 'other_initial_logps': other_initial.tolist(), - 'other_trained_logps': other_trained.tolist(), - 'original_devices': allocated, - 'cross_model_reused_devices': reused, - 'health': health.json(), - 'idle_expired_error': str(expired.value)[:400], - 'lazy_start': True, - 'shared_or_exclusive_isolation': True, - 'idle_reclaimed': True, - 'fresh_recreated': True, - 'passed': True - }) - (Path(os.environ['TWINKLE_DEVICE_RUN_DIR']) / 'progress.json').write_text(json.dumps(evidence, indent=2)) diff --git a/tools/validation/on_demand_training/matrix_validation.py b/tools/validation/on_demand_training/matrix_validation.py deleted file mode 100644 index 387da6371..000000000 --- a/tools/validation/on_demand_training/matrix_validation.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Run the full training matrix on an isolated server; runtimes start from zero on client demand.""" -import argparse -import json -import os -import ray -import socket -import subprocess -import sys -import time -import uuid -import yaml -from pathlib import Path -from ray import serve -from xml.etree import ElementTree - -from twinkle.server.config import ServerConfig -from twinkle.server.launcher.server_launcher import ServerLauncher - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument('--backend', required=True, choices=['transformers', 'megatron']) - parser.add_argument('--model', required=True) - parser.add_argument('--run-dir', required=True) - args = parser.parse_args() - root = Path(__file__).resolve().parents[3] - run_dir = Path(args.run_dir) - run_dir.mkdir(parents=True, exist_ok=True) - namespace = 'matrix_' + uuid.uuid4().hex[:8] - with socket.socket() as listener: - listener.bind(('127.0.0.1', 0)) - port = listener.getsockname()[1] - source = root / 'tests/server/config' / ('server_config_4b_e2e.yaml' if args.backend == 'transformers' else - 'server_config_4b_e2e_megatron.yaml') - raw = yaml.safe_load(source.read_text()) - raw.setdefault('training', {})['shutdown_timeout_s'] = 180 - raw['ray_namespace'] = namespace - raw['http_options'] = {'host': '127.0.0.1', 'port': port} - for app in raw['applications']: - params = app['args'] - if app['import_path'] == 'model': - params['model_id'] = args.model - params['nproc_per_node'] = 2 - params['device_group']['ranks'] = 2 - params['device_mesh'] = {'device_type': 'cuda', 'dp_size': 2} - params['min_replicas'] = 0 - params['max_replicas'] = 2 - params['idle_timeout_s'] = 3600 - params['adapter_config'] = {'adapter_timeout': 3600, 'adapter_max_lifetime': 36000} - params['env_vars'] = {'TWINKLE_FAIL_FAST': '0', 'TWINKLE_TRUST_REMOTE_CODE': '1'} - app.pop('deployments', None) - elif app['import_path'] == 'sampler': - params['model_id'] = args.model - params['engine_args']['attention_config'] = {'flash_attn_version': 3} - config = ServerConfig.model_validate(raw) - (run_dir / 'server.yaml').write_text(yaml.safe_dump(config.to_yaml_dict(), sort_keys=False)) - os.environ.update({ - 'TWINKLE_STATE_RAY_NAMESPACE': namespace, - 'TWINKLE_PERSISTENCE_MODE': 'memory', - 'TWINKLE_DEFAULT_SAVE_DIR': str(run_dir / 'checkpoints'), - 'TWINKLE_MODEL_ID_ALIASES': json.dumps({'Qwen/Qwen3.5-4B': args.model}), - 'TWINKLE_TRUST_REMOTE_CODE': '1', - 'TWINKLE_FAIL_FAST': '0', - 'TWINKLE_TEST_GPU_E2E': '1', - 'TWINKLE_TEST_BACKEND': args.backend, - 'TWINKLE_SERVER_URL': f'http://127.0.0.1:{port}', - 'TWINKLE_TEST_OPERATION_TIMEOUT': '600', - 'PYTHONPATH': str(root / 'src'), - }) - evidence = {'backend': args.backend, 'namespace': namespace, 'cases': [], 'lazy_start': False} - ray.init( - address='local', - namespace=namespace, - num_cpus=24, - num_gpus=3, - _temp_dir='/tmp/twm-' + namespace[-8:], - include_dashboard=False) - launcher = ServerLauncher(config) - try: - launcher._start_serve() - for app in config.applications: - launcher._deploy_application(app) - launcher._deploy_training() - assert ray.get(launcher._training_controller.diagnostics.remote())['instances'] == [] - evidence['lazy_start'] = True - for filename in ['test_sft_e2e.py', 'test_dpo_e2e.py', 'test_grpo_e2e.py', 'test_full_cycle_e2e.py']: - start = time.monotonic() - with (run_dir / (filename + '.log')).open('w') as output: - result = subprocess.run([ - sys.executable, - '-m', - 'pytest', - str(root / 'tests/server/integration' / filename), - '-vv', - '-s', - '--junitxml=' + str(run_dir / (filename + '.xml')), - ], - stdout=output, - stderr=subprocess.STDOUT, - env=os.environ.copy(), - timeout=7200) - junit = run_dir / (filename + '.xml') - cases = ElementTree.parse(junit).getroot().findall('.//testcase') if junit.exists() else [] - skipped = sum(case.find('skipped') is not None for case in cases) - evidence['cases'].append({ - 'file': filename, - 'returncode': result.returncode, - 'tests': len(cases), - 'skipped': skipped, - 'seconds': time.monotonic() - start - }) - (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) - assert result.returncode == 0, f'{filename} failed; see {run_dir / (filename + ".log")}' - assert len(cases) == (1 if filename == 'test_full_cycle_e2e.py' else 2) and skipped == 0 - finally: - try: - controller = getattr(launcher, '_training_controller', None) - if controller is not None: - ray.get(controller.shutdown.remote(), timeout=180) - diagnostic = ray.get(controller.diagnostics.remote(), timeout=30) - assert all(row['state'] == 'STOPPED' for row in diagnostic['instances']) - evidence['training_resources_released'] = True - launcher._shutdown() - finally: - serve.shutdown() - ray.shutdown() - (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) - evidence['passed'] = True - (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) - return 0 - - -if __name__ == '__main__': - raise SystemExit(main()) From 3c96aebb71e314854fd7c42723148912af244288 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Wed, 21 Oct 2026 02:48:46 +0000 Subject: [PATCH 07/46] docs(server): remove links to private training design --- docs/source_en/Usage Guide/Server and Client/Server.md | 2 +- .../\346\234\215\345\212\241\347\253\257.md" | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source_en/Usage Guide/Server and Client/Server.md b/docs/source_en/Usage Guide/Server and Client/Server.md index 4701f24f6..fe09347e9 100644 --- a/docs/source_en/Usage Guide/Server and Client/Server.md +++ b/docs/source_en/Usage Guide/Server and Client/Server.md @@ -484,4 +484,4 @@ holds the cleanup lease. Existing `applications` entries with `import_path: model` remain valid. Set `args.min_replicas: 0` for lazy startup and explicitly set `args.idle_timeout_s` for business inactivity. Legacy `adapter_config.adapter_timeout` remains a heartbeat timeout and `adapter_max_lifetime` retains its absolute lifetime. New `modes` default to 1800 seconds of business inactivity; heartbeats and polling do not renew that deadline; queued and running operations protect training state. Empty instances above the resident minimum are destroyed so other models can reuse their devices. Existing deployment min/max replicas become training instance minimums/limits. -To offer both LoRA and full training for one model, configure `num_devices`, `device_mesh`, `min_replicas`, and `idle_timeout_s` separately under `args.modes.lora/full`. LoRA users share compatible instances; each full-training handle owns an instance. Sampling configuration stays the same. See the [specification](../../../design/on-demand-training.md) and `cookbook/client/server/on_demand_training.yaml`. +To offer both LoRA and full training for one model, configure `num_devices`, `device_mesh`, `min_replicas`, and `idle_timeout_s` separately under `args.modes.lora/full`. LoRA users share compatible instances; each full-training handle owns an instance. Sampling configuration stays the same. See `cookbook/client/server/on_demand_training.yaml` in the repository. diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" index 8125400ca..cc7899561 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" @@ -471,4 +471,4 @@ twinkle-server check-config -c server_config.yaml 原有 `applications` 和 `import_path: model` 配置继续有效。单模式配置增加 `args.min_replicas: 0` 即可懒启动，`args.idle_timeout_s` 显式设置训练句柄的业务闲置时间。旧 `adapter_config.adapter_timeout` 仍按心跳判断会话存活，`adapter_max_lifetime` 保留绝对寿命；不再把旧字段转换成业务闲置时间。新增 `modes` 的业务闲置默认 1800 秒，心跳和轮询不延长该期限，运行中或已排队的任务保护训练状态。最后一个用户释放或过期后，超出保底数量的实例立即销毁，释放的卡可用于其他模型。已有 deployment 的 min/max_replicas 转成训练实例的保底/上限。 -同一模型需要同时提供 LoRA 和全参数训练时，在 `args.modes.lora/full` 内分别设置 `num_devices`、`device_mesh`、`min_replicas` 和 `idle_timeout_s`；不用改成另一套顶层模型目录。LoRA 用户共享兼容实例，全参数用户各占一个实例。采样服务继续使用原配置。详见 [完整规范](../../../design/on-demand-training.md) 与仓库 `cookbook/client/server/on_demand_training.yaml`。 +同一模型需要同时提供 LoRA 和全参数训练时，在 `args.modes.lora/full` 内分别设置 `num_devices`、`device_mesh`、`min_replicas` 和 `idle_timeout_s`；不用改成另一套顶层模型目录。LoRA 用户共享兼容实例，全参数用户各占一个实例。采样服务继续使用原配置。详见仓库 `cookbook/client/server/on_demand_training.yaml`。 From 3f2f4d94dd91fc2c315fcd88fa8a90d5cc433dd0 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Wed, 21 Oct 2026 08:45:13 +0000 Subject: [PATCH 08/46] fix(server): squash PR #325 into on-demand training Integrate 015732eb while preserving controller-owned training creation and sequence ledgers. Adapt full-mode regression tests to the catalog/controller path. Verified 434 targeted tests and the complete pre-commit suite. --- src/twinkle/loss/base.py | 8 + src/twinkle/loss/grpo.py | 91 +++++- .../server/checkpoint/checkpoint_manager.py | 11 +- src/twinkle/server/gateway/openai_handlers.py | 98 +++++- src/twinkle/server/gateway/use_cases.py | 2 + src/twinkle/server/lifecycle/submit.py | 3 +- src/twinkle/server/model/backends/common.py | 32 +- src/twinkle/server/model/tinker_handlers.py | 8 + src/twinkle/server/model/twinkle_handlers.py | 30 +- src/twinkle/server/task_queue/worker.py | 24 +- src/twinkle_client/_future.py | 1 + src/twinkle_client/exceptions.py | 2 + tests/loss/test_grpo_gkd.py | 71 +++++ tests/server/gateway/test_openai_handlers.py | 295 +++++++++++++++++- .../server/lifecycle/test_run_submit_dedup.py | 16 +- .../model/test_full_mode_exclusivity.py | 106 +++++++ .../model/test_preflight_loss_inputs.py | 103 ++++++ .../server/model/test_twinkle_async_inputs.py | 66 ++++ tests/server/utils/test_checkpoint_manager.py | 19 ++ tests/server/utils/test_task_queue_mixin.py | 42 ++- tests/twinkle_client/test_future_layer.py | 9 + 21 files changed, 1006 insertions(+), 31 deletions(-) create mode 100644 tests/server/model/test_full_mode_exclusivity.py create mode 100644 tests/server/model/test_preflight_loss_inputs.py diff --git a/src/twinkle/loss/base.py b/src/twinkle/loss/base.py index 945ed6353..7e8db18a4 100644 --- a/src/twinkle/loss/base.py +++ b/src/twinkle/loss/base.py @@ -12,6 +12,14 @@ class Loss: def __call__(self, inputs: InputFeature, outputs: ModelOutput, **kwargs) -> LossOutput: ... + def validate_inputs(self, inputs, **kwargs) -> None: + """Reject a batch before any rank starts its forward pass. + + Loss implementations whose alignment rules can fail per sample override + this hook. The default accepts every batch. + """ + return None + def micro_batch_scale(self, inputs: list[InputFeature], indices: list[int]) -> float: if len(indices) == len(inputs): return 1.0 diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index 2bf8d39fc..94b9d9d0c 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -152,10 +152,10 @@ def _pad_and_align_to_batch( * Response-only form (e.g. ``old_logps`` from vLLM): length equals the number of trainable positions in ``mask[i]`` and is scattered directly onto those positions. - * Full-sequence form (e.g. ``ref_logps`` from a ref-model forward, - right-padded to ``mask.shape[1]``): length ``>= mask.shape[1]``; - we slice to ``seq_len`` and index by ``mask[i]`` to extract the - trainable positions, then scatter. + * Full-sequence form (e.g. ``ref_logps`` from a ref-model forward): + every trainable mask position must exist in the per-sample row. + Rows need not include the batch's right-padding; we index by + ``mask[i]`` to extract the trainable positions, then scatter. Any other length is a real bug and triggers a hard assert — never silently truncate, since that misaligns IS ratios to the wrong tokens. @@ -196,25 +196,96 @@ def _pad_and_align_to_batch( pos = mask[i].nonzero(as_tuple=True)[0] n_pos = len(pos) n_sample = sample.numel() + last_pos = int(pos[-1].item()) if n_pos else -1 + if not self._alignment_supported(n_pos, n_sample, last_pos): + self._raise_alignment_mismatch(i, n_pos, n_sample, seq_len) if n_sample == 1: result[i, pos] = sample.item() elif n_sample == n_pos: # Response-only form (e.g. old_logps from vLLM). result[i, pos] = sample - elif n_pos == 0 or (n_sample > 0 and pos[-1].item() < n_sample): + else: # Variable-length full-sequence form. The processor right-pads the # batch, but per-sample RL fields from Tinker remain unpadded. They # are valid when every selected mask position exists in this row. result[i, pos] = sample[pos] - else: - raise AssertionError(f'data/mask length mismatch at sample {i}: ' - f'n_pos={n_pos}, n_sample={n_sample}, seq_len={seq_len} ' - '(expected n_sample == n_pos for response-only form, or all masked positions ' - 'to exist in the per-sample full-sequence form)') return result + @staticmethod + def _alignment_supported(n_pos: int, n_sample: int, last_pos: int) -> bool: + """Whether one ragged row fits the response-only or full-sequence form.""" + return (n_sample == 1 or n_sample == n_pos or n_pos == 0 or (n_sample > 0 and last_pos < n_sample)) + + @staticmethod + def _raise_alignment_mismatch(sample_index: int, n_pos: int, n_sample: int, seq_len: int): + raise AssertionError(f'data/mask length mismatch at sample {sample_index}: ' + f'n_pos={n_pos}, n_sample={n_sample}, seq_len={seq_len} ' + '(expected n_sample == n_pos for response-only form, or all masked positions ' + 'to exist in the per-sample full-sequence form)') + + @staticmethod + def _alignment_rows(data, batch_size: int): + """Return one row per sample, or None when the shape is not attributable.""" + import numpy as np + import torch + + if torch.is_tensor(data): + if data.dim() == 1 and data.shape[0] == batch_size: + return [data[index].reshape(1) for index in range(batch_size)] + if data.dim() >= 2 and data.shape[0] == batch_size: + return list(data) + return None + if isinstance(data, np.ndarray): + return GRPOLoss._alignment_rows(torch.as_tensor(data), batch_size) + if isinstance(data, (list, tuple)) and len(data) == batch_size: + return list(data) + return None + + def validate_inputs(self, inputs, *, old_logps=None, ref_logps=None, advantages=None, **kwargs) -> None: + """Reject ragged per-sample lengths before any rank starts its forward. + + ``_pad_and_align_to_batch`` already rejects these shapes, but it runs inside + the loss on each rank independently. Raising there lets one rank return + while the others enter the DDP backward collective, which deadlocks the + replica. Running the same rule on the full, unsliced batch first makes + every rank fail together before any collective is entered. + """ + import torch + + if old_logps is None and ref_logps is None and advantages is None: + return + samples = inputs if isinstance(inputs, list) else [inputs] + if not samples: + return + for data in (old_logps, ref_logps, advantages): + if data is None: + continue + rows = self._alignment_rows(data, len(samples)) + if rows is None: + # The shape does not map one row per sample; the in-loss check + # remains the only applicable validation for this batch. + continue + for index, (sample, row) in enumerate(zip(samples, rows)): + if not isinstance(sample, dict): + continue + labels = sample.get('labels') + if labels is None: + continue + if not torch.is_tensor(labels): + labels = torch.as_tensor(labels) + if labels.dim() == 1: + labels = labels.unsqueeze(0) + mask = self._resolve_loss_mask(sample, labels)[0] + pos = mask.nonzero(as_tuple=True)[0] + row_tensor = row if torch.is_tensor(row) else torch.as_tensor(row) + n_sample = row_tensor.numel() + n_pos = len(pos) + last_pos = int(pos[-1].item()) if n_pos else -1 + if not self._alignment_supported(n_pos, n_sample, last_pos): + self._raise_alignment_mismatch(index, n_pos, n_sample, labels.shape[-1]) + def _resolve_loss_mask(self, inputs: Dict, labels: 'torch.Tensor') -> 'torch.Tensor': """Positions this loss may score: trainable *and* log-prob-bearing. diff --git a/src/twinkle/server/checkpoint/checkpoint_manager.py b/src/twinkle/server/checkpoint/checkpoint_manager.py index d93d992d7..d86fadadd 100644 --- a/src/twinkle/server/checkpoint/checkpoint_manager.py +++ b/src/twinkle/server/checkpoint/checkpoint_manager.py @@ -456,12 +456,21 @@ def delete(self, model_id: str, checkpoint_id: str) -> bool: """Remove a checkpoint and refresh only its kind's last-checkpoint pointer. Sampler links own versions under .versions; legacy directories remain - supported. Missing checkpoints or traversal attempts return False. + supported. Missing checkpoints, traversal attempts, and paths that are + not a real checkpoint all return False. """ # Basic safety check to prevent directory traversal if '..' in checkpoint_id: return False + # Delete only a checkpoint that the manager itself can resolve. This + # rejects kind roots (``weights``), the run root (``.``), aliases whose + # metadata names another checkpoint (``sampler_weights/latest``), and + # any directory that has no checkpoint metadata. + checkpoint = self.get(model_id, checkpoint_id) + if checkpoint is None or checkpoint.checkpoint_id != checkpoint_id: + return False + ckpt_dir = self.get_ckpt_dir(model_id, checkpoint_id) if not ckpt_dir.exists(): diff --git a/src/twinkle/server/gateway/openai_handlers.py b/src/twinkle/server/gateway/openai_handlers.py index aeed51954..cec235b02 100644 --- a/src/twinkle/server/gateway/openai_handlers.py +++ b/src/twinkle/server/gateway/openai_handlers.py @@ -8,7 +8,9 @@ """ from __future__ import annotations +import asyncio import json +import os import uuid from collections.abc import Callable from fastapi import Depends, FastAPI, Request @@ -22,13 +24,68 @@ import httpx +import twinkle.protocol.types as types +from twinkle.server.lifecycle.envelope import envelope_from_record from twinkle.server.utils import get_template_for_model from twinkle.utils.logger import get_logger from . import routes from .openai_bridge import make_error, translate_chat_request, translate_response, translate_stream_chunk +from .use_cases import poll_future logger = get_logger() +_OPENAI_SAMPLE_TIMEOUT_ENV = 'TWINKLE_OPENAI_SAMPLE_TIMEOUT' +_DEFAULT_OPENAI_SAMPLE_TIMEOUT = 1800.0 + + +class SamplerWaitTimeoutError(TimeoutError): + """The sampler task did not reach a terminal state inside the gateway's budget.""" + + +def _sample_wait_timeout() -> float: + return float(os.environ.get(_OPENAI_SAMPLE_TIMEOUT_ENV, str(_DEFAULT_OPENAI_SAMPLE_TIMEOUT))) + + +async def _resolve_sampler_envelope(gateway: GatewayServer, sampler_data: dict[str, Any], + request: Request) -> types.TaskEnvelope | None: + """Resolve the sampler's TaskEnvelope to a terminal state. + + ``/twinkle/sample`` always answers with a TaskEnvelope; a slow sample is + still ``queued``/``running`` when the proxy returns. Polling the shared + state backend is the only way to obtain its result without changing the + upstream wire contract. Returns ``None`` when the future record was lost. + + ``TWINKLE_OPENAI_SAMPLE_TIMEOUT`` bounds the entire terminal-state wait, + including state reads. Client disconnection is checked between long polls. + """ + envelope = types.TaskEnvelope.model_validate(sampler_data) + if envelope.status in types.TERMINAL_STATUSES: + return envelope + + async def wait_for_terminal(envelope: types.TaskEnvelope) -> types.TaskEnvelope | None: + while envelope.status not in types.TERMINAL_STATUSES: + if await request.is_disconnected(): + raise SamplerWaitTimeoutError('Client disconnected while waiting for the sampling task.') + # ``poll_future`` waits out its whole window before reporting a missing + # record, so check once here: a lost record is a terminal answer, not a + # reason to keep the coroutine alive. + if await gateway.state.get_future(envelope.request_id) is None: + return None + outcome = await poll_future(gateway.state, envelope.request_id) + if outcome.record is None: + return None + envelope = envelope_from_record(envelope.request_id, outcome.record) + return envelope + + timeout = _sample_wait_timeout() + try: + return await asyncio.wait_for(wait_for_terminal(envelope), timeout=timeout) + except SamplerWaitTimeoutError: + raise + except asyncio.TimeoutError: + raise SamplerWaitTimeoutError( + f'Sampling task {envelope.request_id} did not finish within {timeout:g}s.') from None + def _register_openai_routes(app: FastAPI, self_fn: Callable[[], GatewayServer]) -> None: """Register OpenAI-compatible routes on the gateway FastAPI app.""" @@ -105,8 +162,43 @@ async def chat_completions( ) sampler_data = json.loads(response.body) + try: + envelope = await _resolve_sampler_envelope(self, sampler_data, request) + except SamplerWaitTimeoutError as e: + return JSONResponse( + status_code=504, + content=make_error(message=str(e), error_type='timeout'), + ) + if envelope is None: + return JSONResponse( + status_code=500, + content=make_error( + message=f'Sampling task {sampler_data.get("request_id")} was lost by the server', + error_type='server_error', + ), + ) + if envelope.status != 'completed': + error = envelope.error + message = error.error if error is not None else f'Sampling task {envelope.status}' + error_type = ('invalid_request_error' + if error is not None and error.category.value == 'user' else 'server_error') + return JSONResponse( + status_code=error.error_code if error is not None else 500, + content=make_error(message=message, error_type=error_type), + ) + + sampler_result = envelope.result or {} + if not sampler_result.get('samples'): + return JSONResponse( + status_code=502, + content=make_error( + message='Sampler completed without returning any samples', + error_type='server_error', + ), + ) + request_id = f'chatcmpl-{uuid.uuid4().hex[:24]}' - openai_response = translate_response(sampler_data, model, request_id) + openai_response = translate_response(sampler_result, model, request_id) resp_headers = {} replica_id = response.headers.get('x-twinkle-replica-id') if replica_id: @@ -227,9 +319,11 @@ async def _ensure_template( return template_cls = get_template_for_model(base_model) + # The sampler's strict schema accepts only template_cls / adapter_name / + # init_kwargs. It already knows its own model_id and injects it when the + # template constructor needs one, so sending model_id here is rejected 422. set_template_body = json.dumps({ 'template_cls': template_cls, - 'model_id': base_model, 'adapter_name': '', }).encode() diff --git a/src/twinkle/server/gateway/use_cases.py b/src/twinkle/server/gateway/use_cases.py index 26691014b..f3c5df5dd 100644 --- a/src/twinkle/server/gateway/use_cases.py +++ b/src/twinkle/server/gateway/use_cases.py @@ -18,6 +18,8 @@ from twinkle.server.checkpoint import create_checkpoint_manager, create_training_run_manager from twinkle.server.lifecycle.poll_config import long_poll_window, retrieve_poll_interval +# Deliberately local: this layer works on canonical records and must not import the +# protocol models (see test_gateway_services_do_not_import_protocol_models). _TERMINAL_STATUSES = frozenset({'completed', 'failed', 'cancelled'}) diff --git a/src/twinkle/server/lifecycle/submit.py b/src/twinkle/server/lifecycle/submit.py index b4da7ecb5..1c6b78937 100644 --- a/src/twinkle/server/lifecycle/submit.py +++ b/src/twinkle/server/lifecycle/submit.py @@ -204,7 +204,8 @@ async def _task(): retention = self.sequence_retention_s prior_request_id = await self.state.claim_training_seq(training_id, stream, seq_id, request_id, retention) else: - ttl = int(self.task_queue_config.effective_execution_timeout) + 60 + # Legacy claims must cover queue wait as well as backend execution. + ttl = int(self.task_queue_config.queue_timeout + self.task_queue_config.effective_execution_timeout) + 60 prior_request_id = await self.state.claim_seq(dedup_key, request_id, ttl) if prior_request_id is not None: return await self._peek_terminal(prior_request_id, fallback_status='pending') diff --git a/src/twinkle/server/model/backends/common.py b/src/twinkle/server/model/backends/common.py index 2ca5f8a74..0c62f3090 100644 --- a/src/twinkle/server/model/backends/common.py +++ b/src/twinkle/server/model/backends/common.py @@ -10,7 +10,8 @@ from tinker import types from typing import Any, List -from twinkle import DeviceMesh +from twinkle import DeviceMesh, remote_function +from twinkle.server.model.tinker_datum import datum_to_input_feature, extract_rl_features_for_loss from twinkle.server.validation.tinker_loss import validate_tinker_loss from twinkle.template import Template @@ -173,6 +174,35 @@ def _tinker_setup_loss(self, loss_fn: str, inputs, adapter_name: str, kwargs: di grpo_beta = kwargs.pop('beta', 0.0) self.set_loss('GRPOLoss', adapter_name=adapter_name, epsilon=epsilon, beta=grpo_beta) + @remote_function(dispatch='all', collect='none') + def preflight_loss_inputs(self, *, inputs, adapter_name: str = None, loss_fn: str = None, **kwargs) -> None: + """Validate the full, unsliced batch on every rank before any forward. + + Each backend's ``forward_backward`` is dispatched with ``slice_dp``, so a + rank that rejects its slice inside the loss has already forwarded while + its peers enter the backward collective, which deadlocks the replica + (BUG-05). Running the loss's own rule on the full batch first, on every + rank, makes them fail together before any collective is entered. + + ``loss_fn`` is the tinker route's loss selector: tinker batches carry + their RL features inside the datums, so the loss has to be set up and + the datums converted before the same rule applies. Twinkle-native + batches already pass ``old_logps``/``ref_logps``/``advantages`` as + keyword arguments and leave ``loss_fn`` unset. + """ + adapter_name = adapter_name or self._get_default_group() + if loss_fn is not None: + loss_values = extract_rl_features_for_loss(inputs) + loss_values.pop('ref_outputs', None) + self._tinker_setup_loss(loss_fn, inputs, adapter_name, kwargs) + inputs = datum_to_input_feature(inputs, self.get_template(adapter_name)) + loss_values.update(kwargs) + else: + loss_values = kwargs + loss = getattr(self.optimizer_group.get(adapter_name), 'loss_instance', None) + if loss is not None: + loss.validate_inputs(inputs, **loss_values) + def _ensure_dpo_metric(self, adapter_name: str, beta: float): """Add DPOMetric for the adapter if not already present. diff --git a/src/twinkle/server/model/tinker_handlers.py b/src/twinkle/server/model/tinker_handlers.py index 4566aafa5..cf32eecda 100644 --- a/src/twinkle/server/model/tinker_handlers.py +++ b/src/twinkle/server/model/tinker_handlers.py @@ -122,6 +122,14 @@ async def _do_forward_backward(): datum_list = body.forward_backward_input.data loss_fn = body.forward_backward_input.loss_fn loss_fn_config = body.forward_backward_input.loss_fn_config or {} + preflight = getattr(self.model, 'preflight_loss_inputs', None) + if preflight is not None and loss_fn == 'importance_sampling' and not validate_tinker_loss( + loss_fn, datum_list, data_world_size=self.data_world_size): + # GRPO datums carry ragged per-sample targets; the loss check inside + # the sliced backend would deadlock the peer ranks (BUG-05), so the + # full batch is validated on every rank first. + await self.call_backend( + preflight, inputs=datum_list, adapter_name=model_adapter, loss_fn=loss_fn, **loss_fn_config) output, loss = await self.call_backend( self.model.tinker_forward_backward, inputs=datum_list, diff --git a/src/twinkle/server/model/twinkle_handlers.py b/src/twinkle/server/model/twinkle_handlers.py index 2e1a98d49..b8537e8b7 100644 --- a/src/twinkle/server/model/twinkle_handlers.py +++ b/src/twinkle/server/model/twinkle_handlers.py @@ -133,11 +133,19 @@ def first_element(data): for key in inputs: if isinstance(inputs[key], list) and isinstance(first_element(inputs[key]), (int, float)): inputs[key] = torch.tensor(inputs[key]) + kwargs = backend_kwargs(body) + resolved_adapter_name = self.resolve_model_adapter_name(adapter_name) + preflight = getattr(self.model, 'preflight_loss_inputs', None) + if preflight is not None and any(key in kwargs for key in ('old_logps', 'ref_logps', 'advantages')): + # Full, unsliced batch on every rank: a loss alignment error must + # fail all ranks before any of them enters its backward collective. + await self.call_backend(preflight, inputs=all_inputs, adapter_name=resolved_adapter_name, **kwargs) + elif preflight is None: + logger.warning( + 'Backend %s has no preflight_loss_inputs; a loss alignment error can still deadlock the DP peers.', + type(self.model).__name__) ret = await self.call_backend( - self.model.forward_backward, - inputs=all_inputs, - adapter_name=self.resolve_model_adapter_name(adapter_name), - **backend_kwargs(body)) + self.model.forward_backward, inputs=all_inputs, adapter_name=resolved_adapter_name, **kwargs) return {'result': ret} return await run_submit( @@ -237,11 +245,17 @@ async def forward_backward_from_data_plane( async def _call(self, body, adapter_name, token): raw_inputs, field_kwargs = await resolve_data_plane_model_inputs(body, self.data_plane) kwargs = merge_forward_kwargs(backend_kwargs(body), field_kwargs) + inputs = to_backend_inputs(raw_inputs) + resolved_adapter_name = self.resolve_model_adapter_name(adapter_name) + preflight = getattr(self.model, 'preflight_loss_inputs', None) + if preflight is not None and any(key in kwargs for key in ('old_logps', 'ref_logps', 'advantages')): + await self.call_backend(preflight, inputs=inputs, adapter_name=resolved_adapter_name, **kwargs) + elif preflight is None: + logger.warning( + 'Backend %s has no preflight_loss_inputs; a loss alignment error can still deadlock the DP peers.', + type(self.model).__name__) ret = await self.call_backend( - self.model.forward_backward, - inputs=to_backend_inputs(raw_inputs), - adapter_name=self.resolve_model_adapter_name(adapter_name), - **kwargs) + self.model.forward_backward, inputs=inputs, adapter_name=resolved_adapter_name, **kwargs) return {'result': ret} return await run_submit( diff --git a/src/twinkle/server/task_queue/worker.py b/src/twinkle/server/task_queue/worker.py index 6dbc4ee89..de5bc42e1 100644 --- a/src/twinkle/server/task_queue/worker.py +++ b/src/twinkle/server/task_queue/worker.py @@ -13,6 +13,7 @@ import traceback from collections import deque from dataclasses import replace +from ray.exceptions import RayTaskError from typing import TYPE_CHECKING, Any, Callable, Deque from twinkle.protocol.types.errors import ErrorCategory @@ -35,6 +36,27 @@ logger = get_logger() +def _exception_summary(exc: BaseException) -> str: + """Build a single-line summary, unwrapping Ray's transport wrapper. + + A RayTaskError's own ``str`` starts with the Ray call header (actor id, + pid, repr); the actionable error (e.g. the loss assertion) is in ``cause``. + Preferring it keeps the summary useful without enlarging the wire field, + while the full traceback still travels separately. + + Only Ray's wrapper is unwrapped: elsewhere ``raise X from e`` is a + deliberate message and must not be replaced by its inner cause. + """ + cause = None + if isinstance(exc, RayTaskError): + cause = getattr(exc, 'cause', None) + if not isinstance(cause, BaseException): + cause = exc.__cause__ + if isinstance(cause, BaseException): + return f'{type(cause).__name__}: {cause}' + return f'{type(exc).__name__}: {exc}' + + def _reason_code_for_server_error(exc: TwinkleServerError) -> str: """Map typed execution failures to protocol-independent reason codes.""" mappings: tuple[tuple[type[TwinkleServerError], str], ...] = ( @@ -497,7 +519,7 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: exec_time = time.monotonic() - exec_start # error is a single-line summary; the full traceback goes only to the # traceback field, never into `error`. - error = f'{type(exc).__name__}: {exc}' + error = _exception_summary(exc) logger.error(f'[ComputeWorker] Task {task.request_id} FAILED after {exec_time:.2f}s, ' f'type={task_type}:\n{traceback.format_exc(limit=3)}') await self._store_task_failed( diff --git a/src/twinkle_client/_future.py b/src/twinkle_client/_future.py index b8fc5a3b4..59570f0c4 100644 --- a/src/twinkle_client/_future.py +++ b/src/twinkle_client/_future.py @@ -99,6 +99,7 @@ def _unwrap(env: TaskEnvelope, model_cls) -> Any: request_id=env.request_id, error_code=p.error_code, details=p.details, + traceback=p.traceback, ) if env.status == 'cancelled': p = env.error diff --git a/src/twinkle_client/exceptions.py b/src/twinkle_client/exceptions.py index dd56d6e35..32f59990c 100644 --- a/src/twinkle_client/exceptions.py +++ b/src/twinkle_client/exceptions.py @@ -80,6 +80,7 @@ def __init__( request_id: str, error_code: int | None = None, details: list[dict[str, Any]] | None = None, + traceback: str | None = None, ) -> None: super().__init__(error) self.error = error @@ -87,6 +88,7 @@ def __init__( self.request_id = request_id self.error_code = error_code self.details = details + self.traceback = traceback class TaskCancelledError(Exception): diff --git a/tests/loss/test_grpo_gkd.py b/tests/loss/test_grpo_gkd.py index 1f9daa707..9623df67e 100644 --- a/tests/loss/test_grpo_gkd.py +++ b/tests/loss/test_grpo_gkd.py @@ -1,5 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Tests for GRPO family losses and GKD loss.""" +import numpy as np import pytest import torch import torch.nn.functional as F @@ -286,3 +287,73 @@ def test_gkd_no_ignored_tokens(self): result = loss_fn({'labels': labels}, {'logits': student_logits}, teacher_logits=teacher_logits) assert torch.isfinite(result['loss']) + + +class TestGRPOValidateInputs: + """Pre-forward validation must match the in-loss alignment rules.""" + + @staticmethod + def _batch(): + # Two trainable positions per row, at indices 2 and 3, so a length-3 + # row cannot be read as the full-sequence form (index 3 is missing). + labels = torch.tensor([[-100, -100, 0, 0], [-100, -100, 0, 0]]) + return [{'labels': labels}, {'labels': labels}] + + def test_accepts_response_only_scalar_and_full_length(self): + loss_fn = GRPOLoss() + inputs = self._batch() + loss_fn.validate_inputs(inputs, old_logps=[[0.1, 0.2], [0.3, 0.4]]) + loss_fn.validate_inputs(inputs, old_logps=[[0.1], [0.2]]) + loss_fn.validate_inputs(inputs, old_logps=torch.zeros(2, 4)) + + def test_rejects_ragged_length_before_forward(self): + loss_fn = GRPOLoss() + inputs = self._batch() + with pytest.raises(AssertionError, match='data/mask length mismatch at sample 1'): + loss_fn.validate_inputs(inputs, old_logps=[[0.1, 0.2], [0.3, 0.4, 0.5]]) + + def test_rejects_ragged_advantages(self): + loss_fn = GRPOLoss() + inputs = self._batch() + with pytest.raises(AssertionError, match='data/mask length mismatch'): + loss_fn.validate_inputs(inputs, advantages=[1.0, [1.0, 2.0, 3.0]]) + + @pytest.mark.parametrize('field', ['old_logps', 'ref_logps', 'advantages']) + @pytest.mark.parametrize('representation', ['list', 'tuple', 'tensor', 'numpy']) + @pytest.mark.parametrize('labels,n_sample,supported', [ + ([-100, -100, 0, 0, -100], 0, False), + ([-100, -100, 0, 0, -100], 1, True), + ([-100, -100, 0, 0, -100], 2, True), + ([-100, -100, 0, 0, -100], 3, False), + ([-100, -100, 0, 0, -100], 4, True), + ([-100, -100, 0, 0, -100], 5, True), + ([-100, -100, 0, 0, -100], 6, True), + ([-100, -100, -100], 0, True), + ([-100, -100, -100], 1, True), + ([-100, -100, -100], 2, True), + ]) + def test_preflight_and_loss_agree_on_alignment(self, labels, n_sample, supported, representation, field): + loss_fn = GRPOLoss(beta=0.1) + inputs = {'labels': torch.tensor(labels)} + outputs = {'logps': torch.zeros(1, len(labels), requires_grad=True)} + data = [[0.1] * n_sample] + if representation == 'tuple': + data = tuple(tuple(row) for row in data) + elif representation == 'tensor': + data = torch.tensor(data) + elif representation == 'numpy': + data = np.asarray(data, dtype=np.float32) + kwargs = {'advantages': [1.0], field: data} + + if not supported: + with pytest.raises(AssertionError, match='data/mask length mismatch at sample 0'): + loss_fn.validate_inputs([inputs], **kwargs) + with pytest.raises(AssertionError, match='data/mask length mismatch at sample 0'): + loss_fn(inputs, outputs, **kwargs) + return + + loss_fn.validate_inputs([inputs], **kwargs) + result = loss_fn(inputs, outputs, **kwargs) + assert torch.isfinite(result['loss']) + result['loss'].backward() + assert torch.isfinite(outputs['logps'].grad).all() diff --git a/tests/server/gateway/test_openai_handlers.py b/tests/server/gateway/test_openai_handlers.py index 048a4ad61..55c80d2ff 100644 --- a/tests/server/gateway/test_openai_handlers.py +++ b/tests/server/gateway/test_openai_handlers.py @@ -6,9 +6,11 @@ """ from __future__ import annotations +import asyncio import json +import httpx import pytest -from fastapi import FastAPI +from fastapi import FastAPI, Request from fastapi.testclient import TestClient from unittest.mock import AsyncMock, MagicMock @@ -50,10 +52,14 @@ def mock_gateway(): class TestChatCompletions: - def _make_sample_response(self, text='Hello!', stop_reason='stop'): - """Build a mock Response matching what proxy_request returns.""" + def _make_sample_response(self, text='Hello!', stop_reason='stop', status='completed'): + """Build a mock Response matching what proxy_request returns. + + ``/twinkle/sample`` always answers with a TaskEnvelope whose completed + result carries the samples; the raw samples dict is not the wire shape. + """ from fastapi import Response - body = json.dumps({ + result = { 'samples': [{ 'sequences': [{ 'stop_reason': stop_reason, @@ -61,9 +67,24 @@ def _make_sample_response(self, text='Hello!', stop_reason='stop'): 'decoded': text, }] }] - }).encode() + } + payload = {'request_id': 'req-mock', 'status': status} + if status == 'completed': + payload['result'] = result + body = json.dumps(payload).encode() return Response(content=body, status_code=200, media_type='application/json') + def _samples_result(self, text='Hello!'): + return { + 'samples': [{ + 'sequences': [{ + 'stop_reason': 'stop', + 'tokens': list(range(len(text))), + 'decoded': text, + }] + }] + } + def test_basic_chat_completion(self, mock_gateway): mock_self, app = mock_gateway mock_self.proxy.proxy_request.return_value = self._make_sample_response('Hi there!') @@ -211,6 +232,270 @@ def test_adapter_resolves_via_metadata(self, mock_gateway): call_kwargs = mock_self.proxy.proxy_request.call_args assert call_kwargs.kwargs.get('base_model') or call_kwargs.args[2] == 'Qwen/Qwen3.5-4B' + def test_non_stream_polls_queued_task_to_completion(self, mock_gateway): + """A queued sampler envelope is polled until completed, then translated.""" + mock_self, app = mock_gateway + mock_self.proxy.proxy_request.return_value = self._make_sample_response(status='queued') + mock_self.state.get_future = AsyncMock(return_value={ + 'status': 'completed', + 'result': self._samples_result('Polled!'), + }) + + client = TestClient(app) + resp = client.post( + '/chat/completions', + json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{ + 'role': 'user', + 'content': 'Hello' + }], + }, + headers={'Authorization': 'Bearer key'}, + ) + + assert resp.status_code == 200 + assert resp.json()['choices'][0]['message']['content'] == 'Polled!' + mock_self.state.get_future.assert_awaited() + + def test_non_stream_failed_task_returns_error(self, mock_gateway): + """A failed sampler envelope becomes an OpenAI error, never an empty 200.""" + from fastapi import Response + mock_self, app = mock_gateway + envelope = { + 'request_id': 'req-failed', + 'status': 'failed', + 'result': None, + 'error': { + 'error': 'data/mask length mismatch at sample 0', + 'category': 'server', + 'error_code': 500, + 'request_id': 'req-failed', + }, + } + mock_self.proxy.proxy_request.return_value = Response( + content=json.dumps(envelope).encode(), status_code=200, media_type='application/json') + + client = TestClient(app) + resp = client.post( + '/chat/completions', + json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{ + 'role': 'user', + 'content': 'Hello' + }], + }, + headers={'Authorization': 'Bearer key'}, + ) + + assert resp.status_code == 500 + assert 'mismatch' in resp.json()['error']['message'] + assert resp.json()['error']['type'] == 'server_error' + + def test_non_stream_wait_is_bounded(self, mock_gateway, monkeypatch): + """A task that never terminates returns 504 instead of polling forever.""" + monkeypatch.setenv('TWINKLE_OPENAI_SAMPLE_TIMEOUT', '0.01') + mock_self, app = mock_gateway + mock_self.proxy.proxy_request.return_value = self._make_sample_response(status='running') + mock_self.state.get_future = AsyncMock(return_value={'status': 'running'}) + + client = TestClient(app) + resp = client.post( + '/chat/completions', + json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{ + 'role': 'user', + 'content': 'Hello' + }], + }, + headers={'Authorization': 'Bearer key'}, + ) + + assert resp.status_code == 504 + assert resp.json()['error']['type'] == 'timeout' + assert 'did not finish' in resp.json()['error']['message'] + + @pytest.mark.asyncio + @pytest.mark.parametrize('blocked_step', ['disconnect_check', 'initial_read', 'long_poll']) + async def test_non_stream_timeout_covers_each_wait_step(self, mock_gateway, monkeypatch, blocked_step): + from twinkle.server.gateway import openai_handlers + + monkeypatch.setenv('TWINKLE_OPENAI_SAMPLE_TIMEOUT', '0.01') + mock_self, app = mock_gateway + mock_self.proxy.proxy_request.return_value = self._make_sample_response(status='running') + mock_self.state.get_future = AsyncMock(return_value={'status': 'running'}) + started = asyncio.Event() + cancelled = asyncio.Event() + + async def blocked(*args, **kwargs): + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + if blocked_step == 'disconnect_check': + monkeypatch.setattr(Request, 'is_disconnected', blocked) + elif blocked_step == 'initial_read': + mock_self.state.get_future.side_effect = blocked + else: + monkeypatch.setattr(openai_handlers, 'poll_future', blocked) + + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url='http://test') as client: + resp = await asyncio.wait_for( + client.post('/chat/completions', json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{'role': 'user', 'content': 'Hello'}], + }), timeout=1.0) + + assert started.is_set() + assert cancelled.is_set() + assert resp.status_code == 504 + assert resp.json()['error']['type'] == 'timeout' + assert 'did not finish' in resp.json()['error']['message'] + + @pytest.mark.asyncio + async def test_non_stream_polls_share_one_total_budget(self, mock_gateway, monkeypatch): + from twinkle.server.gateway import openai_handlers + from twinkle.server.gateway.use_cases import FuturePollResult + + monkeypatch.setenv('TWINKLE_OPENAI_SAMPLE_TIMEOUT', '0.05') + mock_self, app = mock_gateway + mock_self.proxy.proxy_request.return_value = self._make_sample_response(status='running') + mock_self.state.get_future = AsyncMock(return_value={'status': 'running'}) + calls = 0 + + async def slow_poll(*args, **kwargs): + nonlocal calls + calls += 1 + await asyncio.sleep(0.02) + status = 'completed' if calls == 10 else 'running' + return FuturePollResult(record={ + 'status': status, + 'result': self._samples_result('Late!') if status == 'completed' else None, + }, timed_out=status != 'completed') + + monkeypatch.setattr(openai_handlers, 'poll_future', slow_poll) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url='http://test') as client: + resp = await asyncio.wait_for( + client.post('/chat/completions', json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{'role': 'user', 'content': 'Hello'}], + }), timeout=1.0) + + assert resp.status_code == 504 + assert resp.json()['error']['type'] == 'timeout' + assert calls < 10 + + @pytest.mark.asyncio + async def test_non_stream_disconnect_stops_before_state_reads(self, mock_gateway, monkeypatch): + monkeypatch.setattr(Request, 'is_disconnected', AsyncMock(return_value=True)) + mock_self, app = mock_gateway + mock_self.proxy.proxy_request.return_value = self._make_sample_response(status='running') + + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url='http://test') as client: + resp = await client.post('/chat/completions', json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{'role': 'user', 'content': 'Hello'}], + }) + + assert resp.status_code == 504 + assert 'disconnected' in resp.json()['error']['message'] + mock_self.state.get_future.assert_not_awaited() + + def test_non_stream_completed_task_needs_no_wait_budget(self, mock_gateway, monkeypatch): + monkeypatch.setenv('TWINKLE_OPENAI_SAMPLE_TIMEOUT', '0') + mock_self, app = mock_gateway + mock_self.proxy.proxy_request.return_value = self._make_sample_response('Done!') + + with TestClient(app) as client: + resp = client.post('/chat/completions', json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{'role': 'user', 'content': 'Hello'}], + }) + + assert resp.status_code == 200 + assert resp.json()['choices'][0]['message']['content'] == 'Done!' + mock_self.state.get_future.assert_not_awaited() + + def test_non_stream_lost_record_is_a_server_error(self, mock_gateway): + """A vanished future record is the server's loss, not a client 404.""" + mock_self, app = mock_gateway + mock_self.proxy.proxy_request.return_value = self._make_sample_response(status='queued') + mock_self.state.get_future = AsyncMock(return_value=None) + + client = TestClient(app) + resp = client.post( + '/chat/completions', + json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{ + 'role': 'user', + 'content': 'Hello' + }], + }, + headers={'Authorization': 'Bearer key'}, + ) + + assert resp.status_code == 500 + assert 'lost' in resp.json()['error']['message'] + assert resp.json()['error']['type'] == 'server_error' + + def test_template_request_uses_only_declared_sampler_fields(self, mock_gateway): + """The sampler's strict schema rejects top-level model_id; do not send it.""" + from fastapi import Response + mock_self, app = mock_gateway + mock_self.proxy.proxy_request.side_effect = [ + Response(content=b'{"status":"ok"}', status_code=200, media_type='application/json'), + self._make_sample_response('Hi'), + ] + + client = TestClient(app) + resp = client.post( + '/chat/completions', + json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{ + 'role': 'user', + 'content': 'Hello' + }], + }, + headers={'Authorization': 'Bearer key'}, + ) + + assert resp.status_code == 200 + template_call = mock_self.proxy.proxy_request.call_args_list[0] + template_body = json.loads(template_call.kwargs['body_override']) + assert set(template_body) == {'template_cls', 'adapter_name'} + assert 'Qwen/Qwen3.5-4B' in mock_self._template_initialized + + def test_failed_template_call_is_not_cached(self, mock_gateway): + """A rejected set_template must not be marked initialized.""" + from fastapi import Response + mock_self, app = mock_gateway + mock_self.proxy.proxy_request.side_effect = [ + Response(content=b'{}', status_code=422, media_type='application/json'), + self._make_sample_response('Hi'), + ] + + client = TestClient(app) + resp = client.post( + '/chat/completions', + json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{ + 'role': 'user', + 'content': 'Hello' + }], + }, + headers={'Authorization': 'Bearer key'}, + ) + + assert resp.status_code == 200 + assert 'Qwen/Qwen3.5-4B' not in mock_self._template_initialized + # ---------- /chat/completions streaming tests ------------------------------ # diff --git a/tests/server/lifecycle/test_run_submit_dedup.py b/tests/server/lifecycle/test_run_submit_dedup.py index 4d3b3291e..53d1b6925 100644 --- a/tests/server/lifecycle/test_run_submit_dedup.py +++ b/tests/server/lifecycle/test_run_submit_dedup.py @@ -20,11 +20,13 @@ class _FakeState: def __init__(self, record_after_claim): self._record_after_claim = record_after_claim self.claimed = {} + self.claim_ttls = [] self.released = [] async def claim_seq(self, dedup_key, request_id, ttl): # Unseen -> claim it and let the caller proceed (returns None). self.claimed[dedup_key] = request_id + self.claim_ttls.append(ttl) return None async def get_future(self, request_id): @@ -38,7 +40,7 @@ async def release_seq(self, dedup_key): class _FakeManagement: def __init__(self, record_after_claim): self.state = _FakeState(record_after_claim) - self.task_queue_config = SimpleNamespace(effective_execution_timeout=60.0) + self.task_queue_config = SimpleNamespace(queue_timeout=300.0, effective_execution_timeout=1800.0) # A real deployment declares its backend; preflight reads it from here. self.backend = 'transformers' self.data_world_size = 1 @@ -87,6 +89,18 @@ async def test_release_when_never_enqueued(): assert mgmt.state.released == ['seq::sess-1::sess-1-ad::forward_backward::7'], 'claim should be released when nothing enqueued' +@pytest.mark.asyncio +async def test_dedup_ttl_covers_queue_and_execution(): + """The claim must outlive queue wait plus execution, not just execution.""" + mgmt = _FakeManagement(record_after_claim={'status': 'queued'}) + with pytest.raises(RuntimeError): + await run_submit(mgmt, _request(), _body(), task_type='forward_backward', backend_call=_call) + config = mgmt.task_queue_config + assert mgmt.state.claim_ttls == [int(config.queue_timeout + config.effective_execution_timeout) + 60] + # An execution-only window is exactly what let a queued first copy expire. + assert mgmt.state.claim_ttls[0] > int(config.effective_execution_timeout) + 60 + + @pytest.mark.asyncio async def test_dedup_key_is_scoped_per_adapter(): """Two adapters in ONE session must not collide on the same seq_id. diff --git a/tests/server/model/test_full_mode_exclusivity.py b/tests/server/model/test_full_mode_exclusivity.py new file mode 100644 index 000000000..d49238fca --- /dev/null +++ b/tests/server/model/test_full_mode_exclusivity.py @@ -0,0 +1,106 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Full-mode exclusivity must be one atomic state transition, not check-then-register.""" +from __future__ import annotations + +import asyncio + +import pytest +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock +from starlette.requests import Request + +from tests.server.state.fake_backend import FakeBackend +from tests.server.training.test_controller import request, setup +from twinkle.server.config.training import TrainingConfig +from twinkle.server.training.catalog import TrainingCatalog +from twinkle.server.training.router import TrainingRouter + + +async def _setup_full_instance(): + """Full-mode slots belong to the controller, not the retired model create route.""" + controller, state, runtime, _ = await setup(backend=FakeBackend()) + raw = controller.config.model_dump() + raw['models'][0]['modes']['full']['max_replicas'] = 1 + controller.config = TrainingConfig.model_validate(raw) + controller.catalog = TrainingCatalog(controller.config) + runtime.start_gate.clear() + return controller, state, runtime + + +@pytest.mark.asyncio +async def test_concurrent_full_mode_claims_have_exactly_one_winner(): + controller, state, runtime = await _setup_full_instance() + try: + refs = await asyncio.gather(*(controller.create(request(name, mode='full')) + for name in ('tenant-a', 'tenant-b'))) + assert len(controller.instances) == 1 + assert len(next(iter(controller.instances.values())).training_ids) == 1 + runtime.start_gate.set() + await controller.wait_for_jobs() + records = [await state.get_future(ref['request_id']) for ref in refs] + assert sum(record['status'] == 'completed' for record in records) == 1 + assert sum(record['status'] == 'failed' and record['failure']['reason_code'] == 'capacity_unavailable' + for record in records) == 1 + assert await state._backend.get('token_count::owner') == 1 + finally: + runtime.start_gate.set() + await controller.wait_for_jobs() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_same_tenant_may_reclaim_its_own_resource(): + controller, _, runtime = await _setup_full_instance() + try: + refs = await asyncio.gather(*(controller.create(request('same', mode='full')) for _ in range(2))) + assert refs[0] == refs[1] + assert len(controller.instances) == 1 + assert len(next(iter(controller.instances.values())).training_ids) == 1 + finally: + runtime.start_gate.set() + await controller.wait_for_jobs() + await controller.shutdown() + + +def _request(body): + async def receive(): + return {'type': 'http.request', 'body': json.dumps(body).encode(), 'more_body': False} + incoming = Request({'type': 'http', 'method': 'POST', 'path': '/twinkle/add_adapter_to_model', 'headers': []}, receive) + incoming.state.session_id = 'session' + incoming.state.token = 'owner' + incoming.state.request_id = 'req' + return incoming + + +@pytest.mark.asyncio +async def test_concurrent_handler_creates_have_exactly_one_winner(): + """Drive the current catalog handler into real controller reservations.""" + controller, state, runtime = await _setup_full_instance() + router = TrainingRouter.__new__(TrainingRouter) + router.state = state + router.catalog = controller.catalog + router.model = SimpleNamespace(name='small') + router.controller = SimpleNamespace(create=SimpleNamespace(remote=AsyncMock(side_effect=controller.create))) + + async def create_remote(payload): + from twinkle.server.state.training_records import CreateTraining + return await controller.create(CreateTraining.model_validate(payload)) + router.controller.create.remote.side_effect = create_remote + try: + refs = await asyncio.gather(*(router.native_request(_request({ + 'client_create_id': name, 'adapter_name': name, + }), 'add_adapter_to_model') for name in ('adapter-a', 'adapter-b'))) + assert len(controller.instances) == 1 + assert len(next(iter(controller.instances.values())).training_ids) == 1 + runtime.start_gate.set() + await controller.wait_for_jobs() + records = [await state.get_future(ref.request_id) for ref in refs] + assert sum(record['status'] == 'completed' for record in records) == 1 + assert sum(record['status'] == 'failed' and record['failure']['reason_code'] == 'capacity_unavailable' + for record in records) == 1 + assert await state._backend.get('token_count::owner') == 1 + finally: + runtime.start_gate.set() + await controller.wait_for_jobs() + await controller.shutdown() diff --git a/tests/server/model/test_preflight_loss_inputs.py b/tests/server/model/test_preflight_loss_inputs.py new file mode 100644 index 000000000..6c76c9354 --- /dev/null +++ b/tests/server/model/test_preflight_loss_inputs.py @@ -0,0 +1,103 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Both backend entry points share one full-batch preflight implementation.""" +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from twinkle.loss.grpo import GRPOLoss +from twinkle.server.model.backends import common +from twinkle.server.model.backends.common import TwinkleCompatModelBase +from twinkle.server.model.backends.megatron_model import (TwinkleCompatFullMegatronModel, TwinkleCompatMegatronModel) +from twinkle.server.model.backends.transformers_model import (TwinkleCompatFullTransformersModel, + TwinkleCompatTransformersModel) + + +class _FakeBackend(TwinkleCompatModelBase): + """Just the state preflight touches; no Ray actor, no device.""" + + def __init__(self, loss=GRPOLoss, with_loss=True): + self.optimizer_group = { + 'adapter': SimpleNamespace(loss_instance=loss(epsilon=0.2) if with_loss and loss else None), + } + self.set_loss_calls = [] + + def _get_default_group(self): + return 'adapter' + + def set_loss(self, loss_cls, **kwargs): + kwargs.pop('adapter_name', None) + self.set_loss_calls.append((loss_cls, kwargs)) + self.optimizer_group['adapter'].loss_instance = GRPOLoss(**{ + key: value + for key, value in kwargs.items() if key in ('epsilon', 'beta') + }) + + def get_template(self, adapter_name): + return None + + +_BAD_SAMPLE = { + 'labels': [-100, -100, 0, 0], +} +_BAD_OLD_LOGPS = [[0.1, 0.2, 0.3]] + + +@pytest.mark.parametrize('backend', [ + TwinkleCompatTransformersModel, + TwinkleCompatFullTransformersModel, + TwinkleCompatMegatronModel, + TwinkleCompatFullMegatronModel, +]) +def test_every_backend_inherits_the_preflight_entry_point(backend): + """A backend without it silently skips the check and can still deadlock.""" + assert hasattr(backend, 'preflight_loss_inputs') + + +def test_preflight_rejects_a_ragged_batch_on_the_twinkle_route(): + backend = _FakeBackend() + with pytest.raises(AssertionError, match='mismatch'): + TwinkleCompatModelBase.preflight_loss_inputs( + backend, + inputs=[_BAD_SAMPLE], + adapter_name='adapter', + old_logps=_BAD_OLD_LOGPS, + ) + assert backend.set_loss_calls == [] + + +def test_preflight_accepts_a_response_only_row(): + backend = _FakeBackend() + TwinkleCompatModelBase.preflight_loss_inputs( + backend, + inputs=[_BAD_SAMPLE], + adapter_name='adapter', + old_logps=[[0.1, 0.2]], + ) + + +def test_preflight_tinker_route_sets_up_the_loss_before_validating(monkeypatch): + """The tinker payload hides its RL features in the datums, so they must be extracted.""" + monkeypatch.setattr(common, 'datum_to_input_feature', lambda inputs, template: [_BAD_SAMPLE]) + monkeypatch.setattr(common, 'extract_rl_features_for_loss', lambda inputs: {'old_logps': _BAD_OLD_LOGPS, 'ref_outputs': {'logps': []}}) + backend = _FakeBackend(with_loss=False) + + with pytest.raises(AssertionError, match='mismatch'): + TwinkleCompatModelBase.preflight_loss_inputs( + backend, + inputs=[SimpleNamespace(loss_fn_inputs={})], + adapter_name='adapter', + loss_fn='importance_sampling', + epsilon=0.3, + ) + + assert backend.set_loss_calls == [('GRPOLoss', { + 'epsilon': 0.3, + 'beta': 0.0 + })] + + +def test_preflight_is_a_no_op_without_a_loss_instance(): + backend = _FakeBackend(with_loss=False) + TwinkleCompatModelBase.preflight_loss_inputs(backend, inputs=[_BAD_SAMPLE], adapter_name='adapter', old_logps=_BAD_OLD_LOGPS) diff --git a/tests/server/model/test_twinkle_async_inputs.py b/tests/server/model/test_twinkle_async_inputs.py index 9f2ee1e64..21cf17dec 100644 --- a/tests/server/model/test_twinkle_async_inputs.py +++ b/tests/server/model/test_twinkle_async_inputs.py @@ -272,3 +272,69 @@ async def test_checkpoint_consumers_keep_backend_path_and_options(monkeypatch, t else: service.model.upload_to_hub.assert_called_once_with( checkpoint_dir=directory, hub_model_id='org/output', hub_token='hub-token', async_upload=False) + + +class _PreflightManagement(_SchedulingManagement): + + def __init__(self): + super().__init__() + self.preflight_calls = [] + + def preflight_loss_inputs(self, *, inputs, adapter_name, **kwargs): + self.preflight_calls.append((inputs, adapter_name, kwargs)) + self.model_calls.append(('preflight', adapter_name)) + + +@pytest.mark.asyncio +async def test_forward_backward_preflights_loss_inputs_before_backend() -> None: + """The full batch is validated before the backend call can slice it.""" + management = _PreflightManagement() + app = FastAPI() + _register_model_twinkle_routes(app, lambda: management) + route = next(route for route in app.routes if getattr(route, 'path', None) == '/twinkle/forward_backward') + request = Request({'type': 'http', 'headers': []}) + request.state.session_id = 'session' + body = types.ForwardBackwardTaskRequest( + inputs=[{ + 'input_ids': [index, index + 1], + 'labels': [-100, 5], + } for index in range(3)], + adapter_name='adapter', + loss_kwargs={'old_logps': [[-0.1], [-0.2], [-0.3]]}, + ) + + await route.endpoint(request, body, management) + + assert len(management.preflight_calls) == 1 + preflight_inputs, preflight_adapter, preflight_kwargs = management.preflight_calls[0] + # Slice-free: a per-rank slice would carry at most one of the three samples. + assert len(preflight_inputs) == 3 + assert preflight_adapter == 'session-adapter' + assert preflight_kwargs['old_logps'] == [[-0.1], [-0.2], [-0.3]] + assert management.model_calls[0][0] == 'preflight' + assert management.model_calls[-1][0] != 'preflight' + assert management.model_calls[-1][1] == 'session-adapter' + + +@pytest.mark.asyncio +async def test_forward_backward_skips_preflight_without_rl_inputs() -> None: + """Plain SFT batches carry no ragged targets, so the extra actor round trip is skipped.""" + management = _PreflightManagement() + app = FastAPI() + _register_model_twinkle_routes(app, lambda: management) + route = next(route for route in app.routes if getattr(route, 'path', None) == '/twinkle/forward_backward') + request = Request({'type': 'http', 'headers': []}) + request.state.session_id = 'session' + body = types.ForwardBackwardTaskRequest( + inputs=[{ + 'input_ids': [1, 2], + 'labels': [-100, 5], + }], + adapter_name='adapter', + ) + + await route.endpoint(request, body, management) + + assert management.preflight_calls == [] + assert len(management.model_calls) == 1 + assert management.model_calls[0][1] == 'session-adapter' diff --git a/tests/server/utils/test_checkpoint_manager.py b/tests/server/utils/test_checkpoint_manager.py index 88a40e0d0..5303a876b 100644 --- a/tests/server/utils/test_checkpoint_manager.py +++ b/tests/server/utils/test_checkpoint_manager.py @@ -420,3 +420,22 @@ async def test_reserved_sampler_name_is_rejected_before_backend_or_file_io(manag await manager.save_sampler('run', 'latest', save) save.assert_not_awaited() assert not Path(manager.get_save_dir('run', is_sampler=True)).exists() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('manager', ['tinker', 'twinkle'], indirect=True) +async def test_delete_rejects_non_checkpoint_paths(manager): + """Only a concrete checkpoint may be deleted, never a kind root, alias, or the run root.""" + await manager.save_training('run', 'step', _writer('training')) + await manager.save_sampler('run', None, _writer('live')) + run_dir = manager.training_run_manager.get_model_dir('run') + + def snapshot(): + return sorted(path.relative_to(run_dir).as_posix() for path in run_dir.rglob('*')) + + before = snapshot() + for checkpoint_id in ('weights', 'sampler_weights', '.', 'sampler_weights/latest', 'weights/../weights/step'): + assert manager.delete('run', checkpoint_id) is False, checkpoint_id + assert snapshot() == before + assert manager.delete('run', 'weights/step') is True + assert not (run_dir / 'weights/step').exists() diff --git a/tests/server/utils/test_task_queue_mixin.py b/tests/server/utils/test_task_queue_mixin.py index a947065d1..5338f4885 100644 --- a/tests/server/utils/test_task_queue_mixin.py +++ b/tests/server/utils/test_task_queue_mixin.py @@ -5,7 +5,47 @@ from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.task_queue.mixin import TaskQueueMixin from twinkle.server.task_queue.types import UserTaskError -from twinkle.server.task_queue.worker import ComputeWorker +from twinkle.server.task_queue.worker import ComputeWorker, _exception_summary + + +def test_exception_summary_prefers_the_wrapped_cause(): + from ray.exceptions import RayTaskError + + err = RayTaskError( + 'TwinkleCompatTransformersModel.preflight_loss_inputs', + 'ray::TwinkleCompatTransformersModel.preflight_loss_inputs() (pid=1, ip=10.0.0.1)', + AssertionError('data/mask length mismatch at sample 0'), + ) + + summary = _exception_summary(err) + assert summary.startswith('AssertionError:') + assert 'mismatch' in summary + + +def test_exception_summary_keeps_a_deliberate_wrapper_message(): + """``raise X from e`` is a deliberate message; only Ray's wrapper is unwrapped.""" + + class _Rejected(Exception): + pass + + try: + try: + raise ValueError('internal detail') + except ValueError as inner: + raise _Rejected('adapter_name is required') from inner + except _Rejected as outer: + summary = _exception_summary(outer) + + assert summary == '_Rejected: adapter_name is required' + + +def test_exception_summary_falls_back_to_the_ray_error_itself(): + from ray.exceptions import RayTaskError + + err = RayTaskError('Model.step', 'ray::Model.step() (pid=1)', None) + summary = _exception_summary(err) + assert summary.startswith('RayTaskError') + assert 'Model.step' in summary class _DummyState: diff --git a/tests/twinkle_client/test_future_layer.py b/tests/twinkle_client/test_future_layer.py index e078b5388..4569dcbe8 100644 --- a/tests/twinkle_client/test_future_layer.py +++ b/tests/twinkle_client/test_future_layer.py @@ -62,6 +62,15 @@ def test_terminal_submit_failure_raises_taskfailed_with_payload(monkeypatch): assert not isinstance(exc.value, requests.HTTPError) +def test_terminal_failure_keeps_remote_traceback(monkeypatch): + """The client exception must carry the server diagnostic, not only the summary.""" + monkeypatch.setattr(_future, '_post_retrieve', lambda _r, _transport: pytest.fail('no retrieve')) + diagnostic = 'Traceback (most recent call last):\nAssertionError: data/mask length mismatch at sample 0' + with pytest.raises(TaskFailedError) as exc: + _future.resolve(_failed(traceback=diagnostic), model_cls=_Model) + assert exc.value.traceback == diagnostic + + def test_model_cls_none_returns_none_result(monkeypatch): """A method that returned None before still returns None (not swallowed).""" monkeypatch.setattr(_future, '_post_retrieve', lambda _r, _transport: pytest.fail('no retrieve')) From 0be5b9d74306f6efff0caa35026da836e2dd290f Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Wed, 21 Oct 2026 16:07:21 +0000 Subject: [PATCH 09/46] fix(server): harden on-demand lifecycle and preserve YAML options --- .../Usage Guide/Server and Client/Server.md | 2 + .../\346\234\215\345\212\241\347\253\257.md" | 2 + src/twinkle/hub/__init__.py | 14 +- src/twinkle/server/__init__.py | 11 +- src/twinkle/server/config/server_config.py | 3 +- src/twinkle/server/config/training.py | 1 + .../server/config/training_normalization.py | 10 +- src/twinkle/server/deployment.py | 3 +- .../server/launcher/server_launcher.py | 17 ++- src/twinkle/server/service_proxy.py | 5 +- src/twinkle/server/state/__init__.py | 33 +++-- src/twinkle/server/state/backend/base.py | 6 + .../server/state/backend/memory_backend.py | 10 +- .../server/state/backend/redis_backend.py | 6 +- .../server/state/cleanup_coordinator.py | 66 ++++----- src/twinkle/server/state/future_manager.py | 22 ++- src/twinkle/server/state/server_state.py | 4 + src/twinkle/server/task_queue/mixin.py | 96 +++++-------- src/twinkle/server/task_queue/worker.py | 52 +++---- src/twinkle/server/training/actor.py | 65 +++++---- src/twinkle/server/training/driver.py | 38 ++++-- src/twinkle/utils/loader.py | 3 +- tests/server/config/test_training_config.py | 89 ++++++++++++ tests/server/gateway/test_proxy.py | 34 +++++ .../on_demand_training/fragmentation_probe.py | 4 +- .../on_demand_training/http_probe.py | 19 +-- .../integration/test_mock_mode_startup.py | 32 ++++- .../sampler/test_sampling_concurrency.py | 7 + tests/server/state/fake_backend.py | 6 +- tests/server/state/test_future_lifecycle.py | 110 ++++++--------- tests/server/state/test_leader_election.py | 91 ++++++++++-- tests/server/state/test_update_atomic.py | 37 +++++ tests/server/test_dependencies.py | 40 ++++++ .../training/test_actor_initialization.py | 76 ++++++++++- .../server/training/test_failure_recovery.py | 52 +++++++ tests/server/utils/test_task_queue_mixin.py | 129 +++++++++++++++++- 36 files changed, 885 insertions(+), 310 deletions(-) diff --git a/docs/source_en/Usage Guide/Server and Client/Server.md b/docs/source_en/Usage Guide/Server and Client/Server.md index fe09347e9..67a56d4ed 100644 --- a/docs/source_en/Usage Guide/Server and Client/Server.md +++ b/docs/source_en/Usage Guide/Server and Client/Server.md @@ -484,4 +484,6 @@ holds the cleanup lease. Existing `applications` entries with `import_path: model` remain valid. Set `args.min_replicas: 0` for lazy startup and explicitly set `args.idle_timeout_s` for business inactivity. Legacy `adapter_config.adapter_timeout` remains a heartbeat timeout and `adapter_max_lifetime` retains its absolute lifetime. New `modes` default to 1800 seconds of business inactivity; heartbeats and polling do not renew that deadline; queued and running operations protect training state. Empty instances above the resident minimum are destroyed so other models can reuse their devices. Existing deployment min/max replicas become training instance minimums/limits. +`modes` is optional: single-mode training does not require rewriting existing YAML. Model `deployments` options configure the CPU router; autoscaling request targets apply to router replicas. The Controller owns training-instance minimums/limits, and each compute instance has exactly one Serve replica. Non-autoscaling request limits, health/shutdown settings, and `ray_actor_options.runtime_env` also reach compute instances. `args.env_vars` overrides matching deployment environment variables. Unsupported `adapter_config` fields fail validation instead of being silently ignored. + To offer both LoRA and full training for one model, configure `num_devices`, `device_mesh`, `min_replicas`, and `idle_timeout_s` separately under `args.modes.lora/full`. LoRA users share compatible instances; each full-training handle owns an instance. Sampling configuration stays the same. See `cookbook/client/server/on_demand_training.yaml` in the repository. diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" index cc7899561..37d808d8f 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" @@ -471,4 +471,6 @@ twinkle-server check-config -c server_config.yaml 原有 `applications` 和 `import_path: model` 配置继续有效。单模式配置增加 `args.min_replicas: 0` 即可懒启动，`args.idle_timeout_s` 显式设置训练句柄的业务闲置时间。旧 `adapter_config.adapter_timeout` 仍按心跳判断会话存活，`adapter_max_lifetime` 保留绝对寿命；不再把旧字段转换成业务闲置时间。新增 `modes` 的业务闲置默认 1800 秒，心跳和轮询不延长该期限，运行中或已排队的任务保护训练状态。最后一个用户释放或过期后，超出保底数量的实例立即销毁，释放的卡可用于其他模型。已有 deployment 的 min/max_replicas 转成训练实例的保底/上限。 +`modes` 是可选扩展，单模式训练不需要改写原 YAML。原模型 `deployments` 的 Serve 设置用于 CPU 路由部署，其中 autoscaling 的请求目标用于路由副本扩缩容；训练计算实例的 min/max 仍由 Controller 管理，每个计算实例固定一个 Serve replica。非 autoscaling 的请求上限、健康/关闭设置和 `ray_actor_options.runtime_env` 会传给计算实例；`args.env_vars` 优先于 deployment 的同名环境变量。未知或未实现的 `adapter_config` 字段会报错，不再静默忽略。 + 同一模型需要同时提供 LoRA 和全参数训练时，在 `args.modes.lora/full` 内分别设置 `num_devices`、`device_mesh`、`min_replicas` 和 `idle_timeout_s`；不用改成另一套顶层模型目录。LoRA 用户共享兼容实例，全参数用户各占一个实例。采样服务继续使用原配置。详见仓库 `cookbook/client/server/on_demand_training.yaml`。 diff --git a/src/twinkle/hub/__init__.py b/src/twinkle/hub/__init__.py index 4180d5f63..933a2a5c9 100644 --- a/src/twinkle/hub/__init__.py +++ b/src/twinkle/hub/__init__.py @@ -1,3 +1,13 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -from .hub import HFHub, HubOperation, MSHub -from .model_alias import MODEL_ID_ALIASES_ENV, build_model_alias_map, load_model_alias_map, resolve_model_id_alias +import sys + +from twinkle._lazy_module import _LazyModule + +sys.modules[__name__] = _LazyModule( + __name__, + __file__, { + 'hub': ['HFHub', 'HubOperation', 'MSHub'], + 'model_alias': + ['MODEL_ID_ALIASES_ENV', 'build_model_alias_map', 'load_model_alias_map', 'resolve_model_id_alias'], + }, + module_spec=__spec__) diff --git a/src/twinkle/server/__init__.py b/src/twinkle/server/__init__.py index 5bdb1f97c..264552bfc 100644 --- a/src/twinkle/server/__init__.py +++ b/src/twinkle/server/__init__.py @@ -1,7 +1,8 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -from .launcher import ServerLauncher, launch_server +import sys -__all__ = [ - 'ServerLauncher', - 'launch_server', -] +from twinkle._lazy_module import _LazyModule + +sys.modules[__name__] = _LazyModule( + __name__, __file__, {'launcher.server_launcher': ['ServerLauncher', 'launch_server']}, module_spec=__spec__) +sys.modules[__name__].__all__ = ['ServerLauncher', 'launch_server'] diff --git a/src/twinkle/server/config/server_config.py b/src/twinkle/server/config/server_config.py index a87e7f037..31d22b5fb 100644 --- a/src/twinkle/server/config/server_config.py +++ b/src/twinkle/server/config/server_config.py @@ -14,7 +14,6 @@ """ from __future__ import annotations -from functools import cached_property from pathlib import Path from pydantic import BaseModel, ConfigDict, Field, model_validator from typing import Any @@ -86,7 +85,7 @@ def _validate_cross_field(self) -> ServerConfig: raise ValueError(f'training route conflicts with application: {path}') return self - @cached_property + @property def training_config(self) -> TrainingConfig: """Normalize old and extended model args into one runtime catalog.""" return TrainingConfig( diff --git a/src/twinkle/server/config/training.py b/src/twinkle/server/config/training.py index 7c7a10bf8..d4c372b9c 100644 --- a/src/twinkle/server/config/training.py +++ b/src/twinkle/server/config/training.py @@ -87,6 +87,7 @@ class TrainingModelConfig(_TrainingConfig): route_prefix: str | None = None data_plane_url: str | None = None env_vars: dict[str, str] = Field(default_factory=dict) + deploy_options: dict = Field(default_factory=dict) @model_validator(mode='after') def _validate_backend(self): diff --git a/src/twinkle/server/config/training_normalization.py b/src/twinkle/server/config/training_normalization.py index da7487cab..ee38a67c8 100644 --- a/src/twinkle/server/config/training_normalization.py +++ b/src/twinkle/server/config/training_normalization.py @@ -8,6 +8,7 @@ def normalize_model_application(app: ApplicationSpec) -> TrainingModelConfig: args = app.args assert isinstance(args, ModelArgs) + deployment = {} if args.modes is not None: baseline = ModelArgs(model_id=args.model_id, backend=args.backend, device_group={}, device_mesh={}) common = { @@ -33,6 +34,9 @@ def normalize_model_application(app: ApplicationSpec) -> TrainingModelConfig: if len(app.deployments) > 1: raise ValueError(f'{app.name}: a model application must have one deployment configuration') deployment = app.deployments[0] if app.deployments else {} + unsupported = set(args.adapter_config or {}) - {'adapter_timeout', 'adapter_max_lifetime'} + if unsupported: + raise ValueError(f'{app.name}: unsupported adapter_config fields: {sorted(unsupported)}') scaling = deployment.get('autoscaling_config') or {} ranks = args.device_group.get('ranks', args.nproc_per_node) if isinstance(ranks, list): @@ -84,4 +88,8 @@ def normalize_model_application(app: ApplicationSpec) -> TrainingModelConfig: legacy_expiration=legacy_expiration, route_prefix=app.route_prefix.rstrip('/') or '/', data_plane_url=args.data_plane_url, - env_vars=env_vars) + env_vars=env_vars, + deploy_options={ + key: value + for key, value in deployment.items() if key != 'name' + }) diff --git a/src/twinkle/server/deployment.py b/src/twinkle/server/deployment.py index 45ca43fbc..ca38798f8 100644 --- a/src/twinkle/server/deployment.py +++ b/src/twinkle/server/deployment.py @@ -329,4 +329,5 @@ async def _ensure_state_cleanup_started(self) -> None: await self.state.start_cleanup_task() except Exception as e: logger.warning(f'Failed to start ServerState cleanup task: {e}') - self._state_cleanup_started = True + else: + self._state_cleanup_started = True diff --git a/src/twinkle/server/launcher/server_launcher.py b/src/twinkle/server/launcher/server_launcher.py index 342fe469b..7cd6bd125 100644 --- a/src/twinkle/server/launcher/server_launcher.py +++ b/src/twinkle/server/launcher/server_launcher.py @@ -13,6 +13,7 @@ import signal import threading from collections.abc import Callable +from copy import deepcopy from pathlib import Path from typing import Any @@ -348,16 +349,14 @@ def _deploy_training(self) -> None: raise self._training_controller = controller for model in self.config.training_config.models: + options = deepcopy(model.deploy_options) + if 'autoscaling_config' not in options: + options.setdefault('num_replicas', 1) + actor_options = options.setdefault('ray_actor_options', {}) + actor_options.setdefault('num_cpus', 0.1) + actor_options.setdefault('runtime_env', {})['env_vars'] = {**env, **model.env_vars} app = build_training_router(self.config.training_config, model, namespace, - self.config.http_options.model_dump(), { - 'num_replicas': 1, - 'ray_actor_options': { - 'num_cpus': 0.1, - 'runtime_env': { - 'env_vars': env - } - } - }) + self.config.http_options.model_dump(), options) serve.run( app, name=f'training-router-{namespace}-{model.name}', diff --git a/src/twinkle/server/service_proxy.py b/src/twinkle/server/service_proxy.py index 071d2253f..478a7d6cf 100644 --- a/src/twinkle/server/service_proxy.py +++ b/src/twinkle/server/service_proxy.py @@ -9,6 +9,7 @@ from twinkle.protocol.headers import H_MULTIPLEX, H_MULTIPLEX_LEGACY, H_REQUEST_ID, H_REQUEST_ID_LEGACY from twinkle.protocol.types.errors import ErrorCategory, ErrorPayload +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.telemetry.tracing import inject_context from twinkle.utils.logger import get_logger @@ -20,7 +21,9 @@ class ServiceProxy: def __init__(self, http_options: dict[str, Any] | None = None, route_prefix: str = '/api/v1'): self.http_options = http_options or {} self.route_prefix = route_prefix - self.client = httpx.AsyncClient(timeout=None, trust_env=False) + queue = TaskQueueConfig() + budget = queue.queue_timeout + queue.effective_execution_timeout + self.client = httpx.AsyncClient(timeout=httpx.Timeout(budget, connect=10, pool=10), trust_env=False) async def close(self) -> None: await self.client.aclose() diff --git a/src/twinkle/server/state/__init__.py b/src/twinkle/server/state/__init__.py index 5f9a47549..3c270ea30 100644 --- a/src/twinkle/server/state/__init__.py +++ b/src/twinkle/server/state/__init__.py @@ -1,17 +1,26 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -from .backend import create_backend -from .base import BaseManager -from .cleanup_coordinator import ResourceCleanupCoordinator -from .count_publisher import ResourceCountPublisher -from .future_manager import FutureManager -from .model_manager import ModelManager -from .models import FutureRecord, ModelRecord, SamplingSessionRecord, SessionRecord -from .replica_registry import ReplicaRegistry -from .sampling_manager import SamplingSessionManager -from .server_state import ServerState, get_server_state, reset_server_state_cache -from .session_manager import SessionManager +import sys -__all__ = [ +from twinkle._lazy_module import _LazyModule + +sys.modules[__name__] = _LazyModule( + __name__, + __file__, { + 'backend': ['create_backend'], + 'base': ['BaseManager'], + 'cleanup_coordinator': ['ResourceCleanupCoordinator'], + 'count_publisher': ['ResourceCountPublisher'], + 'future_manager': ['FutureManager'], + 'model_manager': ['ModelManager'], + 'models': ['FutureRecord', 'ModelRecord', 'SamplingSessionRecord', 'SessionRecord'], + 'replica_registry': ['ReplicaRegistry'], + 'sampling_manager': ['SamplingSessionManager'], + 'server_state': ['ServerState', 'get_server_state', 'reset_server_state_cache'], + 'session_manager': ['SessionManager'], + }, + module_spec=__spec__) + +sys.modules[__name__].__all__ = [ # Pydantic record models 'SessionRecord', 'ModelRecord', diff --git a/src/twinkle/server/state/backend/base.py b/src/twinkle/server/state/backend/base.py index b65f7d2b2..4ed5bbc49 100644 --- a/src/twinkle/server/state/backend/base.py +++ b/src/twinkle/server/state/backend/base.py @@ -74,6 +74,8 @@ async def update_atomic( key: str, transform: Callable[[Any | None], Any | None], ttl: int | None = None, + *, + return_changed: bool = False, ) -> Any | None: """Read key, call ``transform(current_value)``, write the result atomically. @@ -87,6 +89,10 @@ async def update_atomic( and may raise :class:`ConcurrencyError` after exhausting internal retries (currently 16). + With ``return_changed=True``, return whether THIS call committed a + write, not the stored value. A lost write acknowledgement raises; it + must never be inferred as success from a later read. + ``transform`` must be picklable when running against a Ray-backed backend — pass module-level functions wrapped with ``functools.partial``, not lambdas or local closures. diff --git a/src/twinkle/server/state/backend/memory_backend.py b/src/twinkle/server/state/backend/memory_backend.py index dde709339..a92561264 100644 --- a/src/twinkle/server/state/backend/memory_backend.py +++ b/src/twinkle/server/state/backend/memory_backend.py @@ -84,14 +84,16 @@ async def update_atomic( key: str, transform: Callable[[Any | None], Any | None], ttl: int | None = None, + *, + return_changed: bool = False, ) -> Any | None: current = None if self._is_expired(key) else self._store[key][0] new_value = transform(current) if new_value is None: - return current + return False if return_changed else current expire_at = (time.time() + ttl) if ttl is not None else None self._store[key] = (new_value, expire_at) - return new_value + return True if return_changed else new_value async def mget(self, keys: list[str]) -> list[Any | None]: results: list[Any | None] = [] @@ -168,8 +170,10 @@ async def update_atomic( key: str, transform: Callable[[Any | None], Any | None], ttl: int | None = None, + *, + return_changed: bool = False, ) -> Any | None: - return await self._actor.update_atomic.remote(key, transform, ttl) + return await self._actor.update_atomic.remote(key, transform, ttl, return_changed=return_changed) async def mget(self, keys: list[str]) -> list[Any | None]: return await self._actor.mget.remote(keys) diff --git a/src/twinkle/server/state/backend/redis_backend.py b/src/twinkle/server/state/backend/redis_backend.py index 62c6a7be1..dc180e7d6 100644 --- a/src/twinkle/server/state/backend/redis_backend.py +++ b/src/twinkle/server/state/backend/redis_backend.py @@ -127,6 +127,8 @@ async def update_atomic( key: str, transform: Callable[[Any | None], Any | None], ttl: int | None = None, + *, + return_changed: bool = False, ) -> Any | None: """WATCH/GET/MULTI/SET/EXEC with bounded retries + jittered backoff on WatchError. @@ -149,7 +151,7 @@ async def update_atomic( if new_value is None: await pipe.unwatch() self._metrics.update_atomic_retries.observe(attempt) - return current + return False if return_changed else current pipe.multi() if ttl is not None: pipe.set(real_key, json.dumps(new_value), ex=ttl) @@ -158,7 +160,7 @@ async def update_atomic( executing = True await pipe.execute() self._metrics.update_atomic_retries.observe(attempt) - return new_value + return True if return_changed else new_value except WatchError as exc: # redis-py also wraps a lost watched connection as WatchError. # Once EXEC was sent, its outcome may be unknown, not aborted. diff --git a/src/twinkle/server/state/cleanup_coordinator.py b/src/twinkle/server/state/cleanup_coordinator.py index c6278c248..419846181 100644 --- a/src/twinkle/server/state/cleanup_coordinator.py +++ b/src/twinkle/server/state/cleanup_coordinator.py @@ -71,6 +71,7 @@ def __init__( self._on_become_leader = on_become_leader self._on_lose_leader = on_lose_leader + self._lifecycle_lock = asyncio.Lock() self._cleanup_task: asyncio.Task | None = None self._cleanup_running = False self._leader_id = uuid.uuid4().hex @@ -143,16 +144,14 @@ async def _cleanup_loop(self) -> None: async def _leader_loop(self) -> None: """Acquire and renew the cleanup-leader lease every LEASE_RENEW seconds.""" - await self._try_acquire_or_renew() # Race for leadership at startup while self._leader_running: try: - await asyncio.sleep(LEASE_RENEW) await self._try_acquire_or_renew() except asyncio.CancelledError: break except Exception as e: logger.warning(f'[ServerState Leader] renew error: {e}') - continue + await asyncio.sleep(LEASE_RENEW) async def _try_acquire_or_renew(self) -> None: was_leader = self._is_leader @@ -169,17 +168,8 @@ async def _try_acquire_or_renew(self) -> None: except Exception as e: logger.warning(f'[ServerState Leader] backend error during election: {e}') self._is_leader = False - if was_leader: - # Our renewal failed but our lease value may still be sitting in the - # backend, so a plain ``set_nx`` would keep returning False for up to - # LEASE_TTL and leadership would stall unclaimed. Best-effort delete - # ONLY when we were the leader (never steal a lease another replica - # legitimately holds), swallowing errors so a delete failure cannot - # escape the election loop. The next tick can then re-acquire. - try: - await self._backend.delete(LEADER_KEY) - except Exception: - pass + # A renewal failure does not prove current ownership. Leave the key + # untouched: our lease expires naturally, or a successor owns it. if self._is_leader and not was_leader: logger.info(f'[ServerState] became cleanup leader (id={self._leader_id[:8]})') @@ -199,36 +189,32 @@ async def start(self) -> bool: ``ServerState``) so the implementation and its guard can never drift apart into a double-start. """ - if self._cleanup_running: - return False - # Rebuild in-memory indexes from backend data before the loops start. - await self._managers['models'].rebuild_indexes() - self._cleanup_running = True - self._cleanup_task = asyncio.create_task(self._cleanup_loop()) - self._leader_running = True - self._leader_task = asyncio.create_task(self._leader_loop()) - return True + async with self._lifecycle_lock: + if self._cleanup_running: + return False + await self._managers['models'].rebuild_indexes() + self._cleanup_running = True + self._leader_running = True + self._cleanup_task = asyncio.create_task(self._cleanup_loop()) + self._leader_task = asyncio.create_task(self._leader_loop()) + return True async def stop(self) -> bool: """Stop the background cleanup + leader-election tasks. Returns ``False`` if not running.""" - if not self._cleanup_running: - return False - self._cleanup_running = False - if self._cleanup_task: - self._cleanup_task.cancel() - self._cleanup_task = None - self._leader_running = False - if self._leader_task: - self._leader_task.cancel() - self._leader_task = None - if self._is_leader: - # Release callback registration; the lease itself expires on its own TTL — - # update_atomic can't express "atomic delete", so we accept a short outage - # where the gauge reads 0 between leaders. - if self._on_lose_leader is not None: + async with self._lifecycle_lock: + if not self._cleanup_running: + return False + self._cleanup_running = self._leader_running = False + tasks = [task for task in (self._cleanup_task, self._leader_task) if task is not None] + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + self._cleanup_task = self._leader_task = None + was_leader, self._is_leader = self._is_leader, False + # Do not delete an uncertain lease; its original TTL is the fence. + if was_leader and self._on_lose_leader is not None: await self._on_lose_leader() - self._is_leader = False - return True + return True async def get_cleanup_stats(self) -> dict[str, Any]: """Get current cleanup configuration and resource counts.""" diff --git a/src/twinkle/server/state/future_manager.py b/src/twinkle/server/state/future_manager.py index 54bca4c70..d3aaae9b9 100644 --- a/src/twinkle/server/state/future_manager.py +++ b/src/twinkle/server/state/future_manager.py @@ -14,7 +14,7 @@ # Status sets used by the do-not-regress guard inside the atomic transform. _TERMINAL_STATUSES = frozenset({'completed', 'failed', 'cancelled'}) -_NON_TERMINAL_STATUSES = frozenset({'pending', 'queued', 'running'}) +_STATUS_ORDER = {'pending': 0, 'queued': 1, 'running': 2} def _future_record_transform( @@ -48,6 +48,11 @@ def _future_record_transform( if new_status != existing_status: logger.warning('future %s already terminal as %r; refusing %r', request_id, existing_status, new_status) return None + if new_status in _STATUS_ORDER and existing_status in _STATUS_ORDER: + if _STATUS_ORDER[new_status] <= _STATUS_ORDER[existing_status]: + return None + if new_status == 'cancelled' and existing_status == 'running': + return None if existing is None: record = FutureRecord( @@ -91,6 +96,13 @@ def _future_record_transform( _CANCELLABLE_STATUSES = frozenset({'pending', 'queued'}) +def _start_if_pending_transform(existing: dict | None, *, now: str) -> dict | None: + """Acquire execution only from pending/queued, using the existing record.""" + if existing is None or existing.get('status') not in _CANCELLABLE_STATUSES: + return None + return {**existing, 'status': 'running', 'queue_state': 'active', 'updated_at': now} + + def _cancel_if_not_started_transform(existing: dict | None, *, failure: dict, now: str) -> dict | None: """Atomic transform: cancel iff the task has not started (pending/queued). @@ -168,6 +180,14 @@ async def store_status( ), ) + async def try_start(self, request_id: str) -> bool: + """True only when this attempt atomically acquired execution rights.""" + return await self._backend.update_atomic( + self._make_key(request_id), + functools.partial(_start_if_pending_transform, now=_now_iso()), + return_changed=True, + ) + async def cancel_if_pending(self, request_id: str) -> str | None: """Cancel a task iff it has not started; return the resulting status. diff --git a/src/twinkle/server/state/server_state.py b/src/twinkle/server/state/server_state.py index 150d95f9f..dc9d683f2 100644 --- a/src/twinkle/server/state/server_state.py +++ b/src/twinkle/server/state/server_state.py @@ -498,6 +498,10 @@ async def release_training_seq(self, training_id: str, stream: str, sequence: in await self._backend.update_atomic(f'seq_ledger::{training_id}::{stream}', functools.partial(release_sequence, sequence=sequence, request_id=request_id)) + async def try_start_future(self, request_id: str) -> bool: + """Acquire task execution atomically, without inferring ownership from reads.""" + return await self._future_mgr.try_start(request_id) + async def cancel_future(self, request_id: str) -> dict[str, Any]: """Best-effort cancel: drop the task iff it has not started running. diff --git a/src/twinkle/server/task_queue/mixin.py b/src/twinkle/server/task_queue/mixin.py index ff3da6e79..b887ed20b 100644 --- a/src/twinkle/server/task_queue/mixin.py +++ b/src/twinkle/server/task_queue/mixin.py @@ -9,7 +9,6 @@ import asyncio import time -import traceback import uuid from collections.abc import Callable, Coroutine from typing import TYPE_CHECKING, Any @@ -128,6 +127,7 @@ def _init_task_queue( self._backend_gate = BackendGate(enable_admission_gate=enable_admission_gate) self._event_loop: asyncio.AbstractEventLoop | None = None + self._queue_tasks: set[asyncio.Task] = set() @property def task_queue_config(self) -> TaskQueueConfig: @@ -257,6 +257,7 @@ async def _prepare_queued_task(self, task: QueuedTask) -> None: async def _enqueue_prepared_task(self, task: QueuedTask) -> None: request_id, model_id = task.request_id, task.model_id + self._compute_worker.ensure_accepting() await self.state.store_future_status( request_id, TaskStatus.PENDING.value, @@ -266,10 +267,24 @@ async def _enqueue_prepared_task(self, task: QueuedTask) -> None: absolute_deadline=self._future_deadline(), ) + try: + await self._compute_worker.ensure_started() + except BaseException: + try: + await self.state.store_future_status( + request_id, + TaskStatus.FAILED.value, + model_id, + failure=FutureFailureRecord( + reason_code='orphaned_replica', + attribution='server', + message='Replica stopped before the task could be enqueued.')) + finally: + if task.on_finished is not None: + await task.on_finished(False, False) + raise queue_key = self._queue_key(model_id=model_id, token=task.token) self._compute_worker.ensure_queue_registered(queue_key) - await self._compute_worker.ensure_started() - q = self._compute_worker.get_queue(queue_key) self._compute_worker.track_task(request_id) await q.put(task) @@ -442,67 +457,12 @@ async def schedule_background_task( model_id: str | None = None, task_type: str | None = None, ) -> dict[str, Any]: - """Schedule a fire-and-forget task outside the serial queue. + """Run pure I/O outside serial compute, with the same supervised lifecycle. - The task is launched immediately as an asyncio task. This is suitable - for pure I/O and for an inference engine such as vLLM that performs its - own safe request concurrency and continuous batching. Stateful Model - operations must continue to use :meth:`schedule_task`. - - Status is tracked via state.store_future_status so clients can poll - progress through the same status endpoints as schedule_task(). - - Args: - coro_factory: Zero-argument callable that creates the coroutine. - model_id: Optional model id for result association. - task_type: Label for logging. - - Returns: - {'request_id': str, 'model_id': str | None} + Capacity, timeout, cancellation and shutdown are owned by ComputeWorker; + background jobs must not create a second unsupervised execution path. """ - request_id = f'req_{uuid.uuid4().hex}' - logger.info(f'[TaskQueue] Scheduling background task {request_id}, ' - f'type={task_type or "unknown"}, model_id={model_id}') - - await self.state.store_future_status( - request_id, - TaskStatus.RUNNING.value, - model_id, - queue_state=QueueState.ACTIVE.value, - replica_id=getattr(self, 'replica_id', None), - absolute_deadline=self._future_deadline(), - ) - - async def _run() -> None: - try: - result = await coro_factory() - await self.state.store_future_status( - request_id, - TaskStatus.COMPLETED.value, - model_id, - result=result, - queue_state=QueueState.ACTIVE.value, - ) - logger.info(f'[TaskQueue] Background task {request_id} completed, type={task_type or "unknown"}') - except Exception as exc: - failure = FutureFailureRecord( - reason_code='internal_error', - message=f'{type(exc).__name__}: {exc}'[:1024], - attribution='server', - diagnostic=traceback.format_exc(), - ) - await self.state.store_future_status( - request_id, - TaskStatus.FAILED.value, - model_id, - failure=failure, - queue_state=QueueState.ACTIVE.value, - ) - logger.error(f'[TaskQueue] Background task {request_id} FAILED, type={task_type or "unknown"}:\n' - f'{traceback.format_exc(limit=3)}') - - asyncio.create_task(_run()) - return {'request_id': request_id, 'model_id': model_id} + return await self.schedule_concurrent_task(coro_factory, model_id=model_id, task_type=task_type) async def _fail_queue_tasks_async(self, queue_key: str, reason: str) -> None: await self._compute_worker.fail_queue_tasks(queue_key, reason) @@ -515,13 +475,23 @@ def fail_pending_tasks_for_model(self, model_id: str, reason: str) -> None: return def _schedule() -> None: - asyncio.create_task(self._fail_queue_tasks_async(queue_key, reason)) + task = asyncio.create_task(self._fail_queue_tasks_async(queue_key, reason)) + self._queue_tasks.add(task) + task.add_done_callback(self._queue_task_done) self._event_loop.call_soon_threadsafe(_schedule) + def _queue_task_done(self, task: asyncio.Task) -> None: + self._queue_tasks.discard(task) + if not task.cancelled(): + error = task.exception() + if error is not None: + logger.error('Queue maintenance failed', exc_info=(type(error), error, error.__traceback__)) + async def shutdown_task_queue(self) -> None: """Gracefully shut down the compute queue and release resources.""" await self._rate_limiter.stop_cleanup_task() await self._compute_worker.stop() + await asyncio.gather(*self._queue_tasks, return_exceptions=True) self._backend_gate.shutdown() logger.debug('[TaskQueue] Task queue shutdown complete') diff --git a/src/twinkle/server/task_queue/worker.py b/src/twinkle/server/task_queue/worker.py index de5bc42e1..86aba9e61 100644 --- a/src/twinkle/server/task_queue/worker.py +++ b/src/twinkle/server/task_queue/worker.py @@ -120,14 +120,19 @@ def __init__( self._detached_tasks: dict[asyncio.Task, QueuedTask] = {} self._detached_reserved = 0 + def ensure_accepting(self) -> None: + if self._stopping: + raise ResourceExhaustedError('The replica is shutting down.') + async def ensure_started(self) -> None: - """Ensure the background worker coroutine is running.""" + """Ensure the background worker coroutine is running; never reopen a stopped replica.""" + self.ensure_accepting() if self._started and self._worker_task is not None and not self._worker_task.done(): return async with self._start_lock: + self.ensure_accepting() if self._started and self._worker_task is not None and not self._worker_task.done(): return - self._stopping = False self._worker_task = asyncio.create_task(self._worker_loop()) self._started = True @@ -151,10 +156,11 @@ async def stop(self) -> None: 'Sampler task stopped during replica shutdown.', QueueState.ACTIVE.value, reason_code='orphaned_replica') + for queue_key in list(self.task_queues): + await self.fail_queue_tasks(queue_key, 'Replica stopped before the task could execute.') self._worker_task = None self._started = False - self.task_queues.clear() - self.queue_order.clear() + self._request_ids.clear() def ensure_queue_registered(self, queue_key: str) -> None: """Register a new per-key queue if it does not yet exist.""" @@ -343,7 +349,8 @@ async def submit_detached(self, task: QueuedTask, queue_key: str, *, replica_id: def reserve_detached(self) -> None: """Reserve before the future write, so capacity rejection leaves no record.""" - if self._stopping or self._detached_reserved >= self._config.concurrent_task_limit: + self.ensure_accepting() + if self._detached_reserved >= self._config.concurrent_task_limit: raise ResourceExhaustedError('Concurrent compute capacity is unavailable on this replica.') self._detached_reserved += 1 @@ -393,6 +400,10 @@ async def run(): def done(completed: asyncio.Task) -> None: self._detached_tasks.pop(completed, None) self.release_detached() + if not completed.cancelled(): + error = completed.exception() + if error is not None: + logger.error('Detached task failed', exc_info=(type(error), error, error.__traceback__)) running.add_done_callback(done) @@ -418,19 +429,9 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: handler_attrs = {TOKEN_ID: task.token, MODEL_ID: task.model_id} handler_span_name = f'{self._deployment_name or "deployment"}.{task_type}' try: - await self._state.store_future_status( - task.request_id, - TaskStatus.RUNNING.value, - task.model_id, - queue_state=QueueState.ACTIVE.value, - ) - # A client cancel (or orphan cleanup) can win the race against the - # running write; the terminal record is kept, so do not start work - # nobody will collect. Once running is stored, cancel_if_pending can - # no longer succeed, so this read closes the window. - if await self._is_record_terminal(task.request_id): + if not await self._state.try_start_future(task.request_id): task_status = 'cancelled' - logger.info(f'[ComputeWorker] Task {task.request_id} became terminal before start; skipping.') + logger.info(f'[ComputeWorker] Task {task.request_id} was not acquired; skipping.') return with traced_operation('task_queue.execute', attrs=queue_attrs): logger.debug(f'[ComputeWorker] Task {task.request_id} executing, ' @@ -565,16 +566,6 @@ async def _try_run_one(self) -> bool: await self._fail_timed_out_task(task, queue_wait, q) continue # try the next queue - # A record already in a Terminal_State (e.g. written 'failed' by the - # state-hygiene orphan handling) must not be executed again. - if await self._is_record_terminal(task.request_id): - logger.info(f'[ComputeWorker] Task {task.request_id} already terminal on dequeue; skipping.') - q.task_done() - self._request_ids.discard(task.request_id) - if task.on_finished is not None: - await task.on_finished(False, False) - continue - # Execute the task (serial: stops after the first execution) await self._execute_task(task, queue_key, q) return True @@ -583,11 +574,8 @@ async def _try_run_one(self) -> bool: async def _is_record_terminal(self, request_id: str) -> bool: """True if the future record already holds a Terminal_State.""" - try: - record = await self._state.get_future(request_id) - except Exception: - return False - if not record: + record = await self._state.get_future(request_id) + if record is None: return False return record.get('status') in (TaskStatus.COMPLETED.value, TaskStatus.FAILED.value, TaskStatus.CANCELLED.value) diff --git a/src/twinkle/server/training/actor.py b/src/twinkle/server/training/actor.py index 2e109d1cf..6eb1245f6 100644 --- a/src/twinkle/server/training/actor.py +++ b/src/twinkle/server/training/actor.py @@ -36,6 +36,7 @@ def __init__(self, config, namespace): self._metrics = get_training_metrics() self._monitors = [] self._unreachable_since = {} + self._failures: dict[str, asyncio.Task] = {} self._initialization = asyncio.get_event_loop().create_task(self._initialize()) async def initialize(self): @@ -67,32 +68,42 @@ async def _heartbeat(self): logger.exception('Training state heartbeat failed; admission remains fenced by the state backend') await asyncio.sleep(5) continue - for instance in list(self.core.instances.values()): - if instance.state not in ('READY', 'RESETTING'): - continue - try: - identity = await self.driver.health(instance) - if identity['unhealthy']: - await self.core.fail_instance( - instance.instance_id, - FutureFailureRecord( - reason_code='orphaned_replica', - attribution='server', - message='Training runtime reported a failed worker')) - continue - await self.store.state.touch_replica_last_seen(instance.instance_id) - self._unreachable_since.pop(instance.instance_id, None) - except Exception: - since = self._unreachable_since.setdefault(instance.instance_id, time.monotonic()) - if time.monotonic() - since >= 30: - await self.core.fail_instance( - instance.instance_id, - FutureFailureRecord( - reason_code='orphaned_replica', - attribution='server', - message='Training runtime became unavailable')) + await asyncio.gather(*(self._check_instance_health(instance) + for instance in list(self.core.instances.values()) + if instance.state in ('READY', 'RESETTING'))) await asyncio.sleep(5) + async def _check_instance_health(self, instance): + try: + identity = await self.driver.health(instance) + if identity['unhealthy']: + self._fail_unhealthy(instance, 'Training runtime reported a failed worker') + return + await self.store.state.touch_replica_last_seen(instance.instance_id) + self._unreachable_since.pop(instance.instance_id, None) + except Exception: + since = self._unreachable_since.setdefault(instance.instance_id, time.monotonic()) + if time.monotonic() - since >= 30: + self._fail_unhealthy(instance, 'Training runtime became unavailable') + + def _fail_unhealthy(self, instance, message): + identity = instance.instance_id + if identity in self._failures: + return + task = asyncio.create_task( + self.core.fail_instance( + identity, FutureFailureRecord(reason_code='orphaned_replica', attribution='server', message=message))) + self._failures[identity] = task + + def done(completed): + self._failures.pop(identity, None) + if not completed.cancelled(): + error = completed.exception() + if error is not None: + logger.error('Unhealthy runtime cleanup failed', exc_info=(type(error), error, error.__traceback__)) + + task.add_done_callback(done) + async def _sweep(self): while True: started, status = time.monotonic(), 'success' @@ -168,6 +179,9 @@ async def diagnostics(self, base_model=None): } async def shutdown(self): + if not self._initialization.done(): + self._initialization.cancel() + await asyncio.gather(self._initialization, return_exceptions=True) restart = getattr(self, '_restart_task', None) if restart is not None: restart.cancel() @@ -176,3 +190,6 @@ async def shutdown(self): task.cancel() await asyncio.gather(*self._monitors, return_exceptions=True) await self.core.shutdown() + await asyncio.gather(*self._failures.values(), return_exceptions=True) + await self.core.wait_for_jobs() + await self.driver.wait_for_cleanup() diff --git a/src/twinkle/server/training/driver.py b/src/twinkle/server/training/driver.py index 0341c548a..955dd83d1 100644 --- a/src/twinkle/server/training/driver.py +++ b/src/twinkle/server/training/driver.py @@ -6,6 +6,7 @@ import os import re import time +from copy import deepcopy from twinkle.server.env_propagation import build_propagated_env_vars from twinkle.server.resources.allocation import DevicePool @@ -23,6 +24,7 @@ def __init__(self, config, store, namespace, controller_name): self.config, self.store = config, store self.namespace, self.controller_name = namespace, controller_name self._cleanup_tasks: set[asyncio.Task] = set() + self._deployments: dict[str, asyncio.Task] = {} self._metrics = get_training_metrics() self._catalog = TrainingCatalog(config) self.pool = DevicePool( @@ -73,6 +75,14 @@ async def deploy(self, instance, spec): if 'hybrid' in kwargs: kwargs['hybrid'] = spec.config.hybrid kwargs['train_mode'] = 'hybrid' if kwargs.get('hybrid') else spec.mode + # Each runtime has one identity. Serve autoscaling belongs to the CPU + # router; device-instance replica bounds are enforced by the controller. + options = deepcopy(spec.model.deploy_options) + options.pop('autoscaling_config', None) + options['num_replicas'] = 1 + actor_options = options.setdefault('ray_actor_options', {}) + actor_options.setdefault('num_cpus', 0.1) + actor_options.setdefault('runtime_env', {})['env_vars'] = env app = build_model_app( model_id=spec.model.model_id, backend=spec.model.backend, @@ -85,15 +95,7 @@ async def deploy(self, instance, spec): device_mesh={ **spec.config.device_mesh.model_dump(exclude_none=True), 'device_type': spec.model.device_type }, - deploy_options={ - 'num_replicas': 1, - 'ray_actor_options': { - 'num_cpus': 0.1, - 'runtime_env': { - 'env_vars': env - } - } - }, + deploy_options=options, queue_config=spec.config.queue_config, data_plane_url=spec.model.data_plane_url, training_context=context, @@ -104,6 +106,7 @@ async def deploy(self, instance, spec): deployment = asyncio.create_task( asyncio.to_thread( serve.run, app, name=name, route_prefix=f'/api/v1/training-runtime/{instance.instance_id}')) + self._deployments[instance.instance_id] = deployment try: await asyncio.shield(deployment) except asyncio.CancelledError: @@ -127,6 +130,14 @@ def done(completed): task.add_done_callback(done) raise + finally: + if deployment.done(): + self._deployments.pop(instance.instance_id, None) + + async def wait_for_cleanup(self): + """Join late deployment cleanup before the controller actor is retired.""" + while self._cleanup_tasks: + await asyncio.gather(*tuple(self._cleanup_tasks), return_exceptions=True) def handle(self, instance): from ray import serve @@ -146,6 +157,15 @@ async def shutdown(self, instance: InstanceRecord): import ray from ray import serve from ray.util import list_named_actors + deployment = self._deployments.get(instance.instance_id) + if deployment is not None: + # Cancellation of Serve's caller does not stop its deployment thread. + # Never certify death while that thread can still create resources. + try: + await asyncio.shield(deployment) + except Exception: + pass + self._deployments.pop(instance.instance_id, None) name = self.application_name(instance) owned = [ actor for actor in list_named_actors(all_namespaces=True) diff --git a/src/twinkle/utils/loader.py b/src/twinkle/utils/loader.py index 0ec971cd5..da8321bf1 100644 --- a/src/twinkle/utils/loader.py +++ b/src/twinkle/utils/loader.py @@ -6,7 +6,6 @@ from types import ModuleType from typing import List, Type, TypeVar, Union -from ..hub import HFHub, MSHub from .unsafe import trust_remote_code T = TypeVar('T') @@ -17,6 +16,8 @@ class Plugin: @staticmethod def load_plugin(plugin_id: str, plugin_base: Type[T], **kwargs) -> Type[T]: + from twinkle.hub.hub import HFHub, MSHub + if plugin_id.startswith('hf://'): plugin_dir = HFHub.download_model(plugin_id[len('hf://'):], **kwargs) elif plugin_id.startswith('ms://'): diff --git a/tests/server/config/test_training_config.py b/tests/server/config/test_training_config.py index 0a313079c..8bb8dfc72 100644 --- a/tests/server/config/test_training_config.py +++ b/tests/server/config/test_training_config.py @@ -185,3 +185,92 @@ def builder(**kwargs): launcher._deploy_application(config.applications[1]) assert captured['training_models'] == ['Qwen/small'] assert captured['supported_models'] == ['teacher'] + + +def legacy_application(): + return {'name': 'model', 'route_prefix': '/api/v1/model/small', 'import_path': 'model', 'args': { + 'model_id': 'weights', 'backend': 'mock', 'train_mode': 'lora', 'max_length': 128, + 'device_group': {'ranks': 2, 'device_type': 'cpu'}, 'device_mesh': {'dp_size': 2}, + 'env_vars': {'USER_OPTION': 'args-wins'}, 'adapter_config': {'adapter_timeout': 30}, + }, 'deployments': [{ + 'name': 'ModelManagement', 'max_ongoing_requests': 17, + 'autoscaling_config': {'min_replicas': 1, 'max_replicas': 2, 'target_ongoing_requests': 16}, + 'ray_actor_options': {'num_cpus': 2, 'runtime_env': { + 'env_vars': {'USER_OPTION': 'deployment', 'OTHER_OPTION': 'retained'}, 'pip': ['example==1.0'] + }}, + }]} + + +def test_legacy_options_reach_cpu_router_without_changing_yaml(monkeypatch): + from types import SimpleNamespace + from twinkle.server.launcher.server_launcher import ServerLauncher + from twinkle.server.training.actor import TrainingControllerActor + import ray + from ray import serve + config = ServerConfig(applications=[legacy_application()]) + remote = SimpleNamespace(initialize=SimpleNamespace(remote=lambda: None)) + monkeypatch.setattr(TrainingControllerActor, 'options', + lambda **kwargs: SimpleNamespace(remote=lambda *args: remote)) + monkeypatch.setattr(ray, 'get', lambda *args, **kwargs: None) + captured = {} + + def router(*args): + captured.update(args[-1]) + return 'app' + + monkeypatch.setattr('twinkle.server.training.router.build_training_router', router) + monkeypatch.setattr(serve, 'run', lambda *args, **kwargs: None) + ServerLauncher(config)._deploy_training() + assert captured['max_ongoing_requests'] == 17 + assert captured['autoscaling_config']['target_ongoing_requests'] == 16 + actor = captured['ray_actor_options'] + assert actor['num_cpus'] == 2 and actor['runtime_env']['pip'] == ['example==1.0'] + assert actor['runtime_env']['env_vars']['USER_OPTION'] == 'args-wins' + assert actor['runtime_env']['env_vars']['OTHER_OPTION'] == 'retained' + assert config.to_yaml_dict()['applications'][0]['deployments'][0] == legacy_application()['deployments'][0] + + +@pytest.mark.asyncio +async def test_legacy_options_reach_runtime_with_one_controller_owned_replica(monkeypatch): + import importlib + from twinkle.server.state.training_records import InstanceRecord + from twinkle.server.training.driver import RayInstanceDriver + from ray import serve + config = ServerConfig(applications=[legacy_application()]) + model = config.training_config.models[0] + driver = RayInstanceDriver(config.training_config, None, 'test', 'controller') + instance = InstanceRecord(instance_id='instance', controller_epoch='epoch', base_model=model.name, + mode='lora', startup_deadline=0, devices=['0', '1']) + captured = {} + + def builder(**kwargs): + captured.update(kwargs) + return 'app' + + monkeypatch.setattr('twinkle.server.model.app.build_model_app', builder) + monkeypatch.setattr(importlib.import_module('ray.util.placement_group'), 'get_placement_group', lambda name: 'pg') + monkeypatch.setattr(serve, 'run', lambda *args, **kwargs: None) + await driver.deploy(instance, driver._catalog.get(model.name, 'lora')) + options = captured['deploy_options'] + assert options['num_replicas'] == 1 and 'autoscaling_config' not in options + assert options['max_ongoing_requests'] == 17 + assert options['ray_actor_options']['num_cpus'] == 2 + assert options['ray_actor_options']['runtime_env']['pip'] == ['example==1.0'] + assert captured['max_length'] == 128 + assert captured['nproc_per_node'] == 2 + assert captured['queue_config'] == model.modes.lora.queue_config + assert not driver._deployments + + +def test_mutated_configuration_does_not_reuse_a_stale_catalog(): + config = ServerConfig(applications=[legacy_application()]) + assert config.training_config.models[0].modes.lora.min_replicas == 1 + config.applications[0].args.min_replicas = 0 + assert config.training_config.models[0].modes.lora.min_replicas == 0 + + +def test_unsupported_adapter_parameters_are_not_silently_ignored(): + app = legacy_application() + app['args']['adapter_config']['misspelled_timeout'] = 12 + with pytest.raises(ValidationError, match='misspelled_timeout'): + ServerConfig(applications=[app]) diff --git a/tests/server/gateway/test_proxy.py b/tests/server/gateway/test_proxy.py index a111b5096..caa35deee 100644 --- a/tests/server/gateway/test_proxy.py +++ b/tests/server/gateway/test_proxy.py @@ -7,8 +7,11 @@ """ from __future__ import annotations +import asyncio import json +import httpx + import pytest from starlette.requests import Request from unittest.mock import AsyncMock @@ -51,6 +54,37 @@ def test_prepare_headers_sets_multiplex_from_request_id(): assert 'host' not in {k.lower() for k in headers} +@pytest.mark.asyncio +async def test_proxy_waits_are_bounded_and_timeout_keeps_existing_wire(): + closed = asyncio.Event() + + async def stall(reader, writer): + try: + await reader.read(4096) + await reader.read() + finally: + writer.close() + await writer.wait_closed() + closed.set() + + server = await asyncio.start_server(stall, '127.0.0.1', 0) + port = server.sockets[0].getsockname()[1] + proxy = ServiceProxy(http_options={'host': '127.0.0.1', 'port': port}) + assert all(value is not None and value > 0 for value in proxy.client.timeout.as_dict().values()) + proxy.client.timeout = httpx.Timeout(0.05) + request = _make_request(headers=[(H_REQUEST_ID.encode(), b'req-timeout')]) + try: + reply = await asyncio.wait_for(proxy.proxy_request(request, 'tinker/forward', 'small', 'model', + body_override=b'{}'), 1) + assert reply.status_code == 502 + assert json.loads(reply.body)['request_id'] == 'req-timeout' + finally: + await proxy.close() + server.close() + await server.wait_closed() + await asyncio.wait_for(closed.wait(), 1) + + @pytest.mark.asyncio async def test_proxy_request_502_fallback_returns_error_payload(): proxy = ServiceProxy(http_options={'host': 'localhost', 'port': 8000}, route_prefix='/api/v1') diff --git a/tests/server/integration/on_demand_training/fragmentation_probe.py b/tests/server/integration/on_demand_training/fragmentation_probe.py index f88132ea3..d715d7e8c 100644 --- a/tests/server/integration/on_demand_training/fragmentation_probe.py +++ b/tests/server/integration/on_demand_training/fragmentation_probe.py @@ -24,7 +24,9 @@ async def main(): cluster.add_node(num_cpus=4, num_gpus=4, resources={'left': 1}, include_dashboard=False, dashboard_port=0) cluster.add_node(num_cpus=4, num_gpus=4, resources={'right': 1}) ray.init(address=cluster.address, namespace='fragmentation_' + uuid.uuid4().hex, log_to_driver=False) - pool = DevicePool(ray.get_runtime_context().namespace, allocation_timeout=1, shutdown_timeout=10) + # Keep fragmentation rejection fast, but allow a cold Ray worker to + # import the allocation probe before confirming a successful placement. + pool = DevicePool(ray.get_runtime_context().namespace, allocation_timeout=1, shutdown_timeout=60) spec = TrainingCatalog( TrainingConfig(models=[{ 'name': 'two', diff --git a/tests/server/integration/on_demand_training/http_probe.py b/tests/server/integration/on_demand_training/http_probe.py index ebeda6cc8..b469e7593 100644 --- a/tests/server/integration/on_demand_training/http_probe.py +++ b/tests/server/integration/on_demand_training/http_probe.py @@ -79,16 +79,17 @@ def main(): 'args': { 'model_id': 'other', 'backend': 'mock', - 'device_type': 'cpu', - 'modes': { - 'full': { - 'num_devices': 1, - 'device_mesh': { - 'dp_size': 1 - } - } - } + 'train_mode': 'full', + 'device_group': {'ranks': 1, 'device_type': 'cpu'}, + 'device_mesh': {'dp_size': 1}, + 'min_replicas': 0, + 'max_replicas': 1, }, + 'deployments': [{ + 'max_ongoing_requests': 8, + 'autoscaling_config': {'min_replicas': 1, 'max_replicas': 1, 'target_ongoing_requests': 4}, + 'ray_actor_options': {'num_cpus': 0.1, 'runtime_env': {'env_vars': {'LEGACY_CONFIG_PROBE': '1'}}}, + }], }, { 'name': 'gateway', 'import_path': 'server', diff --git a/tests/server/integration/test_mock_mode_startup.py b/tests/server/integration/test_mock_mode_startup.py index b1c38259a..e377afff2 100644 --- a/tests/server/integration/test_mock_mode_startup.py +++ b/tests/server/integration/test_mock_mode_startup.py @@ -146,6 +146,16 @@ def test_mock_mode_reaches_ready_and_is_deterministic(ray_cluster, monkeypatch) # separately because tinker's polling can stall the test budget on flaky # cross-replica state. if os.environ.get('TWINKLE_TEST_TINKER', '0') == '1': + # The native patch dirties the resident instance; release replaces it. + # A create during that replacement is correctly rejected, not queued. + deadline = time.monotonic() + READY_BUDGET_SECONDS + while time.monotonic() < deadline: + health = _http(f'{base}/api/v1/model/mock-model/healthz').json() + if health.get('healthy'): + break + time.sleep(0.2) + else: + raise AssertionError('Resident replacement did not become ready') _exercise_tinker_client(base) @@ -208,7 +218,11 @@ def _exercise_twinkle_clients(base: str) -> None: # must not publish files to an external account. # --- sampler service --- - sampler = vLLMSampler(model_id='mock-model') + from twinkle_client.http import ClientTransport + from twinkle_client.http.context import capture_transport + # Cold PEFT imports on NAS can exceed the default 90-second HTTP timeout. + sampler_transport = ClientTransport(capture_transport().context, timeout=600) + sampler = vLLMSampler(model_id='mock-model', transport=sampler_transport) # Sampler /add_adapter_to_sampler reuses the model-side ``AddAdapterRequest`` # shape (``config: str``), so send JSON. Skip a real ``LoraConfig`` here # because its ``runtime_config`` member isn't JSON-serializable. @@ -236,6 +250,9 @@ def _exercise_twinkle_clients(base: str) -> None: ) sampler.apply_patch('NoopPatch') + sampler_transport.close() + model.remove_adapter() + twinkle.close() def _exercise_tinker_client(base: str) -> None: @@ -276,6 +293,12 @@ def _exercise_tinker_client(base: str) -> None: sampler_ckpt = training.save_weights_for_sampler(name='step-1').result() # path mode: save_weights_for_sampler(name) must return path != None assert sampler_ckpt.path is not None + from pathlib import Path + from peft import LoraConfig + from twinkle.server.checkpoint import create_checkpoint_manager + checkpoint_manager = create_checkpoint_manager('EMPTY_TOKEN', client_type='tinker') + sampler_dir = Path(checkpoint_manager.resolve_checkpoint_dir(sampler_ckpt.path)).resolve() + LoraConfig(r=4, target_modules='all-linear').save_pretrained(sampler_dir) # Gateway's /asample resolves ``base_model`` from ``body.base_model`` or # ``sampling_session_id``; pass it explicitly because the SDK only sets # ``model_path`` and the gateway doesn't parse ``twinkle://`` URIs. @@ -289,6 +312,11 @@ def _exercise_tinker_client(base: str) -> None: # sampling_session_seq_id mode: save_weights_and_get_sampling_client() # must return path == None and sampling_session_id != None (asserted by SDK internally) sampling_client = training.save_weights_and_get_sampling_client() + # Mock saves contain no tensors. Mark the newly versioned sampler checkpoint + # as LoRA, just as the native smoke does, rather than pretending it is full weights. + for checkpoint in sampler_dir.parent.iterdir(): + if checkpoint.is_dir(): + LoraConfig(r=4, target_modules='all-linear').save_pretrained(checkpoint) sampling_client.sample( prompt=types.ModelInput.from_ints([1, 2, 3]), num_samples=1, @@ -299,3 +327,5 @@ def _exercise_tinker_client(base: str) -> None: # the wire-level handler is /tinker/save_weights either way. ckpt = training.save_state(name='step-2').result() training.load_state(ckpt.path).result() + from tests.server.integration.on_demand_training.on_demand_sdk import unload_tinker_model + unload_tinker_model(client, training.model_id) diff --git a/tests/server/sampler/test_sampling_concurrency.py b/tests/server/sampler/test_sampling_concurrency.py index cf57479b0..88a7e6cb8 100644 --- a/tests/server/sampler/test_sampling_concurrency.py +++ b/tests/server/sampler/test_sampling_concurrency.py @@ -40,6 +40,13 @@ def __init__(self): async def store_future_status(self, request_id, status, model_id=None, **kwargs): self.records[request_id] = {'status': status, **kwargs} + async def try_start_future(self, request_id): + record = self.records.get(request_id) + if record is None or record['status'] not in ('pending', 'queued'): + return False + record['status'] = 'running' + return True + async def get_future(self, request_id): return self.records.get(request_id) diff --git a/tests/server/state/fake_backend.py b/tests/server/state/fake_backend.py index eaac98833..10f5c9e56 100644 --- a/tests/server/state/fake_backend.py +++ b/tests/server/state/fake_backend.py @@ -67,13 +67,15 @@ async def update_atomic( key: str, transform: Callable[[Any | None], Any | None], ttl: int | None = None, + *, + return_changed: bool = False, ) -> Any | None: current = await self.get(key) updated = transform(current) if updated is None: - return current + return False if return_changed else current await self.set(key, updated, ttl) - return updated + return True if return_changed else updated async def flush_all(self) -> None: self._store.clear() diff --git a/tests/server/state/test_future_lifecycle.py b/tests/server/state/test_future_lifecycle.py index 9728780f3..d9c0e63ff 100644 --- a/tests/server/state/test_future_lifecycle.py +++ b/tests/server/state/test_future_lifecycle.py @@ -7,86 +7,20 @@ """ from __future__ import annotations +import asyncio import time -from collections.abc import Callable -from fnmatch import fnmatch -from typing import Any from unittest import mock import pytest -from twinkle.server.state.backend.base import StateBackend +from tests.server.state.fake_backend import FakeBackend from twinkle.server.state.future_manager import FutureManager from twinkle.server.state.models import FutureFailureRecord -class _FakeBackend(StateBackend): - """Dict-backed StateBackend. Single event loop, so ``update_atomic`` is atomic - simply by not awaiting between read and write -- the real backends' guarantee.""" - - def __init__(self) -> None: - self._store: dict[str, tuple[Any, float | None]] = {} - - def _is_expired(self, key: str) -> bool: - entry = self._store.get(key) - if entry is None: - return True - _, expire_at = entry - if expire_at is not None and time.time() >= expire_at: - del self._store[key] - return True - return False - - async def set(self, key: str, value: Any, ttl: int | None = None) -> None: - self._store[key] = (value, (time.time() + ttl) if ttl is not None else None) - - async def get(self, key: str) -> Any | None: - return None if self._is_expired(key) else self._store[key][0] - - async def mget(self, keys: list[str]) -> list[Any | None]: - return [await self.get(key) for key in keys] - - async def delete(self, key: str) -> None: - self._store.pop(key, None) - - async def exists(self, key: str) -> bool: - return not self._is_expired(key) - - async def keys(self, pattern: str) -> list[str]: - return [k for k in list(self._store) if not self._is_expired(k) and fnmatch(k, pattern)] - - async def count(self, pattern: str) -> int: - return len(await self.keys(pattern)) - - async def set_nx(self, key: str, value: Any, ttl: int | None = None) -> bool: - if not self._is_expired(key): - return False - await self.set(key, value, ttl) - return True - - async def update_atomic( - self, - key: str, - transform: Callable[[Any | None], Any | None], - ttl: int | None = None, - ) -> Any | None: - current = await self.get(key) - updated = transform(current) - if updated is None: - return current - await self.set(key, updated, ttl) - return updated - - async def close(self) -> None: - pass - - async def health_check(self) -> bool: - return True - - @pytest.fixture def manager(): - return FutureManager(_FakeBackend(), expiration_timeout=300.0) + return FutureManager(FakeBackend(), expiration_timeout=300.0) async def _store(manager, request_id, status, *, replica_id=None, absolute_deadline=None): @@ -203,7 +137,7 @@ async def test_stored_timestamps_align_with_wall_clock_regardless_of_host_tz(man @pytest.mark.asyncio async def test_claim_seq_dedups_then_release_readmits(): from twinkle.server.state.server_state import ServerState - state = ServerState(backend=_FakeBackend()) + state = ServerState(backend=FakeBackend()) # First claim of a (session, seq_id) is unseen -> None, caller proceeds to enqueue. assert await state.claim_seq('seq::s1::1', 'reqA', ttl=60) is None # A duplicate claim returns the original request_id -> caller returns its envelope. @@ -218,7 +152,7 @@ async def test_claim_seq_dedups_then_release_readmits(): @pytest.mark.asyncio async def test_cancel_drops_pending_but_never_running(): from twinkle.server.state.server_state import ServerState - state = ServerState(backend=_FakeBackend()) + state = ServerState(backend=FakeBackend()) # pending -> cancel drops it to a terminal domain failure. await state.store_future_status('rp', 'pending', 'm1') assert await state.cancel_future('rp') == {'cancelled': True, 'state': 'cancelled'} @@ -247,3 +181,37 @@ def test_cancelled_record_maps_to_error_envelope(): env = envelope_from_record('rc', rec) assert env.status == 'cancelled' assert env.error is not None and env.error.error_code == 499 + + +@pytest.mark.asyncio +async def test_late_queued_and_repeated_running_cannot_regress_or_renew(manager): + await _store(manager, 'race', 'pending') + assert await manager.try_start('race') + before = (await manager.get('race')).model_dump() + await _store(manager, 'race', 'queued') + await _store(manager, 'race', 'running') + assert (await manager.get('race')).model_dump() == before + assert await manager.cancel_if_pending('race') == 'running' + assert not await manager.try_start('race') + assert not await manager.try_start('missing') + + +@pytest.mark.asyncio +@pytest.mark.parametrize('cancel_first', [False, True]) +async def test_start_and_cancel_have_one_atomic_winner(manager, cancel_first): + await _store(manager, 'race', 'queued') + operations = [manager.try_start('race'), manager.cancel_if_pending('race')] + if cancel_first: + operations.reverse() + results = await asyncio.gather(*operations) + started = results[1] if cancel_first else results[0] + status = (await manager.get('race')).status + assert started == (status == 'running') + assert status in ('running', 'cancelled') + assert not await manager.try_start('race') + + +@pytest.mark.asyncio +async def test_concurrent_starters_only_one_obtains_execution(manager): + await _store(manager, 'race', 'pending') + assert sum(await asyncio.gather(*(manager.try_start('race') for _ in range(20)))) == 1 diff --git a/tests/server/state/test_leader_election.py b/tests/server/state/test_leader_election.py index 792989c30..216105e99 100644 --- a/tests/server/state/test_leader_election.py +++ b/tests/server/state/test_leader_election.py @@ -11,10 +11,14 @@ import asyncio import pytest +from types import SimpleNamespace +from unittest.mock import AsyncMock + +from tests.server.state.fake_backend import FakeBackend from twinkle.server.state import ServerState from twinkle.server.state.backend.memory_backend import RayActorBackend -from twinkle.server.state.cleanup_coordinator import LEADER_KEY, LEASE_RENEW +from twinkle.server.state.cleanup_coordinator import LEADER_KEY, LEASE_RENEW, ResourceCleanupCoordinator from twinkle.server.telemetry import MetricsRegistry @@ -100,13 +104,7 @@ async def test_renew_keeps_leader() -> None: @pytest.mark.asyncio async def test_leader_recovers_after_renewal_failure() -> None: - """Regression: when a leader's renewal raises, it releases - the lease best-effort so the very next election tick re-acquires leadership - without waiting LEASE_TTL. - - On unfixed code ``_is_leader`` flips to False while the lease value lingers - in the backend, so ``set_nx`` keeps returning False for up to LEASE_TTL. - """ + """An unconfirmed renewal fences cleanup until the original lease expires.""" backend = RayActorBackend() state = ServerState(backend=backend, cleanup_interval=600.0) @@ -128,11 +126,13 @@ async def flaky_update_atomic(*args, **kwargs): from unittest import mock with mock.patch.object(backend, 'update_atomic', side_effect=flaky_update_atomic): await state._try_acquire_or_renew() - # Renewal failed → no longer leader, and the stale lease was released. assert state._is_leader is False - # The lease key was deleted best-effort, so the next tick re-acquires - # immediately (no LEASE_TTL wait, no lingering self-owned lease). + assert await backend.get(LEADER_KEY) == state._leader_id + await state._try_acquire_or_renew() + assert state._is_leader is False + await backend.set(LEADER_KEY, state._leader_id, ttl=1) + await asyncio.sleep(1.1) await state._try_acquire_or_renew() assert state._is_leader is True @@ -159,6 +159,75 @@ async def test_renewal_failure_does_not_steal_other_leader_lease() -> None: assert await backend.get(LEADER_KEY) == leader_value, 'follower stole the real leader lease' +def coordinator(backend=None): + models = SimpleNamespace(rebuild_indexes=AsyncMock()) + return ResourceCleanupCoordinator(backend or FakeBackend(), {'models': models}, + cleanup_interval=600, expiration_timeout=300, + sweep_processor_quotas=AsyncMock()) + + +@pytest.mark.asyncio +async def test_stale_leader_renewal_failure_preserves_successor(monkeypatch): + backend = FakeBackend() + cleanup = coordinator(backend) + await cleanup._try_acquire_or_renew() + assert cleanup.is_leader + + async def successor_then_fail(*args, **kwargs): + await backend.set(LEADER_KEY, 'successor', ttl=30) + raise ConnectionError('renewal acknowledgement lost') + + monkeypatch.setattr(backend, 'update_atomic', successor_then_fail) + await cleanup._try_acquire_or_renew() + assert not cleanup.is_leader + assert await backend.get(LEADER_KEY) == 'successor' + + +@pytest.mark.asyncio +async def test_cleanup_concurrent_start_rebuilds_once_and_stop_joins(): + cleanup = coordinator() + entered, proceed = asyncio.Event(), asyncio.Event() + + async def rebuild(): + entered.set() + await proceed.wait() + + cleanup._managers['models'].rebuild_indexes.side_effect = rebuild + first = asyncio.create_task(cleanup.start()) + await entered.wait() + second = asyncio.create_task(cleanup.start()) + proceed.set() + assert await asyncio.gather(first, second) == [True, False] + tasks = (cleanup._cleanup_task, cleanup._leader_task) + assert await cleanup.stop() + assert all(task.done() for task in tasks) + assert not await cleanup.stop() + cleanup._managers['models'].rebuild_indexes.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_failed_cleanup_start_can_retry(): + cleanup = coordinator() + cleanup._managers['models'].rebuild_indexes.side_effect = [RuntimeError('offline'), None] + with pytest.raises(RuntimeError, match='offline'): + await cleanup.start() + assert cleanup._cleanup_task is None and cleanup._leader_task is None + assert await cleanup.start() + await cleanup.stop() + + +@pytest.mark.asyncio +async def test_lazy_cleanup_retries_after_failed_start(): + from twinkle.server.deployment import LazyCleanupMixin + client = LazyCleanupMixin() + client.state = SimpleNamespace(start_cleanup_task=AsyncMock(side_effect=[RuntimeError('offline'), True])) + await client._ensure_state_cleanup_started() + assert not client._state_cleanup_started + await client._ensure_state_cleanup_started() + assert client._state_cleanup_started + assert client.state.start_cleanup_task.await_count == 2 + + # ============================================================ # Metrics ObservableGauge (merged from test_metrics_observable_gauge) # ============================================================ diff --git a/tests/server/state/test_update_atomic.py b/tests/server/state/test_update_atomic.py index d03b3a79e..387f69302 100644 --- a/tests/server/state/test_update_atomic.py +++ b/tests/server/state/test_update_atomic.py @@ -129,6 +129,43 @@ async def test_key_prefix_is_hidden_from_logical_keys(backend_pair) -> None: assert await first.keys('session::*') == ['session::one'] +@pytest.mark.asyncio +async def test_atomic_update_reports_this_call_not_the_existing_value(backend_pair): + first, second = backend_pair + assert await first.update_atomic('claim', functools.partial(_replace_with, value=True), + return_changed=True) is True + assert await second.update_atomic('claim', _no_op, return_changed=True) is False + assert await second.get('claim') is True + assert await second.update_atomic('missing', _no_op, return_changed=True) is False + + +@pytest.mark.asyncio +async def test_future_execution_is_claimed_once_across_backend_handles(backend_pair): + from twinkle.server.state.future_manager import FutureManager + first, second = [FutureManager(backend, 300) for backend in backend_pair] + await first.store_status('start', 'pending', 'm') + starts = await asyncio.gather(*(manager.try_start('start') for manager in [first, second] * 8)) + assert sum(starts) == 1 + await second.store_status('start', 'queued', 'm') + assert await first.cancel_if_pending('start') == 'running' + assert (await first.get('start')).status == 'running' + + +@pytest.mark.asyncio +@pytest.mark.parametrize('cancel_first', [False, True]) +async def test_future_start_cancel_race_across_backend_handles(backend_pair, cancel_first): + from twinkle.server.state.future_manager import FutureManager + first, second = [FutureManager(backend, 300) for backend in backend_pair] + await first.store_status('start', 'queued', 'm') + calls = [first.try_start('start'), second.cancel_if_pending('start')] + if cancel_first: + calls.reverse() + results = await asyncio.gather(*calls) + started = results[1] if cancel_first else results[0] + assert started == ((await first.get('start')).status == 'running') + assert not await second.try_start('start') + + # ---------- Memory backend ----------------------------------------------- # diff --git a/tests/server/test_dependencies.py b/tests/server/test_dependencies.py index 3100a71b4..8853f9031 100644 --- a/tests/server/test_dependencies.py +++ b/tests/server/test_dependencies.py @@ -1,6 +1,10 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Executable dependency architecture for every ``twinkle.server`` component.""" import grimp +import subprocess +import sys + +import pytest _ROOT = 'twinkle.server' _ROOT_COMPONENT = '__root__' @@ -29,6 +33,42 @@ def _component(module: str) -> str: return module.removeprefix(f'{_ROOT}.').split('.', 1)[0] +@pytest.mark.parametrize('statement', [ + 'import twinkle.server.exceptions', + 'from twinkle.server.config import ServerConfig', + 'from twinkle.server.state.models import FutureRecord', + 'from twinkle.server.state import SessionRecord', +]) +def test_lightweight_imports_do_not_load_runtime(statement): + code = statement + ''' +import sys +assert not any(name in sys.modules for name in ( + 'torch', 'ray', 'twinkle.server.launcher', 'twinkle.server.training.actor', 'twinkle.hub.hub' +)), sorted(sys.modules) +''' + result = subprocess.run([sys.executable, '-c', code], capture_output=True, text=True, timeout=60) + assert result.returncode == 0, result.stderr + + +def test_hub_first_import_does_not_reenter_plugin_loader(): + code = ''' +from twinkle.hub.hub import HFHub +from twinkle.utils.loader import Plugin +assert HFHub.__module__ == 'twinkle.hub.hub' +assert Plugin.__module__ == 'twinkle.utils.loader' +''' + result = subprocess.run([sys.executable, '-c', code], capture_output=True, text=True, timeout=120) + assert result.returncode == 0, result.stderr + + +def test_public_launcher_is_the_actual_implementation(): + import twinkle.server + from twinkle.server.launcher.server_launcher import ServerLauncher, launch_server + assert twinkle.server.ServerLauncher is ServerLauncher + assert twinkle.server.launch_server is launch_server + assert twinkle.server.__all__ == ['ServerLauncher', 'launch_server'] + + def test_server_components_form_declared_dag(): """Reject cycles, upward imports, same-layer coupling, and unclassified components.""" graph = grimp.build_graph('twinkle', include_external_packages=False) diff --git a/tests/server/training/test_actor_initialization.py b/tests/server/training/test_actor_initialization.py index 211d925e1..10c748341 100644 --- a/tests/server/training/test_actor_initialization.py +++ b/tests/server/training/test_actor_initialization.py @@ -47,7 +47,81 @@ async def test_shutdown_cancels_pending_restart(): cls = TrainingControllerActor.__ray_metadata__.modified_class wait = asyncio.Event() restart = asyncio.create_task(wait.wait()) - actor = SimpleNamespace(core=SimpleNamespace(shutdown=AsyncMock()), _monitors=[], _restart_task=restart) + initialization = asyncio.get_running_loop().create_future() + initialization.set_result(None) + actor = SimpleNamespace(driver=SimpleNamespace(wait_for_cleanup=AsyncMock()), + core=SimpleNamespace(shutdown=AsyncMock(), wait_for_jobs=AsyncMock()), + _monitors=[], _restart_task=restart, _initialization=initialization, _failures={}) await cls.shutdown(actor) assert restart.cancelled() actor.core.shutdown.assert_awaited_once() + actor.driver.wait_for_cleanup.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_heartbeat_keeps_healthy_peer_alive_while_probe_and_close_block(): + cls = TrainingControllerActor.__ray_metadata__.modified_class + good, bad = [SimpleNamespace(instance_id=name, state='READY') for name in ('good', 'bad')] + probe_gate, close_gate = asyncio.Event(), asyncio.Event() + healthy_seen, close_started = asyncio.Event(), asyncio.Event() + + async def health(instance): + if instance is bad: + await probe_gate.wait() + return {'unhealthy': instance is bad} + + async def fail(*args): + close_started.set() + await close_gate.wait() + + async def touch(identity): + if identity == 'good': + healthy_seen.set() + + actor = SimpleNamespace( + store=SimpleNamespace(backend=SimpleNamespace(get=AsyncMock(return_value='active')), + key=lambda *args: 'active', state=SimpleNamespace(touch_replica_last_seen=touch)), + core=SimpleNamespace(epoch='controller', instances={'bad': bad, 'good': good}, fail_instance=fail), + driver=SimpleNamespace(health=health), _unreachable_since={}, _failures={}) + actor._check_instance_health = lambda instance: cls._check_instance_health(actor, instance) + actor._fail_unhealthy = lambda instance, message: cls._fail_unhealthy(actor, instance, message) + monitor = asyncio.create_task(cls._heartbeat(actor)) + try: + await asyncio.wait_for(healthy_seen.wait(), 1) + assert not probe_gate.is_set() + probe_gate.set() + await asyncio.wait_for(close_started.wait(), 1) + # Closing a bad peer does not occupy the heartbeat monitor. + healthy_seen.clear() + await cls._check_instance_health(actor, good) + assert healthy_seen.is_set() + cls._fail_unhealthy(actor, bad, 'duplicate') + assert len(actor._failures) == 1 + finally: + probe_gate.set() + close_gate.set() + monitor.cancel() + await asyncio.gather(monitor, *actor._failures.values(), return_exceptions=True) + assert not actor._failures + + +@pytest.mark.asyncio +async def test_shutdown_cancels_and_joins_initialization_before_closing_core(): + cls = TrainingControllerActor.__ray_metadata__.modified_class + started, exited = asyncio.Event(), asyncio.Event() + + async def initialize(): + started.set() + try: + await asyncio.Event().wait() + finally: + exited.set() + + initialization = asyncio.create_task(initialize()) + await started.wait() + actor = SimpleNamespace(driver=SimpleNamespace(wait_for_cleanup=AsyncMock()), + core=SimpleNamespace(shutdown=AsyncMock(), wait_for_jobs=AsyncMock()), + _initialization=initialization, _monitors=[], _failures={}) + await cls.shutdown(actor) + assert initialization.cancelled() and exited.is_set() + actor.core.shutdown.assert_awaited_once() diff --git a/tests/server/training/test_failure_recovery.py b/tests/server/training/test_failure_recovery.py index bf93488c1..d93caf3fe 100644 --- a/tests/server/training/test_failure_recovery.py +++ b/tests/server/training/test_failure_recovery.py @@ -306,6 +306,58 @@ async def test_explicit_tinker_request_id_is_preserved(): await controller.shutdown() +@pytest.mark.asyncio +async def test_shutdown_waits_for_late_serve_deployment_before_confirming_release(monkeypatch): + import importlib + import threading + from types import SimpleNamespace + from twinkle.server.config.training import TrainingConfig + from twinkle.server.state.training_records import InstanceRecord + from twinkle.server.training.driver import RayInstanceDriver + from ray import serve + config = TrainingConfig(models=[{ + 'name': 'small', 'model_id': 'weights', 'backend': 'mock', 'device_type': 'cpu', + 'modes': {'lora': {'num_devices': 1, 'device_mesh': {'dp_size': 1}}} + }]) + store = SimpleNamespace(state=SimpleNamespace(unregister_replica_capacity=AsyncMock())) + driver = RayInstanceDriver(config, store, 'test-late', 'controller') + instance = InstanceRecord(instance_id='instance', controller_epoch='epoch', base_model='small', + mode='lora', startup_deadline=0) + entered, release = threading.Event(), threading.Event() + remove = AsyncMock() + + def deploy(*args, **kwargs): + entered.set() + if not release.wait(10): + raise TimeoutError('test did not release Serve') + + monkeypatch.setattr('twinkle.server.model.app.build_model_app', lambda **kwargs: 'app') + monkeypatch.setattr(importlib.import_module('ray.util.placement_group'), 'get_placement_group', lambda name: 'pg') + monkeypatch.setattr('ray.util.list_named_actors', lambda **kwargs: []) + monkeypatch.setattr(serve, 'run', deploy) + monkeypatch.setattr(serve, 'delete', lambda name: None) + monkeypatch.setattr(driver.pool, 'remove', remove) + deploying = asyncio.create_task(driver.deploy(instance, driver._catalog.get('small', 'lora'))) + closing = None + try: + assert await asyncio.to_thread(entered.wait, 2) + deploying.cancel() + await asyncio.gather(deploying, return_exceptions=True) + closing = asyncio.create_task(driver.shutdown(instance)) + await asyncio.sleep(0.01) + assert not closing.done() + remove.assert_not_awaited() + store.state.unregister_replica_capacity.assert_not_awaited() + finally: + release.set() + await asyncio.gather(deploying, return_exceptions=True) + if closing is not None: + await closing + await driver.wait_for_cleanup() + assert not driver._deployments and not driver._cleanup_tasks + store.state.unregister_replica_capacity.assert_awaited() + + @pytest.mark.asyncio async def test_capacity_unregister_does_not_retire_managed_model(): controller, state, _, _ = await setup(backend=FakeBackend()) diff --git a/tests/server/utils/test_task_queue_mixin.py b/tests/server/utils/test_task_queue_mixin.py index 5338f4885..23ad79a80 100644 --- a/tests/server/utils/test_task_queue_mixin.py +++ b/tests/server/utils/test_task_queue_mixin.py @@ -65,6 +65,13 @@ async def store_future_status(self, *args, **kwargs): 'queue_state_reason': kwargs.get('queue_state_reason'), } + async def try_start_future(self, request_id): + record = self._latest.get(request_id) + if record is None or record['status'] not in ('pending', 'queued'): + return False + await self.store_future_status(request_id, 'running', None, queue_state='active') + return self._latest[request_id]['status'] == 'running' + async def get_future(self, request_id): return self._latest.get(request_id) @@ -94,6 +101,7 @@ def enable_compute_worker(self): deployment_name=self._deployment_name, ) self._event_loop = None + self._queue_tasks = set() @pytest.mark.asyncio @@ -140,7 +148,7 @@ async def wait_for_with_completion_race(coro, timeout): @pytest.mark.asyncio -async def test_worker_can_restart_after_stopping_while_idle(): +async def test_idle_worker_stop_is_terminal_for_its_replica(): queue = _DummyQueue() queue.enable_compute_worker() worker = queue._compute_worker @@ -152,9 +160,11 @@ async def work(): await worker.ensure_started() await asyncio.sleep(0) # Let the idle worker wait for new work. await asyncio.wait_for(worker.stop(), timeout=1) - env = await queue.submit_and_peek(work, model_id='model1', token='token1') - assert env.status == 'completed' - assert env.result == {'ok': True} + from twinkle.server.exceptions import ResourceExhaustedError + with pytest.raises(ResourceExhaustedError): + await queue.submit_and_peek(work, model_id='model1', token='token1') + assert not queue.state.records + assert not worker._started finally: await worker.stop() @@ -221,6 +231,7 @@ async def test_preflight_accepts_batch_with_per_dp_multiple(): @pytest.mark.asyncio async def test_background_task_tracks_status(): queue = _DummyQueue() + queue.enable_compute_worker() async def work(): return {'ok': True} @@ -229,9 +240,13 @@ async def work(): work, model_id='model1', ) - await asyncio.sleep(0) + for _ in range(100): + if not queue._compute_worker._detached_tasks: + break + await asyncio.sleep(0.001) + await queue._compute_worker.stop() - assert [args[1] for args, _ in queue.state.records] == ['running', 'completed'] + assert [args[1] for args, _ in queue.state.records] == ['pending', 'running', 'completed'] assert queue.state.records[0][1]['absolute_deadline'] > 0 assert queue.state.records[-1][1]['result'] == {'ok': True} @@ -614,6 +629,108 @@ async def write(*args, **kwargs): assert [args[1] for args, _ in queue.state.records] == ['pending', 'failed'] +@pytest.mark.asyncio +async def test_background_capacity_and_shutdown_are_supervised(make_concurrent_queue): + from twinkle.server.exceptions import ResourceExhaustedError + queue = make_concurrent_queue(TaskQueueConfig(concurrent_task_limit=1)) + entered, exited = asyncio.Event(), asyncio.Event() + + async def work(): + entered.set() + try: + await asyncio.Event().wait() + finally: + exited.set() + + ref = await queue.schedule_background_task(work) + await asyncio.wait_for(entered.wait(), 1) + with pytest.raises(ResourceExhaustedError): + await queue.schedule_background_task(work) + assert len(queue.state._latest) == 1 + await queue._compute_worker.stop() + assert exited.is_set() + assert queue.state._latest[ref['request_id']]['status'] == 'failed' + assert queue._compute_worker._detached_reserved == 0 + + +@pytest.mark.asyncio +async def test_stopped_replica_rejects_serial_submission_without_a_future(): + from twinkle.server.exceptions import ResourceExhaustedError + queue = _DummyQueue() + queue.enable_compute_worker() + await queue._compute_worker.stop() + with pytest.raises(ResourceExhaustedError): + await queue.schedule_task(lambda: asyncio.sleep(0)) + assert not queue.state.records + assert not queue._compute_worker._started + + +@pytest.mark.asyncio +async def test_stop_during_pending_write_cannot_restart_the_serial_worker(): + from twinkle.server.exceptions import ResourceExhaustedError + queue = _DummyQueue() + queue.enable_compute_worker() + original = queue.state.store_future_status + + async def write_then_stop(*args, **kwargs): + await original(*args, **kwargs) + if args[1] == 'pending': + await queue._compute_worker.stop() + + queue.state.store_future_status = write_then_stop + with pytest.raises(ResourceExhaustedError): + await queue.schedule_task(lambda: asyncio.sleep(0), request_id='late-serial') + assert queue.state._latest['late-serial']['status'] == 'failed' + assert not queue._compute_worker._started + assert queue._compute_worker.total_queued() == 0 + + +@pytest.mark.asyncio +async def test_shutdown_drains_serial_tasks_that_never_started(): + queue = _DummyQueue() + queue.enable_compute_worker() + executed = [] + + async def work(): + executed.append(True) + + ref = await queue.schedule_task(work) + await queue._compute_worker.stop() + assert not executed + assert queue.state._latest[ref['request_id']]['status'] == 'failed' + assert queue._compute_worker.total_queued() == 0 + assert not queue._compute_worker.has_task(ref['request_id']) + + +@pytest.mark.asyncio +async def test_unknown_start_ack_never_executes_backend(make_concurrent_queue, monkeypatch): + from tests.server.state.fake_backend import FakeBackend + from twinkle.server.state.server_state import ServerState + queue = make_concurrent_queue() + backend = FakeBackend() + queue.state = queue._compute_worker._state = ServerState(backend=backend) + original = backend.update_atomic + executed = [] + + async def lost_ack(*args, **kwargs): + result = await original(*args, **kwargs) + if kwargs.get('return_changed'): + raise ConnectionError('start acknowledgement lost') + return result + + async def work(): + executed.append(True) + + monkeypatch.setattr(backend, 'update_atomic', lost_ack) + ref = await queue.schedule_concurrent_task(work) + for _ in range(100): + if not queue._compute_worker._detached_tasks: + break + await asyncio.sleep(0.001) + assert not executed + assert (await queue.state.get_future(ref['request_id']))['status'] == 'failed' + + @pytest.mark.asyncio async def test_cancel_winning_running_write_race_skips_execution(make_concurrent_queue): queue = make_concurrent_queue() From 911c2f1f0e88ab4c4349ad07e85fc4f8a26468e6 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Thu, 22 Oct 2026 12:43:58 +0000 Subject: [PATCH 10/46] feat(server): share instance lifecycle with on-demand sampling pools --- .../client/server/megatron/server_config.yaml | 53 ++- .../Usage Guide/Server and Client/Server.md | 48 +++ .../\346\234\215\345\212\241\347\253\257.md" | 48 +++ src/twinkle/infra/_ray/resource_manager.py | 12 +- .../sampler/vllm_sampler/vllm_sampler.py | 27 +- src/twinkle/server/config/application_spec.py | 21 ++ src/twinkle/server/config/training.py | 2 + .../server/launcher/server_launcher.py | 18 +- src/twinkle/server/model/managed_runtime.py | 27 +- src/twinkle/server/model/routing.py | 75 ---- src/twinkle/server/resources/instances.py | 110 ++++++ src/twinkle/server/resources/runtime.py | 130 +++++++ src/twinkle/server/sampler/app.py | 53 ++- .../server/sampler/backends/streaming.py | 16 +- src/twinkle/server/sampler/controller.py | 322 +++++++++++++++++ src/twinkle/server/sampler/driver.py | 100 ++++++ src/twinkle/server/sampler/router.py | 135 +++++++ src/twinkle/server/sampler/runtime.py | 147 ++++++++ .../server/sampler/twinkle_handlers.py | 52 ++- src/twinkle/server/task_queue/mixin.py | 12 +- src/twinkle/server/task_queue/types.py | 17 +- src/twinkle/server/task_queue/worker.py | 49 ++- src/twinkle/server/training/actor.py | 46 +-- src/twinkle/server/training/controller.py | 153 +++----- src/twinkle/server/training/driver.py | 143 ++------ src/twinkle_client/auto/agent/server_tools.py | 39 +- tests/server/config/test_server_config.py | 2 +- tests/server/config/test_training_config.py | 49 ++- .../on_demand_training/acceptance_support.py | 3 +- .../on_demand_training/matrix_validation.py | 18 +- .../on_demand_training/sampling_probe.py | 178 +++++++++ tests/server/sampler/test_pool.py | 337 ++++++++++++++++++ tests/server/sampler/test_request_runtime.py | 208 +++++++++++ .../server/sampler/test_stream_guarantees.py | 5 + .../server/static/backend_call_exemptions.py | 8 +- tests/server/test_dependencies.py | 3 +- .../training/test_acceptance_harness.py | 28 ++ .../training/test_actor_initialization.py | 16 +- tests/server/training/test_controller.py | 127 ++++++- .../server/training/test_failure_recovery.py | 49 ++- tests/server/training/test_runtime_queue.py | 149 +++++++- tests/twinkle_client/test_auto_agent_tools.py | 15 + 42 files changed, 2570 insertions(+), 480 deletions(-) delete mode 100644 src/twinkle/server/model/routing.py create mode 100644 src/twinkle/server/resources/instances.py create mode 100644 src/twinkle/server/resources/runtime.py create mode 100644 src/twinkle/server/sampler/controller.py create mode 100644 src/twinkle/server/sampler/driver.py create mode 100644 src/twinkle/server/sampler/router.py create mode 100644 src/twinkle/server/sampler/runtime.py create mode 100644 tests/server/integration/on_demand_training/sampling_probe.py create mode 100644 tests/server/sampler/test_pool.py create mode 100644 tests/server/sampler/test_request_runtime.py diff --git a/cookbook/client/server/megatron/server_config.yaml b/cookbook/client/server/megatron/server_config.yaml index f10de4754..7e509bcb0 100644 --- a/cookbook/client/server/megatron/server_config.yaml +++ b/cookbook/client/server/megatron/server_config.yaml @@ -1,4 +1,8 @@ -# Twinkle Server Configuration - Megatron Backend +# Twinkle Server Configuration - Megatron LoRA Backend +# Candidate for 8 x L20D (~269 GiB each): 2-4 single-GPU sampler instances +# plus up to 2 two-GPU LoRA training instances (3 handles each, 6 total). +# Validate 27B/32K peak memory, throughput and optimizer resume on L20D +# before production rollout; local smaller-model tests do not prove those limits. # proxy_location: determines where the HTTP proxy runs. # "EveryNode" means each Ray node runs its own proxy (good for multi-node). @@ -26,6 +30,13 @@ persistence: mode: redis redis_url: redis://localhost:6380 +# Independent allocation, startup and confirmed-shutdown budgets. +training: + allocation_timeout_s: 3 + startup_timeout_s: 1800 + shutdown_timeout_s: 180 + sweep_interval_s: 5 + # Applications: each entry defines a service component deployed on the server applications: @@ -36,7 +47,7 @@ applications: import_path: server # Python module to import args: server_config: - per_token_model_limit: 3 # Maximum number of models (adapters) per token (server-globally enforced) + per_token_model_limit: 3 # Global per-token limit; shared-token clients share these 3 slots supported_models: - Qwen/Qwen3.8-27B @@ -56,30 +67,33 @@ applications: # 3. Sampler Service - Runs inference / sampling using vLLM engine # Used for generating text from the model (e.g., evaluating LoRA results). - # Config: TP=2 x DP=2 on 4 GPUs, ~27GB weights/GPU, ~37GB for KV cache + LoRA + # Each instance has DP=1 and one GPU. Two are preloaded; scale to four. + # Memory usage includes each engine's preallocated KV cache, not just weights. - name: sampler-Qwen3.8-27B route_prefix: /api/v1/sampler/Qwen/Qwen3.8-27B import_path: sampler args: model_id: "ms://Qwen/Qwen3.8-27B" # ModelScope model identifier - nproc_per_node: 4 # Number of GPU processes per node + nproc_per_node: 1 # One GPU per independently allocated instance + min_replicas: 2 + max_replicas: 4 + target_ongoing_requests: 16 + idle_timeout_s: 600 sampler_type: vllm # Inference engine: vllm engine_args: # vLLM engine-specific settings max_model_len: 32768 # Maximum sequence length the engine supports - gpu_memory_utilization: 0.75 # 80% utilization, ~64GB/GPU, leaves buffer for safety + gpu_memory_utilization: 0.75 # 75% engine memory budget; does not free Ray GPU reservations enable_lora: true # Allow loading LoRA adapters during inference max_loras: 5 # Max allowed loras working on vLLM at the same time - max_lora_rank: 32 # Support up to rank 64 LoRA adapters + max_lora_rank: 32 # Maximum supported LoRA rank enable_tower_connector_lora: true device_group: # Logical device group for the sampler name: sampler - gpus_per_worker: 2 - ranks: 4 # GPU rank indices to use + ranks: 1 # Per-instance device count, not fixed physical GPU IDs device_type: cuda device_mesh: device_type: cuda - dp_size: 2 - tp_size: 2 # 2 TP replicas for multi-tenant throughput + dp_size: 1 # Pool instance count is controlled by args.min/max_replicas queue_config: rps_limit: 20 # Max requests per second tps_limit: 131072 # Max tokens per second @@ -97,8 +111,9 @@ applications: TWINKLE_TRUST_REMOTE_CODE: "0" TWINKLE_LONG_POLL_TIMEOUT: "120" - # 2. Model Service - Hosts the base model for training. - # Config: PP=2 x DP=2 on 4 GPUs, ~27GB weights/GPU, comfortable for LoRA training + # 2. Model Service - LoRA-only, independently queued two-GPU instances. + # Megatron instances require at least two devices; PP=2 shards each model. + # min=0 enables cold startup; max=2 limits total training devices to four. - name: models-Qwen3.8-27B route_prefix: /api/v1/model/Qwen/Qwen3.8-27B import_path: model @@ -106,15 +121,19 @@ applications: backend: megatron # Use Megatron-LM backend model_id: "ms://Qwen/Qwen3.8-27B" # ModelScope model identifier max_length: 32768 # model max length - max_loras: 3 # model max loras - nproc_per_node: 4 # Number of GPU processes per node + train_mode: lora + max_loras: 3 # Handles per instance, not simultaneous backward calls + min_replicas: 0 # Do not reserve training GPUs at startup + max_replicas: 2 # Up to 2 independent two-GPU instances + idle_timeout_s: 600 # Business idle; heartbeats/polling do not renew it + nproc_per_node: 2 # Two workers per training instance device_group: name: model - ranks: 4 + ranks: 2 device_type: cuda device_mesh: device_type: cuda - dp_size: 2 + dp_size: 1 pp_size: 2 queue_config: @@ -122,7 +141,7 @@ applications: tps_limit: 131072 # Max tokens per second max_input_tokens: 32768 adapter_config: - adapter_timeout: 120 # Seconds before idle adapter unload + adapter_timeout: 120 # Session heartbeat timeout, not business idle adapter_max_lifetime: 36000 # Maximum lifetime of an adapter in seconds (e.g., 10 hours) deployments: - name: ModelManagement diff --git a/docs/source_en/Usage Guide/Server and Client/Server.md b/docs/source_en/Usage Guide/Server and Client/Server.md index 67a56d4ed..87d97d3d4 100644 --- a/docs/source_en/Usage Guide/Server and Client/Server.md +++ b/docs/source_en/Usage Guide/Server and Client/Server.md @@ -487,3 +487,51 @@ Existing `applications` entries with `import_path: model` remain valid. Set `arg `modes` is optional: single-mode training does not require rewriting existing YAML. Model `deployments` options configure the CPU router; autoscaling request targets apply to router replicas. The Controller owns training-instance minimums/limits, and each compute instance has exactly one Serve replica. Non-autoscaling request limits, health/shutdown settings, and `ray_actor_options.runtime_env` also reach compute instances. `args.env_vars` overrides matching deployment environment variables. Unsupported `adapter_config` fields fail validation instead of being silently ignored. To offer both LoRA and full training for one model, configure `num_devices`, `device_mesh`, `min_replicas`, and `idle_timeout_s` separately under `args.modes.lora/full`. LoRA users share compatible instances; each full-training handle owns an instance. Sampling configuration stays the same. See `cookbook/client/server/on_demand_training.yaml` in the repository. + +### Single-node, eight-GPU LoRA capacity + +`cookbook/client/server/megatron/server_config.yaml` is a candidate for large-memory L20D GPUs. It retains single-mode YAML: two preloaded single-GPU sampling instances (DP=1, no TP setting), scaling up to four, plus up to two two-GPU training instances (PP=2, DP=1). Megatron configurations require at least two devices per instance before allocation; single-device Megatron is not supported. Only `train_mode: lora` is configured; full training remains available in other framework configurations. Key model args: + +```yaml +train_mode: lora +nproc_per_node: 2 +device_group: {name: model, ranks: 2, device_type: cuda} +device_mesh: {device_type: cuda, dp_size: 1, pp_size: 2} +max_loras: 3 +min_replicas: 0 +max_replicas: 2 +idle_timeout_s: 600 +``` + +- `max_replicas × max_loras = 6` bounds training handles, not throughput. Creation fills existing compatible LoRA slots before starting another instance; active training is not migrated for load balancing. Stateful operations are serial within an instance, with at most two independent training queues. `per_token_model_limit: 3` remains a server-wide per-token quota, not a total-user cap. Clients sharing one token share those three slots; independent users need independent tokens, or a deliberately increased shared-token quota. Set it to 1 separately if the service requires one handle per independent token. +- PP=2 shards each model over two devices; removing DP replication from the original PP=2/DP=2 instance leaves room for two independent instances. Account for 32K activations, LoRA optimizers, communication and temporary workspaces. Smaller-model or different-hardware E2E tests do not certify L20D/27B/32K peak memory or throughput. +- Save checkpoints during training and explicitly call `remove_adapter()` or unload the Tinker model when finished. Heartbeats and polling do not renew business idle. Expired handles require a new create and explicit checkpoint resume; old training is not resumed automatically. +- `min_replicas: 0` permits reclaiming every empty training instance at the cost of cold startup. Both pools use the same device-allocation and instance-lifecycle implementation with independent state. Reclaiming training devices does not directly scale sampling; sampling grows according to its own in-flight load. Lowering vLLM `gpu_memory_utilization` does not release Ray's whole-device reservation. No fractional-GPU or training/sampling colocation is enabled. + +Before rollout, validate the actual 27B model on L20D: a two-GPU PP=2/DP=1 32K workload at the maximum business batch, three resident rank-32 LoRAs, two instances/six handles alongside sampling, optimizer save/resume, tenant isolation and confirmed device reclamation. Measure peak memory, training tokens/s, sampling output tokens/s, queue wait and P95 latency against the original four-GPU PP=2/DP=2 layout using the same work. If the two-GPU layout does not meet business requirements, restore four GPUs per instance (PP=2, DP=2, at most one instance); do not infer throughput from an idle memory snapshot. + +### On-demand sampling instances and dispatch + +A stable CPU Router selects the ready instance with the fewest actual in-flight operations for ordinary sampling, DataRefs and streaming. Clients do not supply an instance ID, and users/LoRAs are not pinned to a device. Per-instance DP/TP describes an engine's parallel shape, not pool replicas. The production candidate uses one GPU and DP=1 per instance, without TP configuration. + +```yaml +nproc_per_node: 1 +device_group: {name: sampler, ranks: 1, device_type: cuda} +device_mesh: {device_type: cuda, dp_size: 1} +min_replicas: 2 +max_replicas: 4 +target_ongoing_requests: 16 +idle_timeout_s: 600 +``` + +Both minimum instances must finish loading before readiness. Demand is `ceil((in-flight requests + 1) / 16)`, capped at four. Starting and quarantined instances count against that cap. Existing ready instances serve during expansion; only idle extras with no actual generation are reclaimed after 600 seconds, down to two. Neither pool preempts or colocates on a GPU; four sampling plus four training GPUs fit the eight-device maximum. + +New instances load only the base model and default template. LoRAs load from each request's `adapter_uri`, `model_path` or persistent sampling-session checkpoint path, into independent disposable caches. No cache synchronization or new cache database is introduced. `engine_args.max_loras: 5` bounds active engine LoRAs, not training handles. Custom templates, named adapter settings and patches are broadcast under pool-wide exclusive protection and replayed before new instances become ready. Partial broadcast failure fences all participating replicas. Controller restart clears the old pool and dynamic control settings; clients must reapply them, while existing sampling-session/checkpoint paths remain unchanged. + +HTTP disconnect and future timeout do not certify generation exit. Completion handles and weight guards remain held until actual exit; unconfirmed death remains `QUARANTINED` without early capacity release. Full-checkpoint sampling retains pool-wide exclusive weight changes. L20D/27B single-GPU 32K memory, four-engine concurrency, output tokens/s and P95 latency require controlled same-hardware validation; local PPU/4B tests cannot certify them. + +### Rollout compatibility and safe states + +Keeping YAML syntax does not keep the old SDK wire unchanged. The Twinkle SDK performs the `training_handles` handshake, generates `client_create_id` and supplies `training_id`; upgrade it together with the server. Tinker keeps its own model-handle wire. + +Ray 2.55 is the minimum dependency: Serve ingress must await asynchronous user constructors; do not use the old CI Ray 2.48 pin. Startup, business idle, task execution and confirmed shutdown have separate budgets. Atomic create start and cancel are mutually exclusive; an unknown state-backend acknowledgement must not authorize execution. `QUARANTINED` means device death is unconfirmed and still counts toward the instance limit; release must not report unconfirmed reclamation as success. Only an actually executed operation renews business idle, not a queue failure or unknown start acknowledgement. diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" index 37d808d8f..367f4a967 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" @@ -474,3 +474,51 @@ twinkle-server check-config -c server_config.yaml `modes` 是可选扩展，单模式训练不需要改写原 YAML。原模型 `deployments` 的 Serve 设置用于 CPU 路由部署，其中 autoscaling 的请求目标用于路由副本扩缩容；训练计算实例的 min/max 仍由 Controller 管理，每个计算实例固定一个 Serve replica。非 autoscaling 的请求上限、健康/关闭设置和 `ray_actor_options.runtime_env` 会传给计算实例；`args.env_vars` 优先于 deployment 的同名环境变量。未知或未实现的 `adapter_config` 字段会报错，不再静默忽略。 同一模型需要同时提供 LoRA 和全参数训练时，在 `args.modes.lora/full` 内分别设置 `num_devices`、`device_mesh`、`min_replicas` 和 `idle_timeout_s`；不用改成另一套顶层模型目录。LoRA 用户共享兼容实例，全参数用户各占一个实例。采样服务继续使用原配置。详见仓库 `cookbook/client/server/on_demand_training.yaml`。 + +### 单机 8 卡、仅 LoRA 的容量配置 + +`cookbook/client/server/megatron/server_config.yaml` 提供大显存 L20D 的候选方案，保持单模式 YAML：采样每实例 1 卡、DP=1，不配置 TP，默认加载 2 实例、最多 4 实例；训练每实例 2 卡（PP=2、DP=1），最多 2 实例。Megatron 配置在分配资源前校验每实例至少 2 卡，不提供单卡路径。仅配置 `train_mode: lora`，不提供 full 模式；框架仍支持其他配置中的 full。关键训练参数： + +```yaml +train_mode: lora +nproc_per_node: 2 +device_group: {name: model, ranks: 2, device_type: cuda} +device_mesh: {device_type: cuda, dp_size: 1, pp_size: 2} +max_loras: 3 +min_replicas: 0 +max_replicas: 2 +idle_timeout_s: 600 +``` + +- `max_replicas × max_loras = 6` 是训练句柄容量上限，不是吞吐保证；创建时优先填满已有实例的 LoRA 槽位，槽位不足才创建新实例，不会按空闲度迁移现有训练。同实例的 forward/backward/step 串行执行，最多有 2 条独立训练队列。`per_token_model_limit: 3` 仍是全服务的每 token 配额，不是总用户上限；不同用户使用独立 token 才能分别获得配额。同一 token 的所有客户端合计仍只有 3 个名额；共享 token 的部署需按业务调大该配额，不能仅增加实例数。若希望避免单个独立 token 占多个名额，可另设为 1。 +- 两卡实例用 PP=2 切分模型，取消原 PP=2/DP=2 实例内的 DP 副本，以容纳两个独立实例。还需考虑 32K 激活、LoRA optimizer、通信及临时 workspace；较小模型或其他硬件的 E2E 不能证明 L20D/27B/32K 的峰值和吞吐。 +- 业务完成后主动 `remove_adapter()` 或卸载 Tinker 模型；运行期应定期保存 checkpoint。心跳/结果轮询不会续业务 idle，句柄过期后必须重新创建并显式 resume；不会自动恢复旧训练。 +- `min_replicas: 0` 能回收所有空训练实例，但首次训练需冷加载。训练和采样通过共用的设备申请与实例启停机制管理资源，两池使用独立状态。释放训练设备预约不主动触发采样扩容，采样按自身在途请求扩容；降低 `gpu_memory_utilization` 也不会释放 Ray 的整卡预约。不启用 fractional GPU 或训练/采样同卡共置。 + +上线前使用同型号 L20D 与实际 27B 模型，依次验证两卡 PP=2/DP=1 的 32K 最大业务 batch、3 个 rank=32 LoRA 同驻留、2 实例共 6 句柄与 sampler 并行、save + optimizer resume、训练之间隔离和最后用户退出后的设备回收。记录峰值显存、训练 tokens/s、采样输出 tokens/s、队列等待及 P95 延迟；与原 PP=2/DP=2 四卡方案使用相同工作量比较。若两卡方案不满足业务要求，可恢复四卡、PP=2/DP=2、最多 1 实例；不凭显存快照承诺缩卡后的吞吐。 + +### Sampler 按需实例与请求分派 + +固定 CPU Router 将普通采样、DataRef 与流式请求分派给真实在途请求最少的 ready 实例，不要求客户端传实例 ID，也不把用户固定绑定某张卡。实例内 DP/TP 描述一个引擎的并行规格，不是池内副本数；线上示例每实例单卡、DP=1，删除 TP 配置。 + +```yaml +nproc_per_node: 1 +device_group: {name: sampler, ranks: 1, device_type: cuda} +device_mesh: {device_type: cuda, dp_size: 1} +min_replicas: 2 +max_replicas: 4 +target_ongoing_requests: 16 +idle_timeout_s: 600 +``` + +保底 2 个实例加载完成后才报告 ready；需求按 `ceil((在途请求数 + 1) / 16)` 计算，上限 4。启动中的实例和死亡未确认的隔离实例均占上限。扩容加载期间现有实例继续服务；闲置 600 秒且没有真实 generation 的超保底实例才回收，缩至 2。两池不抢占、不共享一张卡，训练最多 4 卡加采样最多 4 卡为 8 卡。 + +新实例只加载基础模型和默认 template；LoRA 按请求的 `adapter_uri`、`model_path` 或 sampling session 中的 checkpoint 路径加载，各引擎独立缓存，不同步缓存或新增缓存数据库。`engine_args.max_loras: 5` 是引擎同时活跃 LoRA 数，不等于训练句柄容量。自定义 template、命名 adapter 和 patch 在池级排他保护下广播，新实例 ready 前重放；广播失败将参与实例隔离，避免混用不同配置。Controller 重启清理旧池、动态控制设置失效，客户端须重新设置；已有 sampling session/checkpoint 路径不变。 + +HTTP 断连或 Future 超时不代表 generation 已退出，后台完成句柄和权重保护保持到实际退出；死亡无法确认则保持 `QUARANTINED`，不提前还卡。Full-checkpoint 采样仍使用池级排他权重切换。单卡采样的 L20D/27B/32K 峰值、四实例并发、输出 tokens/s 与 P95 延迟需同型号受控验收，本机 PPU/4B 测试不能替代。 + +### 上线兼容与安全状态 + +配置语法保留不代表 SDK/wire 不变：Twinkle SDK 自动完成 `training_handles` 握手、生成 `client_create_id`、传递 `training_id`；新 server 不接受缺少这些信息的旧 Twinkle SDK，应配套升级。Tinker 仍使用自己的模型句柄与 wire。 + +Ray 依赖下限为 2.55，Serve ingress 必须等待异步用户构造；不要继续使用旧 CI 的 Ray 2.48。资源启动、业务 idle、任务超时和关闭确认是不同预算。创建的原子 start 与 cancel 互斥，状态后端 ACK 不明不得执行；`QUARANTINED` 表示设备死亡尚未确认，仍计入实例上限，release 不应假报释放成功。只有真实执行过的操作续业务 idle，排队失败或 start ACK 不明不续期。 diff --git a/src/twinkle/infra/_ray/resource_manager.py b/src/twinkle/infra/_ray/resource_manager.py index c9013c972..875ba4a08 100644 --- a/src/twinkle/infra/_ray/resource_manager.py +++ b/src/twinkle/infra/_ray/resource_manager.py @@ -22,14 +22,18 @@ class ResourceManager: @classmethod def from_allocation(cls, group: DeviceGroup, placement_group, devices: list[str]): - """Use an allocation owned by the training controller, without reserving again. + """Use an allocation owned by a lifecycle controller, without reserving again. Ranks are local to this allocation; devices are the actual Ray-selected IDs. The controller owns removal, so this view's destroy method is a no-op. """ import ray manager = cls.__new__(cls) - ranks = list(range(len(devices))) + width = group.gpus_per_worker + if len(devices) % width: + raise ValueError('Allocation device count must be divisible by gpus_per_worker') + chunks = [devices[index:index + width] for index in range(0, len(devices), width)] + ranks = list(range(len(chunks))) group.ranks = ranks manager.group_configs = [group] manager.placement_groups = [placement_group] @@ -37,10 +41,10 @@ def from_allocation(cls, group: DeviceGroup, placement_group, devices: list[str] manager._external_allocation = True manager.device_groups = { group.name: [{ - 'gpu_rank': [device], + 'gpu_rank': chunk, 'placement_group': placement_group, 'ray_address': str(ray.get_runtime_context().gcs_address) - } for device in devices] + } for chunk in chunks] } if (group.device_type or '').upper() == 'CPU': manager.device_groups[group.name] = [{ diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index ca9b58840..683f99067 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -143,17 +143,24 @@ async def _drain(): except Exception as exc: q.put(exc) finally: - q.put(_SENTINEL) + try: + await async_gen_coro.aclose() + finally: + q.put(_SENTINEL) - asyncio.run_coroutine_threadsafe(_drain(), self._async_loop) - - while True: - item = q.get() - if item is _SENTINEL: - break - if isinstance(item, Exception): - raise item - yield item + generation = asyncio.run_coroutine_threadsafe(_drain(), self._async_loop) + try: + while True: + item = q.get() + if item is _SENTINEL: + break + if isinstance(item, Exception): + raise item + yield item + finally: + # Queue termination or closing the Python iterator is not evidence + # that the engine-loop task has exited. Join before releasing weights. + generation.result() async def _create_engine_async(self, engine_cls, model_id, engine_kwargs): """Create engine in async context to ensure output_handler starts correctly.""" diff --git a/src/twinkle/server/config/application_spec.py b/src/twinkle/server/config/application_spec.py index 924beece0..a25da604e 100644 --- a/src/twinkle/server/config/application_spec.py +++ b/src/twinkle/server/config/application_spec.py @@ -122,6 +122,27 @@ class SamplerArgs(_ArgsBase): engine_args: dict[str, Any] | None = None queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) data_plane_url: str | None = None + min_replicas: int = Field(default=1, ge=0, strict=True) + max_replicas: int = Field(default=1, gt=0, strict=True) + idle_timeout_s: float = Field(default=600, gt=0, allow_inf_nan=False) + target_ongoing_requests: int = Field(default=16, gt=0, strict=True) + + @model_validator(mode='after') + def _validate_pool(self): + if self.min_replicas > self.max_replicas: + raise ValueError('sampler min_replicas must not exceed max_replicas') + ranks = self.device_group.get('ranks', self.nproc_per_node) + width = self.device_group.get('gpus_per_worker', 1) + if type(ranks) is not int or ranks <= 0 or type(width) is not int or width <= 0 or ranks % width: + raise ValueError('sampler device_group requires a positive device count divisible by gpus_per_worker') + if self.sampler_type != 'mock': + from twinkle.utils.device_mesh import DeviceMesh + if DeviceMesh.from_sizes(**self.device_mesh).world_size != ranks: + raise ValueError('sampler device_mesh world_size must equal its device count') + tp = (self.engine_args or {}).get('tensor_parallel_size', width) + if tp != width: + raise ValueError('sampler tensor_parallel_size must equal gpus_per_worker') + return self class ServerStateArgs(_ArgsBase): diff --git a/src/twinkle/server/config/training.py b/src/twinkle/server/config/training.py index d4c372b9c..964692f0e 100644 --- a/src/twinkle/server/config/training.py +++ b/src/twinkle/server/config/training.py @@ -99,6 +99,8 @@ def _validate_backend(self): for mode in (self.modes.lora, self.modes.full): if mode is None: continue + if self.backend == 'megatron' and mode.num_devices < 2: + raise ValueError('megatron training instances require at least 2 devices') if mode.device_mesh.device_type not in (None, self.device_type): raise ValueError('device_mesh.device_type must match device_type') if mode.strategy is not None and self.backend != 'transformers': diff --git a/src/twinkle/server/launcher/server_launcher.py b/src/twinkle/server/launcher/server_launcher.py index 7cd6bd125..4c6ac6d5c 100644 --- a/src/twinkle/server/launcher/server_launcher.py +++ b/src/twinkle/server/launcher/server_launcher.py @@ -63,6 +63,7 @@ def __init__( self._builders: dict[str, Callable] = {} self._ray_initialized = False self._serve_started = False + self._sampler_controllers = [] def _get_builders(self) -> dict[str, Callable]: """Get (and cache) the builder functions for all app types.""" @@ -151,6 +152,16 @@ def _deploy_application(self, app_spec: ApplicationSpec) -> None: # Model applications are specifications. _deploy_training creates # their CPU routes and the controller owns all device instances. return + if import_path == 'sampler': + from twinkle.server.sampler.router import deploy_sampler_pool + namespace = self.ray_namespace or self.config.ray_namespace or 'twinkle_cluster' + env = build_propagated_env_vars() + for key in ('PYTHONPATH', 'TWINKLE_DEFAULT_SAVE_DIR', 'TMPDIR', 'OMP_NUM_THREADS'): + if key in os.environ: + env.setdefault(key, os.environ[key]) + controller = deploy_sampler_pool(app_spec, namespace, self.config.training.model_dump(), env) + self._sampler_controllers.append(controller) + return # Shallow-dump the typed args to a kwargs dict WITHOUT recursing into # nested models: ``dict(model)`` yields top-level (field, value) pairs # with nested models left as instances, so ``queue_config`` stays a @@ -213,7 +224,7 @@ def launch(self) -> None: """ # Apply Ray Serve patches before initializing Ray apply_ray_serve_patches() - if self.config.training_config.models: + if self.config.training_config.models or any(app.import_path == 'sampler' for app in self.config.applications): os.environ['TWINKLE_STATE_RAY_NAMESPACE'] = ( self.ray_namespace or self.config.ray_namespace or 'twinkle_cluster') @@ -305,7 +316,12 @@ def _shutdown(self) -> None: Telemetry flush runs AFTER ``serve.shutdown()`` and is wrapped so a telemetry-shutdown failure cannot mask the user-facing shutdown path. """ + import ray from ray import serve + for controller in self._sampler_controllers: + ray.get(controller.shutdown.remote()) + ray.kill(controller, no_restart=True) + self._sampler_controllers.clear() if getattr(self, '_training_controller', None) is not None: import ray ray.get(self._training_controller.shutdown.remote()) diff --git a/src/twinkle/server/model/managed_runtime.py b/src/twinkle/server/model/managed_runtime.py index d9f1737d5..5a0d3da8c 100644 --- a/src/twinkle/server/model/managed_runtime.py +++ b/src/twinkle/server/model/managed_runtime.py @@ -6,7 +6,6 @@ import os import uuid from contextvars import ContextVar -from datetime import datetime from twinkle.protocol.serialize import deserialize_object from twinkle.protocol.types import model as native @@ -89,13 +88,6 @@ async def _prepare_queued_task(self, task): task.request_id, renew=task.task_type not in ('get_info', 'get_train_configs')) - async def enqueued(): - try: - await self._training_controller.commit_operation.remote(task.model_id, task.request_id, - context['epoch']) - except Exception: - logger.exception('Training enqueue notification failed for %s', task.request_id) - async def finished(executed, unsafe): if unsafe: self._model_unhealthy = True @@ -107,12 +99,18 @@ async def finished(executed, unsafe): message='Backend execution timed out; runtime is being terminated.')) else: try: + outcome = self._compute_worker.get_task_outcome(task.request_id) await self._training_controller.finish_operation.remote( - task.model_id, task.request_id, context['epoch'], executed=executed) + task.model_id, + task.request_id, + context['epoch'], + executed=executed, + finished_at=outcome['finished_at']) + return True except Exception: logger.exception('Training completion notification failed for %s', task.request_id) - task.on_enqueued, task.on_finished = enqueued, finished + task.on_finished = finished if task.task_type == 'apply_patch': await self._training_controller.mark_dirty.remote(context['instance_id'], context['epoch']) @@ -210,11 +208,4 @@ async def verify_operation(self, request_id): return {'state': 'unsafe'} if self._compute_worker.has_task(request_id): return {'state': 'pending'} - future = await self.state.get_future(request_id) - if future and future['status'] in ('completed', 'failed', 'cancelled'): - return { - 'state': 'finished', - 'executed': future['status'] != 'cancelled', - 'finished_at': datetime.fromisoformat(future['updated_at']).timestamp() - } - return {'state': 'missing'} + return self._compute_worker.get_task_outcome(request_id) or {'state': 'missing'} diff --git a/src/twinkle/server/model/routing.py b/src/twinkle/server/model/routing.py deleted file mode 100644 index 02b7669a3..000000000 --- a/src/twinkle/server/model/routing.py +++ /dev/null @@ -1,75 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -from ray.serve.request_router import FIFOMixin, MultiplexMixin, PendingRequest, ReplicaID, RequestRouter, RunningReplica - -from twinkle.server.state import ServerState, get_server_state -from twinkle.utils.logger import get_logger - -logger = get_logger() - - -class StickyLoraRequestRouter(FIFOMixin, MultiplexMixin, RequestRouter): - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - self.state: ServerState = get_server_state() - - async def choose_replicas( - self, - candidate_replicas: list[RunningReplica], - pending_request: PendingRequest | None = None, - ) -> list[list[RunningReplica]]: - """ - This method chooses the best replica for the request based on - multiplexed and available lora count. The algorithm - works as follows: - - 1. Populate top_ranked_replicas based on available replicas based on - multiplex_id (only one replica is chosen) - 2. Populate and override top_ranked_replicas info based on available lora - slots of the replica. - """ - - # Take the best set of replicas for the multiplexed model - if (pending_request is not None and pending_request.metadata.multiplexed_model_id): - ranked_replicas_multiplex: list[RunningReplica] = (self.rank_replicas_via_multiplex( - replicas=candidate_replicas, - multiplexed_model_id=pending_request.metadata.multiplexed_model_id, - ))[0] - - # If found any replica, return it - if ranked_replicas_multiplex: - logger.debug('[Router] Found replica for multiplexed model !!!') - return [ranked_replicas_multiplex] - - # Dictionary to hold the top-ranked replicas - top_ranked_replicas: dict[ReplicaID, RunningReplica] = {} - - # Filter out replicas that are not available (queue length exceed max ongoing request) - ranked_replicas_locality = self.select_available_replicas(candidates=candidate_replicas) - - for replica in ranked_replicas_locality: - top_ranked_replicas[replica.replica_id] = replica - - # Filter out replicas that exceed max lora count (query from server state) - candidate_ids = [r.replica_id.unique_id for r in top_ranked_replicas.values()] - available_ids = set(await self.state.get_available_replica_ids(candidate_ids)) - if available_ids: - top_ranked_replicas = { - rid: r - for rid, r in top_ranked_replicas.items() if r.replica_id.unique_id in available_ids - } - - if not top_ranked_replicas: - # No replica has remaining LoRA capacity – fall back to all candidates - logger.debug('[Router] No replica has remaining LoRA capacity') - return [candidate_replicas] - - logger.debug('[Router] StickyLoraRequestRouter choosing replica for request') - - # Take the replica with minimum throughput. - min_throughput_replicas = min( - [replica for replica in top_ranked_replicas.values()], - key=lambda r: r.routing_stats.get('throughput', 0), - ) - return [[min_throughput_replicas]] diff --git a/src/twinkle/server/resources/instances.py b/src/twinkle/server/resources/instances.py new file mode 100644 index 000000000..d08045108 --- /dev/null +++ b/src/twinkle/server/resources/instances.py @@ -0,0 +1,110 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shared instance supervision and bounds, with feature-owned records and policies.""" +from __future__ import annotations + +import asyncio +import logging +import time + +logger = logging.getLogger(__name__) + + +class InstanceLifecycle: + """Own tasks, not a second copy of instance records or persisted status.""" + + def __init__(self, label, *, clock=time.monotonic): + self.label = label + self.clock = clock + self.unreachable_since = {} + self.jobs = set() + self.startups = {} + self.closing = {} + self.named_jobs = {} + + def spawn(self, coroutine): + task = asyncio.create_task(coroutine) + self.jobs.add(task) + + def done(completed): + self.jobs.discard(completed) + if not completed.cancelled(): + error = completed.exception() + if error is not None: + logger.error( + '%s background job failed', self.label, exc_info=(type(error), error, error.__traceback__)) + + task.add_done_callback(done) + return task + + def spawn_once(self, key, factory): + task = self.named_jobs.get(key) + if task is None: + task = self.spawn(factory()) + self.named_jobs[key] = task + task.add_done_callback(lambda completed: self.named_jobs.pop(key, None)) + return task + + def start(self, identity, factory): + if identity not in self.startups: + self.startups[identity] = self.spawn(factory()) + return self.startups[identity] + + async def cancel_start(self, identity): + task = self.startups.get(identity) + if task is not None and task is not asyncio.current_task() and not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + async def close(self, identity, stop, timeout, stopped, quarantined): + task = self.closing.get(identity) + if task is None: + + async def close_once(): + try: + await asyncio.wait_for(stop(), timeout) + except Exception as error: + await quarantined(error) + else: + # Only the feature can retire its users/quota and commit STOPPED. + await stopped() + self.unreachable_since.pop(identity, None) + + task = self.spawn(close_once()) + self.closing[identity] = task + + def clear(completed): + if self.closing.get(identity) is completed: + self.closing.pop(identity, None) + + task.add_done_callback(clear) + await asyncio.shield(task) + + async def health(self, identity, probe, *, unreachable_timeout=30): + """Distinguish explicit backend failure from a transient control-plane timeout.""" + try: + result = await probe() + except Exception: + since = self.unreachable_since.setdefault(identity, self.clock()) + return False if self.clock() - since >= unreachable_timeout else None + self.unreachable_since.pop(identity, None) + return not result['unhealthy'] + + @staticmethod + def owned_count(instances): + return sum(instance.state != 'STOPPED' for instance in instances) + + @classmethod + def headroom(cls, instances, maximum): + return float('inf') if maximum is None else max(0, maximum - cls.owned_count(instances)) + + @classmethod + def missing(cls, instances, minimum, maximum, eligible): + instances = tuple(instances) + needed = max(0, minimum - sum(eligible(instance) for instance in instances)) + return min(needed, cls.headroom(instances, maximum)) + + async def join(self): + while self.jobs: + tasks = tuple(self.jobs) + await asyncio.gather(*tasks, return_exceptions=True) + self.jobs.difference_update(task for task in tasks if task.done()) diff --git a/src/twinkle/server/resources/runtime.py b/src/twinkle/server/resources/runtime.py new file mode 100644 index 000000000..48f6958bd --- /dev/null +++ b/src/twinkle/server/resources/runtime.py @@ -0,0 +1,130 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Ray/Serve resource effects shared by training and sampling owners.""" +from __future__ import annotations + +import asyncio +import logging +import os +import time +from copy import deepcopy + +from twinkle.server.env_propagation import build_propagated_env_vars +from .allocation import DevicePool + +logger = logging.getLogger(__name__) + + +class RayRuntime: + """Own deployment threads and certify worker/PG death before releasing capacity.""" + + def __init__(self, namespace, *, allocation_timeout, shutdown_timeout): + self.namespace = namespace + self.pool = DevicePool(namespace, allocation_timeout=allocation_timeout, shutdown_timeout=shutdown_timeout) + self.deployments = {} + self.cleanup_tasks = set() + + def deployment_options(self, options, name, env_vars=None): + options = deepcopy(options) + options.pop('autoscaling_config', None) + options['num_replicas'] = 1 + actor = options.setdefault('ray_actor_options', {}) + actor.setdefault('num_cpus', 0.1) + runtime_env = actor.setdefault('runtime_env', {}) + env = build_propagated_env_vars() + for key in ('PYTHONPATH', 'TWINKLE_DEFAULT_SAVE_DIR', 'TMPDIR', 'OMP_NUM_THREADS'): + if key in os.environ: + env.setdefault(key, os.environ[key]) + env.update(runtime_env.get('env_vars', {})) + env.update(env_vars or {}) + env.update(TWINKLE_RAY_NAMESPACE=self.namespace, TWINKLE_RAY_REGISTRY_NAME=name + '-registry') + runtime_env['env_vars'] = env + return options + + async def deploy(self, identity, name, app, route_prefix, cleanup): + from ray import serve + deployment = asyncio.create_task(asyncio.to_thread(serve.run, app, name=name, route_prefix=route_prefix)) + self.deployments[identity] = deployment + try: + await asyncio.shield(deployment) + except asyncio.CancelledError: + + async def clean_late(): + try: + await deployment + except Exception: + pass + await cleanup() + + task = asyncio.create_task(clean_late()) + self.cleanup_tasks.add(task) + + def done(completed): + self.cleanup_tasks.discard(completed) + if not completed.cancelled(): + error = completed.exception() + if error is not None: + logger.error( + 'Late deployment cleanup failed', exc_info=(type(error), error, error.__traceback__)) + + task.add_done_callback(done) + raise + finally: + if deployment.done(): + self.deployments.pop(identity, None) + + async def wait_for_cleanup(self): + while self.cleanup_tasks: + await asyncio.gather(*tuple(self.cleanup_tasks), return_exceptions=True) + + @staticmethod + def handle(name): + from ray import serve + return serve.get_app_handle(name) + + async def shutdown(self, identity, name, *, observe_removal=None): + import ray + from ray import serve + from ray.util import list_named_actors + deployment = self.deployments.get(identity) + if deployment is not None: + try: + await asyncio.shield(deployment) + except Exception: + pass + self.deployments.pop(identity, None) + # Both features give workers the controller identity, not a Serve replica ID. + actors = [] + for row in list_named_actors(all_namespaces=True): + if row['namespace'] != self.namespace or identity not in row['name']: + continue + try: + handle = ray.get_actor(row['name'], namespace=self.namespace) + except ValueError: + continue + actors.append(handle) + ray.kill(handle, no_restart=True) + await asyncio.to_thread(serve.delete, name) + started, status = time.monotonic(), 'error' + try: + await self.pool.remove(name) + status = 'success' + finally: + if observe_removal is not None: + observe_removal(time.monotonic() - started, status) + for handle in actors: + try: + await asyncio.wait_for(handle.__ray_ready__.remote(), 5) + except ray.exceptions.RayActorError: + continue + raise RuntimeError('Owned worker death could not be confirmed') + + async def recover_orphans(self, decode, shutdown): + """Decode only this owner's names; never sweep another service's resources.""" + from ray import serve + from ray.util.placement_group import placement_group_table + names = set(serve.status().applications) + names.update(row.get('name', '') for row in placement_group_table().values() if row.get('state') != 'REMOVED') + for name in sorted(names): + instance = decode(name) + if instance is not None: + await asyncio.wait_for(shutdown(instance), self.pool.shutdown_timeout) diff --git a/src/twinkle/server/sampler/app.py b/src/twinkle/server/sampler/app.py index 42e3be7a3..c3faf2867 100644 --- a/src/twinkle/server/sampler/app.py +++ b/src/twinkle/server/sampler/app.py @@ -7,6 +7,8 @@ """ from __future__ import annotations +import asyncio +import os from fastapi import FastAPI, Request from ray import serve from typing import Any @@ -15,13 +17,16 @@ from twinkle.server.config.backend_dispatch import BackendSelector from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.deployment import LazyCleanupMixin, bind_deployment, build_deployment_app +from twinkle.server.exceptions import RequestRejectedError from twinkle.server.middleware.auth import get_token_from_request from twinkle.server.runtime import init_twinkle_runtime from twinkle.server.state import ServerState, get_server_state from twinkle.server.task_queue.mixin import TaskQueueMixin +from twinkle.server.task_queue.types import execution_scope from twinkle.server.utils import get_template_for_model, wrap_builder_with_device_group_env from twinkle.utils.logger import get_logger from .concurrency import SamplerConcurrency +from .runtime import SamplerRuntime from .tinker_handlers import _register_tinker_sampler_routes from .twinkle_handlers import _register_twinkle_sampler_routes @@ -88,7 +93,17 @@ async def __init__(self, queue_config: TaskQueueConfig | None = None, data_plane_url: str | None = None, **kwargs): + self.request_runtime = SamplerRuntime(self) self.device_group = DeviceGroup(**device_group) + context = kwargs.pop('sampling_context', None) + if context is not None: + from twinkle.infra._ray.ray_helper import RayHelper + from twinkle.infra._ray.resource_manager import ResourceManager + os.environ['TWINKLE_RAY_NAMESPACE'] = context['namespace'] + os.environ['TWINKLE_RAY_REGISTRY_NAME'] = context['registry_name'] + RayHelper.resource_manager = ResourceManager.from_allocation(self.device_group, + kwargs.pop('placement_group'), + kwargs.pop('allocation_devices')) self.device_mesh = init_twinkle_runtime( is_mock=(sampler_type == 'mock'), nproc_per_node=nproc_per_node, @@ -98,7 +113,7 @@ async def __init__(self, self.sampler_type = sampler_type self.model_id = model_id replica_context = serve.get_replica_context() - self.replica_id = replica_context.replica_id.unique_id + self.replica_id = context['instance_id'] if context is not None else replica_context.replica_id.unique_id sampler_kwargs: dict[str, Any] = { 'model_id': model_id, @@ -117,6 +132,8 @@ async def __init__(self, else: sampler_kwargs.update(kwargs) self.sampler = _construct_sampler_backend(sampler_type, sampler_kwargs, data_plane_url) + actors = getattr(self.sampler, '_actors', ()) + await asyncio.gather(*(actor.__ray_ready__.remote() for actor in actors)) self.state: ServerState = get_server_state() from twinkle.server.data_plane import DataPlaneProxy @@ -147,8 +164,24 @@ async def shutdown(self) -> None: finally: await self.shutdown_task_queue() await self.sampling_guard.wait_closed() + await self.request_runtime.shutdown() await self.data_plane.close() + async def invoke(self, ticket, payload, stream_queue=None): + return await self.request_runtime.invoke(ticket, payload, stream_queue) + + async def wait_operation(self, ticket): + await self.request_runtime.wait_operation(ticket) + + async def install_base_weights(self, path): + await self.request_runtime.install_base_weights(path) + + async def pending_futures(self): + return await self.request_runtime.pending_futures() + + async def runtime_identity(self): + return await self.request_runtime.runtime_identity() + async def check_health(self) -> None: self.sampling_guard.check_available() @@ -161,7 +194,8 @@ async def _ensure_sticky(self): await self._sticky_entry(sticky_key) async def _on_request_start(self, request: Request) -> str: - await self._ensure_sticky() + if execution_scope.get() is None: + await self._ensure_sticky() await self.state.touch_replica_last_seen(self.replica_id) await self._ensure_state_cleanup_started() return get_token_from_request(request) @@ -204,9 +238,20 @@ def build_sampler_app(model_id: str, # Build the FastAPI app + middleware stack + routes via the shared scaffold, # then bind the Ray Serve deployment. The Sampler passes its FastAPI # title/description/version through ``fastapi_kwargs``. + managed = kwargs.get('sampling_context') is not None + def register_routes(app: FastAPI, get_self: Any) -> None: - _register_tinker_sampler_routes(app, get_self) - _register_twinkle_sampler_routes(app, get_self) + if managed: + + def managed_self(): + # Internal compute URLs must not bypass the controller's tickets. + raise RequestRejectedError('Sampling requests must use the stable sampler endpoint') + + dependency = managed_self + else: + dependency = get_self + _register_tinker_sampler_routes(app, dependency) + _register_twinkle_sampler_routes(app, dependency) app = build_deployment_app( 'Sampler', diff --git a/src/twinkle/server/sampler/backends/streaming.py b/src/twinkle/server/sampler/backends/streaming.py index dc61ff7e7..2ca8edf04 100644 --- a/src/twinkle/server/sampler/backends/streaming.py +++ b/src/twinkle/server/sampler/backends/streaming.py @@ -17,10 +17,20 @@ def stream_to_queue(sampler, queue, inputs, sampling_params=None, adapter_name=' Works with any object that exposes a ``sample_stream`` iterator. """ + generation = sampler.sample_stream(inputs, sampling_params, adapter_name, adapter_path) try: - for delta, reason in sampler.sample_stream(inputs, sampling_params, adapter_name, adapter_path): + for delta, reason in generation: queue.put((delta, reason)) except Exception as e: - queue.put(e) + try: + queue.put(e) + except Exception: + pass # The HTTP consumer may already have closed its queue. finally: - queue.put(STREAM_SENTINEL) + close = getattr(generation, 'close', None) + if close is not None: + close() + try: + queue.put(STREAM_SENTINEL) + except Exception: + pass diff --git a/src/twinkle/server/sampler/controller.py b/src/twinkle/server/sampler/controller.py new file mode 100644 index 000000000..e2246abf8 --- /dev/null +++ b/src/twinkle/server/sampler/controller.py @@ -0,0 +1,322 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Sampling pool decisions; generation ownership outlives HTTP and future timeouts.""" +from __future__ import annotations + +import asyncio +import hashlib +import logging +import math +import time +import uuid +from dataclasses import asdict, dataclass, field + +from twinkle.server.config.application_spec import SamplerArgs +from twinkle.server.exceptions import ResourceExhaustedError +from twinkle.server.resources.instances import InstanceLifecycle +from .concurrency import SamplerConcurrency + +logger = logging.getLogger(__name__) +_CONTROL_ENDPOINTS = frozenset( + {'twinkle/set_template', 'twinkle/add_adapter_to_sampler', 'twinkle/apply_patch', 'twinkle/unload_adapter_paths'}) + + +def sampler_controller_name(namespace, pool_id): + return f'sampler-controller-{namespace}-{pool_id}' + + +def sampling_pool_id(route): + return hashlib.sha256(route.encode()).hexdigest()[:16] + + +@dataclass +class SamplingInstance: + identity: str + epoch: str + state: str = 'STARTING' + application: str | None = None + devices: list[str] = field(default_factory=list) + requests: set[str] = field(default_factory=set) + last_finished: float = 0 + failure: str | None = None + + +class SamplerController: + + def __init__(self, args: SamplerArgs, driver, budgets: dict, *, clock=time.monotonic): + self.args, self.driver, self.budgets, self.clock = args, driver, budgets, clock + self.epoch = uuid.uuid4().hex + self.instances = {} + self.lifecycle = InstanceLifecycle('Sampling') + self._actions = [] + self._base_path = None + self._lock = asyncio.Lock() + self._config_lock = asyncio.Lock() + self._guard = SamplerConcurrency() + self._ready = False + self._stopping = False + self._initialization = None + self._monitor_task = None + + def _owned_count(self): + return self.lifecycle.owned_count(self.instances.values()) + + def _ready_instances(self): + return [instance for instance in self.instances.values() if instance.state == 'READY'] + + def _new_instance(self): + if not self.lifecycle.headroom(self.instances.values(), self.args.max_replicas): + raise ResourceExhaustedError('Sampling replica limit reached') + instance = SamplingInstance('sampler_' + uuid.uuid4().hex, self.epoch, last_finished=self.clock()) + self.instances[instance.identity] = instance + self.lifecycle.start(instance.identity, lambda: self._start(instance)) + return instance + + async def initialize(self): + if self._initialization is None: + self._initialization = asyncio.create_task(self._initialize()) + await asyncio.shield(self._initialization) + + async def _initialize(self): + await self.driver.recover_orphans() + try: + for _ in range(self.args.min_replicas): + instance = self._new_instance() + await asyncio.shield(self.lifecycle.startups[instance.identity]) + self._ready = True + except Exception: + self._stopping = True + await asyncio.gather(*(self._close(instance) for instance in list(self.instances.values()))) + raise + + async def _start(self, instance): + + async def startup(): + await self.driver.start(instance) + async with self._config_lock: + if self._stopping or instance.state != 'STARTING': + raise ResourceExhaustedError('Sampling startup no longer owns this instance') + for action in self._actions: + if 'base_path' in action: + await self.driver.load_weights(instance, action['base_path']) + else: + reply = await self._invoke_finished(instance, action) + if reply['status'] >= 400: + raise RuntimeError('Sampler control replay failed') + instance.state = 'READY' + instance.last_finished = self.clock() + + try: + # Replaying control state is part of readiness, inside the same budget. + await asyncio.wait_for(startup(), self.budgets['startup_timeout_s']) + except BaseException: + await self._close(instance) + raise + + async def _select(self): + while True: + async with self._lock: + if self._stopping or not self._ready: + raise ResourceExhaustedError('Sampling pool is not ready') + ready = self._ready_instances() + active = sum(len(instance.requests) for instance in self.instances.values()) + desired = min(self.args.max_replicas, + max(self.args.min_replicas, math.ceil((active + 1) / self.args.target_ongoing_requests))) + if self._owned_count() < desired: + self._new_instance() + if ready: + instance = min(ready, key=lambda item: (len(item.requests), item.identity)) + ticket = uuid.uuid4().hex + instance.requests.add(ticket) + return instance, ticket + pending = [task for task in self.lifecycle.startups.values() if not task.done()] + if not pending: + raise ResourceExhaustedError('No available sampler instance; device allocation was not confirmed') + await asyncio.wait_for(asyncio.shield(pending[0]), self.budgets['startup_timeout_s']) + + async def _invoke_finished(self, instance, payload): + ticket = uuid.uuid4().hex + reply = await self.driver.invoke(instance, ticket, payload) + await self.driver.wait_operation(instance, ticket) + return reply + + async def request(self, payload, stream_queue=None, full_path=None): + await self.initialize() + endpoint = payload['endpoint'] + if endpoint == 'twinkle/create': + return {'status': 200, 'headers': [(b'content-type', b'application/json')], 'body': b'{"status":"ok"}'} + control = endpoint in _CONTROL_ENDPOINTS + reply_ready = asyncio.get_running_loop().create_future() + + async def operation(): + instance = None + ticket = None + try: + if control: + if not self._ready_instances(): + instance, ticket = await self._select() + async with self._config_lock: + replies = [] + replicas = self._ready_instances() + try: + for replica in replicas: + response = await self._invoke_finished(replica, payload) + if response['status'] >= 400: + await asyncio.gather(*(self._close(peer) for peer in replicas)) + reply_ready.set_result(response) + return + replies.append(response) + except BaseException: + # Even an ambiguous ACK can leave this replica mutated. + # Fence every participant before allowing further routing. + for peer in replicas: + peer.state = 'DRAINING' + await asyncio.gather(*(self._close(peer) for peer in replicas)) + raise + if not replies: + raise ResourceExhaustedError('No ready sampler for control update') + if endpoint != 'twinkle/unload_adapter_paths': + self._actions.append(payload) + reply_ready.set_result(replies[0]) + return + if full_path is not None and full_path != self._base_path: + async with self._config_lock: + try: + for replica in self._ready_instances(): + await self.driver.load_weights(replica, full_path) + self._base_path = full_path + self._actions.append({'base_path': full_path}) + except Exception: + for replica in self._ready_instances(): + await self._close(replica) + raise + instance, ticket = await self._select() + response = await self.driver.invoke(instance, ticket, payload, stream_queue) + if not reply_ready.done(): + reply_ready.set_result(response) + # A pending future reply is not evidence that the generation stopped. + await asyncio.wait_for( + self.driver.wait_operation(instance, ticket), self.args.queue_config.absolute_future_ttl(1)) + except BaseException as error: + if not reply_ready.done(): + reply_ready.set_exception(error) + if instance is not None: + await self._close(instance) + raise + finally: + if instance is not None and ticket is not None: + if instance.state != 'QUARANTINED': + instance.requests.discard(ticket) + instance.last_finished = self.clock() + + async def guarded(): + try: + await self._guard.run(operation, exclusive=control or full_path is not None) + except BaseException as error: + if not reply_ready.done(): + reply_ready.set_exception(error) + raise + + self.lifecycle.spawn(guarded()) + return await asyncio.shield(reply_ready) + + async def _close(self, instance): + if instance.state == 'STOPPED': + return + instance.state = 'DRAINING' + + async def quarantined(error): + instance.state = 'QUARANTINED' + instance.failure = str(error) + + async def stopped(): + instance.state = 'STOPPED' + instance.requests.clear() + + await self.lifecycle.close(instance.identity, lambda: self.driver.shutdown(instance), + self.budgets['shutdown_timeout_s'], stopped, quarantined) + + async def sweep(self): + async with self._lock: + ready = sorted(self._ready_instances(), key=lambda instance: instance.last_finished) + removable = max(0, len(ready) - self.args.min_replicas) + victims = [ + instance for instance in ready + if not instance.requests and self.clock() - instance.last_finished >= self.args.idle_timeout_s + ][:removable] + for instance in victims: + instance.state = 'DRAINING' + missing = self.lifecycle.missing(self.instances.values(), self.args.min_replicas, self.args.max_replicas, + lambda instance: instance.state in ('READY', 'STARTING')) + if not self._stopping: + for _ in range(missing): + self._new_instance() + for instance in list(self.instances.values()): + if instance.state in ('DRAINING', 'QUARANTINED'): + self.lifecycle.spawn(self._close(instance)) + elif instance.state == 'STOPPED' and not instance.requests: + self.instances.pop(instance.identity, None) + self.lifecycle.startups.pop(instance.identity, None) + + async def _monitor(self): + + async def health(instance): + healthy = await self.lifecycle.health(instance.identity, lambda: self.driver.health(instance)) + if healthy is False: + await self._close(instance) + + while not self._stopping: + try: + for instance in self._ready_instances(): + self.lifecycle.spawn_once(('health', instance.identity), lambda item=instance: health(item)) + await self.sweep() + except Exception: + logger.exception('Sampler monitor failed; ownership remains reserved') + await asyncio.sleep(self.budgets['sweep_interval_s']) + + async def diagnostics(self): + await self.initialize() + return { + 'ready': self._ready and not self._stopping and len(self._ready_instances()) >= self.args.min_replicas, + 'epoch': self.epoch, + 'instances': [{ + **asdict(instance), 'requests': len(instance.requests) + } for instance in self.instances.values()] + } + + async def shutdown(self): + self._stopping = True + self._ready = False + if self._monitor_task is not None: + self._monitor_task.cancel() + await asyncio.gather(self._monitor_task, return_exceptions=True) + for task in list(self.lifecycle.startups.values()): + if not task.done(): + task.cancel() + await asyncio.gather(*tuple(self.lifecycle.startups.values()), return_exceptions=True) + await asyncio.gather(*(self._close(instance) for instance in list(self.instances.values()))) + self._guard.close() + await self.lifecycle.join() + + +def create_controller_actor(): + import ray + + from .driver import SamplerDriver + + @ray.remote(num_cpus=0.1, max_restarts=-1, max_task_retries=0) + class SamplerControllerActor(SamplerController): + + def __init__(self, values, namespace, pool_id, options, budgets): + args = SamplerArgs.model_validate(values) + super().__init__(args, SamplerDriver(args, namespace, pool_id, options, budgets), budgets) + + async def initialize(self): + await super().initialize() + if not self._stopping and self._monitor_task is None: + self._monitor_task = asyncio.create_task(self._monitor()) + + async def shutdown(self): + await super().shutdown() + await self.driver.runtime.wait_for_cleanup() + + return SamplerControllerActor diff --git a/src/twinkle/server/sampler/driver.py b/src/twinkle/server/sampler/driver.py new file mode 100644 index 000000000..1144072ff --- /dev/null +++ b/src/twinkle/server/sampler/driver.py @@ -0,0 +1,100 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Owned Ray/Serve effects for sampling instances; no training lifecycle dependency.""" +from __future__ import annotations + +import asyncio + +from twinkle.server.config.application_spec import SamplerArgs +from twinkle.server.resources.runtime import RayRuntime +from twinkle.server.state import get_server_state +from twinkle.server.state.models import FutureFailureRecord + + +class SamplerDriver: + + def __init__(self, args: SamplerArgs, namespace: str, pool_id: str, options: dict, budgets: dict): + self.args, self.namespace, self.pool_id = args, namespace, pool_id + self.options, self.budgets = options, budgets + self.runtime = RayRuntime( + namespace, + allocation_timeout=budgets['allocation_timeout_s'], + shutdown_timeout=budgets['shutdown_timeout_s']) + + @property + def prefix(self): + return f'resources-{self.namespace}-sampling_{self.pool_id}_' + + async def start(self, instance): + from ray.util.placement_group import get_placement_group + + from .app import build_sampler_app + group = self.args.device_group + device_type = str(group.get('device_type', 'cpu' if self.args.sampler_type == 'mock' else 'cuda')).lower() + allocation = await self.runtime.pool.allocate( + group.get('ranks', self.args.nproc_per_node), device_type, f'sampling_{self.pool_id}_{instance.epoch}', + instance.identity) + instance.application = allocation.name + instance.devices = list(allocation.devices) + options = self.runtime.deployment_options(self.options, allocation.name) + values = self.args.model_dump( + exclude={'min_replicas', 'max_replicas', 'idle_timeout_s', 'target_ongoing_requests', 'queue_config'}) + values['device_group'].setdefault('name', 'sampler') + values['device_group'].setdefault('ranks', self.args.nproc_per_node) + values['device_group'].setdefault('device_type', device_type) + app = build_sampler_app( + **values, + queue_config=self.args.queue_config, + deploy_options=options, + sampling_context={ + 'namespace': self.namespace, + 'registry_name': allocation.name + '-registry', + 'instance_id': instance.identity + }, + placement_group=get_placement_group(allocation.name), + allocation_devices=list(allocation.devices)) + await self.runtime.deploy(instance.identity, allocation.name, app, + f'/api/v1/sampling-runtime/{instance.identity}', lambda: self.shutdown(instance)) + + def handle(self, instance): + return self.runtime.handle(instance.application) + + async def invoke(self, instance, ticket, payload, stream_queue=None): + return await self.handle(instance).invoke.remote(ticket, payload, stream_queue) + + async def wait_operation(self, instance, ticket): + await self.handle(instance).wait_operation.remote(ticket) + + async def load_weights(self, instance, path): + await self.handle(instance).install_base_weights.remote(path) + + async def health(self, instance): + return await asyncio.wait_for(self.handle(instance).runtime_identity.remote(), 5) + + async def shutdown(self, instance): + name = instance.application or self.runtime.pool.name(f'sampling_{self.pool_id}_{instance.epoch}', + instance.identity) + futures = [] + try: + futures = await asyncio.wait_for(self.handle(instance).pending_futures.remote(), 5) + except Exception: + pass + await self.runtime.shutdown(instance.identity, name) + failure = FutureFailureRecord( + reason_code='orphaned_replica', + attribution='server', + message='Sampler instance stopped; retry with the checkpoint path.') + state = get_server_state() + for request_id in futures: + await state.store_future_status(request_id, 'failed', None, failure=failure) + + async def recover_orphans(self): + import re + from types import SimpleNamespace + pattern = re.compile(re.escape(self.prefix) + r'([0-9a-f]{32})-(sampler_[0-9a-f]{32})$') + + def decode(name): + match = pattern.fullmatch(name) + if match: + return SimpleNamespace(identity=match[2], epoch=match[1], application=name) + + await self.runtime.recover_orphans(decode, self.shutdown) diff --git a/src/twinkle/server/sampler/router.py b/src/twinkle/server/sampler/router.py new file mode 100644 index 000000000..3342c46a7 --- /dev/null +++ b/src/twinkle/server/sampler/router.py @@ -0,0 +1,135 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Stable CPU ingress for an owned sampling pool; existing sampler wire is forwarded intact.""" +from __future__ import annotations + +import asyncio +import json +import uuid +from copy import deepcopy +from fastapi import Depends, Request, Response +from fastapi.responses import StreamingResponse + +from twinkle.server.checkpoint import create_checkpoint_manager +from twinkle.server.deployment import bind_deployment, build_deployment_app +from twinkle.server.middleware.auth import get_token_from_request +from twinkle.server.state import get_server_state +from .controller import create_controller_actor, sampler_controller_name, sampling_pool_id +from .weights import is_full_checkpoint + + +class SamplerRouter: + + async def __init__(self, namespace, pool_id): + import ray + self.state = get_server_state() + self.controller = ray.get_actor(sampler_controller_name(namespace, pool_id), namespace=namespace) + + async def forward(self, request, endpoint): + body = await request.body() + payload = { + 'endpoint': endpoint, + 'body': body, + 'method': request.method, + 'headers': list(request.headers.raw), + 'query_string': request.scope.get('query_string', b'') + } + full_path = None + if endpoint in ('twinkle/sample', 'twinkle/sample_stream', 'twinkle/sample_to_data_plane', 'tinker/asample'): + try: + data = json.loads(body) + except (ValueError, UnicodeDecodeError): + data = {} + if not isinstance(data, dict): + data = {} + # The existing handler remains responsible for its validation/wire errors. + path = data.get('adapter_uri') or data.get('model_path') + if not path and endpoint == 'tinker/asample' and data.get('sampling_session_id'): + session = await self.state.get_sampling_session(data['sampling_session_id']) + path = session.get('model_path') if session else None + if path: + client_type = 'tinker' if endpoint.startswith('tinker/') else 'twinkle' + _, resolved = create_checkpoint_manager( + get_token_from_request(request), client_type=client_type).parse_adapter_uri(path) + if await asyncio.to_thread(is_full_checkpoint, resolved): + full_path = resolved + stream = endpoint == 'twinkle/sample_stream' + queue = None + if stream: + from ray.util.queue import Queue + queue = Queue(maxsize=128) + try: + reply = await self.controller.request.remote(payload, queue, full_path) + except BaseException: + if queue is not None: + queue.shutdown(force=True) + raise + headers = { + key.decode(): value.decode() + for key, value in reply['headers'] if key.lower() not in (b'content-length', b'transfer-encoding') + } + if not stream: + return Response(reply['body'], status_code=reply['status'], headers=headers) + + async def generate(): + try: + while True: + item = await queue.get_async() + if item is None: + return + yield item + finally: + queue.shutdown(force=True) + + return StreamingResponse(generate(), status_code=reply['status'], headers=headers) + + +def build_sampler_router(namespace, pool_id, deploy_options): + + def routes(app, get_self): + + @app.get('/healthz') + async def health(self=Depends(get_self)): + diagnostic = await self.controller.diagnostics.remote() + return {'healthy': diagnostic['ready'], **diagnostic} + + @app.post('/twinkle/{endpoint}') + async def native(request: Request, endpoint: str, self=Depends(get_self)): + return await self.forward(request, f'twinkle/{endpoint}') + + @app.post('/tinker/{endpoint}') + async def tinker(request: Request, endpoint: str, self=Depends(get_self)): + return await self.forward(request, f'tinker/{endpoint}') + + app = build_deployment_app('SamplerRouter', routes) + return bind_deployment( + app, SamplerRouter, deploy_options, deployment_name='SamplerRouter', bind_args=(namespace, pool_id)) + + +def deploy_sampler_pool(spec, namespace, budgets, env): + """Driver-only composition. Each compute instance owns its external allocation.""" + import ray + from ray import serve + options = deepcopy(spec.deployments[0]) if spec.deployments else {} + options.pop('name', None) + options.pop('autoscaling_config', None) + options['num_replicas'] = 1 + actor_options = options.setdefault('ray_actor_options', {}) + actor_options.setdefault('num_cpus', 0.1) + runtime_env = actor_options.setdefault('runtime_env', {}) + runtime_env['env_vars'] = {**env, **runtime_env.get('env_vars', {}), 'TWINKLE_STATE_RAY_NAMESPACE': namespace} + pool_id = sampling_pool_id(spec.route_prefix) + name = sampler_controller_name(namespace, pool_id) + controller = create_controller_actor().options( + name=name, namespace=namespace, runtime_env=runtime_env).remote(spec.args.model_dump(), namespace, pool_id, + options, budgets) + try: + ray.get(controller.initialize.remote()) + app = build_sampler_router(namespace, pool_id, options) + serve.run(app, name=spec.name, route_prefix=spec.route_prefix) + except BaseException: + try: + ray.get(controller.shutdown.remote()) + finally: + ray.kill(controller, no_restart=True) + raise + return controller diff --git a/src/twinkle/server/sampler/runtime.py b/src/twinkle/server/sampler/runtime.py new file mode 100644 index 000000000..bb80f878b --- /dev/null +++ b/src/twinkle/server/sampler/runtime.py @@ -0,0 +1,147 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Request ownership inside one allocated sampling replica, independent of HTTP waiters.""" +from __future__ import annotations + +import asyncio + +from twinkle.server.deployment import build_deployment_app +from twinkle.server.exceptions import ResourceNotFoundError +from twinkle.server.task_queue.types import ExecutionScope, execution_scope + + +class SamplerRuntime: + + def __init__(self, owner): + self.owner = owner + self._operations = {} + self._stream_tasks = set() + self._rpc_app = None + + def start_stream_job(self, factory): + job = asyncio.create_task(factory()) + scope = execution_scope.get() + if scope is not None: + scope.jobs.add(job) + self._stream_tasks.add(job) + + def done(task): + self._stream_tasks.discard(task) + if not task.cancelled(): + task.exception() + + job.add_done_callback(done) + return job + + async def shutdown(self): + await asyncio.gather(*tuple(self._stream_tasks), return_exceptions=True) + + async def invoke(self, ticket, payload, stream_queue=None): + """Return the HTTP reply early while retaining every actual generation job.""" + if ticket in self._operations: + raise ResourceNotFoundError('Sampling request was already dispatched') + if self._rpc_app is None: + from .tinker_handlers import _register_tinker_sampler_routes + from .twinkle_handlers import _register_twinkle_sampler_routes + + def routes(app, get_self): + _register_tinker_sampler_routes(app, lambda: self.owner) + _register_twinkle_sampler_routes(app, lambda: self.owner) + + self._rpc_app = build_deployment_app('Sampler', routes) + loop = asyncio.get_running_loop() + reply = loop.create_future() + record = {'scope': ExecutionScope(), 'task': None} + self._operations[ticket] = record + + async def dispatch(): + token = execution_scope.set(record['scope']) + response = {'status': 500, 'headers': [], 'body': b''} + sent = False + + async def receive(): + nonlocal sent + if not sent: + sent = True + return {'type': 'http.request', 'body': payload['body'], 'more_body': False} + await asyncio.Event().wait() + + async def send(message): + if message['type'] == 'http.response.start': + response['status'] = message['status'] + response['headers'] = message['headers'] + if stream_queue is not None and not reply.done(): + reply.set_result(dict(response)) + elif message['type'] == 'http.response.body': + chunk = message.get('body', b'') + if stream_queue is None: + response['body'] += chunk + elif chunk: + try: + await stream_queue.put_async(chunk) + except Exception: + pass # Keep draining generation after an HTTP consumer disconnects. + + scope = { + 'type': 'http', + 'asgi': { + 'version': '3.0' + }, + 'http_version': '1.1', + 'scheme': 'http', + 'method': payload['method'], + 'path': '/' + payload['endpoint'], + 'raw_path': ('/' + payload['endpoint']).encode(), + 'root_path': '', + 'query_string': payload.get('query_string', b''), + 'headers': payload['headers'], + 'server': ('sampler', 80), + 'client': ('controller', 0) + } + try: + await self._rpc_app(scope, receive, send) + if not reply.done(): + reply.set_result(response) + except BaseException as error: + if not reply.done(): + reply.set_exception(error) + raise + finally: + execution_scope.reset(token) + # ComputeWorker supplied completion handles, not a terminal status read. + await record['scope'].wait() + if stream_queue is not None: + try: + await stream_queue.put_async(None) + except Exception: + pass + + record['task'] = asyncio.create_task(dispatch()) + record['task'].add_done_callback(lambda task: task.exception() if not task.cancelled() else None) + return await asyncio.shield(reply) + + async def wait_operation(self, ticket): + record = self._operations.get(ticket) + if record is None: + raise ResourceNotFoundError('Unknown sampling operation') + await asyncio.shield(record['task']) + self.owner.sampling_guard.check_available() + self._operations.pop(ticket, None) + + async def install_base_weights(self, path): + from .weights import resolve_sampler_weights + await self.owner.sampling_guard.run( + lambda: resolve_sampler_weights(self.owner, path, full_checkpoint=True), exclusive=True) + + async def pending_futures(self): + return list({identity for record in self._operations.values() for identity in record['scope'].request_ids}) + + async def runtime_identity(self): + import ray + actors = getattr(self.owner.sampler, '_actors', ()) + states = await asyncio.to_thread( + lambda: [ray._private.state.actors(actor._actor_id.hex()).get('State') for actor in actors]) + return { + 'unhealthy': + self.owner.sampling_guard._closed or self.owner.sampling_guard._weights_error is not None + or any(state != 'ALIVE' for state in states) + } diff --git a/src/twinkle/server/sampler/twinkle_handlers.py b/src/twinkle/server/sampler/twinkle_handlers.py index b6cf58d3b..b23f65e4a 100644 --- a/src/twinkle/server/sampler/twinkle_handlers.py +++ b/src/twinkle/server/sampler/twinkle_handlers.py @@ -165,9 +165,10 @@ async def _stream_queue(q, sentinel, request_id: str, total_timeout: float, sing async def _await_generation(service: SamplerManagement, submission_id: str): """Poll one admitted generation through the backend boundary.""" collected = False + confirmed = False async def poll(): - nonlocal collected + nonlocal collected, confirmed poll_interval = 0.01 while True: try: @@ -180,6 +181,7 @@ async def poll(): await asyncio.sleep(poll_interval) poll_interval = min(poll_interval * 1.5, 0.25) continue + confirmed = bool(states) and all(state.get('status') in ('completed', 'failed') for state in states) failed = next( (state for state in states if state.get('status') not in ('running', 'completed')), None, @@ -198,7 +200,11 @@ async def poll(): return await poll() finally: if not collected: - + # A submission-cancel ACK can precede the engine-loop finalizer. + # Fence this runtime; its operation may finish only after confirmed + # worker death, not after a transport future reports cancellation. + if not confirmed: + service.sampling_guard.invalidate_weights(RuntimeError('Generation completion was not confirmed')) try: await service.call_backend(service.sampler.cancel_generation, submission_id) except Exception: @@ -368,8 +374,11 @@ async def add_adapter_to_sampler( raise RequestRejectedError('`adapter_name` is required and must be non-empty.') full_adapter_name = resolve_twinkle_adapter_name(request, body.adapter_name) - from peft import LoraConfig - config = LoraConfig(**body.config) if isinstance(body.config, dict) else body.config + def build_config(): + from peft import LoraConfig + return LoraConfig(**body.config) if isinstance(body.config, dict) else body.config + + config = await asyncio.to_thread(build_config) with traced_operation('sampler.add_adapter_to_sampler', attrs={MODEL_ID: self.model_id}): await self.sampling_guard.run( @@ -405,7 +414,7 @@ async def sample_stream( """ token = await self._on_request_start(request) - adapter_path = None + resolved_uri = None adapter_name = body.adapter_name or '' full_adapter_name = resolve_twinkle_adapter_name(request, adapter_name) or '' @@ -413,8 +422,6 @@ async def sample_stream( from twinkle.server.checkpoint import create_checkpoint_manager checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') _, resolved_uri = checkpoint_manager.parse_adapter_uri(body.adapter_uri) - await self.call_backend(self.sampler.reset_prefix_cache) - adapter_path = await resolve_sampler_weights(self, resolved_uri) # Streaming accepts exactly one input; the shared seam enforces that and # returns a single parsed object. Its ValueError maps to the same 400 this @@ -443,14 +450,31 @@ async def _no_actor_generator(): return StreamingResponse(_no_actor_generator(), media_type='application/x-ndjson') q = Queue(maxsize=128) actor = actors[0] - actor.sample_stream_to_queue.remote( - q, - inputs_parsed, - params, - adapter_name=full_adapter_name, - adapter_path=adapter_path, - ) + async def stream_job(): + + async def _stream_generation(adapter_path): + # Retain the actor result, not merely the HTTP queue iterator. + ref = actor.sample_stream_to_queue.remote( + q, inputs_parsed, params, adapter_name=full_adapter_name, adapter_path=adapter_path) + try: + await ref + except BaseException as error: + # An errored bridge cannot certify the inner engine request + # exited. Controller must confirm worker death before reuse. + self.sampling_guard.invalidate_weights(error) + raise + + try: + await sample_with_weights(self, resolved_uri, _stream_generation) + except Exception as error: + try: + await q.put_async(error) + await q.put_async(STREAM_SENTINEL) + except Exception: + pass + + self.request_runtime.start_stream_job(stream_job) return StreamingResponse( _stream_queue( q, diff --git a/src/twinkle/server/task_queue/mixin.py b/src/twinkle/server/task_queue/mixin.py index b887ed20b..6c2dfb7c2 100644 --- a/src/twinkle/server/task_queue/mixin.py +++ b/src/twinkle/server/task_queue/mixin.py @@ -47,9 +47,8 @@ class TaskQueueMixin: concurrent requests itself, such as vLLM sampling. 3. Background task (schedule_background_task): - asyncio.create_task, runs concurrently with compute queue. Use for I/O - or an engine that provides its own safe concurrency and batching. - Status is still tracked; clients can poll the same status endpoints. + Reuses the supervised concurrent path for pure I/O. Status, capacity + and shutdown are owned by the same worker as concurrent compute. Requirements ------------ @@ -248,7 +247,7 @@ async def _schedule_task( await self._enqueue_prepared_task(task) except BaseException: if task.on_finished is not None and await self.state.get_future(request_id) is None: - await task.on_finished(False, False) + await self._compute_worker.finish_task(task, False) raise return {'request_id': request_id, 'model_id': model_id} @@ -280,8 +279,7 @@ async def _enqueue_prepared_task(self, task: QueuedTask) -> None: attribution='server', message='Replica stopped before the task could be enqueued.')) finally: - if task.on_finished is not None: - await task.on_finished(False, False) + await self._compute_worker.finish_task(task, False) raise queue_key = self._queue_key(model_id=model_id, token=task.token) self._compute_worker.ensure_queue_registered(queue_key) @@ -289,8 +287,6 @@ async def _enqueue_prepared_task(self, task: QueuedTask) -> None: self._compute_worker.track_task(request_id) await q.put(task) try: - if task.on_enqueued is not None: - await task.on_enqueued() await self.state.store_future_status( request_id, TaskStatus.QUEUED.value, diff --git a/src/twinkle/server/task_queue/types.py b/src/twinkle/server/task_queue/types.py index 6e522785e..b5f0fd61f 100644 --- a/src/twinkle/server/task_queue/types.py +++ b/src/twinkle/server/task_queue/types.py @@ -11,7 +11,8 @@ import asyncio from collections.abc import Callable, Coroutine -from dataclasses import dataclass +from contextvars import ContextVar +from dataclasses import dataclass, field from enum import Enum # Ray_Get_Timeout is classified the same as asyncio.TimeoutError: 504/Server. @@ -58,6 +59,19 @@ class QueueState(Enum): UNKNOWN = 'unknown' # Unknown or unspecified state +@dataclass +class ExecutionScope: + """Completion witnesses for an operation, independent of its client future.""" + jobs: set[asyncio.Task] = field(default_factory=set) + request_ids: set[str] = field(default_factory=set) + + async def wait(self): + await asyncio.gather(*tuple(self.jobs), return_exceptions=True) + + +execution_scope = ContextVar('execution_scope', default=None) + + @dataclass class QueuedTask: """Dataclass representing a task waiting in the compute queue.""" @@ -69,4 +83,3 @@ class QueuedTask: task_type: str | None created_at: float on_finished: Callable[[bool, bool], Coroutine] | None = None - on_enqueued: Callable[[], Coroutine] | None = None diff --git a/src/twinkle/server/task_queue/worker.py b/src/twinkle/server/task_queue/worker.py index 86aba9e61..a5bfc00ec 100644 --- a/src/twinkle/server/task_queue/worker.py +++ b/src/twinkle/server/task_queue/worker.py @@ -27,7 +27,8 @@ from twinkle.server.telemetry.correlation import MODEL_ID, TOKEN_ID from twinkle.server.telemetry.tracing import traced_operation from twinkle.utils.logger import get_logger -from .types import BACKEND_TIMEOUT_EXCEPTIONS, BackendBusyError, QueuedTask, QueueState, TaskStatus, UserTaskError +from .types import (BACKEND_TIMEOUT_EXCEPTIONS, BackendBusyError, QueuedTask, QueueState, TaskStatus, UserTaskError, + execution_scope) if TYPE_CHECKING: from twinkle.server.state import ServerState @@ -108,6 +109,7 @@ def __init__( self.task_queues: dict[str, asyncio.Queue] = {} self._request_ids: set[str] = set() + self._task_outcomes: dict[str, dict[str, Any]] = {} self.queue_order: Deque[str] = deque() self.new_task_event: asyncio.Event = asyncio.Event() @@ -185,6 +187,33 @@ def has_task(self, request_id: str) -> bool: """True while a task is queued or executing, including a terminal-future race.""" return request_id in self._request_ids + def get_task_outcome(self, request_id: str) -> dict[str, Any] | None: + """Read bounded local execution facts, never infer them from a future status.""" + cutoff = time.time() - self._config.absolute_future_ttl(1) + for identity, outcome in list(self._task_outcomes.items()): + if outcome['finished_at'] <= cutoff: + self._task_outcomes.pop(identity) + outcome = self._task_outcomes.get(request_id) + return dict(outcome) if outcome is not None else None + + async def finish_task(self, + task: QueuedTask, + executed: bool, + unsafe: bool = False, + finished_at: float | None = None) -> None: + if task.on_finished is not None: + self.get_task_outcome(task.request_id) + self._task_outcomes[task.request_id] = { + 'state': 'unsafe' if unsafe else 'finished', + 'executed': executed, + 'finished_at': time.time() if finished_at is None else finished_at, + } + self._request_ids.discard(task.request_id) + if task.on_finished is not None: + acknowledged = await task.on_finished(executed, unsafe) + if acknowledged is True: + self._task_outcomes.pop(task.request_id, None) + def total_queued(self) -> int: """Total number of pending tasks across all per-key queues.""" return sum(q.qsize() for q in self.task_queues.values()) @@ -282,9 +311,7 @@ async def fail_queue_tasks(self, queue_key: str, reason: str) -> None: for task in drained: await self._store_task_failed(task, reason, QueueState.UNKNOWN.value, queue_state_reason=reason) q.task_done() - self._request_ids.discard(task.request_id) - if task.on_finished is not None: - await task.on_finished(False, False) + await self.finish_task(task, False) self.task_queues.pop(queue_key, None) try: @@ -311,9 +338,7 @@ async def _fail_timed_out_task(self, task: QueuedTask, waited: float, q: asyncio await self._store_task_failed(task, error, QueueState.PAUSED_CAPACITY.value, queue_state_reason=error) self._record_tasks_total(task.task_type or 'unknown', 'timeout') q.task_done() - self._request_ids.discard(task.request_id) - if task.on_finished is not None: - await task.on_finished(False, False) + await self.finish_task(task, False) async def _execute_task(self, task: QueuedTask, queue_key: str, q: asyncio.Queue) -> None: """Execute a dequeued task, then release its queue slot.""" @@ -396,6 +421,10 @@ async def run(): running = asyncio.create_task(run()) self._detached_tasks[running] = task + scope = execution_scope.get() + if scope is not None: + scope.jobs.add(running) + scope.request_ids.add(task.request_id) def done(completed: asyncio.Task) -> None: self._detached_tasks.pop(completed, None) @@ -417,6 +446,7 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: task_status = 'completed' exec_time = 0.0 executed = False + finished_at = None # One span per queued task execution, tagged `.` # (e.g. `model.forward`, `sampler.sample`). queue_attrs = { @@ -442,6 +472,7 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: # effective_execution_timeout is always positive (0 -> finite fallback), # so wait_for is always in effect. result = await asyncio.wait_for(coro, timeout=self._config.effective_execution_timeout) + finished_at = time.time() exec_time = time.monotonic() - exec_start logger.info(f'[ComputeWorker] Task {task.request_id} completed in {exec_time:.2f}s, type={task_type}') await self._state.store_future_status( @@ -531,9 +562,7 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: traceback_text=traceback.format_exc(), ) finally: - self._request_ids.discard(task.request_id) - if task.on_finished is not None: - await task.on_finished(executed, executed and task_status in ('timeout', 'cancelled')) + await self.finish_task(task, executed, executed and task_status in ('timeout', 'cancelled'), finished_at) self._record_execution_time(task_type, exec_time) self._record_tasks_total(task_type, task_status) diff --git a/src/twinkle/server/training/actor.py b/src/twinkle/server/training/actor.py index 6eb1245f6..bae531238 100644 --- a/src/twinkle/server/training/actor.py +++ b/src/twinkle/server/training/actor.py @@ -35,8 +35,6 @@ def __init__(self, config, namespace): self.core = TrainingController(self.config, self.store, self.driver) self._metrics = get_training_metrics() self._monitors = [] - self._unreachable_since = {} - self._failures: dict[str, asyncio.Task] = {} self._initialization = asyncio.get_event_loop().create_task(self._initialize()) async def initialize(self): @@ -74,35 +72,21 @@ async def _heartbeat(self): await asyncio.sleep(5) async def _check_instance_health(self, instance): - try: + + async def probe(): identity = await self.driver.health(instance) - if identity['unhealthy']: - self._fail_unhealthy(instance, 'Training runtime reported a failed worker') - return - await self.store.state.touch_replica_last_seen(instance.instance_id) - self._unreachable_since.pop(instance.instance_id, None) - except Exception: - since = self._unreachable_since.setdefault(instance.instance_id, time.monotonic()) - if time.monotonic() - since >= 30: - self._fail_unhealthy(instance, 'Training runtime became unavailable') + if not identity['unhealthy']: + await self.store.state.touch_replica_last_seen(instance.instance_id) + return identity + + healthy = await self.core.lifecycle.health(instance.instance_id, probe) + if healthy is False: + self._fail_unhealthy(instance, 'Training runtime became unavailable or reported a failed worker') def _fail_unhealthy(self, instance, message): - identity = instance.instance_id - if identity in self._failures: - return - task = asyncio.create_task( - self.core.fail_instance( - identity, FutureFailureRecord(reason_code='orphaned_replica', attribution='server', message=message))) - self._failures[identity] = task - - def done(completed): - self._failures.pop(identity, None) - if not completed.cancelled(): - error = completed.exception() - if error is not None: - logger.error('Unhealthy runtime cleanup failed', exc_info=(type(error), error, error.__traceback__)) - - task.add_done_callback(done) + self.core.lifecycle.spawn_once(('health', instance.instance_id), lambda: self.core.fail_instance( + instance.instance_id, + FutureFailureRecord(reason_code='orphaned_replica', attribution='server', message=message))) async def _sweep(self): while True: @@ -151,9 +135,6 @@ async def admit(self, *args, **kwargs): await self.initialize() return await self.core.admit(*args, **kwargs) - async def commit_operation(self, *args): - return await self.core.commit_operation(*args) - async def finish_operation(self, *args, **kwargs): return await self.core.finish_operation(*args, **kwargs) @@ -190,6 +171,5 @@ async def shutdown(self): task.cancel() await asyncio.gather(*self._monitors, return_exceptions=True) await self.core.shutdown() - await asyncio.gather(*self._failures.values(), return_exceptions=True) await self.core.wait_for_jobs() - await self.driver.wait_for_cleanup() + await self.driver.runtime.wait_for_cleanup() diff --git a/src/twinkle/server/training/controller.py b/src/twinkle/server/training/controller.py index 455c04027..3abbae71c 100644 --- a/src/twinkle/server/training/controller.py +++ b/src/twinkle/server/training/controller.py @@ -15,6 +15,7 @@ from twinkle.server.exceptions import (RequestRejectedError, ResourceExhaustedError, ResourceNotFoundError, ResourceQuotaExceededError) from twinkle.server.resources.allocation import CapacityUnavailableError +from twinkle.server.resources.instances import InstanceLifecycle from twinkle.server.state.models import FutureFailureRecord from twinkle.server.state.training_records import (LIVE_INSTANCES, TERMINAL_TRAINING, CreateTraining, InstanceRecord, TrainingRecord) @@ -85,10 +86,8 @@ def __init__(self, self.epoch = 'controller_' + uuid.uuid4().hex self.instances: dict[str, InstanceRecord] = {} self.trainings: dict[str, TrainingRecord] = {} - self._startup: dict[str, asyncio.Task] = {} - self._jobs: set[asyncio.Task] = set() + self.lifecycle = InstanceLifecycle('Training') self._cleaning: set[str] = set() - self._closing: dict[str, asyncio.Task] = {} self._reconciling: set[str] = set() self._failed_instances: dict[str, FutureFailureRecord] = {} self._lock = asyncio.Lock() @@ -141,21 +140,6 @@ def _quarantine_reason(exc: Exception) -> str: return 'pg_removal' return 'other' - def _spawn(self, coroutine) -> asyncio.Task: - task = asyncio.create_task(coroutine) - self._jobs.add(task) - - def done(completed): - self._jobs.discard(completed) - if completed.cancelled(): - return - exc = completed.exception() - if exc is not None: - logger.error('Training background job failed', exc_info=(type(exc), exc, exc.__traceback__)) - - task.add_done_callback(done) - return task - async def initialize(self) -> None: """Discard old state only after owned compute has been confirmed dead.""" old_instances = await self.store.instances() @@ -191,7 +175,7 @@ async def initialize(self) -> None: for _ in range(spec.config.min_replicas): async with self._lock: instance = await self._new_instance(spec, resident=True) - await asyncio.shield(self._startup[instance.instance_id]) + await asyncio.shield(self.lifecycle.startups[instance.instance_id]) except Exception: await self.shutdown() raise @@ -234,6 +218,9 @@ async def _assert_epoch(self) -> None: raise ResourceNotFoundError('Training controller epoch is no longer active') async def _new_instance(self, spec: TrainingSpec, *, resident: bool = False) -> InstanceRecord: + matching = [item for item in self.instances.values() if (item.base_model, item.mode) == spec.key] + if not self.lifecycle.headroom(matching, spec.config.max_replicas): + raise CapacityUnavailableError('Training replica limit reached') instance = InstanceRecord( instance_id='instance_' + uuid.uuid4().hex, controller_epoch=self.epoch, @@ -244,7 +231,7 @@ async def _new_instance(self, spec: TrainingSpec, *, resident: bool = False) -> await self.store.save_instance(instance) self.instances[instance.instance_id] = instance self._report_instance_metrics() - self._startup[instance.instance_id] = self._spawn(self._start_instance(instance)) + self.lifecycle.start(instance.instance_id, lambda: self._start_instance(instance)) return instance async def _start_instance(self, instance: InstanceRecord) -> None: @@ -346,9 +333,8 @@ async def create(self, request: CreateTraining) -> dict[str, str]: if not self.residents_ready_for(spec): raise CapacityUnavailableError( 'Resident training capacity is recovering; retry a new create later.') - count = sum((item.base_model, item.mode) == spec.key and item.state != 'STOPPED' - for item in self.instances.values()) - if spec.config.max_replicas is not None and count >= spec.config.max_replicas: + matching = [item for item in self.instances.values() if (item.base_model, item.mode) == spec.key] + if not self.lifecycle.headroom(matching, spec.config.max_replicas): raise CapacityUnavailableError('Training replica limit reached; retry a new create later.') instance = await self._new_instance(spec) training.instance_id = instance.instance_id @@ -357,7 +343,7 @@ async def create(self, request: CreateTraining) -> dict[str, str]: self.trainings[training_id] = training await self.store.save_training(training) await self.store.save_instance(instance) - self._spawn(self._create_user(training, request)) + self.lifecycle.spawn(self._create_user(training, request)) except Exception as exc: if training is not None: self.trainings.pop(training_id, None) @@ -367,7 +353,7 @@ async def create(self, request: CreateTraining) -> dict[str, str]: instance.training_ids.remove(training_id) await self.store.save_instance(instance) if not instance.training_ids and not instance.resident: - self._spawn(self.fail_instance(instance.instance_id, failure_for(exc))) + self.lifecycle.spawn(self.fail_instance(instance.instance_id, failure_for(exc))) training.state = 'FAILED' training.terminal_reason = failure_for(exc).reason_code await self.store.retire(training) @@ -377,7 +363,7 @@ async def create(self, request: CreateTraining) -> dict[str, str]: async def _create_user(self, training: TrainingRecord, request: CreateTraining) -> None: instance = self.instances[training.instance_id] try: - await asyncio.shield(self._startup[instance.instance_id]) + await asyncio.shield(self.lifecycle.startups[instance.instance_id]) async with self._lock: future = await self.store.state.get_future(training.create_request_id) if future and future['status'] == 'cancelled': @@ -387,7 +373,11 @@ async def _create_user(self, training: TrainingRecord, request: CreateTraining) raise RequestRejectedError('Training creation is no longer pending') training.creation_deadline = self.clock() + self.config.startup_timeout_s await self.store.save_training(training) - await self.store.future(training.create_request_id, 'running', training.training_id) + if not await self.store.state.try_start_future(training.create_request_id): + future = await self.store.state.get_future(training.create_request_id) + if future and future['status'] == 'cancelled': + training.terminal_reason = 'cancelled' + raise RequestRejectedError('Training creation did not acquire execution') remaining = training.creation_deadline - self.clock() await asyncio.wait_for(self.driver.create_training(instance, training, request), max(0, remaining)) # Driver includes training-run metadata commit, before public success. @@ -464,38 +454,18 @@ async def admit(self, training.state = 'CLOSING' training.terminal_reason = 'expired' await self.store.save_training(training) - self._spawn(self._cleanup_user(training)) + self.lifecycle.spawn(self._cleanup_user(training)) raise TrainingExpiredError( 'Training handle expired; create a new handle and explicitly load a checkpoint.') instance = self.instances[training.instance_id] if instance.state != 'READY': raise ResourceNotFoundError('Training runtime is unavailable') if request_id not in training.pending_operation_ids: - training.pending_operation_ids[request_id] = { - 'accepted_at': self.clock(), - 'committed': False, - 'renew': renew - } + training.pending_operation_ids[request_id] = {'accepted_at': self.clock(), 'renew': renew} await self.store.save_training(training) await self.store.own_future(request_id, token) return {'instance_id': instance.instance_id, 'epoch': self.epoch} - async def commit_operation(self, training_id: str, request_id: str, epoch: str) -> None: - async with self._lock: - await self._assert_epoch() - training = self.trainings.get(training_id) - if epoch != self.epoch or training is None: - raise ResourceNotFoundError('Operation admission no longer exists') - if request_id not in training.pending_operation_ids: - # A very short task can finish while the enqueue acknowledgment - # is in flight. Its original terminal future is conclusive. - future = await self.store.state.get_future(request_id) - if future and future['status'] in ('completed', 'failed', 'cancelled'): - return - raise ResourceNotFoundError('Operation admission no longer exists') - training.pending_operation_ids[request_id]['committed'] = True - await self.store.save_training(training) - async def finish_operation(self, training_id: str, request_id: str, @@ -511,16 +481,21 @@ async def finish_operation(self, training = self.trainings.get(training_id) if epoch != self.epoch or training is None: return - ticket = training.pending_operation_ids.pop(request_id, None) + ticket = training.pending_operation_ids.get(request_id) if ticket is None: return + updated = training.model_copy(deep=True) + updated.pending_operation_ids.pop(request_id) if executed and ticket['renew']: age = max(0, time.time() - finished_at) if finished_at is not None else 0 - training.last_activity_at = max(training.last_activity_at, self.clock() - age) - await self.store.save_training(training) + updated.last_activity_at = max(training.last_activity_at, self.clock() - age) + # A failed completion write must leave the original ticket retryable. + await self.store.save_training(updated) + training.pending_operation_ids = updated.pending_operation_ids + training.last_activity_at = updated.last_activity_at cleanup = training.state == 'CLOSING' and not training.pending_operation_ids if cleanup: - self._spawn(self._cleanup_user(training)) + self.lifecycle.spawn(self._cleanup_user(training)) async def route(self, training_id: str, token: str, session_id: str | None = None) -> dict[str, str]: async with self._lock: @@ -538,7 +513,7 @@ async def release(self, training_id: str, token: str, session_id: str | None = N await self._cleanup_user(training) config = self.catalog.get(training.base_model, training.mode).config.queue_config deadline = asyncio.get_running_loop().time() + config.absolute_future_ttl(1) + self.config.shutdown_timeout_s - while training.state == 'CLOSING' or self.instances[training.instance_id].state == 'DRAINING': + while training.state == 'CLOSING' or self.instances[training.instance_id].state in ('DRAINING', 'QUARANTINED'): if self.instances[training.instance_id].state == 'QUARANTINED': raise RuntimeError('Runtime cleanup is quarantined; device release is not confirmed') if asyncio.get_running_loop().time() >= deadline: @@ -628,25 +603,11 @@ async def _cleanup_user(self, training: TrainingRecord) -> None: self._cleaning.discard(training.training_id) async def _close_instance(self, instance: InstanceRecord) -> None: - close_task = self._closing.get(instance.instance_id) - if close_task is None: - close_task = self._spawn(self._close_instance_once(instance)) - self._closing[instance.instance_id] = close_task - - def clear(completed: asyncio.Task) -> None: - if self._closing.get(instance.instance_id) is completed: - self._closing.pop(instance.instance_id, None) - - close_task.add_done_callback(clear) - await asyncio.shield(close_task) - - async def _close_instance_once(self, instance: InstanceRecord) -> None: started = time.monotonic() tags = self._instance_metric_tags(instance) timing_tags = {'backend': tags['backend'], 'mode': tags['mode']} - try: - await asyncio.wait_for(self.driver.shutdown(instance), self.config.shutdown_timeout_s) - except Exception as exc: + + async def quarantined(exc): async with self._lock: quarantined = instance.model_copy(update={'state': 'QUARANTINED', 'failure': str(exc)}) await self.store.save_instance(quarantined) @@ -660,8 +621,18 @@ async def _close_instance_once(self, instance: InstanceRecord) -> None: 'reason': self._quarantine_reason(exc), }) self._metrics.shutdown_seconds.observe(time.monotonic() - started, tags={**timing_tags, 'status': 'error'}) - return - self._metrics.shutdown_seconds.observe(time.monotonic() - started, tags={**timing_tags, 'status': 'success'}) + + async def stopped(): + self._metrics.shutdown_seconds.observe( + time.monotonic() - started, tags={ + **timing_tags, 'status': 'success' + }) + await self._retire_instance(instance) + + await self.lifecycle.close(instance.instance_id, lambda: self.driver.shutdown(instance), + self.config.shutdown_timeout_s, stopped, quarantined) + + async def _retire_instance(self, instance): async with self._lock: terminal_records = [] terminal_at = time.time() @@ -729,10 +700,7 @@ async def fail_instance(self, instance_id: str, failure: FutureFailureRecord) -> training.terminal_reason = failure.reason_code for request_id in (training.create_request_id, *training.pending_operation_ids): await self.store.future(request_id, 'failed', training_id, failure=failure) - startup = self._startup.get(instance_id) - if startup is not None and startup is not asyncio.current_task() and not startup.done(): - startup.cancel() - await asyncio.gather(startup, return_exceptions=True) + await self.lifecycle.cancel_start(instance_id) finally: # Once DRAINING is durable, shutdown must be attempted even if a # later training/future write fails. Sweep will retry if this fails. @@ -766,24 +734,24 @@ async def sweep(self) -> None: and not training.pending_operation_ids and training.training_id not in self._cleaning ] for training in closing: - self._spawn(self._cleanup_user(training)) + self.lifecycle.spawn(self._cleanup_user(training)) for training, request_id in uncertain: - self._spawn(self._reconcile_operation(training, request_id)) + self.lifecycle.spawn(self._reconcile_operation(training, request_id)) for instance in list(self.instances.values()): failure = self._failed_instances.get(instance.instance_id) if failure is not None: - self._spawn(self.fail_instance(instance.instance_id, failure)) + self.lifecycle.spawn(self.fail_instance(instance.instance_id, failure)) elif instance.state in ('DRAINING', 'QUARANTINED'): - self._spawn(self._close_instance(instance)) + self.lifecycle.spawn(self._close_instance(instance)) # Rebuilding one specification must not stop expiry or other models. for spec in self.catalog.specs(): - residents = sum( - instance.resident and instance.instance_id not in self._failed_instances and ( - instance.base_model, instance.mode) == spec.key and instance.state in LIVE_INSTANCES - for instance in self.instances.values()) - for _ in range(spec.config.min_replicas - residents): - async with self._lock: - instance = await self._new_instance(spec, resident=True) + async with self._lock: + matching = [item for item in self.instances.values() if (item.base_model, item.mode) == spec.key] + missing = self.lifecycle.missing( + matching, spec.config.min_replicas, spec.config.max_replicas, lambda instance: instance.resident and + instance.instance_id not in self._failed_instances and instance.state in LIVE_INSTANCES) + for _ in range(missing): + await self._new_instance(spec, resident=True) await self.store.cleanup_receipts() # Terminal diagnostics expire like their backend tombstones. Repeated # model creation must not grow controller memory for the server lifetime. @@ -797,7 +765,7 @@ async def sweep(self) -> None: if (instance.state == 'STOPPED' and instance.stopped_at is not None and instance.stopped_at <= cutoff and not instance.training_ids): self.instances.pop(identity) - self._startup.pop(identity, None) + self.lifecycle.startups.pop(identity, None) self._report_instance_metrics() async def _reconcile_operation(self, training: TrainingRecord, request_id: str) -> None: @@ -807,9 +775,7 @@ async def _reconcile_operation(self, training: TrainingRecord, request_id: str) return result = await asyncio.wait_for( self.driver.verify_operation(instance, request_id), self.config.shutdown_timeout_s) - if result['state'] == 'pending': - await self.commit_operation(training.training_id, request_id, self.epoch) - elif result['state'] == 'finished': + if result['state'] == 'finished': await self.finish_operation( training.training_id, request_id, @@ -833,10 +799,7 @@ async def _reconcile_operation(self, training: TrainingRecord, request_id: str) self._reconciling.discard(request_id) async def wait_for_jobs(self) -> None: - while self._jobs: - jobs = tuple(self._jobs) - await asyncio.gather(*jobs, return_exceptions=True) - self._jobs.difference_update(task for task in jobs if task.done()) + await self.lifecycle.join() async def shutdown(self) -> None: self._ready = False diff --git a/src/twinkle/server/training/driver.py b/src/twinkle/server/training/driver.py index 955dd83d1..e10c8e705 100644 --- a/src/twinkle/server/training/driver.py +++ b/src/twinkle/server/training/driver.py @@ -3,13 +3,10 @@ from __future__ import annotations import asyncio -import os import re import time -from copy import deepcopy -from twinkle.server.env_propagation import build_propagated_env_vars -from twinkle.server.resources.allocation import DevicePool +from twinkle.server.resources.runtime import RayRuntime from twinkle.server.state.training_records import InstanceRecord from twinkle.server.telemetry.metrics import get_training_metrics from twinkle.utils.logger import get_logger @@ -23,22 +20,20 @@ class RayInstanceDriver: def __init__(self, config, store, namespace, controller_name): self.config, self.store = config, store self.namespace, self.controller_name = namespace, controller_name - self._cleanup_tasks: set[asyncio.Task] = set() - self._deployments: dict[str, asyncio.Task] = {} self._metrics = get_training_metrics() self._catalog = TrainingCatalog(config) - self.pool = DevicePool( + self.runtime = RayRuntime( namespace, allocation_timeout=config.allocation_timeout_s, shutdown_timeout=config.shutdown_timeout_s) def application_name(self, instance): - return self.pool.name(instance.controller_epoch, instance.instance_id) + return self.runtime.pool.name(instance.controller_epoch, instance.instance_id) async def allocate(self, instance, spec): started, status = time.monotonic(), 'error' tags = {'backend': spec.model.backend, 'mode': spec.mode, 'device_type': spec.model.device_type} try: - allocation = await self.pool.allocate(spec.config.num_devices, spec.model.device_type, - instance.controller_epoch, instance.instance_id) + allocation = await self.runtime.pool.allocate(spec.config.num_devices, spec.model.device_type, + instance.controller_epoch, instance.instance_id) status = 'success' return { 'application_name': allocation.name, @@ -50,18 +45,11 @@ async def allocate(self, instance, spec): self._metrics.allocation_seconds.observe(time.monotonic() - started, tags={**tags, 'status': status}) async def deploy(self, instance, spec): - from ray import serve from ray.util.placement_group import get_placement_group from twinkle.server.model.app import build_model_app name = self.application_name(instance) - env = build_propagated_env_vars() - env.update(spec.model.env_vars) - for key in ('PYTHONPATH', 'TWINKLE_DEFAULT_SAVE_DIR', 'TMPDIR', 'OMP_NUM_THREADS'): - if key in os.environ: - env.setdefault(key, os.environ[key]) - env.update({'TWINKLE_RAY_REGISTRY_NAME': name + '-registry', 'TWINKLE_RAY_NAMESPACE': self.namespace}) context = { 'instance_id': instance.instance_id, 'epoch': instance.controller_epoch, @@ -75,14 +63,7 @@ async def deploy(self, instance, spec): if 'hybrid' in kwargs: kwargs['hybrid'] = spec.config.hybrid kwargs['train_mode'] = 'hybrid' if kwargs.get('hybrid') else spec.mode - # Each runtime has one identity. Serve autoscaling belongs to the CPU - # router; device-instance replica bounds are enforced by the controller. - options = deepcopy(spec.model.deploy_options) - options.pop('autoscaling_config', None) - options['num_replicas'] = 1 - actor_options = options.setdefault('ray_actor_options', {}) - actor_options.setdefault('num_cpus', 0.1) - actor_options.setdefault('runtime_env', {})['env_vars'] = env + options = self.runtime.deployment_options(spec.model.deploy_options, name, spec.model.env_vars) app = build_model_app( model_id=spec.model.model_id, backend=spec.model.backend, @@ -102,46 +83,11 @@ async def deploy(self, instance, spec): placement_group=get_placement_group(name), allocation_devices=list(instance.devices), **kwargs) - # Shield the actual thread so its late return can still be cleaned. - deployment = asyncio.create_task( - asyncio.to_thread( - serve.run, app, name=name, route_prefix=f'/api/v1/training-runtime/{instance.instance_id}')) - self._deployments[instance.instance_id] = deployment - try: - await asyncio.shield(deployment) - except asyncio.CancelledError: - - async def clean_late(): - try: - await deployment - except Exception: - pass - await self.shutdown(instance) - - task = asyncio.create_task(clean_late()) - self._cleanup_tasks.add(task) - - def done(completed): - self._cleanup_tasks.discard(completed) - if not completed.cancelled(): - exc = completed.exception() - if exc is not None: - logger.error('Late deployment cleanup failed', exc_info=(type(exc), exc, exc.__traceback__)) - - task.add_done_callback(done) - raise - finally: - if deployment.done(): - self._deployments.pop(instance.instance_id, None) - - async def wait_for_cleanup(self): - """Join late deployment cleanup before the controller actor is retired.""" - while self._cleanup_tasks: - await asyncio.gather(*tuple(self._cleanup_tasks), return_exceptions=True) + await self.runtime.deploy(instance.instance_id, name, app, f'/api/v1/training-runtime/{instance.instance_id}', + lambda: self.shutdown(instance)) def handle(self, instance): - from ray import serve - return serve.get_app_handle(instance.application_name or self.application_name(instance)) + return self.runtime.handle(instance.application_name or self.application_name(instance)) async def create_training(self, instance, training, request): await self.handle(instance).initialize_training.remote(training.model_dump(), request.model_dump()) @@ -154,55 +100,18 @@ async def verify_operation(self, instance, request_id): return await self.handle(instance).verify_operation.remote(request_id) async def shutdown(self, instance: InstanceRecord): - import ray - from ray import serve - from ray.util import list_named_actors - deployment = self._deployments.get(instance.instance_id) - if deployment is not None: - # Cancellation of Serve's caller does not stop its deployment thread. - # Never certify death while that thread can still create resources. - try: - await asyncio.shield(deployment) - except Exception: - pass - self._deployments.pop(instance.instance_id, None) - name = self.application_name(instance) - owned = [ - actor for actor in list_named_actors(all_namespaces=True) - if actor['namespace'] == self.namespace and instance.instance_id in actor['name'] - ] - actors = [] - for actor in owned: - try: - handle = ray.get_actor(actor['name'], namespace=self.namespace) - except ValueError: - continue - actors.append(handle) - ray.kill(handle, no_restart=True) - await asyncio.to_thread(serve.delete, instance.application_name or name) - removal_started, removal_status = time.monotonic(), 'error' + name = instance.application_name or self.application_name(instance) try: - await self.pool.remove(name) - removal_status = 'success' - finally: - try: - spec = self._catalog.get(instance.base_model, instance.mode) - except Exception: - tags = {'backend': 'unknown', 'mode': instance.mode, 'device_type': 'unknown'} - else: - tags = {'backend': spec.model.backend, 'mode': spec.mode, 'device_type': spec.model.device_type} - self._metrics.removal_seconds.observe( - time.monotonic() - removal_started, tags={ - **tags, 'status': removal_status - }) - # The PG kills even unnamed remaining children. Named workers must also - # have failed readiness RPCs before the controller retires quota/tickets. - for handle in actors: - try: - await asyncio.wait_for(handle.__ray_ready__.remote(), 5) - except ray.exceptions.RayActorError: - continue - raise RuntimeError('Owned worker death could not be confirmed') + spec = self._catalog.get(instance.base_model, instance.mode) + except Exception: + tags = {'backend': 'unknown', 'mode': instance.mode, 'device_type': 'unknown'} + else: + tags = {'backend': spec.model.backend, 'mode': spec.mode, 'device_type': spec.model.device_type} + + def observe_removal(duration, status): + self._metrics.removal_seconds.observe(duration, tags={**tags, 'status': status}) + + await self.runtime.shutdown(instance.instance_id, name, observe_removal=observe_removal) await self.store.state.unregister_replica_capacity(instance.instance_id) async def health(self, instance): @@ -210,20 +119,18 @@ async def health(self, instance): async def recover_orphans(self): """Ownership names survive loss of the memory state actor.""" - from ray import serve - from ray.util.placement_group import placement_group_table pattern = re.compile(r'resources-' + re.escape(self.namespace) + r'-(controller_[0-9a-f]{32})-(instance_[0-9a-f]{32})$') - names = set(serve.status().applications) - names.update(row.get('name', '') for row in placement_group_table().values() if row.get('state') != 'REMOVED') - for name in sorted(names): + + def decode(name): match = pattern.fullmatch(name) if match: - orphan = InstanceRecord( + return InstanceRecord( controller_epoch=match[1], instance_id=match[2], application_name=name, base_model='', mode='full', startup_deadline=0) - await asyncio.wait_for(self.shutdown(orphan), self.config.shutdown_timeout_s) + + await self.runtime.recover_orphans(decode, self.shutdown) diff --git a/src/twinkle_client/auto/agent/server_tools.py b/src/twinkle_client/auto/agent/server_tools.py index b8f11333f..40e3f8c99 100644 --- a/src/twinkle_client/auto/agent/server_tools.py +++ b/src/twinkle_client/auto/agent/server_tools.py @@ -368,6 +368,8 @@ def _generate_server_config( sampler_list = samplers or [] if backend not in ('transformers', 'megatron'): raise ValueError(f'Unsupported training backend: {backend!r}') + if backend == 'megatron' and train_gpus < 2: + raise ValueError('megatron training instances require at least 2 devices') # Sanitize model name for use in route/names def _short(mid: str) -> str: @@ -424,21 +426,28 @@ def _short(mid: str) -> str: 'env_vars': { 'TWINKLE_TRUST_REMOTE_CODE': '1' }, - 'modes': { - 'lora': { - 'num_devices': train_gpus, - 'device_mesh': { - 'device_type': 'cuda', - 'dp_size': train_gpus - }, - 'max_length': 500000, - 'idle_timeout_s': 600, - 'queue_config': { - 'rps_limit': 100, - 'tps_limit': 100000, - 'max_input_tokens': 500000 - }, - } + 'train_mode': 'lora', + 'nproc_per_node': train_gpus, + 'device_group': { + 'name': 'model', + 'ranks': train_gpus, + 'device_type': 'cuda' + }, + 'device_mesh': { + 'device_type': 'cuda', + 'dp_size': train_gpus + }, + 'min_replicas': 0, + 'max_replicas': 1, + 'max_length': 500000, + 'idle_timeout_s': 600, + 'adapter_config': { + 'adapter_timeout': 600 + }, + 'queue_config': { + 'rps_limit': 100, + 'tps_limit': 100000, + 'max_input_tokens': 500000 }, }, }) diff --git a/tests/server/config/test_server_config.py b/tests/server/config/test_server_config.py index 8e47461b0..3c7e7568e 100644 --- a/tests/server/config/test_server_config.py +++ b/tests/server/config/test_server_config.py @@ -39,7 +39,7 @@ 'model_id': st.just('model-id'), 'backend': st.sampled_from(['mock', 'transformers', 'megatron']), 'device_type': st.just('cuda'), - 'modes': st.just({'lora': {'num_devices': 1, 'device_mesh': {'dp_size': 1}}}), + 'modes': st.just({'lora': {'num_devices': 2, 'device_mesh': {'dp_size': 2}}}), }), }) diff --git a/tests/server/config/test_training_config.py b/tests/server/config/test_training_config.py index 8bb8dfc72..5c9eedc73 100644 --- a/tests/server/config/test_training_config.py +++ b/tests/server/config/test_training_config.py @@ -142,7 +142,11 @@ def test_existing_yaml_examples_keep_their_model_entries(monkeypatch): assert model.model_id == models[0].args.model_id assert model.route_prefix == models[0].route_prefix mode = model.modes.full or model.modes.lora - assert mode.min_replicas == 1 + original = models[0] + expected_min = original.args.min_replicas + if expected_min is None: + expected_min = original.deployments[0]['autoscaling_config']['min_replicas'] + assert mode.min_replicas == expected_min assert ServerConfig.model_validate(config.to_yaml_dict()).training_config == config.training_config @@ -259,7 +263,7 @@ def builder(**kwargs): assert captured['max_length'] == 128 assert captured['nproc_per_node'] == 2 assert captured['queue_config'] == model.modes.lora.queue_config - assert not driver._deployments + assert not driver.runtime.deployments def test_mutated_configuration_does_not_reuse_a_stale_catalog(): @@ -274,3 +278,44 @@ def test_unsupported_adapter_parameters_are_not_silently_ignored(): app['args']['adapter_config']['misspelled_timeout'] = 12 with pytest.raises(ValidationError, match='misspelled_timeout'): ServerConfig(applications=[app]) + + +def test_l20d_candidate_is_lora_only_and_bounds_total_devices(): + config = ServerConfig.from_yaml('cookbook/client/server/megatron/server_config.yaml') + model = config.training_config.models[0] + mode = model.modes.lora + assert model.modes.full is None + assert (mode.num_devices, mode.min_replicas, mode.max_replicas, mode.max_loras) == (2, 0, 2, 3) + assert mode.device_mesh.world_size() == 2 and mode.device_mesh.pp_size == 2 and mode.idle_timeout_s == 600 + sampler = next(app for app in config.applications if app.import_path == 'sampler') + assert sampler.args.device_group['ranks'] == 1 + assert sampler.args.device_mesh['dp_size'] == 1 and 'tp_size' not in sampler.args.device_mesh + assert 'tensor_parallel_size' not in sampler.args.engine_args + assert (sampler.args.min_replicas, sampler.args.max_replicas, sampler.args.target_ongoing_requests) == (2, 4, 16) + assert mode.num_devices * mode.max_replicas + sampler.args.device_group['ranks'] * sampler.args.max_replicas == 8 + assert mode.max_loras * mode.max_replicas == 6 + assert config.applications[-1].args.modes is None + with pytest.raises(RequestRejectedError, match='does not provide full'): + TrainingCatalog(config.training_config).get(model.name, 'full') + + +@pytest.mark.parametrize('mode', ['lora', 'full']) +@pytest.mark.parametrize('legacy', [False, True]) +def test_single_device_megatron_is_rejected_before_allocation(mode, legacy): + if legacy: + app = legacy_application() + app['args'].update(backend='megatron', train_mode=mode, device_group={'ranks': 1, 'device_type': 'cuda'}, + device_mesh={'dp_size': 1}) + else: + model = model_config() + model['backend'] = 'megatron' + model['modes'] = {mode: {'num_devices': 1, 'device_mesh': {'dp_size': 1}}} + app = model_application(model) + with pytest.raises(ValidationError, match='at least 2 devices'): + ServerConfig(applications=[app]) + + +def test_transformers_still_accepts_single_device(): + model = model_config() + model['modes'] = {'lora': {'num_devices': 1, 'device_mesh': {'dp_size': 1}}} + assert ServerConfig(applications=[model_application(model)]).training_config.models[0].modes.lora.num_devices == 1 diff --git a/tests/server/integration/on_demand_training/acceptance_support.py b/tests/server/integration/on_demand_training/acceptance_support.py index b24b5d9f0..f3ae8f5e7 100644 --- a/tests/server/integration/on_demand_training/acceptance_support.py +++ b/tests/server/integration/on_demand_training/acceptance_support.py @@ -146,7 +146,8 @@ def assert_released(controller, namespace, redis_url=None, prefix=None): from ray.util.placement_group import placement_group_table diagnostic = ray.get(controller.diagnostics.remote(), timeout=30) assert all(row['state'] == 'STOPPED' for row in diagnostic['instances']), diagnostic - owned_prefix = 'resources-' + namespace + '-' + # Training teardown must not require a healthy sampling pool to disappear. + owned_prefix = 'resources-' + namespace + '-controller_' assert not [name for name in serve.status().applications if name.startswith(owned_prefix)] assert not [row for row in placement_group_table().values() if row.get('name', '').startswith(owned_prefix) and row.get('state') != 'REMOVED'] diff --git a/tests/server/integration/on_demand_training/matrix_validation.py b/tests/server/integration/on_demand_training/matrix_validation.py index 35a30ee6d..9112a3dff 100644 --- a/tests/server/integration/on_demand_training/matrix_validation.py +++ b/tests/server/integration/on_demand_training/matrix_validation.py @@ -211,6 +211,8 @@ def main(): parser = argparse.ArgumentParser() parser.add_argument('--backend', required=True, choices=['transformers', 'megatron']) parser.add_argument('--model', required=True) + parser.add_argument('--training-devices', type=int, default=2) + parser.add_argument('--training-pp-size', type=int, default=1) parser.add_argument('--run-dir', required=True) parser.add_argument('--persistence', choices=['memory', 'redis'], default='memory') parser.add_argument('--managed-redis', action='store_true', help='Start an owned Redis for restart injection') @@ -231,6 +233,12 @@ def main(): if any(not math.isfinite(value) or value <= 0 for value in (args.startup_timeout_s, args.shutdown_timeout_s)): parser.error('startup and shutdown timeouts must be finite and positive') + if args.training_devices < 1: + parser.error('--training-devices must be positive') + if args.backend == 'megatron' and args.training_devices < 2: + parser.error('megatron training instances require at least 2 devices') + if args.training_pp_size < 1 or args.training_devices % args.training_pp_size: + parser.error('--training-pp-size must be positive and divide --training-devices') if args.recovery_rounds < 0 or args.stress_iterations < 0 or args.soak_seconds < 0: parser.error('validation budgets must be non-negative') root = Path(__file__).resolve().parents[4] @@ -267,7 +275,9 @@ def main(): params = app['args'] if app['import_path'] == 'model': common = { - 'num_devices': 2, 'device_mesh': {'device_type': 'cuda', 'dp_size': 2}, + 'num_devices': args.training_devices, + 'device_mesh': {'device_type': 'cuda', 'dp_size': args.training_devices // args.training_pp_size, + 'pp_size': args.training_pp_size}, 'min_replicas': 0, 'max_replicas': 1, 'idle_timeout_s': 3600, 'max_length': params.get('max_length'), 'queue_config': {**params.get('queue_config', {}), 'execution_timeout': 600}, @@ -320,7 +330,7 @@ def main(): address='local', namespace=namespace, num_cpus=24, - num_gpus=2 if args.recovery_only else 3, + num_gpus=args.training_devices + (0 if args.recovery_only else 1), _temp_dir='/tmp/twm-' + namespace[-8:], include_dashboard=False) launcher = ServerLauncher(config) @@ -373,6 +383,10 @@ def main(): evidence['recovery'] = [] for _ in range(args.recovery_rounds): recovery = validate_controller_restart(base_url, launcher._training_controller, model_name) + for sampler in launcher._sampler_controllers: + assert ray.get(sampler.diagnostics.remote(), timeout=30)['ready'], 'Training restart disturbed sampler' + if launcher._sampler_controllers: + recovery['sampler_remained_ready'] = True evidence['recovery'].append(recovery) evidence['restart_epoch'] = recovery['epoch'] evidence['worker_fault_recovered'] = recovery['worker_fault_recovered'] diff --git a/tests/server/integration/on_demand_training/sampling_probe.py b/tests/server/integration/on_demand_training/sampling_probe.py new file mode 100644 index 000000000..345cce133 --- /dev/null +++ b/tests/server/integration/on_demand_training/sampling_probe.py @@ -0,0 +1,178 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Real Ray/Serve 2->4->2 allocation and death confirmation, without model kernels.""" +import asyncio +import json +import os +import uuid + +import ray +from ray import serve +from ray.util.placement_group import get_placement_group, placement_group_table +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + +from twinkle.server.config.application_spec import SamplerArgs +from twinkle.server.deployment import bind_deployment, build_deployment_app +from twinkle.server.resources.allocation import CapacityUnavailableError, DevicePool +from twinkle.server.sampler.controller import SamplerController +from twinkle.server.sampler.driver import SamplerDriver + + +@ray.remote(num_cpus=0.01, num_gpus=1) +class Worker: + + def device(self): + return ray.get_runtime_context().get_accelerator_ids()['GPU'] + + +class Compute: + + async def __init__(self, namespace, identity, pg): + self.operations = {} + self.worker = Worker.options( + name=identity + '-worker', + namespace=namespace, + scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=pg)).remote() + self.devices = await self.worker.device.remote() + + async def invoke(self, ticket, payload, stream_queue=None): + self.operations[ticket] = asyncio.Event() + return {'status': 200, 'headers': [], 'body': b'{"status":"pending"}'} + + async def wait_operation(self, ticket): + await self.operations[ticket].wait() + self.operations.pop(ticket) + + async def finish(self): + for operation in self.operations.values(): + operation.set() + + async def pending_futures(self): + return [] + + async def runtime_identity(self): + return {'unhealthy': False, 'devices': self.devices} + + +class Driver(SamplerDriver): + + async def start(self, instance): + allocation = await self.runtime.pool.allocate(1, 'cuda', f'sampling_{self.pool_id}_{instance.epoch}', + instance.identity) + instance.application, instance.devices = allocation.name, list(allocation.devices) + app = bind_deployment( + build_deployment_app('AllocationProbe', lambda app, get_self: None), + Compute, + self.runtime.deployment_options({}, allocation.name), + deployment_name='AllocationCompute', + bind_args=(self.namespace, instance.identity, get_placement_group(allocation.name))) + await self.runtime.deploy(instance.identity, allocation.name, app, '/probe/' + instance.identity, + lambda: self.shutdown(instance)) + + +async def run(namespace): + args = SamplerArgs( + model_id='allocation-probe', + sampler_type='mock', + device_group={ + 'ranks': 1, + 'device_type': 'cuda' + }, + device_mesh={'dp_size': 1}, + min_replicas=2, + max_replicas=4, + target_ongoing_requests=1, + idle_timeout_s=1) + budgets = {'allocation_timeout_s': 5, 'startup_timeout_s': 600, 'shutdown_timeout_s': 180, 'sweep_interval_s': 1} + driver = Driver(args, namespace, 'allocation', {}, budgets) + clock = [0] + pool = SamplerController(args, driver, budgets, clock=lambda: clock[0]) + training_pool = DevicePool(namespace, allocation_timeout=5, shutdown_timeout=180) + training = [] + evidence = {'sizes': [], 'device_ids': []} + try: + # Two training-sized allocations plus four sampling instances share exactly eight GPUs. + for index in range(2): + training.append(await training_pool.allocate(2, 'cuda', 'probe_training', str(index))) + await pool.initialize() + evidence['sizes'].append(pool._owned_count()) + payload = {'endpoint': 'twinkle/sample', 'method': 'POST', 'headers': [], 'body': b'{}'} + for _ in range(8): + await pool.request(payload) + await asyncio.gather(*tuple(pool.lifecycle.startups.values())) + evidence['sizes'].append(pool._owned_count()) + all_devices = [device for allocation in training for device in allocation.devices] + for instance in pool._ready_instances(): + identity = await driver.health(instance) + assert identity['devices'] == instance.devices + all_devices.extend(instance.devices) + assert len(all_devices) == len(set(all_devices)) == 8 + evidence['device_ids'] = sorted(all_devices) + try: + await training_pool.allocate(1, 'cuda', 'probe_training', 'overflow') + except CapacityUnavailableError: + evidence['no_overcommit'] = True + else: + raise AssertionError('A ninth device allocation unexpectedly succeeded') + clock[0] = 10 + await pool.sweep() + assert pool._owned_count() == 4 # Public replies did not retire actual operations. + await asyncio.gather(*(driver.handle(instance).finish.remote() for instance in pool._ready_instances())) + await pool.lifecycle.join() + clock[0] = 12 + await pool.sweep() + await pool.lifecycle.join() + evidence['sizes'].append(len(pool._ready_instances())) + assert evidence['sizes'] == [2, 4, 2] + old_names = {instance.application for instance in pool._ready_instances()} + recovered_driver = Driver(args, namespace, 'allocation', {}, budgets) + recovered = SamplerController(args, recovered_driver, budgets, clock=lambda: clock[0]) + try: + await recovered.initialize() + assert recovered.epoch != pool.epoch and len(recovered._ready_instances()) == 2 + assert not old_names.intersection(serve.status().applications) + assert all( + placement_group_table(get_placement_group(allocation.name))['state'] != 'REMOVED' + for allocation in training) + evidence['owned_orphan_recovery'] = True + finally: + await recovered.shutdown() + await recovered_driver.runtime.wait_for_cleanup() + await pool.shutdown() + await driver.runtime.wait_for_cleanup() + names = set(serve.status().applications) + assert not any(name.startswith(driver.prefix) for name in names) + assert not any(row['state'] != 'REMOVED' and row.get('name', '').startswith(driver.prefix) + for row in placement_group_table().values()) + from ray.util import list_named_actors + assert not any(row['namespace'] == namespace and 'sampler_' in row['name'] + for row in list_named_actors(all_namespaces=True)) + evidence['confirmed_release'] = True + return evidence + finally: + await pool.shutdown() + await driver.runtime.wait_for_cleanup() + for allocation in training: + await training_pool.remove(allocation.name) + + +def main(): + namespace = 'sampling_probe_' + uuid.uuid4().hex[:8] + os.environ.update( + TWINKLE_STATE_RAY_NAMESPACE=namespace, TWINKLE_PERSISTENCE_MODE='memory', TWINKLE_GPU_PG_CPUS_PER_PROC='1') + ray.init( + address='local', + namespace=namespace, + num_cpus=24, + num_gpus=8, + _temp_dir='/tmp/tws-' + namespace[-8:], + include_dashboard=False) + try: + serve.start(http_options={'host': '127.0.0.1', 'port': 0}) + print(json.dumps(asyncio.run(run(namespace))), flush=True) + finally: + serve.shutdown() + ray.shutdown() + + +if __name__ == '__main__': + main() diff --git a/tests/server/sampler/test_pool.py b/tests/server/sampler/test_pool.py new file mode 100644 index 000000000..ba47028a1 --- /dev/null +++ b/tests/server/sampler/test_pool.py @@ -0,0 +1,337 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Pool decisions against a controllable driver, without Ray or GPU side effects.""" +import asyncio +import json +from types import SimpleNamespace + +import pytest + +from twinkle.server.config.application_spec import SamplerArgs +from twinkle.server.sampler.controller import SamplerController + + +@pytest.fixture(scope='session', autouse=True) +def _ray_runtime(): + yield + + +class Driver: + + def __init__(self): + self.started = [] + self.stopped = [] + self.calls = [] + self.finish = asyncio.Event() + self.start_gate = asyncio.Event() + self.start_gate.set() + self.fail_start = False + self.fail_close = False + + async def start(self, instance): + self.started.append(instance.identity) + await self.start_gate.wait() + if self.fail_start: + raise RuntimeError('allocation failed') + instance.devices = [str(len(self.started))] + + async def invoke(self, instance, ticket, payload, queue=None): + self.calls.append((instance.identity, payload)) + return {'status': 200, 'headers': [], 'body': b'{"status":"pending"}'} + + async def wait_operation(self, instance, ticket): + await self.finish.wait() + + async def load_weights(self, instance, path): + self.calls.append((instance.identity, {'base_path': path})) + + async def shutdown(self, instance): + if self.fail_close: + raise RuntimeError('death unconfirmed') + self.stopped.append(instance.identity) + + async def recover_orphans(self): + pass + + +def setup(minimum=2, maximum=4, target=1): + args = SamplerArgs( + model_id='mock', + sampler_type='mock', + device_group={ + 'name': 'sampler', + 'ranks': 1, + 'device_type': 'cpu' + }, + device_mesh={'dp_size': 1}, + min_replicas=minimum, + max_replicas=maximum, + target_ongoing_requests=target, + idle_timeout_s=10) + driver = Driver() + clock = SimpleNamespace(now=0) + pool = SamplerController( + args, driver, { + 'startup_timeout_s': 2, + 'shutdown_timeout_s': 1, + 'sweep_interval_s': 1 + }, clock=lambda: clock.now) + return pool, driver, clock + + +def request(endpoint='twinkle/sample', body=None): + return { + 'endpoint': endpoint, + 'body': json.dumps(body or {}).encode(), + 'method': 'POST', + 'headers': [(b'x-request-id', b'test')], + 'query_string': b'' + } + + +async def settle_starts(pool): + await asyncio.gather(*tuple(pool.lifecycle.startups.values()), return_exceptions=True) + + +async def settle_jobs(pool): + await asyncio.gather(*tuple(pool.lifecycle.jobs), return_exceptions=True) + + +@pytest.mark.asyncio +async def test_minimum_expand_and_reclaim_only_idle_extras(): + pool, driver, clock = setup() + await pool.initialize() + assert len(driver.started) == 2 + try: + for _ in range(8): + await pool.request(request()) + await settle_starts(pool) + assert len(driver.started) == 4 and pool._owned_count() == 4 + live = pool._ready_instances() + assert all(instance.requests for instance in live) + assert max(len(instance.requests) for instance in live) - min(len(instance.requests) for instance in live) <= 1 + clock.now = 30 + await pool.sweep() + assert not driver.stopped + driver.finish.set() + await settle_jobs(pool) + clock.now += 11 + await pool.sweep() + await settle_jobs(pool) + assert len(driver.stopped) == 2 and len(pool._ready_instances()) == 2 + finally: + driver.finish.set() + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_waiter_cancellation_preserves_generation_and_defers_control(): + pool, driver, _ = setup() + await pool.initialize() + await pool.request(request()) + control = asyncio.create_task(pool.request(request('twinkle/set_template', {'template_cls': 'Template'}))) + try: + await asyncio.sleep(0) + control.cancel() + await asyncio.gather(control, return_exceptions=True) + assert pool._guard._readers == 1 + assert sum(len(instance.requests) for instance in pool.instances.values()) == 1 + assert len(driver.calls) == 1 + driver.finish.set() + await settle_jobs(pool) + assert any(payload['endpoint'] == 'twinkle/set_template' for _, payload in driver.calls) + assert pool._actions[0]['endpoint'] == 'twinkle/set_template' + finally: + driver.finish.set() + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_new_instance_replays_controls_not_lora_weight_cache(): + pool, driver, _ = setup(minimum=1, target=1) + driver.finish.set() + await pool.initialize() + try: + await pool.request(request('twinkle/set_template', {'template_cls': 'Template'})) + driver.finish.clear() + await pool.request(request(body={'adapter_uri': 'checkpoint-A'})) + await pool.request(request(body={'adapter_uri': 'checkpoint-B'})) + driver.finish.set() + await settle_starts(pool) + await settle_jobs(pool) + assert len(driver.started) == 2 + new = driver.started[-1] + replay = [payload for identity, payload in driver.calls if identity == new] + assert replay and all(payload['endpoint'] == 'twinkle/set_template' for payload in replay) + assert all('adapter_uri' not in json.loads(payload['body']) for payload in pool._actions) + finally: + driver.finish.set() + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_quarantine_counts_against_maximum_until_confirmed_death(): + pool, driver, _ = setup(minimum=1, maximum=1) + await pool.initialize() + instance = pool._ready_instances()[0] + driver.fail_close = True + await pool._close(instance) + try: + assert instance.state == 'QUARANTINED' + for _ in range(3): + await pool.sweep() + await settle_jobs(pool) + assert len(driver.started) == 1 and pool._owned_count() == 1 + driver.fail_close = False + await pool.sweep() + await settle_jobs(pool) + await pool.sweep() + await settle_starts(pool) + assert len(driver.started) == 2 and len(pool._ready_instances()) == 1 + finally: + driver.finish.set() + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_startup_failure_rolls_back_loaded_minimum(): + pool, driver, _ = setup() + original = driver.start + + async def fail_second(instance): + if len(driver.started) == 1: + driver.fail_start = True + await original(instance) + + driver.start = fail_second + with pytest.raises(RuntimeError, match='allocation'): + await pool.initialize() + assert not pool._ready and pool._owned_count() == 0 + + +@pytest.mark.asyncio +async def test_lazy_first_requests_share_one_startup(): + pool, driver, _ = setup(minimum=0, maximum=1) + await pool.initialize() + driver.start_gate.clear() + calls = [asyncio.create_task(pool.request(request())) for _ in range(8)] + try: + for _ in range(10): + await asyncio.sleep(0) + assert len(driver.started) == 1 + driver.start_gate.set() + await asyncio.wait_for(asyncio.gather(*calls), 2) + assert len(driver.started) == 1 + finally: + driver.finish.set() + await pool.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('fault', ['raise', 'http_error']) +async def test_partial_control_broadcast_fences_all_participants(fault): + pool, driver, _ = setup() + driver.finish.set() + await pool.initialize() + invoke = driver.invoke + changed = [] + + async def update(instance, ticket, payload, queue=None): + changed.append(instance.identity) + if len(changed) == 2: + if fault == 'raise': + raise ConnectionError('update ACK unknown') + return {'status': 500, 'headers': [], 'body': b'error'} + return await invoke(instance, ticket, payload, queue) + + driver.invoke = update + try: + if fault == 'raise': + with pytest.raises(ConnectionError): + await pool.request(request('twinkle/set_template')) + else: + reply = await pool.request(request('twinkle/set_template')) + assert reply['status'] == 500 + await settle_jobs(pool) + assert len(driver.stopped) == 2 + assert not pool._ready_instances() and not pool._actions + finally: + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_unknown_generation_exit_keeps_quarantined_capacity_reserved(): + pool, driver, clock = setup(minimum=1, maximum=1) + await pool.initialize() + driver.fail_close = True + + async def unknown(*args): + raise TimeoutError('generation completion unknown') + + driver.wait_operation = unknown + await pool.request(request()) + await settle_jobs(pool) + instance = next(iter(pool.instances.values())) + assert instance.state == 'QUARANTINED' and instance.requests + clock.now = 100 + await pool.sweep() + await settle_jobs(pool) + assert len(driver.started) == 1 and instance.requests + driver.fail_close = False + await pool.sweep() + await settle_jobs(pool) + assert instance.state == 'STOPPED' and not instance.requests + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_training_and_sampling_use_the_same_lifecycle_owner(): + from tests.server.state.fake_backend import FakeBackend + from tests.server.training.test_controller import setup as training_setup + pool, driver, _ = setup() + training, _, _, _ = await training_setup(backend=FakeBackend()) + try: + assert type(training.lifecycle) is type(pool.lifecycle) + assert training.lifecycle is not pool.lifecycle + finally: + driver.finish.set() + await pool.shutdown() + await training.shutdown() + + +@pytest.mark.asyncio +async def test_replaying_config_has_the_same_startup_deadline(): + pool, driver, _ = setup(minimum=0, maximum=1) + await pool.initialize() + pool._actions.append(request('twinkle/set_template')) + pool.budgets['startup_timeout_s'] = 0.05 + instance = pool._new_instance() + with pytest.raises(TimeoutError): + await pool.lifecycle.startups[instance.identity] + assert instance.state == 'STOPPED' and driver.stopped == [instance.identity] + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_shared_health_rule_tolerates_transient_probe_timeout(): + from twinkle.server.resources.instances import InstanceLifecycle + clock = SimpleNamespace(now=0) + lifecycle = InstanceLifecycle('health', clock=lambda: clock.now) + + async def unreachable(): + raise TimeoutError('control-plane probe timed out') + + async def healthy(): + return {'unhealthy': False} + + async def failed(): + return {'unhealthy': True} + + assert await lifecycle.health('instance', unreachable) is None + clock.now = 29 + assert await lifecycle.health('instance', unreachable) is None + clock.now = 30 + assert await lifecycle.health('instance', unreachable) is False + assert await lifecycle.health('instance', healthy) is True + assert not lifecycle.unreachable_since + assert await lifecycle.health('instance', unreachable) is None + assert await lifecycle.health('instance', failed) is False diff --git a/tests/server/sampler/test_request_runtime.py b/tests/server/sampler/test_request_runtime.py new file mode 100644 index 000000000..ef7c5809c --- /dev/null +++ b/tests/server/sampler/test_request_runtime.py @@ -0,0 +1,208 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Actual completion handles outlive future timeout and streaming iterator closure.""" +import ast +import asyncio +import threading +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from tests.server.utils.test_task_queue_mixin import _DummyQueue +from twinkle.server.config.task_queue import TaskQueueConfig +from twinkle.server.sampler.concurrency import SamplerConcurrency +from twinkle.server.task_queue.types import ExecutionScope, execution_scope + + +@pytest.mark.asyncio +async def test_worker_exports_actual_completion_handle_without_private_task_lookup(): + queue = _DummyQueue() + queue._task_queue_config = TaskQueueConfig(execution_timeout=0.05) + queue.enable_compute_worker() + guard = SamplerConcurrency() + entered, exit_gate = asyncio.Event(), asyncio.Event() + scope = ExecutionScope() + + async def backend(): + entered.set() + await exit_gate.wait() + return 'late result' + + token = execution_scope.set(scope) + try: + ref = await queue.schedule_concurrent_task(lambda: guard.run(backend)) + finally: + execution_scope.reset(token) + try: + await asyncio.wait_for(entered.wait(), 1) + completion = asyncio.create_task(scope.wait()) + for _ in range(100): + record = await queue.state.get_future(ref['request_id']) + if record['status'] == 'failed': + break + await asyncio.sleep(0.01) + assert record['status'] == 'failed' + assert not completion.done() and guard._readers == 1 + assert scope.request_ids == {ref['request_id']} and len(scope.jobs) == 1 + exit_gate.set() + await asyncio.wait_for(completion, 1) + assert guard._readers == 0 + finally: + exit_gate.set() + await queue._compute_worker.stop() + + +def test_stream_iterator_close_joins_the_real_engine_loop_task(): + # Execute the repository method without constructing a GPU-backed vLLM engine. + path = Path('src/twinkle/sampler/vllm_sampler/vllm_sampler.py') + tree = ast.parse(path.read_text()) + method = next(node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == '_iter_in_loop') + namespace = {'asyncio': asyncio} + exec(compile(ast.Module(body=[method], type_ignores=[]), str(path), 'exec'), namespace) + loop = asyncio.new_event_loop() + running, finalizing, release, closed = [threading.Event() for _ in range(4)] + + def run_loop(): + asyncio.set_event_loop(loop) + running.set() + loop.run_forever() + + thread = threading.Thread(target=run_loop) + thread.start() + running.wait(1) + + async def generation(): + try: + yield 'first' + finally: + finalizing.set() + await asyncio.to_thread(release.wait) + + iterator = namespace['_iter_in_loop'](SimpleNamespace(_async_loop=loop), generation()) + closer = None + try: + assert next(iterator) == 'first' + assert finalizing.wait(1) + + def close(): + iterator.close() + closed.set() + + closer = threading.Thread(target=close) + closer.start() + assert not closed.wait(0.05) + release.set() + assert closed.wait(2) + finally: + release.set() + if closer: + closer.join(2) + loop.call_soon_threadsafe(loop.stop) + thread.join(2) + loop.close() + + +@pytest.mark.asyncio +async def test_runtime_does_not_certify_an_unconfirmed_cancel_ack(): + from twinkle.server.sampler.runtime import SamplerRuntime + from twinkle.server.task_queue.types import BackendBusyError + guard = SamplerConcurrency() + runtime = SamplerRuntime(SimpleNamespace(sampling_guard=guard)) + task = asyncio.create_task(asyncio.sleep(0)) + await task + runtime._operations['cancelled'] = {'task': task, 'scope': ExecutionScope()} + guard.invalidate_weights(RuntimeError('cancel ACK preceded actual generation exit')) + with pytest.raises(BackendBusyError): + await runtime.wait_operation('cancelled') + assert 'cancelled' in runtime._operations + + +@pytest.mark.asyncio +async def test_cold_lora_config_does_not_block_sampler_health(monkeypatch): + import json + import sys + from twinkle.server.sampler.runtime import SamplerRuntime + entered, release = threading.Event(), threading.Event() + calls = [] + + class ColdConfig: + + def __init__(self, **kwargs): + entered.set() + assert release.wait(2) + + monkeypatch.setitem(sys.modules, 'peft', SimpleNamespace(LoraConfig=ColdConfig)) + + async def call_backend(fn, *args, **kwargs): + return fn(*args, **kwargs) + + def add_adapter(name, config): + calls.append((name, config)) + + owner = SimpleNamespace( + model_id='mock', + sampling_guard=SamplerConcurrency(), + sampler=SimpleNamespace(add_adapter_to_sampler=add_adapter), + call_backend=call_backend) + runtime = SamplerRuntime(owner) + payload = { + 'endpoint': 'twinkle/add_adapter_to_sampler', + 'method': 'POST', + 'headers': [(b'content-type', b'application/json'), (b'x-request-id', b'configuration')], + 'body': json.dumps({ + 'adapter_name': 'a', + 'config': { + 'r': 4 + } + }).encode() + } + invocation = asyncio.create_task(runtime.invoke('configuration', payload)) + try: + began = await asyncio.to_thread(entered.wait, 5) + assert began, invocation.result() if invocation.done() else 'Configuration route did not start' + assert (await asyncio.wait_for(runtime.runtime_identity(), 0.1))['unhealthy'] is False + assert not invocation.done() and not calls + release.set() + assert (await invocation)['status'] == 200 + await runtime.wait_operation('configuration') + assert len(calls) == 1 + finally: + release.set() + await asyncio.gather(invocation, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_managed_compute_http_cannot_bypass_pool_admission(monkeypatch): + import httpx + from fastapi import FastAPI + from twinkle.server.exceptions import RequestRejectedError + from twinkle.server.sampler import app as module + + def scaffold(title, register_routes, **kwargs): + app = FastAPI() + register_routes(app, lambda: SimpleNamespace()) + return app + + monkeypatch.setattr(module, 'build_deployment_app', scaffold) + monkeypatch.setattr(module, 'bind_deployment', lambda app, *args, **kwargs: app) + args = { + 'model_id': 'mock', + 'nproc_per_node': 1, + 'device_group': { + 'name': 'sampler', + 'ranks': 1, + 'device_type': 'cpu' + }, + 'device_mesh': { + 'dp_size': 1 + }, + 'deploy_options': {}, + 'sampler_type': 'mock' + } + managed = module.build_sampler_app(**args, sampling_context={'instance_id': 'owned'}) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=managed), base_url='http://test') as client: + with pytest.raises(RequestRejectedError, match='stable sampler endpoint'): + await client.post('/twinkle/create', json={}) + standalone = module.build_sampler_app(**args) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=standalone), base_url='http://test') as client: + assert (await client.post('/twinkle/create', json={})).status_code == 200 diff --git a/tests/server/sampler/test_stream_guarantees.py b/tests/server/sampler/test_stream_guarantees.py index 60623c970..a98d3d36f 100644 --- a/tests/server/sampler/test_stream_guarantees.py +++ b/tests/server/sampler/test_stream_guarantees.py @@ -89,6 +89,8 @@ async def test_stream_timeout_returns_error_payload_and_closes_queue(): class _GenerationService: def __init__(self) -> None: + from twinkle.server.sampler.concurrency import SamplerConcurrency + self.sampling_guard = SamplerConcurrency() self.cancelled = False self.sampler = SimpleNamespace( get_generation_status=lambda _submission_id: {'status': 'running'}, @@ -111,3 +113,6 @@ async def test_generation_poll_has_total_timeout_and_cancels(): await asyncio.wait_for(_await_generation(service, 'submission'), timeout=0.05) assert service.cancelled is True + from twinkle.server.task_queue.types import BackendBusyError + with pytest.raises(BackendBusyError, match='restarting'): + service.sampling_guard.check_available() diff --git a/tests/server/static/backend_call_exemptions.py b/tests/server/static/backend_call_exemptions.py index 1d0307ffe..881862dd7 100644 --- a/tests/server/static/backend_call_exemptions.py +++ b/tests/server/static/backend_call_exemptions.py @@ -15,9 +15,9 @@ - the ray ``Queue.get`` inside ``sample_stream``'s ``_stream_queue``: it bridges the sampler actor's process boundary and is bounded by a dedicated double-timeout, not by ``call_backend``; -- the ``.sample_stream_to_queue.remote(...)`` call inside ``sample_stream`` - itself: streaming generation must keep producing while the HTTP response streams, - so it cannot use ``call_backend`` as-is and carries its own double timeout. This is +- the ``.sample_stream_to_queue.remote(...)`` call inside ``_stream_generation``: + the HTTP queue has its own deadline, while the actor completion is retained by + ExecutionScope and the weight guard until real engine-loop generation exits. This is a ``remote_function`` bypass that the guard now *detects* (via backend-derived local tracking) and that is *explicitly* accepted here — replacing the previous "No remote_function call is exempt" claim, which was true only because the guard @@ -28,5 +28,5 @@ # (module_relpath under src/twinkle/server, innermost enclosing function name) BACKEND_CALL_EXEMPTIONS: frozenset[tuple[str, str]] = frozenset({ ('sampler/twinkle_handlers.py', '_stream_queue'), - ('sampler/twinkle_handlers.py', 'sample_stream'), + ('sampler/twinkle_handlers.py', '_stream_generation'), }) diff --git a/tests/server/test_dependencies.py b/tests/server/test_dependencies.py index 8853f9031..5150294bf 100644 --- a/tests/server/test_dependencies.py +++ b/tests/server/test_dependencies.py @@ -13,9 +13,10 @@ # the same layer are independent and therefore may not import one another. _LAYERS = ( frozenset({'exceptions'}), - frozenset({'config', 'middleware', 'resources', 'runtime', 'task_errors', 'utils', 'validation'}), + frozenset({'config', 'middleware', 'runtime', 'task_errors', 'utils', 'validation'}), frozenset({'telemetry'}), frozenset({'env_propagation', 'service_proxy', 'state'}), + frozenset({'resources'}), frozenset({'deployment', 'lifecycle', 'session_resource'}), frozenset({'task_queue'}), frozenset({'checkpoint', 'data_plane'}), diff --git a/tests/server/training/test_acceptance_harness.py b/tests/server/training/test_acceptance_harness.py index 92b8b972c..cdbb11d27 100644 --- a/tests/server/training/test_acceptance_harness.py +++ b/tests/server/training/test_acceptance_harness.py @@ -55,7 +55,9 @@ def clock(): def test_failed_training_never_reports_success(monkeypatch): + class Broken(FakeHandle): + def step(self): raise RuntimeError('injected failure') @@ -82,3 +84,29 @@ def unavailable(): finally: client.close() managed.stop() + + +def test_training_release_assertion_excludes_but_does_not_remove_sampling(monkeypatch): + import importlib + import ray + from ray import serve + from types import SimpleNamespace + from tests.server.integration.on_demand_training.acceptance_support import assert_released + namespace = 'owned' + sampling_name = 'resources-owned-sampling_pool_epoch-sampler_id' + applications = {sampling_name: object()} + groups = {'sampler': {'name': sampling_name, 'state': 'CREATED'}} + controller = SimpleNamespace(diagnostics=SimpleNamespace(remote=lambda: {'instances': [{'state': 'STOPPED'}]})) + monkeypatch.setattr(ray, 'get', lambda value, **kwargs: value) + monkeypatch.setattr(serve, 'status', lambda: SimpleNamespace(applications=applications)) + monkeypatch.setattr(importlib.import_module('ray.util.placement_group'), 'placement_group_table', lambda: groups) + assert assert_released(controller, namespace) + assert sampling_name in applications and groups['sampler']['state'] == 'CREATED' + training_name = 'resources-owned-controller_epoch-instance_id' + applications[training_name] = object() + with pytest.raises(AssertionError): + assert_released(controller, namespace) + applications.pop(training_name) + groups['training'] = {'name': training_name, 'state': 'CREATED'} + with pytest.raises(AssertionError): + assert_released(controller, namespace) diff --git a/tests/server/training/test_actor_initialization.py b/tests/server/training/test_actor_initialization.py index 10c748341..0c7c97eeb 100644 --- a/tests/server/training/test_actor_initialization.py +++ b/tests/server/training/test_actor_initialization.py @@ -6,6 +6,7 @@ import pytest +from twinkle.server.resources.instances import InstanceLifecycle from twinkle.server.training.actor import TrainingControllerActor @@ -49,13 +50,13 @@ async def test_shutdown_cancels_pending_restart(): restart = asyncio.create_task(wait.wait()) initialization = asyncio.get_running_loop().create_future() initialization.set_result(None) - actor = SimpleNamespace(driver=SimpleNamespace(wait_for_cleanup=AsyncMock()), + actor = SimpleNamespace(driver=SimpleNamespace(runtime=SimpleNamespace(wait_for_cleanup=AsyncMock())), core=SimpleNamespace(shutdown=AsyncMock(), wait_for_jobs=AsyncMock()), _monitors=[], _restart_task=restart, _initialization=initialization, _failures={}) await cls.shutdown(actor) assert restart.cancelled() actor.core.shutdown.assert_awaited_once() - actor.driver.wait_for_cleanup.assert_awaited_once() + actor.driver.runtime.wait_for_cleanup.assert_awaited_once() @pytest.mark.asyncio @@ -81,7 +82,8 @@ async def touch(identity): actor = SimpleNamespace( store=SimpleNamespace(backend=SimpleNamespace(get=AsyncMock(return_value='active')), key=lambda *args: 'active', state=SimpleNamespace(touch_replica_last_seen=touch)), - core=SimpleNamespace(epoch='controller', instances={'bad': bad, 'good': good}, fail_instance=fail), + core=SimpleNamespace(epoch='controller', instances={'bad': bad, 'good': good}, fail_instance=fail, + lifecycle=InstanceLifecycle('Training')), driver=SimpleNamespace(health=health), _unreachable_since={}, _failures={}) actor._check_instance_health = lambda instance: cls._check_instance_health(actor, instance) actor._fail_unhealthy = lambda instance, message: cls._fail_unhealthy(actor, instance, message) @@ -96,13 +98,13 @@ async def touch(identity): await cls._check_instance_health(actor, good) assert healthy_seen.is_set() cls._fail_unhealthy(actor, bad, 'duplicate') - assert len(actor._failures) == 1 + assert len(actor.core.lifecycle.named_jobs) == 1 finally: probe_gate.set() close_gate.set() monitor.cancel() - await asyncio.gather(monitor, *actor._failures.values(), return_exceptions=True) - assert not actor._failures + await asyncio.gather(monitor, *actor.core.lifecycle.named_jobs.values(), return_exceptions=True) + assert not actor.core.lifecycle.named_jobs @pytest.mark.asyncio @@ -119,7 +121,7 @@ async def initialize(): initialization = asyncio.create_task(initialize()) await started.wait() - actor = SimpleNamespace(driver=SimpleNamespace(wait_for_cleanup=AsyncMock()), + actor = SimpleNamespace(driver=SimpleNamespace(runtime=SimpleNamespace(wait_for_cleanup=AsyncMock())), core=SimpleNamespace(shutdown=AsyncMock(), wait_for_jobs=AsyncMock()), _initialization=initialization, _monitors=[], _failures={}) await cls.shutdown(actor) diff --git a/tests/server/training/test_controller.py b/tests/server/training/test_controller.py index 358c88e15..a8a485f27 100644 --- a/tests/server/training/test_controller.py +++ b/tests/server/training/test_controller.py @@ -199,7 +199,6 @@ async def test_reconciliation_of_committed_completion_does_not_extend_idle_time( controller, state, runtime, clock = await setup(backend=FakeBackend()) identity = await create(controller, state, 'active') await controller.admit(identity, 'owner', 'session', 'lost-finish') - await controller.commit_operation(identity, 'lost-finish', controller.epoch) async def finished(instance, request_id): return {'state': 'finished', 'executed': True, 'finished_at': time.time() - 31} runtime.verify_operation = finished @@ -362,7 +361,6 @@ async def test_heartbeat_does_not_prevent_idle_expiry_or_affect_busy_user(): idle = await create(controller, state, 1) busy = await create(controller, state, 2) await controller.admit(busy, 'owner', 'session', 'op') - await controller.commit_operation(busy, 'op', controller.epoch) clock.now = 11 await state.touch_session('session') await controller.sweep() @@ -387,7 +385,6 @@ async def test_release_drains_accepted_operation_before_cleanup(): controller, state, runtime, _ = await setup() identity = await create(controller, state, 1) await controller.admit(identity, 'owner', 'session', 'op') - await controller.commit_operation(identity, 'op', controller.epoch) releasing = asyncio.create_task(controller.release(identity, 'owner')) while controller.trainings[identity].state == 'ACTIVE': await asyncio.sleep(0) @@ -485,7 +482,6 @@ async def test_unconfirmed_worker_shutdown_holds_devices_and_quota_until_retry() identity = await create(controller, state, 1, mode='full') instance = controller.instances[controller.trainings[identity].instance_id] await controller.admit(identity, 'owner', 'session', 'running') - await controller.commit_operation(identity, 'running', controller.epoch) runtime.fail_shutdown = True await controller.fail_instance(instance.instance_id, FutureFailureRecord( reason_code='execution_timeout', message='RPC timed out', attribution='server')) @@ -651,7 +647,7 @@ async def test_retired_diagnostics_are_bounded_without_discarding_live_training( await controller.sweep() await controller.wait_for_jobs() assert old not in controller.trainings and retired_instance not in controller.instances - assert retired_instance not in controller._startup + assert retired_instance not in controller.lifecycle.startups assert (await controller.route(live, 'owner', 'session'))['instance_id'] in runtime.allocations await controller.shutdown() @@ -677,3 +673,124 @@ async def failing_set(key, value, *args, **kwargs): assert (await state.get_future(ref['request_id']))['status'] == 'failed' assert not runtime.allocations assert await state._backend.get('token_count::owner') == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize('fault', ['cancel_after_read', 'unknown_start_ack']) +async def test_create_requires_its_own_atomic_start_ack(fault, monkeypatch): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, _ = await setup(backend=FakeBackend()) + initialized = [] + original_create = runtime.create_training + + async def initialize(*args): + initialized.append(True) + await original_create(*args) + + runtime.create_training = initialize + if fault == 'cancel_after_read': + save = controller.store.save_training + + async def cancel_during_save(record): + await save(record) + instance = controller.instances.get(record.instance_id) + if record.state == 'PENDING' and instance is not None and instance.state == 'READY': + assert (await state.cancel_future(record.create_request_id))['cancelled'] + + monkeypatch.setattr(controller.store, 'save_training', cancel_during_save) + else: + start = state.try_start_future + + async def lose_ack(request_id): + assert await start(request_id) + raise ConnectionError('start acknowledgement lost') + + monkeypatch.setattr(state, 'try_start_future', lose_ack) + try: + ref = await controller.create(request(fault)) + await controller.wait_for_jobs() + future = await state.get_future(ref['request_id']) + assert future['status'] == ('cancelled' if fault == 'cancel_after_read' else 'failed') + assert not initialized and not runtime.users and not runtime.allocations + assert (await state._backend.get('token_count::owner') or 0) == 0 + finally: + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_quarantined_resident_counts_towards_maximum_until_death_confirmed(): + from tests.server.state.fake_backend import FakeBackend + from twinkle.server.state.models import FutureFailureRecord + controller, _, runtime, _ = await setup(backend=FakeBackend(), min_replicas=1, max_replicas=1) + old = next(iter(controller.instances.values())) + runtime.fail_shutdown = True + try: + await controller.fail_instance(old.instance_id, FutureFailureRecord( + reason_code='orphaned_replica', attribution='server', message='worker lost')) + for _ in range(3): + await controller.sweep() + await controller.wait_for_jobs() + assert old.state == 'QUARANTINED' + assert runtime.starts == 1 and list(runtime.allocations) == [old.instance_id] + runtime.fail_shutdown = False + await controller.sweep() + await controller.wait_for_jobs() + await controller.sweep() + await controller.wait_for_jobs() + live = [instance for instance in controller.instances.values() if instance.state != 'STOPPED'] + assert old.state == 'STOPPED' and len(live) == 1 + assert live[0].resident and live[0].state == 'READY' + assert live[0].instance_id != old.instance_id + finally: + runtime.fail_shutdown = False + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_last_lora_release_does_not_report_unconfirmed_device_release(): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'last-user') + instance = controller.instances[controller.trainings[identity].instance_id] + runtime.fail_shutdown = True + try: + with pytest.raises(RuntimeError, match='quarantined'): + await controller.release(identity, 'owner', 'session') + assert instance.state == 'QUARANTINED' and runtime.allocations + runtime.fail_shutdown = False + await controller.sweep() + await controller.wait_for_jobs() + assert instance.state == 'STOPPED' and not runtime.allocations + finally: + runtime.fail_shutdown = False + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_two_device_lora_capacity_shares_two_instances_and_releases_them(): + from tests.server.state.fake_backend import FakeBackend + from twinkle.server.config import ServerConfig + from twinkle.server.training.catalog import TrainingCatalog + controller, state, runtime, _ = await setup(backend=FakeBackend(), devices=4, quota=3) + controller.config = ServerConfig.from_yaml('cookbook/client/server/megatron/server_config.yaml').training_config + controller.catalog = TrainingCatalog(controller.config) + name = controller.config.models[0].name + try: + requests = [request(index, token=f'owner-{index // 3}').model_copy(update={'base_model': name}) + for index in range(6)] + refs = await asyncio.gather(*(controller.create(incoming) for incoming in requests)) + await controller.wait_for_jobs() + futures = [await state.get_future(ref['request_id']) for ref in refs] + assert all(future['status'] == 'completed' for future in futures) + assert len(controller.instances) == 2 + assert all(len(instance.training_ids) == 3 for instance in controller.instances.values()) + assert sum(runtime.allocations.values()) == 4 and set(runtime.allocations.values()) == {2} + excess = request('excess', token='new-owner').model_copy(update={'base_model': name}) + refused = await controller.create(excess) + assert (await state.get_future(refused['request_id']))['failure']['reason_code'] == 'capacity_unavailable' + for incoming, future in zip(requests, futures): + await controller.release(future['result']['training_id'], incoming.token, incoming.session_id) + assert not runtime.allocations and not runtime.users + assert all(instance.state == 'STOPPED' for instance in controller.instances.values()) + finally: + await controller.shutdown() diff --git a/tests/server/training/test_failure_recovery.py b/tests/server/training/test_failure_recovery.py index d93caf3fe..7466dc741 100644 --- a/tests/server/training/test_failure_recovery.py +++ b/tests/server/training/test_failure_recovery.py @@ -25,8 +25,7 @@ async def wait_idle(queue): @pytest.mark.asyncio -@pytest.mark.parametrize('fault', ['commit', 'queued_status']) -async def test_warm_idle_worker_wakes_after_enqueue_failure(fault, monkeypatch): +async def test_warm_idle_worker_wakes_after_enqueue_failure(monkeypatch): controller, state, _, _ = await setup(backend=FakeBackend()) identity = await create(controller, state, 'queue-failure') instance = controller.instances[controller.trainings[identity].instance_id] @@ -42,24 +41,18 @@ async def operation(): await settle(state, warmup['request_id']) await wait_idle(queue) executed.clear() - if fault == 'commit': - async def broken(*args): - raise RuntimeError('injected commit failure') - queue._training_controller.commit_operation = RPC(broken) - ref = await queue.schedule_task(operation, model_id=identity, token='owner', task_type='step') - else: - previous = state.store_future_status - - async def broken(request_id, status, *args, **kwargs): - if status == 'queued': - raise RuntimeError('injected queued write failure') - return await previous(request_id, status, *args, **kwargs) - - monkeypatch.setattr(state, 'store_future_status', broken) - with pytest.raises(RuntimeError, match='queued write'): - await queue.schedule_task(operation, model_id=identity, token='owner', task_type='step', - request_id='queued-failure') - ref = {'request_id': 'queued-failure'} + previous = state.store_future_status + + async def broken(request_id, status, *args, **kwargs): + if status == 'queued': + raise RuntimeError('injected queued write failure') + return await previous(request_id, status, *args, **kwargs) + + monkeypatch.setattr(state, 'store_future_status', broken) + with pytest.raises(RuntimeError, match='queued write'): + await queue.schedule_task(operation, model_id=identity, token='owner', task_type='step', + request_id='queued-failure') + ref = {'request_id': 'queued-failure'} assert (await settle(state, ref['request_id']))['status'] == 'completed' await wait_idle(queue) assert executed == [True] @@ -107,7 +100,7 @@ async def save(record): await controller.release(identity, 'owner', 'session') assert failed assert identity not in controller._cleaning - assert instance.instance_id not in controller._closing + assert instance.instance_id not in controller.lifecycle.closing assert training.state == 'CLOSING' assert instance.state != 'STOPPED' for _ in range(3): @@ -218,7 +211,7 @@ async def slow_shutdown(target): resume.set() await asyncio.gather(closing, shutdown) assert instance.state == 'STOPPED' - assert instance.instance_id not in controller._closing + assert instance.instance_id not in controller.lifecycle.closing @pytest.mark.asyncio @@ -255,11 +248,11 @@ async def broken(): raise RuntimeError('background sentinel') with caplog.at_level(logging.ERROR): - controller._spawn(broken()) + controller.lifecycle.spawn(broken()) await controller.wait_for_jobs() - records = [record for record in caplog.records if record.message == 'Training background job failed'] + records = [record for record in caplog.records if record.getMessage() == 'Training background job failed'] assert records and records[0].exc_info[1].args == ('background sentinel',) - assert not controller._jobs + assert not controller.lifecycle.jobs await controller.shutdown() @@ -336,7 +329,7 @@ def deploy(*args, **kwargs): monkeypatch.setattr('ray.util.list_named_actors', lambda **kwargs: []) monkeypatch.setattr(serve, 'run', deploy) monkeypatch.setattr(serve, 'delete', lambda name: None) - monkeypatch.setattr(driver.pool, 'remove', remove) + monkeypatch.setattr(driver.runtime.pool, 'remove', remove) deploying = asyncio.create_task(driver.deploy(instance, driver._catalog.get('small', 'lora'))) closing = None try: @@ -353,8 +346,8 @@ def deploy(*args, **kwargs): await asyncio.gather(deploying, return_exceptions=True) if closing is not None: await closing - await driver.wait_for_cleanup() - assert not driver._deployments and not driver._cleanup_tasks + await driver.runtime.wait_for_cleanup() + assert not driver.runtime.deployments and not driver.runtime.cleanup_tasks store.state.unregister_replica_capacity.assert_awaited() diff --git a/tests/server/training/test_runtime_queue.py b/tests/server/training/test_runtime_queue.py index 57f02083f..71ffcd2d6 100644 --- a/tests/server/training/test_runtime_queue.py +++ b/tests/server/training/test_runtime_queue.py @@ -27,7 +27,7 @@ def __init__(self, controller, state, instance): self._model_unhealthy = False self._training_controller = SimpleNamespace(**{ name: RPC(getattr(controller, name)) - for name in ('admit', 'commit_operation', 'finish_operation', 'fail_instance') + for name in ('admit', 'finish_operation', 'fail_instance') }) self._init_task_queue(TaskQueueConfig(execution_timeout=0.1, inline_fast_path_timeout=0.01), enable_admission_gate=True) @@ -103,7 +103,11 @@ async def operation(): @pytest.mark.asyncio -async def test_cancelled_queued_operation_releases_ticket_without_business_renewal(): +async def test_cancelled_queued_operation_releases_ticket_without_business_renewal(monkeypatch): + import time + # Freeze wall time as well as the injected business clock; RPC latency is not activity. + now = time.time() + monkeypatch.setattr('twinkle.server.task_queue.worker.time.time', lambda: now) controller, state, runtime, clock = await setup() identity = await create(controller, state, 1) instance = controller.instances[controller.trainings[identity].instance_id] @@ -189,3 +193,144 @@ async def compute(): finally: await queue.shutdown_task_queue() await controller.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('fault,executed', [('queue_timeout', False), ('unknown_start_ack', False), + ('backend_error', True)]) +async def test_lost_failure_notification_uses_actual_execution_facts(fault, executed, monkeypatch): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, clock = await setup(backend=FakeBackend()) + identity = await create(controller, state, fault) + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + + async def lost(*args, **kwargs): + raise ConnectionError('Completion RPC lost') + + queue._training_controller.finish_operation = RPC(lost) + runtime.verify_operation = lambda instance, request_id: queue.verify_operation(request_id) + ran = [] + + async def operation(): + ran.append(True) + raise ValueError('executed backend error') + + if fault == 'queue_timeout': + queue._compute_worker._config = TaskQueueConfig(queue_timeout=0) + elif fault == 'unknown_start_ack': + start = state.try_start_future + + async def lose_ack(request_id): + assert await start(request_id) + raise ConnectionError('start acknowledgement lost') + + monkeypatch.setattr(state, 'try_start_future', lose_ack) + try: + clock.now = 7 + ref = await queue.schedule_task(operation, model_id=identity, token='owner', task_type='step') + assert (await settle(state, ref['request_id']))['status'] == 'failed' + for _ in range(100): + if not queue._compute_worker.has_task(ref['request_id']): + break + await asyncio.sleep(0.01) + outcome = await queue.verify_operation(ref['request_id']) + assert outcome['state'] == 'finished' and outcome['executed'] is executed + assert bool(ran) is executed + assert controller.trainings[identity].pending_operation_ids + clock.now = 13 # Admission at 7 must age past the five-second reconciliation interval. + await controller.sweep() + await controller.wait_for_jobs() + training = controller.trainings[identity] + assert not training.pending_operation_ids + assert (training.last_activity_at > 6) is executed + if not executed: + clock.now = 14 + await controller.sweep() + await controller.wait_for_jobs() + assert training.state == 'EXPIRED' + finally: + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_execution_facts_expire_without_future_status_guessing(monkeypatch): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'facts-expire') + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + + async def lost(*args, **kwargs): + raise ConnectionError('Completion RPC lost') + + queue._training_controller.finish_operation = RPC(lost) + try: + ref = await queue.schedule_task(lambda: asyncio.sleep(0), model_id=identity, token='owner', task_type='step') + await settle(state, ref['request_id']) + worker = queue._compute_worker + for _ in range(100): + if not worker.has_task(ref['request_id']): + break + await asyncio.sleep(0.01) + outcome = worker.get_task_outcome(ref['request_id']) + assert outcome['executed'] + monkeypatch.setattr('twinkle.server.task_queue.worker.time.time', + lambda: outcome['finished_at'] + worker._config.absolute_future_ttl(1) + 1) + assert (await queue.verify_operation(ref['request_id']))['state'] == 'missing' + assert not worker._task_outcomes + finally: + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_completion_write_failure_retains_ticket_and_execution_fact(monkeypatch): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, clock = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'finish-write') + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + runtime.verify_operation = lambda instance, request_id: queue.verify_operation(request_id) + save = controller.store.save_training + failed = False + + async def fail_completion(record): + nonlocal failed + if record.state == 'ACTIVE' and not record.pending_operation_ids and not failed: + failed = True + raise ConnectionError('completion write failed') + await save(record) + + monkeypatch.setattr(controller.store, 'save_training', fail_completion) + try: + clock.now = 7 + ref = await queue.schedule_task(lambda: asyncio.sleep(0), model_id=identity, token='owner', task_type='step') + await settle(state, ref['request_id']) + for _ in range(100): + if not queue._compute_worker.has_task(ref['request_id']): + break + await asyncio.sleep(0.01) + for _ in range(100): + if failed: + break + await asyncio.sleep(0.01) + training = controller.trainings[identity] + assert failed and ref['request_id'] in training.pending_operation_ids + assert training.last_activity_at == 0 + assert queue._compute_worker.get_task_outcome(ref['request_id'])['executed'] + clock.now = 13 + await controller.sweep() + await controller.wait_for_jobs() + assert not training.pending_operation_ids and 6 < training.last_activity_at <= 13 + followup = await queue.schedule_task(lambda: asyncio.sleep(0), model_id=identity, token='owner', task_type='step') + await settle(state, followup['request_id']) + for _ in range(100): + if queue._compute_worker.get_task_outcome(followup['request_id']) is None: + break + await asyncio.sleep(0.01) + assert queue._compute_worker.get_task_outcome(followup['request_id']) is None + finally: + await queue.shutdown_task_queue() + await controller.shutdown() diff --git a/tests/twinkle_client/test_auto_agent_tools.py b/tests/twinkle_client/test_auto_agent_tools.py index dcfc871e5..c4ea1a817 100644 --- a/tests/twinkle_client/test_auto_agent_tools.py +++ b/tests/twinkle_client/test_auto_agent_tools.py @@ -81,6 +81,13 @@ def test_generated_server_config_matches_launcher_schema(tmp_path, monkeypatch, gateway = next(app for app in config.applications if app.import_path == 'server') assert gateway.deployments[0]['name'] == 'GatewayServer' assert gateway.args.supported_models == ['Qwen/student'] + [s['model_id'] for s in samplers] + application = next(app for app in config.applications if app.import_path == 'model') + assert application.args.modes is None and application.args.train_mode == 'lora' + model = config.training_config.models[0] + assert model.modes.full is None + assert model.modes.lora.min_replicas == 0 and model.modes.lora.max_replicas == 1 + assert model.modes.lora.num_devices == 2 and model.modes.lora.idle_timeout_s == 600 + assert model.legacy_expiration.session_timeout_s == 600 if samplers: sampler = next(app for app in config.applications if app.import_path == 'sampler') assert sampler.args.nproc_per_node == 4 @@ -102,6 +109,14 @@ def test_removed_sampler_is_not_advertised_or_written(tmp_path, monkeypatch): assert not (tmp_path / '.cache' / 'twinkle' / 'server_config.yaml').exists() +def test_single_device_megatron_config_is_not_written(tmp_path, monkeypatch): + from pathlib import Path + monkeypatch.setattr(Path, 'home', lambda: tmp_path) + with pytest.raises(ValueError, match='at least 2 devices'): + ToolExecutor._generate_server_config('Qwen/student', train_gpus=1, backend='megatron') + assert not (tmp_path / '.cache' / 'twinkle' / 'server_config.yaml').exists() + + @pytest.mark.parametrize('backend', ['torch', 'missing']) def test_unsupported_training_backend_is_not_written(tmp_path, monkeypatch, backend): from pathlib import Path From bf37084501a43fe5ed2fd64e6731b9ef04550860 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Thu, 22 Oct 2026 21:42:54 +0800 Subject: [PATCH 11/46] fix: guard sampler mutations and preserve training ownership --- src/twinkle/server/sampler/concurrency.py | 16 +- src/twinkle/server/sampler/controller.py | 139 ++++++++---- .../model/multi_lora_transformers.py | 2 + tests/pluger/test_loader.py | 14 +- tests/server/resources/test_instances.py | 23 ++ tests/server/sampler/test_pool.py | 203 ++++++++++++++++++ .../server/training/test_failure_recovery.py | 19 -- tests/twinkle_client/test_async_components.py | 55 ++++- 8 files changed, 398 insertions(+), 73 deletions(-) create mode 100644 tests/server/resources/test_instances.py diff --git a/src/twinkle/server/sampler/concurrency.py b/src/twinkle/server/sampler/concurrency.py index 71d7e7c68..b3e7accc0 100644 --- a/src/twinkle/server/sampler/concurrency.py +++ b/src/twinkle/server/sampler/concurrency.py @@ -11,7 +11,7 @@ class SamplerConcurrency: """Readers share stable weights; mutations wait for every reader to finish. Jobs retain their lock if an HTTP waiter leaves. Admission and execution - timeouts belong to ComputeWorker, not this lock. + timeouts belong to the caller; the pool may also bound lock admission. """ def __init__(self): @@ -35,13 +35,17 @@ def invalidate_weights(self, error: Exception): def close(self): self._closed = True - async def _execute(self, factory, exclusive: bool): + async def _execute(self, factory, exclusive: bool, queue_timeout): async with self._condition: if exclusive: self._waiting_writers += 1 try: - await self._condition.wait_for(lambda: not self._writer and - (self._readers == 0 if exclusive else self._waiting_writers == 0)) + + def available(): + return not self._writer and (self._readers == 0 if exclusive else self._waiting_writers == 0) + + if not available(): + await asyncio.wait_for(self._condition.wait_for(available), queue_timeout) self.check_available() if exclusive: self._writer = True @@ -61,7 +65,7 @@ async def _execute(self, factory, exclusive: bool): self._readers -= 1 self._condition.notify_all() - async def run(self, factory, *, exclusive: bool = False): + async def run(self, factory, *, exclusive: bool = False, queue_timeout: float | None = None): self.check_available() started = False @@ -70,7 +74,7 @@ async def call(): started = True return await factory() - job = asyncio.create_task(self._execute(call, exclusive)) + job = asyncio.create_task(self._execute(call, exclusive, queue_timeout)) self._jobs.add(job) def finished(task): diff --git a/src/twinkle/server/sampler/controller.py b/src/twinkle/server/sampler/controller.py index e2246abf8..7ba3f5d73 100644 --- a/src/twinkle/server/sampler/controller.py +++ b/src/twinkle/server/sampler/controller.py @@ -4,10 +4,12 @@ import asyncio import hashlib +import json import logging import math import time import uuid +from copy import deepcopy from dataclasses import asdict, dataclass, field from twinkle.server.config.application_spec import SamplerArgs @@ -139,6 +141,97 @@ async def _invoke_finished(self, instance, payload): await self.driver.wait_operation(instance, ticket) return reply + @staticmethod + def _action_key(action): + if 'base_path' in action: + return ('base_path', ) + endpoint = action['endpoint'] + if endpoint == 'twinkle/set_template': + return (endpoint, ) + if endpoint == 'twinkle/add_adapter_to_sampler': + # Adapter names are scoped by the caller's token in the handler. + headers = {key.lower(): value for key, value in action['headers']} + owner = headers.get(b'twinkle-authorization') or headers.get(b'authorization') + return endpoint, owner, json.loads(action['body'])['adapter_name'] + return None + + def _remember(self, action): + """Compact overwritten state only within the segment after the last opaque patch.""" + key = self._action_key(action) + if key is not None: + for index in range(len(self._actions) - 1, -1, -1): + previous = self._actions[index] + if previous.get('endpoint') == 'twinkle/apply_patch': + break + if self._action_key(previous) == key: + del self._actions[index] + self._actions.append(deepcopy(action)) + + async def _release(self, instance, ticket): + async with self._lock: + # Unknown completion retains ownership until shutdown is confirmed. + if instance.state not in ('DRAINING', 'QUARANTINED'): + instance.requests.discard(ticket) + instance.last_finished = self.clock() + + async def _fence(self, replicas): + async with self._lock: + for replica in replicas: + if replica.state != 'STOPPED': + replica.state = 'DRAINING' + await asyncio.gather(*(self._close(replica) for replica in replicas)) + + async def _mutate(self, payload=None, full_path=None): + # Lock order: config -> state. Never wait for startup while holding config, + # because startup also replays the committed configuration under that lock. + seed, seed_ticket = await self._select() + replicas = [] + ticket = uuid.uuid4().hex + try: + async with self._config_lock: + async with self._lock: + if self._stopping: + raise ResourceExhaustedError('Sampling pool is stopping') + replicas = self._ready_instances() + if not replicas: + raise ResourceExhaustedError('No ready sampler for control update') + for replica in replicas: + replica.requests.add(ticket) + + async def broadcast(): + response = None + for replica in replicas: + if full_path is not None: + await self.driver.load_weights(replica, full_path) + else: + response = await self._invoke_finished(replica, payload) + if response['status'] >= 400: + return response + return response + + try: + response = await asyncio.wait_for(broadcast(), self.args.queue_config.effective_execution_timeout) + if response is not None and response['status'] >= 400: + await self._fence(replicas) + return response + # A health probe or shutdown can invalidate a participant while + # its RPC is pending. Such an ACK must not commit pool state. + if any(replica.state != 'READY' for replica in replicas): + raise ResourceExhaustedError('Sampler changed state during control update') + if full_path is not None: + self._base_path = full_path + self._remember({'base_path': full_path}) + elif payload['endpoint'] != 'twinkle/unload_adapter_paths': + self._remember(payload) + return response + except BaseException: + await self._fence(replicas) + raise + finally: + for replica in replicas: + await self._release(replica, ticket) + await self._release(seed, seed_ticket) + async def request(self, payload, stream_queue=None, full_path=None): await self.initialize() endpoint = payload['endpoint'] @@ -152,43 +245,10 @@ async def operation(): ticket = None try: if control: - if not self._ready_instances(): - instance, ticket = await self._select() - async with self._config_lock: - replies = [] - replicas = self._ready_instances() - try: - for replica in replicas: - response = await self._invoke_finished(replica, payload) - if response['status'] >= 400: - await asyncio.gather(*(self._close(peer) for peer in replicas)) - reply_ready.set_result(response) - return - replies.append(response) - except BaseException: - # Even an ambiguous ACK can leave this replica mutated. - # Fence every participant before allowing further routing. - for peer in replicas: - peer.state = 'DRAINING' - await asyncio.gather(*(self._close(peer) for peer in replicas)) - raise - if not replies: - raise ResourceExhaustedError('No ready sampler for control update') - if endpoint != 'twinkle/unload_adapter_paths': - self._actions.append(payload) - reply_ready.set_result(replies[0]) + reply_ready.set_result(await self._mutate(payload=payload)) return if full_path is not None and full_path != self._base_path: - async with self._config_lock: - try: - for replica in self._ready_instances(): - await self.driver.load_weights(replica, full_path) - self._base_path = full_path - self._actions.append({'base_path': full_path}) - except Exception: - for replica in self._ready_instances(): - await self._close(replica) - raise + await self._mutate(full_path=full_path) instance, ticket = await self._select() response = await self.driver.invoke(instance, ticket, payload, stream_queue) if not reply_ready.done(): @@ -204,13 +264,14 @@ async def operation(): raise finally: if instance is not None and ticket is not None: - if instance.state != 'QUARANTINED': - instance.requests.discard(ticket) - instance.last_finished = self.clock() + await self._release(instance, ticket) async def guarded(): try: - await self._guard.run(operation, exclusive=control or full_path is not None) + await self._guard.run( + operation, + exclusive=control or full_path is not None, + queue_timeout=self.args.queue_config.queue_timeout) except BaseException as error: if not reply_ready.done(): reply_ready.set_exception(error) diff --git a/src/twinkle_client/model/multi_lora_transformers.py b/src/twinkle_client/model/multi_lora_transformers.py index b445aca5c..876bb5bd1 100644 --- a/src/twinkle_client/model/multi_lora_transformers.py +++ b/src/twinkle_client/model/multi_lora_transformers.py @@ -148,6 +148,8 @@ def add_adapter_to_model( with self._training_lock: if adapter_name in self._creating: raise RuntimeError('A training with this adapter_name is already being created on this object.') + if adapter_name in self._training_ids: + raise RuntimeError('A training with this adapter_name already exists on this object; release it first.') self._creating.add(adapter_name) try: result = self._submit( diff --git a/tests/pluger/test_loader.py b/tests/pluger/test_loader.py index 76735a697..9d688c4fb 100644 --- a/tests/pluger/test_loader.py +++ b/tests/pluger/test_loader.py @@ -35,7 +35,7 @@ def test_load_plugin_invalid_id(self): def test_load_plugin_safe_mode(self): """Test loading plugin when trust_remote_code is False.""" - with patch('twinkle.utils.loader.MSHub.download_model', return_value='/tmp/fake'): + with patch('twinkle.hub.hub.MSHub.download_model', return_value='/tmp/fake'): with patch('twinkle.utils.loader.trust_remote_code', return_value=False): with pytest.raises(ValueError, match='Twinkle does not support plugin in safe mode'): Plugin.load_plugin('ms://test/plugin', BasePlugin) @@ -43,7 +43,7 @@ def test_load_plugin_safe_mode(self): def test_load_plugin_missing_init_file(self): """Test loading plugin when __init__.py is missing.""" with tempfile.TemporaryDirectory() as tmpdir: - with patch('twinkle.utils.loader.MSHub.download_model', return_value=tmpdir): + with patch('twinkle.hub.hub.MSHub.download_model', return_value=tmpdir): with patch('twinkle.utils.loader.trust_remote_code', return_value=True): with pytest.raises(AssertionError, match='does not exist'): Plugin.load_plugin('ms://test/plugin', BasePlugin) @@ -58,7 +58,7 @@ def test_load_plugin_no_subclass(self): other_class = type('OtherClass', (), {}) other_class.__module__ = str(init_file) - with patch('twinkle.utils.loader.MSHub.download_model', return_value=tmpdir): + with patch('twinkle.hub.hub.MSHub.download_model', return_value=tmpdir): with patch('twinkle.utils.loader.trust_remote_code', return_value=True): with patch('twinkle.utils.loader.importlib.import_module') as mock_import: mock_module = MagicMock() @@ -87,7 +87,7 @@ def test_load_plugin_ms_hub(self): test_plugin_class.__module__ = '__init__' mock_module.TestPlugin = test_plugin_class - with patch('twinkle.utils.loader.MSHub.download_model', return_value=tmpdir): + with patch('twinkle.hub.hub.MSHub.download_model', return_value=tmpdir): with patch('twinkle.utils.loader.trust_remote_code', return_value=True): with patch('twinkle.utils.loader.importlib.import_module', return_value=mock_module): plugin_cls = Plugin.load_plugin('ms://test/plugin', BasePlugin) @@ -111,7 +111,7 @@ def test_load_plugin_hf_hub(self): test_plugin_class.__module__ = '__init__' mock_module.TestPlugin = test_plugin_class - with patch('twinkle.utils.loader.HFHub.download_model', return_value=tmpdir): + with patch('twinkle.hub.hub.HFHub.download_model', return_value=tmpdir): with patch('twinkle.utils.loader.trust_remote_code', return_value=True): with patch('twinkle.utils.loader.importlib.import_module', return_value=mock_module): plugin_cls = Plugin.load_plugin('hf://test/plugin', BasePlugin) @@ -137,7 +137,7 @@ def test_load_plugin_sys_path_management(self): assert tmpdir not in sys.path - with patch('twinkle.utils.loader.MSHub.download_model', return_value=tmpdir): + with patch('twinkle.hub.hub.MSHub.download_model', return_value=tmpdir): with patch('twinkle.utils.loader.trust_remote_code', return_value=True): with patch('twinkle.utils.loader.importlib.import_module', return_value=mock_module): Plugin.load_plugin('ms://test/plugin', BasePlugin) @@ -186,7 +186,7 @@ def test_construct_class_with_plugin_id(self): test_plugin_class.__module__ = '__init__' mock_module.TestPlugin = test_plugin_class - with patch('twinkle.utils.loader.MSHub.download_model', return_value=tmpdir): + with patch('twinkle.hub.hub.MSHub.download_model', return_value=tmpdir): with patch('twinkle.utils.loader.trust_remote_code', return_value=True): with patch('twinkle.utils.loader.importlib.import_module', return_value=mock_module): result = construct_class('ms://test/plugin', BasePlugin, [], name='test') diff --git a/tests/server/resources/test_instances.py b/tests/server/resources/test_instances.py new file mode 100644 index 000000000..42c78503f --- /dev/null +++ b/tests/server/resources/test_instances.py @@ -0,0 +1,23 @@ +import logging +import pytest +from twinkle.server.resources.instances import InstanceLifecycle + + +@pytest.mark.asyncio +async def test_spawn_logs_exception_with_traceback(caplog, monkeypatch): + # The ``twinkle`` logger disables propagation, so attach caplog to the lifecycle logger directly. + target = logging.getLogger('twinkle.server.resources.instances') + monkeypatch.setattr(target, 'handlers', [*target.handlers, caplog.handler]) + lifecycle = InstanceLifecycle('Training') + + async def broken(): + raise RuntimeError('background sentinel') + + with caplog.at_level(logging.ERROR): + lifecycle.spawn(broken()) + await lifecycle.join() + records = [record for record in caplog.records if record.getMessage() == 'Training background job failed'] + assert records and records[0].exc_info[1].args == ('background sentinel',) + assert not lifecycle.jobs + + diff --git a/tests/server/sampler/test_pool.py b/tests/server/sampler/test_pool.py index ba47028a1..847c266c2 100644 --- a/tests/server/sampler/test_pool.py +++ b/tests/server/sampler/test_pool.py @@ -335,3 +335,206 @@ async def failed(): assert not lifecycle.unreachable_since assert await lifecycle.health('instance', unreachable) is None assert await lifecycle.health('instance', failed) is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize('kind', ['control', 'weights']) +async def test_mutation_reserves_every_replica_against_idle_sweep(kind): + pool, driver, clock = setup(minimum=1, maximum=2) + driver.finish.set() + await pool.initialize() + pool._new_instance() + await settle_starts(pool) + clock.now = 20 + entered, resume = asyncio.Event(), asyncio.Event() + invoke, load = driver.invoke, driver.load_weights + + async def update(instance, *args): + entered.set() + await resume.wait() + assert instance.state == 'READY' + return await (invoke(instance, *args) if kind == 'control' else load(instance, *args)) + + if kind == 'control': + driver.invoke = update + task = asyncio.create_task(pool.request(request('twinkle/set_template'))) + else: + driver.load_weights = update + task = asyncio.create_task(pool.request(request(), full_path='B')) + try: + await asyncio.wait_for(entered.wait(), 1) + await pool.sweep() + # Only wait for the close jobs, never for the deliberately blocked mutation. + await asyncio.gather(*tuple(pool.lifecycle.closing.values())) + assert len(pool._ready_instances()) == 2 + assert not driver.stopped + finally: + resume.set() + await asyncio.gather(task, return_exceptions=True) + await pool.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('stage', ['invoke', 'wait_operation', 'load_weights']) +async def test_control_deadline_fences_all_participants(stage): + pool, driver, _ = setup() + driver.finish.set() + pool.args.queue_config.execution_timeout = 0.03 + entered, resume = asyncio.Event(), asyncio.Event() + + async def blocked(*args): + entered.set() + await resume.wait() + + setattr(driver, stage, blocked) + task = asyncio.create_task(pool.request( + request() if stage == 'load_weights' else request('twinkle/set_template'), + full_path='B' if stage == 'load_weights' else None)) + try: + await asyncio.wait_for(entered.wait(), 1) + with pytest.raises(TimeoutError): + await asyncio.wait_for(asyncio.shield(task), 0.3) + assert task.done(), 'Configured execution deadline was not enforced' + await settle_jobs(pool) + assert not pool._ready_instances() + assert len(driver.stopped) == 2 + assert not pool._actions + finally: + resume.set() + await asyncio.gather(task, return_exceptions=True) + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_superseded_checkpoint_is_not_needed_to_start_new_replica(): + pool, driver, _ = setup(minimum=1) + driver.finish.set() + valid = {'A', 'B'} + original = driver.load_weights + + async def load(instance, path): + if path not in valid: + raise FileNotFoundError(path) + await original(instance, path) + + driver.load_weights = load + try: + for path in ('A', 'B'): + await pool.request(request(), full_path=path) + await settle_jobs(pool) + valid.remove('A') + instance = pool._new_instance() + await pool.lifecycle.startups[instance.identity] + assert instance.state == 'READY' + assert [payload['base_path'] for identity, payload in driver.calls + if identity == instance.identity and 'base_path' in payload] == ['B'] + finally: + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_queued_control_expires_without_mutating_or_fencing_generation(): + pool, driver, _ = setup(minimum=1, maximum=1) + pool.args.queue_config.queue_timeout = 0.02 + await pool.request(request()) + try: + with pytest.raises(TimeoutError): + await asyncio.wait_for(pool.request(request('twinkle/set_template')), 1) + assert len(driver.calls) == 1 and not driver.stopped + driver.finish.set() + await settle_jobs(pool) + await pool.request(request('twinkle/set_template')) + assert len(driver.calls) == 2 + finally: + driver.finish.set() + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_cancelled_control_waiter_keeps_replica_occupied_until_completion(): + pool, driver, clock = setup(minimum=1, maximum=2) + await pool.initialize() + pool._new_instance() + await settle_starts(pool) + entered = asyncio.Event() + invoke = driver.invoke + + async def update(*args): + entered.set() + return await invoke(*args) + + driver.invoke = update + task = asyncio.create_task(pool.request(request('twinkle/set_template'))) + try: + await entered.wait() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + clock.now = 30 + await pool.sweep() + assert len(pool._ready_instances()) == 2 and not driver.stopped + driver.finish.set() + await settle_jobs(pool) + clock.now += 11 + await pool.sweep() + await settle_jobs(pool) + assert len(pool._ready_instances()) == 1 + finally: + driver.finish.set() + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_control_timeout_retains_unknown_allocations_then_recovers(): + pool, driver, _ = setup(minimum=2, maximum=2) + pool.args.queue_config.execution_timeout = 0.02 + driver.fail_close = True + try: + with pytest.raises(TimeoutError): + await pool.request(request('twinkle/set_template')) + await settle_jobs(pool) + assert all(i.state == 'QUARANTINED' and i.requests for i in pool.instances.values()) + await pool.sweep() + await settle_jobs(pool) + assert len(driver.started) == 2 + driver.fail_close = False + driver.finish.set() + await pool.sweep() + await settle_jobs(pool) + await pool.sweep() + await settle_starts(pool) + await pool.request(request()) + assert len(driver.started) == 4 + assert len(pool._ready_instances()) == 2 + finally: + driver.fail_close = False + driver.finish.set() + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_replay_compacts_setters_per_owner_without_crossing_patch(): + pool, driver, _ = setup(minimum=1) + driver.finish.set() + try: + for index in range(12): + await pool.request(request('twinkle/set_template', {'template_cls': f'T{index}'})) + for owner in ('one', 'two'): + for rank in (4, 8): + payload = request('twinkle/add_adapter_to_sampler', {'adapter_name': 'same', 'config': {'r': rank}}) + payload['headers'].append((b'twinkle-authorization', owner.encode())) + await pool.request(payload) + assert len(pool._actions) == 3 + await pool.request(request('twinkle/apply_patch', {'patch_cls': 'P'})) + await pool.request(request('twinkle/set_template', {'template_cls': 'AfterPatch'})) + await pool.request(request('twinkle/unload_adapter_paths', {'adapter_paths': ['unused']})) + instance = pool._new_instance() + await pool.lifecycle.startups[instance.identity] + replay = [p for identity, p in driver.calls if identity == instance.identity] + assert [p['endpoint'] for p in replay] == [ + 'twinkle/set_template', 'twinkle/add_adapter_to_sampler', 'twinkle/add_adapter_to_sampler', + 'twinkle/apply_patch', 'twinkle/set_template'] + assert json.loads(replay[0]['body'])['template_cls'] == 'T11' + assert all(json.loads(p['body'])['config']['r'] == 8 for p in replay[1:3]) + assert json.loads(replay[-1]['body'])['template_cls'] == 'AfterPatch' + finally: + await pool.shutdown() diff --git a/tests/server/training/test_failure_recovery.py b/tests/server/training/test_failure_recovery.py index 7466dc741..f310f268e 100644 --- a/tests/server/training/test_failure_recovery.py +++ b/tests/server/training/test_failure_recovery.py @@ -237,25 +237,6 @@ async def test_unready_service_is_503_server_and_reason_sets_have_no_dead_codes( await controller.shutdown() -@pytest.mark.asyncio -async def test_spawn_logs_exception_with_traceback(caplog, monkeypatch): - # The ``twinkle`` logger disables propagation, so attach caplog to the controller logger directly. - target = logging.getLogger('twinkle.server.training.controller') - monkeypatch.setattr(target, 'handlers', [*target.handlers, caplog.handler]) - controller, _, _, _ = await setup(backend=FakeBackend()) - - async def broken(): - raise RuntimeError('background sentinel') - - with caplog.at_level(logging.ERROR): - controller.lifecycle.spawn(broken()) - await controller.wait_for_jobs() - records = [record for record in caplog.records if record.getMessage() == 'Training background job failed'] - assert records and records[0].exc_info[1].args == ('background sentinel',) - assert not controller.lifecycle.jobs - await controller.shutdown() - - @pytest.mark.asyncio async def test_sweep_heartbeat_io_does_not_hold_route_lock(monkeypatch): controller, state, _, _ = await setup(backend=FakeBackend()) diff --git a/tests/twinkle_client/test_async_components.py b/tests/twinkle_client/test_async_components.py index 313986191..fff5c1206 100644 --- a/tests/twinkle_client/test_async_components.py +++ b/tests/twinkle_client/test_async_components.py @@ -247,8 +247,9 @@ def post_model(self, url, body): self.created += 1 return _Response(_completed({'training_id': f'training-{self.created}', 'adapter_name': 'same'})) if url.endswith('/remove_adapter'): - entered.set() - assert resume.wait(3) + if not entered.is_set(): + entered.set() + assert resume.wait(3) return _Response(_completed(None)) transport = Transport() model = MultiLoraTransformersModel('small', transport=transport) @@ -261,6 +262,7 @@ def remove(): releasing.start() try: assert entered.wait(3) + model.remove_adapter('same') model.add_adapter_to_model('same', None) finally: resume.set() @@ -270,3 +272,52 @@ def remove(): assert transport.operations[-1][1]['training_id'] == 'training-2' removal = next(body for endpoint, body in transport.operations if endpoint == 'remove_adapter') assert removal['training_id'] == 'training-1' + + +def test_duplicate_adapter_is_rejected_before_rpc_and_release_allows_recreate(monkeypatch): + import pytest + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + calls = [] + _patch_transport(monkeypatch, calls, lambda url: _completed(None)) + first = MultiLoraTransformersModel('base') + first.add_adapter_to_model('same', None) + count = len(calls) + with pytest.raises(RuntimeError, match='already'): + first.add_adapter_to_model('same', None) + assert len(calls) == count + second = MultiLoraTransformersModel('base') + second.add_adapter_to_model('same', None) + first.remove_adapter('same') + first.add_adapter_to_model('same', None) + + +def test_concurrent_duplicate_and_failed_creation_leave_reusable_name(monkeypatch): + import threading + from concurrent.futures import ThreadPoolExecutor + import pytest + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + calls = [] + _patch_transport(monkeypatch, calls, lambda url: _completed(None)) + model = MultiLoraTransformersModel('base') + original = model._submit + entered, resume = threading.Event(), threading.Event() + + def fail(*args, **kwargs): + entered.set() + assert resume.wait(3) + raise RuntimeError('create failed') + + monkeypatch.setattr(model, '_submit', fail) + with ThreadPoolExecutor(1) as executor: + future = executor.submit(model.add_adapter_to_model, 'same', None) + try: + assert entered.wait(3) + with pytest.raises(RuntimeError, match='already being created'): + model.add_adapter_to_model('same', None) + finally: + resume.set() + with pytest.raises(RuntimeError, match='create failed'): + future.result() + monkeypatch.setattr(model, '_submit', original) + model.add_adapter_to_model('same', None) + assert model._training_ids['same'] == 'training-test' From 747f655d2b81befb2d75014076e5948515afbbe6 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Thu, 22 Oct 2026 21:55:41 +0800 Subject: [PATCH 12/46] refactor: align server responsibilities and test ownership --- .gitignore | 2 - src/twinkle/server/README.md | 44 ++++++ src/twinkle/server/gateway/routes.py | 22 +-- src/twinkle/server/sampler/actor.py | 30 +++++ src/twinkle/server/sampler/controller.py | 24 ---- src/twinkle/server/sampler/router.py | 3 +- .../server/telemetry/training_metrics.py | 40 ++++++ src/twinkle/server/training/controller.py | 43 +----- .../server/{utils => checkpoint}/__init__.py | 0 .../test_checkpoint_manager.py | 0 .../test_task_queue_config.py} | 0 .../resources}/__init__.py | 0 .../allocation_probe.py | 0 .../fragmentation_probe.py | 0 .../resources}/test_allocation.py | 4 +- tests/server/integration/sampler/__init__.py | 0 .../sampling_probe.py | 0 .../model/test_full_mode_exclusivity.py | 2 +- tests/server/resources/__init__.py | 0 tests/server/resources/test_instances.py | 61 +++++++++ tests/server/sampler/test_pool.py | 28 +--- tests/server/sampler/test_request_runtime.py | 2 +- .../{training => state}/test_sequence.py | 2 +- tests/server/task_queue/__init__.py | 0 tests/server/task_queue/helpers.py | 60 +++++++++ .../test_rate_limiter.py | 0 .../test_task_queue_mixin.py | 51 +------ .../server/telemetry/test_training_metrics.py | 6 +- tests/server/test_dependencies.py | 11 ++ .../test_proxy.py => test_service_proxy.py} | 7 +- tests/server/training/helpers.py | 126 ++++++++++++++++++ tests/server/training/test_controller.py | 85 +----------- .../server/training/test_failure_recovery.py | 4 +- tests/server/training/test_router.py | 2 +- .../training/test_runtime_initialization.py | 2 +- tests/server/training/test_runtime_queue.py | 33 +---- .../server/training/test_storage_recovery.py | 2 +- 37 files changed, 400 insertions(+), 296 deletions(-) create mode 100644 src/twinkle/server/README.md create mode 100644 src/twinkle/server/sampler/actor.py create mode 100644 src/twinkle/server/telemetry/training_metrics.py rename tests/server/{utils => checkpoint}/__init__.py (100%) rename tests/server/{utils => checkpoint}/test_checkpoint_manager.py (100%) rename tests/server/{utils/task_queue/test_config.py => config/test_task_queue_config.py} (100%) rename tests/server/{utils/task_queue => integration/resources}/__init__.py (100%) rename tests/server/integration/{on_demand_training => resources}/allocation_probe.py (100%) rename tests/server/integration/{on_demand_training => resources}/fragmentation_probe.py (100%) rename tests/server/{training => integration/resources}/test_allocation.py (84%) create mode 100644 tests/server/integration/sampler/__init__.py rename tests/server/integration/{on_demand_training => sampler}/sampling_probe.py (100%) create mode 100644 tests/server/resources/__init__.py rename tests/server/{training => state}/test_sequence.py (97%) create mode 100644 tests/server/task_queue/__init__.py create mode 100644 tests/server/task_queue/helpers.py rename tests/server/{utils => task_queue}/test_rate_limiter.py (100%) rename tests/server/{utils => task_queue}/test_task_queue_mixin.py (92%) rename tests/server/{gateway/test_proxy.py => test_service_proxy.py} (93%) create mode 100644 tests/server/training/helpers.py diff --git a/.gitignore b/.gitignore index cc52a2f6a..0d55deae1 100644 --- a/.gitignore +++ b/.gitignore @@ -78,8 +78,6 @@ instance/ docs/_build/ # Internal design notes and developer tools are local-only. -/docs/design/ -/tools/ # PyBuilder target/ diff --git a/src/twinkle/server/README.md b/src/twinkle/server/README.md new file mode 100644 index 000000000..917479856 --- /dev/null +++ b/src/twinkle/server/README.md @@ -0,0 +1,44 @@ +# Server responsibilities + +The training and sampling pools share resource mechanisms, but own separate scheduling policies and records. + +```mermaid +flowchart TD + Launcher --> TrainingActor + Launcher --> SamplingActor + TrainingActor --> TrainingController + SamplingActor --> SamplerController + TrainingController --> InstanceLifecycle + SamplerController --> InstanceLifecycle + TrainingController --> TrainingDriver + SamplerController --> SamplerDriver + TrainingDriver --> RayRuntime + SamplerDriver --> RayRuntime + RayRuntime --> DevicePool +``` + +- `training/`: training catalog, handle ownership, placement policy, router and Ray actor composition. +- `model/`: training execution and backend/protocol adapters. `managed_runtime.py` implements worker-side admission and completion, not pool scheduling. +- `sampler/`: sampling policy, runtime execution and backend/protocol adapters. `actor.py` owns Ray composition; `controller.py` accepts an injected driver. +- `resources/`: shared device allocations, Ray/Serve deployment and confirmed shutdown, and background instance supervision. No training or sampling policy imports. +- `state/`: serializable records and persistence. Training records stay here so model execution need not import the training controller. +- `lifecycle/`: request submission, retrieval and Future/Envelope handling. This differs from the instance lifetime supervised in `resources/`. +- `config/`: configuration declarations and normalization. `telemetry/` consumes plain snapshots without importing controllers or their catalogs. + +The three runtime modules have different boundaries: `server/runtime.py` initializes Twinkle and DeviceMesh; `resources/runtime.py` owns Ray/Serve effects; `sampler/runtime.py` tracks request execution inside a sampling instance. + +## Sampling ownership and state reconstruction + +Stable sampling uses shared access; configuration and full-weight updates use exclusive access. Every participating instance is reserved before an update. Idle sweep checks the same reservations under the state lock. The lock order is configuration then state; startup is never awaited while holding the configuration lock. + +Pool admission is bounded by `queue_timeout`. A complete control broadcast or full-weight update uses one `effective_execution_timeout`, including completion confirmation. An ambiguous result fences all participants. Cancellation alone never confirms device release; shutdown failure keeps the instance quarantined and counted against capacity. + +Only successful broadcasts update replay state. Superseded templates, adapter configurations (per owner/name), and full checkpoints are compacted after the last `apply_patch`. Arbitrary patches preserve order and form a compaction barrier: dependencies before a patch must remain available. LoRA weight caches are loaded per request and eviction is not replayed. + +A client object rejects creation of an adapter name it already owns. Release it before creating a replacement; separate objects may use the same label. + +## Tests + +Unit tests follow the component they exercise. Reusable fake drivers and queues live in helper modules, not in imported `test_*.py` modules. Resource allocation probes live under `tests/server/integration/resources`, sampling probes under `tests/server/integration/sampler`, and training acceptance tools under `tests/server/integration/on_demand_training`. + +Control fixtures demonstrate scheduling and release mechanisms, not real model execution. Hardware acceptance must separately record the tested commit, runtime import origins, real model operations, numerical checks and confirmed cleanup. The package DAG and lightweight import checks live in `tests/server/test_dependencies.py`. diff --git a/src/twinkle/server/gateway/routes.py b/src/twinkle/server/gateway/routes.py index 0cececc05..96a8df495 100644 --- a/src/twinkle/server/gateway/routes.py +++ b/src/twinkle/server/gateway/routes.py @@ -1,26 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Route-path constants shared by the gateway's proxy and its OpenAI bridge. +"""Sampler endpoint paths used by the OpenAI bridge.""" -Single source for the ``route_prefix`` convention these three places must agree on; they -previously agreed by having the same string typed out in three files. Placed on -the consumer side (gateway) rather than in ``launcher`` -- launcher *produces* the -``route_prefix``, gateway consumes it -- to avoid a new ``gateway -> launcher`` edge. -""" -from __future__ import annotations - -# Downstream endpoint paths the gateway proxies to (relative to the service's route prefix). TWINKLE_SAMPLE = 'twinkle/sample' TWINKLE_SAMPLE_STREAM = 'twinkle/sample_stream' TWINKLE_SET_TEMPLATE = 'twinkle/set_template' - -TINKER_PREFIX = 'tinker' - - -def tinker_endpoint(endpoint: str) -> str: - """``tinker/`` -- the shape ``proxy_to_model`` / ``proxy_to_sampler`` build.""" - return f'{TINKER_PREFIX}/{endpoint}' - - -def target_url(route_prefix: str, service_type: str, base_model: str, endpoint: str) -> str: - """The one definition of ``{route_prefix}/{service_type}/{base_model}/{endpoint}``.""" - return f'{route_prefix.rstrip("/")}/{service_type}/{base_model}/{endpoint}' diff --git a/src/twinkle/server/sampler/actor.py b/src/twinkle/server/sampler/actor.py new file mode 100644 index 000000000..393e85343 --- /dev/null +++ b/src/twinkle/server/sampler/actor.py @@ -0,0 +1,30 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Ray composition for the sampling controller and its owned driver.""" +import asyncio + +from twinkle.server.config.application_spec import SamplerArgs +from .controller import SamplerController + + +def create_controller_actor(): + import ray + + from .driver import SamplerDriver + + @ray.remote(num_cpus=0.1, max_restarts=-1, max_task_retries=0) + class SamplerControllerActor(SamplerController): + + def __init__(self, values, namespace, pool_id, options, budgets): + args = SamplerArgs.model_validate(values) + super().__init__(args, SamplerDriver(args, namespace, pool_id, options, budgets), budgets) + + async def initialize(self): + await super().initialize() + if not self._stopping and self._monitor_task is None: + self._monitor_task = asyncio.create_task(self._monitor()) + + async def shutdown(self): + await super().shutdown() + await self.driver.runtime.wait_for_cleanup() + + return SamplerControllerActor diff --git a/src/twinkle/server/sampler/controller.py b/src/twinkle/server/sampler/controller.py index 7ba3f5d73..c2f20bb66 100644 --- a/src/twinkle/server/sampler/controller.py +++ b/src/twinkle/server/sampler/controller.py @@ -357,27 +357,3 @@ async def shutdown(self): await asyncio.gather(*(self._close(instance) for instance in list(self.instances.values()))) self._guard.close() await self.lifecycle.join() - - -def create_controller_actor(): - import ray - - from .driver import SamplerDriver - - @ray.remote(num_cpus=0.1, max_restarts=-1, max_task_retries=0) - class SamplerControllerActor(SamplerController): - - def __init__(self, values, namespace, pool_id, options, budgets): - args = SamplerArgs.model_validate(values) - super().__init__(args, SamplerDriver(args, namespace, pool_id, options, budgets), budgets) - - async def initialize(self): - await super().initialize() - if not self._stopping and self._monitor_task is None: - self._monitor_task = asyncio.create_task(self._monitor()) - - async def shutdown(self): - await super().shutdown() - await self.driver.runtime.wait_for_cleanup() - - return SamplerControllerActor diff --git a/src/twinkle/server/sampler/router.py b/src/twinkle/server/sampler/router.py index 3342c46a7..9e767346d 100644 --- a/src/twinkle/server/sampler/router.py +++ b/src/twinkle/server/sampler/router.py @@ -13,7 +13,8 @@ from twinkle.server.deployment import bind_deployment, build_deployment_app from twinkle.server.middleware.auth import get_token_from_request from twinkle.server.state import get_server_state -from .controller import create_controller_actor, sampler_controller_name, sampling_pool_id +from .actor import create_controller_actor +from .controller import sampler_controller_name, sampling_pool_id from .weights import is_full_checkpoint diff --git a/src/twinkle/server/telemetry/training_metrics.py b/src/twinkle/server/telemetry/training_metrics.py new file mode 100644 index 000000000..b6b7fd63a --- /dev/null +++ b/src/twinkle/server/telemetry/training_metrics.py @@ -0,0 +1,40 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Aggregate plain instance snapshots without depending on training state or scheduling.""" + + +class TrainingInstanceMetrics: + + def __init__(self, metrics): + self.metrics = metrics + self._instance_keys = set() + self._resource_keys = set() + + def report(self, snapshots): + counts, resources = {}, {} + for tags, device_type, device_count, group_count in snapshots: + key = tuple(tags[name] for name in ('backend', 'mode', 'state', 'resident')) + counts[key] = counts.get(key, 0) + 1 + resource_key = tags['backend'], tags['mode'], device_type + devices, groups = resources.get(resource_key, (0, 0)) + resources[resource_key] = devices + device_count, groups + group_count + for key in self._instance_keys | counts.keys(): + tags = dict(zip(('backend', 'mode', 'state', 'resident'), key)) + self.metrics.instances_active.set(counts.get(key, 0), tags=tags) + for key in self._resource_keys | resources.keys(): + tags = dict(zip(('backend', 'mode', 'device_type'), key)) + devices, groups = resources.get(key, (0, 0)) + self.metrics.devices_allocated.set(devices, tags=tags) + self.metrics.placement_groups_active.set(groups, tags=tags) + self._instance_keys = set(counts) + self._resource_keys = set(resources) + + +def quarantine_reason(exc: Exception) -> str: + message = str(exc).lower() + if isinstance(exc, TimeoutError): + return 'timeout' + if 'worker death' in message: + return 'worker_death' + if 'placement group' in message: + return 'pg_removal' + return 'other' diff --git a/src/twinkle/server/training/controller.py b/src/twinkle/server/training/controller.py index 3abbae71c..fabbf2eef 100644 --- a/src/twinkle/server/training/controller.py +++ b/src/twinkle/server/training/controller.py @@ -21,6 +21,7 @@ TrainingRecord) from twinkle.server.state.training_store import TrainingStore from twinkle.server.telemetry.metrics import get_training_metrics +from twinkle.server.telemetry.training_metrics import TrainingInstanceMetrics, quarantine_reason from .catalog import TrainingCatalog, TrainingSpec from .errors import TrainingExpiredError @@ -93,8 +94,7 @@ def __init__(self, self._lock = asyncio.Lock() self._ready = False self._metrics = get_training_metrics() - self._reported_metric_keys: set[tuple[str, str, str, str]] = set() - self._reported_resource_metric_keys: set[tuple[str, str, str]] = set() + self._instance_metrics = TrainingInstanceMetrics(self._metrics) def _instance_metric_tags(self, instance: InstanceRecord) -> dict[str, str]: spec = self.catalog.get(instance.base_model, instance.mode) @@ -106,39 +106,10 @@ def _instance_metric_tags(self, instance: InstanceRecord) -> dict[str, str]: } def _report_instance_metrics(self) -> None: - counts: dict[tuple[str, str, str, str], int] = {} - resources: dict[tuple[str, str], tuple[int, int]] = {} - for instance in self.instances.values(): - if instance.state == 'STOPPED': - continue - tags = self._instance_metric_tags(instance) - key = tuple(tags[name] for name in ('backend', 'mode', 'state', 'resident')) - counts[key] = counts.get(key, 0) + 1 - resource_key = tags['backend'], tags['mode'], self.catalog.get(instance.base_model, - instance.mode).model.device_type - devices, groups = resources.get(resource_key, (0, 0)) - resources[resource_key] = (devices + len(instance.devices), groups + int(bool(instance.application_name))) - for key in self._reported_metric_keys | counts.keys(): - tags = dict(zip(('backend', 'mode', 'state', 'resident'), key)) - self._metrics.instances_active.set(counts.get(key, 0), tags=tags) - for key in self._reported_resource_metric_keys | resources.keys(): - tags = dict(zip(('backend', 'mode', 'device_type'), key)) - devices, groups = resources.get(key, (0, 0)) - self._metrics.devices_allocated.set(devices, tags=tags) - self._metrics.placement_groups_active.set(groups, tags=tags) - self._reported_metric_keys = set(counts) - self._reported_resource_metric_keys = set(resources) - - @staticmethod - def _quarantine_reason(exc: Exception) -> str: - message = str(exc).lower() - if isinstance(exc, TimeoutError): - return 'timeout' - if 'worker death' in message: - return 'worker_death' - if 'placement group' in message: - return 'pg_removal' - return 'other' + self._instance_metrics.report((self._instance_metric_tags(instance), + self.catalog.get(instance.base_model, instance.mode).model.device_type, + len(instance.devices), int(bool(instance.application_name))) + for instance in self.instances.values() if instance.state != 'STOPPED') async def initialize(self) -> None: """Discard old state only after owned compute has been confirmed dead.""" @@ -618,7 +589,7 @@ async def quarantined(exc): self._metrics.quarantined_total.inc(tags={ 'backend': tags['backend'], 'mode': tags['mode'], - 'reason': self._quarantine_reason(exc), + 'reason': quarantine_reason(exc), }) self._metrics.shutdown_seconds.observe(time.monotonic() - started, tags={**timing_tags, 'status': 'error'}) diff --git a/tests/server/utils/__init__.py b/tests/server/checkpoint/__init__.py similarity index 100% rename from tests/server/utils/__init__.py rename to tests/server/checkpoint/__init__.py diff --git a/tests/server/utils/test_checkpoint_manager.py b/tests/server/checkpoint/test_checkpoint_manager.py similarity index 100% rename from tests/server/utils/test_checkpoint_manager.py rename to tests/server/checkpoint/test_checkpoint_manager.py diff --git a/tests/server/utils/task_queue/test_config.py b/tests/server/config/test_task_queue_config.py similarity index 100% rename from tests/server/utils/task_queue/test_config.py rename to tests/server/config/test_task_queue_config.py diff --git a/tests/server/utils/task_queue/__init__.py b/tests/server/integration/resources/__init__.py similarity index 100% rename from tests/server/utils/task_queue/__init__.py rename to tests/server/integration/resources/__init__.py diff --git a/tests/server/integration/on_demand_training/allocation_probe.py b/tests/server/integration/resources/allocation_probe.py similarity index 100% rename from tests/server/integration/on_demand_training/allocation_probe.py rename to tests/server/integration/resources/allocation_probe.py diff --git a/tests/server/integration/on_demand_training/fragmentation_probe.py b/tests/server/integration/resources/fragmentation_probe.py similarity index 100% rename from tests/server/integration/on_demand_training/fragmentation_probe.py rename to tests/server/integration/resources/fragmentation_probe.py diff --git a/tests/server/training/test_allocation.py b/tests/server/integration/resources/test_allocation.py similarity index 84% rename from tests/server/training/test_allocation.py rename to tests/server/integration/resources/test_allocation.py index 519347b8c..8ebe04d1c 100644 --- a/tests/server/training/test_allocation.py +++ b/tests/server/integration/resources/test_allocation.py @@ -5,7 +5,7 @@ def test_production_allocator_uses_disjoint_devices_and_reuses_released_group(): - result = subprocess.run([sys.executable, str(Path(__file__).resolve().parents[3] / 'tests/server/integration/on_demand_training/allocation_probe.py')], + result = subprocess.run([sys.executable, str(Path(__file__).with_name('allocation_probe.py'))], capture_output=True, text=True, timeout=120) assert result.returncode == 0, result.stdout + result.stderr evidence = json.loads(result.stdout.strip().splitlines()[-1]) @@ -14,7 +14,7 @@ def test_production_allocator_uses_disjoint_devices_and_reuses_released_group(): def test_fragmented_devices_cannot_be_combined_across_nodes(): - result = subprocess.run([sys.executable, str(Path(__file__).resolve().parents[3] / 'tests/server/integration/on_demand_training/fragmentation_probe.py')], + result = subprocess.run([sys.executable, str(Path(__file__).with_name('fragmentation_probe.py'))], capture_output=True, text=True, timeout=120) assert result.returncode == 0, result.stdout + result.stderr evidence = json.loads(result.stdout.strip().splitlines()[-1]) diff --git a/tests/server/integration/sampler/__init__.py b/tests/server/integration/sampler/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/server/integration/on_demand_training/sampling_probe.py b/tests/server/integration/sampler/sampling_probe.py similarity index 100% rename from tests/server/integration/on_demand_training/sampling_probe.py rename to tests/server/integration/sampler/sampling_probe.py diff --git a/tests/server/model/test_full_mode_exclusivity.py b/tests/server/model/test_full_mode_exclusivity.py index d49238fca..c7a68fedb 100644 --- a/tests/server/model/test_full_mode_exclusivity.py +++ b/tests/server/model/test_full_mode_exclusivity.py @@ -11,7 +11,7 @@ from starlette.requests import Request from tests.server.state.fake_backend import FakeBackend -from tests.server.training.test_controller import request, setup +from tests.server.training.helpers import request, setup from twinkle.server.config.training import TrainingConfig from twinkle.server.training.catalog import TrainingCatalog from twinkle.server.training.router import TrainingRouter diff --git a/tests/server/resources/__init__.py b/tests/server/resources/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/server/resources/test_instances.py b/tests/server/resources/test_instances.py index 42c78503f..97af22bb1 100644 --- a/tests/server/resources/test_instances.py +++ b/tests/server/resources/test_instances.py @@ -1,3 +1,4 @@ +from types import SimpleNamespace import logging import pytest from twinkle.server.resources.instances import InstanceLifecycle @@ -21,3 +22,63 @@ async def broken(): assert not lifecycle.jobs + + +@pytest.mark.asyncio +async def test_concurrent_close_shares_one_stop_and_retries_quarantine(): + import asyncio + lifecycle = InstanceLifecycle('test') + entered, resume = asyncio.Event(), asyncio.Event() + calls, events = [], [] + + async def stop(): + calls.append('stop') + entered.set() + await resume.wait() + if len(calls) == 1: + raise RuntimeError('death unconfirmed') + + async def stopped(): + events.append('stopped') + + async def quarantined(error): + events.append(str(error)) + + first = asyncio.create_task(lifecycle.close('one', stop, 1, stopped, quarantined)) + await entered.wait() + second = asyncio.create_task(lifecycle.close('one', stop, 1, stopped, quarantined)) + await asyncio.sleep(0) + resume.set() + await asyncio.gather(first, second) + assert calls == ['stop'] and events == ['death unconfirmed'] + await lifecycle.close('one', stop, 1, stopped, quarantined) + assert calls == ['stop', 'stop'] and events[-1] == 'stopped' + await lifecycle.join() + assert not lifecycle.jobs and not lifecycle.closing + +@pytest.mark.asyncio +async def test_shared_health_rule_tolerates_transient_probe_timeout(): + from twinkle.server.resources.instances import InstanceLifecycle + clock = SimpleNamespace(now=0) + lifecycle = InstanceLifecycle('health', clock=lambda: clock.now) + + async def unreachable(): + raise TimeoutError('control-plane probe timed out') + + async def healthy(): + return {'unhealthy': False} + + async def failed(): + return {'unhealthy': True} + + assert await lifecycle.health('instance', unreachable) is None + clock.now = 29 + assert await lifecycle.health('instance', unreachable) is None + clock.now = 30 + assert await lifecycle.health('instance', unreachable) is False + assert await lifecycle.health('instance', healthy) is True + assert not lifecycle.unreachable_since + assert await lifecycle.health('instance', unreachable) is None + assert await lifecycle.health('instance', failed) is False + + diff --git a/tests/server/sampler/test_pool.py b/tests/server/sampler/test_pool.py index 847c266c2..86eb8c03d 100644 --- a/tests/server/sampler/test_pool.py +++ b/tests/server/sampler/test_pool.py @@ -286,7 +286,7 @@ async def unknown(*args): @pytest.mark.asyncio async def test_training_and_sampling_use_the_same_lifecycle_owner(): from tests.server.state.fake_backend import FakeBackend - from tests.server.training.test_controller import setup as training_setup + from tests.server.training.helpers import setup as training_setup pool, driver, _ = setup() training, _, _, _ = await training_setup(backend=FakeBackend()) try: @@ -311,32 +311,6 @@ async def test_replaying_config_has_the_same_startup_deadline(): await pool.shutdown() -@pytest.mark.asyncio -async def test_shared_health_rule_tolerates_transient_probe_timeout(): - from twinkle.server.resources.instances import InstanceLifecycle - clock = SimpleNamespace(now=0) - lifecycle = InstanceLifecycle('health', clock=lambda: clock.now) - - async def unreachable(): - raise TimeoutError('control-plane probe timed out') - - async def healthy(): - return {'unhealthy': False} - - async def failed(): - return {'unhealthy': True} - - assert await lifecycle.health('instance', unreachable) is None - clock.now = 29 - assert await lifecycle.health('instance', unreachable) is None - clock.now = 30 - assert await lifecycle.health('instance', unreachable) is False - assert await lifecycle.health('instance', healthy) is True - assert not lifecycle.unreachable_since - assert await lifecycle.health('instance', unreachable) is None - assert await lifecycle.health('instance', failed) is False - - @pytest.mark.asyncio @pytest.mark.parametrize('kind', ['control', 'weights']) async def test_mutation_reserves_every_replica_against_idle_sweep(kind): diff --git a/tests/server/sampler/test_request_runtime.py b/tests/server/sampler/test_request_runtime.py index ef7c5809c..eac13d725 100644 --- a/tests/server/sampler/test_request_runtime.py +++ b/tests/server/sampler/test_request_runtime.py @@ -8,7 +8,7 @@ import pytest -from tests.server.utils.test_task_queue_mixin import _DummyQueue +from tests.server.task_queue.helpers import _DummyQueue from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.sampler.concurrency import SamplerConcurrency from twinkle.server.task_queue.types import ExecutionScope, execution_scope diff --git a/tests/server/training/test_sequence.py b/tests/server/state/test_sequence.py similarity index 97% rename from tests/server/training/test_sequence.py rename to tests/server/state/test_sequence.py index 924258bc6..053ef23ef 100644 --- a/tests/server/training/test_sequence.py +++ b/tests/server/state/test_sequence.py @@ -57,7 +57,7 @@ async def test_out_of_order_sequences_and_preflight_release_remain_retryable(sta @pytest.mark.asyncio async def test_retiring_a_handle_removes_native_and_tinker_stream_ledgers(state): - from tests.server.training.test_controller import create, setup + from tests.server.training.helpers import create, setup controller, state, _, _ = await setup(backend=state._backend) try: identity = await create(controller, state, 'dedup-retirement') diff --git a/tests/server/task_queue/__init__.py b/tests/server/task_queue/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/server/task_queue/helpers.py b/tests/server/task_queue/helpers.py new file mode 100644 index 000000000..4a29e2e4c --- /dev/null +++ b/tests/server/task_queue/helpers.py @@ -0,0 +1,60 @@ +from twinkle.server.config.task_queue import TaskQueueConfig +from twinkle.server.task_queue.mixin import TaskQueueMixin +from twinkle.server.task_queue.worker import ComputeWorker + +class _DummyState: + + def __init__(self): + self.records = [] + self._latest = {} + + async def store_future_status(self, *args, **kwargs): + self.records.append((args, kwargs)) + request_id, status = args[0], args[1] + self._latest[request_id] = { + 'status': status, + 'result': kwargs.get('result'), + 'failure': (kwargs['failure'].model_dump() if kwargs.get('failure') is not None else None), + 'queue_state': kwargs.get('queue_state'), + 'queue_state_reason': kwargs.get('queue_state_reason'), + } + + async def try_start_future(self, request_id): + record = self._latest.get(request_id) + if record is None or record['status'] not in ('pending', 'queued'): + return False + await self.store_future_status(request_id, 'running', None, queue_state='active') + return self._latest[request_id]['status'] == 'running' + + async def get_future(self, request_id): + return self._latest.get(request_id) + + + +class _AllowingRateLimiter: + + async def check_and_record(self, token, input_tokens): + return True, None + + + +class _DummyQueue(TaskQueueMixin): + + def __init__(self): + self.state = _DummyState() + # A generous Inline_Fast_Path window keeps "task settles inside submit" + # deterministic for the trivial in-process coroutines used here. + self._task_queue_config = TaskQueueConfig(inline_fast_path_timeout=5.0) + self._rate_limiter = _AllowingRateLimiter() + self._task_metrics = None + self._deployment_name = 'test' + + def enable_compute_worker(self): + self._compute_worker = ComputeWorker( + state=self.state, + config=self._task_queue_config, + task_metrics=None, + deployment_name=self._deployment_name, + ) + self._event_loop = None + self._queue_tasks = set() diff --git a/tests/server/utils/test_rate_limiter.py b/tests/server/task_queue/test_rate_limiter.py similarity index 100% rename from tests/server/utils/test_rate_limiter.py rename to tests/server/task_queue/test_rate_limiter.py diff --git a/tests/server/utils/test_task_queue_mixin.py b/tests/server/task_queue/test_task_queue_mixin.py similarity index 92% rename from tests/server/utils/test_task_queue_mixin.py rename to tests/server/task_queue/test_task_queue_mixin.py index 23ad79a80..cc049b02f 100644 --- a/tests/server/utils/test_task_queue_mixin.py +++ b/tests/server/task_queue/test_task_queue_mixin.py @@ -1,3 +1,4 @@ +from tests.server.task_queue.helpers import _DummyState, _AllowingRateLimiter, _DummyQueue import asyncio import pytest import pytest_asyncio @@ -48,60 +49,10 @@ def test_exception_summary_falls_back_to_the_ray_error_itself(): assert 'Model.step' in summary -class _DummyState: - def __init__(self): - self.records = [] - self._latest = {} - async def store_future_status(self, *args, **kwargs): - self.records.append((args, kwargs)) - request_id, status = args[0], args[1] - self._latest[request_id] = { - 'status': status, - 'result': kwargs.get('result'), - 'failure': (kwargs['failure'].model_dump() if kwargs.get('failure') is not None else None), - 'queue_state': kwargs.get('queue_state'), - 'queue_state_reason': kwargs.get('queue_state_reason'), - } - async def try_start_future(self, request_id): - record = self._latest.get(request_id) - if record is None or record['status'] not in ('pending', 'queued'): - return False - await self.store_future_status(request_id, 'running', None, queue_state='active') - return self._latest[request_id]['status'] == 'running' - async def get_future(self, request_id): - return self._latest.get(request_id) - - -class _AllowingRateLimiter: - - async def check_and_record(self, token, input_tokens): - return True, None - - -class _DummyQueue(TaskQueueMixin): - - def __init__(self): - self.state = _DummyState() - # A generous Inline_Fast_Path window keeps "task settles inside submit" - # deterministic for the trivial in-process coroutines used here. - self._task_queue_config = TaskQueueConfig(inline_fast_path_timeout=5.0) - self._rate_limiter = _AllowingRateLimiter() - self._task_metrics = None - self._deployment_name = 'test' - - def enable_compute_worker(self): - self._compute_worker = ComputeWorker( - state=self.state, - config=self._task_queue_config, - task_metrics=None, - deployment_name=self._deployment_name, - ) - self._event_loop = None - self._queue_tasks = set() @pytest.mark.asyncio diff --git a/tests/server/telemetry/test_training_metrics.py b/tests/server/telemetry/test_training_metrics.py index ad446cfeb..427db6a86 100644 --- a/tests/server/telemetry/test_training_metrics.py +++ b/tests/server/telemetry/test_training_metrics.py @@ -7,6 +7,7 @@ from twinkle.server.telemetry import metrics from twinkle.server.training.controller import TrainingController +from twinkle.server.telemetry.training_metrics import TrainingInstanceMetrics _ALLOWED_ATTRIBUTES = {'backend', 'mode', 'state', 'status', 'reason', 'device_type', 'resident'} @@ -106,8 +107,7 @@ def test_controller_metrics_clear_old_states_and_released_resources(monkeypatch, controller = TrainingController.__new__(TrainingController) controller.catalog = SimpleNamespace(get=lambda base_model, mode: spec) controller._metrics = registry.training_metrics() - controller._reported_metric_keys = set() - controller._reported_resource_metric_keys = set() + controller._instance_metrics = TrainingInstanceMetrics(controller._metrics) controller.instances = { str(index): SimpleNamespace( instance_id=f'instance_{index}', base_model='private-model', mode='lora', @@ -148,5 +148,3 @@ def totals(name): controller._report_instance_metrics() for name in (instances_name, devices_name, groups_name): assert all(value == 0 for value in totals(name).values()) - assert not controller._reported_metric_keys - assert not controller._reported_resource_metric_keys diff --git a/tests/server/test_dependencies.py b/tests/server/test_dependencies.py index 5150294bf..127f5a7d6 100644 --- a/tests/server/test_dependencies.py +++ b/tests/server/test_dependencies.py @@ -98,3 +98,14 @@ def test_server_components_form_declared_dag(): assert not violations, ( 'Server dependencies must point strictly down the declared DAG; ' f'invalid edges={violations}') + + +@pytest.mark.parametrize('component', ['training', 'sampler']) +def test_controller_import_does_not_load_ray_composition(component): + code = f'''import twinkle.server.{component}.controller +import sys +assert 'ray' not in sys.modules or not sys.modules['ray'].is_initialized() +assert not any('twinkle.server.{component}.' + name in sys.modules for name in ('actor', 'driver', 'app')) +''' + result = subprocess.run([sys.executable, '-c', code], capture_output=True, text=True, timeout=60) + assert result.returncode == 0, result.stderr diff --git a/tests/server/gateway/test_proxy.py b/tests/server/test_service_proxy.py similarity index 93% rename from tests/server/gateway/test_proxy.py rename to tests/server/test_service_proxy.py index caa35deee..e217bedfc 100644 --- a/tests/server/gateway/test_proxy.py +++ b/tests/server/test_service_proxy.py @@ -1,10 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Directed tests for ``gateway/proxy.py``. - -Covers the route-URL construction now sourced from ``gateway.routes``, the -``H_MULTIPLEX`` header compatibility in ``_prepare_headers``, and the 502 ``ErrorPayload`` -fallback when the upstream is unreachable. -""" +"""ServiceProxy URL construction, headers, bounded waits, and upstream error mapping.""" from __future__ import annotations import asyncio diff --git a/tests/server/training/helpers.py b/tests/server/training/helpers.py new file mode 100644 index 000000000..916d47f45 --- /dev/null +++ b/tests/server/training/helpers.py @@ -0,0 +1,126 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations +import asyncio +from twinkle.server.config.training import TrainingConfig +from twinkle.server.state import ServerState +from twinkle.server.state.backend.memory_backend import RayActorBackend +from twinkle.server.training.controller import TrainingController +from twinkle.server.resources.allocation import CapacityUnavailableError +from twinkle.server.state.training_records import CreateTraining +from types import SimpleNamespace +from twinkle.server.config.task_queue import TaskQueueConfig +from twinkle.server.task_queue import TaskQueueMixin +from twinkle.server.model.managed_runtime import ManagedTrainingRuntime + +class Clock: + now = 0.0 + + def __call__(self): + return self.now + + +class Runtime: + """Allocates a finite pool and stores actual per-user state for assertions.""" + + def __init__(self, devices=16): + self.devices = devices + self.allocations = {} + self.users = {} + self.starts = 0 + self.stopped = [] + self.start_gate = asyncio.Event() + self.start_gate.set() + self.fail_start = False + self.fail_shutdown = False + + async def allocate(self, instance, spec): + self.starts += 1 + if self.fail_start: + raise RuntimeError('injected startup failure') + used = sum(self.allocations.values()) + if used + spec.config.num_devices > self.devices: + raise CapacityUnavailableError('No devices available') + self.allocations[instance.instance_id] = spec.config.num_devices + return {} + + async def deploy(self, instance, spec): + await self.start_gate.wait() + + async def create_training(self, instance, training, request): + assert instance.instance_id in self.allocations + assert training.training_id not in self.users + self.users[training.training_id] = {'instance': instance.instance_id, 'label': training.adapter_label, 'steps': 0} + + async def release_training(self, instance, training): + self.users.pop(training.training_id, None) + + async def shutdown(self, instance): + if self.fail_shutdown: + raise RuntimeError('cannot confirm process death') + self.allocations.pop(instance.instance_id, None) + self.users = {key: value for key, value in self.users.items() if value['instance'] != instance.instance_id} + self.stopped.append(instance.instance_id) + + async def verify_operation(self, instance, request_id): + return {'state': 'pending'} + + async def recover_orphans(self): + pass + + +async def setup(*, devices=16, min_replicas=0, max_replicas=None, quota=100, backend=None): + mode = {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}, 'max_loras': 5, + 'idle_timeout_s': 10, 'min_replicas': min_replicas, 'max_replicas': max_replicas} + config = TrainingConfig(models=[{ + 'name': 'small', 'model_id': 'ms://small', 'backend': 'mock', 'device_type': 'cuda', + 'modes': {'lora': mode, 'full': {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}, 'idle_timeout_s': 10}}, + }]) + state = ServerState(backend=backend or RayActorBackend(), per_token_model_limit=quota) + await state.create_session({'session_id': 'session'}) + runtime, clock = Runtime(devices), Clock() + controller = TrainingController(config, state.create_training_store('unit'), runtime, clock=clock) + await controller.initialize() + return controller, state, runtime, clock + + +def request(key, *, mode='lora', label='same', token='owner'): + return CreateTraining(base_model='small', mode=mode, token=token, session_id='session', create_key=str(key), + client_type='twinkle', payload={}, adapter_label=label) + + +async def create(controller, state, key, **kwargs): + ref = await controller.create(request(key, **kwargs)) + await controller.wait_for_jobs() + future = await state.get_future(ref['request_id']) + assert future['status'] == 'completed', future + return future['result']['training_id'] + + +class RPC: + def __init__(self, method): + self.method = method + + def remote(self, *args, **kwargs): + return asyncio.create_task(self.method(*args, **kwargs)) + + +class RuntimeQueue(ManagedTrainingRuntime, TaskQueueMixin): + def __init__(self, controller, state, instance): + self.state, self.replica_id = state, instance.instance_id + self.training_context = {'instance_id': instance.instance_id, 'epoch': controller.epoch} + self._model_unhealthy = False + self._training_controller = SimpleNamespace(**{ + name: RPC(getattr(controller, name)) + for name in ('admit', 'finish_operation', 'fail_instance') + }) + self._init_task_queue(TaskQueueConfig(execution_timeout=0.1, inline_fast_path_timeout=0.01), + enable_admission_gate=True) + + +async def settle(state, request_id): + for _ in range(300): + future = await state.get_future(request_id) + if future and future['status'] in ('completed', 'failed', 'cancelled'): + return future + await asyncio.sleep(0.01) + raise AssertionError('Future did not settle') diff --git a/tests/server/training/test_controller.py b/tests/server/training/test_controller.py index a8a485f27..434d8e2eb 100644 --- a/tests/server/training/test_controller.py +++ b/tests/server/training/test_controller.py @@ -1,6 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Exercise real controller/storage decisions against a controllable runtime.""" from __future__ import annotations +from tests.server.training.helpers import Clock, Runtime, setup, request, create import asyncio @@ -16,62 +17,6 @@ from twinkle.server.state.training_records import CreateTraining -class Clock: - now = 0.0 - - def __call__(self): - return self.now - - -class Runtime: - """Allocates a finite pool and stores actual per-user state for assertions.""" - - def __init__(self, devices=16): - self.devices = devices - self.allocations = {} - self.users = {} - self.starts = 0 - self.stopped = [] - self.start_gate = asyncio.Event() - self.start_gate.set() - self.fail_start = False - self.fail_shutdown = False - - async def allocate(self, instance, spec): - self.starts += 1 - if self.fail_start: - raise RuntimeError('injected startup failure') - used = sum(self.allocations.values()) - if used + spec.config.num_devices > self.devices: - raise CapacityUnavailableError('No devices available') - self.allocations[instance.instance_id] = spec.config.num_devices - return {} - - async def deploy(self, instance, spec): - await self.start_gate.wait() - - async def create_training(self, instance, training, request): - assert instance.instance_id in self.allocations - assert training.training_id not in self.users - self.users[training.training_id] = {'instance': instance.instance_id, 'label': training.adapter_label, 'steps': 0} - - async def release_training(self, instance, training): - self.users.pop(training.training_id, None) - - async def shutdown(self, instance): - if self.fail_shutdown: - raise RuntimeError('cannot confirm process death') - self.allocations.pop(instance.instance_id, None) - self.users = {key: value for key, value in self.users.items() if value['instance'] != instance.instance_id} - self.stopped.append(instance.instance_id) - - async def verify_operation(self, instance, request_id): - return {'state': 'pending'} - - async def recover_orphans(self): - pass - - @pytest.mark.asyncio async def test_initialization_timeout_only_retires_the_failing_tenant(): from tests.server.state.fake_backend import FakeBackend @@ -243,34 +188,6 @@ async def test_legacy_session_and_max_lifetime_bounds_keep_their_original_meanin assert controller.trainings[second].state == 'EXPIRED' -async def setup(*, devices=16, min_replicas=0, max_replicas=None, quota=100, backend=None): - mode = {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}, 'max_loras': 5, - 'idle_timeout_s': 10, 'min_replicas': min_replicas, 'max_replicas': max_replicas} - config = TrainingConfig(models=[{ - 'name': 'small', 'model_id': 'ms://small', 'backend': 'mock', 'device_type': 'cuda', - 'modes': {'lora': mode, 'full': {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}, 'idle_timeout_s': 10}}, - }]) - state = ServerState(backend=backend or RayActorBackend(), per_token_model_limit=quota) - await state.create_session({'session_id': 'session'}) - runtime, clock = Runtime(devices), Clock() - controller = TrainingController(config, state.create_training_store('unit'), runtime, clock=clock) - await controller.initialize() - return controller, state, runtime, clock - - -def request(key, *, mode='lora', label='same', token='owner'): - return CreateTraining(base_model='small', mode=mode, token=token, session_id='session', create_key=str(key), - client_type='twinkle', payload={}, adapter_label=label) - - -async def create(controller, state, key, **kwargs): - ref = await controller.create(request(key, **kwargs)) - await controller.wait_for_jobs() - future = await state.get_future(ref['request_id']) - assert future['status'] == 'completed', future - return future['result']['training_id'] - - @pytest.mark.asyncio async def test_lazy_queries_allocate_nothing(): controller, state, runtime, _ = await setup() diff --git a/tests/server/training/test_failure_recovery.py b/tests/server/training/test_failure_recovery.py index f310f268e..36863fcfd 100644 --- a/tests/server/training/test_failure_recovery.py +++ b/tests/server/training/test_failure_recovery.py @@ -7,8 +7,8 @@ import pytest from tests.server.state.fake_backend import FakeBackend -from tests.server.training.test_controller import create, request, setup -from tests.server.training.test_runtime_queue import RPC, RuntimeQueue, settle +from tests.server.training.helpers import create, request, setup +from tests.server.training.helpers import RPC, RuntimeQueue, settle from twinkle.protocol.types.errors import ErrorCategory from twinkle.server.exceptions import ResourceExhaustedError, ResourceNotFoundError from twinkle.server.state.models import FAILURE_REASON_CODES, FutureFailureRecord diff --git a/tests/server/training/test_router.py b/tests/server/training/test_router.py index 37741d0b7..8f70a5893 100644 --- a/tests/server/training/test_router.py +++ b/tests/server/training/test_router.py @@ -7,7 +7,7 @@ from starlette.requests import Request from tests.server.state.fake_backend import FakeBackend -from tests.server.training.test_runtime_queue import settle +from tests.server.training.helpers import settle from twinkle.hub import HubOperation from twinkle.server.state import ServerState from twinkle.server.training.router import TrainingRouter diff --git a/tests/server/training/test_runtime_initialization.py b/tests/server/training/test_runtime_initialization.py index 27c8ec219..1c91f2aad 100644 --- a/tests/server/training/test_runtime_initialization.py +++ b/tests/server/training/test_runtime_initialization.py @@ -8,7 +8,7 @@ import pytest from tests.server.state.fake_backend import FakeBackend -from tests.server.training.test_controller import create, setup +from tests.server.training.helpers import create, setup from twinkle.server.model.managed_runtime import ManagedTrainingRuntime from twinkle.server.state.training_records import CreateTraining diff --git a/tests/server/training/test_runtime_queue.py b/tests/server/training/test_runtime_queue.py index 71ffcd2d6..664876b86 100644 --- a/tests/server/training/test_runtime_queue.py +++ b/tests/server/training/test_runtime_queue.py @@ -3,43 +3,14 @@ import asyncio import threading from types import SimpleNamespace +from tests.server.training.helpers import RPC, RuntimeQueue, settle import pytest from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.task_queue import TaskQueueMixin from twinkle.server.model.managed_runtime import ManagedTrainingRuntime, _tinker_operation -from tests.server.training.test_controller import create, setup - - -class RPC: - def __init__(self, method): - self.method = method - - def remote(self, *args, **kwargs): - return asyncio.create_task(self.method(*args, **kwargs)) - - -class RuntimeQueue(ManagedTrainingRuntime, TaskQueueMixin): - def __init__(self, controller, state, instance): - self.state, self.replica_id = state, instance.instance_id - self.training_context = {'instance_id': instance.instance_id, 'epoch': controller.epoch} - self._model_unhealthy = False - self._training_controller = SimpleNamespace(**{ - name: RPC(getattr(controller, name)) - for name in ('admit', 'finish_operation', 'fail_instance') - }) - self._init_task_queue(TaskQueueConfig(execution_timeout=0.1, inline_fast_path_timeout=0.01), - enable_admission_gate=True) - - -async def settle(state, request_id): - for _ in range(300): - future = await state.get_future(request_id) - if future and future['status'] in ('completed', 'failed', 'cancelled'): - return future - await asyncio.sleep(0.01) - raise AssertionError('Future did not settle') +from tests.server.training.helpers import create, setup @pytest.mark.asyncio diff --git a/tests/server/training/test_storage_recovery.py b/tests/server/training/test_storage_recovery.py index c2529a6ca..63b2cddd2 100644 --- a/tests/server/training/test_storage_recovery.py +++ b/tests/server/training/test_storage_recovery.py @@ -8,7 +8,7 @@ from twinkle.server.state.backend.memory_backend import RayActorBackend from twinkle.server.state.backend.redis_backend import RedisBackend from twinkle.server.training.controller import TrainingController -from tests.server.training.test_controller import create, request, setup +from tests.server.training.helpers import create, request, setup @pytest.mark.asyncio From 857db4aefcedf0e2969133d46e84f8b873c5cacc Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Thu, 22 Oct 2026 21:56:43 +0800 Subject: [PATCH 13/46] test: exercise real sampler pool recovery and checkpoint rebuilding --- .../on_demand_training/matrix_validation.py | 10 +- .../integration/sampler/test_real_pool_e2e.py | 326 ++++++++++++++++++ 2 files changed, 335 insertions(+), 1 deletion(-) create mode 100644 tests/server/integration/sampler/test_real_pool_e2e.py diff --git a/tests/server/integration/on_demand_training/matrix_validation.py b/tests/server/integration/on_demand_training/matrix_validation.py index 9112a3dff..d4c65226d 100644 --- a/tests/server/integration/on_demand_training/matrix_validation.py +++ b/tests/server/integration/on_demand_training/matrix_validation.py @@ -349,6 +349,13 @@ def main(): elif args.full_cycle_only: filenames = ['test_full_cycle_e2e.py'] for filename in filenames: + collection = subprocess.run([ + sys.executable, '-m', 'pytest', '--noconftest', '--collect-only', '-q', + str(root / 'tests/server/integration' / filename), + ], capture_output=True, text=True, env=os.environ.copy(), timeout=180) + (run_dir / (filename + '.collect.log')).write_text(collection.stdout + collection.stderr) + expected = sum('::test_' in line for line in collection.stdout.splitlines()) + assert collection.returncode == 0 and expected > 0, collection.stdout + collection.stderr start = time.monotonic() with (run_dir / (filename + '.log')).open('w') as output: result = subprocess.run([ @@ -377,7 +384,8 @@ def main(): }) (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) assert result.returncode == 0, f'{filename} failed; see {run_dir / (filename + ".log")}' - assert len(cases) == (1 if filename == 'test_full_cycle_e2e.py' else 2) and skipped == 0 + assert len(cases) == expected and skipped == 0 + assert not any(case.find('failure') is not None or case.find('error') is not None for case in cases) base_url = f'http://127.0.0.1:{port}' model_name = config.training_config.models[0].name evidence['recovery'] = [] diff --git a/tests/server/integration/sampler/test_real_pool_e2e.py b/tests/server/integration/sampler/test_real_pool_e2e.py new file mode 100644 index 000000000..8fc2acca1 --- /dev/null +++ b/tests/server/integration/sampler/test_real_pool_e2e.py @@ -0,0 +1,326 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Real vLLM pool acceptance. Run only on four dedicated devices with the PPU SDK loaded.""" +import asyncio +import json +import math +import os +import shutil +import subprocess +import time +import uuid +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.skipif(os.environ.get('TWINKLE_TEST_REAL_POOL') != '1', + reason='Requires four dedicated devices and a real cached model') + + +@pytest.fixture(scope='session', autouse=True) +def _ray_runtime(): + yield + + +@pytest.fixture(autouse=True) +def _reset_canonical_state_actor(): + yield + + +def payload(endpoint, body): + return {'endpoint': endpoint, 'method': 'POST', 'body': json.dumps(body).encode(), + 'headers': [(b'content-type', b'application/json'), (b'x-request-id', uuid.uuid4().hex.encode()), + (b'authorization', b'Bearer pool-acceptance')], 'query_string': b''} + + +def checkpoint_alias(model, target): + target.mkdir() + for source in model.iterdir(): + if source.is_file(): + (target / source.name).symlink_to(source) + return str(target) + + +def make_lora(model, target): + import torch + from peft import LoraConfig + from safetensors import safe_open + from safetensors.torch import save_file + rank, weights = 8, {} + generator = torch.Generator().manual_seed(312) + for shard in model.glob('*.safetensors'): + with safe_open(str(shard), framework='pt', device='cpu') as source: + for name in source.keys(): + if name.endswith('.self_attn.q_proj.weight'): + output_size, input_size = source.get_slice(name).get_shape() + prefix = 'base_model.model.' + name[:-len('.weight')] + weights[prefix + '.lora_A.weight'] = torch.randn(rank, input_size, generator=generator) * .05 + weights[prefix + '.lora_B.weight'] = torch.randn(output_size, rank, generator=generator) * .05 + assert weights, 'No q_proj weights found; this is not a skipped acceptance case' + target.mkdir() + LoraConfig(r=rank, lora_alpha=rank, target_modules=['q_proj'], bias='none', task_type='CAUSAL_LM', + inference_mode=True, base_model_name_or_path=str(model)).save_pretrained(target) + save_file(weights, str(target / 'adapter_model.safetensors')) + return str(target) + + +async def exercise(model, directory, namespace, evidence): + import numpy as np + import ray + from ray import serve + from ray.util import list_named_actors + from ray.util.placement_group import placement_group_table + from transformers import AutoTokenizer + from twinkle.server.config.application_spec import SamplerArgs + from twinkle.server.sampler.controller import SamplerController + from twinkle.server.sampler.driver import SamplerDriver + from twinkle.server.state import get_server_state + + args = SamplerArgs(model_id=str(model), sampler_type='vllm', + device_group={'name': 'sampler', 'ranks': 1, 'device_type': 'cuda'}, + device_mesh={'device_type': 'cuda', 'dp_size': 1}, + engine_args={'max_model_len': 2048, 'gpu_memory_utilization': .4, 'enable_lora': True, + 'enforce_eager': True, 'logprobs_mode': 'processed_logprobs', + 'attention_config': {'flash_attn_version': 3}}, + min_replicas=2, max_replicas=4, target_ongoing_requests=1, idle_timeout_s=2, + queue_config={'execution_timeout': 600, 'queue_timeout': 600}) + budgets = {'allocation_timeout_s': 60, 'startup_timeout_s': 900, + 'shutdown_timeout_s': 180, 'sweep_interval_s': 1} + driver = SamplerDriver(args, namespace, 'real-pool', {}, budgets) + pool = SamplerController(args, driver, budgets) + pools = [pool] + state = get_server_state() + tokenizer = AutoTokenizer.from_pretrained(str(model), local_files_only=True) + ids = tokenizer.apply_chat_template([{'role': 'user', 'content': 'Count the integers from one to one hundred.'}], + tokenize=True, add_generation_prompt=True, enable_thinking=False, + return_dict=False) + calls = [] + invoke = driver.invoke + + async def observed(instance, ticket, body, queue=None): + calls.append({'instance': instance.identity, 'endpoint': body['endpoint'], 'devices': instance.devices}) + return await invoke(instance, ticket, body, queue) + + driver.invoke = observed + + async def sample(target=pool, path=None, full=None, max_tokens=64): + body = {'inputs': [{'input_ids': ids}], + 'sampling_params': {'max_tokens': max_tokens, 'temperature': 0, 'logprobs': 1, 'ignore_eos': True}} + if path: + body['adapter_uri'] = path + reply = await target.request(payload('twinkle/sample', body), full_path=full) + assert reply['status'] == 200, reply + result = json.loads(reply['body']) + deadline = time.monotonic() + 600 + while result['status'] not in ('completed', 'failed', 'cancelled'): + assert time.monotonic() < deadline, result + await asyncio.sleep(.05) + result = await state.get_future(result['request_id']) + assert result['status'] == 'completed', result + seq = result['result']['samples'][0]['sequences'][0] + assert seq['tokens'] and len(seq['tokens']) == len(seq['logprobs']) + assert all(math.isfinite(position[0][1]) for position in seq['logprobs']) + return seq + + async def settled(target=pool): + await asyncio.gather(*tuple(target.lifecycle.startups.values())) + await target.lifecycle.join() + + def save(): + (directory / 'results.json').write_text(json.dumps(evidence, indent=2)) + + try: + await pool.initialize() + evidence['sizes'] = [len(pool._ready_instances())] + from twinkle.patch import Patch + from twinkle.protocol.serialize import serialize_object + + class VerifyImportOrigin(Patch): + def __call__(self, module, expected_root): + import twinkle + from pathlib import Path + assert Path(twinkle.__file__).resolve().is_relative_to(expected_root), twinkle.__file__ + return module + + source_root = str(Path(__file__).resolve().parents[4]) + origin_reply = await pool.request(payload('twinkle/apply_patch', { + 'patch_cls': serialize_object(VerifyImportOrigin), 'adapter_name': '', + 'init_kwargs': {'expected_root': source_root}})) + assert origin_reply['status'] == 200, origin_reply + evidence['worker_import_origin_verified'] = source_root + save() + await asyncio.gather(*(sample(max_tokens=256) for _ in range(8))) + await settled() + assert len(pool._ready_instances()) == 4 + devices = [device for instance in pool._ready_instances() for device in instance.devices] + assert len(devices) == len(set(devices)) == 4 + calls.clear() + await asyncio.gather(*(sample(max_tokens=128) for _ in range(8))) + await settled() + assert {row['instance'] for row in calls} == {i.identity for i in pool._ready_instances()} + for instance in pool._ready_instances(): + assert not (await driver.health(instance))['unhealthy'] + evidence.update(device_ids=devices, request_placements=list(calls)) + evidence['sizes'].append(4) + save() + + # Pause the ACK after a real backend update, then sweep across the idle window. + entered, resume = asyncio.Event(), asyncio.Event() + original_wait = driver.wait_operation + + async def delayed_ack(instance, ticket): + await original_wait(instance, ticket) + entered.set() + await resume.wait() + + driver.wait_operation = delayed_ack + update = asyncio.create_task(pool.request(payload('twinkle/set_template', {'template_cls': 'Qwen3_5Template'}))) + try: + await asyncio.wait_for(entered.wait(), 120) + await asyncio.sleep(2.1) + await pool.sweep() + assert len(pool._ready_instances()) == 4 + finally: + resume.set() + await update + driver.wait_operation = original_wait + evidence['control_sweep_safe'] = True + + old = checkpoint_alias(model, directory / 'checkpoint-A') + new = checkpoint_alias(model, directory / 'checkpoint-B') + await sample(full=old) + await settled() + original_load = driver.load_weights + entered.clear(); resume.clear() + + async def delayed_load(instance, path): + await original_load(instance, path) + entered.set() + await resume.wait() + + driver.load_weights = delayed_load + update = asyncio.create_task(sample(full=new)) + try: + await asyncio.wait_for(entered.wait(), 180) + await asyncio.sleep(2.1) + await pool.sweep() + assert len(pool._ready_instances()) == 4 + finally: + resume.set() + await update + driver.load_weights = original_load + await settled() + shutil.rmtree(old) # Only the test-owned symlink directory; never the model cache. + evidence['full_update_sweep_safe'] = True + await asyncio.sleep(2.1) + await pool.sweep() + await settled() + evidence['sizes'].append(len(pool._ready_instances())) + assert evidence['sizes'] == [2, 4, 2] + save() + await asyncio.gather(*(sample(max_tokens=128) for _ in range(8))) + await settled() + assert len(pool._ready_instances()) == 4 + evidence['superseded_checkpoint_rebuild'] = True + + lora = make_lora(model, directory / 'lora') + before = await sample(path=lora) + await settled() + response = await pool.request(payload('twinkle/unload_adapter_paths', {'adapter_paths': [lora]})) + assert response['status'] == 200 + after = await sample(path=lora) + await settled() + assert before['tokens'] == after['tokens'] + np.testing.assert_allclose([x[0][1] for x in before['logprobs']], + [x[0][1] for x in after['logprobs']], rtol=0, atol=1e-5) + evidence['lora_eviction_reload'] = True + save() + + # Lose a real control ACK; all four real participants must be stopped. + entered.clear() + async def missing_ack(instance, ticket): + await original_wait(instance, ticket) + entered.set() + await asyncio.Event().wait() + driver.wait_operation = missing_ack + args.queue_config.execution_timeout = 5 + try: + with pytest.raises(TimeoutError): + await pool.request(payload('twinkle/set_template', {'template_cls': 'Qwen3_5Template'})) + assert entered.is_set() + assert not pool._ready_instances() and pool._owned_count() == 0 + finally: + driver.wait_operation = original_wait + args.queue_config.execution_timeout = 600 + await pool.sweep() + await settled() + await sample(full=new) + await settled() + evidence['control_timeout_recovered'] = True + save() + + victim = pool._ready_instances()[0] + workers = [row for row in list_named_actors(all_namespaces=True) + if row['namespace'] == namespace and victim.identity in row['name'] + and row['name'].startswith('sampler-') and not row['name'].endswith('-registry')] + assert workers, 'No owned sampler worker for fault injection' + ray.kill(ray.get_actor(workers[0]['name'], namespace=namespace), no_restart=True) + deadline = time.monotonic() + 60 + while not (await driver.health(victim))['unhealthy']: + assert time.monotonic() < deadline + await asyncio.sleep(.2) + await pool._close(victim) + await pool.sweep() + await settled() + await sample(full=new) + await settled() + evidence['worker_fault_recovered'] = True + save() + + # The new controller's production recovery finds and retires the old pool's resources. + old_names = {i.application for i in pool._ready_instances()} + replacement_driver = SamplerDriver(args, namespace, 'real-pool', {}, budgets) + replacement = SamplerController(args, replacement_driver, budgets) + pools.append(replacement) + await replacement.initialize() + assert replacement.epoch != pool.epoch + assert not old_names.intersection(serve.status().applications) + await sample(target=replacement) + await settled(replacement) + evidence['controller_orphan_recovery'] = True + save() + finally: + for target in reversed(pools): + await target.shutdown() + await target.driver.runtime.wait_for_cleanup() + assert not any(name.startswith(driver.prefix) for name in serve.status().applications) + assert not any(row.get('name', '').startswith(driver.prefix) and row['state'] != 'REMOVED' + for row in placement_group_table().values()) + assert not any(row['namespace'] == namespace and 'sampler_' in row['name'] + for row in list_named_actors(all_namespaces=True)) + evidence['confirmed_release'] = True + save() + evidence['passed'] = True + save() + + +def test_real_pool_lifecycle(): + import ray + import twinkle + from ray import serve + model = Path(os.environ['TWINKLE_TEST_MODEL_PATH']) + directory = Path(os.environ['TWINKLE_TEST_ARTIFACT_DIR']) + directory.mkdir(parents=True, exist_ok=True) + namespace = 'pool_acceptance_' + uuid.uuid4().hex[:8] + os.environ.update(TWINKLE_STATE_RAY_NAMESPACE=namespace, TWINKLE_PERSISTENCE_MODE='memory', + TWINKLE_GPU_PG_CPUS_PER_PROC='1', TWINKLE_TRUST_REMOTE_CODE='1', + TWINKLE_DEFAULT_SAVE_DIR=str(directory / 'checkpoints')) + evidence = {'namespace': namespace, 'model': str(model), 'import': twinkle.__file__, + 'sha': subprocess.check_output(['git', 'rev-parse', 'HEAD'], text=True).strip()} + ray.init(address='local', namespace=namespace, num_cpus=24, num_gpus=4, + _temp_dir='/tmp/twp-' + namespace[-8:], include_dashboard=False, dashboard_port=0) + try: + serve.start(http_options={'host': '127.0.0.1', 'port': 0}) + asyncio.run(exercise(model, directory, namespace, evidence)) + finally: + serve.shutdown() + ray.shutdown() From c2d7394ab480433abb72e977d92dbe7e087fc079 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Thu, 22 Oct 2026 22:19:49 +0800 Subject: [PATCH 14/46] test: verify worker imports through Ray actor diagnostics --- .../integration/sampler/test_real_pool_e2e.py | 39 +++++++++++-------- tests/server/resources/test_instances.py | 2 - 2 files changed, 23 insertions(+), 18 deletions(-) diff --git a/tests/server/integration/sampler/test_real_pool_e2e.py b/tests/server/integration/sampler/test_real_pool_e2e.py index 8fc2acca1..180df47cb 100644 --- a/tests/server/integration/sampler/test_real_pool_e2e.py +++ b/tests/server/integration/sampler/test_real_pool_e2e.py @@ -126,27 +126,31 @@ async def settled(target=pool): await target.lifecycle.join() def save(): + snapshots = evidence.setdefault('device_snapshots', []) + name = f'devices-stage-{len(snapshots):02d}.txt' + (directory / name).write_text(subprocess.check_output(['ppu-smi'], text=True)) + snapshots.append(name) (directory / 'results.json').write_text(json.dumps(evidence, indent=2)) try: await pool.initialize() evidence['sizes'] = [len(pool._ready_instances())] - from twinkle.patch import Patch - from twinkle.protocol.serialize import serialize_object - - class VerifyImportOrigin(Patch): - def __call__(self, module, expected_root): - import twinkle - from pathlib import Path - assert Path(twinkle.__file__).resolve().is_relative_to(expected_root), twinkle.__file__ - return module - source_root = str(Path(__file__).resolve().parents[4]) - origin_reply = await pool.request(payload('twinkle/apply_patch', { - 'patch_cls': serialize_object(VerifyImportOrigin), 'adapter_name': '', - 'init_kwargs': {'expected_root': source_root}})) - assert origin_reply['status'] == 200, origin_reply - evidence['worker_import_origin_verified'] = source_root + + async def verify_worker_origins(): + workers = [row for row in list_named_actors(all_namespaces=True) + if row['namespace'] == namespace and row['name'].startswith('sampler-') + and not row['name'].endswith('-registry')] + assert len(workers) == len(pool._ready_instances()) + origins = {} + for row in workers: + worker = ray.get_actor(row['name'], namespace=namespace) + origin = await worker.__ray_call__.remote(lambda _: __import__('twinkle').__file__) + assert Path(origin).resolve().is_relative_to(source_root), origin + origins[row['name']] = origin + evidence['worker_import_origins'] = origins + + await verify_worker_origins() save() await asyncio.gather(*(sample(max_tokens=256) for _ in range(8))) await settled() @@ -154,8 +158,10 @@ def __call__(self, module, expected_root): devices = [device for instance in pool._ready_instances() for device in instance.devices] assert len(devices) == len(set(devices)) == 4 calls.clear() - await asyncio.gather(*(sample(max_tokens=128) for _ in range(8))) + samples = await asyncio.gather(*(sample(max_tokens=128) for _ in range(8))) await settled() + await verify_worker_origins() + evidence['sample_result'] = samples[0] assert {row['instance'] for row in calls} == {i.identity for i in pool._ready_instances()} for instance in pool._ready_instances(): assert not (await driver.health(instance))['unhealthy'] @@ -232,6 +238,7 @@ async def delayed_load(instance, path): assert before['tokens'] == after['tokens'] np.testing.assert_allclose([x[0][1] for x in before['logprobs']], [x[0][1] for x in after['logprobs']], rtol=0, atol=1e-5) + evidence['lora_results'] = {'before_eviction': before, 'after_reload': after} evidence['lora_eviction_reload'] = True save() diff --git a/tests/server/resources/test_instances.py b/tests/server/resources/test_instances.py index 97af22bb1..2870d7ec5 100644 --- a/tests/server/resources/test_instances.py +++ b/tests/server/resources/test_instances.py @@ -80,5 +80,3 @@ async def failed(): assert not lifecycle.unreachable_since assert await lifecycle.health('instance', unreachable) is None assert await lifecycle.health('instance', failed) is False - - From 6432c31cbacfbd3705881d8a9f9e65e3c8ec3d81 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Thu, 22 Oct 2026 22:44:29 +0800 Subject: [PATCH 15/46] test: preserve sampling request identity while polling futures --- tests/server/integration/sampler/test_real_pool_e2e.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/server/integration/sampler/test_real_pool_e2e.py b/tests/server/integration/sampler/test_real_pool_e2e.py index 180df47cb..8b7345919 100644 --- a/tests/server/integration/sampler/test_real_pool_e2e.py +++ b/tests/server/integration/sampler/test_real_pool_e2e.py @@ -7,6 +7,7 @@ import shutil import subprocess import time +import traceback import uuid from pathlib import Path @@ -110,11 +111,12 @@ async def sample(target=pool, path=None, full=None, max_tokens=64): reply = await target.request(payload('twinkle/sample', body), full_path=full) assert reply['status'] == 200, reply result = json.loads(reply['body']) + request_id = result['request_id'] deadline = time.monotonic() + 600 while result['status'] not in ('completed', 'failed', 'cancelled'): assert time.monotonic() < deadline, result await asyncio.sleep(.05) - result = await state.get_future(result['request_id']) + result = await state.get_future(request_id) assert result['status'] == 'completed', result seq = result['result']['samples'][0]['sequences'][0] assert seq['tokens'] and len(seq['tokens']) == len(seq['logprobs']) @@ -295,6 +297,10 @@ async def missing_ack(instance, ticket): await settled(replacement) evidence['controller_orphan_recovery'] = True save() + except BaseException: + evidence['failure'] = traceback.format_exc() + save() + raise finally: for target in reversed(pools): await target.shutdown() From bbf11bc82615a6d83721087eafd784caa4a23a31 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Thu, 22 Oct 2026 22:56:19 +0800 Subject: [PATCH 16/46] test: verify checkpoint logits and optimizer continuation numerically --- .../on_demand_training/acceptance_support.py | 14 +++++-- .../on_demand_training/matrix_validation.py | 24 +++++++++-- .../server/integration/test_full_cycle_e2e.py | 41 ++++++++----------- .../training/test_acceptance_harness.py | 7 ++++ 4 files changed, 54 insertions(+), 32 deletions(-) diff --git a/tests/server/integration/on_demand_training/acceptance_support.py b/tests/server/integration/on_demand_training/acceptance_support.py index f3ae8f5e7..b23274602 100644 --- a/tests/server/integration/on_demand_training/acceptance_support.py +++ b/tests/server/integration/on_demand_training/acceptance_support.py @@ -8,6 +8,15 @@ from tests.server.integration.on_demand_training.on_demand_sdk import unload_tinker_model +def checked_logps(values): + """Reject empty/default metric values before numerical checkpoint comparisons.""" + import numpy as np + values = np.asarray(values, dtype=float).reshape(-1) + assert values.size and np.isfinite(values).all(), 'Empty or non-finite checkpoint probe' + assert (values < 0).any(), 'Checkpoint probe contains no nontrivial token log probabilities' + return values + + class TrainingHandle: """A small fixed-token workload with optimizer checkpoint support.""" @@ -77,10 +86,7 @@ def evaluate(self): else: result = self.client.forward(self.datums(), 'cross_entropy').result() values = [output['logprobs'].tolist() for output in result.loss_fn_outputs] - import numpy as np - values = np.asarray(values, dtype=float).reshape(-1) - assert values.size and np.isfinite(values).all(), 'Non-finite checkpoint probe' - return values + return checked_logps(values) def save(self, name): if self.protocol == 'twinkle': diff --git a/tests/server/integration/on_demand_training/matrix_validation.py b/tests/server/integration/on_demand_training/matrix_validation.py index d4c65226d..0b2883bc0 100644 --- a/tests/server/integration/on_demand_training/matrix_validation.py +++ b/tests/server/integration/on_demand_training/matrix_validation.py @@ -37,8 +37,16 @@ def validate_controller_restart(base_url, controller, model_name): snapshots = [] try: for handle in old: + # Two calls end the native GA=2 workload on an optimizer boundary. handle.step() - snapshots.append((handle.save('restart-state'), handle.evaluate(), handle.consumed_samples)) + handle.step() + path = handle.save('restart-state') + values, consumed = handle.evaluate(), handle.consumed_samples + handle.step() + handle.step() + continued = handle.evaluate() + assert np.max(np.abs(continued - values)) > 2e-3, 'The reference optimizer did not change logprobs' + snapshots.append((path, values, consumed, continued)) old_epoch = ray.get(controller.diagnostics.remote())['epoch'] ray.kill(controller, no_restart=False) @@ -57,7 +65,8 @@ def restarted(): try: # Keep restored peers on one runtime: checkpoint/optimizer isolation # is checked without reloading the base model for each protocol. - for protocol, (path, values, consumed) in zip(('twinkle', 'tinker'), snapshots): + optimizer_errors = {} + for protocol, (path, values, consumed, continued) in zip(('twinkle', 'tinker'), snapshots): fresh = TrainingHandle(base_url, model_name, protocol, 'restored-' + protocol) restored.append(fresh) progress = fresh.resume(path) @@ -65,6 +74,14 @@ def restarted(): assert progress['consumed_train_samples'] == consumed np.testing.assert_allclose(fresh.evaluate(), values, rtol=0, atol=2e-3) fresh.step() + fresh.step() + resumed = fresh.evaluate() + np.testing.assert_allclose(resumed, continued, rtol=0, atol=2e-3) + optimizer_errors[protocol] = { + 'max_error': float(np.max(np.abs(resumed - continued))), + 'checkpoint_logps': values.tolist(), 'continued_logps': continued.tolist(), + 'resumed_logps': resumed.tolist(), + } active = next(row for row in ray.get(controller.diagnostics.remote())['instances'] if restored[0].identity in row['training_ids']) assert set(active['training_ids']) == {handle.identity for handle in restored} @@ -90,7 +107,8 @@ def restarted(): finally: for handle in old: handle.close(stale=True) - return {'epoch': diagnostic['epoch'], 'worker_fault_recovered': True, 'checkpoint_restored': True} + return {'epoch': diagnostic['epoch'], 'worker_fault_recovered': True, 'checkpoint_restored': True, + 'optimizer_continuation': optimizer_errors} def validate_idempotent_create(base_url, controller, model_name): diff --git a/tests/server/integration/test_full_cycle_e2e.py b/tests/server/integration/test_full_cycle_e2e.py index 12d17de9c..ec7a4c52d 100644 --- a/tests/server/integration/test_full_cycle_e2e.py +++ b/tests/server/integration/test_full_cycle_e2e.py @@ -11,7 +11,7 @@ Phase A — initial training (STEPS_PHASE_A steps) Phase B — keep training STEPS_PHASE_B more steps, save again Phase C — RELOAD VERIFY: load() ckpt_a, forward_only on fixed batch, - assert loss matches Phase A end + assert token log probabilities match Phase A end Phase D — RESUME VERIFY: resume_from_checkpoint(ckpt_b), train more Phase E + F — vLLM LoRA-effect greedy probe @@ -63,7 +63,7 @@ STEPS_PHASE_A = 60 STEPS_PHASE_B = 4 STEPS_PHASE_D = 4 -RELOAD_LOSS_TOLERANCE = 0.05 # |reloaded - original| / original +RELOAD_LOGPROB_ATOL = 2e-3 # Same absolute tolerance as controller-restart acceptance. RESUME_LOSS_BAND = 3.0 # resumed step's loss must be within this factor of Phase-B's last SAMPLE_MAX_TOKENS = 32 @@ -124,21 +124,13 @@ def _train_n_steps(model, dataloader, n: int, *, label: str, start_step: int = 0 return losses -def _record_fixed_batch_loss(model, batch, *, label: str) -> float: - """Run forward_only on a fixed batch and report the loss for reload comparison. - - Uses ``forward_only`` + ``calculate_metric(is_training=False)`` to retrieve - loss. This is compatible with both Transformers and Megatron backends - (Megatron does not support standalone ``calculate_loss``). - """ - model.forward_only(inputs=batch) - metric = model.calculate_metric(is_training=False) - try: - val = float(metric.result.get('loss')) if hasattr(metric.result, 'get') else float(metric.result['loss']) - except (TypeError, KeyError): - val = float(metric.result) - logger.info(f'[{label}] fixed-batch loss = {val:.4f}') - return val +def _record_fixed_batch_logps(model, batch, *, label: str): + """Read actual forward outputs; forward_only need not populate loss metrics.""" + from tests.server.integration.on_demand_training.acceptance_support import checked_logps + values = checked_logps(model.forward_only(inputs=batch).result['logps']) + logger.info('[%s] fixed-batch logps: count=%d min=%.6f max=%.6f', + label, values.size, values.min(), values.max()) + return values # Probe prompts for vLLM LoRA-effect verification. Mixing different prompt @@ -216,7 +208,7 @@ def main() -> int: # Capture a fixed batch we'll re-use after reload to assert state survived. fixed_batch = next(iter(_build_dataset_loader())) - loss_a_fixed = _record_fixed_batch_loss(model_a, fixed_batch, label='A-fixed') + logps_a_fixed = _record_fixed_batch_logps(model_a, fixed_batch, label='A-fixed') ckpt_a_resp = model_a.save( name='e2e-phase-a', @@ -248,11 +240,11 @@ def main() -> int: logger.info('Phase C: reload-verify (load ckpt_a, fixed batch)') logger.info('=' * 60) model_a.load(ckpt_a) - loss_c_fixed = _record_fixed_batch_loss(model_a, fixed_batch, label='C-fixed') - delta = abs(loss_c_fixed - loss_a_fixed) / max(abs(loss_a_fixed), 1e-6) - logger.info(f'Phase C reload delta: |{loss_c_fixed:.4f} - {loss_a_fixed:.4f}| / {loss_a_fixed:.4f} = {delta:.4f}') - assert delta <= RELOAD_LOSS_TOLERANCE, ( - f'Phase C FAILED: reload delta {delta:.4f} > tolerance {RELOAD_LOSS_TOLERANCE}') + import numpy as np + logps_c_fixed = _record_fixed_batch_logps(model_a, fixed_batch, label='C-fixed') + np.testing.assert_allclose(logps_c_fixed, logps_a_fixed, rtol=0, atol=RELOAD_LOGPROB_ATOL) + delta = float(np.max(np.abs(logps_c_fixed - logps_a_fixed))) + logger.info('Phase C reload max logprob error: %.6f', delta) # --------------------------------------------------------------------- # Phase D — RESUME VERIFY @@ -306,8 +298,7 @@ def main() -> int: logger.info('=' * 60) logger.info(' Phase A losses (%d steps): first=%.3f last=%.3f', len(losses_a), losses_a[0][1], losses_a[-1][1]) logger.info(' Phase B losses (%d steps): %s', len(losses_b), [f'{l:.3f}' for _, l in losses_b]) - logger.info(' Phase C reload: |%.4f - %.4f| / %.4f = %.4f (tol %.2f)', loss_c_fixed, loss_a_fixed, - loss_a_fixed, delta, RELOAD_LOSS_TOLERANCE) + logger.info(' Phase C reload max logprob error: %.6f (atol %.4f)', delta, RELOAD_LOGPROB_ATOL) logger.info(' Phase D resume losses (%d steps): %s', len(losses_d), [f'{l:.3f}' for _, l in losses_d]) logger.info(' Phase F LoRA-effect probes (%d/%d diverged):', n_diverged, len(PROBE_PROMPTS)) for prompt, _, _, div in probe_results: diff --git a/tests/server/training/test_acceptance_harness.py b/tests/server/training/test_acceptance_harness.py index cdbb11d27..a24ef6d5d 100644 --- a/tests/server/training/test_acceptance_harness.py +++ b/tests/server/training/test_acceptance_harness.py @@ -110,3 +110,10 @@ def test_training_release_assertion_excludes_but_does_not_remove_sampling(monkey groups['training'] = {'name': training_name, 'state': 'CREATED'} with pytest.raises(AssertionError): assert_released(controller, namespace) + + +@pytest.mark.parametrize('values', [[], [0.0, 0.0], [float('nan')], [float('-inf')]]) +def test_checkpoint_probe_rejects_vacuous_comparisons(values): + from tests.server.integration.on_demand_training.acceptance_support import checked_logps + with pytest.raises(AssertionError): + checked_logps(values) From 17cdb4fd09915793ce226a44a2d6dc74f0ae1838 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Thu, 22 Oct 2026 23:30:19 +0800 Subject: [PATCH 17/46] test: compare variable-length checkpoint log probabilities --- .../integration/on_demand_training/acceptance_support.py | 3 ++- tests/server/training/test_acceptance_harness.py | 5 +++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/server/integration/on_demand_training/acceptance_support.py b/tests/server/integration/on_demand_training/acceptance_support.py index b23274602..c14827c6c 100644 --- a/tests/server/integration/on_demand_training/acceptance_support.py +++ b/tests/server/integration/on_demand_training/acceptance_support.py @@ -11,7 +11,8 @@ def checked_logps(values): """Reject empty/default metric values before numerical checkpoint comparisons.""" import numpy as np - values = np.asarray(values, dtype=float).reshape(-1) + rows = [np.asarray(row, dtype=float).reshape(-1) for row in values] + values = np.concatenate(rows) if rows else np.array([], dtype=float) assert values.size and np.isfinite(values).all(), 'Empty or non-finite checkpoint probe' assert (values < 0).any(), 'Checkpoint probe contains no nontrivial token log probabilities' return values diff --git a/tests/server/training/test_acceptance_harness.py b/tests/server/training/test_acceptance_harness.py index a24ef6d5d..58d46e7de 100644 --- a/tests/server/training/test_acceptance_harness.py +++ b/tests/server/training/test_acceptance_harness.py @@ -117,3 +117,8 @@ def test_checkpoint_probe_rejects_vacuous_comparisons(values): from tests.server.integration.on_demand_training.acceptance_support import checked_logps with pytest.raises(AssertionError): checked_logps(values) + + +def test_checkpoint_probe_accepts_variable_length_sequences(): + from tests.server.integration.on_demand_training.acceptance_support import checked_logps + assert checked_logps([[-0.25, -0.5], [-0.75]]).tolist() == [-0.25, -0.5, -0.75] From a8e0f1e3be37ac720497be3b205ca16a046c30de Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 01:58:46 +0800 Subject: [PATCH 18/46] fix: fence late training release acknowledgements after instance failure --- src/twinkle/server/training/controller.py | 5 + .../test_resident_release_e2e.py | 189 ++++++++++++++++++ .../server/training/test_failure_recovery.py | 68 +++++++ 3 files changed, 262 insertions(+) create mode 100644 tests/server/integration/on_demand_training/test_resident_release_e2e.py diff --git a/src/twinkle/server/training/controller.py b/src/twinkle/server/training/controller.py index fabbf2eef..379bb28a4 100644 --- a/src/twinkle/server/training/controller.py +++ b/src/twinkle/server/training/controller.py @@ -542,6 +542,11 @@ async def _cleanup_user(self, training: TrainingRecord) -> None: diagnostic=traceback.format_exc())) return async with self._lock: + # A health failure may have taken ownership while release/reset + # was in flight. Its late ACK cannot retire quota or revive READY. + if (instance.instance_id in self._failed_instances or instance.state not in ('READY', 'RESETTING') + or training.state != 'CLOSING'): + return terminal = training.model_copy( deep=True, update={ diff --git a/tests/server/integration/on_demand_training/test_resident_release_e2e.py b/tests/server/integration/on_demand_training/test_resident_release_e2e.py new file mode 100644 index 000000000..dfd3e3c4d --- /dev/null +++ b/tests/server/integration/on_demand_training/test_resident_release_e2e.py @@ -0,0 +1,189 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Real resident reset ACK racing worker death; needs two dedicated accelerators.""" +import asyncio +import json +import os +import socket +import subprocess +import sys +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import numpy as np +import pytest + +pytestmark = pytest.mark.skipif(os.environ.get('TWINKLE_TEST_RESIDENT_RELEASE') != '1', + reason='Requires two dedicated devices and a real cached model') + + +@pytest.fixture(scope='session', autouse=True) +def _ray_runtime(): + yield + + +@pytest.fixture(autouse=True) +def _reset_canonical_state_actor(): + yield + + +def release_gate(actor, action): + """Delay only the ACK after the real model reset; never substitute computation.""" + if action == 'install': + entered, resume = asyncio.Event(), asyncio.Event() + original = actor.driver.release_training + + async def delayed(instance, training): + await original(instance, training) + entered.set() + await resume.wait() + + actor.driver.release_training = delayed + actor._release_test_gate = (entered, resume, original) + elif action == 'entered': + return actor._release_test_gate[0].is_set() + elif action == 'resume': + actor._release_test_gate[1].set() + actor.driver.release_training = actor._release_test_gate[2] + + +def worker_origin(actor): + import twinkle + from twinkle.server.training import controller + return {'python': sys.executable, 'twinkle': twinkle.__file__, 'controller': controller.__file__, + 'visible_devices': os.environ.get('CUDA_VISIBLE_DEVICES')} + + +def test_resident_release_ack_after_worker_failure(tmp_path): + import ray + from ray import serve + from ray.util import list_named_actors + from twinkle.server.config import ServerConfig + from twinkle.server.launcher.server_launcher import ServerLauncher + from tests.server.integration.on_demand_training.acceptance_support import TrainingHandle, wait_until, assert_released + from tests.server.integration.on_demand_training.managed_redis import ManagedRedis + + root = Path(__file__).resolve().parents[4] + directory = Path(os.environ.get('TWINKLE_TEST_RUN_DIR', str(tmp_path))) + directory.mkdir(parents=True, exist_ok=True) + backend = os.environ['TWINKLE_TEST_BACKEND'] + model = os.environ['TWINKLE_TEST_MODEL'] + namespace = 'resident_' + uuid.uuid4().hex[:8] + persistence = 'redis' if backend == 'transformers' else 'memory' + managed = ManagedRedis(directory / 'redis') if persistence == 'redis' else None + if managed: + managed.start() + prefix = namespace + '::' + with socket.socket() as listener: + listener.bind(('127.0.0.1', 0)) + port = listener.getsockname()[1] + model_name = 'Qwen/Qwen3.5-4B' + raw = { + 'ray_namespace': namespace, 'http_options': {'host': '127.0.0.1', 'port': port}, + 'training': {'startup_timeout_s': 1800, 'shutdown_timeout_s': 180, 'sweep_interval_s': 1}, + 'persistence': ({'mode': persistence, 'redis_url': managed.url, 'key_prefix': prefix} + if managed else {'mode': persistence}), + 'applications': [ + {'name': 'server', 'route_prefix': '/api/v1', 'import_path': 'server', + 'args': {'supported_models': [model_name]}}, + {'name': 'model', 'route_prefix': '/api/v1/model/' + model_name, 'import_path': 'model', + 'args': {'model_id': model, 'backend': backend, 'device_type': 'cuda', + 'modes': {'full': {'num_devices': 2, 'min_replicas': 1, 'max_replicas': 1, + 'device_mesh': {'dp_size': 2 if backend == 'transformers' else 1, + 'pp_size': 1 if backend == 'transformers' else 2}, + 'max_length': 512, 'idle_timeout_s': 3600, + 'queue_config': {'execution_timeout': 600}}}, + 'env_vars': {'TWINKLE_TRUST_REMOTE_CODE': '1', 'TWINKLE_FAIL_FAST': '0'}}}, + ], + } + (directory / 'server.json').write_text(json.dumps(raw, indent=2)) + os.environ.update(TWINKLE_STATE_RAY_NAMESPACE=namespace, TWINKLE_PERSISTENCE_MODE=persistence, + TWINKLE_DEFAULT_SAVE_DIR=str(directory / 'checkpoints'), TWINKLE_TRUST_REMOTE_CODE='1', + TWINKLE_FAIL_FAST='0', TWINKLE_MODEL_ID_ALIASES=json.dumps({model_name: model}), + PYTHONPATH=os.pathsep.join((str(root / 'src'), str(root)))) + if managed: + os.environ.update(TWINKLE_PERSISTENCE_REDIS_URL=managed.url, TWINKLE_PERSISTENCE_KEY_PREFIX=prefix) + else: + os.environ.pop('TWINKLE_PERSISTENCE_REDIS_URL', None) + os.environ.pop('TWINKLE_PERSISTENCE_KEY_PREFIX', None) + evidence = {'backend': backend, 'namespace': namespace, + 'sha': subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=root, text=True).strip()} + ray.init(address='local', namespace=namespace, num_cpus=24, num_gpus=2, + _temp_dir='/tmp/twr-' + namespace[-8:], include_dashboard=False) + launcher = ServerLauncher(ServerConfig.model_validate(raw)) + handles, controller = [], None + try: + launcher._start_serve() + for app in launcher.config.applications: + launcher._deploy_application(app) + launcher._deploy_training() + controller = launcher._training_controller + url = f'http://127.0.0.1:{port}' + old = TrainingHandle(url, model_name, 'twinkle', 'before-fault', mode='full') + handles.append(old) + baseline = old.evaluate() + old.step() + old.step() + trained = old.evaluate() + assert np.max(np.abs(trained - baseline)) > 2e-3, 'Training did not update real model logprobs' + diagnostic = ray.get(controller.diagnostics.remote()) + instance = next(row for row in diagnostic['instances'] if old.identity in row['training_ids']) + owned = [row for row in list_named_actors(all_namespaces=True) + if row['namespace'] == namespace and instance['instance_id'] in row['name'] + and row['name'].startswith('model-') and not row['name'].endswith('-registry')] + assert len(owned) == 2 + evidence['worker_origins'] = [ray.get(ray.get_actor(row['name'], namespace=namespace).__ray_call__.remote( + worker_origin)) for row in owned] + for origin in evidence['worker_origins']: + assert Path(origin['twinkle']).is_relative_to(root / 'src') + assert origin['python'] == sys.executable + ray.get(controller.__ray_call__.remote(release_gate, 'install')) + with ThreadPoolExecutor(max_workers=1) as executor: + releasing = executor.submit(old.close) + try: + wait_until(lambda: ray.get(controller.__ray_call__.remote(release_gate, 'entered')), + 600, 'Real reset did not finish') + ray.kill(ray.get_actor(owned[0]['name'], namespace=namespace), no_restart=True) + + def stopped(): + rows = ray.get(controller.diagnostics.remote())['instances'] + return next((row for row in rows if row['instance_id'] == instance['instance_id'] + and row['state'] == 'STOPPED'), None) + + wait_until(stopped, 600, 'Health monitor did not retire the failed resident') + finally: + ray.get(controller.__ray_call__.remote(release_gate, 'resume')) + releasing.result(timeout=600) + assert stopped(), 'A late reset ACK resurrected the dead resident' + fresh = TrainingHandle(url, model_name, 'tinker', 'after-fault', mode='full') + handles.append(fresh) + restored = fresh.evaluate() + np.testing.assert_allclose(restored, baseline, rtol=0, atol=2e-3) + fresh.step() + continued = fresh.evaluate() + assert np.max(np.abs(continued - restored)) > 2e-3 + fresh.close() + assert stopped() + rows = ray.get(controller.diagnostics.remote())['instances'] + assert sum(row['state'] == 'READY' for row in rows) == 1 + evidence.update(late_ack_fenced=True, recovered=True, baseline_logps=baseline.tolist(), + trained_logps=trained.tolist(), restored_logps=restored.tolist(), + continued_logps=continued.tolist(), instances=rows) + finally: + if controller is not None: + ray.get(controller.__ray_call__.remote( + lambda actor: release_gate(actor, 'resume') if hasattr(actor, '_release_test_gate') else None)) + for handle in handles: + handle.close(stale=True) + try: + if controller is not None: + ray.get(controller.shutdown.remote(), timeout=240) + assert_released(controller, namespace, managed.url if managed else None, prefix) + evidence['released'] = True + launcher._shutdown() + finally: + serve.shutdown() + ray.shutdown() + if managed: + managed.stop() + (directory / 'result.json').write_text(json.dumps(evidence, indent=2)) diff --git a/tests/server/training/test_failure_recovery.py b/tests/server/training/test_failure_recovery.py index 36863fcfd..f27b9c955 100644 --- a/tests/server/training/test_failure_recovery.py +++ b/tests/server/training/test_failure_recovery.py @@ -345,3 +345,71 @@ async def test_capacity_unregister_does_not_retire_managed_model(): assert await state.get_model_metadata(identity) is None assert await state._backend.get('token_count::owner') == 0 await controller.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('mode', ['lora', 'full']) +@pytest.mark.parametrize('failure_stage', ['stopped', 'quarantined', 'persist_draining']) +async def test_late_release_ack_cannot_retire_a_failed_instance_user(mode, failure_stage, monkeypatch): + from twinkle.server.config.training import TrainingConfig + from twinkle.server.training.catalog import TrainingCatalog + + controller, state, runtime, _ = await setup(backend=FakeBackend()) + config = controller.config.model_dump() + config['models'][0]['modes'][mode].update(min_replicas=1, max_replicas=1) + controller.config = TrainingConfig.model_validate(config) + controller.catalog = TrainingCatalog(controller.config) + await controller.sweep() + await controller.wait_for_jobs() + identity = await create(controller, state, 'late-release', mode=mode) + training = controller.trainings[identity] + instance = controller.instances[training.instance_id] + entered, resume = asyncio.Event(), asyncio.Event() + release = runtime.release_training + + async def delayed_ack(target, user): + await release(target, user) + entered.set() + await resume.wait() + + runtime.release_training = delayed_ack + training.state, training.terminal_reason = 'CLOSING', 'released' + await controller.store.save_training(training) + cleanup = asyncio.create_task(controller._cleanup_user(training)) + save_instance = controller.store.save_instance + + async def unavailable(record): + if record.state == 'DRAINING': + raise ConnectionError('draining write unavailable') + await save_instance(record) + + try: + await asyncio.wait_for(entered.wait(), 1) + runtime.fail_shutdown = failure_stage == 'quarantined' + failure = FutureFailureRecord(reason_code='orphaned_replica', attribution='server', message='Worker died') + if failure_stage == 'persist_draining': + monkeypatch.setattr(controller.store, 'save_instance', unavailable) + with pytest.raises(ConnectionError): + await controller.fail_instance(instance.instance_id, failure) + else: + await controller.fail_instance(instance.instance_id, failure) + state_before_ack = instance.state + resume.set() + await asyncio.wait_for(cleanup, 1) + assert instance.state == state_before_ack + if failure_stage == 'stopped': + assert training.state == 'FAILED' + assert not runtime.allocations + else: + assert training.state == 'CLOSING' + assert identity in instance.training_ids + assert await state.get_model_metadata(identity) is not None + assert await state._backend.get('token_count::owner') == 1 + with pytest.raises(ResourceNotFoundError): + await controller.route(identity, 'owner') + finally: + resume.set() + await asyncio.gather(cleanup, return_exceptions=True) + runtime.fail_shutdown = False + monkeypatch.setattr(controller.store, 'save_instance', save_instance) + await controller.shutdown() From 5708b87814abd31e8ed0cc09e53c50af43641bc0 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 02:12:59 +0800 Subject: [PATCH 19/46] fix: retain future ownership for the lifetime of its result --- src/twinkle/server/state/future_manager.py | 4 +++ src/twinkle/server/state/models.py | 1 + src/twinkle/server/state/server_state.py | 8 ++--- src/twinkle/server/state/training_store.py | 15 ++++++--- src/twinkle/server/task_queue/mixin.py | 4 ++- src/twinkle/server/task_queue/worker.py | 3 +- src/twinkle/server/training/actor.py | 4 +-- src/twinkle/server/training/controller.py | 5 ++- src/twinkle/server/training/router.py | 6 ++-- tests/server/state/test_future_lifecycle.py | 31 +++++++++++++++++++ tests/server/state/test_update_atomic.py | 16 ++++++++++ .../task_queue/test_task_queue_mixin.py | 30 ++++++++++++++++++ tests/server/training/test_router.py | 2 +- 13 files changed, 108 insertions(+), 21 deletions(-) diff --git a/src/twinkle/server/state/future_manager.py b/src/twinkle/server/state/future_manager.py index d3aaae9b9..fbfe316e7 100644 --- a/src/twinkle/server/state/future_manager.py +++ b/src/twinkle/server/state/future_manager.py @@ -31,6 +31,7 @@ def _future_record_transform( replica_id: str | None, absolute_deadline: float | None, now: str, + owner_token: str | None = None, ) -> dict | None: """Atomic transform body for :meth:`FutureManager.store_status`. @@ -56,6 +57,7 @@ def _future_record_transform( if existing is None: record = FutureRecord( + owner_token=owner_token, status=new_status, model_id=model_id, reason=reason, @@ -140,6 +142,7 @@ async def store_status( queue_state_reason: str | None = None, replica_id: str | None = None, absolute_deadline: float | None = None, + owner_token: str | None = None, ) -> None: """Create or update a future record with the latest status. @@ -176,6 +179,7 @@ async def store_status( queue_state_reason=queue_state_reason, replica_id=replica_id, absolute_deadline=absolute_deadline, + owner_token=owner_token, now=now, ), ) diff --git a/src/twinkle/server/state/models.py b/src/twinkle/server/state/models.py index ced4475e1..875699dea 100644 --- a/src/twinkle/server/state/models.py +++ b/src/twinkle/server/state/models.py @@ -89,6 +89,7 @@ class FutureFailureRecord(BaseModel): class FutureRecord(BaseModel): """Represents an async task future / request status.""" + owner_token: str | None = None status: str model_id: str | None = None reason: str | None = None diff --git a/src/twinkle/server/state/server_state.py b/src/twinkle/server/state/server_state.py index dc9d683f2..d38bd8398 100644 --- a/src/twinkle/server/state/server_state.py +++ b/src/twinkle/server/state/server_state.py @@ -183,9 +183,6 @@ async def save_model_record(self, model_id: str, record: ModelRecord) -> None: async def get_training_records(self, namespace: str): return await self._model_mgr.get_training_records(namespace) - async def set_future_owner(self, request_id: str, token: str) -> None: - await self._backend.set(f'future_owner::{request_id}', token, ttl=max(1, int(self.expiration_timeout))) - async def delete_training_sequences(self, training_id: str) -> None: for key in await self._backend.keys(f'seq_ledger::{training_id}::*'): await self._backend.delete(key) @@ -460,7 +457,8 @@ async def get_future(self, request_id: str) -> dict[str, Any] | None: return record.model_dump() if record is not None else None async def assert_future_owner(self, request_id: str, token: str) -> None: - owner = await self._backend.get(f'future_owner::{request_id}') + record = await self._future_mgr.get(request_id) + owner = record.owner_token if record is not None else None if owner is not None and owner != token: raise ResourceNotFoundError('Unknown future') @@ -526,6 +524,7 @@ async def store_future_status( queue_state_reason: str | None = None, replica_id: str | None = None, absolute_deadline: float | None = None, + owner_token: str | None = None, ) -> None: """Store task status with either a success result or domain failure. @@ -557,6 +556,7 @@ async def store_future_status( queue_state_reason=queue_state_reason, replica_id=replica_id, absolute_deadline=absolute_deadline, + owner_token=owner_token, ) # ----- Resource Cleanup ----- diff --git a/src/twinkle/server/state/training_store.py b/src/twinkle/server/state/training_store.py index c814e6a73..4e6a3ca9f 100644 --- a/src/twinkle/server/state/training_store.py +++ b/src/twinkle/server/state/training_store.py @@ -82,12 +82,17 @@ async def future(self, epoch: str | None = None, result: Any = None, failure: FutureFailureRecord | None = None, - deadline: float | None = None) -> None: + deadline: float | None = None, + owner_token: str | None = None) -> None: await self.state.store_future_status( - request_id, status, model_id, result=result, failure=failure, replica_id=epoch, absolute_deadline=deadline) - - async def own_future(self, request_id: str, token: str) -> None: - await self.state.set_future_owner(request_id, token) + request_id, + status, + model_id, + result=result, + failure=failure, + replica_id=epoch, + absolute_deadline=deadline, + owner_token=owner_token) async def begin_epoch(self, epoch: str) -> None: await self.backend.set(self.key('controller', 'active'), {'epoch': epoch, 'started_at': time.time()}) diff --git a/src/twinkle/server/task_queue/mixin.py b/src/twinkle/server/task_queue/mixin.py index 6c2dfb7c2..f1c9b8218 100644 --- a/src/twinkle/server/task_queue/mixin.py +++ b/src/twinkle/server/task_queue/mixin.py @@ -264,6 +264,7 @@ async def _enqueue_prepared_task(self, task: QueuedTask) -> None: queue_state=QueueState.ACTIVE.value, replica_id=getattr(self, 'replica_id', None), absolute_deadline=self._future_deadline(), + owner_token=task.token, ) try: @@ -452,13 +453,14 @@ async def schedule_background_task( coro_factory: Callable[[], Coroutine], model_id: str | None = None, task_type: str | None = None, + token: str | None = None, ) -> dict[str, Any]: """Run pure I/O outside serial compute, with the same supervised lifecycle. Capacity, timeout, cancellation and shutdown are owned by ComputeWorker; background jobs must not create a second unsupervised execution path. """ - return await self.schedule_concurrent_task(coro_factory, model_id=model_id, task_type=task_type) + return await self.schedule_concurrent_task(coro_factory, model_id=model_id, task_type=task_type, token=token) async def _fail_queue_tasks_async(self, queue_key: str, reason: str) -> None: await self._compute_worker.fail_queue_tasks(queue_key, reason) diff --git a/src/twinkle/server/task_queue/worker.py b/src/twinkle/server/task_queue/worker.py index a5bfc00ec..bd8efeef1 100644 --- a/src/twinkle/server/task_queue/worker.py +++ b/src/twinkle/server/task_queue/worker.py @@ -359,7 +359,8 @@ async def submit_detached(self, task: QueuedTask, queue_key: str, *, replica_id: task.model_id, queue_state=QueueState.ACTIVE.value, replica_id=replica_id, - absolute_deadline=absolute_deadline) + absolute_deadline=absolute_deadline, + owner_token=task.token) wrote_status = True self.run_detached(task, queue_key) except BaseException: diff --git a/src/twinkle/server/training/actor.py b/src/twinkle/server/training/actor.py index bae531238..07b2619e6 100644 --- a/src/twinkle/server/training/actor.py +++ b/src/twinkle/server/training/actor.py @@ -116,8 +116,8 @@ async def _recover_lost_state(self): if training.state in ('PENDING', 'ACTIVE', 'CLOSING'): await self.store.save_training(training) for request_id in (training.create_request_id, *training.pending_operation_ids): - await self.store.own_future(request_id, training.token) - await self.store.future(request_id, 'failed', training.training_id, failure=failure) + await self.store.future( + request_id, 'failed', training.training_id, failure=failure, owner_token=training.token) ray.kill(ray.get_actor(controller_name(self.namespace), namespace=self.namespace), no_restart=False) async def register_runtime(self, instance_id, epoch, incarnation): diff --git a/src/twinkle/server/training/controller.py b/src/twinkle/server/training/controller.py index 379bb28a4..01718119f 100644 --- a/src/twinkle/server/training/controller.py +++ b/src/twinkle/server/training/controller.py @@ -277,8 +277,8 @@ async def create(self, request: CreateTraining) -> dict[str, str]: 'pending', training_id, epoch=self.epoch, - deadline=time.time() + 2 * self.config.startup_timeout_s) - await self.store.own_future(request_id, request.token) + deadline=time.time() + 2 * self.config.startup_timeout_s, + owner_token=request.token) await self.store.save_create(identity, fingerprint, request_id, request.session_id) training = None try: @@ -434,7 +434,6 @@ async def admit(self, if request_id not in training.pending_operation_ids: training.pending_operation_ids[request_id] = {'accepted_at': self.clock(), 'renew': renew} await self.store.save_training(training) - await self.store.own_future(request_id, token) return {'instance_id': instance.instance_id, 'epoch': self.epoch} async def finish_operation(self, diff --git a/src/twinkle/server/training/router.py b/src/twinkle/server/training/router.py index 125336da5..6e833695e 100644 --- a/src/twinkle/server/training/router.py +++ b/src/twinkle/server/training/router.py @@ -165,8 +165,7 @@ async def upload(): token=body.hub_token or token, private=True) - ref = await self.schedule_background_task(upload, task_type=endpoint) - await self.store.own_future(ref['request_id'], token) + ref = await self.schedule_background_task(upload, task_type=endpoint, token=token) return envelope_from_record(ref['request_id'], await self.state.get_future(ref['request_id'])) if not body.training_id: raise ProtocolMismatchError('Upgrade Twinkle SDK: training_id is required') @@ -221,8 +220,7 @@ async def release(): await self.controller.release.remote(training_id, token, session) return {'model_id': training_id} if client_type == 'tinker' else None - ref = await self.schedule_background_task(release, task_type='release_training') - await self.store.own_future(ref['request_id'], token) + ref = await self.schedule_background_task(release, task_type='release_training', token=token) if client_type == 'tinker': return {'request_id': ref['request_id']} return envelope_from_record(ref['request_id'], await self.state.get_future(ref['request_id'])) diff --git a/tests/server/state/test_future_lifecycle.py b/tests/server/state/test_future_lifecycle.py index d9c0e63ff..e596391de 100644 --- a/tests/server/state/test_future_lifecycle.py +++ b/tests/server/state/test_future_lifecycle.py @@ -215,3 +215,34 @@ async def test_start_and_cancel_have_one_atomic_winner(manager, cancel_first): async def test_concurrent_starters_only_one_obtains_execution(manager): await _store(manager, 'race', 'pending') assert sum(await asyncio.gather(*(manager.try_start('race') for _ in range(20)))) == 1 + + +@pytest.mark.asyncio +async def test_future_owner_outlives_admission_retention_until_result_cleanup(monkeypatch): + from twinkle.server.exceptions import ResourceNotFoundError + from twinkle.server.state import ServerState + state = ServerState(backend=FakeBackend(), expiration_timeout=10) + await state.store_future_status('owned', 'pending', 'training_test', owner_token='owner') + await state.store_future_status('owned', 'completed', 'training_test', result={'value': 42}) + now = time.time() + monkeypatch.setattr(time, 'time', lambda: now + 11) + assert await state.get_future('owned') is not None + await state.assert_future_owner('owned', 'owner') + with pytest.raises(ResourceNotFoundError): + await state.assert_future_owner('owned', 'other') + await state._future_mgr.cleanup_expired(now + 11) + assert await state.get_future('owned') is None + assert not await state._backend.keys('future*') + + +@pytest.mark.asyncio +async def test_future_owner_is_immutable_and_not_part_of_wire_envelope(): + from twinkle.server.state import ServerState + from twinkle.server.lifecycle.envelope import envelope_from_record + state = ServerState(backend=FakeBackend()) + await state.store_future_status('owned', 'pending', None, owner_token='secret-owner') + await state.store_future_status('owned', 'completed', None, result=42, owner_token='replacement') + record = await state.get_future('owned') + assert record['owner_token'] == 'secret-owner' + assert 'owner_token' not in envelope_from_record('owned', record).model_dump() + assert 'secret-owner' not in envelope_from_record('owned', record).model_dump_json() diff --git a/tests/server/state/test_update_atomic.py b/tests/server/state/test_update_atomic.py index 387f69302..2bcc78f29 100644 --- a/tests/server/state/test_update_atomic.py +++ b/tests/server/state/test_update_atomic.py @@ -295,3 +295,19 @@ def force_watch_conflict(current: Any | None) -> int: 'error_code': 503, 'request_id': 'req-contention', } + + +@pytest.mark.asyncio +async def test_future_owner_is_shared_with_result_and_removed_with_it(backend_pair): + from twinkle.server.exceptions import ResourceNotFoundError + from twinkle.server.state import ServerState + first, second = (ServerState(backend=backend) for backend in backend_pair) + await first.store_future_status('owned', 'pending', None, owner_token='owner') + await second.store_future_status('owned', 'completed', None, result=42) + await first.assert_future_owner('owned', 'owner') + with pytest.raises(ResourceNotFoundError): + await second.assert_future_owner('owned', 'other') + assert (await second.get_future('owned'))['owner_token'] == 'owner' + await first._future_mgr.cleanup_expired(float('inf')) + assert await second.get_future('owned') is None + assert not await backend_pair[0].keys('future*') diff --git a/tests/server/task_queue/test_task_queue_mixin.py b/tests/server/task_queue/test_task_queue_mixin.py index cc049b02f..cf193d202 100644 --- a/tests/server/task_queue/test_task_queue_mixin.py +++ b/tests/server/task_queue/test_task_queue_mixin.py @@ -711,3 +711,33 @@ async def work(): assert queue.state._latest[ref['request_id']]['status'] == 'cancelled' assert queue._compute_worker._detached_reserved == 0 assert not queue._compute_worker._detached_tasks + + +@pytest.mark.asyncio +@pytest.mark.parametrize('method', ['schedule_task', 'schedule_concurrent_task', 'schedule_background_task']) +async def test_admission_persists_future_and_owner_together(method): + from tests.server.state.fake_backend import FakeBackend + from twinkle.server.state import ServerState + queue = TaskQueueMixin() + queue.state = ServerState(backend=FakeBackend()) + queue.replica_id = 'ownership-test' + queue._init_task_queue() + entered, finish = asyncio.Event(), asyncio.Event() + + async def operation(): + entered.set() + await finish.wait() + return 42 + + try: + ref = await getattr(queue, method)(operation, token='owner') + record = await queue.state.get_future(ref['request_id']) + assert record['owner_token'] == 'owner' + assert not await queue.state._backend.keys('future_owner::*') + await asyncio.wait_for(entered.wait(), 1) + from twinkle.server.exceptions import ResourceNotFoundError + with pytest.raises(ResourceNotFoundError): + await queue.state.assert_future_owner(ref['request_id'], 'other') + finally: + finish.set() + await queue.shutdown_task_queue() diff --git a/tests/server/training/test_router.py b/tests/server/training/test_router.py index 8f70a5893..1c610447f 100644 --- a/tests/server/training/test_router.py +++ b/tests/server/training/test_router.py @@ -43,7 +43,7 @@ async def test_upload_after_recycle_is_cpu_only_and_preserves_future_ownership(m future = await settle(router.state, response.request_id) assert future['status'] == 'completed' upload.assert_called_once_with(repo_id='owner/saved', folder_path=str(tmp_path), token='hub-token', private=True) - assert await router.state._backend.get('future_owner::' + response.request_id) == 'owner' + assert (await router.state.get_future(response.request_id))['owner_token'] == 'owner' assert not await router.state._backend.keys('model::*') finally: await router.shutdown_task_queue() From d9e7c6cf2761608c1b9fc543324c27a276fd8e89 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 02:13:57 +0800 Subject: [PATCH 20/46] fix: keep actor shutdown lookups outside the controller event loop --- src/twinkle/server/resources/runtime.py | 9 ++++-- tests/server/resources/test_instances.py | 40 ++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/src/twinkle/server/resources/runtime.py b/src/twinkle/server/resources/runtime.py index 48f6958bd..bf83ef218 100644 --- a/src/twinkle/server/resources/runtime.py +++ b/src/twinkle/server/resources/runtime.py @@ -94,15 +94,18 @@ async def shutdown(self, identity, name, *, observe_removal=None): self.deployments.pop(identity, None) # Both features give workers the controller identity, not a Serve replica ID. actors = [] - for row in list_named_actors(all_namespaces=True): + for row in await asyncio.to_thread(list_named_actors, all_namespaces=True): if row['namespace'] != self.namespace or identity not in row['name']: continue try: - handle = ray.get_actor(row['name'], namespace=self.namespace) + # Ray may import the worker class when resolving its handle. + # Keep imports and GCS calls outside the controller event loop + # so its shutdown deadline and other instances remain live. + handle = await asyncio.to_thread(ray.get_actor, row['name'], namespace=self.namespace) except ValueError: continue actors.append(handle) - ray.kill(handle, no_restart=True) + await asyncio.to_thread(ray.kill, handle, no_restart=True) await asyncio.to_thread(serve.delete, name) started, status = time.monotonic(), 'error' try: diff --git a/tests/server/resources/test_instances.py b/tests/server/resources/test_instances.py index 2870d7ec5..ec1f15907 100644 --- a/tests/server/resources/test_instances.py +++ b/tests/server/resources/test_instances.py @@ -1,3 +1,4 @@ +import asyncio from types import SimpleNamespace import logging import pytest @@ -80,3 +81,42 @@ async def failed(): assert not lifecycle.unreachable_since assert await lifecycle.health('instance', unreachable) is None assert await lifecycle.health('instance', failed) is False + + +@pytest.mark.asyncio +async def test_shutdown_actor_lookup_cannot_block_the_close_deadline(monkeypatch): + import threading + from unittest.mock import AsyncMock + import ray + from ray import serve + from twinkle.server.resources.runtime import RayRuntime + + runtime = RayRuntime('test', allocation_timeout=1, shutdown_timeout=1) + entered, resume, returned = threading.Event(), threading.Event(), threading.Event() + + def lookup(*args, **kwargs): + entered.set() + try: + # A bounded fallback makes the broken implementation fail rather than hang pytest. + if not resume.wait(2): + raise TimeoutError('blocked actor class import') + raise ValueError('actor already gone') + finally: + returned.set() + + monkeypatch.setattr(ray.util, 'list_named_actors', lambda **kwargs: [{'namespace': 'test', 'name': 'owned-worker'}]) + monkeypatch.setattr(ray, 'get_actor', lookup) + monkeypatch.setattr(serve, 'delete', lambda name: None) + runtime.pool.remove = AsyncMock() + stopped, quarantined = AsyncMock(), AsyncMock() + lifecycle = InstanceLifecycle('test') + try: + await lifecycle.close('owned', lambda: runtime.shutdown('owned', 'app'), .02, stopped, quarantined) + assert entered.is_set() + quarantined.assert_awaited_once() + stopped.assert_not_awaited() + assert not returned.is_set(), 'Synchronous actor lookup blocked the event loop past the close deadline' + finally: + resume.set() + await asyncio.to_thread(returned.wait, 3) + await lifecycle.join() From 2b472ac4149c4560dbe422bbbfaee1b07ee144f6 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 03:53:00 +0800 Subject: [PATCH 21/46] fix: synchronize every resident LoRA slot across data parallel workers --- src/twinkle/model/multi_lora.py | 21 +++ .../transformers/multi_lora_transformers.py | 21 ++- .../test_multi_lora_checkpoint_context.py | 1 + .../test_multi_lora_distributed_slots.py | 121 ++++++++++++++++++ 4 files changed, 161 insertions(+), 3 deletions(-) create mode 100644 tests/transformers/test_multi_lora_distributed_slots.py diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index 797c3e628..1e194e543 100644 --- a/src/twinkle/model/multi_lora.py +++ b/src/twinkle/model/multi_lora.py @@ -124,6 +124,27 @@ def reset_adapter_status(self): else: self.module.set_adapter('lora_0') + @contextmanager + def trainable_slots(self): + """Register every preallocated slot with distributed gradient reducers.""" + modules = self.module if isinstance(self.module, list) else [self.module] + masks = [{ + name: param.requires_grad + for name, param in module.named_parameters() if 'lora_' in name + } for module in modules] + try: + for module, mask in zip(modules, masks): + for name, param in module.named_parameters(): + if name in mask: + param.requires_grad_(True) + yield + finally: + # FSDP may replace Parameters while wrapping; restore by name. + for module, mask in zip(modules, masks): + for name, param in module.named_parameters(): + if name in mask: + param.requires_grad_(mask[name]) + def activate_adapter(self, tenant_adapter_name: str, call_enable=False): if not self.has_lora(tenant_adapter_name): raise ValueError(f'Adapter {tenant_adapter_name} does not exist') diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index ea889f4ba..e5a03d491 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/src/twinkle/model/transformers/multi_lora_transformers.py @@ -48,7 +48,10 @@ def __init__( self.device_mesh = device_mesh self.mixed_precision = mixed_precision self._fsdp_config = dict(fsdp_config or {}) - self._ddp_config = ddp_config or {} + self._ddp_config = dict(ddp_config or {}) + if self._ddp_config.get('find_unused_parameters') is False or self._ddp_config.get('static_graph', False): + raise ValueError('Multi-LoRA DDP requires find_unused_parameters=True and static_graph=False.') + self._ddp_config.setdefault('find_unused_parameters', True) self._memory_efficient_init = memory_efficient_init self._router_replay_enabled = bool(kwargs.pop('enable_router_replay', False)) self._router_replay_applied = False @@ -84,7 +87,7 @@ def __init__( # Initialize expert parallel attributes (required by set_optimizer in TransformersModel) self.optimizer_group: Dict[str, OptimizerGroup] = {} self.multi_adapter = MultiLora(max_loras=max_loras, max_r=max_r, max_length=max_length) - self.model.gradient_checkpointing_enable() + self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False}) self.model = self.multi_adapter.patch(self.model, target_modules=target_modules, lora_config=self.lora_config) # PEFT creates LoRA parameters in FP32 even when the base model uses a # lower-precision dtype. Align the preallocated slots before FSDP2 @@ -120,7 +123,15 @@ def unregister_mm_forward_hook(self, optimizer_group: OptimizerGroup): pass def _lazy_wrap_model(self): - return super()._lazy_wrap_model() + if self._model_wrapped: + return + # PEFT freezes inactive slots. DDP must nevertheless register all slots, + # since a later tenant may activate any of them on this resident model. + with self.multi_adapter.trainable_slots(): + super()._lazy_wrap_model() + # Reset must restore the synchronized (and possibly sharded) basis, + # not the rank-local initialization captured before model wrapping. + self.multi_adapter.save_initial_weights() def _maybe_apply_expert_parallel(self): if self._memory_efficient_init: @@ -313,6 +324,7 @@ def _optimizer_param_name_mapping(self, adapter_name: str, optimizer: Optimizer) def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs): adapter_name = kwargs.get('adapter_name') self._check_adapter_valid(adapter_name) + self._lazy_wrap_model() with self._adapter_context(adapter_name), self.multi_adapter.save_context(adapter_name): checkpoint_dir = super().save(name, output_dir, interval, **kwargs) if dist.is_initialized(): @@ -323,6 +335,9 @@ def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs): def load(self, name: str, output_dir: Optional[str] = None, **kwargs): adapter_name = kwargs.get('adapter_name') self._check_adapter_valid(adapter_name) + # Establish the synchronized reset basis before checkpoint weights can + # replace it, including when load is the first model operation. + self._lazy_wrap_model() with self._adapter_context(adapter_name), self.multi_adapter.save_context(adapter_name): load_optimizer = kwargs.get('load_optimizer', False) if output_dir is None: diff --git a/tests/transformers/test_multi_lora_checkpoint_context.py b/tests/transformers/test_multi_lora_checkpoint_context.py index 1b89e3d51..67694a4e4 100644 --- a/tests/transformers/test_multi_lora_checkpoint_context.py +++ b/tests/transformers/test_multi_lora_checkpoint_context.py @@ -21,6 +21,7 @@ def test_checkpoint_activates_target_tenant_before_optimizer_state_initializatio wrapper = object.__new__(MultiLoraTransformersModel) wrapper.__dict__['model'] = model wrapper.multi_adapter = manager + wrapper._model_wrapped = True wrapper.strategy = SimpleNamespace(unwrap_model=lambda value: value) with manager.adapter('first'): params = [param for name, param in model.named_parameters() if '.lora_0.' in name] diff --git a/tests/transformers/test_multi_lora_distributed_slots.py b/tests/transformers/test_multi_lora_distributed_slots.py new file mode 100644 index 000000000..89d22f0fd --- /dev/null +++ b/tests/transformers/test_multi_lora_distributed_slots.py @@ -0,0 +1,121 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Real two-process CPU regression for resident Multi-LoRA slot switching.""" +from datetime import timedelta +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from peft import LoraConfig +from safetensors.torch import save_file + +from twinkle.model.multi_lora import MultiLora +from twinkle.model.transformers.multi_lora_transformers import MultiLoraTransformersModel + + +def _wrapper(module, manager, wrap_model): + wrapper = object.__new__(MultiLoraTransformersModel) + torch.nn.Module.__init__(wrapper) + wrapper.model = module + wrapper.multi_adapter = manager + wrapper._model_wrapped = False + wrapper.optimizer_group = {} + wrapper._enable_expert_parallel = False + wrapper.sp_strategy = None + wrapper.strategy = SimpleNamespace( + capture_pre_ep_state_if_needed=lambda *args, **kwargs: None, + wrap_model=wrap_model, + unwrap_model=lambda model: model.module if isinstance(model, torch.nn.parallel.DistributedDataParallel) else model) + # This linear model does not use expert/sequence parallelism or GDN patches. + for name in ('_maybe_apply_expert_parallel', '_ensure_sp_strategy', '_ensure_gdn_padding_free_patch'): + object.__setattr__(wrapper, name, lambda: None) + return wrapper + + +def _distributed_slots(rank, rendezvous, mode): + dist.init_process_group('gloo', init_method='file://' + rendezvous, rank=rank, world_size=2, + timeout=timedelta(seconds=60)) + try: + torch.set_num_threads(1) + # DDP must synchronize initialization even when workers have different RNG states. + torch.manual_seed(42 + rank) + config = LoraConfig(r=2, lora_alpha=2, target_modules=['0']) + manager = MultiLora(max_loras=2, max_r=2) + module = manager.patch(torch.nn.Sequential(torch.nn.Linear(4, 4)), target_modules=['0'], lora_config=config) + manager.save_initial_weights() + manager.reset_adapter_status() + manager.acquire_lora('first', config) + wrapper = _wrapper(module, manager, lambda model: torch.nn.parallel.DistributedDataParallel( + model, find_unused_parameters=True)) + wrapper._lazy_wrap_model() + assert not any(param.requires_grad for name, param in module.named_parameters() if '.base_layer.' in name) + + def synchronized_weights(tenant): + state = manager.get_state_dict(tenant) + flat = torch.cat([value.detach().flatten() for value in state.values()]) + copies = [torch.empty_like(flat) for _ in range(2)] + dist.all_gather(copies, flat) + torch.testing.assert_close(copies[0], copies[1], rtol=0, atol=1e-7) + return {key: value.detach().clone() for key, value in state.items()} + + for tenant in (('first', 'second') if mode == 'switch' else ('first', 'reused')): + if tenant == 'second': + manager.acquire_lora(tenant, config) + elif tenant == 'reused': + manager.release_lora('first') + manager.acquire_lora(tenant, config) + initial = synchronized_weights(tenant) + with manager.adapter(tenant): + parameters = [p for p in module.parameters() if p.requires_grad] + optimizer = torch.optim.SGD(parameters, lr=.1) + for _ in range(3): + optimizer.zero_grad() + # Different local batches expose missing gradient synchronization. + wrapper.model(torch.ones(2, 4) * (rank + 1)).square().mean().backward() + optimizer.step() + trained = synchronized_weights(tenant) + assert any(not torch.equal(initial[name], value) for name, value in trained.items()) + with manager.adapter(tenant), torch.no_grad(): + fixed = wrapper.model(torch.ones(2, 4)).detach().clone() + # Rank 0 checkpoint must preserve predictions on every DP replica. + checkpoint = [trained if rank == 0 else None] + dist.broadcast_object_list(checkpoint, src=0) + manager.set_state_dict(tenant, checkpoint[0]) + torch.testing.assert_close(wrapper.model(torch.ones(2, 4)), fixed, rtol=0, atol=1e-7) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize('mode', ['switch', 'reuse']) +def test_all_slots_sync_across_dp_ranks_after_switch_and_reuse(tmp_path, mode): + mp.spawn(_distributed_slots, args=(str(tmp_path / 'rendezvous'), mode), nprocs=2, join=True) + + +@pytest.mark.parametrize('options', [{'find_unused_parameters': False}, {'static_graph': True}]) +def test_multi_lora_rejects_static_ddp_reducers(options, monkeypatch): + monkeypatch.setattr(MultiLoraTransformersModel, '_try_init_process_group', lambda _: None) + with pytest.raises(ValueError, match='Multi-LoRA DDP requires'): + MultiLoraTransformersModel(ddp_config=options) + + +def test_checkpoint_as_first_operation_preserves_fresh_slot_reset_basis(tmp_path): + config = LoraConfig(r=2, lora_alpha=2, target_modules=['0']) + manager = MultiLora(max_loras=1, max_r=2) + module = manager.patch(torch.nn.Sequential(torch.nn.Linear(4, 4)), target_modules=['0'], lora_config=config) + manager.save_initial_weights() + manager.acquire_lora('first', config) + initial = {name: value.detach().clone() for name, value in manager.get_state_dict('first').items()} + checkpoint = {name: torch.full_like(value, 2) for name, value in initial.items()} + save_file(checkpoint, str(tmp_path / 'adapter_model.safetensors')) + wrapper = _wrapper(module, manager, lambda model: model) + wrapper.optimizer_group = {'first': SimpleNamespace(optimizer=None)} + wrapper.load(str(tmp_path), adapter_name='first') + for name, value in manager.get_state_dict('first').items(): + torch.testing.assert_close(value, checkpoint[name]) + # Further operations must not replace the initial reset basis with trained weights. + wrapper._lazy_wrap_model() + manager.release_lora('first') + manager.acquire_lora('replacement', config) + for name, value in manager.get_state_dict('replacement').items(): + torch.testing.assert_close(value, initial[name]) From 5ef1e05444609b0842b1090f3841679055212949 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 03:55:40 +0800 Subject: [PATCH 22/46] chore: remove orphaned local ignore comment --- .gitignore | 2 -- 1 file changed, 2 deletions(-) diff --git a/.gitignore b/.gitignore index 0d55deae1..6ba7b5381 100644 --- a/.gitignore +++ b/.gitignore @@ -77,8 +77,6 @@ instance/ # Sphinx documentation docs/_build/ -# Internal design notes and developer tools are local-only. - # PyBuilder target/ From 77b230bf6fbb7fa78a2621e5ca07bd3d9d5428ab Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 05:34:53 +0800 Subject: [PATCH 23/46] fix: reject FSDP view retention for multi-tenant LoRA --- .../model/transformers/multi_lora_transformers.py | 2 ++ src/twinkle/server/README.md | 2 ++ .../test_multi_lora_distributed_slots.py | 12 ++++++++++++ 3 files changed, 16 insertions(+) diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index e5a03d491..017eae3b0 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/src/twinkle/model/transformers/multi_lora_transformers.py @@ -48,6 +48,8 @@ def __init__( self.device_mesh = device_mesh self.mixed_precision = mixed_precision self._fsdp_config = dict(fsdp_config or {}) + if self._fsdp_config.get('reshard_after_forward', True) is not True: + raise ValueError('Multi-LoRA FSDP requires reshard_after_forward=True for safe adapter switching.') self._ddp_config = dict(ddp_config or {}) if self._ddp_config.get('find_unused_parameters') is False or self._ddp_config.get('static_graph', False): raise ValueError('Multi-LoRA DDP requires find_unused_parameters=True and static_graph=False.') diff --git a/src/twinkle/server/README.md b/src/twinkle/server/README.md index 917479856..2f8ab3d25 100644 --- a/src/twinkle/server/README.md +++ b/src/twinkle/server/README.md @@ -27,6 +27,8 @@ flowchart TD The three runtime modules have different boundaries: `server/runtime.py` initializes Twinkle and DeviceMesh; `resources/runtime.py` owns Ray/Serve effects; `sampler/runtime.py` tracks request execution inside a sampling instance. +Multi-LoRA execution registers every preallocated slot with the distributed reducer. DDP requires `find_unused_parameters=True` and `static_graph=False`; FSDP requires `reshard_after_forward=True` so adapter switches act on sharded parameters instead of retained views. Incompatible overrides fail during model initialization, before weights are loaded. + ## Sampling ownership and state reconstruction Stable sampling uses shared access; configuration and full-weight updates use exclusive access. Every participating instance is reserved before an update. Idle sweep checks the same reservations under the state lock. The lock order is configuration then state; startup is never awaited while holding the configuration lock. diff --git a/tests/transformers/test_multi_lora_distributed_slots.py b/tests/transformers/test_multi_lora_distributed_slots.py index 89d22f0fd..499b661fb 100644 --- a/tests/transformers/test_multi_lora_distributed_slots.py +++ b/tests/transformers/test_multi_lora_distributed_slots.py @@ -99,6 +99,18 @@ def test_multi_lora_rejects_static_ddp_reducers(options, monkeypatch): MultiLoraTransformersModel(ddp_config=options) +@pytest.mark.parametrize('reshard_after_forward', [None, False, 2]) +def test_multi_lora_rejects_retained_or_partial_fsdp_views(reshard_after_forward, monkeypatch): + monkeypatch.setattr(MultiLoraTransformersModel, '_try_init_process_group', lambda _: None) + + def unexpected_strategy(*args): + raise AssertionError('Invalid FSDP config reached strategy initialization') + + monkeypatch.setattr(MultiLoraTransformersModel, '_decide_strategy', unexpected_strategy) + with pytest.raises(ValueError, match='Multi-LoRA FSDP requires reshard_after_forward=True'): + MultiLoraTransformersModel(fsdp_config={'reshard_after_forward': reshard_after_forward}) + + def test_checkpoint_as_first_operation_preserves_fresh_slot_reset_basis(tmp_path): config = LoraConfig(r=2, lora_alpha=2, target_modules=['0']) manager = MultiLora(max_loras=1, max_r=2) From a46e25488b6dc2ca4f0a530f76299e51b1785571 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 09:36:07 +0800 Subject: [PATCH 24/46] fix: keep Serve metadata RPCs within controller deadlines --- src/twinkle/server/README.md | 2 + src/twinkle/server/resources/runtime.py | 18 ++++-- src/twinkle/server/sampler/driver.py | 21 ++++--- src/twinkle/server/training/driver.py | 18 +++--- tests/server/resources/test_instances.py | 77 ++++++++++++++++++++++++ 5 files changed, 117 insertions(+), 19 deletions(-) diff --git a/src/twinkle/server/README.md b/src/twinkle/server/README.md index 2f8ab3d25..71933cc6c 100644 --- a/src/twinkle/server/README.md +++ b/src/twinkle/server/README.md @@ -35,6 +35,8 @@ Stable sampling uses shared access; configuration and full-weight updates use ex Pool admission is bounded by `queue_timeout`. A complete control broadcast or full-weight update uses one `effective_execution_timeout`, including completion confirmation. An ambiguous result fences all participants. Cancellation alone never confirms device release; shutdown failure keeps the instance quarantined and counted against capacity. +Serve handle and resource inventory lookups run outside the controller event loop, so synchronous Ray metadata RPCs cannot prevent these deadlines from firing. + Only successful broadcasts update replay state. Superseded templates, adapter configurations (per owner/name), and full checkpoints are compacted after the last `apply_patch`. Arbitrary patches preserve order and form a compaction barrier: dependencies before a patch must remain available. LoRA weight caches are loaded per request and eviction is not replayed. A client object rejects creation of an adapter name it already owns. Release it before creating a replacement; separate objects may use the same label. diff --git a/src/twinkle/server/resources/runtime.py b/src/twinkle/server/resources/runtime.py index bf83ef218..eb844d4f2 100644 --- a/src/twinkle/server/resources/runtime.py +++ b/src/twinkle/server/resources/runtime.py @@ -77,9 +77,15 @@ async def wait_for_cleanup(self): await asyncio.gather(*tuple(self.cleanup_tasks), return_exceptions=True) @staticmethod - def handle(name): - from ray import serve - return serve.get_app_handle(name) + async def handle(name): + + def lookup(): + from ray import serve + return serve.get_app_handle(name) + + # get_app_handle waits synchronously for the Serve controller. Keep + # that RPC inside the caller's cancellable asynchronous deadline. + return await asyncio.to_thread(lookup) async def shutdown(self, identity, name, *, observe_removal=None): import ray @@ -125,8 +131,10 @@ async def recover_orphans(self, decode, shutdown): """Decode only this owner's names; never sweep another service's resources.""" from ray import serve from ray.util.placement_group import placement_group_table - names = set(serve.status().applications) - names.update(row.get('name', '') for row in placement_group_table().values() if row.get('state') != 'REMOVED') + applications, groups = await asyncio.gather( + asyncio.to_thread(serve.status), asyncio.to_thread(placement_group_table)) + names = set(applications.applications) + names.update(row.get('name', '') for row in groups.values() if row.get('state') != 'REMOVED') for name in sorted(names): instance = decode(name) if instance is not None: diff --git a/src/twinkle/server/sampler/driver.py b/src/twinkle/server/sampler/driver.py index 1144072ff..c48e6de08 100644 --- a/src/twinkle/server/sampler/driver.py +++ b/src/twinkle/server/sampler/driver.py @@ -55,27 +55,34 @@ async def start(self, instance): await self.runtime.deploy(instance.identity, allocation.name, app, f'/api/v1/sampling-runtime/{instance.identity}', lambda: self.shutdown(instance)) - def handle(self, instance): - return self.runtime.handle(instance.application) + async def handle(self, instance): + return await self.runtime.handle(instance.application) async def invoke(self, instance, ticket, payload, stream_queue=None): - return await self.handle(instance).invoke.remote(ticket, payload, stream_queue) + handle = await self.handle(instance) + return await handle.invoke.remote(ticket, payload, stream_queue) async def wait_operation(self, instance, ticket): - await self.handle(instance).wait_operation.remote(ticket) + handle = await self.handle(instance) + await handle.wait_operation.remote(ticket) async def load_weights(self, instance, path): - await self.handle(instance).install_base_weights.remote(path) + handle = await self.handle(instance) + await handle.install_base_weights.remote(path) async def health(self, instance): - return await asyncio.wait_for(self.handle(instance).runtime_identity.remote(), 5) + async with asyncio.timeout(5): + handle = await self.handle(instance) + return await handle.runtime_identity.remote() async def shutdown(self, instance): name = instance.application or self.runtime.pool.name(f'sampling_{self.pool_id}_{instance.epoch}', instance.identity) futures = [] try: - futures = await asyncio.wait_for(self.handle(instance).pending_futures.remote(), 5) + async with asyncio.timeout(5): + handle = await self.handle(instance) + futures = await handle.pending_futures.remote() except Exception: pass await self.runtime.shutdown(instance.identity, name) diff --git a/src/twinkle/server/training/driver.py b/src/twinkle/server/training/driver.py index e10c8e705..ed642f1ff 100644 --- a/src/twinkle/server/training/driver.py +++ b/src/twinkle/server/training/driver.py @@ -86,18 +86,20 @@ async def deploy(self, instance, spec): await self.runtime.deploy(instance.instance_id, name, app, f'/api/v1/training-runtime/{instance.instance_id}', lambda: self.shutdown(instance)) - def handle(self, instance): - return self.runtime.handle(instance.application_name or self.application_name(instance)) + async def handle(self, instance): + return await self.runtime.handle(instance.application_name or self.application_name(instance)) async def create_training(self, instance, training, request): - await self.handle(instance).initialize_training.remote(training.model_dump(), request.model_dump()) + handle = await self.handle(instance) + await handle.initialize_training.remote(training.model_dump(), request.model_dump()) async def release_training(self, instance, training): - await self.handle(instance).release_training.remote(training.model_dump(), instance.resident - and not instance.dirty) + handle = await self.handle(instance) + await handle.release_training.remote(training.model_dump(), instance.resident and not instance.dirty) async def verify_operation(self, instance, request_id): - return await self.handle(instance).verify_operation.remote(request_id) + handle = await self.handle(instance) + return await handle.verify_operation.remote(request_id) async def shutdown(self, instance: InstanceRecord): name = instance.application_name or self.application_name(instance) @@ -115,7 +117,9 @@ def observe_removal(duration, status): await self.store.state.unregister_replica_capacity(instance.instance_id) async def health(self, instance): - return await asyncio.wait_for(self.handle(instance).runtime_identity.remote(), 5) + async with asyncio.timeout(5): + handle = await self.handle(instance) + return await handle.runtime_identity.remote() async def recover_orphans(self): """Ownership names survive loss of the memory state actor.""" diff --git a/tests/server/resources/test_instances.py b/tests/server/resources/test_instances.py index ec1f15907..f915d2dda 100644 --- a/tests/server/resources/test_instances.py +++ b/tests/server/resources/test_instances.py @@ -120,3 +120,80 @@ def lookup(*args, **kwargs): resume.set() await asyncio.to_thread(returned.wait, 3) await lifecycle.join() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('kind', ['sampler', 'training']) +async def test_serve_handle_lookup_cannot_block_driver_deadlines(monkeypatch, kind): + import threading + from types import SimpleNamespace + from ray import serve + from twinkle.server.resources.runtime import RayRuntime + from twinkle.server.sampler.driver import SamplerDriver + from twinkle.server.training.driver import RayInstanceDriver + + entered, resume, returned = threading.Event(), threading.Event(), threading.Event() + + def lookup(name): + entered.set() + try: + if not resume.wait(2): + raise TimeoutError('blocked Serve controller RPC') + raise RuntimeError('lookup released after cancellation') + finally: + returned.set() + + monkeypatch.setattr(serve, 'get_app_handle', lookup) + driver_type = SamplerDriver if kind == 'sampler' else RayInstanceDriver + driver = driver_type.__new__(driver_type) + driver.runtime = RayRuntime('test', allocation_timeout=1, shutdown_timeout=1) + instance = SimpleNamespace(application='app', application_name='app') + operation = asyncio.create_task(driver.health(instance)) + try: + assert await asyncio.to_thread(entered.wait, 3) + with pytest.raises(TimeoutError): + await asyncio.wait_for(operation, .02) + assert not returned.is_set(), 'Serve handle lookup blocked the controller event loop' + finally: + resume.set() + await asyncio.gather(operation, return_exceptions=True) + await asyncio.to_thread(returned.wait, 3) + + +@pytest.mark.asyncio +@pytest.mark.parametrize('kind', ['applications', 'placement_groups']) +async def test_orphan_inventory_cannot_block_recovery_deadlines(monkeypatch, kind): + import threading + from types import SimpleNamespace + from ray import serve + from twinkle.server.resources.runtime import RayRuntime + + # Import the module rather than ray.util's same-named public function. + import importlib + placement_module = importlib.import_module('ray.util.placement_group') + entered, resume, returned = threading.Event(), threading.Event(), threading.Event() + + def inventory(): + entered.set() + try: + if not resume.wait(2): + raise TimeoutError('blocked resource inventory RPC') + return SimpleNamespace(applications={}) if kind == 'applications' else {} + finally: + returned.set() + + monkeypatch.setattr( + serve, 'status', inventory if kind == 'applications' else lambda: SimpleNamespace(applications={})) + monkeypatch.setattr( + placement_module, 'placement_group_table', inventory if kind == 'placement_groups' else lambda: {}) + runtime = RayRuntime('test', allocation_timeout=1, shutdown_timeout=1) + operation = asyncio.create_task(runtime.recover_orphans(lambda name: None, None)) + try: + assert await asyncio.to_thread(entered.wait, 3) + with pytest.raises(TimeoutError): + await asyncio.wait_for(operation, .02) + assert not returned.is_set(), 'Resource inventory blocked the controller event loop' + finally: + resume.set() + await asyncio.gather(operation, return_exceptions=True) + await asyncio.to_thread(returned.wait, 3) From d925db33fa95725d93614323658dc8e4afc2cd05 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 09:57:13 +0800 Subject: [PATCH 25/46] fix: retain placement group ownership across allocation timeouts --- src/twinkle/server/README.md | 1 + src/twinkle/server/resources/allocation.py | 63 +++++++++---- src/twinkle/server/sampler/driver.py | 2 +- src/twinkle/server/training/driver.py | 2 +- tests/server/resources/test_allocation.py | 102 +++++++++++++++++++++ 5 files changed, 149 insertions(+), 21 deletions(-) create mode 100644 tests/server/resources/test_allocation.py diff --git a/src/twinkle/server/README.md b/src/twinkle/server/README.md index 71933cc6c..bc36441ae 100644 --- a/src/twinkle/server/README.md +++ b/src/twinkle/server/README.md @@ -36,6 +36,7 @@ Stable sampling uses shared access; configuration and full-weight updates use ex Pool admission is bounded by `queue_timeout`. A complete control broadcast or full-weight update uses one `effective_execution_timeout`, including completion confirmation. An ambiguous result fences all participants. Cancellation alone never confirms device release; shutdown failure keeps the instance quarantined and counted against capacity. Serve handle and resource inventory lookups run outside the controller event loop, so synchronous Ray metadata RPCs cannot prevent these deadlines from firing. +Placement-group creation remains tracked across cancellation; removal waits for the late result before certifying release. Only successful broadcasts update replay state. Superseded templates, adapter configurations (per owner/name), and full checkpoints are compacted after the last `apply_patch`. Arbitrary patches preserve order and form a compaction barrier: dependencies before a patch must remain available. LoRA weight caches are loaded per request and eviction is not replayed. diff --git a/src/twinkle/server/resources/allocation.py b/src/twinkle/server/resources/allocation.py index 0fb27342b..10b4ebfdb 100644 --- a/src/twinkle/server/resources/allocation.py +++ b/src/twinkle/server/resources/allocation.py @@ -40,29 +40,38 @@ def __init__(self, namespace: str, *, allocation_timeout: float, shutdown_timeou self.namespace = namespace self.allocation_timeout = allocation_timeout self.shutdown_timeout = shutdown_timeout + self._creating = {} def name(self, epoch: str, instance_id: str) -> str: return f'resources-{self.namespace}-{epoch}-{instance_id}' - async def allocate(self, num_devices: int, device_type: str, epoch: str, instance_id: str) -> DeviceAllocation: + @staticmethod + def _create_group(resource, count, name): import ray from ray.util.placement_group import placement_group - from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy - - if type(num_devices) is not int or num_devices <= 0: - raise ValueError('num_devices must be a positive integer') - resource = 'CPU' if device_type == 'cpu' else Platform.get_platform(device_type).__name__ - count = num_devices candidates = [node for node in ray.nodes() if node.get('Alive') and node['Resources'].get(resource, 0) >= count] available = ray.available_resources() if not candidates or available.get(resource, 0) < count: raise CapacityUnavailableError(f'No single-node allocation available for {count} {resource} devices') cpus = min(_gpu_placement_group_cpus(int(node['Resources'].get('CPU', 0)), count) for node in candidates) bundle = {resource: count, 'CPU': count if resource == 'CPU' else cpus} + return placement_group([bundle], strategy='STRICT_PACK', name=name) + + async def allocate(self, num_devices: int, device_type: str, epoch: str, instance_id: str) -> DeviceAllocation: + import ray + from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + + if type(num_devices) is not int or num_devices <= 0: + raise ValueError('num_devices must be a positive integer') + resource = 'CPU' if device_type == 'cpu' else Platform.get_platform(device_type).__name__ + count = num_devices name = self.name(epoch, instance_id) - pg = placement_group([bundle], strategy='STRICT_PACK', name=name) + creating = asyncio.create_task(asyncio.to_thread(self._create_group, resource, count, name)) + self._creating[name] = creating try: - await asyncio.wait_for(pg.ready(), self.allocation_timeout) + async with asyncio.timeout(self.allocation_timeout): + pg = await asyncio.shield(creating) + await pg.ready() options = {'num_cpus': 0.01, 'scheduling_strategy': PlacementGroupSchedulingStrategy(placement_group=pg)} if resource == 'GPU': options['num_gpus'] = count @@ -76,22 +85,38 @@ async def allocate(self, num_devices: int, device_type: str, epoch: str, instanc raise RuntimeError('Ray allocation did not return the requested unique device IDs') return DeviceAllocation(name, pg.id.hex(), node, tuple(devices)) except BaseException as exc: + # Cancellation cannot stop a synchronous GCS create. Keep its task + # owned until removal has observed the late result; a close timeout + # leaves the enclosing instance quarantined for a later retry. await self.remove(name) if isinstance(exc, TimeoutError): raise CapacityUnavailableError( 'Single-node allocation could not be confirmed; retry a new create later.') from exc raise + finally: + if creating.done() and self._creating.get(name) is creating: + self._creating.pop(name) async def remove(self, name: str) -> None: from ray.util.placement_group import get_placement_group, placement_group_table, remove_placement_group try: - pg = get_placement_group(name) - except ValueError: - return - remove_placement_group(pg) - deadline = asyncio.get_running_loop().time() + self.shutdown_timeout - while asyncio.get_running_loop().time() < deadline: - if placement_group_table(pg).get('state') == 'REMOVED': - return - await asyncio.sleep(0.05) - raise TimeoutError(f'Placement group removal was not confirmed: {name}') + async with asyncio.timeout(self.shutdown_timeout): + creating = self._creating.get(name) + if creating is not None: + try: + await asyncio.shield(creating) + except Exception: + pass + if self._creating.get(name) is creating: + self._creating.pop(name) + try: + pg = await asyncio.to_thread(get_placement_group, name) + except ValueError: + return + await asyncio.to_thread(remove_placement_group, pg) + while True: + if (await asyncio.to_thread(placement_group_table, pg)).get('state') == 'REMOVED': + return + await asyncio.sleep(0.05) + except TimeoutError as exc: + raise TimeoutError(f'Placement group removal was not confirmed: {name}') from exc diff --git a/src/twinkle/server/sampler/driver.py b/src/twinkle/server/sampler/driver.py index c48e6de08..1435efcc8 100644 --- a/src/twinkle/server/sampler/driver.py +++ b/src/twinkle/server/sampler/driver.py @@ -50,7 +50,7 @@ async def start(self, instance): 'registry_name': allocation.name + '-registry', 'instance_id': instance.identity }, - placement_group=get_placement_group(allocation.name), + placement_group=await asyncio.to_thread(get_placement_group, allocation.name), allocation_devices=list(allocation.devices)) await self.runtime.deploy(instance.identity, allocation.name, app, f'/api/v1/sampling-runtime/{instance.identity}', lambda: self.shutdown(instance)) diff --git a/src/twinkle/server/training/driver.py b/src/twinkle/server/training/driver.py index ed642f1ff..698b210b0 100644 --- a/src/twinkle/server/training/driver.py +++ b/src/twinkle/server/training/driver.py @@ -80,7 +80,7 @@ async def deploy(self, instance, spec): queue_config=spec.config.queue_config, data_plane_url=spec.model.data_plane_url, training_context=context, - placement_group=get_placement_group(name), + placement_group=await asyncio.to_thread(get_placement_group, name), allocation_devices=list(instance.devices), **kwargs) await self.runtime.deploy(instance.instance_id, name, app, f'/api/v1/training-runtime/{instance.instance_id}', diff --git a/tests/server/resources/test_allocation.py b/tests/server/resources/test_allocation.py new file mode 100644 index 000000000..cac7c2e3d --- /dev/null +++ b/tests/server/resources/test_allocation.py @@ -0,0 +1,102 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Resource RPC deadlines must preserve ownership of late placement groups.""" +import asyncio +import threading +from types import SimpleNamespace + +import pytest + +from twinkle.server.resources.allocation import DevicePool + + +@pytest.mark.asyncio +@pytest.mark.parametrize('blocked', ['nodes', 'available_resources', 'placement_group']) +async def test_cancelled_allocation_keeps_late_creation_owned(monkeypatch, blocked): + import importlib + import ray + placement = importlib.import_module('ray.util.placement_group') + entered, resume, returned = threading.Event(), threading.Event(), threading.Event() + created, removed = [], [] + + async def ready(): + return None + + pg = SimpleNamespace(id=SimpleNamespace(hex=lambda: 'pg-id'), ready=ready) + + def call(name, result): + if name == blocked: + entered.set() + try: + if not resume.wait(2): + raise TimeoutError('blocked allocation RPC') + finally: + returned.set() + return result + + def create(*args, **kwargs): + result = call('placement_group', pg) + created.append(result) + return result + + def get(name): + if not created: + raise ValueError('not created yet') + return pg + + monkeypatch.setattr(ray, 'nodes', lambda: call('nodes', [{'Alive': True, 'Resources': {'CPU': 2}}])) + monkeypatch.setattr(ray, 'available_resources', lambda: call('available_resources', {'CPU': 2})) + monkeypatch.setattr(placement, 'placement_group', create) + monkeypatch.setattr(placement, 'get_placement_group', get) + monkeypatch.setattr(placement, 'remove_placement_group', lambda group: removed.append(group)) + monkeypatch.setattr(placement, 'placement_group_table', lambda group: {'state': 'REMOVED'}) + pool = DevicePool('test', allocation_timeout=.02, shutdown_timeout=.02) + operation = asyncio.create_task(pool.allocate(1, 'cpu', 'epoch', 'instance')) + try: + assert await asyncio.to_thread(entered.wait, 3) + with pytest.raises(TimeoutError): + await operation + assert not returned.is_set(), 'Allocation RPC blocked its deadline' + with pytest.raises(TimeoutError): + await asyncio.wait_for(pool.remove(pool.name('epoch', 'instance')), .02) + assert not created, 'The blocked RPC should still be pending' + finally: + resume.set() + await asyncio.gather(operation, return_exceptions=True) + await asyncio.to_thread(returned.wait, 3) + await asyncio.wait_for(pool.remove(pool.name('epoch', 'instance')), 1) + assert created and removed and all(group is pg for group in removed) + + +@pytest.mark.asyncio +@pytest.mark.parametrize('blocked', ['get_placement_group', 'remove_placement_group', 'placement_group_table']) +async def test_removal_rpc_cannot_block_quarantine_deadline(monkeypatch, blocked): + import importlib + placement = importlib.import_module('ray.util.placement_group') + entered, resume, returned = threading.Event(), threading.Event(), threading.Event() + pg = object() + + def call(name, result): + if name == blocked: + entered.set() + try: + if not resume.wait(2): + raise TimeoutError('blocked placement group RPC') + finally: + returned.set() + return result + + monkeypatch.setattr(placement, 'get_placement_group', lambda name: call('get_placement_group', pg)) + monkeypatch.setattr(placement, 'remove_placement_group', lambda group: call('remove_placement_group', None)) + monkeypatch.setattr( + placement, 'placement_group_table', lambda group: call('placement_group_table', {'state': 'REMOVED'})) + pool = DevicePool('test', allocation_timeout=1, shutdown_timeout=.02) + operation = asyncio.create_task(pool.remove('owned')) + try: + assert await asyncio.to_thread(entered.wait, 3) + with pytest.raises(TimeoutError): + await operation + assert not returned.is_set(), 'Placement group RPC blocked the close deadline' + finally: + resume.set() + await asyncio.gather(operation, return_exceptions=True) + await asyncio.to_thread(returned.wait, 3) From b213911d9c742aee23388bf291bd43c6fecde9ca Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 14:25:40 +0800 Subject: [PATCH 26/46] refactor: keep the existing model application YAML format --- .../client/server/on_demand_training.yaml | 41 ------ .../Usage Guide/Server and Client/Server.md | 6 +- .../\346\234\215\345\212\241\347\253\257.md" | 6 +- src/twinkle/server/README.md | 4 + src/twinkle/server/config/application_spec.py | 13 -- src/twinkle/server/config/server_config.py | 2 +- .../server/config/training_normalization.py | 118 +++++++----------- src/twinkle_client/auto/agent/server_tools.py | 1 - tests/server/config/test_server_config.py | 5 +- tests/server/config/test_training_config.py | 63 +++++++--- .../on_demand_training/http_probe.py | 37 ++---- .../on_demand_training/matrix_validation.py | 25 ++-- .../test_resident_release_e2e.py | 13 +- tests/server/training/test_launcher.py | 8 +- tests/twinkle_client/test_auto_agent_tools.py | 2 +- 15 files changed, 147 insertions(+), 197 deletions(-) delete mode 100644 cookbook/client/server/on_demand_training.yaml diff --git a/cookbook/client/server/on_demand_training.yaml b/cookbook/client/server/on_demand_training.yaml deleted file mode 100644 index 60ff1caba..000000000 --- a/cookbook/client/server/on_demand_training.yaml +++ /dev/null @@ -1,41 +0,0 @@ -# On-demand training: applications stays the public configuration entry. -training: - allocation_timeout_s: 3 - startup_timeout_s: 1800 - # PPU actor/PlacementGroup death confirmation can exceed the 60s framework default. - shutdown_timeout_s: 180 - sweep_interval_s: 5 -applications: - - name: model-qwen - route_prefix: /api/v1/model/Qwen/Qwen3-1.7B - import_path: model - args: - model_id: ms://Qwen/Qwen3-1.7B - backend: transformers - device_type: cuda - modes: - lora: - num_devices: 2 - device_mesh: {device_type: cuda, dp_size: 2} - min_replicas: 1 - max_loras: 5 - max_r: 32 - idle_timeout_s: 1800 - target_modules: all-linear - full: - num_devices: 4 - device_mesh: {device_type: cuda, dp_size: 4} - min_replicas: 0 - idle_timeout_s: 1800 - - name: model-small - route_prefix: /api/v1/model/Qwen/Qwen3-0.6B - import_path: model - args: - model_id: ms://Qwen/Qwen3-0.6B - backend: megatron - device_type: cuda - modes: - lora: - num_devices: 2 - device_mesh: {device_type: cuda, dp_size: 2} - max_loras: 5 diff --git a/docs/source_en/Usage Guide/Server and Client/Server.md b/docs/source_en/Usage Guide/Server and Client/Server.md index 87d97d3d4..e6fa99611 100644 --- a/docs/source_en/Usage Guide/Server and Client/Server.md +++ b/docs/source_en/Usage Guide/Server and Client/Server.md @@ -482,11 +482,11 @@ holds the cleanup lease. ## On-demand training and device reclamation -Existing `applications` entries with `import_path: model` remain valid. Set `args.min_replicas: 0` for lazy startup and explicitly set `args.idle_timeout_s` for business inactivity. Legacy `adapter_config.adapter_timeout` remains a heartbeat timeout and `adapter_max_lifetime` retains its absolute lifetime. New `modes` default to 1800 seconds of business inactivity; heartbeats and polling do not renew that deadline; queued and running operations protect training state. Empty instances above the resident minimum are destroyed so other models can reuse their devices. Existing deployment min/max replicas become training instance minimums/limits. +Existing `applications` entries with `import_path: model` remain valid. Set `args.min_replicas: 0` for lazy startup and explicitly set `args.idle_timeout_s` for business inactivity. `adapter_config.adapter_timeout` remains a heartbeat timeout and `adapter_max_lifetime` retains its absolute lifetime. Heartbeats and polling do not renew the business inactivity deadline; queued and running operations protect training state. Empty instances above the resident minimum are destroyed so other models can reuse their devices. Existing deployment min/max replicas become training instance minimums/limits. -`modes` is optional: single-mode training does not require rewriting existing YAML. Model `deployments` options configure the CPU router; autoscaling request targets apply to router replicas. The Controller owns training-instance minimums/limits, and each compute instance has exactly one Serve replica. Non-autoscaling request limits, health/shutdown settings, and `ray_actor_options.runtime_env` also reach compute instances. `args.env_vars` overrides matching deployment environment variables. Unsupported `adapter_config` fields fail validation instead of being silently ignored. +Training uses the existing YAML format: select `args.train_mode: lora`, `hybrid`, or `full`, and configure devices through `device_group` and `device_mesh`. Model `deployments` options configure the CPU router; autoscaling request targets apply to router replicas. The Controller owns training-instance minimums/limits, and each compute instance has exactly one Serve replica. Non-autoscaling request limits, health/shutdown settings, and `ray_actor_options.runtime_env` also reach compute instances. `args.env_vars` overrides matching deployment environment variables. Unsupported `adapter_config` fields fail validation instead of being silently ignored. -To offer both LoRA and full training for one model, configure `num_devices`, `device_mesh`, `min_replicas`, and `idle_timeout_s` separately under `args.modes.lora/full`. LoRA users share compatible instances; each full-training handle owns an instance. Sampling configuration stays the same. See `cookbook/client/server/on_demand_training.yaml` in the repository. +To offer both LoRA and full training for the same weights, declare two `import_path: model` applications with distinct `name` and `route_prefix` values. Set `args.train_mode: lora` on one and `full` on the other; both may use the same `model_id`. Clients select the model alias from the corresponding route. LoRA users share compatible instances; each full-training handle owns an instance. Sampling configuration stays the same. Complete examples are in `cookbook/client/server/transformer/server_config.yaml` and `cookbook/client/server/megatron/server_config.yaml`. ### Single-node, eight-GPU LoRA capacity diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" index 367f4a967..084fc1901 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" @@ -469,11 +469,11 @@ twinkle-server check-config -c server_config.yaml ## 按需训练与设备回收 -原有 `applications` 和 `import_path: model` 配置继续有效。单模式配置增加 `args.min_replicas: 0` 即可懒启动，`args.idle_timeout_s` 显式设置训练句柄的业务闲置时间。旧 `adapter_config.adapter_timeout` 仍按心跳判断会话存活，`adapter_max_lifetime` 保留绝对寿命；不再把旧字段转换成业务闲置时间。新增 `modes` 的业务闲置默认 1800 秒，心跳和轮询不延长该期限，运行中或已排队的任务保护训练状态。最后一个用户释放或过期后，超出保底数量的实例立即销毁，释放的卡可用于其他模型。已有 deployment 的 min/max_replicas 转成训练实例的保底/上限。 +原有 `applications` 和 `import_path: model` 配置继续有效。单模式配置增加 `args.min_replicas: 0` 即可懒启动，`args.idle_timeout_s` 显式设置训练句柄的业务闲置时间。`adapter_config.adapter_timeout` 仍按心跳判断会话存活，`adapter_max_lifetime` 保留绝对寿命；不再把旧字段转换成业务闲置时间。心跳和轮询不延长业务闲置期限，运行中或已排队的任务保护训练状态。最后一个用户释放或过期后，超出保底数量的实例立即销毁，释放的卡可用于其他模型。已有 deployment 的 min/max_replicas 转成训练实例的保底/上限。 -`modes` 是可选扩展，单模式训练不需要改写原 YAML。原模型 `deployments` 的 Serve 设置用于 CPU 路由部署，其中 autoscaling 的请求目标用于路由副本扩缩容；训练计算实例的 min/max 仍由 Controller 管理，每个计算实例固定一个 Serve replica。非 autoscaling 的请求上限、健康/关闭设置和 `ray_actor_options.runtime_env` 会传给计算实例；`args.env_vars` 优先于 deployment 的同名环境变量。未知或未实现的 `adapter_config` 字段会报错，不再静默忽略。 +训练统一使用原有 YAML 格式，通过 `args.train_mode: lora`、`hybrid` 或 `full` 选择训练模式，通过 `device_group` 和 `device_mesh` 配置设备。原模型 `deployments` 的 Serve 设置用于 CPU 路由部署，其中 autoscaling 的请求目标用于路由副本扩缩容；训练计算实例的 min/max 仍由 Controller 管理，每个计算实例固定一个 Serve replica。非 autoscaling 的请求上限、健康/关闭设置和 `ray_actor_options.runtime_env` 会传给计算实例；`args.env_vars` 优先于 deployment 的同名环境变量。未知或未实现的 `adapter_config` 字段会报错，不再静默忽略。 -同一模型需要同时提供 LoRA 和全参数训练时，在 `args.modes.lora/full` 内分别设置 `num_devices`、`device_mesh`、`min_replicas` 和 `idle_timeout_s`；不用改成另一套顶层模型目录。LoRA 用户共享兼容实例，全参数用户各占一个实例。采样服务继续使用原配置。详见仓库 `cookbook/client/server/on_demand_training.yaml`。 +同一权重需要同时提供 LoRA 和全参数训练时，声明两个 `import_path: model` 的 application，使用不同的 `name` 和 `route_prefix`，分别设置 `args.train_mode: lora` 和 `full`；`model_id` 可以指向同一权重。客户端使用各自路由中的模型别名。LoRA 用户共享兼容实例，全参数用户各占一个实例。采样服务继续使用原配置；完整示例见 `cookbook/client/server/transformer/server_config.yaml` 和 `cookbook/client/server/megatron/server_config.yaml`。 ### 单机 8 卡、仅 LoRA 的容量配置 diff --git a/src/twinkle/server/README.md b/src/twinkle/server/README.md index bc36441ae..db5c358e4 100644 --- a/src/twinkle/server/README.md +++ b/src/twinkle/server/README.md @@ -29,6 +29,10 @@ The three runtime modules have different boundaries: `server/runtime.py` initial Multi-LoRA execution registers every preallocated slot with the distributed reducer. DDP requires `find_unused_parameters=True` and `static_graph=False`; FSDP requires `reshard_after_forward=True` so adapter switches act on sharded parameters instead of retained views. Incompatible overrides fail during model initialization, before weights are loaded. +## Training configuration + +The public YAML keeps `applications` with `import_path: model`, `args.train_mode`, `device_group`, and `device_mesh`. Each model application selects one training mode. To expose LoRA and full training for the same weights, use two applications with distinct names and route aliases. `config/training_normalization.py` translates these entries into the internal catalog; internal mode records are not an alternative YAML schema. + ## Sampling ownership and state reconstruction Stable sampling uses shared access; configuration and full-weight updates use exclusive access. Every participating instance is reserved before an update. Idle sweep checks the same reservations under the state lock. The lock order is configuration then state; startup is never awaited while holding the configuration lock. diff --git a/src/twinkle/server/config/application_spec.py b/src/twinkle/server/config/application_spec.py index a25da604e..ab1bbd845 100644 --- a/src/twinkle/server/config/application_spec.py +++ b/src/twinkle/server/config/application_spec.py @@ -17,7 +17,6 @@ from .hybrid import SpectralHybridArgs from .task_queue import TaskQueueConfig -from .training import TrainingModes # Env var keys the launcher sets from the gateway ``server_config`` so that any # Ray worker (model / sampler / processor), not just the gateway, applies the @@ -69,9 +68,7 @@ class ModelArgs(_ArgsBase): device_group: dict[str, Any] device_mesh: dict[str, Any] backend: Literal['mock', 'transformers', 'megatron'] - device_type: Literal['cuda', 'npu', 'mps', 'cpu'] | None = None env_vars: dict[str, str] = Field(default_factory=dict) - modes: TrainingModes | None = None min_replicas: int | None = Field(default=None, ge=0, strict=True) max_replicas: int | None = Field(default=None, gt=0, strict=True) idle_timeout_s: float | None = Field(default=None, gt=0, allow_inf_nan=False) @@ -86,18 +83,8 @@ class ModelArgs(_ArgsBase): hybrid: SpectralHybridArgs | None = None data_plane_url: str | None = None - @model_validator(mode='before') - @classmethod - def _mode_resources(cls, data): - if isinstance(data, dict) and data.get('modes') is not None: - return {'device_group': {}, 'device_mesh': {}, **data} - return data - @model_validator(mode='after') def _validate_spectral_backend(self): - if self.modes is not None and any( - mode.idle_timeout_s is None for mode in (self.modes.lora, self.modes.full) if mode is not None): - raise ValueError('modes require a finite positive idle_timeout_s') if self.strategy is not None and self.backend != 'transformers': raise ValueError('strategy is only supported by the transformers backend') if self.hybrid is not None: diff --git a/src/twinkle/server/config/server_config.py b/src/twinkle/server/config/server_config.py index 31d22b5fb..bceec284d 100644 --- a/src/twinkle/server/config/server_config.py +++ b/src/twinkle/server/config/server_config.py @@ -87,7 +87,7 @@ def _validate_cross_field(self) -> ServerConfig: @property def training_config(self) -> TrainingConfig: - """Normalize old and extended model args into one runtime catalog.""" + """Normalize model application args into the internal runtime catalog.""" return TrainingConfig( **self.training.model_dump(), models=[normalize_model_application(app) for app in self.applications if app.import_path == 'model']) diff --git a/src/twinkle/server/config/training_normalization.py b/src/twinkle/server/config/training_normalization.py index ee38a67c8..5a63fc162 100644 --- a/src/twinkle/server/config/training_normalization.py +++ b/src/twinkle/server/config/training_normalization.py @@ -1,5 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Translate model application configuration without creating a second runtime.""" +"""Normalize the model application YAML into the internal training catalog.""" from twinkle.hub.model_alias import extract_model_alias from .application_spec import ApplicationSpec, ModelArgs from .training import LegacyExpiration, TrainingModelConfig, TrainingModes @@ -8,76 +8,52 @@ def normalize_model_application(app: ApplicationSpec) -> TrainingModelConfig: args = app.args assert isinstance(args, ModelArgs) - deployment = {} - if args.modes is not None: - baseline = ModelArgs(model_id=args.model_id, backend=args.backend, device_group={}, device_mesh={}) - common = { - 'model_id', 'backend', 'modes', 'data_plane_url', 'device_group', 'device_mesh', 'device_type', 'env_vars' - } - legacy_fields = { - key - for key in type(args).model_fields if key not in common and getattr(args, key) != getattr(baseline, key) - } - if legacy_fields or args.device_group or args.device_mesh: - raise ValueError(f'{app.name}: modes cannot be mixed with single-mode training parameters') - if app.deployments: - raise ValueError(f'{app.name}: configure instance replicas inside modes, not deployments') - modes = args.modes - legacy_expiration = None - device_types = {mode.device_mesh.device_type for mode in (modes.lora, modes.full) if mode is not None} - device_types.discard(None) - if len(device_types) > 1: - raise ValueError(f'{app.name}: modes must use the same device type') - device_type = args.device_type or next(iter(device_types), 'cpu' if args.backend == 'mock' else 'cuda').lower() - env_vars = args.env_vars - else: - if len(app.deployments) > 1: - raise ValueError(f'{app.name}: a model application must have one deployment configuration') - deployment = app.deployments[0] if app.deployments else {} - unsupported = set(args.adapter_config or {}) - {'adapter_timeout', 'adapter_max_lifetime'} - if unsupported: - raise ValueError(f'{app.name}: unsupported adapter_config fields: {sorted(unsupported)}') - scaling = deployment.get('autoscaling_config') or {} - ranks = args.device_group.get('ranks', args.nproc_per_node) - if isinstance(ranks, list): - raise ValueError(f'{app.name}: use a device count instead of fixed ranks for pooled training') - device_type = args.device_type or str( - args.device_group.get('device_type', - args.device_mesh.get('device_type', - 'cpu' if args.backend == 'mock' else 'cuda'))).lower() - mesh = {**args.device_mesh, 'device_type': device_type} - mode = { - 'num_devices': - ranks, - 'device_mesh': - mesh, - 'min_replicas': - args.min_replicas if args.min_replicas is not None else scaling.get('min_replicas', - deployment.get('num_replicas', 0)), - 'max_replicas': - args.max_replicas if args.max_replicas is not None else scaling.get('max_replicas', - deployment.get('num_replicas')), - 'idle_timeout_s': - args.idle_timeout_s, - 'strategy': - args.strategy, - 'queue_config': - args.queue_config, - 'max_length': - args.max_length, - } - mode_name = 'full' if args.train_mode == 'full' else 'lora' - if mode_name == 'lora': - mode.update( - max_loras=args.max_loras, max_r=args.max_r, target_modules=args.target_modules, hybrid=args.hybrid) - modes = TrainingModes.model_validate({mode_name: mode}) - legacy_expiration = LegacyExpiration( - session_timeout_s=(args.adapter_config or {}).get('adapter_timeout', 1800), - max_lifetime_s=(args.adapter_config or {}).get('adapter_max_lifetime', 36000)) - env_vars = { - **(deployment.get('ray_actor_options') or {}).get('runtime_env', {}).get('env_vars', {}), - **args.env_vars - } + if len(app.deployments) > 1: + raise ValueError(f'{app.name}: a model application must have one deployment configuration') + deployment = app.deployments[0] if app.deployments else {} + unsupported = set(args.adapter_config or {}) - {'adapter_timeout', 'adapter_max_lifetime'} + if unsupported: + raise ValueError(f'{app.name}: unsupported adapter_config fields: {sorted(unsupported)}') + scaling = deployment.get('autoscaling_config') or {} + ranks = args.device_group.get('ranks', args.nproc_per_node) + if isinstance(ranks, list): + raise ValueError(f'{app.name}: use a device count instead of fixed ranks for pooled training') + device_type = str( + args.device_group.get('device_type', + args.device_mesh.get('device_type', + 'cpu' if args.backend == 'mock' else 'cuda'))).lower() + mesh = {**args.device_mesh, 'device_type': device_type} + mode = { + 'num_devices': + ranks, + 'device_mesh': + mesh, + 'min_replicas': + args.min_replicas if args.min_replicas is not None else scaling.get('min_replicas', + deployment.get('num_replicas', 0)), + 'max_replicas': + args.max_replicas if args.max_replicas is not None else scaling.get('max_replicas', + deployment.get('num_replicas')), + 'idle_timeout_s': + args.idle_timeout_s, + 'strategy': + args.strategy, + 'queue_config': + args.queue_config, + 'max_length': + args.max_length, + } + mode_name = 'full' if args.train_mode == 'full' else 'lora' + if mode_name == 'lora': + mode.update(max_loras=args.max_loras, max_r=args.max_r, target_modules=args.target_modules, hybrid=args.hybrid) + modes = TrainingModes.model_validate({mode_name: mode}) + legacy_expiration = LegacyExpiration( + session_timeout_s=(args.adapter_config or {}).get('adapter_timeout', 1800), + max_lifetime_s=(args.adapter_config or {}).get('adapter_max_lifetime', 36000)) + env_vars = { + **(deployment.get('ray_actor_options') or {}).get('runtime_env', {}).get('env_vars', {}), + **args.env_vars + } name = extract_model_alias(app.route_prefix) or app.name return TrainingModelConfig( name=name, diff --git a/src/twinkle_client/auto/agent/server_tools.py b/src/twinkle_client/auto/agent/server_tools.py index 40e3f8c99..700bddc6e 100644 --- a/src/twinkle_client/auto/agent/server_tools.py +++ b/src/twinkle_client/auto/agent/server_tools.py @@ -422,7 +422,6 @@ def _short(mid: str) -> str: 'args': { 'model_id': f'ms://{model_id}', 'backend': backend, - 'device_type': 'cuda', 'env_vars': { 'TWINKLE_TRUST_REMOTE_CODE': '1' }, diff --git a/tests/server/config/test_server_config.py b/tests/server/config/test_server_config.py index 3c7e7568e..9f507fa61 100644 --- a/tests/server/config/test_server_config.py +++ b/tests/server/config/test_server_config.py @@ -38,8 +38,9 @@ 'args': st.fixed_dictionaries({ 'model_id': st.just('model-id'), 'backend': st.sampled_from(['mock', 'transformers', 'megatron']), - 'device_type': st.just('cuda'), - 'modes': st.just({'lora': {'num_devices': 2, 'device_mesh': {'dp_size': 2}}}), + 'device_group': st.just({'ranks': 2, 'device_type': 'cuda'}), + 'device_mesh': st.just({'dp_size': 2}), + 'train_mode': st.sampled_from(['lora', 'full']), }), }) diff --git a/tests/server/config/test_training_config.py b/tests/server/config/test_training_config.py index 5c9eedc73..9c9a055fa 100644 --- a/tests/server/config/test_training_config.py +++ b/tests/server/config/test_training_config.py @@ -23,24 +23,25 @@ def model_config(): } -def model_application(model=None): - model = dict(model or model_config()) - name = model.pop('name') - return {'name': 'model-' + name, 'route_prefix': '/api/v1/model/' + name, - 'import_path': 'model', 'args': model} +def model_application(): + return {'name': 'model-small', 'route_prefix': '/api/v1/model/Qwen/small', + 'import_path': 'model', 'args': { + 'model_id': 'ms://Qwen/small', 'backend': 'transformers', + 'device_group': {'ranks': 2, 'device_type': 'cuda'}, + 'device_mesh': {'fsdp_size': 2}, + }} def test_catalog_defaults_and_mode_selection(): - config = ServerConfig(applications=[model_application()]) - catalog = TrainingCatalog(config.training_config) + config = TrainingConfig(models=[model_config()]) + catalog = TrainingCatalog(config) shared = catalog.get('Qwen/small', 'lora') exclusive = catalog.get('Qwen/small', 'full') assert shared.capacity == 5 assert exclusive.capacity == 1 assert shared.config.min_replicas == exclusive.config.min_replicas == 0 assert shared.config.idle_timeout_s == 1800 - assert config.training.allocation_timeout_s == 3 - assert config.to_yaml_dict()['applications'][0]['args']['model_id'] == 'ms://Qwen/small' + assert config.allocation_timeout_s == 3 with pytest.raises(RequestRejectedError, match='Unknown training model'): catalog.get('not-in-catalog', 'lora') @@ -163,11 +164,20 @@ def test_legacy_lifecycle_overrides_and_full_instances(): assert config.training_config.models[0].modes.lora is None -def test_extended_modes_do_not_accept_ambiguous_single_mode_settings(): - app = model_application() - app['args']['train_mode'] = 'full' - with pytest.raises(ValidationError, match='cannot be mixed'): - ServerConfig(applications=[app]) +def test_separate_model_applications_select_lora_and_full(): + shared = model_application() + exclusive = copy.deepcopy(shared) + exclusive.update(name='model-small-full', route_prefix='/api/v1/model/Qwen/small-full') + exclusive['args'].update(train_mode='full', min_replicas=0, max_replicas=2, idle_timeout_s=60) + config = ServerConfig(applications=[shared, exclusive]) + catalog = TrainingCatalog(config.training_config) + assert catalog.get('Qwen/small', 'lora').capacity == 5 + full = catalog.get('Qwen/small-full', 'full') + assert full.capacity == 1 + assert (full.config.min_replicas, full.config.max_replicas, full.config.idle_timeout_s) == (0, 2, 60) + with pytest.raises(RequestRejectedError, match='does not provide full'): + catalog.get('Qwen/small', 'full') + assert ServerConfig.model_validate(config.to_yaml_dict()).training_config == config.training_config def test_launcher_keeps_training_and_sampling_capabilities_separate(monkeypatch): @@ -205,6 +215,18 @@ def legacy_application(): }]} +@pytest.mark.parametrize('extra', [ + {'modes': None}, + {'modes': {'lora': {'num_devices': 2, 'device_mesh': {'dp_size': 2}}}}, + {'device_type': 'cpu'}, +]) +def test_model_yaml_rejects_alternative_resource_and_mode_syntax(extra): + app = legacy_application() + app['args'].update(extra) + with pytest.raises(ValidationError, match='Extra inputs are not permitted'): + ServerConfig(applications=[app]) + + def test_legacy_options_reach_cpu_router_without_changing_yaml(monkeypatch): from types import SimpleNamespace from twinkle.server.launcher.server_launcher import ServerLauncher @@ -294,7 +316,6 @@ def test_l20d_candidate_is_lora_only_and_bounds_total_devices(): assert (sampler.args.min_replicas, sampler.args.max_replicas, sampler.args.target_ongoing_requests) == (2, 4, 16) assert mode.num_devices * mode.max_replicas + sampler.args.device_group['ranks'] * sampler.args.max_replicas == 8 assert mode.max_loras * mode.max_replicas == 6 - assert config.applications[-1].args.modes is None with pytest.raises(RequestRejectedError, match='does not provide full'): TrainingCatalog(config.training_config).get(model.name, 'full') @@ -310,12 +331,14 @@ def test_single_device_megatron_is_rejected_before_allocation(mode, legacy): model = model_config() model['backend'] = 'megatron' model['modes'] = {mode: {'num_devices': 1, 'device_mesh': {'dp_size': 1}}} - app = model_application(model) with pytest.raises(ValidationError, match='at least 2 devices'): - ServerConfig(applications=[app]) + if legacy: + ServerConfig(applications=[app]) + else: + TrainingConfig(models=[model]) def test_transformers_still_accepts_single_device(): - model = model_config() - model['modes'] = {'lora': {'num_devices': 1, 'device_mesh': {'dp_size': 1}}} - assert ServerConfig(applications=[model_application(model)]).training_config.models[0].modes.lora.num_devices == 1 + app = model_application() + app['args'].update(device_group={'ranks': 1, 'device_type': 'cuda'}, device_mesh={'dp_size': 1}) + assert ServerConfig(applications=[app]).training_config.models[0].modes.lora.num_devices == 1 diff --git a/tests/server/integration/on_demand_training/http_probe.py b/tests/server/integration/on_demand_training/http_probe.py index b469e7593..5d8f56248 100644 --- a/tests/server/integration/on_demand_training/http_probe.py +++ b/tests/server/integration/on_demand_training/http_probe.py @@ -54,30 +54,18 @@ def main(): 'args': { 'model_id': 'small', 'backend': 'mock', - 'device_type': 'cpu', - 'modes': { - 'lora': { - 'num_devices': 1, - 'device_mesh': { - 'dp_size': 1 - }, - 'max_loras': 2, - 'max_replicas': 1 - }, - 'full': { - 'num_devices': 1, - 'device_mesh': { - 'dp_size': 1 - } - } - } + 'train_mode': 'lora', + 'device_group': {'ranks': 1, 'device_type': 'cpu'}, + 'device_mesh': {'dp_size': 1}, + 'max_loras': 2, + 'max_replicas': 1, }, }, { 'name': 'other', 'import_path': 'model', 'route_prefix': '/api/v1/model/other', 'args': { - 'model_id': 'other', + 'model_id': 'small', 'backend': 'mock', 'train_mode': 'full', 'device_group': {'ranks': 1, 'device_type': 'cpu'}, @@ -129,7 +117,8 @@ def wait_epoch(controller, old_epoch): transport = ClientTransport(ClientContext(base_url=base, api_key='probe-owner', session_id=session)) first = MultiLoraTransformersModel('small', transport=transport) second = MultiLoraTransformersModel('small', transport=transport) - first.add_adapter_to_model('same', None) + from peft import LoraConfig + first.add_adapter_to_model('same', LoraConfig(r=8, target_modules='all-linear')) instance = next( item for item in ray.get(controller.diagnostics.remote())['instances'] if first._training_ids['same'] in item['training_ids']) @@ -142,17 +131,16 @@ def wait_epoch(controller, old_epoch): headers={'Authorization': 'Bearer probe-owner'}, timeout=10) assert bypass.status_code == 400 and 'catalog endpoint' in bypass.text, bypass.text - second.add_adapter_to_model('same', None) + second.add_adapter_to_model('same', LoraConfig(r=8, target_modules='all-linear')) old_id = first._training_ids['same'] assert old_id != second._training_ids['same'] assert json.loads(first.get_train_configs().result) == {} first.zero_grad() first.remove_adapter() - first.add_adapter_to_model('same', None) + first.add_adapter_to_model('same', LoraConfig(r=8, target_modules='all-linear')) assert old_id != first._training_ids['same'] first.remove_adapter() second.remove_adapter() - from peft import LoraConfig native_lora = MultiLoraTransformersModel('small', transport=transport) native_lora.add_adapter_to_model('native-lora', LoraConfig(r=8, target_modules='all-linear')) native_lora.zero_grad() @@ -174,7 +162,7 @@ def wait_epoch(controller, old_epoch): }) assert shared1.forward_backward([datum], 'cross_entropy').result().metrics shared1.optim_step(types.AdamParams(learning_rate=1e-4)).result() - exclusive = service.create_full_training_client(base_model='small') + exclusive = service.create_full_training_client(base_model='other') assert not exclusive.get_info().is_lora assert exclusive.forward_backward([datum], 'cross_entropy').result().metrics exclusive.optim_step(types.AdamParams(learning_rate=1e-4)).result() @@ -208,7 +196,8 @@ def wait_epoch(controller, old_epoch): before = ray.available_resources().get('CPU', 0) for index in range(10): model = MultiLoraTransformersModel('small' if index % 2 == 0 else 'other', transport=transport) - model.add_adapter_to_model('cycle', None) + model.add_adapter_to_model('cycle', LoraConfig(r=8, target_modules='all-linear') + if index % 2 == 0 else None) model.zero_grad() model.remove_adapter() print(json.dumps({'phase': 'cycles', 'completed': index + 1}), flush=True) diff --git a/tests/server/integration/on_demand_training/matrix_validation.py b/tests/server/integration/on_demand_training/matrix_validation.py index 0b2883bc0..5528b0197 100644 --- a/tests/server/integration/on_demand_training/matrix_validation.py +++ b/tests/server/integration/on_demand_training/matrix_validation.py @@ -157,7 +157,7 @@ def submit(_): return True -def validate_capacity(base_url, controller, model_name, capacity): +def validate_capacity(base_url, controller, model_name, full_model_name, capacity): """Assert shared LoRA capacity and full-training exclusivity on real devices.""" from tests.server.integration.on_demand_training.acceptance_support import TrainingHandle, assert_released handles = [] @@ -179,13 +179,13 @@ def validate_capacity(base_url, controller, model_name, capacity): for handle in reversed(handles): handle.close() assert_released(controller, ray.get_runtime_context().namespace) - exclusive = TrainingHandle(base_url, model_name, 'twinkle', 'exclusive', mode='full') + exclusive = TrainingHandle(base_url, full_model_name, 'twinkle', 'exclusive', mode='full') try: exclusive.step() live = [row for row in ray.get(controller.diagnostics.remote())['instances'] if row['state'] != 'STOPPED'] assert len(live) == 1 and live[0]['mode'] == 'full' and len(live[0]['training_ids']) == 1 try: - overflow = TrainingHandle(base_url, model_name, 'tinker', 'exclusive-overflow', mode='full') + overflow = TrainingHandle(base_url, full_model_name, 'tinker', 'exclusive-overflow', mode='full') except Exception as exc: assert 'replica limit' in str(exc).lower(), exc else: @@ -289,11 +289,12 @@ def main(): } if args.persistence == 'redis' else {'mode': 'memory'}) raw['ray_namespace'] = namespace raw['http_options'] = {'host': '127.0.0.1', 'port': port} + full_applications = [] for app in raw['applications']: params = app['args'] if app['import_path'] == 'model': common = { - 'num_devices': args.training_devices, + 'device_group': {'ranks': args.training_devices, 'device_type': 'cuda'}, 'device_mesh': {'device_type': 'cuda', 'dp_size': args.training_devices // args.training_pp_size, 'pp_size': args.training_pp_size}, 'min_replicas': 0, 'max_replicas': 1, 'idle_timeout_s': 3600, @@ -301,14 +302,22 @@ def main(): 'queue_config': {**params.get('queue_config', {}), 'execution_timeout': 600}, } app['args'] = { - 'model_id': args.model, 'backend': args.backend, 'device_type': 'cuda', - 'modes': {'lora': {**common, 'max_loras': 10, 'max_r': 32}, 'full': common}, + **common, 'model_id': args.model, 'backend': args.backend, 'train_mode': 'lora', + 'max_loras': 10, 'max_r': 32, 'env_vars': {'TWINKLE_FAIL_FAST': '0', 'TWINKLE_TRUST_REMOTE_CODE': '1'}, } app.pop('deployments', None) + full_applications.append({ + 'name': app['name'] + '-full', + 'route_prefix': app['route_prefix'] + '-full', + 'import_path': 'model', + 'args': {**common, 'model_id': args.model, 'backend': args.backend, 'train_mode': 'full', + 'env_vars': app['args']['env_vars']}, + }) elif app['import_path'] == 'sampler': params['model_id'] = args.model params['engine_args']['attention_config'] = {'flash_attn_version': 3} + raw['applications'].extend(full_applications) if args.recovery_only: raw['applications'] = [app for app in raw['applications'] if app['import_path'] in ('server', 'model')] config = ServerConfig.model_validate(raw) @@ -320,7 +329,8 @@ def main(): 'TWINKLE_STATE_RAY_NAMESPACE': namespace, 'TWINKLE_PERSISTENCE_MODE': args.persistence, 'TWINKLE_DEFAULT_SAVE_DIR': str(run_dir / 'checkpoints'), - 'TWINKLE_MODEL_ID_ALIASES': json.dumps({'Qwen/Qwen3.5-4B': args.model}), + 'TWINKLE_MODEL_ID_ALIASES': json.dumps({model.name: model.model_id + for model in config.training_config.models}), 'TWINKLE_TRUST_REMOTE_CODE': '1', 'TWINKLE_FAIL_FAST': '0', 'TWINKLE_TEST_GPU_E2E': '1', @@ -425,6 +435,7 @@ def main(): base_url, launcher._training_controller, model_name) evidence['capacity_checked'] = validate_capacity( base_url, launcher._training_controller, model_name, + next(model.name for model in config.training_config.models if model.modes.full is not None), config.training_config.models[0].modes.lora.max_loras) if args.stress_iterations or args.soak_seconds: from tests.server.integration.on_demand_training.acceptance_support import assert_released diff --git a/tests/server/integration/on_demand_training/test_resident_release_e2e.py b/tests/server/integration/on_demand_training/test_resident_release_e2e.py index dfd3e3c4d..39f41f340 100644 --- a/tests/server/integration/on_demand_training/test_resident_release_e2e.py +++ b/tests/server/integration/on_demand_training/test_resident_release_e2e.py @@ -87,12 +87,13 @@ def test_resident_release_ack_after_worker_failure(tmp_path): {'name': 'server', 'route_prefix': '/api/v1', 'import_path': 'server', 'args': {'supported_models': [model_name]}}, {'name': 'model', 'route_prefix': '/api/v1/model/' + model_name, 'import_path': 'model', - 'args': {'model_id': model, 'backend': backend, 'device_type': 'cuda', - 'modes': {'full': {'num_devices': 2, 'min_replicas': 1, 'max_replicas': 1, - 'device_mesh': {'dp_size': 2 if backend == 'transformers' else 1, - 'pp_size': 1 if backend == 'transformers' else 2}, - 'max_length': 512, 'idle_timeout_s': 3600, - 'queue_config': {'execution_timeout': 600}}}, + 'args': {'model_id': model, 'backend': backend, 'train_mode': 'full', + 'device_group': {'ranks': 2, 'device_type': 'cuda'}, + 'min_replicas': 1, 'max_replicas': 1, + 'device_mesh': {'dp_size': 2 if backend == 'transformers' else 1, + 'pp_size': 1 if backend == 'transformers' else 2}, + 'max_length': 512, 'idle_timeout_s': 3600, + 'queue_config': {'execution_timeout': 600}, 'env_vars': {'TWINKLE_TRUST_REMOTE_CODE': '1', 'TWINKLE_FAIL_FAST': '0'}}}, ], } diff --git a/tests/server/training/test_launcher.py b/tests/server/training/test_launcher.py index d212e43e0..cd8ee7613 100644 --- a/tests/server/training/test_launcher.py +++ b/tests/server/training/test_launcher.py @@ -19,9 +19,9 @@ def test_failed_initialization_does_not_leave_an_actor_that_blocks_a_retry(): payload = { 'ray_namespace': namespace, 'http_options': {'host': '127.0.0.1', 'port': port}, 'applications': [{'name': 'retry', 'route_prefix': '/api/v1/model/retry', 'import_path': 'model', - 'args': {'model_id': 'mock', 'backend': 'mock', 'device_type': 'cpu', - 'modes': {'full': {'num_devices': devices, 'device_mesh': {'dp_size': devices}, - 'min_replicas': 1}}}}], + 'args': {'model_id': 'mock', 'backend': 'mock', 'train_mode': 'full', + 'device_group': {'ranks': devices, 'device_type': 'cpu'}, + 'device_mesh': {'dp_size': devices}, 'min_replicas': 1}}], } launcher = ServerLauncher(ServerConfig.model_validate(payload)) launcher._start_serve() @@ -30,7 +30,7 @@ def test_failed_initialization_does_not_leave_an_actor_that_blocks_a_retry(): launcher._deploy_training() with pytest.raises(ValueError): ray.get_actor(controller_name(namespace), namespace=namespace) - payload['applications'][0]['args']['modes']['full']['min_replicas'] = 0 + payload['applications'][0]['args']['min_replicas'] = 0 launcher.config = ServerConfig.model_validate(payload) launcher._deploy_training() diagnostic = ray.get(launcher._training_controller.diagnostics.remote()) diff --git a/tests/twinkle_client/test_auto_agent_tools.py b/tests/twinkle_client/test_auto_agent_tools.py index c4ea1a817..41c5216a5 100644 --- a/tests/twinkle_client/test_auto_agent_tools.py +++ b/tests/twinkle_client/test_auto_agent_tools.py @@ -82,7 +82,7 @@ def test_generated_server_config_matches_launcher_schema(tmp_path, monkeypatch, assert gateway.deployments[0]['name'] == 'GatewayServer' assert gateway.args.supported_models == ['Qwen/student'] + [s['model_id'] for s in samplers] application = next(app for app in config.applications if app.import_path == 'model') - assert application.args.modes is None and application.args.train_mode == 'lora' + assert application.args.train_mode == 'lora' model = config.training_config.models[0] assert model.modes.full is None assert model.modes.lora.min_replicas == 0 and model.modes.lora.max_replicas == 1 From c9926d42d3c2f8c46cf6eb4636a359306678abf1 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 14:48:19 +0800 Subject: [PATCH 27/46] fix: preserve default CPU reservations for standalone registries --- src/twinkle/infra/_ray/ray_helper.py | 3 +- tests/infra/test_ray_registry.py | 62 ++++++++++++++++++++++++++++ 2 files changed, 64 insertions(+), 1 deletion(-) create mode 100644 tests/infra/test_ray_registry.py diff --git a/src/twinkle/infra/_ray/ray_helper.py b/src/twinkle/infra/_ray/ray_helper.py index d676f4661..5dd7aa878 100644 --- a/src/twinkle/infra/_ray/ray_helper.py +++ b/src/twinkle/infra/_ray/ray_helper.py @@ -81,7 +81,8 @@ def clear(self): name=registry_name, namespace=registry_namespace, lifetime='detached', - num_cpus=0.01 if registry_namespace else 1).remote() + # Preserve Ray's default CPU policy outside managed instances. + num_cpus=0.01 if registry_namespace else None).remote() except ValueError: RayHelper._registry = ray.get_actor(registry_name, namespace=registry_namespace) assert RayHelper._registry is not None diff --git a/tests/infra/test_ray_registry.py b/tests/infra/test_ray_registry.py new file mode 100644 index 000000000..f3b4c6f2e --- /dev/null +++ b/tests/infra/test_ray_registry.py @@ -0,0 +1,62 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Registry isolation must not reserve a CPU for ordinary non-server use.""" +import subprocess +import sys +import textwrap + + +def test_registry_cpu_reservations_and_instance_isolation(): + # Own a subprocess and cluster so the CPU assertions do not count actors + # belonging to other tests or reset their process-global RayHelper cache. + code = textwrap.dedent(''' + import os + import time + import uuid + import ray + from twinkle.infra._ray.ray_helper import RayHelper + + namespace = 'registry-' + uuid.uuid4().hex[:8] + ray.init(address='local', num_cpus=2, namespace=namespace, + include_dashboard=False, log_to_driver=False, + _temp_dir='/tmp/twreg-' + uuid.uuid4().hex[:6]) + + def registry(name, managed): + RayHelper._registry = None + os.environ['TWINKLE_RAY_REGISTRY_NAME'] = name + if managed: + os.environ['TWINKLE_RAY_NAMESPACE'] = namespace + else: + os.environ.pop('TWINKLE_RAY_NAMESPACE', None) + RayHelper.init_registry() + actor = RayHelper._registry + ray.get(actor.add_config.remote('owner', name)) + return actor + + def assert_free_cpus(expected): + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + available = ray.available_resources().get('CPU', 0) + if abs(available - expected) < 1e-6: + return + time.sleep(.01) + raise AssertionError((available, expected)) + + try: + ordinary = registry('ordinary', False) + assert_free_cpus(2) + ray.kill(ordinary, no_restart=True) + left = registry('left', True) + right = registry('right', True) + assert_free_cpus(1.98) + assert ray.get(left.get_config.remote('owner')) == 'left' + assert ray.get(right.get_config.remote('owner')) == 'right' + ray.kill(left, no_restart=True) + assert_free_cpus(1.99) + assert ray.get(right.get_config.remote('owner')) == 'right' + ray.kill(right, no_restart=True) + assert_free_cpus(2) + finally: + ray.shutdown() + ''') + result = subprocess.run([sys.executable, '-c', code], capture_output=True, text=True, timeout=120) + assert result.returncode == 0, result.stdout + result.stderr From 45eba79f459a0c706bd6e7f5a43a7a393726b552 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 14:48:33 +0800 Subject: [PATCH 28/46] docs: explain on-demand configuration for service users --- .../Usage Guide/Server and Client/Server.md | 65 +++++++++---------- .../\346\234\215\345\212\241\347\253\257.md" | 65 +++++++++---------- 2 files changed, 64 insertions(+), 66 deletions(-) diff --git a/docs/source_en/Usage Guide/Server and Client/Server.md b/docs/source_en/Usage Guide/Server and Client/Server.md index e6fa99611..4f94d8989 100644 --- a/docs/source_en/Usage Guide/Server and Client/Server.md +++ b/docs/source_en/Usage Guide/Server and Client/Server.md @@ -482,37 +482,44 @@ holds the cleanup lease. ## On-demand training and device reclamation -Existing `applications` entries with `import_path: model` remain valid. Set `args.min_replicas: 0` for lazy startup and explicitly set `args.idle_timeout_s` for business inactivity. `adapter_config.adapter_timeout` remains a heartbeat timeout and `adapter_max_lifetime` retains its absolute lifetime. Heartbeats and polling do not renew the business inactivity deadline; queued and running operations protect training state. Empty instances above the resident minimum are destroyed so other models can reuse their devices. Existing deployment min/max replicas become training instance minimums/limits. +Training services can load models when a request arrives and release devices when they are no longer needed. Keep the existing `applications` format and configure the `args` of an `import_path: model` entry: -Training uses the existing YAML format: select `args.train_mode: lora`, `hybrid`, or `full`, and configure devices through `device_group` and `device_mesh`. Model `deployments` options configure the CPU router; autoscaling request targets apply to router replicas. The Controller owns training-instance minimums/limits, and each compute instance has exactly one Serve replica. Non-autoscaling request limits, health/shutdown settings, and `ray_actor_options.runtime_env` also reach compute instances. `args.env_vars` overrides matching deployment environment variables. Unsupported `adapter_config` fields fail validation instead of being silently ignored. - -To offer both LoRA and full training for the same weights, declare two `import_path: model` applications with distinct `name` and `route_prefix` values. Set `args.train_mode: lora` on one and `full` on the other; both may use the same `model_id`. Clients select the model alias from the corresponding route. LoRA users share compatible instances; each full-training handle owns an instance. Sampling configuration stays the same. Complete examples are in `cookbook/client/server/transformer/server_config.yaml` and `cookbook/client/server/megatron/server_config.yaml`. +```yaml +- name: training-qwen + route_prefix: /api/v1/model/Qwen/Qwen3.5-4B + import_path: model + args: + model_id: ms://Qwen/Qwen3.5-4B + backend: transformers + train_mode: lora + nproc_per_node: 2 + device_group: {name: model, ranks: 2, device_type: cuda} + device_mesh: {device_type: cuda, dp_size: 2} + max_loras: 3 + min_replicas: 0 + max_replicas: 2 + idle_timeout_s: 600 +``` -### Single-node, eight-GPU LoRA capacity +| Setting | How to use it | +|---------|---------------| +| `device_group`, `device_mesh` | Set the devices and parallel layout for each training instance. Adjust them to fit your model and available devices. Megatron requires at least two devices per instance. | +| `min_replicas` | Training instances to keep loaded at startup. With `0`, the first training request waits for model loading. | +| `max_replicas` | Maximum number of training instances. | +| `max_loras` | Maximum LoRA trainings per instance. Used only for LoRA training. | +| `idle_timeout_s` | Seconds of training inactivity before its handle is released. Heartbeats and result polling do not extend this time. Queued or running training operations are protected. | -`cookbook/client/server/megatron/server_config.yaml` is a candidate for large-memory L20D GPUs. It retains single-mode YAML: two preloaded single-GPU sampling instances (DP=1, no TP setting), scaling up to four, plus up to two two-GPU training instances (PP=2, DP=1). Megatron configurations require at least two devices per instance before allocation; single-device Megatron is not supported. Only `train_mode: lora` is configured; full training remains available in other framework configurations. Key model args: +This example allows two instances, each holding three LoRA trainings. Clients are also subject to `server_config.per_token_model_limit`; clients using the same token share that token's quota. -```yaml -train_mode: lora -nproc_per_node: 2 -device_group: {name: model, ranks: 2, device_type: cuda} -device_mesh: {device_type: cuda, dp_size: 1, pp_size: 2} -max_loras: 3 -min_replicas: 0 -max_replicas: 2 -idle_timeout_s: 600 -``` +For full-parameter training, set `train_mode: full` and remove `max_loras`. Each full training uses its own instance. To offer LoRA and full training for the same weights, configure two model applications with different `name` and `route_prefix` values and select the mode in each. Clients choose the model alias from the corresponding route. See [Hybrid LoRA Training](../Hybrid-LoRA-Training.md) for Hybrid LoRA configuration. -- `max_replicas × max_loras = 6` bounds training handles, not throughput. Creation fills existing compatible LoRA slots before starting another instance; active training is not migrated for load balancing. Stateful operations are serial within an instance, with at most two independent training queues. `per_token_model_limit: 3` remains a server-wide per-token quota, not a total-user cap. Clients sharing one token share those three slots; independent users need independent tokens, or a deliberately increased shared-token quota. Set it to 1 separately if the service requires one handle per independent token. -- PP=2 shards each model over two devices; removing DP replication from the original PP=2/DP=2 instance leaves room for two independent instances. Account for 32K activations, LoRA optimizers, communication and temporary workspaces. Smaller-model or different-hardware E2E tests do not certify L20D/27B/32K peak memory or throughput. -- Save checkpoints during training and explicitly call `remove_adapter()` or unload the Tinker model when finished. Heartbeats and polling do not renew business idle. Expired handles require a new create and explicit checkpoint resume; old training is not resumed automatically. -- `min_replicas: 0` permits reclaiming every empty training instance at the cost of cold startup. Both pools use the same device-allocation and instance-lifecycle implementation with independent state. Reclaiming training devices does not directly scale sampling; sampling grows according to its own in-flight load. Lowering vLLM `gpu_memory_utilization` does not release Ray's whole-device reservation. No fractional-GPU or training/sampling colocation is enabled. +Release models when training finishes; Twinkle clients can call `remove_adapter()`. After the last training releases an instance, empty instances above `min_replicas` shut down. Save checkpoints regularly. If a handle expires or the service restarts, create a new training and restore from a checkpoint. Upgrade the Twinkle SDK together with the server. -Before rollout, validate the actual 27B model on L20D: a two-GPU PP=2/DP=1 32K workload at the maximum business batch, three resident rank-32 LoRAs, two instances/six handles alongside sampling, optimizer save/resume, tenant isolation and confirmed device reclamation. Measure peak memory, training tokens/s, sampling output tokens/s, queue wait and P95 latency against the original four-GPU PP=2/DP=2 layout using the same work. If the two-GPU layout does not meet business requirements, restore four GPUs per instance (PP=2, DP=2, at most one instance); do not infer throughput from an idle memory snapshot. +Complete configurations are available in `cookbook/client/server/transformer/server_config.yaml` and `cookbook/client/server/megatron/server_config.yaml`. -### On-demand sampling instances and dispatch +### Automatic sampling scaling -A stable CPU Router selects the ready instance with the fewest actual in-flight operations for ordinary sampling, DataRefs and streaming. Clients do not supply an instance ID, and users/LoRAs are not pinned to a device. Per-instance DP/TP describes an engine's parallel shape, not pool replicas. The production candidate uses one GPU and DP=1 per instance, without TP configuration. +Set instance counts and the scaling target in the `args` of your `import_path: sampler` entry. For example: ```yaml nproc_per_node: 1 @@ -524,14 +531,6 @@ target_ongoing_requests: 16 idle_timeout_s: 600 ``` -Both minimum instances must finish loading before readiness. Demand is `ceil((in-flight requests + 1) / 16)`, capped at four. Starting and quarantined instances count against that cap. Existing ready instances serve during expansion; only idle extras with no actual generation are reclaimed after 600 seconds, down to two. Neither pool preempts or colocates on a GPU; four sampling plus four training GPUs fit the eight-device maximum. - -New instances load only the base model and default template. LoRAs load from each request's `adapter_uri`, `model_path` or persistent sampling-session checkpoint path, into independent disposable caches. No cache synchronization or new cache database is introduced. `engine_args.max_loras: 5` bounds active engine LoRAs, not training handles. Custom templates, named adapter settings and patches are broadcast under pool-wide exclusive protection and replayed before new instances become ready. Partial broadcast failure fences all participating replicas. Controller restart clears the old pool and dynamic control settings; clients must reapply them, while existing sampling-session/checkpoint paths remain unchanged. - -HTTP disconnect and future timeout do not certify generation exit. Completion handles and weight guards remain held until actual exit; unconfirmed death remains `QUARANTINED` without early capacity release. Full-checkpoint sampling retains pool-wide exclusive weight changes. L20D/27B single-GPU 32K memory, four-engine concurrency, output tokens/s and P95 latency require controlled same-hardware validation; local PPU/4B tests cannot certify them. - -### Rollout compatibility and safe states - -Keeping YAML syntax does not keep the old SDK wire unchanged. The Twinkle SDK performs the `training_handles` handshake, generates `client_create_id` and supplies `training_id`; upgrade it together with the server. Tinker keeps its own model-handle wire. +This configuration loads two sampling instances at startup and scales up to four as requests increase. Extra instances shut down after 600 seconds of continuous inactivity, leaving two loaded. `target_ongoing_requests` is the target number of in-flight requests per instance; a smaller value triggers earlier scaling. `device_group` and `device_mesh` determine the devices used by each instance. Allow enough devices for both training and sampling. -Ray 2.55 is the minimum dependency: Serve ingress must await asynchronous user constructors; do not use the old CI Ray 2.48 pin. Startup, business idle, task execution and confirmed shutdown have separate budgets. Atomic create start and cancel are mutually exclusive; an unknown state-backend acknowledgement must not authorize execution. `QUARANTINED` means device death is unconfirmed and still counts toward the instance limit; release must not report unconfirmed reclamation as success. Only an actually executed operation renews business idle, not a queue failure or unknown start acknowledgement. +Clients keep using the existing sampling API without selecting an instance. LoRA checkpoints load as requests need them. After a service restart, reapply custom templates and other runtime settings. diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" index 084fc1901..4711f11bb 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" @@ -469,37 +469,44 @@ twinkle-server check-config -c server_config.yaml ## 按需训练与设备回收 -原有 `applications` 和 `import_path: model` 配置继续有效。单模式配置增加 `args.min_replicas: 0` 即可懒启动，`args.idle_timeout_s` 显式设置训练句柄的业务闲置时间。`adapter_config.adapter_timeout` 仍按心跳判断会话存活，`adapter_max_lifetime` 保留绝对寿命；不再把旧字段转换成业务闲置时间。心跳和轮询不延长业务闲置期限，运行中或已排队的任务保护训练状态。最后一个用户释放或过期后，超出保底数量的实例立即销毁，释放的卡可用于其他模型。已有 deployment 的 min/max_replicas 转成训练实例的保底/上限。 +训练服务可以在收到请求时再加载模型，用完后自动释放设备。继续使用原来的 `applications` 配置即可，在 `import_path: model` 的 `args` 中设置： -训练统一使用原有 YAML 格式，通过 `args.train_mode: lora`、`hybrid` 或 `full` 选择训练模式，通过 `device_group` 和 `device_mesh` 配置设备。原模型 `deployments` 的 Serve 设置用于 CPU 路由部署，其中 autoscaling 的请求目标用于路由副本扩缩容；训练计算实例的 min/max 仍由 Controller 管理，每个计算实例固定一个 Serve replica。非 autoscaling 的请求上限、健康/关闭设置和 `ray_actor_options.runtime_env` 会传给计算实例；`args.env_vars` 优先于 deployment 的同名环境变量。未知或未实现的 `adapter_config` 字段会报错，不再静默忽略。 - -同一权重需要同时提供 LoRA 和全参数训练时，声明两个 `import_path: model` 的 application，使用不同的 `name` 和 `route_prefix`，分别设置 `args.train_mode: lora` 和 `full`；`model_id` 可以指向同一权重。客户端使用各自路由中的模型别名。LoRA 用户共享兼容实例，全参数用户各占一个实例。采样服务继续使用原配置；完整示例见 `cookbook/client/server/transformer/server_config.yaml` 和 `cookbook/client/server/megatron/server_config.yaml`。 +```yaml +- name: training-qwen + route_prefix: /api/v1/model/Qwen/Qwen3.5-4B + import_path: model + args: + model_id: ms://Qwen/Qwen3.5-4B + backend: transformers + train_mode: lora + nproc_per_node: 2 + device_group: {name: model, ranks: 2, device_type: cuda} + device_mesh: {device_type: cuda, dp_size: 2} + max_loras: 3 + min_replicas: 0 + max_replicas: 2 + idle_timeout_s: 600 +``` -### 单机 8 卡、仅 LoRA 的容量配置 +| 配置项 | 如何使用 | +|--------|----------| +| `device_group`、`device_mesh` | 配置每个训练实例使用的设备和并行方式，请按模型大小与可用设备调整。Megatron 每实例至少需要 2 张设备。 | +| `min_replicas` | 启动时保留的训练实例数。设为 `0` 时，首次训练需要等待模型加载。 | +| `max_replicas` | 最多启动多少个训练实例。 | +| `max_loras` | 每个实例最多同时容纳多少个 LoRA 训练，仅用于 LoRA 模式。 | +| `idle_timeout_s` | 训练闲置多久后自动释放句柄，单位为秒。心跳和结果轮询不会延长这个时间，排队中或正在执行的训练不会因此中断。 | -`cookbook/client/server/megatron/server_config.yaml` 提供大显存 L20D 的候选方案，保持单模式 YAML：采样每实例 1 卡、DP=1，不配置 TP，默认加载 2 实例、最多 4 实例；训练每实例 2 卡（PP=2、DP=1），最多 2 实例。Megatron 配置在分配资源前校验每实例至少 2 卡，不提供单卡路径。仅配置 `train_mode: lora`，不提供 full 模式；框架仍支持其他配置中的 full。关键训练参数： +上例最多提供 2 个实例，每个实例容纳 3 个 LoRA 训练。客户端同时还受 `server_config.per_token_model_limit` 限制；多个客户端使用同一个 token 时，共享该 token 的配额。 -```yaml -train_mode: lora -nproc_per_node: 2 -device_group: {name: model, ranks: 2, device_type: cuda} -device_mesh: {device_type: cuda, dp_size: 1, pp_size: 2} -max_loras: 3 -min_replicas: 0 -max_replicas: 2 -idle_timeout_s: 600 -``` +使用全参数训练时，将 `train_mode` 改为 `full`，并删除 `max_loras`。每个全参数训练单独使用一个实例。同一权重需要同时提供 LoRA 和全参数训练时，配置两个 model application，分别设置 `train_mode`，并使用不同的 `name` 和 `route_prefix`。客户端通过对应路由中的模型别名选择服务。Hybrid LoRA 的配置见 [Hybrid LoRA 训练](../Hybrid%20LoRA训练.md)。 -- `max_replicas × max_loras = 6` 是训练句柄容量上限，不是吞吐保证；创建时优先填满已有实例的 LoRA 槽位，槽位不足才创建新实例，不会按空闲度迁移现有训练。同实例的 forward/backward/step 串行执行，最多有 2 条独立训练队列。`per_token_model_limit: 3` 仍是全服务的每 token 配额，不是总用户上限；不同用户使用独立 token 才能分别获得配额。同一 token 的所有客户端合计仍只有 3 个名额；共享 token 的部署需按业务调大该配额，不能仅增加实例数。若希望避免单个独立 token 占多个名额，可另设为 1。 -- 两卡实例用 PP=2 切分模型，取消原 PP=2/DP=2 实例内的 DP 副本，以容纳两个独立实例。还需考虑 32K 激活、LoRA optimizer、通信及临时 workspace；较小模型或其他硬件的 E2E 不能证明 L20D/27B/32K 的峰值和吞吐。 -- 业务完成后主动 `remove_adapter()` 或卸载 Tinker 模型；运行期应定期保存 checkpoint。心跳/结果轮询不会续业务 idle，句柄过期后必须重新创建并显式 resume；不会自动恢复旧训练。 -- `min_replicas: 0` 能回收所有空训练实例，但首次训练需冷加载。训练和采样通过共用的设备申请与实例启停机制管理资源，两池使用独立状态。释放训练设备预约不主动触发采样扩容，采样按自身在途请求扩容；降低 `gpu_memory_utilization` 也不会释放 Ray 的整卡预约。不启用 fractional GPU 或训练/采样同卡共置。 +训练结束后请主动释放模型，Twinkle 客户端可调用 `remove_adapter()`。最后一个训练释放后，超出 `min_replicas` 的空闲实例会关闭。训练期间应定期保存 checkpoint；句柄过期或服务重启后，需要重新创建训练并从 checkpoint 恢复。升级服务端时，也请配套升级 Twinkle SDK。 -上线前使用同型号 L20D 与实际 27B 模型，依次验证两卡 PP=2/DP=1 的 32K 最大业务 batch、3 个 rank=32 LoRA 同驻留、2 实例共 6 句柄与 sampler 并行、save + optimizer resume、训练之间隔离和最后用户退出后的设备回收。记录峰值显存、训练 tokens/s、采样输出 tokens/s、队列等待及 P95 延迟；与原 PP=2/DP=2 四卡方案使用相同工作量比较。若两卡方案不满足业务要求，可恢复四卡、PP=2/DP=2、最多 1 实例；不凭显存快照承诺缩卡后的吞吐。 +完整配置可参考仓库中的 `cookbook/client/server/transformer/server_config.yaml` 和 `cookbook/client/server/megatron/server_config.yaml`。 -### Sampler 按需实例与请求分派 +### 采样服务自动扩缩容 -固定 CPU Router 将普通采样、DataRef 与流式请求分派给真实在途请求最少的 ready 实例，不要求客户端传实例 ID，也不把用户固定绑定某张卡。实例内 DP/TP 描述一个引擎的并行规格，不是池内副本数；线上示例每实例单卡、DP=1，删除 TP 配置。 +在 `import_path: sampler` 的 `args` 中设置实例数量和扩容目标。例如： ```yaml nproc_per_node: 1 @@ -511,14 +518,6 @@ target_ongoing_requests: 16 idle_timeout_s: 600 ``` -保底 2 个实例加载完成后才报告 ready；需求按 `ceil((在途请求数 + 1) / 16)` 计算，上限 4。启动中的实例和死亡未确认的隔离实例均占上限。扩容加载期间现有实例继续服务；闲置 600 秒且没有真实 generation 的超保底实例才回收，缩至 2。两池不抢占、不共享一张卡，训练最多 4 卡加采样最多 4 卡为 8 卡。 - -新实例只加载基础模型和默认 template；LoRA 按请求的 `adapter_uri`、`model_path` 或 sampling session 中的 checkpoint 路径加载，各引擎独立缓存，不同步缓存或新增缓存数据库。`engine_args.max_loras: 5` 是引擎同时活跃 LoRA 数，不等于训练句柄容量。自定义 template、命名 adapter 和 patch 在池级排他保护下广播，新实例 ready 前重放；广播失败将参与实例隔离，避免混用不同配置。Controller 重启清理旧池、动态控制设置失效，客户端须重新设置；已有 sampling session/checkpoint 路径不变。 - -HTTP 断连或 Future 超时不代表 generation 已退出，后台完成句柄和权重保护保持到实际退出；死亡无法确认则保持 `QUARANTINED`，不提前还卡。Full-checkpoint 采样仍使用池级排他权重切换。单卡采样的 L20D/27B/32K 峰值、四实例并发、输出 tokens/s 与 P95 延迟需同型号受控验收，本机 PPU/4B 测试不能替代。 - -### 上线兼容与安全状态 - -配置语法保留不代表 SDK/wire 不变：Twinkle SDK 自动完成 `training_handles` 握手、生成 `client_create_id`、传递 `training_id`；新 server 不接受缺少这些信息的旧 Twinkle SDK，应配套升级。Tinker 仍使用自己的模型句柄与 wire。 +这份配置启动时加载 2 个采样实例，请求增多时最多扩到 4 个；额外实例连续空闲 600 秒后关闭，保留 2 个。`target_ongoing_requests` 是每个实例的目标在途请求数，调小后会更早扩容。每个实例的设备用量由 `device_group` 和 `device_mesh` 决定，请为训练和采样预留足够的设备。 -Ray 依赖下限为 2.55，Serve ingress 必须等待异步用户构造；不要继续使用旧 CI 的 Ray 2.48。资源启动、业务 idle、任务超时和关闭确认是不同预算。创建的原子 start 与 cancel 互斥，状态后端 ACK 不明不得执行；`QUARANTINED` 表示设备死亡尚未确认，仍计入实例上限，release 不应假报释放成功。只有真实执行过的操作续业务 idle，排队失败或 start ACK 不明不续期。 +客户端继续使用原来的采样接口，无需指定实例。采样所需的 LoRA checkpoint 按请求加载；服务重启后，请重新设置自定义 template 等运行时配置。 From 91efeb2b24c16e5fcc0957114caa8548abf5e46a Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 16:49:33 +0800 Subject: [PATCH 29/46] fix(model): register only owned LoRA slots during distributed wrapping --- src/twinkle/model/multi_lora.py | 62 ++++++++++++------- .../test_multi_lora_distributed_slots.py | 22 +++++++ 2 files changed, 61 insertions(+), 23 deletions(-) diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index 1e194e543..960ddf16a 100644 --- a/src/twinkle/model/multi_lora.py +++ b/src/twinkle/model/multi_lora.py @@ -12,7 +12,7 @@ from twinkle import torch_util from twinkle.data_format import InputFeature from twinkle.utils import get_logger -from .multi_lora_target_parameters import TargetParameterLoraManager +from .multi_lora_target_parameters import TargetParameterLoraManager, TargetParameterLoraWrapper logger = get_logger() @@ -124,26 +124,48 @@ def reset_adapter_status(self): else: self.module.set_adapter('lora_0') + def _named_slot_parameters(self, module): + """Identify owned adapter parameters through their actual layer containers.""" + slots = {tenant.adapter_name for tenant in self.loras} + for layer_name, layer in module.named_modules(): + if isinstance(layer, LoraLayer): + containers = layer.adapter_layer_names + elif isinstance(layer, TargetParameterLoraWrapper): + containers = ('lora_A', 'lora_B') + else: + continue + for container_name in containers: + container = getattr(layer, container_name) + for slot in slots & container.keys(): + value = container[slot] + prefix = '.'.join(part for part in (layer_name, container_name, slot) if part) + if isinstance(value, torch.nn.Parameter): + yield prefix, value + else: + for name, parameter in value.named_parameters(): + yield prefix + '.' + name, parameter + @contextmanager def trainable_slots(self): - """Register every preallocated slot with distributed gradient reducers.""" + """Expose all slots while wrapping DDP/FSDP, then restore the active mask. + + Reducers discover parameters once during wrapping. The mask is restored + by adapter path because FSDP can replace the Parameter objects themselves. + """ modules = self.module if isinstance(self.module, list) else [self.module] masks = [{ - name: param.requires_grad - for name, param in module.named_parameters() if 'lora_' in name + name: parameter.requires_grad + for name, parameter in self._named_slot_parameters(module) } for module in modules] try: - for module, mask in zip(modules, masks): - for name, param in module.named_parameters(): - if name in mask: - param.requires_grad_(True) + for module in modules: + for _, parameter in self._named_slot_parameters(module): + parameter.requires_grad_(True) yield finally: - # FSDP may replace Parameters while wrapping; restore by name. for module, mask in zip(modules, masks): - for name, param in module.named_parameters(): - if name in mask: - param.requires_grad_(mask[name]) + for name, parameter in self._named_slot_parameters(module): + parameter.requires_grad_(mask[name]) def activate_adapter(self, tenant_adapter_name: str, call_enable=False): if not self.has_lora(tenant_adapter_name): @@ -580,18 +602,12 @@ def _patch_megatron(_module): # PEFT's add_adapter calls set_adapter(active_adapters) which only keeps the # first adapter's requires_grad=True. We need ALL LoRA params to be trainable # so that MegatronDDP registers them all in its gradient buffers (main_grad). - def _enable_all_lora_grad(_module): - for name, param in _module.named_parameters(): - if 'lora_' in name and not param.requires_grad: - param.requires_grad_(True) - - if isinstance(module, list): - for _m in module: - _enable_all_lora_grad(_m) - else: - _enable_all_lora_grad(module) - self.module = module + modules = module if isinstance(module, list) else [module] + for chunk in modules: + for _, parameter in self._named_slot_parameters(chunk): + parameter.requires_grad_(True) + if not isinstance(module, list): self.lora_layer_names = [name for name, layer in module.named_modules() if isinstance(layer, Linear)] return module diff --git a/tests/transformers/test_multi_lora_distributed_slots.py b/tests/transformers/test_multi_lora_distributed_slots.py index 499b661fb..dd5a92cb8 100644 --- a/tests/transformers/test_multi_lora_distributed_slots.py +++ b/tests/transformers/test_multi_lora_distributed_slots.py @@ -131,3 +131,25 @@ def test_checkpoint_as_first_operation_preserves_fresh_slot_reset_basis(tmp_path manager.acquire_lora('replacement', config) for name, value in manager.get_state_dict('replacement').items(): torch.testing.assert_close(value, initial[name]) + + +def test_slot_registration_never_unfreezes_base_parameters_with_lora_in_their_name(): + from collections import OrderedDict + config = LoraConfig(r=2, target_modules=['lora_projection']) + manager = MultiLora(max_loras=2, max_r=2) + module = manager.patch(torch.nn.Sequential(OrderedDict([ + ('lora_projection', torch.nn.Linear(4, 4)), + ])), target_modules=['lora_projection'], lora_config=config) + base = {name: parameter for name, parameter in module.named_parameters() if '.base_layer.' in name} + assert base + # Respect an explicitly frozen base even when its module name contains lora_. + for parameter in base.values(): + parameter.requires_grad_(False) + manager.reset_adapter_status() + mask = {name: parameter.requires_grad for name, parameter in module.named_parameters()} + with pytest.raises(RuntimeError, match='wrapper failed'): + with manager.trainable_slots(): + assert all(parameter.requires_grad for _, parameter in manager._named_slot_parameters(module)) + assert not any(parameter.requires_grad for parameter in base.values()) + raise RuntimeError('wrapper failed') + assert {name: parameter.requires_grad for name, parameter in module.named_parameters()} == mask From 8c87bbc98d2b1868f507ec61aa819b805c5c7412 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 16:49:35 +0800 Subject: [PATCH 30/46] refactor(server): clarify pool ownership and consolidate runtime contracts --- README.md | 2 +- README_ZH.md | 2 +- cookbook/client/async_rl/server_config.yaml | 4 +- .../client/server/megatron/server_config.yaml | 5 +- .../server/megatron/server_config_4b.yaml | 5 +- .../server/transformer/server_config.yaml | 4 +- .../server_config_hybrid_lora.yaml | 2 - cookbook/client/tinker/dpo.py | 2 +- cookbook/client/tinker/multi_modal.py | 2 +- cookbook/client/tinker/self_cognition.py | 2 +- cookbook/client/tinker/upload_to_hub.py | 2 - cookbook/client/twinkle/upload_to_hub.py | 2 - .../Usage Guide/Introduction-with-Qwen3.5.md | 2 +- docs/source_en/Usage Guide/Quick-Start.md | 2 +- .../Usage Guide/Server and Client/Server.md | 6 - .../Tinker-Compatible-Client.md | 2 +- .../Usage Guide/Train-as-a-Service.md | 2 +- ...00\344\275\263\345\256\236\350\267\265.md" | 2 +- ...53\351\200\237\345\274\200\345\247\213.md" | 2 +- ...71\345\256\242\346\210\267\347\253\257.md" | 2 +- .../\346\234\215\345\212\241\347\253\257.md" | 6 - ...55\347\273\203\346\234\215\345\212\241.md" | 2 +- src/twinkle/protocol/types/base.py | 2 +- src/twinkle/protocol/types/errors.py | 12 +- src/twinkle/protocol/types/lifecycle.py | 25 +- src/twinkle/protocol/types/model.py | 19 +- .../sampler/vllm_sampler/vllm_sampler.py | 2 +- src/twinkle/server/README.md | 59 ++- .../server/checkpoint/checkpoint_manager.py | 4 +- src/twinkle/server/config/__init__.py | 3 +- src/twinkle/server/config/application_spec.py | 44 +- src/twinkle/server/config/backend_dispatch.py | 2 +- src/twinkle/server/config/hybrid.py | 21 - src/twinkle/server/config/resources.py | 55 ++ src/twinkle/server/config/server_config.py | 14 +- src/twinkle/server/config/training.py | 229 ++++++--- .../server/config/training_normalization.py | 71 --- src/twinkle/server/env_propagation.py | 37 -- src/twinkle/server/exceptions.py | 24 + src/twinkle/server/gateway/app.py | 2 +- src/twinkle/server/gateway/openai_handlers.py | 11 +- src/twinkle/server/gateway/routes.py | 6 - src/twinkle/server/gateway/tinker_handlers.py | 42 +- src/twinkle/server/launcher/__init__.py | 1 - .../server/launcher/builder_registry.py | 4 +- .../server/launcher/server_launcher.py | 486 +++++++----------- src/twinkle/server/lifecycle/envelope.py | 36 +- .../{service_proxy.py => lifecycle/proxy.py} | 18 +- src/twinkle/server/lifecycle/submit.py | 86 ++-- src/twinkle/server/model/tinker_handlers.py | 294 ----------- src/twinkle/server/processor/app.py | 2 +- src/twinkle/server/resources/allocation.py | 55 +- src/twinkle/server/resources/runtime.py | 51 +- src/twinkle/server/sampler/actor.py | 30 -- .../server/sampler/backends/__init__.py | 11 - .../server/sampler/backends/streaming.py | 36 -- src/twinkle/server/sampling/__init__.py | 2 + src/twinkle/server/sampling/actor.py | 51 ++ .../{sampler => sampling}/controller.py | 22 +- .../server/{sampler => sampling}/driver.py | 19 +- .../execution}/__init__.py | 0 .../{sampler => sampling/execution}/app.py | 17 +- .../sampling/execution/backends/__init__.py | 2 + .../execution}/backends/mock_sampler.py | 0 .../sampling/execution/backends/streaming.py | 44 ++ .../execution}/backends/vllm_async.py | 0 .../execution}/concurrency.py | 2 +- .../execution}/runtime.py | 45 +- .../execution}/tinker_handlers.py | 8 +- .../execution}/twinkle_handlers.py | 19 +- .../execution}/weights.py | 9 +- .../server/{sampler => sampling}/router.py | 40 +- .../server/session_resource/__init__.py | 3 +- .../server/session_resource/adapter.py | 86 ---- .../server/session_resource/processor.py | 6 +- .../server/state/backend/redis_backend.py | 10 +- src/twinkle/server/state/future_manager.py | 213 ++++---- src/twinkle/server/state/model_manager.py | 4 +- src/twinkle/server/state/models.py | 38 +- src/twinkle/server/state/sequence.py | 27 - src/twinkle/server/state/server_state.py | 71 ++- src/twinkle/server/state/training_records.py | 26 +- src/twinkle/server/state/training_store.py | 94 +++- src/twinkle/server/task_errors.py | 46 ++ src/twinkle/server/task_queue/__init__.py | 5 +- src/twinkle/server/task_queue/backend_gate.py | 6 +- src/twinkle/server/task_queue/mixin.py | 9 +- src/twinkle/server/task_queue/types.py | 68 +-- src/twinkle/server/task_queue/worker.py | 24 +- src/twinkle/server/telemetry/metrics.py | 54 +- .../server/telemetry/training_metrics.py | 40 -- src/twinkle/server/training/actor.py | 4 +- src/twinkle/server/training/catalog.py | 44 -- src/twinkle/server/training/controller.py | 117 ++--- src/twinkle/server/training/driver.py | 41 +- src/twinkle/server/training/errors.py | 10 - .../{model => training/execution}/__init__.py | 0 .../{model => training/execution}/app.py | 26 +- .../execution}/backends/__init__.py | 0 .../execution}/backends/common.py | 2 +- .../execution}/backends/megatron_model.py | 6 +- .../execution}/backends/mock_model.py | 2 +- .../execution}/backends/transformers_model.py | 6 +- .../execution}/data_plane_inputs.py | 0 .../execution/runtime.py} | 86 ++-- .../execution}/tinker_datum.py | 0 .../training/execution/tinker_handlers.py | 301 +++++++++++ .../execution}/twinkle_handlers.py | 4 +- src/twinkle/server/training/router.py | 14 +- src/twinkle_client/auto/agent/server_tools.py | 3 - .../model/multi_lora_transformers.py | 191 ++----- .../checkpoint/test_checkpoint_manager.py | 4 +- tests/server/config/server_config_4b_e2e.yaml | 5 +- .../config/server_config_4b_e2e_megatron.yaml | 4 +- tests/server/config/test_server_config.py | 2 +- tests/server/config/test_training_config.py | 143 +++--- tests/server/contract/client_api_harness.py | 8 +- .../contract/test_client_api_contract.py | 2 +- tests/server/fixtures/server_config_mock.yaml | 2 + .../fixtures/server_config_mock_redis.yaml | 2 + tests/server/integration/e2e_helpers.py | 31 ++ .../on_demand_training/acceptance_support.py | 2 +- .../on_demand_training/http_probe.py | 36 +- .../on_demand_training/matrix_validation.py | 6 +- .../test_resident_release_e2e.py | 2 +- .../integration/resources/allocation_probe.py | 40 +- .../resources/fragmentation_probe.py | 21 +- .../integration/sampler/sampling_probe.py | 4 +- .../integration/sampler/test_real_pool_e2e.py | 4 +- .../server/integration/test_actor_recovery.py | 4 +- .../integration/test_blocking_boundary.py | 2 +- tests/server/integration/test_dpo_e2e.py | 2 +- .../server/integration/test_full_param_e2e.py | 2 +- .../test_sampler_checkpoints_e2e.py | 8 +- tests/server/integration/test_sft_e2e.py | 2 +- .../test_tinker_rl_correctness_e2e.py | 46 +- tests/server/lifecycle/test_envelope.py | 18 +- .../server/lifecycle/test_run_submit_dedup.py | 32 +- tests/server/lifecycle/test_static_guards.py | 28 +- .../model/test_full_mode_exclusivity.py | 8 +- tests/server/model/test_mock_model.py | 4 +- .../model/test_preflight_loss_inputs.py | 8 +- tests/server/model/test_replica_lifecycle.py | 14 +- .../server/model/test_tinker_compat_output.py | 4 +- tests/server/model/test_tinker_handlers.py | 20 +- .../server/model/test_twinkle_async_inputs.py | 58 +-- tests/server/resources/test_instances.py | 43 +- tests/server/sampler/test_mock_sampler.py | 8 +- tests/server/sampler/test_pool.py | 4 +- tests/server/sampler/test_request_runtime.py | 49 +- .../sampler/test_resolve_sampler_weights.py | 4 +- .../sampler/test_sampling_concurrency.py | 14 +- .../server/sampler/test_stream_guarantees.py | 8 +- tests/server/sampler/test_streaming_bridge.py | 45 ++ tests/server/sampler/test_tinker_handlers.py | 4 +- .../server/sampler/test_twinkle_async_rows.py | 6 +- .../server/session_resource/test_contract.py | 7 +- tests/server/start_e2e_server.py | 3 +- tests/server/state/test_managers.py | 17 + tests/server/state/test_redis_read_retry.py | 49 ++ .../test_redis_recovery.py} | 0 .../server/static/backend_call_exemptions.py | 4 +- .../static/test_adapter_name_mapping.py | 6 +- .../static/test_no_direct_backend_call.py | 2 +- .../static/test_no_twinkle_http_exception.py | 4 +- tests/server/task_queue/helpers.py | 3 +- .../task_queue/test_task_queue_mixin.py | 26 +- .../server/telemetry/test_training_metrics.py | 115 ++++- .../test_app_builders_characterization.py | 4 +- tests/server/test_dependencies.py | 14 +- tests/server/test_service_proxy.py | 2 +- tests/server/training/helpers.py | 86 +++- .../training/test_actor_initialization.py | 2 +- tests/server/training/test_controller.py | 53 +- .../server/training/test_failure_recovery.py | 43 +- .../server/training/test_http_integration.py | 2 + tests/server/training/test_launcher.py | 44 ++ tests/server/training/test_router.py | 10 +- .../training/test_runtime_initialization.py | 16 +- tests/server/training/test_runtime_queue.py | 13 +- tests/server/validation/test_preflight.py | 45 +- tests/server/validation/test_request_wire.py | 14 +- .../transformers/test_spectral_hybrid_lora.py | 6 +- .../test_vllm_sampler_tq_generation.py | 2 +- tests/twinkle_client/test_async_components.py | 28 + tests/twinkle_client/test_auto_agent_tools.py | 7 +- tests/twinkle_client/test_request_builder.py | 26 +- tests/twinkle_client/test_types_contract.py | 8 - 188 files changed, 2644 insertions(+), 2603 deletions(-) delete mode 100644 src/twinkle/server/config/hybrid.py create mode 100644 src/twinkle/server/config/resources.py delete mode 100644 src/twinkle/server/config/training_normalization.py delete mode 100644 src/twinkle/server/env_propagation.py delete mode 100644 src/twinkle/server/gateway/routes.py rename src/twinkle/server/{service_proxy.py => lifecycle/proxy.py} (93%) delete mode 100644 src/twinkle/server/model/tinker_handlers.py delete mode 100644 src/twinkle/server/sampler/actor.py delete mode 100644 src/twinkle/server/sampler/backends/__init__.py delete mode 100644 src/twinkle/server/sampler/backends/streaming.py create mode 100644 src/twinkle/server/sampling/__init__.py create mode 100644 src/twinkle/server/sampling/actor.py rename src/twinkle/server/{sampler => sampling}/controller.py (94%) rename src/twinkle/server/{sampler => sampling}/driver.py (87%) rename src/twinkle/server/{sampler => sampling/execution}/__init__.py (100%) rename src/twinkle/server/{sampler => sampling/execution}/app.py (95%) create mode 100644 src/twinkle/server/sampling/execution/backends/__init__.py rename src/twinkle/server/{sampler => sampling/execution}/backends/mock_sampler.py (100%) create mode 100644 src/twinkle/server/sampling/execution/backends/streaming.py rename src/twinkle/server/{sampler => sampling/execution}/backends/vllm_async.py (100%) rename src/twinkle/server/{sampler => sampling/execution}/concurrency.py (98%) rename src/twinkle/server/{sampler => sampling/execution}/runtime.py (79%) rename src/twinkle/server/{sampler => sampling/execution}/tinker_handlers.py (94%) rename src/twinkle/server/{sampler => sampling/execution}/twinkle_handlers.py (96%) rename src/twinkle/server/{sampler => sampling/execution}/weights.py (84%) rename src/twinkle/server/{sampler => sampling}/router.py (72%) delete mode 100644 src/twinkle/server/session_resource/adapter.py delete mode 100644 src/twinkle/server/state/sequence.py delete mode 100644 src/twinkle/server/telemetry/training_metrics.py delete mode 100644 src/twinkle/server/training/catalog.py delete mode 100644 src/twinkle/server/training/errors.py rename src/twinkle/server/{model => training/execution}/__init__.py (100%) rename src/twinkle/server/{model => training/execution}/app.py (91%) rename src/twinkle/server/{model => training/execution}/backends/__init__.py (100%) rename src/twinkle/server/{model => training/execution}/backends/common.py (99%) rename src/twinkle/server/{model => training/execution}/backends/megatron_model.py (95%) rename src/twinkle/server/{model => training/execution}/backends/mock_model.py (99%) rename src/twinkle/server/{model => training/execution}/backends/transformers_model.py (95%) rename src/twinkle/server/{model => training/execution}/data_plane_inputs.py (100%) rename src/twinkle/server/{model/managed_runtime.py => training/execution/runtime.py} (75%) rename src/twinkle/server/{model => training/execution}/tinker_datum.py (100%) create mode 100644 src/twinkle/server/training/execution/tinker_handlers.py rename src/twinkle/server/{model => training/execution}/twinkle_handlers.py (99%) create mode 100644 tests/server/sampler/test_streaming_bridge.py create mode 100644 tests/server/state/test_redis_read_retry.py rename tests/server/{training/test_redis_restart_reads.py => state/test_redis_recovery.py} (100%) diff --git a/README.md b/README.md index 6f4b05a03..00bf1f9f2 100644 --- a/README.md +++ b/README.md @@ -256,7 +256,7 @@ from twinkle import init_tinker_client from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.model.tinker_datum import input_feature_to_datum +from twinkle.server.training.execution.tinker_datum import input_feature_to_datum base_model = 'ms://Qwen/Qwen3.8-27B' base_url='your-base-url' diff --git a/README_ZH.md b/README_ZH.md index 3d5f05fac..3317c5921 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -245,7 +245,7 @@ from twinkle import init_tinker_client from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.model.tinker_datum import input_feature_to_datum +from twinkle.server.training.execution.tinker_datum import input_feature_to_datum base_model = 'ms://Qwen/Qwen3.8-27B' base_url='your-base-url' diff --git a/cookbook/client/async_rl/server_config.yaml b/cookbook/client/async_rl/server_config.yaml index 6e9c46441..668277a89 100644 --- a/cookbook/client/async_rl/server_config.yaml +++ b/cookbook/client/async_rl/server_config.yaml @@ -64,6 +64,8 @@ applications: args: backend: transformers model_id: ${oc.env:TWINKLE_LOCAL_MODEL_PATH} + min_replicas: 1 + max_replicas: 1 max_loras: 8 max_length: 10240 data_plane_url: http://127.0.0.1:8000/api/v1/data-plane @@ -78,8 +80,6 @@ applications: queue_config: rps_limit: 100 tps_limit: 100000 - adapter_config: - adapter_timeout: 30 deployments: - name: ModelManagement autoscaling_config: diff --git a/cookbook/client/server/megatron/server_config.yaml b/cookbook/client/server/megatron/server_config.yaml index 7e509bcb0..93e1be9fc 100644 --- a/cookbook/client/server/megatron/server_config.yaml +++ b/cookbook/client/server/megatron/server_config.yaml @@ -31,7 +31,7 @@ persistence: redis_url: redis://localhost:6380 # Independent allocation, startup and confirmed-shutdown budgets. -training: +resources: allocation_timeout_s: 3 startup_timeout_s: 1800 shutdown_timeout_s: 180 @@ -140,9 +140,6 @@ applications: rps_limit: 20 # Max requests per second tps_limit: 131072 # Max tokens per second max_input_tokens: 32768 - adapter_config: - adapter_timeout: 120 # Session heartbeat timeout, not business idle - adapter_max_lifetime: 36000 # Maximum lifetime of an adapter in seconds (e.g., 10 hours) deployments: - name: ModelManagement autoscaling_config: diff --git a/cookbook/client/server/megatron/server_config_4b.yaml b/cookbook/client/server/megatron/server_config_4b.yaml index 5b6cd9feb..36f20b1df 100644 --- a/cookbook/client/server/megatron/server_config_4b.yaml +++ b/cookbook/client/server/megatron/server_config_4b.yaml @@ -39,6 +39,8 @@ applications: args: backend: megatron model_id: "ms://Qwen/Qwen3.5-4B" # ModelScope model identifier + min_replicas: 1 + max_replicas: 1 max_length: 10240 nproc_per_node: 2 # Number of GPU processes per node device_group: @@ -52,9 +54,6 @@ applications: rps_limit: 100 # Max requests per second tps_limit: 100000 # Max tokens per second for a single user max_input_tokens: 60000 # Maximum input tokens per request - adapter_config: - adapter_timeout: 30 # Seconds before idle adapter unload - adapter_max_lifetime: 36000 # Maximum lifetime of an adapter in seconds (e.g., 10 hours) max_loras: 5 # Maximum number of LoRA adapters per model deployments: - name: ModelManagement diff --git a/cookbook/client/server/transformer/server_config.yaml b/cookbook/client/server/transformer/server_config.yaml index dd6cec311..bf7ff1558 100644 --- a/cookbook/client/server/transformer/server_config.yaml +++ b/cookbook/client/server/transformer/server_config.yaml @@ -55,6 +55,8 @@ applications: args: backend: transformers # Model backend: transformers | megatron model_id: "ms://Qwen/Qwen3.5-4B" # ModelScope model identifier + min_replicas: 1 + max_replicas: 1 max_length: 10240 nproc_per_node: 1 # Number of GPU processes per node device_group: @@ -67,8 +69,6 @@ applications: queue_config: rps_limit: 100 # Max requests per second tps_limit: 100000 # Max tokens per second - adapter_config: - adapter_timeout: 30 # Seconds before idle adapter unload deployments: - name: ModelManagement autoscaling_config: diff --git a/cookbook/client/server/transformer/server_config_hybrid_lora.yaml b/cookbook/client/server/transformer/server_config_hybrid_lora.yaml index 013024cb8..a8b53cee7 100644 --- a/cookbook/client/server/transformer/server_config_hybrid_lora.yaml +++ b/cookbook/client/server/transformer/server_config_hybrid_lora.yaml @@ -55,8 +55,6 @@ applications: queue_config: rps_limit: 100 tps_limit: 100000 - adapter_config: - adapter_timeout: 1800 deployments: - name: ModelManagement autoscaling_config: diff --git a/cookbook/client/tinker/dpo.py b/cookbook/client/tinker/dpo.py index 3f6e25d05..5f57be719 100644 --- a/cookbook/client/tinker/dpo.py +++ b/cookbook/client/tinker/dpo.py @@ -28,7 +28,7 @@ from twinkle.dataset import Dataset, DatasetMeta, LazyDataset from twinkle.dataloader import DataLoader from twinkle.preprocessor import EmojiDPOProcessor -from twinkle.server.model.tinker_datum import input_feature_to_datum +from twinkle.server.training.execution.tinker_datum import input_feature_to_datum logger = get_logger() diff --git a/cookbook/client/tinker/multi_modal.py b/cookbook/client/tinker/multi_modal.py index 594c4bb8c..4541669cd 100644 --- a/cookbook/client/tinker/multi_modal.py +++ b/cookbook/client/tinker/multi_modal.py @@ -30,7 +30,7 @@ from twinkle.preprocessor import Preprocessor from twinkle.dataset import DatasetMeta, LazyDataset from twinkle.dataloader import DataLoader -from twinkle.server.model.tinker_datum import input_feature_to_datum # Key: converts InputFeature -> Datum +from twinkle.server.training.execution.tinker_datum import input_feature_to_datum # Key: converts InputFeature -> Datum from twinkle import get_logger logger = get_logger() diff --git a/cookbook/client/tinker/self_cognition.py b/cookbook/client/tinker/self_cognition.py index 3a1cd75be..6b43c12f4 100644 --- a/cookbook/client/tinker/self_cognition.py +++ b/cookbook/client/tinker/self_cognition.py @@ -16,7 +16,7 @@ from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.model.tinker_datum import input_feature_to_datum +from twinkle.server.training.execution.tinker_datum import input_feature_to_datum # Initialize the Tinker client before importing ServiceClient init_tinker_client() diff --git a/cookbook/client/tinker/upload_to_hub.py b/cookbook/client/tinker/upload_to_hub.py index d35822144..96491aff7 100644 --- a/cookbook/client/tinker/upload_to_hub.py +++ b/cookbook/client/tinker/upload_to_hub.py @@ -9,8 +9,6 @@ # Task_Envelope with a request_id immediately, so the HTTP call never times out. # 2. The client's future layer long-polls /twinkle/retrieve_future and blocks # until the upload reaches a terminal state, raising on failure. -# (`upload_to_hub` keeps its `poll_interval` / `async_upload` arguments for -# signature compatibility; both are deprecated and have no effect.) # # Prerequisites: # - Server must be running (see server.py / server_config.yaml) diff --git a/cookbook/client/twinkle/upload_to_hub.py b/cookbook/client/twinkle/upload_to_hub.py index 7a2bc2fd2..e8e6921cb 100644 --- a/cookbook/client/twinkle/upload_to_hub.py +++ b/cookbook/client/twinkle/upload_to_hub.py @@ -9,8 +9,6 @@ # Task_Envelope with a request_id immediately, so the HTTP call never times out. # 2. The client's future layer long-polls /twinkle/retrieve_future and blocks # until the upload reaches a terminal state, raising on failure. -# (`upload_to_hub` keeps its `poll_interval` / `async_upload` arguments for -# signature compatibility; both are deprecated and have no effect.) # # Prerequisites: # - Server must be running (see server.py / server_config.yaml) diff --git a/docs/source_en/Usage Guide/Introduction-with-Qwen3.5.md b/docs/source_en/Usage Guide/Introduction-with-Qwen3.5.md index e61ed9197..731f6d7be 100644 --- a/docs/source_en/Usage Guide/Introduction-with-Qwen3.5.md +++ b/docs/source_en/Usage Guide/Introduction-with-Qwen3.5.md @@ -458,7 +458,7 @@ from twinkle import init_tinker_client from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.model.tinker_datum import input_feature_to_datum +from twinkle.server.training.execution.tinker_datum import input_feature_to_datum # Initialize Tinker client (must be called before importing ServiceClient) init_tinker_client() diff --git a/docs/source_en/Usage Guide/Quick-Start.md b/docs/source_en/Usage Guide/Quick-Start.md index 2efc140ad..5ea0a8500 100644 --- a/docs/source_en/Usage Guide/Quick-Start.md +++ b/docs/source_en/Usage Guide/Quick-Start.md @@ -761,7 +761,7 @@ from tinker import ServiceClient from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.model.tinker_datum import input_feature_to_datum +from twinkle.server.training.execution.tinker_datum import input_feature_to_datum # The base model to fine-tune / evaluate base_model = 'ms://Qwen/Qwen3.5-4B' diff --git a/docs/source_en/Usage Guide/Server and Client/Server.md b/docs/source_en/Usage Guide/Server and Client/Server.md index 4f94d8989..ecb5f1730 100644 --- a/docs/source_en/Usage Guide/Server and Client/Server.md +++ b/docs/source_en/Usage Guide/Server and Client/Server.md @@ -233,9 +233,6 @@ applications: rps_limit: 100 # Max requests per second tps_limit: 10000 # Max tokens per second per user max_input_tokens: 10000 # Maximum input tokens per request - adapter_config: - adapter_timeout: 30 # Idle adapter timeout unload time (seconds) - adapter_max_lifetime: 36000 # Maximum adapter lifetime (seconds) max_loras: 1 # Maximum number of LoRA adapters per model deployments: - name: ModelManagement @@ -328,9 +325,6 @@ The difference from the Megatron backend is only in the `backend` parameter of t device_mesh: device_type: cuda dp_size: 2 - adapter_config: - adapter_timeout: 1800 # Idle adapter timeout unload time (seconds) - adapter_max_lifetime: 36000 deployments: - name: ModelManagement autoscaling_config: diff --git a/docs/source_en/Usage Guide/Server and Client/Tinker-Compatible-Client.md b/docs/source_en/Usage Guide/Server and Client/Tinker-Compatible-Client.md index 6e98414b3..9e22f83c2 100644 --- a/docs/source_en/Usage Guide/Server and Client/Tinker-Compatible-Client.md +++ b/docs/source_en/Usage Guide/Server and Client/Tinker-Compatible-Client.md @@ -53,7 +53,7 @@ from twinkle import init_tinker_client from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.model.tinker_datum import input_feature_to_datum +from twinkle.server.training.execution.tinker_datum import input_feature_to_datum # Step 1: Initialize Tinker client before importing ServiceClient init_tinker_client() diff --git a/docs/source_en/Usage Guide/Train-as-a-Service.md b/docs/source_en/Usage Guide/Train-as-a-Service.md index c07e072e1..1d27635f3 100644 --- a/docs/source_en/Usage Guide/Train-as-a-Service.md +++ b/docs/source_en/Usage Guide/Train-as-a-Service.md @@ -28,7 +28,7 @@ from twinkle import init_tinker_client from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.model.tinker_datum import input_feature_to_datum +from twinkle.server.training.execution.tinker_datum import input_feature_to_datum base_model = 'ms://Qwen/Qwen3.8-27B' base_url='https://www.modelscope.cn/twinkle' diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Qwen3.5\346\234\200\344\275\263\345\256\236\350\267\265.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Qwen3.5\346\234\200\344\275\263\345\256\236\350\267\265.md" index d6fc45d2b..88d80c0b6 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Qwen3.5\346\234\200\344\275\263\345\256\236\350\267\265.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Qwen3.5\346\234\200\344\275\263\345\256\236\350\267\265.md" @@ -458,7 +458,7 @@ from twinkle import init_tinker_client from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.model.tinker_datum import input_feature_to_datum +from twinkle.server.training.execution.tinker_datum import input_feature_to_datum # 初始化 Tinker 客户端（必须在导入 ServiceClient 之前） init_tinker_client() diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\345\277\253\351\200\237\345\274\200\345\247\213.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\345\277\253\351\200\237\345\274\200\345\247\213.md" index 7ea2d838c..579a2a31e 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\345\277\253\351\200\237\345\274\200\345\247\213.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\345\277\253\351\200\237\345\274\200\345\247\213.md" @@ -762,7 +762,7 @@ from tinker import ServiceClient from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.model.tinker_datum import input_feature_to_datum +from twinkle.server.training.execution.tinker_datum import input_feature_to_datum # The base model to fine-tune / evaluate base_model = 'Qwen/Qwen3.5-4B' diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/Tinker\345\205\274\345\256\271\345\256\242\346\210\267\347\253\257.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/Tinker\345\205\274\345\256\271\345\256\242\346\210\267\347\253\257.md" index ec207e28a..1c787aee6 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/Tinker\345\205\274\345\256\271\345\256\242\346\210\267\347\253\257.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/Tinker\345\205\274\345\256\271\345\256\242\346\210\267\347\253\257.md" @@ -53,7 +53,7 @@ from twinkle import init_tinker_client from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.model.tinker_datum import input_feature_to_datum +from twinkle.server.training.execution.tinker_datum import input_feature_to_datum # Step 1: 在导入 ServiceClient 之前，先初始化 Tinker 客户端 init_tinker_client() diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" index 4711f11bb..5767454d7 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" @@ -233,9 +233,6 @@ applications: rps_limit: 100 # 最大请求速率（每秒） tps_limit: 10000 # 单用户最大 token 速率（每秒） max_input_tokens: 10000 # 每次请求最大输入 token 数 - adapter_config: - adapter_timeout: 30 # 空闲适配器超时卸载时间（秒） - adapter_max_lifetime: 36000 # 适配器最大生命周期（秒） max_loras: 1 # 每个模型最多加载的 LoRA 数量 deployments: - name: ModelManagement @@ -328,9 +325,6 @@ Transformers 后端与 Megatron 后端的区别仅在 Model 服务的 `backend` device_mesh: device_type: cuda dp_size: 2 - adapter_config: - adapter_timeout: 1800 # 空闲适配器超时卸载时间（秒） - adapter_max_lifetime: 36000 deployments: - name: ModelManagement autoscaling_config: diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\350\256\255\347\273\203\346\234\215\345\212\241.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\350\256\255\347\273\203\346\234\215\345\212\241.md" index fb515944a..6a2031e28 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\350\256\255\347\273\203\346\234\215\345\212\241.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\350\256\255\347\273\203\346\234\215\345\212\241.md" @@ -29,7 +29,7 @@ from twinkle import init_tinker_client from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.model.tinker_datum import input_feature_to_datum +from twinkle.server.training.execution.tinker_datum import input_feature_to_datum base_model = 'ms://Qwen/Qwen3.8-27B' base_url='https://www.modelscope.cn/twinkle' diff --git a/src/twinkle/protocol/types/base.py b/src/twinkle/protocol/types/base.py index d428afc81..a3ef46cd7 100644 --- a/src/twinkle/protocol/types/base.py +++ b/src/twinkle/protocol/types/base.py @@ -22,7 +22,7 @@ ``status_code`` -- an execution-time failure is delivered with HTTP 200, so the value is systematically unequal to the response status code. 4. A closed value set on a wire field is declared as ``Literal`` / enum, never a - bare ``str`` (see ``QueueStateLiteral`` in ``errors.py``). + bare ``str`` (see ``QueueState`` in ``errors.py``). Field roles ----------- diff --git a/src/twinkle/protocol/types/errors.py b/src/twinkle/protocol/types/errors.py index b896e4eed..5d3e5e32b 100644 --- a/src/twinkle/protocol/types/errors.py +++ b/src/twinkle/protocol/types/errors.py @@ -20,10 +20,14 @@ class can be lined up against its tinker counterpart): from .base import ResponseModel -# Closed value set, kept in sync with the server-side ``QueueState`` enum values -# (a consistency test asserts the two sets are equal). Wire fields carrying a queue -# state declare this alias, never a bare ``str`` (naming ruling 4). -QueueStateLiteral = Literal['active', 'paused_rate_limit', 'paused_capacity', 'unknown'] + +class QueueState(StrEnum): + """Scheduling hints; these are independent of a task's lifecycle status.""" + + ACTIVE = 'active' + PAUSED_RATE_LIMIT = 'paused_rate_limit' + PAUSED_CAPACITY = 'paused_capacity' + UNKNOWN = 'unknown' class ErrorCategory(StrEnum): diff --git a/src/twinkle/protocol/types/lifecycle.py b/src/twinkle/protocol/types/lifecycle.py index 32555e24b..641c47fbb 100644 --- a/src/twinkle/protocol/types/lifecycle.py +++ b/src/twinkle/protocol/types/lifecycle.py @@ -7,18 +7,25 @@ """ from __future__ import annotations -from typing import Any, Literal, Optional +from enum import StrEnum +from typing import Any, Optional from .base import ResponseModel, StrictRequest -from .errors import ErrorPayload, QueueStateLiteral +from .errors import ErrorPayload, QueueState -# The lifecycle status of a queued task. Kept in sync with the server-side -# ``TaskStatus`` enum values (a consistency test asserts the two sets are equal). -TaskStatus = Literal['pending', 'queued', 'running', 'completed', 'failed', 'cancelled'] -# The two states past which a task never changes again. ``frozenset`` so a caller -# cannot mutate the shared set. -TERMINAL_STATUSES: frozenset[str] = frozenset({'completed', 'failed', 'cancelled'}) +class TaskStatus(StrEnum): + """Canonical task states shared by the wire envelope and server execution.""" + + PENDING = 'pending' + QUEUED = 'queued' + RUNNING = 'running' + COMPLETED = 'completed' + FAILED = 'failed' + CANCELLED = 'cancelled' + + +TERMINAL_STATUSES = frozenset({TaskStatus.COMPLETED.value, TaskStatus.FAILED.value, TaskStatus.CANCELLED.value}) class RetrieveFutureRequest(StrictRequest): @@ -79,5 +86,5 @@ class TaskEnvelope(ResponseModel): status: TaskStatus result: Any | None = None # set iff status == 'completed' error: ErrorPayload | None = None # set iff status == 'failed' - queue_state: QueueStateLiteral | None = None + queue_state: QueueState | None = None queue_state_reason: str | None = None diff --git a/src/twinkle/protocol/types/model.py b/src/twinkle/protocol/types/model.py index 1e267d066..51e5c9ce9 100644 --- a/src/twinkle/protocol/types/model.py +++ b/src/twinkle/protocol/types/model.py @@ -41,7 +41,8 @@ class CreateRequest(StrictRequest): class TrainingRequest(StrictRequest): """A training handle is control metadata, never a backend kwarg.""" - training_id: str | None = None + training_id: str = Field(min_length=1) + adapter_name: str class AdapterRequest(TrainingRequest): @@ -54,7 +55,6 @@ class AdapterRequest(TrainingRequest): which would silently disable that dedup. """ - adapter_name: str seq_id: int | None = None gradient_accumulation_steps: int | None = backend_kwarg(default=None, ge=1) @@ -94,7 +94,6 @@ class ClipGradAndStepRequest(ClipGradNormRequest): class CalculateMetricRequest(TrainingRequest): """Body of ``POST /twinkle/calculate_metric``.""" - adapter_name: str is_training: bool = True @@ -135,7 +134,6 @@ class _InlineForwardBase(TrainingRequest): class ForwardRequest(_InlineForwardBase): """Body of ``POST /twinkle/forward``: keeps the graph, mutates no gradients.""" - adapter_name: str disable_lora: bool | None = backend_kwarg(default=None) @@ -147,14 +145,12 @@ class ForwardOnlyRequest(_InlineForwardBase): nothing is mutated to be idempotent about. """ - adapter_name: str disable_lora: bool | None = backend_kwarg(default=None) class ForwardBackwardTaskRequest(_InlineForwardBase): """Body of ``POST /twinkle/forward_backward``: accumulates gradients.""" - adapter_name: str seq_id: int | None = None sync_gradients: bool | None = backend_kwarg(default=None) loss_scale: float | None = backend_kwarg(default=None) @@ -177,7 +173,6 @@ class DataPlaneForwardRequest(TrainingRequest): # Values are *field paths*, not parameter values, so this is not a passthrough # region: nothing in it is forwarded verbatim. kwarg_fields: dict[str, str] = Field(default_factory=dict) - adapter_name: str seq_id: int | None = None task: str | None = backend_kwarg(default=None) temperature: float | None = backend_kwarg(default=None, gt=0) @@ -199,7 +194,6 @@ class DataPlaneForwardOnlyRequest(TrainingRequest): input_refs: list[DataRef] = Field(min_length=1) input_field: str | None = None kwarg_fields: dict[str, str] = Field(default_factory=dict) - adapter_name: str task: str | None = backend_kwarg(default=None) temperature: float | None = backend_kwarg(default=None, gt=0) return_logits: bool | None = backend_kwarg(default=None) @@ -232,19 +226,16 @@ def validate_output(self) -> DataPlaneForwardOnlyRequest: class SetLossRequest(TrainingRequest): loss_cls: str - adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() class SetOptimizerRequest(TrainingRequest): optimizer_cls: str - adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() class SetLrSchedulerRequest(TrainingRequest): scheduler_cls: str - adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() @@ -258,26 +249,22 @@ class SetTemplateRequest(TrainingRequest): """ template_cls: str - adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() class SetProcessorRequest(TrainingRequest): processor_cls: str - adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() class AddMetricRequest(TrainingRequest): metric_cls: str - adapter_name: str is_training: bool | None = None init_kwargs: dict[str, JsonValue] = passthrough() class ApplyPatchRequest(TrainingRequest): patch_cls: str - adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() @@ -287,7 +274,6 @@ class ApplyPatchRequest(TrainingRequest): class SaveRequest(TrainingRequest): - adapter_name: str name: str | None = None save_optimizer: bool = False is_sampler: bool = False # If True, delete existing sampler weights before saving @@ -296,7 +282,6 @@ class SaveRequest(TrainingRequest): class LoadRequest(TrainingRequest): - adapter_name: str name: str load_optimizer: bool = False no_load_optim: bool | None = backend_only('megatron', default=None) diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index 683f99067..73aebfeee 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -459,7 +459,7 @@ def sample_stream( def sample_stream_to_queue(self, queue, inputs, sampling_params=None, adapter_name='', adapter_path=None): """Push streaming deltas to a cross-process Ray queue.""" - from twinkle.server.sampler.backends import stream_to_queue + from twinkle.server.sampling.execution.backends.streaming import stream_to_queue stream_to_queue(self, queue, inputs, sampling_params, adapter_name, adapter_path) @remote_function(dispatch='all', collect='first') diff --git a/src/twinkle/server/README.md b/src/twinkle/server/README.md index db5c358e4..da288c240 100644 --- a/src/twinkle/server/README.md +++ b/src/twinkle/server/README.md @@ -4,6 +4,10 @@ The training and sampling pools share resource mechanisms, but own separate sche ```mermaid flowchart TD + Client --> TrainingRouter + Client --> SamplerRouter + TrainingRouter --> TrainingActor + SamplerRouter --> SamplingActor Launcher --> TrainingActor Launcher --> SamplingActor TrainingActor --> TrainingController @@ -15,23 +19,27 @@ flowchart TD TrainingDriver --> RayRuntime SamplerDriver --> RayRuntime RayRuntime --> DevicePool + TrainingDriver --> TrainingExecutionApp + SamplerDriver --> SamplingExecutionApp + TrainingExecutionApp --> TrainingRuntime + SamplingExecutionApp --> SamplerRuntime ``` -- `training/`: training catalog, handle ownership, placement policy, router and Ray actor composition. -- `model/`: training execution and backend/protocol adapters. `managed_runtime.py` implements worker-side admission and completion, not pool scheduling. -- `sampler/`: sampling policy, runtime execution and backend/protocol adapters. `actor.py` owns Ray composition; `controller.py` accepts an injected driver. +- `training/`: validated model lookup, handle ownership, placement policy, router and Ray actor composition. Its `execution/` subpackage implements backend/protocol adapters and worker-side admission/completion. +- `sampling/`: sampling placement, scaling, control broadcasts, router and Ray actor composition. Its `execution/` subpackage implements backend/protocol adapters and request completion tracking. +- In both domains, the controller makes business decisions, the driver performs domain-specific setup, the actor composes the controller with its driver, and the router handles public requests. Controllers accept injected drivers; importing a controller does not compose Ray actors. - `resources/`: shared device allocations, Ray/Serve deployment and confirmed shutdown, and background instance supervision. No training or sampling policy imports. - `state/`: serializable records and persistence. Training records stay here so model execution need not import the training controller. - `lifecycle/`: request submission, retrieval and Future/Envelope handling. This differs from the instance lifetime supervised in `resources/`. -- `config/`: configuration declarations and normalization. `telemetry/` consumes plain snapshots without importing controllers or their catalogs. +- `config/`: configuration declarations and normalization. `telemetry/` consumes plain snapshots without importing controllers or domain configuration. -The three runtime modules have different boundaries: `server/runtime.py` initializes Twinkle and DeviceMesh; `resources/runtime.py` owns Ray/Serve effects; `sampler/runtime.py` tracks request execution inside a sampling instance. +The runtime modules have different boundaries: `server/runtime.py` initializes Twinkle and DeviceMesh; `resources/runtime.py` owns Ray/Serve effects; `sampling/execution/runtime.py` tracks sampling request execution, and `training/execution/runtime.py` executes controller-owned training handles. Multi-LoRA execution registers every preallocated slot with the distributed reducer. DDP requires `find_unused_parameters=True` and `static_graph=False`; FSDP requires `reshard_after_forward=True` so adapter switches act on sharded parameters instead of retained views. Incompatible overrides fail during model initialization, before weights are loaded. ## Training configuration -The public YAML keeps `applications` with `import_path: model`, `args.train_mode`, `device_group`, and `device_mesh`. Each model application selects one training mode. To expose LoRA and full training for the same weights, use two applications with distinct names and route aliases. `config/training_normalization.py` translates these entries into the internal catalog; internal mode records are not an alternative YAML schema. +The public YAML keeps `applications` with `import_path: model`, `args.train_mode`, `device_group`, and `device_mesh`. Each model application selects one training mode. To expose LoRA and full training for the same weights, use two applications with distinct names and route aliases. `config/training.py` owns the public arguments and resolves each application to one validated instance specification. Training idle expiration uses only `idle_timeout_s`. Optional top-level `resources` contains the shared allocation/startup/shutdown/monitor budgets, independent of training or sampling policy. ## Sampling ownership and state reconstruction @@ -51,3 +59,42 @@ A client object rejects creation of an adapter name it already owns. Release it Unit tests follow the component they exercise. Reusable fake drivers and queues live in helper modules, not in imported `test_*.py` modules. Resource allocation probes live under `tests/server/integration/resources`, sampling probes under `tests/server/integration/sampler`, and training acceptance tools under `tests/server/integration/on_demand_training`. Control fixtures demonstrate scheduling and release mechanisms, not real model execution. Hardware acceptance must separately record the tested commit, runtime import origins, real model operations, numerical checks and confirmed cleanup. The package DAG and lightweight import checks live in `tests/server/test_dependencies.py`. + +## Public training entrypoints + +Clients still create and release training handles. A stable CPU router now owns those public paths: + +| API | Public path | Behavior | +| --- | --- | --- | +| Twinkle | `twinkle/create` | Return service capabilities; no device allocation | +| Twinkle | `twinkle/add_adapter_to_model` | Allocate a handle and start an instance when needed | +| Twinkle | `twinkle/remove_adapter` | Release the handle; reclaim idle excess capacity | +| Tinker | `tinker/create_model` | Allocate a training handle | +| Tinker | `tinker/unload_model` | Release that handle | + +Forward, backward, optimizer, save/load and retrieve paths remain available through the router. The SDK carries the returned `training_id` (Tinker calls it `model_id`); the router resolves its owning instance and epoch on every request. This replaces Serve multiplexing affinity: a handle stays with its owning instance, while clients may reach any CPU router replica. A worker rejects direct create/unload calls because it cannot allocate or retire controller-owned resources. No public lifecycle endpoint was removed. + +## Ray infrastructure boundary + +`infra/_ray/ResourceManager` builds worker groups and rank/device meshes; `RayHelper` exposes that process-local infrastructure. Pool drivers reuse `ResourceManager.from_allocation`, platform device discovery and the shared CPU sizing policy. `resources/DevicePool` owns dynamic placement groups and waits for confirmed removal; `resources/RayRuntime` owns Serve applications and late deployment cleanup. These async ownership/deadline rules do not belong in the process-global helper. Both pools use the same resource implementation. + +State repositories compose the shared domain managers: `ServerState.models`, `.sessions`, `.futures`, and `.sampling_sessions`. `TrainingStore` owns training persistence and its retry ledger, while the record definitions remain independent of I/O. Constructing a store performs no storage operation. + +## Component composition and execution contracts + +There is no separate Catalog object: `TrainingConfig` already owns the validated model specifications and lookup by model/mode. The launcher normalizes YAML once and builds the two pools. `actor.py` is the Ray composition boundary; `controller.py` owns scheduling and mutable pool state; `driver.py` translates a scheduling decision into resource operations and domain-specific worker setup; `router.py` authenticates/routes public requests; `execution/app.py` builds the compute service and backend. Runtime objects execute requests inside that service. Neither router nor app allocates devices independently. + +A common business Controller base would need switches for training handles, LoRA slots, exclusive full training, sampling load, and broadcast replay. Those policies have different invariants. Both controllers instead compose `InstanceLifecycle`; both drivers compose `RayRuntime`, which composes `DevicePool`. This is the reusable layer. Subdirectories keep each domain's scheduling and execution adjacent without mixing their responsibilities. + +## Task queue flow + +1. Parse/validate the request and check capabilities, quotas and owner admission before writing a Future. +2. Persist `pending` with an immutable deadline. Training work enters an adapter queue; sampling can reserve bounded concurrent execution. +3. Serial work records `queued` and wakes the round-robin worker. An atomic start claim changes an eligible Future to `running`; cancellation can win only before this claim. +4. Execution writes `completed` or `failed`. Queue timeout fails unstarted work; execution timeout fails the client Future. Late queued/running writes cannot regress a terminal record. +5. A client timeout, disconnect, or failed Future does not prove backend completion. The owning task/scope retains capacity and weight protection until actual completion; unknown completion fences the instance until confirmed shutdown. +6. Shutdown stops admission, fails queued work and joins owned execution. Unconfirmed worker/placement-group removal keeps the instance quarantined and counted against capacity. + +`protocol/types` owns serializable request/response models and shared status enums. `task_queue/types.py` owns in-process coroutine factories, asyncio tasks and callbacks; these cannot be wire fields. `ExecutionScope` is passed explicitly from the internal request to admitted jobs, and waits for jobs registered by earlier jobs as well. + +Business identity, retry metadata, resource ownership and completion witnesses must use explicit arguments or owned records. Use `ContextVar` only for genuinely ambient context such as logging/tracing where passing through every API adds no domain meaning. Every temporary context assignment must pair `set()` with `reset(token)` in `finally`; spawning tasks copies context and does not extend resource ownership. diff --git a/src/twinkle/server/checkpoint/checkpoint_manager.py b/src/twinkle/server/checkpoint/checkpoint_manager.py index d86fadadd..136768a72 100644 --- a/src/twinkle/server/checkpoint/checkpoint_manager.py +++ b/src/twinkle/server/checkpoint/checkpoint_manager.py @@ -27,7 +27,7 @@ from twinkle import get_logger from twinkle.hub import HubOperation from twinkle.protocol.types import ResolvedLoadPath -from twinkle.server.task_queue.types import UserTaskError +from twinkle.server.exceptions import RequestRejectedError from twinkle.utils.parallel import FileLock, acquire_lock, release_lock from .paths import CHECKPOINT_INFO_FILENAME, validate_user_path from .training_run_manager import BaseFileManager, BaseTrainingRunManager @@ -370,7 +370,7 @@ async def save_sampler(self, model_id: str, name: str | None, """ explicit_name = self.get_ckpt_name(name) if name else None if explicit_name == 'latest': - raise UserTaskError('Checkpoint name "latest" is reserved for unnamed sampler saves.') + raise RequestRejectedError('Checkpoint name "latest" is reserved for unnamed sampler saves.') version_id = uuid.uuid4().hex checkpoint_name = explicit_name or 'latest' uri_name = explicit_name or f'live_{version_id}' diff --git a/src/twinkle/server/config/__init__.py b/src/twinkle/server/config/__init__.py index 9eb6aa09e..f6cb0ca2e 100644 --- a/src/twinkle/server/config/__init__.py +++ b/src/twinkle/server/config/__init__.py @@ -2,11 +2,12 @@ """Server configuration package — aggregate root and per-deployment specs.""" from .application_spec import (ApplicationSpec, DataPlaneArgs, HttpOptions, ModelArgs, ProcessorArgs, SamplerArgs, - ServerArgs, SpectralHybridArgs) + ServerArgs) from .persistence import PersistenceConfig from .server_config import ServerConfig from .task_queue import TaskQueueConfig from .telemetry import TelemetryConfig +from .training import SpectralHybridArgs __all__ = [ 'ApplicationSpec', diff --git a/src/twinkle/server/config/application_spec.py b/src/twinkle/server/config/application_spec.py index ab1bbd845..5ed19bae3 100644 --- a/src/twinkle/server/config/application_spec.py +++ b/src/twinkle/server/config/application_spec.py @@ -15,8 +15,8 @@ from pydantic import BaseModel, ConfigDict, Field, model_validator from typing import Any, Literal -from .hybrid import SpectralHybridArgs from .task_queue import TaskQueueConfig +from .training import ModelArgs # Env var keys the launcher sets from the gateway ``server_config`` so that any # Ray worker (model / sampler / processor), not just the gateway, applies the @@ -55,46 +55,6 @@ class HttpOptions(BaseModel): # ---------- per-deployment args schemas ------------------------------------ # -class ModelArgs(_ArgsBase): - """Args for the ``model`` deployment. - - The ``backend`` field selects the model implementation (transformers / - megatron / mock) — see ``model/app.py`` for the dispatch site that picks - the runtime class based on this value. - """ - - model_id: str - nproc_per_node: int = 1 - device_group: dict[str, Any] - device_mesh: dict[str, Any] - backend: Literal['mock', 'transformers', 'megatron'] - env_vars: dict[str, str] = Field(default_factory=dict) - min_replicas: int | None = Field(default=None, ge=0, strict=True) - max_replicas: int | None = Field(default=None, gt=0, strict=True) - idle_timeout_s: float | None = Field(default=None, gt=0, allow_inf_nan=False) - train_mode: Literal['lora', 'hybrid', 'full'] = 'lora' - strategy: Literal['accelerate', 'native_fsdp'] | None = None - adapter_config: dict[str, Any] | None = None - queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) - max_loras: int = 5 - max_r: int = Field(default=32, gt=0) - max_length: int | None = None - target_modules: str | list[str] = 'all-linear' - hybrid: SpectralHybridArgs | None = None - data_plane_url: str | None = None - - @model_validator(mode='after') - def _validate_spectral_backend(self): - if self.strategy is not None and self.backend != 'transformers': - raise ValueError('strategy is only supported by the transformers backend') - if self.hybrid is not None: - if self.backend != 'transformers': - raise ValueError('hybrid is only supported by the transformers backend') - if self.train_mode == 'full': - raise ValueError('hybrid is not supported with train_mode="full"') - return self - - class SamplerArgs(_ArgsBase): """Args for the ``sampler`` deployment. @@ -253,7 +213,7 @@ class ApplicationSpec(BaseModel): # ``import_path`` and defaults a missing block to ``{}``), so the field is # required here — the validator runs first and fills it. args: ServerArgs | ModelArgs | SamplerArgs | ProcessorArgs | DataPlaneArgs - deployments: list[dict[str, Any]] = Field(default_factory=list) + deployments: list[dict[str, Any]] = Field(default_factory=list, max_length=1) @model_validator(mode='before') @classmethod diff --git a/src/twinkle/server/config/backend_dispatch.py b/src/twinkle/server/config/backend_dispatch.py index b4da33f90..c7cb70ebd 100644 --- a/src/twinkle/server/config/backend_dispatch.py +++ b/src/twinkle/server/config/backend_dispatch.py @@ -5,7 +5,7 @@ type selector (``mock | vllm | torch``) share one validate-then-dispatch shape. ``BackendSelector`` expresses it once, parameterized by the permitted value set and the per-value construction callbacks. The callbacks stay defined in -``model/app.py`` and ``sampler/app.py`` so the lazy backend imports remain local +``training/execution/app.py`` and ``sampling/execution/app.py`` so the lazy backend imports remain local — importing the mock backend on a CPU-only host must not pull in torch/vllm/megatron. """ diff --git a/src/twinkle/server/config/hybrid.py b/src/twinkle/server/config/hybrid.py deleted file mode 100644 index 1456a8913..000000000 --- a/src/twinkle/server/config/hybrid.py +++ /dev/null @@ -1,21 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Shared Spectral Hybrid configuration.""" -import math -from pydantic import BaseModel, ConfigDict, Field, field_validator - - -class SpectralHybridArgs(BaseModel): - """Strict server-owned Spectral Hybrid allocation and optimizer defaults.""" - - model_config = ConfigDict(extra='forbid') - - allocation_path: str = Field(min_length=1) - default_lr_lora: float = Field(default=2.5e-5, gt=0) - default_lr_fft: float = Field(default=1.0e-6, gt=0) - - @field_validator('default_lr_lora', 'default_lr_fft') - @classmethod - def _finite_learning_rate(cls, value: float) -> float: - if not math.isfinite(value): - raise ValueError('learning rate must be finite') - return value diff --git a/src/twinkle/server/config/resources.py b/src/twinkle/server/config/resources.py new file mode 100644 index 000000000..12a46d31d --- /dev/null +++ b/src/twinkle/server/config/resources.py @@ -0,0 +1,55 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shared instance startup, allocation, shutdown and monitoring budgets.""" +import os +from pydantic import BaseModel, ConfigDict, Field + + +class ResourceOptions(BaseModel): + model_config = ConfigDict(extra='forbid', frozen=True) + + allocation_timeout_s: float = Field(default=3, gt=0, allow_inf_nan=False) + startup_timeout_s: float = Field(default=1800, gt=0, allow_inf_nan=False) + shutdown_timeout_s: float = Field(default=60, gt=0, allow_inf_nan=False) + sweep_interval_s: float = Field(default=5, gt=0, allow_inf_nan=False) + + +SHARED_WORKER_ENV_KEYS: tuple[str, ...] = ( + 'TWINKLE_STATE_RAY_NAMESPACE', + 'TWINKLE_TELEMETRY_ENABLED', + 'TWINKLE_TELEMETRY_DEBUG', + 'TWINKLE_TELEMETRY_SERVICE', + 'TWINKLE_TELEMETRY_ENDPOINT', + 'TWINKLE_TELEMETRY_INTERVAL', + 'TWINKLE_MODEL_ID_ALIASES', +) + + +def build_shared_env_vars() -> dict[str, str]: + return {key: os.environ[key] for key in SHARED_WORKER_ENV_KEYS if key in os.environ} + + +def build_persistence_env_vars() -> dict[str, str]: + from twinkle.server.config.persistence import PERSISTENCE_ENV_KEYS + return {key: os.environ[key] for key in PERSISTENCE_ENV_KEYS if key in os.environ} + + +def build_server_state_env_vars() -> dict[str, str]: + from twinkle.server.config.application_spec import SERVER_STATE_ENV_KEYS + return {key: os.environ[key] for key in SERVER_STATE_ENV_KEYS if key in os.environ} + + +def build_propagated_env_vars() -> dict[str, str]: + merged: dict[str, str] = {} + merged.update(build_shared_env_vars()) + merged.update(build_persistence_env_vars()) + merged.update(build_server_state_env_vars()) + return merged + + +def build_worker_env_vars() -> dict[str, str]: + """Policy plus the explicit process settings shared by launcher and resource owners.""" + env = build_propagated_env_vars() + for key in ('PYTHONPATH', 'TWINKLE_DEFAULT_SAVE_DIR', 'TMPDIR', 'OMP_NUM_THREADS'): + if key in os.environ: + env.setdefault(key, os.environ[key]) + return env diff --git a/src/twinkle/server/config/server_config.py b/src/twinkle/server/config/server_config.py index bceec284d..7de25f094 100644 --- a/src/twinkle/server/config/server_config.py +++ b/src/twinkle/server/config/server_config.py @@ -21,10 +21,10 @@ from twinkle.server.exceptions import ConfigParseError from .application_spec import ApplicationSpec, HttpOptions from .persistence import PersistenceConfig +from .resources import ResourceOptions from .task_queue import TaskQueueConfig from .telemetry import TelemetryConfig -from .training import TrainingConfig, TrainingOptions -from .training_normalization import normalize_model_application +from .training import TrainingConfig, TrainingModelConfig class ServerConfig(BaseModel): @@ -38,7 +38,7 @@ class ServerConfig(BaseModel): telemetry: TelemetryConfig = Field(default_factory=TelemetryConfig) persistence: PersistenceConfig = Field(default_factory=PersistenceConfig) task_queue: TaskQueueConfig = Field(default_factory=TaskQueueConfig) - training: TrainingOptions = Field(default_factory=TrainingOptions) + resources: ResourceOptions = Field(default_factory=ResourceOptions) applications: list[ApplicationSpec] = Field(default_factory=list) # ---- loading ---------------------------------------------------------- # @@ -72,6 +72,8 @@ def from_yaml(cls, path: str | Path) -> ServerConfig: @model_validator(mode='after') def _validate_cross_field(self) -> ServerConfig: + if sum(app.import_path == 'server' for app in self.applications) > 1: + raise ValueError('configure one gateway per server') if self.persistence.mode == 'redis' and not self.persistence.redis_url: raise ValueError("persistence.redis_url is required when persistence.mode == 'redis'", ) if self.training_config.models: @@ -89,8 +91,10 @@ def _validate_cross_field(self) -> ServerConfig: def training_config(self) -> TrainingConfig: """Normalize model application args into the internal runtime catalog.""" return TrainingConfig( - **self.training.model_dump(), - models=[normalize_model_application(app) for app in self.applications if app.import_path == 'model']) + **self.resources.model_dump(), + models=[ + TrainingModelConfig.from_application(app) for app in self.applications if app.import_path == 'model' + ]) # ---- round-trip / serialization -------------------------------------- # diff --git a/src/twinkle/server/config/training.py b/src/twinkle/server/config/training.py index 964692f0e..335fb134f 100644 --- a/src/twinkle/server/config/training.py +++ b/src/twinkle/server/config/training.py @@ -1,21 +1,99 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Administrator-owned training catalog. Parsing never allocates devices.""" +"""Training YAML arguments and the controller's validated runtime configuration.""" from __future__ import annotations -from pydantic import BaseModel, ConfigDict, Field, model_validator -from typing import Literal +import math +from pydantic import BaseModel, ConfigDict, Field, SerializeAsAny, ValidationInfo, field_validator, model_validator +from typing import TYPE_CHECKING, Any, Literal -from .hybrid import SpectralHybridArgs +from twinkle.hub.model_alias import extract_model_alias +from twinkle.server.exceptions import RequestRejectedError +from .resources import ResourceOptions from .task_queue import TaskQueueConfig +if TYPE_CHECKING: + from .application_spec import ApplicationSpec -class _TrainingConfig(BaseModel): - model_config = ConfigDict(extra='forbid', frozen=True) + +class SpectralHybridArgs(BaseModel): + """Strict server-owned Spectral Hybrid allocation and optimizer defaults.""" + + model_config = ConfigDict(extra='forbid') + + allocation_path: str = Field(min_length=1) + default_lr_lora: float = Field(default=2.5e-5, gt=0) + default_lr_fft: float = Field(default=1.0e-6, gt=0) + + @field_validator('default_lr_lora', 'default_lr_fft') + @classmethod + def _finite_learning_rate(cls, value: float) -> float: + if not math.isfinite(value): + raise ValueError('learning rate must be finite') + return value + + +class TrainingArguments(BaseModel): + """Training options shared by model YAML and resolved instance configuration.""" + + model_config = ConfigDict(extra='forbid') + + min_replicas: int = Field(default=0, ge=0, strict=True) + max_replicas: int | None = Field(default=None, gt=0, strict=True) + idle_timeout_s: float = Field(default=1800, gt=0, allow_inf_nan=False) + strategy: Literal['accelerate', 'native_fsdp'] | None = None + queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) + max_length: int | None = Field(default=None, gt=0, strict=True) + + @model_validator(mode='after') + def _validate_replicas(self): + if self.max_replicas is not None and self.min_replicas > self.max_replicas: + raise ValueError('min_replicas must not exceed max_replicas') + return self + + +class LoraArguments(BaseModel): + """LoRA parameters shared by YAML and LoRA instance configuration.""" + + max_loras: int = Field(default=5, gt=0, strict=True) + max_r: int = Field(default=32, gt=0, strict=True) + target_modules: str | list[str] = 'all-linear' + hybrid: SpectralHybridArgs | None = None + + +class ModelArgs(TrainingArguments, LoraArguments): + """Args for the ``model`` deployment. + + The ``backend`` field selects the model implementation (transformers / + megatron / mock) — see ``training/execution/app.py`` for the dispatch site that picks + the runtime class based on this value. + """ + + model_id: str + nproc_per_node: int = 1 + device_group: dict[str, Any] + device_mesh: dict[str, Any] + backend: Literal['mock', 'transformers', 'megatron'] + env_vars: dict[str, str] = Field(default_factory=dict) + train_mode: Literal['lora', 'hybrid', 'full'] = 'lora' + data_plane_url: str | None = None + + @model_validator(mode='after') + def _validate_spectral_backend(self): + if self.strategy is not None and self.backend != 'transformers': + raise ValueError('strategy is only supported by the transformers backend') + if self.hybrid is not None: + if self.backend != 'transformers': + raise ValueError('hybrid is only supported by the transformers backend') + if self.train_mode == 'full': + raise ValueError('hybrid is not supported with train_mode="full"') + return self -class TrainingMesh(_TrainingConfig): +class TrainingMesh(BaseModel): """Sizes accepted by DeviceMesh.from_sizes, including overlapping EP axes.""" + model_config = ConfigDict(extra='forbid', frozen=True) + device_type: str | None = None dp_size: int = Field(default=1, gt=0, strict=True) fsdp_size: int | None = Field(default=None, gt=0, strict=True) @@ -34,61 +112,46 @@ def world_size(self) -> int: return DeviceMesh.from_sizes(**self.model_dump(exclude_none=True)).world_size -class TrainingModeConfig(_TrainingConfig): +class TrainingModeConfig(TrainingArguments): + model_config = ConfigDict(extra='forbid', frozen=True) + num_devices: int = Field(gt=0, strict=True) device_mesh: TrainingMesh - min_replicas: int = Field(default=0, ge=0, strict=True) - max_replicas: int | None = Field(default=None, gt=0, strict=True) - idle_timeout_s: float | None = Field(default=1800, gt=0, allow_inf_nan=False) - strategy: Literal['accelerate', 'native_fsdp'] | None = None - queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) - max_length: int | None = Field(default=None, gt=0, strict=True) @model_validator(mode='after') def _validate_resources(self): - if self.max_replicas is not None and self.min_replicas > self.max_replicas: - raise ValueError('min_replicas must not exceed max_replicas') if self.device_mesh.world_size() != self.num_devices: raise ValueError('device_mesh world_size must equal num_devices') return self -class LoraTrainingConfig(TrainingModeConfig): - max_loras: int = Field(default=5, gt=0, strict=True) - max_r: int = Field(default=32, gt=0, strict=True) - target_modules: str | list[str] = 'all-linear' - hybrid: SpectralHybridArgs | None = None +class LoraTrainingConfig(TrainingModeConfig, LoraArguments): + pass -class TrainingModes(_TrainingConfig): - lora: LoraTrainingConfig | None = None - full: TrainingModeConfig | None = None - - @model_validator(mode='after') - def _not_empty(self): - if self.lora is None and self.full is None: - raise ValueError('at least one training mode is required') - return self - - -class LegacyExpiration(_TrainingConfig): - """Internal compatibility policy derived from existing adapter_config fields.""" - session_timeout_s: float | None = Field(default=1800, gt=0, allow_inf_nan=False) - max_lifetime_s: float | None = Field(default=36000, gt=0, allow_inf_nan=False) +class TrainingModelConfig(BaseModel): + """One application exposes one training mode and one instance specification.""" + model_config = ConfigDict(extra='forbid', frozen=True) -class TrainingModelConfig(_TrainingConfig): name: str = Field(min_length=1) model_id: str = Field(min_length=1) backend: Literal['mock', 'transformers', 'megatron'] device_type: Literal['cuda', 'npu', 'mps', 'cpu'] - modes: TrainingModes - legacy_expiration: LegacyExpiration | None = None + mode: Literal['lora', 'full'] + config: SerializeAsAny[TrainingModeConfig] route_prefix: str | None = None data_plane_url: str | None = None env_vars: dict[str, str] = Field(default_factory=dict) deploy_options: dict = Field(default_factory=dict) + @field_validator('config', mode='before') + @classmethod + def _instance_config(cls, value, info: ValidationInfo): + config_type = LoraTrainingConfig if info.data.get('mode') == 'lora' else TrainingModeConfig + data = value.model_dump() if isinstance(value, TrainingModeConfig) else value + return config_type.model_validate(data) + @model_validator(mode='after') def _validate_backend(self): if (self.name.startswith('/') or any(part in ('', '.', '..') for part in self.name.split('/')) @@ -96,29 +159,68 @@ def _validate_backend(self): raise ValueError('name must be a nonempty relative model path without dot segments') if self.device_type == 'cpu' and self.backend != 'mock': raise ValueError('cpu is supported only by the mock training backend') - for mode in (self.modes.lora, self.modes.full): - if mode is None: - continue - if self.backend == 'megatron' and mode.num_devices < 2: - raise ValueError('megatron training instances require at least 2 devices') - if mode.device_mesh.device_type not in (None, self.device_type): - raise ValueError('device_mesh.device_type must match device_type') - if mode.strategy is not None and self.backend != 'transformers': - raise ValueError('strategy requires the transformers backend') - if self.modes.lora and self.modes.lora.hybrid and self.backend != 'transformers': + if self.backend == 'megatron' and self.config.num_devices < 2: + raise ValueError('megatron training instances require at least 2 devices') + if self.config.device_mesh.device_type not in (None, self.device_type): + raise ValueError('device_mesh.device_type must match device_type') + if self.config.strategy is not None and self.backend != 'transformers': + raise ValueError('strategy requires the transformers backend') + if isinstance(self.config, LoraTrainingConfig) and self.config.hybrid and self.backend != 'transformers': raise ValueError('hybrid requires the transformers backend') return self - -class TrainingOptions(_TrainingConfig): - allocation_timeout_s: float = Field(default=3, gt=0, allow_inf_nan=False) - startup_timeout_s: float = Field(default=1800, gt=0, allow_inf_nan=False) - shutdown_timeout_s: float = Field(default=60, gt=0, allow_inf_nan=False) - sweep_interval_s: float = Field(default=5, gt=0, allow_inf_nan=False) - - -class TrainingConfig(TrainingOptions): - """Internal normalized catalog; public model entries stay in applications.""" + @property + def key(self) -> tuple[str, str]: + return self.name, self.mode + + @property + def capacity(self) -> int: + return self.config.max_loras if isinstance(self.config, LoraTrainingConfig) else 1 + + @classmethod + def from_application(cls, app: ApplicationSpec) -> TrainingModelConfig: + args = app.args + if not isinstance(args, ModelArgs): + raise TypeError('Training configuration requires model application arguments') + if len(app.deployments) > 1: + raise ValueError(f'{app.name}: a model application must have one deployment configuration') + deployment = app.deployments[0] if app.deployments else {} + device_type = str( + args.device_group.get('device_type', + args.device_mesh.get('device_type', + 'cpu' if args.backend == 'mock' else 'cuda'))).lower() + mode = 'full' if args.train_mode == 'full' else 'lora' + fields = set(TrainingArguments.model_fields) + if mode == 'lora': + fields.update(LoraArguments.model_fields) + config = {key: getattr(args, key) for key in fields} + config.update( + num_devices=args.device_group.get('ranks', args.nproc_per_node), + device_mesh={ + **args.device_mesh, 'device_type': device_type + }) + env_vars = { + **(deployment.get('ray_actor_options') or {}).get('runtime_env', {}).get('env_vars', {}), + **args.env_vars, + } + return cls( + name=extract_model_alias(app.route_prefix) or app.name, + model_id=args.model_id, + backend=args.backend, + device_type=device_type, + mode=mode, + config=config, + route_prefix=app.route_prefix.rstrip('/') or '/', + data_plane_url=args.data_plane_url, + env_vars=env_vars, + deploy_options={ + key: value + for key, value in deployment.items() if key != 'name' + }) + + +class TrainingConfig(ResourceOptions): + """Resolved catalog; public model entries stay in applications.""" models: tuple[TrainingModelConfig, ...] = () @@ -128,3 +230,12 @@ def _unique_names(self): if len(names) != len(set(names)): raise ValueError('training model names must be unique') return self + + def get(self, name: str, mode: Literal['lora', 'full']) -> TrainingModelConfig: + """Resolve a declared specification without allocating resources.""" + model = next((model for model in self.models if model.name == name), None) + if model is None: + raise RequestRejectedError(f'Unknown training model: {name}') + if model.mode != mode: + raise RequestRejectedError(f'Training model {name} does not provide {mode} mode') + return model diff --git a/src/twinkle/server/config/training_normalization.py b/src/twinkle/server/config/training_normalization.py deleted file mode 100644 index 5a63fc162..000000000 --- a/src/twinkle/server/config/training_normalization.py +++ /dev/null @@ -1,71 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Normalize the model application YAML into the internal training catalog.""" -from twinkle.hub.model_alias import extract_model_alias -from .application_spec import ApplicationSpec, ModelArgs -from .training import LegacyExpiration, TrainingModelConfig, TrainingModes - - -def normalize_model_application(app: ApplicationSpec) -> TrainingModelConfig: - args = app.args - assert isinstance(args, ModelArgs) - if len(app.deployments) > 1: - raise ValueError(f'{app.name}: a model application must have one deployment configuration') - deployment = app.deployments[0] if app.deployments else {} - unsupported = set(args.adapter_config or {}) - {'adapter_timeout', 'adapter_max_lifetime'} - if unsupported: - raise ValueError(f'{app.name}: unsupported adapter_config fields: {sorted(unsupported)}') - scaling = deployment.get('autoscaling_config') or {} - ranks = args.device_group.get('ranks', args.nproc_per_node) - if isinstance(ranks, list): - raise ValueError(f'{app.name}: use a device count instead of fixed ranks for pooled training') - device_type = str( - args.device_group.get('device_type', - args.device_mesh.get('device_type', - 'cpu' if args.backend == 'mock' else 'cuda'))).lower() - mesh = {**args.device_mesh, 'device_type': device_type} - mode = { - 'num_devices': - ranks, - 'device_mesh': - mesh, - 'min_replicas': - args.min_replicas if args.min_replicas is not None else scaling.get('min_replicas', - deployment.get('num_replicas', 0)), - 'max_replicas': - args.max_replicas if args.max_replicas is not None else scaling.get('max_replicas', - deployment.get('num_replicas')), - 'idle_timeout_s': - args.idle_timeout_s, - 'strategy': - args.strategy, - 'queue_config': - args.queue_config, - 'max_length': - args.max_length, - } - mode_name = 'full' if args.train_mode == 'full' else 'lora' - if mode_name == 'lora': - mode.update(max_loras=args.max_loras, max_r=args.max_r, target_modules=args.target_modules, hybrid=args.hybrid) - modes = TrainingModes.model_validate({mode_name: mode}) - legacy_expiration = LegacyExpiration( - session_timeout_s=(args.adapter_config or {}).get('adapter_timeout', 1800), - max_lifetime_s=(args.adapter_config or {}).get('adapter_max_lifetime', 36000)) - env_vars = { - **(deployment.get('ray_actor_options') or {}).get('runtime_env', {}).get('env_vars', {}), - **args.env_vars - } - name = extract_model_alias(app.route_prefix) or app.name - return TrainingModelConfig( - name=name, - model_id=args.model_id, - backend=args.backend, - device_type=device_type, - modes=modes, - legacy_expiration=legacy_expiration, - route_prefix=app.route_prefix.rstrip('/') or '/', - data_plane_url=args.data_plane_url, - env_vars=env_vars, - deploy_options={ - key: value - for key, value in deployment.items() if key != 'name' - }) diff --git a/src/twinkle/server/env_propagation.py b/src/twinkle/server/env_propagation.py deleted file mode 100644 index 9482acc90..000000000 --- a/src/twinkle/server/env_propagation.py +++ /dev/null @@ -1,37 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Collect environment variables propagated to Ray worker processes.""" -from __future__ import annotations - -import os - -TELEMETRY_ENV_KEYS: tuple[str, ...] = ( - 'TWINKLE_STATE_RAY_NAMESPACE', - 'TWINKLE_TELEMETRY_ENABLED', - 'TWINKLE_TELEMETRY_DEBUG', - 'TWINKLE_TELEMETRY_SERVICE', - 'TWINKLE_TELEMETRY_ENDPOINT', - 'TWINKLE_TELEMETRY_INTERVAL', - 'TWINKLE_MODEL_ID_ALIASES', -) - - -def build_telemetry_env_vars() -> dict[str, str]: - return {key: os.environ[key] for key in TELEMETRY_ENV_KEYS if key in os.environ} - - -def build_persistence_env_vars() -> dict[str, str]: - from twinkle.server.config.persistence import PERSISTENCE_ENV_KEYS - return {key: os.environ[key] for key in PERSISTENCE_ENV_KEYS if key in os.environ} - - -def build_server_state_env_vars() -> dict[str, str]: - from twinkle.server.config.application_spec import SERVER_STATE_ENV_KEYS - return {key: os.environ[key] for key in SERVER_STATE_ENV_KEYS if key in os.environ} - - -def build_propagated_env_vars() -> dict[str, str]: - merged: dict[str, str] = {} - merged.update(build_telemetry_env_vars()) - merged.update(build_persistence_env_vars()) - merged.update(build_server_state_env_vars()) - return merged diff --git a/src/twinkle/server/exceptions.py b/src/twinkle/server/exceptions.py index b8985879d..5fc3cc044 100644 --- a/src/twinkle/server/exceptions.py +++ b/src/twinkle/server/exceptions.py @@ -9,6 +9,8 @@ from __future__ import annotations +from typing import Literal + from twinkle.protocol.types.errors import ErrorCategory @@ -176,3 +178,25 @@ def __init__(self, current_holder: str) -> None: super().__init__('This deployment runs in full-parameter (exclusive) mode and is already ' f'held by another training task ({current_holder}). Only one full-parameter ' 'training task is allowed at a time; retry after it is released.') + + +class TrainingExpiredError(ResourceNotFoundError): + """The training handle expired and must be recreated before explicit resume.""" + + +class ProtocolMismatchError(RequestRejectedError): + """The requested client protocol does not match the training handle.""" + + +class ResourceReleaseError(TwinkleServerError, RuntimeError): + """Owned resource removal failed at a known phase; capacity remains reserved.""" + + error_code = 503 + + def __init__(self, message: str, *, phase: Literal['worker_death', 'pg_removal']): + super().__init__(message) + self.phase = phase + + +class BackendBusyError(ResourceExhaustedError, RuntimeError): + """A timed-out backend still owns execution/weight resources on this replica.""" diff --git a/src/twinkle/server/gateway/app.py b/src/twinkle/server/gateway/app.py index d284f83dd..c68ea7271 100644 --- a/src/twinkle/server/gateway/app.py +++ b/src/twinkle/server/gateway/app.py @@ -13,7 +13,7 @@ import twinkle.protocol.types as types from twinkle.server.deployment import LazyCleanupMixin, bind_deployment, build_deployment_app -from twinkle.server.service_proxy import ServiceProxy +from twinkle.server.lifecycle.proxy import ServiceProxy from twinkle.server.state import get_server_state from twinkle.utils.logger import get_logger from .openai_handlers import _register_openai_routes diff --git a/src/twinkle/server/gateway/openai_handlers.py b/src/twinkle/server/gateway/openai_handlers.py index cec235b02..c62bc4c60 100644 --- a/src/twinkle/server/gateway/openai_handlers.py +++ b/src/twinkle/server/gateway/openai_handlers.py @@ -28,12 +28,15 @@ from twinkle.server.lifecycle.envelope import envelope_from_record from twinkle.server.utils import get_template_for_model from twinkle.utils.logger import get_logger -from . import routes from .openai_bridge import make_error, translate_chat_request, translate_response, translate_stream_chunk from .use_cases import poll_future logger = get_logger() +TWINKLE_SAMPLE = 'twinkle/sample' +TWINKLE_SAMPLE_STREAM = 'twinkle/sample_stream' +TWINKLE_SET_TEMPLATE = 'twinkle/set_template' + _OPENAI_SAMPLE_TIMEOUT_ENV = 'TWINKLE_OPENAI_SAMPLE_TIMEOUT' _DEFAULT_OPENAI_SAMPLE_TIMEOUT = 1800.0 @@ -145,7 +148,7 @@ async def chat_completions( # Non-streaming: proxy to /twinkle/sample, translate response response = await self.proxy.proxy_request( request, - endpoint=routes.TWINKLE_SAMPLE, + endpoint=TWINKLE_SAMPLE, base_model=base_model, service_type='sampler', body_override=body_bytes, @@ -214,7 +217,7 @@ async def _sse_generator(): try: async for line in self.proxy.proxy_request_stream( request, - endpoint=routes.TWINKLE_SAMPLE_STREAM, + endpoint=TWINKLE_SAMPLE_STREAM, base_model=base_model, service_type='sampler', body_override=body_bytes, @@ -330,7 +333,7 @@ async def _ensure_template( try: resp = await gateway.proxy.proxy_request( request, - endpoint=routes.TWINKLE_SET_TEMPLATE, + endpoint=TWINKLE_SET_TEMPLATE, base_model=base_model, service_type='sampler', body_override=set_template_body, diff --git a/src/twinkle/server/gateway/routes.py b/src/twinkle/server/gateway/routes.py deleted file mode 100644 index 96a8df495..000000000 --- a/src/twinkle/server/gateway/routes.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Sampler endpoint paths used by the OpenAI bridge.""" - -TWINKLE_SAMPLE = 'twinkle/sample' -TWINKLE_SAMPLE_STREAM = 'twinkle/sample_stream' -TWINKLE_SET_TEMPLATE = 'twinkle/set_template' diff --git a/src/twinkle/server/gateway/tinker_handlers.py b/src/twinkle/server/gateway/tinker_handlers.py index 1a16e45aa..1277ddd35 100644 --- a/src/twinkle/server/gateway/tinker_handlers.py +++ b/src/twinkle/server/gateway/tinker_handlers.py @@ -17,9 +17,8 @@ from twinkle.hub import HubOperation from twinkle.server.checkpoint import create_checkpoint_manager, create_training_run_manager +from twinkle.server.lifecycle.envelope import error_payload_from_failure from twinkle.server.middleware.auth import get_token_from_request -from twinkle.server.state.models import FutureFailureRecord -from twinkle.server.task_errors import trim_traceback from twinkle.utils.logger import get_logger from .use_cases import delete_checkpoint from .use_cases import get_training_run as get_training_run_use_case @@ -27,45 +26,10 @@ logger = get_logger() -# Keys must equal ``state.models.FAILURE_REASON_CODES`` (guarded by test_envelope). -_TINKER_FAILURE_WIRE: dict[str, tuple[int, str]] = { - 'capacity_unavailable': (409, 'user'), - 'startup_failed': (500, 'server'), - 'startup_timeout': (504, 'server'), - 'invalid_request': (400, 'user'), - 'request_rejected': (400, 'user'), - 'resource_not_found': (404, 'user'), - 'full_mode_busy': (409, 'user'), - 'input_tokens_exceeded': (422, 'user'), - 'batch_size_invalid': (422, 'user'), - 'rate_limit_exceeded': (429, 'user'), - 'resource_quota_exceeded': (429, 'user'), - 'cancelled': (499, 'user'), - 'endpoint_unavailable': (501, 'server'), - 'state_contention': (503, 'server'), - 'backend_gate_unavailable': (503, 'server'), - 'orphaned_replica': (503, 'server'), - 'execution_timeout': (504, 'server'), - 'deadline_exceeded': (500, 'server'), - 'internal_error': (500, 'server'), -} - def _tinker_error_from_failure(stored: Any, *, request_id: str) -> dict[str, Any]: - """Map domain failure state to the existing Tinker-compatible error body.""" - failure = FutureFailureRecord.model_validate(stored) - error_code, category = _TINKER_FAILURE_WIRE.get(failure.reason_code, (500, 'server')) - payload: dict[str, Any] = { - 'error': failure.message[:1024], - 'category': category, - 'error_code': error_code, - 'request_id': request_id, - } - if failure.details is not None: - payload['details'] = failure.details - if failure.diagnostic and category == 'server': - payload['traceback'] = trim_traceback(failure.diagnostic) - return payload + """Serialize the shared failure presentation with Tinker's field names.""" + return error_payload_from_failure(stored, request_id=request_id).model_dump(mode='json', exclude_none=True) def _register_gateway_tinker_routes(app: FastAPI, self_fn: Callable[[], GatewayServer]) -> None: diff --git a/src/twinkle/server/launcher/__init__.py b/src/twinkle/server/launcher/__init__.py index 0a955b8bc..83769bd32 100644 --- a/src/twinkle/server/launcher/__init__.py +++ b/src/twinkle/server/launcher/__init__.py @@ -8,7 +8,6 @@ - ``server_launcher`` — the ``ServerLauncher`` class and the ``launch_server`` entry point (the real launcher logic); - ``builder_registry`` — ``import_path`` → deployment-builder resolution; -- ``env_propagation`` — telemetry / persistence env-var collection for workers. This ``__init__`` is a thin aggregator: it only re-exports the public surface so the dotted paths ``twinkle.server.launcher.ServerLauncher`` / ``launch_server`` diff --git a/src/twinkle/server/launcher/builder_registry.py b/src/twinkle/server/launcher/builder_registry.py index 9c81d07e0..2b9395b6c 100644 --- a/src/twinkle/server/launcher/builder_registry.py +++ b/src/twinkle/server/launcher/builder_registry.py @@ -28,9 +28,9 @@ def get_builders() -> dict[str, Callable]: """ from twinkle.server.data_plane import build_data_plane_app from twinkle.server.gateway import build_gateway_app - from twinkle.server.model import build_model_app from twinkle.server.processor import build_processor_app - from twinkle.server.sampler import build_sampler_app + from twinkle.server.sampling.execution import build_sampler_app + from twinkle.server.training.execution import build_model_app return { 'build_gateway_app': build_gateway_app, diff --git a/src/twinkle/server/launcher/server_launcher.py b/src/twinkle/server/launcher/server_launcher.py index 4c6ac6d5c..1268e1adb 100644 --- a/src/twinkle/server/launcher/server_launcher.py +++ b/src/twinkle/server/launcher/server_launcher.py @@ -1,11 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""``ServerLauncher`` and the ``launch_server`` entry point. - -The real launcher module. It lives here (rather than in the package -``__init__.py``) so the package's ``__init__`` stays a thin aggregator; the -public dotted paths ``twinkle.server.launcher.ServerLauncher`` / -``launch_server`` are preserved by the ``__init__`` re-export. -""" +"""Initialize server configuration, deploy applications, and own their shutdown.""" from __future__ import annotations import json @@ -22,320 +16,141 @@ from twinkle.patch.ray_serve import apply_ray_serve_patches, get_runtime_env_for_patches from twinkle.server.config import ServerConfig from twinkle.server.config.application_spec import ApplicationSpec -from twinkle.server.env_propagation import build_propagated_env_vars +from twinkle.server.config.resources import build_propagated_env_vars, build_worker_env_vars from .builder_registry import get_builders, resolve_builder logger = get_logger() class ServerLauncher: - """ - Unified server launcher. - - This class handles Ray/Serve initialization and application deployment. + """Own startup and cleanup; business scheduling stays in the controllers.""" - Attributes: - config: The server configuration dictionary - ray_namespace: The Ray namespace for the cluster - """ - - def __init__( - self, - config: ServerConfig, - ray_namespace: str | None = None, - ): - """ - Initialize the server launcher. - - Args: - config: A validated :class:`ServerConfig` instance. Raw dicts are - rejected — operators must build ``ServerConfig`` via - ``ServerConfig.from_yaml`` or its constructor so cross-field - validation runs before the launcher consumes anything. - ray_namespace: Ray namespace (default: 'twinkle_cluster') - """ + def __init__(self, config: ServerConfig, ray_namespace: str | None = None): if not isinstance(config, ServerConfig): raise TypeError('ServerLauncher requires a typed ServerConfig instance; ' f'got {type(config).__name__}. Build one with ' 'ServerConfig.from_yaml(path) or ServerConfig(...).') - self.config: ServerConfig = config - self.ray_namespace = ray_namespace + self.config = config + self.ray_namespace = ray_namespace or config.ray_namespace or 'twinkle_cluster' self._builders: dict[str, Callable] = {} self._ray_initialized = False self._serve_started = False + self._training_controller = None self._sampler_controllers = [] def _get_builders(self) -> dict[str, Callable]: - """Get (and cache) the builder functions for all app types.""" if not self._builders: self._builders = get_builders() return self._builders def _resolve_builder(self, import_path: str) -> Callable: - """Resolve an import_path to a builder function. - - Raises: - ValueError: If the import_path cannot be resolved. - """ return resolve_builder(import_path, self._get_builders()) + def _configure_environment(self) -> None: + if any(app.import_path in ('model', 'sampler') for app in self.config.applications): + os.environ['TWINKLE_STATE_RAY_NAMESPACE'] = self.ray_namespace + os.environ.update(self.config.persistence.to_env_vars()) + gateways = [app for app in self.config.applications if app.import_path == 'server'] + if gateways: + os.environ.update(gateways[0].args.server_config.to_env_vars()) + aliases = build_model_alias_map(self.config.applications) + if aliases: + os.environ[MODEL_ID_ALIASES_ENV] = json.dumps(aliases, ensure_ascii=False) + else: + os.environ.pop(MODEL_ID_ALIASES_ENV, None) + self._configure_telemetry() + + def _configure_telemetry(self) -> None: + telemetry = self.config.telemetry + os.environ['TWINKLE_TELEMETRY_ENABLED'] = str(int(telemetry.enabled)) + if not telemetry.enabled: + return + from twinkle.server.telemetry import init_telemetry + init_telemetry(telemetry) + os.environ.update({ + 'TWINKLE_TELEMETRY_DEBUG': str(int(telemetry.debug)), + 'TWINKLE_TELEMETRY_SERVICE': telemetry.service_name, + 'TWINKLE_TELEMETRY_ENDPOINT': telemetry.otlp_endpoint, + 'TWINKLE_TELEMETRY_INTERVAL': str(telemetry.export_interval_ms), + }) + def _init_ray(self) -> None: - """Initialize Ray if not already initialized.""" if self._ray_initialized: return - import ray - - namespace = self.ray_namespace or self.config.ray_namespace or 'twinkle_cluster' - if not ray.is_initialized(): - # Use runtime_env to apply patches in worker processes - # This is required because Ray Serve's ProxyActor runs in separate processes runtime_env = get_runtime_env_for_patches() - # Propagate telemetry + persistence env vars to all Ray workers - propagated_env_vars = build_propagated_env_vars() - if propagated_env_vars: - merged_env_vars = dict(runtime_env.get('env_vars') or {}) - merged_env_vars.update(propagated_env_vars) - runtime_env['env_vars'] = merged_env_vars - # Connect to existing cluster if available, otherwise start local instance - ray.init( - address='auto', - namespace=namespace, - runtime_env=runtime_env, - ) - logger.info(f'Ray initialized with namespace={namespace}') - + runtime_env['env_vars'] = {**runtime_env.get('env_vars', {}), **build_propagated_env_vars()} + ray.init(address='auto', namespace=self.ray_namespace, runtime_env=runtime_env) self._ray_initialized = True def _start_serve(self) -> None: - """Start Ray Serve with http_options from config.""" if self._serve_started: return - from ray import serve - - try: - from ray.serve.context import _get_global_client - _get_global_client() - # Serve is running, shut it down before re-starting - serve.shutdown() - except Exception: - # Serve not running — nothing to shut down - pass - - http_options = self.config.http_options.model_dump() - serve_kwargs: dict[str, Any] = {'http_options': http_options} - # ``proxy_location`` controls where the Ray Serve HTTP proxy runs - # (``EveryNode`` / ``HeadOnly`` / ``Disabled``). The example configs - # set this field, so honour it here instead of silently ignoring. + options = {'http_options': self.config.http_options.model_dump()} if self.config.proxy_location: - serve_kwargs['proxy_location'] = self.config.proxy_location - serve.start(**serve_kwargs) - logger.info(f'Ray Serve started with http_options={http_options}, ' - f'proxy_location={self.config.proxy_location!r}') - + options['proxy_location'] = self.config.proxy_location + serve.start(**options) self._serve_started = True - def _deploy_application(self, app_spec: ApplicationSpec) -> None: - """Deploy a single application. - - Args: - app_spec: Validated :class:`ApplicationSpec` from the typed config. - """ + def _deployment_options(self, options: dict, env: dict | None = None) -> dict: + options = deepcopy(options) + options.pop('name', None) + actor = options.setdefault('ray_actor_options', {}) + runtime_env = actor.setdefault('runtime_env', {}) + runtime_env['env_vars'] = {**build_worker_env_vars(), **runtime_env.get('env_vars', {}), **(env or {})} + return options + + def _gateway_arguments(self) -> dict[str, Any]: + args = {'http_options': self.config.http_options.model_dump()} + models = self.config.training_config.models + if models: + args['training_models'] = [model.name for model in models] + args['supported_models'] = list( + dict.fromkeys( + alias for app in self.config.applications + if app.import_path == 'sampler' and (alias := extract_model_alias(app.route_prefix, 'sampler')))) + return args + + def _deploy_sampler(self, app: ApplicationSpec) -> None: + import ray from ray import serve - name = app_spec.name - route_prefix = app_spec.route_prefix - import_path = app_spec.import_path - if import_path == 'model': - # Model applications are specifications. _deploy_training creates - # their CPU routes and the controller owns all device instances. - return - if import_path == 'sampler': - from twinkle.server.sampler.router import deploy_sampler_pool - namespace = self.ray_namespace or self.config.ray_namespace or 'twinkle_cluster' - env = build_propagated_env_vars() - for key in ('PYTHONPATH', 'TWINKLE_DEFAULT_SAVE_DIR', 'TMPDIR', 'OMP_NUM_THREADS'): - if key in os.environ: - env.setdefault(key, os.environ[key]) - controller = deploy_sampler_pool(app_spec, namespace, self.config.training.model_dump(), env) - self._sampler_controllers.append(controller) - return - # Shallow-dump the typed args to a kwargs dict WITHOUT recursing into - # nested models: ``dict(model)`` yields top-level (field, value) pairs - # with nested models left as instances, so ``queue_config`` stays a - # typed ``TaskQueueConfig`` and is not re-serialized to a dict only to be - # revived via ``from_dict`` inside the builder (the prior triple - # validation). ``None`` values are dropped to match ``exclude_none``. - args = {k: v for k, v in dict(app_spec.args).items() if v is not None} - deployments = list(app_spec.deployments or []) - - logger.info(f'Starting {name} at {route_prefix}...') - - builder = self._resolve_builder(import_path) - - deploy_options = {} - if deployments: - if len(deployments) > 1: - logger.warning(f'Application "{name}" has {len(deployments)} deployments configured, ' - f'but only the first deployment will be used.') - deploy_config = deployments[0] - if isinstance(deploy_config, dict): - deploy_options = {k: v for k, v in deploy_config.items() if k != 'name'} - - # Inject telemetry + persistence env vars into the deployment's - # runtime_env so that Ray Serve replicas (worker processes) can - # initialize telemetry and resolve the configured persistence backend - # regardless of deployment startup order. - # User-specified env_vars take precedence over our defaults. - propagated_env_vars = build_propagated_env_vars() - if propagated_env_vars: - ray_actor_options = dict(deploy_options.get('ray_actor_options') or {}) - runtime_env = dict(ray_actor_options.get('runtime_env') or {}) - env_vars = dict(runtime_env.get('env_vars') or {}) - for k, v in propagated_env_vars.items(): - env_vars.setdefault(k, v) - runtime_env['env_vars'] = env_vars - ray_actor_options['runtime_env'] = runtime_env - deploy_options['ray_actor_options'] = ray_actor_options - - # Pass http_options to server apps for internal proxy routing - if import_path == 'server': - args.setdefault('http_options', self.config.http_options.model_dump()) - if self.config.training_config.models: - args['training_models'] = [model.name for model in self.config.training_config.models] - sampler_apps = [app for app in self.config.applications if app.import_path == 'sampler'] - args['supported_models'] = list( - dict.fromkeys( - alias for app in sampler_apps if (alias := extract_model_alias(app.route_prefix, 'sampler')))) + from twinkle.server.sampling.actor import create_sampler_controller + from twinkle.server.sampling.router import build_sampler_router + controller, options, pool_id = create_sampler_controller(app, self.ray_namespace, + self.config.resources.model_dump(), + build_worker_env_vars()) + self._sampler_controllers.append(controller) + ray.get(controller.initialize.remote()) + serve.run( + build_sampler_router(self.ray_namespace, pool_id, options), name=app.name, route_prefix=app.route_prefix) - app = builder(deploy_options=deploy_options, **args) - - serve.run(app, name=name, route_prefix=route_prefix) - logger.info(f'Deployed {name} at {route_prefix}') - - def launch(self) -> None: - """Launch the server with all configured applications. - - Blocks the calling thread to keep the server running. Installs signal - handlers for SIGINT/SIGTERM so that ``serve.shutdown()`` is called on - termination instead of leaving orphaned deployments. - """ - # Apply Ray Serve patches before initializing Ray - apply_ray_serve_patches() - if self.config.training_config.models or any(app.import_path == 'sampler' for app in self.config.applications): - os.environ['TWINKLE_STATE_RAY_NAMESPACE'] = ( - self.ray_namespace or self.config.ray_namespace or 'twinkle_cluster') - - # Initialize telemetry if configured - telemetry = self.config.telemetry - if telemetry.enabled: - from twinkle.server.telemetry import init_telemetry - init_telemetry(telemetry) - # Export config to env vars for Ray worker processes - os.environ['TWINKLE_TELEMETRY_ENABLED'] = '1' - os.environ['TWINKLE_TELEMETRY_DEBUG'] = '1' if telemetry.debug else '0' - os.environ['TWINKLE_TELEMETRY_SERVICE'] = telemetry.service_name - os.environ['TWINKLE_TELEMETRY_ENDPOINT'] = telemetry.otlp_endpoint - os.environ['TWINKLE_TELEMETRY_INTERVAL'] = str(telemetry.export_interval_ms) - - # Export top-level persistence to env vars so any worker - # (not just Gateway) can build the same backend on first call to - # get_server_state(). - persistence = self.config.persistence - for k, v in persistence.to_env_vars().items(): - os.environ[k] = v - logger.info(f'Persistence backend configured: mode={persistence.mode}') - - # Export the gateway ``server`` application's ServerState policy (quota / - # expiry / cleanup / metrics interval) to env vars for the same reason: - # so every worker's first ``get_server_state()`` applies the configured - # values instead of the hardcoded defaults. Without this the model worker - # that enforces ``per_token_model_limit`` runs on the default (30), - # silently ignoring the YAML value. - server_specs = [a for a in self.config.applications if a.import_path == 'server'] - if len(server_specs) > 1: - logger.warning(f'{len(server_specs)} "server" applications declared; using the first ' - 'for ServerState policy env propagation.') - if server_specs: - server_state_env = server_specs[0].args.server_config.to_env_vars() - for k, v in server_state_env.items(): - os.environ[k] = v - if server_state_env: - logger.info(f'ServerState policy exported to worker env: {server_state_env}') - - model_alias_map = build_model_alias_map(self.config.applications) - model_alias_map.update({model.name: model.model_id for model in self.config.training_config.models}) - if model_alias_map: - os.environ[MODEL_ID_ALIASES_ENV] = json.dumps(model_alias_map, ensure_ascii=False) - else: - os.environ.pop(MODEL_ID_ALIASES_ENV, None) - - self._init_ray() - self._start_serve() - - applications = self.config.applications - if not applications and not self.config.training_config.models: - logger.warning('No applications configured') + def _deploy_application(self, app_spec: ApplicationSpec) -> None: + """Model declarations go to the training controller; other apps deploy here.""" + if app_spec.import_path == 'model': + return + if app_spec.import_path == 'sampler': + self._deploy_sampler(app_spec) return - - for app_spec in applications: - if app_spec.import_path != 'model': - self._deploy_application(app_spec) - - if self.config.training_config.models: - self._deploy_training() - - host = self.config.http_options.host - port = self.config.http_options.port - - logger.info('All applications started!') - for app_spec in applications: - logger.info(' Endpoint: http://%s:%s%s', host, port, app_spec.route_prefix) - - # Graceful shutdown via signal handling - shutdown_event = threading.Event() - - def _handle_signal(signum, frame): - sig_name = signal.Signals(signum).name - logger.info(f'Received {sig_name}, shutting down gracefully...') - shutdown_event.set() - - signal.signal(signal.SIGINT, _handle_signal) - signal.signal(signal.SIGTERM, _handle_signal) - - # Block until a termination signal is received - shutdown_event.wait() - - self._shutdown() - - def _shutdown(self) -> None: - """Tear down Serve and flush telemetry on graceful termination. - - Telemetry flush runs AFTER ``serve.shutdown()`` and is wrapped so a - telemetry-shutdown failure cannot mask the user-facing shutdown path. - """ - import ray from ray import serve - for controller in self._sampler_controllers: - ray.get(controller.shutdown.remote()) - ray.kill(controller, no_restart=True) - self._sampler_controllers.clear() - if getattr(self, '_training_controller', None) is not None: - import ray - ray.get(self._training_controller.shutdown.remote()) - ray.kill(self._training_controller, no_restart=True) - try: - serve.shutdown() - logger.info('Ray Serve shut down successfully') - except Exception: - logger.warning('Error during Ray Serve shutdown', exc_info=True) - - # Flush buffered OTLP batches (traces / metrics / logs) on the driver - # after Serve is down. Errors are swallowed by the helper. - from twinkle.server.telemetry import flush_telemetry_safely - flush_telemetry_safely() + args = {key: value for key, value in dict(app_spec.args).items() if value is not None} + if app_spec.import_path == 'server': + defaults = self._gateway_arguments() + for key, value in defaults.items(): + args.setdefault(key, value) + deployment = app_spec.deployments[0] if app_spec.deployments else {} + options = self._deployment_options(deployment) + builder = self._resolve_builder(app_spec.import_path) + serve.run(builder(deploy_options=options, **args), name=app_spec.name, route_prefix=app_spec.route_prefix) + + def _training_router_options(self, model) -> dict: + options = self._deployment_options(model.deploy_options, model.env_vars) + if 'autoscaling_config' not in options: + options.setdefault('num_replicas', 1) + options['ray_actor_options'].setdefault('num_cpus', 0.1) + return options def _deploy_training(self) -> None: import ray @@ -343,40 +158,101 @@ def _deploy_training(self) -> None: from twinkle.server.training.actor import TrainingControllerActor, controller_name from twinkle.server.training.router import build_training_router - - namespace = self.ray_namespace or self.config.ray_namespace or 'twinkle_cluster' - env = build_propagated_env_vars() - for key in ('PYTHONPATH', 'TWINKLE_DEFAULT_SAVE_DIR', 'TMPDIR', 'OMP_NUM_THREADS'): - if key in os.environ: - env.setdefault(key, os.environ[key]) + config = self.config.training_config controller = TrainingControllerActor.options( - name=controller_name(namespace), namespace=namespace, runtime_env={ - 'env_vars': env - }).remote(self.config.training_config.model_dump(), namespace) + name=controller_name(self.ray_namespace), + namespace=self.ray_namespace, + runtime_env={ + 'env_vars': build_worker_env_vars() + }).remote(config.model_dump(), self.ray_namespace) + self._training_controller = controller try: ray.get(controller.initialize.remote()) + for model in config.models: + app = build_training_router(config, model, self.ray_namespace, self.config.http_options.model_dump(), + self._training_router_options(model)) + serve.run( + app, + name=f'training-router-{self.ray_namespace}-{model.name}', + route_prefix=model.route_prefix or f'/api/v1/model/{model.name}') except Exception: + self._run_cleanup([('training controller rollback', self._stop_training_controller)], raise_errors=False) + raise + + @staticmethod + def _stop_controller(controller) -> None: + import ray + ray.get(controller.shutdown.remote()) + ray.kill(controller, no_restart=True) + + def _stop_training_controller(self) -> None: + if self._training_controller is not None: + self._stop_controller(self._training_controller) + self._training_controller = None + + def _stop_sampler_controller(self, controller) -> None: + self._stop_controller(controller) + self._sampler_controllers.remove(controller) + + @staticmethod + def _run_cleanup(actions: list[tuple[str, Callable]], *, raise_errors: bool) -> None: + """Attempt every independent cleanup; rollback preserves the startup error.""" + failures = [] + for name, action in actions: try: - ray.get(controller.shutdown.remote()) - except Exception: - logger.exception('Training controller cleanup failed during initialization rollback') - finally: - ray.kill(controller, no_restart=True) + action() + except Exception as exc: + failures.append((name, exc)) + logger.exception('%s failed', name) + if failures and raise_errors: + names = ', '.join(name for name, _ in failures) + raise RuntimeError(f'Server cleanup is incomplete: {names}') from failures[0][1] + + def _shutdown(self, *, raise_errors: bool = True) -> None: + from functools import partial + from ray import serve + + from twinkle.server.telemetry import flush_telemetry_safely + actions = [('training controller shutdown', self._stop_training_controller)] + actions.extend(('sampler controller shutdown', partial(self._stop_sampler_controller, controller)) + for controller in tuple(self._sampler_controllers)) + if self._serve_started: + actions.append(('Ray Serve shutdown', serve.shutdown)) + actions.append(('telemetry flush', flush_telemetry_safely)) + self._run_cleanup(actions, raise_errors=raise_errors) + self._serve_started = False + + @staticmethod + def _wait_for_shutdown() -> None: + stopped = threading.Event() + + def signal_shutdown(signum, frame): + logger.info('Received %s', signal.Signals(signum).name) + stopped.set() + + signal.signal(signal.SIGINT, signal_shutdown) + signal.signal(signal.SIGTERM, signal_shutdown) + stopped.wait() + + def launch(self) -> None: + if not self.config.applications: + logger.warning('No applications configured') + return + apply_ray_serve_patches() + self._configure_environment() + self._init_ray() + try: + self._start_serve() + for app in self.config.applications: + self._deploy_application(app) + if self.config.training_config.models: + self._deploy_training() + logger.info('All applications started') + self._wait_for_shutdown() + except Exception: + self._shutdown(raise_errors=False) raise - self._training_controller = controller - for model in self.config.training_config.models: - options = deepcopy(model.deploy_options) - if 'autoscaling_config' not in options: - options.setdefault('num_replicas', 1) - actor_options = options.setdefault('ray_actor_options', {}) - actor_options.setdefault('num_cpus', 0.1) - actor_options.setdefault('runtime_env', {})['env_vars'] = {**env, **model.env_vars} - app = build_training_router(self.config.training_config, model, namespace, - self.config.http_options.model_dump(), options) - serve.run( - app, - name=f'training-router-{namespace}-{model.name}', - route_prefix=model.route_prefix or f'/api/v1/model/{model.name}') + self._shutdown() def launch_server( diff --git a/src/twinkle/server/lifecycle/envelope.py b/src/twinkle/server/lifecycle/envelope.py index f089e29ec..dc193aadf 100644 --- a/src/twinkle/server/lifecycle/envelope.py +++ b/src/twinkle/server/lifecycle/envelope.py @@ -13,44 +13,18 @@ from twinkle.protocol.types.errors import ErrorCategory, ErrorPayload from twinkle.protocol.types.lifecycle import TaskEnvelope from twinkle.server.state.models import FutureFailureRecord -from twinkle.server.task_errors import trim_traceback - -# Keys must equal ``state.models.FAILURE_REASON_CODES`` (guarded by test_envelope). -_FAILURE_WIRE: dict[str, tuple[int, ErrorCategory]] = { - 'capacity_unavailable': (409, ErrorCategory.User), - 'startup_failed': (500, ErrorCategory.Server), - 'startup_timeout': (504, ErrorCategory.Server), - 'invalid_request': (400, ErrorCategory.User), - 'request_rejected': (400, ErrorCategory.User), - 'resource_not_found': (404, ErrorCategory.User), - 'full_mode_busy': (409, ErrorCategory.User), - 'input_tokens_exceeded': (422, ErrorCategory.User), - 'batch_size_invalid': (422, ErrorCategory.User), - 'rate_limit_exceeded': (429, ErrorCategory.User), - 'resource_quota_exceeded': (429, ErrorCategory.User), - 'cancelled': (499, ErrorCategory.User), - 'endpoint_unavailable': (501, ErrorCategory.Server), - 'state_contention': (503, ErrorCategory.Server), - 'backend_gate_unavailable': (503, ErrorCategory.Server), - 'orphaned_replica': (503, ErrorCategory.Server), - 'execution_timeout': (504, ErrorCategory.Server), - 'deadline_exceeded': (500, ErrorCategory.Server), - 'internal_error': (500, ErrorCategory.Server), -} +from twinkle.server.task_errors import FailureResponse, trim_traceback def error_payload_from_failure(stored: Any, *, request_id: str) -> ErrorPayload: """Map one protocol-independent failure record to Twinkle's wire model.""" failure = FutureFailureRecord.model_validate(stored) - error_code, category = _FAILURE_WIRE.get( - failure.reason_code, - (500, ErrorCategory.Server), - ) - diagnostic = failure.diagnostic if category is ErrorCategory.Server else None + response = FailureResponse.for_reason(failure.reason_code) + diagnostic = failure.diagnostic if response.category is ErrorCategory.Server else None return ErrorPayload( error=failure.message[:1024], - category=category, - error_code=error_code, + category=response.category, + error_code=response.error_code, request_id=request_id, traceback=trim_traceback(diagnostic) if diagnostic else None, details=failure.details, diff --git a/src/twinkle/server/service_proxy.py b/src/twinkle/server/lifecycle/proxy.py similarity index 93% rename from src/twinkle/server/service_proxy.py rename to src/twinkle/server/lifecycle/proxy.py index 478a7d6cf..cee04c878 100644 --- a/src/twinkle/server/service_proxy.py +++ b/src/twinkle/server/lifecycle/proxy.py @@ -5,7 +5,7 @@ import httpx from fastapi import Request, Response from fastapi.responses import JSONResponse -from typing import Any +from typing import Any, TypedDict from twinkle.protocol.headers import H_MULTIPLEX, H_MULTIPLEX_LEGACY, H_REQUEST_ID, H_REQUEST_ID_LEGACY from twinkle.protocol.types.errors import ErrorCategory, ErrorPayload @@ -16,6 +16,22 @@ logger = get_logger() +class ForwardedRequest(TypedDict): + """HTTP request fields carried intact across an internal Ray RPC.""" + + endpoint: str + body: bytes + method: str + headers: list[tuple[bytes, bytes]] + query_string: bytes + + +class ForwardedReply(TypedDict): + status: int + headers: list[tuple[bytes, bytes]] + body: bytes + + class ServiceProxy: def __init__(self, http_options: dict[str, Any] | None = None, route_prefix: str = '/api/v1'): diff --git a/src/twinkle/server/lifecycle/submit.py b/src/twinkle/server/lifecycle/submit.py index 1c6b78937..5c432281b 100644 --- a/src/twinkle/server/lifecycle/submit.py +++ b/src/twinkle/server/lifecycle/submit.py @@ -16,6 +16,8 @@ from twinkle.protocol.types.base import FieldRole, fields_with_role from twinkle.protocol.types.data import export_batch from twinkle.protocol.types.lifecycle import TaskEnvelope +from twinkle.protocol.types.model import (AdapterRequest, DataPlaneForwardRequest, ForwardBackwardTaskRequest, + TrainingRequest) from twinkle.server.middleware.auth import get_session_id_from_request from twinkle.server.validation import assert_request_supported @@ -85,11 +87,11 @@ def backend_kwargs(body: Any) -> dict[str, Any]: model_cls = type(body) kwargs: dict[str, Any] = {} for name in fields_with_role(model_cls, FieldRole.BackendKwarg): - value = getattr(body, name, None) + value = getattr(body, name) if value is not None: kwargs[name] = value for name in fields_with_role(model_cls, FieldRole.Passthrough): - region = getattr(body, name, None) or {} + region = getattr(body, name) or {} overlap = set(region) & set(kwargs) if overlap: raise ValueError(f'{name} collides with declared backend parameters: {", ".join(sorted(overlap))}') @@ -123,7 +125,7 @@ def resolve_twinkle_adapter_name(request: Request, adapter_name: str | None) -> async def run_submit( self: QueuedDeployment, request: Request, - body: Any, + body: TrainingRequest, *, task_type: str, backend_call: Callable[..., Coroutine], @@ -152,13 +154,10 @@ async def run_submit( signature-patching wrapper once deepened the route graph past CPython's C-stack recursion guard during ``serve.ingress``). - Four queued endpoints deliberately do NOT route through here and call - ``submit_and_peek`` / ``submit_background_and_peek`` directly: model - ``add_adapter_to_model`` (creates the adapter, so the resource assertion cannot - apply and it owns the train_mode/full-mode checks), model ``upload_to_hub`` (pure - I/O, background task), and sampler ``sample`` / ``sample_to_data_plane`` (no adapter - semantics). Zero-write and admission still hold for them because both live inside - ``schedule_task`` -> ``_perform_preflight_checks``, not in this shell. + Upload uses the background I/O path; sampling uses concurrent admission. + Their scheduling entrypoints apply the same preflight and Future ownership rules. + Training creation/release belongs to the public Router and Controller, rather + than this compute-operation shell. ``backend_call(self, body, adapter_name, token)`` runs the endpoint-specific call and returns the JSON-safe task result. ``metrics(self, body)`` supplies scheduling @@ -168,7 +167,7 @@ async def run_submit( not implemented by every backend. """ token = await self._on_request_start(request) - adapter_name = getattr(body, 'training_id', None) or resolve_twinkle_adapter_name(request, body.adapter_name) + adapter_name = body.training_id # ---- Preflight: decidable from the body plus this deployment's backend, so it # runs before any state write and before the enqueue. ---- @@ -192,39 +191,40 @@ async def _task(): # forward_backward would be dropped as a duplicate AND handed the first adapter's # loss. That is a silent wrong-result bug, i.e. the exact failure this dedup # exists to prevent, one level up. ---- - request_id = f'req_{uuid.uuid4().hex}' - seq_id = getattr(body, 'seq_id', None) - dedup_key = None + sequenced = isinstance(body, (AdapterRequest, ForwardBackwardTaskRequest, DataPlaneForwardRequest)) + return await submit_training_task( + self, + _task, + training_id=body.training_id, + stream=f'native::{get_session_id_from_request(request)}::{task_type}', + seq_id=body.seq_id if sequenced else None, + peek=True, + token=token, + task_type=task_type, + **schedule_kwargs) + + +async def submit_training_task(self, + factory, + *, + training_id: str, + stream: str, + seq_id: int | None, + peek: bool = False, + **kwargs): + """Submit with explicit request identity; retain claims once any future was written.""" + request_id = kwargs.pop('request_id', None) or f'req_{uuid.uuid4().hex}' if seq_id is not None: - session_id = get_session_id_from_request(request) or request.state.request_id - dedup_key = f'seq::{session_id}::{adapter_name or "-"}::{task_type}::{seq_id}' - training_id = getattr(body, 'training_id', None) - if training_id is not None: - stream = f'native::{session_id}::{task_type}' - retention = self.sequence_retention_s - prior_request_id = await self.state.claim_training_seq(training_id, stream, seq_id, request_id, retention) - else: - # Legacy claims must cover queue wait as well as backend execution. - ttl = int(self.task_queue_config.queue_timeout + self.task_queue_config.effective_execution_timeout) + 60 - prior_request_id = await self.state.claim_seq(dedup_key, request_id, ttl) - if prior_request_id is not None: - return await self._peek_terminal(prior_request_id, fallback_status='pending') - - # ---- Decision_Boundary: preflight (in schedule_task) then peek ---- + prior = await self.state.claim_training_seq(training_id, stream, seq_id, request_id, self.sequence_retention_s) + if prior is not None: + if peek: + return await self._peek_terminal(prior, fallback_status='pending') + return {'request_id': prior, 'model_id': training_id} + submit = self.submit_and_peek if peek else self.schedule_task try: - return await self.submit_and_peek( - _task, model_id=adapter_name, token=token, task_type=task_type, request_id=request_id, **schedule_kwargs) + return await submit(factory, model_id=training_id, request_id=request_id, **kwargs) except Exception: - # Release the seq claim only when the task never made it onto the queue -- - # decided by whether a future record exists, NOT by the exception type. A - # preflight rejection raises before any record is written, so releasing lets - # a retry re-enqueue. But a failure *after* enqueue (e.g. a transient state - # error inside the peek) leaves a live task that will still run; releasing - # there would let a retry enqueue a duplicate -> the exact double-apply this - # dedup prevents. When unsure (record exists), keep the claim. - if dedup_key is not None and await self.state.get_future(request_id) is None: - if training_id is not None: - await self.state.release_training_seq(training_id, stream, seq_id, request_id) - else: - await self.state.release_seq(dedup_key) + # A failed response after enqueue must not allow the same update to run twice. + if seq_id is not None and await self.state.get_future(request_id) is None: + await self.state.release_training_seq(training_id, stream, seq_id, request_id) raise diff --git a/src/twinkle/server/model/tinker_handlers.py b/src/twinkle/server/model/tinker_handlers.py deleted file mode 100644 index cf32eecda..000000000 --- a/src/twinkle/server/model/tinker_handlers.py +++ /dev/null @@ -1,294 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Tinker-compatible routes for the Model deployment. - -Registered by ``_register_model_tinker_routes(app, self_fn)`` -- module-level route -registration closing over ``self_fn`` via ``Depends``, not a mixin: there is no -inheritance relationship with the deployment class. All endpoints are prefixed -/tinker/... and use schedule_task() returning UntypedAPIFuture. ``self_fn`` is injected -via FastAPI Depends to obtain the ModelManagement instance at request time. -""" -from __future__ import annotations - -import traceback -from collections.abc import Callable -from fastapi import Depends, FastAPI, Request -from functools import partial -from tinker import types -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from .app import ModelManagement - -from twinkle.server.checkpoint import create_checkpoint_manager, create_training_run_manager -from twinkle.server.exceptions import RequestRejectedError -from twinkle.server.validation.tinker_loss import validate_tinker_loss -from twinkle.utils.logger import get_logger - -logger = get_logger() - - -def _register_model_tinker_routes(app: FastAPI, self_fn: Callable[[], ModelManagement]) -> None: - """Register all /tinker/* routes on the given FastAPI app. - - self_fn is a zero-argument callable that returns the current ModelManagement - replica instance. It is wired in via Depends so it is resolved lazily at request time. - """ - - @app.post('/tinker/create_model') - async def create_model( - request: Request, - body: types.CreateModelRequest, - self: ModelManagement = Depends(self_fn), - ) -> types.UntypedAPIFuture: - raise RequestRejectedError('Training lifecycle operations must use the catalog endpoint') - - @app.post('/tinker/get_info') - async def get_info(request: Request, body: types.GetInfoRequest, - self: ModelManagement = Depends(self_fn)) -> types.GetInfoResponse: - token = await self._on_request_start(request) - training_run_manager = create_training_run_manager(token, client_type='tinker') - metadata = training_run_manager.get(str(body.model_id)) - model_name = metadata.base_model if metadata else self.base_model - lora_rank = None - is_lora = False - if metadata and hasattr(metadata, 'lora_rank') and metadata.lora_rank: - lora_rank = metadata.lora_rank - is_lora = metadata.is_lora - return types.GetInfoResponse( - model_data=types.ModelData(model_name=model_name), - model_id=body.model_id, - is_lora=is_lora, - lora_rank=lora_rank, - model_name=model_name, - ) - - @app.post('/tinker/unload_model') - async def unload_model( - request: Request, - body: types.UnloadModelRequest, - self: ModelManagement = Depends(self_fn), - ) -> types.UntypedAPIFuture: - raise RequestRejectedError('Training lifecycle operations must use the catalog endpoint') - - @app.post('/tinker/forward') - async def forward(request: Request, body: types.ForwardRequest, - self: ModelManagement = Depends(self_fn)) -> types.UntypedAPIFuture: - token = await self._on_request_start(request) - - async def _do_forward(): - try: - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - self.assert_resource_exists(adapter_name) - model_adapter = self.resolve_model_adapter_name(adapter_name) - datum_list = body.forward_input.data - loss_fn_config = body.forward_input.loss_fn_config or {} - output, loss = await self.call_backend( - self.model.tinker_forward_only, inputs=datum_list, adapter_name=model_adapter, **loss_fn_config) - return types.ForwardBackwardOutput( - loss_fn_output_type='CrossEntropyLossReturn', - loss_fn_outputs=output, - metrics={'loss:mean': loss}, - ) - except Exception: - logger.error(traceback.format_exc()) - raise - - datum_list = body.forward_input.data - input_tokens = sum(len(d.model_input.to_ints()) for d in datum_list) - batch_size = len(datum_list) - return await self.schedule_task( - _do_forward, - model_id=body.model_id, - token=token, - input_tokens=input_tokens, - batch_size=batch_size, - data_world_size=self.data_world_size, - task_type='forward', - ) - - @app.post('/tinker/forward_backward') - async def forward_backward( - request: Request, - body: types.ForwardBackwardRequest, - self: ModelManagement = Depends(self_fn), - ) -> types.UntypedAPIFuture: - token = await self._on_request_start(request) - - async def _do_forward_backward(): - try: - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - self.assert_resource_exists(adapter_name) - model_adapter = self.resolve_model_adapter_name(adapter_name) - datum_list = body.forward_backward_input.data - loss_fn = body.forward_backward_input.loss_fn - loss_fn_config = body.forward_backward_input.loss_fn_config or {} - preflight = getattr(self.model, 'preflight_loss_inputs', None) - if preflight is not None and loss_fn == 'importance_sampling' and not validate_tinker_loss( - loss_fn, datum_list, data_world_size=self.data_world_size): - # GRPO datums carry ragged per-sample targets; the loss check inside - # the sliced backend would deadlock the peer ranks (BUG-05), so the - # full batch is validated on every rank first. - await self.call_backend( - preflight, inputs=datum_list, adapter_name=model_adapter, loss_fn=loss_fn, **loss_fn_config) - output, loss = await self.call_backend( - self.model.tinker_forward_backward, - inputs=datum_list, - adapter_name=model_adapter, - loss_fn=loss_fn, - **loss_fn_config) - output_type = ('ImportanceSamplingLossReturn' - if loss_fn == 'importance_sampling' else 'CrossEntropyLossReturn') - self.set_resource_state(adapter_name, 'grad_ready', True) - return types.ForwardBackwardOutput( - loss_fn_output_type=output_type, - loss_fn_outputs=output, - metrics={'loss:mean': loss}, - ) - except Exception: - logger.error(traceback.format_exc()) - raise - - datum_list = body.forward_backward_input.data - input_tokens = sum(len(d.model_input.to_ints()) for d in datum_list) - batch_size = len(datum_list) - loss_fn = body.forward_backward_input.loss_fn - is_dpo = validate_tinker_loss(loss_fn, datum_list, data_world_size=self.data_world_size) - return await self.schedule_task( - _do_forward_backward, - model_id=body.model_id, - token=token, - input_tokens=input_tokens, - batch_size=batch_size, - data_world_size=self.data_world_size, - batch_size_multiple=2 if is_dpo else None, - task_type='forward_backward', - ) - - @app.post('/tinker/optim_step') - async def optim_step( - request: Request, - body: types.OptimStepRequest, - self: ModelManagement = Depends(self_fn), - ) -> types.UntypedAPIFuture: - token = await self._on_request_start(request) - - async def _do_optim(): - try: - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - self.assert_resource_exists(adapter_name) - model_adapter = self.resolve_model_adapter_name(adapter_name) - if not self.get_resource_state(adapter_name, 'grad_ready', False): - raise RuntimeError(f'No accumulated gradients for adapter={adapter_name}; ' - 'call forward_backward before optim_step') - await self.call_backend( - self.model.tinker_step, adam_params=body.adam_params, adapter_name=model_adapter) - self.set_resource_state(adapter_name, 'grad_ready', False) - metrics = await self.call_backend( - self.model.tinker_calculate_metric, is_training=True, adapter_name=model_adapter) - return types.OptimStepResponse(metrics=metrics) - except Exception: - logger.error(traceback.format_exc()) - raise - - return await self.schedule_task(_do_optim, model_id=body.model_id, token=token, task_type='optim_step') - - @app.post('/tinker/save_weights') - async def save_weights( - request: Request, - body: types.SaveWeightsRequest, - self: ModelManagement = Depends(self_fn), - ) -> types.UntypedAPIFuture: - token = await self._on_request_start(request) - - async def _do_save(): - try: - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - self.assert_resource_exists(adapter_name) - checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') - save_weights = partial( - self.call_backend, - self.model.save, - adapter_name=self.resolve_model_adapter_name(adapter_name), - save_optimizer=True) - tinker_path, _ = await checkpoint_manager.save_training(body.model_id, body.path, save_weights) - return types.SaveWeightsResponse(path=tinker_path, type='save_weights') - except Exception: - logger.error(traceback.format_exc()) - raise - - return await self.schedule_task(_do_save, model_id=body.model_id, token=token, task_type='save_weights') - - @app.post('/tinker/save_weights_for_sampler') - async def save_weights_for_sampler( - request: Request, - body: types.SaveWeightsForSamplerRequest, - self: ModelManagement = Depends(self_fn), - ) -> types.UntypedAPIFuture: - token = await self._on_request_start(request) - - async def _do_save_for_sampler(): - try: - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - self.assert_resource_exists(adapter_name) - checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') - - save_weights = partial( - self.call_backend, - self.model.save, - adapter_name=self.resolve_model_adapter_name(adapter_name), - save_optimizer=False) - - tinker_path, _ = await checkpoint_manager.save_sampler(body.model_id, body.path, save_weights) - payload = body.model_dump() - payload['model_path'] = tinker_path - metadata = await self.state.get_model_metadata(body.model_id) or {} - if metadata.get('base_model'): - payload['base_model'] = metadata['base_model'] - sampling_session_id = await self.state.create_sampling_session(payload) - # Tinker SDK distinguishes two modes by whether sampling_session_seq_id is set: - # 1. save_weights_for_sampler(name): expects path != None - # 2. save_weights_and_get_sampling_client(): expects path == None, sampling_session_id != None - if body.sampling_session_seq_id is not None: - return types.SaveWeightsForSamplerResponseInternal( - path=None, sampling_session_id=sampling_session_id) - else: - return types.SaveWeightsForSamplerResponseInternal( - path=tinker_path, sampling_session_id=sampling_session_id) - except Exception: - logger.error(traceback.format_exc()) - raise - - return await self.schedule_task( - _do_save_for_sampler, model_id=body.model_id, token=token, task_type='save_weights_for_sampler') - - @app.post('/tinker/load_weights') - async def load_weights( - request: Request, - body: types.LoadWeightsRequest, - self: ModelManagement = Depends(self_fn), - ) -> types.UntypedAPIFuture: - token = await self._on_request_start(request) - - async def _do_load(): - try: - assert self.model is not None, 'Model not loaded, please load model first' - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - self.assert_resource_exists(adapter_name) - # Path resolution (token isolation, twinkle-vs-external path shapes) is a - # server storage policy and belongs in the handler, not in the GPU-side - # backend actor. The backend receives an already-resolved location. - checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') - resolved = checkpoint_manager.resolve_load_path(body.path) - await self.call_backend( - self.model.tinker_load, - checkpoint_name=resolved.checkpoint_name, - output_dir=resolved.checkpoint_dir if resolved.is_twinkle_path else None, - load_optimizer=body.optimizer, - adapter_name=self.resolve_model_adapter_name(adapter_name)) - self.set_resource_state(adapter_name, 'grad_ready', False) - return types.LoadWeightsResponse(path=body.path, type='load_weights') - except Exception: - logger.error(traceback.format_exc()) - raise - - return await self.schedule_task(_do_load, model_id=body.model_id, token=token, task_type='load_weights') diff --git a/src/twinkle/server/processor/app.py b/src/twinkle/server/processor/app.py index b471cede8..c7a7de209 100644 --- a/src/twinkle/server/processor/app.py +++ b/src/twinkle/server/processor/app.py @@ -5,7 +5,7 @@ Provides a Ray Serve deployment for managing distributed processors (datasets, dataloaders, preprocessors, rewards, templates, weight loaders, etc.). -Follows the same structural pattern as model/app.py: +Follows the same structural pattern as training/execution/app.py: - ProcessorManagement is a top-level class inheriting ProcessorManagerMixin - Routes are registered in build_processor_app() via _register_processor_routes() - serve.ingress(app)(ProcessorManagement) applied before deployment diff --git a/src/twinkle/server/resources/allocation.py b/src/twinkle/server/resources/allocation.py index 10b4ebfdb..3b20680a1 100644 --- a/src/twinkle/server/resources/allocation.py +++ b/src/twinkle/server/resources/allocation.py @@ -7,7 +7,7 @@ from twinkle import Platform from twinkle.infra._ray.resource_manager import _gpu_placement_group_cpus -from twinkle.server.exceptions import RequestRejectedError +from twinkle.server.exceptions import RequestRejectedError, ResourceReleaseError class CapacityUnavailableError(RequestRejectedError): @@ -97,26 +97,41 @@ async def allocate(self, num_devices: int, device_type: str, epoch: str, instanc if creating.done() and self._creating.get(name) is creating: self._creating.pop(name) + async def _finish_creation(self, name): + creating = self._creating.get(name) + if creating is None: + return + # A failed create has no result to free. A pending create must stay tracked + # across cancellation, so do not remove it until this await has completed. + try: + await asyncio.shield(creating) + except Exception: + pass # A failed create has no allocation result; cancellation remains unsafe. + if self._creating.get(name) is creating: + self._creating.pop(name) + + @staticmethod + async def _find_group(name): + from ray.util.placement_group import get_placement_group + try: + return await asyncio.to_thread(get_placement_group, name) + except ValueError: + return None # A missing named group is already released. + + async def _remove(self, name): + from ray.util.placement_group import placement_group_table, remove_placement_group + await self._finish_creation(name) + pg = await self._find_group(name) + if pg is None: + return + await asyncio.to_thread(remove_placement_group, pg) + while (await asyncio.to_thread(placement_group_table, pg)).get('state') != 'REMOVED': + await asyncio.sleep(0.05) + async def remove(self, name: str) -> None: - from ray.util.placement_group import get_placement_group, placement_group_table, remove_placement_group try: - async with asyncio.timeout(self.shutdown_timeout): - creating = self._creating.get(name) - if creating is not None: - try: - await asyncio.shield(creating) - except Exception: - pass - if self._creating.get(name) is creating: - self._creating.pop(name) - try: - pg = await asyncio.to_thread(get_placement_group, name) - except ValueError: - return - await asyncio.to_thread(remove_placement_group, pg) - while True: - if (await asyncio.to_thread(placement_group_table, pg)).get('state') == 'REMOVED': - return - await asyncio.sleep(0.05) + await asyncio.wait_for(self._remove(name), self.shutdown_timeout) except TimeoutError as exc: raise TimeoutError(f'Placement group removal was not confirmed: {name}') from exc + except Exception as exc: + raise ResourceReleaseError(f'Placement group removal failed: {name}', phase='pg_removal') from exc diff --git a/src/twinkle/server/resources/runtime.py b/src/twinkle/server/resources/runtime.py index eb844d4f2..5eba8b920 100644 --- a/src/twinkle/server/resources/runtime.py +++ b/src/twinkle/server/resources/runtime.py @@ -4,12 +4,13 @@ import asyncio import logging -import os import time from copy import deepcopy -from twinkle.server.env_propagation import build_propagated_env_vars +from twinkle.server.config.resources import build_worker_env_vars +from twinkle.server.exceptions import ResourceReleaseError from .allocation import DevicePool +from .instances import InstanceLifecycle logger = logging.getLogger(__name__) @@ -21,7 +22,7 @@ def __init__(self, namespace, *, allocation_timeout, shutdown_timeout): self.namespace = namespace self.pool = DevicePool(namespace, allocation_timeout=allocation_timeout, shutdown_timeout=shutdown_timeout) self.deployments = {} - self.cleanup_tasks = set() + self.lifecycle = InstanceLifecycle('Ray deployment cleanup') def deployment_options(self, options, name, env_vars=None): options = deepcopy(options) @@ -30,10 +31,7 @@ def deployment_options(self, options, name, env_vars=None): actor = options.setdefault('ray_actor_options', {}) actor.setdefault('num_cpus', 0.1) runtime_env = actor.setdefault('runtime_env', {}) - env = build_propagated_env_vars() - for key in ('PYTHONPATH', 'TWINKLE_DEFAULT_SAVE_DIR', 'TMPDIR', 'OMP_NUM_THREADS'): - if key in os.environ: - env.setdefault(key, os.environ[key]) + env = build_worker_env_vars() env.update(runtime_env.get('env_vars', {})) env.update(env_vars or {}) env.update(TWINKLE_RAY_NAMESPACE=self.namespace, TWINKLE_RAY_REGISTRY_NAME=name + '-registry') @@ -49,32 +47,17 @@ async def deploy(self, identity, name, app, route_prefix, cleanup): except asyncio.CancelledError: async def clean_late(): - try: - await deployment - except Exception: - pass + await asyncio.gather(deployment, return_exceptions=True) await cleanup() - task = asyncio.create_task(clean_late()) - self.cleanup_tasks.add(task) - - def done(completed): - self.cleanup_tasks.discard(completed) - if not completed.cancelled(): - error = completed.exception() - if error is not None: - logger.error( - 'Late deployment cleanup failed', exc_info=(type(error), error, error.__traceback__)) - - task.add_done_callback(done) + self.lifecycle.spawn(clean_late()) raise finally: if deployment.done(): self.deployments.pop(identity, None) async def wait_for_cleanup(self): - while self.cleanup_tasks: - await asyncio.gather(*tuple(self.cleanup_tasks), return_exceptions=True) + await self.lifecycle.join() @staticmethod async def handle(name): @@ -93,10 +76,7 @@ async def shutdown(self, identity, name, *, observe_removal=None): from ray.util import list_named_actors deployment = self.deployments.get(identity) if deployment is not None: - try: - await asyncio.shield(deployment) - except Exception: - pass + await asyncio.gather(asyncio.shield(deployment), return_exceptions=True) self.deployments.pop(identity, None) # Both features give workers the controller identity, not a Serve replica ID. actors = [] @@ -113,6 +93,13 @@ async def shutdown(self, identity, name, *, observe_removal=None): actors.append(handle) await asyncio.to_thread(ray.kill, handle, no_restart=True) await asyncio.to_thread(serve.delete, name) + for handle in actors: + try: + await asyncio.wait_for(handle.__ray_ready__.remote(), 5) + except ray.exceptions.RayActorError: + continue + raise ResourceReleaseError('Owned worker death could not be confirmed', phase='worker_death') + started, status = time.monotonic(), 'error' try: await self.pool.remove(name) @@ -120,12 +107,6 @@ async def shutdown(self, identity, name, *, observe_removal=None): finally: if observe_removal is not None: observe_removal(time.monotonic() - started, status) - for handle in actors: - try: - await asyncio.wait_for(handle.__ray_ready__.remote(), 5) - except ray.exceptions.RayActorError: - continue - raise RuntimeError('Owned worker death could not be confirmed') async def recover_orphans(self, decode, shutdown): """Decode only this owner's names; never sweep another service's resources.""" diff --git a/src/twinkle/server/sampler/actor.py b/src/twinkle/server/sampler/actor.py deleted file mode 100644 index 393e85343..000000000 --- a/src/twinkle/server/sampler/actor.py +++ /dev/null @@ -1,30 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Ray composition for the sampling controller and its owned driver.""" -import asyncio - -from twinkle.server.config.application_spec import SamplerArgs -from .controller import SamplerController - - -def create_controller_actor(): - import ray - - from .driver import SamplerDriver - - @ray.remote(num_cpus=0.1, max_restarts=-1, max_task_retries=0) - class SamplerControllerActor(SamplerController): - - def __init__(self, values, namespace, pool_id, options, budgets): - args = SamplerArgs.model_validate(values) - super().__init__(args, SamplerDriver(args, namespace, pool_id, options, budgets), budgets) - - async def initialize(self): - await super().initialize() - if not self._stopping and self._monitor_task is None: - self._monitor_task = asyncio.create_task(self._monitor()) - - async def shutdown(self): - await super().shutdown() - await self.driver.runtime.wait_for_cleanup() - - return SamplerControllerActor diff --git a/src/twinkle/server/sampler/backends/__init__.py b/src/twinkle/server/sampler/backends/__init__.py deleted file mode 100644 index c2caf8e9e..000000000 --- a/src/twinkle/server/sampler/backends/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Sampler backend implementations. - -The cross-process streaming bridge lives in ``streaming`` and is re-exported here -so a sibling backend module can import it from ``.streaming`` without importing -this package root, while external consumers keep using -``from twinkle.server.sampler.backends import stream_to_queue``. -""" -from .streaming import STREAM_SENTINEL, stream_to_queue - -__all__ = ['STREAM_SENTINEL', 'stream_to_queue'] diff --git a/src/twinkle/server/sampler/backends/streaming.py b/src/twinkle/server/sampler/backends/streaming.py deleted file mode 100644 index 2ca8edf04..000000000 --- a/src/twinkle/server/sampler/backends/streaming.py +++ /dev/null @@ -1,36 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Cross-process streaming bridge for sampler backends. - -Lives in its own module (rather than the ``backends`` package ``__init__``) so -that a backend implementation (e.g. ``mock_sampler``) can import it without -importing its own package root — which would be a package-initialisation-order -dependency. ``backends/__init__`` re-exports these names, so external consumers -(``from twinkle.server.sampler.backends import stream_to_queue``) are unchanged. -""" -from __future__ import annotations - -STREAM_SENTINEL = '__STREAM_END__' - - -def stream_to_queue(sampler, queue, inputs, sampling_params=None, adapter_name='', adapter_path=None): - """Push streaming deltas from *sampler* to a cross-process Ray queue. - - Works with any object that exposes a ``sample_stream`` iterator. - """ - generation = sampler.sample_stream(inputs, sampling_params, adapter_name, adapter_path) - try: - for delta, reason in generation: - queue.put((delta, reason)) - except Exception as e: - try: - queue.put(e) - except Exception: - pass # The HTTP consumer may already have closed its queue. - finally: - close = getattr(generation, 'close', None) - if close is not None: - close() - try: - queue.put(STREAM_SENTINEL) - except Exception: - pass diff --git a/src/twinkle/server/sampling/__init__.py b/src/twinkle/server/sampling/__init__.py new file mode 100644 index 000000000..f088c2670 --- /dev/null +++ b/src/twinkle/server/sampling/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Sampling orchestration with backend execution in the execution subpackage.""" diff --git a/src/twinkle/server/sampling/actor.py b/src/twinkle/server/sampling/actor.py new file mode 100644 index 000000000..7d571e94b --- /dev/null +++ b/src/twinkle/server/sampling/actor.py @@ -0,0 +1,51 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Ray composition for the sampling controller and its owned driver.""" +import asyncio +from copy import deepcopy + +from twinkle.server.config.application_spec import SamplerArgs +from .controller import SamplerController + + +def create_controller_actor(): + import ray + + from .driver import SamplerDriver + + @ray.remote(num_cpus=0.1, max_restarts=-1, max_task_retries=0) + class SamplerControllerActor(SamplerController): + + def __init__(self, values, namespace, pool_id, options, budgets): + args = SamplerArgs.model_validate(values) + super().__init__(args, SamplerDriver(args, namespace, pool_id, options, budgets), budgets) + + async def initialize(self): + await super().initialize() + if not self._stopping and self._monitor_task is None: + self._monitor_task = asyncio.create_task(self._monitor()) + + async def shutdown(self): + await super().shutdown() + await self.driver.runtime.wait_for_cleanup() + + return SamplerControllerActor + + +def create_sampler_controller(spec, namespace, budgets, env): + """Driver-only composition. Each compute instance owns its external allocation.""" + from .controller import sampler_controller_name, sampling_pool_id + + options = deepcopy(spec.deployments[0]) if spec.deployments else {} + options.pop('name', None) + options.pop('autoscaling_config', None) + options['num_replicas'] = 1 + actor_options = options.setdefault('ray_actor_options', {}) + actor_options.setdefault('num_cpus', 0.1) + runtime_env = actor_options.setdefault('runtime_env', {}) + runtime_env['env_vars'] = {**env, **runtime_env.get('env_vars', {}), 'TWINKLE_STATE_RAY_NAMESPACE': namespace} + pool_id = sampling_pool_id(spec.route_prefix) + name = sampler_controller_name(namespace, pool_id) + controller = create_controller_actor().options( + name=name, namespace=namespace, runtime_env=runtime_env).remote(spec.args.model_dump(), namespace, pool_id, + options, budgets) + return controller, options, pool_id diff --git a/src/twinkle/server/sampler/controller.py b/src/twinkle/server/sampling/controller.py similarity index 94% rename from src/twinkle/server/sampler/controller.py rename to src/twinkle/server/sampling/controller.py index c2f20bb66..f79cb6106 100644 --- a/src/twinkle/server/sampler/controller.py +++ b/src/twinkle/server/sampling/controller.py @@ -11,11 +11,16 @@ import uuid from copy import deepcopy from dataclasses import asdict, dataclass, field +from typing import TYPE_CHECKING, Literal + +if TYPE_CHECKING: + from twinkle.server.lifecycle.proxy import ForwardedRequest, ForwardedReply from twinkle.server.config.application_spec import SamplerArgs +from twinkle.server.config.resources import ResourceOptions from twinkle.server.exceptions import ResourceExhaustedError from twinkle.server.resources.instances import InstanceLifecycle -from .concurrency import SamplerConcurrency +from twinkle.server.sampling.execution.concurrency import SamplerConcurrency logger = logging.getLogger(__name__) _CONTROL_ENDPOINTS = frozenset( @@ -34,7 +39,7 @@ def sampling_pool_id(route): class SamplingInstance: identity: str epoch: str - state: str = 'STARTING' + state: Literal['STARTING', 'READY', 'DRAINING', 'QUARANTINED', 'STOPPED'] = 'STARTING' application: str | None = None devices: list[str] = field(default_factory=list) requests: set[str] = field(default_factory=set) @@ -45,7 +50,8 @@ class SamplingInstance: class SamplerController: def __init__(self, args: SamplerArgs, driver, budgets: dict, *, clock=time.monotonic): - self.args, self.driver, self.budgets, self.clock = args, driver, budgets, clock + self.args, self.driver, self.clock = args, driver, clock + self.budgets = ResourceOptions.model_validate(budgets) self.epoch = uuid.uuid4().hex self.instances = {} self.lifecycle = InstanceLifecycle('Sampling') @@ -109,7 +115,7 @@ async def startup(): try: # Replaying control state is part of readiness, inside the same budget. - await asyncio.wait_for(startup(), self.budgets['startup_timeout_s']) + await asyncio.wait_for(startup(), self.budgets.startup_timeout_s) except BaseException: await self._close(instance) raise @@ -133,7 +139,7 @@ async def _select(self): pending = [task for task in self.lifecycle.startups.values() if not task.done()] if not pending: raise ResourceExhaustedError('No available sampler instance; device allocation was not confirmed') - await asyncio.wait_for(asyncio.shield(pending[0]), self.budgets['startup_timeout_s']) + await asyncio.wait_for(asyncio.shield(pending[0]), self.budgets.startup_timeout_s) async def _invoke_finished(self, instance, payload): ticket = uuid.uuid4().hex @@ -232,7 +238,7 @@ async def broadcast(): await self._release(replica, ticket) await self._release(seed, seed_ticket) - async def request(self, payload, stream_queue=None, full_path=None): + async def request(self, payload: ForwardedRequest, stream_queue=None, full_path=None) -> ForwardedReply: await self.initialize() endpoint = payload['endpoint'] if endpoint == 'twinkle/create': @@ -294,7 +300,7 @@ async def stopped(): instance.requests.clear() await self.lifecycle.close(instance.identity, lambda: self.driver.shutdown(instance), - self.budgets['shutdown_timeout_s'], stopped, quarantined) + self.budgets.shutdown_timeout_s, stopped, quarantined) async def sweep(self): async with self._lock: @@ -332,7 +338,7 @@ async def health(instance): await self.sweep() except Exception: logger.exception('Sampler monitor failed; ownership remains reserved') - await asyncio.sleep(self.budgets['sweep_interval_s']) + await asyncio.sleep(self.budgets.sweep_interval_s) async def diagnostics(self): await self.initialize() diff --git a/src/twinkle/server/sampler/driver.py b/src/twinkle/server/sampling/driver.py similarity index 87% rename from src/twinkle/server/sampler/driver.py rename to src/twinkle/server/sampling/driver.py index 1435efcc8..cad3e85c2 100644 --- a/src/twinkle/server/sampler/driver.py +++ b/src/twinkle/server/sampling/driver.py @@ -5,20 +5,22 @@ import asyncio from twinkle.server.config.application_spec import SamplerArgs +from twinkle.server.config.resources import ResourceOptions from twinkle.server.resources.runtime import RayRuntime from twinkle.server.state import get_server_state -from twinkle.server.state.models import FutureFailureRecord +from twinkle.server.state.models import FutureFailureRecord, RuntimeContext class SamplerDriver: def __init__(self, args: SamplerArgs, namespace: str, pool_id: str, options: dict, budgets: dict): self.args, self.namespace, self.pool_id = args, namespace, pool_id - self.options, self.budgets = options, budgets + self.options = options + self.budgets = ResourceOptions.model_validate(budgets) self.runtime = RayRuntime( namespace, - allocation_timeout=budgets['allocation_timeout_s'], - shutdown_timeout=budgets['shutdown_timeout_s']) + allocation_timeout=self.budgets.allocation_timeout_s, + shutdown_timeout=self.budgets.shutdown_timeout_s) @property def prefix(self): @@ -27,7 +29,7 @@ def prefix(self): async def start(self, instance): from ray.util.placement_group import get_placement_group - from .app import build_sampler_app + from twinkle.server.sampling.execution.app import build_sampler_app group = self.args.device_group device_type = str(group.get('device_type', 'cpu' if self.args.sampler_type == 'mock' else 'cuda')).lower() allocation = await self.runtime.pool.allocate( @@ -45,11 +47,8 @@ async def start(self, instance): **values, queue_config=self.args.queue_config, deploy_options=options, - sampling_context={ - 'namespace': self.namespace, - 'registry_name': allocation.name + '-registry', - 'instance_id': instance.identity - }, + sampling_context=RuntimeContext( + namespace=self.namespace, registry_name=allocation.name + '-registry', instance_id=instance.identity), placement_group=await asyncio.to_thread(get_placement_group, allocation.name), allocation_devices=list(allocation.devices)) await self.runtime.deploy(instance.identity, allocation.name, app, diff --git a/src/twinkle/server/sampler/__init__.py b/src/twinkle/server/sampling/execution/__init__.py similarity index 100% rename from src/twinkle/server/sampler/__init__.py rename to src/twinkle/server/sampling/execution/__init__.py diff --git a/src/twinkle/server/sampler/app.py b/src/twinkle/server/sampling/execution/app.py similarity index 95% rename from src/twinkle/server/sampler/app.py rename to src/twinkle/server/sampling/execution/app.py index c3faf2867..00aac7956 100644 --- a/src/twinkle/server/sampler/app.py +++ b/src/twinkle/server/sampling/execution/app.py @@ -21,8 +21,8 @@ from twinkle.server.middleware.auth import get_token_from_request from twinkle.server.runtime import init_twinkle_runtime from twinkle.server.state import ServerState, get_server_state +from twinkle.server.state.models import RuntimeContext from twinkle.server.task_queue.mixin import TaskQueueMixin -from twinkle.server.task_queue.types import execution_scope from twinkle.server.utils import get_template_for_model, wrap_builder_with_device_group_env from twinkle.utils.logger import get_logger from .concurrency import SamplerConcurrency @@ -97,10 +97,11 @@ async def __init__(self, self.device_group = DeviceGroup(**device_group) context = kwargs.pop('sampling_context', None) if context is not None: + context = RuntimeContext.model_validate(context) from twinkle.infra._ray.ray_helper import RayHelper from twinkle.infra._ray.resource_manager import ResourceManager - os.environ['TWINKLE_RAY_NAMESPACE'] = context['namespace'] - os.environ['TWINKLE_RAY_REGISTRY_NAME'] = context['registry_name'] + os.environ['TWINKLE_RAY_NAMESPACE'] = context.namespace + os.environ['TWINKLE_RAY_REGISTRY_NAME'] = context.registry_name RayHelper.resource_manager = ResourceManager.from_allocation(self.device_group, kwargs.pop('placement_group'), kwargs.pop('allocation_devices')) @@ -185,17 +186,7 @@ async def runtime_identity(self): async def check_health(self) -> None: self.sampling_guard.check_available() - @serve.multiplexed(max_num_models_per_replica=5) - async def _sticky_entry(self, sticky_key: str): - return sticky_key - - async def _ensure_sticky(self): - sticky_key = serve.get_multiplexed_model_id() - await self._sticky_entry(sticky_key) - async def _on_request_start(self, request: Request) -> str: - if execution_scope.get() is None: - await self._ensure_sticky() await self.state.touch_replica_last_seen(self.replica_id) await self._ensure_state_cleanup_started() return get_token_from_request(request) diff --git a/src/twinkle/server/sampling/execution/backends/__init__.py b/src/twinkle/server/sampling/execution/backends/__init__.py new file mode 100644 index 000000000..a734d4690 --- /dev/null +++ b/src/twinkle/server/sampling/execution/backends/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Sampling backend implementations; helpers are imported from their defining modules.""" diff --git a/src/twinkle/server/sampler/backends/mock_sampler.py b/src/twinkle/server/sampling/execution/backends/mock_sampler.py similarity index 100% rename from src/twinkle/server/sampler/backends/mock_sampler.py rename to src/twinkle/server/sampling/execution/backends/mock_sampler.py diff --git a/src/twinkle/server/sampling/execution/backends/streaming.py b/src/twinkle/server/sampling/execution/backends/streaming.py new file mode 100644 index 000000000..5d07e1b59 --- /dev/null +++ b/src/twinkle/server/sampling/execution/backends/streaming.py @@ -0,0 +1,44 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Lightweight queue bridge shared by real and mock sampling backends. + +Kept independent of the HTTP runtime so backend workers do not import FastAPI +or depend on the mock backend. Both consumers import this defining module. +""" +from __future__ import annotations + +STREAM_SENTINEL = '__STREAM_END__' + + +class StreamPublisher: + """A disconnected consumer cannot end the backend generation lifetime.""" + + def __init__(self, queue): + self.queue = queue + self.connected = True + + def put(self, item): + if not self.connected: + return + try: + self.queue.put(item) + except Exception: + # Ray queue transport can disappear after the HTTP reader closes. + self.connected = False + + +def stream_to_queue(sampler, queue, inputs, sampling_params=None, adapter_name='', adapter_path=None): + """Drain and close the generator even when its HTTP consumer disconnects.""" + from contextlib import ExitStack + + publisher = StreamPublisher(queue) + generation = sampler.sample_stream(inputs, sampling_params, adapter_name, adapter_path) + with ExitStack() as cleanup: + cleanup.callback(publisher.put, STREAM_SENTINEL) + close = getattr(generation, 'close', None) + if close is not None: + cleanup.callback(close) + try: + for delta, reason in generation: + publisher.put((delta, reason)) + except Exception as error: + publisher.put(error) diff --git a/src/twinkle/server/sampler/backends/vllm_async.py b/src/twinkle/server/sampling/execution/backends/vllm_async.py similarity index 100% rename from src/twinkle/server/sampler/backends/vllm_async.py rename to src/twinkle/server/sampling/execution/backends/vllm_async.py diff --git a/src/twinkle/server/sampler/concurrency.py b/src/twinkle/server/sampling/execution/concurrency.py similarity index 98% rename from src/twinkle/server/sampler/concurrency.py rename to src/twinkle/server/sampling/execution/concurrency.py index b3e7accc0..0acb32232 100644 --- a/src/twinkle/server/sampler/concurrency.py +++ b/src/twinkle/server/sampling/execution/concurrency.py @@ -4,7 +4,7 @@ import asyncio -from twinkle.server.task_queue.types import BackendBusyError +from twinkle.server.exceptions import BackendBusyError class SamplerConcurrency: diff --git a/src/twinkle/server/sampler/runtime.py b/src/twinkle/server/sampling/execution/runtime.py similarity index 79% rename from src/twinkle/server/sampler/runtime.py rename to src/twinkle/server/sampling/execution/runtime.py index bb80f878b..5549cdeec 100644 --- a/src/twinkle/server/sampler/runtime.py +++ b/src/twinkle/server/sampling/execution/runtime.py @@ -3,25 +3,37 @@ from __future__ import annotations import asyncio +from dataclasses import dataclass from twinkle.server.deployment import build_deployment_app from twinkle.server.exceptions import ResourceNotFoundError -from twinkle.server.task_queue.types import ExecutionScope, execution_scope +from twinkle.server.lifecycle.proxy import ForwardedReply, ForwardedRequest +from twinkle.server.task_queue.types import ExecutionScope + + +def request_execution_scope(request) -> ExecutionScope | None: + """Internal ASGI metadata, absent for direct standalone endpoint calls.""" + return request.scope.get('state', {}).get('execution_scope') + + +@dataclass +class SamplingOperation: + scope: ExecutionScope + task: asyncio.Task | None = None class SamplerRuntime: def __init__(self, owner): self.owner = owner - self._operations = {} + self._operations: dict[str, SamplingOperation] = {} self._stream_tasks = set() self._rpc_app = None - def start_stream_job(self, factory): + def start_stream_job(self, factory, *, execution_scope: ExecutionScope | None = None): job = asyncio.create_task(factory()) - scope = execution_scope.get() - if scope is not None: - scope.jobs.add(job) + if execution_scope is not None: + execution_scope.track(job) self._stream_tasks.add(job) def done(task): @@ -35,7 +47,7 @@ def done(task): async def shutdown(self): await asyncio.gather(*tuple(self._stream_tasks), return_exceptions=True) - async def invoke(self, ticket, payload, stream_queue=None): + async def invoke(self, ticket: str, payload: ForwardedRequest, stream_queue=None) -> ForwardedReply: """Return the HTTP reply early while retaining every actual generation job.""" if ticket in self._operations: raise ResourceNotFoundError('Sampling request was already dispatched') @@ -50,12 +62,11 @@ def routes(app, get_self): self._rpc_app = build_deployment_app('Sampler', routes) loop = asyncio.get_running_loop() reply = loop.create_future() - record = {'scope': ExecutionScope(), 'task': None} + record = SamplingOperation(ExecutionScope()) self._operations[ticket] = record async def dispatch(): - token = execution_scope.set(record['scope']) - response = {'status': 500, 'headers': [], 'body': b''} + response: ForwardedReply = {'status': 500, 'headers': [], 'body': b''} sent = False async def receive(): @@ -83,6 +94,9 @@ async def send(message): scope = { 'type': 'http', + 'state': { + 'execution_scope': record.scope + }, 'asgi': { 'version': '3.0' }, @@ -106,24 +120,23 @@ async def send(message): reply.set_exception(error) raise finally: - execution_scope.reset(token) # ComputeWorker supplied completion handles, not a terminal status read. - await record['scope'].wait() + await record.scope.wait() if stream_queue is not None: try: await stream_queue.put_async(None) except Exception: pass - record['task'] = asyncio.create_task(dispatch()) - record['task'].add_done_callback(lambda task: task.exception() if not task.cancelled() else None) + record.task = asyncio.create_task(dispatch()) + record.task.add_done_callback(lambda task: task.exception() if not task.cancelled() else None) return await asyncio.shield(reply) async def wait_operation(self, ticket): record = self._operations.get(ticket) if record is None: raise ResourceNotFoundError('Unknown sampling operation') - await asyncio.shield(record['task']) + await asyncio.shield(record.task) self.owner.sampling_guard.check_available() self._operations.pop(ticket, None) @@ -133,7 +146,7 @@ async def install_base_weights(self, path): lambda: resolve_sampler_weights(self.owner, path, full_checkpoint=True), exclusive=True) async def pending_futures(self): - return list({identity for record in self._operations.values() for identity in record['scope'].request_ids}) + return list({identity for record in self._operations.values() for identity in record.scope.request_ids}) async def runtime_identity(self): import ray diff --git a/src/twinkle/server/sampler/tinker_handlers.py b/src/twinkle/server/sampling/execution/tinker_handlers.py similarity index 94% rename from src/twinkle/server/sampler/tinker_handlers.py rename to src/twinkle/server/sampling/execution/tinker_handlers.py index 719ab0621..ac585bad7 100644 --- a/src/twinkle/server/sampler/tinker_handlers.py +++ b/src/twinkle/server/sampling/execution/tinker_handlers.py @@ -20,9 +20,10 @@ from twinkle.data_format import SamplingParams from twinkle.server.checkpoint import create_checkpoint_manager -from twinkle.server.sampler.weights import sample_with_weights -from twinkle.server.task_queue.types import UserTaskError +from twinkle.server.exceptions import RequestRejectedError +from twinkle.server.sampling.execution.weights import sample_with_weights from twinkle.utils.logger import get_logger +from .runtime import request_execution_scope logger = get_logger() @@ -84,7 +85,7 @@ async def _do_sample(): # Base-model sampling is valid when no model_path was provided. if adapter_uri and not os.path.exists(adapter_uri): - raise UserTaskError(f'Adapter URI {model_path} does not exist. Please check the model_path.') + raise RequestRejectedError(f'Adapter URI {model_path} does not exist. Please check the model_path.') # Convert tinker SamplingParams to twinkle SamplingParams if needed sampling_params = None @@ -153,4 +154,5 @@ async def generate(lora_path): token=token, input_tokens=input_tokens, task_type='sample', + execution_scope=request_execution_scope(request), ) diff --git a/src/twinkle/server/sampler/twinkle_handlers.py b/src/twinkle/server/sampling/execution/twinkle_handlers.py similarity index 96% rename from src/twinkle/server/sampler/twinkle_handlers.py rename to src/twinkle/server/sampling/execution/twinkle_handlers.py index b23f65e4a..92fc8dc5f 100644 --- a/src/twinkle/server/sampler/twinkle_handlers.py +++ b/src/twinkle/server/sampling/execution/twinkle_handlers.py @@ -28,11 +28,12 @@ from twinkle.protocol.types import sampler as sampler_types from twinkle.server.exceptions import EndpointUnavailableError, RequestRejectedError from twinkle.server.lifecycle.submit import backend_kwargs, resolve_twinkle_adapter_name, to_backend_inputs -from twinkle.server.sampler.weights import resolve_sampler_weights, sample_with_weights +from twinkle.server.sampling.execution.weights import resolve_sampler_weights, sample_with_weights from twinkle.server.task_errors import task_error_payload from twinkle.server.telemetry.correlation import MODEL_ID from twinkle.server.telemetry.tracing import traced_operation from twinkle.utils.logger import get_logger +from .runtime import request_execution_scope logger = get_logger() @@ -268,7 +269,12 @@ async def generate(adapter_path): # Calculate metrics for queue scheduling. The body is wire-validated, so the # entries are models and ``input_ids`` is absent or a list of ints. input_tokens = sum(len(getattr(entry, 'input_ids', None) or ()) for entry in body.inputs) - return await self.submit_concurrent_and_peek(_task, token=token, input_tokens=input_tokens, task_type='sample') + return await self.submit_concurrent_and_peek( + _task, + token=token, + input_tokens=input_tokens, + task_type='sample', + execution_scope=request_execution_scope(request)) @app.post('/twinkle/sample_to_data_plane', response_model=types.TaskEnvelope) async def sample_to_data_plane( @@ -325,7 +331,10 @@ async def generate(): return ref.model_dump() return await self.submit_concurrent_and_peek( - _generate_and_store, model_id=full_adapter_name or None, task_type='sample_to_data_plane') + _generate_and_store, + model_id=full_adapter_name or None, + task_type='sample_to_data_plane', + execution_scope=request_execution_scope(request)) @app.post('/twinkle/unload_adapter_paths') async def unload_adapter_paths( @@ -437,7 +446,7 @@ async def sample_stream( from ray.util.queue import Queue - from .backends import STREAM_SENTINEL + from .backends.streaming import STREAM_SENTINEL request_id = f'req_{uuid.uuid4().hex}' actors = self.sampler._actors @@ -474,7 +483,7 @@ async def _stream_generation(adapter_path): except Exception: pass - self.request_runtime.start_stream_job(stream_job) + self.request_runtime.start_stream_job(stream_job, execution_scope=request_execution_scope(request)) return StreamingResponse( _stream_queue( q, diff --git a/src/twinkle/server/sampler/weights.py b/src/twinkle/server/sampling/execution/weights.py similarity index 84% rename from src/twinkle/server/sampler/weights.py rename to src/twinkle/server/sampling/execution/weights.py index f2d1077f7..dac788542 100644 --- a/src/twinkle/server/sampler/weights.py +++ b/src/twinkle/server/sampling/execution/weights.py @@ -7,7 +7,8 @@ from safetensors import SafetensorError from twinkle.exceptions import FullWeightLoadError -from twinkle.server.task_queue.types import BACKEND_TIMEOUT_EXCEPTIONS, UserTaskError +from twinkle.server.exceptions import RequestRejectedError +from twinkle.server.task_queue.backend_gate import BACKEND_TIMEOUT_EXCEPTIONS def is_full_checkpoint(resolved_uri: str | None) -> bool: @@ -16,10 +17,10 @@ def is_full_checkpoint(resolved_uri: str | None) -> bool: if os.path.isfile(os.path.join(resolved_uri, 'adapter_config.json')): return False if os.path.isabs(resolved_uri) and not os.path.isdir(resolved_uri): - raise UserTaskError(f'Checkpoint directory does not exist: {resolved_uri}') + raise RequestRejectedError(f'Checkpoint directory does not exist: {resolved_uri}') if os.path.isdir(resolved_uri) and not (os.path.isfile(os.path.join(resolved_uri, 'model.safetensors.index.json')) or glob.glob(os.path.join(resolved_uri, '*.safetensors'))): - raise UserTaskError(f'No .safetensors weights found under {resolved_uri}') + raise RequestRejectedError(f'No .safetensors weights found under {resolved_uri}') return True @@ -38,7 +39,7 @@ async def resolve_sampler_weights(service, resolved_uri: str | None, *, full_che service.sampling_guard.invalidate_weights(error) raise except (ValueError, OSError, KeyError, SafetensorError) as error: - raise UserTaskError(str(error)) from error + raise RequestRejectedError(str(error)) from error return None diff --git a/src/twinkle/server/sampler/router.py b/src/twinkle/server/sampling/router.py similarity index 72% rename from src/twinkle/server/sampler/router.py rename to src/twinkle/server/sampling/router.py index 9e767346d..eabb3a015 100644 --- a/src/twinkle/server/sampler/router.py +++ b/src/twinkle/server/sampling/router.py @@ -4,18 +4,16 @@ import asyncio import json -import uuid -from copy import deepcopy from fastapi import Depends, Request, Response from fastapi.responses import StreamingResponse from twinkle.server.checkpoint import create_checkpoint_manager from twinkle.server.deployment import bind_deployment, build_deployment_app +from twinkle.server.lifecycle.proxy import ForwardedRequest from twinkle.server.middleware.auth import get_token_from_request +from twinkle.server.sampling.execution.weights import is_full_checkpoint from twinkle.server.state import get_server_state -from .actor import create_controller_actor -from .controller import sampler_controller_name, sampling_pool_id -from .weights import is_full_checkpoint +from .controller import sampler_controller_name class SamplerRouter: @@ -27,7 +25,7 @@ async def __init__(self, namespace, pool_id): async def forward(self, request, endpoint): body = await request.body() - payload = { + payload: ForwardedRequest = { 'endpoint': endpoint, 'body': body, 'method': request.method, @@ -104,33 +102,3 @@ async def tinker(request: Request, endpoint: str, self=Depends(get_self)): app = build_deployment_app('SamplerRouter', routes) return bind_deployment( app, SamplerRouter, deploy_options, deployment_name='SamplerRouter', bind_args=(namespace, pool_id)) - - -def deploy_sampler_pool(spec, namespace, budgets, env): - """Driver-only composition. Each compute instance owns its external allocation.""" - import ray - from ray import serve - options = deepcopy(spec.deployments[0]) if spec.deployments else {} - options.pop('name', None) - options.pop('autoscaling_config', None) - options['num_replicas'] = 1 - actor_options = options.setdefault('ray_actor_options', {}) - actor_options.setdefault('num_cpus', 0.1) - runtime_env = actor_options.setdefault('runtime_env', {}) - runtime_env['env_vars'] = {**env, **runtime_env.get('env_vars', {}), 'TWINKLE_STATE_RAY_NAMESPACE': namespace} - pool_id = sampling_pool_id(spec.route_prefix) - name = sampler_controller_name(namespace, pool_id) - controller = create_controller_actor().options( - name=name, namespace=namespace, runtime_env=runtime_env).remote(spec.args.model_dump(), namespace, pool_id, - options, budgets) - try: - ray.get(controller.initialize.remote()) - app = build_sampler_router(namespace, pool_id, options) - serve.run(app, name=spec.name, route_prefix=spec.route_prefix) - except BaseException: - try: - ray.get(controller.shutdown.remote()) - finally: - ray.kill(controller, no_restart=True) - raise - return controller diff --git a/src/twinkle/server/session_resource/__init__.py b/src/twinkle/server/session_resource/__init__.py index fdb7ebdc1..7384532cb 100644 --- a/src/twinkle/server/session_resource/__init__.py +++ b/src/twinkle/server/session_resource/__init__.py @@ -7,8 +7,7 @@ session-driven expiration of session-bound resources. """ -from .adapter import AdapterManagerMixin from .base import SessionResourceMixin from .processor import ProcessorManagerMixin -__all__ = ['AdapterManagerMixin', 'ProcessorManagerMixin', 'SessionResourceMixin'] +__all__ = ['ProcessorManagerMixin', 'SessionResourceMixin'] diff --git a/src/twinkle/server/session_resource/adapter.py b/src/twinkle/server/session_resource/adapter.py deleted file mode 100644 index 392a05bf1..000000000 --- a/src/twinkle/server/session_resource/adapter.py +++ /dev/null @@ -1,86 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -""" -Adapter Lifecycle Manager Mixin for Twinkle Server. - -This module provides adapter lifecycle management as a mixin class that can be -inherited directly by services. It tracks adapter activity and provides interfaces -for registration, heartbeat updates, and expiration handling. - -By inheriting this mixin, services can override the _on_adapter_expired() method -to handle expired adapters without using callbacks or polling. -""" -from __future__ import annotations - -from abc import abstractmethod - -from twinkle.utils.logger import get_logger -from .base import SessionResourceMixin - -logger = get_logger() - - -class AdapterManagerMixin(SessionResourceMixin): - """Mixin for adapter lifecycle management with session-based expiration. - - This mixin tracks adapter activity and automatically expires adapters - when their associated session expires. - - Inheriting classes should: - 1. Call _init_adapter_manager() in __init__ - 2. Override _on_adapter_expired() to customize expiration handling - - The inactivity timeout / max lifetime are stored on the base mixin as - ``_resource_timeout`` / ``_resource_max_lifetime`` (set via - ``_init_adapter_manager``). - """ - - # Set resource type for logging - _resource_type = 'Adapter' - - def _init_adapter_manager( - self, - adapter_timeout: float = 1800.0, - adapter_max_lifetime: float = 36000.0, - ) -> None: - """Initialize the adapter manager. - - This should be called in the __init__ of the inheriting class. - - Args: - adapter_timeout: Timeout in seconds used to check whether a session is still alive. - Default is 1800.0 (30 minutes). - adapter_max_lifetime: Maximum lifetime in seconds for an adapter regardless of session - liveness. Adapters older than this are treated as expired. Default is 36000.0 (10 hours). - """ - self._init_resource_manager( - resource_timeout=adapter_timeout, - resource_max_lifetime=adapter_max_lifetime, - ) - - async def _on_resource_expired(self, resource_id: str) -> None: - """Base-class expiry hook; forwards to the domain hook ``_on_adapter_expired``. - - ``_on_adapter_expired`` is the supported extension point: the adapter-domain - name is kept deliberately so subclass authors override a method named for - adapters rather than the generic base-class resource hook. - """ - await self._on_adapter_expired(resource_id) - - @abstractmethod - async def _on_adapter_expired(self, adapter_name: str) -> None: - """Hook method called when an adapter expires.""" - ... - - @staticmethod - def get_adapter_name(adapter_name: str) -> str: - """Get the adapter name for a request. - - This is a passthrough method for consistency with the original API. - - Args: - adapter_name: The adapter name (typically model_id) - - Returns: - The adapter name to use - """ - return adapter_name diff --git a/src/twinkle/server/session_resource/processor.py b/src/twinkle/server/session_resource/processor.py index 69a6192e3..17c398aa5 100644 --- a/src/twinkle/server/session_resource/processor.py +++ b/src/twinkle/server/session_resource/processor.py @@ -2,7 +2,7 @@ """ Processor Lifecycle Manager Mixin for Twinkle Server. -Mirrors AdapterManagerMixin but adds a global per-token processor limit. +Adds a global per-token processor limit to session resources. Sessions are tracked via session ID; processors expire when their session expires. """ from __future__ import annotations @@ -18,7 +18,7 @@ class ProcessorManagerMixin(SessionResourceMixin): """Mixin for processor lifecycle management with session-based expiration. - Mirrors AdapterManagerMixin with an additional per-token processor limit. + Manages processor sessions with a global per-token processor limit. Inheriting classes should: 1. Call _init_processor_manager() in __init__ @@ -81,8 +81,6 @@ async def _on_resource_expired(self, resource_id: str) -> None: ``_on_processor_expired`` is the supported extension point: the processor-domain name is kept deliberately so subclass authors override a method named for processors rather than the generic base-class hook. It is - ``async`` to match the sibling ``AdapterManagerMixin._on_adapter_expired`` - contract, so both resource kinds expose the same extension-point shape. """ await self._on_processor_expired(resource_id) diff --git a/src/twinkle/server/state/backend/redis_backend.py b/src/twinkle/server/state/backend/redis_backend.py index dc180e7d6..53b4d78ce 100644 --- a/src/twinkle/server/state/backend/redis_backend.py +++ b/src/twinkle/server/state/backend/redis_backend.py @@ -60,7 +60,7 @@ async def set(self, key: str, value: Any, ttl: int | None = None) -> None: else: await self._client.set(real_key, data) - async def _read_with_reconnect(self, operation, *args): + async def _retry_read(self, operation, *args): """Retry only read commands after a stale pooled connection is discarded. Writes/EXEC are deliberately excluded: a lost acknowledgement does not @@ -80,7 +80,7 @@ async def _read_with_reconnect(self, operation, *args): async def get(self, key: str) -> Any | None: """Retrieve value, return None if not found or expired.""" real_key = self._make_key(key) - raw = await self._read_with_reconnect(self._client.get, real_key) + raw = await self._retry_read(self._client.get, real_key) if raw is None: return None return json.loads(raw) @@ -93,7 +93,7 @@ async def delete(self, key: str) -> None: async def exists(self, key: str) -> bool: """Check if key exists and is not expired.""" real_key = self._make_key(key) - return bool(await self._read_with_reconnect(self._client.exists, real_key)) + return bool(await self._retry_read(self._client.exists, real_key)) async def keys(self, pattern: str) -> list[str]: """Return all key names matching the pattern using SCAN. @@ -191,7 +191,7 @@ async def mget(self, keys: list[str]) -> list[Any | None]: if not keys: return [] real_keys = [self._make_key(k) for k in keys] - raw_values = await self._read_with_reconnect(self._client.mget, real_keys) + raw_values = await self._retry_read(self._client.mget, real_keys) return [json.loads(v) if v is not None else None for v in raw_values] async def close(self) -> None: @@ -201,7 +201,7 @@ async def close(self) -> None: async def health_check(self) -> bool: """Check if Redis is healthy and available.""" try: - healthy = bool(await self._read_with_reconnect(self._client.ping)) + healthy = bool(await self._retry_read(self._client.ping)) except Exception: healthy = False self._metrics.redis_health.set(1 if healthy else 0) diff --git a/src/twinkle/server/state/future_manager.py b/src/twinkle/server/state/future_manager.py index fbfe316e7..4e36d28b4 100644 --- a/src/twinkle/server/state/future_manager.py +++ b/src/twinkle/server/state/future_manager.py @@ -5,6 +5,7 @@ import time from typing import Any +from twinkle.protocol.types.lifecycle import TERMINAL_STATUSES from twinkle.utils.logger import get_logger from .backend.base import StateBackend from .base import BaseManager @@ -12,114 +13,114 @@ logger = get_logger() -# Status sets used by the do-not-regress guard inside the atomic transform. -_TERMINAL_STATUSES = frozenset({'completed', 'failed', 'cancelled'}) -_STATUS_ORDER = {'pending': 0, 'queued': 1, 'running': 2} - - -def _future_record_transform( - existing: dict | None, - *, - request_id: str, - new_status: str, - model_id: str | None, - reason: str | None, - result: Any, - failure: dict[str, Any] | None, - queue_state: str | None, - queue_state_reason: str | None, - replica_id: str | None, - absolute_deadline: float | None, - now: str, - owner_token: str | None = None, -) -> dict | None: - """Atomic transform body for :meth:`FutureManager.store_status`. - - Module-level so it remains picklable when forwarded across the Ray actor - boundary (closures and lambdas cannot be). - - A record already in a terminal state is never overwritten (returns ``None``, - which ``update_atomic`` treats as "keep current value"). A write of a - *different* terminal state is logged; a write of the *same* terminal state is - dropped silently (State_Backend idempotent retries produce these and they - indicate no defect). - """ - existing_status = existing.get('status') if existing is not None else None - if existing_status in _TERMINAL_STATUSES: - if new_status != existing_status: - logger.warning('future %s already terminal as %r; refusing %r', request_id, existing_status, new_status) - return None - if new_status in _STATUS_ORDER and existing_status in _STATUS_ORDER: - if _STATUS_ORDER[new_status] <= _STATUS_ORDER[existing_status]: - return None - if new_status == 'cancelled' and existing_status == 'running': - return None - - if existing is None: - record = FutureRecord( - owner_token=owner_token, - status=new_status, - model_id=model_id, - reason=reason, - result=result, - failure=failure, - queue_state=queue_state, - queue_state_reason=queue_state_reason, - replica_id=replica_id, - absolute_deadline=absolute_deadline, - created_at=now, - updated_at=now, - ) - return record.model_dump() - - updated = dict(existing) - updated['status'] = new_status - updated['model_id'] = model_id - updated['updated_at'] = now - # replica_id is set at creation and is deliberately NOT overwritten here. - if reason is not None: - updated['reason'] = reason - if new_status in _TERMINAL_STATUSES: - if new_status in ('failed', 'cancelled'): - updated['result'] = None - updated['failure'] = failure - else: - updated['failure'] = None - updated['result'] = result - elif result is not None: - updated['result'] = result - if queue_state is not None: - updated['queue_state'] = queue_state - if queue_state_reason is not None: - updated['queue_state_reason'] = queue_state_reason - return updated +class FutureTransitions: + """Pure, picklable atomic transitions; the manager owns I/O and retention.""" -_CANCELLABLE_STATUSES = frozenset({'pending', 'queued'}) + terminal = TERMINAL_STATUSES + cancellable = frozenset({'pending', 'queued'}) + order = {'pending': 0, 'queued': 1, 'running': 2} + @staticmethod + def update( + existing: dict | None, + *, + request_id: str, + new_status: str, + model_id: str | None, + reason: str | None, + result: Any, + failure: dict[str, Any] | None, + queue_state: str | None, + queue_state_reason: str | None, + replica_id: str | None, + absolute_deadline: float | None, + now: str, + owner_token: str | None = None, + ) -> dict | None: + """Atomic transform body for :meth:`FutureManager.store_status`. -def _start_if_pending_transform(existing: dict | None, *, now: str) -> dict | None: - """Acquire execution only from pending/queued, using the existing record.""" - if existing is None or existing.get('status') not in _CANCELLABLE_STATUSES: - return None - return {**existing, 'status': 'running', 'queue_state': 'active', 'updated_at': now} + A static method remains picklable across the Ray actor boundary. + A record already in a terminal state is never overwritten (returns ``None``, + which ``update_atomic`` treats as "keep current value"). A write of a + *different* terminal state is logged; a write of the *same* terminal state is + dropped silently (State_Backend idempotent retries produce these and they + indicate no defect). + """ + existing_status = existing.get('status') if existing is not None else None + if existing_status in FutureTransitions.terminal: + if new_status != existing_status: + logger.warning('future %s already terminal as %r; refusing %r', request_id, existing_status, new_status) + return None + if new_status in FutureTransitions.order and existing_status in FutureTransitions.order: + if FutureTransitions.order[new_status] <= FutureTransitions.order[existing_status]: + return None + if new_status == 'cancelled' and existing_status == 'running': + return None -def _cancel_if_not_started_transform(existing: dict | None, *, failure: dict, now: str) -> dict | None: - """Atomic transform: cancel iff the task has not started (pending/queued). + if existing is None: + record = FutureRecord( + owner_token=owner_token, + status=new_status, + model_id=model_id, + reason=reason, + result=result, + failure=failure, + queue_state=queue_state, + queue_state_reason=queue_state_reason, + replica_id=replica_id, + absolute_deadline=absolute_deadline, + created_at=now, + updated_at=now, + ) + return record.model_dump() + + updated = dict(existing) + updated['status'] = new_status + updated['model_id'] = model_id + updated['updated_at'] = now + # replica_id is set at creation and is deliberately NOT overwritten here. + if reason is not None: + updated['reason'] = reason + if new_status in FutureTransitions.terminal: + if new_status in ('failed', 'cancelled'): + updated['result'] = None + updated['failure'] = failure + else: + updated['failure'] = None + updated['result'] = result + elif result is not None: + updated['result'] = result + if queue_state is not None: + updated['queue_state'] = queue_state + if queue_state_reason is not None: + updated['queue_state_reason'] = queue_state_reason + return updated + + @staticmethod + def start(existing: dict | None, *, now: str) -> dict | None: + """Acquire execution only from pending/queued, using the existing record.""" + if existing is None or existing.get('status') not in FutureTransitions.cancellable: + return None + return {**existing, 'status': 'running', 'queue_state': 'active', 'updated_at': now} + + @staticmethod + def cancel(existing: dict | None, *, failure: dict, now: str) -> dict | None: + """Atomic transform: cancel iff the task has not started (pending/queued). - Returns ``None`` (no change) for running/terminal/missing records so a task - already executing is never interrupted -- cancel is best-effort on the queue. - """ - status = existing.get('status') if existing is not None else None - if status not in _CANCELLABLE_STATUSES: - return None - updated = dict(existing) - updated['status'] = 'cancelled' - updated['result'] = None - updated['failure'] = failure - updated['updated_at'] = now - return updated + Returns ``None`` (no change) for running/terminal/missing records so a task + already executing is never interrupted -- cancel is best-effort on the queue. + """ + status = existing.get('status') if existing is not None else None + if status not in FutureTransitions.cancellable: + return None + updated = dict(existing) + updated['status'] = 'cancelled' + updated['result'] = None + updated['failure'] = failure + updated['updated_at'] = now + return updated class FutureManager(BaseManager[FutureRecord]): @@ -168,7 +169,7 @@ async def store_status( await self._backend.update_atomic( self._make_key(request_id), functools.partial( - _future_record_transform, + FutureTransitions.update, request_id=request_id, new_status=status, model_id=model_id, @@ -188,7 +189,7 @@ async def try_start(self, request_id: str) -> bool: """True only when this attempt atomically acquired execution rights.""" return await self._backend.update_atomic( self._make_key(request_id), - functools.partial(_start_if_pending_transform, now=_now_iso()), + functools.partial(FutureTransitions.start, now=_now_iso()), return_changed=True, ) @@ -207,7 +208,7 @@ async def cancel_if_pending(self, request_id: str) -> str | None: ) result = await self._backend.update_atomic( self._make_key(request_id), - functools.partial(_cancel_if_not_started_transform, failure=failure.model_dump(), now=_now_iso()), + functools.partial(FutureTransitions.cancel, failure=failure.model_dump(), now=_now_iso()), ) return result.get('status') if result else None @@ -243,14 +244,14 @@ async def cleanup_expired( now = time.time() expired_ids: list[str] = [] for request_id, record in all_records.items(): - if record.status in _TERMINAL_STATUSES: + if record.status in FutureTransitions.terminal: timestamp = self._parse_timestamp(record.updated_at or record.created_at) if timestamp < cutoff_time: expired_ids.append(request_id) continue # Non-terminal records are never deleted -- only ever written ``failed``. - # replica_id None (pre-upgrade) => ownership unknown => treated as alive. + # Unowned background tasks have no replica; only explicit owners can be declared dead. replica_id = record.replica_id replica_alive = (replica_id is None or alive_replica_ids is None or replica_id in alive_replica_ids) if not replica_alive: diff --git a/src/twinkle/server/state/model_manager.py b/src/twinkle/server/state/model_manager.py index de8f872c7..39c327e72 100644 --- a/src/twinkle/server/state/model_manager.py +++ b/src/twinkle/server/state/model_manager.py @@ -93,8 +93,8 @@ async def rebuild_indexes(self) -> None: async def get_capacity_info(self) -> dict[str, int]: """Return global LoRA capacity across all registered replicas.""" - replicas = await self._replicas.get_all() - replicas = {identity: capacity for identity, capacity in replicas.items() if capacity > 0} + capacities = await self._replicas.get_all() + replicas = {identity: capacity for identity, capacity in capacities.items() if capacity > 0} loaded_per_replica = await self._loaded_per_replica() total_max = sum(replicas.values()) total_used = sum(loaded_per_replica.get(rid, 0) for rid in replicas) diff --git a/src/twinkle/server/state/models.py b/src/twinkle/server/state/models.py index 875699dea..d84760d38 100644 --- a/src/twinkle/server/state/models.py +++ b/src/twinkle/server/state/models.py @@ -3,7 +3,7 @@ import time from datetime import datetime, timezone -from pydantic import BaseModel, Field, model_validator +from pydantic import BaseModel, ConfigDict, Field, model_validator from typing import Any, Literal @@ -14,6 +14,16 @@ def _now_iso() -> str: return datetime.now(timezone.utc).isoformat() +class RuntimeContext(BaseModel): + """Allocation identity passed to a worker; never a client request payload.""" + + model_config = ConfigDict(frozen=True, extra='forbid') + + namespace: str + registry_name: str + instance_id: str + + class SessionRecord(BaseModel): """Represents a client session.""" @@ -60,32 +70,6 @@ class FutureFailureRecord(BaseModel): diagnostic: str | None = None -# Canonical set of protocol-independent failure reason codes. The Twinkle and Tinker -# gateways each map this same key set to their own wire vocabularies; keeping the set -# in one place lets a consistency test catch a map that drifts out of coverage. -FAILURE_REASON_CODES: frozenset[str] = frozenset({ - 'capacity_unavailable', - 'startup_failed', - 'startup_timeout', - 'invalid_request', - 'request_rejected', - 'resource_not_found', - 'full_mode_busy', - 'input_tokens_exceeded', - 'batch_size_invalid', - 'rate_limit_exceeded', - 'resource_quota_exceeded', - 'cancelled', - 'endpoint_unavailable', - 'state_contention', - 'backend_gate_unavailable', - 'orphaned_replica', - 'execution_timeout', - 'deadline_exceeded', - 'internal_error', -}) - - class FutureRecord(BaseModel): """Represents an async task future / request status.""" diff --git a/src/twinkle/server/state/sequence.py b/src/twinkle/server/state/sequence.py deleted file mode 100644 index 78efab55b..000000000 --- a/src/twinkle/server/state/sequence.py +++ /dev/null @@ -1,27 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Bounded training retry receipts with a persistent expired-sequence fence.""" - - -def claim_sequence(current, *, sequence, request_id, now, retention): - floor = current.get('floor') if current else None - receipts = {} - for key, receipt in (current or {}).get('receipts', {}).items(): - if receipt['expires_at'] <= now: - floor = int(key) if floor is None else max(floor, int(key)) - else: - receipts[key] = receipt - key = str(sequence) - # Unexpired receipts remain retryable even if another sequence expired first. - if key not in receipts and (floor is None or sequence > floor): - receipts[key] = {'request_id': request_id, 'expires_at': now + retention} - return {'floor': floor, 'receipts': receipts} - - -def release_sequence(current, *, sequence, request_id): - if current is None: - return None - receipts = dict(current['receipts']) - key = str(sequence) - if receipts.get(key, {}).get('request_id') == request_id: - receipts.pop(key) - return {'floor': current['floor'], 'receipts': receipts} diff --git a/src/twinkle/server/state/server_state.py b/src/twinkle/server/state/server_state.py index d38bd8398..01cf54607 100644 --- a/src/twinkle/server/state/server_state.py +++ b/src/twinkle/server/state/server_state.py @@ -10,7 +10,7 @@ from typing import Any from twinkle.server.config.persistence import PersistenceConfig -from twinkle.server.exceptions import RequestRejectedError, ResourceNotFoundError, ResourceQuotaExceededError +from twinkle.server.exceptions import ResourceNotFoundError, ResourceQuotaExceededError from twinkle.server.telemetry.correlation import (BASE_MODEL, MODEL_ID, REPLICA_ID, SAMPLING_SESSION_ID, SESSION_ID, TOKEN_ID) from twinkle.server.telemetry.tracing import traced_operation @@ -24,6 +24,7 @@ from .models import FutureFailureRecord, ModelRecord, SamplingSessionRecord, SessionRecord from .sampling_manager import SamplingSessionManager from .session_manager import SessionManager +from .training_store import TrainingSequences, TrainingStore logger = get_logger() @@ -131,6 +132,7 @@ def __init__( self._model_mgr = ModelManager(self._backend, expiration_timeout, per_token_model_limit) self._sampling_mgr = SamplingSessionManager(self._backend, expiration_timeout) self._future_mgr = FutureManager(self._backend, expiration_timeout) + self.training_sequences = TrainingSequences(self._backend) self.expiration_timeout = expiration_timeout self.cleanup_interval = cleanup_interval @@ -173,29 +175,32 @@ async def _on_lose_leader(self) -> None: async def get_capacity_info(self) -> dict[str, int]: return await self._model_mgr.get_capacity_info() - def create_training_store(self, namespace: str): - from .training_store import TrainingStore - return TrainingStore(self, namespace, self._backend) - - async def save_model_record(self, model_id: str, record: ModelRecord) -> None: - await self._model_mgr.save_record(model_id, record) - - async def get_training_records(self, namespace: str): - return await self._model_mgr.get_training_records(namespace) + @property + def models(self) -> ModelManager: + return self._model_mgr - async def delete_training_sequences(self, training_id: str) -> None: - for key in await self._backend.keys(f'seq_ledger::{training_id}::*'): - await self._backend.delete(key) + @property + def sessions(self) -> SessionManager: + return self._session_mgr - async def unregister_replica_capacity(self, replica_id: str) -> None: - await self._model_mgr.unregister_replica_capacity(replica_id) + @property + def futures(self) -> FutureManager: + return self._future_mgr - async def register_model_record(self, model_id: str, record: ModelRecord) -> None: - """Reserve quota for a complete model record, including managed training state.""" - await self._model_mgr.add(model_id, record) + @property + def sampling_sessions(self) -> SamplingSessionManager: + return self._sampling_mgr - async def rebuild_model_indexes(self) -> None: - await self._model_mgr.rebuild_indexes() + def create_training_store(self, namespace: str) -> TrainingStore: + """Compose a namespaced repository; construction performs no I/O.""" + return TrainingStore( + namespace, + self._backend, + models=self.models, + sessions=self.sessions, + futures=self.futures, + sequences=self.training_sequences, + expiration_timeout=self.expiration_timeout) # ----- Session Management ----- @@ -481,20 +486,10 @@ async def release_seq(self, dedup_key: str) -> None: async def claim_training_seq(self, training_id: str, stream: str, sequence: int, request_id: str, retention: float) -> str | None: - from .sequence import claim_sequence - ledger = await self._backend.update_atomic( - f'seq_ledger::{training_id}::{stream}', - functools.partial( - claim_sequence, sequence=sequence, request_id=request_id, now=time.time(), retention=retention)) - receipt = ledger['receipts'].get(str(sequence)) - if receipt is None: - raise RequestRejectedError('Sequence retry window expired; the operation will not be executed again') - return None if receipt['request_id'] == request_id else receipt['request_id'] + return await self.training_sequences.claim(training_id, stream, sequence, request_id, retention) async def release_training_seq(self, training_id: str, stream: str, sequence: int, request_id: str) -> None: - from .sequence import release_sequence - await self._backend.update_atomic(f'seq_ledger::{training_id}::{stream}', - functools.partial(release_sequence, sequence=sequence, request_id=request_id)) + await self.training_sequences.release(training_id, stream, sequence, request_id) async def try_start_future(self, request_id: str) -> bool: """Acquire task execution atomically, without inferring ownership from reads.""" @@ -669,7 +664,7 @@ def get_server_state(cache_key: str = 'twinkle_server_state', # launcher-propagated env (so a non-gateway worker honours the operator's # YAML instead of the hardcoded default), else the module default. from twinkle.server.config.application_spec import ServerStateArgs - env_policy = ServerStateArgs.from_env() + env_policy = ServerStateArgs.from_env() or ServerStateArgs() def _resolve(explicit, env_value, default): if explicit is not None: @@ -678,13 +673,11 @@ def _resolve(explicit, env_value, default): return env_value return default - expiration_timeout = _resolve(expiration_timeout, getattr(env_policy, 'expiration_timeout', None), - _DEFAULT_EXPIRATION_TIMEOUT) - cleanup_interval = _resolve(cleanup_interval, getattr(env_policy, 'cleanup_interval', None), - _DEFAULT_CLEANUP_INTERVAL) - per_token_model_limit = _resolve(per_token_model_limit, getattr(env_policy, 'per_token_model_limit', None), + expiration_timeout = _resolve(expiration_timeout, env_policy.expiration_timeout, _DEFAULT_EXPIRATION_TIMEOUT) + cleanup_interval = _resolve(cleanup_interval, env_policy.cleanup_interval, _DEFAULT_CLEANUP_INTERVAL) + per_token_model_limit = _resolve(per_token_model_limit, env_policy.per_token_model_limit, _DEFAULT_PER_TOKEN_MODEL_LIMIT) - metrics_update_interval = _resolve(metrics_update_interval, getattr(env_policy, 'metrics_update_interval', None), + metrics_update_interval = _resolve(metrics_update_interval, env_policy.metrics_update_interval, _DEFAULT_METRICS_UPDATE_INTERVAL) state = ServerState( diff --git a/src/twinkle/server/state/training_records.py b/src/twinkle/server/state/training_records.py index e89eb3d3d..87e7aa693 100644 --- a/src/twinkle/server/state/training_records.py +++ b/src/twinkle/server/state/training_records.py @@ -4,8 +4,9 @@ from pydantic import BaseModel, Field from typing import Any, Literal +from typing_extensions import TypedDict -from .models import ModelRecord +from .models import ModelRecord, RuntimeContext TrainingState = Literal['PENDING', 'ACTIVE', 'CLOSING', 'RELEASED', 'EXPIRED', 'FAILED', 'CANCELLED'] InstanceState = Literal['ALLOCATING', 'STARTING', 'READY', 'RESETTING', 'DRAINING', 'STOPPED', 'QUARANTINED'] @@ -13,6 +14,22 @@ LIVE_INSTANCES = frozenset({'ALLOCATING', 'STARTING', 'READY', 'RESETTING'}) +class TrainingOperation(TypedDict): + """Admission facts retained until actual backend completion is confirmed.""" + + accepted_at: float + renew: bool + + +class RuntimeOwner(TypedDict): + instance_id: str + epoch: str + + +class TrainingRoute(RuntimeOwner): + base_model: str + + class TrainingRecord(ModelRecord): training_id: str controller_epoch: str @@ -23,7 +40,7 @@ class TrainingRecord(ModelRecord): create_request_id: str creation_deadline: float last_activity_at: float - pending_operation_ids: dict[str, dict[str, Any]] = Field(default_factory=dict) + pending_operation_ids: dict[str, TrainingOperation] = Field(default_factory=dict) terminal_reason: str | None = None terminal_at: float | None = None @@ -56,3 +73,8 @@ class CreateTraining(BaseModel): client_type: Literal['twinkle', 'tinker'] payload: dict[str, Any] adapter_label: str | None = None + + +class TrainingRuntimeContext(RuntimeContext): + controller_name: str + epoch: str diff --git a/src/twinkle/server/state/training_store.py b/src/twinkle/server/state/training_store.py index 4e6a3ca9f..0920095d3 100644 --- a/src/twinkle/server/state/training_store.py +++ b/src/twinkle/server/state/training_store.py @@ -1,25 +1,29 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Controller-owned persistence; other key-spaces are accessed through ServerState.""" +"""Training persistence and retry receipts, composed from shared state managers.""" from __future__ import annotations +import functools import hashlib import json import time -from typing import TYPE_CHECKING, Any +from typing import Any +from twinkle.server.exceptions import RequestRejectedError from .backend.base import StateBackend +from .future_manager import FutureManager +from .model_manager import ModelManager from .models import FutureFailureRecord +from .session_manager import SessionManager from .training_records import CreateTraining, InstanceRecord, TrainingRecord -if TYPE_CHECKING: - from .server_state import ServerState - class TrainingStore: - def __init__(self, state: ServerState, namespace: str, backend: StateBackend): - self.state = state + def __init__(self, namespace: str, backend: StateBackend, *, models: ModelManager, sessions: SessionManager, + futures: FutureManager, sequences: TrainingSequences, expiration_timeout: float): self.backend = backend + self.models, self.sessions, self.futures, self.sequences = models, sessions, futures, sequences + self.expiration_timeout = expiration_timeout self.prefix = f'training::{namespace}::' def key(self, kind: str, identity: str) -> str: @@ -32,22 +36,20 @@ def create_identity(request: CreateTraining) -> tuple[str, str]: return hashlib.sha256(identity.encode()).hexdigest(), hashlib.sha256(fingerprint.encode()).hexdigest() async def reserve(self, record: TrainingRecord) -> None: - await self.state.register_model_record(record.training_id, record) + await self.models.add(record.training_id, record) async def save_training(self, record: TrainingRecord) -> None: - await self.state.save_model_record(record.training_id, record) + await self.models.save_record(record.training_id, record) async def retire(self, record: TrainingRecord) -> None: await self.backend.set( - self.key('terminal', record.training_id), - record.model_dump(), - ttl=max(1, int(self.state.expiration_timeout))) - await self.state.delete_training_sequences(record.training_id) + self.key('terminal', record.training_id), record.model_dump(), ttl=max(1, int(self.expiration_timeout))) + await self.sequences.delete(record.training_id) # Release the quota row last so a failed earlier step can be idempotently retried. - await self.state.unload_model(record.training_id) + await self.models.remove(record.training_id) async def save_instance(self, record: InstanceRecord) -> None: - ttl = max(1, int(self.state.expiration_timeout)) if record.state == 'STOPPED' else None + ttl = max(1, int(self.expiration_timeout)) if record.state == 'STOPPED' else None await self.backend.set(self.key('instance', record.instance_id), record.model_dump(), ttl=ttl) async def instances(self) -> list[InstanceRecord]: @@ -55,7 +57,7 @@ async def instances(self) -> list[InstanceRecord]: return [InstanceRecord.model_validate(value) for value in await self.backend.mget(keys) if value is not None] async def trainings(self) -> list[TrainingRecord]: - return await self.state.get_training_records(self.prefix) + return await self.models.get_training_records(self.prefix) async def get_create(self, identity: str) -> dict[str, str] | None: return await self.backend.get(self.key('create', identity)) @@ -71,7 +73,7 @@ async def save_create(self, identity: str, fingerprint: str, request_id: str, se async def cleanup_receipts(self) -> None: for key in await self.backend.keys(self.key('create', '*')): receipt = await self.backend.get(key) - if receipt and await self.state.get_session_last_heartbeat(receipt['session_id']) is None: + if receipt and await self.sessions.get_last_heartbeat(receipt['session_id']) is None: await self.backend.delete(key) async def future(self, @@ -84,7 +86,7 @@ async def future(self, failure: FutureFailureRecord | None = None, deadline: float | None = None, owner_token: str | None = None) -> None: - await self.state.store_future_status( + await self.futures.store_status( request_id, status, model_id, @@ -96,3 +98,59 @@ async def future(self, async def begin_epoch(self, epoch: str) -> None: await self.backend.set(self.key('controller', 'active'), {'epoch': epoch, 'started_at': time.time()}) + + +class TrainingSequences: + """Atomic retry receipts and an expired-sequence fence, scoped to a training handle. + + Expiration compacts receipts but retains the sequence floor until that handle is + retired. An old update therefore cannot be applied twice after its receipt ages out. + """ + + def __init__(self, backend: StateBackend): + self.backend = backend + + async def claim(self, training_id: str, stream: str, sequence: int, request_id: str, + retention: float) -> str | None: + ledger = await self.backend.update_atomic( + f'seq_ledger::{training_id}::{stream}', + functools.partial( + self.claim_transform, sequence=sequence, request_id=request_id, now=time.time(), retention=retention)) + receipt = ledger['receipts'].get(str(sequence)) + if receipt is None: + raise RequestRejectedError('Sequence retry window expired; the operation will not be executed again') + return None if receipt['request_id'] == request_id else receipt['request_id'] + + async def release(self, training_id: str, stream: str, sequence: int, request_id: str) -> None: + await self.backend.update_atomic( + f'seq_ledger::{training_id}::{stream}', + functools.partial(self.release_transform, sequence=sequence, request_id=request_id)) + + async def delete(self, training_id: str) -> None: + for key in await self.backend.keys(f'seq_ledger::{training_id}::*'): + await self.backend.delete(key) + + @staticmethod + def claim_transform(current, *, sequence, request_id, now, retention): + floor = current.get('floor') if current else None + receipts = {} + for key, receipt in (current or {}).get('receipts', {}).items(): + if receipt['expires_at'] <= now: + floor = int(key) if floor is None else max(floor, int(key)) + else: + receipts[key] = receipt + key = str(sequence) + # Unexpired receipts remain retryable even if another sequence expired first. + if key not in receipts and (floor is None or sequence > floor): + receipts[key] = {'request_id': request_id, 'expires_at': now + retention} + return {'floor': floor, 'receipts': receipts} + + @staticmethod + def release_transform(current, *, sequence, request_id): + if current is None: + return None + receipts = dict(current['receipts']) + key = str(sequence) + if receipts.get(key, {}).get('request_id') == request_id: + receipts.pop(key) + return {'floor': current['floor'], 'receipts': receipts} diff --git a/src/twinkle/server/task_errors.py b/src/twinkle/server/task_errors.py index cd94d7b70..3c98140eb 100644 --- a/src/twinkle/server/task_errors.py +++ b/src/twinkle/server/task_errors.py @@ -2,10 +2,56 @@ """Helpers for constructing protocol-layer ``ErrorPayload`` values.""" from __future__ import annotations +from dataclasses import dataclass +from functools import lru_cache from typing import Any from twinkle.protocol.types.errors import ErrorCategory, ErrorPayload + +@dataclass(frozen=True) +class FailureResponse: + """How a domain failure is presented inside a response, not its HTTP status.""" + + error_code: int + category: ErrorCategory + + @classmethod + def for_reason(cls, reason: str) -> FailureResponse: + policies = cls._policies() + return policies.get(reason, policies['internal_error']) + + @classmethod + def reasons(cls) -> frozenset[str]: + return frozenset(cls._policies()) + + @classmethod + @lru_cache(maxsize=1) + def _policies(cls) -> dict[str, FailureResponse]: + return { + 'capacity_unavailable': cls(error_code=409, category=ErrorCategory.User), + 'startup_failed': cls(error_code=500, category=ErrorCategory.Server), + 'startup_timeout': cls(error_code=504, category=ErrorCategory.Server), + 'invalid_request': cls(error_code=400, category=ErrorCategory.User), + 'request_rejected': cls(error_code=400, category=ErrorCategory.User), + 'resource_not_found': cls(error_code=404, category=ErrorCategory.User), + 'full_mode_busy': cls(error_code=409, category=ErrorCategory.User), + 'input_tokens_exceeded': cls(error_code=422, category=ErrorCategory.User), + 'batch_size_invalid': cls(error_code=422, category=ErrorCategory.User), + 'rate_limit_exceeded': cls(error_code=429, category=ErrorCategory.User), + 'resource_quota_exceeded': cls(error_code=429, category=ErrorCategory.User), + 'cancelled': cls(error_code=499, category=ErrorCategory.User), + 'endpoint_unavailable': cls(error_code=501, category=ErrorCategory.Server), + 'state_contention': cls(error_code=503, category=ErrorCategory.Server), + 'backend_gate_unavailable': cls(error_code=503, category=ErrorCategory.Server), + 'orphaned_replica': cls(error_code=503, category=ErrorCategory.Server), + 'backend_cancelled': cls(error_code=503, category=ErrorCategory.Server), + 'execution_timeout': cls(error_code=504, category=ErrorCategory.Server), + 'deadline_exceeded': cls(error_code=500, category=ErrorCategory.Server), + 'internal_error': cls(error_code=500, category=ErrorCategory.Server), + } + + _ERROR_MAX = 1024 _TRACEBACK_MAX = 65536 _TRUNCATION_MARKER = '...[traceback truncated, tail kept]...\n' diff --git a/src/twinkle/server/task_queue/__init__.py b/src/twinkle/server/task_queue/__init__.py index 6d7dd81be..7f4805db5 100644 --- a/src/twinkle/server/task_queue/__init__.py +++ b/src/twinkle/server/task_queue/__init__.py @@ -8,15 +8,16 @@ - TaskQueueMixin - mixin with schedule_task / schedule_background_task - RateLimiter - sliding-window rate limiter """ +from twinkle.protocol.types.errors import QueueState +from twinkle.protocol.types.lifecycle import TaskStatus from .mixin import TaskQueueMixin from .rate_limiter import RateLimiter -from .types import QueuedTask, QueueState, TaskStatus, UserTaskError +from .types import QueuedTask from .worker import ComputeWorker __all__ = [ 'TaskStatus', 'QueueState', - 'UserTaskError', 'QueuedTask', 'TaskQueueMixin', 'ComputeWorker', diff --git a/src/twinkle/server/task_queue/backend_gate.py b/src/twinkle/server/task_queue/backend_gate.py index 82d17c1e0..2570a1e10 100644 --- a/src/twinkle/server/task_queue/backend_gate.py +++ b/src/twinkle/server/task_queue/backend_gate.py @@ -20,9 +20,13 @@ import functools from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor +from ray.exceptions import GetTimeoutError from typing import Any -from .types import BackendBusyError +from twinkle.server.exceptions import BackendBusyError + +# Backend execution has the same deadline semantics for asyncio and Ray waits. +BACKEND_TIMEOUT_EXCEPTIONS = (TimeoutError, GetTimeoutError) class BackendGate: diff --git a/src/twinkle/server/task_queue/mixin.py b/src/twinkle/server/task_queue/mixin.py index f1c9b8218..eae14de97 100644 --- a/src/twinkle/server/task_queue/mixin.py +++ b/src/twinkle/server/task_queue/mixin.py @@ -13,7 +13,8 @@ from collections.abc import Callable, Coroutine from typing import TYPE_CHECKING, Any -from twinkle.protocol.types.lifecycle import TERMINAL_STATUSES, TaskEnvelope +from twinkle.protocol.types.errors import QueueState +from twinkle.protocol.types.lifecycle import TERMINAL_STATUSES, TaskEnvelope, TaskStatus from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.exceptions import BatchSizeError, ConfigError, InputTokensExceededError, RateLimitExceededError from twinkle.server.lifecycle.envelope import envelope_from_record @@ -23,7 +24,7 @@ from twinkle.utils.logger import get_logger from .backend_gate import BackendGate from .rate_limiter import RateLimiter -from .types import QueuedTask, QueueState, TaskStatus +from .types import ExecutionScope, QueuedTask from .worker import ComputeWorker if TYPE_CHECKING: @@ -352,6 +353,7 @@ async def schedule_concurrent_task( input_tokens: int = 0, task_type: str | None = None, request_id: str | None = None, + execution_scope: ExecutionScope | None = None, ) -> dict[str, Any]: """Schedule a compute task that runs now instead of in serial order. @@ -376,7 +378,8 @@ async def schedule_concurrent_task( token=token, input_tokens=input_tokens, task_type=task_type, - created_at=time.monotonic()), + created_at=time.monotonic(), + execution_scope=execution_scope), self._queue_key(model_id=model_id, token=token), replica_id=getattr(self, 'replica_id', None), absolute_deadline=self._future_deadline()) diff --git a/src/twinkle/server/task_queue/types.py b/src/twinkle/server/task_queue/types.py index b5f0fd61f..abfb23809 100644 --- a/src/twinkle/server/task_queue/types.py +++ b/src/twinkle/server/task_queue/types.py @@ -1,62 +1,14 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -""" -Core type definitions for the task queue system. +"""In-process execution records, never HTTP/Ray wire models. -Provides: -- TaskStatus: Enum for tracking task lifecycle states -- QueueState: Enum for tinker client compatibility (retry behavior hints) -- QueuedTask: Dataclass representing a queued work item +Coroutine factories, asyncio tasks, and callbacks belong here. Public task and +queue state enums live in twinkle.protocol.types and are shared by SDK/server. """ from __future__ import annotations import asyncio from collections.abc import Callable, Coroutine -from contextvars import ContextVar from dataclasses import dataclass, field -from enum import Enum - -# Ray_Get_Timeout is classified the same as asyncio.TimeoutError: 504/Server. -try: - from ray.exceptions import GetTimeoutError as _RayGetTimeout - BACKEND_TIMEOUT_EXCEPTIONS: tuple[type[BaseException], ...] = (asyncio.TimeoutError, _RayGetTimeout) -except Exception: # pragma: no cover - ray always present in server runtime - BACKEND_TIMEOUT_EXCEPTIONS = (asyncio.TimeoutError, ) - - -class TaskStatus(Enum): - """Task lifecycle status.""" - PENDING = 'pending' # Task created, waiting to be processed - QUEUED = 'queued' # Task in queue waiting for execution - RUNNING = 'running' # Task currently executing - COMPLETED = 'completed' # Task completed successfully - FAILED = 'failed' # Task failed with error - CANCELLED = 'cancelled' # Task cancelled by the client before it started running - - -class UserTaskError(ValueError): - """A queued operation rejected because of caller input or usage.""" - - -class BackendBusyError(RuntimeError): - """Raised when the per-replica Admission_Gate is held by a leaked backend call. - - A new backend call arriving while the gate is closed (its holder is a call that - already exceeded ``asyncio.wait_for`` but whose executor thread has not yet - returned) fails fast with this error instead of queueing behind it. The worker - maps it to ``ErrorPayload(category='server', error_code=503)``. - """ - - -class QueueState(Enum): - """Queue state for tinker client compatibility. - - These states are returned to the tinker client to indicate the current - state of the task queue and help the client adjust its retry behavior. - """ - ACTIVE = 'active' # Queue is actively processing tasks - PAUSED_RATE_LIMIT = 'paused_rate_limit' # Queue paused due to rate limiting - PAUSED_CAPACITY = 'paused_capacity' # Queue paused due to capacity limits - UNKNOWN = 'unknown' # Unknown or unspecified state @dataclass @@ -65,11 +17,16 @@ class ExecutionScope: jobs: set[asyncio.Task] = field(default_factory=set) request_ids: set[str] = field(default_factory=set) - async def wait(self): - await asyncio.gather(*tuple(self.jobs), return_exceptions=True) - + def track(self, job: asyncio.Task, request_id: str | None = None) -> None: + self.jobs.add(job) + if request_id is not None: + self.request_ids.add(request_id) -execution_scope = ContextVar('execution_scope', default=None) + async def wait(self): + observed = set() + while pending := self.jobs - observed: + await asyncio.gather(*pending, return_exceptions=True) + observed.update(pending) @dataclass @@ -83,3 +40,4 @@ class QueuedTask: task_type: str | None created_at: float on_finished: Callable[[bool, bool], Coroutine] | None = None + execution_scope: ExecutionScope | None = None diff --git a/src/twinkle/server/task_queue/worker.py b/src/twinkle/server/task_queue/worker.py index bd8efeef1..a19cac88e 100644 --- a/src/twinkle/server/task_queue/worker.py +++ b/src/twinkle/server/task_queue/worker.py @@ -16,19 +16,20 @@ from ray.exceptions import RayTaskError from typing import TYPE_CHECKING, Any, Callable, Deque -from twinkle.protocol.types.errors import ErrorCategory +from twinkle.protocol.types.errors import ErrorCategory, QueueState +from twinkle.protocol.types.lifecycle import TaskStatus from twinkle.server.config.task_queue import TaskQueueConfig -from twinkle.server.exceptions import (BatchSizeError, EndpointUnavailableError, FullModeBusyError, +from twinkle.server.exceptions import (BackendBusyError, BatchSizeError, EndpointUnavailableError, FullModeBusyError, InputTokensExceededError, RateLimitExceededError, RequestRejectedError, ResourceExhaustedError, ResourceNotFoundError, ResourceQuotaExceededError, StateBackendError, TrainModeMismatchError, TwinkleServerError) from twinkle.server.state.backend.base import ConcurrencyError from twinkle.server.state.models import FutureFailureRecord +from twinkle.server.task_queue.backend_gate import BACKEND_TIMEOUT_EXCEPTIONS from twinkle.server.telemetry.correlation import MODEL_ID, TOKEN_ID from twinkle.server.telemetry.tracing import traced_operation from twinkle.utils.logger import get_logger -from .types import (BACKEND_TIMEOUT_EXCEPTIONS, BackendBusyError, QueuedTask, QueueState, TaskStatus, UserTaskError, - execution_scope) +from .types import QueuedTask if TYPE_CHECKING: from twinkle.server.state import ServerState @@ -422,10 +423,9 @@ async def run(): running = asyncio.create_task(run()) self._detached_tasks[running] = task - scope = execution_scope.get() + scope = task.execution_scope if scope is not None: - scope.jobs.add(running) - scope.request_ids.add(task.request_id) + scope.track(running, task.request_id) def done(completed: asyncio.Task) -> None: self._detached_tasks.pop(completed, None) @@ -509,16 +509,6 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: await self._on_backend_timeout() except Exception: logger.error(f'[ComputeWorker] backend-timeout probe failed:\n{traceback.format_exc(limit=3)}') - except UserTaskError as exc: - task_status = 'failed' - exec_time = time.monotonic() - exec_start - await self._store_task_failed( - task, - f'{type(exc).__name__}: {exc}', - QueueState.UNKNOWN.value, - reason_code='request_rejected', - attribution='user', - ) except BackendBusyError as exc: task_status = 'failed' exec_time = time.monotonic() - exec_start diff --git a/src/twinkle/server/telemetry/metrics.py b/src/twinkle/server/telemetry/metrics.py index 1c287ab9e..0564fa473 100644 --- a/src/twinkle/server/telemetry/metrics.py +++ b/src/twinkle/server/telemetry/metrics.py @@ -17,9 +17,11 @@ from __future__ import annotations +from dataclasses import dataclass from pydantic import BaseModel, ConfigDict from typing import Any +from twinkle.server.exceptions import ResourceReleaseError from .provider import get_meter try: @@ -224,7 +226,8 @@ def __init__(self) -> None: description='Placement group removal duration', unit='s') self.training_quarantined_total = meter.create_counter( - 'twinkle.training.quarantined.total', description='Training instances quarantined during cleanup') + 'twinkle.training.quarantined.total', + description='Transitions to quarantine during cleanup; retries in quarantine are excluded') self.training_sweep_total = meter.create_counter( 'twinkle.training.sweep.total', description='Training lifecycle sweep runs by status') self.training_sweep_seconds = meter.create_histogram( @@ -373,3 +376,52 @@ def get_training_metrics() -> TrainingMetrics: def get_state_backend_metrics() -> StateBackendMetrics: """Return persistent state backend metric adapters.""" return MetricsRegistry.get().state_backend_metrics() + + +@dataclass(frozen=True) +class TrainingInstanceSnapshot: + """Known resource ownership, including draining/quarantined instances.""" + + backend: str + mode: str + state: str + resident: str + device_type: str + device_count: int + placement_group_count: int + + +class TrainingInstanceMetrics: + """Aggregate snapshots and clear retired label sets; never own scheduling state.""" + + def __init__(self, metrics: TrainingMetrics): + self.metrics = metrics + self._instance_keys = set() + self._resource_keys = set() + + def report(self, snapshots): + counts, resources = {}, {} + for snapshot in snapshots: + key = snapshot.backend, snapshot.mode, snapshot.state, snapshot.resident + counts[key] = counts.get(key, 0) + 1 + resource_key = snapshot.backend, snapshot.mode, snapshot.device_type + devices, groups = resources.get(resource_key, (0, 0)) + resources[resource_key] = devices + snapshot.device_count, groups + snapshot.placement_group_count + for key in self._instance_keys | counts.keys(): + tags = dict(zip(('backend', 'mode', 'state', 'resident'), key)) + self.metrics.instances_active.set(counts.get(key, 0), tags=tags) + for key in self._resource_keys | resources.keys(): + tags = dict(zip(('backend', 'mode', 'device_type'), key)) + devices, groups = resources.get(key, (0, 0)) + self.metrics.devices_allocated.set(devices, tags=tags) + self.metrics.placement_groups_active.set(groups, tags=tags) + self._instance_keys = set(counts) + self._resource_keys = set(resources) + + @staticmethod + def quarantine_reason(exc: Exception) -> str: + if isinstance(exc, TimeoutError): + return 'timeout' + if isinstance(exc, ResourceReleaseError): + return exc.phase + return 'other' diff --git a/src/twinkle/server/telemetry/training_metrics.py b/src/twinkle/server/telemetry/training_metrics.py deleted file mode 100644 index b6b7fd63a..000000000 --- a/src/twinkle/server/telemetry/training_metrics.py +++ /dev/null @@ -1,40 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Aggregate plain instance snapshots without depending on training state or scheduling.""" - - -class TrainingInstanceMetrics: - - def __init__(self, metrics): - self.metrics = metrics - self._instance_keys = set() - self._resource_keys = set() - - def report(self, snapshots): - counts, resources = {}, {} - for tags, device_type, device_count, group_count in snapshots: - key = tuple(tags[name] for name in ('backend', 'mode', 'state', 'resident')) - counts[key] = counts.get(key, 0) + 1 - resource_key = tags['backend'], tags['mode'], device_type - devices, groups = resources.get(resource_key, (0, 0)) - resources[resource_key] = devices + device_count, groups + group_count - for key in self._instance_keys | counts.keys(): - tags = dict(zip(('backend', 'mode', 'state', 'resident'), key)) - self.metrics.instances_active.set(counts.get(key, 0), tags=tags) - for key in self._resource_keys | resources.keys(): - tags = dict(zip(('backend', 'mode', 'device_type'), key)) - devices, groups = resources.get(key, (0, 0)) - self.metrics.devices_allocated.set(devices, tags=tags) - self.metrics.placement_groups_active.set(groups, tags=tags) - self._instance_keys = set(counts) - self._resource_keys = set(resources) - - -def quarantine_reason(exc: Exception) -> str: - message = str(exc).lower() - if isinstance(exc, TimeoutError): - return 'timeout' - if 'worker death' in message: - return 'worker_death' - if 'placement group' in message: - return 'pg_removal' - return 'other' diff --git a/src/twinkle/server/training/actor.py b/src/twinkle/server/training/actor.py index 07b2619e6..4e554f76b 100644 --- a/src/twinkle/server/training/actor.py +++ b/src/twinkle/server/training/actor.py @@ -61,7 +61,7 @@ async def _heartbeat(self): if active is None: await self._recover_lost_state() return - await self.store.state.touch_replica_last_seen(self.core.epoch) + await self.store.models.touch_replica_last_seen(self.core.epoch) except Exception: logger.exception('Training state heartbeat failed; admission remains fenced by the state backend') await asyncio.sleep(5) @@ -76,7 +76,7 @@ async def _check_instance_health(self, instance): async def probe(): identity = await self.driver.health(instance) if not identity['unhealthy']: - await self.store.state.touch_replica_last_seen(instance.instance_id) + await self.store.models.touch_replica_last_seen(instance.instance_id) return identity healthy = await self.core.lifecycle.health(instance.instance_id, probe) diff --git a/src/twinkle/server/training/catalog.py b/src/twinkle/server/training/catalog.py deleted file mode 100644 index 1be0548e6..000000000 --- a/src/twinkle/server/training/catalog.py +++ /dev/null @@ -1,44 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Pure lookup and request validation, separate from runtime allocation.""" -from __future__ import annotations - -from dataclasses import dataclass -from typing import Literal - -from twinkle.server.config.training import LoraTrainingConfig, TrainingConfig, TrainingModeConfig, TrainingModelConfig -from twinkle.server.exceptions import RequestRejectedError - - -@dataclass(frozen=True) -class TrainingSpec: - model: TrainingModelConfig - mode: Literal['lora', 'full'] - config: TrainingModeConfig - - @property - def key(self) -> tuple[str, str]: - return self.model.name, self.mode - - @property - def capacity(self) -> int: - return self.config.max_loras if isinstance(self.config, LoraTrainingConfig) else 1 - - -class TrainingCatalog: - - def __init__(self, config: TrainingConfig): - self._models = {model.name: model for model in config.models} - - def get(self, name: str, mode: Literal['lora', 'full']) -> TrainingSpec: - model = self._models.get(name) - if model is None: - raise RequestRejectedError(f'Unknown training model: {name}') - config = getattr(model.modes, mode) - if config is None: - raise RequestRejectedError(f'Training model {name} does not provide {mode} mode') - return TrainingSpec(model, mode, config) - - def specs(self) -> tuple[TrainingSpec, ...]: - return tuple( - TrainingSpec(model, mode, config) for model in self._models.values() for mode in ('lora', 'full') - if (config := getattr(model.modes, mode)) is not None) diff --git a/src/twinkle/server/training/controller.py b/src/twinkle/server/training/controller.py index 01718119f..1e89d5f8e 100644 --- a/src/twinkle/server/training/controller.py +++ b/src/twinkle/server/training/controller.py @@ -8,32 +8,28 @@ import traceback import uuid from collections.abc import Callable -from datetime import datetime from typing import Any, Protocol -from twinkle.server.config.training import TrainingConfig +from twinkle.server.config.training import TrainingConfig, TrainingModelConfig from twinkle.server.exceptions import (RequestRejectedError, ResourceExhaustedError, ResourceNotFoundError, - ResourceQuotaExceededError) + ResourceQuotaExceededError, TrainingExpiredError) from twinkle.server.resources.allocation import CapacityUnavailableError from twinkle.server.resources.instances import InstanceLifecycle from twinkle.server.state.models import FutureFailureRecord from twinkle.server.state.training_records import (LIVE_INSTANCES, TERMINAL_TRAINING, CreateTraining, InstanceRecord, - TrainingRecord) + RuntimeOwner, TrainingRecord, TrainingRoute) from twinkle.server.state.training_store import TrainingStore -from twinkle.server.telemetry.metrics import get_training_metrics -from twinkle.server.telemetry.training_metrics import TrainingInstanceMetrics, quarantine_reason -from .catalog import TrainingCatalog, TrainingSpec -from .errors import TrainingExpiredError +from twinkle.server.telemetry.metrics import TrainingInstanceMetrics, TrainingInstanceSnapshot, get_training_metrics logger = logging.getLogger(__name__) class InstanceDriver(Protocol): - async def allocate(self, instance: InstanceRecord, spec: TrainingSpec) -> dict[str, Any]: + async def allocate(self, instance: InstanceRecord, spec: TrainingModelConfig) -> dict[str, Any]: ... - async def deploy(self, instance: InstanceRecord, spec: TrainingSpec) -> None: + async def deploy(self, instance: InstanceRecord, spec: TrainingModelConfig) -> None: ... async def create_training(self, instance: InstanceRecord, training: TrainingRecord, @@ -80,7 +76,6 @@ def __init__(self, *, clock: Callable[[], float] = time.monotonic): self.config = config - self.catalog = TrainingCatalog(config) self.store = store self.driver = driver self.clock = clock @@ -97,19 +92,27 @@ def __init__(self, self._instance_metrics = TrainingInstanceMetrics(self._metrics) def _instance_metric_tags(self, instance: InstanceRecord) -> dict[str, str]: - spec = self.catalog.get(instance.base_model, instance.mode) + spec = self.config.get(instance.base_model, instance.mode) return { - 'backend': spec.model.backend, + 'backend': spec.backend, 'mode': instance.mode, 'state': instance.state, 'resident': str(instance.resident).lower(), } def _report_instance_metrics(self) -> None: - self._instance_metrics.report((self._instance_metric_tags(instance), - self.catalog.get(instance.base_model, instance.mode).model.device_type, - len(instance.devices), int(bool(instance.application_name))) - for instance in self.instances.values() if instance.state != 'STOPPED') + snapshots = [] + for instance in self.instances.values(): + if instance.state == 'STOPPED': + continue + model = self.config.get(instance.base_model, instance.mode) + snapshots.append( + TrainingInstanceSnapshot( + **self._instance_metric_tags(instance), + device_type=model.device_type, + device_count=len(instance.devices), + placement_group_count=int(bool(instance.application_name)))) + self._instance_metrics.report(snapshots) async def initialize(self) -> None: """Discard old state only after owned compute has been confirmed dead.""" @@ -136,13 +139,13 @@ async def initialize(self) -> None: training.terminal_reason = 'orphaned_replica' training.terminal_at = time.time() await self.store.retire(training) - await self.store.state.rebuild_model_indexes() + await self.store.models.rebuild_indexes() if old_controller: - await self.store.state.unregister_replica(old_controller['epoch']) - await self.store.state.register_replica(self.epoch, 0) - await self.store.state.touch_replica_last_seen(self.epoch) + await self.store.models.unregister_replica(old_controller['epoch']) + await self.store.models.register_replica(self.epoch, 0) + await self.store.models.touch_replica_last_seen(self.epoch) try: - for spec in self.catalog.specs(): + for spec in self.config.models: for _ in range(spec.config.min_replicas): async with self._lock: instance = await self._new_instance(spec, resident=True) @@ -154,9 +157,9 @@ async def initialize(self) -> None: @property def residents_ready(self) -> bool: - return all(self.residents_ready_for(spec) for spec in self.catalog.specs()) + return all(self.residents_ready_for(spec) for spec in self.config.models) - def residents_ready_for(self, spec: TrainingSpec) -> bool: + def residents_ready_for(self, spec: TrainingModelConfig) -> bool: return sum( instance.resident and instance.instance_id not in self._failed_instances and ( instance.base_model, instance.mode) == spec.key and instance.state in ('READY', 'RESETTING') @@ -188,14 +191,14 @@ async def _assert_epoch(self) -> None: if not active or active['epoch'] != self.epoch: raise ResourceNotFoundError('Training controller epoch is no longer active') - async def _new_instance(self, spec: TrainingSpec, *, resident: bool = False) -> InstanceRecord: + async def _new_instance(self, spec: TrainingModelConfig, *, resident: bool = False) -> InstanceRecord: matching = [item for item in self.instances.values() if (item.base_model, item.mode) == spec.key] if not self.lifecycle.headroom(matching, spec.config.max_replicas): raise CapacityUnavailableError('Training replica limit reached') instance = InstanceRecord( instance_id='instance_' + uuid.uuid4().hex, controller_epoch=self.epoch, - base_model=spec.model.name, + base_model=spec.name, mode=spec.mode, resident=resident, startup_deadline=self.clock() + self.config.startup_timeout_s) @@ -206,9 +209,9 @@ async def _new_instance(self, spec: TrainingSpec, *, resident: bool = False) -> return instance async def _start_instance(self, instance: InstanceRecord) -> None: - spec = self.catalog.get(instance.base_model, instance.mode) + spec = self.config.get(instance.base_model, instance.mode) started, status = time.monotonic(), 'error' - tags = {'backend': spec.model.backend, 'mode': spec.mode} + tags = {'backend': spec.backend, 'mode': spec.mode} try: async def startup(): @@ -246,7 +249,7 @@ async def startup(): finally: self._metrics.startup_seconds.observe(time.monotonic() - started, tags={**tags, 'status': status}) - def _select(self, spec: TrainingSpec) -> InstanceRecord | None: + def _select(self, spec: TrainingModelConfig) -> InstanceRecord | None: candidates = [ instance for instance in self.instances.values() if instance.instance_id not in self._failed_instances and (instance.base_model, instance.mode) == spec.key @@ -259,7 +262,7 @@ def _select(self, spec: TrainingSpec) -> InstanceRecord | None: async def create(self, request: CreateTraining) -> dict[str, str]: # Catalog/mode errors are rejected before a future or any resource exists. - spec = self.catalog.get(request.base_model, request.mode) + spec = self.config.get(request.base_model, request.mode) identity, fingerprint = self.store.create_identity(request) async with self._lock: await self._assert_epoch() @@ -336,17 +339,17 @@ async def _create_user(self, training: TrainingRecord, request: CreateTraining) try: await asyncio.shield(self.lifecycle.startups[instance.instance_id]) async with self._lock: - future = await self.store.state.get_future(training.create_request_id) - if future and future['status'] == 'cancelled': + future = await self.store.futures.get(training.create_request_id) + if future and future.status == 'cancelled': training.state = 'CLOSING' training.terminal_reason = 'cancelled' if training.state != 'PENDING' or instance.state != 'READY': raise RequestRejectedError('Training creation is no longer pending') training.creation_deadline = self.clock() + self.config.startup_timeout_s await self.store.save_training(training) - if not await self.store.state.try_start_future(training.create_request_id): - future = await self.store.state.get_future(training.create_request_id) - if future and future['status'] == 'cancelled': + if not await self.store.futures.try_start(training.create_request_id): + future = await self.store.futures.get(training.create_request_id) + if future and future.status == 'cancelled': training.terminal_reason = 'cancelled' raise RequestRejectedError('Training creation did not acquire execution') remaining = training.creation_deadline - self.clock() @@ -380,23 +383,14 @@ async def _create_user(self, training: TrainingRecord, request: CreateTraining) await self._cleanup_user(training) async def _expired(self, training: TrainingRecord) -> bool: - heartbeat = await self.store.state.get_session_last_heartbeat(training.session_id) + heartbeat = await self.store.sessions.get_last_heartbeat(training.session_id) return self._expired_with_heartbeat(training, heartbeat) def _expired_with_heartbeat(self, training: TrainingRecord, heartbeat: float | None) -> bool: - spec = self.catalog.get(training.base_model, training.mode) - config, legacy = spec.config, spec.model.legacy_expiration + spec = self.config.get(training.base_model, training.mode) if heartbeat is None: return True - if legacy is not None: - now = time.time() - if legacy.session_timeout_s is not None and now - heartbeat >= legacy.session_timeout_s: - return True - if legacy.max_lifetime_s is not None: - created_at = datetime.fromisoformat(training.created_at).timestamp() - if now - created_at >= legacy.max_lifetime_s: - return True - return config.idle_timeout_s is not None and self.clock() - training.last_activity_at >= config.idle_timeout_s + return self.clock() - training.last_activity_at >= spec.config.idle_timeout_s async def _authorized(self, training_id: str, token: str, session_id: str | None) -> TrainingRecord: await self._assert_epoch() @@ -408,7 +402,7 @@ async def _authorized(self, training_id: str, token: str, session_id: str | None raise TrainingExpiredError('Training handle expired; create a new handle and explicitly load a checkpoint.') if training.state != 'ACTIVE' or training.instance_id in self._failed_instances: raise ResourceNotFoundError('Training handle is no longer active') - if await self.store.state.get_session_last_heartbeat(training.session_id) is None: + if await self.store.sessions.get_last_heartbeat(training.session_id) is None: raise TrainingExpiredError('Training session expired; create a new training handle') return training @@ -418,7 +412,7 @@ async def admit(self, session_id: str | None, request_id: str, *, - renew: bool = True) -> dict[str, str]: + renew: bool = True) -> RuntimeOwner: async with self._lock: training = await self._authorized(training_id, token, session_id) if not training.pending_operation_ids and await self._expired(training): @@ -467,7 +461,7 @@ async def finish_operation(self, if cleanup: self.lifecycle.spawn(self._cleanup_user(training)) - async def route(self, training_id: str, token: str, session_id: str | None = None) -> dict[str, str]: + async def route(self, training_id: str, token: str, session_id: str | None = None) -> TrainingRoute: async with self._lock: training = await self._authorized(training_id, token, session_id) return {'instance_id': training.instance_id, 'base_model': training.base_model, 'epoch': self.epoch} @@ -481,7 +475,7 @@ async def release(self, training_id: str, token: str, session_id: str | None = N cleanup = not training.pending_operation_ids if cleanup: await self._cleanup_user(training) - config = self.catalog.get(training.base_model, training.mode).config.queue_config + config = self.config.get(training.base_model, training.mode).config.queue_config deadline = asyncio.get_running_loop().time() + config.absolute_future_ttl(1) + self.config.shutdown_timeout_s while training.state == 'CLOSING' or self.instances[training.instance_id].state in ('DRAINING', 'QUARANTINED'): if self.instances[training.instance_id].state == 'QUARANTINED': @@ -519,7 +513,7 @@ async def _cleanup_user(self, training: TrainingRecord) -> None: await self._close_instance(instance) return try: - spec = self.catalog.get(training.base_model, training.mode) + spec = self.config.get(training.base_model, training.mode) await asyncio.wait_for( self.driver.release_training(instance, training), spec.config.queue_config.absolute_future_ttl(spec.config.num_devices)) @@ -584,17 +578,20 @@ async def _close_instance(self, instance: InstanceRecord) -> None: async def quarantined(exc): async with self._lock: + entered_quarantine = instance.state != 'QUARANTINED' quarantined = instance.model_copy(update={'state': 'QUARANTINED', 'failure': str(exc)}) await self.store.save_instance(quarantined) instance.state = 'QUARANTINED' instance.failure = str(exc) self._report_instance_metrics() tags = self._instance_metric_tags(instance) - self._metrics.quarantined_total.inc(tags={ - 'backend': tags['backend'], - 'mode': tags['mode'], - 'reason': quarantine_reason(exc), - }) + if entered_quarantine: + self._metrics.quarantined_total.inc( + tags={ + 'backend': tags['backend'], + 'mode': tags['mode'], + 'reason': TrainingInstanceMetrics.quarantine_reason(exc), + }) self._metrics.shutdown_seconds.observe(time.monotonic() - started, tags={**timing_tags, 'status': 'error'}) async def stopped(): @@ -687,7 +684,7 @@ async def sweep(self) -> None: training.session_id for training in self.trainings.values() if training.state == 'ACTIVE' and not training.pending_operation_ids }) - heartbeats = await self.store.state.get_session_last_heartbeats(session_ids) + heartbeats = await self.store.sessions.get_last_heartbeats(session_ids) async with self._lock: await self._assert_epoch() for training in self.trainings.values(): @@ -719,7 +716,7 @@ async def sweep(self) -> None: elif instance.state in ('DRAINING', 'QUARANTINED'): self.lifecycle.spawn(self._close_instance(instance)) # Rebuilding one specification must not stop expiry or other models. - for spec in self.catalog.specs(): + for spec in self.config.models: async with self._lock: matching = [item for item in self.instances.values() if (item.base_model, item.mode) == spec.key] missing = self.lifecycle.missing( @@ -731,7 +728,7 @@ async def sweep(self) -> None: # Terminal diagnostics expire like their backend tombstones. Repeated # model creation must not grow controller memory for the server lifetime. async with self._lock: - cutoff = time.time() - self.store.state.expiration_timeout + cutoff = time.time() - self.store.expiration_timeout for identity, training in list(self.trainings.items()): if training.state in TERMINAL_TRAINING and training.terminal_at is not None: if training.terminal_at <= cutoff: diff --git a/src/twinkle/server/training/driver.py b/src/twinkle/server/training/driver.py index 698b210b0..4106d68d9 100644 --- a/src/twinkle/server/training/driver.py +++ b/src/twinkle/server/training/driver.py @@ -7,10 +7,9 @@ import time from twinkle.server.resources.runtime import RayRuntime -from twinkle.server.state.training_records import InstanceRecord +from twinkle.server.state.training_records import InstanceRecord, TrainingRuntimeContext from twinkle.server.telemetry.metrics import get_training_metrics from twinkle.utils.logger import get_logger -from .catalog import TrainingCatalog logger = get_logger() @@ -21,7 +20,6 @@ def __init__(self, config, store, namespace, controller_name): self.config, self.store = config, store self.namespace, self.controller_name = namespace, controller_name self._metrics = get_training_metrics() - self._catalog = TrainingCatalog(config) self.runtime = RayRuntime( namespace, allocation_timeout=config.allocation_timeout_s, shutdown_timeout=config.shutdown_timeout_s) @@ -30,9 +28,9 @@ def application_name(self, instance): async def allocate(self, instance, spec): started, status = time.monotonic(), 'error' - tags = {'backend': spec.model.backend, 'mode': spec.mode, 'device_type': spec.model.device_type} + tags = {'backend': spec.backend, 'mode': spec.mode, 'device_type': spec.device_type} try: - allocation = await self.runtime.pool.allocate(spec.config.num_devices, spec.model.device_type, + allocation = await self.runtime.pool.allocate(spec.config.num_devices, spec.device_type, instance.controller_epoch, instance.instance_id) status = 'success' return { @@ -47,38 +45,37 @@ async def allocate(self, instance, spec): async def deploy(self, instance, spec): from ray.util.placement_group import get_placement_group - from twinkle.server.model.app import build_model_app + from twinkle.server.training.execution.app import build_model_app name = self.application_name(instance) - context = { - 'instance_id': instance.instance_id, - 'epoch': instance.controller_epoch, - 'namespace': self.namespace, - 'controller_name': self.controller_name, - 'registry_name': name + '-registry' - } + context = TrainingRuntimeContext( + instance_id=instance.instance_id, + epoch=instance.controller_epoch, + namespace=self.namespace, + controller_name=self.controller_name, + registry_name=name + '-registry') kwargs = spec.config.model_dump( exclude={'num_devices', 'device_mesh', 'min_replicas', 'max_replicas', 'idle_timeout_s', 'queue_config'}, exclude_none=True) if 'hybrid' in kwargs: kwargs['hybrid'] = spec.config.hybrid kwargs['train_mode'] = 'hybrid' if kwargs.get('hybrid') else spec.mode - options = self.runtime.deployment_options(spec.model.deploy_options, name, spec.model.env_vars) + options = self.runtime.deployment_options(spec.deploy_options, name, spec.env_vars) app = build_model_app( - model_id=spec.model.model_id, - backend=spec.model.backend, + model_id=spec.model_id, + backend=spec.backend, nproc_per_node=spec.config.num_devices, device_group={ 'name': 'model', 'ranks': spec.config.num_devices, - 'device_type': spec.model.device_type + 'device_type': spec.device_type }, device_mesh={ - **spec.config.device_mesh.model_dump(exclude_none=True), 'device_type': spec.model.device_type + **spec.config.device_mesh.model_dump(exclude_none=True), 'device_type': spec.device_type }, deploy_options=options, queue_config=spec.config.queue_config, - data_plane_url=spec.model.data_plane_url, + data_plane_url=spec.data_plane_url, training_context=context, placement_group=await asyncio.to_thread(get_placement_group, name), allocation_devices=list(instance.devices), @@ -104,17 +101,17 @@ async def verify_operation(self, instance, request_id): async def shutdown(self, instance: InstanceRecord): name = instance.application_name or self.application_name(instance) try: - spec = self._catalog.get(instance.base_model, instance.mode) + spec = self.config.get(instance.base_model, instance.mode) except Exception: tags = {'backend': 'unknown', 'mode': instance.mode, 'device_type': 'unknown'} else: - tags = {'backend': spec.model.backend, 'mode': spec.mode, 'device_type': spec.model.device_type} + tags = {'backend': spec.backend, 'mode': spec.mode, 'device_type': spec.device_type} def observe_removal(duration, status): self._metrics.removal_seconds.observe(duration, tags={**tags, 'status': status}) await self.runtime.shutdown(instance.instance_id, name, observe_removal=observe_removal) - await self.store.state.unregister_replica_capacity(instance.instance_id) + await self.store.models.unregister_replica_capacity(instance.instance_id) async def health(self, instance): async with asyncio.timeout(5): diff --git a/src/twinkle/server/training/errors.py b/src/twinkle/server/training/errors.py deleted file mode 100644 index d0d508bc6..000000000 --- a/src/twinkle/server/training/errors.py +++ /dev/null @@ -1,10 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -from twinkle.server.exceptions import RequestRejectedError, ResourceNotFoundError - - -class TrainingExpiredError(ResourceNotFoundError): - pass - - -class ProtocolMismatchError(RequestRejectedError): - pass diff --git a/src/twinkle/server/model/__init__.py b/src/twinkle/server/training/execution/__init__.py similarity index 100% rename from src/twinkle/server/model/__init__.py rename to src/twinkle/server/training/execution/__init__.py diff --git a/src/twinkle/server/model/app.py b/src/twinkle/server/training/execution/app.py similarity index 91% rename from src/twinkle/server/model/app.py rename to src/twinkle/server/training/execution/app.py index dfdb7e5b7..5e4a389da 100644 --- a/src/twinkle/server/model/app.py +++ b/src/twinkle/server/training/execution/app.py @@ -16,12 +16,12 @@ from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.deployment import bind_deployment, build_deployment_app from twinkle.server.runtime import init_twinkle_runtime -from twinkle.server.session_resource import AdapterManagerMixin from twinkle.server.state import ServerState, get_server_state +from twinkle.server.state.training_records import TrainingRuntimeContext from twinkle.server.task_queue import TaskQueueMixin from twinkle.server.utils import wrap_builder_with_device_group_env from twinkle.utils.logger import get_logger -from .managed_runtime import ManagedTrainingRuntime +from .runtime import TrainingRuntime from .tinker_handlers import _register_model_tinker_routes from .twinkle_handlers import _register_model_twinkle_routes @@ -79,14 +79,14 @@ def _make_megatron_model(kw: dict[str, Any]) -> Any: ) -class ModelManagement(ManagedTrainingRuntime, TaskQueueMixin, AdapterManagerMixin): +class ModelManagement(TrainingRuntime, TaskQueueMixin): """Unified model management service. Handles: - Base model and multiple LoRA adapters (multi-user) - Tinker training operations via /tinker/* endpoints (async/polling) - Twinkle training operations via /twinkle/* endpoints (synchronous) - - AdapterManagerMixin provides local registration only; expiry is handled by the training controller. + - Worker-local gradient state is keyed by the controller-issued training handle. - Per-user rate limiting via TaskQueueMixin """ @@ -101,16 +101,16 @@ async def __init__(self, **kwargs): self.backend = backend self.device_group = DeviceGroup(**device_group) - self.training_context = kwargs.pop('training_context') - await self.initialize_managed_runtime(self.training_context, kwargs.pop('placement_group'), - kwargs.pop('allocation_devices')) + self.training_context = TrainingRuntimeContext.model_validate(kwargs.pop('training_context')) + await self.initialize_training_runtime(self.training_context, kwargs.pop('placement_group'), + kwargs.pop('allocation_devices')) self.device_mesh = init_twinkle_runtime( is_mock=(backend == 'mock'), nproc_per_node=nproc_per_node, device_group=self.device_group, device_mesh_dict=device_mesh, ) - self.replica_id = self.training_context['instance_id'] + self.replica_id = self.training_context.instance_id self.max_loras = kwargs.get('max_loras', 5) self.base_model = model_id # ``train_mode`` selects LoRA multi-tenant (default) or exclusive @@ -147,7 +147,7 @@ async def __init__(self, collect_width=len(actors) if actors else 1, ) self.model._ray_get_timeout = self.task_queue_config.effective_execution_timeout - self._init_adapter_manager() + self._training_states = {} await self._register_replica_on_startup() @property @@ -168,7 +168,7 @@ async def _register_replica_on_startup(self) -> None: self._replica_registered = True async def _on_request_start(self, request: Request) -> str: - return await self.managed_request_start(request) + return await self.training_request_start(request) async def shutdown(self) -> None: """Explicit async cleanup — called via FastAPI shutdown event.""" @@ -215,9 +215,6 @@ def resolve_model_adapter_name(self, adapter_name: str | None) -> str | None: return '' return adapter_name - async def _on_adapter_expired(self, adapter_name: str) -> None: - raise RuntimeError('Managed training expiry must be handled by the training controller') - def build_model_app(model_id: str, nproc_per_node: int, @@ -253,8 +250,7 @@ def build_model_app(model_id: str, # Build the FastAPI app + middleware stack + routes via the shared scaffold, # then bind the Ray Serve deployment. The Model passes its ``shutdown()`` - # teardown via ``on_shutdown`` and its sticky-LoRA router via - # ``request_router_config``. + # teardown via ``on_shutdown``; the training controller owns handle routing. def register_routes(app: FastAPI, get_self: Any) -> None: _register_model_tinker_routes(app, get_self) _register_model_twinkle_routes(app, get_self) diff --git a/src/twinkle/server/model/backends/__init__.py b/src/twinkle/server/training/execution/backends/__init__.py similarity index 100% rename from src/twinkle/server/model/backends/__init__.py rename to src/twinkle/server/training/execution/backends/__init__.py diff --git a/src/twinkle/server/model/backends/common.py b/src/twinkle/server/training/execution/backends/common.py similarity index 99% rename from src/twinkle/server/model/backends/common.py rename to src/twinkle/server/training/execution/backends/common.py index 0c62f3090..024a0bfa1 100644 --- a/src/twinkle/server/model/backends/common.py +++ b/src/twinkle/server/training/execution/backends/common.py @@ -11,7 +11,7 @@ from typing import Any, List from twinkle import DeviceMesh, remote_function -from twinkle.server.model.tinker_datum import datum_to_input_feature, extract_rl_features_for_loss +from twinkle.server.training.execution.tinker_datum import datum_to_input_feature, extract_rl_features_for_loss from twinkle.server.validation.tinker_loss import validate_tinker_loss from twinkle.template import Template diff --git a/src/twinkle/server/model/backends/megatron_model.py b/src/twinkle/server/training/execution/backends/megatron_model.py similarity index 95% rename from src/twinkle/server/model/backends/megatron_model.py rename to src/twinkle/server/training/execution/backends/megatron_model.py index b06228248..dcad8cb29 100644 --- a/src/twinkle/server/model/backends/megatron_model.py +++ b/src/twinkle/server/training/execution/backends/megatron_model.py @@ -16,9 +16,9 @@ from twinkle.data_format import InputFeature, Trajectory from twinkle.infra import collect_tensor_dict from twinkle.model.megatron import MegatronModel, MultiLoraMegatronModel -from twinkle.server.model.backends.common import (TwinkleCompatModelBase, clean_metrics, - collect_forward_backward_results, to_cpu_safe_output) -from twinkle.server.model.tinker_datum import datum_to_input_feature, extract_rl_features_for_loss +from twinkle.server.training.execution.backends.common import (TwinkleCompatModelBase, clean_metrics, + collect_forward_backward_results, to_cpu_safe_output) +from twinkle.server.training.execution.tinker_datum import datum_to_input_feature, extract_rl_features_for_loss from twinkle.utils.nccl_safe import nccl_safe_megatron diff --git a/src/twinkle/server/model/backends/mock_model.py b/src/twinkle/server/training/execution/backends/mock_model.py similarity index 99% rename from src/twinkle/server/model/backends/mock_model.py rename to src/twinkle/server/training/execution/backends/mock_model.py index d204090a8..b2685f3ca 100644 --- a/src/twinkle/server/model/backends/mock_model.py +++ b/src/twinkle/server/training/execution/backends/mock_model.py @@ -181,7 +181,7 @@ def save(self, name: str = 'latest', output_dir: str | None = None, **kwargs: An The real backend creates ``//`` on disk and returns its path. Tinker /asample later does ``os.path.exists(...)`` on a symlink - pointing here (see ``sampler/tinker_handlers.py:72``), so the directory + pointing here (see ``sampling/execution/tinker_handlers.py``), so the directory has to really exist for the mock path to clear. """ target = os.path.join(output_dir or '/tmp/twinkle_mock_ckpt', name) diff --git a/src/twinkle/server/model/backends/transformers_model.py b/src/twinkle/server/training/execution/backends/transformers_model.py similarity index 95% rename from src/twinkle/server/model/backends/transformers_model.py rename to src/twinkle/server/training/execution/backends/transformers_model.py index db42aa7d0..38639ada1 100644 --- a/src/twinkle/server/model/backends/transformers_model.py +++ b/src/twinkle/server/training/execution/backends/transformers_model.py @@ -19,9 +19,9 @@ from twinkle.infra import collect_tensor_dict from twinkle.model import MultiLoraTransformersModel, SpectralHybridTransformersModel from twinkle.model.transformers import TransformersModel -from twinkle.server.model.backends.common import (TwinkleCompatModelBase, clean_metrics, - collect_forward_backward_results, to_cpu_safe_output) -from twinkle.server.model.tinker_datum import datum_to_input_feature, extract_rl_features_for_loss +from twinkle.server.training.execution.backends.common import (TwinkleCompatModelBase, clean_metrics, + collect_forward_backward_results, to_cpu_safe_output) +from twinkle.server.training.execution.tinker_datum import datum_to_input_feature, extract_rl_features_for_loss class _TransformersTinkerCompatMixin(TwinkleCompatModelBase): diff --git a/src/twinkle/server/model/data_plane_inputs.py b/src/twinkle/server/training/execution/data_plane_inputs.py similarity index 100% rename from src/twinkle/server/model/data_plane_inputs.py rename to src/twinkle/server/training/execution/data_plane_inputs.py diff --git a/src/twinkle/server/model/managed_runtime.py b/src/twinkle/server/training/execution/runtime.py similarity index 75% rename from src/twinkle/server/model/managed_runtime.py rename to src/twinkle/server/training/execution/runtime.py index 5a0d3da8c..ff535734c 100644 --- a/src/twinkle/server/model/managed_runtime.py +++ b/src/twinkle/server/training/execution/runtime.py @@ -5,7 +5,7 @@ import asyncio import os import uuid -from contextvars import ContextVar +from dataclasses import dataclass from twinkle.protocol.serialize import deserialize_object from twinkle.protocol.types import model as native @@ -16,33 +16,52 @@ from twinkle.server.lifecycle.submit import backend_kwargs from twinkle.server.middleware.auth import get_session_id_from_request, get_token_from_request from twinkle.server.state.models import FutureFailureRecord -from twinkle.server.state.training_records import CreateTraining, TrainingRecord +from twinkle.server.state.training_records import CreateTraining, TrainingRecord, TrainingRuntimeContext from twinkle.server.utils import get_template_for_model from twinkle.utils.logger import get_logger -_tinker_operation = ContextVar('managed_tinker_operation', default=None) logger = get_logger() -class ManagedTrainingRuntime: +@dataclass +class TrainingExecutionState: + """Worker-local state; ownership and expiration belong to the controller.""" + + grad_ready: bool = False + + +class TrainingRuntime: + """Execute controller-owned training handles on one allocated backend.""" + + training_context: TrainingRuntimeContext + + _training_states: dict[str, TrainingExecutionState] + + def get_training_state(self, training_id: str) -> TrainingExecutionState: + state = self._training_states.get(training_id) + if state is None: + raise ResourceNotFoundError('Training is not registered on this runtime') + return state + + def assert_resource_exists(self, training_id: str) -> None: + self.get_training_state(training_id) async def wait_for_training_workers(self): actors = getattr(self.model, '_actors', None) or [] await asyncio.gather(*(actor.__ray_ready__.remote() for actor in actors)) - async def initialize_managed_runtime(self, context, placement_group, devices): + async def initialize_training_runtime(self, context, placement_group, devices): import ray from twinkle.infra._ray.ray_helper import RayHelper from twinkle.infra._ray.resource_manager import ResourceManager - self._training_controller = ray.get_actor(context['controller_name'], namespace=context['namespace']) - await self._training_controller.register_runtime.remote(context['instance_id'], context['epoch'], - uuid.uuid4().hex) - os.environ['TWINKLE_RAY_REGISTRY_NAME'] = context['registry_name'] - os.environ['TWINKLE_RAY_NAMESPACE'] = context['namespace'] + self._training_controller = ray.get_actor(context.controller_name, namespace=context.namespace) + await self._training_controller.register_runtime.remote(context.instance_id, context.epoch, uuid.uuid4().hex) + os.environ['TWINKLE_RAY_REGISTRY_NAME'] = context.registry_name + os.environ['TWINKLE_RAY_NAMESPACE'] = context.namespace RayHelper.resource_manager = ResourceManager.from_allocation(self.device_group, placement_group, devices) - async def managed_request_start(self, request): + async def training_request_start(self, request): token = get_token_from_request(request) if self._model_unhealthy: raise ResourceNotFoundError('Training runtime is unavailable') @@ -52,31 +71,10 @@ async def managed_request_start(self, request): if '/twinkle/' in request.url.path and not session: raise ResourceNotFoundError('Training requests require the creating session') route = await self._training_controller.route.remote(training_id, token, session) - if route['instance_id'] != self.training_context['instance_id'] or route['epoch'] != self.training_context[ - 'epoch']: + if route['instance_id'] != self.training_context.instance_id or route['epoch'] != self.training_context.epoch: raise ResourceNotFoundError('Runtime does not own this training handle') - _tinker_operation.set((training_id, body.get('seq_id')) if '/tinker/' in request.url.path else None) return token - async def schedule_task(self, coro_factory, **kwargs): - operation = _tinker_operation.get() - if operation is None or operation[1] is None: - return await super().schedule_task(coro_factory, **kwargs) - training_id, seq_id = operation - task_type = kwargs.get('task_type') - stream = f'tinker::{task_type}' - request_id = kwargs.pop('request_id', None) or 'req_' + uuid.uuid4().hex - retention = self.sequence_retention_s - previous = await self.state.claim_training_seq(training_id, stream, seq_id, request_id, retention) - if previous is not None: - return {'request_id': previous, 'model_id': training_id} - try: - return await super().schedule_task(coro_factory, request_id=request_id, **kwargs) - except Exception: - if await self.state.get_future(request_id) is None: - await self.state.release_training_seq(training_id, stream, seq_id, request_id) - raise - async def _prepare_queued_task(self, task): if task.model_id is None: return @@ -92,7 +90,7 @@ async def finished(executed, unsafe): if unsafe: self._model_unhealthy = True self._training_controller.fail_instance.remote( - context['instance_id'], + context.instance_id, FutureFailureRecord( reason_code='execution_timeout', attribution='server', @@ -103,7 +101,7 @@ async def finished(executed, unsafe): await self._training_controller.finish_operation.remote( task.model_id, task.request_id, - context['epoch'], + context.epoch, executed=executed, finished_at=outcome['finished_at']) return True @@ -112,9 +110,9 @@ async def finished(executed, unsafe): task.on_finished = finished if task.task_type == 'apply_patch': - await self._training_controller.mark_dirty.remote(context['instance_id'], context['epoch']) + await self._training_controller.mark_dirty.remote(context.instance_id, context.epoch) - async def _managed_serial(self, factory, task_type): + async def _run_lifecycle_operation(self, factory, task_type): ref = await self.schedule_task(factory, task_type=task_type) deadline = asyncio.get_running_loop().time() + self.task_queue_config.absolute_future_ttl(1) while True: @@ -134,8 +132,7 @@ async def initialize_training(self, training_data, request_data): async def initialize(): adapter = training.training_id backend_adapter = self.resolve_model_adapter_name(adapter) - self.register_resource(adapter, training.token, training.session_id) - self.set_resource_state(adapter, 'grad_ready', False) + self._training_states[adapter] = TrainingExecutionState() if request.client_type == 'twinkle': body = native.AddAdapterRequest.model_validate(request.payload) # Dynamic imports (e.g. PEFT) and deserialization must not block the health-probe event loop. @@ -173,24 +170,23 @@ def save_training_metadata(): await asyncio.to_thread(save_training_metadata) - await self._managed_serial(initialize, 'initialize_training') + await self._run_lifecycle_operation(initialize, 'initialize_training') async def release_training(self, training_data, reset_full): training = TrainingRecord.model_validate(training_data) async def release(): adapter = training.training_id - if self.get_resource_info(adapter) is None: + if adapter not in self._training_states: return if self.is_full_mode: if reset_full: await self.call_backend(self.model.reload_initial_weights, reset_training_state=True) else: await self.call_backend(self.model.remove_adapter, adapter) - self.clear_resource_state(adapter) - self.unregister_resource(adapter) + self._training_states.pop(adapter) - await self._managed_serial(release, 'release_training') + await self._run_lifecycle_operation(release, 'release_training') async def runtime_identity(self): import ray @@ -199,7 +195,7 @@ async def runtime_identity(self): lambda: [ray._private.state.actors(actor._actor_id.hex()).get('State') for actor in actors]) return { 'pid': os.getpid(), - 'instance_id': self.training_context['instance_id'], + 'instance_id': self.training_context.instance_id, 'unhealthy': self._model_unhealthy or any(state != 'ALIVE' for state in states) } diff --git a/src/twinkle/server/model/tinker_datum.py b/src/twinkle/server/training/execution/tinker_datum.py similarity index 100% rename from src/twinkle/server/model/tinker_datum.py rename to src/twinkle/server/training/execution/tinker_datum.py diff --git a/src/twinkle/server/training/execution/tinker_handlers.py b/src/twinkle/server/training/execution/tinker_handlers.py new file mode 100644 index 000000000..b1fd47f69 --- /dev/null +++ b/src/twinkle/server/training/execution/tinker_handlers.py @@ -0,0 +1,301 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Tinker-compatible routes for the Model deployment. + +Registered by ``_register_model_tinker_routes(app, self_fn)`` -- module-level route +registration closing over ``self_fn`` via ``Depends``, not a mixin: there is no +inheritance relationship with the deployment class. All endpoints are prefixed +/tinker/... and use schedule_task() returning UntypedAPIFuture. ``self_fn`` is injected +via FastAPI Depends to obtain the ModelManagement instance at request time. +""" +from __future__ import annotations + +from collections.abc import Callable +from fastapi import Depends, FastAPI, Request +from functools import partial +from tinker import types +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .app import ModelManagement + +from twinkle.server.checkpoint import create_checkpoint_manager, create_training_run_manager +from twinkle.server.exceptions import RequestRejectedError +from twinkle.server.lifecycle.submit import submit_training_task +from twinkle.server.validation.tinker_loss import validate_tinker_loss +from twinkle.utils.logger import get_logger + +logger = get_logger() + + +def _register_model_tinker_routes(app: FastAPI, self_fn: Callable[[], ModelManagement]) -> None: + """Register all /tinker/* routes on the given FastAPI app. + + self_fn is a zero-argument callable that returns the current ModelManagement + replica instance. It is wired in via Depends so it is resolved lazily at request time. + """ + + @app.post('/tinker/create_model') + async def create_model( + request: Request, + body: types.CreateModelRequest, + self: ModelManagement = Depends(self_fn), + ) -> types.UntypedAPIFuture: + raise RequestRejectedError('Training lifecycle operations must use the catalog endpoint') + + @app.post('/tinker/get_info') + async def get_info(request: Request, body: types.GetInfoRequest, + self: ModelManagement = Depends(self_fn)) -> types.GetInfoResponse: + token = await self._on_request_start(request) + training_run_manager = create_training_run_manager(token, client_type='tinker') + metadata = training_run_manager.get(str(body.model_id)) + model_name = metadata.base_model if metadata else self.base_model + lora_rank = None + is_lora = False + if metadata and hasattr(metadata, 'lora_rank') and metadata.lora_rank: + lora_rank = metadata.lora_rank + is_lora = metadata.is_lora + return types.GetInfoResponse( + model_data=types.ModelData(model_name=model_name), + model_id=body.model_id, + is_lora=is_lora, + lora_rank=lora_rank, + model_name=model_name, + ) + + @app.post('/tinker/unload_model') + async def unload_model( + request: Request, + body: types.UnloadModelRequest, + self: ModelManagement = Depends(self_fn), + ) -> types.UntypedAPIFuture: + raise RequestRejectedError('Training lifecycle operations must use the catalog endpoint') + + @app.post('/tinker/forward') + async def forward(request: Request, body: types.ForwardRequest, + self: ModelManagement = Depends(self_fn)) -> types.UntypedAPIFuture: + token = await self._on_request_start(request) + + async def _do_forward(): + adapter_name = body.model_id + self.assert_resource_exists(adapter_name) + model_adapter = self.resolve_model_adapter_name(adapter_name) + datum_list = body.forward_input.data + loss_fn_config = body.forward_input.loss_fn_config or {} + output, loss = await self.call_backend( + self.model.tinker_forward_only, inputs=datum_list, adapter_name=model_adapter, **loss_fn_config) + return types.ForwardBackwardOutput( + loss_fn_output_type='CrossEntropyLossReturn', + loss_fn_outputs=output, + metrics={'loss:mean': loss}, + ) + + datum_list = body.forward_input.data + input_tokens = sum(len(d.model_input.to_ints()) for d in datum_list) + batch_size = len(datum_list) + return await submit_training_task( + self, + _do_forward, + training_id=body.model_id, + seq_id=body.seq_id, + stream='tinker::forward', + token=token, + input_tokens=input_tokens, + batch_size=batch_size, + data_world_size=self.data_world_size, + task_type='forward', + ) + + @app.post('/tinker/forward_backward') + async def forward_backward( + request: Request, + body: types.ForwardBackwardRequest, + self: ModelManagement = Depends(self_fn), + ) -> types.UntypedAPIFuture: + token = await self._on_request_start(request) + + async def _do_forward_backward(): + adapter_name = body.model_id + self.assert_resource_exists(adapter_name) + model_adapter = self.resolve_model_adapter_name(adapter_name) + datum_list = body.forward_backward_input.data + loss_fn = body.forward_backward_input.loss_fn + loss_fn_config = body.forward_backward_input.loss_fn_config or {} + preflight = getattr(self.model, 'preflight_loss_inputs', None) + if preflight is not None and loss_fn == 'importance_sampling' and not validate_tinker_loss( + loss_fn, datum_list, data_world_size=self.data_world_size): + # GRPO datums carry ragged per-sample targets; the loss check inside + # the sliced backend would deadlock the peer ranks (BUG-05), so the + # full batch is validated on every rank first. + await self.call_backend( + preflight, inputs=datum_list, adapter_name=model_adapter, loss_fn=loss_fn, **loss_fn_config) + output, loss = await self.call_backend( + self.model.tinker_forward_backward, + inputs=datum_list, + adapter_name=model_adapter, + loss_fn=loss_fn, + **loss_fn_config) + output_type = ('ImportanceSamplingLossReturn' + if loss_fn == 'importance_sampling' else 'CrossEntropyLossReturn') + self.get_training_state(adapter_name).grad_ready = True + return types.ForwardBackwardOutput( + loss_fn_output_type=output_type, + loss_fn_outputs=output, + metrics={'loss:mean': loss}, + ) + + datum_list = body.forward_backward_input.data + input_tokens = sum(len(d.model_input.to_ints()) for d in datum_list) + batch_size = len(datum_list) + loss_fn = body.forward_backward_input.loss_fn + is_dpo = validate_tinker_loss(loss_fn, datum_list, data_world_size=self.data_world_size) + return await submit_training_task( + self, + _do_forward_backward, + training_id=body.model_id, + seq_id=body.seq_id, + stream='tinker::forward_backward', + token=token, + input_tokens=input_tokens, + batch_size=batch_size, + data_world_size=self.data_world_size, + batch_size_multiple=2 if is_dpo else None, + task_type='forward_backward', + ) + + @app.post('/tinker/optim_step') + async def optim_step( + request: Request, + body: types.OptimStepRequest, + self: ModelManagement = Depends(self_fn), + ) -> types.UntypedAPIFuture: + token = await self._on_request_start(request) + + async def _do_optim(): + adapter_name = body.model_id + self.assert_resource_exists(adapter_name) + model_adapter = self.resolve_model_adapter_name(adapter_name) + if not self.get_training_state(adapter_name).grad_ready: + raise RuntimeError(f'No accumulated gradients for adapter={adapter_name}; ' + 'call forward_backward before optim_step') + await self.call_backend(self.model.tinker_step, adam_params=body.adam_params, adapter_name=model_adapter) + self.get_training_state(adapter_name).grad_ready = False + metrics = await self.call_backend( + self.model.tinker_calculate_metric, is_training=True, adapter_name=model_adapter) + return types.OptimStepResponse(metrics=metrics) + + return await submit_training_task( + self, + _do_optim, + training_id=body.model_id, + seq_id=body.seq_id, + stream='tinker::optim_step', + token=token, + task_type='optim_step') + + @app.post('/tinker/save_weights') + async def save_weights( + request: Request, + body: types.SaveWeightsRequest, + self: ModelManagement = Depends(self_fn), + ) -> types.UntypedAPIFuture: + token = await self._on_request_start(request) + + async def _do_save(): + adapter_name = body.model_id + self.assert_resource_exists(adapter_name) + checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') + save_weights = partial( + self.call_backend, + self.model.save, + adapter_name=self.resolve_model_adapter_name(adapter_name), + save_optimizer=True) + tinker_path, _ = await checkpoint_manager.save_training(body.model_id, body.path, save_weights) + return types.SaveWeightsResponse(path=tinker_path, type='save_weights') + + return await submit_training_task( + self, + _do_save, + training_id=body.model_id, + seq_id=body.seq_id, + stream='tinker::save_weights', + token=token, + task_type='save_weights') + + @app.post('/tinker/save_weights_for_sampler') + async def save_weights_for_sampler( + request: Request, + body: types.SaveWeightsForSamplerRequest, + self: ModelManagement = Depends(self_fn), + ) -> types.UntypedAPIFuture: + token = await self._on_request_start(request) + + async def _do_save_for_sampler(): + adapter_name = body.model_id + self.assert_resource_exists(adapter_name) + checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') + + save_weights = partial( + self.call_backend, + self.model.save, + adapter_name=self.resolve_model_adapter_name(adapter_name), + save_optimizer=False) + + tinker_path, _ = await checkpoint_manager.save_sampler(body.model_id, body.path, save_weights) + payload = body.model_dump() + payload['model_path'] = tinker_path + metadata = await self.state.get_model_metadata(body.model_id) or {} + if metadata.get('base_model'): + payload['base_model'] = metadata['base_model'] + sampling_session_id = await self.state.create_sampling_session(payload) + # Tinker SDK distinguishes two modes by whether sampling_session_seq_id is set: + # 1. save_weights_for_sampler(name): expects path != None + # 2. save_weights_and_get_sampling_client(): expects path == None, sampling_session_id != None + if body.sampling_session_seq_id is not None: + return types.SaveWeightsForSamplerResponseInternal(path=None, sampling_session_id=sampling_session_id) + else: + return types.SaveWeightsForSamplerResponseInternal( + path=tinker_path, sampling_session_id=sampling_session_id) + + return await submit_training_task( + self, + _do_save_for_sampler, + training_id=body.model_id, + seq_id=body.seq_id, + stream='tinker::save_weights_for_sampler', + token=token, + task_type='save_weights_for_sampler') + + @app.post('/tinker/load_weights') + async def load_weights( + request: Request, + body: types.LoadWeightsRequest, + self: ModelManagement = Depends(self_fn), + ) -> types.UntypedAPIFuture: + token = await self._on_request_start(request) + + async def _do_load(): + assert self.model is not None, 'Model not loaded, please load model first' + adapter_name = body.model_id + self.assert_resource_exists(adapter_name) + # Path resolution (token isolation, twinkle-vs-external path shapes) is a + # server storage policy and belongs in the handler, not in the GPU-side + # backend actor. The backend receives an already-resolved location. + checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') + resolved = checkpoint_manager.resolve_load_path(body.path) + await self.call_backend( + self.model.tinker_load, + checkpoint_name=resolved.checkpoint_name, + output_dir=resolved.checkpoint_dir if resolved.is_twinkle_path else None, + load_optimizer=body.optimizer, + adapter_name=self.resolve_model_adapter_name(adapter_name)) + self.get_training_state(adapter_name).grad_ready = False + return types.LoadWeightsResponse(path=body.path, type='load_weights') + + return await submit_training_task( + self, + _do_load, + training_id=body.model_id, + seq_id=body.seq_id, + stream='tinker::load_weights', + token=token, + task_type='load_weights') diff --git a/src/twinkle/server/model/twinkle_handlers.py b/src/twinkle/server/training/execution/twinkle_handlers.py similarity index 99% rename from src/twinkle/server/model/twinkle_handlers.py rename to src/twinkle/server/training/execution/twinkle_handlers.py index b8537e8b7..4bd630689 100644 --- a/src/twinkle/server/model/twinkle_handlers.py +++ b/src/twinkle/server/training/execution/twinkle_handlers.py @@ -25,8 +25,8 @@ from twinkle.server.checkpoint import create_checkpoint_manager from twinkle.server.exceptions import RequestRejectedError from twinkle.server.lifecycle.submit import backend_kwargs, input_metrics, run_submit, to_backend_inputs -from twinkle.server.model.data_plane_inputs import (data_plane_request_shape, merge_forward_kwargs, - resolve_data_plane_model_inputs, select_output_rows) +from twinkle.server.training.execution.data_plane_inputs import (data_plane_request_shape, merge_forward_kwargs, + resolve_data_plane_model_inputs, select_output_rows) from twinkle.server.validation import BackendCapability from twinkle.utils.logger import get_logger diff --git a/src/twinkle/server/training/router.py b/src/twinkle/server/training/router.py index 6e833695e..23f4d8bcb 100644 --- a/src/twinkle/server/training/router.py +++ b/src/twinkle/server/training/router.py @@ -16,17 +16,15 @@ from twinkle.protocol.types import model as native from twinkle.server.checkpoint import _resolve_client_save_dir, create_checkpoint_manager, validate_user_path from twinkle.server.deployment import LazyCleanupMixin, bind_deployment, build_deployment_app -from twinkle.server.exceptions import RequestRejectedError, ResourceNotFoundError +from twinkle.server.exceptions import ProtocolMismatchError, RequestRejectedError, ResourceNotFoundError from twinkle.server.lifecycle.envelope import envelope_from_record +from twinkle.server.lifecycle.proxy import ServiceProxy from twinkle.server.middleware.auth import get_session_id_from_request, get_token_from_request -from twinkle.server.service_proxy import ServiceProxy from twinkle.server.state import get_server_state from twinkle.server.state.training_records import CreateTraining from twinkle.server.task_queue import TaskQueueMixin from twinkle.server.validation import assert_request_supported from .actor import controller_name -from .catalog import TrainingCatalog -from .errors import ProtocolMismatchError NATIVE_REQUESTS = { 'forward': native.ForwardRequest, @@ -75,7 +73,7 @@ class TrainingRouter(TaskQueueMixin, LazyCleanupMixin): async def __init__(self, config, model, namespace, http_options): - self.catalog = TrainingCatalog(config) + self.config = config self.model = model self.backend = model.backend self.state = get_server_state() @@ -126,7 +124,7 @@ async def native_request(self, request, endpoint): except (ValueError, TypeError) as exc: raise RequestRejectedError('Expected a serialized LoRA config') from exc mode = 'full' if config is None else 'lora' - spec = self.catalog.get(self.model.name, mode) + spec = self.config.get(self.model.name, mode) if config is not None: rank = config.get('r') if isinstance(config, dict) else None if (not isinstance(config, dict) or config.get('_TWINKLE_TYPE_') != 'LoraConfig' @@ -167,8 +165,6 @@ async def upload(): ref = await self.schedule_background_task(upload, task_type=endpoint, token=token) return envelope_from_record(ref['request_id'], await self.state.get_future(ref['request_id'])) - if not body.training_id: - raise ProtocolMismatchError('Upgrade Twinkle SDK: training_id is required') session = get_session_id_from_request(request) await self._check_session(session) if endpoint == 'remove_adapter': @@ -193,7 +189,7 @@ async def tinker_request(self, request, endpoint): if body.base_model != self.model.name: raise RequestRejectedError('Model does not match the catalog endpoint') mode = 'lora' if body.lora_config is not None else 'full' - spec = self.catalog.get(self.model.name, mode) + spec = self.config.get(self.model.name, mode) if body.lora_config is not None and not 0 < body.lora_config.rank <= spec.config.max_r: raise RequestRejectedError('LoRA rank exceeds the training specification') await self._check_session(body.session_id) diff --git a/src/twinkle_client/auto/agent/server_tools.py b/src/twinkle_client/auto/agent/server_tools.py index 700bddc6e..c2b031cc6 100644 --- a/src/twinkle_client/auto/agent/server_tools.py +++ b/src/twinkle_client/auto/agent/server_tools.py @@ -440,9 +440,6 @@ def _short(mid: str) -> str: 'max_replicas': 1, 'max_length': 500000, 'idle_timeout_s': 600, - 'adapter_config': { - 'adapter_timeout': 600 - }, 'queue_config': { 'rps_limit': 100, 'tps_limit': 100000, diff --git a/src/twinkle_client/model/multi_lora_transformers.py b/src/twinkle_client/model/multi_lora_transformers.py index 876bb5bd1..05981a861 100644 --- a/src/twinkle_client/model/multi_lora_transformers.py +++ b/src/twinkle_client/model/multi_lora_transformers.py @@ -6,7 +6,7 @@ import uuid from collections.abc import Mapping from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from peft import LoraConfig @@ -34,7 +34,7 @@ class MultiLoraTransformersModel: """Client wrapper for TwinkleModel that calls server HTTP endpoints. This client manages adapters and sends training/inference requests to the model server. - The server-side session (managed by TwinkleClient) keeps the model alive. + The session authenticates requests; executed training work renews the idle deadline. Every method builds its endpoint's request model rather than a dict, so a misspelled or wrongly-typed argument fails here -- in the caller's own stack trace, @@ -48,11 +48,9 @@ def __init__( model_id: str, *, transport: ClientTransport | None = None, - **kwargs, ): """Initialize a model wrapper bound to one immutable request identity.""" self._transport = capture_transport(transport) - kwargs.pop('data_plane_url', None) if '://' in model_id: model_id = model_id.split('://')[1] @@ -63,7 +61,7 @@ def __init__( self._creating: set[str] = set() self._training_lock = threading.Lock() # Per-client monotonic sequence for idempotent dedup of stateful training ops: - # the server dedups on (session_id, seq_id) so a retried grad/step call is + # the server dedups on (training_id, session, operation, seq_id), so each update is # applied at most once. Reserved once per call and reused on retry. self._seq_counter = itertools.count(1) self._seq_lock = threading.Lock() @@ -93,17 +91,34 @@ def _ensure_created(self) -> None: raise RuntimeError('Server does not support training handles; upgrade Twinkle server and SDK together.') self._created = True + _SEQUENCED_ENDPOINTS = frozenset({ + 'forward_backward', + 'forward_backward_from_data_plane', + 'backward', + 'step', + 'lr_step', + 'clip_grad_and_step', + }) + + def _training_values(self, values: dict[str, Any]) -> dict[str, Any]: + """Choose the label and its handle together, under one lock.""" + with self._training_lock: + name = values.get('adapter_name', self.adapter_name) + training_id = values.get('training_id') or self._training_ids.get(name) + if training_id is None: + raise RuntimeError('Create this training with add_adapter_to_model on this object before using it.') + return {**values, 'adapter_name': name, 'training_id': training_id} + def _submit(self, endpoint: str, model_cls, response_cls, **values): - """Build, send, and resolve one twinkle-native request.""" - self._ensure_created() + """Validate once, then send with a stable handle and retry identity.""" if issubclass(model_cls, model_types.TrainingRequest): - name = values.get('adapter_name') - with self._training_lock: - training_id = values.get('training_id') or self._training_ids.get(name) - if training_id is None: - raise RuntimeError('Create this training with add_adapter_to_model on this object before using it.') - values['training_id'] = training_id + values = self._training_values(values) + if endpoint in self._SEQUENCED_ENDPOINTS: + if 'seq_id' in values: + raise TypeError('seq_id is managed by the client') + values['seq_id'] = self._next_seq_id() body = build_request(model_cls, **values) + self._ensure_created() response = self._transport.post_model(f'{self.server_url}/{endpoint}', body) return self._await_task(response, response_cls) @@ -186,23 +201,12 @@ def remove_adapter(self, adapter_name: str | None = None) -> None: def forward(self, inputs: Any, **kwargs) -> model_types.ForwardResponse: """Execute forward pass on inline model inputs.""" - return self._submit( - 'forward', - model_types.ForwardRequest, - model_types.ForwardResponse, - inputs=inputs, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + return self._submit('forward', model_types.ForwardRequest, model_types.ForwardResponse, inputs=inputs, **kwargs) def forward_only(self, inputs: Any, **kwargs) -> model_types.ForwardResponse: """Execute forward pass without gradient computation on inline inputs.""" return self._submit( - 'forward_only', - model_types.ForwardOnlyRequest, - model_types.ForwardResponse, - inputs=inputs, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + 'forward_only', model_types.ForwardOnlyRequest, model_types.ForwardResponse, inputs=inputs, **kwargs) def forward_backward(self, inputs: Any, **kwargs) -> model_types.ForwardBackwardResponse: """Execute combined forward and backward pass on inline inputs.""" @@ -211,37 +215,20 @@ def forward_backward(self, inputs: Any, **kwargs) -> model_types.ForwardBackward model_types.ForwardBackwardTaskRequest, model_types.ForwardBackwardResponse, inputs=inputs, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - seq_id=self._next_seq_id(), **kwargs) def calculate_loss(self, **kwargs) -> model_types.CalculateLossResponse: """Calculate loss from model outputs.""" - return self._submit( - 'calculate_loss', - model_types.AdapterRequest, - model_types.CalculateLossResponse, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + return self._submit('calculate_loss', model_types.AdapterRequest, model_types.CalculateLossResponse, **kwargs) def get_train_configs(self, **kwargs) -> model_types.GetTrainConfigsResponse: """Get training configs.""" - return self._submit( - 'get_train_configs', - model_types.AdapterRequest, - model_types.GetTrainConfigsResponse, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + return self._submit('get_train_configs', model_types.AdapterRequest, model_types.GetTrainConfigsResponse, + **kwargs) def backward(self, **kwargs) -> None: """Execute backward pass.""" - self._submit( - 'backward', - model_types.AdapterRequest, - None, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - seq_id=self._next_seq_id(), - **kwargs) + self._submit('backward', model_types.AdapterRequest, None, **kwargs) # ------------------------------------------------------------------ # # Data-plane forward family @@ -261,7 +248,6 @@ def forward_from_data_plane( model_types.DataPlaneForwardRequest, model_types.ForwardResponse, input_refs=_data_refs(inputs), - adapter_name=kwargs.pop('adapter_name', self.adapter_name), input_field=input_field, kwarg_fields=kwarg_fields or {}, **kwargs) @@ -282,7 +268,6 @@ def forward_only_from_data_plane( model_types.DataPlaneForwardOnlyRequest, model_types.ForwardResponse, input_refs=_data_refs(inputs), - adapter_name=kwargs.pop('adapter_name', self.adapter_name), input_field=input_field, kwarg_fields=kwarg_fields or {}, output_ref=output_ref.model_dump() if output_ref is not None else None, @@ -306,10 +291,8 @@ def forward_backward_from_data_plane( model_types.DataPlaneForwardRequest, model_types.ForwardBackwardResponse, input_refs=_data_refs(inputs), - adapter_name=kwargs.pop('adapter_name', self.adapter_name), input_field=input_field, kwarg_fields=kwarg_fields or {}, - seq_id=self._next_seq_id(), **kwargs) # ------------------------------------------------------------------ # @@ -318,32 +301,15 @@ def forward_backward_from_data_plane( def step(self, **kwargs) -> None: """Execute optimizer step.""" - self._submit( - 'step', - model_types.StepRequest, - None, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - seq_id=self._next_seq_id(), - **kwargs) + self._submit('step', model_types.StepRequest, None, **kwargs) def zero_grad(self, **kwargs) -> None: """Zero out gradients.""" - self._submit( - 'zero_grad', - model_types.AdapterRequest, - None, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('zero_grad', model_types.AdapterRequest, None, **kwargs) def lr_step(self, **kwargs) -> None: """Execute learning rate scheduler step.""" - self._submit( - 'lr_step', - model_types.LrStepRequest, - None, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - seq_id=self._next_seq_id(), - **kwargs) + self._submit('lr_step', model_types.LrStepRequest, None, **kwargs) def clip_grad_norm(self, max_grad_norm: float = 1.0, @@ -354,7 +320,6 @@ def clip_grad_norm(self, 'clip_grad_norm', model_types.ClipGradNormRequest, model_types.ClipGradNormResponse, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), max_grad_norm=max_grad_norm, norm_type=norm_type, **kwargs) @@ -365,10 +330,8 @@ def clip_grad_and_step(self, max_grad_norm: float = 1.0, norm_type: int = 2, **k 'clip_grad_and_step', model_types.ClipGradAndStepRequest, None, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), max_grad_norm=max_grad_norm, norm_type=norm_type, - seq_id=self._next_seq_id(), **kwargs) # ------------------------------------------------------------------ # @@ -377,33 +340,15 @@ def clip_grad_and_step(self, max_grad_norm: float = 1.0, norm_type: int = 2, **k def set_loss(self, loss_cls: str, **kwargs) -> None: """Set the loss function.""" - self._submit( - 'set_loss', - model_types.SetLossRequest, - None, - loss_cls=loss_cls, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('set_loss', model_types.SetLossRequest, None, loss_cls=loss_cls, **kwargs) def set_optimizer(self, optimizer_cls: str, **kwargs) -> None: """Set the optimizer.""" - self._submit( - 'set_optimizer', - model_types.SetOptimizerRequest, - None, - optimizer_cls=optimizer_cls, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('set_optimizer', model_types.SetOptimizerRequest, None, optimizer_cls=optimizer_cls, **kwargs) def set_lr_scheduler(self, scheduler_cls: str, **kwargs) -> None: """Set the learning rate scheduler.""" - self._submit( - 'set_lr_scheduler', - model_types.SetLrSchedulerRequest, - None, - scheduler_cls=scheduler_cls, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('set_lr_scheduler', model_types.SetLrSchedulerRequest, None, scheduler_cls=scheduler_cls, **kwargs) def set_template(self, template_cls: str, **kwargs) -> None: """Set the template for data processing. @@ -413,44 +358,20 @@ def set_template(self, template_cls: str, **kwargs) -> None: effect. A caller that passes it explicitly still reaches the template constructor through the passthrough region. """ - self._submit( - 'set_template', - model_types.SetTemplateRequest, - None, - template_cls=template_cls, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('set_template', model_types.SetTemplateRequest, None, template_cls=template_cls, **kwargs) def set_processor(self, processor_cls: str, **kwargs) -> None: """Set the input processor.""" - self._submit( - 'set_processor', - model_types.SetProcessorRequest, - None, - processor_cls=processor_cls, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('set_processor', model_types.SetProcessorRequest, None, processor_cls=processor_cls, **kwargs) def add_metric(self, metric_cls: str, is_training: bool | None = None, **kwargs) -> None: """Add a metric to the model.""" self._submit( - 'add_metric', - model_types.AddMetricRequest, - None, - metric_cls=metric_cls, - is_training=is_training, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + 'add_metric', model_types.AddMetricRequest, None, metric_cls=metric_cls, is_training=is_training, **kwargs) def apply_patch(self, patch_cls: str, **kwargs) -> None: """Apply a patch to the model.""" - self._submit( - 'apply_patch', - model_types.ApplyPatchRequest, - None, - patch_cls=patch_cls, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('apply_patch', model_types.ApplyPatchRequest, None, patch_cls=patch_cls, **kwargs) def calculate_metric(self, is_training: bool = True, **kwargs) -> model_types.CalculateMetricResponse: """Calculate metrics from model outputs.""" @@ -459,7 +380,6 @@ def calculate_metric(self, is_training: bool = True, **kwargs) -> model_types.Ca model_types.CalculateMetricRequest, model_types.CalculateMetricResponse, is_training=is_training, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), **kwargs) # ------------------------------------------------------------------ # @@ -468,23 +388,11 @@ def calculate_metric(self, is_training: bool = True, **kwargs) -> model_types.Ca def save(self, name: str, **kwargs) -> model_types.SaveResponse: """Save model checkpoint.""" - return self._submit( - 'save', - model_types.SaveRequest, - model_types.SaveResponse, - name=name, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + return self._submit('save', model_types.SaveRequest, model_types.SaveResponse, name=name, **kwargs) def load(self, name: str, **kwargs) -> None: """Load model checkpoint.""" - self._submit( - 'load', - model_types.LoadRequest, - None, - name=name, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('load', model_types.LoadRequest, None, name=name, **kwargs) def resume_from_checkpoint(self, name: str, *, resume_only_model: bool = False, **kwargs) -> dict[str, Any]: """Resume weights (and optionally optimizer state) from a checkpoint.""" @@ -493,7 +401,6 @@ def resume_from_checkpoint(self, name: str, *, resume_only_model: bool = False, model_types.ResumeFromCheckpointRequest, model_types.TrainingProgressResponse, name=name, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), resume_only_model=resume_only_model, **kwargs) return progress.result @@ -503,8 +410,6 @@ def upload_to_hub( checkpoint_dir: str, hub_model_id: str, hub_token: str | None = None, - async_upload: bool = True, - poll_interval: float = 5.0, ) -> None: """Upload model checkpoint to hub. @@ -515,10 +420,6 @@ def upload_to_hub( checkpoint_dir: The directory path of the checkpoint to upload. hub_model_id: The hub model id. hub_token: The hub token (optional). - async_upload: Deprecated, has no effect. The server always runs the - upload in the background and the client waits via the future layer. - poll_interval: Deprecated, has no effect. Pacing is now owned by the - server-side long-poll of the Retrieve_Endpoint. """ logger.info('[upload_to_hub] submitting upload, waiting for completion...') self._submit( diff --git a/tests/server/checkpoint/test_checkpoint_manager.py b/tests/server/checkpoint/test_checkpoint_manager.py index 5303a876b..b342d2935 100644 --- a/tests/server/checkpoint/test_checkpoint_manager.py +++ b/tests/server/checkpoint/test_checkpoint_manager.py @@ -414,9 +414,9 @@ def write(path, data): @pytest.mark.asyncio async def test_reserved_sampler_name_is_rejected_before_backend_or_file_io(manager): from unittest.mock import AsyncMock - from twinkle.server.task_queue.types import UserTaskError + from twinkle.server.exceptions import RequestRejectedError save = AsyncMock() - with pytest.raises(UserTaskError, match='reserved'): + with pytest.raises(RequestRejectedError, match='reserved'): await manager.save_sampler('run', 'latest', save) save.assert_not_awaited() assert not Path(manager.get_save_dir('run', is_sampler=True)).exists() diff --git a/tests/server/config/server_config_4b_e2e.yaml b/tests/server/config/server_config_4b_e2e.yaml index 6fc4d658f..07cfeca25 100644 --- a/tests/server/config/server_config_4b_e2e.yaml +++ b/tests/server/config/server_config_4b_e2e.yaml @@ -32,6 +32,8 @@ applications: args: backend: transformers model_id: "ms://Qwen/Qwen3.5-4B" + min_replicas: 1 + max_replicas: 1 max_length: 10240 max_loras: 10 nproc_per_node: 2 @@ -45,9 +47,6 @@ applications: queue_config: rps_limit: 100 tps_limit: 100000 - adapter_config: - adapter_timeout: 30 - adapter_max_lifetime: 36000 deployments: - name: ModelManagement autoscaling_config: diff --git a/tests/server/config/server_config_4b_e2e_megatron.yaml b/tests/server/config/server_config_4b_e2e_megatron.yaml index ab71f6bb5..d235a0aa3 100644 --- a/tests/server/config/server_config_4b_e2e_megatron.yaml +++ b/tests/server/config/server_config_4b_e2e_megatron.yaml @@ -32,6 +32,8 @@ applications: args: backend: megatron model_id: "ms://Qwen/Qwen3.5-4B" + min_replicas: 1 + max_replicas: 1 max_length: 10240 max_loras: 10 nproc_per_node: 4 @@ -46,8 +48,6 @@ applications: queue_config: rps_limit: 100 tps_limit: 100000 - adapter_config: - adapter_timeout: 30 deployments: - name: ModelManagement autoscaling_config: diff --git a/tests/server/config/test_server_config.py b/tests/server/config/test_server_config.py index 9f507fa61..4ce5d5704 100644 --- a/tests/server/config/test_server_config.py +++ b/tests/server/config/test_server_config.py @@ -21,7 +21,7 @@ from twinkle.server.config import ApplicationSpec, ServerConfig from twinkle.server.exceptions import ConfigParseError from twinkle.server.launcher import ServerLauncher -from twinkle.server.sampler.app import SAMPLER_SELECTOR, build_sampler_app +from twinkle.server.sampling.execution.app import SAMPLER_SELECTOR, build_sampler_app # ---------- minimal valid config strategy ---------------------------------- # diff --git a/tests/server/config/test_training_config.py b/tests/server/config/test_training_config.py index 9c9a055fa..35ecba6f0 100644 --- a/tests/server/config/test_training_config.py +++ b/tests/server/config/test_training_config.py @@ -10,16 +10,14 @@ from twinkle.server.config import ServerConfig from twinkle.server.config.training import TrainingConfig from twinkle.server.exceptions import RequestRejectedError -from twinkle.server.training.catalog import TrainingCatalog -def model_config(): +def model_config(mode='lora'): + devices = 2 if mode == 'lora' else 4 return { - 'name': 'Qwen/small', 'model_id': 'ms://Qwen/small', 'backend': 'transformers', 'device_type': 'cuda', - 'modes': { - 'lora': {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}}, - 'full': {'num_devices': 4, 'device_mesh': {'fsdp_size': 4}}, - }, + 'name': 'Qwen/small' if mode == 'lora' else 'Qwen/small-full', + 'model_id': 'ms://Qwen/small', 'backend': 'transformers', 'device_type': 'cuda', + 'mode': mode, 'config': {'num_devices': devices, 'device_mesh': {'fsdp_size': devices}}, } @@ -33,10 +31,10 @@ def model_application(): def test_catalog_defaults_and_mode_selection(): - config = TrainingConfig(models=[model_config()]) - catalog = TrainingCatalog(config) + config = TrainingConfig(models=[model_config(), model_config('full')]) + catalog = config shared = catalog.get('Qwen/small', 'lora') - exclusive = catalog.get('Qwen/small', 'full') + exclusive = catalog.get('Qwen/small-full', 'full') assert shared.capacity == 5 assert exclusive.capacity == 1 assert shared.config.min_replicas == exclusive.config.min_replicas == 0 @@ -46,22 +44,18 @@ def test_catalog_defaults_and_mode_selection(): catalog.get('not-in-catalog', 'lora') -def test_existing_yaml_retains_heartbeat_timeout_and_max_lifetime(): +def test_existing_yaml_uses_explicit_instance_limits_and_business_idle(): from pathlib import Path paths = sorted(Path('cookbook/client/server').glob('**/server_config*.yaml')) assert paths for path in paths: config = ServerConfig.from_yaml(path) - for app in config.applications: - if app.import_path != 'model': - continue - normalized = config.training_config.models[ - [item.model_id for item in config.training_config.models].index(app.args.model_id)] - mode = normalized.modes.full if app.args.train_mode == 'full' else normalized.modes.lora - legacy = app.args.adapter_config or {} - assert mode.idle_timeout_s == app.args.idle_timeout_s - assert normalized.legacy_expiration.session_timeout_s == legacy.get('adapter_timeout', 1800) - assert normalized.legacy_expiration.max_lifetime_s == legacy.get('adapter_max_lifetime', 36000) + applications = [app for app in config.applications if app.import_path == 'model'] + for app, model in zip(applications, config.training_config.models): + assert model.config.idle_timeout_s == app.args.idle_timeout_s + assert model.config.min_replicas == app.args.min_replicas + assert model.config.max_replicas == app.args.max_replicas + assert 'adapter_config' not in type(app.args).model_fields @pytest.mark.parametrize('field,value', [ @@ -71,7 +65,7 @@ def test_existing_yaml_retains_heartbeat_timeout_and_max_lifetime(): ]) def test_invalid_mode_cannot_reach_allocation(field, value): model = model_config() - model['modes']['lora'][field] = value + model['config'][field] = value with pytest.raises(ValidationError): TrainingConfig(models=[model]) @@ -79,41 +73,42 @@ def test_invalid_mode_cannot_reach_allocation(field, value): def test_ep_and_vpp_do_not_multiply_physical_device_count(): model = model_config() model['backend'] = 'megatron' - model['modes']['lora'] = { + model['config'] = { 'num_devices': 8, 'device_mesh': {'tp_size': 2, 'dp_size': 2, 'pp_size': 2, 'ep_size': 2, 'vpp_size': 2}, } - assert TrainingConfig(models=[model]).models[0].modes.lora.num_devices == 8 + assert TrainingConfig(models=[model]).models[0].config.num_devices == 8 -@pytest.mark.parametrize('mutation', ['duplicate', 'no_modes', 'full_lora_field', 'hybrid_megatron', 'mesh_type', 'min_max']) +@pytest.mark.parametrize('mutation', ['duplicate', 'missing_mode', 'full_lora_field', 'hybrid_megatron', 'mesh_type', 'min_max']) def test_cross_field_constraints(mutation): model = model_config() models = [model] if mutation == 'duplicate': models.append(copy.deepcopy(model)) - elif mutation == 'no_modes': - model['modes'] = {} + elif mutation == 'missing_mode': + del model['mode'] elif mutation == 'full_lora_field': - model['modes']['full']['max_loras'] = 2 + model = model_config('full') + model['config']['max_loras'] = 2 + models = [model] elif mutation == 'hybrid_megatron': model['backend'] = 'megatron' - model['modes']['lora']['hybrid'] = {'allocation_path': '/tmp/allocation.json'} + model['config']['hybrid'] = {'allocation_path': '/tmp/allocation.json'} elif mutation == 'mesh_type': - model['modes']['lora']['device_mesh']['device_type'] = 'npu' + model['config']['device_mesh']['device_type'] = 'npu' elif mutation == 'min_max': - model['modes']['lora'].update(min_replicas=2, max_replicas=1) + model['config'].update(min_replicas=2, max_replicas=1) with pytest.raises(ValidationError): TrainingConfig(models=models) def test_full_only_is_explicit_and_lora_is_unavailable(): - model = model_config() - del model['modes']['lora'] - catalog = TrainingCatalog(TrainingConfig(models=[model])) - assert catalog.get('Qwen/small', 'full').capacity == 1 + model = model_config('full') + catalog = TrainingConfig(models=[model]) + assert catalog.get('Qwen/small-full', 'full').capacity == 1 with pytest.raises(RequestRejectedError, match='does not provide lora'): - catalog.get('Qwen/small', 'lora') + catalog.get('Qwen/small-full', 'lora') def test_training_namespace_conflicts_are_rejected(): @@ -142,26 +137,23 @@ def test_existing_yaml_examples_keep_their_model_entries(monkeypatch): model = config.training_config.models[0] assert model.model_id == models[0].args.model_id assert model.route_prefix == models[0].route_prefix - mode = model.modes.full or model.modes.lora + mode = model.config original = models[0] expected_min = original.args.min_replicas - if expected_min is None: - expected_min = original.deployments[0]['autoscaling_config']['min_replicas'] assert mode.min_replicas == expected_min assert ServerConfig.model_validate(config.to_yaml_dict()).training_config == config.training_config -def test_legacy_lifecycle_overrides_and_full_instances(): +def test_configured_lifecycle_overrides_and_full_instances(): app = {'name': 'model', 'route_prefix': '/api/v1/model/small', 'import_path': 'model', 'args': { 'model_id': 'weights', 'backend': 'mock', 'train_mode': 'full', 'device_group': {'ranks': 2, 'device_type': 'CPU'}, 'device_mesh': {'dp_size': 2}, 'min_replicas': 0, 'max_replicas': 3, 'idle_timeout_s': 17, - 'adapter_config': {'adapter_timeout': 30}, }, 'deployments': [{'num_replicas': 1}]} config = ServerConfig(applications=[app]) - mode = config.training_config.models[0].modes.full + mode = config.training_config.models[0].config assert (mode.num_devices, mode.min_replicas, mode.max_replicas, mode.idle_timeout_s) == (2, 0, 3, 17) - assert config.training_config.models[0].modes.lora is None + assert config.training_config.models[0].mode == 'full' def test_separate_model_applications_select_lora_and_full(): @@ -170,7 +162,7 @@ def test_separate_model_applications_select_lora_and_full(): exclusive.update(name='model-small-full', route_prefix='/api/v1/model/Qwen/small-full') exclusive['args'].update(train_mode='full', min_replicas=0, max_replicas=2, idle_timeout_s=60) config = ServerConfig(applications=[shared, exclusive]) - catalog = TrainingCatalog(config.training_config) + catalog = config.training_config assert catalog.get('Qwen/small', 'lora').capacity == 5 full = catalog.get('Qwen/small-full', 'full') assert full.capacity == 1 @@ -201,11 +193,11 @@ def builder(**kwargs): assert captured['supported_models'] == ['teacher'] -def legacy_application(): +def configured_application(): return {'name': 'model', 'route_prefix': '/api/v1/model/small', 'import_path': 'model', 'args': { 'model_id': 'weights', 'backend': 'mock', 'train_mode': 'lora', 'max_length': 128, 'device_group': {'ranks': 2, 'device_type': 'cpu'}, 'device_mesh': {'dp_size': 2}, - 'env_vars': {'USER_OPTION': 'args-wins'}, 'adapter_config': {'adapter_timeout': 30}, + 'env_vars': {'USER_OPTION': 'args-wins'}, 'min_replicas': 1, 'max_replicas': 2, }, 'deployments': [{ 'name': 'ModelManagement', 'max_ongoing_requests': 17, 'autoscaling_config': {'min_replicas': 1, 'max_replicas': 2, 'target_ongoing_requests': 16}, @@ -219,21 +211,22 @@ def legacy_application(): {'modes': None}, {'modes': {'lora': {'num_devices': 2, 'device_mesh': {'dp_size': 2}}}}, {'device_type': 'cpu'}, + {'adapter_config': {'adapter_timeout': 30}}, ]) def test_model_yaml_rejects_alternative_resource_and_mode_syntax(extra): - app = legacy_application() + app = configured_application() app['args'].update(extra) with pytest.raises(ValidationError, match='Extra inputs are not permitted'): ServerConfig(applications=[app]) -def test_legacy_options_reach_cpu_router_without_changing_yaml(monkeypatch): +def test_configured_options_reach_cpu_router_without_changing_yaml(monkeypatch): from types import SimpleNamespace from twinkle.server.launcher.server_launcher import ServerLauncher from twinkle.server.training.actor import TrainingControllerActor import ray from ray import serve - config = ServerConfig(applications=[legacy_application()]) + config = ServerConfig(applications=[configured_application()]) remote = SimpleNamespace(initialize=SimpleNamespace(remote=lambda: None)) monkeypatch.setattr(TrainingControllerActor, 'options', lambda **kwargs: SimpleNamespace(remote=lambda *args: remote)) @@ -253,16 +246,16 @@ def router(*args): assert actor['num_cpus'] == 2 and actor['runtime_env']['pip'] == ['example==1.0'] assert actor['runtime_env']['env_vars']['USER_OPTION'] == 'args-wins' assert actor['runtime_env']['env_vars']['OTHER_OPTION'] == 'retained' - assert config.to_yaml_dict()['applications'][0]['deployments'][0] == legacy_application()['deployments'][0] + assert config.to_yaml_dict()['applications'][0]['deployments'][0] == configured_application()['deployments'][0] @pytest.mark.asyncio -async def test_legacy_options_reach_runtime_with_one_controller_owned_replica(monkeypatch): +async def test_configured_options_reach_runtime_with_one_controller_owned_replica(monkeypatch): import importlib from twinkle.server.state.training_records import InstanceRecord from twinkle.server.training.driver import RayInstanceDriver from ray import serve - config = ServerConfig(applications=[legacy_application()]) + config = ServerConfig(applications=[configured_application()]) model = config.training_config.models[0] driver = RayInstanceDriver(config.training_config, None, 'test', 'controller') instance = InstanceRecord(instance_id='instance', controller_epoch='epoch', base_model=model.name, @@ -273,10 +266,10 @@ def builder(**kwargs): captured.update(kwargs) return 'app' - monkeypatch.setattr('twinkle.server.model.app.build_model_app', builder) + monkeypatch.setattr('twinkle.server.training.execution.app.build_model_app', builder) monkeypatch.setattr(importlib.import_module('ray.util.placement_group'), 'get_placement_group', lambda name: 'pg') monkeypatch.setattr(serve, 'run', lambda *args, **kwargs: None) - await driver.deploy(instance, driver._catalog.get(model.name, 'lora')) + await driver.deploy(instance, driver.config.get(model.name, 'lora')) options = captured['deploy_options'] assert options['num_replicas'] == 1 and 'autoscaling_config' not in options assert options['max_ongoing_requests'] == 17 @@ -284,29 +277,31 @@ def builder(**kwargs): assert options['ray_actor_options']['runtime_env']['pip'] == ['example==1.0'] assert captured['max_length'] == 128 assert captured['nproc_per_node'] == 2 - assert captured['queue_config'] == model.modes.lora.queue_config + assert captured['queue_config'] == model.config.queue_config assert not driver.runtime.deployments def test_mutated_configuration_does_not_reuse_a_stale_catalog(): - config = ServerConfig(applications=[legacy_application()]) - assert config.training_config.models[0].modes.lora.min_replicas == 1 + config = ServerConfig(applications=[configured_application()]) + assert config.training_config.models[0].config.min_replicas == 1 config.applications[0].args.min_replicas = 0 - assert config.training_config.models[0].modes.lora.min_replicas == 0 + assert config.training_config.models[0].config.min_replicas == 0 -def test_unsupported_adapter_parameters_are_not_silently_ignored(): - app = legacy_application() - app['args']['adapter_config']['misspelled_timeout'] = 12 - with pytest.raises(ValidationError, match='misspelled_timeout'): - ServerConfig(applications=[app]) +def test_router_replica_limits_do_not_implicitly_set_training_instance_limits(): + app = configured_application() + app['args'].pop('min_replicas') + app['args'].pop('max_replicas') + model = ServerConfig(applications=[app]).training_config.models[0] + assert model.config.min_replicas == 0 and model.config.max_replicas is None + assert model.deploy_options['autoscaling_config']['min_replicas'] == 1 def test_l20d_candidate_is_lora_only_and_bounds_total_devices(): config = ServerConfig.from_yaml('cookbook/client/server/megatron/server_config.yaml') model = config.training_config.models[0] - mode = model.modes.lora - assert model.modes.full is None + mode = model.config + assert model.mode == 'lora' assert (mode.num_devices, mode.min_replicas, mode.max_replicas, mode.max_loras) == (2, 0, 2, 3) assert mode.device_mesh.world_size() == 2 and mode.device_mesh.pp_size == 2 and mode.idle_timeout_s == 600 sampler = next(app for app in config.applications if app.import_path == 'sampler') @@ -317,22 +312,22 @@ def test_l20d_candidate_is_lora_only_and_bounds_total_devices(): assert mode.num_devices * mode.max_replicas + sampler.args.device_group['ranks'] * sampler.args.max_replicas == 8 assert mode.max_loras * mode.max_replicas == 6 with pytest.raises(RequestRejectedError, match='does not provide full'): - TrainingCatalog(config.training_config).get(model.name, 'full') + config.training_config.get(model.name, 'full') @pytest.mark.parametrize('mode', ['lora', 'full']) -@pytest.mark.parametrize('legacy', [False, True]) -def test_single_device_megatron_is_rejected_before_allocation(mode, legacy): - if legacy: - app = legacy_application() +@pytest.mark.parametrize('yaml_entrypoint', [False, True]) +def test_single_device_megatron_is_rejected_before_allocation(mode, yaml_entrypoint): + if yaml_entrypoint: + app = configured_application() app['args'].update(backend='megatron', train_mode=mode, device_group={'ranks': 1, 'device_type': 'cuda'}, device_mesh={'dp_size': 1}) else: - model = model_config() + model = model_config(mode) model['backend'] = 'megatron' - model['modes'] = {mode: {'num_devices': 1, 'device_mesh': {'dp_size': 1}}} + model['config'] = {'num_devices': 1, 'device_mesh': {'dp_size': 1}} with pytest.raises(ValidationError, match='at least 2 devices'): - if legacy: + if yaml_entrypoint: ServerConfig(applications=[app]) else: TrainingConfig(models=[model]) @@ -341,4 +336,4 @@ def test_single_device_megatron_is_rejected_before_allocation(mode, legacy): def test_transformers_still_accepts_single_device(): app = model_application() app['args'].update(device_group={'ranks': 1, 'device_type': 'cuda'}, device_mesh={'dp_size': 1}) - assert ServerConfig(applications=[app]).training_config.models[0].modes.lora.num_devices == 1 + assert ServerConfig(applications=[app]).training_config.models[0].config.num_devices == 1 diff --git a/tests/server/contract/client_api_harness.py b/tests/server/contract/client_api_harness.py index ab82e0582..479b36049 100644 --- a/tests/server/contract/client_api_harness.py +++ b/tests/server/contract/client_api_harness.py @@ -66,8 +66,8 @@ def build_gateway_app() -> FastAPI: def build_model_app() -> FastAPI: - from twinkle.server.model.tinker_handlers import _register_model_tinker_routes - from twinkle.server.model.twinkle_handlers import _register_model_twinkle_routes + from twinkle.server.training.execution.tinker_handlers import _register_model_tinker_routes + from twinkle.server.training.execution.twinkle_handlers import _register_model_twinkle_routes app = FastAPI() _register_model_tinker_routes(app, _noop_self) @@ -76,8 +76,8 @@ def build_model_app() -> FastAPI: def build_sampler_app() -> FastAPI: - from twinkle.server.sampler.tinker_handlers import _register_tinker_sampler_routes - from twinkle.server.sampler.twinkle_handlers import _register_twinkle_sampler_routes + from twinkle.server.sampling.execution.tinker_handlers import _register_tinker_sampler_routes + from twinkle.server.sampling.execution.twinkle_handlers import _register_twinkle_sampler_routes app = FastAPI() _register_tinker_sampler_routes(app, _noop_self) diff --git a/tests/server/contract/test_client_api_contract.py b/tests/server/contract/test_client_api_contract.py index f543ee436..89d9fe00a 100644 --- a/tests/server/contract/test_client_api_contract.py +++ b/tests/server/contract/test_client_api_contract.py @@ -66,7 +66,7 @@ def test_new_client_types_importable(): for symbol in ('StrictRequest', 'ResponseModel', 'DataModel', 'backend_only'): assert hasattr(base, symbol) - for symbol in ('ErrorPayload', 'ErrorCategory', 'QueueStateLiteral'): + for symbol in ('ErrorPayload', 'ErrorCategory', 'QueueState'): assert hasattr(errors, symbol) diff --git a/tests/server/fixtures/server_config_mock.yaml b/tests/server/fixtures/server_config_mock.yaml index 67303ed5c..2a7711062 100644 --- a/tests/server/fixtures/server_config_mock.yaml +++ b/tests/server/fixtures/server_config_mock.yaml @@ -42,6 +42,8 @@ applications: args: backend: mock model_id: mock-model + min_replicas: 1 + max_replicas: 1 nproc_per_node: 1 device_group: name: model diff --git a/tests/server/fixtures/server_config_mock_redis.yaml b/tests/server/fixtures/server_config_mock_redis.yaml index 4c7297fe3..59c2ab436 100644 --- a/tests/server/fixtures/server_config_mock_redis.yaml +++ b/tests/server/fixtures/server_config_mock_redis.yaml @@ -41,6 +41,8 @@ applications: args: backend: mock model_id: mock-model + min_replicas: 1 + max_replicas: 1 nproc_per_node: 1 device_group: name: model diff --git a/tests/server/integration/e2e_helpers.py b/tests/server/integration/e2e_helpers.py index 2eca95f0d..8ac1c6905 100644 --- a/tests/server/integration/e2e_helpers.py +++ b/tests/server/integration/e2e_helpers.py @@ -290,3 +290,34 @@ def assert_metrics_valid(metrics: Any, label: str) -> None: assert math.isfinite(val), ( f'[{label}] Metric {key}={val} is not finite!') log(f'[{label}] Metrics valid: {metrics}') + + +def sampled_datum(prompt_tokens, sample, advantage): + from tinker import types + + tokens = prompt_tokens + list(sample.tokens) + prompt_positions = len(prompt_tokens) - 1 + size = len(tokens) - 1 + + def tensor(values, dtype): + return types.TensorData(data=values, dtype=dtype, shape=[size]) + + return types.Datum( + model_input=types.ModelInput.from_ints(tokens[:-1]), + loss_fn_inputs={ + 'target_tokens': tensor(tokens[1:], 'int64'), + 'weights': tensor([0.0] * prompt_positions + [1.0] * len(sample.tokens), 'float32'), + 'logprobs': tensor([0.0] * prompt_positions + list(sample.logprobs), 'float32'), + 'advantages': tensor([0.0] * prompt_positions + [advantage] * len(sample.tokens), 'float32'), + }, + ) + + +def completion_logprobs(client, datum, timeout): + result = client.forward([datum], 'cross_entropy').result(timeout=timeout) + values = np.asarray(result.loss_fn_outputs[0]['logprobs'].tolist(), dtype=np.float64) + weights = np.asarray(datum.loss_fn_inputs['weights'].tolist()) + assert len(values) >= len(weights) + values = values[:len(weights)][weights != 0] + assert values.size and np.isfinite(values).all() + return values diff --git a/tests/server/integration/on_demand_training/acceptance_support.py b/tests/server/integration/on_demand_training/acceptance_support.py index c14827c6c..07abc6b5b 100644 --- a/tests/server/integration/on_demand_training/acceptance_support.py +++ b/tests/server/integration/on_demand_training/acceptance_support.py @@ -65,7 +65,7 @@ def configure(self): self.client.set_processor('InputProcessor', padding_side='right') def datums(self): - from twinkle.server.model.tinker_datum import input_feature_to_datum + from twinkle.server.training.execution.tinker_datum import input_feature_to_datum return [input_feature_to_datum(feature) for feature in self.inputs] def step(self): diff --git a/tests/server/integration/on_demand_training/http_probe.py b/tests/server/integration/on_demand_training/http_probe.py index 5d8f56248..13efd3f6c 100644 --- a/tests/server/integration/on_demand_training/http_probe.py +++ b/tests/server/integration/on_demand_training/http_probe.py @@ -43,7 +43,7 @@ def main(): 'host': '127.0.0.1', 'port': port }, - 'training': { + 'resources': { 'startup_timeout_s': 600, 'shutdown_timeout_s': 180 }, @@ -78,6 +78,13 @@ def main(): 'autoscaling_config': {'min_replicas': 1, 'max_replicas': 1, 'target_ongoing_requests': 4}, 'ray_actor_options': {'num_cpus': 0.1, 'runtime_env': {'env_vars': {'LEGACY_CONFIG_PROBE': '1'}}}, }], + }, { + 'name': 'routed', 'import_path': 'model', 'route_prefix': '/api/v1/model/routed', + 'args': { + 'model_id': 'small', 'backend': 'mock', 'train_mode': 'lora', + 'device_group': {'ranks': 1, 'device_type': 'cpu'}, 'device_mesh': {'dp_size': 1}, + 'max_loras': 1, 'max_replicas': 2, + }, }, { 'name': 'gateway', 'import_path': 'server', @@ -118,6 +125,24 @@ def wait_epoch(controller, old_epoch): first = MultiLoraTransformersModel('small', transport=transport) second = MultiLoraTransformersModel('small', transport=transport) from peft import LoraConfig + + def assert_separate_instances(identities): + diagnostic = ray.get(controller.diagnostics.remote()) + owners = {identity: item['instance_id'] for item in diagnostic['instances'] + for identity in item['training_ids'] if identity in identities} + assert set(owners) == set(identities) and len(set(owners.values())) == 2, owners + return owners + + routed = [MultiLoraTransformersModel('routed', transport=transport) for _ in range(2)] + for client in routed: + client.add_adapter_to_model('same', LoraConfig(r=8, target_modules='all-linear')) + native_owners = assert_separate_instances([client._training_ids['same'] for client in routed]) + for _ in range(3): + for client in reversed(routed): + assert json.loads(client.get_train_configs().result) == {} + client.zero_grad() + for client in routed: + client.remove_adapter() first.add_adapter_to_model('same', LoraConfig(r=8, target_modules='all-linear')) instance = next( item for item in ray.get(controller.diagnostics.remote())['instances'] @@ -148,6 +173,13 @@ def wait_epoch(controller, old_epoch): init_tinker_client() from tinker import ServiceClient service = ServiceClient(base_url=base, api_key='probe-tinker') + routed_tinker = [service.create_lora_training_client(base_model='routed', rank=8) for _ in range(2)] + tinker_owners = assert_separate_instances([client.model_id for client in routed_tinker]) + for _ in range(3): + for client in reversed(routed_tinker): + assert client.get_info().model_id == client.model_id + for client in routed_tinker: + assert unload_tinker_model(service, client.model_id).model_id == client.model_id shared1 = service.create_lora_training_client(base_model='small', rank=8) shared2 = service.create_lora_training_client(base_model='small', rank=8) assert shared1.model_id != shared2.model_id @@ -220,6 +252,8 @@ def wait_epoch(controller, old_epoch): json.dumps({ 'passed': True, 'native_recreation': True, + 'native_stable_routing': native_owners, + 'tinker_stable_routing': tinker_owners, 'native_lora': True, 'tinker_shared': True, 'tinker_full': True, diff --git a/tests/server/integration/on_demand_training/matrix_validation.py b/tests/server/integration/on_demand_training/matrix_validation.py index 5528b0197..59bfcc9e6 100644 --- a/tests/server/integration/on_demand_training/matrix_validation.py +++ b/tests/server/integration/on_demand_training/matrix_validation.py @@ -277,7 +277,7 @@ def main(): source = root / 'tests/server/config' / ('server_config_4b_e2e.yaml' if args.backend == 'transformers' else 'server_config_4b_e2e_megatron.yaml') raw = yaml.safe_load(source.read_text()) - raw.setdefault('training', {}).update({ + raw.setdefault('resources', {}).update({ 'startup_timeout_s': args.startup_timeout_s, 'shutdown_timeout_s': args.shutdown_timeout_s, }) @@ -435,8 +435,8 @@ def main(): base_url, launcher._training_controller, model_name) evidence['capacity_checked'] = validate_capacity( base_url, launcher._training_controller, model_name, - next(model.name for model in config.training_config.models if model.modes.full is not None), - config.training_config.models[0].modes.lora.max_loras) + next(model.name for model in config.training_config.models if model.mode == 'full'), + config.training_config.models[0].config.max_loras) if args.stress_iterations or args.soak_seconds: from tests.server.integration.on_demand_training.acceptance_support import assert_released controller = launcher._training_controller diff --git a/tests/server/integration/on_demand_training/test_resident_release_e2e.py b/tests/server/integration/on_demand_training/test_resident_release_e2e.py index 39f41f340..c6c830729 100644 --- a/tests/server/integration/on_demand_training/test_resident_release_e2e.py +++ b/tests/server/integration/on_demand_training/test_resident_release_e2e.py @@ -80,7 +80,7 @@ def test_resident_release_ack_after_worker_failure(tmp_path): model_name = 'Qwen/Qwen3.5-4B' raw = { 'ray_namespace': namespace, 'http_options': {'host': '127.0.0.1', 'port': port}, - 'training': {'startup_timeout_s': 1800, 'shutdown_timeout_s': 180, 'sweep_interval_s': 1}, + 'resources': {'startup_timeout_s': 1800, 'shutdown_timeout_s': 180, 'sweep_interval_s': 1}, 'persistence': ({'mode': persistence, 'redis_url': managed.url, 'key_prefix': prefix} if managed else {'mode': persistence}), 'applications': [ diff --git a/tests/server/integration/resources/allocation_probe.py b/tests/server/integration/resources/allocation_probe.py index d3cc874ec..9d057d215 100644 --- a/tests/server/integration/resources/allocation_probe.py +++ b/tests/server/integration/resources/allocation_probe.py @@ -10,7 +10,6 @@ from twinkle.infra._ray.resource_manager import ResourceManager from twinkle.server.config.training import TrainingConfig from twinkle.server.resources.allocation import CapacityUnavailableError, DevicePool -from twinkle.server.training.catalog import TrainingCatalog async def main(): @@ -30,28 +29,15 @@ async def main(): allocations = [] try: for index, count in enumerate((2, 2, 4)): - config = TrainingConfig(models=[{ - 'name': str(index), - 'model_id': str(index), - 'backend': 'mock', - 'device_type': 'cuda', - 'modes': { - 'full': { - 'num_devices': count, - 'device_mesh': { - 'fsdp_size': count - } - } - }, - }]) - spec = TrainingCatalog(config).get(str(index), 'full') - allocation = await pool.allocate(spec.config.num_devices, spec.model.device_type, 'epoch', str(index)) + config = TrainingConfig(models=[{'name': str(index), 'model_id': str(index), 'backend': 'mock', 'device_type': 'cuda', 'mode': 'full', 'config': {'num_devices': count, 'device_mesh': {'fsdp_size': count}}}]) + spec = config.get(str(index), 'full') + allocation = await pool.allocate(spec.config.num_devices, spec.device_type, 'epoch', str(index)) allocations.append(allocation) devices = [device for allocation in allocations for device in allocation.devices] assert len(set(devices)) == 8, devices assert len({allocation.node_id for allocation in allocations}) == 1 try: - await pool.allocate(spec.config.num_devices, spec.model.device_type, 'epoch', 'overflow') + await pool.allocate(spec.config.num_devices, spec.device_type, 'epoch', 'overflow') raise AssertionError('over-allocation was admitted') except CapacityUnavailableError: pass @@ -66,22 +52,8 @@ async def main(): assert get_placement_group(allocations[0].name).id == owned_pg.id freed = allocations[0] await pool.remove(freed.name) - two = TrainingCatalog( - TrainingConfig(models=[{ - 'name': 'replacement', - 'model_id': 'replacement', - 'backend': 'mock', - 'device_type': 'cuda', - 'modes': { - 'full': { - 'num_devices': 2, - 'device_mesh': { - 'fsdp_size': 2 - } - } - }, - }])).get('replacement', 'full') - replacement = await pool.allocate(two.config.num_devices, two.model.device_type, 'epoch', 'replacement') + two = TrainingConfig(models=[{'name': 'replacement', 'model_id': 'replacement', 'backend': 'mock', 'device_type': 'cuda', 'mode': 'full', 'config': {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}}}]).get('replacement', 'full') + replacement = await pool.allocate(two.config.num_devices, two.device_type, 'epoch', 'replacement') allocations.append(replacement) assert set(replacement.devices) == set(freed.devices) print(json.dumps({'passed': True, 'devices': devices, 'reused': list(replacement.devices)})) diff --git a/tests/server/integration/resources/fragmentation_probe.py b/tests/server/integration/resources/fragmentation_probe.py index d715d7e8c..f684c51c3 100644 --- a/tests/server/integration/resources/fragmentation_probe.py +++ b/tests/server/integration/resources/fragmentation_probe.py @@ -10,7 +10,6 @@ from twinkle.server.config.training import TrainingConfig from twinkle.server.resources.allocation import CapacityUnavailableError, DevicePool -from twinkle.server.training.catalog import TrainingCatalog async def main(): @@ -27,28 +26,14 @@ async def main(): # Keep fragmentation rejection fast, but allow a cold Ray worker to # import the allocation probe before confirming a successful placement. pool = DevicePool(ray.get_runtime_context().namespace, allocation_timeout=1, shutdown_timeout=60) - spec = TrainingCatalog( - TrainingConfig(models=[{ - 'name': 'two', - 'model_id': 'two', - 'backend': 'mock', - 'device_type': 'cuda', - 'modes': { - 'full': { - 'num_devices': 2, - 'device_mesh': { - 'fsdp_size': 2 - } - } - }, - }])).get('two', 'full') + spec = TrainingConfig(models=[{'name': 'two', 'model_id': 'two', 'backend': 'mock', 'device_type': 'cuda', 'mode': 'full', 'config': {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}}}]).get('two', 'full') for marker in ('left', 'right'): pg = placement_group([{'GPU': 3, marker: 0.01}], strategy='STRICT_PACK') occupied.append(pg) await asyncio.wait_for(pg.ready(), 20) assert ray.available_resources()['GPU'] == 2 try: - await pool.allocate(spec.config.num_devices, spec.model.device_type, 'epoch', 'fragmented') + await pool.allocate(spec.config.num_devices, spec.device_type, 'epoch', 'fragmented') except CapacityUnavailableError: pass else: @@ -63,7 +48,7 @@ async def main(): await asyncio.sleep(0.05) else: raise AssertionError('Owned GPU reservation was not released') - allocated = await pool.allocate(spec.config.num_devices, spec.model.device_type, 'epoch', 'available') + allocated = await pool.allocate(spec.config.num_devices, spec.device_type, 'epoch', 'available') assert len(set(allocated.devices)) == 2 print( json.dumps({ diff --git a/tests/server/integration/sampler/sampling_probe.py b/tests/server/integration/sampler/sampling_probe.py index 345cce133..8eb1a5f35 100644 --- a/tests/server/integration/sampler/sampling_probe.py +++ b/tests/server/integration/sampler/sampling_probe.py @@ -13,8 +13,8 @@ from twinkle.server.config.application_spec import SamplerArgs from twinkle.server.deployment import bind_deployment, build_deployment_app from twinkle.server.resources.allocation import CapacityUnavailableError, DevicePool -from twinkle.server.sampler.controller import SamplerController -from twinkle.server.sampler.driver import SamplerDriver +from twinkle.server.sampling.controller import SamplerController +from twinkle.server.sampling.driver import SamplerDriver @ray.remote(num_cpus=0.01, num_gpus=1) diff --git a/tests/server/integration/sampler/test_real_pool_e2e.py b/tests/server/integration/sampler/test_real_pool_e2e.py index 8b7345919..49e7d052b 100644 --- a/tests/server/integration/sampler/test_real_pool_e2e.py +++ b/tests/server/integration/sampler/test_real_pool_e2e.py @@ -72,8 +72,8 @@ async def exercise(model, directory, namespace, evidence): from ray.util.placement_group import placement_group_table from transformers import AutoTokenizer from twinkle.server.config.application_spec import SamplerArgs - from twinkle.server.sampler.controller import SamplerController - from twinkle.server.sampler.driver import SamplerDriver + from twinkle.server.sampling.controller import SamplerController + from twinkle.server.sampling.driver import SamplerDriver from twinkle.server.state import get_server_state args = SamplerArgs(model_id=str(model), sampler_type='vllm', diff --git a/tests/server/integration/test_actor_recovery.py b/tests/server/integration/test_actor_recovery.py index dfc03a74b..26aa3725a 100644 --- a/tests/server/integration/test_actor_recovery.py +++ b/tests/server/integration/test_actor_recovery.py @@ -9,8 +9,8 @@ import pytest from fastapi import FastAPI -from twinkle.server.model.app import ModelManagement -from twinkle.server.model.twinkle_handlers import _register_model_twinkle_routes +from twinkle.server.training.execution.app import ModelManagement +from twinkle.server.training.execution.twinkle_handlers import _register_model_twinkle_routes class _MockModel: diff --git a/tests/server/integration/test_blocking_boundary.py b/tests/server/integration/test_blocking_boundary.py index 7d35e229d..7fb895993 100644 --- a/tests/server/integration/test_blocking_boundary.py +++ b/tests/server/integration/test_blocking_boundary.py @@ -21,7 +21,7 @@ from twinkle.server.task_queue.backend_gate import BackendGate # noqa: E402 from twinkle.server.task_queue.mixin import TaskQueueMixin # noqa: E402 -from twinkle.server.task_queue.types import BackendBusyError # noqa: E402 +from twinkle.server.exceptions import BackendBusyError # noqa: E402 class _Harness(TaskQueueMixin): diff --git a/tests/server/integration/test_dpo_e2e.py b/tests/server/integration/test_dpo_e2e.py index 7a69cbcab..6597d5d55 100644 --- a/tests/server/integration/test_dpo_e2e.py +++ b/tests/server/integration/test_dpo_e2e.py @@ -182,7 +182,7 @@ def test_dpo_tinker(): """ from tinker import types from twinkle.dataloader import DataLoader - from twinkle.server.model.tinker_datum import input_feature_to_datum + from twinkle.server.training.execution.tinker_datum import input_feature_to_datum backend = get_backend() log(f'=== test_dpo_tinker [backend={backend}] ===') diff --git a/tests/server/integration/test_full_param_e2e.py b/tests/server/integration/test_full_param_e2e.py index 32d605191..c6f54f0c6 100644 --- a/tests/server/integration/test_full_param_e2e.py +++ b/tests/server/integration/test_full_param_e2e.py @@ -46,7 +46,7 @@ from twinkle.dataloader import DataLoader # noqa: E402 from twinkle.dataset import Dataset, DatasetMeta # noqa: E402 from twinkle.preprocessor import SelfCognitionProcessor # noqa: E402 -from twinkle.server.model.tinker_datum import input_feature_to_datum # noqa: E402 +from twinkle.server.training.execution.tinker_datum import input_feature_to_datum # noqa: E402 init_tinker_client() diff --git a/tests/server/integration/test_sampler_checkpoints_e2e.py b/tests/server/integration/test_sampler_checkpoints_e2e.py index e95858118..4affee9a3 100644 --- a/tests/server/integration/test_sampler_checkpoints_e2e.py +++ b/tests/server/integration/test_sampler_checkpoints_e2e.py @@ -8,7 +8,7 @@ import numpy as np import pytest -from tests.server.integration.test_tinker_rl_correctness_e2e import _completion_logprobs, _datum +from tests.server.integration.e2e_helpers import completion_logprobs, sampled_datum pytestmark = pytest.mark.skipif(os.environ.get('TWINKLE_TEST_GPU_E2E') != '1', reason='Requires a real trainer and vLLM sampler') @@ -58,15 +58,15 @@ def same(actual, expected): sampler_x = service.create_sampling_client(base_model=model_id, model_path=x) first = sample(sampler_x) same(sample(sampler_x), first) - datum = _datum(prompt, first, 1.0) - before = _completion_logprobs(client, datum, timeout) + datum = sampled_datum(prompt, first, 1.0) + before = completion_logprobs(client, datum, timeout) losses = [] for _ in range(3): result = client.forward_backward([datum], 'cross_entropy').result(timeout=timeout) losses.append(float(result.metrics['loss:mean'])) assert np.isfinite(losses[-1]) client.optim_step(types.AdamParams(learning_rate=1e-4)).result(timeout=timeout) - after = _completion_logprobs(client, datum, timeout) + after = completion_logprobs(client, datum, timeout) assert float(np.mean(after - before)) > 1e-4 y = client.save_weights_for_sampler(name='y').result(timeout=timeout).path diff --git a/tests/server/integration/test_sft_e2e.py b/tests/server/integration/test_sft_e2e.py index 16a014b8c..cff8a2a74 100644 --- a/tests/server/integration/test_sft_e2e.py +++ b/tests/server/integration/test_sft_e2e.py @@ -156,7 +156,7 @@ def test_sft_tinker(): """ from tinker import types from twinkle.dataloader import DataLoader - from twinkle.server.model.tinker_datum import input_feature_to_datum + from twinkle.server.training.execution.tinker_datum import input_feature_to_datum backend = get_backend() log(f'=== test_sft_tinker [backend={backend}] ===') diff --git a/tests/server/integration/test_tinker_rl_correctness_e2e.py b/tests/server/integration/test_tinker_rl_correctness_e2e.py index 925841dbe..5f3fc5a66 100644 --- a/tests/server/integration/test_tinker_rl_correctness_e2e.py +++ b/tests/server/integration/test_tinker_rl_correctness_e2e.py @@ -14,6 +14,8 @@ import numpy as np import pytest +from tests.server.integration.e2e_helpers import completion_logprobs, sampled_datum + pytestmark = pytest.mark.skipif( os.environ.get('TWINKLE_TEST_GPU_E2E', '0') != '1', reason='Requires a running training + vLLM server on real accelerators', @@ -83,50 +85,20 @@ def live_clients(): return service, model_id, prompt_tokens, samples[0], timeout -def _datum(prompt_tokens, sample, advantage): - from tinker import types - - tokens = prompt_tokens + list(sample.tokens) - prompt_positions = len(prompt_tokens) - 1 - size = len(tokens) - 1 - - def tensor(values, dtype): - return types.TensorData(data=values, dtype=dtype, shape=[size]) - - return types.Datum( - model_input=types.ModelInput.from_ints(tokens[:-1]), - loss_fn_inputs={ - 'target_tokens': tensor(tokens[1:], 'int64'), - 'weights': tensor([0.0] * prompt_positions + [1.0] * len(sample.tokens), 'float32'), - 'logprobs': tensor([0.0] * prompt_positions + list(sample.logprobs), 'float32'), - 'advantages': tensor([0.0] * prompt_positions + [advantage] * len(sample.tokens), 'float32'), - }, - ) - - -def _completion_logprobs(client, datum, timeout): - result = client.forward([datum], 'cross_entropy').result(timeout=timeout) - values = np.asarray(result.loss_fn_outputs[0]['logprobs'].tolist(), dtype=np.float64) - weights = np.asarray(datum.loss_fn_inputs['weights'].tolist()) - assert len(values) >= len(weights) - values = values[:len(weights)][weights != 0] - assert values.size and np.isfinite(values).all() - return values - def test_live_unsupported_losses_are_rejected_without_changing_policy(live_clients): from tinker import BadRequestError service, model_id, prompt, sample, timeout = live_clients client = service.create_lora_training_client(base_model=model_id, rank=16, seed=2026) - datum = _datum(prompt, sample, -1.0) - before = _completion_logprobs(client, datum, timeout) + datum = sampled_datum(prompt, sample, -1.0) + before = completion_logprobs(client, datum, timeout) for loss_fn in ('ppo', 'cispo', 'dro'): with pytest.raises(BadRequestError) as exc: client.forward_backward([datum], loss_fn).result(timeout=timeout) assert exc.value.status_code == 400 assert f"Unsupported Tinker loss_fn '{loss_fn}'" in str(exc.value) - after = _completion_logprobs(client, datum, timeout) + after = completion_logprobs(client, datum, timeout) np.testing.assert_allclose(after, before, rtol=0, atol=1e-5) _report(check='unsupported_loss', loss_fn=loss_fn, http_status=400, policy_unchanged=True) @@ -138,15 +110,15 @@ def test_live_importance_sampling_updates_policy_in_reward_direction(live_client service, model_id, prompt, sample, timeout = live_clients client = service.create_lora_training_client(base_model=model_id, rank=16, seed=2026) - datum = _datum(prompt, sample, advantage) - before = _completion_logprobs(client, datum, timeout) + datum = sampled_datum(prompt, sample, advantage) + before = completion_logprobs(client, datum, timeout) sampler_difference = float(np.mean(np.abs(before - np.asarray(sample.logprobs)))) assert sampler_difference < 0.05, f'Sampler/trainer logprob mismatch: {sampler_difference}' result = client.forward_backward([datum] * batch_size, 'importance_sampling').result(timeout=timeout) assert len(result.loss_fn_outputs) == batch_size assert np.isfinite(float(result.metrics['loss:mean'])) client.optim_step(types.AdamParams(learning_rate=1e-4)).result(timeout=timeout) - after = _completion_logprobs(client, datum, timeout) + after = completion_logprobs(client, datum, timeout) delta = float(np.mean(after) - np.mean(before)) _report( check='reward_direction', @@ -166,7 +138,7 @@ def test_live_dpo_still_requires_pairs(live_clients): service, model_id, prompt, sample, timeout = live_clients client = service.create_lora_training_client(base_model=model_id, rank=16, seed=2026) - datum = _datum(prompt, sample, 1.0) + datum = sampled_datum(prompt, sample, 1.0) reference = client.forward([datum], 'cross_entropy').result(timeout=timeout) datum.loss_fn_inputs['ref_logps'] = reference.loss_fn_outputs[0]['logprobs'] with pytest.raises(UnprocessableEntityError) as exc: diff --git a/tests/server/lifecycle/test_envelope.py b/tests/server/lifecycle/test_envelope.py index fdc588bd8..1dd368a94 100644 --- a/tests/server/lifecycle/test_envelope.py +++ b/tests/server/lifecycle/test_envelope.py @@ -12,12 +12,18 @@ def test_wire_maps_cover_exactly_the_canonical_reason_codes(): """The Twinkle and Tinker failure maps must both stay in lockstep with the canonical reason-code set, so a newly added domain reason cannot silently fall back to 500 on one protocol face.""" - from twinkle.server.gateway.tinker_handlers import _TINKER_FAILURE_WIRE - from twinkle.server.lifecycle.envelope import _FAILURE_WIRE - from twinkle.server.state.models import FAILURE_REASON_CODES - - assert set(_FAILURE_WIRE) == FAILURE_REASON_CODES - assert set(_TINKER_FAILURE_WIRE) == FAILURE_REASON_CODES + from twinkle.server.gateway.tinker_handlers import _tinker_error_from_failure + from twinkle.server.lifecycle.envelope import error_payload_from_failure + from twinkle.server.task_errors import FailureResponse + + for reason in FailureResponse.reasons(): + policy = FailureResponse.for_reason(reason) + failure = {'reason_code': reason, 'message': 'failed', 'attribution': policy.category.value} + native = error_payload_from_failure(failure, request_id='request') + assert native.error_code == policy.error_code + assert native.category is policy.category + assert _tinker_error_from_failure(failure, request_id='request') == native.model_dump( + mode='json', exclude_none=True) def test_completed_with_none_result_is_a_success_not_a_failure(): diff --git a/tests/server/lifecycle/test_run_submit_dedup.py b/tests/server/lifecycle/test_run_submit_dedup.py index 53d1b6925..b5f7fe710 100644 --- a/tests/server/lifecycle/test_run_submit_dedup.py +++ b/tests/server/lifecycle/test_run_submit_dedup.py @@ -23,24 +23,25 @@ def __init__(self, record_after_claim): self.claim_ttls = [] self.released = [] - async def claim_seq(self, dedup_key, request_id, ttl): - # Unseen -> claim it and let the caller proceed (returns None). - self.claimed[dedup_key] = request_id - self.claim_ttls.append(ttl) - return None + async def claim_training_seq(self, training_id, stream, seq_id, request_id, retention): + key = (training_id, stream, seq_id) + prior = self.claimed.get(key) + self.claimed[key] = prior or request_id + self.claim_ttls.append(retention) + return prior async def get_future(self, request_id): # Simulates whether a PENDING/QUEUED record was written (task enqueued). return self._record_after_claim - async def release_seq(self, dedup_key): - self.released.append(dedup_key) + async def release_training_seq(self, training_id, stream, seq_id, request_id): + self.released.append((training_id, stream, seq_id)) class _FakeManagement: def __init__(self, record_after_claim): self.state = _FakeState(record_after_claim) - self.task_queue_config = SimpleNamespace(queue_timeout=300.0, effective_execution_timeout=1800.0) + self.sequence_retention_s = 3600 # A real deployment declares its backend; preflight reads it from here. self.backend = 'transformers' self.data_world_size = 1 @@ -63,7 +64,7 @@ def _body(adapter_name: str = 'ad', seq_id: int = 7) -> ForwardBackwardTaskReque ``run_submit`` now reads field roles off the body to build the backend kwargs and to run preflight, so a ``SimpleNamespace`` would exercise a shape production never sees. """ - return ForwardBackwardTaskRequest(inputs=[{'input_ids': [1, 2]}], adapter_name=adapter_name, seq_id=seq_id) + return ForwardBackwardTaskRequest(inputs=[{'input_ids': [1, 2]}], adapter_name=adapter_name, seq_id=seq_id, training_id='training-' + adapter_name) async def _call(self, body, adapter_name, token): # pragma: no cover - never invoked @@ -77,7 +78,7 @@ async def test_release_kept_when_task_already_enqueued(): with pytest.raises(RuntimeError): await run_submit(mgmt, _request(), _body(), task_type='forward_backward', backend_call=_call) assert mgmt.state.released == [], 'claim wrongly released for an already-enqueued task' - assert 'seq::sess-1::sess-1-ad::forward_backward::7' in mgmt.state.claimed + assert ('training-ad', 'native::sess-1::forward_backward', 7) in mgmt.state.claimed @pytest.mark.asyncio @@ -86,19 +87,16 @@ async def test_release_when_never_enqueued(): mgmt = _FakeManagement(record_after_claim=None) with pytest.raises(RuntimeError): await run_submit(mgmt, _request(), _body(), task_type='forward_backward', backend_call=_call) - assert mgmt.state.released == ['seq::sess-1::sess-1-ad::forward_backward::7'], 'claim should be released when nothing enqueued' + assert mgmt.state.released == [('training-ad', 'native::sess-1::forward_backward', 7)], 'claim should be released when nothing enqueued' @pytest.mark.asyncio -async def test_dedup_ttl_covers_queue_and_execution(): +async def test_dedup_uses_training_sequence_retention(): """The claim must outlive queue wait plus execution, not just execution.""" mgmt = _FakeManagement(record_after_claim={'status': 'queued'}) with pytest.raises(RuntimeError): await run_submit(mgmt, _request(), _body(), task_type='forward_backward', backend_call=_call) - config = mgmt.task_queue_config - assert mgmt.state.claim_ttls == [int(config.queue_timeout + config.effective_execution_timeout) + 60] - # An execution-only window is exactly what let a queued first copy expire. - assert mgmt.state.claim_ttls[0] > int(config.effective_execution_timeout) + 60 + assert mgmt.state.claim_ttls == [mgmt.sequence_retention_s] @pytest.mark.asyncio @@ -121,5 +119,5 @@ async def test_dedup_key_is_scoped_per_adapter(): backend_call=_call) claimed = set(mgmt.state.claimed) - assert claimed == {'seq::sess-1::sess-1-lora-A::forward_backward::1', 'seq::sess-1::sess-1-lora-B::forward_backward::1'}, ( + assert claimed == {('training-lora-A', 'native::sess-1::forward_backward', 1), ('training-lora-B', 'native::sess-1::forward_backward', 1)}, ( f'adapters collided on one dedup key: {claimed}') diff --git a/tests/server/lifecycle/test_static_guards.py b/tests/server/lifecycle/test_static_guards.py index 730c4ef7f..6278cd2b8 100644 --- a/tests/server/lifecycle/test_static_guards.py +++ b/tests/server/lifecycle/test_static_guards.py @@ -5,7 +5,7 @@ - ``TaskEnvelope`` has exactly one construction site. - The client-side invariant: the client HTTP timeout is <= 120 and strictly greater than the server Long_Poll_Window. -- The task status set has two independent declarations that must not drift. +- Task status uses one shared protocol declaration. These exist because the spec states most of its guarantees in prose. A prose claim that nothing checks decays into a false claim -- as happened with "a consistency test asserts @@ -71,28 +71,12 @@ def test_task_envelope_has_exactly_one_construction_site(): assert sites, 'expected to find the construction sites inside envelope.py' -def test_server_task_status_enum_matches_client_literal(): - """The two independent declarations of the task status set must not drift. +def test_task_status_is_shared_by_queue_and_protocol(): + from twinkle.protocol.types.lifecycle import TaskStatus, TERMINAL_STATUSES + from twinkle.server.task_queue import TaskStatus as QueueTaskStatus - ``twinkle.protocol.types.lifecycle.TaskStatus`` (a Literal on the wire model) and the - server's ``TaskStatus`` enum are declared separately. ``envelope_from_record`` copies - ``record['status']`` straight into ``TaskEnvelope.status``, so a value the server can - write but the Literal does not list would fail pydantic validation *while serialising - the response* -- i.e. a 500 from retrieve for a task that actually finished. - - The client module's comment claimed such a test existed; it did not. This is it. - """ - from typing import get_args - - from twinkle.server.task_queue.types import TaskStatus as ServerTaskStatus - from twinkle.protocol.types.lifecycle import TERMINAL_STATUSES - from twinkle.protocol.types.lifecycle import TaskStatus as WireTaskStatus - - server_values = {member.value for member in ServerTaskStatus} - wire_values = set(get_args(WireTaskStatus)) - assert server_values == wire_values, (f'task status sets drifted: server-only={server_values - wire_values}, ' - f'wire-only={wire_values - server_values}') - assert TERMINAL_STATUSES <= wire_values, 'TERMINAL_STATUSES must be a subset of the declared statuses' + assert QueueTaskStatus is TaskStatus + assert TERMINAL_STATUSES == {TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED} def test_client_future_layer_is_not_imported_by_the_server(): diff --git a/tests/server/model/test_full_mode_exclusivity.py b/tests/server/model/test_full_mode_exclusivity.py index c7a68fedb..70888e335 100644 --- a/tests/server/model/test_full_mode_exclusivity.py +++ b/tests/server/model/test_full_mode_exclusivity.py @@ -13,7 +13,6 @@ from tests.server.state.fake_backend import FakeBackend from tests.server.training.helpers import request, setup from twinkle.server.config.training import TrainingConfig -from twinkle.server.training.catalog import TrainingCatalog from twinkle.server.training.router import TrainingRouter @@ -21,9 +20,8 @@ async def _setup_full_instance(): """Full-mode slots belong to the controller, not the retired model create route.""" controller, state, runtime, _ = await setup(backend=FakeBackend()) raw = controller.config.model_dump() - raw['models'][0]['modes']['full']['max_replicas'] = 1 + raw['models'][1]['config']['max_replicas'] = 1 controller.config = TrainingConfig.model_validate(raw) - controller.catalog = TrainingCatalog(controller.config) runtime.start_gate.clear() return controller, state, runtime @@ -79,8 +77,8 @@ async def test_concurrent_handler_creates_have_exactly_one_winner(): controller, state, runtime = await _setup_full_instance() router = TrainingRouter.__new__(TrainingRouter) router.state = state - router.catalog = controller.catalog - router.model = SimpleNamespace(name='small') + router.config = controller.config + router.model = SimpleNamespace(name='small-full') router.controller = SimpleNamespace(create=SimpleNamespace(remote=AsyncMock(side_effect=controller.create))) async def create_remote(payload): diff --git a/tests/server/model/test_mock_model.py b/tests/server/model/test_mock_model.py index 5564d4344..8dde9e6ab 100644 --- a/tests/server/model/test_mock_model.py +++ b/tests/server/model/test_mock_model.py @@ -15,8 +15,8 @@ from hypothesis import strategies as st from twinkle.server.exceptions import ConfigError -from twinkle.server.model.app import MODEL_SELECTOR -from twinkle.server.model.backends.mock_model import TwinkleCompatMockModel +from twinkle.server.training.execution.app import MODEL_SELECTOR +from twinkle.server.training.execution.backends.mock_model import TwinkleCompatMockModel _MODEL_BACKENDS = tuple(MODEL_SELECTOR.builders) diff --git a/tests/server/model/test_preflight_loss_inputs.py b/tests/server/model/test_preflight_loss_inputs.py index 6c76c9354..6a6df4946 100644 --- a/tests/server/model/test_preflight_loss_inputs.py +++ b/tests/server/model/test_preflight_loss_inputs.py @@ -7,10 +7,10 @@ import pytest from twinkle.loss.grpo import GRPOLoss -from twinkle.server.model.backends import common -from twinkle.server.model.backends.common import TwinkleCompatModelBase -from twinkle.server.model.backends.megatron_model import (TwinkleCompatFullMegatronModel, TwinkleCompatMegatronModel) -from twinkle.server.model.backends.transformers_model import (TwinkleCompatFullTransformersModel, +from twinkle.server.training.execution.backends import common +from twinkle.server.training.execution.backends.common import TwinkleCompatModelBase +from twinkle.server.training.execution.backends.megatron_model import (TwinkleCompatFullMegatronModel, TwinkleCompatMegatronModel) +from twinkle.server.training.execution.backends.transformers_model import (TwinkleCompatFullTransformersModel, TwinkleCompatTransformersModel) diff --git a/tests/server/model/test_replica_lifecycle.py b/tests/server/model/test_replica_lifecycle.py index 0b1160085..2849f427f 100644 --- a/tests/server/model/test_replica_lifecycle.py +++ b/tests/server/model/test_replica_lifecycle.py @@ -5,7 +5,7 @@ import pytest -from twinkle.server.model.app import MODEL_SELECTOR, ModelManagement +from twinkle.server.training.execution.app import MODEL_SELECTOR, ModelManagement class _CapacityState: @@ -67,14 +67,13 @@ async def test_async_constructor_registers_replica_before_ready() -> None: manager = ModelManagement.__new__(ModelManagement) manager._task_queue_config = SimpleNamespace(effective_execution_timeout=1800.0) - with patch('twinkle.server.model.app.DeviceGroup', return_value=SimpleNamespace(name='group')), \ - patch('twinkle.server.model.app.init_twinkle_runtime', return_value=None), \ + with patch('twinkle.server.training.execution.app.DeviceGroup', return_value=SimpleNamespace(name='group')), \ + patch('twinkle.server.training.execution.app.init_twinkle_runtime', return_value=None), \ patch.object(MODEL_SELECTOR, 'construct', return_value=MagicMock()), \ - patch('twinkle.server.model.app.get_server_state', return_value=state), \ + patch('twinkle.server.training.execution.app.get_server_state', return_value=state), \ patch('twinkle.server.data_plane.DataPlaneProxy', return_value=MagicMock()), \ patch.object(ModelManagement, '_init_task_queue'), \ - patch.object(ModelManagement, '_init_adapter_manager'), \ - patch.object(ModelManagement, 'initialize_managed_runtime', new_callable=AsyncMock), \ + patch.object(ModelManagement, 'initialize_training_runtime', new_callable=AsyncMock), \ patch.object(ModelManagement, 'wait_for_training_workers', new_callable=AsyncMock): await ModelManagement.__init__( manager, @@ -84,7 +83,8 @@ async def test_async_constructor_registers_replica_before_ready() -> None: device_mesh={}, backend='mock', max_loras=3, - training_context={'instance_id': 'replica-1'}, + training_context={'instance_id': 'replica-1', 'epoch': 'epoch', 'namespace': 'test', + 'registry_name': 'test', 'controller_name': 'test'}, placement_group=None, allocation_devices=[], ) diff --git a/tests/server/model/test_tinker_compat_output.py b/tests/server/model/test_tinker_compat_output.py index c6e95ec50..f4d76bbcc 100644 --- a/tests/server/model/test_tinker_compat_output.py +++ b/tests/server/model/test_tinker_compat_output.py @@ -3,8 +3,8 @@ from tinker import types from unittest.mock import Mock -from twinkle.server.model.tinker_datum import extract_rl_features_for_loss -from twinkle.server.model.backends.common import TwinkleCompatModelBase +from twinkle.server.training.execution.tinker_datum import extract_rl_features_for_loss +from twinkle.server.training.execution.backends.common import TwinkleCompatModelBase from twinkle.server.exceptions import RequestRejectedError diff --git a/tests/server/model/test_tinker_handlers.py b/tests/server/model/test_tinker_handlers.py index 9d116ed51..55d4a8e84 100644 --- a/tests/server/model/test_tinker_handlers.py +++ b/tests/server/model/test_tinker_handlers.py @@ -5,12 +5,13 @@ from tinker import types from unittest.mock import ANY, AsyncMock, MagicMock -from twinkle.server.model.app import ModelManagement -from twinkle.server.model.tinker_handlers import _register_model_tinker_routes +from twinkle.server.training.execution.app import ModelManagement +from twinkle.server.training.execution.tinker_handlers import _register_model_tinker_routes from twinkle.server.deployment import twinkle_server_error_handler from twinkle.server.exceptions import BatchSizeError, TwinkleServerError from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.task_queue.mixin import TaskQueueMixin +from twinkle.server.training.execution.runtime import TrainingExecutionState class _DummyManagement: @@ -190,6 +191,7 @@ class _SaveWeightsDummyManagement: is_full_mode = ModelManagement.is_full_mode def __init__(self): + self.training_state = TrainingExecutionState(grad_ready=True) self.model = MagicMock() self.state = MagicMock() self.state.get_model_metadata = AsyncMock(return_value={'base_model': 'test-model'}) @@ -198,8 +200,8 @@ def __init__(self): async def _on_request_start(self, request): return 'token1' - def get_adapter_name(self, adapter_name=None): - return adapter_name + def get_training_state(self, training_id): + return self.training_state resolve_model_adapter_name = ModelManagement.resolve_model_adapter_name @@ -218,7 +220,7 @@ async def call_backend(self, fn, /, *args, **kwargs): ids=['lora', 'full-model', 'backend-failure']) async def test_training_save_records_metadata_only_after_backend_success(monkeypatch, tmp_path, fail_backend, full_mode): from pathlib import Path - import twinkle.server.model.tinker_handlers as handlers + import twinkle.server.training.execution.tinker_handlers as handlers from twinkle.server.checkpoint.tinker import TinkerCheckpointManager, TinkerTrainingRunManager monkeypatch.setattr('twinkle.server.checkpoint.training_run_manager.TWINKLE_DEFAULT_SAVE_DIR', str(tmp_path)) @@ -260,7 +262,7 @@ def save(name, output_dir, **kwargs): @pytest.mark.asyncio @pytest.mark.parametrize('path,output_dir', [('org/model', None), ('twinkle://model1/weights/step', 'local')]) async def test_load_weights_resolves_location_and_preserves_optimizer_flag(monkeypatch, tmp_path, path, output_dir): - import twinkle.server.model.tinker_handlers as handlers + import twinkle.server.training.execution.tinker_handlers as handlers from twinkle.server.checkpoint.tinker import TinkerCheckpointManager, TinkerTrainingRunManager monkeypatch.setattr('twinkle.server.checkpoint.training_run_manager.TWINKLE_DEFAULT_SAVE_DIR', str(tmp_path)) manager = TinkerCheckpointManager('token1', TinkerTrainingRunManager('token1')) @@ -269,7 +271,6 @@ async def test_load_weights_resolves_location_and_preserves_optimizer_flag(monke output_dir = manager.get_save_dir('model1') monkeypatch.setattr(handlers, 'create_checkpoint_manager', lambda *args, **kwargs: manager) management = _SaveWeightsDummyManagement() - management.set_resource_state = MagicMock() app = FastAPI() _register_model_tinker_routes(app, lambda: management) endpoint = next(r.endpoint for r in app.routes if getattr(r, 'path', '') == '/tinker/load_weights') @@ -278,7 +279,7 @@ async def test_load_weights_resolves_location_and_preserves_optimizer_flag(monke management.model.tinker_load.assert_called_once_with( checkpoint_name='step' if output_dir else path, output_dir=output_dir, load_optimizer=True, adapter_name='model1') - management.set_resource_state.assert_called_once_with('model1', 'grad_ready', False) + assert management.training_state.grad_ready is False assert response.path == path and response.type == 'load_weights' @@ -289,7 +290,6 @@ async def test_tinker_loss_metric_survives_sdk_reduction(backward): management = _SaveWeightsDummyManagement() management.data_world_size = 1 - management.set_resource_state = MagicMock() outputs = [{'logprobs': types.TensorData(data=[-0.25, -0.5], dtype='float32', shape=[2])}] management.model.tinker_forward_only.return_value = (outputs, 0.375) management.model.tinker_forward_backward.return_value = (outputs, 0.375) @@ -321,7 +321,7 @@ async def test_tinker_loss_metric_survives_sdk_reduction(backward): async def test_sampler_save_materializes_weights_and_returns_sdk_mode(monkeypatch, tmp_path, name, sequence_id, full_mode): from pathlib import Path from twinkle.server.checkpoint import create_checkpoint_manager - import twinkle.server.model.tinker_handlers as handlers + import twinkle.server.training.execution.tinker_handlers as handlers monkeypatch.setattr('twinkle.server.checkpoint.training_run_manager.TWINKLE_DEFAULT_SAVE_DIR', str(tmp_path)) manager = create_checkpoint_manager('token1', client_type='tinker') manager.training_run_manager._write_info('model1', {'base_model': 'test-model'}) diff --git a/tests/server/model/test_twinkle_async_inputs.py b/tests/server/model/test_twinkle_async_inputs.py index 21cf17dec..8a97bc8c4 100644 --- a/tests/server/model/test_twinkle_async_inputs.py +++ b/tests/server/model/test_twinkle_async_inputs.py @@ -5,8 +5,8 @@ from starlette.requests import Request import twinkle.protocol.types as types -from twinkle.server.model.data_plane_inputs import model_result_rows -from twinkle.server.model.twinkle_handlers import _register_model_twinkle_routes +from twinkle.server.training.execution.data_plane_inputs import model_result_rows +from twinkle.server.training.execution.twinkle_handlers import _register_model_twinkle_routes def test_model_result_rows_keeps_one_output_row_per_sample() -> None: @@ -101,7 +101,7 @@ async def test_forward_backward_resolves_multiple_data_refs_and_field_kwargs() - kwarg_fields={ 'old_logps': 'sampled_logprobs', 'advantages': 'advantage', - }, + }, training_id='training-a', ) await route.endpoint(request, body, management) @@ -109,7 +109,7 @@ async def test_forward_backward_resolves_multiple_data_refs_and_field_kwargs() - assert management.scheduled[-1]['batch_size'] == 8 assert management.scheduled[-1]['data_world_size'] == 2 inputs, adapter_name, forwarded_kwargs = management.model_calls[-1] - assert adapter_name == 'session-adapter' + assert adapter_name == 'training-a' assert [row['input_ids'] for row in inputs] == [[index] for index in range(8)] assert forwarded_kwargs['old_logps'] == [[-0.1]] * 4 + [[-0.2]] * 4 assert forwarded_kwargs['advantages'] == [1.0] * 4 + [-1.0] * 4 @@ -139,13 +139,13 @@ async def test_forward_backward_binds_nested_dpo_ref_logps_without_coercion() -> body = types.DataPlaneForwardRequest( adapter_name='adapter', input_refs=[types.DataRef(ref_id='dpo', size=2, num_tokens=6)], - kwarg_fields={'ref_outputs.logps': 'ref_logps'}, + kwarg_fields={'ref_outputs.logps': 'ref_logps'}, training_id='training-a', ) await route.endpoint(request, body, management) inputs, adapter_name, forwarded_kwargs = management.model_calls[-1] - assert adapter_name == 'session-adapter' + assert adapter_name == 'training-a' assert [row['input_ids'] for row in inputs] == [[1, 2, 3], [1, 4, 5]] assert forwarded_kwargs['ref_outputs']['logps'] == [ [-0.1, -0.2, -0.3], @@ -161,8 +161,8 @@ async def test_save_forwards_backend_options_and_returns_public_path( monkeypatch, tmp_path, is_sampler, save_optimizer, full_mode): from pathlib import Path from twinkle.server.checkpoint.twinkle import TwinkleCheckpointManager, TwinkleTrainingRunManager - from twinkle.server.model.app import ModelManagement - import twinkle.server.model.twinkle_handlers as handlers + from twinkle.server.training.execution.app import ModelManagement + import twinkle.server.training.execution.twinkle_handlers as handlers class Management(_SchedulingManagement): is_full_mode = ModelManagement.is_full_mode @@ -177,7 +177,7 @@ def save(self, name, output_dir, **kwargs): monkeypatch.setattr('twinkle.server.checkpoint.training_run_manager.TWINKLE_DEFAULT_SAVE_DIR', str(tmp_path)) runs = TwinkleTrainingRunManager('token') - runs._write_info('session-adapter', {'base_model': 'test-model', 'is_lora': True}) + runs._write_info('training-a', {'base_model': 'test-model', 'is_lora': True}) manager = TwinkleCheckpointManager('token', runs) monkeypatch.setattr(handlers, 'create_checkpoint_manager', lambda *args, **kwargs: manager) service = Management() @@ -189,21 +189,21 @@ def save(self, name, output_dir, **kwargs): request.state.session_id = 'session' body = types.SaveRequest(adapter_name='adapter', name='step/one', is_sampler=is_sampler, - save_optimizer=save_optimizer, consumed_train_samples=7) + save_optimizer=save_optimizer, consumed_train_samples=7, training_id='training-a') response = await endpoint(request, body, service) weights_type = 'sampler_weights' if is_sampler else 'weights' - directory = manager.get_ckpt_dir('session-adapter', f'{weights_type}/step_one') - assert response.result == {'twinkle_path': f'twinkle://session-adapter/{weights_type}/step_one', + directory = manager.get_ckpt_dir('training-a', f'{weights_type}/step_one') + assert response.result == {'twinkle_path': f'twinkle://training-a/{weights_type}/step_one', 'checkpoint_dir': str(directory)} assert (directory / 'weights').read_text() == '1' - assert service.model_calls == [('step_one', {'adapter_name': '' if full_mode else 'session-adapter', + assert service.model_calls == [('step_one', {'adapter_name': '' if full_mode else 'training-a', 'save_optimizer': save_optimizer, 'consumed_train_samples': 7})] @pytest.mark.asyncio async def test_training_save_failure_does_not_publish_checkpoint(monkeypatch, tmp_path): from twinkle.server.checkpoint.twinkle import TwinkleCheckpointManager, TwinkleTrainingRunManager - import twinkle.server.model.twinkle_handlers as handlers + import twinkle.server.training.execution.twinkle_handlers as handlers class Management(_SchedulingManagement): def save(self, **kwargs): @@ -211,9 +211,9 @@ def save(self, **kwargs): monkeypatch.setattr('twinkle.server.checkpoint.training_run_manager.TWINKLE_DEFAULT_SAVE_DIR', str(tmp_path)) runs = TwinkleTrainingRunManager('token') - runs._write_info('session-adapter', {'base_model': 'test-model'}) + runs._write_info('training-a', {'base_model': 'test-model'}) manager = TwinkleCheckpointManager('token', runs) - before = runs._read_info('session-adapter') + before = runs._read_info('training-a') monkeypatch.setattr(handlers, 'create_checkpoint_manager', lambda *args, **kwargs: manager) service = Management() app = FastAPI() @@ -221,12 +221,12 @@ def save(self, **kwargs): endpoint = next(r.endpoint for r in app.routes if getattr(r, 'path', '') == '/twinkle/save') request = Request({'type': 'http', 'headers': []}) request.state.session_id = 'session' - body = types.SaveRequest(adapter_name='adapter', name='failed') + body = types.SaveRequest(adapter_name='adapter', name='failed', training_id='training-a') with pytest.raises(RuntimeError, match='backend save failed'): await endpoint(request, body, service) - assert manager.list_checkpoints('session-adapter').checkpoints == [] - assert runs._read_info('session-adapter') == before + assert manager.list_checkpoints('training-a').checkpoints == [] + assert runs._read_info('training-a') == before @pytest.mark.asyncio @@ -234,7 +234,7 @@ def save(self, **kwargs): @pytest.mark.parametrize('local', [False, True], ids=['external', 'checkpoint-uri']) async def test_checkpoint_consumers_keep_backend_path_and_options(monkeypatch, tmp_path, operation, local): from unittest.mock import MagicMock - import twinkle.server.model.twinkle_handlers as handlers + import twinkle.server.training.execution.twinkle_handlers as handlers from twinkle.server.checkpoint.twinkle import TwinkleCheckpointManager, TwinkleTrainingRunManager monkeypatch.setattr('twinkle.server.checkpoint.training_run_manager.TWINKLE_DEFAULT_SAVE_DIR', str(tmp_path)) @@ -253,9 +253,9 @@ async def test_checkpoint_consumers_keep_backend_path_and_options(monkeypatch, t request.state.session_id = 'session' if operation == 'load': - body = types.LoadRequest(adapter_name='adapter', name=path, load_optimizer=True, strict=False) + body = types.LoadRequest(adapter_name='adapter', name=path, load_optimizer=True, strict=False, training_id='training-a') elif operation == 'resume_from_checkpoint': - body = types.ResumeFromCheckpointRequest(adapter_name='adapter', name=path, resume_only_model=True) + body = types.ResumeFromCheckpointRequest(adapter_name='adapter', name=path, resume_only_model=True, training_id='training-a') service.model.resume_from_checkpoint.return_value = {'step': 7} else: body = types.UploadToHubRequest(checkpoint_dir=path, hub_model_id='org/output', hub_token='hub-token') @@ -264,10 +264,10 @@ async def test_checkpoint_consumers_keep_backend_path_and_options(monkeypatch, t if operation == 'load': service.model.load.assert_called_once_with( name='step' if local else path, output_dir=manager.get_save_dir('run') if local else None, - load_optimizer=True, adapter_name='session-adapter', token='token', strict=False) + load_optimizer=True, adapter_name='training-a', token='token', strict=False) elif operation == 'resume_from_checkpoint': service.model.resume_from_checkpoint.assert_called_once_with( - directory, resume_only_model=True, adapter_name='session-adapter') + directory, resume_only_model=True, adapter_name='training-a') assert response.result == {'result': {'step': 7}} else: service.model.upload_to_hub.assert_called_once_with( @@ -300,7 +300,7 @@ async def test_forward_backward_preflights_loss_inputs_before_backend() -> None: 'labels': [-100, 5], } for index in range(3)], adapter_name='adapter', - loss_kwargs={'old_logps': [[-0.1], [-0.2], [-0.3]]}, + loss_kwargs={'old_logps': [[-0.1], [-0.2], [-0.3]]}, training_id='training-a', ) await route.endpoint(request, body, management) @@ -309,11 +309,11 @@ async def test_forward_backward_preflights_loss_inputs_before_backend() -> None: preflight_inputs, preflight_adapter, preflight_kwargs = management.preflight_calls[0] # Slice-free: a per-rank slice would carry at most one of the three samples. assert len(preflight_inputs) == 3 - assert preflight_adapter == 'session-adapter' + assert preflight_adapter == 'training-a' assert preflight_kwargs['old_logps'] == [[-0.1], [-0.2], [-0.3]] assert management.model_calls[0][0] == 'preflight' assert management.model_calls[-1][0] != 'preflight' - assert management.model_calls[-1][1] == 'session-adapter' + assert management.model_calls[-1][1] == 'training-a' @pytest.mark.asyncio @@ -330,11 +330,11 @@ async def test_forward_backward_skips_preflight_without_rl_inputs() -> None: 'input_ids': [1, 2], 'labels': [-100, 5], }], - adapter_name='adapter', + adapter_name='adapter', training_id='training-a', ) await route.endpoint(request, body, management) assert management.preflight_calls == [] assert len(management.model_calls) == 1 - assert management.model_calls[0][1] == 'session-adapter' + assert management.model_calls[0][1] == 'training-a' diff --git a/tests/server/resources/test_instances.py b/tests/server/resources/test_instances.py index f915d2dda..3abe66281 100644 --- a/tests/server/resources/test_instances.py +++ b/tests/server/resources/test_instances.py @@ -129,7 +129,7 @@ async def test_serve_handle_lookup_cannot_block_driver_deadlines(monkeypatch, ki from types import SimpleNamespace from ray import serve from twinkle.server.resources.runtime import RayRuntime - from twinkle.server.sampler.driver import SamplerDriver + from twinkle.server.sampling.driver import SamplerDriver from twinkle.server.training.driver import RayInstanceDriver entered, resume, returned = threading.Event(), threading.Event(), threading.Event() @@ -197,3 +197,44 @@ def inventory(): resume.set() await asyncio.gather(operation, return_exceptions=True) await asyncio.to_thread(returned.wait, 3) + + +@pytest.mark.asyncio +@pytest.mark.parametrize('outcome', ['unconfirmed_worker', 'removed', 'removal_failed']) +async def test_worker_death_precedes_group_release_and_removal_timing(monkeypatch, outcome): + from types import SimpleNamespace + from unittest.mock import AsyncMock, Mock + import ray + from ray import serve + from twinkle.server.exceptions import ResourceReleaseError + from twinkle.server.resources.runtime import RayRuntime + + async def ready(): + return None # Worker can still answer: death has not been confirmed. + + handle = SimpleNamespace(__ray_ready__=SimpleNamespace(remote=ready)) + rows = [{'namespace': 'test', 'name': 'owned-worker'}] if outcome == 'unconfirmed_worker' else [] + monkeypatch.setattr(ray.util, 'list_named_actors', lambda **kwargs: rows) + monkeypatch.setattr(ray, 'get_actor', lambda *args, **kwargs: handle) + monkeypatch.setattr(ray, 'kill', lambda *args, **kwargs: None) + monkeypatch.setattr(serve, 'delete', lambda name: None) + runtime = RayRuntime('test', allocation_timeout=1, shutdown_timeout=1) + runtime.pool.remove = AsyncMock(side_effect=RuntimeError('PG removal failed') + if outcome == 'removal_failed' else None) + observe = Mock() + if outcome == 'unconfirmed_worker': + with pytest.raises(ResourceReleaseError) as failure: + await runtime.shutdown('owned', 'app', observe_removal=observe) + assert failure.value.phase == 'worker_death' + runtime.pool.remove.assert_not_awaited() + observe.assert_not_called() + else: + if outcome == 'removal_failed': + with pytest.raises(RuntimeError, match='PG removal failed'): + await runtime.shutdown('owned', 'app', observe_removal=observe) + else: + await runtime.shutdown('owned', 'app', observe_removal=observe) + runtime.pool.remove.assert_awaited_once_with('app') + assert observe.call_count == 1 + duration, status = observe.call_args.args + assert duration >= 0 and status == ('error' if outcome == 'removal_failed' else 'success') diff --git a/tests/server/sampler/test_mock_sampler.py b/tests/server/sampler/test_mock_sampler.py index f39828428..d9ca81066 100644 --- a/tests/server/sampler/test_mock_sampler.py +++ b/tests/server/sampler/test_mock_sampler.py @@ -17,8 +17,8 @@ from twinkle.data_format import InputFeature, SamplingParams from twinkle.server.exceptions import ConfigError -from twinkle.server.sampler.app import SAMPLER_SELECTOR, _construct_sampler_backend -from twinkle.server.sampler.backends.mock_sampler import MockSampler +from twinkle.server.sampling.execution.app import SAMPLER_SELECTOR, _construct_sampler_backend +from twinkle.server.sampling.execution.backends.mock_sampler import MockSampler _SAMPLER_TYPES = tuple(SAMPLER_SELECTOR.builders) @@ -121,7 +121,7 @@ def test_explicit_async_vllm_uses_non_blocking_sampler(monkeypatch) -> None: # in the `async-rl` extra, so skip when it is absent instead of failing a sampler test # for a data-path dependency it does not use. pytest.importorskip('transfer_queue') - from twinkle.server.sampler.backends import vllm_async as module + from twinkle.server.sampling.execution.backends import vllm_async as module captured = {} @@ -268,7 +268,7 @@ def test_server_concurrency_does_not_change_native_async_rl_defaults(monkeypatch import twinkle.infra as infra from twinkle.infra._ray import RayHelper from twinkle.sampler.vllm_sampler.vllm_sampler import SAMPLER_MAX_CONCURRENCY - from twinkle.server.sampler.backends.vllm_async import ServerVLLMSampler + from twinkle.server.sampling.execution.backends.vllm_async import ServerVLLMSampler from twinkle.utils.device_mesh import DeviceMesh from twinkle_agentic.async_rl.vllm_sampler_tq import VLLMSamplerTQ diff --git a/tests/server/sampler/test_pool.py b/tests/server/sampler/test_pool.py index 86eb8c03d..4f177fc99 100644 --- a/tests/server/sampler/test_pool.py +++ b/tests/server/sampler/test_pool.py @@ -7,7 +7,7 @@ import pytest from twinkle.server.config.application_spec import SamplerArgs -from twinkle.server.sampler.controller import SamplerController +from twinkle.server.sampling.controller import SamplerController @pytest.fixture(scope='session', autouse=True) @@ -303,7 +303,7 @@ async def test_replaying_config_has_the_same_startup_deadline(): pool, driver, _ = setup(minimum=0, maximum=1) await pool.initialize() pool._actions.append(request('twinkle/set_template')) - pool.budgets['startup_timeout_s'] = 0.05 + pool.budgets = pool.budgets.model_copy(update={'startup_timeout_s': 0.05}) instance = pool._new_instance() with pytest.raises(TimeoutError): await pool.lifecycle.startups[instance.identity] diff --git a/tests/server/sampler/test_request_runtime.py b/tests/server/sampler/test_request_runtime.py index eac13d725..b940aaabd 100644 --- a/tests/server/sampler/test_request_runtime.py +++ b/tests/server/sampler/test_request_runtime.py @@ -10,8 +10,8 @@ from tests.server.task_queue.helpers import _DummyQueue from twinkle.server.config.task_queue import TaskQueueConfig -from twinkle.server.sampler.concurrency import SamplerConcurrency -from twinkle.server.task_queue.types import ExecutionScope, execution_scope +from twinkle.server.sampling.execution.concurrency import SamplerConcurrency +from twinkle.server.task_queue.types import ExecutionScope @pytest.mark.asyncio @@ -28,11 +28,7 @@ async def backend(): await exit_gate.wait() return 'late result' - token = execution_scope.set(scope) - try: - ref = await queue.schedule_concurrent_task(lambda: guard.run(backend)) - finally: - execution_scope.reset(token) + ref = await queue.schedule_concurrent_task(lambda: guard.run(backend), execution_scope=scope) try: await asyncio.wait_for(entered.wait(), 1) completion = asyncio.create_task(scope.wait()) @@ -104,13 +100,14 @@ def close(): @pytest.mark.asyncio async def test_runtime_does_not_certify_an_unconfirmed_cancel_ack(): - from twinkle.server.sampler.runtime import SamplerRuntime - from twinkle.server.task_queue.types import BackendBusyError + from twinkle.server.sampling.execution.runtime import SamplerRuntime + from twinkle.server.exceptions import BackendBusyError guard = SamplerConcurrency() runtime = SamplerRuntime(SimpleNamespace(sampling_guard=guard)) task = asyncio.create_task(asyncio.sleep(0)) await task - runtime._operations['cancelled'] = {'task': task, 'scope': ExecutionScope()} + from twinkle.server.sampling.execution.runtime import SamplingOperation + runtime._operations['cancelled'] = SamplingOperation(ExecutionScope(), task) guard.invalidate_weights(RuntimeError('cancel ACK preceded actual generation exit')) with pytest.raises(BackendBusyError): await runtime.wait_operation('cancelled') @@ -121,7 +118,7 @@ async def test_runtime_does_not_certify_an_unconfirmed_cancel_ack(): async def test_cold_lora_config_does_not_block_sampler_health(monkeypatch): import json import sys - from twinkle.server.sampler.runtime import SamplerRuntime + from twinkle.server.sampling.execution.runtime import SamplerRuntime entered, release = threading.Event(), threading.Event() calls = [] @@ -176,7 +173,7 @@ async def test_managed_compute_http_cannot_bypass_pool_admission(monkeypatch): import httpx from fastapi import FastAPI from twinkle.server.exceptions import RequestRejectedError - from twinkle.server.sampler import app as module + from twinkle.server.sampling.execution import app as module def scaffold(title, register_routes, **kwargs): app = FastAPI() @@ -206,3 +203,31 @@ def scaffold(title, register_routes, **kwargs): standalone = module.build_sampler_app(**args) async with httpx.AsyncClient(transport=httpx.ASGITransport(app=standalone), base_url='http://test') as client: assert (await client.post('/twinkle/create', json={})).status_code == 200 + + +@pytest.mark.asyncio +async def test_scope_waits_for_jobs_registered_by_an_earlier_job(): + scope = ExecutionScope() + child_started, child_release, parent_release = asyncio.Event(), asyncio.Event(), asyncio.Event() + + async def child(): + child_started.set() + await child_release.wait() + + async def parent(): + await parent_release.wait() + scope.track(asyncio.create_task(child())) + + scope.track(asyncio.create_task(parent())) + waiter = asyncio.create_task(scope.wait()) + try: + await asyncio.sleep(0) + parent_release.set() + await child_started.wait() + with pytest.raises(TimeoutError): + await asyncio.wait_for(asyncio.shield(waiter), 0.05) + child_release.set() + await waiter + finally: + child_release.set() + await asyncio.gather(waiter, return_exceptions=True) diff --git a/tests/server/sampler/test_resolve_sampler_weights.py b/tests/server/sampler/test_resolve_sampler_weights.py index 7bbf9f25d..4f01d9b96 100644 --- a/tests/server/sampler/test_resolve_sampler_weights.py +++ b/tests/server/sampler/test_resolve_sampler_weights.py @@ -4,7 +4,7 @@ import pytest -from twinkle.server.sampler.weights import resolve_sampler_weights +from twinkle.server.sampling.execution.weights import resolve_sampler_weights class _FakeSampler: @@ -17,7 +17,7 @@ class _FakeService: """Minimal stand-in exposing the two attributes the helper touches.""" def __init__(self): - from twinkle.server.sampler.concurrency import SamplerConcurrency + from twinkle.server.sampling.execution.concurrency import SamplerConcurrency self.sampling_guard = SamplerConcurrency() self.sampler = _FakeSampler() self.full_weight_loads: list[str] = [] diff --git a/tests/server/sampler/test_sampling_concurrency.py b/tests/server/sampler/test_sampling_concurrency.py index 88a7e6cb8..781a17556 100644 --- a/tests/server/sampler/test_sampling_concurrency.py +++ b/tests/server/sampler/test_sampling_concurrency.py @@ -16,13 +16,13 @@ from twinkle.data_format import SampledSequence, SampleResponse from twinkle.exceptions import FullWeightLoadError from twinkle.protocol.types import SampleRequest -from twinkle.server.sampler.app import SamplerManagement -from twinkle.server.sampler.concurrency import SamplerConcurrency -from twinkle.server.sampler.tinker_handlers import _register_tinker_sampler_routes -from twinkle.server.sampler.twinkle_handlers import _register_twinkle_sampler_routes +from twinkle.server.sampling.execution.app import SamplerManagement +from twinkle.server.sampling.execution.concurrency import SamplerConcurrency +from twinkle.server.sampling.execution.tinker_handlers import _register_tinker_sampler_routes +from twinkle.server.sampling.execution.twinkle_handlers import _register_twinkle_sampler_routes from twinkle.server.task_queue.backend_gate import BackendGate from twinkle.server.config.task_queue import TaskQueueConfig -from twinkle.server.task_queue.types import BackendBusyError +from twinkle.server.exceptions import BackendBusyError from twinkle.server.task_queue.worker import ComputeWorker @@ -166,7 +166,7 @@ def create(**kwargs): @pytest.mark.asyncio async def test_standalone_checkpoint_load_waits_for_ordinary_sampling(make_service, tmp_path): - from twinkle.server.sampler.weights import resolve_sampler_weights + from twinkle.server.sampling.execution.weights import resolve_sampler_weights checkpoint = tmp_path / '2' checkpoint.mkdir() (checkpoint / 'model.safetensors').touch() @@ -401,7 +401,7 @@ def partial_load(path): @pytest.mark.asyncio async def test_completed_generation_failure_does_not_quarantine_replica(make_service): - from twinkle.server.sampler.twinkle_handlers import _await_generation + from twinkle.server.sampling.execution.twinkle_handlers import _await_generation service = make_service() service.sampler.get_generation_status = lambda _: {'status': 'failed', 'error': 'invalid input'} service.sampler.cancel_generation = lambda _: {'status': 'completed'} diff --git a/tests/server/sampler/test_stream_guarantees.py b/tests/server/sampler/test_stream_guarantees.py index a98d3d36f..54c6e7498 100644 --- a/tests/server/sampler/test_stream_guarantees.py +++ b/tests/server/sampler/test_stream_guarantees.py @@ -9,8 +9,8 @@ import pytest from fastapi import FastAPI -from twinkle.server.sampler.app import SamplerManagement -from twinkle.server.sampler.twinkle_handlers import _await_generation, _register_twinkle_sampler_routes, _stream_queue +from twinkle.server.sampling.execution.app import SamplerManagement +from twinkle.server.sampling.execution.twinkle_handlers import _await_generation, _register_twinkle_sampler_routes, _stream_queue class _BlockingQueue: @@ -89,7 +89,7 @@ async def test_stream_timeout_returns_error_payload_and_closes_queue(): class _GenerationService: def __init__(self) -> None: - from twinkle.server.sampler.concurrency import SamplerConcurrency + from twinkle.server.sampling.execution.concurrency import SamplerConcurrency self.sampling_guard = SamplerConcurrency() self.cancelled = False self.sampler = SimpleNamespace( @@ -113,6 +113,6 @@ async def test_generation_poll_has_total_timeout_and_cancels(): await asyncio.wait_for(_await_generation(service, 'submission'), timeout=0.05) assert service.cancelled is True - from twinkle.server.task_queue.types import BackendBusyError + from twinkle.server.exceptions import BackendBusyError with pytest.raises(BackendBusyError, match='restarting'): service.sampling_guard.check_available() diff --git a/tests/server/sampler/test_streaming_bridge.py b/tests/server/sampler/test_streaming_bridge.py new file mode 100644 index 000000000..96cc3c900 --- /dev/null +++ b/tests/server/sampler/test_streaming_bridge.py @@ -0,0 +1,45 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Consumer disconnection never shortens backend ownership.""" +from types import SimpleNamespace + +import pytest + +from twinkle.server.sampling.execution.backends.streaming import STREAM_SENTINEL, stream_to_queue + + +@pytest.mark.parametrize('disconnect', [False, True]) +def test_stream_drains_and_closes_after_consumer_disconnect(disconnect): + consumed, closed, published = [], [], [] + + def generate(*args): + try: + for index in range(3): + consumed.append(index) + yield index, None + finally: + closed.append(True) + + def put(value): + if disconnect: + raise RuntimeError('queue actor is gone') + published.append(value) + + stream_to_queue(SimpleNamespace(sample_stream=generate), SimpleNamespace(put=put), []) + assert consumed == [0, 1, 2] and closed == [True] + assert published == ([] if disconnect else [(0, None), (1, None), (2, None), STREAM_SENTINEL]) + + +def test_backend_failure_is_reported_and_iterator_closed(): + closed, published = [], [] + + def generate(*args): + try: + yield 1, None + raise ValueError('generation failed') + finally: + closed.append(True) + + stream_to_queue(SimpleNamespace(sample_stream=generate), SimpleNamespace(put=published.append), []) + assert closed == [True] + assert isinstance(published[-2], ValueError) + assert published[-1] == STREAM_SENTINEL diff --git a/tests/server/sampler/test_tinker_handlers.py b/tests/server/sampler/test_tinker_handlers.py index 2f4a0aad8..956b085d2 100644 --- a/tests/server/sampler/test_tinker_handlers.py +++ b/tests/server/sampler/test_tinker_handlers.py @@ -5,8 +5,8 @@ from tinker import types from twinkle.data_format import SampledSequence, SampleResponse -from twinkle.server.sampler.tinker_handlers import _register_tinker_sampler_routes -from twinkle.server.sampler.concurrency import SamplerConcurrency +from twinkle.server.sampling.execution.tinker_handlers import _register_tinker_sampler_routes +from twinkle.server.sampling.execution.concurrency import SamplerConcurrency class _DummyState: diff --git a/tests/server/sampler/test_twinkle_async_rows.py b/tests/server/sampler/test_twinkle_async_rows.py index aa1877474..a323e64f9 100644 --- a/tests/server/sampler/test_twinkle_async_rows.py +++ b/tests/server/sampler/test_twinkle_async_rows.py @@ -8,7 +8,7 @@ import twinkle.protocol.types as types from twinkle.data_format import SampledSequence, SampleResponse -from twinkle.server.sampler.twinkle_handlers import ( +from twinkle.server.sampling.execution.twinkle_handlers import ( _register_twinkle_sampler_routes, _build_rollout_rows_and_tags, ) @@ -64,7 +64,7 @@ def test_async_sampler_rejects_group_id_count_mismatch() -> None: class _SamplerManagement: def __init__(self): - from twinkle.server.sampler.concurrency import SamplerConcurrency + from twinkle.server.sampling.execution.concurrency import SamplerConcurrency self.sampling_guard = SamplerConcurrency() self.sampler = self self.data_plane = self @@ -76,7 +76,7 @@ def __init__(self): async def _on_request_start(self, _request): return 'token' - async def submit_concurrent_and_peek(self, coro_factory, *, model_id=None, task_type=None): + async def submit_concurrent_and_peek(self, coro_factory, *, model_id=None, task_type=None, execution_scope=None): self.scheduled.append({'model_id': model_id, 'task_type': task_type}) result = await coro_factory() from twinkle.protocol.types.lifecycle import TaskEnvelope diff --git a/tests/server/session_resource/test_contract.py b/tests/server/session_resource/test_contract.py index 477b509cf..843b7d611 100644 --- a/tests/server/session_resource/test_contract.py +++ b/tests/server/session_resource/test_contract.py @@ -7,7 +7,6 @@ import pytest from twinkle.server.exceptions import RequestRejectedError -from twinkle.server.session_resource.adapter import AdapterManagerMixin from twinkle.server.session_resource.base import SessionResourceMixin from twinkle.server.session_resource.processor import ProcessorManagerMixin @@ -39,16 +38,12 @@ class _MissingBaseHook(SessionResourceMixin): pass -class _MissingAdapterHook(AdapterManagerMixin): - pass - - class _MissingProcessorHook(ProcessorManagerMixin): pass def test_all_resource_expiry_hooks_are_abstract() -> None: - for cls in (_MissingBaseHook, _MissingAdapterHook, _MissingProcessorHook): + for cls in (_MissingBaseHook, _MissingProcessorHook): with pytest.raises(TypeError): cls() diff --git a/tests/server/start_e2e_server.py b/tests/server/start_e2e_server.py index 1e98c914a..337e387f2 100644 --- a/tests/server/start_e2e_server.py +++ b/tests/server/start_e2e_server.py @@ -42,8 +42,7 @@ "max_replicas": 1, # Full-parameter training holds the deployment exclusively; use long # timeouts so the tenant is not expired mid-test. - "adapter_config.adapter_timeout": 3600, - "adapter_config.adapter_max_lifetime": 36000, + "idle_timeout_s": 3600, } VARIANTS = { "lora": (TRANSFORMERS_BASE, {}), diff --git a/tests/server/state/test_managers.py b/tests/server/state/test_managers.py index 242dc3508..6c2dfd47b 100644 --- a/tests/server/state/test_managers.py +++ b/tests/server/state/test_managers.py @@ -534,3 +534,20 @@ async def get_all_then_touch(): session_alive = await state.get_session_last_heartbeat('s-race') is not None model_alive = await state.get_model_metadata('m-race') is not None assert session_alive == model_alive + + +@pytest.mark.parametrize('environment_limit', [None, '7']) +def test_state_initializes_without_or_with_a_partial_environment_policy(monkeypatch, environment_limit): + """A missing optional policy must resolve to defaults before field access.""" + import uuid + from tests.server.state.fake_backend import FakeBackend + from twinkle.server.config.application_spec import SERVER_STATE_ENV_KEYS + from twinkle.server.state.server_state import get_server_state + + for key in SERVER_STATE_ENV_KEYS: + monkeypatch.delenv(key, raising=False) + if environment_limit is not None: + monkeypatch.setenv('TWINKLE_SERVER_STATE_PER_TOKEN_MODEL_LIMIT', environment_limit) + state = get_server_state(cache_key='policy-' + uuid.uuid4().hex, backend=FakeBackend()) + assert state.expiration_timeout == 86400 + assert state.cleanup_interval == 3600 diff --git a/tests/server/state/test_redis_read_retry.py b/tests/server/state/test_redis_read_retry.py new file mode 100644 index 000000000..c5cac36e3 --- /dev/null +++ b/tests/server/state/test_redis_read_retry.py @@ -0,0 +1,49 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Read retry is bounded, request-local, and never replays writes.""" +import asyncio +from unittest.mock import AsyncMock + +import pytest +from redis.exceptions import ConnectionError + +from twinkle.server.state.backend.redis_backend import RedisBackend + + +@pytest.mark.asyncio +async def test_stale_read_retries_without_blocking_another_reader(monkeypatch): + backend = RedisBackend('redis://127.0.0.1:1') + entered, resume = asyncio.Event(), asyncio.Event() + attempts = [] + + async def read(key): + attempts.append(key) + if key == 'stale' and attempts.count(key) == 1: + entered.set() + await resume.wait() + raise ConnectionError('stale connection') + return '"' + key + '"' + + backend._client.get = read + monkeypatch.setattr('twinkle.server.state.backend.redis_backend.asyncio.sleep', AsyncMock()) + try: + first = asyncio.create_task(backend.get('stale')) + await entered.wait() + assert await backend.get('healthy') == 'healthy' + resume.set() + assert await first == 'stale' + assert attempts == ['stale', 'healthy', 'stale'] + finally: + resume.set() + await backend.close() + + +@pytest.mark.asyncio +async def test_cancellation_is_not_retried(): + backend = RedisBackend('redis://127.0.0.1:1') + backend._client.get = AsyncMock(side_effect=asyncio.CancelledError()) + try: + with pytest.raises(asyncio.CancelledError): + await backend.get('key') + backend._client.get.assert_awaited_once() + finally: + await backend.close() diff --git a/tests/server/training/test_redis_restart_reads.py b/tests/server/state/test_redis_recovery.py similarity index 100% rename from tests/server/training/test_redis_restart_reads.py rename to tests/server/state/test_redis_recovery.py diff --git a/tests/server/static/backend_call_exemptions.py b/tests/server/static/backend_call_exemptions.py index 881862dd7..385901d3c 100644 --- a/tests/server/static/backend_call_exemptions.py +++ b/tests/server/static/backend_call_exemptions.py @@ -27,6 +27,6 @@ # (module_relpath under src/twinkle/server, innermost enclosing function name) BACKEND_CALL_EXEMPTIONS: frozenset[tuple[str, str]] = frozenset({ - ('sampler/twinkle_handlers.py', '_stream_queue'), - ('sampler/twinkle_handlers.py', '_stream_generation'), + ('sampling/execution/twinkle_handlers.py', '_stream_queue'), + ('sampling/execution/twinkle_handlers.py', '_stream_generation'), }) diff --git a/tests/server/static/test_adapter_name_mapping.py b/tests/server/static/test_adapter_name_mapping.py index 6856f099e..a150fc53f 100644 --- a/tests/server/static/test_adapter_name_mapping.py +++ b/tests/server/static/test_adapter_name_mapping.py @@ -2,14 +2,14 @@ """Static check: model backend calls map the tenant adapter name (F002 / P002). ``run_submit`` hands each ``backend_call`` a *tenant-scoped* adapter name -(``owner_id-``). Before that reaches a backend model method it must be +(``training_id``). Before that reaches a backend model method it must be translated by ``ModelManagement.resolve_model_adapter_name`` — which returns ``''`` in full-parameter mode (the empty-string default optimizer group) and the name unchanged in LoRA mode. The inline forward/backward/step endpoints did this; the ``*_from_data_plane`` endpoints originally passed the raw tenant name, so a full-mode data-plane request drove the wrong optimizer group. -This guard fails if any ``self.call_backend(...)`` in ``model/twinkle_handlers.py`` +This guard fails if any ``self.call_backend(...)`` in ``training/execution/twinkle_handlers.py`` passes ``adapter_name`` as the bare local (i.e. unmapped). Passing the raw name *positionally* (as ``add_adapter_to_model`` does when creating a tenant adapter) is intentionally not a keyword and is therefore not flagged. @@ -21,7 +21,7 @@ import twinkle -_HANDLERS = pathlib.Path(twinkle.__file__).resolve().parent / 'server' / 'model' / 'twinkle_handlers.py' +_HANDLERS = pathlib.Path(twinkle.__file__).resolve().parent / 'server' / 'training' / 'execution' / 'twinkle_handlers.py' def _is_call_backend(func: ast.AST) -> bool: diff --git a/tests/server/static/test_no_direct_backend_call.py b/tests/server/static/test_no_direct_backend_call.py index 89573d371..f48a66326 100644 --- a/tests/server/static/test_no_direct_backend_call.py +++ b/tests/server/static/test_no_direct_backend_call.py @@ -123,7 +123,7 @@ def test_no_direct_backend_call_in_server(): def test_exemptions_are_read_from_shared_file(): - assert ('sampler/twinkle_handlers.py', '_stream_queue') in BACKEND_CALL_EXEMPTIONS + assert ('sampling/execution/twinkle_handlers.py', '_stream_queue') in BACKEND_CALL_EXEMPTIONS def test_checker_detects_indirect_backend_calls(): diff --git a/tests/server/static/test_no_twinkle_http_exception.py b/tests/server/static/test_no_twinkle_http_exception.py index 28a67a392..5092415e5 100644 --- a/tests/server/static/test_no_twinkle_http_exception.py +++ b/tests/server/static/test_no_twinkle_http_exception.py @@ -6,8 +6,8 @@ _SERVER = Path(__file__).resolve().parents[3] / 'src' / 'twinkle' / 'server' _TWINKLE_HANDLERS = ( _SERVER / 'gateway' / 'twinkle_handlers.py', - _SERVER / 'model' / 'twinkle_handlers.py', - _SERVER / 'sampler' / 'twinkle_handlers.py', + _SERVER / 'training' / 'execution' / 'twinkle_handlers.py', + _SERVER / 'sampling' / 'execution' / 'twinkle_handlers.py', _SERVER / 'processor' / 'twinkle_handlers.py', ) diff --git a/tests/server/task_queue/helpers.py b/tests/server/task_queue/helpers.py index 4a29e2e4c..b3e9bb5de 100644 --- a/tests/server/task_queue/helpers.py +++ b/tests/server/task_queue/helpers.py @@ -2,6 +2,7 @@ from twinkle.server.task_queue.mixin import TaskQueueMixin from twinkle.server.task_queue.worker import ComputeWorker + class _DummyState: def __init__(self): @@ -30,14 +31,12 @@ async def get_future(self, request_id): return self._latest.get(request_id) - class _AllowingRateLimiter: async def check_and_record(self, token, input_tokens): return True, None - class _DummyQueue(TaskQueueMixin): def __init__(self): diff --git a/tests/server/task_queue/test_task_queue_mixin.py b/tests/server/task_queue/test_task_queue_mixin.py index cf193d202..4fae0d1bd 100644 --- a/tests/server/task_queue/test_task_queue_mixin.py +++ b/tests/server/task_queue/test_task_queue_mixin.py @@ -5,7 +5,7 @@ from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.task_queue.mixin import TaskQueueMixin -from twinkle.server.task_queue.types import UserTaskError +from twinkle.server.exceptions import RequestRejectedError from twinkle.server.task_queue.worker import ComputeWorker, _exception_summary @@ -285,7 +285,7 @@ async def test_user_task_error_is_stored_as_user_failure(): queue.enable_compute_worker() async def work(): - raise UserTaskError('invalid request') + raise RequestRejectedError('invalid request') try: await queue.schedule_task(work, model_id='model1', token='token1') @@ -421,7 +421,7 @@ async def test_concurrent_task_keeps_queue_failure_mapping(make_concurrent_queue queue = make_concurrent_queue() async def work(): - raise UserTaskError('invalid request') + raise RequestRejectedError('invalid request') await queue.schedule_concurrent_task(work, token='token1') for _ in range(100): @@ -741,3 +741,23 @@ async def operation(): finally: finish.set() await queue.shutdown_task_queue() + + +@pytest.mark.asyncio +async def test_backend_cancellation_is_an_explicit_server_failure(make_concurrent_queue): + from twinkle.server.lifecycle.envelope import envelope_from_record + from twinkle.server.task_queue.types import ExecutionScope + queue = make_concurrent_queue() + scope = ExecutionScope() + + async def work(): + raise asyncio.CancelledError() + + ref = await queue.schedule_concurrent_task(work, token='token1', execution_scope=scope) + await scope.wait() + record = await queue.state.get_future(ref['request_id']) + assert record['status'] == 'failed' + assert record['failure']['reason_code'] == 'backend_cancelled' + envelope = envelope_from_record(ref['request_id'], record) + assert envelope.error.error_code == 503 + assert envelope.error.category == 'server' diff --git a/tests/server/telemetry/test_training_metrics.py b/tests/server/telemetry/test_training_metrics.py index 427db6a86..157f2447a 100644 --- a/tests/server/telemetry/test_training_metrics.py +++ b/tests/server/telemetry/test_training_metrics.py @@ -7,7 +7,7 @@ from twinkle.server.telemetry import metrics from twinkle.server.training.controller import TrainingController -from twinkle.server.telemetry.training_metrics import TrainingInstanceMetrics +from twinkle.server.telemetry.metrics import TrainingInstanceMetrics _ALLOWED_ATTRIBUTES = {'backend', 'mode', 'state', 'status', 'reason', 'device_type', 'resident'} @@ -103,9 +103,9 @@ def test_training_metric_adapters_rebind_after_registry_reset(monkeypatch): @pytest.mark.parametrize('backend', ['transformers', 'megatron']) def test_controller_metrics_clear_old_states_and_released_resources(monkeypatch, backend): registry, meter = build_registry(monkeypatch) - spec = SimpleNamespace(model=SimpleNamespace(backend=backend, device_type='cuda')) + spec = SimpleNamespace(backend=backend, device_type='cuda') controller = TrainingController.__new__(TrainingController) - controller.catalog = SimpleNamespace(get=lambda base_model, mode: spec) + controller.config = SimpleNamespace(get=lambda base_model, mode: spec) controller._metrics = registry.training_metrics() controller._instance_metrics = TrainingInstanceMetrics(controller._metrics) controller.instances = { @@ -148,3 +148,112 @@ def totals(name): controller._report_instance_metrics() for name in (instances_name, devices_name, groups_name): assert all(value == 0 for value in totals(name).values()) + + +@pytest.mark.asyncio +async def test_repeated_failed_cleanup_is_one_quarantine_transition(): + from unittest.mock import MagicMock + from tests.server.state.fake_backend import FakeBackend + from tests.server.training.helpers import create, setup + + controller, state, driver, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'metric-retry') + instruments = MagicMock() + controller._metrics = instruments + driver.fail_shutdown = True + try: + with pytest.raises(RuntimeError, match='quarantined'): + await controller.release(identity, 'owner', 'session') + for _ in range(2): + await controller.sweep() + await controller.wait_for_jobs() + assert instruments.quarantined_total.inc.call_count == 1 + assert instruments.shutdown_seconds.observe.call_count == 3 + finally: + driver.fail_shutdown = False + await controller.shutdown() + + +@pytest.mark.parametrize('phase', ['worker_death', 'pg_removal']) +def test_quarantine_classification_uses_typed_phase(phase): + from twinkle.server.exceptions import ResourceReleaseError + assert TrainingInstanceMetrics.quarantine_reason(ResourceReleaseError('unconfirmed', phase=phase)) == phase + assert TrainingInstanceMetrics.quarantine_reason(RuntimeError('user text: worker death')) == 'other' + + +@pytest.mark.asyncio +@pytest.mark.parametrize('fail_start', [False, True]) +async def test_startup_and_shutdown_histograms_follow_real_attempts(fail_start): + from unittest.mock import MagicMock + from tests.server.state.fake_backend import FakeBackend + from tests.server.training.helpers import setup, request + + controller, state, driver, _ = await setup(backend=FakeBackend()) + controller._metrics = MagicMock() + driver.fail_start = fail_start + try: + ref = await controller.create(request('timed-start')) + await controller.wait_for_jobs() + future = await state.get_future(ref['request_id']) + assert future['status'] == ('failed' if fail_start else 'completed') + call = controller._metrics.startup_seconds.observe.call_args + assert call.args[0] >= 0 + assert call.kwargs['tags'] == {'backend': 'mock', 'mode': 'lora', + 'status': 'error' if fail_start else 'success'} + if not fail_start: + await controller.release(future['result']['training_id'], 'owner', 'session') + call = controller._metrics.shutdown_seconds.observe.call_args + assert call.args[0] >= 0 and call.kwargs['tags']['status'] == 'success' + finally: + await controller.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('failed', [False, True]) +async def test_sweep_instruments_one_iteration_and_its_outcome(monkeypatch, failed): + import asyncio + from unittest.mock import AsyncMock, MagicMock + from twinkle.server.training.actor import TrainingControllerActor + + async def stop_after_iteration(delay): + raise asyncio.CancelledError + + actor = SimpleNamespace(core=SimpleNamespace(sweep=AsyncMock( + side_effect=RuntimeError('state unavailable') if failed else None)), + config=SimpleNamespace(sweep_interval_s=5), _metrics=MagicMock()) + monkeypatch.setattr('twinkle.server.training.actor.asyncio.sleep', stop_after_iteration) + with pytest.raises(asyncio.CancelledError): + await TrainingControllerActor.__ray_metadata__.modified_class._sweep(actor) + tags = {'status': 'error' if failed else 'success'} + actor._metrics.sweep_total.inc.assert_called_once_with(tags=tags) + call = actor._metrics.sweep_seconds.observe.call_args + assert call.args[0] >= 0 and call.kwargs['tags'] == tags + + +@pytest.mark.asyncio +@pytest.mark.parametrize('failed', [False, True]) +async def test_driver_records_allocation_attempts(failed): + from unittest.mock import AsyncMock, MagicMock + from twinkle.server.config.training import TrainingConfig + from twinkle.server.resources.allocation import DeviceAllocation + from twinkle.server.training.driver import RayInstanceDriver + + config = TrainingConfig(models=[{'name': 'tiny', 'model_id': 'tiny', 'backend': 'mock', + 'device_type': 'cpu', 'mode': 'full', 'config': {'num_devices': 1, 'device_mesh': {'device_type': 'cpu'}}}]) + store = SimpleNamespace(models=SimpleNamespace(unregister_replica_capacity=AsyncMock())) + driver = RayInstanceDriver(config, store, 'test', 'controller') + driver._metrics = MagicMock() + allocation = DeviceAllocation('owned', 'pg', 'node', ('0',)) + driver.runtime.pool.allocate = AsyncMock(return_value=allocation, + side_effect=RuntimeError('allocation failed') if failed else None) + instance = SimpleNamespace(instance_id='instance', controller_epoch='epoch', application_name='owned', + base_model='tiny', mode='full') + if failed: + with pytest.raises(RuntimeError, match='allocation failed'): + await driver.allocate(instance, config.models[0]) + else: + assert (await driver.allocate(instance, config.models[0]))['devices'] == ['0'] + call = driver._metrics.allocation_seconds.observe.call_args + assert call.args[0] >= 0 + assert call.kwargs['tags'] == {'backend': 'mock', 'mode': 'full', 'device_type': 'cpu', + 'status': 'error' if failed else 'success'} diff --git a/tests/server/test_app_builders_characterization.py b/tests/server/test_app_builders_characterization.py index 55c42b4e6..a1f5fb3a9 100644 --- a/tests/server/test_app_builders_characterization.py +++ b/tests/server/test_app_builders_characterization.py @@ -133,7 +133,7 @@ def test_gateway_builder_characterization(monkeypatch) -> None: def test_model_builder_characterization(monkeypatch) -> None: - from twinkle.server.model import app as model_mod + from twinkle.server.training.execution import app as model_mod res = _capture_builder( monkeypatch, @@ -152,7 +152,7 @@ def test_model_builder_characterization(monkeypatch) -> None: def test_sampler_builder_characterization(monkeypatch) -> None: - from twinkle.server.sampler import app as sampler_mod + from twinkle.server.sampling.execution import app as sampler_mod res = _capture_builder( monkeypatch, diff --git a/tests/server/test_dependencies.py b/tests/server/test_dependencies.py index 127f5a7d6..575b4dd98 100644 --- a/tests/server/test_dependencies.py +++ b/tests/server/test_dependencies.py @@ -15,13 +15,13 @@ frozenset({'exceptions'}), frozenset({'config', 'middleware', 'runtime', 'task_errors', 'utils', 'validation'}), frozenset({'telemetry'}), - frozenset({'env_propagation', 'service_proxy', 'state'}), + frozenset({'state'}), frozenset({'resources'}), frozenset({'deployment', 'lifecycle', 'session_resource'}), frozenset({'task_queue'}), frozenset({'checkpoint', 'data_plane'}), - frozenset({'gateway', 'model', 'processor'}), - frozenset({'sampler', 'training'}), + frozenset({'gateway', 'training.execution', 'sampling.execution', 'processor'}), + frozenset({'sampling', 'training'}), frozenset({'launcher'}), frozenset({'cli'}), frozenset({'__main__', _ROOT_COMPONENT}), @@ -31,7 +31,11 @@ def _component(module: str) -> str: if module == _ROOT: return _ROOT_COMPONENT - return module.removeprefix(f'{_ROOT}.').split('.', 1)[0] + relative = module.removeprefix(f'{_ROOT}.') + for execution in ('training.execution', 'sampling.execution'): + if relative == execution or relative.startswith(execution + '.'): + return execution + return relative.split('.', 1)[0] @pytest.mark.parametrize('statement', [ @@ -100,7 +104,7 @@ def test_server_components_form_declared_dag(): f'invalid edges={violations}') -@pytest.mark.parametrize('component', ['training', 'sampler']) +@pytest.mark.parametrize('component', ['training', 'sampling']) def test_controller_import_does_not_load_ray_composition(component): code = f'''import twinkle.server.{component}.controller import sys diff --git a/tests/server/test_service_proxy.py b/tests/server/test_service_proxy.py index e217bedfc..dc9eb4359 100644 --- a/tests/server/test_service_proxy.py +++ b/tests/server/test_service_proxy.py @@ -11,7 +11,7 @@ from starlette.requests import Request from unittest.mock import AsyncMock -from twinkle.server.service_proxy import ServiceProxy +from twinkle.server.lifecycle.proxy import ServiceProxy from twinkle.protocol.headers import H_MULTIPLEX, H_MULTIPLEX_LEGACY, H_REQUEST_ID diff --git a/tests/server/training/helpers.py b/tests/server/training/helpers.py index 916d47f45..e55d27c14 100644 --- a/tests/server/training/helpers.py +++ b/tests/server/training/helpers.py @@ -1,16 +1,19 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from __future__ import annotations + import asyncio +from types import SimpleNamespace + +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.config.training import TrainingConfig +from twinkle.server.resources.allocation import CapacityUnavailableError from twinkle.server.state import ServerState from twinkle.server.state.backend.memory_backend import RayActorBackend -from twinkle.server.training.controller import TrainingController -from twinkle.server.resources.allocation import CapacityUnavailableError from twinkle.server.state.training_records import CreateTraining -from types import SimpleNamespace -from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.task_queue import TaskQueueMixin -from twinkle.server.model.managed_runtime import ManagedTrainingRuntime +from twinkle.server.training.controller import TrainingController +from twinkle.server.training.execution.runtime import TrainingRuntime + class Clock: now = 0.0 @@ -49,7 +52,11 @@ async def deploy(self, instance, spec): async def create_training(self, instance, training, request): assert instance.instance_id in self.allocations assert training.training_id not in self.users - self.users[training.training_id] = {'instance': instance.instance_id, 'label': training.adapter_label, 'steps': 0} + self.users[training.training_id] = { + 'instance': instance.instance_id, + 'label': training.adapter_label, + 'steps': 0 + } async def release_training(self, instance, training): self.users.pop(training.training_id, None) @@ -69,11 +76,36 @@ async def recover_orphans(self): async def setup(*, devices=16, min_replicas=0, max_replicas=None, quota=100, backend=None): - mode = {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}, 'max_loras': 5, - 'idle_timeout_s': 10, 'min_replicas': min_replicas, 'max_replicas': max_replicas} + mode = { + 'num_devices': 2, + 'device_mesh': { + 'fsdp_size': 2 + }, + 'max_loras': 5, + 'idle_timeout_s': 10, + 'min_replicas': min_replicas, + 'max_replicas': max_replicas + } config = TrainingConfig(models=[{ - 'name': 'small', 'model_id': 'ms://small', 'backend': 'mock', 'device_type': 'cuda', - 'modes': {'lora': mode, 'full': {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}, 'idle_timeout_s': 10}}, + 'name': 'small', + 'model_id': 'ms://small', + 'backend': 'mock', + 'device_type': 'cuda', + 'mode': 'lora', + 'config': mode, + }, { + 'name': 'small-full', + 'model_id': 'ms://small', + 'backend': 'mock', + 'device_type': 'cuda', + 'mode': 'full', + 'config': { + 'num_devices': 2, + 'device_mesh': { + 'fsdp_size': 2 + }, + 'idle_timeout_s': 10 + }, }]) state = ServerState(backend=backend or RayActorBackend(), per_token_model_limit=quota) await state.create_session({'session_id': 'session'}) @@ -84,8 +116,15 @@ async def setup(*, devices=16, min_replicas=0, max_replicas=None, quota=100, bac def request(key, *, mode='lora', label='same', token='owner'): - return CreateTraining(base_model='small', mode=mode, token=token, session_id='session', create_key=str(key), - client_type='twinkle', payload={}, adapter_label=label) + return CreateTraining( + base_model='small-full' if mode == 'full' else 'small', + mode=mode, + token=token, + session_id='session', + create_key=str(key), + client_type='twinkle', + payload={}, + adapter_label=label) async def create(controller, state, key, **kwargs): @@ -97,6 +136,7 @@ async def create(controller, state, key, **kwargs): class RPC: + def __init__(self, method): self.method = method @@ -104,17 +144,23 @@ def remote(self, *args, **kwargs): return asyncio.create_task(self.method(*args, **kwargs)) -class RuntimeQueue(ManagedTrainingRuntime, TaskQueueMixin): +class RuntimeQueue(TrainingRuntime, TaskQueueMixin): + def __init__(self, controller, state, instance): self.state, self.replica_id = state, instance.instance_id - self.training_context = {'instance_id': instance.instance_id, 'epoch': controller.epoch} + from twinkle.server.state.training_records import TrainingRuntimeContext + self.training_context = TrainingRuntimeContext( + instance_id=instance.instance_id, + epoch=controller.epoch, + namespace='test', + registry_name='test', + controller_name='test') self._model_unhealthy = False - self._training_controller = SimpleNamespace(**{ - name: RPC(getattr(controller, name)) - for name in ('admit', 'finish_operation', 'fail_instance') - }) - self._init_task_queue(TaskQueueConfig(execution_timeout=0.1, inline_fast_path_timeout=0.01), - enable_admission_gate=True) + self._training_controller = SimpleNamespace( + **{name: RPC(getattr(controller, name)) + for name in ('admit', 'finish_operation', 'fail_instance')}) + self._init_task_queue( + TaskQueueConfig(execution_timeout=0.1, inline_fast_path_timeout=0.01), enable_admission_gate=True) async def settle(state, request_id): diff --git a/tests/server/training/test_actor_initialization.py b/tests/server/training/test_actor_initialization.py index 0c7c97eeb..c5a81b43e 100644 --- a/tests/server/training/test_actor_initialization.py +++ b/tests/server/training/test_actor_initialization.py @@ -81,7 +81,7 @@ async def touch(identity): actor = SimpleNamespace( store=SimpleNamespace(backend=SimpleNamespace(get=AsyncMock(return_value='active')), - key=lambda *args: 'active', state=SimpleNamespace(touch_replica_last_seen=touch)), + key=lambda *args: 'active', models=SimpleNamespace(touch_replica_last_seen=touch)), core=SimpleNamespace(epoch='controller', instances={'bad': bad, 'good': good}, fail_instance=fail, lifecycle=InstanceLifecycle('Training')), driver=SimpleNamespace(health=health), _unreachable_since={}, _failures={}) diff --git a/tests/server/training/test_controller.py b/tests/server/training/test_controller.py index 434d8e2eb..df9e4fe4b 100644 --- a/tests/server/training/test_controller.py +++ b/tests/server/training/test_controller.py @@ -12,7 +12,7 @@ from twinkle.server.state import ServerState from twinkle.server.state.backend.memory_backend import RayActorBackend from twinkle.server.training.controller import TrainingController -from twinkle.server.training.errors import TrainingExpiredError +from twinkle.server.exceptions import TrainingExpiredError from twinkle.server.resources.allocation import CapacityUnavailableError from twinkle.server.state.training_records import CreateTraining @@ -43,10 +43,8 @@ async def test_late_initializer_keeps_its_slot_reserved_until_serial_cleanup(): active = await create(controller, state, 'active') raw = controller.config.model_dump() raw['startup_timeout_s'] = 0.03 - raw['models'][0]['modes']['lora']['max_loras'] = 2 + raw['models'][0]['config']['max_loras'] = 2 controller.config = TrainingConfig.model_validate(raw) - from twinkle.server.training.catalog import TrainingCatalog - controller.catalog = TrainingCatalog(controller.config) entered, unblock = asyncio.Event(), asyncio.Event() original_create, original_release = runtime.create_training, runtime.release_training initializer = None @@ -105,13 +103,11 @@ async def initialize(instance, training, req): @pytest.mark.asyncio async def test_resident_rebuild_does_not_block_expiry_or_another_mode(): from tests.server.state.fake_backend import FakeBackend - from twinkle.server.training.catalog import TrainingCatalog controller, state, runtime, clock = await setup(backend=FakeBackend()) expired = await create(controller, state, 'expired-full', mode='full') raw = controller.config.model_dump() - raw['models'][0]['modes']['lora']['min_replicas'] = 1 + raw['models'][0]['config']['min_replicas'] = 1 controller.config = TrainingConfig.model_validate(raw) - controller.catalog = TrainingCatalog(controller.config) runtime.start_gate.clear() await asyncio.wait_for(controller.sweep(), 0.1) for _ in range(20): @@ -158,41 +154,28 @@ async def finished(instance, request_id): @pytest.mark.asyncio -async def test_legacy_session_and_max_lifetime_bounds_keep_their_original_meaning(): - import time - from datetime import datetime, timezone +async def test_old_creation_time_and_live_heartbeat_do_not_override_business_idle(): from tests.server.state.fake_backend import FakeBackend - from twinkle.server.training.catalog import TrainingCatalog controller, state, runtime, clock = await setup(backend=FakeBackend()) - raw = controller.config.model_dump() - raw['models'][0]['modes']['lora']['idle_timeout_s'] = None - raw['models'][0]['legacy_expiration'] = {'session_timeout_s': 30, 'max_lifetime_s': 36000} - controller.config = TrainingConfig.model_validate(raw) - controller.catalog = TrainingCatalog(controller.config) - identity = await create(controller, state, 'legacy') - clock.now = 1801 + identity = await create(controller, state, 'idle-only') + controller.trainings[identity].created_at = '2000-01-01T00:00:00+00:00' + clock.now = 9 await state.touch_session('session') await controller.sweep() - await controller.wait_for_jobs() assert controller.trainings[identity].state == 'ACTIVE' - controller.trainings[identity].created_at = datetime.fromtimestamp(time.time() - 36001, timezone.utc).isoformat() + clock.now = 10 + await state.touch_session('session') await controller.sweep() await controller.wait_for_jobs() assert controller.trainings[identity].state == 'EXPIRED' - second = await create(controller, state, 'stale-heartbeat') - session = await state._backend.get('session::session') - session['last_heartbeat'] = time.time() - 31 - await state._backend.set('session::session', session) - await controller.sweep() - await controller.wait_for_jobs() - assert controller.trainings[second].state == 'EXPIRED' + assert not runtime.allocations @pytest.mark.asyncio async def test_lazy_queries_allocate_nothing(): controller, state, runtime, _ = await setup() assert runtime.starts == 0 - assert controller.catalog.get('small', 'full').capacity == 1 + assert controller.config.get('small-full', 'full').capacity == 1 with pytest.raises(ResourceNotFoundError): await controller.route('unknown', 'owner') assert not runtime.allocations @@ -472,11 +455,9 @@ async def test_info_does_not_renew_and_expiry_wins_before_next_sweep(): async def test_full_resident_cannot_accept_user_while_resetting(): controller, state, runtime, _ = await setup() original = controller.config.model_dump() - original['models'][0]['modes']['full']['min_replicas'] = 1 - original['models'][0]['modes']['full']['max_replicas'] = 1 + original['models'][1]['config']['min_replicas'] = 1 + original['models'][1]['config']['max_replicas'] = 1 controller.config = TrainingConfig.model_validate(original) - from twinkle.server.training.catalog import TrainingCatalog - controller.catalog = TrainingCatalog(controller.config) await controller.sweep() await controller.wait_for_jobs() identity = await create(controller, state, 1, mode='full') @@ -543,7 +524,7 @@ async def late(instance, training, request): async def test_failed_resident_startup_rolls_back_previously_loaded_residents(): controller, state, runtime, clock = await setup(devices=2) raw = controller.config.model_dump() - raw['models'][0]['modes']['lora']['min_replicas'] = 2 + raw['models'][0]['config']['min_replicas'] = 2 replacement = TrainingController(TrainingConfig.model_validate(raw), state.create_training_store('unit'), runtime, clock=clock) with pytest.raises(CapacityUnavailableError): await replacement.initialize() @@ -616,13 +597,13 @@ async def cancel_during_save(record): monkeypatch.setattr(controller.store, 'save_training', cancel_during_save) else: - start = state.try_start_future + start = state.futures.try_start async def lose_ack(request_id): assert await start(request_id) raise ConnectionError('start acknowledgement lost') - monkeypatch.setattr(state, 'try_start_future', lose_ack) + monkeypatch.setattr(state.futures, 'try_start', lose_ack) try: ref = await controller.create(request(fault)) await controller.wait_for_jobs() @@ -687,10 +668,8 @@ async def test_last_lora_release_does_not_report_unconfirmed_device_release(): async def test_two_device_lora_capacity_shares_two_instances_and_releases_them(): from tests.server.state.fake_backend import FakeBackend from twinkle.server.config import ServerConfig - from twinkle.server.training.catalog import TrainingCatalog controller, state, runtime, _ = await setup(backend=FakeBackend(), devices=4, quota=3) controller.config = ServerConfig.from_yaml('cookbook/client/server/megatron/server_config.yaml').training_config - controller.catalog = TrainingCatalog(controller.config) name = controller.config.models[0].name try: requests = [request(index, token=f'owner-{index // 3}').model_copy(update={'base_model': name}) diff --git a/tests/server/training/test_failure_recovery.py b/tests/server/training/test_failure_recovery.py index f27b9c955..62da32968 100644 --- a/tests/server/training/test_failure_recovery.py +++ b/tests/server/training/test_failure_recovery.py @@ -11,7 +11,7 @@ from tests.server.training.helpers import RPC, RuntimeQueue, settle from twinkle.protocol.types.errors import ErrorCategory from twinkle.server.exceptions import ResourceExhaustedError, ResourceNotFoundError -from twinkle.server.state.models import FAILURE_REASON_CODES, FutureFailureRecord +from twinkle.server.state.models import FutureFailureRecord async def wait_idle(queue): @@ -41,14 +41,14 @@ async def operation(): await settle(state, warmup['request_id']) await wait_idle(queue) executed.clear() - previous = state.store_future_status + previous = state.futures.store_status async def broken(request_id, status, *args, **kwargs): if status == 'queued': raise RuntimeError('injected queued write failure') return await previous(request_id, status, *args, **kwargs) - monkeypatch.setattr(state, 'store_future_status', broken) + monkeypatch.setattr(state.futures, 'store_status', broken) with pytest.raises(RuntimeError, match='queued write'): await queue.schedule_task(operation, model_id=identity, token='owner', task_type='step', request_id='queued-failure') @@ -227,13 +227,12 @@ async def test_cleanup_error_is_not_classified_as_startup_failure(): @pytest.mark.asyncio -async def test_unready_service_is_503_server_and_reason_sets_have_no_dead_codes(): +async def test_unready_service_is_503_server(): controller, _, _, _ = await setup(backend=FakeBackend()) controller._ready = False with pytest.raises(ResourceExhaustedError) as error: await controller.create(request('unready')) assert error.value.error_code == 503 and error.value.category is ErrorCategory.Server - assert not {'protocol_mismatch', 'training_expired'} & FAILURE_REASON_CODES await controller.shutdown() @@ -242,14 +241,14 @@ async def test_sweep_heartbeat_io_does_not_hold_route_lock(monkeypatch): controller, state, _, _ = await setup(backend=FakeBackend()) identity = await create(controller, state, 'slow-heartbeat') entered, resume = asyncio.Event(), asyncio.Event() - original = state.get_session_last_heartbeats + original = state.sessions.get_last_heartbeats async def slow(identities): entered.set() await resume.wait() return await original(identities) - monkeypatch.setattr(state, 'get_session_last_heartbeats', slow) + monkeypatch.setattr(state.sessions, 'get_last_heartbeats', slow) sweeping = asyncio.create_task(controller.sweep()) try: await entered.wait() @@ -263,19 +262,18 @@ async def slow(identities): @pytest.mark.asyncio async def test_explicit_tinker_request_id_is_preserved(): - from twinkle.server.model.managed_runtime import _tinker_operation + from twinkle.server.lifecycle.submit import submit_training_task controller, state, _, _ = await setup(backend=FakeBackend()) identity = await create(controller, state, 'explicit-id') instance = controller.instances[controller.trainings[identity].instance_id] queue = RuntimeQueue(controller, state, instance) - context = _tinker_operation.set((identity, 10)) try: - ref = await queue.schedule_task(AsyncMock(return_value='ok'), model_id=identity, token='owner', + ref = await submit_training_task(queue, AsyncMock(return_value='ok'), training_id=identity, seq_id=10, + stream='tinker::step', token='owner', task_type='step', request_id='explicit-id') assert ref['request_id'] == 'explicit-id' assert (await settle(state, 'explicit-id'))['status'] == 'completed' finally: - _tinker_operation.reset(context) await queue.shutdown_task_queue() await controller.shutdown() @@ -289,11 +287,8 @@ async def test_shutdown_waits_for_late_serve_deployment_before_confirming_releas from twinkle.server.state.training_records import InstanceRecord from twinkle.server.training.driver import RayInstanceDriver from ray import serve - config = TrainingConfig(models=[{ - 'name': 'small', 'model_id': 'weights', 'backend': 'mock', 'device_type': 'cpu', - 'modes': {'lora': {'num_devices': 1, 'device_mesh': {'dp_size': 1}}} - }]) - store = SimpleNamespace(state=SimpleNamespace(unregister_replica_capacity=AsyncMock())) + config = TrainingConfig(models=[{'name': 'small', 'model_id': 'weights', 'backend': 'mock', 'device_type': 'cpu', 'mode': 'lora', 'config': {'num_devices': 1, 'device_mesh': {'dp_size': 1}}}]) + store = SimpleNamespace(models=SimpleNamespace(unregister_replica_capacity=AsyncMock())) driver = RayInstanceDriver(config, store, 'test-late', 'controller') instance = InstanceRecord(instance_id='instance', controller_epoch='epoch', base_model='small', mode='lora', startup_deadline=0) @@ -305,13 +300,13 @@ def deploy(*args, **kwargs): if not release.wait(10): raise TimeoutError('test did not release Serve') - monkeypatch.setattr('twinkle.server.model.app.build_model_app', lambda **kwargs: 'app') + monkeypatch.setattr('twinkle.server.training.execution.app.build_model_app', lambda **kwargs: 'app') monkeypatch.setattr(importlib.import_module('ray.util.placement_group'), 'get_placement_group', lambda name: 'pg') monkeypatch.setattr('ray.util.list_named_actors', lambda **kwargs: []) monkeypatch.setattr(serve, 'run', deploy) monkeypatch.setattr(serve, 'delete', lambda name: None) monkeypatch.setattr(driver.runtime.pool, 'remove', remove) - deploying = asyncio.create_task(driver.deploy(instance, driver._catalog.get('small', 'lora'))) + deploying = asyncio.create_task(driver.deploy(instance, driver.config.get('small', 'lora'))) closing = None try: assert await asyncio.to_thread(entered.wait, 2) @@ -321,15 +316,15 @@ def deploy(*args, **kwargs): await asyncio.sleep(0.01) assert not closing.done() remove.assert_not_awaited() - store.state.unregister_replica_capacity.assert_not_awaited() + store.models.unregister_replica_capacity.assert_not_awaited() finally: release.set() await asyncio.gather(deploying, return_exceptions=True) if closing is not None: await closing await driver.runtime.wait_for_cleanup() - assert not driver.runtime.deployments and not driver.runtime.cleanup_tasks - store.state.unregister_replica_capacity.assert_awaited() + assert not driver.runtime.deployments and not driver.runtime.lifecycle.jobs + store.models.unregister_replica_capacity.assert_awaited() @pytest.mark.asyncio @@ -338,7 +333,7 @@ async def test_capacity_unregister_does_not_retire_managed_model(): identity = await create(controller, state, 'capacity-only') instance = controller.instances[controller.trainings[identity].instance_id] await state.register_replica(instance.instance_id, 5) - await state.unregister_replica_capacity(instance.instance_id) + await state.models.unregister_replica_capacity(instance.instance_id) assert await state.get_model_metadata(identity) is not None assert await state._backend.get('token_count::owner') == 1 await controller.release(identity, 'owner') @@ -352,13 +347,11 @@ async def test_capacity_unregister_does_not_retire_managed_model(): @pytest.mark.parametrize('failure_stage', ['stopped', 'quarantined', 'persist_draining']) async def test_late_release_ack_cannot_retire_a_failed_instance_user(mode, failure_stage, monkeypatch): from twinkle.server.config.training import TrainingConfig - from twinkle.server.training.catalog import TrainingCatalog controller, state, runtime, _ = await setup(backend=FakeBackend()) config = controller.config.model_dump() - config['models'][0]['modes'][mode].update(min_replicas=1, max_replicas=1) + next(model['config'] for model in config['models'] if model['mode'] == mode).update(min_replicas=1, max_replicas=1) controller.config = TrainingConfig.model_validate(config) - controller.catalog = TrainingCatalog(controller.config) await controller.sweep() await controller.wait_for_jobs() identity = await create(controller, state, 'late-release', mode=mode) diff --git a/tests/server/training/test_http_integration.py b/tests/server/training/test_http_integration.py index b20f14b29..61b0bb54e 100644 --- a/tests/server/training/test_http_integration.py +++ b/tests/server/training/test_http_integration.py @@ -26,6 +26,8 @@ def test_installed_native_and_tinker_sdks_use_managed_http_routes(): assert result.returncode == 0, result.stdout + result.stderr evidence = next(json.loads(line) for line in result.stdout.splitlines() if line.startswith('{"passed":')) assert evidence['passed'] and evidence['native_recreation'] and evidence['tinker_shared'] + assert len(set(evidence['native_stable_routing'].values())) == 2 + assert len(set(evidence['tinker_stable_routing'].values())) == 2 assert evidence['native_lora'] assert evidence['runtime_lifecycle_rejected'] assert evidence['tinker_full'] diff --git a/tests/server/training/test_launcher.py b/tests/server/training/test_launcher.py index cd8ee7613..289fd9b48 100644 --- a/tests/server/training/test_launcher.py +++ b/tests/server/training/test_launcher.py @@ -37,3 +37,47 @@ def test_failed_initialization_does_not_leave_an_actor_that_blocks_a_retry(): assert diagnostic['ready'] and diagnostic['instances'] == [] finally: launcher._shutdown() + + +def test_cleanup_attempts_all_owners_and_reports_the_original_failure(): + calls = [] + error = ConnectionError('shutdown ACK unknown') + + def failed(): + calls.append('failed') + raise error + + def healthy(): + calls.append('healthy') + + with pytest.raises(RuntimeError, match='cleanup is incomplete') as caught: + ServerLauncher._run_cleanup([('first', failed), ('second', healthy)], raise_errors=True) + assert calls == ['failed', 'healthy'] + assert caught.value.__cause__ is error + + +def test_startup_rollback_preserves_startup_error_and_cleans_all_owners(monkeypatch): + config = ServerConfig.model_validate({'applications': [ + {'name': 'gateway', 'import_path': 'server', 'route_prefix': '/api/v1'}]}) + launcher = ServerLauncher(config) + calls = [] + error = ValueError('invalid application') + monkeypatch.setattr('twinkle.server.launcher.server_launcher.apply_ray_serve_patches', lambda: None) + monkeypatch.setattr(launcher, '_configure_environment', lambda: None) + monkeypatch.setattr(launcher, '_init_ray', lambda: None) + monkeypatch.setattr(launcher, '_start_serve', lambda: None) + + def deploy(app): + raise error + + def failed_cleanup(): + calls.append('training') + raise ConnectionError('cleanup pending') + + monkeypatch.setattr(launcher, '_deploy_application', deploy) + monkeypatch.setattr(launcher, '_stop_training_controller', failed_cleanup) + monkeypatch.setattr('twinkle.server.telemetry.flush_telemetry_safely', lambda: calls.append('telemetry')) + with pytest.raises(ValueError) as caught: + launcher.launch() + assert caught.value is error + assert calls == ['training', 'telemetry'] diff --git a/tests/server/training/test_router.py b/tests/server/training/test_router.py index 1c610447f..da5271f40 100644 --- a/tests/server/training/test_router.py +++ b/tests/server/training/test_router.py @@ -11,7 +11,6 @@ from twinkle.hub import HubOperation from twinkle.server.state import ServerState from twinkle.server.training.router import TrainingRouter -from twinkle.server.training.catalog import TrainingCatalog from twinkle.server.config.training import TrainingConfig @@ -59,10 +58,7 @@ async def test_native_lora_admission_does_not_import_peft(monkeypatch, rank): router.model = SimpleNamespace(name='small') router.state = ServerState(backend=FakeBackend()) await router.state.create_session({'session_id': 'session'}) - router.catalog = TrainingCatalog(TrainingConfig(models=[{ - 'name': 'small', 'model_id': 'weights', 'backend': 'transformers', 'device_type': 'cuda', - 'modes': {'lora': {'num_devices': 1, 'device_mesh': {'dp_size': 1}}}, - }])) + router.config = TrainingConfig(models=[{'name': 'small', 'model_id': 'weights', 'backend': 'transformers', 'device_type': 'cuda', 'mode': 'lora', 'config': {'num_devices': 1, 'device_mesh': {'dp_size': 1}}}]) async def create(payload): await router.state.store_future_status('created', 'pending', None) return {'request_id': 'created'} @@ -88,8 +84,8 @@ def cpu_import(name, *args, **kwargs): def test_router_schemas_match_model_runtime_routes(): from fastapi import FastAPI - from twinkle.server.model.tinker_handlers import _register_model_tinker_routes - from twinkle.server.model.twinkle_handlers import _register_model_twinkle_routes + from twinkle.server.training.execution.tinker_handlers import _register_model_tinker_routes + from twinkle.server.training.execution.twinkle_handlers import _register_model_twinkle_routes from twinkle.server.training.router import NATIVE_REQUESTS, TINKER_REQUESTS app = FastAPI() diff --git a/tests/server/training/test_runtime_initialization.py b/tests/server/training/test_runtime_initialization.py index 1c91f2aad..94b25c0c9 100644 --- a/tests/server/training/test_runtime_initialization.py +++ b/tests/server/training/test_runtime_initialization.py @@ -9,14 +9,14 @@ from tests.server.state.fake_backend import FakeBackend from tests.server.training.helpers import create, setup -from twinkle.server.model.managed_runtime import ManagedTrainingRuntime +from twinkle.server.training.execution.runtime import TrainingRuntime from twinkle.server.state.training_records import CreateTraining @pytest.mark.asyncio async def test_configuration_deserialization_does_not_block_event_loop(monkeypatch): - controller, _, _, _ = await setup(backend=FakeBackend()) - identity = await create(controller, controller.store.state, 'configuration-import', mode='full') + controller, state, _, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'configuration-import', mode='full') training = controller.trainings[identity] entered, resume, finished = threading.Event(), threading.Event(), threading.Event() @@ -26,18 +26,18 @@ def deserialize(config): finished.set() return None - monkeypatch.setattr('twinkle.server.model.managed_runtime.deserialize_object', deserialize) + monkeypatch.setattr('twinkle.server.training.execution.runtime.deserialize_object', deserialize) manager = SimpleNamespace(save=Mock()) - monkeypatch.setattr('twinkle.server.model.managed_runtime.create_training_run_manager', lambda *args, **kw: manager) + monkeypatch.setattr('twinkle.server.training.execution.runtime.create_training_run_manager', lambda *args, **kw: manager) async def serial(factory, task_type): return await factory() - runtime = SimpleNamespace(_managed_serial=serial, resolve_model_adapter_name=lambda _: '', - register_resource=Mock(), set_resource_state=Mock(), is_full_mode=True) + runtime = SimpleNamespace(_run_lifecycle_operation=serial, resolve_model_adapter_name=lambda _: '', + _training_states={}, is_full_mode=True) request = CreateTraining(base_model='small', mode='full', token='owner', session_id='session', create_key='key', client_type='twinkle', payload={'adapter_name': 'same', 'config': None}) - initializing = asyncio.create_task(ManagedTrainingRuntime.initialize_training( + initializing = asyncio.create_task(TrainingRuntime.initialize_training( runtime, training.model_dump(), request.model_dump())) try: assert await asyncio.to_thread(entered.wait, 2) diff --git a/tests/server/training/test_runtime_queue.py b/tests/server/training/test_runtime_queue.py index 664876b86..cf772b711 100644 --- a/tests/server/training/test_runtime_queue.py +++ b/tests/server/training/test_runtime_queue.py @@ -9,7 +9,8 @@ from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.task_queue import TaskQueueMixin -from twinkle.server.model.managed_runtime import ManagedTrainingRuntime, _tinker_operation +from twinkle.server.training.execution.runtime import TrainingRuntime +from twinkle.server.lifecycle.submit import submit_training_task from tests.server.training.helpers import create, setup @@ -56,19 +57,19 @@ async def test_tinker_seq_retry_runs_once_and_different_operations_do_not_collid async def operation(): counter.append('executed') return {'ok': True} - context = _tinker_operation.set((identity, 1)) try: first, retry = await asyncio.gather(*( - queue.schedule_task(operation, model_id=identity, token='owner', task_type='forward_backward') + submit_training_task(queue, operation, training_id=identity, seq_id=1, + stream='tinker::forward_backward', token='owner', task_type='forward_backward') for _ in range(2))) assert first == retry assert (await settle(state, first['request_id']))['status'] == 'completed' - step = await queue.schedule_task(operation, model_id=identity, token='owner', task_type='optim_step') + step = await submit_training_task(queue, operation, training_id=identity, seq_id=1, + stream='tinker::optim_step', token='owner', task_type='optim_step') assert step['request_id'] != first['request_id'] await settle(state, step['request_id']) assert len(counter) == 2 finally: - _tinker_operation.reset(context) await queue.shutdown_task_queue() await controller.shutdown() @@ -124,7 +125,7 @@ def __init__(self): def is_loaded(self): return self.loaded worker = Worker.remote() - runtime = ManagedTrainingRuntime() + runtime = TrainingRuntime() runtime.model = SimpleNamespace(_actors=[worker]) try: await runtime.wait_for_training_workers() diff --git a/tests/server/validation/test_preflight.py b/tests/server/validation/test_preflight.py index 2912b1e8e..f6ffe539c 100644 --- a/tests/server/validation/test_preflight.py +++ b/tests/server/validation/test_preflight.py @@ -28,22 +28,23 @@ def __init__(self, backend: str = 'transformers'): self.futures: dict[str, dict] = {} self.backend_calls: list = [] self.claimed: list = [] + self.sequence_retention_s = 3600 self._task_queue_config = SimpleNamespace(effective_execution_timeout=60.0) self.state = SimpleNamespace( - claim_seq=self._claim_seq, + claim_training_seq=self._claim_seq, get_future=self._get_future, - release_seq=self._release_seq, + release_training_seq=self._release_seq, ) - async def _claim_seq(self, dedup_key, request_id, ttl): - self.claimed.append(dedup_key) + async def _claim_seq(self, training_id, stream, sequence, request_id, retention): + self.claimed.append((training_id, stream, sequence)) return None async def _get_future(self, request_id): return self.futures.get(request_id) - async def _release_seq(self, dedup_key): - self.claimed.remove(dedup_key) + async def _release_seq(self, training_id, stream, sequence, request_id): + self.claimed.remove((training_id, stream, sequence)) async def _on_request_start(self, request): return 'token' @@ -82,7 +83,7 @@ def test_megatron_rejects_the_split_gradient_endpoints_with_501(): for capability in (BackendCapability.Forward, BackendCapability.Backward, BackendCapability.CalculateLoss): with pytest.raises(EndpointUnavailableError) as raised: assert_request_supported( - _Deployment('megatron'), model_types.AdapterRequest(adapter_name='a'), capability=capability) + _Deployment('megatron'), model_types.AdapterRequest(adapter_name='a', training_id='training-a'), capability=capability) assert raised.value.error_code == 501 assert 'forward_backward' in str(raised.value), 'the alternative endpoints must be named' @@ -90,14 +91,14 @@ def test_megatron_rejects_the_split_gradient_endpoints_with_501(): def test_transformers_serves_the_split_gradient_endpoints(): assert_request_supported( _Deployment('transformers'), - model_types.AdapterRequest(adapter_name='a'), + model_types.AdapterRequest(adapter_name='a', training_id='training-a'), capability=BackendCapability.Forward) def test_mock_backend_serves_every_endpoint(): """The mock is a test double; restricting it would only break tests.""" assert_request_supported( - _Deployment('mock'), model_types.AdapterRequest(adapter_name='a'), capability=BackendCapability.Forward) + _Deployment('mock'), model_types.AdapterRequest(adapter_name='a', training_id='training-a'), capability=BackendCapability.Forward) @pytest.mark.asyncio @@ -107,7 +108,7 @@ async def test_capability_rejection_writes_no_future_and_calls_no_backend(): await run_submit( deployment, _request(), - model_types.AdapterRequest(adapter_name='a'), + model_types.AdapterRequest(adapter_name='a', training_id='training-a'), task_type='backward', backend_call=_call, capability=BackendCapability.Backward) @@ -123,7 +124,7 @@ async def test_capability_rejection_writes_no_future_and_calls_no_backend(): def test_a_megatron_only_parameter_is_rejected_on_transformers(): with pytest.raises(RequestRejectedError) as raised: - assert_request_supported(_Deployment('transformers'), model_types.SaveRequest(adapter_name='a', merge_lora=True)) + assert_request_supported(_Deployment('transformers'), model_types.SaveRequest(adapter_name='a', merge_lora=True, training_id='training-a')) assert raised.value.error_code == 422 assert 'merge_lora' in str(raised.value) assert 'megatron' in str(raised.value) @@ -132,7 +133,7 @@ def test_a_megatron_only_parameter_is_rejected_on_transformers(): def test_a_transformers_only_parameter_is_rejected_on_megatron(): with pytest.raises(RequestRejectedError): assert_request_supported( - _Deployment('megatron'), model_types.LoadRequest(adapter_name='a', name='ckpt', strict=True)) + _Deployment('megatron'), model_types.LoadRequest(adapter_name='a', name='ckpt', strict=True, training_id='training-a')) def test_an_unset_backend_only_parameter_is_not_rejected(): @@ -142,8 +143,8 @@ def test_an_unset_backend_only_parameter_is_not_rejected(): backend's own default, it would look sent on every request and the other half of the fleet would reject everything. """ - assert_request_supported(_Deployment('transformers'), model_types.SaveRequest(adapter_name='a')) - assert_request_supported(_Deployment('megatron'), model_types.LoadRequest(adapter_name='a', name='ckpt')) + assert_request_supported(_Deployment('transformers'), model_types.SaveRequest(adapter_name='a', training_id='training-a')) + assert_request_supported(_Deployment('megatron'), model_types.LoadRequest(adapter_name='a', name='ckpt', training_id='training-a')) def test_every_restricted_field_is_optional_with_a_none_default(): @@ -178,14 +179,14 @@ def test_a_real_parameter_read_from_kwargs_is_not_rejected(): assert_request_supported( _Deployment(), model_types.SetProcessorRequest( - processor_cls='InputProcessor', adapter_name='a', init_kwargs={'padding_side': 'right'})) + processor_cls='InputProcessor', adapter_name='a', init_kwargs={'padding_side': 'right'}, training_id='training-a')) def test_an_unrecognised_plugin_argument_is_forwarded(): """Plugins accept ``**kwargs``, so "unknown" cannot mean "wrong".""" assert_request_supported( _Deployment(), - model_types.SetLossRequest(loss_cls='DPOLoss', adapter_name='a', init_kwargs={'my_custom_knob': 1})) + model_types.SetLossRequest(loss_cls='DPOLoss', adapter_name='a', init_kwargs={'my_custom_knob': 1}, training_id='training-a')) def test_no_plugin_download_happens_during_validation(monkeypatch): @@ -195,7 +196,7 @@ def test_no_plugin_download_happens_during_validation(monkeypatch): lambda *a, **k: pytest.fail('validation must not download a plugin')) assert_request_supported( _Deployment(), - model_types.SetLossRequest(loss_cls='ms://someone/MyLoss', adapter_name='a', init_kwargs={'beta': 0.1})) + model_types.SetLossRequest(loss_cls='ms://someone/MyLoss', adapter_name='a', init_kwargs={'beta': 0.1}, training_id='training-a')) # --------------------------------------------------------------------------- # @@ -210,32 +211,32 @@ def test_control_fields_are_never_forwarded_to_the_backend(): into a backend signature. """ body = model_types.ForwardBackwardTaskRequest( - inputs=[{'input_ids': [1, 2]}], adapter_name='a', seq_id=3, task='embedding') + inputs=[{'input_ids': [1, 2]}], adapter_name='a', seq_id=3, task='embedding', training_id='training-a') assert backend_kwargs(body) == {'task': 'embedding'} def test_unset_backend_parameters_are_not_forwarded(): - body = model_types.ForwardRequest(inputs=[{'input_ids': [1]}], adapter_name='a') + body = model_types.ForwardRequest(inputs=[{'input_ids': [1]}], adapter_name='a', training_id='training-a') assert backend_kwargs(body) == {} def test_passthrough_contents_are_flattened(): body = model_types.ForwardBackwardTaskRequest( - inputs=[{'input_ids': [1]}], adapter_name='a', loss_kwargs={'advantages': [0.5]}) + inputs=[{'input_ids': [1]}], adapter_name='a', loss_kwargs={'advantages': [0.5]}, training_id='training-a') assert backend_kwargs(body) == {'advantages': [0.5]} def test_a_passthrough_key_shadowing_a_declared_parameter_is_an_error(): """Silently letting one win would make the effective value depend on merge order.""" body = model_types.ForwardRequest(inputs=[{'input_ids': [1]}], adapter_name='a', task='causal_lm', - loss_kwargs={'task': 'embedding'}) + loss_kwargs={'task': 'embedding'}, training_id='training-a') with pytest.raises(ValueError, match='collides'): backend_kwargs(body) def test_plugin_identifiers_are_not_forwarded_as_kwargs(): """The handler passes ``loss_cls`` positionally; forwarding it too would duplicate it.""" - body = model_types.SetLossRequest(loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': 0.1}) + body = model_types.SetLossRequest(loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': 0.1}, training_id='training-a') assert backend_kwargs(body) == {'beta': 0.1} diff --git a/tests/server/validation/test_request_wire.py b/tests/server/validation/test_request_wire.py index 73759d25d..5e0ab0acf 100644 --- a/tests/server/validation/test_request_wire.py +++ b/tests/server/validation/test_request_wire.py @@ -55,7 +55,7 @@ def _post(client: TestClient, path: str, body: dict): def _valid_forward(**overrides) -> dict: - return {'inputs': [{'input_ids': [1, 2, 3]}], 'adapter_name': 'a', **overrides} + return {'training_id': 'training-a', 'inputs': [{'input_ids': [1, 2, 3]}], 'adapter_name': 'a', **overrides} def test_a_valid_body_reaches_the_handler(client): @@ -140,6 +140,7 @@ def test_a_dynamic_plugin_argument_is_accepted_inside_its_region(client): response = _post(client, '/set_loss', { 'loss_cls': 'DPOLoss', 'adapter_name': 'a', + 'training_id': 'training-a', 'init_kwargs': {'beta': 0.1, 'anything_at_all': [1, 2]}, }) assert response.status_code == 200 @@ -154,7 +155,7 @@ def test_a_non_json_value_fails_in_the_client_before_any_request(): """ from pydantic import ValidationError with pytest.raises(ValidationError) as raised: - model_types.SetLossRequest(loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': object()}) + model_types.SetLossRequest(loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': object()}, training_id='training-a') assert 'init_kwargs' in str(raised.value), 'the error must name the field that holds the bad value' @@ -246,3 +247,12 @@ def test_sampler_add_adapter_accepts_a_dict_config_where_model_rejects_it(): if __name__ == '__main__': raise SystemExit(pytest.main([__file__, '-v'])) + + +def test_missing_training_handle_is_rejected_before_handler(client): + body = _valid_forward() + body.pop('training_id') + response = _post(client, '/forward', body) + assert response.status_code == 422 + assert any(detail['field'] == 'training_id' for detail in response.json()['details']) + assert 'reached_handler' not in response.json() diff --git a/tests/transformers/test_spectral_hybrid_lora.py b/tests/transformers/test_spectral_hybrid_lora.py index 9e9714249..828e8bd8d 100644 --- a/tests/transformers/test_spectral_hybrid_lora.py +++ b/tests/transformers/test_spectral_hybrid_lora.py @@ -507,8 +507,8 @@ def test_hybrid_rejects_client_owned_fft_and_target_parameter_config(): def test_transformers_server_selects_hybrid_model_only_for_hybrid_train_mode(monkeypatch): - from twinkle.server.model import app - from twinkle.server.model.backends import transformers_model + from twinkle.server.training.execution import app + from twinkle.server.training.execution.backends import transformers_model class Regular: @@ -535,7 +535,7 @@ class Hybrid(Regular): def test_hybrid_server_model_exposes_device_mesh_to_remote_class(): from twinkle.infra import _get_device_mesh_param_name - from twinkle.server.model.backends.transformers_model import TwinkleCompatSpectralHybridTransformersModel + from twinkle.server.training.execution.backends.transformers_model import TwinkleCompatSpectralHybridTransformersModel assert _get_device_mesh_param_name(TwinkleCompatSpectralHybridTransformersModel.__init__) == 'device_mesh' diff --git a/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py b/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py index 9567138c4..90bf4e6a1 100644 --- a/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py +++ b/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py @@ -10,7 +10,7 @@ from twinkle import DeviceMesh from twinkle.data_format import SampledSequence, SampleResponse, SamplingParams from twinkle.infra import _dispatch_args -from twinkle.server.sampler.twinkle_handlers import _await_generation +from twinkle.server.sampling.execution.twinkle_handlers import _await_generation from twinkle_agentic.async_rl import LoraContext from twinkle_agentic.async_rl.types import PartitionAdmission, PromptGroup, RolloutPolicy from twinkle_agentic.async_rl.vllm_sampler_tq import (VLLMSamplerTQ, _dispatch_generation, _GeneratedSample, diff --git a/tests/twinkle_client/test_async_components.py b/tests/twinkle_client/test_async_components.py index fff5c1206..30adcb789 100644 --- a/tests/twinkle_client/test_async_components.py +++ b/tests/twinkle_client/test_async_components.py @@ -321,3 +321,31 @@ def fail(*args, **kwargs): monkeypatch.setattr(model, '_submit', original) model.add_adapter_to_model('same', None) assert model._training_ids['same'] == 'training-test' + + +def test_unknown_training_rejected_before_handshake(monkeypatch): + import pytest + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + calls = [] + _patch_transport(monkeypatch, calls, lambda url: _completed({})) + model = MultiLoraTransformersModel('base') + with pytest.raises(RuntimeError, match='add_adapter_to_model'): + model.step() + assert not calls + + +def test_sequence_and_handle_selection_are_object_local(monkeypatch): + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + calls = [] + _patch_transport(monkeypatch, calls, lambda url: _completed({})) + first, second = MultiLoraTransformersModel('base'), MultiLoraTransformersModel('base') + for model in (first, second): + model._created = True + model.adapter_name = 'default' + first._training_ids = {'default': 'first-default', 'other': 'first-other'} + second._training_ids = {'default': 'second-default'} + first.step() + first.step(adapter_name='other') + second.step() + assert [(body['adapter_name'], body['training_id'], body['seq_id']) for _, body in calls] == [ + ('default', 'first-default', 1), ('other', 'first-other', 2), ('default', 'second-default', 1)] diff --git a/tests/twinkle_client/test_auto_agent_tools.py b/tests/twinkle_client/test_auto_agent_tools.py index 41c5216a5..747d463e1 100644 --- a/tests/twinkle_client/test_auto_agent_tools.py +++ b/tests/twinkle_client/test_auto_agent_tools.py @@ -84,10 +84,9 @@ def test_generated_server_config_matches_launcher_schema(tmp_path, monkeypatch, application = next(app for app in config.applications if app.import_path == 'model') assert application.args.train_mode == 'lora' model = config.training_config.models[0] - assert model.modes.full is None - assert model.modes.lora.min_replicas == 0 and model.modes.lora.max_replicas == 1 - assert model.modes.lora.num_devices == 2 and model.modes.lora.idle_timeout_s == 600 - assert model.legacy_expiration.session_timeout_s == 600 + assert model.mode == 'lora' + assert model.config.min_replicas == 0 and model.config.max_replicas == 1 + assert model.config.num_devices == 2 and model.config.idle_timeout_s == 600 if samplers: sampler = next(app for app in config.applications if app.import_path == 'sampler') assert sampler.args.nproc_per_node == 4 diff --git a/tests/twinkle_client/test_request_builder.py b/tests/twinkle_client/test_request_builder.py index 0c7a15bce..f629f108c 100644 --- a/tests/twinkle_client/test_request_builder.py +++ b/tests/twinkle_client/test_request_builder.py @@ -32,7 +32,7 @@ def no_network(monkeypatch): def test_a_declared_name_goes_to_its_field(): - body = build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', task='embedding') + body = build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', task='embedding', training_id='training-a') assert body.task == 'embedding' assert body.loss_kwargs == {} @@ -46,7 +46,7 @@ def test_an_undeclared_name_goes_to_the_single_passthrough_region(): }], adapter_name='a', advantages=[0.5], - old_logps=[[-1.0]]) + old_logps=[[-1.0]], training_id='training-a') assert body.loss_kwargs == {'advantages': [0.5], 'old_logps': [[-1.0]]} @@ -71,7 +71,7 @@ class _TwoRegions(StrictRequest): def test_a_model_without_a_region_rejects_an_unknown_name(): with pytest.raises(TwinkleClientValidationError) as raised: - build_request(model_types.SaveRequest, adapter_name='a', name='ckpt', typo=1) + build_request(model_types.SaveRequest, adapter_name='a', name='ckpt', typo=1, training_id='training-a') assert 'typo' in str(raised.value) @@ -81,7 +81,7 @@ def test_an_explicit_region_and_a_routed_key_are_merged(): loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': 0.1}, - loss_type='sigmoid') + loss_type='sigmoid', training_id='training-a') assert body.init_kwargs == {'beta': 0.1, 'loss_type': 'sigmoid'} @@ -89,7 +89,7 @@ def test_a_key_passed_twice_is_an_error(): """Silently letting one win would make the effective value depend on merge order.""" with pytest.raises(TwinkleClientValidationError, match='both directly and inside'): build_request( - model_types.SetLossRequest, loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': 0.1}, beta=0.2) + model_types.SetLossRequest, loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': 0.1}, beta=0.2, training_id='training-a') def test_an_omitted_optional_argument_is_not_routed_into_the_region(): @@ -98,7 +98,7 @@ def test_an_omitted_optional_argument_is_not_routed_into_the_region(): Routing an explicit ``None`` would hand the backend a null argument it never received before, changing behaviour for callers who simply did not pass anything. """ - body = build_request(model_types.SetLossRequest, loss_cls='DPOLoss', adapter_name='a', unused=None) + body = build_request(model_types.SetLossRequest, loss_cls='DPOLoss', adapter_name='a', unused=None, training_id='training-a') assert body.init_kwargs == {} @@ -109,28 +109,28 @@ def test_an_omitted_optional_argument_is_not_routed_into_the_region(): def test_a_wrongly_typed_field_fails_in_process(): with pytest.raises(ValidationError): - build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', temperature='hot') + build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', temperature='hot', training_id='training-a') def test_an_out_of_range_value_fails_in_process(): with pytest.raises(ValidationError): - build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', temperature=0) + build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', temperature=0, training_id='training-a') def test_malformed_inputs_fail_in_process(): """The client shares the server's schema, so it catches this without asking.""" with pytest.raises(ValidationError): - build_request(model_types.ForwardRequest, inputs=[{'input_ids': [True]}], adapter_name='a') + build_request(model_types.ForwardRequest, inputs=[{'input_ids': [True]}], adapter_name='a', training_id='training-a') def test_a_missing_required_field_fails_in_process(): with pytest.raises(ValidationError): - build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}]) + build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], training_id='training-a') def test_forward_only_requires_an_adapter_context(): with pytest.raises(ValidationError): - build_request(model_types.ForwardOnlyRequest, inputs=[{'input_ids': [1]}]) + build_request(model_types.ForwardOnlyRequest, inputs=[{'input_ids': [1]}], training_id='training-a') # --------------------------------------------------------------------------- # @@ -141,9 +141,9 @@ def test_forward_only_requires_an_adapter_context(): def test_unset_optionals_stay_off_the_wire(): """Absent and "not requested" must look the same, or the server needs its own defaults.""" import json - body = build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a') + body = build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', training_id='training-a') payload = json.loads(request_json(body)) - assert payload == {'inputs': [{'input_ids': [1]}], 'adapter_name': 'a', 'loss_kwargs': {}} + assert payload == {'training_id': 'training-a', 'inputs': [{'input_ids': [1]}], 'adapter_name': 'a', 'loss_kwargs': {}} def test_a_lora_config_is_serialized_to_the_form_the_server_decodes(): diff --git a/tests/twinkle_client/test_types_contract.py b/tests/twinkle_client/test_types_contract.py index 7dbaa25d1..868fd673c 100644 --- a/tests/twinkle_client/test_types_contract.py +++ b/tests/twinkle_client/test_types_contract.py @@ -1,7 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Contract-base consistency and naming-disambiguation tests. -- ``QueueStateLiteral`` value set equals the server ``QueueState`` enum. - naming disambiguation guard. The two SDKs already share public names. The contract freezes that legacy set and @@ -15,9 +14,7 @@ import typing import twinkle -from twinkle.server.task_queue.types import QueueState from twinkle.protocol.types import model as model_types -from twinkle.protocol.types.errors import QueueStateLiteral from twinkle.protocol.types.server import GetServerCapabilitiesResponse _TWINKLE_SRC = pathlib.Path(twinkle.__file__).resolve().parent @@ -43,11 +40,6 @@ }) -def test_queue_state_literal_matches_server_enum(): - literal_values = set(typing.get_args(QueueStateLiteral)) - enum_values = {state.value for state in QueueState} - assert literal_values == enum_values, (f'QueueStateLiteral {literal_values} != QueueState {enum_values}') - def test_old_capabilities_response_gets_conservative_defaults(): response = GetServerCapabilitiesResponse.model_validate({'supported_models': []}) From 7d66859446cd0ce1df58a656f56dd140e4155e5e Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 16:57:25 +0800 Subject: [PATCH 31/46] test(server): remove obsolete fail-fast settings from acceptance scripts --- .../integration/on_demand_training/matrix_validation.py | 3 +-- .../on_demand_training/test_resident_release_e2e.py | 4 ++-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/server/integration/on_demand_training/matrix_validation.py b/tests/server/integration/on_demand_training/matrix_validation.py index 59bfcc9e6..b889d2d4b 100644 --- a/tests/server/integration/on_demand_training/matrix_validation.py +++ b/tests/server/integration/on_demand_training/matrix_validation.py @@ -304,7 +304,7 @@ def main(): app['args'] = { **common, 'model_id': args.model, 'backend': args.backend, 'train_mode': 'lora', 'max_loras': 10, 'max_r': 32, - 'env_vars': {'TWINKLE_FAIL_FAST': '0', 'TWINKLE_TRUST_REMOTE_CODE': '1'}, + 'env_vars': {'TWINKLE_TRUST_REMOTE_CODE': '1'}, } app.pop('deployments', None) full_applications.append({ @@ -332,7 +332,6 @@ def main(): 'TWINKLE_MODEL_ID_ALIASES': json.dumps({model.name: model.model_id for model in config.training_config.models}), 'TWINKLE_TRUST_REMOTE_CODE': '1', - 'TWINKLE_FAIL_FAST': '0', 'TWINKLE_TEST_GPU_E2E': '1', 'TWINKLE_TEST_BACKEND': args.backend, 'TWINKLE_SERVER_URL': f'http://127.0.0.1:{port}', diff --git a/tests/server/integration/on_demand_training/test_resident_release_e2e.py b/tests/server/integration/on_demand_training/test_resident_release_e2e.py index c6c830729..198045607 100644 --- a/tests/server/integration/on_demand_training/test_resident_release_e2e.py +++ b/tests/server/integration/on_demand_training/test_resident_release_e2e.py @@ -94,13 +94,13 @@ def test_resident_release_ack_after_worker_failure(tmp_path): 'pp_size': 1 if backend == 'transformers' else 2}, 'max_length': 512, 'idle_timeout_s': 3600, 'queue_config': {'execution_timeout': 600}, - 'env_vars': {'TWINKLE_TRUST_REMOTE_CODE': '1', 'TWINKLE_FAIL_FAST': '0'}}}, + 'env_vars': {'TWINKLE_TRUST_REMOTE_CODE': '1'}}}, ], } (directory / 'server.json').write_text(json.dumps(raw, indent=2)) os.environ.update(TWINKLE_STATE_RAY_NAMESPACE=namespace, TWINKLE_PERSISTENCE_MODE=persistence, TWINKLE_DEFAULT_SAVE_DIR=str(directory / 'checkpoints'), TWINKLE_TRUST_REMOTE_CODE='1', - TWINKLE_FAIL_FAST='0', TWINKLE_MODEL_ID_ALIASES=json.dumps({model_name: model}), + TWINKLE_MODEL_ID_ALIASES=json.dumps({model_name: model}), PYTHONPATH=os.pathsep.join((str(root / 'src'), str(root)))) if managed: os.environ.update(TWINKLE_PERSISTENCE_REDIS_URL=managed.url, TWINKLE_PERSISTENCE_KEY_PREFIX=prefix) From 2f7f99765fee6642cbdeeb2ea857752b19f6adce Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 17:21:37 +0800 Subject: [PATCH 32/46] fix(model): re-enable FFT execution when activating hybrid slots --- src/twinkle/model/transformers/hybrid/fft_slots.py | 4 ++++ tests/transformers/test_spectral_hybrid_lora.py | 6 +++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/src/twinkle/model/transformers/hybrid/fft_slots.py b/src/twinkle/model/transformers/hybrid/fft_slots.py index 39cde9579..38d675d2b 100644 --- a/src/twinkle/model/transformers/hybrid/fft_slots.py +++ b/src/twinkle/model/transformers/hybrid/fft_slots.py @@ -120,6 +120,10 @@ def activate_fft_slot(self, adapter_name: str) -> None: fft_adapter_name = self._fft_adapter_name(adapter_name) if self.is_hybrid(adapter_name) else None selected_adapters = [fft_adapter_name] if fft_adapter_name is not None else [] for wrapper in self._iter_fft_wrappers(): + # Selecting a regular LoRA can disable auxiliary wrappers. Restoring + # an FFT slot name alone does not restore its execution state. + if fft_adapter_name is not None and wrapper.disable_adapters: + wrapper.enable_adapters(True) if wrapper.active_adapters != selected_adapters: wrapper.set_adapter(selected_adapters) self._enable_preallocated_fft_grad(wrapper) diff --git a/tests/transformers/test_spectral_hybrid_lora.py b/tests/transformers/test_spectral_hybrid_lora.py index 828e8bd8d..a9a7bd62f 100644 --- a/tests/transformers/test_spectral_hybrid_lora.py +++ b/tests/transformers/test_spectral_hybrid_lora.py @@ -871,11 +871,15 @@ def test_disable_lora_restores_slots_after_operation_fails(): assert all(parameter.requires_grad for name, parameter in model.named_parameters() if 'lora_' in name) -def test_multi_tenant_hybrid_isolation_and_non_destructive_merge(): +@pytest.mark.parametrize('fft_disabled', [False, True]) +def test_multi_tenant_hybrid_isolation_and_non_destructive_merge(fft_disabled): base, model, manager, spectral = _make_multi_tenant_hybrid() inputs = torch.randn(2, 8) with manager.adapter('regular'): regular_before = model(inputs).detach().clone() + if fft_disabled: + # A normal LoRA selection can disable auxiliary FFT wrappers in PEFT. + spectral._get_fft_wrapper('layers.0.self_attn.q_proj').enable_adapters(False) hybrid = manager.find_lora_by_tenant('hybrid') modules = dict(model.named_modules()) From ee4a64abea319bb0f25da746031df3430f810580 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 18:07:53 +0800 Subject: [PATCH 33/46] test(server): budget HTTP acceptance for repeated cold workers --- tests/server/training/test_http_integration.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/server/training/test_http_integration.py b/tests/server/training/test_http_integration.py index 61b0bb54e..3d73695e4 100644 --- a/tests/server/training/test_http_integration.py +++ b/tests/server/training/test_http_integration.py @@ -11,11 +11,14 @@ def test_installed_native_and_tinker_sdks_use_managed_http_routes(): import ray ray.shutdown() try: + # Both SDKs, recovery and ten release/recreate cycles start many fresh + # workers. Allow cold imports from a shared environment; each service + # operation still enforces its own startup/execution deadline. result = subprocess.run([ sys.executable, str(Path(__file__).resolve().parents[3] / 'tests/server/integration/on_demand_training/http_probe.py') - ], capture_output=True, text=True, timeout=1800) + ], capture_output=True, text=True, timeout=2 * 60 * 60) finally: ray.init( ignore_reinit_error=True, From 73cc8cf372274cb23c24578080683dcdd1e36cc8 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 20:15:46 +0800 Subject: [PATCH 34/46] fix(server): close startup, release and shutdown gaps --- .../server/launcher/server_launcher.py | 11 ++-- src/twinkle/server/resources/allocation.py | 5 +- src/twinkle/server/sampling/execution/app.py | 6 ++- .../model/multi_lora_transformers.py | 9 +++- tests/server/resources/test_allocation.py | 22 +++++++- tests/server/sampler/test_request_runtime.py | 34 ++++++++++++ tests/server/training/test_launcher.py | 52 +++++++++++++++++++ tests/twinkle_client/test_async_components.py | 39 ++++++++++++++ 8 files changed, 167 insertions(+), 11 deletions(-) diff --git a/src/twinkle/server/launcher/server_launcher.py b/src/twinkle/server/launcher/server_launcher.py index 1268e1adb..0d9b13b91 100644 --- a/src/twinkle/server/launcher/server_launcher.py +++ b/src/twinkle/server/launcher/server_launcher.py @@ -179,11 +179,14 @@ def _deploy_training(self) -> None: self._run_cleanup([('training controller rollback', self._stop_training_controller)], raise_errors=False) raise - @staticmethod - def _stop_controller(controller) -> None: + def _stop_controller(self, controller) -> None: import ray - ray.get(controller.shutdown.remote()) - ray.kill(controller, no_restart=True) + try: + ray.get(controller.shutdown.remote(), timeout=self.config.resources.shutdown_timeout_s) + finally: + # Killing the owner prevents restart but does not certify its devices + # were released. Preserve a failed/unknown shutdown for _run_cleanup. + ray.kill(controller, no_restart=True) def _stop_training_controller(self) -> None: if self._training_controller is not None: diff --git a/src/twinkle/server/resources/allocation.py b/src/twinkle/server/resources/allocation.py index 3b20680a1..1a43d82df 100644 --- a/src/twinkle/server/resources/allocation.py +++ b/src/twinkle/server/resources/allocation.py @@ -50,9 +50,10 @@ def _create_group(resource, count, name): import ray from ray.util.placement_group import placement_group candidates = [node for node in ray.nodes() if node.get('Alive') and node['Resources'].get(resource, 0) >= count] - available = ray.available_resources() - if not candidates or available.get(resource, 0) < count: + if not candidates: raise CapacityUnavailableError(f'No single-node allocation available for {count} {resource} devices') + # The aggregate free-resource snapshot can lag a confirmed group removal. + # pg.ready() is the allocation witness and is bounded by allocate(). cpus = min(_gpu_placement_group_cpus(int(node['Resources'].get('CPU', 0)), count) for node in candidates) bundle = {resource: count, 'CPU': count if resource == 'CPU' else cpus} return placement_group([bundle], strategy='STRICT_PACK', name=name) diff --git a/src/twinkle/server/sampling/execution/app.py b/src/twinkle/server/sampling/execution/app.py index 00aac7956..ee3964e59 100644 --- a/src/twinkle/server/sampling/execution/app.py +++ b/src/twinkle/server/sampling/execution/app.py @@ -113,8 +113,10 @@ async def __init__(self, ) self.sampler_type = sampler_type self.model_id = model_id - replica_context = serve.get_replica_context() - self.replica_id = context['instance_id'] if context is not None else replica_context.replica_id.unique_id + if context is not None: + self.replica_id = context.instance_id + else: + self.replica_id = serve.get_replica_context().replica_id.unique_id sampler_kwargs: dict[str, Any] = { 'model_id': model_id, diff --git a/src/twinkle_client/model/multi_lora_transformers.py b/src/twinkle_client/model/multi_lora_transformers.py index 05981a861..5b2ef2dc5 100644 --- a/src/twinkle_client/model/multi_lora_transformers.py +++ b/src/twinkle_client/model/multi_lora_transformers.py @@ -14,6 +14,7 @@ from twinkle.protocol.types import model as model_types from twinkle.protocol.types.component import DataRef from twinkle_client._request_builder import build_request +from twinkle_client.exceptions import TwinkleHTTPError from twinkle_client.http import ClientTransport from twinkle_client.http.context import capture_transport @@ -188,7 +189,13 @@ def remove_adapter(self, adapter_name: str | None = None) -> None: name = adapter_name or self.adapter_name with self._training_lock: training_id = self._training_ids.get(name) - self._submit('remove_adapter', model_types.AdapterRequest, None, adapter_name=name, training_id=training_id) + try: + self._submit('remove_adapter', model_types.AdapterRequest, None, adapter_name=name, training_id=training_id) + except TwinkleHTTPError as exc: + if exc.status_code != 404: + raise + # An expired or lost handle is unusable. Forget its local binding so + # this label can be recreated; unknown cleanup outcomes still raise. with self._training_lock: if self._training_ids.get(name) == training_id: self._training_ids.pop(name, None) diff --git a/tests/server/resources/test_allocation.py b/tests/server/resources/test_allocation.py index cac7c2e3d..9e89ef403 100644 --- a/tests/server/resources/test_allocation.py +++ b/tests/server/resources/test_allocation.py @@ -9,8 +9,27 @@ from twinkle.server.resources.allocation import DevicePool +def test_stale_resource_summary_does_not_reject_a_schedulable_group(monkeypatch): + import importlib + import ray + placement = importlib.import_module('ray.util.placement_group') + monkeypatch.setattr(ray, 'nodes', lambda: [{'Alive': True, 'Resources': {'CPU': 4, 'GPU': 2}}]) + monkeypatch.setattr(ray, 'available_resources', lambda: {'CPU': 4}) + created = [] + group = object() + + def create(bundles, **kwargs): + created.append((bundles, kwargs)) + return group + + monkeypatch.setattr(placement, 'placement_group', create) + assert DevicePool._create_group('GPU', 2, 'replacement') is group + assert created[0][0][0]['GPU'] == 2 + assert created[0][1] == {'strategy': 'STRICT_PACK', 'name': 'replacement'} + + @pytest.mark.asyncio -@pytest.mark.parametrize('blocked', ['nodes', 'available_resources', 'placement_group']) +@pytest.mark.parametrize('blocked', ['nodes', 'placement_group']) async def test_cancelled_allocation_keeps_late_creation_owned(monkeypatch, blocked): import importlib import ray @@ -44,7 +63,6 @@ def get(name): return pg monkeypatch.setattr(ray, 'nodes', lambda: call('nodes', [{'Alive': True, 'Resources': {'CPU': 2}}])) - monkeypatch.setattr(ray, 'available_resources', lambda: call('available_resources', {'CPU': 2})) monkeypatch.setattr(placement, 'placement_group', create) monkeypatch.setattr(placement, 'get_placement_group', get) monkeypatch.setattr(placement, 'remove_placement_group', lambda group: removed.append(group)) diff --git a/tests/server/sampler/test_request_runtime.py b/tests/server/sampler/test_request_runtime.py index b940aaabd..e45f34aae 100644 --- a/tests/server/sampler/test_request_runtime.py +++ b/tests/server/sampler/test_request_runtime.py @@ -14,6 +14,40 @@ from twinkle.server.task_queue.types import ExecutionScope +@pytest.mark.asyncio +async def test_allocated_sampler_initializes_with_typed_owner_context(monkeypatch): + from unittest.mock import Mock + from twinkle.infra._ray.ray_helper import RayHelper + from twinkle.infra._ray.resource_manager import ResourceManager + from twinkle.server.sampling.execution import app as module + from twinkle.server.state.models import RuntimeContext + + context = RuntimeContext(namespace='owned', registry_name='registry', instance_id='sampler-owned') + allocation = object() + manager = object() + allocate = Mock(return_value=manager) + sampler = SimpleNamespace() + monkeypatch.setenv('TWINKLE_RAY_NAMESPACE', 'before') + monkeypatch.setenv('TWINKLE_RAY_REGISTRY_NAME', 'before') + monkeypatch.setattr(RayHelper, 'resource_manager', None) + monkeypatch.setattr(ResourceManager, 'from_allocation', allocate) + monkeypatch.setattr(module, 'init_twinkle_runtime', lambda **kwargs: None) + monkeypatch.setattr(module, '_construct_sampler_backend', lambda *args: sampler) + monkeypatch.setattr(module.SamplerManagement, '_init_task_queue', lambda *args, **kwargs: None) + monkeypatch.setattr(module, 'get_server_state', lambda: SimpleNamespace()) + monkeypatch.setattr(module.serve, 'get_replica_context', lambda: SimpleNamespace( + replica_id=SimpleNamespace(unique_id='serve-replica'))) + owner = module.SamplerManagement.__new__(module.SamplerManagement) + await owner.__init__( + model_id='mock', nproc_per_node=1, device_group={'name': 'sampler', 'ranks': 1, 'device_type': 'cpu'}, + device_mesh={'dp_size': 1}, sampler_type='mock', sampling_context=context, + placement_group=allocation, allocation_devices=['0']) + assert owner.replica_id == context.instance_id + assert RayHelper.resource_manager is manager + allocate.assert_called_once_with(owner.device_group, allocation, ['0']) + await owner.data_plane.close() + + @pytest.mark.asyncio async def test_worker_exports_actual_completion_handle_without_private_task_lookup(): queue = _DummyQueue() diff --git a/tests/server/training/test_launcher.py b/tests/server/training/test_launcher.py index 289fd9b48..630e47cb9 100644 --- a/tests/server/training/test_launcher.py +++ b/tests/server/training/test_launcher.py @@ -1,5 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import socket +import asyncio +import threading import uuid import pytest @@ -10,6 +12,56 @@ from twinkle.server.training.actor import controller_name +def test_unresponsive_controller_does_not_block_other_cleanup(): + @ray.remote(num_cpus=0) + class Controller: + def __init__(self): + self.entered = asyncio.Event() + self.resume = asyncio.Event() + + async def shutdown(self): + self.entered.set() + await self.resume.wait() + + async def wait_until_entered(self): + await self.entered.wait() + + async def unblock(self): + self.resume.set() + + controller = Controller.remote() + launcher = ServerLauncher(ServerConfig(resources={'shutdown_timeout_s': 0.3})) + launcher._training_controller = controller + finished = threading.Event() + errors, cleaned = [], [] + + def stop(): + try: + launcher._run_cleanup([ + ('blocked controller', launcher._stop_training_controller), + ('other owner', lambda: cleaned.append(True))], raise_errors=True) + except Exception as error: + errors.append(error) + finally: + finished.set() + + worker = threading.Thread(target=stop) + # Ensure actor startup is outside the deadline under test. + ray.get(controller.__ray_ready__.remote(), timeout=10) + entered = controller.wait_until_entered.remote() + worker.start() + try: + ray.get(entered, timeout=10) + assert finished.wait(2), 'Controller shutdown exceeded its budget and blocked the next owner' + assert cleaned == [True] + assert len(errors) == 1 and isinstance(errors[0].__cause__, ray.exceptions.GetTimeoutError) + finally: + if not finished.is_set(): + ray.get(controller.unblock.remote(), timeout=5) + worker.join(5) + ray.kill(controller, no_restart=True) + + def test_failed_initialization_does_not_leave_an_actor_that_blocks_a_retry(): namespace = 'launcher-retry-' + uuid.uuid4().hex with socket.socket() as listener: diff --git a/tests/twinkle_client/test_async_components.py b/tests/twinkle_client/test_async_components.py index 30adcb789..8c6629f81 100644 --- a/tests/twinkle_client/test_async_components.py +++ b/tests/twinkle_client/test_async_components.py @@ -334,6 +334,45 @@ def test_unknown_training_rejected_before_handshake(monkeypatch): assert not calls +def test_expired_handle_can_be_removed_and_recreated_on_same_object(monkeypatch): + from twinkle_client.exceptions import TwinkleHTTPError + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + calls = [] + _patch_transport(monkeypatch, calls, lambda url: _completed(None)) + model = MultiLoraTransformersModel('base') + model.add_adapter_to_model('same', None) + submit = model._submit + + def expired(*args, **kwargs): + raise TwinkleHTTPError('Training handle expired', status_code=404, error_code=404) + + monkeypatch.setattr(model, '_submit', expired) + model.remove_adapter('same') + monkeypatch.setattr(model, '_submit', submit) + model.add_adapter_to_model('same', None) + model.zero_grad() + assert calls[-1][1]['training_id'] == model._training_ids['same'] + + +def test_failed_release_keeps_handle_when_outcome_is_unknown(monkeypatch): + import pytest + from twinkle_client.exceptions import TwinkleHTTPError + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + calls = [] + _patch_transport(monkeypatch, calls, lambda url: _completed(None)) + model = MultiLoraTransformersModel('base') + model.add_adapter_to_model('same', None) + identity = model._training_ids['same'] + for error in (TwinkleHTTPError('Unavailable', status_code=503), TimeoutError('ACK unknown')): + def failed(*args, **kwargs): + raise error + monkeypatch.setattr(model, '_submit', failed) + with pytest.raises(type(error)): + model.remove_adapter('same') + assert model._training_ids['same'] == identity + assert model.adapter_name == 'same' + + def test_sequence_and_handle_selection_are_object_local(monkeypatch): from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel calls = [] From e13e1b77b3e39aaf2839db7a682c58dfa35ce107 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 20:16:25 +0800 Subject: [PATCH 35/46] refactor(server): preserve public datum imports and remove extra docs --- README.md | 2 +- README_ZH.md | 2 +- cookbook/client/tinker/dpo.py | 2 +- cookbook/client/tinker/multi_modal.py | 2 +- cookbook/client/tinker/self_cognition.py | 2 +- .../Usage Guide/Introduction-with-Qwen3.5.md | 2 +- docs/source_en/Usage Guide/Quick-Start.md | 2 +- .../Tinker-Compatible-Client.md | 2 +- .../Usage Guide/Train-as-a-Service.md | 2 +- ...00\344\275\263\345\256\236\350\267\265.md" | 2 +- ...53\351\200\237\345\274\200\345\247\213.md" | 2 +- ...71\345\256\242\346\210\267\347\253\257.md" | 2 +- ...55\347\273\203\346\234\215\345\212\241.md" | 2 +- src/twinkle/server/README.md | 100 ------------------ src/twinkle/server/model/__init__.py | 1 + .../execution => model}/tinker_datum.py | 0 .../server/session_resource/processor.py | 2 +- .../training/execution/backends/common.py | 2 +- .../execution/backends/megatron_model.py | 2 +- .../execution/backends/transformers_model.py | 2 +- .../on_demand_training/acceptance_support.py | 2 +- tests/server/integration/test_dpo_e2e.py | 2 +- .../server/integration/test_full_param_e2e.py | 2 +- tests/server/integration/test_sft_e2e.py | 2 +- .../lifecycle/test_retrieve_endpoint.py | 2 +- .../server/model/test_tinker_compat_output.py | 2 +- tests/server/start_e2e_server.py | 1 - .../task_queue/test_task_queue_mixin.py | 6 -- tests/server/test_dependencies.py | 2 +- 29 files changed, 25 insertions(+), 131 deletions(-) delete mode 100644 src/twinkle/server/README.md create mode 100644 src/twinkle/server/model/__init__.py rename src/twinkle/server/{training/execution => model}/tinker_datum.py (100%) diff --git a/README.md b/README.md index 00bf1f9f2..6f4b05a03 100644 --- a/README.md +++ b/README.md @@ -256,7 +256,7 @@ from twinkle import init_tinker_client from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.training.execution.tinker_datum import input_feature_to_datum +from twinkle.server.model.tinker_datum import input_feature_to_datum base_model = 'ms://Qwen/Qwen3.8-27B' base_url='your-base-url' diff --git a/README_ZH.md b/README_ZH.md index 3317c5921..3d5f05fac 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -245,7 +245,7 @@ from twinkle import init_tinker_client from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.training.execution.tinker_datum import input_feature_to_datum +from twinkle.server.model.tinker_datum import input_feature_to_datum base_model = 'ms://Qwen/Qwen3.8-27B' base_url='your-base-url' diff --git a/cookbook/client/tinker/dpo.py b/cookbook/client/tinker/dpo.py index 5f57be719..3f6e25d05 100644 --- a/cookbook/client/tinker/dpo.py +++ b/cookbook/client/tinker/dpo.py @@ -28,7 +28,7 @@ from twinkle.dataset import Dataset, DatasetMeta, LazyDataset from twinkle.dataloader import DataLoader from twinkle.preprocessor import EmojiDPOProcessor -from twinkle.server.training.execution.tinker_datum import input_feature_to_datum +from twinkle.server.model.tinker_datum import input_feature_to_datum logger = get_logger() diff --git a/cookbook/client/tinker/multi_modal.py b/cookbook/client/tinker/multi_modal.py index 4541669cd..594c4bb8c 100644 --- a/cookbook/client/tinker/multi_modal.py +++ b/cookbook/client/tinker/multi_modal.py @@ -30,7 +30,7 @@ from twinkle.preprocessor import Preprocessor from twinkle.dataset import DatasetMeta, LazyDataset from twinkle.dataloader import DataLoader -from twinkle.server.training.execution.tinker_datum import input_feature_to_datum # Key: converts InputFeature -> Datum +from twinkle.server.model.tinker_datum import input_feature_to_datum # Key: converts InputFeature -> Datum from twinkle import get_logger logger = get_logger() diff --git a/cookbook/client/tinker/self_cognition.py b/cookbook/client/tinker/self_cognition.py index 6b43c12f4..3a1cd75be 100644 --- a/cookbook/client/tinker/self_cognition.py +++ b/cookbook/client/tinker/self_cognition.py @@ -16,7 +16,7 @@ from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.training.execution.tinker_datum import input_feature_to_datum +from twinkle.server.model.tinker_datum import input_feature_to_datum # Initialize the Tinker client before importing ServiceClient init_tinker_client() diff --git a/docs/source_en/Usage Guide/Introduction-with-Qwen3.5.md b/docs/source_en/Usage Guide/Introduction-with-Qwen3.5.md index 731f6d7be..e61ed9197 100644 --- a/docs/source_en/Usage Guide/Introduction-with-Qwen3.5.md +++ b/docs/source_en/Usage Guide/Introduction-with-Qwen3.5.md @@ -458,7 +458,7 @@ from twinkle import init_tinker_client from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.training.execution.tinker_datum import input_feature_to_datum +from twinkle.server.model.tinker_datum import input_feature_to_datum # Initialize Tinker client (must be called before importing ServiceClient) init_tinker_client() diff --git a/docs/source_en/Usage Guide/Quick-Start.md b/docs/source_en/Usage Guide/Quick-Start.md index 5ea0a8500..2efc140ad 100644 --- a/docs/source_en/Usage Guide/Quick-Start.md +++ b/docs/source_en/Usage Guide/Quick-Start.md @@ -761,7 +761,7 @@ from tinker import ServiceClient from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.training.execution.tinker_datum import input_feature_to_datum +from twinkle.server.model.tinker_datum import input_feature_to_datum # The base model to fine-tune / evaluate base_model = 'ms://Qwen/Qwen3.5-4B' diff --git a/docs/source_en/Usage Guide/Server and Client/Tinker-Compatible-Client.md b/docs/source_en/Usage Guide/Server and Client/Tinker-Compatible-Client.md index 9e22f83c2..6e98414b3 100644 --- a/docs/source_en/Usage Guide/Server and Client/Tinker-Compatible-Client.md +++ b/docs/source_en/Usage Guide/Server and Client/Tinker-Compatible-Client.md @@ -53,7 +53,7 @@ from twinkle import init_tinker_client from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.training.execution.tinker_datum import input_feature_to_datum +from twinkle.server.model.tinker_datum import input_feature_to_datum # Step 1: Initialize Tinker client before importing ServiceClient init_tinker_client() diff --git a/docs/source_en/Usage Guide/Train-as-a-Service.md b/docs/source_en/Usage Guide/Train-as-a-Service.md index 1d27635f3..c07e072e1 100644 --- a/docs/source_en/Usage Guide/Train-as-a-Service.md +++ b/docs/source_en/Usage Guide/Train-as-a-Service.md @@ -28,7 +28,7 @@ from twinkle import init_tinker_client from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.training.execution.tinker_datum import input_feature_to_datum +from twinkle.server.model.tinker_datum import input_feature_to_datum base_model = 'ms://Qwen/Qwen3.8-27B' base_url='https://www.modelscope.cn/twinkle' diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Qwen3.5\346\234\200\344\275\263\345\256\236\350\267\265.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Qwen3.5\346\234\200\344\275\263\345\256\236\350\267\265.md" index 88d80c0b6..d6fc45d2b 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Qwen3.5\346\234\200\344\275\263\345\256\236\350\267\265.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Qwen3.5\346\234\200\344\275\263\345\256\236\350\267\265.md" @@ -458,7 +458,7 @@ from twinkle import init_tinker_client from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.training.execution.tinker_datum import input_feature_to_datum +from twinkle.server.model.tinker_datum import input_feature_to_datum # 初始化 Tinker 客户端（必须在导入 ServiceClient 之前） init_tinker_client() diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\345\277\253\351\200\237\345\274\200\345\247\213.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\345\277\253\351\200\237\345\274\200\345\247\213.md" index 579a2a31e..7ea2d838c 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\345\277\253\351\200\237\345\274\200\345\247\213.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\345\277\253\351\200\237\345\274\200\345\247\213.md" @@ -762,7 +762,7 @@ from tinker import ServiceClient from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.training.execution.tinker_datum import input_feature_to_datum +from twinkle.server.model.tinker_datum import input_feature_to_datum # The base model to fine-tune / evaluate base_model = 'Qwen/Qwen3.5-4B' diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/Tinker\345\205\274\345\256\271\345\256\242\346\210\267\347\253\257.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/Tinker\345\205\274\345\256\271\345\256\242\346\210\267\347\253\257.md" index 1c787aee6..ec207e28a 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/Tinker\345\205\274\345\256\271\345\256\242\346\210\267\347\253\257.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/Tinker\345\205\274\345\256\271\345\256\242\346\210\267\347\253\257.md" @@ -53,7 +53,7 @@ from twinkle import init_tinker_client from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.training.execution.tinker_datum import input_feature_to_datum +from twinkle.server.model.tinker_datum import input_feature_to_datum # Step 1: 在导入 ServiceClient 之前，先初始化 Tinker 客户端 init_tinker_client() diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\350\256\255\347\273\203\346\234\215\345\212\241.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\350\256\255\347\273\203\346\234\215\345\212\241.md" index 6a2031e28..fb515944a 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\350\256\255\347\273\203\346\234\215\345\212\241.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\350\256\255\347\273\203\346\234\215\345\212\241.md" @@ -29,7 +29,7 @@ from twinkle import init_tinker_client from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.training.execution.tinker_datum import input_feature_to_datum +from twinkle.server.model.tinker_datum import input_feature_to_datum base_model = 'ms://Qwen/Qwen3.8-27B' base_url='https://www.modelscope.cn/twinkle' diff --git a/src/twinkle/server/README.md b/src/twinkle/server/README.md deleted file mode 100644 index da288c240..000000000 --- a/src/twinkle/server/README.md +++ /dev/null @@ -1,100 +0,0 @@ -# Server responsibilities - -The training and sampling pools share resource mechanisms, but own separate scheduling policies and records. - -```mermaid -flowchart TD - Client --> TrainingRouter - Client --> SamplerRouter - TrainingRouter --> TrainingActor - SamplerRouter --> SamplingActor - Launcher --> TrainingActor - Launcher --> SamplingActor - TrainingActor --> TrainingController - SamplingActor --> SamplerController - TrainingController --> InstanceLifecycle - SamplerController --> InstanceLifecycle - TrainingController --> TrainingDriver - SamplerController --> SamplerDriver - TrainingDriver --> RayRuntime - SamplerDriver --> RayRuntime - RayRuntime --> DevicePool - TrainingDriver --> TrainingExecutionApp - SamplerDriver --> SamplingExecutionApp - TrainingExecutionApp --> TrainingRuntime - SamplingExecutionApp --> SamplerRuntime -``` - -- `training/`: validated model lookup, handle ownership, placement policy, router and Ray actor composition. Its `execution/` subpackage implements backend/protocol adapters and worker-side admission/completion. -- `sampling/`: sampling placement, scaling, control broadcasts, router and Ray actor composition. Its `execution/` subpackage implements backend/protocol adapters and request completion tracking. -- In both domains, the controller makes business decisions, the driver performs domain-specific setup, the actor composes the controller with its driver, and the router handles public requests. Controllers accept injected drivers; importing a controller does not compose Ray actors. -- `resources/`: shared device allocations, Ray/Serve deployment and confirmed shutdown, and background instance supervision. No training or sampling policy imports. -- `state/`: serializable records and persistence. Training records stay here so model execution need not import the training controller. -- `lifecycle/`: request submission, retrieval and Future/Envelope handling. This differs from the instance lifetime supervised in `resources/`. -- `config/`: configuration declarations and normalization. `telemetry/` consumes plain snapshots without importing controllers or domain configuration. - -The runtime modules have different boundaries: `server/runtime.py` initializes Twinkle and DeviceMesh; `resources/runtime.py` owns Ray/Serve effects; `sampling/execution/runtime.py` tracks sampling request execution, and `training/execution/runtime.py` executes controller-owned training handles. - -Multi-LoRA execution registers every preallocated slot with the distributed reducer. DDP requires `find_unused_parameters=True` and `static_graph=False`; FSDP requires `reshard_after_forward=True` so adapter switches act on sharded parameters instead of retained views. Incompatible overrides fail during model initialization, before weights are loaded. - -## Training configuration - -The public YAML keeps `applications` with `import_path: model`, `args.train_mode`, `device_group`, and `device_mesh`. Each model application selects one training mode. To expose LoRA and full training for the same weights, use two applications with distinct names and route aliases. `config/training.py` owns the public arguments and resolves each application to one validated instance specification. Training idle expiration uses only `idle_timeout_s`. Optional top-level `resources` contains the shared allocation/startup/shutdown/monitor budgets, independent of training or sampling policy. - -## Sampling ownership and state reconstruction - -Stable sampling uses shared access; configuration and full-weight updates use exclusive access. Every participating instance is reserved before an update. Idle sweep checks the same reservations under the state lock. The lock order is configuration then state; startup is never awaited while holding the configuration lock. - -Pool admission is bounded by `queue_timeout`. A complete control broadcast or full-weight update uses one `effective_execution_timeout`, including completion confirmation. An ambiguous result fences all participants. Cancellation alone never confirms device release; shutdown failure keeps the instance quarantined and counted against capacity. - -Serve handle and resource inventory lookups run outside the controller event loop, so synchronous Ray metadata RPCs cannot prevent these deadlines from firing. -Placement-group creation remains tracked across cancellation; removal waits for the late result before certifying release. - -Only successful broadcasts update replay state. Superseded templates, adapter configurations (per owner/name), and full checkpoints are compacted after the last `apply_patch`. Arbitrary patches preserve order and form a compaction barrier: dependencies before a patch must remain available. LoRA weight caches are loaded per request and eviction is not replayed. - -A client object rejects creation of an adapter name it already owns. Release it before creating a replacement; separate objects may use the same label. - -## Tests - -Unit tests follow the component they exercise. Reusable fake drivers and queues live in helper modules, not in imported `test_*.py` modules. Resource allocation probes live under `tests/server/integration/resources`, sampling probes under `tests/server/integration/sampler`, and training acceptance tools under `tests/server/integration/on_demand_training`. - -Control fixtures demonstrate scheduling and release mechanisms, not real model execution. Hardware acceptance must separately record the tested commit, runtime import origins, real model operations, numerical checks and confirmed cleanup. The package DAG and lightweight import checks live in `tests/server/test_dependencies.py`. - -## Public training entrypoints - -Clients still create and release training handles. A stable CPU router now owns those public paths: - -| API | Public path | Behavior | -| --- | --- | --- | -| Twinkle | `twinkle/create` | Return service capabilities; no device allocation | -| Twinkle | `twinkle/add_adapter_to_model` | Allocate a handle and start an instance when needed | -| Twinkle | `twinkle/remove_adapter` | Release the handle; reclaim idle excess capacity | -| Tinker | `tinker/create_model` | Allocate a training handle | -| Tinker | `tinker/unload_model` | Release that handle | - -Forward, backward, optimizer, save/load and retrieve paths remain available through the router. The SDK carries the returned `training_id` (Tinker calls it `model_id`); the router resolves its owning instance and epoch on every request. This replaces Serve multiplexing affinity: a handle stays with its owning instance, while clients may reach any CPU router replica. A worker rejects direct create/unload calls because it cannot allocate or retire controller-owned resources. No public lifecycle endpoint was removed. - -## Ray infrastructure boundary - -`infra/_ray/ResourceManager` builds worker groups and rank/device meshes; `RayHelper` exposes that process-local infrastructure. Pool drivers reuse `ResourceManager.from_allocation`, platform device discovery and the shared CPU sizing policy. `resources/DevicePool` owns dynamic placement groups and waits for confirmed removal; `resources/RayRuntime` owns Serve applications and late deployment cleanup. These async ownership/deadline rules do not belong in the process-global helper. Both pools use the same resource implementation. - -State repositories compose the shared domain managers: `ServerState.models`, `.sessions`, `.futures`, and `.sampling_sessions`. `TrainingStore` owns training persistence and its retry ledger, while the record definitions remain independent of I/O. Constructing a store performs no storage operation. - -## Component composition and execution contracts - -There is no separate Catalog object: `TrainingConfig` already owns the validated model specifications and lookup by model/mode. The launcher normalizes YAML once and builds the two pools. `actor.py` is the Ray composition boundary; `controller.py` owns scheduling and mutable pool state; `driver.py` translates a scheduling decision into resource operations and domain-specific worker setup; `router.py` authenticates/routes public requests; `execution/app.py` builds the compute service and backend. Runtime objects execute requests inside that service. Neither router nor app allocates devices independently. - -A common business Controller base would need switches for training handles, LoRA slots, exclusive full training, sampling load, and broadcast replay. Those policies have different invariants. Both controllers instead compose `InstanceLifecycle`; both drivers compose `RayRuntime`, which composes `DevicePool`. This is the reusable layer. Subdirectories keep each domain's scheduling and execution adjacent without mixing their responsibilities. - -## Task queue flow - -1. Parse/validate the request and check capabilities, quotas and owner admission before writing a Future. -2. Persist `pending` with an immutable deadline. Training work enters an adapter queue; sampling can reserve bounded concurrent execution. -3. Serial work records `queued` and wakes the round-robin worker. An atomic start claim changes an eligible Future to `running`; cancellation can win only before this claim. -4. Execution writes `completed` or `failed`. Queue timeout fails unstarted work; execution timeout fails the client Future. Late queued/running writes cannot regress a terminal record. -5. A client timeout, disconnect, or failed Future does not prove backend completion. The owning task/scope retains capacity and weight protection until actual completion; unknown completion fences the instance until confirmed shutdown. -6. Shutdown stops admission, fails queued work and joins owned execution. Unconfirmed worker/placement-group removal keeps the instance quarantined and counted against capacity. - -`protocol/types` owns serializable request/response models and shared status enums. `task_queue/types.py` owns in-process coroutine factories, asyncio tasks and callbacks; these cannot be wire fields. `ExecutionScope` is passed explicitly from the internal request to admitted jobs, and waits for jobs registered by earlier jobs as well. - -Business identity, retry metadata, resource ownership and completion witnesses must use explicit arguments or owned records. Use `ContextVar` only for genuinely ambient context such as logging/tracing where passing through every API adds no domain meaning. Every temporary context assignment must pair `set()` with `reset(token)` in `finally`; spawning tasks copies context and does not extend resource ownership. diff --git a/src/twinkle/server/model/__init__.py b/src/twinkle/server/model/__init__.py new file mode 100644 index 000000000..85b3e739d --- /dev/null +++ b/src/twinkle/server/model/__init__.py @@ -0,0 +1 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. diff --git a/src/twinkle/server/training/execution/tinker_datum.py b/src/twinkle/server/model/tinker_datum.py similarity index 100% rename from src/twinkle/server/training/execution/tinker_datum.py rename to src/twinkle/server/model/tinker_datum.py diff --git a/src/twinkle/server/session_resource/processor.py b/src/twinkle/server/session_resource/processor.py index 17c398aa5..dd2a7ff10 100644 --- a/src/twinkle/server/session_resource/processor.py +++ b/src/twinkle/server/session_resource/processor.py @@ -80,7 +80,7 @@ async def _on_resource_expired(self, resource_id: str) -> None: ``_on_processor_expired`` is the supported extension point: the processor-domain name is kept deliberately so subclass authors override a - method named for processors rather than the generic base-class hook. It is + method named for processors rather than the generic base-class hook. """ await self._on_processor_expired(resource_id) diff --git a/src/twinkle/server/training/execution/backends/common.py b/src/twinkle/server/training/execution/backends/common.py index 024a0bfa1..0c62f3090 100644 --- a/src/twinkle/server/training/execution/backends/common.py +++ b/src/twinkle/server/training/execution/backends/common.py @@ -11,7 +11,7 @@ from typing import Any, List from twinkle import DeviceMesh, remote_function -from twinkle.server.training.execution.tinker_datum import datum_to_input_feature, extract_rl_features_for_loss +from twinkle.server.model.tinker_datum import datum_to_input_feature, extract_rl_features_for_loss from twinkle.server.validation.tinker_loss import validate_tinker_loss from twinkle.template import Template diff --git a/src/twinkle/server/training/execution/backends/megatron_model.py b/src/twinkle/server/training/execution/backends/megatron_model.py index dcad8cb29..42605f427 100644 --- a/src/twinkle/server/training/execution/backends/megatron_model.py +++ b/src/twinkle/server/training/execution/backends/megatron_model.py @@ -16,9 +16,9 @@ from twinkle.data_format import InputFeature, Trajectory from twinkle.infra import collect_tensor_dict from twinkle.model.megatron import MegatronModel, MultiLoraMegatronModel +from twinkle.server.model.tinker_datum import datum_to_input_feature, extract_rl_features_for_loss from twinkle.server.training.execution.backends.common import (TwinkleCompatModelBase, clean_metrics, collect_forward_backward_results, to_cpu_safe_output) -from twinkle.server.training.execution.tinker_datum import datum_to_input_feature, extract_rl_features_for_loss from twinkle.utils.nccl_safe import nccl_safe_megatron diff --git a/src/twinkle/server/training/execution/backends/transformers_model.py b/src/twinkle/server/training/execution/backends/transformers_model.py index 38639ada1..abf6c31dc 100644 --- a/src/twinkle/server/training/execution/backends/transformers_model.py +++ b/src/twinkle/server/training/execution/backends/transformers_model.py @@ -19,9 +19,9 @@ from twinkle.infra import collect_tensor_dict from twinkle.model import MultiLoraTransformersModel, SpectralHybridTransformersModel from twinkle.model.transformers import TransformersModel +from twinkle.server.model.tinker_datum import datum_to_input_feature, extract_rl_features_for_loss from twinkle.server.training.execution.backends.common import (TwinkleCompatModelBase, clean_metrics, collect_forward_backward_results, to_cpu_safe_output) -from twinkle.server.training.execution.tinker_datum import datum_to_input_feature, extract_rl_features_for_loss class _TransformersTinkerCompatMixin(TwinkleCompatModelBase): diff --git a/tests/server/integration/on_demand_training/acceptance_support.py b/tests/server/integration/on_demand_training/acceptance_support.py index 07abc6b5b..c14827c6c 100644 --- a/tests/server/integration/on_demand_training/acceptance_support.py +++ b/tests/server/integration/on_demand_training/acceptance_support.py @@ -65,7 +65,7 @@ def configure(self): self.client.set_processor('InputProcessor', padding_side='right') def datums(self): - from twinkle.server.training.execution.tinker_datum import input_feature_to_datum + from twinkle.server.model.tinker_datum import input_feature_to_datum return [input_feature_to_datum(feature) for feature in self.inputs] def step(self): diff --git a/tests/server/integration/test_dpo_e2e.py b/tests/server/integration/test_dpo_e2e.py index 6597d5d55..7a69cbcab 100644 --- a/tests/server/integration/test_dpo_e2e.py +++ b/tests/server/integration/test_dpo_e2e.py @@ -182,7 +182,7 @@ def test_dpo_tinker(): """ from tinker import types from twinkle.dataloader import DataLoader - from twinkle.server.training.execution.tinker_datum import input_feature_to_datum + from twinkle.server.model.tinker_datum import input_feature_to_datum backend = get_backend() log(f'=== test_dpo_tinker [backend={backend}] ===') diff --git a/tests/server/integration/test_full_param_e2e.py b/tests/server/integration/test_full_param_e2e.py index c6f54f0c6..32d605191 100644 --- a/tests/server/integration/test_full_param_e2e.py +++ b/tests/server/integration/test_full_param_e2e.py @@ -46,7 +46,7 @@ from twinkle.dataloader import DataLoader # noqa: E402 from twinkle.dataset import Dataset, DatasetMeta # noqa: E402 from twinkle.preprocessor import SelfCognitionProcessor # noqa: E402 -from twinkle.server.training.execution.tinker_datum import input_feature_to_datum # noqa: E402 +from twinkle.server.model.tinker_datum import input_feature_to_datum # noqa: E402 init_tinker_client() diff --git a/tests/server/integration/test_sft_e2e.py b/tests/server/integration/test_sft_e2e.py index cff8a2a74..16a014b8c 100644 --- a/tests/server/integration/test_sft_e2e.py +++ b/tests/server/integration/test_sft_e2e.py @@ -156,7 +156,7 @@ def test_sft_tinker(): """ from tinker import types from twinkle.dataloader import DataLoader - from twinkle.server.training.execution.tinker_datum import input_feature_to_datum + from twinkle.server.model.tinker_datum import input_feature_to_datum backend = get_backend() log(f'=== test_sft_tinker [backend={backend}] ===') diff --git a/tests/server/lifecycle/test_retrieve_endpoint.py b/tests/server/lifecycle/test_retrieve_endpoint.py index cde575825..704ea2f71 100644 --- a/tests/server/lifecycle/test_retrieve_endpoint.py +++ b/tests/server/lifecycle/test_retrieve_endpoint.py @@ -17,10 +17,10 @@ class _State: + """A fake ServerState whose get_future returns a fixed record (or None).""" async def assert_future_owner(self, request_id, token): return None - """A fake ServerState whose get_future returns a fixed record (or None).""" def __init__(self, record): self._record = record diff --git a/tests/server/model/test_tinker_compat_output.py b/tests/server/model/test_tinker_compat_output.py index f4d76bbcc..6a5f393eb 100644 --- a/tests/server/model/test_tinker_compat_output.py +++ b/tests/server/model/test_tinker_compat_output.py @@ -3,7 +3,7 @@ from tinker import types from unittest.mock import Mock -from twinkle.server.training.execution.tinker_datum import extract_rl_features_for_loss +from twinkle.server.model.tinker_datum import extract_rl_features_for_loss from twinkle.server.training.execution.backends.common import TwinkleCompatModelBase from twinkle.server.exceptions import RequestRejectedError diff --git a/tests/server/start_e2e_server.py b/tests/server/start_e2e_server.py index 337e387f2..21995f969 100644 --- a/tests/server/start_e2e_server.py +++ b/tests/server/start_e2e_server.py @@ -38,7 +38,6 @@ MEGATRON_BASE = "tests/server/config/server_config_4b_e2e_megatron.yaml" _FULL = { "train_mode": "full", - "idle_timeout_s": 3600, "max_replicas": 1, # Full-parameter training holds the deployment exclusively; use long # timeouts so the tenant is not expired mid-test. diff --git a/tests/server/task_queue/test_task_queue_mixin.py b/tests/server/task_queue/test_task_queue_mixin.py index 4fae0d1bd..e2f111e66 100644 --- a/tests/server/task_queue/test_task_queue_mixin.py +++ b/tests/server/task_queue/test_task_queue_mixin.py @@ -49,12 +49,6 @@ def test_exception_summary_falls_back_to_the_ray_error_itself(): assert 'Model.step' in summary - - - - - - @pytest.mark.asyncio async def test_worker_stop_exits_when_wait_for_swallows_cancellation(monkeypatch): queue = _DummyQueue() diff --git a/tests/server/test_dependencies.py b/tests/server/test_dependencies.py index 575b4dd98..09c0a4f49 100644 --- a/tests/server/test_dependencies.py +++ b/tests/server/test_dependencies.py @@ -13,7 +13,7 @@ # the same layer are independent and therefore may not import one another. _LAYERS = ( frozenset({'exceptions'}), - frozenset({'config', 'middleware', 'runtime', 'task_errors', 'utils', 'validation'}), + frozenset({'config', 'middleware', 'model', 'runtime', 'task_errors', 'utils', 'validation'}), frozenset({'telemetry'}), frozenset({'state'}), frozenset({'resources'}), From c48d044d50c68b8731553505943be29d866d4e44 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 20:34:40 +0800 Subject: [PATCH 36/46] fix(model): preserve frozen base layers during FSDP slot activation --- src/twinkle/model/multi_lora.py | 11 ++++++++--- src/twinkle/model/transformers/hybrid/model.py | 6 +++--- .../transformers/multi_lora_transformers.py | 16 +++------------- .../test_multi_lora_distributed_slots.py | 6 ++++++ 4 files changed, 20 insertions(+), 19 deletions(-) diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index 960ddf16a..c71fd3f38 100644 --- a/src/twinkle/model/multi_lora.py +++ b/src/twinkle/model/multi_lora.py @@ -145,6 +145,13 @@ def _named_slot_parameters(self, module): for name, parameter in value.named_parameters(): yield prefix + '.' + name, parameter + def enable_slot_gradients(self): + """Keep owned adapter slots visible to reducers without unfreezing the base.""" + modules = self.module if isinstance(self.module, list) else [self.module] + for module in modules: + for _, parameter in self._named_slot_parameters(module): + parameter.requires_grad_(True) + @contextmanager def trainable_slots(self): """Expose all slots while wrapping DDP/FSDP, then restore the active mask. @@ -158,9 +165,7 @@ def trainable_slots(self): for name, parameter in self._named_slot_parameters(module) } for module in modules] try: - for module in modules: - for _, parameter in self._named_slot_parameters(module): - parameter.requires_grad_(True) + self.enable_slot_gradients() yield finally: for module, mask in zip(modules, masks): diff --git a/src/twinkle/model/transformers/hybrid/model.py b/src/twinkle/model/transformers/hybrid/model.py index 525f97243..3ccae3a71 100644 --- a/src/twinkle/model/transformers/hybrid/model.py +++ b/src/twinkle/model/transformers/hybrid/model.py @@ -64,7 +64,7 @@ def _adapter_context(self, adapter_name: str, disable_lora: bool = False): return with super()._adapter_context(adapter_name) as slot_name: - self._restore_preallocated_lora_grad() + self.multi_adapter.enable_slot_gradients() self.fft_slots.activate_fft_slot(adapter_name) yield slot_name @@ -72,14 +72,14 @@ def _adapter_context(self, adapter_name: str, disable_lora: bool = False): def _disabled_adapter_context(self, adapter_name: str): try: with super()._adapter_context(adapter_name, disable_lora=True) as slot_name: - self._restore_preallocated_lora_grad() + self.multi_adapter.enable_slot_gradients() self.fft_slots.deactivate_fft_slots() try: yield slot_name finally: self.fft_slots.activate_fft_slot(adapter_name) finally: - self._restore_preallocated_lora_grad() + self.multi_adapter.enable_slot_gradients() @remote_function() def add_adapter_to_model(self, adapter_name: str, config_or_dir: Union[PeftConfig, str], **kwargs): diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index bd90a6100..9c8fbb6b8 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/src/twinkle/model/transformers/multi_lora_transformers.py @@ -153,24 +153,14 @@ def _ensure_target_parameter_lora_installed(self, config: LoraConfig) -> None: @contextmanager def _adapter_context(self, adapter_name: str, disable_lora: bool = False): - """Activate one LoRA tenant for a model operation.""" + """Activate a tenant while keeping FSDP2's trainable parameter set stable.""" try: with self.multi_adapter.adapter(adapter_name, disable_lora=disable_lora) as slot_name: - self._restore_preallocated_lora_grad() + self.multi_adapter.enable_slot_gradients() yield slot_name finally: # Leaving a disable_lora context re-enables adapter layers, which flips the flags again. - self._restore_preallocated_lora_grad() - - def _restore_preallocated_lora_grad(self) -> None: - # PEFT's set_adapter / disable_adapter_layers mark every inactive slot requires_grad=False, - # but FSDP2 captures each parameter's trainable state at lazy init. If that state drifts, - # gradient reduction for a slot other than the one active at wrap time breaks (e.g. fp32 - # grads assigned to bf16 params on NPU). Inactive slots still get no gradient because the - # patched LoRA forward only runs active adapters. - for name, parameter in self.multi_adapter.module.named_parameters(): - if 'lora_' in name and not parameter.requires_grad: - parameter.requires_grad_(True) + self.multi_adapter.enable_slot_gradients() @remote_function(dispatch='slice_dp', collect=collect_tensor_dict) def forward(self, *, inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], **kwargs): diff --git a/tests/transformers/test_multi_lora_distributed_slots.py b/tests/transformers/test_multi_lora_distributed_slots.py index dd5a92cb8..ce177078a 100644 --- a/tests/transformers/test_multi_lora_distributed_slots.py +++ b/tests/transformers/test_multi_lora_distributed_slots.py @@ -153,3 +153,9 @@ def test_slot_registration_never_unfreezes_base_parameters_with_lora_in_their_na assert not any(parameter.requires_grad for parameter in base.values()) raise RuntimeError('wrapper failed') assert {name: parameter.requires_grad for name, parameter in module.named_parameters()} == mask + manager.acquire_lora('first', config) + wrapper = _wrapper(module, manager, lambda model: model) + for disable_lora in (False, True): + with wrapper._adapter_context('first', disable_lora=disable_lora): + assert not any(parameter.requires_grad for parameter in base.values()) + assert not any(parameter.requires_grad for parameter in base.values()) From f46a42a46e39a8a5e6ec835e4d5f6bda1d188971 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 20:34:42 +0800 Subject: [PATCH 37/46] fix(server): keep unknown cleanup outcomes visible and retryable --- src/twinkle/server/launcher/server_launcher.py | 11 +++++------ src/twinkle_client/model/multi_lora_transformers.py | 12 ++++++++---- tests/server/training/test_launcher.py | 4 ++++ tests/twinkle_client/test_async_components.py | 11 ++++++++--- 4 files changed, 25 insertions(+), 13 deletions(-) diff --git a/src/twinkle/server/launcher/server_launcher.py b/src/twinkle/server/launcher/server_launcher.py index 0d9b13b91..d34122de5 100644 --- a/src/twinkle/server/launcher/server_launcher.py +++ b/src/twinkle/server/launcher/server_launcher.py @@ -181,12 +181,11 @@ def _deploy_training(self) -> None: def _stop_controller(self, controller) -> None: import ray - try: - ray.get(controller.shutdown.remote(), timeout=self.config.resources.shutdown_timeout_s) - finally: - # Killing the owner prevents restart but does not certify its devices - # were released. Preserve a failed/unknown shutdown for _run_cleanup. - ray.kill(controller, no_restart=True) + + # Keep an unconfirmed owner available for cleanup retry; a timeout alone + # does not certify its workers or devices were released. + ray.get(controller.shutdown.remote(), timeout=self.config.resources.shutdown_timeout_s) + ray.kill(controller, no_restart=True) def _stop_training_controller(self) -> None: if self._training_controller is not None: diff --git a/src/twinkle_client/model/multi_lora_transformers.py b/src/twinkle_client/model/multi_lora_transformers.py index 5b2ef2dc5..efbdd2d60 100644 --- a/src/twinkle_client/model/multi_lora_transformers.py +++ b/src/twinkle_client/model/multi_lora_transformers.py @@ -192,10 +192,14 @@ def remove_adapter(self, adapter_name: str | None = None) -> None: try: self._submit('remove_adapter', model_types.AdapterRequest, None, adapter_name=name, training_id=training_id) except TwinkleHTTPError as exc: - if exc.status_code != 404: - raise - # An expired or lost handle is unusable. Forget its local binding so - # this label can be recreated; unknown cleanup outcomes still raise. + if exc.status_code == 404: + self._forget_training(name, training_id) + # A missing handle is unusable, but does not certify device release. + raise + self._forget_training(name, training_id) + + def _forget_training(self, name: str, training_id: str) -> None: + """A late response must never erase a newer same-name binding.""" with self._training_lock: if self._training_ids.get(name) == training_id: self._training_ids.pop(name, None) diff --git a/tests/server/training/test_launcher.py b/tests/server/training/test_launcher.py index 630e47cb9..8ee8b40be 100644 --- a/tests/server/training/test_launcher.py +++ b/tests/server/training/test_launcher.py @@ -55,6 +55,10 @@ def stop(): assert finished.wait(2), 'Controller shutdown exceeded its budget and blocked the next owner' assert cleaned == [True] assert len(errors) == 1 and isinstance(errors[0].__cause__, ray.exceptions.GetTimeoutError) + assert launcher._training_controller is controller + ray.get(controller.unblock.remote(), timeout=5) + launcher._stop_training_controller() + assert launcher._training_controller is None finally: if not finished.is_set(): ray.get(controller.unblock.remote(), timeout=5) diff --git a/tests/twinkle_client/test_async_components.py b/tests/twinkle_client/test_async_components.py index 8c6629f81..431ed66d7 100644 --- a/tests/twinkle_client/test_async_components.py +++ b/tests/twinkle_client/test_async_components.py @@ -256,8 +256,10 @@ def post_model(self, url, body): model.add_adapter_to_model('same', None) errors = [] def remove(): - try:model.remove_adapter() - except Exception as exc:errors.append(exc) + try: + model.remove_adapter() + except Exception as exc: + errors.append(exc) releasing = threading.Thread(target=remove) releasing.start() try: @@ -335,6 +337,7 @@ def test_unknown_training_rejected_before_handshake(monkeypatch): def test_expired_handle_can_be_removed_and_recreated_on_same_object(monkeypatch): + import pytest from twinkle_client.exceptions import TwinkleHTTPError from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel calls = [] @@ -347,7 +350,9 @@ def expired(*args, **kwargs): raise TwinkleHTTPError('Training handle expired', status_code=404, error_code=404) monkeypatch.setattr(model, '_submit', expired) - model.remove_adapter('same') + with pytest.raises(TwinkleHTTPError) as error: + model.remove_adapter('same') + assert error.value.status_code == 404 monkeypatch.setattr(model, '_submit', submit) model.add_adapter_to_model('same', None) model.zero_grad() From 84e31141fdf6889f4729356bdf9cccdfcc7afafc Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 20:34:44 +0800 Subject: [PATCH 38/46] test(server): simplify probes and exercise actual CPU sampling services --- .../integration/resources/test_allocation.py | 8 + .../integration/sampler/sampling_probe.py | 190 ++++++------------ tests/server/lifecycle/test_static_guards.py | 9 - tests/server/resources/test_instances.py | 3 +- tests/server/sampler/test_pool.py | 15 -- .../server/training/test_storage_recovery.py | 9 +- tests/twinkle_client/test_types_contract.py | 1 - 7 files changed, 75 insertions(+), 160 deletions(-) diff --git a/tests/server/integration/resources/test_allocation.py b/tests/server/integration/resources/test_allocation.py index 8ebe04d1c..54d8f8b11 100644 --- a/tests/server/integration/resources/test_allocation.py +++ b/tests/server/integration/resources/test_allocation.py @@ -21,3 +21,11 @@ def test_fragmented_devices_cannot_be_combined_across_nodes(): assert evidence['passed'] and evidence['nodes'] == 2 assert evidence['fragmented_rejected'] and evidence['pending_removed'] assert len(evidence['devices']) == 2 + + +def test_sampler_uses_real_cpu_services_for_scaling_and_release(): + script = Path(__file__).resolve().parents[1] / 'sampler' / 'sampling_probe.py' + result = subprocess.run([sys.executable, str(script)], capture_output=True, text=True, timeout=1800) + assert result.returncode == 0, result.stdout + result.stderr + evidence = [json.loads(line) for line in result.stdout.splitlines() if line.startswith('{"sizes":')] + assert evidence == [{'sizes': [2, 4, 2], 'requests': 8, 'confirmed_release': True}] diff --git a/tests/server/integration/sampler/sampling_probe.py b/tests/server/integration/sampler/sampling_probe.py index 8eb1a5f35..f335c0a3e 100644 --- a/tests/server/integration/sampler/sampling_probe.py +++ b/tests/server/integration/sampler/sampling_probe.py @@ -1,5 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Real Ray/Serve 2->4->2 allocation and death confirmation, without model kernels.""" +"""Exercise the actual sampler driver and mock compute service on an owned CPU cluster.""" import asyncio import json import os @@ -7,165 +7,95 @@ import ray from ray import serve -from ray.util.placement_group import get_placement_group, placement_group_table -from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy +from ray.util import list_named_actors +from ray.util.placement_group import placement_group_table from twinkle.server.config.application_spec import SamplerArgs -from twinkle.server.deployment import bind_deployment, build_deployment_app -from twinkle.server.resources.allocation import CapacityUnavailableError, DevicePool from twinkle.server.sampling.controller import SamplerController from twinkle.server.sampling.driver import SamplerDriver -@ray.remote(num_cpus=0.01, num_gpus=1) -class Worker: - - def device(self): - return ray.get_runtime_context().get_accelerator_ids()['GPU'] - - -class Compute: - - async def __init__(self, namespace, identity, pg): - self.operations = {} - self.worker = Worker.options( - name=identity + '-worker', - namespace=namespace, - scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=pg)).remote() - self.devices = await self.worker.device.remote() - - async def invoke(self, ticket, payload, stream_queue=None): - self.operations[ticket] = asyncio.Event() - return {'status': 200, 'headers': [], 'body': b'{"status":"pending"}'} - - async def wait_operation(self, ticket): - await self.operations[ticket].wait() - self.operations.pop(ticket) - - async def finish(self): - for operation in self.operations.values(): - operation.set() - - async def pending_futures(self): - return [] - - async def runtime_identity(self): - return {'unhealthy': False, 'devices': self.devices} - - -class Driver(SamplerDriver): - - async def start(self, instance): - allocation = await self.runtime.pool.allocate(1, 'cuda', f'sampling_{self.pool_id}_{instance.epoch}', - instance.identity) - instance.application, instance.devices = allocation.name, list(allocation.devices) - app = bind_deployment( - build_deployment_app('AllocationProbe', lambda app, get_self: None), - Compute, - self.runtime.deployment_options({}, allocation.name), - deployment_name='AllocationCompute', - bind_args=(self.namespace, instance.identity, get_placement_group(allocation.name))) - await self.runtime.deploy(instance.identity, allocation.name, app, '/probe/' + instance.identity, - lambda: self.shutdown(instance)) - - async def run(namespace): args = SamplerArgs( - model_id='allocation-probe', - sampler_type='mock', - device_group={ - 'ranks': 1, - 'device_type': 'cuda' - }, - device_mesh={'dp_size': 1}, - min_replicas=2, - max_replicas=4, - target_ongoing_requests=1, - idle_timeout_s=1) - budgets = {'allocation_timeout_s': 5, 'startup_timeout_s': 600, 'shutdown_timeout_s': 180, 'sweep_interval_s': 1} - driver = Driver(args, namespace, 'allocation', {}, budgets) + model_id='mock', sampler_type='mock', + device_group={'name': 'sampler', 'ranks': 1, 'device_type': 'cpu'}, + device_mesh={'device_type': 'cpu', 'dp_size': 1}, + min_replicas=2, max_replicas=4, target_ongoing_requests=1, idle_timeout_s=1) + budgets = {'allocation_timeout_s': 10, 'startup_timeout_s': 600, + 'shutdown_timeout_s': 60, 'sweep_interval_s': 1} + driver = SamplerDriver(args, namespace, 'cpu-allocation', {}, budgets) clock = [0] pool = SamplerController(args, driver, budgets, clock=lambda: clock[0]) - training_pool = DevicePool(namespace, allocation_timeout=5, shutdown_timeout=180) - training = [] - evidence = {'sizes': [], 'device_ids': []} + # Hold replies after the actual compute route runs. The pool must retain + # their tickets until the caller can confirm backend completion. + invoke = driver.invoke + replies = 0 + dispatched, release = asyncio.Event(), asyncio.Event() + + async def delayed_reply(*args, **kwargs): + nonlocal replies + response = await invoke(*args, **kwargs) + replies += 1 + if replies == 8: + dispatched.set() + await release.wait() + return response + + driver.invoke = delayed_reply + calls = [] try: - # Two training-sized allocations plus four sampling instances share exactly eight GPUs. - for index in range(2): - training.append(await training_pool.allocate(2, 'cuda', 'probe_training', str(index))) await pool.initialize() - evidence['sizes'].append(pool._owned_count()) - payload = {'endpoint': 'twinkle/sample', 'method': 'POST', 'headers': [], 'body': b'{}'} - for _ in range(8): - await pool.request(payload) - await asyncio.gather(*tuple(pool.lifecycle.startups.values())) - evidence['sizes'].append(pool._owned_count()) - all_devices = [device for allocation in training for device in allocation.devices] - for instance in pool._ready_instances(): - identity = await driver.health(instance) - assert identity['devices'] == instance.devices - all_devices.extend(instance.devices) - assert len(all_devices) == len(set(all_devices)) == 8 - evidence['device_ids'] = sorted(all_devices) - try: - await training_pool.allocate(1, 'cuda', 'probe_training', 'overflow') - except CapacityUnavailableError: - evidence['no_overcommit'] = True - else: - raise AssertionError('A ninth device allocation unexpectedly succeeded') + sizes = [len(pool._ready_instances())] + payload = {'endpoint': 'twinkle/sample', 'method': 'POST', + 'headers': [(b'content-type', b'application/json'), + (b'authorization', b'Bearer cpu-probe'), + (b'x-request-id', b'cpu-probe')], 'query_string': b'', + 'body': json.dumps({'inputs': [{'input_ids': [1, 2]}], + 'sampling_params': {'max_tokens': 2}}).encode()} + calls = [asyncio.create_task(pool.request(payload)) for _ in range(8)] + await asyncio.wait_for(dispatched.wait(), 600) + await asyncio.gather(*tuple(pool.lifecycle.startups.values())) + sizes.append(len(pool._ready_instances())) clock[0] = 10 await pool.sweep() - assert pool._owned_count() == 4 # Public replies did not retire actual operations. - await asyncio.gather(*(driver.handle(instance).finish.remote() for instance in pool._ready_instances())) + assert pool._owned_count() == 4 + release.set() + responses = await asyncio.gather(*calls) + assert all(response['status'] == 200 for response in responses), responses await pool.lifecycle.join() + from twinkle.server.state import get_server_state + state = get_server_state() + for response in responses: + record = await state.get_future(json.loads(response['body'])['request_id']) + assert record['status'] == 'completed', record + assert len(record['result']['samples'][0]['sequences'][0]['tokens']) == 2 + health = await asyncio.gather(*(driver.health(instance) for instance in pool._ready_instances())) + assert all(not result['unhealthy'] for result in health) clock[0] = 12 await pool.sweep() await pool.lifecycle.join() - evidence['sizes'].append(len(pool._ready_instances())) - assert evidence['sizes'] == [2, 4, 2] - old_names = {instance.application for instance in pool._ready_instances()} - recovered_driver = Driver(args, namespace, 'allocation', {}, budgets) - recovered = SamplerController(args, recovered_driver, budgets, clock=lambda: clock[0]) - try: - await recovered.initialize() - assert recovered.epoch != pool.epoch and len(recovered._ready_instances()) == 2 - assert not old_names.intersection(serve.status().applications) - assert all( - placement_group_table(get_placement_group(allocation.name))['state'] != 'REMOVED' - for allocation in training) - evidence['owned_orphan_recovery'] = True - finally: - await recovered.shutdown() - await recovered_driver.runtime.wait_for_cleanup() + sizes.append(len(pool._ready_instances())) + assert sizes == [2, 4, 2], sizes await pool.shutdown() await driver.runtime.wait_for_cleanup() - names = set(serve.status().applications) - assert not any(name.startswith(driver.prefix) for name in names) - assert not any(row['state'] != 'REMOVED' and row.get('name', '').startswith(driver.prefix) + assert not any(name.startswith(driver.prefix) for name in serve.status().applications) + assert not any(row.get('name', '').startswith(driver.prefix) and row['state'] != 'REMOVED' for row in placement_group_table().values()) - from ray.util import list_named_actors assert not any(row['namespace'] == namespace and 'sampler_' in row['name'] for row in list_named_actors(all_namespaces=True)) - evidence['confirmed_release'] = True - return evidence + return {'sizes': sizes, 'requests': len(responses), 'confirmed_release': True} finally: + release.set() + await asyncio.gather(*calls, return_exceptions=True) await pool.shutdown() await driver.runtime.wait_for_cleanup() - for allocation in training: - await training_pool.remove(allocation.name) def main(): namespace = 'sampling_probe_' + uuid.uuid4().hex[:8] - os.environ.update( - TWINKLE_STATE_RAY_NAMESPACE=namespace, TWINKLE_PERSISTENCE_MODE='memory', TWINKLE_GPU_PG_CPUS_PER_PROC='1') - ray.init( - address='local', - namespace=namespace, - num_cpus=24, - num_gpus=8, - _temp_dir='/tmp/tws-' + namespace[-8:], - include_dashboard=False) + os.environ.update(TWINKLE_STATE_RAY_NAMESPACE=namespace, TWINKLE_PERSISTENCE_MODE='memory') + ray.init(address='local', namespace=namespace, num_cpus=24, num_gpus=0, + _temp_dir='/tmp/tws-' + namespace[-8:], include_dashboard=False) try: serve.start(http_options={'host': '127.0.0.1', 'port': 0}) print(json.dumps(asyncio.run(run(namespace))), flush=True) diff --git a/tests/server/lifecycle/test_static_guards.py b/tests/server/lifecycle/test_static_guards.py index 6278cd2b8..3e635a294 100644 --- a/tests/server/lifecycle/test_static_guards.py +++ b/tests/server/lifecycle/test_static_guards.py @@ -5,7 +5,6 @@ - ``TaskEnvelope`` has exactly one construction site. - The client-side invariant: the client HTTP timeout is <= 120 and strictly greater than the server Long_Poll_Window. -- Task status uses one shared protocol declaration. These exist because the spec states most of its guarantees in prose. A prose claim that nothing checks decays into a false claim -- as happened with "a consistency test asserts @@ -71,14 +70,6 @@ def test_task_envelope_has_exactly_one_construction_site(): assert sites, 'expected to find the construction sites inside envelope.py' -def test_task_status_is_shared_by_queue_and_protocol(): - from twinkle.protocol.types.lifecycle import TaskStatus, TERMINAL_STATUSES - from twinkle.server.task_queue import TaskStatus as QueueTaskStatus - - assert QueueTaskStatus is TaskStatus - assert TERMINAL_STATUSES == {TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED} - - def test_client_future_layer_is_not_imported_by_the_server(): """``_future.py`` carries an underscore because the dependency runs one way only. diff --git a/tests/server/resources/test_instances.py b/tests/server/resources/test_instances.py index 3abe66281..f7eaed80d 100644 --- a/tests/server/resources/test_instances.py +++ b/tests/server/resources/test_instances.py @@ -20,11 +20,10 @@ async def broken(): await lifecycle.join() records = [record for record in caplog.records if record.getMessage() == 'Training background job failed'] assert records and records[0].exc_info[1].args == ('background sentinel',) + assert records[0].exc_info[2] is not None assert not lifecycle.jobs - - @pytest.mark.asyncio async def test_concurrent_close_shares_one_stop_and_retries_quarantine(): import asyncio diff --git a/tests/server/sampler/test_pool.py b/tests/server/sampler/test_pool.py index 4f177fc99..e5209aad0 100644 --- a/tests/server/sampler/test_pool.py +++ b/tests/server/sampler/test_pool.py @@ -283,21 +283,6 @@ async def unknown(*args): await pool.shutdown() -@pytest.mark.asyncio -async def test_training_and_sampling_use_the_same_lifecycle_owner(): - from tests.server.state.fake_backend import FakeBackend - from tests.server.training.helpers import setup as training_setup - pool, driver, _ = setup() - training, _, _, _ = await training_setup(backend=FakeBackend()) - try: - assert type(training.lifecycle) is type(pool.lifecycle) - assert training.lifecycle is not pool.lifecycle - finally: - driver.finish.set() - await pool.shutdown() - await training.shutdown() - - @pytest.mark.asyncio async def test_replaying_config_has_the_same_startup_deadline(): pool, driver, _ = setup(minimum=0, maximum=1) diff --git a/tests/server/training/test_storage_recovery.py b/tests/server/training/test_storage_recovery.py index 63b2cddd2..7b939f61f 100644 --- a/tests/server/training/test_storage_recovery.py +++ b/tests/server/training/test_storage_recovery.py @@ -52,10 +52,13 @@ async def test_managed_memory_actor_restart_reconnects_existing_handles(): ray.kill(backend._actor, no_restart=False) for _ in range(100): try: - if await backend.health_check() and await backend.get('before') is None:break - except ray.exceptions.RayError:pass + if await backend.health_check() and await backend.get('before') is None: + break + except ray.exceptions.RayError: + pass await asyncio.sleep(0.1) - else:raise AssertionError('State actor did not restart') + else: + raise AssertionError('State actor did not restart') assert await backend.get('before') is None await backend.set('after', 'new') assert await backend.get('after') == 'new' diff --git a/tests/twinkle_client/test_types_contract.py b/tests/twinkle_client/test_types_contract.py index 868fd673c..a2d35bbed 100644 --- a/tests/twinkle_client/test_types_contract.py +++ b/tests/twinkle_client/test_types_contract.py @@ -11,7 +11,6 @@ import ast import pathlib -import typing import twinkle from twinkle.protocol.types import model as model_types From 890802724b18cc7200362b0a20f834a9667b1115 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 21:02:16 +0800 Subject: [PATCH 39/46] refactor(server): nest model and sampler instance pools --- .../sampler/vllm_sampler/vllm_sampler.py | 2 +- src/twinkle/server/config/backend_dispatch.py | 2 +- src/twinkle/server/config/training.py | 2 +- src/twinkle/server/launcher/builder_registry.py | 4 ++-- src/twinkle/server/launcher/server_launcher.py | 8 ++++---- src/twinkle/server/model/__init__.py | 16 ++++++++++++++++ .../server/{training/execution => model}/app.py | 0 .../execution => model}/backends/__init__.py | 0 .../execution => model}/backends/common.py | 0 .../backends/megatron_model.py | 4 ++-- .../execution => model}/backends/mock_model.py | 2 +- .../backends/transformers_model.py | 4 ++-- .../execution => model}/data_plane_inputs.py | 0 src/twinkle/server/model/pool/__init__.py | 2 ++ .../server/{training => model/pool}/actor.py | 0 .../{training => model/pool}/controller.py | 0 .../server/{training => model/pool}/driver.py | 2 +- .../server/{training => model/pool}/router.py | 0 .../{training/execution => model}/runtime.py | 0 .../execution => model}/tinker_handlers.py | 0 .../execution => model}/twinkle_handlers.py | 4 ++-- src/twinkle/server/processor/app.py | 2 +- .../{sampling/execution => sampler}/__init__.py | 0 .../{sampling/execution => sampler}/app.py | 0 .../execution => sampler}/backends/__init__.py | 0 .../backends/mock_sampler.py | 0 .../execution => sampler}/backends/streaming.py | 0 .../backends/vllm_async.py | 0 .../execution => sampler}/concurrency.py | 0 src/twinkle/server/sampler/pool/__init__.py | 2 ++ .../server/{sampling => sampler/pool}/actor.py | 0 .../{sampling => sampler/pool}/controller.py | 2 +- .../server/{sampling => sampler/pool}/driver.py | 2 +- .../server/{sampling => sampler/pool}/router.py | 2 +- .../{sampling/execution => sampler}/runtime.py | 0 .../execution => sampler}/tinker_handlers.py | 2 +- .../execution => sampler}/twinkle_handlers.py | 2 +- .../{sampling/execution => sampler}/weights.py | 0 src/twinkle/server/sampling/__init__.py | 2 -- src/twinkle/server/training/__init__.py | 2 -- .../server/training/execution/__init__.py | 17 ----------------- tests/server/config/test_server_config.py | 2 +- tests/server/config/test_training_config.py | 8 ++++---- tests/server/contract/client_api_harness.py | 8 ++++---- .../test_resident_release_e2e.py | 2 +- .../integration/sampler/sampling_probe.py | 4 ++-- .../integration/sampler/test_real_pool_e2e.py | 4 ++-- tests/server/integration/test_actor_recovery.py | 4 ++-- .../server/model/test_full_mode_exclusivity.py | 2 +- tests/server/model/test_mock_model.py | 4 ++-- .../server/model/test_preflight_loss_inputs.py | 8 ++++---- tests/server/model/test_replica_lifecycle.py | 8 ++++---- tests/server/model/test_tinker_compat_output.py | 2 +- tests/server/model/test_tinker_handlers.py | 12 ++++++------ tests/server/model/test_twinkle_async_inputs.py | 12 ++++++------ tests/server/resources/test_instances.py | 4 ++-- tests/server/sampler/test_mock_sampler.py | 8 ++++---- tests/server/sampler/test_pool.py | 2 +- tests/server/sampler/test_request_runtime.py | 12 ++++++------ .../sampler/test_resolve_sampler_weights.py | 4 ++-- .../server/sampler/test_sampling_concurrency.py | 12 ++++++------ tests/server/sampler/test_stream_guarantees.py | 6 +++--- tests/server/sampler/test_streaming_bridge.py | 2 +- tests/server/sampler/test_tinker_handlers.py | 4 ++-- tests/server/sampler/test_twinkle_async_rows.py | 4 ++-- tests/server/static/backend_call_exemptions.py | 4 ++-- .../server/static/test_adapter_name_mapping.py | 4 ++-- .../static/test_no_direct_backend_call.py | 2 +- .../static/test_no_twinkle_http_exception.py | 4 ++-- tests/server/telemetry/test_training_metrics.py | 8 ++++---- .../test_app_builders_characterization.py | 4 ++-- tests/server/test_dependencies.py | 16 ++++++++-------- tests/server/training/helpers.py | 4 ++-- .../training/test_actor_initialization.py | 8 ++++---- tests/server/training/test_controller.py | 2 +- tests/server/training/test_failure_recovery.py | 4 ++-- tests/server/training/test_launcher.py | 2 +- tests/server/training/test_router.py | 8 ++++---- .../training/test_runtime_initialization.py | 6 +++--- tests/server/training/test_runtime_queue.py | 2 +- tests/server/training/test_storage_recovery.py | 2 +- tests/transformers/test_spectral_hybrid_lora.py | 6 +++--- .../test_vllm_sampler_tq_generation.py | 2 +- 83 files changed, 156 insertions(+), 157 deletions(-) rename src/twinkle/server/{training/execution => model}/app.py (100%) rename src/twinkle/server/{training/execution => model}/backends/__init__.py (100%) rename src/twinkle/server/{training/execution => model}/backends/common.py (100%) rename src/twinkle/server/{training/execution => model}/backends/megatron_model.py (97%) rename src/twinkle/server/{training/execution => model}/backends/mock_model.py (99%) rename src/twinkle/server/{training/execution => model}/backends/transformers_model.py (96%) rename src/twinkle/server/{training/execution => model}/data_plane_inputs.py (100%) create mode 100644 src/twinkle/server/model/pool/__init__.py rename src/twinkle/server/{training => model/pool}/actor.py (100%) rename src/twinkle/server/{training => model/pool}/controller.py (100%) rename src/twinkle/server/{training => model/pool}/driver.py (98%) rename src/twinkle/server/{training => model/pool}/router.py (100%) rename src/twinkle/server/{training/execution => model}/runtime.py (100%) rename src/twinkle/server/{training/execution => model}/tinker_handlers.py (100%) rename src/twinkle/server/{training/execution => model}/twinkle_handlers.py (99%) rename src/twinkle/server/{sampling/execution => sampler}/__init__.py (100%) rename src/twinkle/server/{sampling/execution => sampler}/app.py (100%) rename src/twinkle/server/{sampling/execution => sampler}/backends/__init__.py (100%) rename src/twinkle/server/{sampling/execution => sampler}/backends/mock_sampler.py (100%) rename src/twinkle/server/{sampling/execution => sampler}/backends/streaming.py (100%) rename src/twinkle/server/{sampling/execution => sampler}/backends/vllm_async.py (100%) rename src/twinkle/server/{sampling/execution => sampler}/concurrency.py (100%) create mode 100644 src/twinkle/server/sampler/pool/__init__.py rename src/twinkle/server/{sampling => sampler/pool}/actor.py (100%) rename src/twinkle/server/{sampling => sampler/pool}/controller.py (99%) rename src/twinkle/server/{sampling => sampler/pool}/driver.py (98%) rename src/twinkle/server/{sampling => sampler/pool}/router.py (98%) rename src/twinkle/server/{sampling/execution => sampler}/runtime.py (100%) rename src/twinkle/server/{sampling/execution => sampler}/tinker_handlers.py (98%) rename src/twinkle/server/{sampling/execution => sampler}/twinkle_handlers.py (99%) rename src/twinkle/server/{sampling/execution => sampler}/weights.py (100%) delete mode 100644 src/twinkle/server/sampling/__init__.py delete mode 100644 src/twinkle/server/training/__init__.py delete mode 100644 src/twinkle/server/training/execution/__init__.py diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index 73aebfeee..073fb0e9e 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -459,7 +459,7 @@ def sample_stream( def sample_stream_to_queue(self, queue, inputs, sampling_params=None, adapter_name='', adapter_path=None): """Push streaming deltas to a cross-process Ray queue.""" - from twinkle.server.sampling.execution.backends.streaming import stream_to_queue + from twinkle.server.sampler.backends.streaming import stream_to_queue stream_to_queue(self, queue, inputs, sampling_params, adapter_name, adapter_path) @remote_function(dispatch='all', collect='first') diff --git a/src/twinkle/server/config/backend_dispatch.py b/src/twinkle/server/config/backend_dispatch.py index c7cb70ebd..b4da33f90 100644 --- a/src/twinkle/server/config/backend_dispatch.py +++ b/src/twinkle/server/config/backend_dispatch.py @@ -5,7 +5,7 @@ type selector (``mock | vllm | torch``) share one validate-then-dispatch shape. ``BackendSelector`` expresses it once, parameterized by the permitted value set and the per-value construction callbacks. The callbacks stay defined in -``training/execution/app.py`` and ``sampling/execution/app.py`` so the lazy backend imports remain local +``model/app.py`` and ``sampler/app.py`` so the lazy backend imports remain local — importing the mock backend on a CPU-only host must not pull in torch/vllm/megatron. """ diff --git a/src/twinkle/server/config/training.py b/src/twinkle/server/config/training.py index 335fb134f..e09b0626b 100644 --- a/src/twinkle/server/config/training.py +++ b/src/twinkle/server/config/training.py @@ -64,7 +64,7 @@ class ModelArgs(TrainingArguments, LoraArguments): """Args for the ``model`` deployment. The ``backend`` field selects the model implementation (transformers / - megatron / mock) — see ``training/execution/app.py`` for the dispatch site that picks + megatron / mock) — see ``model/app.py`` for the dispatch site that picks the runtime class based on this value. """ diff --git a/src/twinkle/server/launcher/builder_registry.py b/src/twinkle/server/launcher/builder_registry.py index 2b9395b6c..9c81d07e0 100644 --- a/src/twinkle/server/launcher/builder_registry.py +++ b/src/twinkle/server/launcher/builder_registry.py @@ -28,9 +28,9 @@ def get_builders() -> dict[str, Callable]: """ from twinkle.server.data_plane import build_data_plane_app from twinkle.server.gateway import build_gateway_app + from twinkle.server.model import build_model_app from twinkle.server.processor import build_processor_app - from twinkle.server.sampling.execution import build_sampler_app - from twinkle.server.training.execution import build_model_app + from twinkle.server.sampler import build_sampler_app return { 'build_gateway_app': build_gateway_app, diff --git a/src/twinkle/server/launcher/server_launcher.py b/src/twinkle/server/launcher/server_launcher.py index d34122de5..445de7ad2 100644 --- a/src/twinkle/server/launcher/server_launcher.py +++ b/src/twinkle/server/launcher/server_launcher.py @@ -117,8 +117,8 @@ def _deploy_sampler(self, app: ApplicationSpec) -> None: import ray from ray import serve - from twinkle.server.sampling.actor import create_sampler_controller - from twinkle.server.sampling.router import build_sampler_router + from twinkle.server.sampler.pool.actor import create_sampler_controller + from twinkle.server.sampler.pool.router import build_sampler_router controller, options, pool_id = create_sampler_controller(app, self.ray_namespace, self.config.resources.model_dump(), build_worker_env_vars()) @@ -156,8 +156,8 @@ def _deploy_training(self) -> None: import ray from ray import serve - from twinkle.server.training.actor import TrainingControllerActor, controller_name - from twinkle.server.training.router import build_training_router + from twinkle.server.model.pool.actor import TrainingControllerActor, controller_name + from twinkle.server.model.pool.router import build_training_router config = self.config.training_config controller = TrainingControllerActor.options( name=controller_name(self.ray_namespace), diff --git a/src/twinkle/server/model/__init__.py b/src/twinkle/server/model/__init__.py index 85b3e739d..1b784496b 100644 --- a/src/twinkle/server/model/__init__.py +++ b/src/twinkle/server/model/__init__.py @@ -1 +1,17 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +"""Model deployment package. + +``build_model_app`` is exposed lazily via ``__getattr__`` so importing the +package doesn't eagerly pull in ``app.py`` and its transitive dependencies. +""" +from __future__ import annotations + +__all__ = ['build_model_app'] + + +def __getattr__(name: str): + if name == 'build_model_app': + from .app import build_model_app + + return build_model_app + raise AttributeError(f'module {__name__!r} has no attribute {name!r}') diff --git a/src/twinkle/server/training/execution/app.py b/src/twinkle/server/model/app.py similarity index 100% rename from src/twinkle/server/training/execution/app.py rename to src/twinkle/server/model/app.py diff --git a/src/twinkle/server/training/execution/backends/__init__.py b/src/twinkle/server/model/backends/__init__.py similarity index 100% rename from src/twinkle/server/training/execution/backends/__init__.py rename to src/twinkle/server/model/backends/__init__.py diff --git a/src/twinkle/server/training/execution/backends/common.py b/src/twinkle/server/model/backends/common.py similarity index 100% rename from src/twinkle/server/training/execution/backends/common.py rename to src/twinkle/server/model/backends/common.py diff --git a/src/twinkle/server/training/execution/backends/megatron_model.py b/src/twinkle/server/model/backends/megatron_model.py similarity index 97% rename from src/twinkle/server/training/execution/backends/megatron_model.py rename to src/twinkle/server/model/backends/megatron_model.py index 42605f427..b06228248 100644 --- a/src/twinkle/server/training/execution/backends/megatron_model.py +++ b/src/twinkle/server/model/backends/megatron_model.py @@ -16,9 +16,9 @@ from twinkle.data_format import InputFeature, Trajectory from twinkle.infra import collect_tensor_dict from twinkle.model.megatron import MegatronModel, MultiLoraMegatronModel +from twinkle.server.model.backends.common import (TwinkleCompatModelBase, clean_metrics, + collect_forward_backward_results, to_cpu_safe_output) from twinkle.server.model.tinker_datum import datum_to_input_feature, extract_rl_features_for_loss -from twinkle.server.training.execution.backends.common import (TwinkleCompatModelBase, clean_metrics, - collect_forward_backward_results, to_cpu_safe_output) from twinkle.utils.nccl_safe import nccl_safe_megatron diff --git a/src/twinkle/server/training/execution/backends/mock_model.py b/src/twinkle/server/model/backends/mock_model.py similarity index 99% rename from src/twinkle/server/training/execution/backends/mock_model.py rename to src/twinkle/server/model/backends/mock_model.py index b2685f3ca..7654c9eeb 100644 --- a/src/twinkle/server/training/execution/backends/mock_model.py +++ b/src/twinkle/server/model/backends/mock_model.py @@ -181,7 +181,7 @@ def save(self, name: str = 'latest', output_dir: str | None = None, **kwargs: An The real backend creates ``//`` on disk and returns its path. Tinker /asample later does ``os.path.exists(...)`` on a symlink - pointing here (see ``sampling/execution/tinker_handlers.py``), so the directory + pointing here (see ``sampler/tinker_handlers.py``), so the directory has to really exist for the mock path to clear. """ target = os.path.join(output_dir or '/tmp/twinkle_mock_ckpt', name) diff --git a/src/twinkle/server/training/execution/backends/transformers_model.py b/src/twinkle/server/model/backends/transformers_model.py similarity index 96% rename from src/twinkle/server/training/execution/backends/transformers_model.py rename to src/twinkle/server/model/backends/transformers_model.py index abf6c31dc..db42aa7d0 100644 --- a/src/twinkle/server/training/execution/backends/transformers_model.py +++ b/src/twinkle/server/model/backends/transformers_model.py @@ -19,9 +19,9 @@ from twinkle.infra import collect_tensor_dict from twinkle.model import MultiLoraTransformersModel, SpectralHybridTransformersModel from twinkle.model.transformers import TransformersModel +from twinkle.server.model.backends.common import (TwinkleCompatModelBase, clean_metrics, + collect_forward_backward_results, to_cpu_safe_output) from twinkle.server.model.tinker_datum import datum_to_input_feature, extract_rl_features_for_loss -from twinkle.server.training.execution.backends.common import (TwinkleCompatModelBase, clean_metrics, - collect_forward_backward_results, to_cpu_safe_output) class _TransformersTinkerCompatMixin(TwinkleCompatModelBase): diff --git a/src/twinkle/server/training/execution/data_plane_inputs.py b/src/twinkle/server/model/data_plane_inputs.py similarity index 100% rename from src/twinkle/server/training/execution/data_plane_inputs.py rename to src/twinkle/server/model/data_plane_inputs.py diff --git a/src/twinkle/server/model/pool/__init__.py b/src/twinkle/server/model/pool/__init__.py new file mode 100644 index 000000000..720c35471 --- /dev/null +++ b/src/twinkle/server/model/pool/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Training instance pool: routing, lifecycle decisions and Ray composition.""" diff --git a/src/twinkle/server/training/actor.py b/src/twinkle/server/model/pool/actor.py similarity index 100% rename from src/twinkle/server/training/actor.py rename to src/twinkle/server/model/pool/actor.py diff --git a/src/twinkle/server/training/controller.py b/src/twinkle/server/model/pool/controller.py similarity index 100% rename from src/twinkle/server/training/controller.py rename to src/twinkle/server/model/pool/controller.py diff --git a/src/twinkle/server/training/driver.py b/src/twinkle/server/model/pool/driver.py similarity index 98% rename from src/twinkle/server/training/driver.py rename to src/twinkle/server/model/pool/driver.py index 4106d68d9..34cbd8b6c 100644 --- a/src/twinkle/server/training/driver.py +++ b/src/twinkle/server/model/pool/driver.py @@ -45,7 +45,7 @@ async def allocate(self, instance, spec): async def deploy(self, instance, spec): from ray.util.placement_group import get_placement_group - from twinkle.server.training.execution.app import build_model_app + from twinkle.server.model.app import build_model_app name = self.application_name(instance) context = TrainingRuntimeContext( diff --git a/src/twinkle/server/training/router.py b/src/twinkle/server/model/pool/router.py similarity index 100% rename from src/twinkle/server/training/router.py rename to src/twinkle/server/model/pool/router.py diff --git a/src/twinkle/server/training/execution/runtime.py b/src/twinkle/server/model/runtime.py similarity index 100% rename from src/twinkle/server/training/execution/runtime.py rename to src/twinkle/server/model/runtime.py diff --git a/src/twinkle/server/training/execution/tinker_handlers.py b/src/twinkle/server/model/tinker_handlers.py similarity index 100% rename from src/twinkle/server/training/execution/tinker_handlers.py rename to src/twinkle/server/model/tinker_handlers.py diff --git a/src/twinkle/server/training/execution/twinkle_handlers.py b/src/twinkle/server/model/twinkle_handlers.py similarity index 99% rename from src/twinkle/server/training/execution/twinkle_handlers.py rename to src/twinkle/server/model/twinkle_handlers.py index 4bd630689..b8537e8b7 100644 --- a/src/twinkle/server/training/execution/twinkle_handlers.py +++ b/src/twinkle/server/model/twinkle_handlers.py @@ -25,8 +25,8 @@ from twinkle.server.checkpoint import create_checkpoint_manager from twinkle.server.exceptions import RequestRejectedError from twinkle.server.lifecycle.submit import backend_kwargs, input_metrics, run_submit, to_backend_inputs -from twinkle.server.training.execution.data_plane_inputs import (data_plane_request_shape, merge_forward_kwargs, - resolve_data_plane_model_inputs, select_output_rows) +from twinkle.server.model.data_plane_inputs import (data_plane_request_shape, merge_forward_kwargs, + resolve_data_plane_model_inputs, select_output_rows) from twinkle.server.validation import BackendCapability from twinkle.utils.logger import get_logger diff --git a/src/twinkle/server/processor/app.py b/src/twinkle/server/processor/app.py index c7a7de209..b471cede8 100644 --- a/src/twinkle/server/processor/app.py +++ b/src/twinkle/server/processor/app.py @@ -5,7 +5,7 @@ Provides a Ray Serve deployment for managing distributed processors (datasets, dataloaders, preprocessors, rewards, templates, weight loaders, etc.). -Follows the same structural pattern as training/execution/app.py: +Follows the same structural pattern as model/app.py: - ProcessorManagement is a top-level class inheriting ProcessorManagerMixin - Routes are registered in build_processor_app() via _register_processor_routes() - serve.ingress(app)(ProcessorManagement) applied before deployment diff --git a/src/twinkle/server/sampling/execution/__init__.py b/src/twinkle/server/sampler/__init__.py similarity index 100% rename from src/twinkle/server/sampling/execution/__init__.py rename to src/twinkle/server/sampler/__init__.py diff --git a/src/twinkle/server/sampling/execution/app.py b/src/twinkle/server/sampler/app.py similarity index 100% rename from src/twinkle/server/sampling/execution/app.py rename to src/twinkle/server/sampler/app.py diff --git a/src/twinkle/server/sampling/execution/backends/__init__.py b/src/twinkle/server/sampler/backends/__init__.py similarity index 100% rename from src/twinkle/server/sampling/execution/backends/__init__.py rename to src/twinkle/server/sampler/backends/__init__.py diff --git a/src/twinkle/server/sampling/execution/backends/mock_sampler.py b/src/twinkle/server/sampler/backends/mock_sampler.py similarity index 100% rename from src/twinkle/server/sampling/execution/backends/mock_sampler.py rename to src/twinkle/server/sampler/backends/mock_sampler.py diff --git a/src/twinkle/server/sampling/execution/backends/streaming.py b/src/twinkle/server/sampler/backends/streaming.py similarity index 100% rename from src/twinkle/server/sampling/execution/backends/streaming.py rename to src/twinkle/server/sampler/backends/streaming.py diff --git a/src/twinkle/server/sampling/execution/backends/vllm_async.py b/src/twinkle/server/sampler/backends/vllm_async.py similarity index 100% rename from src/twinkle/server/sampling/execution/backends/vllm_async.py rename to src/twinkle/server/sampler/backends/vllm_async.py diff --git a/src/twinkle/server/sampling/execution/concurrency.py b/src/twinkle/server/sampler/concurrency.py similarity index 100% rename from src/twinkle/server/sampling/execution/concurrency.py rename to src/twinkle/server/sampler/concurrency.py diff --git a/src/twinkle/server/sampler/pool/__init__.py b/src/twinkle/server/sampler/pool/__init__.py new file mode 100644 index 000000000..78ee5693d --- /dev/null +++ b/src/twinkle/server/sampler/pool/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Sampling instance pool: routing, lifecycle decisions and Ray composition.""" diff --git a/src/twinkle/server/sampling/actor.py b/src/twinkle/server/sampler/pool/actor.py similarity index 100% rename from src/twinkle/server/sampling/actor.py rename to src/twinkle/server/sampler/pool/actor.py diff --git a/src/twinkle/server/sampling/controller.py b/src/twinkle/server/sampler/pool/controller.py similarity index 99% rename from src/twinkle/server/sampling/controller.py rename to src/twinkle/server/sampler/pool/controller.py index f79cb6106..45f8f5aca 100644 --- a/src/twinkle/server/sampling/controller.py +++ b/src/twinkle/server/sampler/pool/controller.py @@ -20,7 +20,7 @@ from twinkle.server.config.resources import ResourceOptions from twinkle.server.exceptions import ResourceExhaustedError from twinkle.server.resources.instances import InstanceLifecycle -from twinkle.server.sampling.execution.concurrency import SamplerConcurrency +from twinkle.server.sampler.concurrency import SamplerConcurrency logger = logging.getLogger(__name__) _CONTROL_ENDPOINTS = frozenset( diff --git a/src/twinkle/server/sampling/driver.py b/src/twinkle/server/sampler/pool/driver.py similarity index 98% rename from src/twinkle/server/sampling/driver.py rename to src/twinkle/server/sampler/pool/driver.py index cad3e85c2..e09652f27 100644 --- a/src/twinkle/server/sampling/driver.py +++ b/src/twinkle/server/sampler/pool/driver.py @@ -29,7 +29,7 @@ def prefix(self): async def start(self, instance): from ray.util.placement_group import get_placement_group - from twinkle.server.sampling.execution.app import build_sampler_app + from twinkle.server.sampler.app import build_sampler_app group = self.args.device_group device_type = str(group.get('device_type', 'cpu' if self.args.sampler_type == 'mock' else 'cuda')).lower() allocation = await self.runtime.pool.allocate( diff --git a/src/twinkle/server/sampling/router.py b/src/twinkle/server/sampler/pool/router.py similarity index 98% rename from src/twinkle/server/sampling/router.py rename to src/twinkle/server/sampler/pool/router.py index eabb3a015..138c055d1 100644 --- a/src/twinkle/server/sampling/router.py +++ b/src/twinkle/server/sampler/pool/router.py @@ -11,7 +11,7 @@ from twinkle.server.deployment import bind_deployment, build_deployment_app from twinkle.server.lifecycle.proxy import ForwardedRequest from twinkle.server.middleware.auth import get_token_from_request -from twinkle.server.sampling.execution.weights import is_full_checkpoint +from twinkle.server.sampler.weights import is_full_checkpoint from twinkle.server.state import get_server_state from .controller import sampler_controller_name diff --git a/src/twinkle/server/sampling/execution/runtime.py b/src/twinkle/server/sampler/runtime.py similarity index 100% rename from src/twinkle/server/sampling/execution/runtime.py rename to src/twinkle/server/sampler/runtime.py diff --git a/src/twinkle/server/sampling/execution/tinker_handlers.py b/src/twinkle/server/sampler/tinker_handlers.py similarity index 98% rename from src/twinkle/server/sampling/execution/tinker_handlers.py rename to src/twinkle/server/sampler/tinker_handlers.py index ac585bad7..b6a6543ed 100644 --- a/src/twinkle/server/sampling/execution/tinker_handlers.py +++ b/src/twinkle/server/sampler/tinker_handlers.py @@ -21,7 +21,7 @@ from twinkle.data_format import SamplingParams from twinkle.server.checkpoint import create_checkpoint_manager from twinkle.server.exceptions import RequestRejectedError -from twinkle.server.sampling.execution.weights import sample_with_weights +from twinkle.server.sampler.weights import sample_with_weights from twinkle.utils.logger import get_logger from .runtime import request_execution_scope diff --git a/src/twinkle/server/sampling/execution/twinkle_handlers.py b/src/twinkle/server/sampler/twinkle_handlers.py similarity index 99% rename from src/twinkle/server/sampling/execution/twinkle_handlers.py rename to src/twinkle/server/sampler/twinkle_handlers.py index 92fc8dc5f..a37f6475c 100644 --- a/src/twinkle/server/sampling/execution/twinkle_handlers.py +++ b/src/twinkle/server/sampler/twinkle_handlers.py @@ -28,7 +28,7 @@ from twinkle.protocol.types import sampler as sampler_types from twinkle.server.exceptions import EndpointUnavailableError, RequestRejectedError from twinkle.server.lifecycle.submit import backend_kwargs, resolve_twinkle_adapter_name, to_backend_inputs -from twinkle.server.sampling.execution.weights import resolve_sampler_weights, sample_with_weights +from twinkle.server.sampler.weights import resolve_sampler_weights, sample_with_weights from twinkle.server.task_errors import task_error_payload from twinkle.server.telemetry.correlation import MODEL_ID from twinkle.server.telemetry.tracing import traced_operation diff --git a/src/twinkle/server/sampling/execution/weights.py b/src/twinkle/server/sampler/weights.py similarity index 100% rename from src/twinkle/server/sampling/execution/weights.py rename to src/twinkle/server/sampler/weights.py diff --git a/src/twinkle/server/sampling/__init__.py b/src/twinkle/server/sampling/__init__.py deleted file mode 100644 index f088c2670..000000000 --- a/src/twinkle/server/sampling/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Sampling orchestration with backend execution in the execution subpackage.""" diff --git a/src/twinkle/server/training/__init__.py b/src/twinkle/server/training/__init__.py deleted file mode 100644 index 33dc7caf2..000000000 --- a/src/twinkle/server/training/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""On-demand training: catalog, lifecycle controller and runtime ownership.""" diff --git a/src/twinkle/server/training/execution/__init__.py b/src/twinkle/server/training/execution/__init__.py deleted file mode 100644 index 1b784496b..000000000 --- a/src/twinkle/server/training/execution/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Model deployment package. - -``build_model_app`` is exposed lazily via ``__getattr__`` so importing the -package doesn't eagerly pull in ``app.py`` and its transitive dependencies. -""" -from __future__ import annotations - -__all__ = ['build_model_app'] - - -def __getattr__(name: str): - if name == 'build_model_app': - from .app import build_model_app - - return build_model_app - raise AttributeError(f'module {__name__!r} has no attribute {name!r}') diff --git a/tests/server/config/test_server_config.py b/tests/server/config/test_server_config.py index 4ce5d5704..9f507fa61 100644 --- a/tests/server/config/test_server_config.py +++ b/tests/server/config/test_server_config.py @@ -21,7 +21,7 @@ from twinkle.server.config import ApplicationSpec, ServerConfig from twinkle.server.exceptions import ConfigParseError from twinkle.server.launcher import ServerLauncher -from twinkle.server.sampling.execution.app import SAMPLER_SELECTOR, build_sampler_app +from twinkle.server.sampler.app import SAMPLER_SELECTOR, build_sampler_app # ---------- minimal valid config strategy ---------------------------------- # diff --git a/tests/server/config/test_training_config.py b/tests/server/config/test_training_config.py index 35ecba6f0..11eb9bf5a 100644 --- a/tests/server/config/test_training_config.py +++ b/tests/server/config/test_training_config.py @@ -223,7 +223,7 @@ def test_model_yaml_rejects_alternative_resource_and_mode_syntax(extra): def test_configured_options_reach_cpu_router_without_changing_yaml(monkeypatch): from types import SimpleNamespace from twinkle.server.launcher.server_launcher import ServerLauncher - from twinkle.server.training.actor import TrainingControllerActor + from twinkle.server.model.pool.actor import TrainingControllerActor import ray from ray import serve config = ServerConfig(applications=[configured_application()]) @@ -237,7 +237,7 @@ def router(*args): captured.update(args[-1]) return 'app' - monkeypatch.setattr('twinkle.server.training.router.build_training_router', router) + monkeypatch.setattr('twinkle.server.model.pool.router.build_training_router', router) monkeypatch.setattr(serve, 'run', lambda *args, **kwargs: None) ServerLauncher(config)._deploy_training() assert captured['max_ongoing_requests'] == 17 @@ -253,7 +253,7 @@ def router(*args): async def test_configured_options_reach_runtime_with_one_controller_owned_replica(monkeypatch): import importlib from twinkle.server.state.training_records import InstanceRecord - from twinkle.server.training.driver import RayInstanceDriver + from twinkle.server.model.pool.driver import RayInstanceDriver from ray import serve config = ServerConfig(applications=[configured_application()]) model = config.training_config.models[0] @@ -266,7 +266,7 @@ def builder(**kwargs): captured.update(kwargs) return 'app' - monkeypatch.setattr('twinkle.server.training.execution.app.build_model_app', builder) + monkeypatch.setattr('twinkle.server.model.app.build_model_app', builder) monkeypatch.setattr(importlib.import_module('ray.util.placement_group'), 'get_placement_group', lambda name: 'pg') monkeypatch.setattr(serve, 'run', lambda *args, **kwargs: None) await driver.deploy(instance, driver.config.get(model.name, 'lora')) diff --git a/tests/server/contract/client_api_harness.py b/tests/server/contract/client_api_harness.py index 479b36049..ab82e0582 100644 --- a/tests/server/contract/client_api_harness.py +++ b/tests/server/contract/client_api_harness.py @@ -66,8 +66,8 @@ def build_gateway_app() -> FastAPI: def build_model_app() -> FastAPI: - from twinkle.server.training.execution.tinker_handlers import _register_model_tinker_routes - from twinkle.server.training.execution.twinkle_handlers import _register_model_twinkle_routes + from twinkle.server.model.tinker_handlers import _register_model_tinker_routes + from twinkle.server.model.twinkle_handlers import _register_model_twinkle_routes app = FastAPI() _register_model_tinker_routes(app, _noop_self) @@ -76,8 +76,8 @@ def build_model_app() -> FastAPI: def build_sampler_app() -> FastAPI: - from twinkle.server.sampling.execution.tinker_handlers import _register_tinker_sampler_routes - from twinkle.server.sampling.execution.twinkle_handlers import _register_twinkle_sampler_routes + from twinkle.server.sampler.tinker_handlers import _register_tinker_sampler_routes + from twinkle.server.sampler.twinkle_handlers import _register_twinkle_sampler_routes app = FastAPI() _register_tinker_sampler_routes(app, _noop_self) diff --git a/tests/server/integration/on_demand_training/test_resident_release_e2e.py b/tests/server/integration/on_demand_training/test_resident_release_e2e.py index 198045607..a7879000f 100644 --- a/tests/server/integration/on_demand_training/test_resident_release_e2e.py +++ b/tests/server/integration/on_demand_training/test_resident_release_e2e.py @@ -49,7 +49,7 @@ async def delayed(instance, training): def worker_origin(actor): import twinkle - from twinkle.server.training import controller + from twinkle.server.model.pool import controller return {'python': sys.executable, 'twinkle': twinkle.__file__, 'controller': controller.__file__, 'visible_devices': os.environ.get('CUDA_VISIBLE_DEVICES')} diff --git a/tests/server/integration/sampler/sampling_probe.py b/tests/server/integration/sampler/sampling_probe.py index f335c0a3e..6e32c6ba4 100644 --- a/tests/server/integration/sampler/sampling_probe.py +++ b/tests/server/integration/sampler/sampling_probe.py @@ -11,8 +11,8 @@ from ray.util.placement_group import placement_group_table from twinkle.server.config.application_spec import SamplerArgs -from twinkle.server.sampling.controller import SamplerController -from twinkle.server.sampling.driver import SamplerDriver +from twinkle.server.sampler.pool.controller import SamplerController +from twinkle.server.sampler.pool.driver import SamplerDriver async def run(namespace): diff --git a/tests/server/integration/sampler/test_real_pool_e2e.py b/tests/server/integration/sampler/test_real_pool_e2e.py index 49e7d052b..36b5d35fb 100644 --- a/tests/server/integration/sampler/test_real_pool_e2e.py +++ b/tests/server/integration/sampler/test_real_pool_e2e.py @@ -72,8 +72,8 @@ async def exercise(model, directory, namespace, evidence): from ray.util.placement_group import placement_group_table from transformers import AutoTokenizer from twinkle.server.config.application_spec import SamplerArgs - from twinkle.server.sampling.controller import SamplerController - from twinkle.server.sampling.driver import SamplerDriver + from twinkle.server.sampler.pool.controller import SamplerController + from twinkle.server.sampler.pool.driver import SamplerDriver from twinkle.server.state import get_server_state args = SamplerArgs(model_id=str(model), sampler_type='vllm', diff --git a/tests/server/integration/test_actor_recovery.py b/tests/server/integration/test_actor_recovery.py index 26aa3725a..dfc03a74b 100644 --- a/tests/server/integration/test_actor_recovery.py +++ b/tests/server/integration/test_actor_recovery.py @@ -9,8 +9,8 @@ import pytest from fastapi import FastAPI -from twinkle.server.training.execution.app import ModelManagement -from twinkle.server.training.execution.twinkle_handlers import _register_model_twinkle_routes +from twinkle.server.model.app import ModelManagement +from twinkle.server.model.twinkle_handlers import _register_model_twinkle_routes class _MockModel: diff --git a/tests/server/model/test_full_mode_exclusivity.py b/tests/server/model/test_full_mode_exclusivity.py index 70888e335..2c6ad3635 100644 --- a/tests/server/model/test_full_mode_exclusivity.py +++ b/tests/server/model/test_full_mode_exclusivity.py @@ -13,7 +13,7 @@ from tests.server.state.fake_backend import FakeBackend from tests.server.training.helpers import request, setup from twinkle.server.config.training import TrainingConfig -from twinkle.server.training.router import TrainingRouter +from twinkle.server.model.pool.router import TrainingRouter async def _setup_full_instance(): diff --git a/tests/server/model/test_mock_model.py b/tests/server/model/test_mock_model.py index 8dde9e6ab..5564d4344 100644 --- a/tests/server/model/test_mock_model.py +++ b/tests/server/model/test_mock_model.py @@ -15,8 +15,8 @@ from hypothesis import strategies as st from twinkle.server.exceptions import ConfigError -from twinkle.server.training.execution.app import MODEL_SELECTOR -from twinkle.server.training.execution.backends.mock_model import TwinkleCompatMockModel +from twinkle.server.model.app import MODEL_SELECTOR +from twinkle.server.model.backends.mock_model import TwinkleCompatMockModel _MODEL_BACKENDS = tuple(MODEL_SELECTOR.builders) diff --git a/tests/server/model/test_preflight_loss_inputs.py b/tests/server/model/test_preflight_loss_inputs.py index 6a6df4946..6c76c9354 100644 --- a/tests/server/model/test_preflight_loss_inputs.py +++ b/tests/server/model/test_preflight_loss_inputs.py @@ -7,10 +7,10 @@ import pytest from twinkle.loss.grpo import GRPOLoss -from twinkle.server.training.execution.backends import common -from twinkle.server.training.execution.backends.common import TwinkleCompatModelBase -from twinkle.server.training.execution.backends.megatron_model import (TwinkleCompatFullMegatronModel, TwinkleCompatMegatronModel) -from twinkle.server.training.execution.backends.transformers_model import (TwinkleCompatFullTransformersModel, +from twinkle.server.model.backends import common +from twinkle.server.model.backends.common import TwinkleCompatModelBase +from twinkle.server.model.backends.megatron_model import (TwinkleCompatFullMegatronModel, TwinkleCompatMegatronModel) +from twinkle.server.model.backends.transformers_model import (TwinkleCompatFullTransformersModel, TwinkleCompatTransformersModel) diff --git a/tests/server/model/test_replica_lifecycle.py b/tests/server/model/test_replica_lifecycle.py index 2849f427f..63b9ccaa8 100644 --- a/tests/server/model/test_replica_lifecycle.py +++ b/tests/server/model/test_replica_lifecycle.py @@ -5,7 +5,7 @@ import pytest -from twinkle.server.training.execution.app import MODEL_SELECTOR, ModelManagement +from twinkle.server.model.app import MODEL_SELECTOR, ModelManagement class _CapacityState: @@ -67,10 +67,10 @@ async def test_async_constructor_registers_replica_before_ready() -> None: manager = ModelManagement.__new__(ModelManagement) manager._task_queue_config = SimpleNamespace(effective_execution_timeout=1800.0) - with patch('twinkle.server.training.execution.app.DeviceGroup', return_value=SimpleNamespace(name='group')), \ - patch('twinkle.server.training.execution.app.init_twinkle_runtime', return_value=None), \ + with patch('twinkle.server.model.app.DeviceGroup', return_value=SimpleNamespace(name='group')), \ + patch('twinkle.server.model.app.init_twinkle_runtime', return_value=None), \ patch.object(MODEL_SELECTOR, 'construct', return_value=MagicMock()), \ - patch('twinkle.server.training.execution.app.get_server_state', return_value=state), \ + patch('twinkle.server.model.app.get_server_state', return_value=state), \ patch('twinkle.server.data_plane.DataPlaneProxy', return_value=MagicMock()), \ patch.object(ModelManagement, '_init_task_queue'), \ patch.object(ModelManagement, 'initialize_training_runtime', new_callable=AsyncMock), \ diff --git a/tests/server/model/test_tinker_compat_output.py b/tests/server/model/test_tinker_compat_output.py index 6a5f393eb..c6e95ec50 100644 --- a/tests/server/model/test_tinker_compat_output.py +++ b/tests/server/model/test_tinker_compat_output.py @@ -4,7 +4,7 @@ from unittest.mock import Mock from twinkle.server.model.tinker_datum import extract_rl_features_for_loss -from twinkle.server.training.execution.backends.common import TwinkleCompatModelBase +from twinkle.server.model.backends.common import TwinkleCompatModelBase from twinkle.server.exceptions import RequestRejectedError diff --git a/tests/server/model/test_tinker_handlers.py b/tests/server/model/test_tinker_handlers.py index 55d4a8e84..2df46bb7e 100644 --- a/tests/server/model/test_tinker_handlers.py +++ b/tests/server/model/test_tinker_handlers.py @@ -5,13 +5,13 @@ from tinker import types from unittest.mock import ANY, AsyncMock, MagicMock -from twinkle.server.training.execution.app import ModelManagement -from twinkle.server.training.execution.tinker_handlers import _register_model_tinker_routes +from twinkle.server.model.app import ModelManagement +from twinkle.server.model.tinker_handlers import _register_model_tinker_routes from twinkle.server.deployment import twinkle_server_error_handler from twinkle.server.exceptions import BatchSizeError, TwinkleServerError from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.task_queue.mixin import TaskQueueMixin -from twinkle.server.training.execution.runtime import TrainingExecutionState +from twinkle.server.model.runtime import TrainingExecutionState class _DummyManagement: @@ -220,7 +220,7 @@ async def call_backend(self, fn, /, *args, **kwargs): ids=['lora', 'full-model', 'backend-failure']) async def test_training_save_records_metadata_only_after_backend_success(monkeypatch, tmp_path, fail_backend, full_mode): from pathlib import Path - import twinkle.server.training.execution.tinker_handlers as handlers + import twinkle.server.model.tinker_handlers as handlers from twinkle.server.checkpoint.tinker import TinkerCheckpointManager, TinkerTrainingRunManager monkeypatch.setattr('twinkle.server.checkpoint.training_run_manager.TWINKLE_DEFAULT_SAVE_DIR', str(tmp_path)) @@ -262,7 +262,7 @@ def save(name, output_dir, **kwargs): @pytest.mark.asyncio @pytest.mark.parametrize('path,output_dir', [('org/model', None), ('twinkle://model1/weights/step', 'local')]) async def test_load_weights_resolves_location_and_preserves_optimizer_flag(monkeypatch, tmp_path, path, output_dir): - import twinkle.server.training.execution.tinker_handlers as handlers + import twinkle.server.model.tinker_handlers as handlers from twinkle.server.checkpoint.tinker import TinkerCheckpointManager, TinkerTrainingRunManager monkeypatch.setattr('twinkle.server.checkpoint.training_run_manager.TWINKLE_DEFAULT_SAVE_DIR', str(tmp_path)) manager = TinkerCheckpointManager('token1', TinkerTrainingRunManager('token1')) @@ -321,7 +321,7 @@ async def test_tinker_loss_metric_survives_sdk_reduction(backward): async def test_sampler_save_materializes_weights_and_returns_sdk_mode(monkeypatch, tmp_path, name, sequence_id, full_mode): from pathlib import Path from twinkle.server.checkpoint import create_checkpoint_manager - import twinkle.server.training.execution.tinker_handlers as handlers + import twinkle.server.model.tinker_handlers as handlers monkeypatch.setattr('twinkle.server.checkpoint.training_run_manager.TWINKLE_DEFAULT_SAVE_DIR', str(tmp_path)) manager = create_checkpoint_manager('token1', client_type='tinker') manager.training_run_manager._write_info('model1', {'base_model': 'test-model'}) diff --git a/tests/server/model/test_twinkle_async_inputs.py b/tests/server/model/test_twinkle_async_inputs.py index 8a97bc8c4..38cefbe85 100644 --- a/tests/server/model/test_twinkle_async_inputs.py +++ b/tests/server/model/test_twinkle_async_inputs.py @@ -5,8 +5,8 @@ from starlette.requests import Request import twinkle.protocol.types as types -from twinkle.server.training.execution.data_plane_inputs import model_result_rows -from twinkle.server.training.execution.twinkle_handlers import _register_model_twinkle_routes +from twinkle.server.model.data_plane_inputs import model_result_rows +from twinkle.server.model.twinkle_handlers import _register_model_twinkle_routes def test_model_result_rows_keeps_one_output_row_per_sample() -> None: @@ -161,8 +161,8 @@ async def test_save_forwards_backend_options_and_returns_public_path( monkeypatch, tmp_path, is_sampler, save_optimizer, full_mode): from pathlib import Path from twinkle.server.checkpoint.twinkle import TwinkleCheckpointManager, TwinkleTrainingRunManager - from twinkle.server.training.execution.app import ModelManagement - import twinkle.server.training.execution.twinkle_handlers as handlers + from twinkle.server.model.app import ModelManagement + import twinkle.server.model.twinkle_handlers as handlers class Management(_SchedulingManagement): is_full_mode = ModelManagement.is_full_mode @@ -203,7 +203,7 @@ def save(self, name, output_dir, **kwargs): @pytest.mark.asyncio async def test_training_save_failure_does_not_publish_checkpoint(monkeypatch, tmp_path): from twinkle.server.checkpoint.twinkle import TwinkleCheckpointManager, TwinkleTrainingRunManager - import twinkle.server.training.execution.twinkle_handlers as handlers + import twinkle.server.model.twinkle_handlers as handlers class Management(_SchedulingManagement): def save(self, **kwargs): @@ -234,7 +234,7 @@ def save(self, **kwargs): @pytest.mark.parametrize('local', [False, True], ids=['external', 'checkpoint-uri']) async def test_checkpoint_consumers_keep_backend_path_and_options(monkeypatch, tmp_path, operation, local): from unittest.mock import MagicMock - import twinkle.server.training.execution.twinkle_handlers as handlers + import twinkle.server.model.twinkle_handlers as handlers from twinkle.server.checkpoint.twinkle import TwinkleCheckpointManager, TwinkleTrainingRunManager monkeypatch.setattr('twinkle.server.checkpoint.training_run_manager.TWINKLE_DEFAULT_SAVE_DIR', str(tmp_path)) diff --git a/tests/server/resources/test_instances.py b/tests/server/resources/test_instances.py index f7eaed80d..2fae3d280 100644 --- a/tests/server/resources/test_instances.py +++ b/tests/server/resources/test_instances.py @@ -128,8 +128,8 @@ async def test_serve_handle_lookup_cannot_block_driver_deadlines(monkeypatch, ki from types import SimpleNamespace from ray import serve from twinkle.server.resources.runtime import RayRuntime - from twinkle.server.sampling.driver import SamplerDriver - from twinkle.server.training.driver import RayInstanceDriver + from twinkle.server.sampler.pool.driver import SamplerDriver + from twinkle.server.model.pool.driver import RayInstanceDriver entered, resume, returned = threading.Event(), threading.Event(), threading.Event() diff --git a/tests/server/sampler/test_mock_sampler.py b/tests/server/sampler/test_mock_sampler.py index d9ca81066..f39828428 100644 --- a/tests/server/sampler/test_mock_sampler.py +++ b/tests/server/sampler/test_mock_sampler.py @@ -17,8 +17,8 @@ from twinkle.data_format import InputFeature, SamplingParams from twinkle.server.exceptions import ConfigError -from twinkle.server.sampling.execution.app import SAMPLER_SELECTOR, _construct_sampler_backend -from twinkle.server.sampling.execution.backends.mock_sampler import MockSampler +from twinkle.server.sampler.app import SAMPLER_SELECTOR, _construct_sampler_backend +from twinkle.server.sampler.backends.mock_sampler import MockSampler _SAMPLER_TYPES = tuple(SAMPLER_SELECTOR.builders) @@ -121,7 +121,7 @@ def test_explicit_async_vllm_uses_non_blocking_sampler(monkeypatch) -> None: # in the `async-rl` extra, so skip when it is absent instead of failing a sampler test # for a data-path dependency it does not use. pytest.importorskip('transfer_queue') - from twinkle.server.sampling.execution.backends import vllm_async as module + from twinkle.server.sampler.backends import vllm_async as module captured = {} @@ -268,7 +268,7 @@ def test_server_concurrency_does_not_change_native_async_rl_defaults(monkeypatch import twinkle.infra as infra from twinkle.infra._ray import RayHelper from twinkle.sampler.vllm_sampler.vllm_sampler import SAMPLER_MAX_CONCURRENCY - from twinkle.server.sampling.execution.backends.vllm_async import ServerVLLMSampler + from twinkle.server.sampler.backends.vllm_async import ServerVLLMSampler from twinkle.utils.device_mesh import DeviceMesh from twinkle_agentic.async_rl.vllm_sampler_tq import VLLMSamplerTQ diff --git a/tests/server/sampler/test_pool.py b/tests/server/sampler/test_pool.py index e5209aad0..6fb9e3b67 100644 --- a/tests/server/sampler/test_pool.py +++ b/tests/server/sampler/test_pool.py @@ -7,7 +7,7 @@ import pytest from twinkle.server.config.application_spec import SamplerArgs -from twinkle.server.sampling.controller import SamplerController +from twinkle.server.sampler.pool.controller import SamplerController @pytest.fixture(scope='session', autouse=True) diff --git a/tests/server/sampler/test_request_runtime.py b/tests/server/sampler/test_request_runtime.py index e45f34aae..61c8095bd 100644 --- a/tests/server/sampler/test_request_runtime.py +++ b/tests/server/sampler/test_request_runtime.py @@ -10,7 +10,7 @@ from tests.server.task_queue.helpers import _DummyQueue from twinkle.server.config.task_queue import TaskQueueConfig -from twinkle.server.sampling.execution.concurrency import SamplerConcurrency +from twinkle.server.sampler.concurrency import SamplerConcurrency from twinkle.server.task_queue.types import ExecutionScope @@ -19,7 +19,7 @@ async def test_allocated_sampler_initializes_with_typed_owner_context(monkeypatc from unittest.mock import Mock from twinkle.infra._ray.ray_helper import RayHelper from twinkle.infra._ray.resource_manager import ResourceManager - from twinkle.server.sampling.execution import app as module + from twinkle.server.sampler import app as module from twinkle.server.state.models import RuntimeContext context = RuntimeContext(namespace='owned', registry_name='registry', instance_id='sampler-owned') @@ -134,13 +134,13 @@ def close(): @pytest.mark.asyncio async def test_runtime_does_not_certify_an_unconfirmed_cancel_ack(): - from twinkle.server.sampling.execution.runtime import SamplerRuntime + from twinkle.server.sampler.runtime import SamplerRuntime from twinkle.server.exceptions import BackendBusyError guard = SamplerConcurrency() runtime = SamplerRuntime(SimpleNamespace(sampling_guard=guard)) task = asyncio.create_task(asyncio.sleep(0)) await task - from twinkle.server.sampling.execution.runtime import SamplingOperation + from twinkle.server.sampler.runtime import SamplingOperation runtime._operations['cancelled'] = SamplingOperation(ExecutionScope(), task) guard.invalidate_weights(RuntimeError('cancel ACK preceded actual generation exit')) with pytest.raises(BackendBusyError): @@ -152,7 +152,7 @@ async def test_runtime_does_not_certify_an_unconfirmed_cancel_ack(): async def test_cold_lora_config_does_not_block_sampler_health(monkeypatch): import json import sys - from twinkle.server.sampling.execution.runtime import SamplerRuntime + from twinkle.server.sampler.runtime import SamplerRuntime entered, release = threading.Event(), threading.Event() calls = [] @@ -207,7 +207,7 @@ async def test_managed_compute_http_cannot_bypass_pool_admission(monkeypatch): import httpx from fastapi import FastAPI from twinkle.server.exceptions import RequestRejectedError - from twinkle.server.sampling.execution import app as module + from twinkle.server.sampler import app as module def scaffold(title, register_routes, **kwargs): app = FastAPI() diff --git a/tests/server/sampler/test_resolve_sampler_weights.py b/tests/server/sampler/test_resolve_sampler_weights.py index 4f01d9b96..7bbf9f25d 100644 --- a/tests/server/sampler/test_resolve_sampler_weights.py +++ b/tests/server/sampler/test_resolve_sampler_weights.py @@ -4,7 +4,7 @@ import pytest -from twinkle.server.sampling.execution.weights import resolve_sampler_weights +from twinkle.server.sampler.weights import resolve_sampler_weights class _FakeSampler: @@ -17,7 +17,7 @@ class _FakeService: """Minimal stand-in exposing the two attributes the helper touches.""" def __init__(self): - from twinkle.server.sampling.execution.concurrency import SamplerConcurrency + from twinkle.server.sampler.concurrency import SamplerConcurrency self.sampling_guard = SamplerConcurrency() self.sampler = _FakeSampler() self.full_weight_loads: list[str] = [] diff --git a/tests/server/sampler/test_sampling_concurrency.py b/tests/server/sampler/test_sampling_concurrency.py index 781a17556..14b3e07e1 100644 --- a/tests/server/sampler/test_sampling_concurrency.py +++ b/tests/server/sampler/test_sampling_concurrency.py @@ -16,10 +16,10 @@ from twinkle.data_format import SampledSequence, SampleResponse from twinkle.exceptions import FullWeightLoadError from twinkle.protocol.types import SampleRequest -from twinkle.server.sampling.execution.app import SamplerManagement -from twinkle.server.sampling.execution.concurrency import SamplerConcurrency -from twinkle.server.sampling.execution.tinker_handlers import _register_tinker_sampler_routes -from twinkle.server.sampling.execution.twinkle_handlers import _register_twinkle_sampler_routes +from twinkle.server.sampler.app import SamplerManagement +from twinkle.server.sampler.concurrency import SamplerConcurrency +from twinkle.server.sampler.tinker_handlers import _register_tinker_sampler_routes +from twinkle.server.sampler.twinkle_handlers import _register_twinkle_sampler_routes from twinkle.server.task_queue.backend_gate import BackendGate from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.exceptions import BackendBusyError @@ -166,7 +166,7 @@ def create(**kwargs): @pytest.mark.asyncio async def test_standalone_checkpoint_load_waits_for_ordinary_sampling(make_service, tmp_path): - from twinkle.server.sampling.execution.weights import resolve_sampler_weights + from twinkle.server.sampler.weights import resolve_sampler_weights checkpoint = tmp_path / '2' checkpoint.mkdir() (checkpoint / 'model.safetensors').touch() @@ -401,7 +401,7 @@ def partial_load(path): @pytest.mark.asyncio async def test_completed_generation_failure_does_not_quarantine_replica(make_service): - from twinkle.server.sampling.execution.twinkle_handlers import _await_generation + from twinkle.server.sampler.twinkle_handlers import _await_generation service = make_service() service.sampler.get_generation_status = lambda _: {'status': 'failed', 'error': 'invalid input'} service.sampler.cancel_generation = lambda _: {'status': 'completed'} diff --git a/tests/server/sampler/test_stream_guarantees.py b/tests/server/sampler/test_stream_guarantees.py index 54c6e7498..ed10414d0 100644 --- a/tests/server/sampler/test_stream_guarantees.py +++ b/tests/server/sampler/test_stream_guarantees.py @@ -9,8 +9,8 @@ import pytest from fastapi import FastAPI -from twinkle.server.sampling.execution.app import SamplerManagement -from twinkle.server.sampling.execution.twinkle_handlers import _await_generation, _register_twinkle_sampler_routes, _stream_queue +from twinkle.server.sampler.app import SamplerManagement +from twinkle.server.sampler.twinkle_handlers import _await_generation, _register_twinkle_sampler_routes, _stream_queue class _BlockingQueue: @@ -89,7 +89,7 @@ async def test_stream_timeout_returns_error_payload_and_closes_queue(): class _GenerationService: def __init__(self) -> None: - from twinkle.server.sampling.execution.concurrency import SamplerConcurrency + from twinkle.server.sampler.concurrency import SamplerConcurrency self.sampling_guard = SamplerConcurrency() self.cancelled = False self.sampler = SimpleNamespace( diff --git a/tests/server/sampler/test_streaming_bridge.py b/tests/server/sampler/test_streaming_bridge.py index 96cc3c900..bba725b4f 100644 --- a/tests/server/sampler/test_streaming_bridge.py +++ b/tests/server/sampler/test_streaming_bridge.py @@ -4,7 +4,7 @@ import pytest -from twinkle.server.sampling.execution.backends.streaming import STREAM_SENTINEL, stream_to_queue +from twinkle.server.sampler.backends.streaming import STREAM_SENTINEL, stream_to_queue @pytest.mark.parametrize('disconnect', [False, True]) diff --git a/tests/server/sampler/test_tinker_handlers.py b/tests/server/sampler/test_tinker_handlers.py index 956b085d2..2f4a0aad8 100644 --- a/tests/server/sampler/test_tinker_handlers.py +++ b/tests/server/sampler/test_tinker_handlers.py @@ -5,8 +5,8 @@ from tinker import types from twinkle.data_format import SampledSequence, SampleResponse -from twinkle.server.sampling.execution.tinker_handlers import _register_tinker_sampler_routes -from twinkle.server.sampling.execution.concurrency import SamplerConcurrency +from twinkle.server.sampler.tinker_handlers import _register_tinker_sampler_routes +from twinkle.server.sampler.concurrency import SamplerConcurrency class _DummyState: diff --git a/tests/server/sampler/test_twinkle_async_rows.py b/tests/server/sampler/test_twinkle_async_rows.py index a323e64f9..407ba483a 100644 --- a/tests/server/sampler/test_twinkle_async_rows.py +++ b/tests/server/sampler/test_twinkle_async_rows.py @@ -8,7 +8,7 @@ import twinkle.protocol.types as types from twinkle.data_format import SampledSequence, SampleResponse -from twinkle.server.sampling.execution.twinkle_handlers import ( +from twinkle.server.sampler.twinkle_handlers import ( _register_twinkle_sampler_routes, _build_rollout_rows_and_tags, ) @@ -64,7 +64,7 @@ def test_async_sampler_rejects_group_id_count_mismatch() -> None: class _SamplerManagement: def __init__(self): - from twinkle.server.sampling.execution.concurrency import SamplerConcurrency + from twinkle.server.sampler.concurrency import SamplerConcurrency self.sampling_guard = SamplerConcurrency() self.sampler = self self.data_plane = self diff --git a/tests/server/static/backend_call_exemptions.py b/tests/server/static/backend_call_exemptions.py index 385901d3c..881862dd7 100644 --- a/tests/server/static/backend_call_exemptions.py +++ b/tests/server/static/backend_call_exemptions.py @@ -27,6 +27,6 @@ # (module_relpath under src/twinkle/server, innermost enclosing function name) BACKEND_CALL_EXEMPTIONS: frozenset[tuple[str, str]] = frozenset({ - ('sampling/execution/twinkle_handlers.py', '_stream_queue'), - ('sampling/execution/twinkle_handlers.py', '_stream_generation'), + ('sampler/twinkle_handlers.py', '_stream_queue'), + ('sampler/twinkle_handlers.py', '_stream_generation'), }) diff --git a/tests/server/static/test_adapter_name_mapping.py b/tests/server/static/test_adapter_name_mapping.py index a150fc53f..708bef399 100644 --- a/tests/server/static/test_adapter_name_mapping.py +++ b/tests/server/static/test_adapter_name_mapping.py @@ -9,7 +9,7 @@ the ``*_from_data_plane`` endpoints originally passed the raw tenant name, so a full-mode data-plane request drove the wrong optimizer group. -This guard fails if any ``self.call_backend(...)`` in ``training/execution/twinkle_handlers.py`` +This guard fails if any ``self.call_backend(...)`` in ``model/twinkle_handlers.py`` passes ``adapter_name`` as the bare local (i.e. unmapped). Passing the raw name *positionally* (as ``add_adapter_to_model`` does when creating a tenant adapter) is intentionally not a keyword and is therefore not flagged. @@ -21,7 +21,7 @@ import twinkle -_HANDLERS = pathlib.Path(twinkle.__file__).resolve().parent / 'server' / 'training' / 'execution' / 'twinkle_handlers.py' +_HANDLERS = pathlib.Path(twinkle.__file__).resolve().parent / 'server' / 'model' / 'twinkle_handlers.py' def _is_call_backend(func: ast.AST) -> bool: diff --git a/tests/server/static/test_no_direct_backend_call.py b/tests/server/static/test_no_direct_backend_call.py index f48a66326..89573d371 100644 --- a/tests/server/static/test_no_direct_backend_call.py +++ b/tests/server/static/test_no_direct_backend_call.py @@ -123,7 +123,7 @@ def test_no_direct_backend_call_in_server(): def test_exemptions_are_read_from_shared_file(): - assert ('sampling/execution/twinkle_handlers.py', '_stream_queue') in BACKEND_CALL_EXEMPTIONS + assert ('sampler/twinkle_handlers.py', '_stream_queue') in BACKEND_CALL_EXEMPTIONS def test_checker_detects_indirect_backend_calls(): diff --git a/tests/server/static/test_no_twinkle_http_exception.py b/tests/server/static/test_no_twinkle_http_exception.py index 5092415e5..28a67a392 100644 --- a/tests/server/static/test_no_twinkle_http_exception.py +++ b/tests/server/static/test_no_twinkle_http_exception.py @@ -6,8 +6,8 @@ _SERVER = Path(__file__).resolve().parents[3] / 'src' / 'twinkle' / 'server' _TWINKLE_HANDLERS = ( _SERVER / 'gateway' / 'twinkle_handlers.py', - _SERVER / 'training' / 'execution' / 'twinkle_handlers.py', - _SERVER / 'sampling' / 'execution' / 'twinkle_handlers.py', + _SERVER / 'model' / 'twinkle_handlers.py', + _SERVER / 'sampler' / 'twinkle_handlers.py', _SERVER / 'processor' / 'twinkle_handlers.py', ) diff --git a/tests/server/telemetry/test_training_metrics.py b/tests/server/telemetry/test_training_metrics.py index 157f2447a..7f4214593 100644 --- a/tests/server/telemetry/test_training_metrics.py +++ b/tests/server/telemetry/test_training_metrics.py @@ -6,7 +6,7 @@ import pytest from twinkle.server.telemetry import metrics -from twinkle.server.training.controller import TrainingController +from twinkle.server.model.pool.controller import TrainingController from twinkle.server.telemetry.metrics import TrainingInstanceMetrics _ALLOWED_ATTRIBUTES = {'backend', 'mode', 'state', 'status', 'reason', 'device_type', 'resident'} @@ -213,7 +213,7 @@ async def test_startup_and_shutdown_histograms_follow_real_attempts(fail_start): async def test_sweep_instruments_one_iteration_and_its_outcome(monkeypatch, failed): import asyncio from unittest.mock import AsyncMock, MagicMock - from twinkle.server.training.actor import TrainingControllerActor + from twinkle.server.model.pool.actor import TrainingControllerActor async def stop_after_iteration(delay): raise asyncio.CancelledError @@ -221,7 +221,7 @@ async def stop_after_iteration(delay): actor = SimpleNamespace(core=SimpleNamespace(sweep=AsyncMock( side_effect=RuntimeError('state unavailable') if failed else None)), config=SimpleNamespace(sweep_interval_s=5), _metrics=MagicMock()) - monkeypatch.setattr('twinkle.server.training.actor.asyncio.sleep', stop_after_iteration) + monkeypatch.setattr('twinkle.server.model.pool.actor.asyncio.sleep', stop_after_iteration) with pytest.raises(asyncio.CancelledError): await TrainingControllerActor.__ray_metadata__.modified_class._sweep(actor) tags = {'status': 'error' if failed else 'success'} @@ -236,7 +236,7 @@ async def test_driver_records_allocation_attempts(failed): from unittest.mock import AsyncMock, MagicMock from twinkle.server.config.training import TrainingConfig from twinkle.server.resources.allocation import DeviceAllocation - from twinkle.server.training.driver import RayInstanceDriver + from twinkle.server.model.pool.driver import RayInstanceDriver config = TrainingConfig(models=[{'name': 'tiny', 'model_id': 'tiny', 'backend': 'mock', 'device_type': 'cpu', 'mode': 'full', 'config': {'num_devices': 1, 'device_mesh': {'device_type': 'cpu'}}}]) diff --git a/tests/server/test_app_builders_characterization.py b/tests/server/test_app_builders_characterization.py index a1f5fb3a9..55c42b4e6 100644 --- a/tests/server/test_app_builders_characterization.py +++ b/tests/server/test_app_builders_characterization.py @@ -133,7 +133,7 @@ def test_gateway_builder_characterization(monkeypatch) -> None: def test_model_builder_characterization(monkeypatch) -> None: - from twinkle.server.training.execution import app as model_mod + from twinkle.server.model import app as model_mod res = _capture_builder( monkeypatch, @@ -152,7 +152,7 @@ def test_model_builder_characterization(monkeypatch) -> None: def test_sampler_builder_characterization(monkeypatch) -> None: - from twinkle.server.sampling.execution import app as sampler_mod + from twinkle.server.sampler import app as sampler_mod res = _capture_builder( monkeypatch, diff --git a/tests/server/test_dependencies.py b/tests/server/test_dependencies.py index 09c0a4f49..03d795da0 100644 --- a/tests/server/test_dependencies.py +++ b/tests/server/test_dependencies.py @@ -13,15 +13,15 @@ # the same layer are independent and therefore may not import one another. _LAYERS = ( frozenset({'exceptions'}), - frozenset({'config', 'middleware', 'model', 'runtime', 'task_errors', 'utils', 'validation'}), + frozenset({'config', 'middleware', 'model.tinker_datum', 'runtime', 'task_errors', 'utils', 'validation'}), frozenset({'telemetry'}), frozenset({'state'}), frozenset({'resources'}), frozenset({'deployment', 'lifecycle', 'session_resource'}), frozenset({'task_queue'}), frozenset({'checkpoint', 'data_plane'}), - frozenset({'gateway', 'training.execution', 'sampling.execution', 'processor'}), - frozenset({'sampling', 'training'}), + frozenset({'gateway', 'model', 'sampler', 'processor'}), + frozenset({'model.pool', 'sampler.pool'}), frozenset({'launcher'}), frozenset({'cli'}), frozenset({'__main__', _ROOT_COMPONENT}), @@ -32,9 +32,9 @@ def _component(module: str) -> str: if module == _ROOT: return _ROOT_COMPONENT relative = module.removeprefix(f'{_ROOT}.') - for execution in ('training.execution', 'sampling.execution'): - if relative == execution or relative.startswith(execution + '.'): - return execution + for component in ('model.pool', 'sampler.pool', 'model.tinker_datum'): + if relative == component or relative.startswith(component + '.'): + return component return relative.split('.', 1)[0] @@ -48,7 +48,7 @@ def test_lightweight_imports_do_not_load_runtime(statement): code = statement + ''' import sys assert not any(name in sys.modules for name in ( - 'torch', 'ray', 'twinkle.server.launcher', 'twinkle.server.training.actor', 'twinkle.hub.hub' + 'torch', 'ray', 'twinkle.server.launcher', 'twinkle.server.model.pool.actor', 'twinkle.hub.hub' )), sorted(sys.modules) ''' result = subprocess.run([sys.executable, '-c', code], capture_output=True, text=True, timeout=60) @@ -104,7 +104,7 @@ def test_server_components_form_declared_dag(): f'invalid edges={violations}') -@pytest.mark.parametrize('component', ['training', 'sampling']) +@pytest.mark.parametrize('component', ['model.pool', 'sampler.pool']) def test_controller_import_does_not_load_ray_composition(component): code = f'''import twinkle.server.{component}.controller import sys diff --git a/tests/server/training/helpers.py b/tests/server/training/helpers.py index e55d27c14..3b4f6aa74 100644 --- a/tests/server/training/helpers.py +++ b/tests/server/training/helpers.py @@ -11,8 +11,8 @@ from twinkle.server.state.backend.memory_backend import RayActorBackend from twinkle.server.state.training_records import CreateTraining from twinkle.server.task_queue import TaskQueueMixin -from twinkle.server.training.controller import TrainingController -from twinkle.server.training.execution.runtime import TrainingRuntime +from twinkle.server.model.pool.controller import TrainingController +from twinkle.server.model.runtime import TrainingRuntime class Clock: diff --git a/tests/server/training/test_actor_initialization.py b/tests/server/training/test_actor_initialization.py index c5a81b43e..dd9a3f530 100644 --- a/tests/server/training/test_actor_initialization.py +++ b/tests/server/training/test_actor_initialization.py @@ -7,7 +7,7 @@ import pytest from twinkle.server.resources.instances import InstanceLifecycle -from twinkle.server.training.actor import TrainingControllerActor +from twinkle.server.model.pool.actor import TrainingControllerActor @pytest.mark.asyncio @@ -25,9 +25,9 @@ async def sleep(delay): sleeping.set() await gate.wait() - monkeypatch.setattr('twinkle.server.training.actor.asyncio.sleep', sleep) - monkeypatch.setattr('twinkle.server.training.actor.ray.get_actor', get_actor) - monkeypatch.setattr('twinkle.server.training.actor.ray.kill', kill) + monkeypatch.setattr('twinkle.server.model.pool.actor.asyncio.sleep', sleep) + monkeypatch.setattr('twinkle.server.model.pool.actor.ray.get_actor', get_actor) + monkeypatch.setattr('twinkle.server.model.pool.actor.ray.kill', kill) actor = SimpleNamespace(core=SimpleNamespace(initialize=AsyncMock(side_effect=RuntimeError('allocation failed'))), config=SimpleNamespace(sweep_interval_s=7), namespace='initialization-test', _monitors=[]) actor._restart_after_failure = lambda: cls._restart_after_failure(actor) diff --git a/tests/server/training/test_controller.py b/tests/server/training/test_controller.py index df9e4fe4b..60d3098bd 100644 --- a/tests/server/training/test_controller.py +++ b/tests/server/training/test_controller.py @@ -11,7 +11,7 @@ from twinkle.server.exceptions import RequestRejectedError, ResourceNotFoundError from twinkle.server.state import ServerState from twinkle.server.state.backend.memory_backend import RayActorBackend -from twinkle.server.training.controller import TrainingController +from twinkle.server.model.pool.controller import TrainingController from twinkle.server.exceptions import TrainingExpiredError from twinkle.server.resources.allocation import CapacityUnavailableError from twinkle.server.state.training_records import CreateTraining diff --git a/tests/server/training/test_failure_recovery.py b/tests/server/training/test_failure_recovery.py index 62da32968..115a3c8fc 100644 --- a/tests/server/training/test_failure_recovery.py +++ b/tests/server/training/test_failure_recovery.py @@ -285,7 +285,7 @@ async def test_shutdown_waits_for_late_serve_deployment_before_confirming_releas from types import SimpleNamespace from twinkle.server.config.training import TrainingConfig from twinkle.server.state.training_records import InstanceRecord - from twinkle.server.training.driver import RayInstanceDriver + from twinkle.server.model.pool.driver import RayInstanceDriver from ray import serve config = TrainingConfig(models=[{'name': 'small', 'model_id': 'weights', 'backend': 'mock', 'device_type': 'cpu', 'mode': 'lora', 'config': {'num_devices': 1, 'device_mesh': {'dp_size': 1}}}]) store = SimpleNamespace(models=SimpleNamespace(unregister_replica_capacity=AsyncMock())) @@ -300,7 +300,7 @@ def deploy(*args, **kwargs): if not release.wait(10): raise TimeoutError('test did not release Serve') - monkeypatch.setattr('twinkle.server.training.execution.app.build_model_app', lambda **kwargs: 'app') + monkeypatch.setattr('twinkle.server.model.app.build_model_app', lambda **kwargs: 'app') monkeypatch.setattr(importlib.import_module('ray.util.placement_group'), 'get_placement_group', lambda name: 'pg') monkeypatch.setattr('ray.util.list_named_actors', lambda **kwargs: []) monkeypatch.setattr(serve, 'run', deploy) diff --git a/tests/server/training/test_launcher.py b/tests/server/training/test_launcher.py index 8ee8b40be..7ed539517 100644 --- a/tests/server/training/test_launcher.py +++ b/tests/server/training/test_launcher.py @@ -9,7 +9,7 @@ from twinkle.server.config import ServerConfig from twinkle.server.launcher.server_launcher import ServerLauncher -from twinkle.server.training.actor import controller_name +from twinkle.server.model.pool.actor import controller_name def test_unresponsive_controller_does_not_block_other_cleanup(): diff --git a/tests/server/training/test_router.py b/tests/server/training/test_router.py index da5271f40..52e4b1c7c 100644 --- a/tests/server/training/test_router.py +++ b/tests/server/training/test_router.py @@ -10,7 +10,7 @@ from tests.server.training.helpers import settle from twinkle.hub import HubOperation from twinkle.server.state import ServerState -from twinkle.server.training.router import TrainingRouter +from twinkle.server.model.pool.router import TrainingRouter from twinkle.server.config.training import TrainingConfig @@ -84,9 +84,9 @@ def cpu_import(name, *args, **kwargs): def test_router_schemas_match_model_runtime_routes(): from fastapi import FastAPI - from twinkle.server.training.execution.tinker_handlers import _register_model_tinker_routes - from twinkle.server.training.execution.twinkle_handlers import _register_model_twinkle_routes - from twinkle.server.training.router import NATIVE_REQUESTS, TINKER_REQUESTS + from twinkle.server.model.tinker_handlers import _register_model_tinker_routes + from twinkle.server.model.twinkle_handlers import _register_model_twinkle_routes + from twinkle.server.model.pool.router import NATIVE_REQUESTS, TINKER_REQUESTS app = FastAPI() _register_model_tinker_routes(app, lambda: None) diff --git a/tests/server/training/test_runtime_initialization.py b/tests/server/training/test_runtime_initialization.py index 94b25c0c9..b136d4fc8 100644 --- a/tests/server/training/test_runtime_initialization.py +++ b/tests/server/training/test_runtime_initialization.py @@ -9,7 +9,7 @@ from tests.server.state.fake_backend import FakeBackend from tests.server.training.helpers import create, setup -from twinkle.server.training.execution.runtime import TrainingRuntime +from twinkle.server.model.runtime import TrainingRuntime from twinkle.server.state.training_records import CreateTraining @@ -26,9 +26,9 @@ def deserialize(config): finished.set() return None - monkeypatch.setattr('twinkle.server.training.execution.runtime.deserialize_object', deserialize) + monkeypatch.setattr('twinkle.server.model.runtime.deserialize_object', deserialize) manager = SimpleNamespace(save=Mock()) - monkeypatch.setattr('twinkle.server.training.execution.runtime.create_training_run_manager', lambda *args, **kw: manager) + monkeypatch.setattr('twinkle.server.model.runtime.create_training_run_manager', lambda *args, **kw: manager) async def serial(factory, task_type): return await factory() diff --git a/tests/server/training/test_runtime_queue.py b/tests/server/training/test_runtime_queue.py index cf772b711..04cc0927c 100644 --- a/tests/server/training/test_runtime_queue.py +++ b/tests/server/training/test_runtime_queue.py @@ -9,7 +9,7 @@ from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.task_queue import TaskQueueMixin -from twinkle.server.training.execution.runtime import TrainingRuntime +from twinkle.server.model.runtime import TrainingRuntime from twinkle.server.lifecycle.submit import submit_training_task from tests.server.training.helpers import create, setup diff --git a/tests/server/training/test_storage_recovery.py b/tests/server/training/test_storage_recovery.py index 7b939f61f..a0eefdb9f 100644 --- a/tests/server/training/test_storage_recovery.py +++ b/tests/server/training/test_storage_recovery.py @@ -7,7 +7,7 @@ from twinkle.server.state.backend.memory_backend import RayActorBackend from twinkle.server.state.backend.redis_backend import RedisBackend -from twinkle.server.training.controller import TrainingController +from twinkle.server.model.pool.controller import TrainingController from tests.server.training.helpers import create, request, setup diff --git a/tests/transformers/test_spectral_hybrid_lora.py b/tests/transformers/test_spectral_hybrid_lora.py index a9a7bd62f..5ff8ae167 100644 --- a/tests/transformers/test_spectral_hybrid_lora.py +++ b/tests/transformers/test_spectral_hybrid_lora.py @@ -507,8 +507,8 @@ def test_hybrid_rejects_client_owned_fft_and_target_parameter_config(): def test_transformers_server_selects_hybrid_model_only_for_hybrid_train_mode(monkeypatch): - from twinkle.server.training.execution import app - from twinkle.server.training.execution.backends import transformers_model + from twinkle.server.model import app + from twinkle.server.model.backends import transformers_model class Regular: @@ -535,7 +535,7 @@ class Hybrid(Regular): def test_hybrid_server_model_exposes_device_mesh_to_remote_class(): from twinkle.infra import _get_device_mesh_param_name - from twinkle.server.training.execution.backends.transformers_model import TwinkleCompatSpectralHybridTransformersModel + from twinkle.server.model.backends.transformers_model import TwinkleCompatSpectralHybridTransformersModel assert _get_device_mesh_param_name(TwinkleCompatSpectralHybridTransformersModel.__init__) == 'device_mesh' diff --git a/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py b/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py index 90bf4e6a1..9567138c4 100644 --- a/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py +++ b/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py @@ -10,7 +10,7 @@ from twinkle import DeviceMesh from twinkle.data_format import SampledSequence, SampleResponse, SamplingParams from twinkle.infra import _dispatch_args -from twinkle.server.sampling.execution.twinkle_handlers import _await_generation +from twinkle.server.sampler.twinkle_handlers import _await_generation from twinkle_agentic.async_rl import LoraContext from twinkle_agentic.async_rl.types import PartitionAdmission, PromptGroup, RolloutPolicy from twinkle_agentic.async_rl.vllm_sampler_tq import (VLLMSamplerTQ, _dispatch_generation, _GeneratedSample, From 5569d5d3f6be32c13e6cac1e99c9f6a933eedadf Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 21:04:14 +0800 Subject: [PATCH 40/46] fix(server): preserve adapter heartbeat and lifetime policies --- cookbook/client/async_rl/server_config.yaml | 2 + .../client/server/megatron/server_config.yaml | 3 + .../server/megatron/server_config_4b.yaml | 3 + .../server/transformer/server_config.yaml | 2 + .../server_config_hybrid_lora.yaml | 2 + .../\346\234\215\345\212\241\347\253\257.md" | 7 +++ src/twinkle/server/config/training.py | 10 +++ src/twinkle/server/model/pool/controller.py | 21 ++++++- src/twinkle/server/model/pool/driver.py | 5 +- tests/server/config/server_config_4b_e2e.yaml | 3 + .../config/server_config_4b_e2e_megatron.yaml | 2 + tests/server/config/test_training_config.py | 28 ++++++++- tests/server/training/helpers.py | 4 +- tests/server/training/test_controller.py | 61 ++++++++++++++++++- tests/server/training/test_runtime_queue.py | 59 ++++++++++++++++++ 15 files changed, 203 insertions(+), 9 deletions(-) diff --git a/cookbook/client/async_rl/server_config.yaml b/cookbook/client/async_rl/server_config.yaml index 668277a89..fe86dd552 100644 --- a/cookbook/client/async_rl/server_config.yaml +++ b/cookbook/client/async_rl/server_config.yaml @@ -66,6 +66,8 @@ applications: model_id: ${oc.env:TWINKLE_LOCAL_MODEL_PATH} min_replicas: 1 max_replicas: 1 + adapter_config: + adapter_timeout: 30 max_loras: 8 max_length: 10240 data_plane_url: http://127.0.0.1:8000/api/v1/data-plane diff --git a/cookbook/client/server/megatron/server_config.yaml b/cookbook/client/server/megatron/server_config.yaml index 93e1be9fc..ebae1bff0 100644 --- a/cookbook/client/server/megatron/server_config.yaml +++ b/cookbook/client/server/megatron/server_config.yaml @@ -122,6 +122,9 @@ applications: model_id: "ms://Qwen/Qwen3.8-27B" # ModelScope model identifier max_length: 32768 # model max length train_mode: lora + adapter_config: + adapter_timeout: 120 # Seconds without a session heartbeat before release + adapter_max_lifetime: 36000 # Maximum lifetime of an adapter in seconds (e.g., 10 hours) max_loras: 3 # Handles per instance, not simultaneous backward calls min_replicas: 0 # Do not reserve training GPUs at startup max_replicas: 2 # Up to 2 independent two-GPU instances diff --git a/cookbook/client/server/megatron/server_config_4b.yaml b/cookbook/client/server/megatron/server_config_4b.yaml index 36f20b1df..a0de5caec 100644 --- a/cookbook/client/server/megatron/server_config_4b.yaml +++ b/cookbook/client/server/megatron/server_config_4b.yaml @@ -54,6 +54,9 @@ applications: rps_limit: 100 # Max requests per second tps_limit: 100000 # Max tokens per second for a single user max_input_tokens: 60000 # Maximum input tokens per request + adapter_config: + adapter_timeout: 30 # Seconds without a session heartbeat before release + adapter_max_lifetime: 36000 # Maximum lifetime of an adapter in seconds (e.g., 10 hours) max_loras: 5 # Maximum number of LoRA adapters per model deployments: - name: ModelManagement diff --git a/cookbook/client/server/transformer/server_config.yaml b/cookbook/client/server/transformer/server_config.yaml index bf7ff1558..e9a8bb5e7 100644 --- a/cookbook/client/server/transformer/server_config.yaml +++ b/cookbook/client/server/transformer/server_config.yaml @@ -69,6 +69,8 @@ applications: queue_config: rps_limit: 100 # Max requests per second tps_limit: 100000 # Max tokens per second + adapter_config: + adapter_timeout: 30 # Seconds without a session heartbeat before release deployments: - name: ModelManagement autoscaling_config: diff --git a/cookbook/client/server/transformer/server_config_hybrid_lora.yaml b/cookbook/client/server/transformer/server_config_hybrid_lora.yaml index a8b53cee7..013024cb8 100644 --- a/cookbook/client/server/transformer/server_config_hybrid_lora.yaml +++ b/cookbook/client/server/transformer/server_config_hybrid_lora.yaml @@ -55,6 +55,8 @@ applications: queue_config: rps_limit: 100 tps_limit: 100000 + adapter_config: + adapter_timeout: 1800 deployments: - name: ModelManagement autoscaling_config: diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" index 5767454d7..7b51eb47e 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" @@ -480,6 +480,9 @@ twinkle-server check-config -c server_config.yaml min_replicas: 0 max_replicas: 2 idle_timeout_s: 600 + adapter_config: + adapter_timeout: 30 + adapter_max_lifetime: 36000 ``` | 配置项 | 如何使用 | @@ -489,6 +492,10 @@ twinkle-server check-config -c server_config.yaml | `max_replicas` | 最多启动多少个训练实例。 | | `max_loras` | 每个实例最多同时容纳多少个 LoRA 训练，仅用于 LoRA 模式。 | | `idle_timeout_s` | 训练闲置多久后自动释放句柄，单位为秒。心跳和结果轮询不会延长这个时间，排队中或正在执行的训练不会因此中断。 | +| `adapter_config.adapter_timeout` | 客户端多久没有发送会话心跳后释放训练句柄，单位为秒，默认 `1800`。SDK 会自动发送心跳；设为 `null` 可关闭这项检查，但会话删除后句柄仍会失效。 | +| `adapter_config.adapter_max_lifetime` | 训练句柄从创建起最长保留多久，单位为秒，默认 `36000`（10 小时）。设为 `null` 可取消这个上限。 | + +这三个时间分别控制业务闲置、心跳中断和最长寿命，互不替代。已经排队或正在执行的任务会先完成，再回收句柄。长时间训练请定期保存 checkpoint；句柄过期后需要重新创建，并显式恢复 checkpoint。 上例最多提供 2 个实例，每个实例容纳 3 个 LoRA 训练。客户端同时还受 `server_config.per_token_model_limit` 限制；多个客户端使用同一个 token 时，共享该 token 的配额。 diff --git a/src/twinkle/server/config/training.py b/src/twinkle/server/config/training.py index e09b0626b..5e2346582 100644 --- a/src/twinkle/server/config/training.py +++ b/src/twinkle/server/config/training.py @@ -32,6 +32,15 @@ def _finite_learning_rate(cls, value: float) -> float: return value +class AdapterConfig(BaseModel): + """Session liveness and absolute lifetime of a training handle, in seconds.""" + + model_config = ConfigDict(extra='forbid', frozen=True) + + adapter_timeout: float | None = Field(default=1800, gt=0, allow_inf_nan=False) + adapter_max_lifetime: float | None = Field(default=36000, gt=0, allow_inf_nan=False) + + class TrainingArguments(BaseModel): """Training options shared by model YAML and resolved instance configuration.""" @@ -40,6 +49,7 @@ class TrainingArguments(BaseModel): min_replicas: int = Field(default=0, ge=0, strict=True) max_replicas: int | None = Field(default=None, gt=0, strict=True) idle_timeout_s: float = Field(default=1800, gt=0, allow_inf_nan=False) + adapter_config: AdapterConfig = Field(default_factory=AdapterConfig) strategy: Literal['accelerate', 'native_fsdp'] | None = None queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) max_length: int | None = Field(default=None, gt=0, strict=True) diff --git a/src/twinkle/server/model/pool/controller.py b/src/twinkle/server/model/pool/controller.py index 1e89d5f8e..88caf868f 100644 --- a/src/twinkle/server/model/pool/controller.py +++ b/src/twinkle/server/model/pool/controller.py @@ -8,6 +8,7 @@ import traceback import uuid from collections.abc import Callable +from datetime import datetime from typing import Any, Protocol from twinkle.server.config.training import TrainingConfig, TrainingModelConfig @@ -386,9 +387,20 @@ async def _expired(self, training: TrainingRecord) -> bool: heartbeat = await self.store.sessions.get_last_heartbeat(training.session_id) return self._expired_with_heartbeat(training, heartbeat) + def _session_expired(self, training: TrainingRecord, heartbeat: float | None) -> bool: + timeout = self.config.get(training.base_model, training.mode).config.adapter_config.adapter_timeout + return heartbeat is None or (timeout is not None and time.time() - heartbeat >= timeout) + + def _lifetime_expired(self, training: TrainingRecord) -> bool: + lifetime = self.config.get(training.base_model, training.mode).config.adapter_config.adapter_max_lifetime + if lifetime is None: + return False + created_at = datetime.fromisoformat(training.created_at).timestamp() + return time.time() - created_at >= lifetime + def _expired_with_heartbeat(self, training: TrainingRecord, heartbeat: float | None) -> bool: spec = self.config.get(training.base_model, training.mode) - if heartbeat is None: + if self._session_expired(training, heartbeat) or self._lifetime_expired(training): return True return self.clock() - training.last_activity_at >= spec.config.idle_timeout_s @@ -402,7 +414,8 @@ async def _authorized(self, training_id: str, token: str, session_id: str | None raise TrainingExpiredError('Training handle expired; create a new handle and explicitly load a checkpoint.') if training.state != 'ACTIVE' or training.instance_id in self._failed_instances: raise ResourceNotFoundError('Training handle is no longer active') - if await self.store.sessions.get_last_heartbeat(training.session_id) is None: + heartbeat = await self.store.sessions.get_last_heartbeat(training.session_id) + if self._session_expired(training, heartbeat): raise TrainingExpiredError('Training session expired; create a new training handle') return training @@ -415,7 +428,9 @@ async def admit(self, renew: bool = True) -> RuntimeOwner: async with self._lock: training = await self._authorized(training_id, token, session_id) - if not training.pending_operation_ids and await self._expired(training): + expired = self._lifetime_expired(training) or (not training.pending_operation_ids + and await self._expired(training)) + if request_id not in training.pending_operation_ids and expired: training.state = 'CLOSING' training.terminal_reason = 'expired' await self.store.save_training(training) diff --git a/src/twinkle/server/model/pool/driver.py b/src/twinkle/server/model/pool/driver.py index 34cbd8b6c..cb9071f70 100644 --- a/src/twinkle/server/model/pool/driver.py +++ b/src/twinkle/server/model/pool/driver.py @@ -55,7 +55,10 @@ async def deploy(self, instance, spec): controller_name=self.controller_name, registry_name=name + '-registry') kwargs = spec.config.model_dump( - exclude={'num_devices', 'device_mesh', 'min_replicas', 'max_replicas', 'idle_timeout_s', 'queue_config'}, + exclude={ + 'num_devices', 'device_mesh', 'min_replicas', 'max_replicas', 'idle_timeout_s', 'adapter_config', + 'queue_config' + }, exclude_none=True) if 'hybrid' in kwargs: kwargs['hybrid'] = spec.config.hybrid diff --git a/tests/server/config/server_config_4b_e2e.yaml b/tests/server/config/server_config_4b_e2e.yaml index 07cfeca25..e4b06ea67 100644 --- a/tests/server/config/server_config_4b_e2e.yaml +++ b/tests/server/config/server_config_4b_e2e.yaml @@ -47,6 +47,9 @@ applications: queue_config: rps_limit: 100 tps_limit: 100000 + adapter_config: + adapter_timeout: 30 + adapter_max_lifetime: 36000 deployments: - name: ModelManagement autoscaling_config: diff --git a/tests/server/config/server_config_4b_e2e_megatron.yaml b/tests/server/config/server_config_4b_e2e_megatron.yaml index d235a0aa3..1c83ddac9 100644 --- a/tests/server/config/server_config_4b_e2e_megatron.yaml +++ b/tests/server/config/server_config_4b_e2e_megatron.yaml @@ -48,6 +48,8 @@ applications: queue_config: rps_limit: 100 tps_limit: 100000 + adapter_config: + adapter_timeout: 30 deployments: - name: ModelManagement autoscaling_config: diff --git a/tests/server/config/test_training_config.py b/tests/server/config/test_training_config.py index 11eb9bf5a..a0bdb268b 100644 --- a/tests/server/config/test_training_config.py +++ b/tests/server/config/test_training_config.py @@ -55,7 +55,7 @@ def test_existing_yaml_uses_explicit_instance_limits_and_business_idle(): assert model.config.idle_timeout_s == app.args.idle_timeout_s assert model.config.min_replicas == app.args.min_replicas assert model.config.max_replicas == app.args.max_replicas - assert 'adapter_config' not in type(app.args).model_fields + assert model.config.adapter_config == app.args.adapter_config @pytest.mark.parametrize('field,value', [ @@ -211,7 +211,6 @@ def configured_application(): {'modes': None}, {'modes': {'lora': {'num_devices': 2, 'device_mesh': {'dp_size': 2}}}}, {'device_type': 'cpu'}, - {'adapter_config': {'adapter_timeout': 30}}, ]) def test_model_yaml_rejects_alternative_resource_and_mode_syntax(extra): app = configured_application() @@ -220,6 +219,30 @@ def test_model_yaml_rejects_alternative_resource_and_mode_syntax(extra): ServerConfig(applications=[app]) +@pytest.mark.parametrize('mode', ['lora', 'full']) +def test_adapter_lifecycle_yaml_reaches_controller_and_round_trips(mode): + app = model_application() + policy = {'adapter_timeout': 30, 'adapter_max_lifetime': 36000} + app['args'].update(train_mode=mode, adapter_config=policy, idle_timeout_s=600) + config = ServerConfig(applications=[app]) + resolved = config.training_config.models[0].config + assert resolved.adapter_config.model_dump() == policy + assert resolved.idle_timeout_s == 600 + restored = ServerConfig.model_validate(config.to_yaml_dict()) + assert restored.training_config == config.training_config + + +@pytest.mark.parametrize('policy', [ + {'adapter_timeout': 0}, {'adapter_timeout': math.inf}, {'adapter_timeout': math.nan}, + {'adapter_max_lifetime': -1}, {'adapter_max_lifetime': math.inf}, {'unknown': 30}, +]) +def test_invalid_adapter_lifecycle_is_rejected_before_allocation(policy): + app = model_application() + app['args']['adapter_config'] = policy + with pytest.raises(ValidationError): + ServerConfig(applications=[app]) + + def test_configured_options_reach_cpu_router_without_changing_yaml(monkeypatch): from types import SimpleNamespace from twinkle.server.launcher.server_launcher import ServerLauncher @@ -278,6 +301,7 @@ def builder(**kwargs): assert captured['max_length'] == 128 assert captured['nproc_per_node'] == 2 assert captured['queue_config'] == model.config.queue_config + assert 'adapter_config' not in captured, 'Controller lifecycle policy must not reach model constructors' assert not driver.runtime.deployments diff --git a/tests/server/training/helpers.py b/tests/server/training/helpers.py index 3b4f6aa74..ba4a5fb9d 100644 --- a/tests/server/training/helpers.py +++ b/tests/server/training/helpers.py @@ -75,7 +75,7 @@ async def recover_orphans(self): pass -async def setup(*, devices=16, min_replicas=0, max_replicas=None, quota=100, backend=None): +async def setup(*, devices=16, min_replicas=0, max_replicas=None, quota=100, backend=None, adapter_config=None): mode = { 'num_devices': 2, 'device_mesh': { @@ -86,6 +86,8 @@ async def setup(*, devices=16, min_replicas=0, max_replicas=None, quota=100, bac 'min_replicas': min_replicas, 'max_replicas': max_replicas } + if adapter_config is not None: + mode['adapter_config'] = adapter_config config = TrainingConfig(models=[{ 'name': 'small', 'model_id': 'ms://small', diff --git a/tests/server/training/test_controller.py b/tests/server/training/test_controller.py index 60d3098bd..f11b8c9ea 100644 --- a/tests/server/training/test_controller.py +++ b/tests/server/training/test_controller.py @@ -154,9 +154,10 @@ async def finished(instance, request_id): @pytest.mark.asyncio -async def test_old_creation_time_and_live_heartbeat_do_not_override_business_idle(): +async def test_disabled_max_lifetime_and_live_heartbeat_do_not_override_business_idle(): from tests.server.state.fake_backend import FakeBackend - controller, state, runtime, clock = await setup(backend=FakeBackend()) + controller, state, runtime, clock = await setup( + backend=FakeBackend(), adapter_config={'adapter_max_lifetime': None}) identity = await create(controller, state, 'idle-only') controller.trainings[identity].created_at = '2000-01-01T00:00:00+00:00' clock.now = 9 @@ -171,6 +172,62 @@ async def test_old_creation_time_and_live_heartbeat_do_not_override_business_idl assert not runtime.allocations +@pytest.mark.asyncio +@pytest.mark.parametrize('reason', ['heartbeat', 'max_lifetime']) +@pytest.mark.parametrize('busy', [False, True]) +async def test_adapter_lifecycle_expires_independently_and_waits_for_running_work(reason, busy): + import time + from datetime import datetime, timedelta, timezone + from tests.server.state.fake_backend import FakeBackend + + controller, state, runtime, _ = await setup( + backend=FakeBackend(), adapter_config={'adapter_timeout': 30, 'adapter_max_lifetime': 60}) + identity = await create(controller, state, 'lifecycle') + if busy: + await controller.admit(identity, 'owner', 'session', 'running') + if reason == 'heartbeat': + session = await state.sessions.get('session') + session.last_heartbeat = time.time() - 31 + await state.sessions.add('session', session) + else: + controller.trainings[identity].created_at = (datetime.now(timezone.utc) - timedelta(seconds=61)).isoformat() + await controller.sweep() + await controller.wait_for_jobs() + if busy: + assert controller.trainings[identity].state == 'ACTIVE' + assert identity in runtime.users + if reason == 'heartbeat': + with pytest.raises(TrainingExpiredError): + await controller.admit(identity, 'owner', 'session', 'new-operation') + await controller.finish_operation(identity, 'running', controller.epoch, executed=True) + await controller.sweep() + await controller.wait_for_jobs() + assert controller.trainings[identity].state == 'EXPIRED' + assert identity not in runtime.users + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_max_lifetime_stops_new_work_while_admitted_work_finishes(): + from datetime import datetime, timedelta, timezone + from tests.server.state.fake_backend import FakeBackend + + controller, state, runtime, _ = await setup( + backend=FakeBackend(), adapter_config={'adapter_max_lifetime': 60}) + identity = await create(controller, state, 'lifetime-admission') + await controller.admit(identity, 'owner', 'session', 'running') + training = controller.trainings[identity] + training.created_at = (datetime.now(timezone.utc) - timedelta(seconds=61)).isoformat() + with pytest.raises(TrainingExpiredError): + await controller.admit(identity, 'owner', 'session', 'new-operation') + assert set(training.pending_operation_ids) == {'running'} + assert identity in runtime.users + await controller.finish_operation(identity, 'running', controller.epoch, executed=True) + await controller.wait_for_jobs() + assert training.state == 'EXPIRED' + assert not runtime.allocations + + @pytest.mark.asyncio async def test_lazy_queries_allocate_nothing(): controller, state, runtime, _ = await setup() diff --git a/tests/server/training/test_runtime_queue.py b/tests/server/training/test_runtime_queue.py index 04cc0927c..ac89f4b25 100644 --- a/tests/server/training/test_runtime_queue.py +++ b/tests/server/training/test_runtime_queue.py @@ -135,6 +135,65 @@ def is_loaded(self): ray.kill(worker, no_restart=True) +@pytest.mark.asyncio +@pytest.mark.parametrize('reason', ['heartbeat', 'max_lifetime']) +async def test_adapter_expiry_preserves_already_queued_work(reason): + import time + from datetime import datetime, timedelta, timezone + from tests.server.state.fake_backend import FakeBackend + + controller, state, runtime, _ = await setup( + backend=FakeBackend(), adapter_config={'adapter_timeout': 30, 'adapter_max_lifetime': 60}) + identity = await create(controller, state, 'queued-expiry') + training = controller.trainings[identity] + queue = RuntimeQueue(controller, state, controller.instances[training.instance_id]) + queue._task_queue_config = TaskQueueConfig(execution_timeout=5) + queue._compute_worker._config = queue._task_queue_config + entered, resume, finished = asyncio.Event(), asyncio.Event(), asyncio.Event() + executions = [] + + async def first(): + entered.set() + await resume.wait() + executions.append('first') + + async def second(): + executions.append('second') + + async def notify(*args, **kwargs): + await controller.finish_operation(*args, **kwargs) + if not training.pending_operation_ids: + finished.set() + return True + + queue._training_controller.finish_operation = RPC(notify) + try: + active = await queue.schedule_task(first, model_id=identity, token='owner', task_type='forward_backward') + await asyncio.wait_for(entered.wait(), 2) + queued = await queue.schedule_task(second, model_id=identity, token='owner', task_type='step') + if reason == 'heartbeat': + session = await state.sessions.get('session') + session.last_heartbeat = time.time() - 31 + await state.sessions.add('session', session) + else: + training.created_at = (datetime.now(timezone.utc) - timedelta(seconds=61)).isoformat() + await controller.sweep() + await controller.wait_for_jobs() + assert identity in runtime.users and len(training.pending_operation_ids) == 2 + resume.set() + assert (await settle(state, active['request_id']))['status'] == 'completed' + assert (await settle(state, queued['request_id']))['status'] == 'completed' + await asyncio.wait_for(finished.wait(), 2) + await controller.sweep() + await controller.wait_for_jobs() + assert executions == ['first', 'second'] + assert training.state == 'EXPIRED' and not runtime.allocations + finally: + resume.set() + await queue.shutdown_task_queue() + await controller.shutdown() + + @pytest.mark.asyncio async def test_lost_finish_notification_is_reconciled_without_stopping_compute(): controller, state, runtime, clock = await setup() From 85a293673882a4047204522b396eff7652bbcc9f Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 22:19:32 +0800 Subject: [PATCH 41/46] fix(server): require confirmed resource release before shutdown succeeds --- src/twinkle/server/model/pool/actor.py | 1 - src/twinkle/server/model/pool/controller.py | 13 ++++- src/twinkle/server/resources/instances.py | 11 ++++ src/twinkle/server/resources/runtime.py | 4 +- src/twinkle/server/sampler/pool/controller.py | 1 + tests/server/resources/test_instances.py | 20 ++++++-- tests/server/sampler/test_pool.py | 16 ++++++ tests/server/training/test_controller.py | 50 +++++++++++++++++++ 8 files changed, 108 insertions(+), 8 deletions(-) diff --git a/src/twinkle/server/model/pool/actor.py b/src/twinkle/server/model/pool/actor.py index 4e554f76b..d4f8f0afa 100644 --- a/src/twinkle/server/model/pool/actor.py +++ b/src/twinkle/server/model/pool/actor.py @@ -171,5 +171,4 @@ async def shutdown(self): task.cancel() await asyncio.gather(*self._monitors, return_exceptions=True) await self.core.shutdown() - await self.core.wait_for_jobs() await self.driver.runtime.wait_for_cleanup() diff --git a/src/twinkle/server/model/pool/controller.py b/src/twinkle/server/model/pool/controller.py index 88caf868f..6ef2908a5 100644 --- a/src/twinkle/server/model/pool/controller.py +++ b/src/twinkle/server/model/pool/controller.py @@ -790,8 +790,17 @@ async def wait_for_jobs(self) -> None: async def shutdown(self) -> None: self._ready = False - for instance in list(self.instances.values()): - await self.fail_instance( + outcomes = await asyncio.gather( + *(self.fail_instance( instance.instance_id, FutureFailureRecord( reason_code='orphaned_replica', message='Training service is shutting down', attribution='server')) + for instance in list(self.instances.values())), + return_exceptions=True) + await self.wait_for_jobs() + errors = [outcome for outcome in outcomes if isinstance(outcome, BaseException)] + for error in errors: + logger.error('Training instance shutdown failed', exc_info=(type(error), error, error.__traceback__)) + if errors: + raise errors[0] + self.lifecycle.require_released(self.instances.values()) diff --git a/src/twinkle/server/resources/instances.py b/src/twinkle/server/resources/instances.py index d08045108..8bcdad118 100644 --- a/src/twinkle/server/resources/instances.py +++ b/src/twinkle/server/resources/instances.py @@ -63,6 +63,11 @@ async def close_once(): try: await asyncio.wait_for(stop(), timeout) except Exception as error: + logger.warning( + '%s instance %s shutdown is unconfirmed', + self.label, + identity, + exc_info=(type(error), error, error.__traceback__)) await quarantined(error) else: # Only the feature can retire its users/quota and commit STOPPED. @@ -93,6 +98,12 @@ async def health(self, identity, probe, *, unreachable_timeout=30): def owned_count(instances): return sum(instance.state != 'STOPPED' for instance in instances) + def require_released(self, instances): + """A shutdown ACK must certify release before its owner can be removed.""" + owned = self.owned_count(instances) + if owned: + raise RuntimeError(f'{self.label} shutdown is unconfirmed: {owned} instances remain owned') + @classmethod def headroom(cls, instances, maximum): return float('inf') if maximum is None else max(0, maximum - cls.owned_count(instances)) diff --git a/src/twinkle/server/resources/runtime.py b/src/twinkle/server/resources/runtime.py index 5eba8b920..fbc251d58 100644 --- a/src/twinkle/server/resources/runtime.py +++ b/src/twinkle/server/resources/runtime.py @@ -96,8 +96,10 @@ async def shutdown(self, identity, name, *, observe_removal=None): for handle in actors: try: await asyncio.wait_for(handle.__ray_ready__.remote(), 5) - except ray.exceptions.RayActorError: + except ray.exceptions.ActorDiedError: continue + except ray.exceptions.RayActorError as exc: + raise ResourceReleaseError('Owned worker death could not be confirmed', phase='worker_death') from exc raise ResourceReleaseError('Owned worker death could not be confirmed', phase='worker_death') started, status = time.monotonic(), 'error' diff --git a/src/twinkle/server/sampler/pool/controller.py b/src/twinkle/server/sampler/pool/controller.py index 45f8f5aca..1eb60431c 100644 --- a/src/twinkle/server/sampler/pool/controller.py +++ b/src/twinkle/server/sampler/pool/controller.py @@ -363,3 +363,4 @@ async def shutdown(self): await asyncio.gather(*(self._close(instance) for instance in list(self.instances.values()))) self._guard.close() await self.lifecycle.join() + self.lifecycle.require_released(self.instances.values()) diff --git a/tests/server/resources/test_instances.py b/tests/server/resources/test_instances.py index 2fae3d280..7b346b534 100644 --- a/tests/server/resources/test_instances.py +++ b/tests/server/resources/test_instances.py @@ -25,8 +25,10 @@ async def broken(): @pytest.mark.asyncio -async def test_concurrent_close_shares_one_stop_and_retries_quarantine(): +async def test_concurrent_close_shares_one_stop_and_retries_quarantine(caplog, monkeypatch): import asyncio + target = logging.getLogger('twinkle.server.resources.instances') + monkeypatch.setattr(target, 'handlers', [*target.handlers, caplog.handler]) lifecycle = InstanceLifecycle('test') entered, resume = asyncio.Event(), asyncio.Event() calls, events = [], [] @@ -51,6 +53,9 @@ async def quarantined(error): resume.set() await asyncio.gather(first, second) assert calls == ['stop'] and events == ['death unconfirmed'] + records = [record for record in caplog.records if record.getMessage() == 'test instance one shutdown is unconfirmed'] + assert len(records) == 1 and records[0].exc_info[1].args == ('death unconfirmed',) + assert records[0].exc_info[2] is not None await lifecycle.close('one', stop, 1, stopped, quarantined) assert calls == ['stop', 'stop'] and events[-1] == 'stopped' await lifecycle.join() @@ -199,7 +204,9 @@ def inventory(): @pytest.mark.asyncio -@pytest.mark.parametrize('outcome', ['unconfirmed_worker', 'removed', 'removal_failed']) +@pytest.mark.parametrize('outcome', [ + 'unconfirmed_worker', 'unavailable_worker', 'dead_worker', 'removed', 'removal_failed' +]) async def test_worker_death_precedes_group_release_and_removal_timing(monkeypatch, outcome): from types import SimpleNamespace from unittest.mock import AsyncMock, Mock @@ -209,10 +216,15 @@ async def test_worker_death_precedes_group_release_and_removal_timing(monkeypatc from twinkle.server.resources.runtime import RayRuntime async def ready(): + if outcome == 'unavailable_worker': + raise ray.exceptions.ActorUnavailableError('Control-plane connection lost', None) + if outcome == 'dead_worker': + raise ray.exceptions.ActorDiedError() return None # Worker can still answer: death has not been confirmed. handle = SimpleNamespace(__ray_ready__=SimpleNamespace(remote=ready)) - rows = [{'namespace': 'test', 'name': 'owned-worker'}] if outcome == 'unconfirmed_worker' else [] + rows = ([{'namespace': 'test', 'name': 'owned-worker'}] + if outcome in ('unconfirmed_worker', 'unavailable_worker', 'dead_worker') else []) monkeypatch.setattr(ray.util, 'list_named_actors', lambda **kwargs: rows) monkeypatch.setattr(ray, 'get_actor', lambda *args, **kwargs: handle) monkeypatch.setattr(ray, 'kill', lambda *args, **kwargs: None) @@ -221,7 +233,7 @@ async def ready(): runtime.pool.remove = AsyncMock(side_effect=RuntimeError('PG removal failed') if outcome == 'removal_failed' else None) observe = Mock() - if outcome == 'unconfirmed_worker': + if outcome in ('unconfirmed_worker', 'unavailable_worker'): with pytest.raises(ResourceReleaseError) as failure: await runtime.shutdown('owned', 'app', observe_removal=observe) assert failure.value.phase == 'worker_death' diff --git a/tests/server/sampler/test_pool.py b/tests/server/sampler/test_pool.py index 6fb9e3b67..12592bef4 100644 --- a/tests/server/sampler/test_pool.py +++ b/tests/server/sampler/test_pool.py @@ -96,6 +96,22 @@ async def settle_jobs(pool): await asyncio.gather(*tuple(pool.lifecycle.jobs), return_exceptions=True) +@pytest.mark.asyncio +async def test_shutdown_rejects_unconfirmed_release_and_allows_retry(): + pool, driver, _ = setup() + await pool.initialize() + driver.fail_close = True + try: + with pytest.raises(RuntimeError, match='shutdown.*unconfirmed'): + await pool.shutdown() + assert pool._owned_count() == 2 and not driver.stopped + assert all(instance.state == 'QUARANTINED' for instance in pool.instances.values()) + finally: + driver.fail_close = False + await pool.shutdown() + assert pool._owned_count() == 0 and len(driver.stopped) == 2 + + @pytest.mark.asyncio async def test_minimum_expand_and_reclaim_only_idle_extras(): pool, driver, clock = setup() diff --git a/tests/server/training/test_controller.py b/tests/server/training/test_controller.py index f11b8c9ea..f4c3babad 100644 --- a/tests/server/training/test_controller.py +++ b/tests/server/training/test_controller.py @@ -228,6 +228,56 @@ async def test_max_lifetime_stops_new_work_while_admitted_work_finishes(): assert not runtime.allocations +@pytest.mark.asyncio +async def test_shutdown_rejects_unconfirmed_release_and_allows_retry(): + from tests.server.state.fake_backend import FakeBackend + + controller, state, runtime, _ = await setup(backend=FakeBackend()) + await create(controller, state, 'shutdown-ack') + runtime.fail_shutdown = True + try: + with pytest.raises(RuntimeError, match='shutdown.*unconfirmed'): + await controller.shutdown() + assert runtime.allocations + assert all(instance.state == 'QUARANTINED' for instance in controller.instances.values()) + finally: + runtime.fail_shutdown = False + await controller.shutdown() + assert not runtime.allocations + assert all(instance.state == 'STOPPED' for instance in controller.instances.values()) + + +@pytest.mark.asyncio +async def test_shutdown_attempts_healthy_peer_after_an_instance_metadata_failure(): + from tests.server.state.fake_backend import FakeBackend + + controller, state, runtime, _ = await setup(backend=FakeBackend()) + first = await create(controller, state, 'shutdown-first') + peer = await create(controller, state, 'shutdown-peer', mode='full') + first_instance = controller.instances[controller.trainings[first].instance_id] + peer_instance = controller.instances[controller.trainings[peer].instance_id] + save = controller.store.save_instance + + async def unavailable(record): + if record.instance_id == first_instance.instance_id and record.state == 'DRAINING': + raise RuntimeError('Instance metadata unavailable') + await save(record) + + controller.store.save_instance = unavailable + try: + with pytest.raises(RuntimeError, match='Instance metadata unavailable'): + await controller.shutdown() + assert peer_instance.state == 'STOPPED' + assert peer_instance.instance_id not in runtime.allocations + assert first_instance.instance_id in runtime.allocations + with pytest.raises(ResourceNotFoundError): + await controller.route(first, 'owner') + finally: + controller.store.save_instance = save + await controller.shutdown() + assert not runtime.allocations + + @pytest.mark.asyncio async def test_lazy_queries_allocate_nothing(): controller, state, runtime, _ = await setup() From a0a6aa8b6a13f88231877fe09bfc88216522e227 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 23 Oct 2026 22:46:40 +0800 Subject: [PATCH 42/46] test(server): isolate lifecycle logging from parent handlers --- tests/server/resources/test_instances.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/server/resources/test_instances.py b/tests/server/resources/test_instances.py index 7b346b534..556bf2879 100644 --- a/tests/server/resources/test_instances.py +++ b/tests/server/resources/test_instances.py @@ -7,15 +7,16 @@ @pytest.mark.asyncio async def test_spawn_logs_exception_with_traceback(caplog, monkeypatch): - # The ``twinkle`` logger disables propagation, so attach caplog to the lifecycle logger directly. + # Capture one path regardless of the parent logger's runtime configuration. target = logging.getLogger('twinkle.server.resources.instances') - monkeypatch.setattr(target, 'handlers', [*target.handlers, caplog.handler]) + monkeypatch.setattr(target, 'handlers', [caplog.handler]) + monkeypatch.setattr(target, 'propagate', False) lifecycle = InstanceLifecycle('Training') async def broken(): raise RuntimeError('background sentinel') - with caplog.at_level(logging.ERROR): + with caplog.at_level(logging.ERROR, logger=target.name): lifecycle.spawn(broken()) await lifecycle.join() records = [record for record in caplog.records if record.getMessage() == 'Training background job failed'] @@ -28,7 +29,9 @@ async def broken(): async def test_concurrent_close_shares_one_stop_and_retries_quarantine(caplog, monkeypatch): import asyncio target = logging.getLogger('twinkle.server.resources.instances') - monkeypatch.setattr(target, 'handlers', [*target.handlers, caplog.handler]) + monkeypatch.setattr(target, 'handlers', [caplog.handler]) + monkeypatch.setattr(target, 'propagate', False) + caplog.set_level(logging.WARNING, logger=target.name) lifecycle = InstanceLifecycle('test') entered, resume = asyncio.Event(), asyncio.Event() calls, events = [], [] From a02b7595f78d597efb2bf2ccdf85b5f7c3266f52 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Sat, 24 Oct 2026 00:08:23 +0800 Subject: [PATCH 43/46] test(server): start validation components only when needed --- .../on_demand_training/http_probe.py | 72 ++++++++++++------- .../on_demand_training/matrix_validation.py | 26 ++++++- .../server/training/test_http_integration.py | 14 ++-- 3 files changed, 78 insertions(+), 34 deletions(-) diff --git a/tests/server/integration/on_demand_training/http_probe.py b/tests/server/integration/on_demand_training/http_probe.py index 13efd3f6c..82e87b34c 100644 --- a/tests/server/integration/on_demand_training/http_probe.py +++ b/tests/server/integration/on_demand_training/http_probe.py @@ -8,18 +8,22 @@ import tempfile import time import uuid +from pathlib import Path from on_demand_sdk import unload_tinker_model from ray import serve from twinkle.server.config import ServerConfig from twinkle.server.launcher.server_launcher import ServerLauncher -from twinkle_client import init_tinker_client -from twinkle_client.http import ClientTransport -from twinkle_client.http.context import ClientContext +from twinkle_client import init_tinker_client, init_twinkle_client from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel def main(): + started = time.monotonic() + + def phase(name, **details): + print(json.dumps({'phase': name, 'elapsed_s': round(time.monotonic() - started, 2), **details}), flush=True) + namespace = 'http_' + uuid.uuid4().hex[:8] with socket.socket() as listener: listener.bind(('127.0.0.1', 0)) @@ -59,6 +63,7 @@ def main(): 'device_mesh': {'dp_size': 1}, 'max_loras': 2, 'max_replicas': 1, + 'adapter_config': {'adapter_timeout': 30}, }, }, { 'name': 'other', @@ -72,6 +77,7 @@ def main(): 'device_mesh': {'dp_size': 1}, 'min_replicas': 0, 'max_replicas': 1, + 'adapter_config': {'adapter_timeout': 30}, }, 'deployments': [{ 'max_ongoing_requests': 8, @@ -84,6 +90,7 @@ def main(): 'model_id': 'small', 'backend': 'mock', 'train_mode': 'lora', 'device_group': {'ranks': 1, 'device_type': 'cpu'}, 'device_mesh': {'dp_size': 1}, 'max_loras': 1, 'max_replicas': 2, + 'adapter_config': {'adapter_timeout': 30}, }, }, { 'name': 'gateway', @@ -99,6 +106,7 @@ def main(): launcher = ServerLauncher(config) launcher._start_serve() base = f'http://127.0.0.1:{port}' + native = fresh_native = service = None def wait_epoch(controller, old_epoch): deadline = time.monotonic() + 90 @@ -120,8 +128,9 @@ def wait_epoch(controller, old_epoch): launcher._deploy_training() controller = launcher._training_controller assert ray.get(controller.diagnostics.remote())['instances'] == [] - session = requests.post(base + '/api/v1/twinkle/create_session', json={}, timeout=10).json()['session_id'] - transport = ClientTransport(ClientContext(base_url=base, api_key='probe-owner', session_id=session)) + phase('server_ready') + native = init_twinkle_client(base_url=base, api_key='probe-owner') + transport = native.transport first = MultiLoraTransformersModel('small', transport=transport) second = MultiLoraTransformersModel('small', transport=transport) from peft import LoraConfig @@ -143,6 +152,7 @@ def assert_separate_instances(identities): client.zero_grad() for client in routed: client.remove_adapter() + phase('native_routing_verified') first.add_adapter_to_model('same', LoraConfig(r=8, target_modules='all-linear')) instance = next( item for item in ray.get(controller.diagnostics.remote())['instances'] @@ -170,6 +180,7 @@ def assert_separate_instances(identities): native_lora.add_adapter_to_model('native-lora', LoraConfig(r=8, target_modules='all-linear')) native_lora.zero_grad() native_lora.remove_adapter() + phase('native_lifecycle_verified') init_tinker_client() from tinker import ServiceClient service = ServiceClient(base_url=base, api_key='probe-tinker') @@ -180,6 +191,7 @@ def assert_separate_instances(identities): assert client.get_info().model_id == client.model_id for client in routed_tinker: assert unload_tinker_model(service, client.model_id).model_id == client.model_id + phase('tinker_routing_verified') shared1 = service.create_lora_training_client(base_model='small', rank=8) shared2 = service.create_lora_training_client(base_model='small', rank=8) assert shared1.model_id != shared2.model_id @@ -213,6 +225,7 @@ def assert_separate_instances(identities): ray.kill(controller, no_restart=False) diagnostic = wait_epoch(controller, old_epoch) assert not [item for item in diagnostic['instances'] if item['state'] != 'STOPPED'] + phase('controller_recovered') stale = requests.post( base + '/api/v1/get_info', json={ @@ -226,45 +239,52 @@ def assert_separate_instances(identities): timeout=10) assert stale.status_code == 404, stale.text before = ray.available_resources().get('CPU', 0) - for index in range(10): + for index in range(3): model = MultiLoraTransformersModel('small' if index % 2 == 0 else 'other', transport=transport) model.add_adapter_to_model('cycle', LoraConfig(r=8, target_modules='all-linear') if index % 2 == 0 else None) model.zero_grad() model.remove_adapter() - print(json.dumps({'phase': 'cycles', 'completed': index + 1}), flush=True) + phase('cycles', completed=index + 1) assert ray.available_resources().get('CPU', 0) >= before - 0.02 final = ray.get(controller.diagnostics.remote()) assert all(item['state'] == 'STOPPED' for item in final['instances']) live = MultiLoraTransformersModel('other', transport=transport) live.add_adapter_to_model('state-loss', None) + native.close() ray.kill(ray.get_actor('twinkle_state_actor', namespace=namespace), no_restart=False) after_loss = wait_epoch(controller, final['epoch']) assert not [item for item in after_loss['instances'] if item['state'] != 'STOPPED'] - session = requests.post(base + '/api/v1/twinkle/create_session', json={}, timeout=10).json()['session_id'] - fresh_transport = ClientTransport(ClientContext(base_url=base, api_key='probe-owner', session_id=session)) + fresh_native = init_twinkle_client(base_url=base, api_key='probe-owner') + fresh_transport = fresh_native.transport fresh = MultiLoraTransformersModel('other', transport=fresh_transport) fresh.add_adapter_to_model('state-loss', None) assert fresh._training_ids['state-loss'] != live._training_ids['state-loss'] fresh.remove_adapter() - fresh_transport.close() - print( - json.dumps({ - 'passed': True, - 'native_recreation': True, - 'native_stable_routing': native_owners, - 'tinker_stable_routing': tinker_owners, - 'native_lora': True, - 'tinker_shared': True, - 'tinker_full': True, - 'restart': True, - 'state_loss': True, - 'runtime_lifecycle_rejected': True, - 'cycles': 10, - 'instances': len(final['instances']) - })) - transport.close() + phase('state_loss_recovered') + evidence = { + 'passed': True, + 'native_recreation': True, + 'native_stable_routing': native_owners, + 'tinker_stable_routing': tinker_owners, + 'native_lora': True, + 'tinker_shared': True, + 'tinker_full': True, + 'restart': True, + 'state_loss': True, + 'runtime_lifecycle_rejected': True, + 'cycles': 3, + 'instances': len(final['instances']) + } + if result_path := os.environ.get('TWINKLE_TEST_HTTP_RESULT'): + Path(result_path).write_text(json.dumps(evidence)) + print(json.dumps(evidence), flush=True) finally: + for client in (native, fresh_native): + if client is not None: + client.close() + if service is not None: + service.holder.close() if getattr(launcher, '_training_controller', None) is not None: ray.get(launcher._training_controller.shutdown.remote(), timeout=180) ray.kill(launcher._training_controller, no_restart=True) diff --git a/tests/server/integration/on_demand_training/matrix_validation.py b/tests/server/integration/on_demand_training/matrix_validation.py index b889d2d4b..f4fada660 100644 --- a/tests/server/integration/on_demand_training/matrix_validation.py +++ b/tests/server/integration/on_demand_training/matrix_validation.py @@ -353,6 +353,7 @@ def main(): 'cases': [], 'lazy_start': False, } + started = time.monotonic() ray.init( address='local', namespace=namespace, @@ -360,14 +361,20 @@ def main(): num_gpus=args.training_devices + (0 if args.recovery_only else 1), _temp_dir='/tmp/twm-' + namespace[-8:], include_dashboard=False) + evidence['ray_startup_seconds'] = time.monotonic() - started launcher = ServerLauncher(config) try: + started = time.monotonic() launcher._start_serve() for app in config.applications: - launcher._deploy_application(app) + if app.import_path != 'sampler': + launcher._deploy_application(app) launcher._deploy_training() assert ray.get(launcher._training_controller.diagnostics.remote())['instances'] == [] evidence['lazy_start'] = True + evidence['server_startup_seconds'] = time.monotonic() - started + print(json.dumps({'phase': 'training_server_ready', 'seconds': evidence['server_startup_seconds'], + 'sampler_running': False}), flush=True) filenames = ['test_sft_e2e.py', 'test_dpo_e2e.py', 'test_grpo_e2e.py'] if not args.skip_full_cycle: filenames.append('test_full_cycle_e2e.py') @@ -383,7 +390,19 @@ def main(): (run_dir / (filename + '.collect.log')).write_text(collection.stdout + collection.stderr) expected = sum('::test_' in line for line in collection.stdout.splitlines()) assert collection.returncode == 0 and expected > 0, collection.stdout + collection.stderr + if filename in ('test_grpo_e2e.py', 'test_full_cycle_e2e.py') and not launcher._sampler_controllers: + started = time.monotonic() + print(json.dumps({'phase': 'sampler_starting', 'needed_by': filename}), flush=True) + for app in config.applications: + if app.import_path == 'sampler': + launcher._deploy_application(app) + evidence['sampler_startup_seconds'] = time.monotonic() - started + evidence['sampler_started_for'] = filename + print(json.dumps({'phase': 'sampler_ready', 'seconds': evidence['sampler_startup_seconds'], + 'needed_by': filename}), flush=True) start = time.monotonic() + print(json.dumps({'phase': 'case_start', 'file': filename, + 'sampler_running': bool(launcher._sampler_controllers)}), flush=True) with (run_dir / (filename + '.log')).open('w') as output: result = subprocess.run([ sys.executable, @@ -407,6 +426,7 @@ def main(): 'returncode': result.returncode, 'tests': len(cases), 'skipped': skipped, + 'sampler_running': bool(launcher._sampler_controllers), 'seconds': time.monotonic() - start }) (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) @@ -456,6 +476,7 @@ def main(): progress_path=run_dir / (name + '.json'))) (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) finally: + cleanup_started = time.monotonic() try: controller = getattr(launcher, '_training_controller', None) if controller is not None: @@ -479,9 +500,10 @@ def main(): client.delete(*keys) finally: client.close() - (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) if managed is not None: managed.stop() + evidence['cleanup_seconds'] = time.monotonic() - cleanup_started + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) evidence['passed'] = True (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) return 0 diff --git a/tests/server/training/test_http_integration.py b/tests/server/training/test_http_integration.py index 3d73695e4..164a1562d 100644 --- a/tests/server/training/test_http_integration.py +++ b/tests/server/training/test_http_integration.py @@ -1,24 +1,26 @@ import json +import os import subprocess import sys import uuid from pathlib import Path -def test_installed_native_and_tinker_sdks_use_managed_http_routes(): +def test_installed_native_and_tinker_sdks_use_managed_http_routes(tmp_path): # This probe owns a separate Ray cluster. Stop the session fixture's cluster # first so two local GCS/raylet pairs do not evict each other on a slow host. import ray ray.shutdown() try: - # Both SDKs, recovery and ten release/recreate cycles start many fresh + # Both SDKs, recovery and release/recreate cycles start fresh # workers. Allow cold imports from a shared environment; each service # operation still enforces its own startup/execution deadline. + evidence_path = tmp_path / 'http-result.json' result = subprocess.run([ sys.executable, str(Path(__file__).resolve().parents[3] / 'tests/server/integration/on_demand_training/http_probe.py') - ], capture_output=True, text=True, timeout=2 * 60 * 60) + ], env={**os.environ, 'TWINKLE_TEST_HTTP_RESULT': str(evidence_path)}, timeout=2 * 60 * 60) finally: ray.init( ignore_reinit_error=True, @@ -26,12 +28,12 @@ def test_installed_native_and_tinker_sdks_use_managed_http_routes(): log_to_driver=False, namespace='twinkle_test', _temp_dir='/tmp/twinkle-test-' + uuid.uuid4().hex[:8]) - assert result.returncode == 0, result.stdout + result.stderr - evidence = next(json.loads(line) for line in result.stdout.splitlines() if line.startswith('{"passed":')) + assert result.returncode == 0, 'HTTP probe failed; see subprocess output above' + evidence = json.loads(evidence_path.read_text()) assert evidence['passed'] and evidence['native_recreation'] and evidence['tinker_shared'] assert len(set(evidence['native_stable_routing'].values())) == 2 assert len(set(evidence['tinker_stable_routing'].values())) == 2 assert evidence['native_lora'] assert evidence['runtime_lifecycle_rejected'] assert evidence['tinker_full'] - assert evidence['restart'] and evidence['state_loss'] and evidence['cycles'] == 10 + assert evidence['restart'] and evidence['state_loss'] and evidence['cycles'] == 3 From ffc9c0f93417376e415e9f48ea8df5821a4435c8 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Mon, 26 Oct 2026 11:45:04 +0800 Subject: [PATCH 44/46] fix(model): stabilize metric kwargs, optimizer groups and FFT slots --- src/twinkle/model/optimizer_group.py | 16 +++---- .../model/transformers/hybrid/fft_slots.py | 4 ++ .../model/transformers/transformers.py | 5 +- .../test_optimizer_group_regression.py | 48 +++++++++++++++++++ .../transformers/test_spectral_hybrid_lora.py | 6 ++- 5 files changed, 69 insertions(+), 10 deletions(-) create mode 100644 tests/transformers/test_optimizer_group_regression.py diff --git a/src/twinkle/model/optimizer_group.py b/src/twinkle/model/optimizer_group.py index 150e694bc..59b4b331b 100644 --- a/src/twinkle/model/optimizer_group.py +++ b/src/twinkle/model/optimizer_group.py @@ -64,15 +64,15 @@ def accumulate_metrics(self, is_training): """Accumulate metrics for train/eval status. Override in subclass if needed.""" status = self.train_status if is_training else self.eval_status if len(status.metrics) > 0 and status.inputs is not None and status.outputs is not None: + metric_kwargs = { + **status.forward_kwargs, + 'lr': self._get_lr(), + 'step': self.cur_step - 1, + 'gradient_accumulation_steps': self.gradient_accumulation_steps, + 'grad_norm': self._last_grad_norm, + } for metric in status.metrics: - metric.accumulate( - status.inputs, - status.outputs, - lr=self._get_lr(), - step=self.cur_step - 1, - gradient_accumulation_steps=self.gradient_accumulation_steps, - grad_norm=self._last_grad_norm, - **status.forward_kwargs) + metric.accumulate(status.inputs, status.outputs, **metric_kwargs) def calculate_metrics(self, is_training): """Calculate and return metrics.""" diff --git a/src/twinkle/model/transformers/hybrid/fft_slots.py b/src/twinkle/model/transformers/hybrid/fft_slots.py index 39cde9579..38d675d2b 100644 --- a/src/twinkle/model/transformers/hybrid/fft_slots.py +++ b/src/twinkle/model/transformers/hybrid/fft_slots.py @@ -120,6 +120,10 @@ def activate_fft_slot(self, adapter_name: str) -> None: fft_adapter_name = self._fft_adapter_name(adapter_name) if self.is_hybrid(adapter_name) else None selected_adapters = [fft_adapter_name] if fft_adapter_name is not None else [] for wrapper in self._iter_fft_wrappers(): + # Selecting a regular LoRA can disable auxiliary wrappers. Restoring + # an FFT slot name alone does not restore its execution state. + if fft_adapter_name is not None and wrapper.disable_adapters: + wrapper.enable_adapters(True) if wrapper.active_adapters != selected_adapters: wrapper.set_adapter(selected_adapters) self._enable_preallocated_fft_grad(wrapper) diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index eb90d4a9f..e240c765c 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -1127,7 +1127,10 @@ def get_parameter_names(model, forbidden_layer_types, forbidden_layer_names=None return result forbidden_name_patterns = [r'bias', r'layernorm', r'rmsnorm', r'(?:^|\.)norm(?:$|\.)', r'_norm(?:$|\.)'] - decay_parameters = get_parameter_names(self.model, [torch.nn.LayerNorm], forbidden_name_patterns) + # Match names from the same unwrapped model used by _get_trainable_parameters. + # DDP's module prefix otherwise changes optimizer groups for later tenants. + decay_parameters = get_parameter_names( + self.strategy.unwrap_model(self.model), [torch.nn.LayerNorm], forbidden_name_patterns) params = self._get_trainable_parameters(adapter_name) decay_param_names = [n for n, p in params.items() if (n in decay_parameters and p.requires_grad)] no_decay_param_names = [n for n, p in params.items() if (n not in decay_parameters and p.requires_grad)] diff --git a/tests/transformers/test_optimizer_group_regression.py b/tests/transformers/test_optimizer_group_regression.py new file mode 100644 index 000000000..23922c3a1 --- /dev/null +++ b/tests/transformers/test_optimizer_group_regression.py @@ -0,0 +1,48 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Optimizer groups must not depend on DDP wrapping or repeated GA kwargs.""" +from types import SimpleNamespace +from unittest.mock import Mock + +import torch + +from twinkle.model.optimizer_group import BaseOptimizerGroup, TrainStatus +from twinkle.model.transformers.transformers import TransformersModel + + +def test_forward_kwargs_do_not_duplicate_metric_control_parameters(): + metric = Mock() + kwargs = {'gradient_accumulation_steps': 2, 'lr': 99, 'step': 99, 'grad_norm': 99, 'loss_reduction': 'sum'} + status = TrainStatus(inputs=[{}], outputs={'loss': 1.0}, metrics=[metric], forward_kwargs=kwargs) + group = BaseOptimizerGroup(train_status=status, cur_step=3, gradient_accumulation_steps=2) + group.accumulate_metrics(True) + metric.accumulate.assert_called_once_with( + status.inputs, status.outputs, gradient_accumulation_steps=2, lr=[], step=2, grad_norm=0.0, + loss_reduction='sum') + assert kwargs['lr'] == 99 + + +def test_parameter_groups_and_optimizer_restore_are_stable_after_wrapping(): + model = torch.nn.Sequential(torch.nn.Linear(4, 4), torch.nn.LayerNorm(4)) + + class Wrapper(torch.nn.Module): + def __init__(self, module): + super().__init__() + self.module = module + + harness = SimpleNamespace( + model=model, + strategy=SimpleNamespace(unwrap_model=lambda wrapped: getattr(wrapped, 'module', wrapped)), + _get_trainable_parameters=lambda adapter: dict(model.named_parameters()), + ) + original_groups = TransformersModel._create_param_group(harness, '', weight_decay=0.1) + optimizer = torch.optim.Adam(original_groups, lr=1e-4) + model(torch.ones(2, 4)).sum().backward() + optimizer.step() + state = optimizer.state_dict() + harness.model = Wrapper(model) + late_groups = TransformersModel._create_param_group(harness, '', weight_decay=0.1) + assert [g['param_names'] for g in late_groups] == [g['param_names'] for g in original_groups] + assert original_groups[0]['param_names'] == ['0.weight'] + restored = torch.optim.Adam(late_groups, lr=1e-4) + restored.load_state_dict(state) + assert len(restored.state) == len(optimizer.state) diff --git a/tests/transformers/test_spectral_hybrid_lora.py b/tests/transformers/test_spectral_hybrid_lora.py index 9e9714249..5ff8ae167 100644 --- a/tests/transformers/test_spectral_hybrid_lora.py +++ b/tests/transformers/test_spectral_hybrid_lora.py @@ -871,11 +871,15 @@ def test_disable_lora_restores_slots_after_operation_fails(): assert all(parameter.requires_grad for name, parameter in model.named_parameters() if 'lora_' in name) -def test_multi_tenant_hybrid_isolation_and_non_destructive_merge(): +@pytest.mark.parametrize('fft_disabled', [False, True]) +def test_multi_tenant_hybrid_isolation_and_non_destructive_merge(fft_disabled): base, model, manager, spectral = _make_multi_tenant_hybrid() inputs = torch.randn(2, 8) with manager.adapter('regular'): regular_before = model(inputs).detach().clone() + if fft_disabled: + # A normal LoRA selection can disable auxiliary FFT wrappers in PEFT. + spectral._get_fft_wrapper('layers.0.self_attn.q_proj').enable_adapters(False) hybrid = manager.find_lora_by_tenant('hybrid') modules = dict(model.named_modules()) From 84410936ea6c901896d877d925207054f60bea15 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Mon, 26 Oct 2026 11:45:39 +0800 Subject: [PATCH 45/46] refactor(client): minimize on-demand server config changes --- src/twinkle_client/auto/agent/server_tools.py | 87 ++++++++++++------- tests/twinkle_client/test_auto_agent_tools.py | 3 + 2 files changed, 58 insertions(+), 32 deletions(-) diff --git a/src/twinkle_client/auto/agent/server_tools.py b/src/twinkle_client/auto/agent/server_tools.py index c2b031cc6..2324dc547 100644 --- a/src/twinkle_client/auto/agent/server_tools.py +++ b/src/twinkle_client/auto/agent/server_tools.py @@ -375,6 +375,8 @@ def _generate_server_config( def _short(mid: str) -> str: return mid.split('/')[-1] if '/' in mid else mid + model_short = _short(model_id) + # Collect all model IDs for supported_models all_model_ids = [model_id] + [s['model_id'] for s in sampler_list] @@ -409,44 +411,65 @@ def _short(mid: str) -> str: }], }) - # 2. Build static GPU-requiring sampler applications, + # 2. Build GPU-requiring applications (model + samplers), # then sort by GPU count DESCENDING before appending. # Largest PG deploys first → it has the fewest node choices → # avoids GPU scheduling deadlock on single-machine multi-node. gpu_apps: list[tuple[int, dict]] = [] # (gpu_count, app_config) - applications.append({ - 'name': f'model-{_short(model_id)}', - 'route_prefix': f'/api/v1/model/{model_id}', - 'import_path': 'model', - 'args': { - 'model_id': f'ms://{model_id}', - 'backend': backend, - 'env_vars': { - 'TWINKLE_TRUST_REMOTE_CODE': '1' - }, - 'train_mode': 'lora', - 'nproc_per_node': train_gpus, - 'device_group': { - 'name': 'model', - 'ranks': train_gpus, - 'device_type': 'cuda' - }, - 'device_mesh': { - 'device_type': 'cuda', - 'dp_size': train_gpus - }, - 'min_replicas': 0, - 'max_replicas': 1, - 'max_length': 500000, - 'idle_timeout_s': 600, - 'queue_config': { - 'rps_limit': 100, - 'tps_limit': 100000, - 'max_input_tokens': 500000 + # 2a. Training model worker (student) + gpu_apps.append(( + train_gpus, + { + 'name': + f'models-{model_short}', + 'route_prefix': + f'/api/v1/model/{model_id}', + 'import_path': + 'model', + 'args': { + 'backend': backend, + 'model_id': f'ms://{model_id}', + 'min_replicas': 0, + 'max_replicas': 1, + 'idle_timeout_s': 600, + 'max_length': 500000, # total tokens per forward pass (must match max_input_tokens) + 'nproc_per_node': train_gpus, + 'device_group': { + 'name': 'model', + 'ranks': train_gpus, + 'device_type': 'cuda', + }, + 'device_mesh': { + 'device_type': 'cuda', + 'dp_size': train_gpus, + }, + 'queue_config': { + 'rps_limit': 100, + 'tps_limit': 100000, + 'max_input_tokens': 500000, + }, + 'adapter_config': { + 'adapter_timeout': 600, + }, }, - }, - }) + 'deployments': [{ + 'name': 'ModelManagement', + 'autoscaling_config': { + 'min_replicas': 1, + 'max_replicas': 1, + 'target_ongoing_requests': 16, + }, + 'ray_actor_options': { + 'num_cpus': 0.1, + 'runtime_env': { + 'env_vars': { + 'TWINKLE_TRUST_REMOTE_CODE': '1' + }, + }, + }, + }], + })) # 2b. Sampler/teacher models sampler_name_count: dict[str, int] = {} diff --git a/tests/twinkle_client/test_auto_agent_tools.py b/tests/twinkle_client/test_auto_agent_tools.py index 747d463e1..9d5ecc59d 100644 --- a/tests/twinkle_client/test_auto_agent_tools.py +++ b/tests/twinkle_client/test_auto_agent_tools.py @@ -82,8 +82,11 @@ def test_generated_server_config_matches_launcher_schema(tmp_path, monkeypatch, assert gateway.deployments[0]['name'] == 'GatewayServer' assert gateway.args.supported_models == ['Qwen/student'] + [s['model_id'] for s in samplers] application = next(app for app in config.applications if app.import_path == 'model') + assert application.name == 'models-student' assert application.args.train_mode == 'lora' + assert application.args.adapter_config.adapter_timeout == 600 model = config.training_config.models[0] + assert model.env_vars['TWINKLE_TRUST_REMOTE_CODE'] == '1' assert model.mode == 'lora' assert model.config.min_replicas == 0 and model.config.max_replicas == 1 assert model.config.num_devices == 2 and model.config.idle_timeout_s == 600 From 4b100918655a0bfa4437c00768c2722222f859f1 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Mon, 26 Oct 2026 13:53:10 +0800 Subject: [PATCH 46/46] fix: isolate remaining model state and stream completion repairs --- src/twinkle/model/megatron/megatron.py | 3 + src/twinkle/model/multi_lora.py | 66 +++++-- .../model/transformers/hybrid/model.py | 6 +- .../transformers/multi_lora_transformers.py | 39 +++-- .../model/transformers/transformers.py | 9 + .../sampler/vllm_sampler/vllm_sampler.py | 27 +-- tests/sampler/test_vllm_stream_lifetime.py | 57 +++++++ .../test_multi_lora_checkpoint_context.py | 1 + .../test_multi_lora_distributed_slots.py | 161 ++++++++++++++++++ .../test_resident_training_reset.py | 93 ++++++++++ 10 files changed, 421 insertions(+), 41 deletions(-) create mode 100644 tests/sampler/test_vllm_stream_lifetime.py create mode 100644 tests/transformers/test_multi_lora_distributed_slots.py create mode 100644 tests/transformers/test_resident_training_reset.py diff --git a/src/twinkle/model/megatron/megatron.py b/src/twinkle/model/megatron/megatron.py index a782f9270..d47ee5896 100644 --- a/src/twinkle/model/megatron/megatron.py +++ b/src/twinkle/model/megatron/megatron.py @@ -1010,6 +1010,9 @@ def reload_initial_weights(self, **kwargs): for _model in self.strategy.unwrap_model(self.model): for param in _model.parameters(): param.grad = None + if kwargs.get('reset_training_state', False): + self.optimizer_group = {_default_adapter_name: self._construct_default_optimizer_group()} + self.optimizer_group[_default_adapter_name].adapter_name = _default_adapter_name if dist.is_initialized(): dist.barrier() diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index 797c3e628..c71fd3f38 100644 --- a/src/twinkle/model/multi_lora.py +++ b/src/twinkle/model/multi_lora.py @@ -12,7 +12,7 @@ from twinkle import torch_util from twinkle.data_format import InputFeature from twinkle.utils import get_logger -from .multi_lora_target_parameters import TargetParameterLoraManager +from .multi_lora_target_parameters import TargetParameterLoraManager, TargetParameterLoraWrapper logger = get_logger() @@ -124,6 +124,54 @@ def reset_adapter_status(self): else: self.module.set_adapter('lora_0') + def _named_slot_parameters(self, module): + """Identify owned adapter parameters through their actual layer containers.""" + slots = {tenant.adapter_name for tenant in self.loras} + for layer_name, layer in module.named_modules(): + if isinstance(layer, LoraLayer): + containers = layer.adapter_layer_names + elif isinstance(layer, TargetParameterLoraWrapper): + containers = ('lora_A', 'lora_B') + else: + continue + for container_name in containers: + container = getattr(layer, container_name) + for slot in slots & container.keys(): + value = container[slot] + prefix = '.'.join(part for part in (layer_name, container_name, slot) if part) + if isinstance(value, torch.nn.Parameter): + yield prefix, value + else: + for name, parameter in value.named_parameters(): + yield prefix + '.' + name, parameter + + def enable_slot_gradients(self): + """Keep owned adapter slots visible to reducers without unfreezing the base.""" + modules = self.module if isinstance(self.module, list) else [self.module] + for module in modules: + for _, parameter in self._named_slot_parameters(module): + parameter.requires_grad_(True) + + @contextmanager + def trainable_slots(self): + """Expose all slots while wrapping DDP/FSDP, then restore the active mask. + + Reducers discover parameters once during wrapping. The mask is restored + by adapter path because FSDP can replace the Parameter objects themselves. + """ + modules = self.module if isinstance(self.module, list) else [self.module] + masks = [{ + name: parameter.requires_grad + for name, parameter in self._named_slot_parameters(module) + } for module in modules] + try: + self.enable_slot_gradients() + yield + finally: + for module, mask in zip(modules, masks): + for name, parameter in self._named_slot_parameters(module): + parameter.requires_grad_(mask[name]) + def activate_adapter(self, tenant_adapter_name: str, call_enable=False): if not self.has_lora(tenant_adapter_name): raise ValueError(f'Adapter {tenant_adapter_name} does not exist') @@ -559,18 +607,12 @@ def _patch_megatron(_module): # PEFT's add_adapter calls set_adapter(active_adapters) which only keeps the # first adapter's requires_grad=True. We need ALL LoRA params to be trainable # so that MegatronDDP registers them all in its gradient buffers (main_grad). - def _enable_all_lora_grad(_module): - for name, param in _module.named_parameters(): - if 'lora_' in name and not param.requires_grad: - param.requires_grad_(True) - - if isinstance(module, list): - for _m in module: - _enable_all_lora_grad(_m) - else: - _enable_all_lora_grad(module) - self.module = module + modules = module if isinstance(module, list) else [module] + for chunk in modules: + for _, parameter in self._named_slot_parameters(chunk): + parameter.requires_grad_(True) + if not isinstance(module, list): self.lora_layer_names = [name for name, layer in module.named_modules() if isinstance(layer, Linear)] return module diff --git a/src/twinkle/model/transformers/hybrid/model.py b/src/twinkle/model/transformers/hybrid/model.py index 525f97243..3ccae3a71 100644 --- a/src/twinkle/model/transformers/hybrid/model.py +++ b/src/twinkle/model/transformers/hybrid/model.py @@ -64,7 +64,7 @@ def _adapter_context(self, adapter_name: str, disable_lora: bool = False): return with super()._adapter_context(adapter_name) as slot_name: - self._restore_preallocated_lora_grad() + self.multi_adapter.enable_slot_gradients() self.fft_slots.activate_fft_slot(adapter_name) yield slot_name @@ -72,14 +72,14 @@ def _adapter_context(self, adapter_name: str, disable_lora: bool = False): def _disabled_adapter_context(self, adapter_name: str): try: with super()._adapter_context(adapter_name, disable_lora=True) as slot_name: - self._restore_preallocated_lora_grad() + self.multi_adapter.enable_slot_gradients() self.fft_slots.deactivate_fft_slots() try: yield slot_name finally: self.fft_slots.activate_fft_slot(adapter_name) finally: - self._restore_preallocated_lora_grad() + self.multi_adapter.enable_slot_gradients() @remote_function() def add_adapter_to_model(self, adapter_name: str, config_or_dir: Union[PeftConfig, str], **kwargs): diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index b367af582..9c8fbb6b8 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/src/twinkle/model/transformers/multi_lora_transformers.py @@ -48,7 +48,12 @@ def __init__( self.device_mesh = device_mesh self.mixed_precision = mixed_precision self._fsdp_config = dict(fsdp_config or {}) - self._ddp_config = ddp_config or {} + if self._fsdp_config.get('reshard_after_forward', True) is not True: + raise ValueError('Multi-LoRA FSDP requires reshard_after_forward=True for safe adapter switching.') + self._ddp_config = dict(ddp_config or {}) + if self._ddp_config.get('find_unused_parameters') is False or self._ddp_config.get('static_graph', False): + raise ValueError('Multi-LoRA DDP requires find_unused_parameters=True and static_graph=False.') + self._ddp_config.setdefault('find_unused_parameters', True) self._memory_efficient_init = memory_efficient_init self._router_replay_enabled = bool(kwargs.pop('enable_router_replay', False)) self._router_replay_applied = False @@ -84,7 +89,7 @@ def __init__( # Initialize expert parallel attributes (required by set_optimizer in TransformersModel) self.optimizer_group: Dict[str, OptimizerGroup] = {} self.multi_adapter = MultiLora(max_loras=max_loras, max_r=max_r, max_length=max_length) - self.model.gradient_checkpointing_enable() + self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False}) self.model = self.multi_adapter.patch(self.model, target_modules=target_modules, lora_config=self.lora_config) # PEFT creates LoRA parameters in FP32 even when the base model uses a # lower-precision dtype. Align the preallocated slots before FSDP2 @@ -120,7 +125,15 @@ def unregister_mm_forward_hook(self, optimizer_group: OptimizerGroup): pass def _lazy_wrap_model(self): - return super()._lazy_wrap_model() + if self._model_wrapped: + return + # PEFT freezes inactive slots. DDP must nevertheless register all slots, + # since a later tenant may activate any of them on this resident model. + with self.multi_adapter.trainable_slots(): + super()._lazy_wrap_model() + # Reset must restore the synchronized (and possibly sharded) basis, + # not the rank-local initialization captured before model wrapping. + self.multi_adapter.save_initial_weights() def _maybe_apply_expert_parallel(self): if self._memory_efficient_init: @@ -140,24 +153,14 @@ def _ensure_target_parameter_lora_installed(self, config: LoraConfig) -> None: @contextmanager def _adapter_context(self, adapter_name: str, disable_lora: bool = False): - """Activate one LoRA tenant for a model operation.""" + """Activate a tenant while keeping FSDP2's trainable parameter set stable.""" try: with self.multi_adapter.adapter(adapter_name, disable_lora=disable_lora) as slot_name: - self._restore_preallocated_lora_grad() + self.multi_adapter.enable_slot_gradients() yield slot_name finally: # Leaving a disable_lora context re-enables adapter layers, which flips the flags again. - self._restore_preallocated_lora_grad() - - def _restore_preallocated_lora_grad(self) -> None: - # PEFT's set_adapter / disable_adapter_layers mark every inactive slot requires_grad=False, - # but FSDP2 captures each parameter's trainable state at lazy init. If that state drifts, - # gradient reduction for a slot other than the one active at wrap time breaks (e.g. fp32 - # grads assigned to bf16 params on NPU). Inactive slots still get no gradient because the - # patched LoRA forward only runs active adapters. - for name, parameter in self.multi_adapter.module.named_parameters(): - if 'lora_' in name and not parameter.requires_grad: - parameter.requires_grad_(True) + self.multi_adapter.enable_slot_gradients() @remote_function(dispatch='slice_dp', collect=collect_tensor_dict) def forward(self, *, inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], **kwargs): @@ -328,6 +331,7 @@ def _optimizer_param_name_mapping(self, adapter_name: str, optimizer: Optimizer) def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs): adapter_name = kwargs.get('adapter_name') self._check_adapter_valid(adapter_name) + self._lazy_wrap_model() with self._adapter_context(adapter_name), self.multi_adapter.save_context(adapter_name): checkpoint_dir = super().save(name, output_dir, interval, **kwargs) if dist.is_initialized(): @@ -338,6 +342,9 @@ def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs): def load(self, name: str, output_dir: Optional[str] = None, **kwargs): adapter_name = kwargs.get('adapter_name') self._check_adapter_valid(adapter_name) + # Establish the synchronized reset basis before checkpoint weights can + # replace it, including when load is the first model operation. + self._lazy_wrap_model() with self._adapter_context(adapter_name), self.multi_adapter.save_context(adapter_name): load_optimizer = kwargs.get('load_optimizer', False) if output_dir is None: diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index e240c765c..5a8a38561 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -1504,6 +1504,15 @@ def reload_initial_weights(self, **kwargs): (exclusive) model, so the next tenant starts from clean pretrained weights instead of the previous tenant's trained weights. """ + if kwargs.get('reset_training_state', False): + for optimizer_group in self.optimizer_group.values(): + self.unregister_mm_forward_hook(optimizer_group) + self.optimizer_group = {_default_adapter_name: self._construct_default_optimizer_group()} + self.optimizer_group[_default_adapter_name].adapter_name = _default_adapter_name + self.active_group = _default_adapter_name + if self._model_wrapped: + self.register_mm_forward_hook(self.optimizer_group[_default_adapter_name]) + if not self.model_id: logger.warning('reload_initial_weights skipped: model_id is not set (blank model).') return diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index ca9b58840..683f99067 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -143,17 +143,24 @@ async def _drain(): except Exception as exc: q.put(exc) finally: - q.put(_SENTINEL) + try: + await async_gen_coro.aclose() + finally: + q.put(_SENTINEL) - asyncio.run_coroutine_threadsafe(_drain(), self._async_loop) - - while True: - item = q.get() - if item is _SENTINEL: - break - if isinstance(item, Exception): - raise item - yield item + generation = asyncio.run_coroutine_threadsafe(_drain(), self._async_loop) + try: + while True: + item = q.get() + if item is _SENTINEL: + break + if isinstance(item, Exception): + raise item + yield item + finally: + # Queue termination or closing the Python iterator is not evidence + # that the engine-loop task has exited. Join before releasing weights. + generation.result() async def _create_engine_async(self, engine_cls, model_id, engine_kwargs): """Create engine in async context to ensure output_handler starts correctly.""" diff --git a/tests/sampler/test_vllm_stream_lifetime.py b/tests/sampler/test_vllm_stream_lifetime.py new file mode 100644 index 000000000..d7542bafd --- /dev/null +++ b/tests/sampler/test_vllm_stream_lifetime.py @@ -0,0 +1,57 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""The vLLM stream iterator owns its engine task until actual completion.""" +import ast +import asyncio +import threading +from pathlib import Path +from types import SimpleNamespace + + +def test_stream_iterator_close_joins_the_real_engine_loop_task(): + # Execute the repository method without constructing a GPU-backed vLLM engine. + path = Path('src/twinkle/sampler/vllm_sampler/vllm_sampler.py') + tree = ast.parse(path.read_text()) + method = next(node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == '_iter_in_loop') + namespace = {'asyncio': asyncio} + exec(compile(ast.Module(body=[method], type_ignores=[]), str(path), 'exec'), namespace) + loop = asyncio.new_event_loop() + running, finalizing, release, closed = [threading.Event() for _ in range(4)] + + def run_loop(): + asyncio.set_event_loop(loop) + running.set() + loop.run_forever() + + thread = threading.Thread(target=run_loop) + thread.start() + running.wait(1) + + async def generation(): + try: + yield 'first' + finally: + finalizing.set() + await asyncio.to_thread(release.wait) + + iterator = namespace['_iter_in_loop'](SimpleNamespace(_async_loop=loop), generation()) + closer = None + try: + assert next(iterator) == 'first' + assert finalizing.wait(1) + + def close(): + iterator.close() + closed.set() + + closer = threading.Thread(target=close) + closer.start() + assert not closed.wait(0.05) + release.set() + assert closed.wait(2) + finally: + release.set() + if closer: + closer.join(2) + loop.call_soon_threadsafe(loop.stop) + thread.join(2) + loop.close() diff --git a/tests/transformers/test_multi_lora_checkpoint_context.py b/tests/transformers/test_multi_lora_checkpoint_context.py index 80f4af241..39100bba6 100644 --- a/tests/transformers/test_multi_lora_checkpoint_context.py +++ b/tests/transformers/test_multi_lora_checkpoint_context.py @@ -21,6 +21,7 @@ def test_checkpoint_activates_target_tenant_before_optimizer_state_initializatio wrapper = object.__new__(MultiLoraTransformersModel) wrapper.__dict__['model'] = model wrapper.multi_adapter = manager + wrapper._model_wrapped = True wrapper.strategy = SimpleNamespace(unwrap_model=lambda value: value) with manager.adapter('first'): params = [param for name, param in model.named_parameters() if '.lora_0.' in name] diff --git a/tests/transformers/test_multi_lora_distributed_slots.py b/tests/transformers/test_multi_lora_distributed_slots.py new file mode 100644 index 000000000..ce177078a --- /dev/null +++ b/tests/transformers/test_multi_lora_distributed_slots.py @@ -0,0 +1,161 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Real two-process CPU regression for resident Multi-LoRA slot switching.""" +from datetime import timedelta +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from peft import LoraConfig +from safetensors.torch import save_file + +from twinkle.model.multi_lora import MultiLora +from twinkle.model.transformers.multi_lora_transformers import MultiLoraTransformersModel + + +def _wrapper(module, manager, wrap_model): + wrapper = object.__new__(MultiLoraTransformersModel) + torch.nn.Module.__init__(wrapper) + wrapper.model = module + wrapper.multi_adapter = manager + wrapper._model_wrapped = False + wrapper.optimizer_group = {} + wrapper._enable_expert_parallel = False + wrapper.sp_strategy = None + wrapper.strategy = SimpleNamespace( + capture_pre_ep_state_if_needed=lambda *args, **kwargs: None, + wrap_model=wrap_model, + unwrap_model=lambda model: model.module if isinstance(model, torch.nn.parallel.DistributedDataParallel) else model) + # This linear model does not use expert/sequence parallelism or GDN patches. + for name in ('_maybe_apply_expert_parallel', '_ensure_sp_strategy', '_ensure_gdn_padding_free_patch'): + object.__setattr__(wrapper, name, lambda: None) + return wrapper + + +def _distributed_slots(rank, rendezvous, mode): + dist.init_process_group('gloo', init_method='file://' + rendezvous, rank=rank, world_size=2, + timeout=timedelta(seconds=60)) + try: + torch.set_num_threads(1) + # DDP must synchronize initialization even when workers have different RNG states. + torch.manual_seed(42 + rank) + config = LoraConfig(r=2, lora_alpha=2, target_modules=['0']) + manager = MultiLora(max_loras=2, max_r=2) + module = manager.patch(torch.nn.Sequential(torch.nn.Linear(4, 4)), target_modules=['0'], lora_config=config) + manager.save_initial_weights() + manager.reset_adapter_status() + manager.acquire_lora('first', config) + wrapper = _wrapper(module, manager, lambda model: torch.nn.parallel.DistributedDataParallel( + model, find_unused_parameters=True)) + wrapper._lazy_wrap_model() + assert not any(param.requires_grad for name, param in module.named_parameters() if '.base_layer.' in name) + + def synchronized_weights(tenant): + state = manager.get_state_dict(tenant) + flat = torch.cat([value.detach().flatten() for value in state.values()]) + copies = [torch.empty_like(flat) for _ in range(2)] + dist.all_gather(copies, flat) + torch.testing.assert_close(copies[0], copies[1], rtol=0, atol=1e-7) + return {key: value.detach().clone() for key, value in state.items()} + + for tenant in (('first', 'second') if mode == 'switch' else ('first', 'reused')): + if tenant == 'second': + manager.acquire_lora(tenant, config) + elif tenant == 'reused': + manager.release_lora('first') + manager.acquire_lora(tenant, config) + initial = synchronized_weights(tenant) + with manager.adapter(tenant): + parameters = [p for p in module.parameters() if p.requires_grad] + optimizer = torch.optim.SGD(parameters, lr=.1) + for _ in range(3): + optimizer.zero_grad() + # Different local batches expose missing gradient synchronization. + wrapper.model(torch.ones(2, 4) * (rank + 1)).square().mean().backward() + optimizer.step() + trained = synchronized_weights(tenant) + assert any(not torch.equal(initial[name], value) for name, value in trained.items()) + with manager.adapter(tenant), torch.no_grad(): + fixed = wrapper.model(torch.ones(2, 4)).detach().clone() + # Rank 0 checkpoint must preserve predictions on every DP replica. + checkpoint = [trained if rank == 0 else None] + dist.broadcast_object_list(checkpoint, src=0) + manager.set_state_dict(tenant, checkpoint[0]) + torch.testing.assert_close(wrapper.model(torch.ones(2, 4)), fixed, rtol=0, atol=1e-7) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize('mode', ['switch', 'reuse']) +def test_all_slots_sync_across_dp_ranks_after_switch_and_reuse(tmp_path, mode): + mp.spawn(_distributed_slots, args=(str(tmp_path / 'rendezvous'), mode), nprocs=2, join=True) + + +@pytest.mark.parametrize('options', [{'find_unused_parameters': False}, {'static_graph': True}]) +def test_multi_lora_rejects_static_ddp_reducers(options, monkeypatch): + monkeypatch.setattr(MultiLoraTransformersModel, '_try_init_process_group', lambda _: None) + with pytest.raises(ValueError, match='Multi-LoRA DDP requires'): + MultiLoraTransformersModel(ddp_config=options) + + +@pytest.mark.parametrize('reshard_after_forward', [None, False, 2]) +def test_multi_lora_rejects_retained_or_partial_fsdp_views(reshard_after_forward, monkeypatch): + monkeypatch.setattr(MultiLoraTransformersModel, '_try_init_process_group', lambda _: None) + + def unexpected_strategy(*args): + raise AssertionError('Invalid FSDP config reached strategy initialization') + + monkeypatch.setattr(MultiLoraTransformersModel, '_decide_strategy', unexpected_strategy) + with pytest.raises(ValueError, match='Multi-LoRA FSDP requires reshard_after_forward=True'): + MultiLoraTransformersModel(fsdp_config={'reshard_after_forward': reshard_after_forward}) + + +def test_checkpoint_as_first_operation_preserves_fresh_slot_reset_basis(tmp_path): + config = LoraConfig(r=2, lora_alpha=2, target_modules=['0']) + manager = MultiLora(max_loras=1, max_r=2) + module = manager.patch(torch.nn.Sequential(torch.nn.Linear(4, 4)), target_modules=['0'], lora_config=config) + manager.save_initial_weights() + manager.acquire_lora('first', config) + initial = {name: value.detach().clone() for name, value in manager.get_state_dict('first').items()} + checkpoint = {name: torch.full_like(value, 2) for name, value in initial.items()} + save_file(checkpoint, str(tmp_path / 'adapter_model.safetensors')) + wrapper = _wrapper(module, manager, lambda model: model) + wrapper.optimizer_group = {'first': SimpleNamespace(optimizer=None)} + wrapper.load(str(tmp_path), adapter_name='first') + for name, value in manager.get_state_dict('first').items(): + torch.testing.assert_close(value, checkpoint[name]) + # Further operations must not replace the initial reset basis with trained weights. + wrapper._lazy_wrap_model() + manager.release_lora('first') + manager.acquire_lora('replacement', config) + for name, value in manager.get_state_dict('replacement').items(): + torch.testing.assert_close(value, initial[name]) + + +def test_slot_registration_never_unfreezes_base_parameters_with_lora_in_their_name(): + from collections import OrderedDict + config = LoraConfig(r=2, target_modules=['lora_projection']) + manager = MultiLora(max_loras=2, max_r=2) + module = manager.patch(torch.nn.Sequential(OrderedDict([ + ('lora_projection', torch.nn.Linear(4, 4)), + ])), target_modules=['lora_projection'], lora_config=config) + base = {name: parameter for name, parameter in module.named_parameters() if '.base_layer.' in name} + assert base + # Respect an explicitly frozen base even when its module name contains lora_. + for parameter in base.values(): + parameter.requires_grad_(False) + manager.reset_adapter_status() + mask = {name: parameter.requires_grad for name, parameter in module.named_parameters()} + with pytest.raises(RuntimeError, match='wrapper failed'): + with manager.trainable_slots(): + assert all(parameter.requires_grad for _, parameter in manager._named_slot_parameters(module)) + assert not any(parameter.requires_grad for parameter in base.values()) + raise RuntimeError('wrapper failed') + assert {name: parameter.requires_grad for name, parameter in module.named_parameters()} == mask + manager.acquire_lora('first', config) + wrapper = _wrapper(module, manager, lambda model: model) + for disable_lora in (False, True): + with wrapper._adapter_context('first', disable_lora=disable_lora): + assert not any(parameter.requires_grad for parameter in base.values()) + assert not any(parameter.requires_grad for parameter in base.values()) diff --git a/tests/transformers/test_resident_training_reset.py b/tests/transformers/test_resident_training_reset.py new file mode 100644 index 000000000..710ed55ba --- /dev/null +++ b/tests/transformers/test_resident_training_reset.py @@ -0,0 +1,93 @@ +import ast +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from twinkle.model.transformers.transformers import TransformersModel + + +def test_resident_reset_removes_previous_templates_forward_hook(): + class ScalingTemplate: + def __init__(self, scale): + self.scale = scale + + def pre_forward_hook(self, model, args, kwargs): + return (args[0] * self.scale,), kwargs + + def group(scale): + return SimpleNamespace(template=ScalingTemplate(scale), _handler=None) + + model = torch.nn.Linear(1, 1, bias=False) + model.weight.data.fill_(1) + wrapper = object.__new__(TransformersModel) + wrapper.__dict__['model'] = model + wrapper.model_id = 'local-reset-test' + wrapper._model_wrapped = True + wrapper.strategy = SimpleNamespace(unwrap_model=lambda value: value, + load_full_state_dict=lambda value, state: value.load_state_dict(state)) + wrapper.optimizer_group = {'': group(3)} + object.__setattr__(wrapper, '_construct_default_optimizer_group', lambda: group(1)) + wrapper.register_mm_forward_hook(wrapper.optimizer_group['']) + assert model(torch.ones(1, 1)).item() == 3 + initial = {'weight': torch.ones(1, 1)} + with patch('twinkle.model.transformers.transformers._read_hf_state_dict', return_value=initial): + for _ in range(2): + model.weight.grad = torch.ones_like(model.weight) + wrapper.reload_initial_weights(reset_training_state=True) + assert model(torch.ones(1, 1)).item() == 1 + assert model.weight.grad is None + assert len(model._forward_pre_hooks) == 1 + + +def test_blank_model_still_resets_training_state(): + wrapper = object.__new__(TransformersModel) + wrapper.model_id = '' + wrapper._model_wrapped = False + previous = SimpleNamespace(template=None, _handler=None, optimizer=object()) + fresh = SimpleNamespace(template=None, _handler=None, optimizer=None) + wrapper.optimizer_group = {'': previous} + object.__setattr__(wrapper, '_construct_default_optimizer_group', lambda: fresh) + wrapper.reload_initial_weights(reset_training_state=True) + assert wrapper.optimizer_group == {'': fresh} + assert wrapper.active_group == '' and fresh.adapter_name == '' + + +@pytest.mark.parametrize('reset_training_state', [False, True]) +def test_megatron_reload_resets_optimizer_state_only_when_requested(reset_training_state): + # Exercise the repository method without importing the optional Megatron stack. + path = Path('src/twinkle/model/megatron/megatron.py') + model_class = next(node for node in ast.parse(path.read_text()).body + if isinstance(node, ast.ClassDef) and node.name == 'MegatronModel') + method = next(node for node in model_class.body + if isinstance(node, ast.FunctionDef) and node.name == 'reload_initial_weights') + method.decorator_list = [] + namespace = {'_default_adapter_name': '', 'dist': SimpleNamespace(is_initialized=lambda: False)} + exec(compile(ast.Module(body=[method], type_ignores=[]), str(path), 'exec'), namespace) + model = torch.nn.Linear(1, 1, bias=False) + model.weight.grad = torch.ones_like(model.weight) + original_weights = {'weight': torch.ones_like(model.weight)} + + def load_weights(chunks, checkpoint, *, peft_format): + assert checkpoint == 'initial-checkpoint' and peft_format is False + for chunk in chunks: + chunk.load_state_dict(original_weights) + + previous = SimpleNamespace(optimizer=object(), cur_step=7) + fresh = SimpleNamespace(optimizer=None, cur_step=0) + wrapper = SimpleNamespace( + model=[model], _model_path='initial-checkpoint', + optimizer_group={'': previous}, + strategy=SimpleNamespace(unwrap_model=lambda value: value, + bridge=SimpleNamespace(load_weights=load_weights)), + _is_model_ddp_wrapped=lambda: False, + _construct_default_optimizer_group=lambda: fresh, + ) + namespace['reload_initial_weights'](wrapper, reset_training_state=reset_training_state) + assert model.weight.grad is None + torch.testing.assert_close(model.weight, original_weights['weight']) + assert wrapper.optimizer_group[''] is (fresh if reset_training_state else previous) + if reset_training_state: + assert fresh.adapter_name == '' and fresh.optimizer is None and fresh.cur_step == 0