From fe5a088260d70759a3b63b6ee9808569763122b2 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Wed, 14 Oct 2026 14:53:27 +0800 Subject: [PATCH 01/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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)