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/cookbook/client/async_rl/server_config.yaml b/cookbook/client/async_rl/server_config.yaml index 6e9c46441..fe86dd552 100644 --- a/cookbook/client/async_rl/server_config.yaml +++ b/cookbook/client/async_rl/server_config.yaml @@ -64,6 +64,10 @@ applications: args: backend: transformers model_id: ${oc.env:TWINKLE_LOCAL_MODEL_PATH} + min_replicas: 1 + max_replicas: 1 + adapter_config: + adapter_timeout: 30 max_loras: 8 max_length: 10240 data_plane_url: http://127.0.0.1:8000/api/v1/data-plane @@ -78,8 +82,6 @@ applications: queue_config: rps_limit: 100 tps_limit: 100000 - adapter_config: - adapter_timeout: 30 deployments: - name: ModelManagement autoscaling_config: diff --git a/cookbook/client/server/megatron/server_config.yaml b/cookbook/client/server/megatron/server_config.yaml index f10de4754..ebae1bff0 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. +resources: + 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,24 +121,28 @@ 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 + adapter_config: + adapter_timeout: 120 # Seconds without a session heartbeat before release + adapter_max_lifetime: 36000 # Maximum lifetime of an adapter in seconds (e.g., 10 hours) + max_loras: 3 # Handles per instance, not simultaneous backward calls + min_replicas: 0 # Do not reserve training GPUs at startup + max_replicas: 2 # Up to 2 independent two-GPU instances + 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: rps_limit: 20 # Max requests per second tps_limit: 131072 # Max tokens per second max_input_tokens: 32768 - adapter_config: - adapter_timeout: 120 # Seconds before idle adapter unload - adapter_max_lifetime: 36000 # Maximum lifetime of an adapter in seconds (e.g., 10 hours) deployments: - name: ModelManagement autoscaling_config: diff --git a/cookbook/client/server/megatron/server_config_4b.yaml b/cookbook/client/server/megatron/server_config_4b.yaml index 5b6cd9feb..a0de5caec 100644 --- a/cookbook/client/server/megatron/server_config_4b.yaml +++ b/cookbook/client/server/megatron/server_config_4b.yaml @@ -39,6 +39,8 @@ applications: args: backend: megatron model_id: "ms://Qwen/Qwen3.5-4B" # ModelScope model identifier + min_replicas: 1 + max_replicas: 1 max_length: 10240 nproc_per_node: 2 # Number of GPU processes per node device_group: @@ -53,7 +55,7 @@ applications: tps_limit: 100000 # Max tokens per second for a single user max_input_tokens: 60000 # Maximum input tokens per request adapter_config: - adapter_timeout: 30 # Seconds before idle adapter unload + adapter_timeout: 30 # Seconds without a session heartbeat before release adapter_max_lifetime: 36000 # Maximum lifetime of an adapter in seconds (e.g., 10 hours) max_loras: 5 # Maximum number of LoRA adapters per model deployments: diff --git a/cookbook/client/server/transformer/server_config.yaml b/cookbook/client/server/transformer/server_config.yaml index dd6cec311..e9a8bb5e7 100644 --- a/cookbook/client/server/transformer/server_config.yaml +++ b/cookbook/client/server/transformer/server_config.yaml @@ -55,6 +55,8 @@ applications: args: backend: transformers # Model backend: transformers | megatron model_id: "ms://Qwen/Qwen3.5-4B" # ModelScope model identifier + min_replicas: 1 + max_replicas: 1 max_length: 10240 nproc_per_node: 1 # Number of GPU processes per node device_group: @@ -68,7 +70,7 @@ applications: rps_limit: 100 # Max requests per second tps_limit: 100000 # Max tokens per second adapter_config: - adapter_timeout: 30 # Seconds before idle adapter unload + adapter_timeout: 30 # Seconds without a session heartbeat before release deployments: - name: ModelManagement autoscaling_config: diff --git a/cookbook/client/tinker/upload_to_hub.py b/cookbook/client/tinker/upload_to_hub.py index d35822144..96491aff7 100644 --- a/cookbook/client/tinker/upload_to_hub.py +++ b/cookbook/client/tinker/upload_to_hub.py @@ -9,8 +9,6 @@ # Task_Envelope with a request_id immediately, so the HTTP call never times out. # 2. The client's future layer long-polls /twinkle/retrieve_future and blocks # until the upload reaches a terminal state, raising on failure. -# (`upload_to_hub` keeps its `poll_interval` / `async_upload` arguments for -# signature compatibility; both are deprecated and have no effect.) # # Prerequisites: # - Server must be running (see server.py / server_config.yaml) diff --git a/cookbook/client/twinkle/upload_to_hub.py b/cookbook/client/twinkle/upload_to_hub.py index 7a2bc2fd2..e8e6921cb 100644 --- a/cookbook/client/twinkle/upload_to_hub.py +++ b/cookbook/client/twinkle/upload_to_hub.py @@ -9,8 +9,6 @@ # Task_Envelope with a request_id immediately, so the HTTP call never times out. # 2. The client's future layer long-polls /twinkle/retrieve_future and blocks # until the upload reaches a terminal state, raising on failure. -# (`upload_to_hub` keeps its `poll_interval` / `async_upload` arguments for -# signature compatibility; both are deprecated and have no effect.) # # Prerequisites: # - Server must be running (see server.py / server_config.yaml) diff --git a/docs/source_en/Usage Guide/Server and Client/Server.md b/docs/source_en/Usage Guide/Server and Client/Server.md index dd0064aba..ecb5f1730 100644 --- a/docs/source_en/Usage Guide/Server and Client/Server.md +++ b/docs/source_en/Usage Guide/Server and Client/Server.md @@ -233,9 +233,6 @@ applications: rps_limit: 100 # Max requests per second tps_limit: 10000 # Max tokens per second per user max_input_tokens: 10000 # Maximum input tokens per request - adapter_config: - adapter_timeout: 30 # Idle adapter timeout unload time (seconds) - adapter_max_lifetime: 36000 # Maximum adapter lifetime (seconds) max_loras: 1 # Maximum number of LoRA adapters per model deployments: - name: ModelManagement @@ -328,9 +325,6 @@ The difference from the Megatron backend is only in the `backend` parameter of t device_mesh: device_type: cuda dp_size: 2 - adapter_config: - adapter_timeout: 1800 # Idle adapter timeout unload time (seconds) - adapter_max_lifetime: 36000 deployments: - name: ModelManagement autoscaling_config: @@ -479,3 +473,58 @@ 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 + +Training services can load models when a request arrives and release devices when they are no longer needed. Keep the existing `applications` format and configure the `args` of an `import_path: model` entry: + +```yaml +- name: training-qwen + route_prefix: /api/v1/model/Qwen/Qwen3.5-4B + import_path: model + args: + model_id: ms://Qwen/Qwen3.5-4B + backend: transformers + train_mode: lora + nproc_per_node: 2 + device_group: {name: model, ranks: 2, device_type: cuda} + device_mesh: {device_type: cuda, dp_size: 2} + max_loras: 3 + min_replicas: 0 + max_replicas: 2 + idle_timeout_s: 600 +``` + +| Setting | How to use it | +|---------|---------------| +| `device_group`, `device_mesh` | Set the devices and parallel layout for each training instance. Adjust them to fit your model and available devices. Megatron requires at least two devices per instance. | +| `min_replicas` | Training instances to keep loaded at startup. With `0`, the first training request waits for model loading. | +| `max_replicas` | Maximum number of training instances. | +| `max_loras` | Maximum LoRA trainings per instance. Used only for LoRA training. | +| `idle_timeout_s` | Seconds of training inactivity before its handle is released. Heartbeats and result polling do not extend this time. Queued or running training operations are protected. | + +This example allows two instances, each holding three LoRA trainings. Clients are also subject to `server_config.per_token_model_limit`; clients using the same token share that token's quota. + +For full-parameter training, set `train_mode: full` and remove `max_loras`. Each full training uses its own instance. To offer LoRA and full training for the same weights, configure two model applications with different `name` and `route_prefix` values and select the mode in each. Clients choose the model alias from the corresponding route. See [Hybrid LoRA Training](../Hybrid-LoRA-Training.md) for Hybrid LoRA configuration. + +Release models when training finishes; Twinkle clients can call `remove_adapter()`. After the last training releases an instance, empty instances above `min_replicas` shut down. Save checkpoints regularly. If a handle expires or the service restarts, create a new training and restore from a checkpoint. Upgrade the Twinkle SDK together with the server. + +Complete configurations are available in `cookbook/client/server/transformer/server_config.yaml` and `cookbook/client/server/megatron/server_config.yaml`. + +### Automatic sampling scaling + +Set instance counts and the scaling target in the `args` of your `import_path: sampler` entry. For example: + +```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 +``` + +This configuration loads two sampling instances at startup and scales up to four as requests increase. Extra instances shut down after 600 seconds of continuous inactivity, leaving two loaded. `target_ongoing_requests` is the target number of in-flight requests per instance; a smaller value triggers earlier scaling. `device_group` and `device_mesh` determine the devices used by each instance. Allow enough devices for both training and sampling. + +Clients keep using the existing sampling API without selecting an instance. LoRA checkpoints load as requests need them. After a service restart, reapply custom templates and other runtime settings. diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" index 81fcacfe5..7b51eb47e 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" @@ -233,9 +233,6 @@ applications: rps_limit: 100 # 最大请求速率（每秒） tps_limit: 10000 # 单用户最大 token 速率（每秒） max_input_tokens: 10000 # 每次请求最大输入 token 数 - adapter_config: - adapter_timeout: 30 # 空闲适配器超时卸载时间（秒） - adapter_max_lifetime: 36000 # 适配器最大生命周期（秒） max_loras: 1 # 每个模型最多加载的 LoRA 数量 deployments: - name: ModelManagement @@ -328,9 +325,6 @@ Transformers 后端与 Megatron 后端的区别仅在 Model 服务的 `backend` device_mesh: device_type: cuda dp_size: 2 - adapter_config: - adapter_timeout: 1800 # 空闲适配器超时卸载时间（秒） - adapter_max_lifetime: 36000 deployments: - name: ModelManagement autoscaling_config: @@ -466,3 +460,65 @@ 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` 中设置： + +```yaml +- name: training-qwen + route_prefix: /api/v1/model/Qwen/Qwen3.5-4B + import_path: model + args: + model_id: ms://Qwen/Qwen3.5-4B + backend: transformers + train_mode: lora + nproc_per_node: 2 + device_group: {name: model, ranks: 2, device_type: cuda} + device_mesh: {device_type: cuda, dp_size: 2} + max_loras: 3 + min_replicas: 0 + max_replicas: 2 + idle_timeout_s: 600 + adapter_config: + adapter_timeout: 30 + adapter_max_lifetime: 36000 +``` + +| 配置项 | 如何使用 | +|--------|----------| +| `device_group`、`device_mesh` | 配置每个训练实例使用的设备和并行方式，请按模型大小与可用设备调整。Megatron 每实例至少需要 2 张设备。 | +| `min_replicas` | 启动时保留的训练实例数。设为 `0` 时，首次训练需要等待模型加载。 | +| `max_replicas` | 最多启动多少个训练实例。 | +| `max_loras` | 每个实例最多同时容纳多少个 LoRA 训练，仅用于 LoRA 模式。 | +| `idle_timeout_s` | 训练闲置多久后自动释放句柄，单位为秒。心跳和结果轮询不会延长这个时间，排队中或正在执行的训练不会因此中断。 | +| `adapter_config.adapter_timeout` | 客户端多久没有发送会话心跳后释放训练句柄，单位为秒，默认 `1800`。SDK 会自动发送心跳；设为 `null` 可关闭这项检查，但会话删除后句柄仍会失效。 | +| `adapter_config.adapter_max_lifetime` | 训练句柄从创建起最长保留多久，单位为秒，默认 `36000`（10 小时）。设为 `null` 可取消这个上限。 | + +这三个时间分别控制业务闲置、心跳中断和最长寿命，互不替代。已经排队或正在执行的任务会先完成，再回收句柄。长时间训练请定期保存 checkpoint；句柄过期后需要重新创建，并显式恢复 checkpoint。 + +上例最多提供 2 个实例，每个实例容纳 3 个 LoRA 训练。客户端同时还受 `server_config.per_token_model_limit` 限制；多个客户端使用同一个 token 时，共享该 token 的配额。 + +使用全参数训练时，将 `train_mode` 改为 `full`，并删除 `max_loras`。每个全参数训练单独使用一个实例。同一权重需要同时提供 LoRA 和全参数训练时，配置两个 model application，分别设置 `train_mode`，并使用不同的 `name` 和 `route_prefix`。客户端通过对应路由中的模型别名选择服务。Hybrid LoRA 的配置见 [Hybrid LoRA 训练](../Hybrid%20LoRA训练.md)。 + +训练结束后请主动释放模型，Twinkle 客户端可调用 `remove_adapter()`。最后一个训练释放后，超出 `min_replicas` 的空闲实例会关闭。训练期间应定期保存 checkpoint；句柄过期或服务重启后，需要重新创建训练并从 checkpoint 恢复。升级服务端时，也请配套升级 Twinkle SDK。 + +完整配置可参考仓库中的 `cookbook/client/server/transformer/server_config.yaml` 和 `cookbook/client/server/megatron/server_config.yaml`。 + +### 采样服务自动扩缩容 + +在 `import_path: sampler` 的 `args` 中设置实例数量和扩容目标。例如： + +```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 个采样实例，请求增多时最多扩到 4 个；额外实例连续空闲 600 秒后关闭，保留 2 个。`target_ongoing_requests` 是每个实例的目标在途请求数，调小后会更早扩容。每个实例的设备用量由 `device_group` 和 `device_mesh` 决定，请为训练和采样预留足够的设备。 + +客户端继续使用原来的采样接口，无需指定实例。采样所需的 LoRA checkpoint 按请求加载；服务重启后，请重新设置自定义 template 等运行时配置。 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/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/infra/_ray/ray_helper.py b/src/twinkle/infra/_ray/ray_helper.py index 4cc5f6a66..5dd7aa878 100644 --- a/src/twinkle/infra/_ray/ray_helper.py +++ b/src/twinkle/infra/_ray/ray_helper.py @@ -71,13 +71,20 @@ 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', + # Preserve Ray's default CPU policy outside managed instances. + num_cpus=0.01 if registry_namespace else None).remote() except ValueError: - RayHelper._registry = ray.get_actor('config_registry') + RayHelper._registry = ray.get_actor(registry_name, namespace=registry_namespace) assert RayHelper._registry is not None @staticmethod @@ -359,7 +366,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 +390,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 +426,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 +434,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..875ba4a08 100644 --- a/src/twinkle/infra/_ray/resource_manager.py +++ b/src/twinkle/infra/_ray/resource_manager.py @@ -20,6 +20,39 @@ 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 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) + 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] + manager.cpu_placement_groups = [] + manager._external_allocation = True + manager.device_groups = { + group.name: [{ + 'gpu_rank': chunk, + 'placement_group': placement_group, + 'ray_address': str(ray.get_runtime_context().gcs_address) + } for chunk in chunks] + } + 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 +329,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 a782f9270..d47ee5896 100644 --- a/src/twinkle/model/megatron/megatron.py +++ b/src/twinkle/model/megatron/megatron.py @@ -1010,6 +1010,9 @@ def reload_initial_weights(self, **kwargs): for _model in self.strategy.unwrap_model(self.model): for param in _model.parameters(): param.grad = None + if kwargs.get('reset_training_state', False): + self.optimizer_group = {_default_adapter_name: self._construct_default_optimizer_group()} + self.optimizer_group[_default_adapter_name].adapter_name = _default_adapter_name if dist.is_initialized(): dist.barrier() diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index 797c3e628..c71fd3f38 100644 --- a/src/twinkle/model/multi_lora.py +++ b/src/twinkle/model/multi_lora.py @@ -12,7 +12,7 @@ from twinkle import torch_util from twinkle.data_format import InputFeature from twinkle.utils import get_logger -from .multi_lora_target_parameters import TargetParameterLoraManager +from .multi_lora_target_parameters import TargetParameterLoraManager, TargetParameterLoraWrapper logger = get_logger() @@ -124,6 +124,54 @@ def reset_adapter_status(self): else: self.module.set_adapter('lora_0') + def _named_slot_parameters(self, module): + """Identify owned adapter parameters through their actual layer containers.""" + slots = {tenant.adapter_name for tenant in self.loras} + for layer_name, layer in module.named_modules(): + if isinstance(layer, LoraLayer): + containers = layer.adapter_layer_names + elif isinstance(layer, TargetParameterLoraWrapper): + containers = ('lora_A', 'lora_B') + else: + continue + for container_name in containers: + container = getattr(layer, container_name) + for slot in slots & container.keys(): + value = container[slot] + prefix = '.'.join(part for part in (layer_name, container_name, slot) if part) + if isinstance(value, torch.nn.Parameter): + yield prefix, value + else: + for name, parameter in value.named_parameters(): + yield prefix + '.' + name, parameter + + def enable_slot_gradients(self): + """Keep owned adapter slots visible to reducers without unfreezing the base.""" + modules = self.module if isinstance(self.module, list) else [self.module] + for module in modules: + for _, parameter in self._named_slot_parameters(module): + parameter.requires_grad_(True) + + @contextmanager + def trainable_slots(self): + """Expose all slots while wrapping DDP/FSDP, then restore the active mask. + + Reducers discover parameters once during wrapping. The mask is restored + by adapter path because FSDP can replace the Parameter objects themselves. + """ + modules = self.module if isinstance(self.module, list) else [self.module] + masks = [{ + name: parameter.requires_grad + for name, parameter in self._named_slot_parameters(module) + } for module in modules] + try: + self.enable_slot_gradients() + yield + finally: + for module, mask in zip(modules, masks): + for name, parameter in self._named_slot_parameters(module): + parameter.requires_grad_(mask[name]) + def activate_adapter(self, tenant_adapter_name: str, call_enable=False): if not self.has_lora(tenant_adapter_name): raise ValueError(f'Adapter {tenant_adapter_name} does not exist') @@ -559,18 +607,12 @@ def _patch_megatron(_module): # PEFT's add_adapter calls set_adapter(active_adapters) which only keeps the # first adapter's requires_grad=True. We need ALL LoRA params to be trainable # so that MegatronDDP registers them all in its gradient buffers (main_grad). - def _enable_all_lora_grad(_module): - for name, param in _module.named_parameters(): - if 'lora_' in name and not param.requires_grad: - param.requires_grad_(True) - - if isinstance(module, list): - for _m in module: - _enable_all_lora_grad(_m) - else: - _enable_all_lora_grad(module) - self.module = module + modules = module if isinstance(module, list) else [module] + for chunk in modules: + for _, parameter in self._named_slot_parameters(chunk): + parameter.requires_grad_(True) + if not isinstance(module, list): self.lora_layer_names = [name for name, layer in module.named_modules() if isinstance(layer, Linear)] return module diff --git a/src/twinkle/model/optimizer_group.py b/src/twinkle/model/optimizer_group.py index 150e694bc..59b4b331b 100644 --- a/src/twinkle/model/optimizer_group.py +++ b/src/twinkle/model/optimizer_group.py @@ -64,15 +64,15 @@ def accumulate_metrics(self, is_training): """Accumulate metrics for train/eval status. Override in subclass if needed.""" status = self.train_status if is_training else self.eval_status if len(status.metrics) > 0 and status.inputs is not None and status.outputs is not None: + metric_kwargs = { + **status.forward_kwargs, + 'lr': self._get_lr(), + 'step': self.cur_step - 1, + 'gradient_accumulation_steps': self.gradient_accumulation_steps, + 'grad_norm': self._last_grad_norm, + } for metric in status.metrics: - metric.accumulate( - status.inputs, - status.outputs, - lr=self._get_lr(), - step=self.cur_step - 1, - gradient_accumulation_steps=self.gradient_accumulation_steps, - grad_norm=self._last_grad_norm, - **status.forward_kwargs) + metric.accumulate(status.inputs, status.outputs, **metric_kwargs) def calculate_metrics(self, is_training): """Calculate and return metrics.""" diff --git a/src/twinkle/model/transformers/hybrid/fft_slots.py b/src/twinkle/model/transformers/hybrid/fft_slots.py index 39cde9579..38d675d2b 100644 --- a/src/twinkle/model/transformers/hybrid/fft_slots.py +++ b/src/twinkle/model/transformers/hybrid/fft_slots.py @@ -120,6 +120,10 @@ def activate_fft_slot(self, adapter_name: str) -> None: fft_adapter_name = self._fft_adapter_name(adapter_name) if self.is_hybrid(adapter_name) else None selected_adapters = [fft_adapter_name] if fft_adapter_name is not None else [] for wrapper in self._iter_fft_wrappers(): + # Selecting a regular LoRA can disable auxiliary wrappers. Restoring + # an FFT slot name alone does not restore its execution state. + if fft_adapter_name is not None and wrapper.disable_adapters: + wrapper.enable_adapters(True) if wrapper.active_adapters != selected_adapters: wrapper.set_adapter(selected_adapters) self._enable_preallocated_fft_grad(wrapper) diff --git a/src/twinkle/model/transformers/hybrid/model.py b/src/twinkle/model/transformers/hybrid/model.py index 525f97243..3ccae3a71 100644 --- a/src/twinkle/model/transformers/hybrid/model.py +++ b/src/twinkle/model/transformers/hybrid/model.py @@ -64,7 +64,7 @@ def _adapter_context(self, adapter_name: str, disable_lora: bool = False): return with super()._adapter_context(adapter_name) as slot_name: - self._restore_preallocated_lora_grad() + self.multi_adapter.enable_slot_gradients() self.fft_slots.activate_fft_slot(adapter_name) yield slot_name @@ -72,14 +72,14 @@ def _adapter_context(self, adapter_name: str, disable_lora: bool = False): def _disabled_adapter_context(self, adapter_name: str): try: with super()._adapter_context(adapter_name, disable_lora=True) as slot_name: - self._restore_preallocated_lora_grad() + self.multi_adapter.enable_slot_gradients() self.fft_slots.deactivate_fft_slots() try: yield slot_name finally: self.fft_slots.activate_fft_slot(adapter_name) finally: - self._restore_preallocated_lora_grad() + self.multi_adapter.enable_slot_gradients() @remote_function() def add_adapter_to_model(self, adapter_name: str, config_or_dir: Union[PeftConfig, str], **kwargs): diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index b367af582..9c8fbb6b8 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/src/twinkle/model/transformers/multi_lora_transformers.py @@ -48,7 +48,12 @@ def __init__( self.device_mesh = device_mesh self.mixed_precision = mixed_precision self._fsdp_config = dict(fsdp_config or {}) - self._ddp_config = ddp_config or {} + if self._fsdp_config.get('reshard_after_forward', True) is not True: + raise ValueError('Multi-LoRA FSDP requires reshard_after_forward=True for safe adapter switching.') + self._ddp_config = dict(ddp_config or {}) + if self._ddp_config.get('find_unused_parameters') is False or self._ddp_config.get('static_graph', False): + raise ValueError('Multi-LoRA DDP requires find_unused_parameters=True and static_graph=False.') + self._ddp_config.setdefault('find_unused_parameters', True) self._memory_efficient_init = memory_efficient_init self._router_replay_enabled = bool(kwargs.pop('enable_router_replay', False)) self._router_replay_applied = False @@ -84,7 +89,7 @@ def __init__( # Initialize expert parallel attributes (required by set_optimizer in TransformersModel) self.optimizer_group: Dict[str, OptimizerGroup] = {} self.multi_adapter = MultiLora(max_loras=max_loras, max_r=max_r, max_length=max_length) - self.model.gradient_checkpointing_enable() + self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False}) self.model = self.multi_adapter.patch(self.model, target_modules=target_modules, lora_config=self.lora_config) # PEFT creates LoRA parameters in FP32 even when the base model uses a # lower-precision dtype. Align the preallocated slots before FSDP2 @@ -120,7 +125,15 @@ def unregister_mm_forward_hook(self, optimizer_group: OptimizerGroup): pass def _lazy_wrap_model(self): - return super()._lazy_wrap_model() + if self._model_wrapped: + return + # PEFT freezes inactive slots. DDP must nevertheless register all slots, + # since a later tenant may activate any of them on this resident model. + with self.multi_adapter.trainable_slots(): + super()._lazy_wrap_model() + # Reset must restore the synchronized (and possibly sharded) basis, + # not the rank-local initialization captured before model wrapping. + self.multi_adapter.save_initial_weights() def _maybe_apply_expert_parallel(self): if self._memory_efficient_init: @@ -140,24 +153,14 @@ def _ensure_target_parameter_lora_installed(self, config: LoraConfig) -> None: @contextmanager def _adapter_context(self, adapter_name: str, disable_lora: bool = False): - """Activate one LoRA tenant for a model operation.""" + """Activate a tenant while keeping FSDP2's trainable parameter set stable.""" try: with self.multi_adapter.adapter(adapter_name, disable_lora=disable_lora) as slot_name: - self._restore_preallocated_lora_grad() + self.multi_adapter.enable_slot_gradients() yield slot_name finally: # Leaving a disable_lora context re-enables adapter layers, which flips the flags again. - self._restore_preallocated_lora_grad() - - def _restore_preallocated_lora_grad(self) -> None: - # PEFT's set_adapter / disable_adapter_layers mark every inactive slot requires_grad=False, - # but FSDP2 captures each parameter's trainable state at lazy init. If that state drifts, - # gradient reduction for a slot other than the one active at wrap time breaks (e.g. fp32 - # grads assigned to bf16 params on NPU). Inactive slots still get no gradient because the - # patched LoRA forward only runs active adapters. - for name, parameter in self.multi_adapter.module.named_parameters(): - if 'lora_' in name and not parameter.requires_grad: - parameter.requires_grad_(True) + self.multi_adapter.enable_slot_gradients() @remote_function(dispatch='slice_dp', collect=collect_tensor_dict) def forward(self, *, inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], **kwargs): @@ -328,6 +331,7 @@ def _optimizer_param_name_mapping(self, adapter_name: str, optimizer: Optimizer) def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs): adapter_name = kwargs.get('adapter_name') self._check_adapter_valid(adapter_name) + self._lazy_wrap_model() with self._adapter_context(adapter_name), self.multi_adapter.save_context(adapter_name): checkpoint_dir = super().save(name, output_dir, interval, **kwargs) if dist.is_initialized(): @@ -338,6 +342,9 @@ def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs): def load(self, name: str, output_dir: Optional[str] = None, **kwargs): adapter_name = kwargs.get('adapter_name') self._check_adapter_valid(adapter_name) + # Establish the synchronized reset basis before checkpoint weights can + # replace it, including when load is the first model operation. + self._lazy_wrap_model() with self._adapter_context(adapter_name), self.multi_adapter.save_context(adapter_name): load_optimizer = kwargs.get('load_optimizer', False) if output_dir is None: diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index eb90d4a9f..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)] @@ -1501,6 +1504,15 @@ def reload_initial_weights(self, **kwargs): (exclusive) model, so the next tenant starts from clean pretrained weights instead of the previous tenant's trained weights. """ + if kwargs.get('reset_training_state', False): + for optimizer_group in self.optimizer_group.values(): + self.unregister_mm_forward_hook(optimizer_group) + self.optimizer_group = {_default_adapter_name: self._construct_default_optimizer_group()} + self.optimizer_group[_default_adapter_name].adapter_name = _default_adapter_name + self.active_group = _default_adapter_name + if self._model_wrapped: + self.register_mm_forward_hook(self.optimizer_group[_default_adapter_name]) + if not self.model_id: logger.warning('reload_initial_weights skipped: model_id is not set (blank model).') return diff --git a/src/twinkle/protocol/types/base.py b/src/twinkle/protocol/types/base.py index d428afc81..a3ef46cd7 100644 --- a/src/twinkle/protocol/types/base.py +++ b/src/twinkle/protocol/types/base.py @@ -22,7 +22,7 @@ ``status_code`` -- an execution-time failure is delivered with HTTP 200, so the value is systematically unequal to the response status code. 4. A closed value set on a wire field is declared as ``Literal`` / enum, never a - bare ``str`` (see ``QueueStateLiteral`` in ``errors.py``). + bare ``str`` (see ``QueueState`` in ``errors.py``). Field roles ----------- diff --git a/src/twinkle/protocol/types/errors.py b/src/twinkle/protocol/types/errors.py index b896e4eed..5d3e5e32b 100644 --- a/src/twinkle/protocol/types/errors.py +++ b/src/twinkle/protocol/types/errors.py @@ -20,10 +20,14 @@ class can be lined up against its tinker counterpart): from .base import ResponseModel -# Closed value set, kept in sync with the server-side ``QueueState`` enum values -# (a consistency test asserts the two sets are equal). Wire fields carrying a queue -# state declare this alias, never a bare ``str`` (naming ruling 4). -QueueStateLiteral = Literal['active', 'paused_rate_limit', 'paused_capacity', 'unknown'] + +class QueueState(StrEnum): + """Scheduling hints; these are independent of a task's lifecycle status.""" + + ACTIVE = 'active' + PAUSED_RATE_LIMIT = 'paused_rate_limit' + PAUSED_CAPACITY = 'paused_capacity' + UNKNOWN = 'unknown' class ErrorCategory(StrEnum): diff --git a/src/twinkle/protocol/types/lifecycle.py b/src/twinkle/protocol/types/lifecycle.py index 32555e24b..641c47fbb 100644 --- a/src/twinkle/protocol/types/lifecycle.py +++ b/src/twinkle/protocol/types/lifecycle.py @@ -7,18 +7,25 @@ """ from __future__ import annotations -from typing import Any, Literal, Optional +from enum import StrEnum +from typing import Any, Optional from .base import ResponseModel, StrictRequest -from .errors import ErrorPayload, QueueStateLiteral +from .errors import ErrorPayload, QueueState -# The lifecycle status of a queued task. Kept in sync with the server-side -# ``TaskStatus`` enum values (a consistency test asserts the two sets are equal). -TaskStatus = Literal['pending', 'queued', 'running', 'completed', 'failed', 'cancelled'] -# The two states past which a task never changes again. ``frozenset`` so a caller -# cannot mutate the shared set. -TERMINAL_STATUSES: frozenset[str] = frozenset({'completed', 'failed', 'cancelled'}) +class TaskStatus(StrEnum): + """Canonical task states shared by the wire envelope and server execution.""" + + PENDING = 'pending' + QUEUED = 'queued' + RUNNING = 'running' + COMPLETED = 'completed' + FAILED = 'failed' + CANCELLED = 'cancelled' + + +TERMINAL_STATUSES = frozenset({TaskStatus.COMPLETED.value, TaskStatus.FAILED.value, TaskStatus.CANCELLED.value}) class RetrieveFutureRequest(StrictRequest): @@ -79,5 +86,5 @@ class TaskEnvelope(ResponseModel): status: TaskStatus result: Any | None = None # set iff status == 'completed' error: ErrorPayload | None = None # set iff status == 'failed' - queue_state: QueueStateLiteral | None = None + queue_state: QueueState | None = None queue_state_reason: str | None = None diff --git a/src/twinkle/protocol/types/model.py b/src/twinkle/protocol/types/model.py index f2e6bd976..51e5c9ce9 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,14 @@ class CreateRequest(StrictRequest): # --------------------------------------------------------------------------- # -class AdapterRequest(StrictRequest): +class TrainingRequest(StrictRequest): + """A training handle is control metadata, never a backend kwarg.""" + + training_id: str = Field(min_length=1) + adapter_name: str + + +class AdapterRequest(TrainingRequest): """The shared shape of an adapter-scoped operation. ``seq_id`` is an idempotency key, not a backend parameter: the submit shell @@ -46,7 +55,6 @@ class AdapterRequest(StrictRequest): which would silently disable that dedup. """ - adapter_name: str seq_id: int | None = None gradient_accumulation_steps: int | None = backend_kwarg(default=None, ge=1) @@ -83,10 +91,9 @@ 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 is_training: bool = True @@ -105,7 +112,7 @@ class CalculateMetricRequest(StrictRequest): # --------------------------------------------------------------------------- # -class _InlineForwardBase(StrictRequest): +class _InlineForwardBase(TrainingRequest): """Fields common to the three inline forward endpoints.""" inputs: WireInputBatch @@ -127,7 +134,6 @@ class _InlineForwardBase(StrictRequest): class ForwardRequest(_InlineForwardBase): """Body of ``POST /twinkle/forward``: keeps the graph, mutates no gradients.""" - adapter_name: str disable_lora: bool | None = backend_kwarg(default=None) @@ -139,14 +145,12 @@ class ForwardOnlyRequest(_InlineForwardBase): nothing is mutated to be idempotent about. """ - adapter_name: str disable_lora: bool | None = backend_kwarg(default=None) class ForwardBackwardTaskRequest(_InlineForwardBase): """Body of ``POST /twinkle/forward_backward``: accumulates gradients.""" - adapter_name: str seq_id: int | None = None sync_gradients: bool | None = backend_kwarg(default=None) loss_scale: float | None = backend_kwarg(default=None) @@ -161,7 +165,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) @@ -169,7 +173,6 @@ class DataPlaneForwardRequest(StrictRequest): # Values are *field paths*, not parameter values, so this is not a passthrough # region: nothing in it is forwarded verbatim. kwarg_fields: dict[str, str] = Field(default_factory=dict) - adapter_name: str seq_id: int | None = None task: str | None = backend_kwarg(default=None) temperature: float | None = backend_kwarg(default=None, gt=0) @@ -180,7 +183,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 @@ -191,7 +194,6 @@ class DataPlaneForwardOnlyRequest(StrictRequest): input_refs: list[DataRef] = Field(min_length=1) input_field: str | None = None kwarg_fields: dict[str, str] = Field(default_factory=dict) - adapter_name: str task: str | None = backend_kwarg(default=None) temperature: float | None = backend_kwarg(default=None, gt=0) return_logits: bool | None = backend_kwarg(default=None) @@ -222,25 +224,22 @@ 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 @@ -250,26 +249,22 @@ class SetTemplateRequest(StrictRequest): """ template_cls: str - adapter_name: str 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,8 +273,7 @@ class ApplyPatchRequest(StrictRequest): # --------------------------------------------------------------------------- # -class SaveRequest(StrictRequest): - adapter_name: str +class SaveRequest(TrainingRequest): name: str | None = None save_optimizer: bool = False is_sampler: bool = False # If True, delete existing sampler weights before saving @@ -287,8 +281,7 @@ class SaveRequest(StrictRequest): merge_lora: bool | None = backend_only('megatron', default=None) -class LoadRequest(StrictRequest): - adapter_name: str +class LoadRequest(TrainingRequest): name: str load_optimizer: bool = False no_load_optim: bool | None = backend_only('megatron', default=None) @@ -296,7 +289,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 +298,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 +412,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/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index ca9b58840..073fb0e9e 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.""" @@ -452,7 +459,7 @@ def sample_stream( def sample_stream_to_queue(self, queue, inputs, sampling_params=None, adapter_name='', adapter_path=None): """Push streaming deltas to a cross-process Ray queue.""" - from twinkle.server.sampler.backends import stream_to_queue + from twinkle.server.sampler.backends.streaming import stream_to_queue stream_to_queue(self, queue, inputs, sampling_params, adapter_name, adapter_path) @remote_function(dispatch='all', collect='first') diff --git a/src/twinkle/server/__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/checkpoint/checkpoint_manager.py b/src/twinkle/server/checkpoint/checkpoint_manager.py index d86fadadd..136768a72 100644 --- a/src/twinkle/server/checkpoint/checkpoint_manager.py +++ b/src/twinkle/server/checkpoint/checkpoint_manager.py @@ -27,7 +27,7 @@ from twinkle import get_logger from twinkle.hub import HubOperation from twinkle.protocol.types import ResolvedLoadPath -from twinkle.server.task_queue.types import UserTaskError +from twinkle.server.exceptions import RequestRejectedError from twinkle.utils.parallel import FileLock, acquire_lock, release_lock from .paths import CHECKPOINT_INFO_FILENAME, validate_user_path from .training_run_manager import BaseFileManager, BaseTrainingRunManager @@ -370,7 +370,7 @@ async def save_sampler(self, model_id: str, name: str | None, """ explicit_name = self.get_ckpt_name(name) if name else None if explicit_name == 'latest': - raise UserTaskError('Checkpoint name "latest" is reserved for unnamed sampler saves.') + raise RequestRejectedError('Checkpoint name "latest" is reserved for unnamed sampler saves.') version_id = uuid.uuid4().hex checkpoint_name = explicit_name or 'latest' uri_name = explicit_name or f'live_{version_id}' diff --git a/src/twinkle/server/config/__init__.py b/src/twinkle/server/config/__init__.py index 3e3fe2f71..f6cb0ca2e 100644 --- a/src/twinkle/server/config/__init__.py +++ b/src/twinkle/server/config/__init__.py @@ -2,10 +2,12 @@ """Server configuration package — aggregate root and per-deployment specs.""" from .application_spec import (ApplicationSpec, DataPlaneArgs, HttpOptions, ModelArgs, ProcessorArgs, SamplerArgs, - ServerArgs, SpectralHybridArgs) + ServerArgs) from .persistence import PersistenceConfig from .server_config import ServerConfig +from .task_queue import TaskQueueConfig from .telemetry import TelemetryConfig +from .training import SpectralHybridArgs __all__ = [ 'ApplicationSpec', @@ -18,5 +20,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 de5c8bc3a..5ed19bae3 100644 --- a/src/twinkle/server/config/application_spec.py +++ b/src/twinkle/server/config/application_spec.py @@ -11,12 +11,12 @@ """ 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 .task_queue import TaskQueueConfig +from .training import ModelArgs # Env var keys the launcher sets from the gateway ``server_config`` so that any # Ray worker (model / sampler / processor), not just the gateway, applies the @@ -55,57 +55,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. - - The ``backend`` field selects the model implementation (transformers / - megatron / mock) — see ``model/app.py`` for the dispatch site that picks - the runtime class based on this value. - """ - - model_id: str - nproc_per_node: int = 1 - device_group: dict[str, Any] - device_mesh: dict[str, Any] - backend: Literal['mock', 'transformers', 'megatron'] - train_mode: Literal['lora', 'hybrid', 'full'] = 'lora' - strategy: Literal['accelerate', 'native_fsdp'] | None = None - adapter_config: dict[str, Any] | None = None - queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) - max_loras: int = 5 - max_r: int = Field(default=32, gt=0) - max_length: int | None = None - target_modules: str | list[str] = 'all-linear' - hybrid: SpectralHybridArgs | None = None - data_plane_url: str | None = None - - @model_validator(mode='after') - def _validate_spectral_backend(self): - if self.strategy is not None and self.backend != 'transformers': - raise ValueError('strategy is only supported by the transformers backend') - if self.hybrid is not None: - if self.backend != 'transformers': - raise ValueError('hybrid is only supported by the transformers backend') - if self.train_mode == 'full': - raise ValueError('hybrid is not supported with train_mode="full"') - return self - - class SamplerArgs(_ArgsBase): """Args for the ``sampler`` deployment. @@ -120,6 +69,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): @@ -243,7 +213,7 @@ class ApplicationSpec(BaseModel): # ``import_path`` and defaults a missing block to ``{}``), so the field is # required here — the validator runs first and fills it. args: ServerArgs | ModelArgs | SamplerArgs | ProcessorArgs | DataPlaneArgs - deployments: list[dict[str, Any]] = Field(default_factory=list) + deployments: list[dict[str, Any]] = Field(default_factory=list, max_length=1) @model_validator(mode='before') @classmethod @@ -269,23 +239,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/resources.py b/src/twinkle/server/config/resources.py new file mode 100644 index 000000000..12a46d31d --- /dev/null +++ b/src/twinkle/server/config/resources.py @@ -0,0 +1,55 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shared instance startup, allocation, shutdown and monitoring budgets.""" +import os +from pydantic import BaseModel, ConfigDict, Field + + +class ResourceOptions(BaseModel): + model_config = ConfigDict(extra='forbid', frozen=True) + + allocation_timeout_s: float = Field(default=3, gt=0, allow_inf_nan=False) + startup_timeout_s: float = Field(default=1800, gt=0, allow_inf_nan=False) + shutdown_timeout_s: float = Field(default=60, gt=0, allow_inf_nan=False) + sweep_interval_s: float = Field(default=5, gt=0, allow_inf_nan=False) + + +SHARED_WORKER_ENV_KEYS: tuple[str, ...] = ( + 'TWINKLE_STATE_RAY_NAMESPACE', + 'TWINKLE_TELEMETRY_ENABLED', + 'TWINKLE_TELEMETRY_DEBUG', + 'TWINKLE_TELEMETRY_SERVICE', + 'TWINKLE_TELEMETRY_ENDPOINT', + 'TWINKLE_TELEMETRY_INTERVAL', + 'TWINKLE_MODEL_ID_ALIASES', +) + + +def build_shared_env_vars() -> dict[str, str]: + return {key: os.environ[key] for key in SHARED_WORKER_ENV_KEYS if key in os.environ} + + +def build_persistence_env_vars() -> dict[str, str]: + from twinkle.server.config.persistence import PERSISTENCE_ENV_KEYS + return {key: os.environ[key] for key in PERSISTENCE_ENV_KEYS if key in os.environ} + + +def build_server_state_env_vars() -> dict[str, str]: + from twinkle.server.config.application_spec import SERVER_STATE_ENV_KEYS + return {key: os.environ[key] for key in SERVER_STATE_ENV_KEYS if key in os.environ} + + +def build_propagated_env_vars() -> dict[str, str]: + merged: dict[str, str] = {} + merged.update(build_shared_env_vars()) + merged.update(build_persistence_env_vars()) + merged.update(build_server_state_env_vars()) + return merged + + +def build_worker_env_vars() -> dict[str, str]: + """Policy plus the explicit process settings shared by launcher and resource owners.""" + env = build_propagated_env_vars() + for key in ('PYTHONPATH', 'TWINKLE_DEFAULT_SAVE_DIR', 'TMPDIR', 'OMP_NUM_THREADS'): + if key in os.environ: + env.setdefault(key, os.environ[key]) + return env diff --git a/src/twinkle/server/config/server_config.py b/src/twinkle/server/config/server_config.py index 82bac4d44..7de25f094 100644 --- a/src/twinkle/server/config/server_config.py +++ b/src/twinkle/server/config/server_config.py @@ -19,10 +19,12 @@ 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 .resources import ResourceOptions +from .task_queue import TaskQueueConfig from .telemetry import TelemetryConfig +from .training import TrainingConfig, TrainingModelConfig class ServerConfig(BaseModel): @@ -36,6 +38,7 @@ class ServerConfig(BaseModel): telemetry: TelemetryConfig = Field(default_factory=TelemetryConfig) persistence: PersistenceConfig = Field(default_factory=PersistenceConfig) task_queue: TaskQueueConfig = Field(default_factory=TaskQueueConfig) + resources: ResourceOptions = Field(default_factory=ResourceOptions) applications: list[ApplicationSpec] = Field(default_factory=list) # ---- loading ---------------------------------------------------------- # @@ -69,10 +72,30 @@ def from_yaml(cls, path: str | Path) -> ServerConfig: @model_validator(mode='after') def _validate_cross_field(self) -> ServerConfig: + if sum(app.import_path == 'server' for app in self.applications) > 1: + raise ValueError('configure one gateway per server') if self.persistence.mode == 'redis' and not self.persistence.redis_url: raise ValueError("persistence.redis_url is required when persistence.mode == 'redis'", ) + if self.training_config.models: + 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 + @property + def training_config(self) -> TrainingConfig: + """Normalize model application args into the internal runtime catalog.""" + return TrainingConfig( + **self.resources.model_dump(), + models=[ + TrainingModelConfig.from_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/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 new file mode 100644 index 000000000..5e2346582 --- /dev/null +++ b/src/twinkle/server/config/training.py @@ -0,0 +1,251 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Training YAML arguments and the controller's validated runtime configuration.""" +from __future__ import annotations + +import math +from pydantic import BaseModel, ConfigDict, Field, SerializeAsAny, ValidationInfo, field_validator, model_validator +from typing import TYPE_CHECKING, Any, Literal + +from twinkle.hub.model_alias import extract_model_alias +from twinkle.server.exceptions import RequestRejectedError +from .resources import ResourceOptions +from .task_queue import TaskQueueConfig + +if TYPE_CHECKING: + from .application_spec import ApplicationSpec + + +class SpectralHybridArgs(BaseModel): + """Strict server-owned Spectral Hybrid allocation and optimizer defaults.""" + + model_config = ConfigDict(extra='forbid') + + allocation_path: str = Field(min_length=1) + default_lr_lora: float = Field(default=2.5e-5, gt=0) + default_lr_fft: float = Field(default=1.0e-6, gt=0) + + @field_validator('default_lr_lora', 'default_lr_fft') + @classmethod + def _finite_learning_rate(cls, value: float) -> float: + if not math.isfinite(value): + raise ValueError('learning rate must be finite') + return value + + +class AdapterConfig(BaseModel): + """Session liveness and absolute lifetime of a training handle, in seconds.""" + + model_config = ConfigDict(extra='forbid', frozen=True) + + adapter_timeout: float | None = Field(default=1800, gt=0, allow_inf_nan=False) + adapter_max_lifetime: float | None = Field(default=36000, gt=0, allow_inf_nan=False) + + +class TrainingArguments(BaseModel): + """Training options shared by model YAML and resolved instance configuration.""" + + model_config = ConfigDict(extra='forbid') + + min_replicas: int = Field(default=0, ge=0, strict=True) + max_replicas: int | None = Field(default=None, gt=0, strict=True) + idle_timeout_s: float = Field(default=1800, gt=0, allow_inf_nan=False) + adapter_config: AdapterConfig = Field(default_factory=AdapterConfig) + strategy: Literal['accelerate', 'native_fsdp'] | None = None + queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) + max_length: int | None = Field(default=None, gt=0, strict=True) + + @model_validator(mode='after') + def _validate_replicas(self): + if self.max_replicas is not None and self.min_replicas > self.max_replicas: + raise ValueError('min_replicas must not exceed max_replicas') + return self + + +class LoraArguments(BaseModel): + """LoRA parameters shared by YAML and LoRA instance configuration.""" + + max_loras: int = Field(default=5, gt=0, strict=True) + max_r: int = Field(default=32, gt=0, strict=True) + target_modules: str | list[str] = 'all-linear' + hybrid: SpectralHybridArgs | None = None + + +class ModelArgs(TrainingArguments, LoraArguments): + """Args for the ``model`` deployment. + + The ``backend`` field selects the model implementation (transformers / + megatron / mock) — see ``model/app.py`` for the dispatch site that picks + the runtime class based on this value. + """ + + model_id: str + nproc_per_node: int = 1 + device_group: dict[str, Any] + device_mesh: dict[str, Any] + backend: Literal['mock', 'transformers', 'megatron'] + env_vars: dict[str, str] = Field(default_factory=dict) + train_mode: Literal['lora', 'hybrid', 'full'] = 'lora' + data_plane_url: str | None = None + + @model_validator(mode='after') + def _validate_spectral_backend(self): + if self.strategy is not None and self.backend != 'transformers': + raise ValueError('strategy is only supported by the transformers backend') + if self.hybrid is not None: + if self.backend != 'transformers': + raise ValueError('hybrid is only supported by the transformers backend') + if self.train_mode == 'full': + raise ValueError('hybrid is not supported with train_mode="full"') + return self + + +class TrainingMesh(BaseModel): + """Sizes accepted by DeviceMesh.from_sizes, including overlapping EP axes.""" + + model_config = ConfigDict(extra='forbid', frozen=True) + + device_type: str | None = None + dp_size: int = Field(default=1, gt=0, strict=True) + fsdp_size: int | None = Field(default=None, gt=0, strict=True) + 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(TrainingArguments): + model_config = ConfigDict(extra='forbid', frozen=True) + + num_devices: int = Field(gt=0, strict=True) + device_mesh: TrainingMesh + + @model_validator(mode='after') + def _validate_resources(self): + if self.device_mesh.world_size() != self.num_devices: + raise ValueError('device_mesh world_size must equal num_devices') + return self + + +class LoraTrainingConfig(TrainingModeConfig, LoraArguments): + pass + + +class TrainingModelConfig(BaseModel): + """One application exposes one training mode and one instance specification.""" + + model_config = ConfigDict(extra='forbid', frozen=True) + + 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'] + mode: Literal['lora', 'full'] + config: SerializeAsAny[TrainingModeConfig] + route_prefix: str | None = None + data_plane_url: str | None = None + env_vars: dict[str, str] = Field(default_factory=dict) + deploy_options: dict = Field(default_factory=dict) + + @field_validator('config', mode='before') + @classmethod + def _instance_config(cls, value, info: ValidationInfo): + config_type = LoraTrainingConfig if info.data.get('mode') == 'lora' else TrainingModeConfig + data = value.model_dump() if isinstance(value, TrainingModeConfig) else value + return config_type.model_validate(data) + + @model_validator(mode='after') + def _validate_backend(self): + if (self.name.startswith('/') or any(part in ('', '.', '..') for part in self.name.split('/')) + 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') + if self.backend == 'megatron' and self.config.num_devices < 2: + raise ValueError('megatron training instances require at least 2 devices') + if self.config.device_mesh.device_type not in (None, self.device_type): + raise ValueError('device_mesh.device_type must match device_type') + if self.config.strategy is not None and self.backend != 'transformers': + raise ValueError('strategy requires the transformers backend') + if isinstance(self.config, LoraTrainingConfig) and self.config.hybrid and self.backend != 'transformers': + raise ValueError('hybrid requires the transformers backend') + return self + + @property + def key(self) -> tuple[str, str]: + return self.name, self.mode + + @property + def capacity(self) -> int: + return self.config.max_loras if isinstance(self.config, LoraTrainingConfig) else 1 + + @classmethod + def from_application(cls, app: ApplicationSpec) -> TrainingModelConfig: + args = app.args + if not isinstance(args, ModelArgs): + raise TypeError('Training configuration requires model application arguments') + if len(app.deployments) > 1: + raise ValueError(f'{app.name}: a model application must have one deployment configuration') + deployment = app.deployments[0] if app.deployments else {} + device_type = str( + args.device_group.get('device_type', + args.device_mesh.get('device_type', + 'cpu' if args.backend == 'mock' else 'cuda'))).lower() + mode = 'full' if args.train_mode == 'full' else 'lora' + fields = set(TrainingArguments.model_fields) + if mode == 'lora': + fields.update(LoraArguments.model_fields) + config = {key: getattr(args, key) for key in fields} + config.update( + num_devices=args.device_group.get('ranks', args.nproc_per_node), + device_mesh={ + **args.device_mesh, 'device_type': device_type + }) + env_vars = { + **(deployment.get('ray_actor_options') or {}).get('runtime_env', {}).get('env_vars', {}), + **args.env_vars, + } + return cls( + name=extract_model_alias(app.route_prefix) or app.name, + model_id=args.model_id, + backend=args.backend, + device_type=device_type, + mode=mode, + config=config, + route_prefix=app.route_prefix.rstrip('/') or '/', + data_plane_url=args.data_plane_url, + env_vars=env_vars, + deploy_options={ + key: value + for key, value in deployment.items() if key != 'name' + }) + + +class TrainingConfig(ResourceOptions): + """Resolved catalog; public model entries stay in applications.""" + + models: tuple[TrainingModelConfig, ...] = () + + @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 + + def get(self, name: str, mode: Literal['lora', 'full']) -> TrainingModelConfig: + """Resolve a declared specification without allocating resources.""" + model = next((model for model in self.models if model.name == name), None) + if model is None: + raise RequestRejectedError(f'Unknown training model: {name}') + if model.mode != mode: + raise RequestRejectedError(f'Training model {name} does not provide {mode} mode') + return model diff --git a/src/twinkle/server/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/exceptions.py b/src/twinkle/server/exceptions.py index b8985879d..5fc3cc044 100644 --- a/src/twinkle/server/exceptions.py +++ b/src/twinkle/server/exceptions.py @@ -9,6 +9,8 @@ from __future__ import annotations +from typing import Literal + from twinkle.protocol.types.errors import ErrorCategory @@ -176,3 +178,25 @@ def __init__(self, current_holder: str) -> None: super().__init__('This deployment runs in full-parameter (exclusive) mode and is already ' f'held by another training task ({current_holder}). Only one full-parameter ' 'training task is allowed at a time; retry after it is released.') + + +class TrainingExpiredError(ResourceNotFoundError): + """The training handle expired and must be recreated before explicit resume.""" + + +class ProtocolMismatchError(RequestRejectedError): + """The requested client protocol does not match the training handle.""" + + +class ResourceReleaseError(TwinkleServerError, RuntimeError): + """Owned resource removal failed at a known phase; capacity remains reserved.""" + + error_code = 503 + + def __init__(self, message: str, *, phase: Literal['worker_death', 'pg_removal']): + super().__init__(message) + self.phase = phase + + +class BackendBusyError(ResourceExhaustedError, RuntimeError): + """A timed-out backend still owns execution/weight resources on this replica.""" diff --git a/src/twinkle/server/gateway/app.py b/src/twinkle/server/gateway/app.py index 0c54f506b..c68ea7271 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.lifecycle.proxy import ServiceProxy from twinkle.server.state import get_server_state from twinkle.utils.logger import get_logger from .openai_handlers import _register_openai_routes -from .proxy import ServiceProxy from .tinker_handlers import _register_gateway_tinker_routes from .twinkle_handlers import _register_gateway_twinkle_routes @@ -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/openai_handlers.py b/src/twinkle/server/gateway/openai_handlers.py index cec235b02..c62bc4c60 100644 --- a/src/twinkle/server/gateway/openai_handlers.py +++ b/src/twinkle/server/gateway/openai_handlers.py @@ -28,12 +28,15 @@ from twinkle.server.lifecycle.envelope import envelope_from_record from twinkle.server.utils import get_template_for_model from twinkle.utils.logger import get_logger -from . import routes from .openai_bridge import make_error, translate_chat_request, translate_response, translate_stream_chunk from .use_cases import poll_future logger = get_logger() +TWINKLE_SAMPLE = 'twinkle/sample' +TWINKLE_SAMPLE_STREAM = 'twinkle/sample_stream' +TWINKLE_SET_TEMPLATE = 'twinkle/set_template' + _OPENAI_SAMPLE_TIMEOUT_ENV = 'TWINKLE_OPENAI_SAMPLE_TIMEOUT' _DEFAULT_OPENAI_SAMPLE_TIMEOUT = 1800.0 @@ -145,7 +148,7 @@ async def chat_completions( # Non-streaming: proxy to /twinkle/sample, translate response response = await self.proxy.proxy_request( request, - endpoint=routes.TWINKLE_SAMPLE, + endpoint=TWINKLE_SAMPLE, base_model=base_model, service_type='sampler', body_override=body_bytes, @@ -214,7 +217,7 @@ async def _sse_generator(): try: async for line in self.proxy.proxy_request_stream( request, - endpoint=routes.TWINKLE_SAMPLE_STREAM, + endpoint=TWINKLE_SAMPLE_STREAM, base_model=base_model, service_type='sampler', body_override=body_bytes, @@ -330,7 +333,7 @@ async def _ensure_template( try: resp = await gateway.proxy.proxy_request( request, - endpoint=routes.TWINKLE_SET_TEMPLATE, + endpoint=TWINKLE_SET_TEMPLATE, base_model=base_model, service_type='sampler', body_override=set_template_body, diff --git a/src/twinkle/server/gateway/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/routes.py b/src/twinkle/server/gateway/routes.py deleted file mode 100644 index 0cececc05..000000000 --- a/src/twinkle/server/gateway/routes.py +++ /dev/null @@ -1,26 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Route-path constants shared by the gateway's proxy and its 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/gateway/tinker_handlers.py b/src/twinkle/server/gateway/tinker_handlers.py index 93809d780..1277ddd35 100644 --- a/src/twinkle/server/gateway/tinker_handlers.py +++ b/src/twinkle/server/gateway/tinker_handlers.py @@ -17,9 +17,8 @@ from twinkle.hub import HubOperation from twinkle.server.checkpoint import create_checkpoint_manager, create_training_run_manager +from twinkle.server.lifecycle.envelope import error_payload_from_failure from twinkle.server.middleware.auth import get_token_from_request -from twinkle.server.state.models import FutureFailureRecord -from twinkle.server.task_errors import trim_traceback from twinkle.utils.logger import get_logger from .use_cases import delete_checkpoint from .use_cases import get_training_run as get_training_run_use_case @@ -27,42 +26,10 @@ logger = get_logger() -# Keys must equal ``state.models.FAILURE_REASON_CODES`` (guarded by test_envelope). -_TINKER_FAILURE_WIRE: dict[str, tuple[int, str]] = { - 'invalid_request': (400, 'user'), - 'request_rejected': (400, 'user'), - 'resource_not_found': (404, 'user'), - 'full_mode_busy': (409, 'user'), - 'input_tokens_exceeded': (422, 'user'), - 'batch_size_invalid': (422, 'user'), - 'rate_limit_exceeded': (429, 'user'), - 'resource_quota_exceeded': (429, 'user'), - 'cancelled': (499, 'user'), - 'endpoint_unavailable': (501, 'server'), - 'state_contention': (503, 'server'), - 'backend_gate_unavailable': (503, 'server'), - 'orphaned_replica': (503, 'server'), - 'execution_timeout': (504, 'server'), - 'deadline_exceeded': (500, 'server'), - 'internal_error': (500, 'server'), -} - def _tinker_error_from_failure(stored: Any, *, request_id: str) -> dict[str, Any]: - """Map domain failure state to the existing Tinker-compatible error body.""" - failure = FutureFailureRecord.model_validate(stored) - error_code, category = _TINKER_FAILURE_WIRE.get(failure.reason_code, (500, 'server')) - payload: dict[str, Any] = { - 'error': failure.message[:1024], - 'category': category, - 'error_code': error_code, - 'request_id': request_id, - } - if failure.details is not None: - payload['details'] = failure.details - if failure.diagnostic and category == 'server': - payload['traceback'] = trim_traceback(failure.diagnostic) - return payload + """Serialize the shared failure presentation with Tinker's field names.""" + return error_payload_from_failure(stored, request_id=request_id).model_dump(mode='json', exclude_none=True) def _register_gateway_tinker_routes(app: FastAPI, self_fn: Callable[[], GatewayServer]) -> None: @@ -83,7 +50,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 +88,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/__init__.py b/src/twinkle/server/launcher/__init__.py index 0a955b8bc..83769bd32 100644 --- a/src/twinkle/server/launcher/__init__.py +++ b/src/twinkle/server/launcher/__init__.py @@ -8,7 +8,6 @@ - ``server_launcher`` — the ``ServerLauncher`` class and the ``launch_server`` entry point (the real launcher logic); - ``builder_registry`` — ``import_path`` → deployment-builder resolution; -- ``env_propagation`` — telemetry / persistence env-var collection for workers. This ``__init__`` is a thin aggregator: it only re-exports the public surface so the dotted paths ``twinkle.server.launcher.ServerLauncher`` / ``launch_server`` diff --git a/src/twinkle/server/launcher/env_propagation.py b/src/twinkle/server/launcher/env_propagation.py deleted file mode 100644 index c478444d9..000000000 --- a/src/twinkle/server/launcher/env_propagation.py +++ /dev/null @@ -1,52 +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_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 75fd0bd52..445de7ad2 100644 --- a/src/twinkle/server/launcher/server_launcher.py +++ b/src/twinkle/server/launcher/server_launcher.py @@ -1,11 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""``ServerLauncher`` and the ``launch_server`` entry point. - -The real launcher module. It lives here (rather than in the package -``__init__.py``) so the package's ``__init__`` stays a thin aggregator; the -public dotted paths ``twinkle.server.launcher.ServerLauncher`` / -``launch_server`` are preserved by the ``__init__`` re-export. -""" +"""Initialize server configuration, deploy applications, and own their shutdown.""" from __future__ import annotations import json @@ -13,290 +7,254 @@ import signal import threading from collections.abc import Callable +from copy import deepcopy from pathlib import Path 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 +from twinkle.server.config.resources import build_propagated_env_vars, build_worker_env_vars from .builder_registry import get_builders, resolve_builder -from .env_propagation import build_propagated_env_vars logger = get_logger() class ServerLauncher: - """ - Unified server launcher. - - This class handles Ray/Serve initialization and application deployment. - - Attributes: - config: The server configuration dictionary - ray_namespace: The Ray namespace for the cluster - """ + """Own startup and cleanup; business scheduling stays in the controllers.""" - def __init__( - self, - config: ServerConfig, - ray_namespace: str | None = None, - ): - """ - Initialize the server launcher. - - Args: - config: A validated :class:`ServerConfig` instance. Raw dicts are - rejected — operators must build ``ServerConfig`` via - ``ServerConfig.from_yaml`` or its constructor so cross-field - validation runs before the launcher consumes anything. - ray_namespace: Ray namespace (default: 'twinkle_cluster') - """ + def __init__(self, config: ServerConfig, ray_namespace: str | None = None): if not isinstance(config, ServerConfig): raise TypeError('ServerLauncher requires a typed ServerConfig instance; ' f'got {type(config).__name__}. Build one with ' 'ServerConfig.from_yaml(path) or ServerConfig(...).') - self.config: ServerConfig = config - self.ray_namespace = ray_namespace + self.config = config + self.ray_namespace = ray_namespace or config.ray_namespace or 'twinkle_cluster' self._builders: dict[str, Callable] = {} self._ray_initialized = False self._serve_started = False + self._training_controller = None + self._sampler_controllers = [] def _get_builders(self) -> dict[str, Callable]: - """Get (and cache) the builder functions for all app types.""" if not self._builders: self._builders = get_builders() return self._builders def _resolve_builder(self, import_path: str) -> Callable: - """Resolve an import_path to a builder function. - - Raises: - ValueError: If the import_path cannot be resolved. - """ return resolve_builder(import_path, self._get_builders()) + def _configure_environment(self) -> None: + if any(app.import_path in ('model', 'sampler') for app in self.config.applications): + os.environ['TWINKLE_STATE_RAY_NAMESPACE'] = self.ray_namespace + os.environ.update(self.config.persistence.to_env_vars()) + gateways = [app for app in self.config.applications if app.import_path == 'server'] + if gateways: + os.environ.update(gateways[0].args.server_config.to_env_vars()) + aliases = build_model_alias_map(self.config.applications) + if aliases: + os.environ[MODEL_ID_ALIASES_ENV] = json.dumps(aliases, ensure_ascii=False) + else: + os.environ.pop(MODEL_ID_ALIASES_ENV, None) + self._configure_telemetry() + + def _configure_telemetry(self) -> None: + telemetry = self.config.telemetry + os.environ['TWINKLE_TELEMETRY_ENABLED'] = str(int(telemetry.enabled)) + if not telemetry.enabled: + return + from twinkle.server.telemetry import init_telemetry + init_telemetry(telemetry) + os.environ.update({ + 'TWINKLE_TELEMETRY_DEBUG': str(int(telemetry.debug)), + 'TWINKLE_TELEMETRY_SERVICE': telemetry.service_name, + 'TWINKLE_TELEMETRY_ENDPOINT': telemetry.otlp_endpoint, + 'TWINKLE_TELEMETRY_INTERVAL': str(telemetry.export_interval_ms), + }) + def _init_ray(self) -> None: - """Initialize Ray if not already initialized.""" if self._ray_initialized: return - import ray - - namespace = self.ray_namespace or self.config.ray_namespace or 'twinkle_cluster' - if not ray.is_initialized(): - # Use runtime_env to apply patches in worker processes - # This is required because Ray Serve's ProxyActor runs in separate processes runtime_env = get_runtime_env_for_patches() - # Propagate telemetry + persistence env vars to all Ray workers - propagated_env_vars = build_propagated_env_vars() - if propagated_env_vars: - merged_env_vars = dict(runtime_env.get('env_vars') or {}) - merged_env_vars.update(propagated_env_vars) - runtime_env['env_vars'] = merged_env_vars - # Connect to existing cluster if available, otherwise start local instance - ray.init( - address='auto', - namespace=namespace, - runtime_env=runtime_env, - ) - logger.info(f'Ray initialized with namespace={namespace}') - + runtime_env['env_vars'] = {**runtime_env.get('env_vars', {}), **build_propagated_env_vars()} + ray.init(address='auto', namespace=self.ray_namespace, runtime_env=runtime_env) self._ray_initialized = True def _start_serve(self) -> None: - """Start Ray Serve with http_options from config.""" if self._serve_started: return - from ray import serve - - try: - from ray.serve.context import _get_global_client - _get_global_client() - # Serve is running, shut it down before re-starting - serve.shutdown() - except Exception: - # Serve not running — nothing to shut down - pass - - http_options = self.config.http_options.model_dump() - serve_kwargs: dict[str, Any] = {'http_options': http_options} - # ``proxy_location`` controls where the Ray Serve HTTP proxy runs - # (``EveryNode`` / ``HeadOnly`` / ``Disabled``). The example configs - # set this field, so honour it here instead of silently ignoring. + options = {'http_options': self.config.http_options.model_dump()} if self.config.proxy_location: - serve_kwargs['proxy_location'] = self.config.proxy_location - serve.start(**serve_kwargs) - logger.info(f'Ray Serve started with http_options={http_options}, ' - f'proxy_location={self.config.proxy_location!r}') - + options['proxy_location'] = self.config.proxy_location + serve.start(**options) self._serve_started = True - def _deploy_application(self, app_spec: ApplicationSpec) -> None: - """Deploy a single application. - - Args: - app_spec: Validated :class:`ApplicationSpec` from the typed config. - """ + def _deployment_options(self, options: dict, env: dict | None = None) -> dict: + options = deepcopy(options) + options.pop('name', None) + actor = options.setdefault('ray_actor_options', {}) + runtime_env = actor.setdefault('runtime_env', {}) + runtime_env['env_vars'] = {**build_worker_env_vars(), **runtime_env.get('env_vars', {}), **(env or {})} + return options + + def _gateway_arguments(self) -> dict[str, Any]: + args = {'http_options': self.config.http_options.model_dump()} + models = self.config.training_config.models + if models: + args['training_models'] = [model.name for model in models] + args['supported_models'] = list( + dict.fromkeys( + alias for app in self.config.applications + if app.import_path == 'sampler' and (alias := extract_model_alias(app.route_prefix, 'sampler')))) + return args + + def _deploy_sampler(self, app: ApplicationSpec) -> None: + import ray from ray import serve - name = app_spec.name - route_prefix = app_spec.route_prefix - import_path = app_spec.import_path - # Shallow-dump the typed args to a kwargs dict WITHOUT recursing into - # nested models: ``dict(model)`` yields top-level (field, value) pairs - # with nested models left as instances, so ``queue_config`` stays a - # typed ``TaskQueueConfig`` and is not re-serialized to a dict only to be - # revived via ``from_dict`` inside the builder (the prior triple - # validation). ``None`` values are dropped to match ``exclude_none``. - args = {k: v for k, v in dict(app_spec.args).items() if v is not None} - deployments = list(app_spec.deployments or []) - - logger.info(f'Starting {name} at {route_prefix}...') - - builder = self._resolve_builder(import_path) - - deploy_options = {} - if deployments: - if len(deployments) > 1: - logger.warning(f'Application "{name}" has {len(deployments)} deployments configured, ' - f'but only the first deployment will be used.') - deploy_config = deployments[0] - if isinstance(deploy_config, dict): - deploy_options = {k: v for k, v in deploy_config.items() if k != 'name'} - - # Inject telemetry + persistence env vars into the deployment's - # runtime_env so that Ray Serve replicas (worker processes) can - # initialize telemetry and resolve the configured persistence backend - # regardless of deployment startup order. - # User-specified env_vars take precedence over our defaults. - propagated_env_vars = build_propagated_env_vars() - if propagated_env_vars: - ray_actor_options = dict(deploy_options.get('ray_actor_options') or {}) - runtime_env = dict(ray_actor_options.get('runtime_env') or {}) - env_vars = dict(runtime_env.get('env_vars') or {}) - for k, v in propagated_env_vars.items(): - env_vars.setdefault(k, v) - runtime_env['env_vars'] = env_vars - ray_actor_options['runtime_env'] = runtime_env - deploy_options['ray_actor_options'] = ray_actor_options - - # Pass http_options to server apps for internal proxy routing - if import_path == 'server': - args.setdefault('http_options', self.config.http_options.model_dump()) - - app = builder(deploy_options=deploy_options, **args) - - serve.run(app, name=name, route_prefix=route_prefix) - logger.info(f'Deployed {name} at {route_prefix}') - - def launch(self) -> None: - """Launch the server with all configured applications. + from twinkle.server.sampler.pool.actor import create_sampler_controller + from twinkle.server.sampler.pool.router import build_sampler_router + controller, options, pool_id = create_sampler_controller(app, self.ray_namespace, + self.config.resources.model_dump(), + build_worker_env_vars()) + self._sampler_controllers.append(controller) + ray.get(controller.initialize.remote()) + serve.run( + build_sampler_router(self.ray_namespace, pool_id, options), name=app.name, route_prefix=app.route_prefix) - Blocks the calling thread to keep the server running. Installs signal - handlers for SIGINT/SIGTERM so that ``serve.shutdown()`` is called on - termination instead of leaving orphaned deployments. - """ - # Apply Ray Serve patches before initializing Ray - apply_ray_serve_patches() - - # Initialize telemetry if configured - telemetry = self.config.telemetry - if telemetry.enabled: - from twinkle.server.telemetry import init_telemetry - init_telemetry(telemetry) - # Export config to env vars for Ray worker processes - os.environ['TWINKLE_TELEMETRY_ENABLED'] = '1' - os.environ['TWINKLE_TELEMETRY_DEBUG'] = '1' if telemetry.debug else '0' - os.environ['TWINKLE_TELEMETRY_SERVICE'] = telemetry.service_name - os.environ['TWINKLE_TELEMETRY_ENDPOINT'] = telemetry.otlp_endpoint - os.environ['TWINKLE_TELEMETRY_INTERVAL'] = str(telemetry.export_interval_ms) - - # Export top-level persistence to env vars so any worker - # (not just Gateway) can build the same backend on first call to - # get_server_state(). - persistence = self.config.persistence - for k, v in persistence.to_env_vars().items(): - os.environ[k] = v - logger.info(f'Persistence backend configured: mode={persistence.mode}') - - # Export the gateway ``server`` application's ServerState policy (quota / - # expiry / cleanup / metrics interval) to env vars for the same reason: - # so every worker's first ``get_server_state()`` applies the configured - # values instead of the hardcoded defaults. Without this the model worker - # that enforces ``per_token_model_limit`` runs on the default (30), - # silently ignoring the YAML value. - server_specs = [a for a in self.config.applications if a.import_path == 'server'] - if len(server_specs) > 1: - logger.warning(f'{len(server_specs)} "server" applications declared; using the first ' - 'for ServerState policy env propagation.') - if server_specs: - server_state_env = server_specs[0].args.server_config.to_env_vars() - for k, v in server_state_env.items(): - os.environ[k] = v - if server_state_env: - logger.info(f'ServerState policy exported to worker env: {server_state_env}') - - model_alias_map = build_model_alias_map(self.config.applications) - if model_alias_map: - os.environ[MODEL_ID_ALIASES_ENV] = json.dumps(model_alias_map, ensure_ascii=False) - else: - os.environ.pop(MODEL_ID_ALIASES_ENV, None) - - self._init_ray() - self._start_serve() - - applications = self.config.applications - if not applications: - logger.warning('No applications configured') + def _deploy_application(self, app_spec: ApplicationSpec) -> None: + """Model declarations go to the training controller; other apps deploy here.""" + if app_spec.import_path == 'model': return + if app_spec.import_path == 'sampler': + self._deploy_sampler(app_spec) + return + from ray import serve + args = {key: value for key, value in dict(app_spec.args).items() if value is not None} + if app_spec.import_path == 'server': + defaults = self._gateway_arguments() + for key, value in defaults.items(): + args.setdefault(key, value) + deployment = app_spec.deployments[0] if app_spec.deployments else {} + options = self._deployment_options(deployment) + builder = self._resolve_builder(app_spec.import_path) + serve.run(builder(deploy_options=options, **args), name=app_spec.name, route_prefix=app_spec.route_prefix) + + def _training_router_options(self, model) -> dict: + options = self._deployment_options(model.deploy_options, model.env_vars) + if 'autoscaling_config' not in options: + options.setdefault('num_replicas', 1) + options['ray_actor_options'].setdefault('num_cpus', 0.1) + return options + + def _deploy_training(self) -> None: + import ray + from ray import serve - for app_spec in applications: - self._deploy_application(app_spec) - - host = self.config.http_options.host - port = self.config.http_options.port - - logger.info('All applications started!') - for app_spec in applications: - logger.info(' Endpoint: http://%s:%s%s', host, port, app_spec.route_prefix) + from twinkle.server.model.pool.actor import TrainingControllerActor, controller_name + from twinkle.server.model.pool.router import build_training_router + config = self.config.training_config + controller = TrainingControllerActor.options( + name=controller_name(self.ray_namespace), + namespace=self.ray_namespace, + runtime_env={ + 'env_vars': build_worker_env_vars() + }).remote(config.model_dump(), self.ray_namespace) + self._training_controller = controller + try: + ray.get(controller.initialize.remote()) + for model in config.models: + app = build_training_router(config, model, self.ray_namespace, self.config.http_options.model_dump(), + self._training_router_options(model)) + serve.run( + app, + name=f'training-router-{self.ray_namespace}-{model.name}', + route_prefix=model.route_prefix or f'/api/v1/model/{model.name}') + except Exception: + self._run_cleanup([('training controller rollback', self._stop_training_controller)], raise_errors=False) + raise - # Graceful shutdown via signal handling - shutdown_event = threading.Event() + def _stop_controller(self, controller) -> None: + import ray - def _handle_signal(signum, frame): - sig_name = signal.Signals(signum).name - logger.info(f'Received {sig_name}, shutting down gracefully...') - shutdown_event.set() + # Keep an unconfirmed owner available for cleanup retry; a timeout alone + # does not certify its workers or devices were released. + ray.get(controller.shutdown.remote(), timeout=self.config.resources.shutdown_timeout_s) + ray.kill(controller, no_restart=True) + + def _stop_training_controller(self) -> None: + if self._training_controller is not None: + self._stop_controller(self._training_controller) + self._training_controller = None + + def _stop_sampler_controller(self, controller) -> None: + self._stop_controller(controller) + self._sampler_controllers.remove(controller) + + @staticmethod + def _run_cleanup(actions: list[tuple[str, Callable]], *, raise_errors: bool) -> None: + """Attempt every independent cleanup; rollback preserves the startup error.""" + failures = [] + for name, action in actions: + try: + action() + except Exception as exc: + failures.append((name, exc)) + logger.exception('%s failed', name) + if failures and raise_errors: + names = ', '.join(name for name, _ in failures) + raise RuntimeError(f'Server cleanup is incomplete: {names}') from failures[0][1] + + def _shutdown(self, *, raise_errors: bool = True) -> None: + from functools import partial + from ray import serve - signal.signal(signal.SIGINT, _handle_signal) - signal.signal(signal.SIGTERM, _handle_signal) + from twinkle.server.telemetry import flush_telemetry_safely + actions = [('training controller shutdown', self._stop_training_controller)] + actions.extend(('sampler controller shutdown', partial(self._stop_sampler_controller, controller)) + for controller in tuple(self._sampler_controllers)) + if self._serve_started: + actions.append(('Ray Serve shutdown', serve.shutdown)) + actions.append(('telemetry flush', flush_telemetry_safely)) + self._run_cleanup(actions, raise_errors=raise_errors) + self._serve_started = False - # Block until a termination signal is received - shutdown_event.wait() + @staticmethod + def _wait_for_shutdown() -> None: + stopped = threading.Event() - self._shutdown() + def signal_shutdown(signum, frame): + logger.info('Received %s', signal.Signals(signum).name) + stopped.set() - def _shutdown(self) -> None: - """Tear down Serve and flush telemetry on graceful termination. + signal.signal(signal.SIGINT, signal_shutdown) + signal.signal(signal.SIGTERM, signal_shutdown) + stopped.wait() - Telemetry flush runs AFTER ``serve.shutdown()`` and is wrapped so a - telemetry-shutdown failure cannot mask the user-facing shutdown path. - """ - from ray import serve + def launch(self) -> None: + if not self.config.applications: + logger.warning('No applications configured') + return + apply_ray_serve_patches() + self._configure_environment() + self._init_ray() try: - serve.shutdown() - logger.info('Ray Serve shut down successfully') + self._start_serve() + for app in self.config.applications: + self._deploy_application(app) + if self.config.training_config.models: + self._deploy_training() + logger.info('All applications started') + self._wait_for_shutdown() except Exception: - logger.warning('Error during Ray Serve shutdown', exc_info=True) - - # Flush buffered OTLP batches (traces / metrics / logs) on the driver - # after Serve is down. Errors are swallowed by the helper. - from twinkle.server.telemetry import flush_telemetry_safely - flush_telemetry_safely() + self._shutdown(raise_errors=False) + raise + self._shutdown() def launch_server( diff --git a/src/twinkle/server/lifecycle/envelope.py b/src/twinkle/server/lifecycle/envelope.py index 25343c9b0..dc193aadf 100644 --- a/src/twinkle/server/lifecycle/envelope.py +++ b/src/twinkle/server/lifecycle/envelope.py @@ -13,41 +13,18 @@ from twinkle.protocol.types.errors import ErrorCategory, ErrorPayload from twinkle.protocol.types.lifecycle import TaskEnvelope from twinkle.server.state.models import FutureFailureRecord -from twinkle.server.task_errors import trim_traceback - -# Keys must equal ``state.models.FAILURE_REASON_CODES`` (guarded by test_envelope). -_FAILURE_WIRE: dict[str, tuple[int, ErrorCategory]] = { - 'invalid_request': (400, ErrorCategory.User), - 'request_rejected': (400, ErrorCategory.User), - 'resource_not_found': (404, ErrorCategory.User), - 'full_mode_busy': (409, ErrorCategory.User), - 'input_tokens_exceeded': (422, ErrorCategory.User), - 'batch_size_invalid': (422, ErrorCategory.User), - 'rate_limit_exceeded': (429, ErrorCategory.User), - 'resource_quota_exceeded': (429, ErrorCategory.User), - 'cancelled': (499, ErrorCategory.User), - 'endpoint_unavailable': (501, ErrorCategory.Server), - 'state_contention': (503, ErrorCategory.Server), - 'backend_gate_unavailable': (503, ErrorCategory.Server), - 'orphaned_replica': (503, ErrorCategory.Server), - 'execution_timeout': (504, ErrorCategory.Server), - 'deadline_exceeded': (500, ErrorCategory.Server), - 'internal_error': (500, ErrorCategory.Server), -} +from twinkle.server.task_errors import FailureResponse, trim_traceback def error_payload_from_failure(stored: Any, *, request_id: str) -> ErrorPayload: """Map one protocol-independent failure record to Twinkle's wire model.""" failure = FutureFailureRecord.model_validate(stored) - error_code, category = _FAILURE_WIRE.get( - failure.reason_code, - (500, ErrorCategory.Server), - ) - diagnostic = failure.diagnostic if category is ErrorCategory.Server else None + response = FailureResponse.for_reason(failure.reason_code) + diagnostic = failure.diagnostic if response.category is ErrorCategory.Server else None return ErrorPayload( error=failure.message[:1024], - category=category, - error_code=error_code, + category=response.category, + error_code=response.error_code, request_id=request_id, traceback=trim_traceback(diagnostic) if diagnostic else None, details=failure.details, diff --git a/src/twinkle/server/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/lifecycle/proxy.py b/src/twinkle/server/lifecycle/proxy.py new file mode 100644 index 000000000..cee04c878 --- /dev/null +++ b/src/twinkle/server/lifecycle/proxy.py @@ -0,0 +1,124 @@ +# 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, TypedDict + +from twinkle.protocol.headers import H_MULTIPLEX, H_MULTIPLEX_LEGACY, H_REQUEST_ID, H_REQUEST_ID_LEGACY +from twinkle.protocol.types.errors import ErrorCategory, ErrorPayload +from twinkle.server.config.task_queue import TaskQueueConfig +from twinkle.server.telemetry.tracing import inject_context +from twinkle.utils.logger import get_logger + +logger = get_logger() + + +class ForwardedRequest(TypedDict): + """HTTP request fields carried intact across an internal Ray RPC.""" + + endpoint: str + body: bytes + method: str + headers: list[tuple[bytes, bytes]] + query_string: bytes + + +class ForwardedReply(TypedDict): + status: int + headers: list[tuple[bytes, bytes]] + body: bytes + + +class ServiceProxy: + + def __init__(self, http_options: dict[str, Any] | None = None, route_prefix: str = '/api/v1'): + self.http_options = http_options or {} + self.route_prefix = route_prefix + 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() + + 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/lifecycle/submit.py b/src/twinkle/server/lifecycle/submit.py index 947dfb88e..5c432281b 100644 --- a/src/twinkle/server/lifecycle/submit.py +++ b/src/twinkle/server/lifecycle/submit.py @@ -16,6 +16,8 @@ from twinkle.protocol.types.base import FieldRole, fields_with_role from twinkle.protocol.types.data import export_batch from twinkle.protocol.types.lifecycle import TaskEnvelope +from twinkle.protocol.types.model import (AdapterRequest, DataPlaneForwardRequest, ForwardBackwardTaskRequest, + TrainingRequest) from twinkle.server.middleware.auth import get_session_id_from_request from twinkle.server.validation import assert_request_supported @@ -85,11 +87,11 @@ def backend_kwargs(body: Any) -> dict[str, Any]: model_cls = type(body) kwargs: dict[str, Any] = {} for name in fields_with_role(model_cls, FieldRole.BackendKwarg): - value = getattr(body, name, None) + value = getattr(body, name) if value is not None: kwargs[name] = value for name in fields_with_role(model_cls, FieldRole.Passthrough): - region = getattr(body, name, None) or {} + region = getattr(body, name) or {} overlap = set(region) & set(kwargs) if overlap: raise ValueError(f'{name} collides with declared backend parameters: {", ".join(sorted(overlap))}') @@ -123,7 +125,7 @@ def resolve_twinkle_adapter_name(request: Request, adapter_name: str | None) -> async def run_submit( self: QueuedDeployment, request: Request, - body: Any, + body: TrainingRequest, *, task_type: str, backend_call: Callable[..., Coroutine], @@ -152,13 +154,10 @@ async def run_submit( signature-patching wrapper once deepened the route graph past CPython's C-stack recursion guard during ``serve.ingress``). - Four queued endpoints deliberately do NOT route through here and call - ``submit_and_peek`` / ``submit_background_and_peek`` directly: model - ``add_adapter_to_model`` (creates the adapter, so the resource assertion cannot - apply and it owns the train_mode/full-mode checks), model ``upload_to_hub`` (pure - I/O, background task), and sampler ``sample`` / ``sample_to_data_plane`` (no adapter - semantics). Zero-write and admission still hold for them because both live inside - ``schedule_task`` -> ``_perform_preflight_checks``, not in this shell. + Upload uses the background I/O path; sampling uses concurrent admission. + Their scheduling entrypoints apply the same preflight and Future ownership rules. + Training creation/release belongs to the public Router and Controller, rather + than this compute-operation shell. ``backend_call(self, body, adapter_name, token)`` runs the endpoint-specific call and returns the JSON-safe task result. ``metrics(self, body)`` supplies scheduling @@ -168,7 +167,7 @@ async def run_submit( not implemented by every backend. """ token = await self._on_request_start(request) - adapter_name = resolve_twinkle_adapter_name(request, body.adapter_name) + adapter_name = body.training_id # ---- Preflight: decidable from the body plus this deployment's backend, so it # runs before any state write and before the enqueue. ---- @@ -192,33 +191,40 @@ async def _task(): # forward_backward would be dropped as a duplicate AND handed the first adapter's # loss. That is a silent wrong-result bug, i.e. the exact failure this dedup # exists to prevent, one level up. ---- - request_id = f'req_{uuid.uuid4().hex}' - seq_id = getattr(body, 'seq_id', None) - dedup_key = None + sequenced = isinstance(body, (AdapterRequest, ForwardBackwardTaskRequest, DataPlaneForwardRequest)) + return await submit_training_task( + self, + _task, + training_id=body.training_id, + stream=f'native::{get_session_id_from_request(request)}::{task_type}', + seq_id=body.seq_id if sequenced else None, + peek=True, + token=token, + task_type=task_type, + **schedule_kwargs) + + +async def submit_training_task(self, + factory, + *, + training_id: str, + stream: str, + seq_id: int | None, + peek: bool = False, + **kwargs): + """Submit with explicit request identity; retain claims once any future was written.""" + request_id = kwargs.pop('request_id', None) or f'req_{uuid.uuid4().hex}' if seq_id is not None: - session_id = get_session_id_from_request(request) or request.state.request_id - dedup_key = f'seq::{session_id}::{adapter_name or "-"}::{seq_id}' - # The claim must outlive the whole legal lifetime of the original task: - # it may wait in the queue and then execute. A TTL that covers only the - # execution bound lets a retry re-enqueue while the original is still - # queued, i.e. applies one gradient update twice. - ttl = int(self.task_queue_config.queue_timeout + self.task_queue_config.effective_execution_timeout) + 60 - prior_request_id = await self.state.claim_seq(dedup_key, request_id, ttl) - if prior_request_id is not None: - return await self._peek_terminal(prior_request_id, fallback_status='pending') - - # ---- Decision_Boundary: preflight (in schedule_task) then peek ---- + prior = await self.state.claim_training_seq(training_id, stream, seq_id, request_id, self.sequence_retention_s) + if prior is not None: + if peek: + return await self._peek_terminal(prior, fallback_status='pending') + return {'request_id': prior, 'model_id': training_id} + submit = self.submit_and_peek if peek else self.schedule_task try: - return await self.submit_and_peek( - _task, model_id=adapter_name, token=token, task_type=task_type, request_id=request_id, **schedule_kwargs) + return await submit(factory, model_id=training_id, request_id=request_id, **kwargs) except Exception: - # Release the seq claim only when the task never made it onto the queue -- - # decided by whether a future record exists, NOT by the exception type. A - # preflight rejection raises before any record is written, so releasing lets - # a retry re-enqueue. But a failure *after* enqueue (e.g. a transient state - # error inside the peek) leaves a live task that will still run; releasing - # there would let a retry enqueue a duplicate -> the exact double-apply this - # dedup prevents. When unsure (record exists), keep the claim. - if dedup_key is not None and await self.state.get_future(request_id) is None: - await self.state.release_seq(dedup_key) + # A failed response after enqueue must not allow the same update to run twice. + if seq_id is not None and await self.state.get_future(request_id) is None: + await self.state.release_training_seq(training_id, stream, seq_id, request_id) raise diff --git a/src/twinkle/server/model/app.py b/src/twinkle/server/model/app.py index b98fbb2fb..5e4a389da 100644 --- a/src/twinkle/server/model/app.py +++ b/src/twinkle/server/model/app.py @@ -8,23 +8,20 @@ from __future__ import annotations import asyncio -from fastapi import FastAPI, HTTPException, Request -from ray import serve -from ray.serve.config import RequestRouterConfig +from fastapi import FastAPI, Request 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.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.state.training_records import TrainingRuntimeContext +from twinkle.server.task_queue import TaskQueueMixin from twinkle.server.utils import wrap_builder_with_device_group_env from twinkle.utils.logger import get_logger +from .runtime import TrainingRuntime from .tinker_handlers import _register_model_tinker_routes from .twinkle_handlers import _register_model_twinkle_routes @@ -82,14 +79,14 @@ def _make_megatron_model(kw: dict[str, Any]) -> Any: ) -class ModelManagement(LazyCleanupMixin, TaskQueueMixin, AdapterManagerMixin): +class ModelManagement(TrainingRuntime, TaskQueueMixin): """Unified model management service. Handles: - Base model and multiple LoRA adapters (multi-user) - Tinker training operations via /tinker/* endpoints (async/polling) - Twinkle training operations via /twinkle/* endpoints (synchronous) - - Adapter lifecycle via AdapterManagerMixin + - Worker-local gradient state is keyed by the controller-issued training handle. - Per-user rate limiting via TaskQueueMixin """ @@ -99,19 +96,21 @@ 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): self.backend = backend self.device_group = DeviceGroup(**device_group) + self.training_context = TrainingRuntimeContext.model_validate(kwargs.pop('training_context')) + await self.initialize_training_runtime(self.training_context, kwargs.pop('placement_group'), + kwargs.pop('allocation_devices')) self.device_mesh = init_twinkle_runtime( is_mock=(backend == 'mock'), nproc_per_node=nproc_per_node, device_group=self.device_group, device_mesh_dict=device_mesh, ) - self.replica_id = 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 +129,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 @@ -147,9 +147,8 @@ 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._training_states = {} await self._register_replica_on_startup() - # Note: countdown task is started lazily in _ensure_sticky() @property def data_world_size(self) -> int: @@ -164,34 +163,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.training_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 +179,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: @@ -215,28 +194,11 @@ 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: - """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,36 +215,6 @@ 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) - - def claim_full_mode_resource(self, adapter_name: str, token: str, session_id: str) -> None: - """Atomically take the exclusive full-mode resource for one tenant. - - The check and the registration are synchronous neighbours: no coroutine - can interleave between them, so two concurrent creates cannot both pass - the availability check before either is recorded. - """ - self.assert_full_mode_available(adapter_name) - self.register_resource(adapter_name, token, session_id) - - 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) - def build_model_app(model_id: str, nproc_per_node: int, @@ -290,7 +222,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): @@ -307,7 +238,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 @@ -320,8 +250,7 @@ def build_model_app(model_id: str, # Build the FastAPI app + middleware stack + routes via the shared scaffold, # then bind the Ray Serve deployment. The Model passes its ``shutdown()`` - # teardown via ``on_shutdown`` and its sticky-LoRA router via - # ``request_router_config``. + # teardown via ``on_shutdown``; the training controller owns handle routing. def register_routes(app: FastAPI, get_self: Any) -> None: _register_model_tinker_routes(app, get_self) _register_model_twinkle_routes(app, get_self) @@ -336,9 +265,7 @@ 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_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/model/backends/mock_model.py b/src/twinkle/server/model/backends/mock_model.py index d204090a8..7654c9eeb 100644 --- a/src/twinkle/server/model/backends/mock_model.py +++ b/src/twinkle/server/model/backends/mock_model.py @@ -181,7 +181,7 @@ def save(self, name: str = 'latest', output_dir: str | None = None, **kwargs: An The real backend creates ``//`` on disk and returns its path. Tinker /asample later does ``os.path.exists(...)`` on a symlink - pointing here (see ``sampler/tinker_handlers.py:72``), so the directory + pointing here (see ``sampler/tinker_handlers.py``), so the directory has to really exist for the mock path to clear. """ target = os.path.join(output_dir or '/tmp/twinkle_mock_ckpt', name) diff --git a/src/twinkle/server/model/pool/__init__.py b/src/twinkle/server/model/pool/__init__.py new file mode 100644 index 000000000..720c35471 --- /dev/null +++ b/src/twinkle/server/model/pool/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Training instance pool: routing, lifecycle decisions and Ray composition.""" diff --git a/src/twinkle/server/model/pool/actor.py b/src/twinkle/server/model/pool/actor.py new file mode 100644 index 000000000..d4f8f0afa --- /dev/null +++ b/src/twinkle/server/model/pool/actor.py @@ -0,0 +1,174 @@ +# 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.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 + +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): + 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._initialization = asyncio.get_event_loop().create_task(self._initialize()) + + async def initialize(self): + await asyncio.shield(self._initialization) + + async def _initialize(self): + 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: + active = await self.store.backend.get(self.store.key('controller', 'active')) + if active is None: + await self._recover_lost_state() + return + await self.store.models.touch_replica_last_seen(self.core.epoch) + except Exception: + logger.exception('Training state heartbeat failed; admission remains fenced by the state backend') + await asyncio.sleep(5) + continue + 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): + + async def probe(): + identity = await self.driver.health(instance) + if not identity['unhealthy']: + await self.store.models.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): + 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: + 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): + # 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.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): + 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 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): + 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() + await asyncio.gather(restart, return_exceptions=True) + for task in self._monitors: + task.cancel() + await asyncio.gather(*self._monitors, return_exceptions=True) + await self.core.shutdown() + await self.driver.runtime.wait_for_cleanup() diff --git a/src/twinkle/server/model/pool/controller.py b/src/twinkle/server/model/pool/controller.py new file mode 100644 index 000000000..6ef2908a5 --- /dev/null +++ b/src/twinkle/server/model/pool/controller.py @@ -0,0 +1,806 @@ +# 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, TrainingModelConfig +from twinkle.server.exceptions import (RequestRejectedError, ResourceExhaustedError, ResourceNotFoundError, + ResourceQuotaExceededError, TrainingExpiredError) +from twinkle.server.resources.allocation import CapacityUnavailableError +from twinkle.server.resources.instances import InstanceLifecycle +from twinkle.server.state.models import FutureFailureRecord +from twinkle.server.state.training_records import (LIVE_INSTANCES, TERMINAL_TRAINING, CreateTraining, InstanceRecord, + RuntimeOwner, TrainingRecord, TrainingRoute) +from twinkle.server.state.training_store import TrainingStore +from twinkle.server.telemetry.metrics import TrainingInstanceMetrics, TrainingInstanceSnapshot, get_training_metrics + +logger = logging.getLogger(__name__) + + +class InstanceDriver(Protocol): + + async def allocate(self, instance: InstanceRecord, spec: TrainingModelConfig) -> dict[str, Any]: + ... + + async def deploy(self, instance: InstanceRecord, spec: TrainingModelConfig) -> None: + ... + + 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.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.lifecycle = InstanceLifecycle('Training') + self._cleaning: set[str] = set() + self._reconciling: set[str] = set() + self._failed_instances: dict[str, FutureFailureRecord] = {} + self._lock = asyncio.Lock() + self._ready = False + self._metrics = get_training_metrics() + self._instance_metrics = TrainingInstanceMetrics(self._metrics) + + def _instance_metric_tags(self, instance: InstanceRecord) -> dict[str, str]: + spec = self.config.get(instance.base_model, instance.mode) + return { + 'backend': spec.backend, + 'mode': instance.mode, + 'state': instance.state, + 'resident': str(instance.resident).lower(), + } + + def _report_instance_metrics(self) -> None: + snapshots = [] + for instance in self.instances.values(): + if instance.state == 'STOPPED': + continue + model = self.config.get(instance.base_model, instance.mode) + snapshots.append( + TrainingInstanceSnapshot( + **self._instance_metric_tags(instance), + device_type=model.device_type, + device_count=len(instance.devices), + placement_group_count=int(bool(instance.application_name)))) + self._instance_metrics.report(snapshots) + + async def initialize(self) -> None: + """Discard old state only after owned compute has been confirmed dead.""" + 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.models.rebuild_indexes() + if old_controller: + await self.store.models.unregister_replica(old_controller['epoch']) + await self.store.models.register_replica(self.epoch, 0) + await self.store.models.touch_replica_last_seen(self.epoch) + try: + for spec in self.config.models: + for _ in range(spec.config.min_replicas): + async with self._lock: + instance = await self._new_instance(spec, resident=True) + await asyncio.shield(self.lifecycle.startups[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.config.models) + + def residents_ready_for(self, spec: TrainingModelConfig) -> bool: + return sum( + instance.resident and instance.instance_id not in self._failed_instances and ( + instance.base_model, instance.mode) == spec.key and instance.state in ('READY', 'RESETTING') + 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) + self._report_instance_metrics() + + 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: TrainingModelConfig, *, resident: bool = False) -> InstanceRecord: + matching = [item for item in self.instances.values() if (item.base_model, item.mode) == spec.key] + if not self.lifecycle.headroom(matching, spec.config.max_replicas): + raise CapacityUnavailableError('Training replica limit reached') + instance = InstanceRecord( + instance_id='instance_' + uuid.uuid4().hex, + controller_epoch=self.epoch, + base_model=spec.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._report_instance_metrics() + self.lifecycle.start(instance.instance_id, lambda: self._start_instance(instance)) + return instance + + async def _start_instance(self, instance: InstanceRecord) -> None: + spec = self.config.get(instance.base_model, instance.mode) + started, status = time.monotonic(), 'error' + tags = {'backend': spec.backend, 'mode': spec.mode} + try: + + 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) + 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) + + 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') + 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: TrainingModelConfig) -> InstanceRecord | None: + candidates = [ + instance for instance in self.instances.values() + if instance.instance_id not in self._failed_instances and (instance.base_model, instance.mode) == spec.key + 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.config.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 ResourceExhaustedError('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, + owner_token=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.') + 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 + 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.lifecycle.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.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) + 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.lifecycle.startups[instance.instance_id]) + async with self._lock: + future = await self.store.futures.get(training.create_request_id) + if future and future.status == 'cancelled': + training.state = 'CLOSING' + training.terminal_reason = 'cancelled' + if training.state != 'PENDING' or instance.state != 'READY': + raise RequestRejectedError('Training creation is no longer pending') + training.creation_deadline = self.clock() + self.config.startup_timeout_s + await self.store.save_training(training) + if not await self.store.futures.try_start(training.create_request_id): + future = await self.store.futures.get(training.create_request_id) + if future and future.status == 'cancelled': + training.terminal_reason = 'cancelled' + raise RequestRejectedError('Training creation did not acquire execution') + remaining = training.creation_deadline - self.clock() + 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: + heartbeat = await self.store.sessions.get_last_heartbeat(training.session_id) + return self._expired_with_heartbeat(training, heartbeat) + + def _session_expired(self, training: TrainingRecord, heartbeat: float | None) -> bool: + timeout = self.config.get(training.base_model, training.mode).config.adapter_config.adapter_timeout + return heartbeat is None or (timeout is not None and time.time() - heartbeat >= timeout) + + def _lifetime_expired(self, training: TrainingRecord) -> bool: + lifetime = self.config.get(training.base_model, training.mode).config.adapter_config.adapter_max_lifetime + if lifetime is None: + return False + created_at = datetime.fromisoformat(training.created_at).timestamp() + return time.time() - created_at >= lifetime + + def _expired_with_heartbeat(self, training: TrainingRecord, heartbeat: float | None) -> bool: + spec = self.config.get(training.base_model, training.mode) + if self._session_expired(training, heartbeat) or self._lifetime_expired(training): + return True + return self.clock() - training.last_activity_at >= spec.config.idle_timeout_s + + async def _authorized(self, training_id: str, token: str, session_id: str | None) -> TrainingRecord: + await self._assert_epoch() + 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' or training.instance_id in self._failed_instances: + raise ResourceNotFoundError('Training handle is no longer active') + heartbeat = await self.store.sessions.get_last_heartbeat(training.session_id) + if self._session_expired(training, heartbeat): + raise TrainingExpiredError('Training session expired; create a new training handle') + return training + + async def admit(self, + training_id: str, + token: str, + session_id: str | None, + request_id: str, + *, + renew: bool = True) -> RuntimeOwner: + async with self._lock: + training = await self._authorized(training_id, token, session_id) + expired = self._lifetime_expired(training) or (not training.pending_operation_ids + and await self._expired(training)) + if request_id not in training.pending_operation_ids and expired: + training.state = 'CLOSING' + training.terminal_reason = 'expired' + await self.store.save_training(training) + 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(), 'renew': renew} + await self.store.save_training(training) + return {'instance_id': instance.instance_id, 'epoch': self.epoch} + + 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.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 + 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.lifecycle.spawn(self._cleanup_user(training)) + + async def route(self, training_id: str, token: str, session_id: str | None = None) -> TrainingRoute: + async with self._lock: + training = await self._authorized(training_id, token, session_id) + return {'instance_id': training.instance_id, 'base_model': training.base_model, 'epoch': self.epoch} + + 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.config.get(training.base_model, training.mode).config.queue_config + deadline = asyncio.get_running_loop().time() + config.absolute_future_ttl(1) + self.config.shutdown_timeout_s + while training.state == 'CLOSING' or self.instances[training.instance_id].state in ('DRAINING', 'QUARANTINED'): + if self.instances[training.instance_id].state == 'QUARANTINED': + 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] + claimed = False + resetting = False + closing_full = False + try: + # 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 resetting or closing_full: + self._report_instance_metrics() + if closing_full: + await self._close_instance(instance) + return + try: + spec = self.config.get(training.base_model, training.mode) + await asyncio.wait_for( + self.driver.release_training(instance, training), + spec.config.queue_config.absolute_future_ttl(spec.config.num_devices)) + 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: + # 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={ + '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 + self._report_instance_metrics() + if close: + await self._close_instance(instance) + finally: + if claimed: + self._cleaning.discard(training.training_id) + + async def _close_instance(self, instance: InstanceRecord) -> None: + started = time.monotonic() + tags = self._instance_metric_tags(instance) + timing_tags = {'backend': tags['backend'], 'mode': tags['mode']} + + async def quarantined(exc): + async with self._lock: + entered_quarantine = instance.state != 'QUARANTINED' + quarantined = instance.model_copy(update={'state': 'QUARANTINED', 'failure': str(exc)}) + await self.store.save_instance(quarantined) + instance.state = 'QUARANTINED' + instance.failure = str(exc) + self._report_instance_metrics() + tags = self._instance_metric_tags(instance) + if entered_quarantine: + self._metrics.quarantined_total.inc( + tags={ + 'backend': tags['backend'], + 'mode': tags['mode'], + 'reason': TrainingInstanceMetrics.quarantine_reason(exc), + }) + self._metrics.shutdown_seconds.observe(time.monotonic() - started, tags={**timing_tags, 'status': 'error'}) + + async def stopped(): + 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() + 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 = 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) + self._report_instance_metrics() + + async def fail_instance(self, instance_id: str, failure: FutureFailureRecord) -> None: + async with self._lock: + instance = self.instances[instance_id] + if instance.state == 'STOPPED': + self._failed_instances.pop(instance_id, None) + return + # 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 + self._report_instance_metrics() + 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) + 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. + 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.sessions.get_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 + or training.session_id not in heartbeats): + continue + 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' + 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.lifecycle.spawn(self._cleanup_user(training)) + for training, request_id in uncertain: + 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.lifecycle.spawn(self.fail_instance(instance.instance_id, failure)) + elif instance.state in ('DRAINING', 'QUARANTINED'): + self.lifecycle.spawn(self._close_instance(instance)) + # Rebuilding one specification must not stop expiry or other models. + for spec in self.config.models: + async with self._lock: + matching = [item for item in self.instances.values() if (item.base_model, item.mode) == spec.key] + missing = self.lifecycle.missing( + 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. + async with self._lock: + cutoff = time.time() - self.store.expiration_timeout + for identity, training in list(self.trainings.items()): + if training.state in TERMINAL_TRAINING and training.terminal_at is not None: + if training.terminal_at <= cutoff: + 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.lifecycle.startups.pop(identity, None) + self._report_instance_metrics() + + 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'] == '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: + await self.lifecycle.join() + + async def shutdown(self) -> None: + self._ready = False + outcomes = await asyncio.gather( + *(self.fail_instance( + instance.instance_id, + FutureFailureRecord( + reason_code='orphaned_replica', message='Training service is shutting down', attribution='server')) + for instance in list(self.instances.values())), + return_exceptions=True) + await self.wait_for_jobs() + errors = [outcome for outcome in outcomes if isinstance(outcome, BaseException)] + for error in errors: + logger.error('Training instance shutdown failed', exc_info=(type(error), error, error.__traceback__)) + if errors: + raise errors[0] + self.lifecycle.require_released(self.instances.values()) diff --git a/src/twinkle/server/model/pool/driver.py b/src/twinkle/server/model/pool/driver.py new file mode 100644 index 000000000..cb9071f70 --- /dev/null +++ b/src/twinkle/server/model/pool/driver.py @@ -0,0 +1,140 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Ray/Serve effects behind the controller's small runtime interface.""" +from __future__ import annotations + +import asyncio +import re +import time + +from twinkle.server.resources.runtime import RayRuntime +from twinkle.server.state.training_records import InstanceRecord, TrainingRuntimeContext +from twinkle.server.telemetry.metrics import get_training_metrics +from twinkle.utils.logger import get_logger + +logger = get_logger() + + +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._metrics = get_training_metrics() + self.runtime = RayRuntime( + namespace, allocation_timeout=config.allocation_timeout_s, shutdown_timeout=config.shutdown_timeout_s) + + def application_name(self, instance): + 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.backend, 'mode': spec.mode, 'device_type': spec.device_type} + try: + allocation = await self.runtime.pool.allocate(spec.config.num_devices, spec.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.util.placement_group import get_placement_group + + from twinkle.server.model.app import build_model_app + + name = self.application_name(instance) + context = TrainingRuntimeContext( + instance_id=instance.instance_id, + epoch=instance.controller_epoch, + namespace=self.namespace, + controller_name=self.controller_name, + registry_name=name + '-registry') + kwargs = spec.config.model_dump( + exclude={ + 'num_devices', 'device_mesh', 'min_replicas', 'max_replicas', 'idle_timeout_s', 'adapter_config', + 'queue_config' + }, + exclude_none=True) + if 'hybrid' in kwargs: + kwargs['hybrid'] = spec.config.hybrid + kwargs['train_mode'] = 'hybrid' if kwargs.get('hybrid') else spec.mode + options = self.runtime.deployment_options(spec.deploy_options, name, spec.env_vars) + app = build_model_app( + model_id=spec.model_id, + backend=spec.backend, + nproc_per_node=spec.config.num_devices, + device_group={ + 'name': 'model', + 'ranks': spec.config.num_devices, + 'device_type': spec.device_type + }, + device_mesh={ + **spec.config.device_mesh.model_dump(exclude_none=True), 'device_type': spec.device_type + }, + deploy_options=options, + queue_config=spec.config.queue_config, + data_plane_url=spec.data_plane_url, + training_context=context, + 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}', + lambda: self.shutdown(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): + handle = await self.handle(instance) + await handle.initialize_training.remote(training.model_dump(), request.model_dump()) + + async def release_training(self, instance, training): + 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): + 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) + try: + spec = self.config.get(instance.base_model, instance.mode) + except Exception: + tags = {'backend': 'unknown', 'mode': instance.mode, 'device_type': 'unknown'} + else: + tags = {'backend': spec.backend, 'mode': spec.mode, 'device_type': spec.device_type} + + def observe_removal(duration, status): + self._metrics.removal_seconds.observe(duration, tags={**tags, 'status': status}) + + await self.runtime.shutdown(instance.instance_id, name, observe_removal=observe_removal) + await self.store.models.unregister_replica_capacity(instance.instance_id) + + async def health(self, instance): + 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.""" + pattern = re.compile(r'resources-' + re.escape(self.namespace) + + r'-(controller_[0-9a-f]{32})-(instance_[0-9a-f]{32})$') + + def decode(name): + match = pattern.fullmatch(name) + if match: + return InstanceRecord( + controller_epoch=match[1], + instance_id=match[2], + application_name=name, + base_model='', + mode='full', + startup_deadline=0) + + await self.runtime.recover_orphans(decode, self.shutdown) diff --git a/src/twinkle/server/model/pool/router.py b/src/twinkle/server/model/pool/router.py new file mode 100644 index 000000000..23f4d8bcb --- /dev/null +++ b/src/twinkle/server/model/pool/router.py @@ -0,0 +1,256 @@ +# 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 ProtocolMismatchError, RequestRejectedError, ResourceNotFoundError +from twinkle.server.lifecycle.envelope import envelope_from_record +from twinkle.server.lifecycle.proxy import ServiceProxy +from twinkle.server.middleware.auth import get_session_id_from_request, get_token_from_request +from twinkle.server.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 + +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.config = config + self.model = model + self.backend = model.backend + self.state = get_server_state() + 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 + 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.config.get(self.model.name, mode) + if config is not None: + rank = config.get('r') if isinstance(config, dict) else None + if (not isinstance(config, dict) or config.get('_TWINKLE_TYPE_') != 'LoraConfig' + 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, token=token) + return envelope_from_record(ref['request_id'], await self.state.get_future(ref['request_id'])) + 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.config.get(self.model.name, mode) + if body.lora_config is not None and not 0 < body.lora_config.rank <= spec.config.max_r: + raise RequestRejectedError('LoRA rank exceeds the training specification') + await self._check_session(body.session_id) + 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', 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'])) + + +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/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/model/runtime.py b/src/twinkle/server/model/runtime.py new file mode 100644 index 000000000..ff535734c --- /dev/null +++ b/src/twinkle/server/model/runtime.py @@ -0,0 +1,207 @@ +# 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 dataclasses import dataclass + +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.state.training_records import CreateTraining, TrainingRecord, TrainingRuntimeContext +from twinkle.server.utils import get_template_for_model +from twinkle.utils.logger import get_logger + +logger = get_logger() + + +@dataclass +class TrainingExecutionState: + """Worker-local state; ownership and expiration belong to the controller.""" + + grad_ready: bool = False + + +class TrainingRuntime: + """Execute controller-owned training handles on one allocated backend.""" + + training_context: TrainingRuntimeContext + + _training_states: dict[str, TrainingExecutionState] + + def get_training_state(self, training_id: str) -> TrainingExecutionState: + state = self._training_states.get(training_id) + if state is None: + raise ResourceNotFoundError('Training is not registered on this runtime') + return state + + def assert_resource_exists(self, training_id: str) -> None: + self.get_training_state(training_id) + + async def wait_for_training_workers(self): + actors = getattr(self.model, '_actors', None) or [] + await asyncio.gather(*(actor.__ray_ready__.remote() for actor in actors)) + + async def initialize_training_runtime(self, context, placement_group, devices): + import ray + + from twinkle.infra._ray.ray_helper import RayHelper + from twinkle.infra._ray.resource_manager import ResourceManager + self._training_controller = ray.get_actor(context.controller_name, namespace=context.namespace) + await self._training_controller.register_runtime.remote(context.instance_id, context.epoch, uuid.uuid4().hex) + os.environ['TWINKLE_RAY_REGISTRY_NAME'] = context.registry_name + os.environ['TWINKLE_RAY_NAMESPACE'] = context.namespace + RayHelper.resource_manager = ResourceManager.from_allocation(self.device_group, placement_group, devices) + + async def training_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') + return token + + 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 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: + 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, + finished_at=outcome['finished_at']) + return True + except Exception: + logger.exception('Training completion notification failed for %s', task.request_id) + + task.on_finished = finished + if task.task_type == 'apply_patch': + await self._training_controller.mark_dirty.remote(context.instance_id, context.epoch) + + async def _run_lifecycle_operation(self, factory, task_type): + ref = await self.schedule_task(factory, task_type=task_type) + deadline = asyncio.get_running_loop().time() + self.task_queue_config.absolute_future_ttl(1) + while True: + 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._training_states[adapter] = TrainingExecutionState() + if request.client_type == 'twinkle': + body = native.AddAdapterRequest.model_validate(request.payload) + # Dynamic imports (e.g. PEFT) and deserialization must not block the health-probe event loop. + 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( + 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 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=config) + 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) + + 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._run_lifecycle_operation(initialize, 'initialize_training') + + async def release_training(self, training_data, reset_full): + training = TrainingRecord.model_validate(training_data) + + async def release(): + adapter = training.training_id + if adapter not in self._training_states: + return + if self.is_full_mode: + if reset_full: + await self.call_backend(self.model.reload_initial_weights, reset_training_state=True) + else: + await self.call_backend(self.model.remove_adapter, adapter) + self._training_states.pop(adapter) + + await self._run_lifecycle_operation(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'} + return self._compute_worker.get_task_outcome(request_id) or {'state': 'missing'} diff --git a/src/twinkle/server/model/tinker_handlers.py b/src/twinkle/server/model/tinker_handlers.py index 60c864a43..b1fd47f69 100644 --- a/src/twinkle/server/model/tinker_handlers.py +++ b/src/twinkle/server/model/tinker_handlers.py @@ -9,7 +9,6 @@ """ from __future__ import annotations -import traceback from collections.abc import Callable from fastapi import Depends, FastAPI, Request from functools import partial @@ -20,9 +19,8 @@ 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.lifecycle.submit import submit_training_task 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, @@ -138,29 +76,28 @@ async def forward(request: Request, body: types.ForwardRequest, token = await self._on_request_start(request) async def _do_forward(): - try: - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - self.assert_resource_exists(adapter_name) - model_adapter = self.resolve_model_adapter_name(adapter_name) - datum_list = body.forward_input.data - loss_fn_config = body.forward_input.loss_fn_config or {} - output, loss = await self.call_backend( - self.model.tinker_forward_only, inputs=datum_list, adapter_name=model_adapter, **loss_fn_config) - return types.ForwardBackwardOutput( - loss_fn_output_type='CrossEntropyLossReturn', - loss_fn_outputs=output, - metrics={'loss:mean': loss}, - ) - except Exception: - logger.error(traceback.format_exc()) - raise + adapter_name = body.model_id + self.assert_resource_exists(adapter_name) + model_adapter = self.resolve_model_adapter_name(adapter_name) + datum_list = body.forward_input.data + loss_fn_config = body.forward_input.loss_fn_config or {} + output, loss = await self.call_backend( + self.model.tinker_forward_only, inputs=datum_list, adapter_name=model_adapter, **loss_fn_config) + return types.ForwardBackwardOutput( + loss_fn_output_type='CrossEntropyLossReturn', + loss_fn_outputs=output, + metrics={'loss:mean': loss}, + ) datum_list = body.forward_input.data input_tokens = sum(len(d.model_input.to_ints()) for d in datum_list) batch_size = len(datum_list) - return await self.schedule_task( + return await submit_training_task( + self, _do_forward, - model_id=body.model_id, + training_id=body.model_id, + seq_id=body.seq_id, + stream='tinker::forward', token=token, input_tokens=input_tokens, batch_size=batch_size, @@ -177,47 +114,46 @@ async def forward_backward( token = await self._on_request_start(request) async def _do_forward_backward(): - try: - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - self.assert_resource_exists(adapter_name) - model_adapter = self.resolve_model_adapter_name(adapter_name) - datum_list = body.forward_backward_input.data - loss_fn = body.forward_backward_input.loss_fn - loss_fn_config = body.forward_backward_input.loss_fn_config or {} - preflight = getattr(self.model, 'preflight_loss_inputs', None) - if preflight is not None and loss_fn == 'importance_sampling' and not validate_tinker_loss( - loss_fn, datum_list, data_world_size=self.data_world_size): - # GRPO datums carry ragged per-sample targets; the loss check inside - # the sliced backend would deadlock the peer ranks (BUG-05), so the - # full batch is validated on every rank first. - await self.call_backend( - preflight, inputs=datum_list, adapter_name=model_adapter, loss_fn=loss_fn, **loss_fn_config) - output, loss = await self.call_backend( - self.model.tinker_forward_backward, - inputs=datum_list, - adapter_name=model_adapter, - loss_fn=loss_fn, - **loss_fn_config) - output_type = ('ImportanceSamplingLossReturn' - if loss_fn == 'importance_sampling' else 'CrossEntropyLossReturn') - self.set_resource_state(adapter_name, 'grad_ready', True) - return types.ForwardBackwardOutput( - loss_fn_output_type=output_type, - loss_fn_outputs=output, - metrics={'loss:mean': loss}, - ) - except Exception: - logger.error(traceback.format_exc()) - raise + adapter_name = body.model_id + self.assert_resource_exists(adapter_name) + model_adapter = self.resolve_model_adapter_name(adapter_name) + datum_list = body.forward_backward_input.data + loss_fn = body.forward_backward_input.loss_fn + loss_fn_config = body.forward_backward_input.loss_fn_config or {} + preflight = getattr(self.model, 'preflight_loss_inputs', None) + if preflight is not None and loss_fn == 'importance_sampling' and not validate_tinker_loss( + loss_fn, datum_list, data_world_size=self.data_world_size): + # GRPO datums carry ragged per-sample targets; the loss check inside + # the sliced backend would deadlock the peer ranks (BUG-05), so the + # full batch is validated on every rank first. + await self.call_backend( + preflight, inputs=datum_list, adapter_name=model_adapter, loss_fn=loss_fn, **loss_fn_config) + output, loss = await self.call_backend( + self.model.tinker_forward_backward, + inputs=datum_list, + adapter_name=model_adapter, + loss_fn=loss_fn, + **loss_fn_config) + output_type = ('ImportanceSamplingLossReturn' + if loss_fn == 'importance_sampling' else 'CrossEntropyLossReturn') + self.get_training_state(adapter_name).grad_ready = True + return types.ForwardBackwardOutput( + loss_fn_output_type=output_type, + loss_fn_outputs=output, + metrics={'loss:mean': loss}, + ) datum_list = body.forward_backward_input.data input_tokens = sum(len(d.model_input.to_ints()) for d in datum_list) batch_size = len(datum_list) loss_fn = body.forward_backward_input.loss_fn is_dpo = validate_tinker_loss(loss_fn, datum_list, data_world_size=self.data_world_size) - return await self.schedule_task( + return await submit_training_task( + self, _do_forward_backward, - model_id=body.model_id, + training_id=body.model_id, + seq_id=body.seq_id, + stream='tinker::forward_backward', token=token, input_tokens=input_tokens, batch_size=batch_size, @@ -235,24 +171,26 @@ async def optim_step( token = await self._on_request_start(request) async def _do_optim(): - try: - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - self.assert_resource_exists(adapter_name) - model_adapter = self.resolve_model_adapter_name(adapter_name) - if not self.get_resource_state(adapter_name, 'grad_ready', False): - raise RuntimeError(f'No accumulated gradients for adapter={adapter_name}; ' - 'call forward_backward before optim_step') - await self.call_backend( - self.model.tinker_step, adam_params=body.adam_params, adapter_name=model_adapter) - self.set_resource_state(adapter_name, 'grad_ready', False) - metrics = await self.call_backend( - self.model.tinker_calculate_metric, is_training=True, adapter_name=model_adapter) - return types.OptimStepResponse(metrics=metrics) - except Exception: - logger.error(traceback.format_exc()) - raise - - return await self.schedule_task(_do_optim, model_id=body.model_id, token=token, task_type='optim_step') + adapter_name = body.model_id + self.assert_resource_exists(adapter_name) + model_adapter = self.resolve_model_adapter_name(adapter_name) + if not self.get_training_state(adapter_name).grad_ready: + raise RuntimeError(f'No accumulated gradients for adapter={adapter_name}; ' + 'call forward_backward before optim_step') + await self.call_backend(self.model.tinker_step, adam_params=body.adam_params, adapter_name=model_adapter) + self.get_training_state(adapter_name).grad_ready = False + metrics = await self.call_backend( + self.model.tinker_calculate_metric, is_training=True, adapter_name=model_adapter) + return types.OptimStepResponse(metrics=metrics) + + return await submit_training_task( + self, + _do_optim, + training_id=body.model_id, + seq_id=body.seq_id, + stream='tinker::optim_step', + token=token, + task_type='optim_step') @app.post('/tinker/save_weights') async def save_weights( @@ -263,22 +201,25 @@ async def save_weights( token = await self._on_request_start(request) async def _do_save(): - try: - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - self.assert_resource_exists(adapter_name) - checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') - save_weights = partial( - self.call_backend, - self.model.save, - adapter_name=self.resolve_model_adapter_name(adapter_name), - save_optimizer=True) - tinker_path, _ = await checkpoint_manager.save_training(body.model_id, body.path, save_weights) - return types.SaveWeightsResponse(path=tinker_path, type='save_weights') - except Exception: - logger.error(traceback.format_exc()) - raise - - return await self.schedule_task(_do_save, model_id=body.model_id, token=token, task_type='save_weights') + adapter_name = body.model_id + self.assert_resource_exists(adapter_name) + checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') + save_weights = partial( + self.call_backend, + self.model.save, + adapter_name=self.resolve_model_adapter_name(adapter_name), + save_optimizer=True) + tinker_path, _ = await checkpoint_manager.save_training(body.model_id, body.path, save_weights) + return types.SaveWeightsResponse(path=tinker_path, type='save_weights') + + return await submit_training_task( + self, + _do_save, + training_id=body.model_id, + seq_id=body.seq_id, + stream='tinker::save_weights', + token=token, + task_type='save_weights') @app.post('/tinker/save_weights_for_sampler') async def save_weights_for_sampler( @@ -289,39 +230,40 @@ async def save_weights_for_sampler( token = await self._on_request_start(request) async def _do_save_for_sampler(): - try: - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - self.assert_resource_exists(adapter_name) - checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') - - save_weights = partial( - self.call_backend, - self.model.save, - adapter_name=self.resolve_model_adapter_name(adapter_name), - save_optimizer=False) - - tinker_path, _ = await checkpoint_manager.save_sampler(body.model_id, body.path, save_weights) - payload = body.model_dump() - payload['model_path'] = tinker_path - metadata = await self.state.get_model_metadata(body.model_id) or {} - if metadata.get('base_model'): - payload['base_model'] = metadata['base_model'] - sampling_session_id = await self.state.create_sampling_session(payload) - # Tinker SDK distinguishes two modes by whether sampling_session_seq_id is set: - # 1. save_weights_for_sampler(name): expects path != None - # 2. save_weights_and_get_sampling_client(): expects path == None, sampling_session_id != None - if body.sampling_session_seq_id is not None: - return types.SaveWeightsForSamplerResponseInternal( - path=None, sampling_session_id=sampling_session_id) - else: - return types.SaveWeightsForSamplerResponseInternal( - path=tinker_path, sampling_session_id=sampling_session_id) - except Exception: - logger.error(traceback.format_exc()) - raise - - return await self.schedule_task( - _do_save_for_sampler, model_id=body.model_id, token=token, task_type='save_weights_for_sampler') + adapter_name = body.model_id + self.assert_resource_exists(adapter_name) + checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') + + save_weights = partial( + self.call_backend, + self.model.save, + adapter_name=self.resolve_model_adapter_name(adapter_name), + save_optimizer=False) + + tinker_path, _ = await checkpoint_manager.save_sampler(body.model_id, body.path, save_weights) + payload = body.model_dump() + payload['model_path'] = tinker_path + metadata = await self.state.get_model_metadata(body.model_id) or {} + if metadata.get('base_model'): + payload['base_model'] = metadata['base_model'] + sampling_session_id = await self.state.create_sampling_session(payload) + # Tinker SDK distinguishes two modes by whether sampling_session_seq_id is set: + # 1. save_weights_for_sampler(name): expects path != None + # 2. save_weights_and_get_sampling_client(): expects path == None, sampling_session_id != None + if body.sampling_session_seq_id is not None: + return types.SaveWeightsForSamplerResponseInternal(path=None, sampling_session_id=sampling_session_id) + else: + return types.SaveWeightsForSamplerResponseInternal( + path=tinker_path, sampling_session_id=sampling_session_id) + + return await submit_training_task( + self, + _do_save_for_sampler, + training_id=body.model_id, + seq_id=body.seq_id, + stream='tinker::save_weights_for_sampler', + token=token, + task_type='save_weights_for_sampler') @app.post('/tinker/load_weights') async def load_weights( @@ -332,25 +274,28 @@ async def load_weights( token = await self._on_request_start(request) async def _do_load(): - try: - assert self.model is not None, 'Model not loaded, please load model first' - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - self.assert_resource_exists(adapter_name) - # Path resolution (token isolation, twinkle-vs-external path shapes) is a - # server storage policy and belongs in the handler, not in the GPU-side - # backend actor. The backend receives an already-resolved location. - checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') - resolved = checkpoint_manager.resolve_load_path(body.path) - await self.call_backend( - self.model.tinker_load, - checkpoint_name=resolved.checkpoint_name, - output_dir=resolved.checkpoint_dir if resolved.is_twinkle_path else None, - load_optimizer=body.optimizer, - adapter_name=self.resolve_model_adapter_name(adapter_name)) - self.set_resource_state(adapter_name, 'grad_ready', False) - return types.LoadWeightsResponse(path=body.path, type='load_weights') - except Exception: - logger.error(traceback.format_exc()) - raise - - return await self.schedule_task(_do_load, model_id=body.model_id, token=token, task_type='load_weights') + assert self.model is not None, 'Model not loaded, please load model first' + adapter_name = body.model_id + self.assert_resource_exists(adapter_name) + # Path resolution (token isolation, twinkle-vs-external path shapes) is a + # server storage policy and belongs in the handler, not in the GPU-side + # backend actor. The backend receives an already-resolved location. + checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') + resolved = checkpoint_manager.resolve_load_path(body.path) + await self.call_backend( + self.model.tinker_load, + checkpoint_name=resolved.checkpoint_name, + output_dir=resolved.checkpoint_dir if resolved.is_twinkle_path else None, + load_optimizer=body.optimizer, + adapter_name=self.resolve_model_adapter_name(adapter_name)) + self.get_training_state(adapter_name).grad_ready = False + return types.LoadWeightsResponse(path=body.path, type='load_weights') + + return await submit_training_task( + self, + _do_load, + training_id=body.model_id, + seq_id=body.seq_id, + stream='tinker::load_weights', + token=token, + task_type='load_weights') diff --git a/src/twinkle/server/model/twinkle_handlers.py b/src/twinkle/server/model/twinkle_handlers.py index b6d155b1c..b8537e8b7 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 @@ -528,14 +524,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( @@ -552,71 +541,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: - # Re-check at creation time: the admission check above runs - # before the async state write, so two concurrent creates can - # both pass it. The claim below is one synchronous step, which - # makes exactly one of them win. - if self.is_full_mode: - self.claim_full_mode_resource(adapter_name, token, session_id) - else: - 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/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/resources/allocation.py b/src/twinkle/server/resources/allocation.py new file mode 100644 index 000000000..1a43d82df --- /dev/null +++ b/src/twinkle/server/resources/allocation.py @@ -0,0 +1,138 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Single-node device allocation; only the resource 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 twinkle.server.exceptions import RequestRejectedError, ResourceReleaseError + + +class CapacityUnavailableError(RequestRejectedError): + error_code = 409 + + +@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 + self._creating = {} + + def name(self, epoch: str, instance_id: str) -> str: + return f'resources-{self.namespace}-{epoch}-{instance_id}' + + @staticmethod + def _create_group(resource, count, name): + import ray + from ray.util.placement_group import placement_group + candidates = [node for node in ray.nodes() if node.get('Alive') and node['Resources'].get(resource, 0) >= count] + if not candidates: + raise CapacityUnavailableError(f'No single-node allocation available for {count} {resource} devices') + # The aggregate free-resource snapshot can lag a confirmed group removal. + # pg.ready() is the allocation witness and is bounded by allocate(). + cpus = min(_gpu_placement_group_cpus(int(node['Resources'].get('CPU', 0)), count) for node in candidates) + bundle = {resource: count, 'CPU': count if resource == 'CPU' else cpus} + return placement_group([bundle], strategy='STRICT_PACK', name=name) + + 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) + creating = asyncio.create_task(asyncio.to_thread(self._create_group, resource, count, name)) + self._creating[name] = creating + try: + 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 + elif resource != 'CPU': + options['resources'] = {resource: count} + node, devices = await asyncio.wait_for( + ray.remote(_probe_allocation).options(**options).remote(resource, 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: + # 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 _finish_creation(self, name): + creating = self._creating.get(name) + if creating is None: + return + # A failed create has no result to free. A pending create must stay tracked + # across cancellation, so do not remove it until this await has completed. + try: + await asyncio.shield(creating) + except Exception: + pass # A failed create has no allocation result; cancellation remains unsafe. + if self._creating.get(name) is creating: + self._creating.pop(name) + + @staticmethod + async def _find_group(name): + from ray.util.placement_group import get_placement_group + try: + return await asyncio.to_thread(get_placement_group, name) + except ValueError: + return None # A missing named group is already released. + + async def _remove(self, name): + from ray.util.placement_group import placement_group_table, remove_placement_group + await self._finish_creation(name) + pg = await self._find_group(name) + if pg is None: + return + await asyncio.to_thread(remove_placement_group, pg) + while (await asyncio.to_thread(placement_group_table, pg)).get('state') != 'REMOVED': + await asyncio.sleep(0.05) + + async def remove(self, name: str) -> None: + try: + await asyncio.wait_for(self._remove(name), self.shutdown_timeout) + except TimeoutError as exc: + raise TimeoutError(f'Placement group removal was not confirmed: {name}') from exc + except Exception as exc: + raise ResourceReleaseError(f'Placement group removal failed: {name}', phase='pg_removal') from exc diff --git a/src/twinkle/server/resources/instances.py b/src/twinkle/server/resources/instances.py new file mode 100644 index 000000000..8bcdad118 --- /dev/null +++ b/src/twinkle/server/resources/instances.py @@ -0,0 +1,121 @@ +# 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: + logger.warning( + '%s instance %s shutdown is unconfirmed', + self.label, + identity, + exc_info=(type(error), error, error.__traceback__)) + await quarantined(error) + else: + # Only the feature can retire its users/quota and commit STOPPED. + 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) + + def require_released(self, instances): + """A shutdown ACK must certify release before its owner can be removed.""" + owned = self.owned_count(instances) + if owned: + raise RuntimeError(f'{self.label} shutdown is unconfirmed: {owned} instances remain owned') + + @classmethod + def headroom(cls, instances, maximum): + return float('inf') if maximum is None else max(0, maximum - cls.owned_count(instances)) + + @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..fbc251d58 --- /dev/null +++ b/src/twinkle/server/resources/runtime.py @@ -0,0 +1,124 @@ +# 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 time +from copy import deepcopy + +from twinkle.server.config.resources import build_worker_env_vars +from twinkle.server.exceptions import ResourceReleaseError +from .allocation import DevicePool +from .instances import InstanceLifecycle + +logger = logging.getLogger(__name__) + + +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.lifecycle = InstanceLifecycle('Ray deployment cleanup') + + 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_worker_env_vars() + env.update(runtime_env.get('env_vars', {})) + env.update(env_vars or {}) + env.update(TWINKLE_RAY_NAMESPACE=self.namespace, TWINKLE_RAY_REGISTRY_NAME=name + '-registry') + 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(): + await asyncio.gather(deployment, return_exceptions=True) + await cleanup() + + self.lifecycle.spawn(clean_late()) + raise + finally: + if deployment.done(): + self.deployments.pop(identity, None) + + async def wait_for_cleanup(self): + await self.lifecycle.join() + + @staticmethod + 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 + from ray import serve + from ray.util import list_named_actors + deployment = self.deployments.get(identity) + if deployment is not None: + await asyncio.gather(asyncio.shield(deployment), return_exceptions=True) + self.deployments.pop(identity, None) + # Both features give workers the controller identity, not a Serve replica ID. + actors = [] + 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: + # 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) + await asyncio.to_thread(ray.kill, handle, no_restart=True) + await asyncio.to_thread(serve.delete, name) + for handle in actors: + try: + await asyncio.wait_for(handle.__ray_ready__.remote(), 5) + except ray.exceptions.ActorDiedError: + continue + except ray.exceptions.RayActorError as exc: + raise ResourceReleaseError('Owned worker death could not be confirmed', phase='worker_death') from exc + raise ResourceReleaseError('Owned worker death could not be confirmed', phase='worker_death') + + started, status = time.monotonic(), 'error' + try: + await self.pool.remove(name) + status = 'success' + finally: + if observe_removal is not None: + observe_removal(time.monotonic() - started, status) + + 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 + 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: + 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 f0ee357df..ee3964e59 100644 --- a/src/twinkle/server/sampler/app.py +++ b/src/twinkle/server/sampler/app.py @@ -7,21 +7,26 @@ """ from __future__ import annotations +import asyncio +import os 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 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.config import TaskQueueConfig +from twinkle.server.state.models import RuntimeContext 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 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,18 @@ 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: + context = RuntimeContext.model_validate(context) + from twinkle.infra._ray.ray_helper import RayHelper + from twinkle.infra._ray.resource_manager import ResourceManager + os.environ['TWINKLE_RAY_NAMESPACE'] = context.namespace + os.environ['TWINKLE_RAY_REGISTRY_NAME'] = context.registry_name + 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, @@ -97,8 +113,10 @@ async def __init__(self, ) self.sampler_type = sampler_type self.model_id = model_id - replica_context = serve.get_replica_context() - self.replica_id = replica_context.replica_id.unique_id + if context is not None: + self.replica_id = context.instance_id + else: + self.replica_id = serve.get_replica_context().replica_id.unique_id sampler_kwargs: dict[str, Any] = { 'model_id': model_id, @@ -117,6 +135,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,21 +167,28 @@ 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 check_health(self) -> None: - self.sampling_guard.check_available() + 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() - @serve.multiplexed(max_num_models_per_replica=5) - async def _sticky_entry(self, sticky_key: str): - return sticky_key + async def runtime_identity(self): + return await self.request_runtime.runtime_identity() - async def _ensure_sticky(self): - sticky_key = serve.get_multiplexed_model_id() - await self._sticky_entry(sticky_key) + async def check_health(self) -> None: + self.sampling_guard.check_available() 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() return get_token_from_request(request) @@ -204,9 +231,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/__init__.py b/src/twinkle/server/sampler/backends/__init__.py index c2caf8e9e..a734d4690 100644 --- a/src/twinkle/server/sampler/backends/__init__.py +++ b/src/twinkle/server/sampler/backends/__init__.py @@ -1,11 +1,2 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Sampler backend implementations. - -The cross-process streaming bridge lives in ``streaming`` and is re-exported here -so a sibling backend module can import it from ``.streaming`` without importing -this package root, while external consumers keep using -``from twinkle.server.sampler.backends import stream_to_queue``. -""" -from .streaming import STREAM_SENTINEL, stream_to_queue - -__all__ = ['STREAM_SENTINEL', 'stream_to_queue'] +"""Sampling backend implementations; helpers are imported from their defining modules.""" diff --git a/src/twinkle/server/sampler/backends/streaming.py b/src/twinkle/server/sampler/backends/streaming.py index dc61ff7e7..5d07e1b59 100644 --- a/src/twinkle/server/sampler/backends/streaming.py +++ b/src/twinkle/server/sampler/backends/streaming.py @@ -1,26 +1,44 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Cross-process streaming bridge for sampler backends. +"""Lightweight queue bridge shared by real and mock sampling backends. -Lives in its own module (rather than the ``backends`` package ``__init__``) so -that a backend implementation (e.g. ``mock_sampler``) can import it without -importing its own package root — which would be a package-initialisation-order -dependency. ``backends/__init__`` re-exports these names, so external consumers -(``from twinkle.server.sampler.backends import stream_to_queue``) are unchanged. +Kept independent of the HTTP runtime so backend workers do not import FastAPI +or depend on the mock backend. Both consumers import this defining module. """ from __future__ import annotations STREAM_SENTINEL = '__STREAM_END__' +class StreamPublisher: + """A disconnected consumer cannot end the backend generation lifetime.""" + + def __init__(self, queue): + self.queue = queue + self.connected = True + + def put(self, item): + if not self.connected: + return + try: + self.queue.put(item) + except Exception: + # Ray queue transport can disappear after the HTTP reader closes. + self.connected = False + + def stream_to_queue(sampler, queue, inputs, sampling_params=None, adapter_name='', adapter_path=None): - """Push streaming deltas from *sampler* to a cross-process Ray queue. - - Works with any object that exposes a ``sample_stream`` iterator. - """ - try: - for delta, reason in sampler.sample_stream(inputs, sampling_params, adapter_name, adapter_path): - queue.put((delta, reason)) - except Exception as e: - queue.put(e) - finally: - queue.put(STREAM_SENTINEL) + """Drain and close the generator even when its HTTP consumer disconnects.""" + from contextlib import ExitStack + + publisher = StreamPublisher(queue) + generation = sampler.sample_stream(inputs, sampling_params, adapter_name, adapter_path) + with ExitStack() as cleanup: + cleanup.callback(publisher.put, STREAM_SENTINEL) + close = getattr(generation, 'close', None) + if close is not None: + cleanup.callback(close) + try: + for delta, reason in generation: + publisher.put((delta, reason)) + except Exception as error: + publisher.put(error) diff --git a/src/twinkle/server/sampler/concurrency.py b/src/twinkle/server/sampler/concurrency.py index 71d7e7c68..0acb32232 100644 --- a/src/twinkle/server/sampler/concurrency.py +++ b/src/twinkle/server/sampler/concurrency.py @@ -4,14 +4,14 @@ import asyncio -from twinkle.server.task_queue.types import BackendBusyError +from twinkle.server.exceptions import BackendBusyError 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/pool/__init__.py b/src/twinkle/server/sampler/pool/__init__.py new file mode 100644 index 000000000..78ee5693d --- /dev/null +++ b/src/twinkle/server/sampler/pool/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Sampling instance pool: routing, lifecycle decisions and Ray composition.""" diff --git a/src/twinkle/server/sampler/pool/actor.py b/src/twinkle/server/sampler/pool/actor.py new file mode 100644 index 000000000..7d571e94b --- /dev/null +++ b/src/twinkle/server/sampler/pool/actor.py @@ -0,0 +1,51 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Ray composition for the sampling controller and its owned driver.""" +import asyncio +from copy import deepcopy + +from twinkle.server.config.application_spec import SamplerArgs +from .controller import SamplerController + + +def create_controller_actor(): + import ray + + from .driver import SamplerDriver + + @ray.remote(num_cpus=0.1, max_restarts=-1, max_task_retries=0) + class SamplerControllerActor(SamplerController): + + def __init__(self, values, namespace, pool_id, options, budgets): + args = SamplerArgs.model_validate(values) + super().__init__(args, SamplerDriver(args, namespace, pool_id, options, budgets), budgets) + + async def initialize(self): + await super().initialize() + if not self._stopping and self._monitor_task is None: + self._monitor_task = asyncio.create_task(self._monitor()) + + async def shutdown(self): + await super().shutdown() + await self.driver.runtime.wait_for_cleanup() + + return SamplerControllerActor + + +def create_sampler_controller(spec, namespace, budgets, env): + """Driver-only composition. Each compute instance owns its external allocation.""" + from .controller import sampler_controller_name, sampling_pool_id + + options = deepcopy(spec.deployments[0]) if spec.deployments else {} + options.pop('name', None) + options.pop('autoscaling_config', None) + options['num_replicas'] = 1 + actor_options = options.setdefault('ray_actor_options', {}) + actor_options.setdefault('num_cpus', 0.1) + runtime_env = actor_options.setdefault('runtime_env', {}) + runtime_env['env_vars'] = {**env, **runtime_env.get('env_vars', {}), 'TWINKLE_STATE_RAY_NAMESPACE': namespace} + pool_id = sampling_pool_id(spec.route_prefix) + name = sampler_controller_name(namespace, pool_id) + controller = create_controller_actor().options( + name=name, namespace=namespace, runtime_env=runtime_env).remote(spec.args.model_dump(), namespace, pool_id, + options, budgets) + return controller, options, pool_id diff --git a/src/twinkle/server/sampler/pool/controller.py b/src/twinkle/server/sampler/pool/controller.py new file mode 100644 index 000000000..1eb60431c --- /dev/null +++ b/src/twinkle/server/sampler/pool/controller.py @@ -0,0 +1,366 @@ +# 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 json +import logging +import math +import time +import uuid +from copy import deepcopy +from dataclasses import asdict, dataclass, field +from typing import TYPE_CHECKING, Literal + +if TYPE_CHECKING: + from twinkle.server.lifecycle.proxy import ForwardedRequest, ForwardedReply + +from twinkle.server.config.application_spec import SamplerArgs +from twinkle.server.config.resources import ResourceOptions +from twinkle.server.exceptions import ResourceExhaustedError +from twinkle.server.resources.instances import InstanceLifecycle +from twinkle.server.sampler.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: Literal['STARTING', 'READY', 'DRAINING', 'QUARANTINED', 'STOPPED'] = 'STARTING' + application: str | None = None + devices: list[str] = field(default_factory=list) + requests: set[str] = field(default_factory=set) + 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.clock = args, driver, clock + self.budgets = ResourceOptions.model_validate(budgets) + 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 + + @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: ForwardedRequest, stream_queue=None, full_path=None) -> ForwardedReply: + 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: + reply_ready.set_result(await self._mutate(payload=payload)) + return + if full_path is not None and full_path != self._base_path: + 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(): + 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: + await self._release(instance, ticket) + + async def guarded(): + try: + 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) + 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() + self.lifecycle.require_released(self.instances.values()) diff --git a/src/twinkle/server/sampler/pool/driver.py b/src/twinkle/server/sampler/pool/driver.py new file mode 100644 index 000000000..e09652f27 --- /dev/null +++ b/src/twinkle/server/sampler/pool/driver.py @@ -0,0 +1,106 @@ +# 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.config.resources import ResourceOptions +from twinkle.server.resources.runtime import RayRuntime +from twinkle.server.state import get_server_state +from twinkle.server.state.models import FutureFailureRecord, RuntimeContext + + +class SamplerDriver: + + def __init__(self, args: SamplerArgs, namespace: str, pool_id: str, options: dict, budgets: dict): + self.args, self.namespace, self.pool_id = args, namespace, pool_id + self.options = options + self.budgets = ResourceOptions.model_validate(budgets) + self.runtime = RayRuntime( + namespace, + allocation_timeout=self.budgets.allocation_timeout_s, + shutdown_timeout=self.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 twinkle.server.sampler.app import build_sampler_app + group = self.args.device_group + device_type = str(group.get('device_type', 'cpu' if self.args.sampler_type == 'mock' else 'cuda')).lower() + allocation = await self.runtime.pool.allocate( + 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=RuntimeContext( + namespace=self.namespace, registry_name=allocation.name + '-registry', instance_id=instance.identity), + placement_group=await asyncio.to_thread(get_placement_group, allocation.name), + allocation_devices=list(allocation.devices)) + await self.runtime.deploy(instance.identity, allocation.name, app, + f'/api/v1/sampling-runtime/{instance.identity}', lambda: self.shutdown(instance)) + + async def handle(self, instance): + return await self.runtime.handle(instance.application) + + async def invoke(self, instance, ticket, payload, stream_queue=None): + handle = await self.handle(instance) + return await handle.invoke.remote(ticket, payload, stream_queue) + + async def wait_operation(self, instance, ticket): + handle = await self.handle(instance) + await handle.wait_operation.remote(ticket) + + async def load_weights(self, instance, path): + handle = await self.handle(instance) + await handle.install_base_weights.remote(path) + + async def health(self, instance): + 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: + 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) + 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/pool/router.py b/src/twinkle/server/sampler/pool/router.py new file mode 100644 index 000000000..138c055d1 --- /dev/null +++ b/src/twinkle/server/sampler/pool/router.py @@ -0,0 +1,104 @@ +# 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 +from fastapi import Depends, Request, Response +from fastapi.responses import StreamingResponse + +from twinkle.server.checkpoint import create_checkpoint_manager +from twinkle.server.deployment import bind_deployment, build_deployment_app +from twinkle.server.lifecycle.proxy import ForwardedRequest +from twinkle.server.middleware.auth import get_token_from_request +from twinkle.server.sampler.weights import is_full_checkpoint +from twinkle.server.state import get_server_state +from .controller import sampler_controller_name + + +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: ForwardedRequest = { + '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)) diff --git a/src/twinkle/server/sampler/runtime.py b/src/twinkle/server/sampler/runtime.py new file mode 100644 index 000000000..5549cdeec --- /dev/null +++ b/src/twinkle/server/sampler/runtime.py @@ -0,0 +1,160 @@ +# 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 dataclasses import dataclass + +from twinkle.server.deployment import build_deployment_app +from twinkle.server.exceptions import ResourceNotFoundError +from twinkle.server.lifecycle.proxy import ForwardedReply, ForwardedRequest +from twinkle.server.task_queue.types import ExecutionScope + + +def request_execution_scope(request) -> ExecutionScope | None: + """Internal ASGI metadata, absent for direct standalone endpoint calls.""" + return request.scope.get('state', {}).get('execution_scope') + + +@dataclass +class SamplingOperation: + scope: ExecutionScope + task: asyncio.Task | None = None + + +class SamplerRuntime: + + def __init__(self, owner): + self.owner = owner + self._operations: dict[str, SamplingOperation] = {} + self._stream_tasks = set() + self._rpc_app = None + + def start_stream_job(self, factory, *, execution_scope: ExecutionScope | None = None): + job = asyncio.create_task(factory()) + if execution_scope is not None: + execution_scope.track(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: str, payload: ForwardedRequest, stream_queue=None) -> ForwardedReply: + """Return the HTTP reply early while retaining every actual generation job.""" + if ticket in self._operations: + raise ResourceNotFoundError('Sampling request was already dispatched') + 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 = SamplingOperation(ExecutionScope()) + self._operations[ticket] = record + + async def dispatch(): + response: ForwardedReply = {'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', + 'state': { + 'execution_scope': record.scope + }, + '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: + # 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/tinker_handlers.py b/src/twinkle/server/sampler/tinker_handlers.py index 719ab0621..b6a6543ed 100644 --- a/src/twinkle/server/sampler/tinker_handlers.py +++ b/src/twinkle/server/sampler/tinker_handlers.py @@ -20,9 +20,10 @@ from twinkle.data_format import SamplingParams from twinkle.server.checkpoint import create_checkpoint_manager +from twinkle.server.exceptions import RequestRejectedError from twinkle.server.sampler.weights import sample_with_weights -from twinkle.server.task_queue.types import UserTaskError from twinkle.utils.logger import get_logger +from .runtime import request_execution_scope logger = get_logger() @@ -84,7 +85,7 @@ async def _do_sample(): # Base-model sampling is valid when no model_path was provided. if adapter_uri and not os.path.exists(adapter_uri): - raise UserTaskError(f'Adapter URI {model_path} does not exist. Please check the model_path.') + raise RequestRejectedError(f'Adapter URI {model_path} does not exist. Please check the model_path.') # Convert tinker SamplingParams to twinkle SamplingParams if needed sampling_params = None @@ -153,4 +154,5 @@ async def generate(lora_path): token=token, input_tokens=input_tokens, task_type='sample', + execution_scope=request_execution_scope(request), ) diff --git a/src/twinkle/server/sampler/twinkle_handlers.py b/src/twinkle/server/sampler/twinkle_handlers.py index b6cf58d3b..a37f6475c 100644 --- a/src/twinkle/server/sampler/twinkle_handlers.py +++ b/src/twinkle/server/sampler/twinkle_handlers.py @@ -33,6 +33,7 @@ from twinkle.server.telemetry.correlation import MODEL_ID from twinkle.server.telemetry.tracing import traced_operation from twinkle.utils.logger import get_logger +from .runtime import request_execution_scope logger = get_logger() @@ -165,9 +166,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 +182,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 +201,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: @@ -262,7 +269,12 @@ async def generate(adapter_path): # Calculate metrics for queue scheduling. The body is wire-validated, so the # entries are models and ``input_ids`` is absent or a list of ints. input_tokens = sum(len(getattr(entry, 'input_ids', None) or ()) for entry in body.inputs) - return await self.submit_concurrent_and_peek(_task, token=token, input_tokens=input_tokens, task_type='sample') + return await self.submit_concurrent_and_peek( + _task, + token=token, + input_tokens=input_tokens, + task_type='sample', + execution_scope=request_execution_scope(request)) @app.post('/twinkle/sample_to_data_plane', response_model=types.TaskEnvelope) async def sample_to_data_plane( @@ -319,7 +331,10 @@ async def generate(): return ref.model_dump() return await self.submit_concurrent_and_peek( - _generate_and_store, model_id=full_adapter_name or None, task_type='sample_to_data_plane') + _generate_and_store, + model_id=full_adapter_name or None, + task_type='sample_to_data_plane', + execution_scope=request_execution_scope(request)) @app.post('/twinkle/unload_adapter_paths') async def unload_adapter_paths( @@ -368,8 +383,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 +423,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 +431,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 @@ -430,7 +446,7 @@ async def sample_stream( from ray.util.queue import Queue - from .backends import STREAM_SENTINEL + from .backends.streaming import STREAM_SENTINEL request_id = f'req_{uuid.uuid4().hex}' actors = self.sampler._actors @@ -443,14 +459,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, execution_scope=request_execution_scope(request)) return StreamingResponse( _stream_queue( q, diff --git a/src/twinkle/server/sampler/weights.py b/src/twinkle/server/sampler/weights.py index f2d1077f7..dac788542 100644 --- a/src/twinkle/server/sampler/weights.py +++ b/src/twinkle/server/sampler/weights.py @@ -7,7 +7,8 @@ from safetensors import SafetensorError from twinkle.exceptions import FullWeightLoadError -from twinkle.server.task_queue.types import BACKEND_TIMEOUT_EXCEPTIONS, UserTaskError +from twinkle.server.exceptions import RequestRejectedError +from twinkle.server.task_queue.backend_gate import BACKEND_TIMEOUT_EXCEPTIONS def is_full_checkpoint(resolved_uri: str | None) -> bool: @@ -16,10 +17,10 @@ def is_full_checkpoint(resolved_uri: str | None) -> bool: if os.path.isfile(os.path.join(resolved_uri, 'adapter_config.json')): return False if os.path.isabs(resolved_uri) and not os.path.isdir(resolved_uri): - raise UserTaskError(f'Checkpoint directory does not exist: {resolved_uri}') + raise RequestRejectedError(f'Checkpoint directory does not exist: {resolved_uri}') if os.path.isdir(resolved_uri) and not (os.path.isfile(os.path.join(resolved_uri, 'model.safetensors.index.json')) or glob.glob(os.path.join(resolved_uri, '*.safetensors'))): - raise UserTaskError(f'No .safetensors weights found under {resolved_uri}') + raise RequestRejectedError(f'No .safetensors weights found under {resolved_uri}') return True @@ -38,7 +39,7 @@ async def resolve_sampler_weights(service, resolved_uri: str | None, *, full_che service.sampling_guard.invalidate_weights(error) raise except (ValueError, OSError, KeyError, SafetensorError) as error: - raise UserTaskError(str(error)) from error + raise RequestRejectedError(str(error)) from error return None diff --git a/src/twinkle/server/session_resource/__init__.py b/src/twinkle/server/session_resource/__init__.py index fdb7ebdc1..7384532cb 100644 --- a/src/twinkle/server/session_resource/__init__.py +++ b/src/twinkle/server/session_resource/__init__.py @@ -7,8 +7,7 @@ session-driven expiration of session-bound resources. """ -from .adapter import AdapterManagerMixin from .base import SessionResourceMixin from .processor import ProcessorManagerMixin -__all__ = ['AdapterManagerMixin', 'ProcessorManagerMixin', 'SessionResourceMixin'] +__all__ = ['ProcessorManagerMixin', 'SessionResourceMixin'] diff --git a/src/twinkle/server/session_resource/adapter.py b/src/twinkle/server/session_resource/adapter.py deleted file mode 100644 index 392a05bf1..000000000 --- a/src/twinkle/server/session_resource/adapter.py +++ /dev/null @@ -1,86 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -""" -Adapter Lifecycle Manager Mixin for Twinkle Server. - -This module provides adapter lifecycle management as a mixin class that can be -inherited directly by services. It tracks adapter activity and provides interfaces -for registration, heartbeat updates, and expiration handling. - -By inheriting this mixin, services can override the _on_adapter_expired() method -to handle expired adapters without using callbacks or polling. -""" -from __future__ import annotations - -from abc import abstractmethod - -from twinkle.utils.logger import get_logger -from .base import SessionResourceMixin - -logger = get_logger() - - -class AdapterManagerMixin(SessionResourceMixin): - """Mixin for adapter lifecycle management with session-based expiration. - - This mixin tracks adapter activity and automatically expires adapters - when their associated session expires. - - Inheriting classes should: - 1. Call _init_adapter_manager() in __init__ - 2. Override _on_adapter_expired() to customize expiration handling - - The inactivity timeout / max lifetime are stored on the base mixin as - ``_resource_timeout`` / ``_resource_max_lifetime`` (set via - ``_init_adapter_manager``). - """ - - # Set resource type for logging - _resource_type = 'Adapter' - - def _init_adapter_manager( - self, - adapter_timeout: float = 1800.0, - adapter_max_lifetime: float = 36000.0, - ) -> None: - """Initialize the adapter manager. - - This should be called in the __init__ of the inheriting class. - - Args: - adapter_timeout: Timeout in seconds used to check whether a session is still alive. - Default is 1800.0 (30 minutes). - adapter_max_lifetime: Maximum lifetime in seconds for an adapter regardless of session - liveness. Adapters older than this are treated as expired. Default is 36000.0 (10 hours). - """ - self._init_resource_manager( - resource_timeout=adapter_timeout, - resource_max_lifetime=adapter_max_lifetime, - ) - - async def _on_resource_expired(self, resource_id: str) -> None: - """Base-class expiry hook; forwards to the domain hook ``_on_adapter_expired``. - - ``_on_adapter_expired`` is the supported extension point: the adapter-domain - name is kept deliberately so subclass authors override a method named for - adapters rather than the generic base-class resource hook. - """ - await self._on_adapter_expired(resource_id) - - @abstractmethod - async def _on_adapter_expired(self, adapter_name: str) -> None: - """Hook method called when an adapter expires.""" - ... - - @staticmethod - def get_adapter_name(adapter_name: str) -> str: - """Get the adapter name for a request. - - This is a passthrough method for consistency with the original API. - - Args: - adapter_name: The adapter name (typically model_id) - - Returns: - The adapter name to use - """ - return adapter_name diff --git a/src/twinkle/server/session_resource/processor.py b/src/twinkle/server/session_resource/processor.py index 69a6192e3..dd2a7ff10 100644 --- a/src/twinkle/server/session_resource/processor.py +++ b/src/twinkle/server/session_resource/processor.py @@ -2,7 +2,7 @@ """ Processor Lifecycle Manager Mixin for Twinkle Server. -Mirrors AdapterManagerMixin but adds a global per-token processor limit. +Adds a global per-token processor limit to session resources. Sessions are tracked via session ID; processors expire when their session expires. """ from __future__ import annotations @@ -18,7 +18,7 @@ class ProcessorManagerMixin(SessionResourceMixin): """Mixin for processor lifecycle management with session-based expiration. - Mirrors AdapterManagerMixin with an additional per-token processor limit. + Manages processor sessions with a global per-token processor limit. Inheriting classes should: 1. Call _init_processor_manager() in __init__ @@ -80,9 +80,7 @@ async def _on_resource_expired(self, resource_id: str) -> None: ``_on_processor_expired`` is the supported extension point: the processor-domain name is kept deliberately so subclass authors override a - method named for processors rather than the generic base-class hook. It is - ``async`` to match the sibling ``AdapterManagerMixin._on_adapter_expired`` - contract, so both resource kinds expose the same extension-point shape. + method named for processors rather than the generic base-class hook. """ await self._on_processor_expired(resource_id) 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/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..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] = [] @@ -121,21 +123,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) @@ -163,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 e6769551a..53b4d78ce 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 _retry_read(self, operation, *args): + """Retry only read commands after a stale pooled connection is discarded. + + Writes/EXEC are deliberately excluded: a lost acknowledgement does not + 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._retry_read(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._retry_read(self._client.exists, real_key)) async def keys(self, pattern: str) -> list[str]: """Return all key names matching the pattern using SCAN. @@ -106,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. @@ -116,8 +139,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 +150,40 @@ async def update_atomic( new_value = transform(current) if new_value is None: await pipe.unwatch() - return current + self._metrics.update_atomic_retries.observe(attempt) + return False if return_changed else 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() - return new_value - except WatchError: + self._metrics.update_atomic_retries.observe(attempt) + 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. + 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 +191,7 @@ async def mget(self, keys: list[str]) -> list[Any | None]: if not keys: return [] real_keys = [self._make_key(k) for k in keys] - raw_values = await self._client.mget(real_keys) + raw_values = await self._retry_read(self._client.mget, real_keys) return [json.loads(v) if v is not None else None for v in raw_values] async def close(self) -> None: @@ -157,6 +201,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._retry_read(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/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..4e36d28b4 100644 --- a/src/twinkle/server/state/future_manager.py +++ b/src/twinkle/server/state/future_manager.py @@ -5,6 +5,7 @@ import time from typing import Any +from twinkle.protocol.types.lifecycle import TERMINAL_STATUSES from twinkle.utils.logger import get_logger from .backend.base import StateBackend from .base import BaseManager @@ -12,100 +13,114 @@ logger = get_logger() -# 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'}) - - -def _future_record_transform( - existing: dict | None, - *, - request_id: str, - new_status: str, - model_id: str | None, - reason: str | None, - result: Any, - failure: dict[str, Any] | None, - queue_state: str | None, - queue_state_reason: str | None, - replica_id: str | None, - absolute_deadline: float | None, - now: str, -) -> dict | None: - """Atomic transform body for :meth:`FutureManager.store_status`. - - Module-level so it remains picklable when forwarded across the Ray actor - boundary (closures and lambdas cannot be). - - A record already in a terminal state is never overwritten (returns ``None``, - which ``update_atomic`` treats as "keep current value"). A write of a - *different* terminal state is logged; a write of the *same* terminal state is - dropped silently (State_Backend idempotent retries produce these and they - indicate no defect). - """ - existing_status = existing.get('status') if existing is not None else None - if existing_status in _TERMINAL_STATUSES: - if new_status != existing_status: - logger.warning('future %s already terminal as %r; refusing %r', request_id, existing_status, new_status) - return None - - if existing is None: - record = FutureRecord( - status=new_status, - model_id=model_id, - reason=reason, - result=result, - failure=failure, - queue_state=queue_state, - queue_state_reason=queue_state_reason, - replica_id=replica_id, - absolute_deadline=absolute_deadline, - created_at=now, - updated_at=now, - ) - return record.model_dump() - - updated = dict(existing) - updated['status'] = new_status - updated['model_id'] = model_id - updated['updated_at'] = now - # replica_id is set at creation and is deliberately NOT overwritten here. - if reason is not None: - updated['reason'] = reason - if new_status in _TERMINAL_STATUSES: - if new_status in ('failed', 'cancelled'): - updated['result'] = None - updated['failure'] = failure - else: - updated['failure'] = None - updated['result'] = result - elif result is not None: - updated['result'] = result - if queue_state is not None: - updated['queue_state'] = queue_state - if queue_state_reason is not None: - updated['queue_state_reason'] = queue_state_reason - return updated - -_CANCELLABLE_STATUSES = frozenset({'pending', 'queued'}) +class FutureTransitions: + """Pure, picklable atomic transitions; the manager owns I/O and retention.""" + terminal = TERMINAL_STATUSES + cancellable = frozenset({'pending', 'queued'}) + order = {'pending': 0, 'queued': 1, 'running': 2} -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). - - Returns ``None`` (no change) for running/terminal/missing records so a task - already executing is never interrupted -- cancel is best-effort on the queue. - """ - status = existing.get('status') if existing is not None else None - if status not in _CANCELLABLE_STATUSES: - return None - updated = dict(existing) - updated['status'] = 'cancelled' - updated['result'] = None - updated['failure'] = failure - updated['updated_at'] = now - return updated + @staticmethod + def update( + existing: dict | None, + *, + request_id: str, + new_status: str, + model_id: str | None, + reason: str | None, + result: Any, + failure: dict[str, Any] | None, + queue_state: str | None, + queue_state_reason: str | None, + replica_id: str | None, + absolute_deadline: float | None, + now: str, + owner_token: str | None = None, + ) -> dict | None: + """Atomic transform body for :meth:`FutureManager.store_status`. + + A static method remains picklable across the Ray actor boundary. + + A record already in a terminal state is never overwritten (returns ``None``, + which ``update_atomic`` treats as "keep current value"). A write of a + *different* terminal state is logged; a write of the *same* terminal state is + dropped silently (State_Backend idempotent retries produce these and they + indicate no defect). + """ + existing_status = existing.get('status') if existing is not None else None + if existing_status in FutureTransitions.terminal: + if new_status != existing_status: + logger.warning('future %s already terminal as %r; refusing %r', request_id, existing_status, new_status) + return None + if new_status in FutureTransitions.order and existing_status in FutureTransitions.order: + if FutureTransitions.order[new_status] <= FutureTransitions.order[existing_status]: + return None + if new_status == 'cancelled' and existing_status == 'running': + return None + + if existing is None: + record = FutureRecord( + owner_token=owner_token, + status=new_status, + model_id=model_id, + reason=reason, + result=result, + failure=failure, + queue_state=queue_state, + queue_state_reason=queue_state_reason, + replica_id=replica_id, + absolute_deadline=absolute_deadline, + created_at=now, + updated_at=now, + ) + return record.model_dump() + + updated = dict(existing) + updated['status'] = new_status + updated['model_id'] = model_id + updated['updated_at'] = now + # replica_id is set at creation and is deliberately NOT overwritten here. + if reason is not None: + updated['reason'] = reason + if new_status in FutureTransitions.terminal: + if new_status in ('failed', 'cancelled'): + updated['result'] = None + updated['failure'] = failure + else: + updated['failure'] = None + updated['result'] = result + elif result is not None: + updated['result'] = result + if queue_state is not None: + updated['queue_state'] = queue_state + if queue_state_reason is not None: + updated['queue_state_reason'] = queue_state_reason + return updated + + @staticmethod + def start(existing: dict | None, *, now: str) -> dict | None: + """Acquire execution only from pending/queued, using the existing record.""" + if existing is None or existing.get('status') not in FutureTransitions.cancellable: + return None + return {**existing, 'status': 'running', 'queue_state': 'active', 'updated_at': now} + + @staticmethod + def cancel(existing: dict | None, *, failure: dict, now: str) -> dict | None: + """Atomic transform: cancel iff the task has not started (pending/queued). + + Returns ``None`` (no change) for running/terminal/missing records so a task + already executing is never interrupted -- cancel is best-effort on the queue. + """ + status = existing.get('status') if existing is not None else None + if status not in FutureTransitions.cancellable: + return None + updated = dict(existing) + updated['status'] = 'cancelled' + updated['result'] = None + updated['failure'] = failure + updated['updated_at'] = now + return updated class FutureManager(BaseManager[FutureRecord]): @@ -128,6 +143,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. @@ -153,7 +169,7 @@ async def store_status( await self._backend.update_atomic( self._make_key(request_id), functools.partial( - _future_record_transform, + FutureTransitions.update, request_id=request_id, new_status=status, model_id=model_id, @@ -164,10 +180,19 @@ async def store_status( queue_state_reason=queue_state_reason, replica_id=replica_id, absolute_deadline=absolute_deadline, + owner_token=owner_token, now=now, ), ) + 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(FutureTransitions.start, 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. @@ -183,7 +208,7 @@ async def cancel_if_pending(self, request_id: str) -> str | None: ) result = await self._backend.update_atomic( self._make_key(request_id), - functools.partial(_cancel_if_not_started_transform, failure=failure.model_dump(), now=_now_iso()), + functools.partial(FutureTransitions.cancel, failure=failure.model_dump(), now=_now_iso()), ) return result.get('status') if result else None @@ -219,14 +244,14 @@ async def cleanup_expired( now = time.time() expired_ids: list[str] = [] for request_id, record in all_records.items(): - if record.status in _TERMINAL_STATUSES: + if record.status in FutureTransitions.terminal: timestamp = self._parse_timestamp(record.updated_at or record.created_at) if timestamp < cutoff_time: expired_ids.append(request_id) continue # Non-terminal records are never deleted -- only ever written ``failed``. - # replica_id None (pre-upgrade) => ownership unknown => treated as alive. + # Unowned background tasks have no replica; only explicit owners can be declared dead. replica_id = record.replica_id replica_alive = (replica_id is None or alive_replica_ids is None or replica_id in alive_replica_ids) if not replica_alive: diff --git a/src/twinkle/server/state/model_manager.py b/src/twinkle/server/state/model_manager.py index aaf520fca..39c327e72 100644 --- a/src/twinkle/server/state/model_manager.py +++ b/src/twinkle/server/state/model_manager.py @@ -93,7 +93,8 @@ async def rebuild_indexes(self) -> None: async def get_capacity_info(self) -> dict[str, int]: """Return global LoRA capacity across all registered replicas.""" - replicas = await self._replicas.get_all() + capacities = await self._replicas.get_all() + replicas = {identity: capacity for identity, capacity in capacities.items() if capacity > 0} loaded_per_replica = await self._loaded_per_replica() total_max = sum(replicas.values()) total_used = sum(loaded_per_replica.get(rid, 0) for rid in replicas) @@ -116,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) @@ -220,6 +237,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..d84760d38 100644 --- a/src/twinkle/server/state/models.py +++ b/src/twinkle/server/state/models.py @@ -3,7 +3,7 @@ import time from datetime import datetime, timezone -from pydantic import BaseModel, Field, model_validator +from pydantic import BaseModel, ConfigDict, Field, model_validator from typing import Any, Literal @@ -14,6 +14,16 @@ def _now_iso() -> str: return datetime.now(timezone.utc).isoformat() +class RuntimeContext(BaseModel): + """Allocation identity passed to a worker; never a client request payload.""" + + model_config = ConfigDict(frozen=True, extra='forbid') + + namespace: str + registry_name: str + instance_id: str + + class SessionRecord(BaseModel): """Represents a client session.""" @@ -35,6 +45,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): @@ -57,32 +70,10 @@ class FutureFailureRecord(BaseModel): diagnostic: str | None = None -# Canonical set of protocol-independent failure reason codes. The Twinkle and Tinker -# gateways each map this same key set to their own wire vocabularies; keeping the set -# in one place lets a consistency test catch a map that drifts out of coverage. -FAILURE_REASON_CODES: frozenset[str] = frozenset({ - 'invalid_request', - 'request_rejected', - 'resource_not_found', - 'full_mode_busy', - 'input_tokens_exceeded', - 'batch_size_invalid', - 'rate_limit_exceeded', - 'resource_quota_exceeded', - 'cancelled', - 'endpoint_unavailable', - 'state_contention', - 'backend_gate_unavailable', - 'orphaned_replica', - 'execution_timeout', - 'deadline_exceeded', - 'internal_error', -}) - - class FutureRecord(BaseModel): """Represents an async task future / request status.""" + 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 dd926aa57..01cf54607 100644 --- a/src/twinkle/server/state/server_state.py +++ b/src/twinkle/server/state/server_state.py @@ -10,7 +10,7 @@ from typing import Any from twinkle.server.config.persistence import PersistenceConfig -from twinkle.server.exceptions import ResourceQuotaExceededError +from twinkle.server.exceptions import ResourceNotFoundError, ResourceQuotaExceededError from twinkle.server.telemetry.correlation import (BASE_MODEL, MODEL_ID, REPLICA_ID, SAMPLING_SESSION_ID, SESSION_ID, TOKEN_ID) from twinkle.server.telemetry.tracing import traced_operation @@ -24,6 +24,7 @@ from .models import FutureFailureRecord, ModelRecord, SamplingSessionRecord, SessionRecord from .sampling_manager import SamplingSessionManager from .session_manager import SessionManager +from .training_store import TrainingSequences, TrainingStore logger = get_logger() @@ -131,6 +132,7 @@ def __init__( self._model_mgr = ModelManager(self._backend, expiration_timeout, per_token_model_limit) self._sampling_mgr = SamplingSessionManager(self._backend, expiration_timeout) self._future_mgr = FutureManager(self._backend, expiration_timeout) + self.training_sequences = TrainingSequences(self._backend) self.expiration_timeout = expiration_timeout self.cleanup_interval = cleanup_interval @@ -173,6 +175,33 @@ 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 models(self) -> ModelManager: + return self._model_mgr + + @property + def sessions(self) -> SessionManager: + return self._session_mgr + + @property + def futures(self) -> FutureManager: + return self._future_mgr + + @property + def sampling_sessions(self) -> SamplingSessionManager: + return self._sampling_mgr + + def create_training_store(self, namespace: str) -> TrainingStore: + """Compose a namespaced repository; construction performs no I/O.""" + return TrainingStore( + namespace, + self._backend, + models=self.models, + sessions=self.sessions, + futures=self.futures, + sequences=self.training_sequences, + expiration_timeout=self.expiration_timeout) + # ----- Session Management ----- async def create_session(self, payload: dict[str, Any]) -> str: @@ -213,6 +242,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, @@ -429,7 +461,13 @@ 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: + 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') + + 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 +484,17 @@ 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: + return await self.training_sequences.claim(training_id, stream, sequence, request_id, retention) + + async def release_training_seq(self, training_id: str, stream: str, sequence: int, request_id: str) -> None: + await self.training_sequences.release(training_id, stream, sequence, request_id) + + async def try_start_future(self, request_id: str) -> bool: + """Acquire task execution atomically, without inferring ownership from reads.""" + 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. @@ -470,6 +519,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. @@ -501,6 +551,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 ----- @@ -613,7 +664,7 @@ def get_server_state(cache_key: str = 'twinkle_server_state', # launcher-propagated env (so a non-gateway worker honours the operator's # YAML instead of the hardcoded default), else the module default. from twinkle.server.config.application_spec import ServerStateArgs - env_policy = ServerStateArgs.from_env() + env_policy = ServerStateArgs.from_env() or ServerStateArgs() def _resolve(explicit, env_value, default): if explicit is not None: @@ -622,13 +673,11 @@ def _resolve(explicit, env_value, default): return env_value return default - expiration_timeout = _resolve(expiration_timeout, getattr(env_policy, 'expiration_timeout', None), - _DEFAULT_EXPIRATION_TIMEOUT) - cleanup_interval = _resolve(cleanup_interval, getattr(env_policy, 'cleanup_interval', None), - _DEFAULT_CLEANUP_INTERVAL) - per_token_model_limit = _resolve(per_token_model_limit, getattr(env_policy, 'per_token_model_limit', None), + expiration_timeout = _resolve(expiration_timeout, env_policy.expiration_timeout, _DEFAULT_EXPIRATION_TIMEOUT) + cleanup_interval = _resolve(cleanup_interval, env_policy.cleanup_interval, _DEFAULT_CLEANUP_INTERVAL) + per_token_model_limit = _resolve(per_token_model_limit, env_policy.per_token_model_limit, _DEFAULT_PER_TOKEN_MODEL_LIMIT) - metrics_update_interval = _resolve(metrics_update_interval, getattr(env_policy, 'metrics_update_interval', None), + metrics_update_interval = _resolve(metrics_update_interval, env_policy.metrics_update_interval, _DEFAULT_METRICS_UPDATE_INTERVAL) state = ServerState( diff --git a/src/twinkle/server/state/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/state/training_records.py b/src/twinkle/server/state/training_records.py new file mode 100644 index 000000000..87e7aa693 --- /dev/null +++ b/src/twinkle/server/state/training_records.py @@ -0,0 +1,80 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""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 typing_extensions import TypedDict + +from .models import ModelRecord, RuntimeContext + +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 TrainingOperation(TypedDict): + """Admission facts retained until actual backend completion is confirmed.""" + + accepted_at: float + renew: bool + + +class RuntimeOwner(TypedDict): + instance_id: str + epoch: str + + +class TrainingRoute(RuntimeOwner): + base_model: str + + +class TrainingRecord(ModelRecord): + training_id: str + controller_epoch: str + 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, TrainingOperation] = 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 + + +class TrainingRuntimeContext(RuntimeContext): + controller_name: str + epoch: str diff --git a/src/twinkle/server/state/training_store.py b/src/twinkle/server/state/training_store.py new file mode 100644 index 000000000..0920095d3 --- /dev/null +++ b/src/twinkle/server/state/training_store.py @@ -0,0 +1,156 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Training persistence and retry receipts, composed from shared state managers.""" +from __future__ import annotations + +import functools +import hashlib +import json +import time +from typing import Any + +from twinkle.server.exceptions import RequestRejectedError +from .backend.base import StateBackend +from .future_manager import FutureManager +from .model_manager import ModelManager +from .models import FutureFailureRecord +from .session_manager import SessionManager +from .training_records import CreateTraining, InstanceRecord, TrainingRecord + + +class TrainingStore: + + def __init__(self, namespace: str, backend: StateBackend, *, models: ModelManager, sessions: SessionManager, + futures: FutureManager, sequences: TrainingSequences, expiration_timeout: float): + self.backend = backend + self.models, self.sessions, self.futures, self.sequences = models, sessions, futures, sequences + self.expiration_timeout = expiration_timeout + self.prefix = f'training::{namespace}::' + + def key(self, kind: str, identity: str) -> str: + 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: + await self.models.add(record.training_id, record) + + async def save_training(self, record: TrainingRecord) -> None: + await self.models.save_record(record.training_id, record) + + async def retire(self, record: TrainingRecord) -> None: + await self.backend.set( + self.key('terminal', record.training_id), record.model_dump(), ttl=max(1, int(self.expiration_timeout))) + await self.sequences.delete(record.training_id) + # Release the quota row last so a failed earlier step can be idempotently retried. + await self.models.remove(record.training_id) + + async def save_instance(self, record: InstanceRecord) -> None: + ttl = max(1, int(self.expiration_timeout)) if record.state == 'STOPPED' else None + await self.backend.set(self.key('instance', record.instance_id), record.model_dump(), ttl=ttl) + + async def instances(self) -> list[InstanceRecord]: + 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]: + return await self.models.get_training_records(self.prefix) + + async def get_create(self, identity: str) -> dict[str, str] | None: + return await self.backend.get(self.key('create', identity)) + + async def save_create(self, identity: str, fingerprint: str, request_id: str, session_id: str) -> None: + 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.sessions.get_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, + owner_token: str | None = None) -> None: + await self.futures.store_status( + 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()}) + + +class TrainingSequences: + """Atomic retry receipts and an expired-sequence fence, scoped to a training handle. + + Expiration compacts receipts but retains the sequence floor until that handle is + retired. An old update therefore cannot be applied twice after its receipt ages out. + """ + + def __init__(self, backend: StateBackend): + self.backend = backend + + async def claim(self, training_id: str, stream: str, sequence: int, request_id: str, + retention: float) -> str | None: + ledger = await self.backend.update_atomic( + f'seq_ledger::{training_id}::{stream}', + functools.partial( + self.claim_transform, sequence=sequence, request_id=request_id, now=time.time(), retention=retention)) + receipt = ledger['receipts'].get(str(sequence)) + if receipt is None: + raise RequestRejectedError('Sequence retry window expired; the operation will not be executed again') + return None if receipt['request_id'] == request_id else receipt['request_id'] + + async def release(self, training_id: str, stream: str, sequence: int, request_id: str) -> None: + await self.backend.update_atomic( + f'seq_ledger::{training_id}::{stream}', + functools.partial(self.release_transform, sequence=sequence, request_id=request_id)) + + async def delete(self, training_id: str) -> None: + for key in await self.backend.keys(f'seq_ledger::{training_id}::*'): + await self.backend.delete(key) + + @staticmethod + def claim_transform(current, *, sequence, request_id, now, retention): + floor = current.get('floor') if current else None + receipts = {} + for key, receipt in (current or {}).get('receipts', {}).items(): + if receipt['expires_at'] <= now: + floor = int(key) if floor is None else max(floor, int(key)) + else: + receipts[key] = receipt + key = str(sequence) + # Unexpired receipts remain retryable even if another sequence expired first. + if key not in receipts and (floor is None or sequence > floor): + receipts[key] = {'request_id': request_id, 'expires_at': now + retention} + return {'floor': floor, 'receipts': receipts} + + @staticmethod + def release_transform(current, *, sequence, request_id): + if current is None: + return None + receipts = dict(current['receipts']) + key = str(sequence) + if receipts.get(key, {}).get('request_id') == request_id: + receipts.pop(key) + return {'floor': current['floor'], 'receipts': receipts} diff --git a/src/twinkle/server/task_errors.py b/src/twinkle/server/task_errors.py index cd94d7b70..3c98140eb 100644 --- a/src/twinkle/server/task_errors.py +++ b/src/twinkle/server/task_errors.py @@ -2,10 +2,56 @@ """Helpers for constructing protocol-layer ``ErrorPayload`` values.""" from __future__ import annotations +from dataclasses import dataclass +from functools import lru_cache from typing import Any from twinkle.protocol.types.errors import ErrorCategory, ErrorPayload + +@dataclass(frozen=True) +class FailureResponse: + """How a domain failure is presented inside a response, not its HTTP status.""" + + error_code: int + category: ErrorCategory + + @classmethod + def for_reason(cls, reason: str) -> FailureResponse: + policies = cls._policies() + return policies.get(reason, policies['internal_error']) + + @classmethod + def reasons(cls) -> frozenset[str]: + return frozenset(cls._policies()) + + @classmethod + @lru_cache(maxsize=1) + def _policies(cls) -> dict[str, FailureResponse]: + return { + 'capacity_unavailable': cls(error_code=409, category=ErrorCategory.User), + 'startup_failed': cls(error_code=500, category=ErrorCategory.Server), + 'startup_timeout': cls(error_code=504, category=ErrorCategory.Server), + 'invalid_request': cls(error_code=400, category=ErrorCategory.User), + 'request_rejected': cls(error_code=400, category=ErrorCategory.User), + 'resource_not_found': cls(error_code=404, category=ErrorCategory.User), + 'full_mode_busy': cls(error_code=409, category=ErrorCategory.User), + 'input_tokens_exceeded': cls(error_code=422, category=ErrorCategory.User), + 'batch_size_invalid': cls(error_code=422, category=ErrorCategory.User), + 'rate_limit_exceeded': cls(error_code=429, category=ErrorCategory.User), + 'resource_quota_exceeded': cls(error_code=429, category=ErrorCategory.User), + 'cancelled': cls(error_code=499, category=ErrorCategory.User), + 'endpoint_unavailable': cls(error_code=501, category=ErrorCategory.Server), + 'state_contention': cls(error_code=503, category=ErrorCategory.Server), + 'backend_gate_unavailable': cls(error_code=503, category=ErrorCategory.Server), + 'orphaned_replica': cls(error_code=503, category=ErrorCategory.Server), + 'backend_cancelled': cls(error_code=503, category=ErrorCategory.Server), + 'execution_timeout': cls(error_code=504, category=ErrorCategory.Server), + 'deadline_exceeded': cls(error_code=500, category=ErrorCategory.Server), + 'internal_error': cls(error_code=500, category=ErrorCategory.Server), + } + + _ERROR_MAX = 1024 _TRACEBACK_MAX = 65536 _TRUNCATION_MARKER = '...[traceback truncated, tail kept]...\n' diff --git a/src/twinkle/server/task_queue/__init__.py b/src/twinkle/server/task_queue/__init__.py index bf5391611..7f4805db5 100644 --- a/src/twinkle/server/task_queue/__init__.py +++ b/src/twinkle/server/task_queue/__init__.py @@ -5,22 +5,20 @@ 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 twinkle.protocol.types.errors import QueueState +from twinkle.protocol.types.lifecycle import TaskStatus from .mixin import TaskQueueMixin from .rate_limiter import RateLimiter -from .types import QueuedTask, QueueState, TaskStatus, UserTaskError +from .types import QueuedTask from .worker import ComputeWorker __all__ = [ 'TaskStatus', 'QueueState', - 'UserTaskError', 'QueuedTask', - 'TaskQueueConfig', 'TaskQueueMixin', 'ComputeWorker', 'RateLimiter', diff --git a/src/twinkle/server/task_queue/backend_gate.py b/src/twinkle/server/task_queue/backend_gate.py index 82d17c1e0..2570a1e10 100644 --- a/src/twinkle/server/task_queue/backend_gate.py +++ b/src/twinkle/server/task_queue/backend_gate.py @@ -20,9 +20,13 @@ import functools from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor +from ray.exceptions import GetTimeoutError from typing import Any -from .types import BackendBusyError +from twinkle.server.exceptions import BackendBusyError + +# Backend execution has the same deadline semantics for asyncio and Ray waits. +BACKEND_TIMEOUT_EXCEPTIONS = (TimeoutError, GetTimeoutError) class BackendGate: diff --git a/src/twinkle/server/task_queue/mixin.py b/src/twinkle/server/task_queue/mixin.py index f26ec9c5c..eae14de97 100644 --- a/src/twinkle/server/task_queue/mixin.py +++ b/src/twinkle/server/task_queue/mixin.py @@ -9,12 +9,13 @@ import asyncio import time -import traceback import uuid from collections.abc import Callable, Coroutine from typing import TYPE_CHECKING, Any -from twinkle.protocol.types.lifecycle import TERMINAL_STATUSES, TaskEnvelope +from twinkle.protocol.types.errors import QueueState +from twinkle.protocol.types.lifecycle import TERMINAL_STATUSES, TaskEnvelope, TaskStatus +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.exceptions import BatchSizeError, ConfigError, InputTokensExceededError, RateLimitExceededError from twinkle.server.lifecycle.envelope import envelope_from_record from twinkle.server.lifecycle.poll_config import long_poll_window @@ -22,9 +23,8 @@ 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 .types import ExecutionScope, QueuedTask from .worker import ComputeWorker if TYPE_CHECKING: @@ -48,9 +48,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 ------------ @@ -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: @@ -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,33 @@ 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 self._compute_worker.finish_task(task, 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 + self._compute_worker.ensure_accepting() await self.state.store_future_status( request_id, TaskStatus.PENDING.value, @@ -237,36 +265,43 @@ async def _schedule_task( queue_state=QueueState.ACTIVE.value, replica_id=getattr(self, 'replica_id', None), absolute_deadline=self._future_deadline(), + owner_token=task.token, ) - queue_key = self._queue_key(model_id=model_id, token=token) + 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: + 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) - 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(), - )) - 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"}, ' - f'model_id={model_id}, queue_key={queue_key}, ' - f'queue_depth={q.qsize()}, input_tokens={input_tokens}') - - self._compute_worker.notify_new_task() - - return {'request_id': request_id, 'model_id': model_id} + self._compute_worker.track_task(request_id) + await q.put(task) + try: + 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, @@ -318,6 +353,7 @@ async def schedule_concurrent_task( input_tokens: int = 0, task_type: str | None = None, request_id: str | None = None, + execution_scope: ExecutionScope | None = None, ) -> dict[str, Any]: """Schedule a compute task that runs now instead of in serial order. @@ -342,7 +378,8 @@ async def schedule_concurrent_task( token=token, input_tokens=input_tokens, task_type=task_type, - created_at=time.monotonic()), + created_at=time.monotonic(), + execution_scope=execution_scope), self._queue_key(model_id=model_id, token=token), replica_id=getattr(self, 'replica_id', None), absolute_deadline=self._future_deadline()) @@ -419,68 +456,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]: - """Schedule a fire-and-forget task outside the serial queue. - - 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. + """Run pure I/O outside serial compute, with the same supervised lifecycle. - 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, 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) @@ -493,13 +476,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/types.py b/src/twinkle/server/task_queue/types.py index 936a433a3..abfb23809 100644 --- a/src/twinkle/server/task_queue/types.py +++ b/src/twinkle/server/task_queue/types.py @@ -1,61 +1,32 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -""" -Core type definitions for the task queue system. +"""In-process execution records, never HTTP/Ray wire models. -Provides: -- TaskStatus: Enum for tracking task lifecycle states -- QueueState: Enum for tinker client compatibility (retry behavior hints) -- QueuedTask: Dataclass representing a queued work item +Coroutine factories, asyncio tasks, and callbacks belong here. Public task and +queue state enums live in twinkle.protocol.types and are shared by SDK/server. """ from __future__ import annotations import asyncio from collections.abc import Callable, Coroutine -from dataclasses import dataclass -from enum import Enum - -# Ray_Get_Timeout is classified the same as asyncio.TimeoutError: 504/Server. -try: - from ray.exceptions import GetTimeoutError as _RayGetTimeout - BACKEND_TIMEOUT_EXCEPTIONS: tuple[type[BaseException], ...] = (asyncio.TimeoutError, _RayGetTimeout) -except Exception: # pragma: no cover - ray always present in server runtime - BACKEND_TIMEOUT_EXCEPTIONS = (asyncio.TimeoutError, ) - - -class TaskStatus(Enum): - """Task lifecycle status.""" - PENDING = 'pending' # Task created, waiting to be processed - QUEUED = 'queued' # Task in queue waiting for execution - RUNNING = 'running' # Task currently executing - COMPLETED = 'completed' # Task completed successfully - FAILED = 'failed' # Task failed with error - CANCELLED = 'cancelled' # Task cancelled by the client before it started running - +from dataclasses import dataclass, field -class UserTaskError(ValueError): - """A queued operation rejected because of caller input or usage.""" - - -class BackendBusyError(RuntimeError): - """Raised when the per-replica Admission_Gate is held by a leaked backend call. - - A new backend call arriving while the gate is closed (its holder is a call that - already exceeded ``asyncio.wait_for`` but whose executor thread has not yet - returned) fails fast with this error instead of queueing behind it. The worker - maps it to ``ErrorPayload(category='server', error_code=503)``. - """ +@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) -class QueueState(Enum): - """Queue state for tinker client compatibility. + def track(self, job: asyncio.Task, request_id: str | None = None) -> None: + self.jobs.add(job) + if request_id is not None: + self.request_ids.add(request_id) - These states are returned to the tinker client to indicate the current - state of the task queue and help the client adjust its retry behavior. - """ - ACTIVE = 'active' # Queue is actively processing tasks - PAUSED_RATE_LIMIT = 'paused_rate_limit' # Queue paused due to rate limiting - PAUSED_CAPACITY = 'paused_capacity' # Queue paused due to capacity limits - UNKNOWN = 'unknown' # Unknown or unspecified state + async def wait(self): + observed = set() + while pending := self.jobs - observed: + await asyncio.gather(*pending, return_exceptions=True) + observed.update(pending) @dataclass @@ -68,3 +39,5 @@ class QueuedTask: input_tokens: int task_type: str | None created_at: float + on_finished: Callable[[bool, bool], Coroutine] | None = None + execution_scope: ExecutionScope | None = None diff --git a/src/twinkle/server/task_queue/worker.py b/src/twinkle/server/task_queue/worker.py index cf684fe79..a19cac88e 100644 --- a/src/twinkle/server/task_queue/worker.py +++ b/src/twinkle/server/task_queue/worker.py @@ -16,18 +16,20 @@ from ray.exceptions import RayTaskError from typing import TYPE_CHECKING, Any, Callable, Deque -from twinkle.protocol.types.errors import ErrorCategory -from twinkle.server.exceptions import (BatchSizeError, EndpointUnavailableError, FullModeBusyError, +from twinkle.protocol.types.errors import ErrorCategory, QueueState +from twinkle.protocol.types.lifecycle import TaskStatus +from twinkle.server.config.task_queue import TaskQueueConfig +from twinkle.server.exceptions import (BackendBusyError, BatchSizeError, EndpointUnavailableError, FullModeBusyError, InputTokensExceededError, RateLimitExceededError, RequestRejectedError, ResourceExhaustedError, ResourceNotFoundError, ResourceQuotaExceededError, StateBackendError, TrainModeMismatchError, TwinkleServerError) from twinkle.server.state.backend.base import ConcurrencyError from twinkle.server.state.models import FutureFailureRecord +from twinkle.server.task_queue.backend_gate import BACKEND_TIMEOUT_EXCEPTIONS from twinkle.server.telemetry.correlation import MODEL_ID, TOKEN_ID from twinkle.server.telemetry.tracing import traced_operation from twinkle.utils.logger import get_logger -from .config import TaskQueueConfig -from .types import BACKEND_TIMEOUT_EXCEPTIONS, BackendBusyError, QueuedTask, QueueState, TaskStatus, UserTaskError +from .types import QueuedTask if TYPE_CHECKING: from twinkle.server.state import ServerState @@ -107,6 +109,8 @@ def __init__( self._on_backend_timeout = on_backend_timeout 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() @@ -119,14 +123,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 @@ -150,10 +159,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.""" @@ -171,6 +181,40 @@ 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 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()) @@ -268,6 +312,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() + await self.finish_task(task, False) self.task_queues.pop(queue_key, None) try: @@ -294,6 +339,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() + 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.""" @@ -314,7 +360,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: @@ -329,7 +376,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 @@ -375,10 +423,17 @@ async def run(): running = asyncio.create_task(run()) self._detached_tasks[running] = task + scope = task.execution_scope + if scope is not None: + scope.track(running, task.request_id) 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) @@ -391,6 +446,8 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: exec_start = time.monotonic() 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 = { @@ -403,28 +460,20 @@ 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, ' 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. 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( @@ -460,16 +509,6 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: await self._on_backend_timeout() except Exception: logger.error(f'[ComputeWorker] backend-timeout probe failed:\n{traceback.format_exc(limit=3)}') - except UserTaskError as exc: - task_status = 'failed' - exec_time = time.monotonic() - exec_start - await self._store_task_failed( - task, - f'{type(exc).__name__}: {exc}', - QueueState.UNKNOWN.value, - reason_code='request_rejected', - attribution='user', - ) except BackendBusyError as exc: task_status = 'failed' exec_time = time.monotonic() - exec_start @@ -514,6 +553,7 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: traceback_text=traceback.format_exc(), ) finally: + 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) @@ -546,13 +586,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() - continue - # Execute the task (serial: stops after the first execution) await self._execute_task(task, queue_key, q) return True @@ -561,11 +594,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/telemetry/metrics.py b/src/twinkle/server/telemetry/metrics.py index 586277708..0564fa473 100644 --- a/src/twinkle/server/telemetry/metrics.py +++ b/src/twinkle/server/telemetry/metrics.py @@ -17,9 +17,11 @@ from __future__ import annotations +from dataclasses import dataclass from pydantic import BaseModel, ConfigDict from typing import Any +from twinkle.server.exceptions import ResourceReleaseError from .provider import get_meter try: @@ -111,6 +113,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 +206,46 @@ 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='Transitions to quarantine during cleanup; retries in quarantine are excluded') + self.training_sweep_total = meter.create_counter( + 'twinkle.training.sweep.total', description='Training lifecycle sweep runs by status') + self.training_sweep_seconds = meter.create_histogram( + '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 +261,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 +298,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 +366,62 @@ 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() + + +@dataclass(frozen=True) +class TrainingInstanceSnapshot: + """Known resource ownership, including draining/quarantined instances.""" + + backend: str + mode: str + state: str + resident: str + device_type: str + device_count: int + placement_group_count: int + + +class TrainingInstanceMetrics: + """Aggregate snapshots and clear retired label sets; never own scheduling state.""" + + def __init__(self, metrics: TrainingMetrics): + self.metrics = metrics + self._instance_keys = set() + self._resource_keys = set() + + def report(self, snapshots): + counts, resources = {}, {} + for snapshot in snapshots: + key = snapshot.backend, snapshot.mode, snapshot.state, snapshot.resident + counts[key] = counts.get(key, 0) + 1 + resource_key = snapshot.backend, snapshot.mode, snapshot.device_type + devices, groups = resources.get(resource_key, (0, 0)) + resources[resource_key] = devices + snapshot.device_count, groups + snapshot.placement_group_count + for key in self._instance_keys | counts.keys(): + tags = dict(zip(('backend', 'mode', 'state', 'resident'), key)) + self.metrics.instances_active.set(counts.get(key, 0), tags=tags) + for key in self._resource_keys | resources.keys(): + tags = dict(zip(('backend', 'mode', 'device_type'), key)) + devices, groups = resources.get(key, (0, 0)) + self.metrics.devices_allocated.set(devices, tags=tags) + self.metrics.placement_groups_active.set(groups, tags=tags) + self._instance_keys = set(counts) + self._resource_keys = set(resources) + + @staticmethod + def quarantine_reason(exc: Exception) -> str: + if isinstance(exc, TimeoutError): + return 'timeout' + if isinstance(exc, ResourceReleaseError): + return exc.phase + return 'other' diff --git a/src/twinkle/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/src/twinkle_client/auto/agent/server_tools.py b/src/twinkle_client/auto/agent/server_tools.py index 09786c274..c2b031cc6 100644 --- a/src/twinkle_client/auto/agent/server_tools.py +++ b/src/twinkle_client/auto/agent/server_tools.py @@ -368,13 +368,13 @@ 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: 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 +409,44 @@ 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, + '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' - }, - }, - }, - }], - })) + 'train_mode': 'lora', + 'nproc_per_node': train_gpus, + 'device_group': { + 'name': 'model', + 'ranks': train_gpus, + 'device_type': 'cuda' + }, + 'device_mesh': { + 'device_type': 'cuda', + 'dp_size': train_gpus + }, + 'min_replicas': 0, + 'max_replicas': 1, + 'max_length': 500000, + 'idle_timeout_s': 600, + 'queue_config': { + 'rps_limit': 100, + 'tps_limit': 100000, + 'max_input_tokens': 500000 + }, + }, + }) # 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..efbdd2d60 100644 --- a/src/twinkle_client/model/multi_lora_transformers.py +++ b/src/twinkle_client/model/multi_lora_transformers.py @@ -3,9 +3,10 @@ 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 +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from peft import LoraConfig @@ -13,6 +14,7 @@ from twinkle.protocol.types import model as model_types from twinkle.protocol.types.component import DataRef from twinkle_client._request_builder import build_request +from twinkle_client.exceptions import TwinkleHTTPError from twinkle_client.http import ClientTransport from twinkle_client.http.context import capture_transport @@ -33,7 +35,7 @@ class MultiLoraTransformersModel: """Client wrapper for TwinkleModel that calls server HTTP endpoints. This client manages adapters and sends training/inference requests to the model server. - The server-side session (managed by TwinkleClient) keeps the model alive. + The session authenticates requests; executed training work renews the idle deadline. Every method builds its endpoint's request model rather than a dict, so a misspelled or wrongly-typed argument fails here -- in the caller's own stack trace, @@ -47,19 +49,20 @@ def __init__( model_id: str, *, transport: ClientTransport | None = None, - **kwargs, ): """Initialize a model wrapper bound to one immutable request identity.""" self._transport = capture_transport(transport) - kwargs.pop('data_plane_url', None) if '://' in model_id: model_id = model_id.split('://')[1] 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 + # the server dedups on (training_id, session, operation, seq_id), so each update is # applied at most once. Reserved once per call and reused on retry. self._seq_counter = itertools.count(1) self._seq_lock = threading.Lock() @@ -84,13 +87,39 @@ 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 + _SEQUENCED_ENDPOINTS = frozenset({ + 'forward_backward', + 'forward_backward_from_data_plane', + 'backward', + 'step', + 'lr_step', + 'clip_grad_and_step', + }) + + def _training_values(self, values: dict[str, Any]) -> dict[str, Any]: + """Choose the label and its handle together, under one lock.""" + with self._training_lock: + name = values.get('adapter_name', self.adapter_name) + training_id = values.get('training_id') or self._training_ids.get(name) + if training_id is None: + raise RuntimeError('Create this training with add_adapter_to_model on this object before using it.') + return {**values, 'adapter_name': name, 'training_id': training_id} + def _submit(self, endpoint: str, model_cls, response_cls, **values): - """Build, send, and resolve one twinkle-native request.""" - self._ensure_created() + """Validate once, then send with a stable handle and retry identity.""" + if issubclass(model_cls, model_types.TrainingRequest): + values = self._training_values(values) + if endpoint in self._SEQUENCED_ENDPOINTS: + if 'seq_id' in values: + raise TypeError('seq_id is managed by the client') + values['seq_id'] = self._next_seq_id() body = build_request(model_cls, **values) + self._ensure_created() response = self._transport.post_model(f'{self.server_url}/{endpoint}', body) return self._await_task(response, response_cls) @@ -132,22 +161,50 @@ 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.') + 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( + '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) + try: + self._submit('remove_adapter', model_types.AdapterRequest, None, adapter_name=name, training_id=training_id) + except TwinkleHTTPError as exc: + if exc.status_code == 404: + self._forget_training(name, training_id) + # A missing handle is unusable, but does not certify device release. + raise + self._forget_training(name, training_id) + + def _forget_training(self, name: str, training_id: str) -> None: + """A late response must never erase a newer same-name binding.""" + with self._training_lock: + if self._training_ids.get(name) == training_id: + self._training_ids.pop(name, None) + if name == self.adapter_name: + self.adapter_name = None # ------------------------------------------------------------------ # # Inline forward family @@ -155,23 +212,12 @@ def remove_adapter(self, adapter_name: str | None = None) -> None: def forward(self, inputs: Any, **kwargs) -> model_types.ForwardResponse: """Execute forward pass on inline model inputs.""" - return self._submit( - 'forward', - model_types.ForwardRequest, - model_types.ForwardResponse, - inputs=inputs, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + return self._submit('forward', model_types.ForwardRequest, model_types.ForwardResponse, inputs=inputs, **kwargs) def forward_only(self, inputs: Any, **kwargs) -> model_types.ForwardResponse: """Execute forward pass without gradient computation on inline inputs.""" return self._submit( - 'forward_only', - model_types.ForwardOnlyRequest, - model_types.ForwardResponse, - inputs=inputs, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + 'forward_only', model_types.ForwardOnlyRequest, model_types.ForwardResponse, inputs=inputs, **kwargs) def forward_backward(self, inputs: Any, **kwargs) -> model_types.ForwardBackwardResponse: """Execute combined forward and backward pass on inline inputs.""" @@ -180,37 +226,20 @@ def forward_backward(self, inputs: Any, **kwargs) -> model_types.ForwardBackward model_types.ForwardBackwardTaskRequest, model_types.ForwardBackwardResponse, inputs=inputs, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - seq_id=self._next_seq_id(), **kwargs) def calculate_loss(self, **kwargs) -> model_types.CalculateLossResponse: """Calculate loss from model outputs.""" - return self._submit( - 'calculate_loss', - model_types.AdapterRequest, - model_types.CalculateLossResponse, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + return self._submit('calculate_loss', model_types.AdapterRequest, model_types.CalculateLossResponse, **kwargs) def get_train_configs(self, **kwargs) -> model_types.GetTrainConfigsResponse: """Get training configs.""" - return self._submit( - 'get_train_configs', - model_types.AdapterRequest, - model_types.GetTrainConfigsResponse, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + return self._submit('get_train_configs', model_types.AdapterRequest, model_types.GetTrainConfigsResponse, + **kwargs) def backward(self, **kwargs) -> None: """Execute backward pass.""" - self._submit( - 'backward', - model_types.AdapterRequest, - None, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - seq_id=self._next_seq_id(), - **kwargs) + self._submit('backward', model_types.AdapterRequest, None, **kwargs) # ------------------------------------------------------------------ # # Data-plane forward family @@ -230,7 +259,6 @@ def forward_from_data_plane( model_types.DataPlaneForwardRequest, model_types.ForwardResponse, input_refs=_data_refs(inputs), - adapter_name=kwargs.pop('adapter_name', self.adapter_name), input_field=input_field, kwarg_fields=kwarg_fields or {}, **kwargs) @@ -251,7 +279,6 @@ def forward_only_from_data_plane( model_types.DataPlaneForwardOnlyRequest, model_types.ForwardResponse, input_refs=_data_refs(inputs), - adapter_name=kwargs.pop('adapter_name', self.adapter_name), input_field=input_field, kwarg_fields=kwarg_fields or {}, output_ref=output_ref.model_dump() if output_ref is not None else None, @@ -275,10 +302,8 @@ def forward_backward_from_data_plane( model_types.DataPlaneForwardRequest, model_types.ForwardBackwardResponse, input_refs=_data_refs(inputs), - adapter_name=kwargs.pop('adapter_name', self.adapter_name), input_field=input_field, kwarg_fields=kwarg_fields or {}, - seq_id=self._next_seq_id(), **kwargs) # ------------------------------------------------------------------ # @@ -287,32 +312,15 @@ def forward_backward_from_data_plane( def step(self, **kwargs) -> None: """Execute optimizer step.""" - self._submit( - 'step', - model_types.StepRequest, - None, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - seq_id=self._next_seq_id(), - **kwargs) + self._submit('step', model_types.StepRequest, None, **kwargs) def zero_grad(self, **kwargs) -> None: """Zero out gradients.""" - self._submit( - 'zero_grad', - model_types.AdapterRequest, - None, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('zero_grad', model_types.AdapterRequest, None, **kwargs) def lr_step(self, **kwargs) -> None: """Execute learning rate scheduler step.""" - self._submit( - 'lr_step', - model_types.LrStepRequest, - None, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - seq_id=self._next_seq_id(), - **kwargs) + self._submit('lr_step', model_types.LrStepRequest, None, **kwargs) def clip_grad_norm(self, max_grad_norm: float = 1.0, @@ -323,7 +331,6 @@ def clip_grad_norm(self, 'clip_grad_norm', model_types.ClipGradNormRequest, model_types.ClipGradNormResponse, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), max_grad_norm=max_grad_norm, norm_type=norm_type, **kwargs) @@ -334,10 +341,8 @@ def clip_grad_and_step(self, max_grad_norm: float = 1.0, norm_type: int = 2, **k 'clip_grad_and_step', model_types.ClipGradAndStepRequest, None, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), max_grad_norm=max_grad_norm, norm_type=norm_type, - seq_id=self._next_seq_id(), **kwargs) # ------------------------------------------------------------------ # @@ -346,33 +351,15 @@ def clip_grad_and_step(self, max_grad_norm: float = 1.0, norm_type: int = 2, **k def set_loss(self, loss_cls: str, **kwargs) -> None: """Set the loss function.""" - self._submit( - 'set_loss', - model_types.SetLossRequest, - None, - loss_cls=loss_cls, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('set_loss', model_types.SetLossRequest, None, loss_cls=loss_cls, **kwargs) def set_optimizer(self, optimizer_cls: str, **kwargs) -> None: """Set the optimizer.""" - self._submit( - 'set_optimizer', - model_types.SetOptimizerRequest, - None, - optimizer_cls=optimizer_cls, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('set_optimizer', model_types.SetOptimizerRequest, None, optimizer_cls=optimizer_cls, **kwargs) def set_lr_scheduler(self, scheduler_cls: str, **kwargs) -> None: """Set the learning rate scheduler.""" - self._submit( - 'set_lr_scheduler', - model_types.SetLrSchedulerRequest, - None, - scheduler_cls=scheduler_cls, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('set_lr_scheduler', model_types.SetLrSchedulerRequest, None, scheduler_cls=scheduler_cls, **kwargs) def set_template(self, template_cls: str, **kwargs) -> None: """Set the template for data processing. @@ -382,44 +369,20 @@ def set_template(self, template_cls: str, **kwargs) -> None: effect. A caller that passes it explicitly still reaches the template constructor through the passthrough region. """ - self._submit( - 'set_template', - model_types.SetTemplateRequest, - None, - template_cls=template_cls, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('set_template', model_types.SetTemplateRequest, None, template_cls=template_cls, **kwargs) def set_processor(self, processor_cls: str, **kwargs) -> None: """Set the input processor.""" - self._submit( - 'set_processor', - model_types.SetProcessorRequest, - None, - processor_cls=processor_cls, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('set_processor', model_types.SetProcessorRequest, None, processor_cls=processor_cls, **kwargs) def add_metric(self, metric_cls: str, is_training: bool | None = None, **kwargs) -> None: """Add a metric to the model.""" self._submit( - 'add_metric', - model_types.AddMetricRequest, - None, - metric_cls=metric_cls, - is_training=is_training, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + 'add_metric', model_types.AddMetricRequest, None, metric_cls=metric_cls, is_training=is_training, **kwargs) def apply_patch(self, patch_cls: str, **kwargs) -> None: """Apply a patch to the model.""" - self._submit( - 'apply_patch', - model_types.ApplyPatchRequest, - None, - patch_cls=patch_cls, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('apply_patch', model_types.ApplyPatchRequest, None, patch_cls=patch_cls, **kwargs) def calculate_metric(self, is_training: bool = True, **kwargs) -> model_types.CalculateMetricResponse: """Calculate metrics from model outputs.""" @@ -428,7 +391,6 @@ def calculate_metric(self, is_training: bool = True, **kwargs) -> model_types.Ca model_types.CalculateMetricRequest, model_types.CalculateMetricResponse, is_training=is_training, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), **kwargs) # ------------------------------------------------------------------ # @@ -437,23 +399,11 @@ def calculate_metric(self, is_training: bool = True, **kwargs) -> model_types.Ca def save(self, name: str, **kwargs) -> model_types.SaveResponse: """Save model checkpoint.""" - return self._submit( - 'save', - model_types.SaveRequest, - model_types.SaveResponse, - name=name, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + return self._submit('save', model_types.SaveRequest, model_types.SaveResponse, name=name, **kwargs) def load(self, name: str, **kwargs) -> None: """Load model checkpoint.""" - self._submit( - 'load', - model_types.LoadRequest, - None, - name=name, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('load', model_types.LoadRequest, None, name=name, **kwargs) def resume_from_checkpoint(self, name: str, *, resume_only_model: bool = False, **kwargs) -> dict[str, Any]: """Resume weights (and optionally optimizer state) from a checkpoint.""" @@ -462,7 +412,6 @@ def resume_from_checkpoint(self, name: str, *, resume_only_model: bool = False, model_types.ResumeFromCheckpointRequest, model_types.TrainingProgressResponse, name=name, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), resume_only_model=resume_only_model, **kwargs) return progress.result @@ -472,8 +421,6 @@ def upload_to_hub( checkpoint_dir: str, hub_model_id: str, hub_token: str | None = None, - async_upload: bool = True, - poll_interval: float = 5.0, ) -> None: """Upload model checkpoint to hub. @@ -484,10 +431,6 @@ def upload_to_hub( checkpoint_dir: The directory path of the checkpoint to upload. hub_model_id: The hub model id. hub_token: The hub token (optional). - async_upload: Deprecated, has no effect. The server always runs the - upload in the background and the client waits via the future layer. - poll_interval: Deprecated, has no effect. Pacing is now owned by the - server-side long-poll of the Retrieve_Endpoint. """ logger.info('[upload_to_hub] submitting upload, waiting for completion...') self._submit( diff --git a/tests/infra/test_ray_registry.py b/tests/infra/test_ray_registry.py new file mode 100644 index 000000000..f3b4c6f2e --- /dev/null +++ b/tests/infra/test_ray_registry.py @@ -0,0 +1,62 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Registry isolation must not reserve a CPU for ordinary non-server use.""" +import subprocess +import sys +import textwrap + + +def test_registry_cpu_reservations_and_instance_isolation(): + # Own a subprocess and cluster so the CPU assertions do not count actors + # belonging to other tests or reset their process-global RayHelper cache. + code = textwrap.dedent(''' + import os + import time + import uuid + import ray + from twinkle.infra._ray.ray_helper import RayHelper + + namespace = 'registry-' + uuid.uuid4().hex[:8] + ray.init(address='local', num_cpus=2, namespace=namespace, + include_dashboard=False, log_to_driver=False, + _temp_dir='/tmp/twreg-' + uuid.uuid4().hex[:6]) + + def registry(name, managed): + RayHelper._registry = None + os.environ['TWINKLE_RAY_REGISTRY_NAME'] = name + if managed: + os.environ['TWINKLE_RAY_NAMESPACE'] = namespace + else: + os.environ.pop('TWINKLE_RAY_NAMESPACE', None) + RayHelper.init_registry() + actor = RayHelper._registry + ray.get(actor.add_config.remote('owner', name)) + return actor + + def assert_free_cpus(expected): + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + available = ray.available_resources().get('CPU', 0) + if abs(available - expected) < 1e-6: + return + time.sleep(.01) + raise AssertionError((available, expected)) + + try: + ordinary = registry('ordinary', False) + assert_free_cpus(2) + ray.kill(ordinary, no_restart=True) + left = registry('left', True) + right = registry('right', True) + assert_free_cpus(1.98) + assert ray.get(left.get_config.remote('owner')) == 'left' + assert ray.get(right.get_config.remote('owner')) == 'right' + ray.kill(left, no_restart=True) + assert_free_cpus(1.99) + assert ray.get(right.get_config.remote('owner')) == 'right' + ray.kill(right, no_restart=True) + assert_free_cpus(2) + finally: + ray.shutdown() + ''') + result = subprocess.run([sys.executable, '-c', code], capture_output=True, text=True, timeout=120) + assert result.returncode == 0, result.stdout + result.stderr 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/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 99% rename from tests/server/utils/test_checkpoint_manager.py rename to tests/server/checkpoint/test_checkpoint_manager.py index e4d27eb85..b342d2935 100644 --- a/tests/server/utils/test_checkpoint_manager.py +++ b/tests/server/checkpoint/test_checkpoint_manager.py @@ -414,9 +414,9 @@ def write(path, data): @pytest.mark.asyncio async def test_reserved_sampler_name_is_rejected_before_backend_or_file_io(manager): from unittest.mock import AsyncMock - from twinkle.server.task_queue.types import UserTaskError + from twinkle.server.exceptions import RequestRejectedError save = AsyncMock() - with pytest.raises(UserTaskError, match='reserved'): + with pytest.raises(RequestRejectedError, match='reserved'): await manager.save_sampler('run', 'latest', save) save.assert_not_awaited() assert not Path(manager.get_save_dir('run', is_sampler=True)).exists() @@ -439,4 +439,3 @@ def snapshot(): assert snapshot() == before assert manager.delete('run', 'weights/step') is True assert not (run_dir / 'weights/step').exists() - 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/server_config_4b_e2e.yaml b/tests/server/config/server_config_4b_e2e.yaml index 6fc4d658f..e4b06ea67 100644 --- a/tests/server/config/server_config_4b_e2e.yaml +++ b/tests/server/config/server_config_4b_e2e.yaml @@ -32,6 +32,8 @@ applications: args: backend: transformers model_id: "ms://Qwen/Qwen3.5-4B" + min_replicas: 1 + max_replicas: 1 max_length: 10240 max_loras: 10 nproc_per_node: 2 diff --git a/tests/server/config/server_config_4b_e2e_megatron.yaml b/tests/server/config/server_config_4b_e2e_megatron.yaml index ab71f6bb5..1c83ddac9 100644 --- a/tests/server/config/server_config_4b_e2e_megatron.yaml +++ b/tests/server/config/server_config_4b_e2e_megatron.yaml @@ -32,6 +32,8 @@ applications: args: backend: megatron model_id: "ms://Qwen/Qwen3.5-4B" + min_replicas: 1 + max_replicas: 1 max_length: 10240 max_loras: 10 nproc_per_node: 4 diff --git a/tests/server/config/test_server_config.py b/tests/server/config/test_server_config.py index 927089ff2..9f507fa61 100644 --- a/tests/server/config/test_server_config.py +++ b/tests/server/config/test_server_config.py @@ -33,38 +33,26 @@ }), ) -_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_group': st.just({'ranks': 2, 'device_type': 'cuda'}), + 'device_mesh': st.just({'dp_size': 2}), + 'train_mode': st.sampled_from(['lora', 'full']), }), }) _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 +129,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/utils/task_queue/test_config.py b/tests/server/config/test_task_queue_config.py similarity index 98% rename from tests/server/utils/task_queue/test_config.py rename to tests/server/config/test_task_queue_config.py index 60787f567..511d6bde5 100644 --- a/tests/server/utils/task_queue/test_config.py +++ b/tests/server/config/test_task_queue_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/config/test_training_config.py b/tests/server/config/test_training_config.py new file mode 100644 index 000000000..a0bdb268b --- /dev/null +++ b/tests/server/config/test_training_config.py @@ -0,0 +1,363 @@ +# 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 + + +def model_config(mode='lora'): + devices = 2 if mode == 'lora' else 4 + return { + 'name': 'Qwen/small' if mode == 'lora' else 'Qwen/small-full', + 'model_id': 'ms://Qwen/small', 'backend': 'transformers', 'device_type': 'cuda', + 'mode': mode, 'config': {'num_devices': devices, 'device_mesh': {'fsdp_size': devices}}, + } + + +def model_application(): + return {'name': 'model-small', 'route_prefix': '/api/v1/model/Qwen/small', + 'import_path': 'model', 'args': { + 'model_id': 'ms://Qwen/small', 'backend': 'transformers', + 'device_group': {'ranks': 2, 'device_type': 'cuda'}, + 'device_mesh': {'fsdp_size': 2}, + }} + + +def test_catalog_defaults_and_mode_selection(): + config = TrainingConfig(models=[model_config(), model_config('full')]) + catalog = config + shared = catalog.get('Qwen/small', 'lora') + exclusive = catalog.get('Qwen/small-full', 'full') + assert shared.capacity == 5 + assert exclusive.capacity == 1 + assert shared.config.min_replicas == exclusive.config.min_replicas == 0 + assert shared.config.idle_timeout_s == 1800 + assert config.allocation_timeout_s == 3 + with pytest.raises(RequestRejectedError, match='Unknown training model'): + catalog.get('not-in-catalog', 'lora') + + +def test_existing_yaml_uses_explicit_instance_limits_and_business_idle(): + from pathlib import Path + paths = sorted(Path('cookbook/client/server').glob('**/server_config*.yaml')) + assert paths + for path in paths: + config = ServerConfig.from_yaml(path) + applications = [app for app in config.applications if app.import_path == 'model'] + for app, model in zip(applications, config.training_config.models): + assert model.config.idle_timeout_s == app.args.idle_timeout_s + assert model.config.min_replicas == app.args.min_replicas + assert model.config.max_replicas == app.args.max_replicas + assert model.config.adapter_config == app.args.adapter_config + + +@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['config'][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['config'] = { + 'num_devices': 8, + 'device_mesh': {'tp_size': 2, 'dp_size': 2, 'pp_size': 2, 'ep_size': 2, 'vpp_size': 2}, + } + assert TrainingConfig(models=[model]).models[0].config.num_devices == 8 + + +@pytest.mark.parametrize('mutation', ['duplicate', 'missing_mode', 'full_lora_field', 'hybrid_megatron', 'mesh_type', 'min_max']) +def test_cross_field_constraints(mutation): + model = model_config() + models = [model] + if mutation == 'duplicate': + models.append(copy.deepcopy(model)) + elif mutation == 'missing_mode': + del model['mode'] + elif mutation == 'full_lora_field': + model = model_config('full') + model['config']['max_loras'] = 2 + models = [model] + elif mutation == 'hybrid_megatron': + model['backend'] = 'megatron' + model['config']['hybrid'] = {'allocation_path': '/tmp/allocation.json'} + elif mutation == 'mesh_type': + model['config']['device_mesh']['device_type'] = 'npu' + elif mutation == 'min_max': + model['config'].update(min_replicas=2, max_replicas=1) + with pytest.raises(ValidationError): + TrainingConfig(models=models) + + +def test_full_only_is_explicit_and_lora_is_unavailable(): + model = model_config('full') + catalog = TrainingConfig(models=[model]) + assert catalog.get('Qwen/small-full', 'full').capacity == 1 + with pytest.raises(RequestRejectedError, match='does not provide lora'): + catalog.get('Qwen/small-full', '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.config + original = models[0] + expected_min = original.args.min_replicas + assert mode.min_replicas == expected_min + assert ServerConfig.model_validate(config.to_yaml_dict()).training_config == config.training_config + + +def test_configured_lifecycle_overrides_and_full_instances(): + app = {'name': 'model', 'route_prefix': '/api/v1/model/small', 'import_path': 'model', 'args': { + 'model_id': 'weights', 'backend': 'mock', 'train_mode': 'full', + 'device_group': {'ranks': 2, 'device_type': 'CPU'}, 'device_mesh': {'dp_size': 2}, + 'min_replicas': 0, 'max_replicas': 3, 'idle_timeout_s': 17, + }, 'deployments': [{'num_replicas': 1}]} + config = ServerConfig(applications=[app]) + mode = config.training_config.models[0].config + assert (mode.num_devices, mode.min_replicas, mode.max_replicas, mode.idle_timeout_s) == (2, 0, 3, 17) + assert config.training_config.models[0].mode == 'full' + + +def test_separate_model_applications_select_lora_and_full(): + shared = model_application() + exclusive = copy.deepcopy(shared) + exclusive.update(name='model-small-full', route_prefix='/api/v1/model/Qwen/small-full') + exclusive['args'].update(train_mode='full', min_replicas=0, max_replicas=2, idle_timeout_s=60) + config = ServerConfig(applications=[shared, exclusive]) + catalog = config.training_config + assert catalog.get('Qwen/small', 'lora').capacity == 5 + full = catalog.get('Qwen/small-full', 'full') + assert full.capacity == 1 + assert (full.config.min_replicas, full.config.max_replicas, full.config.idle_timeout_s) == (0, 2, 60) + with pytest.raises(RequestRejectedError, match='does not provide full'): + catalog.get('Qwen/small', 'full') + assert ServerConfig.model_validate(config.to_yaml_dict()).training_config == config.training_config + + +def test_launcher_keeps_training_and_sampling_capabilities_separate(monkeypatch): + 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'] + + +def configured_application(): + return {'name': 'model', 'route_prefix': '/api/v1/model/small', 'import_path': 'model', 'args': { + 'model_id': 'weights', 'backend': 'mock', 'train_mode': 'lora', 'max_length': 128, + 'device_group': {'ranks': 2, 'device_type': 'cpu'}, 'device_mesh': {'dp_size': 2}, + 'env_vars': {'USER_OPTION': 'args-wins'}, 'min_replicas': 1, 'max_replicas': 2, + }, 'deployments': [{ + 'name': 'ModelManagement', 'max_ongoing_requests': 17, + 'autoscaling_config': {'min_replicas': 1, 'max_replicas': 2, 'target_ongoing_requests': 16}, + 'ray_actor_options': {'num_cpus': 2, 'runtime_env': { + 'env_vars': {'USER_OPTION': 'deployment', 'OTHER_OPTION': 'retained'}, 'pip': ['example==1.0'] + }}, + }]} + + +@pytest.mark.parametrize('extra', [ + {'modes': None}, + {'modes': {'lora': {'num_devices': 2, 'device_mesh': {'dp_size': 2}}}}, + {'device_type': 'cpu'}, +]) +def test_model_yaml_rejects_alternative_resource_and_mode_syntax(extra): + app = configured_application() + app['args'].update(extra) + with pytest.raises(ValidationError, match='Extra inputs are not permitted'): + ServerConfig(applications=[app]) + + +@pytest.mark.parametrize('mode', ['lora', 'full']) +def test_adapter_lifecycle_yaml_reaches_controller_and_round_trips(mode): + app = model_application() + policy = {'adapter_timeout': 30, 'adapter_max_lifetime': 36000} + app['args'].update(train_mode=mode, adapter_config=policy, idle_timeout_s=600) + config = ServerConfig(applications=[app]) + resolved = config.training_config.models[0].config + assert resolved.adapter_config.model_dump() == policy + assert resolved.idle_timeout_s == 600 + restored = ServerConfig.model_validate(config.to_yaml_dict()) + assert restored.training_config == config.training_config + + +@pytest.mark.parametrize('policy', [ + {'adapter_timeout': 0}, {'adapter_timeout': math.inf}, {'adapter_timeout': math.nan}, + {'adapter_max_lifetime': -1}, {'adapter_max_lifetime': math.inf}, {'unknown': 30}, +]) +def test_invalid_adapter_lifecycle_is_rejected_before_allocation(policy): + app = model_application() + app['args']['adapter_config'] = policy + with pytest.raises(ValidationError): + ServerConfig(applications=[app]) + + +def test_configured_options_reach_cpu_router_without_changing_yaml(monkeypatch): + from types import SimpleNamespace + from twinkle.server.launcher.server_launcher import ServerLauncher + from twinkle.server.model.pool.actor import TrainingControllerActor + import ray + from ray import serve + config = ServerConfig(applications=[configured_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.model.pool.router.build_training_router', router) + monkeypatch.setattr(serve, 'run', lambda *args, **kwargs: None) + ServerLauncher(config)._deploy_training() + assert captured['max_ongoing_requests'] == 17 + 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] == configured_application()['deployments'][0] + + +@pytest.mark.asyncio +async def test_configured_options_reach_runtime_with_one_controller_owned_replica(monkeypatch): + import importlib + from twinkle.server.state.training_records import InstanceRecord + from twinkle.server.model.pool.driver import RayInstanceDriver + from ray import serve + config = ServerConfig(applications=[configured_application()]) + model = config.training_config.models[0] + driver = RayInstanceDriver(config.training_config, None, 'test', 'controller') + instance = InstanceRecord(instance_id='instance', controller_epoch='epoch', base_model=model.name, + 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.config.get(model.name, 'lora')) + options = captured['deploy_options'] + assert options['num_replicas'] == 1 and 'autoscaling_config' not in options + assert options['max_ongoing_requests'] == 17 + 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.config.queue_config + assert 'adapter_config' not in captured, 'Controller lifecycle policy must not reach model constructors' + assert not driver.runtime.deployments + + +def test_mutated_configuration_does_not_reuse_a_stale_catalog(): + config = ServerConfig(applications=[configured_application()]) + assert config.training_config.models[0].config.min_replicas == 1 + config.applications[0].args.min_replicas = 0 + assert config.training_config.models[0].config.min_replicas == 0 + + +def test_router_replica_limits_do_not_implicitly_set_training_instance_limits(): + app = configured_application() + app['args'].pop('min_replicas') + app['args'].pop('max_replicas') + model = ServerConfig(applications=[app]).training_config.models[0] + assert model.config.min_replicas == 0 and model.config.max_replicas is None + assert model.deploy_options['autoscaling_config']['min_replicas'] == 1 + + +def test_l20d_candidate_is_lora_only_and_bounds_total_devices(): + config = ServerConfig.from_yaml('cookbook/client/server/megatron/server_config.yaml') + model = config.training_config.models[0] + mode = model.config + assert model.mode == 'lora' + assert (mode.num_devices, mode.min_replicas, mode.max_replicas, mode.max_loras) == (2, 0, 2, 3) + assert mode.device_mesh.world_size() == 2 and mode.device_mesh.pp_size == 2 and mode.idle_timeout_s == 600 + sampler = next(app for app in config.applications if app.import_path == 'sampler') + 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 + with pytest.raises(RequestRejectedError, match='does not provide full'): + config.training_config.get(model.name, 'full') + + +@pytest.mark.parametrize('mode', ['lora', 'full']) +@pytest.mark.parametrize('yaml_entrypoint', [False, True]) +def test_single_device_megatron_is_rejected_before_allocation(mode, yaml_entrypoint): + if yaml_entrypoint: + app = configured_application() + app['args'].update(backend='megatron', train_mode=mode, device_group={'ranks': 1, 'device_type': 'cuda'}, + device_mesh={'dp_size': 1}) + else: + model = model_config(mode) + model['backend'] = 'megatron' + model['config'] = {'num_devices': 1, 'device_mesh': {'dp_size': 1}} + with pytest.raises(ValidationError, match='at least 2 devices'): + if yaml_entrypoint: + ServerConfig(applications=[app]) + else: + TrainingConfig(models=[model]) + + +def test_transformers_still_accepts_single_device(): + app = model_application() + app['args'].update(device_group={'ranks': 1, 'device_type': 'cuda'}, device_mesh={'dp_size': 1}) + assert ServerConfig(applications=[app]).training_config.models[0].config.num_devices == 1 diff --git a/tests/server/contract/test_client_api_contract.py b/tests/server/contract/test_client_api_contract.py index f543ee436..89d9fe00a 100644 --- a/tests/server/contract/test_client_api_contract.py +++ b/tests/server/contract/test_client_api_contract.py @@ -66,7 +66,7 @@ def test_new_client_types_importable(): for symbol in ('StrictRequest', 'ResponseModel', 'DataModel', 'backend_only'): assert hasattr(base, symbol) - for symbol in ('ErrorPayload', 'ErrorCategory', 'QueueStateLiteral'): + for symbol in ('ErrorPayload', 'ErrorCategory', 'QueueState'): assert hasattr(errors, symbol) diff --git a/tests/server/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/fixtures/server_config_mock.yaml b/tests/server/fixtures/server_config_mock.yaml index 67303ed5c..2a7711062 100644 --- a/tests/server/fixtures/server_config_mock.yaml +++ b/tests/server/fixtures/server_config_mock.yaml @@ -42,6 +42,8 @@ applications: args: backend: mock model_id: mock-model + min_replicas: 1 + max_replicas: 1 nproc_per_node: 1 device_group: name: model diff --git a/tests/server/fixtures/server_config_mock_redis.yaml b/tests/server/fixtures/server_config_mock_redis.yaml index 4c7297fe3..59c2ab436 100644 --- a/tests/server/fixtures/server_config_mock_redis.yaml +++ b/tests/server/fixtures/server_config_mock_redis.yaml @@ -41,6 +41,8 @@ applications: args: backend: mock model_id: mock-model + min_replicas: 1 + max_replicas: 1 nproc_per_node: 1 device_group: name: model diff --git a/tests/server/integration/e2e_helpers.py b/tests/server/integration/e2e_helpers.py index 2eca95f0d..8ac1c6905 100644 --- a/tests/server/integration/e2e_helpers.py +++ b/tests/server/integration/e2e_helpers.py @@ -290,3 +290,34 @@ def assert_metrics_valid(metrics: Any, label: str) -> None: assert math.isfinite(val), ( f'[{label}] Metric {key}={val} is not finite!') log(f'[{label}] Metrics valid: {metrics}') + + +def sampled_datum(prompt_tokens, sample, advantage): + from tinker import types + + tokens = prompt_tokens + list(sample.tokens) + prompt_positions = len(prompt_tokens) - 1 + size = len(tokens) - 1 + + def tensor(values, dtype): + return types.TensorData(data=values, dtype=dtype, shape=[size]) + + return types.Datum( + model_input=types.ModelInput.from_ints(tokens[:-1]), + loss_fn_inputs={ + 'target_tokens': tensor(tokens[1:], 'int64'), + 'weights': tensor([0.0] * prompt_positions + [1.0] * len(sample.tokens), 'float32'), + 'logprobs': tensor([0.0] * prompt_positions + list(sample.logprobs), 'float32'), + 'advantages': tensor([0.0] * prompt_positions + [advantage] * len(sample.tokens), 'float32'), + }, + ) + + +def completion_logprobs(client, datum, timeout): + result = client.forward([datum], 'cross_entropy').result(timeout=timeout) + values = np.asarray(result.loss_fn_outputs[0]['logprobs'].tolist(), dtype=np.float64) + weights = np.asarray(datum.loss_fn_inputs['weights'].tolist()) + assert len(values) >= len(weights) + values = values[:len(weights)][weights != 0] + assert values.size and np.isfinite(values).all() + return values diff --git a/tests/server/integration/on_demand_training/__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..c14827c6c --- /dev/null +++ b/tests/server/integration/on_demand_training/acceptance_support.py @@ -0,0 +1,171 @@ +# 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 + + +def checked_logps(values): + """Reject empty/default metric values before numerical checkpoint comparisons.""" + import numpy as np + 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 + + +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] + return checked_logps(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 + # 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'] + 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/tests/server/integration/on_demand_training/http_probe.py b/tests/server/integration/on_demand_training/http_probe.py new file mode 100644 index 000000000..13efd3f6c --- /dev/null +++ b/tests/server/integration/on_demand_training/http_probe.py @@ -0,0 +1,276 @@ +# 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=8, + _temp_dir='/tmp/twh-' + namespace[-8:], + include_dashboard=False, + log_to_driver=False) + config = ServerConfig.model_validate({ + 'ray_namespace': + namespace, + 'http_options': { + 'host': '127.0.0.1', + 'port': port + }, + 'resources': { + 'startup_timeout_s': 600, + 'shutdown_timeout_s': 180 + }, + 'applications': [{ + 'name': 'small', + 'import_path': 'model', + 'route_prefix': '/api/v1/model/small', + 'args': { + 'model_id': 'small', + 'backend': 'mock', + 'train_mode': 'lora', + 'device_group': {'ranks': 1, 'device_type': 'cpu'}, + 'device_mesh': {'dp_size': 1}, + 'max_loras': 2, + 'max_replicas': 1, + }, + }, { + 'name': 'other', + 'import_path': 'model', + 'route_prefix': '/api/v1/model/other', + 'args': { + 'model_id': 'small', + 'backend': 'mock', + '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': 'routed', 'import_path': 'model', 'route_prefix': '/api/v1/model/routed', + 'args': { + 'model_id': 'small', 'backend': 'mock', 'train_mode': 'lora', + 'device_group': {'ranks': 1, 'device_type': 'cpu'}, 'device_mesh': {'dp_size': 1}, + 'max_loras': 1, 'max_replicas': 2, + }, + }, { + 'name': 'gateway', + 'import_path': 'server', + '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) + from peft import LoraConfig + + def assert_separate_instances(identities): + diagnostic = ray.get(controller.diagnostics.remote()) + owners = {identity: item['instance_id'] for item in diagnostic['instances'] + for identity in item['training_ids'] if identity in identities} + assert set(owners) == set(identities) and len(set(owners.values())) == 2, owners + return owners + + routed = [MultiLoraTransformersModel('routed', transport=transport) for _ in range(2)] + for client in routed: + client.add_adapter_to_model('same', LoraConfig(r=8, target_modules='all-linear')) + native_owners = assert_separate_instances([client._training_ids['same'] for client in routed]) + for _ in range(3): + for client in reversed(routed): + assert json.loads(client.get_train_configs().result) == {} + client.zero_grad() + for client in routed: + client.remove_adapter() + first.add_adapter_to_model('same', LoraConfig(r=8, target_modules='all-linear')) + instance = next( + item for item in ray.get(controller.diagnostics.remote())['instances'] + 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', LoraConfig(r=8, target_modules='all-linear')) + old_id = first._training_ids['same'] + assert old_id != second._training_ids['same'] + assert json.loads(first.get_train_configs().result) == {} + first.zero_grad() + first.remove_adapter() + first.add_adapter_to_model('same', LoraConfig(r=8, target_modules='all-linear')) + assert old_id != first._training_ids['same'] + first.remove_adapter() + second.remove_adapter() + 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') + routed_tinker = [service.create_lora_training_client(base_model='routed', rank=8) for _ in range(2)] + tinker_owners = assert_separate_instances([client.model_id for client in routed_tinker]) + for _ in range(3): + for client in reversed(routed_tinker): + assert client.get_info().model_id == client.model_id + for client in routed_tinker: + assert unload_tinker_model(service, client.model_id).model_id == client.model_id + shared1 = service.create_lora_training_client(base_model='small', rank=8) + shared2 = service.create_lora_training_client(base_model='small', rank=8) + assert shared1.model_id != shared2.model_id + 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='other') + assert not exclusive.get_info().is_lora + assert exclusive.forward_backward([datum], 'cross_entropy').result().metrics + exclusive.optim_step(types.AdamParams(learning_rate=1e-4)).result() + 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', LoraConfig(r=8, target_modules='all-linear') + if index % 2 == 0 else None) + model.zero_grad() + model.remove_adapter() + print(json.dumps({'phase': 'cycles', 'completed': index + 1}), flush=True) + 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_stable_routing': native_owners, + 'tinker_stable_routing': tinker_owners, + 'native_lora': True, + 'tinker_shared': True, + 'tinker_full': True, + 'restart': True, + 'state_loss': True, + 'runtime_lifecycle_rejected': True, + 'cycles': 10, + 'instances': len(final['instances']) + })) + transport.close() + 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) + serve.shutdown() + ray.shutdown() + + +if __name__ == '__main__': + main() 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..b889d2d4b --- /dev/null +++ b/tests/server/integration/on_demand_training/matrix_validation.py @@ -0,0 +1,491 @@ +# 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: + # Two calls end the native GA=2 workload on an optimizer boundary. + handle.step() + 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) + + 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. + 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) + if protocol == 'twinkle': + 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} + 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, + 'optimizer_continuation': optimizer_errors} + + +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, full_model_name, capacity): + """Assert shared LoRA capacity and full-training exclusivity on real devices.""" + from tests.server.integration.on_demand_training.acceptance_support import TrainingHandle, assert_released + handles = [] + 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, full_model_name, 'twinkle', 'exclusive', mode='full') + try: + exclusive.step() + live = [row for row in ray.get(controller.diagnostics.remote())['instances'] if row['state'] != 'STOPPED'] + assert len(live) == 1 and live[0]['mode'] == 'full' and len(live[0]['training_ids']) == 1 + try: + overflow = TrainingHandle(base_url, full_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('--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') + 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.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] + 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('resources', {}).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} + full_applications = [] + for app in raw['applications']: + params = app['args'] + if app['import_path'] == 'model': + common = { + 'device_group': {'ranks': args.training_devices, 'device_type': 'cuda'}, + 'device_mesh': {'device_type': 'cuda', 'dp_size': args.training_devices // args.training_pp_size, + 'pp_size': args.training_pp_size}, + 'min_replicas': 0, 'max_replicas': 1, 'idle_timeout_s': 3600, + 'max_length': params.get('max_length'), + 'queue_config': {**params.get('queue_config', {}), 'execution_timeout': 600}, + } + app['args'] = { + **common, 'model_id': args.model, 'backend': args.backend, 'train_mode': 'lora', + 'max_loras': 10, 'max_r': 32, + 'env_vars': {'TWINKLE_TRUST_REMOTE_CODE': '1'}, + } + app.pop('deployments', None) + full_applications.append({ + 'name': app['name'] + '-full', + 'route_prefix': app['route_prefix'] + '-full', + 'import_path': 'model', + 'args': {**common, 'model_id': args.model, 'backend': args.backend, 'train_mode': 'full', + 'env_vars': app['args']['env_vars']}, + }) + elif app['import_path'] == 'sampler': + params['model_id'] = args.model + params['engine_args']['attention_config'] = {'flash_attn_version': 3} + raw['applications'].extend(full_applications) + if args.recovery_only: + raw['applications'] = [app for app in raw['applications'] if app['import_path'] in ('server', 'model')] + config = ServerConfig.model_validate(raw) + 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({model.name: model.model_id + for model in config.training_config.models}), + 'TWINKLE_TRUST_REMOTE_CODE': '1', + '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=args.training_devices + (0 if args.recovery_only else 1), + _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: + 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([ + 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) == 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'] = [] + 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'] + 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, + next(model.name for model in config.training_config.models if model.mode == 'full'), + config.training_config.models[0].config.max_loras) + if args.stress_iterations or args.soak_seconds: + from tests.server.integration.on_demand_training.acceptance_support import assert_released + controller = launcher._training_controller + 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/tests/server/integration/on_demand_training/on_demand_sdk.py b/tests/server/integration/on_demand_training/on_demand_sdk.py new file mode 100644 index 000000000..d8e30f8aa --- /dev/null +++ b/tests/server/integration/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) 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/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..a7879000f --- /dev/null +++ b/tests/server/integration/on_demand_training/test_resident_release_e2e.py @@ -0,0 +1,190 @@ +# 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.model.pool 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}, + 'resources': {'startup_timeout_s': 1800, 'shutdown_timeout_s': 180, 'sweep_interval_s': 1}, + 'persistence': ({'mode': persistence, 'redis_url': managed.url, 'key_prefix': prefix} + if managed else {'mode': persistence}), + 'applications': [ + {'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, 'train_mode': 'full', + 'device_group': {'ranks': 2, 'device_type': 'cuda'}, + 'min_replicas': 1, 'max_replicas': 1, + 'device_mesh': {'dp_size': 2 if backend == 'transformers' else 1, + 'pp_size': 1 if backend == 'transformers' else 2}, + 'max_length': 512, 'idle_timeout_s': 3600, + 'queue_config': {'execution_timeout': 600}, + 'env_vars': {'TWINKLE_TRUST_REMOTE_CODE': '1'}}}, + ], + } + (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_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/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/resources/allocation_probe.py b/tests/server/integration/resources/allocation_probe.py new file mode 100644 index 000000000..9d057d215 --- /dev/null +++ b/tests/server/integration/resources/allocation_probe.py @@ -0,0 +1,69 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Run in its own process: test production allocation against logical Ray GPUs.""" +import asyncio +import json +import os +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.resources.allocation import CapacityUnavailableError, DevicePool + + +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=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: + for index, count in enumerate((2, 2, 4)): + config = TrainingConfig(models=[{'name': str(index), 'model_id': str(index), 'backend': 'mock', 'device_type': 'cuda', 'mode': 'full', 'config': {'num_devices': count, 'device_mesh': {'fsdp_size': count}}}]) + spec = config.get(str(index), 'full') + allocation = await pool.allocate(spec.config.num_devices, spec.device_type, 'epoch', str(index)) + allocations.append(allocation) + devices = [device for allocation in allocations for device in allocation.devices] + assert len(set(devices)) == 8, devices + assert len({allocation.node_id for allocation in allocations}) == 1 + try: + await pool.allocate(spec.config.num_devices, spec.device_type, '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') + 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 get_placement_group(allocations[0].name).id == owned_pg.id + freed = allocations[0] + await pool.remove(freed.name) + two = TrainingConfig(models=[{'name': 'replacement', 'model_id': 'replacement', 'backend': 'mock', 'device_type': 'cuda', 'mode': 'full', 'config': {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}}}]).get('replacement', 'full') + replacement = await pool.allocate(two.config.num_devices, two.device_type, 'epoch', 'replacement') + allocations.append(replacement) + assert set(replacement.devices) == set(freed.devices) + print(json.dumps({'passed': True, 'devices': devices, 'reused': list(replacement.devices)})) + finally: + try: + for allocation in allocations: + await pool.remove(allocation.name) + finally: + ray.shutdown() + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/tests/server/integration/resources/fragmentation_probe.py b/tests/server/integration/resources/fragmentation_probe.py new file mode 100644 index 000000000..f684c51c3 --- /dev/null +++ b/tests/server/integration/resources/fragmentation_probe.py @@ -0,0 +1,73 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""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 +from ray.util.placement_group import placement_group, placement_group_table, remove_placement_group + +from twinkle.server.config.training import TrainingConfig +from twinkle.server.resources.allocation import CapacityUnavailableError, DevicePool + + +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, 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) + # 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 = TrainingConfig(models=[{'name': 'two', 'model_id': 'two', 'backend': 'mock', 'device_type': 'cuda', 'mode': 'full', 'config': {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}}}]).get('two', 'full') + for marker in ('left', 'right'): + pg = placement_group([{'GPU': 3, marker: 0.01}], strategy='STRICT_PACK') + occupied.append(pg) + await asyncio.wait_for(pg.ready(), 20) + assert ray.available_resources()['GPU'] == 2 + try: + await pool.allocate(spec.config.num_devices, spec.device_type, '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.config.num_devices, spec.device_type, '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/tests/server/integration/resources/test_allocation.py b/tests/server/integration/resources/test_allocation.py new file mode 100644 index 000000000..54d8f8b11 --- /dev/null +++ b/tests/server/integration/resources/test_allocation.py @@ -0,0 +1,31 @@ +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__).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]) + 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__).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]) + assert evidence['passed'] and evidence['nodes'] == 2 + assert evidence['fragmented_rejected'] and evidence['pending_removed'] + assert len(evidence['devices']) == 2 + + +def test_sampler_uses_real_cpu_services_for_scaling_and_release(): + script = Path(__file__).resolve().parents[1] / 'sampler' / 'sampling_probe.py' + result = subprocess.run([sys.executable, str(script)], capture_output=True, text=True, timeout=1800) + assert result.returncode == 0, result.stdout + result.stderr + evidence = [json.loads(line) for line in result.stdout.splitlines() if line.startswith('{"sizes":')] + assert evidence == [{'sizes': [2, 4, 2], 'requests': 8, 'confirmed_release': True}] diff --git a/tests/server/integration/sampler/__init__.py b/tests/server/integration/sampler/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/server/integration/sampler/sampling_probe.py b/tests/server/integration/sampler/sampling_probe.py new file mode 100644 index 000000000..6e32c6ba4 --- /dev/null +++ b/tests/server/integration/sampler/sampling_probe.py @@ -0,0 +1,108 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Exercise the actual sampler driver and mock compute service on an owned CPU cluster.""" +import asyncio +import json +import os +import uuid + +import ray +from ray import serve +from ray.util import list_named_actors +from ray.util.placement_group import placement_group_table + +from twinkle.server.config.application_spec import SamplerArgs +from twinkle.server.sampler.pool.controller import SamplerController +from twinkle.server.sampler.pool.driver import SamplerDriver + + +async def run(namespace): + args = SamplerArgs( + model_id='mock', sampler_type='mock', + device_group={'name': 'sampler', 'ranks': 1, 'device_type': 'cpu'}, + device_mesh={'device_type': 'cpu', 'dp_size': 1}, + min_replicas=2, max_replicas=4, target_ongoing_requests=1, idle_timeout_s=1) + budgets = {'allocation_timeout_s': 10, 'startup_timeout_s': 600, + 'shutdown_timeout_s': 60, 'sweep_interval_s': 1} + driver = SamplerDriver(args, namespace, 'cpu-allocation', {}, budgets) + clock = [0] + pool = SamplerController(args, driver, budgets, clock=lambda: clock[0]) + # Hold replies after the actual compute route runs. The pool must retain + # their tickets until the caller can confirm backend completion. + invoke = driver.invoke + replies = 0 + dispatched, release = asyncio.Event(), asyncio.Event() + + async def delayed_reply(*args, **kwargs): + nonlocal replies + response = await invoke(*args, **kwargs) + replies += 1 + if replies == 8: + dispatched.set() + await release.wait() + return response + + driver.invoke = delayed_reply + calls = [] + try: + await pool.initialize() + sizes = [len(pool._ready_instances())] + payload = {'endpoint': 'twinkle/sample', 'method': 'POST', + 'headers': [(b'content-type', b'application/json'), + (b'authorization', b'Bearer cpu-probe'), + (b'x-request-id', b'cpu-probe')], 'query_string': b'', + 'body': json.dumps({'inputs': [{'input_ids': [1, 2]}], + 'sampling_params': {'max_tokens': 2}}).encode()} + calls = [asyncio.create_task(pool.request(payload)) for _ in range(8)] + await asyncio.wait_for(dispatched.wait(), 600) + await asyncio.gather(*tuple(pool.lifecycle.startups.values())) + sizes.append(len(pool._ready_instances())) + clock[0] = 10 + await pool.sweep() + assert pool._owned_count() == 4 + release.set() + responses = await asyncio.gather(*calls) + assert all(response['status'] == 200 for response in responses), responses + await pool.lifecycle.join() + from twinkle.server.state import get_server_state + state = get_server_state() + for response in responses: + record = await state.get_future(json.loads(response['body'])['request_id']) + assert record['status'] == 'completed', record + assert len(record['result']['samples'][0]['sequences'][0]['tokens']) == 2 + health = await asyncio.gather(*(driver.health(instance) for instance in pool._ready_instances())) + assert all(not result['unhealthy'] for result in health) + clock[0] = 12 + await pool.sweep() + await pool.lifecycle.join() + sizes.append(len(pool._ready_instances())) + assert sizes == [2, 4, 2], sizes + await pool.shutdown() + await 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)) + return {'sizes': sizes, 'requests': len(responses), 'confirmed_release': True} + finally: + release.set() + await asyncio.gather(*calls, return_exceptions=True) + await pool.shutdown() + await driver.runtime.wait_for_cleanup() + + +def main(): + namespace = 'sampling_probe_' + uuid.uuid4().hex[:8] + os.environ.update(TWINKLE_STATE_RAY_NAMESPACE=namespace, TWINKLE_PERSISTENCE_MODE='memory') + ray.init(address='local', namespace=namespace, num_cpus=24, num_gpus=0, + _temp_dir='/tmp/tws-' + namespace[-8:], include_dashboard=False) + try: + serve.start(http_options={'host': '127.0.0.1', 'port': 0}) + print(json.dumps(asyncio.run(run(namespace))), flush=True) + finally: + serve.shutdown() + ray.shutdown() + + +if __name__ == '__main__': + main() 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..36b5d35fb --- /dev/null +++ b/tests/server/integration/sampler/test_real_pool_e2e.py @@ -0,0 +1,339 @@ +# 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 traceback +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.pool.controller import SamplerController + from twinkle.server.sampler.pool.driver import SamplerDriver + from twinkle.server.state import get_server_state + + args = SamplerArgs(model_id=str(model), sampler_type='vllm', + 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']) + 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(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(): + 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())] + source_root = str(Path(__file__).resolve().parents[4]) + + 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() + 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() + 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'] + 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_results'] = {'before_eviction': before, 'after_reload': after} + 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() + except BaseException: + evidence['failure'] = traceback.format_exc() + save() + raise + 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() 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_blocking_boundary.py b/tests/server/integration/test_blocking_boundary.py index 7d35e229d..7fb895993 100644 --- a/tests/server/integration/test_blocking_boundary.py +++ b/tests/server/integration/test_blocking_boundary.py @@ -21,7 +21,7 @@ from twinkle.server.task_queue.backend_gate import BackendGate # noqa: E402 from twinkle.server.task_queue.mixin import TaskQueueMixin # noqa: E402 -from twinkle.server.task_queue.types import BackendBusyError # noqa: E402 +from twinkle.server.exceptions import BackendBusyError # noqa: E402 class _Harness(TaskQueueMixin): diff --git a/tests/server/integration/test_full_cycle_e2e.py b/tests/server/integration/test_full_cycle_e2e.py index 6fa31f006..ec7a4c52d 100644 --- a/tests/server/integration/test_full_cycle_e2e.py +++ b/tests/server/integration/test_full_cycle_e2e.py @@ -6,13 +6,12 @@ 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 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 @@ -58,13 +57,13 @@ 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 -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 @@ -125,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 @@ -217,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', @@ -249,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 @@ -307,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/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..e377afff2 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 @@ -207,6 +146,16 @@ def test_mock_mode_reaches_ready_under_30s_and_is_deterministic(ray_cluster) -> # 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) @@ -257,20 +206,23 @@ 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') - 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') + 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. @@ -287,8 +239,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] @@ -299,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: @@ -339,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. @@ -352,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, @@ -362,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/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/integration/test_sampler_checkpoints_e2e.py b/tests/server/integration/test_sampler_checkpoints_e2e.py index e95858118..4affee9a3 100644 --- a/tests/server/integration/test_sampler_checkpoints_e2e.py +++ b/tests/server/integration/test_sampler_checkpoints_e2e.py @@ -8,7 +8,7 @@ import numpy as np import pytest -from tests.server.integration.test_tinker_rl_correctness_e2e import _completion_logprobs, _datum +from tests.server.integration.e2e_helpers import completion_logprobs, sampled_datum pytestmark = pytest.mark.skipif(os.environ.get('TWINKLE_TEST_GPU_E2E') != '1', reason='Requires a real trainer and vLLM sampler') @@ -58,15 +58,15 @@ def same(actual, expected): sampler_x = service.create_sampling_client(base_model=model_id, model_path=x) first = sample(sampler_x) same(sample(sampler_x), first) - datum = _datum(prompt, first, 1.0) - before = _completion_logprobs(client, datum, timeout) + datum = sampled_datum(prompt, first, 1.0) + before = completion_logprobs(client, datum, timeout) losses = [] for _ in range(3): result = client.forward_backward([datum], 'cross_entropy').result(timeout=timeout) losses.append(float(result.metrics['loss:mean'])) assert np.isfinite(losses[-1]) client.optim_step(types.AdamParams(learning_rate=1e-4)).result(timeout=timeout) - after = _completion_logprobs(client, datum, timeout) + after = completion_logprobs(client, datum, timeout) assert float(np.mean(after - before)) > 1e-4 y = client.save_weights_for_sampler(name='y').result(timeout=timeout).path diff --git a/tests/server/integration/test_tinker_rl_correctness_e2e.py b/tests/server/integration/test_tinker_rl_correctness_e2e.py index 925841dbe..5f3fc5a66 100644 --- a/tests/server/integration/test_tinker_rl_correctness_e2e.py +++ b/tests/server/integration/test_tinker_rl_correctness_e2e.py @@ -14,6 +14,8 @@ import numpy as np import pytest +from tests.server.integration.e2e_helpers import completion_logprobs, sampled_datum + pytestmark = pytest.mark.skipif( os.environ.get('TWINKLE_TEST_GPU_E2E', '0') != '1', reason='Requires a running training + vLLM server on real accelerators', @@ -83,50 +85,20 @@ def live_clients(): return service, model_id, prompt_tokens, samples[0], timeout -def _datum(prompt_tokens, sample, advantage): - from tinker import types - - tokens = prompt_tokens + list(sample.tokens) - prompt_positions = len(prompt_tokens) - 1 - size = len(tokens) - 1 - - def tensor(values, dtype): - return types.TensorData(data=values, dtype=dtype, shape=[size]) - - return types.Datum( - model_input=types.ModelInput.from_ints(tokens[:-1]), - loss_fn_inputs={ - 'target_tokens': tensor(tokens[1:], 'int64'), - 'weights': tensor([0.0] * prompt_positions + [1.0] * len(sample.tokens), 'float32'), - 'logprobs': tensor([0.0] * prompt_positions + list(sample.logprobs), 'float32'), - 'advantages': tensor([0.0] * prompt_positions + [advantage] * len(sample.tokens), 'float32'), - }, - ) - - -def _completion_logprobs(client, datum, timeout): - result = client.forward([datum], 'cross_entropy').result(timeout=timeout) - values = np.asarray(result.loss_fn_outputs[0]['logprobs'].tolist(), dtype=np.float64) - weights = np.asarray(datum.loss_fn_inputs['weights'].tolist()) - assert len(values) >= len(weights) - values = values[:len(weights)][weights != 0] - assert values.size and np.isfinite(values).all() - return values - def test_live_unsupported_losses_are_rejected_without_changing_policy(live_clients): from tinker import BadRequestError service, model_id, prompt, sample, timeout = live_clients client = service.create_lora_training_client(base_model=model_id, rank=16, seed=2026) - datum = _datum(prompt, sample, -1.0) - before = _completion_logprobs(client, datum, timeout) + datum = sampled_datum(prompt, sample, -1.0) + before = completion_logprobs(client, datum, timeout) for loss_fn in ('ppo', 'cispo', 'dro'): with pytest.raises(BadRequestError) as exc: client.forward_backward([datum], loss_fn).result(timeout=timeout) assert exc.value.status_code == 400 assert f"Unsupported Tinker loss_fn '{loss_fn}'" in str(exc.value) - after = _completion_logprobs(client, datum, timeout) + after = completion_logprobs(client, datum, timeout) np.testing.assert_allclose(after, before, rtol=0, atol=1e-5) _report(check='unsupported_loss', loss_fn=loss_fn, http_status=400, policy_unchanged=True) @@ -138,15 +110,15 @@ def test_live_importance_sampling_updates_policy_in_reward_direction(live_client service, model_id, prompt, sample, timeout = live_clients client = service.create_lora_training_client(base_model=model_id, rank=16, seed=2026) - datum = _datum(prompt, sample, advantage) - before = _completion_logprobs(client, datum, timeout) + datum = sampled_datum(prompt, sample, advantage) + before = completion_logprobs(client, datum, timeout) sampler_difference = float(np.mean(np.abs(before - np.asarray(sample.logprobs)))) assert sampler_difference < 0.05, f'Sampler/trainer logprob mismatch: {sampler_difference}' result = client.forward_backward([datum] * batch_size, 'importance_sampling').result(timeout=timeout) assert len(result.loss_fn_outputs) == batch_size assert np.isfinite(float(result.metrics['loss:mean'])) client.optim_step(types.AdamParams(learning_rate=1e-4)).result(timeout=timeout) - after = _completion_logprobs(client, datum, timeout) + after = completion_logprobs(client, datum, timeout) delta = float(np.mean(after) - np.mean(before)) _report( check='reward_direction', @@ -166,7 +138,7 @@ def test_live_dpo_still_requires_pairs(live_clients): service, model_id, prompt, sample, timeout = live_clients client = service.create_lora_training_client(base_model=model_id, rank=16, seed=2026) - datum = _datum(prompt, sample, 1.0) + datum = sampled_datum(prompt, sample, 1.0) reference = client.forward([datum], 'cross_entropy').result(timeout=timeout) datum.loss_fn_inputs['ref_logps'] = reference.loss_fn_outputs[0]['logprobs'] with pytest.raises(UnprocessableEntityError) as exc: diff --git a/tests/server/lifecycle/test_envelope.py b/tests/server/lifecycle/test_envelope.py index fdc588bd8..1dd368a94 100644 --- a/tests/server/lifecycle/test_envelope.py +++ b/tests/server/lifecycle/test_envelope.py @@ -12,12 +12,18 @@ def test_wire_maps_cover_exactly_the_canonical_reason_codes(): """The Twinkle and Tinker failure maps must both stay in lockstep with the canonical reason-code set, so a newly added domain reason cannot silently fall back to 500 on one protocol face.""" - from twinkle.server.gateway.tinker_handlers import _TINKER_FAILURE_WIRE - from twinkle.server.lifecycle.envelope import _FAILURE_WIRE - from twinkle.server.state.models import FAILURE_REASON_CODES - - assert set(_FAILURE_WIRE) == FAILURE_REASON_CODES - assert set(_TINKER_FAILURE_WIRE) == FAILURE_REASON_CODES + from twinkle.server.gateway.tinker_handlers import _tinker_error_from_failure + from twinkle.server.lifecycle.envelope import error_payload_from_failure + from twinkle.server.task_errors import FailureResponse + + for reason in FailureResponse.reasons(): + policy = FailureResponse.for_reason(reason) + failure = {'reason_code': reason, 'message': 'failed', 'attribution': policy.category.value} + native = error_payload_from_failure(failure, request_id='request') + assert native.error_code == policy.error_code + assert native.category is policy.category + assert _tinker_error_from_failure(failure, request_id='request') == native.model_dump( + mode='json', exclude_none=True) def test_completed_with_none_result_is_a_success_not_a_failure(): diff --git a/tests/server/lifecycle/test_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_retrieve_endpoint.py b/tests/server/lifecycle/test_retrieve_endpoint.py index 84b8767e4..704ea2f71 100644 --- a/tests/server/lifecycle/test_retrieve_endpoint.py +++ b/tests/server/lifecycle/test_retrieve_endpoint.py @@ -19,6 +19,9 @@ class _State: """A fake ServerState whose get_future returns a fixed record (or None).""" + async def assert_future_owner(self, request_id, token): + return None + def __init__(self, record): self._record = record diff --git a/tests/server/lifecycle/test_run_submit_dedup.py b/tests/server/lifecycle/test_run_submit_dedup.py index 81d9c840b..b5f7fe710 100644 --- a/tests/server/lifecycle/test_run_submit_dedup.py +++ b/tests/server/lifecycle/test_run_submit_dedup.py @@ -23,24 +23,25 @@ def __init__(self, record_after_claim): self.claim_ttls = [] self.released = [] - async def claim_seq(self, dedup_key, request_id, ttl): - # Unseen -> claim it and let the caller proceed (returns None). - self.claimed[dedup_key] = request_id - self.claim_ttls.append(ttl) - return None + async def claim_training_seq(self, training_id, stream, seq_id, request_id, retention): + key = (training_id, stream, seq_id) + prior = self.claimed.get(key) + self.claimed[key] = prior or request_id + self.claim_ttls.append(retention) + return prior async def get_future(self, request_id): # Simulates whether a PENDING/QUEUED record was written (task enqueued). return self._record_after_claim - async def release_seq(self, dedup_key): - self.released.append(dedup_key) + async def release_training_seq(self, training_id, stream, seq_id, request_id): + self.released.append((training_id, stream, seq_id)) class _FakeManagement: def __init__(self, record_after_claim): self.state = _FakeState(record_after_claim) - self.task_queue_config = SimpleNamespace(queue_timeout=300.0, effective_execution_timeout=1800.0) + self.sequence_retention_s = 3600 # A real deployment declares its backend; preflight reads it from here. self.backend = 'transformers' self.data_world_size = 1 @@ -63,7 +64,7 @@ def _body(adapter_name: str = 'ad', seq_id: int = 7) -> ForwardBackwardTaskReque ``run_submit`` now reads field roles off the body to build the backend kwargs and to run preflight, so a ``SimpleNamespace`` would exercise a shape production never sees. """ - return ForwardBackwardTaskRequest(inputs=[{'input_ids': [1, 2]}], adapter_name=adapter_name, seq_id=seq_id) + return ForwardBackwardTaskRequest(inputs=[{'input_ids': [1, 2]}], adapter_name=adapter_name, seq_id=seq_id, training_id='training-' + adapter_name) async def _call(self, body, adapter_name, token): # pragma: no cover - never invoked @@ -77,7 +78,7 @@ async def test_release_kept_when_task_already_enqueued(): with pytest.raises(RuntimeError): await run_submit(mgmt, _request(), _body(), task_type='forward_backward', backend_call=_call) assert mgmt.state.released == [], 'claim wrongly released for an already-enqueued task' - assert 'seq::sess-1::sess-1-ad::7' in mgmt.state.claimed + assert ('training-ad', 'native::sess-1::forward_backward', 7) in mgmt.state.claimed @pytest.mark.asyncio @@ -86,19 +87,16 @@ async def test_release_when_never_enqueued(): mgmt = _FakeManagement(record_after_claim=None) with pytest.raises(RuntimeError): await run_submit(mgmt, _request(), _body(), task_type='forward_backward', backend_call=_call) - assert mgmt.state.released == ['seq::sess-1::sess-1-ad::7'], 'claim should be released when nothing enqueued' + assert mgmt.state.released == [('training-ad', 'native::sess-1::forward_backward', 7)], 'claim should be released when nothing enqueued' @pytest.mark.asyncio -async def test_dedup_ttl_covers_queue_and_execution(): +async def test_dedup_uses_training_sequence_retention(): """The claim must outlive queue wait plus execution, not just execution.""" mgmt = _FakeManagement(record_after_claim={'status': 'queued'}) with pytest.raises(RuntimeError): await run_submit(mgmt, _request(), _body(), task_type='forward_backward', backend_call=_call) - config = mgmt.task_queue_config - assert mgmt.state.claim_ttls == [int(config.queue_timeout + config.effective_execution_timeout) + 60] - # An execution-only window is exactly what let a queued first copy expire. - assert mgmt.state.claim_ttls[0] > int(config.effective_execution_timeout) + 60 + assert mgmt.state.claim_ttls == [mgmt.sequence_retention_s] @pytest.mark.asyncio @@ -121,5 +119,5 @@ async def test_dedup_key_is_scoped_per_adapter(): backend_call=_call) claimed = set(mgmt.state.claimed) - assert claimed == {'seq::sess-1::sess-1-lora-A::1', 'seq::sess-1::sess-1-lora-B::1'}, ( + assert claimed == {('training-lora-A', 'native::sess-1::forward_backward', 1), ('training-lora-B', 'native::sess-1::forward_backward', 1)}, ( f'adapters collided on one dedup key: {claimed}') diff --git a/tests/server/lifecycle/test_static_guards.py b/tests/server/lifecycle/test_static_guards.py index 730c4ef7f..3e635a294 100644 --- a/tests/server/lifecycle/test_static_guards.py +++ b/tests/server/lifecycle/test_static_guards.py @@ -5,7 +5,6 @@ - ``TaskEnvelope`` has exactly one construction site. - The client-side invariant: the client HTTP timeout is <= 120 and strictly greater than the server Long_Poll_Window. -- The task status set has two independent declarations that must not drift. These exist because the spec states most of its guarantees in prose. A prose claim that nothing checks decays into a false claim -- as happened with "a consistency test asserts @@ -71,30 +70,6 @@ def test_task_envelope_has_exactly_one_construction_site(): assert sites, 'expected to find the construction sites inside envelope.py' -def test_server_task_status_enum_matches_client_literal(): - """The two independent declarations of the task status set must not drift. - - ``twinkle.protocol.types.lifecycle.TaskStatus`` (a Literal on the wire model) and the - server's ``TaskStatus`` enum are declared separately. ``envelope_from_record`` copies - ``record['status']`` straight into ``TaskEnvelope.status``, so a value the server can - write but the Literal does not list would fail pydantic validation *while serialising - the response* -- i.e. a 500 from retrieve for a task that actually finished. - - The client module's comment claimed such a test existed; it did not. This is it. - """ - from typing import get_args - - from twinkle.server.task_queue.types import TaskStatus as ServerTaskStatus - from twinkle.protocol.types.lifecycle import TERMINAL_STATUSES - from twinkle.protocol.types.lifecycle import TaskStatus as WireTaskStatus - - server_values = {member.value for member in ServerTaskStatus} - wire_values = set(get_args(WireTaskStatus)) - assert server_values == wire_values, (f'task status sets drifted: server-only={server_values - wire_values}, ' - f'wire-only={wire_values - server_values}') - assert TERMINAL_STATUSES <= wire_values, 'TERMINAL_STATUSES must be a subset of the declared statuses' - - def test_client_future_layer_is_not_imported_by_the_server(): """``_future.py`` carries an underscore because the dependency runs one way only. diff --git a/tests/server/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/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_full_mode_exclusivity.py b/tests/server/model/test_full_mode_exclusivity.py index a8d207b10..2c6ad3635 100644 --- a/tests/server/model/test_full_mode_exclusivity.py +++ b/tests/server/model/test_full_mode_exclusivity.py @@ -5,120 +5,100 @@ import asyncio import pytest -from fastapi import FastAPI +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock from starlette.requests import Request -import twinkle.protocol.types as types -from twinkle.server.exceptions import FullModeBusyError -from twinkle.server.model import twinkle_handlers -from twinkle.server.model.app import ModelManagement -from twinkle.server.model.twinkle_handlers import _register_model_twinkle_routes -from twinkle.server.session_resource.base import SessionResourceMixin +from tests.server.state.fake_backend import FakeBackend +from tests.server.training.helpers import request, setup +from twinkle.server.config.training import TrainingConfig +from twinkle.server.model.pool.router import TrainingRouter -class _ResourceStub: - """The resource-registry slice of ModelManagement, without Ray Serve.""" - - is_full_mode = True - _resource_type = 'Model' - find_active_resource = SessionResourceMixin.find_active_resource - register_resource = SessionResourceMixin.register_resource - unregister_resource = SessionResourceMixin.unregister_resource - set_resource_state = SessionResourceMixin.set_resource_state - _validate_registration = SessionResourceMixin._validate_registration - _create_resource_record = SessionResourceMixin._create_resource_record - assert_full_mode_available = ModelManagement.assert_full_mode_available - claim_full_mode_resource = ModelManagement.claim_full_mode_resource - - def __init__(self): - self._resource_records = {} +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'][1]['config']['max_replicas'] = 1 + controller.config = TrainingConfig.model_validate(raw) + runtime.start_gate.clear() + return controller, state, runtime @pytest.mark.asyncio async def test_concurrent_full_mode_claims_have_exactly_one_winner(): - stub = _ResourceStub() - - async def claim(name): - # The old check-then-register split had this await between the two steps. - await asyncio.sleep(0) - stub.claim_full_mode_resource(name, 'token', 'session') - - results = await asyncio.gather(claim('tenant-a'), claim('tenant-b'), return_exceptions=True) - assert sum(isinstance(result, FullModeBusyError) for result in results) == 1 - assert len(stub._resource_records) == 1 - - -def test_same_tenant_may_reclaim_its_own_resource(): - stub = _ResourceStub() - stub.claim_full_mode_resource('tenant-a', 'token', 'session') - stub.claim_full_mode_resource('tenant-a', 'token', 'session') - - with pytest.raises(FullModeBusyError): - stub.claim_full_mode_resource('tenant-b', 'token', 'session') - - -class _AsyncState: - - def __init__(self): - self.registered = [] - self.unloaded = [] - - async def register_model(self, config, *, token, model_id, replica_id, session_id): - # The exact window the old check-then-register split raced through. - await asyncio.sleep(0) - self.registered.append(model_id) - - async def unload_model(self, model_id): - self.unloaded.append(model_id) - - -class _RunManager: - - def save(self, adapter_name, run_config): - return None - - -class _FullModeManagement(_ResourceStub): - """The handler-level slice of ModelManagement needed by add_adapter_to_model.""" - - base_model = 'base-model' - replica_id = 'replica' - - def __init__(self): - super().__init__() - self.state = _AsyncState() - - async def _on_request_start(self, request): - return 'token' - - async def submit_and_peek(self, factory, **kwargs): - from twinkle.protocol.types.lifecycle import TaskEnvelope - return TaskEnvelope(request_id='req', status='completed', result=await factory()) - - -def _request(): - request = Request({'type': 'http', 'headers': []}) - request.state.session_id = 'session' - request.state.request_id = 'req' - return request + 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_concurrent_handler_creates_have_exactly_one_winner(monkeypatch): - """Two real handler calls, interleaved inside register_model, must not both win.""" - management = _FullModeManagement() - monkeypatch.setattr(twinkle_handlers, 'create_training_run_manager', lambda *args, **kwargs: _RunManager()) - app = FastAPI() - _register_model_twinkle_routes(app, lambda: management) - endpoint = next(route.endpoint for route in app.routes - if getattr(route, 'path', None) == '/twinkle/add_adapter_to_model') +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 - results = await asyncio.gather( - *(endpoint(_request(), types.AddAdapterRequest(adapter_name=name), management) - for name in ('adapter-a', 'adapter-b')), - return_exceptions=True) - assert sum(isinstance(result, FullModeBusyError) for result in results) == 1 - assert len(management._resource_records) == 1 - # The loser released the model record it wrote before the claim failed. - assert len(management.state.unloaded) == 1 +@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.config = controller.config + router.model = SimpleNamespace(name='small-full') + router.controller = SimpleNamespace(create=SimpleNamespace(remote=AsyncMock(side_effect=controller.create))) + + async def create_remote(payload): + 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_replica_lifecycle.py b/tests/server/model/test_replica_lifecycle.py index 84b961391..63b9ccaa8 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,26 @@ 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, 'initialize_training_runtime', new_callable=AsyncMock), \ + patch.object(ModelManagement, 'wait_for_training_workers', new_callable=AsyncMock): await ModelManagement.__init__( manager, model_id='model', @@ -80,6 +83,10 @@ async def test_async_constructor_registers_replica_before_ready() -> None: device_mesh={}, backend='mock', max_loras=3, + training_context={'instance_id': 'replica-1', 'epoch': 'epoch', 'namespace': 'test', + 'registry_name': 'test', 'controller_name': 'test'}, + placement_group=None, + allocation_devices=[], ) state.register_replica.assert_awaited_once_with('replica-1', 3) diff --git a/tests/server/model/test_tinker_handlers.py b/tests/server/model/test_tinker_handlers.py index 6daa6d566..2df46bb7e 100644 --- a/tests/server/model/test_tinker_handlers.py +++ b/tests/server/model/test_tinker_handlers.py @@ -9,8 +9,9 @@ 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 +from twinkle.server.model.runtime import TrainingExecutionState class _DummyManagement: @@ -190,6 +191,7 @@ class _SaveWeightsDummyManagement: is_full_mode = ModelManagement.is_full_mode def __init__(self): + self.training_state = TrainingExecutionState(grad_ready=True) self.model = MagicMock() self.state = MagicMock() self.state.get_model_metadata = AsyncMock(return_value={'base_model': 'test-model'}) @@ -198,8 +200,8 @@ def __init__(self): async def _on_request_start(self, request): return 'token1' - def get_adapter_name(self, adapter_name=None): - return adapter_name + def get_training_state(self, training_id): + return self.training_state resolve_model_adapter_name = ModelManagement.resolve_model_adapter_name @@ -269,7 +271,6 @@ async def test_load_weights_resolves_location_and_preserves_optimizer_flag(monke output_dir = manager.get_save_dir('model1') monkeypatch.setattr(handlers, 'create_checkpoint_manager', lambda *args, **kwargs: manager) management = _SaveWeightsDummyManagement() - management.set_resource_state = MagicMock() app = FastAPI() _register_model_tinker_routes(app, lambda: management) endpoint = next(r.endpoint for r in app.routes if getattr(r, 'path', '') == '/tinker/load_weights') @@ -278,7 +279,7 @@ async def test_load_weights_resolves_location_and_preserves_optimizer_flag(monke management.model.tinker_load.assert_called_once_with( checkpoint_name='step' if output_dir else path, output_dir=output_dir, load_optimizer=True, adapter_name='model1') - management.set_resource_state.assert_called_once_with('model1', 'grad_ready', False) + assert management.training_state.grad_ready is False assert response.path == path and response.type == 'load_weights' @@ -289,7 +290,6 @@ async def test_tinker_loss_metric_survives_sdk_reduction(backward): management = _SaveWeightsDummyManagement() management.data_world_size = 1 - management.set_resource_state = MagicMock() outputs = [{'logprobs': types.TensorData(data=[-0.25, -0.5], dtype='float32', shape=[2])}] management.model.tinker_forward_only.return_value = (outputs, 0.375) management.model.tinker_forward_backward.return_value = (outputs, 0.375) diff --git a/tests/server/model/test_twinkle_async_inputs.py b/tests/server/model/test_twinkle_async_inputs.py index 21cf17dec..38cefbe85 100644 --- a/tests/server/model/test_twinkle_async_inputs.py +++ b/tests/server/model/test_twinkle_async_inputs.py @@ -101,7 +101,7 @@ async def test_forward_backward_resolves_multiple_data_refs_and_field_kwargs() - kwarg_fields={ 'old_logps': 'sampled_logprobs', 'advantages': 'advantage', - }, + }, training_id='training-a', ) await route.endpoint(request, body, management) @@ -109,7 +109,7 @@ async def test_forward_backward_resolves_multiple_data_refs_and_field_kwargs() - assert management.scheduled[-1]['batch_size'] == 8 assert management.scheduled[-1]['data_world_size'] == 2 inputs, adapter_name, forwarded_kwargs = management.model_calls[-1] - assert adapter_name == 'session-adapter' + assert adapter_name == 'training-a' assert [row['input_ids'] for row in inputs] == [[index] for index in range(8)] assert forwarded_kwargs['old_logps'] == [[-0.1]] * 4 + [[-0.2]] * 4 assert forwarded_kwargs['advantages'] == [1.0] * 4 + [-1.0] * 4 @@ -139,13 +139,13 @@ async def test_forward_backward_binds_nested_dpo_ref_logps_without_coercion() -> body = types.DataPlaneForwardRequest( adapter_name='adapter', input_refs=[types.DataRef(ref_id='dpo', size=2, num_tokens=6)], - kwarg_fields={'ref_outputs.logps': 'ref_logps'}, + kwarg_fields={'ref_outputs.logps': 'ref_logps'}, training_id='training-a', ) await route.endpoint(request, body, management) inputs, adapter_name, forwarded_kwargs = management.model_calls[-1] - assert adapter_name == 'session-adapter' + assert adapter_name == 'training-a' assert [row['input_ids'] for row in inputs] == [[1, 2, 3], [1, 4, 5]] assert forwarded_kwargs['ref_outputs']['logps'] == [ [-0.1, -0.2, -0.3], @@ -177,7 +177,7 @@ def save(self, name, output_dir, **kwargs): monkeypatch.setattr('twinkle.server.checkpoint.training_run_manager.TWINKLE_DEFAULT_SAVE_DIR', str(tmp_path)) runs = TwinkleTrainingRunManager('token') - runs._write_info('session-adapter', {'base_model': 'test-model', 'is_lora': True}) + runs._write_info('training-a', {'base_model': 'test-model', 'is_lora': True}) manager = TwinkleCheckpointManager('token', runs) monkeypatch.setattr(handlers, 'create_checkpoint_manager', lambda *args, **kwargs: manager) service = Management() @@ -189,14 +189,14 @@ def save(self, name, output_dir, **kwargs): request.state.session_id = 'session' body = types.SaveRequest(adapter_name='adapter', name='step/one', is_sampler=is_sampler, - save_optimizer=save_optimizer, consumed_train_samples=7) + save_optimizer=save_optimizer, consumed_train_samples=7, training_id='training-a') response = await endpoint(request, body, service) weights_type = 'sampler_weights' if is_sampler else 'weights' - directory = manager.get_ckpt_dir('session-adapter', f'{weights_type}/step_one') - assert response.result == {'twinkle_path': f'twinkle://session-adapter/{weights_type}/step_one', + directory = manager.get_ckpt_dir('training-a', f'{weights_type}/step_one') + assert response.result == {'twinkle_path': f'twinkle://training-a/{weights_type}/step_one', 'checkpoint_dir': str(directory)} assert (directory / 'weights').read_text() == '1' - assert service.model_calls == [('step_one', {'adapter_name': '' if full_mode else 'session-adapter', + assert service.model_calls == [('step_one', {'adapter_name': '' if full_mode else 'training-a', 'save_optimizer': save_optimizer, 'consumed_train_samples': 7})] @@ -211,9 +211,9 @@ def save(self, **kwargs): monkeypatch.setattr('twinkle.server.checkpoint.training_run_manager.TWINKLE_DEFAULT_SAVE_DIR', str(tmp_path)) runs = TwinkleTrainingRunManager('token') - runs._write_info('session-adapter', {'base_model': 'test-model'}) + runs._write_info('training-a', {'base_model': 'test-model'}) manager = TwinkleCheckpointManager('token', runs) - before = runs._read_info('session-adapter') + before = runs._read_info('training-a') monkeypatch.setattr(handlers, 'create_checkpoint_manager', lambda *args, **kwargs: manager) service = Management() app = FastAPI() @@ -221,12 +221,12 @@ def save(self, **kwargs): endpoint = next(r.endpoint for r in app.routes if getattr(r, 'path', '') == '/twinkle/save') request = Request({'type': 'http', 'headers': []}) request.state.session_id = 'session' - body = types.SaveRequest(adapter_name='adapter', name='failed') + body = types.SaveRequest(adapter_name='adapter', name='failed', training_id='training-a') with pytest.raises(RuntimeError, match='backend save failed'): await endpoint(request, body, service) - assert manager.list_checkpoints('session-adapter').checkpoints == [] - assert runs._read_info('session-adapter') == before + assert manager.list_checkpoints('training-a').checkpoints == [] + assert runs._read_info('training-a') == before @pytest.mark.asyncio @@ -253,9 +253,9 @@ async def test_checkpoint_consumers_keep_backend_path_and_options(monkeypatch, t request.state.session_id = 'session' if operation == 'load': - body = types.LoadRequest(adapter_name='adapter', name=path, load_optimizer=True, strict=False) + body = types.LoadRequest(adapter_name='adapter', name=path, load_optimizer=True, strict=False, training_id='training-a') elif operation == 'resume_from_checkpoint': - body = types.ResumeFromCheckpointRequest(adapter_name='adapter', name=path, resume_only_model=True) + body = types.ResumeFromCheckpointRequest(adapter_name='adapter', name=path, resume_only_model=True, training_id='training-a') service.model.resume_from_checkpoint.return_value = {'step': 7} else: body = types.UploadToHubRequest(checkpoint_dir=path, hub_model_id='org/output', hub_token='hub-token') @@ -264,10 +264,10 @@ async def test_checkpoint_consumers_keep_backend_path_and_options(monkeypatch, t if operation == 'load': service.model.load.assert_called_once_with( name='step' if local else path, output_dir=manager.get_save_dir('run') if local else None, - load_optimizer=True, adapter_name='session-adapter', token='token', strict=False) + load_optimizer=True, adapter_name='training-a', token='token', strict=False) elif operation == 'resume_from_checkpoint': service.model.resume_from_checkpoint.assert_called_once_with( - directory, resume_only_model=True, adapter_name='session-adapter') + directory, resume_only_model=True, adapter_name='training-a') assert response.result == {'result': {'step': 7}} else: service.model.upload_to_hub.assert_called_once_with( @@ -300,7 +300,7 @@ async def test_forward_backward_preflights_loss_inputs_before_backend() -> None: 'labels': [-100, 5], } for index in range(3)], adapter_name='adapter', - loss_kwargs={'old_logps': [[-0.1], [-0.2], [-0.3]]}, + loss_kwargs={'old_logps': [[-0.1], [-0.2], [-0.3]]}, training_id='training-a', ) await route.endpoint(request, body, management) @@ -309,11 +309,11 @@ async def test_forward_backward_preflights_loss_inputs_before_backend() -> None: preflight_inputs, preflight_adapter, preflight_kwargs = management.preflight_calls[0] # Slice-free: a per-rank slice would carry at most one of the three samples. assert len(preflight_inputs) == 3 - assert preflight_adapter == 'session-adapter' + assert preflight_adapter == 'training-a' assert preflight_kwargs['old_logps'] == [[-0.1], [-0.2], [-0.3]] assert management.model_calls[0][0] == 'preflight' assert management.model_calls[-1][0] != 'preflight' - assert management.model_calls[-1][1] == 'session-adapter' + assert management.model_calls[-1][1] == 'training-a' @pytest.mark.asyncio @@ -330,11 +330,11 @@ async def test_forward_backward_skips_preflight_without_rl_inputs() -> None: 'input_ids': [1, 2], 'labels': [-100, 5], }], - adapter_name='adapter', + adapter_name='adapter', training_id='training-a', ) await route.endpoint(request, body, management) assert management.preflight_calls == [] assert len(management.model_calls) == 1 - assert management.model_calls[0][1] == 'session-adapter' + assert management.model_calls[0][1] == 'training-a' diff --git a/tests/server/resources/__init__.py b/tests/server/resources/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/server/resources/test_allocation.py b/tests/server/resources/test_allocation.py new file mode 100644 index 000000000..9e89ef403 --- /dev/null +++ b/tests/server/resources/test_allocation.py @@ -0,0 +1,120 @@ +# 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 + + +def test_stale_resource_summary_does_not_reject_a_schedulable_group(monkeypatch): + import importlib + import ray + placement = importlib.import_module('ray.util.placement_group') + monkeypatch.setattr(ray, 'nodes', lambda: [{'Alive': True, 'Resources': {'CPU': 4, 'GPU': 2}}]) + monkeypatch.setattr(ray, 'available_resources', lambda: {'CPU': 4}) + created = [] + group = object() + + def create(bundles, **kwargs): + created.append((bundles, kwargs)) + return group + + monkeypatch.setattr(placement, 'placement_group', create) + assert DevicePool._create_group('GPU', 2, 'replacement') is group + assert created[0][0][0]['GPU'] == 2 + assert created[0][1] == {'strategy': 'STRICT_PACK', 'name': 'replacement'} + + +@pytest.mark.asyncio +@pytest.mark.parametrize('blocked', ['nodes', '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(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) diff --git a/tests/server/resources/test_instances.py b/tests/server/resources/test_instances.py new file mode 100644 index 000000000..556bf2879 --- /dev/null +++ b/tests/server/resources/test_instances.py @@ -0,0 +1,254 @@ +import asyncio +from types import SimpleNamespace +import logging +import pytest +from twinkle.server.resources.instances import InstanceLifecycle + + +@pytest.mark.asyncio +async def test_spawn_logs_exception_with_traceback(caplog, monkeypatch): + # Capture one path regardless of the parent logger's runtime configuration. + target = logging.getLogger('twinkle.server.resources.instances') + monkeypatch.setattr(target, 'handlers', [caplog.handler]) + monkeypatch.setattr(target, 'propagate', False) + lifecycle = InstanceLifecycle('Training') + + async def broken(): + raise RuntimeError('background sentinel') + + with caplog.at_level(logging.ERROR, logger=target.name): + 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 records[0].exc_info[2] is not None + assert not lifecycle.jobs + + +@pytest.mark.asyncio +async def test_concurrent_close_shares_one_stop_and_retries_quarantine(caplog, monkeypatch): + import asyncio + target = logging.getLogger('twinkle.server.resources.instances') + monkeypatch.setattr(target, 'handlers', [caplog.handler]) + monkeypatch.setattr(target, 'propagate', False) + caplog.set_level(logging.WARNING, logger=target.name) + lifecycle = InstanceLifecycle('test') + entered, resume = asyncio.Event(), asyncio.Event() + calls, events = [], [] + + 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'] + records = [record for record in caplog.records if record.getMessage() == 'test instance one shutdown is unconfirmed'] + assert len(records) == 1 and records[0].exc_info[1].args == ('death unconfirmed',) + assert records[0].exc_info[2] is not None + await lifecycle.close('one', stop, 1, stopped, quarantined) + assert calls == ['stop', 'stop'] and events[-1] == 'stopped' + await lifecycle.join() + 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 + + +@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() + + +@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.pool.driver import SamplerDriver + from twinkle.server.model.pool.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) + + +@pytest.mark.asyncio +@pytest.mark.parametrize('outcome', [ + 'unconfirmed_worker', 'unavailable_worker', 'dead_worker', 'removed', 'removal_failed' +]) +async def test_worker_death_precedes_group_release_and_removal_timing(monkeypatch, outcome): + from types import SimpleNamespace + from unittest.mock import AsyncMock, Mock + import ray + from ray import serve + from twinkle.server.exceptions import ResourceReleaseError + from twinkle.server.resources.runtime import RayRuntime + + async def ready(): + if outcome == 'unavailable_worker': + raise ray.exceptions.ActorUnavailableError('Control-plane connection lost', None) + if outcome == 'dead_worker': + raise ray.exceptions.ActorDiedError() + return None # Worker can still answer: death has not been confirmed. + + handle = SimpleNamespace(__ray_ready__=SimpleNamespace(remote=ready)) + rows = ([{'namespace': 'test', 'name': 'owned-worker'}] + if outcome in ('unconfirmed_worker', 'unavailable_worker', 'dead_worker') else []) + monkeypatch.setattr(ray.util, 'list_named_actors', lambda **kwargs: rows) + monkeypatch.setattr(ray, 'get_actor', lambda *args, **kwargs: handle) + monkeypatch.setattr(ray, 'kill', lambda *args, **kwargs: None) + monkeypatch.setattr(serve, 'delete', lambda name: None) + runtime = RayRuntime('test', allocation_timeout=1, shutdown_timeout=1) + runtime.pool.remove = AsyncMock(side_effect=RuntimeError('PG removal failed') + if outcome == 'removal_failed' else None) + observe = Mock() + if outcome in ('unconfirmed_worker', 'unavailable_worker'): + with pytest.raises(ResourceReleaseError) as failure: + await runtime.shutdown('owned', 'app', observe_removal=observe) + assert failure.value.phase == 'worker_death' + runtime.pool.remove.assert_not_awaited() + observe.assert_not_called() + else: + if outcome == 'removal_failed': + with pytest.raises(RuntimeError, match='PG removal failed'): + await runtime.shutdown('owned', 'app', observe_removal=observe) + else: + await runtime.shutdown('owned', 'app', observe_removal=observe) + runtime.pool.remove.assert_awaited_once_with('app') + assert observe.call_count == 1 + duration, status = observe.call_args.args + assert duration >= 0 and status == ('error' if outcome == 'removal_failed' else 'success') diff --git a/tests/server/sampler/test_pool.py b/tests/server/sampler/test_pool.py new file mode 100644 index 000000000..12592bef4 --- /dev/null +++ b/tests/server/sampler/test_pool.py @@ -0,0 +1,515 @@ +# 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.pool.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_shutdown_rejects_unconfirmed_release_and_allows_retry(): + pool, driver, _ = setup() + await pool.initialize() + driver.fail_close = True + try: + with pytest.raises(RuntimeError, match='shutdown.*unconfirmed'): + await pool.shutdown() + assert pool._owned_count() == 2 and not driver.stopped + assert all(instance.state == 'QUARANTINED' for instance in pool.instances.values()) + finally: + driver.fail_close = False + await pool.shutdown() + assert pool._owned_count() == 0 and len(driver.stopped) == 2 + + +@pytest.mark.asyncio +async def test_minimum_expand_and_reclaim_only_idle_extras(): + pool, driver, clock = setup() + 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_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 = pool.budgets.model_copy(update={'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 +@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/sampler/test_request_runtime.py b/tests/server/sampler/test_request_runtime.py new file mode 100644 index 000000000..61c8095bd --- /dev/null +++ b/tests/server/sampler/test_request_runtime.py @@ -0,0 +1,267 @@ +# 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.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 + + +@pytest.mark.asyncio +async def test_allocated_sampler_initializes_with_typed_owner_context(monkeypatch): + from unittest.mock import Mock + from twinkle.infra._ray.ray_helper import RayHelper + from twinkle.infra._ray.resource_manager import ResourceManager + from twinkle.server.sampler import app as module + from twinkle.server.state.models import RuntimeContext + + context = RuntimeContext(namespace='owned', registry_name='registry', instance_id='sampler-owned') + allocation = object() + manager = object() + allocate = Mock(return_value=manager) + sampler = SimpleNamespace() + monkeypatch.setenv('TWINKLE_RAY_NAMESPACE', 'before') + monkeypatch.setenv('TWINKLE_RAY_REGISTRY_NAME', 'before') + monkeypatch.setattr(RayHelper, 'resource_manager', None) + monkeypatch.setattr(ResourceManager, 'from_allocation', allocate) + monkeypatch.setattr(module, 'init_twinkle_runtime', lambda **kwargs: None) + monkeypatch.setattr(module, '_construct_sampler_backend', lambda *args: sampler) + monkeypatch.setattr(module.SamplerManagement, '_init_task_queue', lambda *args, **kwargs: None) + monkeypatch.setattr(module, 'get_server_state', lambda: SimpleNamespace()) + monkeypatch.setattr(module.serve, 'get_replica_context', lambda: SimpleNamespace( + replica_id=SimpleNamespace(unique_id='serve-replica'))) + owner = module.SamplerManagement.__new__(module.SamplerManagement) + await owner.__init__( + model_id='mock', nproc_per_node=1, device_group={'name': 'sampler', 'ranks': 1, 'device_type': 'cpu'}, + device_mesh={'dp_size': 1}, sampler_type='mock', sampling_context=context, + placement_group=allocation, allocation_devices=['0']) + assert owner.replica_id == context.instance_id + assert RayHelper.resource_manager is manager + allocate.assert_called_once_with(owner.device_group, allocation, ['0']) + await owner.data_plane.close() + + +@pytest.mark.asyncio +async def test_worker_exports_actual_completion_handle_without_private_task_lookup(): + queue = _DummyQueue() + 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' + + ref = await queue.schedule_concurrent_task(lambda: guard.run(backend), execution_scope=scope) + try: + await asyncio.wait_for(entered.wait(), 1) + completion = asyncio.create_task(scope.wait()) + 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.exceptions import BackendBusyError + guard = SamplerConcurrency() + runtime = SamplerRuntime(SimpleNamespace(sampling_guard=guard)) + task = asyncio.create_task(asyncio.sleep(0)) + await task + from twinkle.server.sampler.runtime import SamplingOperation + runtime._operations['cancelled'] = SamplingOperation(ExecutionScope(), task) + guard.invalidate_weights(RuntimeError('cancel ACK preceded actual generation exit')) + with pytest.raises(BackendBusyError): + 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 + + +@pytest.mark.asyncio +async def test_scope_waits_for_jobs_registered_by_an_earlier_job(): + scope = ExecutionScope() + child_started, child_release, parent_release = asyncio.Event(), asyncio.Event(), asyncio.Event() + + async def child(): + child_started.set() + await child_release.wait() + + async def parent(): + await parent_release.wait() + scope.track(asyncio.create_task(child())) + + scope.track(asyncio.create_task(parent())) + waiter = asyncio.create_task(scope.wait()) + try: + await asyncio.sleep(0) + parent_release.set() + await child_started.wait() + with pytest.raises(TimeoutError): + await asyncio.wait_for(asyncio.shield(waiter), 0.05) + child_release.set() + await waiter + finally: + child_release.set() + await asyncio.gather(waiter, return_exceptions=True) diff --git a/tests/server/sampler/test_sampling_concurrency.py b/tests/server/sampler/test_sampling_concurrency.py index f1697ec9c..14b3e07e1 100644 --- a/tests/server/sampler/test_sampling_concurrency.py +++ b/tests/server/sampler/test_sampling_concurrency.py @@ -21,8 +21,8 @@ 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.task_queue.types import BackendBusyError +from twinkle.server.config.task_queue import TaskQueueConfig +from twinkle.server.exceptions import BackendBusyError from twinkle.server.task_queue.worker import ComputeWorker @@ -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/sampler/test_stream_guarantees.py b/tests/server/sampler/test_stream_guarantees.py index 60623c970..ed10414d0 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.exceptions import BackendBusyError + with pytest.raises(BackendBusyError, match='restarting'): + service.sampling_guard.check_available() diff --git a/tests/server/sampler/test_streaming_bridge.py b/tests/server/sampler/test_streaming_bridge.py new file mode 100644 index 000000000..bba725b4f --- /dev/null +++ b/tests/server/sampler/test_streaming_bridge.py @@ -0,0 +1,45 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Consumer disconnection never shortens backend ownership.""" +from types import SimpleNamespace + +import pytest + +from twinkle.server.sampler.backends.streaming import STREAM_SENTINEL, stream_to_queue + + +@pytest.mark.parametrize('disconnect', [False, True]) +def test_stream_drains_and_closes_after_consumer_disconnect(disconnect): + consumed, closed, published = [], [], [] + + def generate(*args): + try: + for index in range(3): + consumed.append(index) + yield index, None + finally: + closed.append(True) + + def put(value): + if disconnect: + raise RuntimeError('queue actor is gone') + published.append(value) + + stream_to_queue(SimpleNamespace(sample_stream=generate), SimpleNamespace(put=put), []) + assert consumed == [0, 1, 2] and closed == [True] + assert published == ([] if disconnect else [(0, None), (1, None), (2, None), STREAM_SENTINEL]) + + +def test_backend_failure_is_reported_and_iterator_closed(): + closed, published = [], [] + + def generate(*args): + try: + yield 1, None + raise ValueError('generation failed') + finally: + closed.append(True) + + stream_to_queue(SimpleNamespace(sample_stream=generate), SimpleNamespace(put=published.append), []) + assert closed == [True] + assert isinstance(published[-2], ValueError) + assert published[-1] == STREAM_SENTINEL diff --git a/tests/server/sampler/test_twinkle_async_rows.py b/tests/server/sampler/test_twinkle_async_rows.py index aa1877474..407ba483a 100644 --- a/tests/server/sampler/test_twinkle_async_rows.py +++ b/tests/server/sampler/test_twinkle_async_rows.py @@ -76,7 +76,7 @@ def __init__(self): async def _on_request_start(self, _request): return 'token' - async def submit_concurrent_and_peek(self, coro_factory, *, model_id=None, task_type=None): + async def submit_concurrent_and_peek(self, coro_factory, *, model_id=None, task_type=None, execution_scope=None): self.scheduled.append({'model_id': model_id, 'task_type': task_type}) result = await coro_factory() from twinkle.protocol.types.lifecycle import TaskEnvelope diff --git a/tests/server/session_resource/test_contract.py b/tests/server/session_resource/test_contract.py index 477b509cf..843b7d611 100644 --- a/tests/server/session_resource/test_contract.py +++ b/tests/server/session_resource/test_contract.py @@ -7,7 +7,6 @@ import pytest from twinkle.server.exceptions import RequestRejectedError -from twinkle.server.session_resource.adapter import AdapterManagerMixin from twinkle.server.session_resource.base import SessionResourceMixin from twinkle.server.session_resource.processor import ProcessorManagerMixin @@ -39,16 +38,12 @@ class _MissingBaseHook(SessionResourceMixin): pass -class _MissingAdapterHook(AdapterManagerMixin): - pass - - class _MissingProcessorHook(ProcessorManagerMixin): pass def test_all_resource_expiry_hooks_are_abstract() -> None: - for cls in (_MissingBaseHook, _MissingAdapterHook, _MissingProcessorHook): + for cls in (_MissingBaseHook, _MissingProcessorHook): with pytest.raises(TypeError): cls() diff --git a/tests/server/start_e2e_server.py b/tests/server/start_e2e_server.py index 4af2e2b85..21995f969 100644 --- a/tests/server/start_e2e_server.py +++ b/tests/server/start_e2e_server.py @@ -38,10 +38,10 @@ MEGATRON_BASE = "tests/server/config/server_config_4b_e2e_megatron.yaml" _FULL = { "train_mode": "full", + "max_replicas": 1, # Full-parameter training holds the deployment exclusively; use long # timeouts so the tenant is not expired mid-test. - "adapter_config.adapter_timeout": 3600, - "adapter_config.adapter_max_lifetime": 36000, + "idle_timeout_s": 3600, } VARIANTS = { "lora": (TRANSFORMERS_BASE, {}), diff --git a/tests/server/state/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..e596391de 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,68 @@ 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 + + +@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_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_managers.py b/tests/server/state/test_managers.py index 242dc3508..6c2dfd47b 100644 --- a/tests/server/state/test_managers.py +++ b/tests/server/state/test_managers.py @@ -534,3 +534,20 @@ async def get_all_then_touch(): session_alive = await state.get_session_last_heartbeat('s-race') is not None model_alive = await state.get_model_metadata('m-race') is not None assert session_alive == model_alive + + +@pytest.mark.parametrize('environment_limit', [None, '7']) +def test_state_initializes_without_or_with_a_partial_environment_policy(monkeypatch, environment_limit): + """A missing optional policy must resolve to defaults before field access.""" + import uuid + from tests.server.state.fake_backend import FakeBackend + from twinkle.server.config.application_spec import SERVER_STATE_ENV_KEYS + from twinkle.server.state.server_state import get_server_state + + for key in SERVER_STATE_ENV_KEYS: + monkeypatch.delenv(key, raising=False) + if environment_limit is not None: + monkeypatch.setenv('TWINKLE_SERVER_STATE_PER_TOKEN_MODEL_LIMIT', environment_limit) + state = get_server_state(cache_key='policy-' + uuid.uuid4().hex, backend=FakeBackend()) + assert state.expiration_timeout == 86400 + assert state.cleanup_interval == 3600 diff --git a/tests/server/state/test_redis_read_retry.py b/tests/server/state/test_redis_read_retry.py new file mode 100644 index 000000000..c5cac36e3 --- /dev/null +++ b/tests/server/state/test_redis_read_retry.py @@ -0,0 +1,49 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Read retry is bounded, request-local, and never replays writes.""" +import asyncio +from unittest.mock import AsyncMock + +import pytest +from redis.exceptions import ConnectionError + +from twinkle.server.state.backend.redis_backend import RedisBackend + + +@pytest.mark.asyncio +async def test_stale_read_retries_without_blocking_another_reader(monkeypatch): + backend = RedisBackend('redis://127.0.0.1:1') + entered, resume = asyncio.Event(), asyncio.Event() + attempts = [] + + async def read(key): + attempts.append(key) + if key == 'stale' and attempts.count(key) == 1: + entered.set() + await resume.wait() + raise ConnectionError('stale connection') + return '"' + key + '"' + + backend._client.get = read + monkeypatch.setattr('twinkle.server.state.backend.redis_backend.asyncio.sleep', AsyncMock()) + try: + first = asyncio.create_task(backend.get('stale')) + await entered.wait() + assert await backend.get('healthy') == 'healthy' + resume.set() + assert await first == 'stale' + assert attempts == ['stale', 'healthy', 'stale'] + finally: + resume.set() + await backend.close() + + +@pytest.mark.asyncio +async def test_cancellation_is_not_retried(): + backend = RedisBackend('redis://127.0.0.1:1') + backend._client.get = AsyncMock(side_effect=asyncio.CancelledError()) + try: + with pytest.raises(asyncio.CancelledError): + await backend.get('key') + backend._client.get.assert_awaited_once() + finally: + await backend.close() diff --git a/tests/server/state/test_redis_recovery.py b/tests/server/state/test_redis_recovery.py new file mode 100644 index 000000000..bfde03295 --- /dev/null +++ b/tests/server/state/test_redis_recovery.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/state/test_sequence.py b/tests/server/state/test_sequence.py new file mode 100644 index 000000000..053ef23ef --- /dev/null +++ b/tests/server/state/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.helpers 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/state/test_update_atomic.py b/tests/server/state/test_update_atomic.py index d03b3a79e..2bcc78f29 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 ----------------------------------------------- # @@ -258,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/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/static/test_adapter_name_mapping.py b/tests/server/static/test_adapter_name_mapping.py index 6856f099e..708bef399 100644 --- a/tests/server/static/test_adapter_name_mapping.py +++ b/tests/server/static/test_adapter_name_mapping.py @@ -2,7 +2,7 @@ """Static check: model backend calls map the tenant adapter name (F002 / P002). ``run_submit`` hands each ``backend_call`` a *tenant-scoped* adapter name -(``owner_id-``). Before that reaches a backend model method it must be +(``training_id``). Before that reaches a backend model method it must be translated by ``ModelManagement.resolve_model_adapter_name`` — which returns ``''`` in full-parameter mode (the empty-string default optimizer group) and the name unchanged in LoRA mode. The inline forward/backward/step endpoints did this; 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..b3e9bb5de --- /dev/null +++ b/tests/server/task_queue/helpers.py @@ -0,0 +1,59 @@ +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 74% rename from tests/server/utils/test_task_queue_mixin.py rename to tests/server/task_queue/test_task_queue_mixin.py index a0316c632..e2f111e66 100644 --- a/tests/server/utils/test_task_queue_mixin.py +++ b/tests/server/task_queue/test_task_queue_mixin.py @@ -1,10 +1,11 @@ +from tests.server.task_queue.helpers import _DummyState, _AllowingRateLimiter, _DummyQueue import asyncio 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.exceptions import RequestRejectedError from twinkle.server.task_queue.worker import ComputeWorker, _exception_summary @@ -48,54 +49,6 @@ 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 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 - - @pytest.mark.asyncio async def test_worker_stop_exits_when_wait_for_swallows_cancellation(monkeypatch): queue = _DummyQueue() @@ -140,7 +93,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 +105,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 +176,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 +185,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} @@ -319,7 +279,7 @@ async def test_user_task_error_is_stored_as_user_failure(): queue.enable_compute_worker() async def work(): - raise UserTaskError('invalid request') + raise RequestRejectedError('invalid request') try: await queue.schedule_task(work, model_id='model1', token='token1') @@ -455,7 +415,7 @@ async def test_concurrent_task_keeps_queue_failure_mapping(make_concurrent_queue queue = make_concurrent_queue() async def work(): - raise UserTaskError('invalid request') + raise RequestRejectedError('invalid request') await queue.schedule_concurrent_task(work, token='token1') for _ in range(100): @@ -614,6 +574,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() @@ -643,3 +705,53 @@ 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() + + +@pytest.mark.asyncio +async def test_backend_cancellation_is_an_explicit_server_failure(make_concurrent_queue): + from twinkle.server.lifecycle.envelope import envelope_from_record + from twinkle.server.task_queue.types import ExecutionScope + queue = make_concurrent_queue() + scope = ExecutionScope() + + async def work(): + raise asyncio.CancelledError() + + ref = await queue.schedule_concurrent_task(work, token='token1', execution_scope=scope) + await scope.wait() + record = await queue.state.get_future(ref['request_id']) + assert record['status'] == 'failed' + assert record['failure']['reason_code'] == 'backend_cancelled' + envelope = envelope_from_record(ref['request_id'], record) + assert envelope.error.error_code == 503 + assert envelope.error.category == 'server' diff --git a/tests/server/telemetry/test_training_metrics.py b/tests/server/telemetry/test_training_metrics.py new file mode 100644 index 000000000..7f4214593 --- /dev/null +++ b/tests/server/telemetry/test_training_metrics.py @@ -0,0 +1,259 @@ +# 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.model.pool.controller import TrainingController +from twinkle.server.telemetry.metrics import TrainingInstanceMetrics + +_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(backend=backend, device_type='cuda') + controller = TrainingController.__new__(TrainingController) + controller.config = SimpleNamespace(get=lambda base_model, mode: spec) + controller._metrics = registry.training_metrics() + controller._instance_metrics = TrainingInstanceMetrics(controller._metrics) + 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()) + + +@pytest.mark.asyncio +async def test_repeated_failed_cleanup_is_one_quarantine_transition(): + from unittest.mock import MagicMock + from tests.server.state.fake_backend import FakeBackend + from tests.server.training.helpers import create, setup + + controller, state, driver, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'metric-retry') + instruments = MagicMock() + controller._metrics = instruments + driver.fail_shutdown = True + try: + with pytest.raises(RuntimeError, match='quarantined'): + await controller.release(identity, 'owner', 'session') + for _ in range(2): + await controller.sweep() + await controller.wait_for_jobs() + assert instruments.quarantined_total.inc.call_count == 1 + assert instruments.shutdown_seconds.observe.call_count == 3 + finally: + driver.fail_shutdown = False + await controller.shutdown() + + +@pytest.mark.parametrize('phase', ['worker_death', 'pg_removal']) +def test_quarantine_classification_uses_typed_phase(phase): + from twinkle.server.exceptions import ResourceReleaseError + assert TrainingInstanceMetrics.quarantine_reason(ResourceReleaseError('unconfirmed', phase=phase)) == phase + assert TrainingInstanceMetrics.quarantine_reason(RuntimeError('user text: worker death')) == 'other' + + +@pytest.mark.asyncio +@pytest.mark.parametrize('fail_start', [False, True]) +async def test_startup_and_shutdown_histograms_follow_real_attempts(fail_start): + from unittest.mock import MagicMock + from tests.server.state.fake_backend import FakeBackend + from tests.server.training.helpers import setup, request + + controller, state, driver, _ = await setup(backend=FakeBackend()) + controller._metrics = MagicMock() + driver.fail_start = fail_start + try: + ref = await controller.create(request('timed-start')) + await controller.wait_for_jobs() + future = await state.get_future(ref['request_id']) + assert future['status'] == ('failed' if fail_start else 'completed') + call = controller._metrics.startup_seconds.observe.call_args + assert call.args[0] >= 0 + assert call.kwargs['tags'] == {'backend': 'mock', 'mode': 'lora', + 'status': 'error' if fail_start else 'success'} + if not fail_start: + await controller.release(future['result']['training_id'], 'owner', 'session') + call = controller._metrics.shutdown_seconds.observe.call_args + assert call.args[0] >= 0 and call.kwargs['tags']['status'] == 'success' + finally: + await controller.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('failed', [False, True]) +async def test_sweep_instruments_one_iteration_and_its_outcome(monkeypatch, failed): + import asyncio + from unittest.mock import AsyncMock, MagicMock + from twinkle.server.model.pool.actor import TrainingControllerActor + + async def stop_after_iteration(delay): + raise asyncio.CancelledError + + actor = SimpleNamespace(core=SimpleNamespace(sweep=AsyncMock( + side_effect=RuntimeError('state unavailable') if failed else None)), + config=SimpleNamespace(sweep_interval_s=5), _metrics=MagicMock()) + monkeypatch.setattr('twinkle.server.model.pool.actor.asyncio.sleep', stop_after_iteration) + with pytest.raises(asyncio.CancelledError): + await TrainingControllerActor.__ray_metadata__.modified_class._sweep(actor) + tags = {'status': 'error' if failed else 'success'} + actor._metrics.sweep_total.inc.assert_called_once_with(tags=tags) + call = actor._metrics.sweep_seconds.observe.call_args + assert call.args[0] >= 0 and call.kwargs['tags'] == tags + + +@pytest.mark.asyncio +@pytest.mark.parametrize('failed', [False, True]) +async def test_driver_records_allocation_attempts(failed): + from unittest.mock import AsyncMock, MagicMock + from twinkle.server.config.training import TrainingConfig + from twinkle.server.resources.allocation import DeviceAllocation + from twinkle.server.model.pool.driver import RayInstanceDriver + + config = TrainingConfig(models=[{'name': 'tiny', 'model_id': 'tiny', 'backend': 'mock', + 'device_type': 'cpu', 'mode': 'full', 'config': {'num_devices': 1, 'device_mesh': {'device_type': 'cpu'}}}]) + store = SimpleNamespace(models=SimpleNamespace(unregister_replica_capacity=AsyncMock())) + driver = RayInstanceDriver(config, store, 'test', 'controller') + driver._metrics = MagicMock() + allocation = DeviceAllocation('owned', 'pg', 'node', ('0',)) + driver.runtime.pool.allocate = AsyncMock(return_value=allocation, + side_effect=RuntimeError('allocation failed') if failed else None) + instance = SimpleNamespace(instance_id='instance', controller_epoch='epoch', application_name='owned', + base_model='tiny', mode='full') + if failed: + with pytest.raises(RuntimeError, match='allocation failed'): + await driver.allocate(instance, config.models[0]) + else: + assert (await driver.allocate(instance, config.models[0]))['devices'] == ['0'] + call = driver._metrics.allocation_seconds.observe.call_args + assert call.args[0] >= 0 + assert call.kwargs['tags'] == {'backend': 'mock', 'mode': 'full', 'device_type': 'cpu', + 'status': 'error' if failed else 'success'} diff --git a/tests/server/test_dependencies.py b/tests/server/test_dependencies.py new file mode 100644 index 000000000..03d795da0 --- /dev/null +++ b/tests/server/test_dependencies.py @@ -0,0 +1,115 @@ +# 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__' + +# 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', 'model.tinker_datum', 'runtime', 'task_errors', 'utils', 'validation'}), + frozenset({'telemetry'}), + frozenset({'state'}), + frozenset({'resources'}), + frozenset({'deployment', 'lifecycle', 'session_resource'}), + frozenset({'task_queue'}), + frozenset({'checkpoint', 'data_plane'}), + frozenset({'gateway', 'model', 'sampler', 'processor'}), + frozenset({'model.pool', 'sampler.pool'}), + frozenset({'launcher'}), + frozenset({'cli'}), + frozenset({'__main__', _ROOT_COMPONENT}), +) + + +def _component(module: str) -> str: + if module == _ROOT: + return _ROOT_COMPONENT + relative = module.removeprefix(f'{_ROOT}.') + for component in ('model.pool', 'sampler.pool', 'model.tinker_datum'): + if relative == component or relative.startswith(component + '.'): + return component + return relative.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.model.pool.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) + 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}') + + +@pytest.mark.parametrize('component', ['model.pool', 'sampler.pool']) +def test_controller_import_does_not_load_ray_composition(component): + code = f'''import twinkle.server.{component}.controller +import sys +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/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/gateway/test_proxy.py b/tests/server/test_service_proxy.py similarity index 61% rename from tests/server/gateway/test_proxy.py rename to tests/server/test_service_proxy.py index 9b03e4a98..dc9eb4359 100644 --- a/tests/server/gateway/test_proxy.py +++ b/tests/server/test_service_proxy.py @@ -1,19 +1,17 @@ # 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 import json +import httpx + import pytest from starlette.requests import Request from unittest.mock import AsyncMock -from twinkle.server.gateway.proxy import ServiceProxy +from twinkle.server.lifecycle.proxy import ServiceProxy from twinkle.protocol.headers import H_MULTIPLEX, H_MULTIPLEX_LEGACY, H_REQUEST_ID @@ -51,6 +49,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/training/helpers.py b/tests/server/training/helpers.py new file mode 100644 index 000000000..ba4a5fb9d --- /dev/null +++ b/tests/server/training/helpers.py @@ -0,0 +1,174 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +from twinkle.server.config.task_queue import TaskQueueConfig +from twinkle.server.config.training import TrainingConfig +from twinkle.server.resources.allocation import CapacityUnavailableError +from twinkle.server.state import ServerState +from twinkle.server.state.backend.memory_backend import RayActorBackend +from twinkle.server.state.training_records import CreateTraining +from twinkle.server.task_queue import TaskQueueMixin +from twinkle.server.model.pool.controller import TrainingController +from twinkle.server.model.runtime import TrainingRuntime + + +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, adapter_config=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 + } + if adapter_config is not None: + mode['adapter_config'] = adapter_config + config = TrainingConfig(models=[{ + 'name': 'small', + 'model_id': 'ms://small', + 'backend': 'mock', + 'device_type': 'cuda', + 'mode': 'lora', + 'config': mode, + }, { + 'name': 'small-full', + 'model_id': 'ms://small', + 'backend': 'mock', + 'device_type': 'cuda', + 'mode': 'full', + 'config': { + 'num_devices': 2, + 'device_mesh': { + 'fsdp_size': 2 + }, + 'idle_timeout_s': 10 + }, + }]) + state = ServerState(backend=backend or RayActorBackend(), per_token_model_limit=quota) + await state.create_session({'session_id': 'session'}) + 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-full' if mode == 'full' else 'small', + mode=mode, + token=token, + session_id='session', + create_key=str(key), + client_type='twinkle', + payload={}, + adapter_label=label) + + +async def create(controller, state, key, **kwargs): + 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(TrainingRuntime, TaskQueueMixin): + + def __init__(self, controller, state, instance): + self.state, self.replica_id = state, instance.instance_id + from twinkle.server.state.training_records import TrainingRuntimeContext + self.training_context = TrainingRuntimeContext( + instance_id=instance.instance_id, + epoch=controller.epoch, + namespace='test', + registry_name='test', + controller_name='test') + self._model_unhealthy = False + self._training_controller = SimpleNamespace( + **{name: RPC(getattr(controller, name)) + for name in ('admit', 'finish_operation', 'fail_instance')}) + self._init_task_queue( + TaskQueueConfig(execution_timeout=0.1, inline_fast_path_timeout=0.01), enable_admission_gate=True) + + +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_acceptance_harness.py b/tests/server/training/test_acceptance_harness.py new file mode 100644 index 000000000..58d46e7de --- /dev/null +++ b/tests/server/training/test_acceptance_harness.py @@ -0,0 +1,124 @@ +# 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() + + +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) + + +@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) + + +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] diff --git a/tests/server/training/test_actor_initialization.py b/tests/server/training/test_actor_initialization.py new file mode 100644 index 000000000..dd9a3f530 --- /dev/null +++ b/tests/server/training/test_actor_initialization.py @@ -0,0 +1,129 @@ +# 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.resources.instances import InstanceLifecycle +from twinkle.server.model.pool.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.model.pool.actor.asyncio.sleep', sleep) + monkeypatch.setattr('twinkle.server.model.pool.actor.ray.get_actor', get_actor) + monkeypatch.setattr('twinkle.server.model.pool.actor.ray.kill', kill) + actor = SimpleNamespace(core=SimpleNamespace(initialize=AsyncMock(side_effect=RuntimeError('allocation failed'))), + config=SimpleNamespace(sweep_interval_s=7), namespace='initialization-test', _monitors=[]) + actor._restart_after_failure = lambda: cls._restart_after_failure(actor) + 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()) + initialization = asyncio.get_running_loop().create_future() + initialization.set_result(None) + 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.runtime.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', models=SimpleNamespace(touch_replica_last_seen=touch)), + core=SimpleNamespace(epoch='controller', instances={'bad': bad, 'good': good}, fail_instance=fail, + lifecycle=InstanceLifecycle('Training')), + driver=SimpleNamespace(health=health), _unreachable_since={}, _failures={}) + 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.core.lifecycle.named_jobs) == 1 + finally: + probe_gate.set() + close_gate.set() + monitor.cancel() + await asyncio.gather(monitor, *actor.core.lifecycle.named_jobs.values(), return_exceptions=True) + assert not actor.core.lifecycle.named_jobs + + +@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(runtime=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_controller.py b/tests/server/training/test_controller.py new file mode 100644 index 000000000..f4c3babad --- /dev/null +++ b/tests/server/training/test_controller.py @@ -0,0 +1,799 @@ +# 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 + +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.model.pool.controller import TrainingController +from twinkle.server.exceptions import TrainingExpiredError +from twinkle.server.resources.allocation import CapacityUnavailableError +from twinkle.server.state.training_records import CreateTraining + + +@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]['config']['max_loras'] = 2 + controller.config = TrainingConfig.model_validate(raw) + 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 + 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]['config']['min_replicas'] = 1 + controller.config = TrainingConfig.model_validate(raw) + 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') + 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_disabled_max_lifetime_and_live_heartbeat_do_not_override_business_idle(): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, clock = await setup( + backend=FakeBackend(), adapter_config={'adapter_max_lifetime': None}) + identity = await create(controller, state, 'idle-only') + controller.trainings[identity].created_at = '2000-01-01T00:00:00+00:00' + clock.now = 9 + await state.touch_session('session') + await controller.sweep() + assert controller.trainings[identity].state == 'ACTIVE' + clock.now = 10 + await state.touch_session('session') + await controller.sweep() + await controller.wait_for_jobs() + assert controller.trainings[identity].state == 'EXPIRED' + assert not runtime.allocations + + +@pytest.mark.asyncio +@pytest.mark.parametrize('reason', ['heartbeat', 'max_lifetime']) +@pytest.mark.parametrize('busy', [False, True]) +async def test_adapter_lifecycle_expires_independently_and_waits_for_running_work(reason, busy): + import time + from datetime import datetime, timedelta, timezone + from tests.server.state.fake_backend import FakeBackend + + controller, state, runtime, _ = await setup( + backend=FakeBackend(), adapter_config={'adapter_timeout': 30, 'adapter_max_lifetime': 60}) + identity = await create(controller, state, 'lifecycle') + if busy: + await controller.admit(identity, 'owner', 'session', 'running') + if reason == 'heartbeat': + session = await state.sessions.get('session') + session.last_heartbeat = time.time() - 31 + await state.sessions.add('session', session) + else: + controller.trainings[identity].created_at = (datetime.now(timezone.utc) - timedelta(seconds=61)).isoformat() + await controller.sweep() + await controller.wait_for_jobs() + if busy: + assert controller.trainings[identity].state == 'ACTIVE' + assert identity in runtime.users + if reason == 'heartbeat': + with pytest.raises(TrainingExpiredError): + await controller.admit(identity, 'owner', 'session', 'new-operation') + await controller.finish_operation(identity, 'running', controller.epoch, executed=True) + await controller.sweep() + await controller.wait_for_jobs() + assert controller.trainings[identity].state == 'EXPIRED' + assert identity not in runtime.users + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_max_lifetime_stops_new_work_while_admitted_work_finishes(): + from datetime import datetime, timedelta, timezone + from tests.server.state.fake_backend import FakeBackend + + controller, state, runtime, _ = await setup( + backend=FakeBackend(), adapter_config={'adapter_max_lifetime': 60}) + identity = await create(controller, state, 'lifetime-admission') + await controller.admit(identity, 'owner', 'session', 'running') + training = controller.trainings[identity] + training.created_at = (datetime.now(timezone.utc) - timedelta(seconds=61)).isoformat() + with pytest.raises(TrainingExpiredError): + await controller.admit(identity, 'owner', 'session', 'new-operation') + assert set(training.pending_operation_ids) == {'running'} + assert identity in runtime.users + await controller.finish_operation(identity, 'running', controller.epoch, executed=True) + await controller.wait_for_jobs() + assert training.state == 'EXPIRED' + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_shutdown_rejects_unconfirmed_release_and_allows_retry(): + from tests.server.state.fake_backend import FakeBackend + + controller, state, runtime, _ = await setup(backend=FakeBackend()) + await create(controller, state, 'shutdown-ack') + runtime.fail_shutdown = True + try: + with pytest.raises(RuntimeError, match='shutdown.*unconfirmed'): + await controller.shutdown() + assert runtime.allocations + assert all(instance.state == 'QUARANTINED' for instance in controller.instances.values()) + finally: + runtime.fail_shutdown = False + await controller.shutdown() + assert not runtime.allocations + assert all(instance.state == 'STOPPED' for instance in controller.instances.values()) + + +@pytest.mark.asyncio +async def test_shutdown_attempts_healthy_peer_after_an_instance_metadata_failure(): + from tests.server.state.fake_backend import FakeBackend + + controller, state, runtime, _ = await setup(backend=FakeBackend()) + first = await create(controller, state, 'shutdown-first') + peer = await create(controller, state, 'shutdown-peer', mode='full') + first_instance = controller.instances[controller.trainings[first].instance_id] + peer_instance = controller.instances[controller.trainings[peer].instance_id] + save = controller.store.save_instance + + async def unavailable(record): + if record.instance_id == first_instance.instance_id and record.state == 'DRAINING': + raise RuntimeError('Instance metadata unavailable') + await save(record) + + controller.store.save_instance = unavailable + try: + with pytest.raises(RuntimeError, match='Instance metadata unavailable'): + await controller.shutdown() + assert peer_instance.state == 'STOPPED' + assert peer_instance.instance_id not in runtime.allocations + assert first_instance.instance_id in runtime.allocations + with pytest.raises(ResourceNotFoundError): + await controller.route(first, 'owner') + finally: + controller.store.save_instance = save + await controller.shutdown() + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_lazy_queries_allocate_nothing(): + controller, state, runtime, _ = await setup() + assert runtime.starts == 0 + assert controller.config.get('small-full', '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') + 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') + 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, state.create_training_store('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') + 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, 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'): + 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'][1]['config']['min_replicas'] = 1 + original['models'][1]['config']['max_replicas'] = 1 + controller.config = TrainingConfig.model_validate(original) + 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.deploy + async def late(instance, spec): + result = await original(instance, spec) + clock.now = instance.startup_deadline + 1 + return result + runtime.deploy = 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]['config']['min_replicas'] = 2 + replacement = TrainingController(TrainingConfig.model_validate(raw), state.create_training_store('unit'), runtime, clock=clock) + with pytest.raises(CapacityUnavailableError): + await replacement.initialize() + 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.lifecycle.startups + 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 + + +@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.futures.try_start + + async def lose_ack(request_id): + assert await start(request_id) + raise ConnectionError('start acknowledgement lost') + + monkeypatch.setattr(state.futures, 'try_start', 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 + 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 + 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 new file mode 100644 index 000000000..115a3c8fc --- /dev/null +++ b/tests/server/training/test_failure_recovery.py @@ -0,0 +1,408 @@ +# 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.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 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 +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] + 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() + previous = state.futures.store_status + + async def broken(request_id, status, *args, **kwargs): + if status == 'queued': + raise RuntimeError('injected queued write failure') + return await previous(request_id, status, *args, **kwargs) + + monkeypatch.setattr(state.futures, 'store_status', broken) + with pytest.raises(RuntimeError, match='queued write'): + await queue.schedule_task(operation, model_id=identity, token='owner', task_type='step', + request_id='queued-failure') + 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.lifecycle.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_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.lifecycle.closing + + +@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(): + 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 + 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.sessions.get_last_heartbeats + + async def slow(identities): + entered.set() + await resume.wait() + return await original(identities) + + monkeypatch.setattr(state.sessions, 'get_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.lifecycle.submit import submit_training_task + controller, state, _, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'explicit-id') + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + try: + ref = await submit_training_task(queue, AsyncMock(return_value='ok'), training_id=identity, seq_id=10, + stream='tinker::step', token='owner', + task_type='step', request_id='explicit-id') + assert ref['request_id'] == 'explicit-id' + assert (await settle(state, 'explicit-id'))['status'] == 'completed' + finally: + await queue.shutdown_task_queue() + 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.model.pool.driver import RayInstanceDriver + from ray import serve + config = TrainingConfig(models=[{'name': 'small', 'model_id': 'weights', 'backend': 'mock', 'device_type': 'cpu', 'mode': 'lora', 'config': {'num_devices': 1, 'device_mesh': {'dp_size': 1}}}]) + store = SimpleNamespace(models=SimpleNamespace(unregister_replica_capacity=AsyncMock())) + 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.runtime.pool, 'remove', remove) + deploying = asyncio.create_task(driver.deploy(instance, driver.config.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.models.unregister_replica_capacity.assert_not_awaited() + finally: + release.set() + await asyncio.gather(deploying, return_exceptions=True) + if closing is not None: + await closing + await driver.runtime.wait_for_cleanup() + assert not driver.runtime.deployments and not driver.runtime.lifecycle.jobs + store.models.unregister_replica_capacity.assert_awaited() + + +@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.models.unregister_replica_capacity(instance.instance_id) + assert await state.get_model_metadata(identity) is not None + assert await state._backend.get('token_count::owner') == 1 + await controller.release(identity, 'owner') + 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 + + controller, state, runtime, _ = await setup(backend=FakeBackend()) + config = controller.config.model_dump() + next(model['config'] for model in config['models'] if model['mode'] == mode).update(min_replicas=1, max_replicas=1) + controller.config = TrainingConfig.model_validate(config) + 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() diff --git a/tests/server/training/test_http_integration.py b/tests/server/training/test_http_integration.py new file mode 100644 index 000000000..3d73695e4 --- /dev/null +++ b/tests/server/training/test_http_integration.py @@ -0,0 +1,37 @@ +import json +import subprocess +import sys +import uuid +from pathlib import Path + + +def test_installed_native_and_tinker_sdks_use_managed_http_routes(): + # This probe owns a separate Ray cluster. Stop the session fixture's cluster + # first so two local GCS/raylet pairs do not evict each other on a slow host. + import ray + ray.shutdown() + try: + # Both SDKs, recovery and ten release/recreate cycles start many fresh + # workers. Allow cold imports from a shared environment; each service + # operation still enforces its own startup/execution deadline. + result = subprocess.run([ + sys.executable, + str(Path(__file__).resolve().parents[3] + / 'tests/server/integration/on_demand_training/http_probe.py') + ], capture_output=True, text=True, timeout=2 * 60 * 60) + 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'] + assert len(set(evidence['native_stable_routing'].values())) == 2 + assert len(set(evidence['tinker_stable_routing'].values())) == 2 + assert evidence['native_lora'] + assert evidence['runtime_lifecycle_rejected'] + assert evidence['tinker_full'] + assert evidence['restart'] and evidence['state_loss'] and evidence['cycles'] == 10 diff --git a/tests/server/training/test_launcher.py b/tests/server/training/test_launcher.py new file mode 100644 index 000000000..7ed539517 --- /dev/null +++ b/tests/server/training/test_launcher.py @@ -0,0 +1,139 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import socket +import asyncio +import threading +import uuid + +import pytest +import ray + +from twinkle.server.config import ServerConfig +from twinkle.server.launcher.server_launcher import ServerLauncher +from twinkle.server.model.pool.actor import controller_name + + +def test_unresponsive_controller_does_not_block_other_cleanup(): + @ray.remote(num_cpus=0) + class Controller: + def __init__(self): + self.entered = asyncio.Event() + self.resume = asyncio.Event() + + async def shutdown(self): + self.entered.set() + await self.resume.wait() + + async def wait_until_entered(self): + await self.entered.wait() + + async def unblock(self): + self.resume.set() + + controller = Controller.remote() + launcher = ServerLauncher(ServerConfig(resources={'shutdown_timeout_s': 0.3})) + launcher._training_controller = controller + finished = threading.Event() + errors, cleaned = [], [] + + def stop(): + try: + launcher._run_cleanup([ + ('blocked controller', launcher._stop_training_controller), + ('other owner', lambda: cleaned.append(True))], raise_errors=True) + except Exception as error: + errors.append(error) + finally: + finished.set() + + worker = threading.Thread(target=stop) + # Ensure actor startup is outside the deadline under test. + ray.get(controller.__ray_ready__.remote(), timeout=10) + entered = controller.wait_until_entered.remote() + worker.start() + try: + ray.get(entered, timeout=10) + assert finished.wait(2), 'Controller shutdown exceeded its budget and blocked the next owner' + assert cleaned == [True] + assert len(errors) == 1 and isinstance(errors[0].__cause__, ray.exceptions.GetTimeoutError) + assert launcher._training_controller is controller + ray.get(controller.unblock.remote(), timeout=5) + launcher._stop_training_controller() + assert launcher._training_controller is None + finally: + if not finished.is_set(): + ray.get(controller.unblock.remote(), timeout=5) + worker.join(5) + ray.kill(controller, no_restart=True) + + +def test_failed_initialization_does_not_leave_an_actor_that_blocks_a_retry(): + namespace = 'launcher-retry-' + uuid.uuid4().hex + with socket.socket() as listener: + 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', 'train_mode': 'full', + 'device_group': {'ranks': devices, 'device_type': 'cpu'}, + 'device_mesh': {'dp_size': devices}, 'min_replicas': 1}}], + } + launcher = ServerLauncher(ServerConfig.model_validate(payload)) + launcher._start_serve() + 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']['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() + + +def test_cleanup_attempts_all_owners_and_reports_the_original_failure(): + calls = [] + error = ConnectionError('shutdown ACK unknown') + + def failed(): + calls.append('failed') + raise error + + def healthy(): + calls.append('healthy') + + with pytest.raises(RuntimeError, match='cleanup is incomplete') as caught: + ServerLauncher._run_cleanup([('first', failed), ('second', healthy)], raise_errors=True) + assert calls == ['failed', 'healthy'] + assert caught.value.__cause__ is error + + +def test_startup_rollback_preserves_startup_error_and_cleans_all_owners(monkeypatch): + config = ServerConfig.model_validate({'applications': [ + {'name': 'gateway', 'import_path': 'server', 'route_prefix': '/api/v1'}]}) + launcher = ServerLauncher(config) + calls = [] + error = ValueError('invalid application') + monkeypatch.setattr('twinkle.server.launcher.server_launcher.apply_ray_serve_patches', lambda: None) + monkeypatch.setattr(launcher, '_configure_environment', lambda: None) + monkeypatch.setattr(launcher, '_init_ray', lambda: None) + monkeypatch.setattr(launcher, '_start_serve', lambda: None) + + def deploy(app): + raise error + + def failed_cleanup(): + calls.append('training') + raise ConnectionError('cleanup pending') + + monkeypatch.setattr(launcher, '_deploy_application', deploy) + monkeypatch.setattr(launcher, '_stop_training_controller', failed_cleanup) + monkeypatch.setattr('twinkle.server.telemetry.flush_telemetry_safely', lambda: calls.append('telemetry')) + with pytest.raises(ValueError) as caught: + launcher.launch() + assert caught.value is error + assert calls == ['training', 'telemetry'] diff --git a/tests/server/training/test_router.py b/tests/server/training/test_router.py new file mode 100644 index 000000000..52e4b1c7c --- /dev/null +++ b/tests/server/training/test_router.py @@ -0,0 +1,97 @@ +# 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.helpers import settle +from twinkle.hub import HubOperation +from twinkle.server.state import ServerState +from twinkle.server.model.pool.router import TrainingRouter +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 = 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) + 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.get_future(response.request_id))['owner_token'] == '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.config = TrainingConfig(models=[{'name': 'small', 'model_id': 'weights', 'backend': 'transformers', 'device_type': 'cuda', 'mode': 'lora', 'config': {'num_devices': 1, 'device_mesh': {'dp_size': 1}}}]) + async def create(payload): + await router.state.store_future_status('created', 'pending', None) + return {'request_id': 'created'} + 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() + + +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.model.pool.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..b136d4fc8 --- /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.helpers import create, setup +from twinkle.server.model.runtime import TrainingRuntime +from twinkle.server.state.training_records import CreateTraining + + +@pytest.mark.asyncio +async def test_configuration_deserialization_does_not_block_event_loop(monkeypatch): + controller, state, _, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'configuration-import', mode='full') + training = controller.trainings[identity] + entered, resume, finished = threading.Event(), threading.Event(), threading.Event() + + def deserialize(config): + entered.set() + resume.wait(3) + finished.set() + return None + + monkeypatch.setattr('twinkle.server.model.runtime.deserialize_object', deserialize) + manager = SimpleNamespace(save=Mock()) + monkeypatch.setattr('twinkle.server.model.runtime.create_training_run_manager', lambda *args, **kw: manager) + + async def serial(factory, task_type): + return await factory() + + runtime = SimpleNamespace(_run_lifecycle_operation=serial, resolve_model_adapter_name=lambda _: '', + _training_states={}, is_full_mode=True) + request = CreateTraining(base_model='small', mode='full', token='owner', session_id='session', create_key='key', + client_type='twinkle', payload={'adapter_name': 'same', 'config': None}) + initializing = asyncio.create_task(TrainingRuntime.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 new file mode 100644 index 000000000..ac89f4b25 --- /dev/null +++ b/tests/server/training/test_runtime_queue.py @@ -0,0 +1,367 @@ +# 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 +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.runtime import TrainingRuntime +from twinkle.server.lifecycle.submit import submit_training_task +from tests.server.training.helpers import create, setup + + +@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() + await controller.wait_for_jobs() + 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} + try: + first, retry = await asyncio.gather(*( + submit_training_task(queue, operation, training_id=identity, seq_id=1, + stream='tinker::forward_backward', token='owner', task_type='forward_backward') + for _ in range(2))) + assert first == retry + assert (await settle(state, first['request_id']))['status'] == 'completed' + step = await submit_training_task(queue, operation, training_id=identity, seq_id=1, + stream='tinker::optim_step', token='owner', task_type='optim_step') + assert step['request_id'] != first['request_id'] + await settle(state, step['request_id']) + assert len(counter) == 2 + finally: + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +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] + 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 = TrainingRuntime() + 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 +@pytest.mark.parametrize('reason', ['heartbeat', 'max_lifetime']) +async def test_adapter_expiry_preserves_already_queued_work(reason): + import time + from datetime import datetime, timedelta, timezone + from tests.server.state.fake_backend import FakeBackend + + controller, state, runtime, _ = await setup( + backend=FakeBackend(), adapter_config={'adapter_timeout': 30, 'adapter_max_lifetime': 60}) + identity = await create(controller, state, 'queued-expiry') + training = controller.trainings[identity] + queue = RuntimeQueue(controller, state, controller.instances[training.instance_id]) + queue._task_queue_config = TaskQueueConfig(execution_timeout=5) + queue._compute_worker._config = queue._task_queue_config + entered, resume, finished = asyncio.Event(), asyncio.Event(), asyncio.Event() + executions = [] + + async def first(): + entered.set() + await resume.wait() + executions.append('first') + + async def second(): + executions.append('second') + + async def notify(*args, **kwargs): + await controller.finish_operation(*args, **kwargs) + if not training.pending_operation_ids: + finished.set() + return True + + queue._training_controller.finish_operation = RPC(notify) + try: + active = await queue.schedule_task(first, model_id=identity, token='owner', task_type='forward_backward') + await asyncio.wait_for(entered.wait(), 2) + queued = await queue.schedule_task(second, model_id=identity, token='owner', task_type='step') + if reason == 'heartbeat': + session = await state.sessions.get('session') + session.last_heartbeat = time.time() - 31 + await state.sessions.add('session', session) + else: + training.created_at = (datetime.now(timezone.utc) - timedelta(seconds=61)).isoformat() + await controller.sweep() + await controller.wait_for_jobs() + assert identity in runtime.users and len(training.pending_operation_ids) == 2 + resume.set() + assert (await settle(state, active['request_id']))['status'] == 'completed' + assert (await settle(state, queued['request_id']))['status'] == 'completed' + await asyncio.wait_for(finished.wait(), 2) + await controller.sweep() + await controller.wait_for_jobs() + assert executions == ['first', 'second'] + assert training.state == 'EXPIRED' and not runtime.allocations + finally: + resume.set() + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_lost_finish_notification_is_reconciled_without_stopping_compute(): + controller, state, runtime, clock = await setup() + 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() + + +@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/server/training/test_storage_recovery.py b/tests/server/training/test_storage_recovery.py new file mode 100644 index 000000000..a0eefdb9f --- /dev/null +++ b/tests/server/training/test_storage_recovery.py @@ -0,0 +1,66 @@ +# 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.model.pool.controller import TrainingController +from tests.server.training.helpers 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, 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 + 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/server/validation/test_preflight.py b/tests/server/validation/test_preflight.py index 2912b1e8e..f6ffe539c 100644 --- a/tests/server/validation/test_preflight.py +++ b/tests/server/validation/test_preflight.py @@ -28,22 +28,23 @@ def __init__(self, backend: str = 'transformers'): self.futures: dict[str, dict] = {} self.backend_calls: list = [] self.claimed: list = [] + self.sequence_retention_s = 3600 self._task_queue_config = SimpleNamespace(effective_execution_timeout=60.0) self.state = SimpleNamespace( - claim_seq=self._claim_seq, + claim_training_seq=self._claim_seq, get_future=self._get_future, - release_seq=self._release_seq, + release_training_seq=self._release_seq, ) - async def _claim_seq(self, dedup_key, request_id, ttl): - self.claimed.append(dedup_key) + async def _claim_seq(self, training_id, stream, sequence, request_id, retention): + self.claimed.append((training_id, stream, sequence)) return None async def _get_future(self, request_id): return self.futures.get(request_id) - async def _release_seq(self, dedup_key): - self.claimed.remove(dedup_key) + async def _release_seq(self, training_id, stream, sequence, request_id): + self.claimed.remove((training_id, stream, sequence)) async def _on_request_start(self, request): return 'token' @@ -82,7 +83,7 @@ def test_megatron_rejects_the_split_gradient_endpoints_with_501(): for capability in (BackendCapability.Forward, BackendCapability.Backward, BackendCapability.CalculateLoss): with pytest.raises(EndpointUnavailableError) as raised: assert_request_supported( - _Deployment('megatron'), model_types.AdapterRequest(adapter_name='a'), capability=capability) + _Deployment('megatron'), model_types.AdapterRequest(adapter_name='a', training_id='training-a'), capability=capability) assert raised.value.error_code == 501 assert 'forward_backward' in str(raised.value), 'the alternative endpoints must be named' @@ -90,14 +91,14 @@ def test_megatron_rejects_the_split_gradient_endpoints_with_501(): def test_transformers_serves_the_split_gradient_endpoints(): assert_request_supported( _Deployment('transformers'), - model_types.AdapterRequest(adapter_name='a'), + model_types.AdapterRequest(adapter_name='a', training_id='training-a'), capability=BackendCapability.Forward) def test_mock_backend_serves_every_endpoint(): """The mock is a test double; restricting it would only break tests.""" assert_request_supported( - _Deployment('mock'), model_types.AdapterRequest(adapter_name='a'), capability=BackendCapability.Forward) + _Deployment('mock'), model_types.AdapterRequest(adapter_name='a', training_id='training-a'), capability=BackendCapability.Forward) @pytest.mark.asyncio @@ -107,7 +108,7 @@ async def test_capability_rejection_writes_no_future_and_calls_no_backend(): await run_submit( deployment, _request(), - model_types.AdapterRequest(adapter_name='a'), + model_types.AdapterRequest(adapter_name='a', training_id='training-a'), task_type='backward', backend_call=_call, capability=BackendCapability.Backward) @@ -123,7 +124,7 @@ async def test_capability_rejection_writes_no_future_and_calls_no_backend(): def test_a_megatron_only_parameter_is_rejected_on_transformers(): with pytest.raises(RequestRejectedError) as raised: - assert_request_supported(_Deployment('transformers'), model_types.SaveRequest(adapter_name='a', merge_lora=True)) + assert_request_supported(_Deployment('transformers'), model_types.SaveRequest(adapter_name='a', merge_lora=True, training_id='training-a')) assert raised.value.error_code == 422 assert 'merge_lora' in str(raised.value) assert 'megatron' in str(raised.value) @@ -132,7 +133,7 @@ def test_a_megatron_only_parameter_is_rejected_on_transformers(): def test_a_transformers_only_parameter_is_rejected_on_megatron(): with pytest.raises(RequestRejectedError): assert_request_supported( - _Deployment('megatron'), model_types.LoadRequest(adapter_name='a', name='ckpt', strict=True)) + _Deployment('megatron'), model_types.LoadRequest(adapter_name='a', name='ckpt', strict=True, training_id='training-a')) def test_an_unset_backend_only_parameter_is_not_rejected(): @@ -142,8 +143,8 @@ def test_an_unset_backend_only_parameter_is_not_rejected(): backend's own default, it would look sent on every request and the other half of the fleet would reject everything. """ - assert_request_supported(_Deployment('transformers'), model_types.SaveRequest(adapter_name='a')) - assert_request_supported(_Deployment('megatron'), model_types.LoadRequest(adapter_name='a', name='ckpt')) + assert_request_supported(_Deployment('transformers'), model_types.SaveRequest(adapter_name='a', training_id='training-a')) + assert_request_supported(_Deployment('megatron'), model_types.LoadRequest(adapter_name='a', name='ckpt', training_id='training-a')) def test_every_restricted_field_is_optional_with_a_none_default(): @@ -178,14 +179,14 @@ def test_a_real_parameter_read_from_kwargs_is_not_rejected(): assert_request_supported( _Deployment(), model_types.SetProcessorRequest( - processor_cls='InputProcessor', adapter_name='a', init_kwargs={'padding_side': 'right'})) + processor_cls='InputProcessor', adapter_name='a', init_kwargs={'padding_side': 'right'}, training_id='training-a')) def test_an_unrecognised_plugin_argument_is_forwarded(): """Plugins accept ``**kwargs``, so "unknown" cannot mean "wrong".""" assert_request_supported( _Deployment(), - model_types.SetLossRequest(loss_cls='DPOLoss', adapter_name='a', init_kwargs={'my_custom_knob': 1})) + model_types.SetLossRequest(loss_cls='DPOLoss', adapter_name='a', init_kwargs={'my_custom_knob': 1}, training_id='training-a')) def test_no_plugin_download_happens_during_validation(monkeypatch): @@ -195,7 +196,7 @@ def test_no_plugin_download_happens_during_validation(monkeypatch): lambda *a, **k: pytest.fail('validation must not download a plugin')) assert_request_supported( _Deployment(), - model_types.SetLossRequest(loss_cls='ms://someone/MyLoss', adapter_name='a', init_kwargs={'beta': 0.1})) + model_types.SetLossRequest(loss_cls='ms://someone/MyLoss', adapter_name='a', init_kwargs={'beta': 0.1}, training_id='training-a')) # --------------------------------------------------------------------------- # @@ -210,32 +211,32 @@ def test_control_fields_are_never_forwarded_to_the_backend(): into a backend signature. """ body = model_types.ForwardBackwardTaskRequest( - inputs=[{'input_ids': [1, 2]}], adapter_name='a', seq_id=3, task='embedding') + inputs=[{'input_ids': [1, 2]}], adapter_name='a', seq_id=3, task='embedding', training_id='training-a') assert backend_kwargs(body) == {'task': 'embedding'} def test_unset_backend_parameters_are_not_forwarded(): - body = model_types.ForwardRequest(inputs=[{'input_ids': [1]}], adapter_name='a') + body = model_types.ForwardRequest(inputs=[{'input_ids': [1]}], adapter_name='a', training_id='training-a') assert backend_kwargs(body) == {} def test_passthrough_contents_are_flattened(): body = model_types.ForwardBackwardTaskRequest( - inputs=[{'input_ids': [1]}], adapter_name='a', loss_kwargs={'advantages': [0.5]}) + inputs=[{'input_ids': [1]}], adapter_name='a', loss_kwargs={'advantages': [0.5]}, training_id='training-a') assert backend_kwargs(body) == {'advantages': [0.5]} def test_a_passthrough_key_shadowing_a_declared_parameter_is_an_error(): """Silently letting one win would make the effective value depend on merge order.""" body = model_types.ForwardRequest(inputs=[{'input_ids': [1]}], adapter_name='a', task='causal_lm', - loss_kwargs={'task': 'embedding'}) + loss_kwargs={'task': 'embedding'}, training_id='training-a') with pytest.raises(ValueError, match='collides'): backend_kwargs(body) def test_plugin_identifiers_are_not_forwarded_as_kwargs(): """The handler passes ``loss_cls`` positionally; forwarding it too would duplicate it.""" - body = model_types.SetLossRequest(loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': 0.1}) + body = model_types.SetLossRequest(loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': 0.1}, training_id='training-a') assert backend_kwargs(body) == {'beta': 0.1} diff --git a/tests/server/validation/test_request_wire.py b/tests/server/validation/test_request_wire.py index 73759d25d..5e0ab0acf 100644 --- a/tests/server/validation/test_request_wire.py +++ b/tests/server/validation/test_request_wire.py @@ -55,7 +55,7 @@ def _post(client: TestClient, path: str, body: dict): def _valid_forward(**overrides) -> dict: - return {'inputs': [{'input_ids': [1, 2, 3]}], 'adapter_name': 'a', **overrides} + return {'training_id': 'training-a', 'inputs': [{'input_ids': [1, 2, 3]}], 'adapter_name': 'a', **overrides} def test_a_valid_body_reaches_the_handler(client): @@ -140,6 +140,7 @@ def test_a_dynamic_plugin_argument_is_accepted_inside_its_region(client): response = _post(client, '/set_loss', { 'loss_cls': 'DPOLoss', 'adapter_name': 'a', + 'training_id': 'training-a', 'init_kwargs': {'beta': 0.1, 'anything_at_all': [1, 2]}, }) assert response.status_code == 200 @@ -154,7 +155,7 @@ def test_a_non_json_value_fails_in_the_client_before_any_request(): """ from pydantic import ValidationError with pytest.raises(ValidationError) as raised: - model_types.SetLossRequest(loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': object()}) + model_types.SetLossRequest(loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': object()}, training_id='training-a') assert 'init_kwargs' in str(raised.value), 'the error must name the field that holds the bad value' @@ -246,3 +247,12 @@ def test_sampler_add_adapter_accepts_a_dict_config_where_model_rejects_it(): if __name__ == '__main__': raise SystemExit(pytest.main([__file__, '-v'])) + + +def test_missing_training_handle_is_rejected_before_handler(client): + body = _valid_forward() + body.pop('training_id') + response = _post(client, '/forward', body) + assert response.status_code == 422 + assert any(detail['field'] == 'training_id' for detail in response.json()['details']) + assert 'reached_handler' not in response.json() diff --git a/tests/transformers/test_multi_lora_checkpoint_context.py b/tests/transformers/test_multi_lora_checkpoint_context.py index 80f4af241..39100bba6 100644 --- a/tests/transformers/test_multi_lora_checkpoint_context.py +++ b/tests/transformers/test_multi_lora_checkpoint_context.py @@ -21,6 +21,7 @@ def test_checkpoint_activates_target_tenant_before_optimizer_state_initializatio wrapper = object.__new__(MultiLoraTransformersModel) wrapper.__dict__['model'] = model wrapper.multi_adapter = manager + wrapper._model_wrapped = True wrapper.strategy = SimpleNamespace(unwrap_model=lambda value: value) with manager.adapter('first'): params = [param for name, param in model.named_parameters() if '.lora_0.' in name] diff --git a/tests/transformers/test_multi_lora_distributed_slots.py b/tests/transformers/test_multi_lora_distributed_slots.py new file mode 100644 index 000000000..ce177078a --- /dev/null +++ b/tests/transformers/test_multi_lora_distributed_slots.py @@ -0,0 +1,161 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Real two-process CPU regression for resident Multi-LoRA slot switching.""" +from datetime import timedelta +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from peft import LoraConfig +from safetensors.torch import save_file + +from twinkle.model.multi_lora import MultiLora +from twinkle.model.transformers.multi_lora_transformers import MultiLoraTransformersModel + + +def _wrapper(module, manager, wrap_model): + wrapper = object.__new__(MultiLoraTransformersModel) + torch.nn.Module.__init__(wrapper) + wrapper.model = module + wrapper.multi_adapter = manager + wrapper._model_wrapped = False + wrapper.optimizer_group = {} + wrapper._enable_expert_parallel = False + wrapper.sp_strategy = None + wrapper.strategy = SimpleNamespace( + capture_pre_ep_state_if_needed=lambda *args, **kwargs: None, + wrap_model=wrap_model, + unwrap_model=lambda model: model.module if isinstance(model, torch.nn.parallel.DistributedDataParallel) else model) + # This linear model does not use expert/sequence parallelism or GDN patches. + for name in ('_maybe_apply_expert_parallel', '_ensure_sp_strategy', '_ensure_gdn_padding_free_patch'): + object.__setattr__(wrapper, name, lambda: None) + return wrapper + + +def _distributed_slots(rank, rendezvous, mode): + dist.init_process_group('gloo', init_method='file://' + rendezvous, rank=rank, world_size=2, + timeout=timedelta(seconds=60)) + try: + torch.set_num_threads(1) + # DDP must synchronize initialization even when workers have different RNG states. + torch.manual_seed(42 + rank) + config = LoraConfig(r=2, lora_alpha=2, target_modules=['0']) + manager = MultiLora(max_loras=2, max_r=2) + module = manager.patch(torch.nn.Sequential(torch.nn.Linear(4, 4)), target_modules=['0'], lora_config=config) + manager.save_initial_weights() + manager.reset_adapter_status() + manager.acquire_lora('first', config) + wrapper = _wrapper(module, manager, lambda model: torch.nn.parallel.DistributedDataParallel( + model, find_unused_parameters=True)) + wrapper._lazy_wrap_model() + assert not any(param.requires_grad for name, param in module.named_parameters() if '.base_layer.' in name) + + def synchronized_weights(tenant): + state = manager.get_state_dict(tenant) + flat = torch.cat([value.detach().flatten() for value in state.values()]) + copies = [torch.empty_like(flat) for _ in range(2)] + dist.all_gather(copies, flat) + torch.testing.assert_close(copies[0], copies[1], rtol=0, atol=1e-7) + return {key: value.detach().clone() for key, value in state.items()} + + for tenant in (('first', 'second') if mode == 'switch' else ('first', 'reused')): + if tenant == 'second': + manager.acquire_lora(tenant, config) + elif tenant == 'reused': + manager.release_lora('first') + manager.acquire_lora(tenant, config) + initial = synchronized_weights(tenant) + with manager.adapter(tenant): + parameters = [p for p in module.parameters() if p.requires_grad] + optimizer = torch.optim.SGD(parameters, lr=.1) + for _ in range(3): + optimizer.zero_grad() + # Different local batches expose missing gradient synchronization. + wrapper.model(torch.ones(2, 4) * (rank + 1)).square().mean().backward() + optimizer.step() + trained = synchronized_weights(tenant) + assert any(not torch.equal(initial[name], value) for name, value in trained.items()) + with manager.adapter(tenant), torch.no_grad(): + fixed = wrapper.model(torch.ones(2, 4)).detach().clone() + # Rank 0 checkpoint must preserve predictions on every DP replica. + checkpoint = [trained if rank == 0 else None] + dist.broadcast_object_list(checkpoint, src=0) + manager.set_state_dict(tenant, checkpoint[0]) + torch.testing.assert_close(wrapper.model(torch.ones(2, 4)), fixed, rtol=0, atol=1e-7) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize('mode', ['switch', 'reuse']) +def test_all_slots_sync_across_dp_ranks_after_switch_and_reuse(tmp_path, mode): + mp.spawn(_distributed_slots, args=(str(tmp_path / 'rendezvous'), mode), nprocs=2, join=True) + + +@pytest.mark.parametrize('options', [{'find_unused_parameters': False}, {'static_graph': True}]) +def test_multi_lora_rejects_static_ddp_reducers(options, monkeypatch): + monkeypatch.setattr(MultiLoraTransformersModel, '_try_init_process_group', lambda _: None) + with pytest.raises(ValueError, match='Multi-LoRA DDP requires'): + MultiLoraTransformersModel(ddp_config=options) + + +@pytest.mark.parametrize('reshard_after_forward', [None, False, 2]) +def test_multi_lora_rejects_retained_or_partial_fsdp_views(reshard_after_forward, monkeypatch): + monkeypatch.setattr(MultiLoraTransformersModel, '_try_init_process_group', lambda _: None) + + def unexpected_strategy(*args): + raise AssertionError('Invalid FSDP config reached strategy initialization') + + monkeypatch.setattr(MultiLoraTransformersModel, '_decide_strategy', unexpected_strategy) + with pytest.raises(ValueError, match='Multi-LoRA FSDP requires reshard_after_forward=True'): + MultiLoraTransformersModel(fsdp_config={'reshard_after_forward': reshard_after_forward}) + + +def test_checkpoint_as_first_operation_preserves_fresh_slot_reset_basis(tmp_path): + config = LoraConfig(r=2, lora_alpha=2, target_modules=['0']) + manager = MultiLora(max_loras=1, max_r=2) + module = manager.patch(torch.nn.Sequential(torch.nn.Linear(4, 4)), target_modules=['0'], lora_config=config) + manager.save_initial_weights() + manager.acquire_lora('first', config) + initial = {name: value.detach().clone() for name, value in manager.get_state_dict('first').items()} + checkpoint = {name: torch.full_like(value, 2) for name, value in initial.items()} + save_file(checkpoint, str(tmp_path / 'adapter_model.safetensors')) + wrapper = _wrapper(module, manager, lambda model: model) + wrapper.optimizer_group = {'first': SimpleNamespace(optimizer=None)} + wrapper.load(str(tmp_path), adapter_name='first') + for name, value in manager.get_state_dict('first').items(): + torch.testing.assert_close(value, checkpoint[name]) + # Further operations must not replace the initial reset basis with trained weights. + wrapper._lazy_wrap_model() + manager.release_lora('first') + manager.acquire_lora('replacement', config) + for name, value in manager.get_state_dict('replacement').items(): + torch.testing.assert_close(value, initial[name]) + + +def test_slot_registration_never_unfreezes_base_parameters_with_lora_in_their_name(): + from collections import OrderedDict + config = LoraConfig(r=2, target_modules=['lora_projection']) + manager = MultiLora(max_loras=2, max_r=2) + module = manager.patch(torch.nn.Sequential(OrderedDict([ + ('lora_projection', torch.nn.Linear(4, 4)), + ])), target_modules=['lora_projection'], lora_config=config) + base = {name: parameter for name, parameter in module.named_parameters() if '.base_layer.' in name} + assert base + # Respect an explicitly frozen base even when its module name contains lora_. + for parameter in base.values(): + parameter.requires_grad_(False) + manager.reset_adapter_status() + mask = {name: parameter.requires_grad for name, parameter in module.named_parameters()} + with pytest.raises(RuntimeError, match='wrapper failed'): + with manager.trainable_slots(): + assert all(parameter.requires_grad for _, parameter in manager._named_slot_parameters(module)) + assert not any(parameter.requires_grad for parameter in base.values()) + raise RuntimeError('wrapper failed') + assert {name: parameter.requires_grad for name, parameter in module.named_parameters()} == mask + manager.acquire_lora('first', config) + wrapper = _wrapper(module, manager, lambda model: model) + for disable_lora in (False, True): + with wrapper._adapter_context('first', disable_lora=disable_lora): + assert not any(parameter.requires_grad for parameter in base.values()) + assert not any(parameter.requires_grad for parameter in base.values()) diff --git a/tests/transformers/test_optimizer_group_regression.py b/tests/transformers/test_optimizer_group_regression.py new file mode 100644 index 000000000..23922c3a1 --- /dev/null +++ b/tests/transformers/test_optimizer_group_regression.py @@ -0,0 +1,48 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Optimizer groups must not depend on DDP wrapping or repeated GA kwargs.""" +from types import SimpleNamespace +from unittest.mock import Mock + +import torch + +from twinkle.model.optimizer_group import BaseOptimizerGroup, TrainStatus +from twinkle.model.transformers.transformers import TransformersModel + + +def test_forward_kwargs_do_not_duplicate_metric_control_parameters(): + metric = Mock() + kwargs = {'gradient_accumulation_steps': 2, 'lr': 99, 'step': 99, 'grad_norm': 99, 'loss_reduction': 'sum'} + status = TrainStatus(inputs=[{}], outputs={'loss': 1.0}, metrics=[metric], forward_kwargs=kwargs) + group = BaseOptimizerGroup(train_status=status, cur_step=3, gradient_accumulation_steps=2) + group.accumulate_metrics(True) + metric.accumulate.assert_called_once_with( + status.inputs, status.outputs, gradient_accumulation_steps=2, lr=[], step=2, grad_norm=0.0, + loss_reduction='sum') + assert kwargs['lr'] == 99 + + +def test_parameter_groups_and_optimizer_restore_are_stable_after_wrapping(): + model = torch.nn.Sequential(torch.nn.Linear(4, 4), torch.nn.LayerNorm(4)) + + class Wrapper(torch.nn.Module): + def __init__(self, module): + super().__init__() + self.module = module + + harness = SimpleNamespace( + model=model, + strategy=SimpleNamespace(unwrap_model=lambda wrapped: getattr(wrapped, 'module', wrapped)), + _get_trainable_parameters=lambda adapter: dict(model.named_parameters()), + ) + original_groups = TransformersModel._create_param_group(harness, '', weight_decay=0.1) + optimizer = torch.optim.Adam(original_groups, lr=1e-4) + model(torch.ones(2, 4)).sum().backward() + optimizer.step() + state = optimizer.state_dict() + harness.model = Wrapper(model) + late_groups = TransformersModel._create_param_group(harness, '', weight_decay=0.1) + assert [g['param_names'] for g in late_groups] == [g['param_names'] for g in original_groups] + assert original_groups[0]['param_names'] == ['0.weight'] + restored = torch.optim.Adam(late_groups, lr=1e-4) + restored.load_state_dict(state) + assert len(restored.state) == len(optimizer.state) diff --git a/tests/transformers/test_resident_training_reset.py b/tests/transformers/test_resident_training_reset.py new file mode 100644 index 000000000..ed1846c5a --- /dev/null +++ b/tests/transformers/test_resident_training_reset.py @@ -0,0 +1,52 @@ +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 + + +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/tests/transformers/test_spectral_hybrid_lora.py b/tests/transformers/test_spectral_hybrid_lora.py index 9e9714249..5ff8ae167 100644 --- a/tests/transformers/test_spectral_hybrid_lora.py +++ b/tests/transformers/test_spectral_hybrid_lora.py @@ -871,11 +871,15 @@ def test_disable_lora_restores_slots_after_operation_fails(): assert all(parameter.requires_grad for name, parameter in model.named_parameters() if 'lora_' in name) -def test_multi_tenant_hybrid_isolation_and_non_destructive_merge(): +@pytest.mark.parametrize('fft_disabled', [False, True]) +def test_multi_tenant_hybrid_isolation_and_non_destructive_merge(fft_disabled): base, model, manager, spectral = _make_multi_tenant_hybrid() inputs = torch.randn(2, 8) with manager.adapter('regular'): regular_before = model(inputs).detach().clone() + if fft_disabled: + # A normal LoRA selection can disable auxiliary FFT wrappers in PEFT. + spectral._get_fft_wrapper('layers.0.self_attn.q_proj').enable_adapters(False) hybrid = manager.find_lora_by_tenant('hybrid') modules = dict(model.named_modules()) diff --git a/tests/twinkle_client/test_async_components.py b/tests/twinkle_client/test_async_components.py index 0a092627c..431ed66d7 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,169 @@ 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'): + if not entered.is_set(): + 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.remove_adapter('same') + 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' + + +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' + + +def test_unknown_training_rejected_before_handshake(monkeypatch): + import pytest + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + calls = [] + _patch_transport(monkeypatch, calls, lambda url: _completed({})) + model = MultiLoraTransformersModel('base') + with pytest.raises(RuntimeError, match='add_adapter_to_model'): + model.step() + assert not calls + + +def test_expired_handle_can_be_removed_and_recreated_on_same_object(monkeypatch): + import pytest + from twinkle_client.exceptions import TwinkleHTTPError + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + calls = [] + _patch_transport(monkeypatch, calls, lambda url: _completed(None)) + model = MultiLoraTransformersModel('base') + model.add_adapter_to_model('same', None) + submit = model._submit + + def expired(*args, **kwargs): + raise TwinkleHTTPError('Training handle expired', status_code=404, error_code=404) + + monkeypatch.setattr(model, '_submit', expired) + with pytest.raises(TwinkleHTTPError) as error: + model.remove_adapter('same') + assert error.value.status_code == 404 + monkeypatch.setattr(model, '_submit', submit) + model.add_adapter_to_model('same', None) + model.zero_grad() + assert calls[-1][1]['training_id'] == model._training_ids['same'] + + +def test_failed_release_keeps_handle_when_outcome_is_unknown(monkeypatch): + import pytest + from twinkle_client.exceptions import TwinkleHTTPError + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + calls = [] + _patch_transport(monkeypatch, calls, lambda url: _completed(None)) + model = MultiLoraTransformersModel('base') + model.add_adapter_to_model('same', None) + identity = model._training_ids['same'] + for error in (TwinkleHTTPError('Unavailable', status_code=503), TimeoutError('ACK unknown')): + def failed(*args, **kwargs): + raise error + monkeypatch.setattr(model, '_submit', failed) + with pytest.raises(type(error)): + model.remove_adapter('same') + assert model._training_ids['same'] == identity + assert model.adapter_name == 'same' + + +def test_sequence_and_handle_selection_are_object_local(monkeypatch): + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + calls = [] + _patch_transport(monkeypatch, calls, lambda url: _completed({})) + first, second = MultiLoraTransformersModel('base'), MultiLoraTransformersModel('base') + for model in (first, second): + model._created = True + model.adapter_name = 'default' + first._training_ids = {'default': 'first-default', 'other': 'first-other'} + second._training_ids = {'default': 'second-default'} + first.step() + first.step(adapter_name='other') + second.step() + assert [(body['adapter_name'], body['training_id'], body['seq_id']) for _, body in calls] == [ + ('default', 'first-default', 1), ('other', 'first-other', 2), ('default', 'second-default', 1)] diff --git a/tests/twinkle_client/test_auto_agent_tools.py b/tests/twinkle_client/test_auto_agent_tools.py index dcfc871e5..747d463e1 100644 --- a/tests/twinkle_client/test_auto_agent_tools.py +++ b/tests/twinkle_client/test_auto_agent_tools.py @@ -81,6 +81,12 @@ 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.train_mode == 'lora' + model = config.training_config.models[0] + assert model.mode == 'lora' + assert model.config.min_replicas == 0 and model.config.max_replicas == 1 + assert model.config.num_devices == 2 and model.config.idle_timeout_s == 600 if samplers: sampler = next(app for app in config.applications if app.import_path == 'sampler') assert sampler.args.nproc_per_node == 4 @@ -102,6 +108,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 diff --git a/tests/twinkle_client/test_request_builder.py b/tests/twinkle_client/test_request_builder.py index 0c7a15bce..f629f108c 100644 --- a/tests/twinkle_client/test_request_builder.py +++ b/tests/twinkle_client/test_request_builder.py @@ -32,7 +32,7 @@ def no_network(monkeypatch): def test_a_declared_name_goes_to_its_field(): - body = build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', task='embedding') + body = build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', task='embedding', training_id='training-a') assert body.task == 'embedding' assert body.loss_kwargs == {} @@ -46,7 +46,7 @@ def test_an_undeclared_name_goes_to_the_single_passthrough_region(): }], adapter_name='a', advantages=[0.5], - old_logps=[[-1.0]]) + old_logps=[[-1.0]], training_id='training-a') assert body.loss_kwargs == {'advantages': [0.5], 'old_logps': [[-1.0]]} @@ -71,7 +71,7 @@ class _TwoRegions(StrictRequest): def test_a_model_without_a_region_rejects_an_unknown_name(): with pytest.raises(TwinkleClientValidationError) as raised: - build_request(model_types.SaveRequest, adapter_name='a', name='ckpt', typo=1) + build_request(model_types.SaveRequest, adapter_name='a', name='ckpt', typo=1, training_id='training-a') assert 'typo' in str(raised.value) @@ -81,7 +81,7 @@ def test_an_explicit_region_and_a_routed_key_are_merged(): loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': 0.1}, - loss_type='sigmoid') + loss_type='sigmoid', training_id='training-a') assert body.init_kwargs == {'beta': 0.1, 'loss_type': 'sigmoid'} @@ -89,7 +89,7 @@ def test_a_key_passed_twice_is_an_error(): """Silently letting one win would make the effective value depend on merge order.""" with pytest.raises(TwinkleClientValidationError, match='both directly and inside'): build_request( - model_types.SetLossRequest, loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': 0.1}, beta=0.2) + model_types.SetLossRequest, loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': 0.1}, beta=0.2, training_id='training-a') def test_an_omitted_optional_argument_is_not_routed_into_the_region(): @@ -98,7 +98,7 @@ def test_an_omitted_optional_argument_is_not_routed_into_the_region(): Routing an explicit ``None`` would hand the backend a null argument it never received before, changing behaviour for callers who simply did not pass anything. """ - body = build_request(model_types.SetLossRequest, loss_cls='DPOLoss', adapter_name='a', unused=None) + body = build_request(model_types.SetLossRequest, loss_cls='DPOLoss', adapter_name='a', unused=None, training_id='training-a') assert body.init_kwargs == {} @@ -109,28 +109,28 @@ def test_an_omitted_optional_argument_is_not_routed_into_the_region(): def test_a_wrongly_typed_field_fails_in_process(): with pytest.raises(ValidationError): - build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', temperature='hot') + build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', temperature='hot', training_id='training-a') def test_an_out_of_range_value_fails_in_process(): with pytest.raises(ValidationError): - build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', temperature=0) + build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', temperature=0, training_id='training-a') def test_malformed_inputs_fail_in_process(): """The client shares the server's schema, so it catches this without asking.""" with pytest.raises(ValidationError): - build_request(model_types.ForwardRequest, inputs=[{'input_ids': [True]}], adapter_name='a') + build_request(model_types.ForwardRequest, inputs=[{'input_ids': [True]}], adapter_name='a', training_id='training-a') def test_a_missing_required_field_fails_in_process(): with pytest.raises(ValidationError): - build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}]) + build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], training_id='training-a') def test_forward_only_requires_an_adapter_context(): with pytest.raises(ValidationError): - build_request(model_types.ForwardOnlyRequest, inputs=[{'input_ids': [1]}]) + build_request(model_types.ForwardOnlyRequest, inputs=[{'input_ids': [1]}], training_id='training-a') # --------------------------------------------------------------------------- # @@ -141,9 +141,9 @@ def test_forward_only_requires_an_adapter_context(): def test_unset_optionals_stay_off_the_wire(): """Absent and "not requested" must look the same, or the server needs its own defaults.""" import json - body = build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a') + body = build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', training_id='training-a') payload = json.loads(request_json(body)) - assert payload == {'inputs': [{'input_ids': [1]}], 'adapter_name': 'a', 'loss_kwargs': {}} + assert payload == {'training_id': 'training-a', 'inputs': [{'input_ids': [1]}], 'adapter_name': 'a', 'loss_kwargs': {}} def test_a_lora_config_is_serialized_to_the_form_the_server_decodes(): diff --git a/tests/twinkle_client/test_types_contract.py b/tests/twinkle_client/test_types_contract.py index 7dbaa25d1..a2d35bbed 100644 --- a/tests/twinkle_client/test_types_contract.py +++ b/tests/twinkle_client/test_types_contract.py @@ -1,7 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Contract-base consistency and naming-disambiguation tests. -- ``QueueStateLiteral`` value set equals the server ``QueueState`` enum. - naming disambiguation guard. The two SDKs already share public names. The contract freezes that legacy set and @@ -12,12 +11,9 @@ import ast import pathlib -import typing import twinkle -from twinkle.server.task_queue.types import QueueState from twinkle.protocol.types import model as model_types -from twinkle.protocol.types.errors import QueueStateLiteral from twinkle.protocol.types.server import GetServerCapabilitiesResponse _TWINKLE_SRC = pathlib.Path(twinkle.__file__).resolve().parent @@ -43,11 +39,6 @@ }) -def test_queue_state_literal_matches_server_enum(): - literal_values = set(typing.get_args(QueueStateLiteral)) - enum_values = {state.value for state in QueueState} - assert literal_values == enum_values, (f'QueueStateLiteral {literal_values} != QueueState {enum_values}') - def test_old_capabilities_response_gets_conservative_defaults(): response = GetServerCapabilitiesResponse.model_validate({'supported_models': []})