diff --git a/.dev_scripts/ci_container_test.sh b/.dev_scripts/ci_container_test.sh index 722e977a4..04cd9d4ac 100644 --- a/.dev_scripts/ci_container_test.sh +++ b/.dev_scripts/ci_container_test.sh @@ -26,7 +26,7 @@ if [ "$MODELSCOPE_SDK_DEBUG" == "True" ]; then pip uninstall autoawq -y pip uninstall lmdeploy -y pip uninstall tensorflow -y - pip install ray==2.48 + pip install ray==2.55.1 pip install optimum # test with install diff --git a/.gitignore b/.gitignore index 6ba7b5381..cc52a2f6a 100644 --- a/.gitignore +++ b/.gitignore @@ -77,6 +77,10 @@ instance/ # Sphinx documentation docs/_build/ +# Internal design notes and developer tools are local-only. +/docs/design/ +/tools/ + # PyBuilder target/ diff --git a/cookbook/client/server/on_demand_training.yaml b/cookbook/client/server/on_demand_training.yaml new file mode 100644 index 000000000..60ff1caba --- /dev/null +++ b/cookbook/client/server/on_demand_training.yaml @@ -0,0 +1,41 @@ +# On-demand training: applications stays the public configuration entry. +training: + allocation_timeout_s: 3 + startup_timeout_s: 1800 + # PPU actor/PlacementGroup death confirmation can exceed the 60s framework default. + shutdown_timeout_s: 180 + sweep_interval_s: 5 +applications: + - name: model-qwen + route_prefix: /api/v1/model/Qwen/Qwen3-1.7B + import_path: model + args: + model_id: ms://Qwen/Qwen3-1.7B + backend: transformers + device_type: cuda + modes: + lora: + num_devices: 2 + device_mesh: {device_type: cuda, dp_size: 2} + min_replicas: 1 + max_loras: 5 + max_r: 32 + idle_timeout_s: 1800 + target_modules: all-linear + full: + num_devices: 4 + device_mesh: {device_type: cuda, dp_size: 4} + min_replicas: 0 + idle_timeout_s: 1800 + - name: model-small + route_prefix: /api/v1/model/Qwen/Qwen3-0.6B + import_path: model + args: + model_id: ms://Qwen/Qwen3-0.6B + backend: megatron + device_type: cuda + modes: + lora: + num_devices: 2 + device_mesh: {device_type: cuda, dp_size: 2} + max_loras: 5 diff --git a/docs/source_en/Usage Guide/Server and Client/Server.md b/docs/source_en/Usage Guide/Server and Client/Server.md index dd0064aba..67a56d4ed 100644 --- a/docs/source_en/Usage Guide/Server and Client/Server.md +++ b/docs/source_en/Usage Guide/Server and Client/Server.md @@ -479,3 +479,11 @@ raises. During that window newly arriving tasks fail fast with a `server`/503 er Each persisted future stores its immutable `absolute_deadline` when it is created. Cleanup therefore reaches the same decision regardless of which deployment process holds the cleanup lease. + +## On-demand training and device reclamation + +Existing `applications` entries with `import_path: model` remain valid. Set `args.min_replicas: 0` for lazy startup and explicitly set `args.idle_timeout_s` for business inactivity. Legacy `adapter_config.adapter_timeout` remains a heartbeat timeout and `adapter_max_lifetime` retains its absolute lifetime. New `modes` default to 1800 seconds of business inactivity; heartbeats and polling do not renew that deadline; queued and running operations protect training state. Empty instances above the resident minimum are destroyed so other models can reuse their devices. Existing deployment min/max replicas become training instance minimums/limits. + +`modes` is optional: single-mode training does not require rewriting existing YAML. Model `deployments` options configure the CPU router; autoscaling request targets apply to router replicas. The Controller owns training-instance minimums/limits, and each compute instance has exactly one Serve replica. Non-autoscaling request limits, health/shutdown settings, and `ray_actor_options.runtime_env` also reach compute instances. `args.env_vars` overrides matching deployment environment variables. Unsupported `adapter_config` fields fail validation instead of being silently ignored. + +To offer both LoRA and full training for one model, configure `num_devices`, `device_mesh`, `min_replicas`, and `idle_timeout_s` separately under `args.modes.lora/full`. LoRA users share compatible instances; each full-training handle owns an instance. Sampling configuration stays the same. See `cookbook/client/server/on_demand_training.yaml` in the repository. diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" index 81fcacfe5..37d808d8f 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" @@ -466,3 +466,11 @@ twinkle-server check-config -c server_config.yaml 任务记录进入终态后，per-replica 准入闸门额外保持关闭的最长时长为 `max(0, Collect_Width × R − T)`：此时记录已是终态，但泄漏的 executor 线程可能仍持有闸门，直到其 `ray.get` 返回或抛出。在该窗口内新到达的任务会以 `server`/503 错误快速失败。 每条持久化 future 在创建时写入不可变的 `absolute_deadline`，因此无论哪个 deployment 进程持有 cleanup lease，清理结果都由任务自身契约决定。 + +## 按需训练与设备回收 + +原有 `applications` 和 `import_path: model` 配置继续有效。单模式配置增加 `args.min_replicas: 0` 即可懒启动，`args.idle_timeout_s` 显式设置训练句柄的业务闲置时间。旧 `adapter_config.adapter_timeout` 仍按心跳判断会话存活，`adapter_max_lifetime` 保留绝对寿命；不再把旧字段转换成业务闲置时间。新增 `modes` 的业务闲置默认 1800 秒，心跳和轮询不延长该期限，运行中或已排队的任务保护训练状态。最后一个用户释放或过期后，超出保底数量的实例立即销毁，释放的卡可用于其他模型。已有 deployment 的 min/max_replicas 转成训练实例的保底/上限。 + +`modes` 是可选扩展，单模式训练不需要改写原 YAML。原模型 `deployments` 的 Serve 设置用于 CPU 路由部署，其中 autoscaling 的请求目标用于路由副本扩缩容；训练计算实例的 min/max 仍由 Controller 管理，每个计算实例固定一个 Serve replica。非 autoscaling 的请求上限、健康/关闭设置和 `ray_actor_options.runtime_env` 会传给计算实例；`args.env_vars` 优先于 deployment 的同名环境变量。未知或未实现的 `adapter_config` 字段会报错，不再静默忽略。 + +同一模型需要同时提供 LoRA 和全参数训练时，在 `args.modes.lora/full` 内分别设置 `num_devices`、`device_mesh`、`min_replicas` 和 `idle_timeout_s`；不用改成另一套顶层模型目录。LoRA 用户共享兼容实例，全参数用户各占一个实例。采样服务继续使用原配置。详见仓库 `cookbook/client/server/on_demand_training.yaml`。 diff --git a/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..d676f4661 100644 --- a/src/twinkle/infra/_ray/ray_helper.py +++ b/src/twinkle/infra/_ray/ray_helper.py @@ -71,13 +71,19 @@ def get_config(self, key: str): def clear(self): self.config.clear() + registry_name = os.environ.get('TWINKLE_RAY_REGISTRY_NAME', 'config_registry') + registry_namespace = os.environ.get('TWINKLE_RAY_NAMESPACE') try: - RayHelper._registry = ray.get_actor('config_registry') + RayHelper._registry = ray.get_actor(registry_name, namespace=registry_namespace) except ValueError: try: - RayHelper._registry = WorkerRegistry.options(name='config_registry', lifetime='detached').remote() + RayHelper._registry = WorkerRegistry.options( + name=registry_name, + namespace=registry_namespace, + lifetime='detached', + num_cpus=0.01 if registry_namespace else 1).remote() except ValueError: - RayHelper._registry = ray.get_actor('config_registry') + RayHelper._registry = ray.get_actor(registry_name, namespace=registry_namespace) assert RayHelper._registry is not None @staticmethod @@ -359,7 +365,7 @@ def create_workers(worker_cls: Type[T], 'scheduling_strategy': PlacementGroupSchedulingStrategy(placement_group=deploy_pg['placement_group']), 'name': worker_name, - 'namespace': 'default', + 'namespace': os.environ.get('TWINKLE_RAY_NAMESPACE', 'default'), 'runtime_env': runtime_env, 'num_cpus': 0.01, } @@ -383,6 +389,8 @@ def create_workers(worker_cls: Type[T], if max_concurrency is not None: worker_options['max_concurrency'] = max_concurrency + if os.environ.get('TWINKLE_RAY_NAMESPACE'): + worker_options.update(max_restarts=0, max_task_retries=0) worker = worker_cls.options(**worker_options).remote(*args, **kwargs) workers.append(worker) else: @@ -417,7 +425,7 @@ def create_workers(worker_cls: Type[T], 'scheduling_strategy': PlacementGroupSchedulingStrategy(placement_group=deploy_pg['placement_group']), 'name': worker_name, - 'namespace': 'default', + 'namespace': os.environ.get('TWINKLE_RAY_NAMESPACE', 'default'), 'runtime_env': runtime_env, 'num_cpus': 0.01, } @@ -425,6 +433,8 @@ def create_workers(worker_cls: Type[T], if max_concurrency is not None: worker_options['max_concurrency'] = max_concurrency + if os.environ.get('TWINKLE_RAY_NAMESPACE'): + worker_options.update(max_restarts=0, max_task_retries=0) worker = worker_cls.options(**worker_options).remote(*args, **kwargs) workers.append(worker) return workers diff --git a/src/twinkle/infra/_ray/resource_manager.py b/src/twinkle/infra/_ray/resource_manager.py index 9b8ba0cef..c9013c972 100644 --- a/src/twinkle/infra/_ray/resource_manager.py +++ b/src/twinkle/infra/_ray/resource_manager.py @@ -20,6 +20,35 @@ def _gpu_placement_group_cpus(node_cpu: int, nproc_per_node: int) -> int: class ResourceManager: + @classmethod + def from_allocation(cls, group: DeviceGroup, placement_group, devices: list[str]): + """Use an allocation owned by the training controller, without reserving again. + + Ranks are local to this allocation; devices are the actual Ray-selected IDs. + The controller owns removal, so this view's destroy method is a no-op. + """ + import ray + manager = cls.__new__(cls) + ranks = list(range(len(devices))) + group.ranks = ranks + manager.group_configs = [group] + manager.placement_groups = [placement_group] + manager.cpu_placement_groups = [] + manager._external_allocation = True + manager.device_groups = { + group.name: [{ + 'gpu_rank': [device], + 'placement_group': placement_group, + 'ray_address': str(ray.get_runtime_context().gcs_address) + } for device in devices] + } + if (group.device_type or '').upper() == 'CPU': + manager.device_groups[group.name] = [{ + 'placement_group': placement_group, + 'ray_address': str(ray.get_runtime_context().gcs_address) + } for _ in devices] + return manager + def __init__(self, nproc_per_node: int, ncpu_proc_per_node: int, groups: List[DeviceGroup]): # CPU placement group default strategy: # - Old approach: use node_cpu//4 as CPU bundle per node, even if only 1~2 CPU processes are needed, @@ -296,6 +325,8 @@ def get_group(self, group: str): return self.device_groups[group] def destroy_placement_group(self): + if getattr(self, '_external_allocation', False): + return import ray for pg in self.placement_groups: ray.util.remove_placement_group(pg) diff --git a/src/twinkle/model/megatron/megatron.py b/src/twinkle/model/megatron/megatron.py index 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/optimizer_group.py b/src/twinkle/model/optimizer_group.py index 150e694bc..59b4b331b 100644 --- a/src/twinkle/model/optimizer_group.py +++ b/src/twinkle/model/optimizer_group.py @@ -64,15 +64,15 @@ def accumulate_metrics(self, is_training): """Accumulate metrics for train/eval status. Override in subclass if needed.""" status = self.train_status if is_training else self.eval_status if len(status.metrics) > 0 and status.inputs is not None and status.outputs is not None: + metric_kwargs = { + **status.forward_kwargs, + 'lr': self._get_lr(), + 'step': self.cur_step - 1, + 'gradient_accumulation_steps': self.gradient_accumulation_steps, + 'grad_norm': self._last_grad_norm, + } for metric in status.metrics: - metric.accumulate( - status.inputs, - status.outputs, - lr=self._get_lr(), - step=self.cur_step - 1, - gradient_accumulation_steps=self.gradient_accumulation_steps, - grad_norm=self._last_grad_norm, - **status.forward_kwargs) + metric.accumulate(status.inputs, status.outputs, **metric_kwargs) def calculate_metrics(self, is_training): """Calculate and return metrics.""" diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 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/model.py b/src/twinkle/protocol/types/model.py index f2e6bd976..1e267d066 100644 --- a/src/twinkle/protocol/types/model.py +++ b/src/twinkle/protocol/types/model.py @@ -28,7 +28,9 @@ class CreateRequest(StrictRequest): - """Body of ``POST /twinkle/create``: a session-establishing no-op.""" + """Body of ``POST /twinkle/create``: CPU-only capability handshake.""" + + training_handles: bool = False # --------------------------------------------------------------------------- # @@ -36,7 +38,13 @@ class CreateRequest(StrictRequest): # --------------------------------------------------------------------------- # -class AdapterRequest(StrictRequest): +class TrainingRequest(StrictRequest): + """A training handle is control metadata, never a backend kwarg.""" + + training_id: str | None = None + + +class AdapterRequest(TrainingRequest): """The shared shape of an adapter-scoped operation. ``seq_id`` is an idempotency key, not a backend parameter: the submit shell @@ -83,7 +91,7 @@ class ClipGradAndStepRequest(ClipGradNormRequest): optim_params: dict[str, JsonValue] | None = backend_kwarg(default=None) -class CalculateMetricRequest(StrictRequest): +class CalculateMetricRequest(TrainingRequest): """Body of ``POST /twinkle/calculate_metric``.""" adapter_name: str @@ -105,7 +113,7 @@ class CalculateMetricRequest(StrictRequest): # --------------------------------------------------------------------------- # -class _InlineForwardBase(StrictRequest): +class _InlineForwardBase(TrainingRequest): """Fields common to the three inline forward endpoints.""" inputs: WireInputBatch @@ -161,7 +169,7 @@ class ForwardBackwardTaskRequest(_InlineForwardBase): # --------------------------------------------------------------------------- # -class DataPlaneForwardRequest(StrictRequest): +class DataPlaneForwardRequest(TrainingRequest): """Body of the ``*_from_data_plane`` forward endpoints.""" input_refs: list[DataRef] = Field(min_length=1) @@ -180,7 +188,7 @@ class DataPlaneForwardRequest(StrictRequest): loss_kwargs: dict[str, JsonValue] = passthrough() -class DataPlaneForwardOnlyRequest(StrictRequest): +class DataPlaneForwardOnlyRequest(TrainingRequest): """Body of ``POST /twinkle/forward_only_from_data_plane``. This endpoint is read-only, so it deliberately has no ``seq_id`` idempotency @@ -222,25 +230,25 @@ def validate_output(self) -> DataPlaneForwardOnlyRequest: # --------------------------------------------------------------------------- # -class SetLossRequest(StrictRequest): +class SetLossRequest(TrainingRequest): loss_cls: str adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() -class SetOptimizerRequest(StrictRequest): +class SetOptimizerRequest(TrainingRequest): optimizer_cls: str adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() -class SetLrSchedulerRequest(StrictRequest): +class SetLrSchedulerRequest(TrainingRequest): scheduler_cls: str adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() -class SetTemplateRequest(StrictRequest): +class SetTemplateRequest(TrainingRequest): """Body of ``POST /twinkle/set_template``. No top-level ``model_id``: the backend always overrides it with its own @@ -254,20 +262,20 @@ class SetTemplateRequest(StrictRequest): init_kwargs: dict[str, JsonValue] = passthrough() -class SetProcessorRequest(StrictRequest): +class SetProcessorRequest(TrainingRequest): processor_cls: str adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() -class AddMetricRequest(StrictRequest): +class AddMetricRequest(TrainingRequest): metric_cls: str adapter_name: str is_training: bool | None = None init_kwargs: dict[str, JsonValue] = passthrough() -class ApplyPatchRequest(StrictRequest): +class ApplyPatchRequest(TrainingRequest): patch_cls: str adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() @@ -278,7 +286,7 @@ class ApplyPatchRequest(StrictRequest): # --------------------------------------------------------------------------- # -class SaveRequest(StrictRequest): +class SaveRequest(TrainingRequest): adapter_name: str name: str | None = None save_optimizer: bool = False @@ -287,7 +295,7 @@ class SaveRequest(StrictRequest): merge_lora: bool | None = backend_only('megatron', default=None) -class LoadRequest(StrictRequest): +class LoadRequest(TrainingRequest): adapter_name: str name: str load_optimizer: bool = False @@ -296,7 +304,7 @@ class LoadRequest(StrictRequest): strict: bool | None = backend_only('transformers', default=None) -class ResumeFromCheckpointRequest(StrictRequest): +class ResumeFromCheckpointRequest(TrainingRequest): """Body of ``POST /twinkle/resume_from_checkpoint``.""" name: str @@ -305,6 +313,7 @@ class ResumeFromCheckpointRequest(StrictRequest): class AddAdapterRequest(StrictRequest): + client_create_id: str | None = None adapter_name: str # ``config`` is None for full-parameter training (no LoRA adapter) and a # serialized LoraConfig string for LoRA training. @@ -418,6 +427,12 @@ class TrainingProgressResponse(ResponseModel): # --- Other responses --- +class TrainingHandleResponse(ResponseModel): + training_id: str + adapter_name: str + + class CreateResponse(ResponseModel): """Response for /create endpoint.""" status: str = 'ok' + training_handles: bool = False diff --git a/src/twinkle/server/__init__.py b/src/twinkle/server/__init__.py index 5bdb1f97c..264552bfc 100644 --- a/src/twinkle/server/__init__.py +++ b/src/twinkle/server/__init__.py @@ -1,7 +1,8 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -from .launcher import ServerLauncher, launch_server +import sys -__all__ = [ - 'ServerLauncher', - 'launch_server', -] +from twinkle._lazy_module import _LazyModule + +sys.modules[__name__] = _LazyModule( + __name__, __file__, {'launcher.server_launcher': ['ServerLauncher', 'launch_server']}, module_spec=__spec__) +sys.modules[__name__].__all__ = ['ServerLauncher', 'launch_server'] diff --git a/src/twinkle/server/config/__init__.py b/src/twinkle/server/config/__init__.py index 3e3fe2f71..9eb6aa09e 100644 --- a/src/twinkle/server/config/__init__.py +++ b/src/twinkle/server/config/__init__.py @@ -5,6 +5,7 @@ ServerArgs, SpectralHybridArgs) from .persistence import PersistenceConfig from .server_config import ServerConfig +from .task_queue import TaskQueueConfig from .telemetry import TelemetryConfig __all__ = [ @@ -18,5 +19,6 @@ 'ServerArgs', 'SpectralHybridArgs', 'ServerConfig', + 'TaskQueueConfig', 'TelemetryConfig', ] diff --git a/src/twinkle/server/config/application_spec.py b/src/twinkle/server/config/application_spec.py index de5c8bc3a..924beece0 100644 --- a/src/twinkle/server/config/application_spec.py +++ b/src/twinkle/server/config/application_spec.py @@ -11,12 +11,13 @@ """ from __future__ import annotations -import math import os -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic import BaseModel, ConfigDict, Field, model_validator from typing import Any, Literal -from twinkle.server.task_queue.config import TaskQueueConfig +from .hybrid import SpectralHybridArgs +from .task_queue import TaskQueueConfig +from .training import TrainingModes # Env var keys the launcher sets from the gateway ``server_config`` so that any # Ray worker (model / sampler / processor), not just the gateway, applies the @@ -55,21 +56,6 @@ class HttpOptions(BaseModel): # ---------- per-deployment args schemas ------------------------------------ # -class SpectralHybridArgs(_ArgsBase): - """Strict server-owned Spectral Hybrid allocation and optimizer defaults.""" - - allocation_path: str = Field(min_length=1) - default_lr_lora: float = Field(default=2.5e-5, gt=0) - default_lr_fft: float = Field(default=1.0e-6, gt=0) - - @field_validator('default_lr_lora', 'default_lr_fft') - @classmethod - def _finite_learning_rate(cls, value: float) -> float: - if not math.isfinite(value): - raise ValueError('learning rate must be finite') - return value - - class ModelArgs(_ArgsBase): """Args for the ``model`` deployment. @@ -83,6 +69,12 @@ class ModelArgs(_ArgsBase): device_group: dict[str, Any] device_mesh: dict[str, Any] backend: Literal['mock', 'transformers', 'megatron'] + device_type: Literal['cuda', 'npu', 'mps', 'cpu'] | None = None + env_vars: dict[str, str] = Field(default_factory=dict) + modes: TrainingModes | None = None + min_replicas: int | None = Field(default=None, ge=0, strict=True) + max_replicas: int | None = Field(default=None, gt=0, strict=True) + idle_timeout_s: float | None = Field(default=None, gt=0, allow_inf_nan=False) train_mode: Literal['lora', 'hybrid', 'full'] = 'lora' strategy: Literal['accelerate', 'native_fsdp'] | None = None adapter_config: dict[str, Any] | None = None @@ -94,8 +86,18 @@ class ModelArgs(_ArgsBase): hybrid: SpectralHybridArgs | None = None data_plane_url: str | None = None + @model_validator(mode='before') + @classmethod + def _mode_resources(cls, data): + if isinstance(data, dict) and data.get('modes') is not None: + return {'device_group': {}, 'device_mesh': {}, **data} + return data + @model_validator(mode='after') def _validate_spectral_backend(self): + if self.modes is not None and any( + mode.idle_timeout_s is None for mode in (self.modes.lora, self.modes.full) if mode is not None): + raise ValueError('modes require a finite positive idle_timeout_s') if self.strategy is not None and self.backend != 'transformers': raise ValueError('strategy is only supported by the transformers backend') if self.hybrid is not None: @@ -269,23 +271,3 @@ def _coerce_args_to_schema(cls, data: Any) -> Any: # ``schema.model_validate`` rejects a non-dict itself with a clean # error, so no separate non-dict guard is needed here. return {**data, 'args': schema.model_validate(raw_args)} - - @model_validator(mode='after') - def _validate_full_mode_single_replica(self) -> ApplicationSpec: - """``train_mode: full`` requires exactly one replica. - - The exclusive-tenant lock lives in per-replica memory - (``ModelManagement._resource_records``), so more than one replica would - silently allow one full-parameter tenant per replica, each rewriting - its own copy of the base weights. Reject that at config-load time. - """ - if not (isinstance(self.args, ModelArgs) and self.args.train_mode == 'full'): - return self - for dep in self.deployments: - num_replicas = dep.get('num_replicas') - max_replicas = (dep.get('autoscaling_config') or {}).get('max_replicas') - if (num_replicas or 1) > 1 or (max_replicas or 1) > 1: - raise ValueError(f"Application '{self.name}': train_mode='full' is an exclusive single-tenant " - 'mode and requires a single replica; set num_replicas/autoscaling_config.' - 'max_replicas to 1.') - return self diff --git a/src/twinkle/server/config/hybrid.py b/src/twinkle/server/config/hybrid.py new file mode 100644 index 000000000..1456a8913 --- /dev/null +++ b/src/twinkle/server/config/hybrid.py @@ -0,0 +1,21 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shared Spectral Hybrid configuration.""" +import math +from pydantic import BaseModel, ConfigDict, Field, field_validator + + +class SpectralHybridArgs(BaseModel): + """Strict server-owned Spectral Hybrid allocation and optimizer defaults.""" + + model_config = ConfigDict(extra='forbid') + + allocation_path: str = Field(min_length=1) + default_lr_lora: float = Field(default=2.5e-5, gt=0) + default_lr_fft: float = Field(default=1.0e-6, gt=0) + + @field_validator('default_lr_lora', 'default_lr_fft') + @classmethod + def _finite_learning_rate(cls, value: float) -> float: + if not math.isfinite(value): + raise ValueError('learning rate must be finite') + return value diff --git a/src/twinkle/server/config/server_config.py b/src/twinkle/server/config/server_config.py index 82bac4d44..31d22b5fb 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 .task_queue import TaskQueueConfig from .telemetry import TelemetryConfig +from .training import TrainingConfig, TrainingOptions +from .training_normalization import normalize_model_application 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) + training: TrainingOptions = Field(default_factory=TrainingOptions) applications: list[ApplicationSpec] = Field(default_factory=list) # ---- loading ---------------------------------------------------------- # @@ -71,8 +74,24 @@ def from_yaml(cls, path: str | Path) -> ServerConfig: def _validate_cross_field(self) -> ServerConfig: if self.persistence.mode == 'redis' and not self.persistence.redis_url: raise ValueError("persistence.redis_url is required when persistence.mode == 'redis'", ) + if self.training_config.models: + static = [app for app in self.applications if app.import_path != 'model'] + names = {app.name for app in static} + for model in self.training_config.models: + if model.name in names: + raise ValueError(f'training model name conflicts with application: {model.name}') + path = model.route_prefix + if any(app.route_prefix.rstrip('/') == path for app in static): + raise ValueError(f'training route conflicts with application: {path}') return self + @property + def training_config(self) -> TrainingConfig: + """Normalize old and extended model args into one runtime catalog.""" + return TrainingConfig( + **self.training.model_dump(), + models=[normalize_model_application(app) for app in self.applications if app.import_path == 'model']) + # ---- round-trip / serialization -------------------------------------- # def to_yaml_dict(self) -> dict[str, Any]: diff --git a/src/twinkle/server/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..d4c372b9c --- /dev/null +++ b/src/twinkle/server/config/training.py @@ -0,0 +1,128 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Administrator-owned training catalog. Parsing never allocates devices.""" +from __future__ import annotations + +from pydantic import BaseModel, ConfigDict, Field, model_validator +from typing import Literal + +from .hybrid import SpectralHybridArgs +from .task_queue import TaskQueueConfig + + +class _TrainingConfig(BaseModel): + model_config = ConfigDict(extra='forbid', frozen=True) + + +class TrainingMesh(_TrainingConfig): + """Sizes accepted by DeviceMesh.from_sizes, including overlapping EP axes.""" + + device_type: str | None = None + dp_size: int = Field(default=1, gt=0, strict=True) + fsdp_size: int | None = Field(default=None, gt=0, strict=True) + tp_size: int | None = Field(default=None, gt=0, strict=True) + pp_size: int | None = Field(default=None, gt=0, strict=True) + cp_size: int | None = Field(default=None, gt=0, strict=True) + ep_size: int | None = Field(default=None, gt=0, strict=True) + etp_size: int = Field(default=1, gt=0, strict=True) + ep_fsdp_size: int | None = Field(default=None, gt=0, strict=True) + vpp_size: int | None = Field(default=None, gt=0, strict=True) + ulysses_size: int | None = Field(default=None, gt=0, strict=True) + sequence_parallel: bool = False + + def world_size(self) -> int: + from twinkle.utils.device_mesh import DeviceMesh + return DeviceMesh.from_sizes(**self.model_dump(exclude_none=True)).world_size + + +class TrainingModeConfig(_TrainingConfig): + num_devices: int = Field(gt=0, strict=True) + device_mesh: TrainingMesh + min_replicas: int = Field(default=0, ge=0, strict=True) + max_replicas: int | None = Field(default=None, gt=0, strict=True) + idle_timeout_s: float | None = Field(default=1800, gt=0, allow_inf_nan=False) + strategy: Literal['accelerate', 'native_fsdp'] | None = None + queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) + max_length: int | None = Field(default=None, gt=0, strict=True) + + @model_validator(mode='after') + def _validate_resources(self): + if self.max_replicas is not None and self.min_replicas > self.max_replicas: + raise ValueError('min_replicas must not exceed max_replicas') + if self.device_mesh.world_size() != self.num_devices: + raise ValueError('device_mesh world_size must equal num_devices') + return self + + +class LoraTrainingConfig(TrainingModeConfig): + max_loras: int = Field(default=5, gt=0, strict=True) + max_r: int = Field(default=32, gt=0, strict=True) + target_modules: str | list[str] = 'all-linear' + hybrid: SpectralHybridArgs | None = None + + +class TrainingModes(_TrainingConfig): + lora: LoraTrainingConfig | None = None + full: TrainingModeConfig | None = None + + @model_validator(mode='after') + def _not_empty(self): + if self.lora is None and self.full is None: + raise ValueError('at least one training mode is required') + return self + + +class LegacyExpiration(_TrainingConfig): + """Internal compatibility policy derived from existing adapter_config fields.""" + session_timeout_s: float | None = Field(default=1800, gt=0, allow_inf_nan=False) + max_lifetime_s: float | None = Field(default=36000, gt=0, allow_inf_nan=False) + + +class TrainingModelConfig(_TrainingConfig): + name: str = Field(min_length=1) + model_id: str = Field(min_length=1) + backend: Literal['mock', 'transformers', 'megatron'] + device_type: Literal['cuda', 'npu', 'mps', 'cpu'] + modes: TrainingModes + legacy_expiration: LegacyExpiration | None = None + route_prefix: str | None = None + data_plane_url: str | None = None + env_vars: dict[str, str] = Field(default_factory=dict) + deploy_options: dict = Field(default_factory=dict) + + @model_validator(mode='after') + def _validate_backend(self): + if (self.name.startswith('/') or any(part in ('', '.', '..') for part in self.name.split('/')) + or any(char.isspace() or char in '?#%\\' for char in self.name)): + raise ValueError('name must be a nonempty relative model path without dot segments') + if self.device_type == 'cpu' and self.backend != 'mock': + raise ValueError('cpu is supported only by the mock training backend') + for mode in (self.modes.lora, self.modes.full): + if mode is None: + continue + if mode.device_mesh.device_type not in (None, self.device_type): + raise ValueError('device_mesh.device_type must match device_type') + if mode.strategy is not None and self.backend != 'transformers': + raise ValueError('strategy requires the transformers backend') + if self.modes.lora and self.modes.lora.hybrid and self.backend != 'transformers': + raise ValueError('hybrid requires the transformers backend') + return self + + +class TrainingOptions(_TrainingConfig): + allocation_timeout_s: float = Field(default=3, gt=0, allow_inf_nan=False) + startup_timeout_s: float = Field(default=1800, gt=0, allow_inf_nan=False) + shutdown_timeout_s: float = Field(default=60, gt=0, allow_inf_nan=False) + sweep_interval_s: float = Field(default=5, gt=0, allow_inf_nan=False) + + +class TrainingConfig(TrainingOptions): + """Internal normalized catalog; public model entries stay in applications.""" + + models: tuple[TrainingModelConfig, ...] = () + + @model_validator(mode='after') + def _unique_names(self): + names = [model.name for model in self.models] + if len(names) != len(set(names)): + raise ValueError('training model names must be unique') + return self diff --git a/src/twinkle/server/config/training_normalization.py b/src/twinkle/server/config/training_normalization.py new file mode 100644 index 000000000..ee38a67c8 --- /dev/null +++ b/src/twinkle/server/config/training_normalization.py @@ -0,0 +1,95 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Translate model application configuration without creating a second runtime.""" +from twinkle.hub.model_alias import extract_model_alias +from .application_spec import ApplicationSpec, ModelArgs +from .training import LegacyExpiration, TrainingModelConfig, TrainingModes + + +def normalize_model_application(app: ApplicationSpec) -> TrainingModelConfig: + args = app.args + assert isinstance(args, ModelArgs) + deployment = {} + if args.modes is not None: + baseline = ModelArgs(model_id=args.model_id, backend=args.backend, device_group={}, device_mesh={}) + common = { + 'model_id', 'backend', 'modes', 'data_plane_url', 'device_group', 'device_mesh', 'device_type', 'env_vars' + } + legacy_fields = { + key + for key in type(args).model_fields if key not in common and getattr(args, key) != getattr(baseline, key) + } + if legacy_fields or args.device_group or args.device_mesh: + raise ValueError(f'{app.name}: modes cannot be mixed with single-mode training parameters') + if app.deployments: + raise ValueError(f'{app.name}: configure instance replicas inside modes, not deployments') + modes = args.modes + legacy_expiration = None + device_types = {mode.device_mesh.device_type for mode in (modes.lora, modes.full) if mode is not None} + device_types.discard(None) + if len(device_types) > 1: + raise ValueError(f'{app.name}: modes must use the same device type') + device_type = args.device_type or next(iter(device_types), 'cpu' if args.backend == 'mock' else 'cuda').lower() + env_vars = args.env_vars + else: + if len(app.deployments) > 1: + raise ValueError(f'{app.name}: a model application must have one deployment configuration') + deployment = app.deployments[0] if app.deployments else {} + unsupported = set(args.adapter_config or {}) - {'adapter_timeout', 'adapter_max_lifetime'} + if unsupported: + raise ValueError(f'{app.name}: unsupported adapter_config fields: {sorted(unsupported)}') + scaling = deployment.get('autoscaling_config') or {} + ranks = args.device_group.get('ranks', args.nproc_per_node) + if isinstance(ranks, list): + raise ValueError(f'{app.name}: use a device count instead of fixed ranks for pooled training') + device_type = args.device_type or str( + args.device_group.get('device_type', + args.device_mesh.get('device_type', + 'cpu' if args.backend == 'mock' else 'cuda'))).lower() + mesh = {**args.device_mesh, 'device_type': device_type} + mode = { + 'num_devices': + ranks, + 'device_mesh': + mesh, + 'min_replicas': + args.min_replicas if args.min_replicas is not None else scaling.get('min_replicas', + deployment.get('num_replicas', 0)), + 'max_replicas': + args.max_replicas if args.max_replicas is not None else scaling.get('max_replicas', + deployment.get('num_replicas')), + 'idle_timeout_s': + args.idle_timeout_s, + 'strategy': + args.strategy, + 'queue_config': + args.queue_config, + 'max_length': + args.max_length, + } + mode_name = 'full' if args.train_mode == 'full' else 'lora' + if mode_name == 'lora': + mode.update( + max_loras=args.max_loras, max_r=args.max_r, target_modules=args.target_modules, hybrid=args.hybrid) + modes = TrainingModes.model_validate({mode_name: mode}) + legacy_expiration = LegacyExpiration( + session_timeout_s=(args.adapter_config or {}).get('adapter_timeout', 1800), + max_lifetime_s=(args.adapter_config or {}).get('adapter_max_lifetime', 36000)) + env_vars = { + **(deployment.get('ray_actor_options') or {}).get('runtime_env', {}).get('env_vars', {}), + **args.env_vars + } + name = extract_model_alias(app.route_prefix) or app.name + return TrainingModelConfig( + name=name, + model_id=args.model_id, + backend=args.backend, + device_type=device_type, + modes=modes, + legacy_expiration=legacy_expiration, + route_prefix=app.route_prefix.rstrip('/') or '/', + data_plane_url=args.data_plane_url, + env_vars=env_vars, + deploy_options={ + key: value + for key, value in deployment.items() if key != 'name' + }) diff --git a/src/twinkle/server/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/env_propagation.py b/src/twinkle/server/env_propagation.py new file mode 100644 index 000000000..9482acc90 --- /dev/null +++ b/src/twinkle/server/env_propagation.py @@ -0,0 +1,37 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Collect environment variables propagated to Ray worker processes.""" +from __future__ import annotations + +import os + +TELEMETRY_ENV_KEYS: tuple[str, ...] = ( + 'TWINKLE_STATE_RAY_NAMESPACE', + 'TWINKLE_TELEMETRY_ENABLED', + 'TWINKLE_TELEMETRY_DEBUG', + 'TWINKLE_TELEMETRY_SERVICE', + 'TWINKLE_TELEMETRY_ENDPOINT', + 'TWINKLE_TELEMETRY_INTERVAL', + 'TWINKLE_MODEL_ID_ALIASES', +) + + +def build_telemetry_env_vars() -> dict[str, str]: + return {key: os.environ[key] for key in TELEMETRY_ENV_KEYS if key in os.environ} + + +def build_persistence_env_vars() -> dict[str, str]: + from twinkle.server.config.persistence import PERSISTENCE_ENV_KEYS + return {key: os.environ[key] for key in PERSISTENCE_ENV_KEYS if key in os.environ} + + +def build_server_state_env_vars() -> dict[str, str]: + from twinkle.server.config.application_spec import SERVER_STATE_ENV_KEYS + return {key: os.environ[key] for key in SERVER_STATE_ENV_KEYS if key in os.environ} + + +def build_propagated_env_vars() -> dict[str, str]: + merged: dict[str, str] = {} + merged.update(build_telemetry_env_vars()) + merged.update(build_persistence_env_vars()) + merged.update(build_server_state_env_vars()) + return merged diff --git a/src/twinkle/server/gateway/app.py b/src/twinkle/server/gateway/app.py index 0c54f506b..d284f83dd 100644 --- a/src/twinkle/server/gateway/app.py +++ b/src/twinkle/server/gateway/app.py @@ -13,10 +13,10 @@ import twinkle.protocol.types as types from twinkle.server.deployment import LazyCleanupMixin, bind_deployment, build_deployment_app +from twinkle.server.service_proxy import ServiceProxy from twinkle.server.state import get_server_state from twinkle.utils.logger import get_logger from .openai_handlers import _register_openai_routes -from .proxy import ServiceProxy from .tinker_handlers import _register_gateway_tinker_routes from .twinkle_handlers import _register_gateway_twinkle_routes @@ -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/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 = '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/api/v1', - ): - self.http_options = http_options or {} - self.route_prefix = route_prefix - # Disable proxy env vars to avoid external routing - self.client = httpx.AsyncClient(timeout=None, trust_env=False) - - async def close(self) -> None: - """Close the underlying httpx.AsyncClient to release connections.""" - await self.client.aclose() - - def _build_target_url(self, service_type: str, base_model: str, endpoint: str) -> str: - """Build the target URL for internal service routing. - - Args: - service_type: Either 'model' or 'sampler' - base_model: The base model name for routing - endpoint: The target endpoint name (already includes tinker/ or twinkle/ prefix) - - Returns: - Complete target URL for the internal service - """ - host = self.http_options.get('host', 'localhost') - port = self.http_options.get('port', 8000) - - if host == '0.0.0.0': - host = 'localhost' - - base_url = f'http://{host}:{port}' - return f'{base_url}{routes.target_url(self.route_prefix, service_type, base_model, endpoint)}' - - def _prepare_headers(self, request_headers) -> dict[str, str]: - """Prepare headers for proxying by removing problematic headers.""" - logger.debug('prepare_headers request_headers=%s', request_headers) - headers = dict(request_headers) - headers.pop('host', None) - headers.pop('content-length', None) - request_id = request_headers.get(H_REQUEST_ID) or request_headers.get(H_REQUEST_ID_LEGACY) - if request_id is not None and not request_headers.get(H_MULTIPLEX_LEGACY): - headers[H_MULTIPLEX_LEGACY] = request_id - if request_id is not None: - headers.setdefault(H_REQUEST_ID, request_id) - headers.setdefault(H_MULTIPLEX, request_id) - return headers - - async def proxy_request( - self, - request: Request, - endpoint: str, - base_model: str, - service_type: str, - *, - body_override: bytes | None = None, - extra_headers: dict[str, str] | None = None, - ) -> Response: - """Generic proxy method to forward requests to model or sampler services. - - Args: - request: The incoming FastAPI request - endpoint: The target endpoint path (e.g., 'tinker/create_model') - base_model: The base model name for routing - service_type: Either 'model' or 'sampler' - body_override: If set, use this as the request body instead of reading from request - extra_headers: Additional headers to merge into the forwarded request - - Returns: - Proxied response from the target service - """ - body_bytes = body_override if body_override is not None else await request.body() - target_url = self._build_target_url(service_type, base_model, endpoint) - headers = self._prepare_headers(request.headers) - if extra_headers: - headers.update(extra_headers) - - # Inject current trace context into outgoing headers for distributed tracing. - # When telemetry is not initialized, this is a noop. - inject_context(headers) - - try: - logger.debug( - 'proxy_request service=%s endpoint=%s target_url=%s request_id=%s', - service_type, - endpoint, - target_url, - headers.get(H_MULTIPLEX_LEGACY), - ) - - response = await self.client.request( - method=request.method, - url=target_url, - content=body_bytes, - headers=headers, - params=request.query_params, - ) - - logger.debug( - 'proxy_response status=%s body_preview=%s', - response.status_code, - response.text[:200], - ) - - return Response( - content=response.content, - status_code=response.status_code, - headers=dict(response.headers), - media_type=response.headers.get('content-type'), - ) - except Exception as e: - logger.error('Proxy error: %s', str(e), exc_info=True) - # The gateway could not reach the upstream deployment. Return the - # unified ``ErrorPayload`` (502/Server) instead of a plain-text body - # so every gateway failure has the same wire shape. Upstream error - # responses are passed through unchanged above, preserving their own - # ErrorPayload body. - request_id = request.headers.get(H_REQUEST_ID) or request.headers.get(H_REQUEST_ID_LEGACY) or '' - payload = ErrorPayload( - error=f'Proxy Error: {str(e)}', - category=ErrorCategory.Server, - error_code=502, - request_id=request_id, - ) - return JSONResponse(status_code=502, content=payload.model_dump(mode='json', exclude_none=True)) - - async def proxy_request_stream( - self, - request: Request, - endpoint: str, - base_model: str, - service_type: str, - *, - body_override: bytes | None = None, - extra_headers: dict[str, str] | None = None, - ): - """Proxy with httpx streaming — returns an async generator of response lines. - - Yields: - str lines from the upstream response (newline-delimited JSON). - - Raises: - httpx.HTTPStatusError: If the upstream returns a non-2xx status. - """ - body_bytes = body_override if body_override is not None else await request.body() - target_url = self._build_target_url(service_type, base_model, endpoint) - headers = self._prepare_headers(request.headers) - if extra_headers: - headers.update(extra_headers) - inject_context(headers) - - logger.debug( - 'proxy_request_stream service=%s endpoint=%s target_url=%s', - service_type, - endpoint, - target_url, - ) - - async with self.client.stream( - method='POST', - url=target_url, - content=body_bytes, - headers=headers, - ) as response: - if response.status_code >= 400: - body = await response.aread() - raise httpx.HTTPStatusError( - f'Upstream returned {response.status_code}', - request=response.request, - response=httpx.Response(status_code=response.status_code, content=body), - ) - async for line in response.aiter_lines(): - if line: - yield line - - async def proxy_to_model(self, request: Request, endpoint: str, base_model: str) -> Response: - """Proxy request to model's tinker endpoint (/tinker/). - - Args: - request: The incoming FastAPI request - endpoint: The tinker endpoint name (e.g., 'create_model', 'forward') - base_model: The base model name for routing - """ - return await self.proxy_request(request, routes.tinker_endpoint(endpoint), base_model, 'model') - - async def proxy_to_sampler(self, request: Request, endpoint: str, base_model: str) -> Response: - """Proxy request to sampler's tinker endpoint (/tinker/). - - Args: - request: The incoming FastAPI request - endpoint: The tinker endpoint name (e.g., 'asample') - base_model: The base model name for routing - """ - return await self.proxy_request(request, routes.tinker_endpoint(endpoint), base_model, 'sampler') diff --git a/src/twinkle/server/gateway/tinker_handlers.py b/src/twinkle/server/gateway/tinker_handlers.py index 93809d780..1a16e45aa 100644 --- a/src/twinkle/server/gateway/tinker_handlers.py +++ b/src/twinkle/server/gateway/tinker_handlers.py @@ -29,6 +29,9 @@ # Keys must equal ``state.models.FAILURE_REASON_CODES`` (guarded by test_envelope). _TINKER_FAILURE_WIRE: dict[str, tuple[int, str]] = { + 'capacity_unavailable': (409, 'user'), + 'startup_failed': (500, 'server'), + 'startup_timeout': (504, 'server'), 'invalid_request': (400, 'user'), 'request_rejected': (400, 'user'), 'resource_not_found': (404, 'user'), @@ -83,7 +86,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 +124,7 @@ async def retrieve_future(request: Request, self: GatewayServer = Depends(self_fn)) -> Any: """Retrieve the result of an async task with long polling.""" request_id = body.request_id + await self.state.assert_future_owner(request_id, get_token_from_request(request)) outcome = await poll_future(self.state, request_id) record = outcome.record if outcome.timed_out: diff --git a/src/twinkle/server/gateway/twinkle_handlers.py b/src/twinkle/server/gateway/twinkle_handlers.py index 3df44ffc8..91962c8d6 100644 --- a/src/twinkle/server/gateway/twinkle_handlers.py +++ b/src/twinkle/server/gateway/twinkle_handlers.py @@ -57,11 +57,14 @@ async def healthz_deep( results = {} all_healthy = True - for model in self.supported_models: + for model in self.training_models: model_name = model.model_name try: resp = await self.proxy.proxy_request(request, 'healthz', model_name, 'model') healthy = (resp.status_code == 200) + if healthy: + import json + healthy = json.loads(resp.body).get('healthy', True) if not healthy: all_healthy = False results[model_name] = { @@ -86,7 +89,7 @@ async def get_server_capabilities( self: GatewayServer = Depends(self_fn), ) -> types.GetServerCapabilitiesResponse: return types.GetServerCapabilitiesResponse( - supported_models=self.supported_models, + supported_models=self.training_models, protocol_version=1, features=types.ClientFeatures( task_envelope=True, @@ -139,6 +142,7 @@ async def retrieve_future( finish. See ``poll_config`` for the numbers. """ request_id = body.request_id + await self.state.assert_future_owner(request_id, get_token_from_request(request)) outcome = await poll_future(self.state, request_id) if outcome.record is None: raise ResourceNotFoundError(f'request_id {request_id} not found or expired') @@ -156,6 +160,7 @@ async def cancel_future( running or already-terminal task is reported but never interrupted, so cancel can never corrupt in-flight GPU/optimizer state. """ + await self.state.assert_future_owner(body.request_id, get_token_from_request(request)) result = await self.state.cancel_future(body.request_id) return types.CancelResponse(**result) diff --git a/src/twinkle/server/launcher/env_propagation.py b/src/twinkle/server/launcher/env_propagation.py 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..7cd6bd125 100644 --- a/src/twinkle/server/launcher/server_launcher.py +++ b/src/twinkle/server/launcher/server_launcher.py @@ -13,16 +13,17 @@ 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.env_propagation import build_propagated_env_vars from .builder_registry import get_builders, resolve_builder -from .env_propagation import build_propagated_env_vars logger = get_logger() @@ -146,6 +147,10 @@ def _deploy_application(self, app_spec: ApplicationSpec) -> None: name = app_spec.name route_prefix = app_spec.route_prefix import_path = app_spec.import_path + if import_path == 'model': + # Model applications are specifications. _deploy_training creates + # their CPU routes and the controller owns all device instances. + return # Shallow-dump the typed args to a kwargs dict WITHOUT recursing into # nested models: ``dict(model)`` yields top-level (field, value) pairs # with nested models left as instances, so ``queue_config`` stays a @@ -187,6 +192,12 @@ def _deploy_application(self, app_spec: ApplicationSpec) -> None: # Pass http_options to server apps for internal proxy routing if import_path == 'server': args.setdefault('http_options', self.config.http_options.model_dump()) + if self.config.training_config.models: + args['training_models'] = [model.name for model in self.config.training_config.models] + sampler_apps = [app for app in self.config.applications if app.import_path == 'sampler'] + args['supported_models'] = list( + dict.fromkeys( + alias for app in sampler_apps if (alias := extract_model_alias(app.route_prefix, 'sampler')))) app = builder(deploy_options=deploy_options, **args) @@ -202,6 +213,9 @@ def launch(self) -> None: """ # Apply Ray Serve patches before initializing Ray apply_ray_serve_patches() + if self.config.training_config.models: + os.environ['TWINKLE_STATE_RAY_NAMESPACE'] = ( + self.ray_namespace or self.config.ray_namespace or 'twinkle_cluster') # Initialize telemetry if configured telemetry = self.config.telemetry @@ -241,6 +255,7 @@ def launch(self) -> None: logger.info(f'ServerState policy exported to worker env: {server_state_env}') model_alias_map = build_model_alias_map(self.config.applications) + model_alias_map.update({model.name: model.model_id for model in self.config.training_config.models}) if model_alias_map: os.environ[MODEL_ID_ALIASES_ENV] = json.dumps(model_alias_map, ensure_ascii=False) else: @@ -250,12 +265,16 @@ def launch(self) -> None: self._start_serve() applications = self.config.applications - if not applications: + if not applications and not self.config.training_config.models: logger.warning('No applications configured') return for app_spec in applications: - self._deploy_application(app_spec) + if app_spec.import_path != 'model': + self._deploy_application(app_spec) + + if self.config.training_config.models: + self._deploy_training() host = self.config.http_options.host port = self.config.http_options.port @@ -287,6 +306,10 @@ def _shutdown(self) -> None: telemetry-shutdown failure cannot mask the user-facing shutdown path. """ from ray import serve + if getattr(self, '_training_controller', None) is not None: + import ray + ray.get(self._training_controller.shutdown.remote()) + ray.kill(self._training_controller, no_restart=True) try: serve.shutdown() logger.info('Ray Serve shut down successfully') @@ -298,6 +321,47 @@ def _shutdown(self) -> None: from twinkle.server.telemetry import flush_telemetry_safely flush_telemetry_safely() + def _deploy_training(self) -> None: + import ray + from ray import serve + + from twinkle.server.training.actor import TrainingControllerActor, controller_name + from twinkle.server.training.router import build_training_router + + namespace = self.ray_namespace or self.config.ray_namespace or 'twinkle_cluster' + env = build_propagated_env_vars() + for key in ('PYTHONPATH', 'TWINKLE_DEFAULT_SAVE_DIR', 'TMPDIR', 'OMP_NUM_THREADS'): + if key in os.environ: + env.setdefault(key, os.environ[key]) + controller = TrainingControllerActor.options( + name=controller_name(namespace), namespace=namespace, runtime_env={ + 'env_vars': env + }).remote(self.config.training_config.model_dump(), namespace) + try: + ray.get(controller.initialize.remote()) + except Exception: + try: + ray.get(controller.shutdown.remote()) + except Exception: + logger.exception('Training controller cleanup failed during initialization rollback') + finally: + ray.kill(controller, no_restart=True) + raise + self._training_controller = controller + for model in self.config.training_config.models: + options = deepcopy(model.deploy_options) + if 'autoscaling_config' not in options: + options.setdefault('num_replicas', 1) + actor_options = options.setdefault('ray_actor_options', {}) + actor_options.setdefault('num_cpus', 0.1) + actor_options.setdefault('runtime_env', {})['env_vars'] = {**env, **model.env_vars} + app = build_training_router(self.config.training_config, model, namespace, + self.config.http_options.model_dump(), options) + serve.run( + app, + name=f'training-router-{namespace}-{model.name}', + route_prefix=model.route_prefix or f'/api/v1/model/{model.name}') + def launch_server( config: ServerConfig | None = None, diff --git a/src/twinkle/server/lifecycle/envelope.py b/src/twinkle/server/lifecycle/envelope.py index 25343c9b0..f089e29ec 100644 --- a/src/twinkle/server/lifecycle/envelope.py +++ b/src/twinkle/server/lifecycle/envelope.py @@ -17,6 +17,9 @@ # Keys must equal ``state.models.FAILURE_REASON_CODES`` (guarded by test_envelope). _FAILURE_WIRE: dict[str, tuple[int, ErrorCategory]] = { + 'capacity_unavailable': (409, ErrorCategory.User), + 'startup_failed': (500, ErrorCategory.Server), + 'startup_timeout': (504, ErrorCategory.Server), 'invalid_request': (400, ErrorCategory.User), 'request_rejected': (400, ErrorCategory.User), 'resource_not_found': (404, ErrorCategory.User), 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/submit.py b/src/twinkle/server/lifecycle/submit.py index 947dfb88e..1c6b78937 100644 --- a/src/twinkle/server/lifecycle/submit.py +++ b/src/twinkle/server/lifecycle/submit.py @@ -168,7 +168,7 @@ async def run_submit( not implemented by every backend. """ token = await self._on_request_start(request) - adapter_name = resolve_twinkle_adapter_name(request, body.adapter_name) + adapter_name = getattr(body, 'training_id', None) or resolve_twinkle_adapter_name(request, body.adapter_name) # ---- Preflight: decidable from the body plus this deployment's backend, so it # runs before any state write and before the enqueue. ---- @@ -197,13 +197,16 @@ async def _task(): dedup_key = None if seq_id is not None: session_id = get_session_id_from_request(request) or request.state.request_id - dedup_key = f'seq::{session_id}::{adapter_name or "-"}::{seq_id}' - # 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) + dedup_key = f'seq::{session_id}::{adapter_name or "-"}::{task_type}::{seq_id}' + training_id = getattr(body, 'training_id', None) + if training_id is not None: + stream = f'native::{session_id}::{task_type}' + retention = self.sequence_retention_s + prior_request_id = await self.state.claim_training_seq(training_id, stream, seq_id, request_id, retention) + else: + # Legacy claims must cover queue wait as well as backend execution. + ttl = int(self.task_queue_config.queue_timeout + self.task_queue_config.effective_execution_timeout) + 60 + prior_request_id = await self.state.claim_seq(dedup_key, request_id, ttl) if prior_request_id is not None: return await self._peek_terminal(prior_request_id, fallback_status='pending') @@ -220,5 +223,8 @@ async def _task(): # there would let a retry enqueue a duplicate -> the exact double-apply this # dedup prevents. When unsure (record exists), keep the claim. if dedup_key is not None and await self.state.get_future(request_id) is None: - await self.state.release_seq(dedup_key) + if training_id is not None: + await self.state.release_training_seq(training_id, stream, seq_id, request_id) + else: + await self.state.release_seq(dedup_key) raise diff --git a/src/twinkle/server/model/app.py b/src/twinkle/server/model/app.py index b98fbb2fb..dfdb7e5b7 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.task_queue import TaskQueueMixin from twinkle.server.utils import wrap_builder_with_device_group_env from twinkle.utils.logger import get_logger +from .managed_runtime import ManagedTrainingRuntime from .tinker_handlers import _register_model_tinker_routes from .twinkle_handlers import _register_model_twinkle_routes @@ -82,14 +79,14 @@ def _make_megatron_model(kw: dict[str, Any]) -> Any: ) -class ModelManagement(LazyCleanupMixin, TaskQueueMixin, AdapterManagerMixin): +class ModelManagement(ManagedTrainingRuntime, TaskQueueMixin, AdapterManagerMixin): """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 + - AdapterManagerMixin provides local registration only; expiry is handled by the training controller. - 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 = kwargs.pop('training_context') + await self.initialize_managed_runtime(self.training_context, kwargs.pop('placement_group'), + kwargs.pop('allocation_devices')) self.device_mesh = init_twinkle_runtime( is_mock=(backend == 'mock'), nproc_per_node=nproc_per_node, device_group=self.device_group, device_mesh_dict=device_mesh, ) - self.replica_id = serve.get_replica_context().replica_id.unique_id + self.replica_id = self.training_context['instance_id'] self.max_loras = kwargs.get('max_loras', 5) self.base_model = model_id # ``train_mode`` selects LoRA multi-tenant (default) or exclusive @@ -130,6 +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._init_adapter_manager() 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.managed_request_start(request) async def shutdown(self) -> None: """Explicit async cleanup — called via FastAPI shutdown event.""" - try: - await self.state.unregister_replica(self.replica_id) - except Exception: - pass await self.shutdown_task_queue() await self.data_plane.close() @@ -199,8 +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,35 +215,8 @@ def resolve_model_adapter_name(self, adapter_name: str | None) -> str | None: return '' return adapter_name - def assert_full_mode_available(self, adapter_name: str | None = None) -> None: - """In full mode, ensure no other tenant currently holds this deployment. - - Call before registering any state so a rejected request leaves nothing - behind. ``adapter_name`` (when given) is excluded from the check so a - tenant re-issuing against its own resource is not blocked. - - Raises: - FullModeBusyError: another (non-expiring) tenant already holds it. - """ - if not self.is_full_mode: - return - holder = self.find_active_resource(exclude=adapter_name) - if holder is not None: - raise FullModeBusyError(holder) - - 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) + raise RuntimeError('Managed training expiry must be handled by the training controller') def build_model_app(model_id: str, @@ -290,7 +225,6 @@ def build_model_app(model_id: str, device_mesh: dict[str, Any], deploy_options: dict[str, Any], backend: str, - adapter_config: dict[str, Any] | None = None, queue_config: TaskQueueConfig | None = None, data_plane_url: str | None = None, **kwargs): @@ -307,7 +241,6 @@ def build_model_app(model_id: str, backend: Model backend selector — ``mock`` | ``transformers`` | ``megatron``. Validated up front; bad values raise :class:`ConfigError` before any side effect. - adapter_config: Adapter lifecycle config (timeout, per-token limits) queue_config: Task queue configuration (rate limiting, etc.) **kwargs: Additional model initialization arguments @@ -336,9 +269,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/managed_runtime.py b/src/twinkle/server/model/managed_runtime.py new file mode 100644 index 000000000..d9f1737d5 --- /dev/null +++ b/src/twinkle/server/model/managed_runtime.py @@ -0,0 +1,220 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Managed training operations on the existing model queue and backend gate.""" +from __future__ import annotations + +import asyncio +import os +import uuid +from contextvars import ContextVar +from datetime import datetime + +from twinkle.protocol.serialize import deserialize_object +from twinkle.protocol.types import model as native +from twinkle.protocol.types.training import CreateModelRequest +from twinkle.protocol.types.training import LoraConfig as RunLoraConfig +from twinkle.server.checkpoint import create_training_run_manager +from twinkle.server.exceptions import ResourceNotFoundError +from twinkle.server.lifecycle.submit import backend_kwargs +from twinkle.server.middleware.auth import get_session_id_from_request, get_token_from_request +from twinkle.server.state.models import FutureFailureRecord +from twinkle.server.state.training_records import CreateTraining, TrainingRecord +from twinkle.server.utils import get_template_for_model +from twinkle.utils.logger import get_logger + +_tinker_operation = ContextVar('managed_tinker_operation', default=None) +logger = get_logger() + + +class ManagedTrainingRuntime: + + async def wait_for_training_workers(self): + actors = getattr(self.model, '_actors', None) or [] + await asyncio.gather(*(actor.__ray_ready__.remote() for actor in actors)) + + async def initialize_managed_runtime(self, context, placement_group, devices): + import ray + + from twinkle.infra._ray.ray_helper import RayHelper + from twinkle.infra._ray.resource_manager import ResourceManager + self._training_controller = ray.get_actor(context['controller_name'], namespace=context['namespace']) + await self._training_controller.register_runtime.remote(context['instance_id'], context['epoch'], + uuid.uuid4().hex) + os.environ['TWINKLE_RAY_REGISTRY_NAME'] = context['registry_name'] + os.environ['TWINKLE_RAY_NAMESPACE'] = context['namespace'] + RayHelper.resource_manager = ResourceManager.from_allocation(self.device_group, placement_group, devices) + + async def managed_request_start(self, request): + token = get_token_from_request(request) + if self._model_unhealthy: + raise ResourceNotFoundError('Training runtime is unavailable') + body = await request.json() + training_id = body.get('training_id') or body.get('model_id') + session = get_session_id_from_request(request) if '/twinkle/' in request.url.path else None + if '/twinkle/' in request.url.path and not session: + raise ResourceNotFoundError('Training requests require the creating session') + route = await self._training_controller.route.remote(training_id, token, session) + if route['instance_id'] != self.training_context['instance_id'] or route['epoch'] != self.training_context[ + 'epoch']: + raise ResourceNotFoundError('Runtime does not own this training handle') + _tinker_operation.set((training_id, body.get('seq_id')) if '/tinker/' in request.url.path else None) + return token + + async def schedule_task(self, coro_factory, **kwargs): + operation = _tinker_operation.get() + if operation is None or operation[1] is None: + return await super().schedule_task(coro_factory, **kwargs) + training_id, seq_id = operation + task_type = kwargs.get('task_type') + stream = f'tinker::{task_type}' + request_id = kwargs.pop('request_id', None) or 'req_' + uuid.uuid4().hex + retention = self.sequence_retention_s + previous = await self.state.claim_training_seq(training_id, stream, seq_id, request_id, retention) + if previous is not None: + return {'request_id': previous, 'model_id': training_id} + try: + return await super().schedule_task(coro_factory, request_id=request_id, **kwargs) + except Exception: + if await self.state.get_future(request_id) is None: + await self.state.release_training_seq(training_id, stream, seq_id, request_id) + raise + + async def _prepare_queued_task(self, task): + if task.model_id is None: + return + context = self.training_context + await self._training_controller.admit.remote( + task.model_id, + task.token, + None, + task.request_id, + renew=task.task_type not in ('get_info', 'get_train_configs')) + + async def enqueued(): + try: + await self._training_controller.commit_operation.remote(task.model_id, task.request_id, + context['epoch']) + except Exception: + logger.exception('Training enqueue notification failed for %s', task.request_id) + + async def finished(executed, unsafe): + if unsafe: + self._model_unhealthy = True + self._training_controller.fail_instance.remote( + context['instance_id'], + FutureFailureRecord( + reason_code='execution_timeout', + attribution='server', + message='Backend execution timed out; runtime is being terminated.')) + else: + try: + await self._training_controller.finish_operation.remote( + task.model_id, task.request_id, context['epoch'], executed=executed) + except Exception: + logger.exception('Training completion notification failed for %s', task.request_id) + + task.on_enqueued, task.on_finished = enqueued, finished + if task.task_type == 'apply_patch': + await self._training_controller.mark_dirty.remote(context['instance_id'], context['epoch']) + + async def _managed_serial(self, factory, task_type): + ref = await self.schedule_task(factory, task_type=task_type) + deadline = asyncio.get_running_loop().time() + self.task_queue_config.absolute_future_ttl(1) + while True: + future = await self.state.get_future(ref['request_id']) + if future and future['status'] in ('failed', 'cancelled'): + raise RuntimeError(future['failure']['message']) + if future and future['status'] == 'completed': + return future['result'] + if asyncio.get_running_loop().time() >= deadline: + raise TimeoutError('Training lifecycle operation has not finished') + await asyncio.sleep(self.task_queue_config.inline_fast_path_timeout) + + async def initialize_training(self, training_data, request_data): + training = TrainingRecord.model_validate(training_data) + request = CreateTraining.model_validate(request_data) + + async def initialize(): + adapter = training.training_id + backend_adapter = self.resolve_model_adapter_name(adapter) + self.register_resource(adapter, training.token, training.session_id) + self.set_resource_state(adapter, 'grad_ready', False) + if request.client_type == 'twinkle': + body = native.AddAdapterRequest.model_validate(request.payload) + # 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._managed_serial(initialize, 'initialize_training') + + async def release_training(self, training_data, reset_full): + training = TrainingRecord.model_validate(training_data) + + async def release(): + adapter = training.training_id + if self.get_resource_info(adapter) is None: + return + if self.is_full_mode: + if reset_full: + await self.call_backend(self.model.reload_initial_weights, reset_training_state=True) + else: + await self.call_backend(self.model.remove_adapter, adapter) + self.clear_resource_state(adapter) + self.unregister_resource(adapter) + + await self._managed_serial(release, 'release_training') + + async def runtime_identity(self): + import ray + actors = getattr(self.model, '_actors', None) or [] + states = await asyncio.to_thread( + lambda: [ray._private.state.actors(actor._actor_id.hex()).get('State') for actor in actors]) + return { + 'pid': os.getpid(), + 'instance_id': self.training_context['instance_id'], + 'unhealthy': self._model_unhealthy or any(state != 'ALIVE' for state in states) + } + + async def verify_operation(self, request_id): + if self._model_unhealthy: + return {'state': 'unsafe'} + if self._compute_worker.has_task(request_id): + return {'state': 'pending'} + future = await self.state.get_future(request_id) + if future and future['status'] in ('completed', 'failed', 'cancelled'): + return { + 'state': 'finished', + 'executed': future['status'] != 'cancelled', + 'finished_at': datetime.fromisoformat(future['updated_at']).timestamp() + } + return {'state': 'missing'} diff --git a/src/twinkle/server/model/tinker_handlers.py b/src/twinkle/server/model/tinker_handlers.py index 60c864a43..cf32eecda 100644 --- a/src/twinkle/server/model/tinker_handlers.py +++ b/src/twinkle/server/model/tinker_handlers.py @@ -20,9 +20,7 @@ from .app import ModelManagement from twinkle.server.checkpoint import create_checkpoint_manager, create_training_run_manager -from twinkle.server.exceptions import FullModeBusyError -from twinkle.server.task_queue.types import UserTaskError -from twinkle.server.utils import get_template_for_model +from twinkle.server.exceptions import RequestRejectedError from twinkle.server.validation.tinker_loss import validate_tinker_loss from twinkle.utils.logger import get_logger @@ -42,60 +40,7 @@ async def create_model( body: types.CreateModelRequest, self: ModelManagement = Depends(self_fn), ) -> types.UntypedAPIFuture: - token = await self._on_request_start(request) - - async def _create_adapter(): - _model_id = None - try: - # Validate lora_config against the deployment's train_mode up front. - if self.is_full_mode and body.lora_config: - raise UserTaskError('This deployment runs in full-parameter (exclusive) mode; do not pass ' - 'lora_config. Use create_full_training_client (or omit lora_config).') - if (not self.is_full_mode) and (not body.lora_config): - raise UserTaskError('This deployment runs in LoRA mode; lora_config is required. ' - 'Use create_lora_training_client.') - # Exclusive full-parameter training: reject early (before touching - # state) if another tenant already holds the deployment. - if self.is_full_mode: - self.assert_full_mode_available() - - _model_id = await self.state.register_model( - body.model_dump(), token=token, replica_id=self.replica_id, session_id=body.session_id) - adapter_name = self.get_adapter_name(adapter_name=_model_id) - model_adapter = self.resolve_model_adapter_name(adapter_name) - # Select template based on model type - template = get_template_for_model(self.base_model) - if self.is_full_mode: - self.register_resource(adapter_name, token, session_id=body.session_id) - await self.call_backend( - self.model.set_template, template, adapter_name=model_adapter, model_id=self.base_model) - await self.call_backend(self.model.set_processor, 'InputProcessor', adapter_name=model_adapter) - await self.call_backend(self.model.set_optimizer, 'Adam', adapter_name=model_adapter) - self.set_resource_state(adapter_name, 'grad_ready', False) - else: - from peft import LoraConfig - lora_cfg = LoraConfig(r=body.lora_config.rank, target_modules='all-linear') - self.register_resource(adapter_name, token, session_id=body.session_id) - await self.call_backend( - self.model.add_adapter_to_model, adapter_name=adapter_name, config_or_dir=lora_cfg) - await self.call_backend( - self.model.set_template, template, adapter_name=adapter_name, model_id=self.base_model) - await self.call_backend(self.model.set_processor, 'InputProcessor', adapter_name=adapter_name) - await self.call_backend(self.model.set_optimizer, 'Adam', adapter_name=adapter_name) - self.set_resource_state(adapter_name, 'grad_ready', False) - training_run_manager = create_training_run_manager(token, client_type='tinker') - training_run_manager.save(_model_id, body) - return types.CreateModelResponse(model_id=_model_id) - except FullModeBusyError as e: - raise UserTaskError(str(e)) from e - except Exception: - if _model_id: - adapter_name = self.get_adapter_name(adapter_name=_model_id) - await self._cleanup_adapter(adapter_name) - logger.error(traceback.format_exc()) - raise - - return await self.schedule_task(_create_adapter, token=token, task_type='create_model') + raise RequestRejectedError('Training lifecycle operations must use the catalog endpoint') @app.post('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/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('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/tinker/forward') async def forward(request: Request, body: types.ForwardRequest, diff --git a/src/twinkle/server/model/twinkle_handlers.py b/src/twinkle/server/model/twinkle_handlers.py index 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('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/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('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/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('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/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..0fb27342b --- /dev/null +++ b/src/twinkle/server/resources/allocation.py @@ -0,0 +1,97 @@ +# 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 + + +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 + + def name(self, epoch: str, instance_id: str) -> str: + return f'resources-{self.namespace}-{epoch}-{instance_id}' + + async def allocate(self, num_devices: int, device_type: str, epoch: str, instance_id: str) -> DeviceAllocation: + import ray + from ray.util.placement_group import placement_group + from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + + if type(num_devices) is not int or num_devices <= 0: + raise ValueError('num_devices must be a positive integer') + resource = 'CPU' if device_type == 'cpu' else Platform.get_platform(device_type).__name__ + count = num_devices + candidates = [node for node in ray.nodes() if node.get('Alive') and node['Resources'].get(resource, 0) >= count] + available = ray.available_resources() + if not candidates or available.get(resource, 0) < count: + raise CapacityUnavailableError(f'No single-node allocation available for {count} {resource} devices') + cpus = min(_gpu_placement_group_cpus(int(node['Resources'].get('CPU', 0)), count) for node in candidates) + bundle = {resource: count, 'CPU': count if resource == 'CPU' else cpus} + name = self.name(epoch, instance_id) + pg = placement_group([bundle], strategy='STRICT_PACK', name=name) + try: + await asyncio.wait_for(pg.ready(), self.allocation_timeout) + options = {'num_cpus': 0.01, 'scheduling_strategy': PlacementGroupSchedulingStrategy(placement_group=pg)} + if resource == 'GPU': + options['num_gpus'] = count + elif resource != 'CPU': + options['resources'] = {resource: count} + node, devices = await asyncio.wait_for( + ray.remote(_probe_allocation).options(**options).remote(resource, device_type), self.shutdown_timeout) + if resource == 'CPU': + devices = [str(index) for index in range(count)] + if len(devices) != count or len(set(devices)) != count: + raise RuntimeError('Ray allocation did not return the requested unique device IDs') + return DeviceAllocation(name, pg.id.hex(), node, tuple(devices)) + except BaseException as exc: + await self.remove(name) + if isinstance(exc, TimeoutError): + raise CapacityUnavailableError( + 'Single-node allocation could not be confirmed; retry a new create later.') from exc + raise + + async def remove(self, name: str) -> None: + from ray.util.placement_group import get_placement_group, placement_group_table, remove_placement_group + try: + pg = get_placement_group(name) + except ValueError: + return + remove_placement_group(pg) + deadline = asyncio.get_running_loop().time() + self.shutdown_timeout + while asyncio.get_running_loop().time() < deadline: + if placement_group_table(pg).get('state') == 'REMOVED': + return + await asyncio.sleep(0.05) + raise TimeoutError(f'Placement group removal was not confirmed: {name}') diff --git a/src/twinkle/server/sampler/app.py b/src/twinkle/server/sampler/app.py index f0ee357df..42e3be7a3 100644 --- a/src/twinkle/server/sampler/app.py +++ b/src/twinkle/server/sampler/app.py @@ -13,11 +13,11 @@ from twinkle import DeviceGroup from twinkle.server.config.backend_dispatch import BackendSelector +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.deployment import LazyCleanupMixin, bind_deployment, build_deployment_app from twinkle.server.middleware.auth import get_token_from_request from twinkle.server.runtime import init_twinkle_runtime from twinkle.server.state import ServerState, get_server_state -from twinkle.server.task_queue.config import TaskQueueConfig from twinkle.server.task_queue.mixin import TaskQueueMixin from twinkle.server.utils import get_template_for_model, wrap_builder_with_device_group_env from twinkle.utils.logger import get_logger diff --git a/src/twinkle/server/service_proxy.py b/src/twinkle/server/service_proxy.py new file mode 100644 index 000000000..478a7d6cf --- /dev/null +++ b/src/twinkle/server/service_proxy.py @@ -0,0 +1,108 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shared HTTP forwarding between services; independent of gateway and training controller.""" +from __future__ import annotations + +import httpx +from fastapi import Request, Response +from fastapi.responses import JSONResponse +from typing import Any + +from twinkle.protocol.headers import H_MULTIPLEX, H_MULTIPLEX_LEGACY, H_REQUEST_ID, H_REQUEST_ID_LEGACY +from twinkle.protocol.types.errors import ErrorCategory, ErrorPayload +from twinkle.server.config.task_queue import TaskQueueConfig +from twinkle.server.telemetry.tracing import inject_context +from twinkle.utils.logger import get_logger + +logger = get_logger() + + +class ServiceProxy: + + def __init__(self, http_options: dict[str, Any] | None = None, route_prefix: str = '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/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/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..dc180e7d6 100644 --- a/src/twinkle/server/state/backend/redis_backend.py +++ b/src/twinkle/server/state/backend/redis_backend.py @@ -6,10 +6,13 @@ from collections.abc import Callable from typing import Any +from twinkle.server.telemetry.metrics import get_state_backend_metrics from .base import ConcurrencyError, StateBackend try: import redis.asyncio as aioredis + from redis.exceptions import ConnectionError as RedisConnectionError + from redis.exceptions import TimeoutError as RedisTimeoutError from redis.exceptions import WatchError _REDIS_AVAILABLE = True @@ -38,6 +41,7 @@ def __init__(self, redis_url: str, key_prefix: str = '') -> None: raise ImportError('redis package required. Install with: pip install redis') self._client = aioredis.from_url(redis_url, decode_responses=True) self._prefix = key_prefix + self._metrics = get_state_backend_metrics() def _make_key(self, key: str) -> str: """Add namespace prefix to key.""" @@ -56,10 +60,27 @@ async def set(self, key: str, value: Any, ttl: int | None = None) -> None: else: await self._client.set(real_key, data) + async def _read_with_reconnect(self, operation, *args): + """Retry only read commands after a stale pooled connection is discarded. + + Writes/EXEC are deliberately excluded: a lost acknowledgement does not + prove that a mutation failed, so replay could apply it twice. + """ + for attempt in range(3): + try: + result = await operation(*args) + self._metrics.redis_health.set(1) + return result + except (RedisConnectionError, RedisTimeoutError): + self._metrics.redis_health.set(0) + if attempt == 2: + raise + await asyncio.sleep(0.05 * (attempt + 1)) + async def get(self, key: str) -> Any | None: """Retrieve value, return None if not found or expired.""" real_key = self._make_key(key) - raw = await self._client.get(real_key) + raw = await self._read_with_reconnect(self._client.get, real_key) if raw is None: return None return json.loads(raw) @@ -72,7 +93,7 @@ async def delete(self, key: str) -> None: async def exists(self, key: str) -> bool: """Check if key exists and is not expired.""" real_key = self._make_key(key) - return bool(await self._client.exists(real_key)) + return bool(await self._read_with_reconnect(self._client.exists, real_key)) async def keys(self, pattern: str) -> list[str]: """Return all key names matching the pattern using SCAN. @@ -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._read_with_reconnect(self._client.mget, real_keys) return [json.loads(v) if v is not None else None for v in raw_values] async def close(self) -> None: @@ -157,6 +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._read_with_reconnect(self._client.ping)) except Exception: - return False + healthy = False + self._metrics.redis_health.set(1 if healthy else 0) + return healthy diff --git a/src/twinkle/server/state/cleanup_coordinator.py b/src/twinkle/server/state/cleanup_coordinator.py index c6278c248..419846181 100644 --- a/src/twinkle/server/state/cleanup_coordinator.py +++ b/src/twinkle/server/state/cleanup_coordinator.py @@ -71,6 +71,7 @@ def __init__( self._on_become_leader = on_become_leader self._on_lose_leader = on_lose_leader + self._lifecycle_lock = asyncio.Lock() self._cleanup_task: asyncio.Task | None = None self._cleanup_running = False self._leader_id = uuid.uuid4().hex @@ -143,16 +144,14 @@ async def _cleanup_loop(self) -> None: async def _leader_loop(self) -> None: """Acquire and renew the cleanup-leader lease every LEASE_RENEW seconds.""" - await self._try_acquire_or_renew() # Race for leadership at startup while self._leader_running: try: - await asyncio.sleep(LEASE_RENEW) await self._try_acquire_or_renew() except asyncio.CancelledError: break except Exception as e: logger.warning(f'[ServerState Leader] renew error: {e}') - continue + await asyncio.sleep(LEASE_RENEW) async def _try_acquire_or_renew(self) -> None: was_leader = self._is_leader @@ -169,17 +168,8 @@ async def _try_acquire_or_renew(self) -> None: except Exception as e: logger.warning(f'[ServerState Leader] backend error during election: {e}') self._is_leader = False - if was_leader: - # Our renewal failed but our lease value may still be sitting in the - # backend, so a plain ``set_nx`` would keep returning False for up to - # LEASE_TTL and leadership would stall unclaimed. Best-effort delete - # ONLY when we were the leader (never steal a lease another replica - # legitimately holds), swallowing errors so a delete failure cannot - # escape the election loop. The next tick can then re-acquire. - try: - await self._backend.delete(LEADER_KEY) - except Exception: - pass + # A renewal failure does not prove current ownership. Leave the key + # untouched: our lease expires naturally, or a successor owns it. if self._is_leader and not was_leader: logger.info(f'[ServerState] became cleanup leader (id={self._leader_id[:8]})') @@ -199,36 +189,32 @@ async def start(self) -> bool: ``ServerState``) so the implementation and its guard can never drift apart into a double-start. """ - if self._cleanup_running: - return False - # Rebuild in-memory indexes from backend data before the loops start. - await self._managers['models'].rebuild_indexes() - self._cleanup_running = True - self._cleanup_task = asyncio.create_task(self._cleanup_loop()) - self._leader_running = True - self._leader_task = asyncio.create_task(self._leader_loop()) - return True + async with self._lifecycle_lock: + if self._cleanup_running: + return False + await self._managers['models'].rebuild_indexes() + self._cleanup_running = True + self._leader_running = True + self._cleanup_task = asyncio.create_task(self._cleanup_loop()) + self._leader_task = asyncio.create_task(self._leader_loop()) + return True async def stop(self) -> bool: """Stop the background cleanup + leader-election tasks. Returns ``False`` if not running.""" - if not self._cleanup_running: - return False - self._cleanup_running = False - if self._cleanup_task: - self._cleanup_task.cancel() - self._cleanup_task = None - self._leader_running = False - if self._leader_task: - self._leader_task.cancel() - self._leader_task = None - if self._is_leader: - # Release callback registration; the lease itself expires on its own TTL — - # update_atomic can't express "atomic delete", so we accept a short outage - # where the gauge reads 0 between leaders. - if self._on_lose_leader is not None: + async with self._lifecycle_lock: + if not self._cleanup_running: + return False + self._cleanup_running = self._leader_running = False + tasks = [task for task in (self._cleanup_task, self._leader_task) if task is not None] + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + self._cleanup_task = self._leader_task = None + was_leader, self._is_leader = self._is_leader, False + # Do not delete an uncertain lease; its original TTL is the fence. + if was_leader and self._on_lose_leader is not None: await self._on_lose_leader() - self._is_leader = False - return True + return True async def get_cleanup_stats(self) -> dict[str, Any]: """Get current cleanup configuration and resource counts.""" diff --git a/src/twinkle/server/state/future_manager.py b/src/twinkle/server/state/future_manager.py index 54bca4c70..d3aaae9b9 100644 --- a/src/twinkle/server/state/future_manager.py +++ b/src/twinkle/server/state/future_manager.py @@ -14,7 +14,7 @@ # Status sets used by the do-not-regress guard inside the atomic transform. _TERMINAL_STATUSES = frozenset({'completed', 'failed', 'cancelled'}) -_NON_TERMINAL_STATUSES = frozenset({'pending', 'queued', 'running'}) +_STATUS_ORDER = {'pending': 0, 'queued': 1, 'running': 2} def _future_record_transform( @@ -48,6 +48,11 @@ def _future_record_transform( if new_status != existing_status: logger.warning('future %s already terminal as %r; refusing %r', request_id, existing_status, new_status) return None + if new_status in _STATUS_ORDER and existing_status in _STATUS_ORDER: + if _STATUS_ORDER[new_status] <= _STATUS_ORDER[existing_status]: + return None + if new_status == 'cancelled' and existing_status == 'running': + return None if existing is None: record = FutureRecord( @@ -91,6 +96,13 @@ def _future_record_transform( _CANCELLABLE_STATUSES = frozenset({'pending', 'queued'}) +def _start_if_pending_transform(existing: dict | None, *, now: str) -> dict | None: + """Acquire execution only from pending/queued, using the existing record.""" + if existing is None or existing.get('status') not in _CANCELLABLE_STATUSES: + return None + return {**existing, 'status': 'running', 'queue_state': 'active', 'updated_at': now} + + def _cancel_if_not_started_transform(existing: dict | None, *, failure: dict, now: str) -> dict | None: """Atomic transform: cancel iff the task has not started (pending/queued). @@ -168,6 +180,14 @@ async def store_status( ), ) + async def try_start(self, request_id: str) -> bool: + """True only when this attempt atomically acquired execution rights.""" + return await self._backend.update_atomic( + self._make_key(request_id), + functools.partial(_start_if_pending_transform, now=_now_iso()), + return_changed=True, + ) + async def cancel_if_pending(self, request_id: str) -> str | None: """Cancel a task iff it has not started; return the resulting status. diff --git a/src/twinkle/server/state/model_manager.py b/src/twinkle/server/state/model_manager.py index aaf520fca..de8f872c7 100644 --- a/src/twinkle/server/state/model_manager.py +++ b/src/twinkle/server/state/model_manager.py @@ -94,6 +94,7 @@ async def rebuild_indexes(self) -> None: async def get_capacity_info(self) -> dict[str, int]: """Return global LoRA capacity across all registered replicas.""" replicas = await self._replicas.get_all() + replicas = {identity: capacity for identity, capacity in replicas.items() if capacity > 0} loaded_per_replica = await self._loaded_per_replica() total_max = sum(replicas.values()) total_used = sum(loaded_per_replica.get(rid, 0) for rid in replicas) @@ -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..ced4475e1 100644 --- a/src/twinkle/server/state/models.py +++ b/src/twinkle/server/state/models.py @@ -35,6 +35,9 @@ class ModelRecord(BaseModel): lora_config: Any = None replica_id: str | None = None created_at: str = Field(default_factory=_now_iso) + controller_epoch: str | None = None + training_namespace: str | None = None + instance_id: str | None = None class SamplingSessionRecord(BaseModel): @@ -61,6 +64,9 @@ class FutureFailureRecord(BaseModel): # gateways each map this same key set to their own wire vocabularies; keeping the set # in one place lets a consistency test catch a map that drifts out of coverage. FAILURE_REASON_CODES: frozenset[str] = frozenset({ + 'capacity_unavailable', + 'startup_failed', + 'startup_timeout', 'invalid_request', 'request_rejected', 'resource_not_found', diff --git a/src/twinkle/server/state/sequence.py b/src/twinkle/server/state/sequence.py new file mode 100644 index 000000000..78efab55b --- /dev/null +++ b/src/twinkle/server/state/sequence.py @@ -0,0 +1,27 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Bounded training retry receipts with a persistent expired-sequence fence.""" + + +def claim_sequence(current, *, sequence, request_id, now, retention): + floor = current.get('floor') if current else None + receipts = {} + for key, receipt in (current or {}).get('receipts', {}).items(): + if receipt['expires_at'] <= now: + floor = int(key) if floor is None else max(floor, int(key)) + else: + receipts[key] = receipt + key = str(sequence) + # Unexpired receipts remain retryable even if another sequence expired first. + if key not in receipts and (floor is None or sequence > floor): + receipts[key] = {'request_id': request_id, 'expires_at': now + retention} + return {'floor': floor, 'receipts': receipts} + + +def release_sequence(current, *, sequence, request_id): + if current is None: + return None + receipts = dict(current['receipts']) + key = str(sequence) + if receipts.get(key, {}).get('request_id') == request_id: + receipts.pop(key) + return {'floor': current['floor'], 'receipts': receipts} diff --git a/src/twinkle/server/state/server_state.py b/src/twinkle/server/state/server_state.py index dd926aa57..dc9d683f2 100644 --- a/src/twinkle/server/state/server_state.py +++ b/src/twinkle/server/state/server_state.py @@ -10,7 +10,7 @@ from typing import Any from twinkle.server.config.persistence import PersistenceConfig -from twinkle.server.exceptions import ResourceQuotaExceededError +from twinkle.server.exceptions import RequestRejectedError, ResourceNotFoundError, ResourceQuotaExceededError from twinkle.server.telemetry.correlation import (BASE_MODEL, MODEL_ID, REPLICA_ID, SAMPLING_SESSION_ID, SESSION_ID, TOKEN_ID) from twinkle.server.telemetry.tracing import traced_operation @@ -173,6 +173,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() + def create_training_store(self, namespace: str): + from .training_store import TrainingStore + return TrainingStore(self, namespace, self._backend) + + async def save_model_record(self, model_id: str, record: ModelRecord) -> None: + await self._model_mgr.save_record(model_id, record) + + async def get_training_records(self, namespace: str): + return await self._model_mgr.get_training_records(namespace) + + async def set_future_owner(self, request_id: str, token: str) -> None: + await self._backend.set(f'future_owner::{request_id}', token, ttl=max(1, int(self.expiration_timeout))) + + async def delete_training_sequences(self, training_id: str) -> None: + for key in await self._backend.keys(f'seq_ledger::{training_id}::*'): + await self._backend.delete(key) + + async def unregister_replica_capacity(self, replica_id: str) -> None: + await self._model_mgr.unregister_replica_capacity(replica_id) + + async def register_model_record(self, model_id: str, record: ModelRecord) -> None: + """Reserve quota for a complete model record, including managed training state.""" + await self._model_mgr.add(model_id, record) + + async def rebuild_model_indexes(self) -> None: + await self._model_mgr.rebuild_indexes() + # ----- Session Management ----- async def create_session(self, payload: dict[str, Any]) -> str: @@ -213,6 +240,9 @@ async def get_session_last_heartbeat(self, session_id: str) -> float | None: """ return await self._session_mgr.get_last_heartbeat(session_id) + async def get_session_last_heartbeats(self, session_ids: list[str]) -> dict[str, float | None]: + return await self._session_mgr.get_last_heartbeats(session_ids) + # ----- Model Registration ----- async def register_model(self, @@ -429,7 +459,12 @@ async def get_future(self, request_id: str) -> dict[str, Any] | None: record = await self._future_mgr.get(request_id) return record.model_dump() if record is not None else None - async def claim_seq(self, dedup_key: str, request_id: str, ttl: int) -> str | None: + async def assert_future_owner(self, request_id: str, token: str) -> None: + owner = await self._backend.get(f'future_owner::{request_id}') + if owner is not None and owner != token: + raise ResourceNotFoundError('Unknown future') + + async def claim_seq(self, dedup_key: str, request_id: str, ttl: int | None) -> str | None: """Idempotency claim for a client seq_id. Atomically records ``dedup_key -> request_id`` if unseen and returns ``None`` @@ -446,6 +481,27 @@ async def release_seq(self, dedup_key: str) -> None: preflight rejected it) so a retry can be admitted rather than see a phantom.""" await self._backend.delete(dedup_key) + async def claim_training_seq(self, training_id: str, stream: str, sequence: int, request_id: str, + retention: float) -> str | None: + from .sequence import claim_sequence + ledger = await self._backend.update_atomic( + f'seq_ledger::{training_id}::{stream}', + functools.partial( + claim_sequence, sequence=sequence, request_id=request_id, now=time.time(), retention=retention)) + receipt = ledger['receipts'].get(str(sequence)) + if receipt is None: + raise RequestRejectedError('Sequence retry window expired; the operation will not be executed again') + return None if receipt['request_id'] == request_id else receipt['request_id'] + + async def release_training_seq(self, training_id: str, stream: str, sequence: int, request_id: str) -> None: + from .sequence import release_sequence + await self._backend.update_atomic(f'seq_ledger::{training_id}::{stream}', + functools.partial(release_sequence, sequence=sequence, request_id=request_id)) + + async def try_start_future(self, request_id: str) -> bool: + """Acquire task execution atomically, without inferring ownership from reads.""" + return await self._future_mgr.try_start(request_id) + async def cancel_future(self, request_id: str) -> dict[str, Any]: """Best-effort cancel: drop the task iff it has not started running. diff --git a/src/twinkle/server/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..e89eb3d3d --- /dev/null +++ b/src/twinkle/server/state/training_records.py @@ -0,0 +1,58 @@ +# 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 .models import ModelRecord + +TrainingState = Literal['PENDING', 'ACTIVE', 'CLOSING', 'RELEASED', 'EXPIRED', 'FAILED', 'CANCELLED'] +InstanceState = Literal['ALLOCATING', 'STARTING', 'READY', 'RESETTING', 'DRAINING', 'STOPPED', 'QUARANTINED'] +TERMINAL_TRAINING = frozenset({'RELEASED', 'EXPIRED', 'FAILED', 'CANCELLED'}) +LIVE_INSTANCES = frozenset({'ALLOCATING', 'STARTING', 'READY', 'RESETTING'}) + + +class TrainingRecord(ModelRecord): + training_id: str + controller_epoch: str + instance_id: str + mode: Literal['lora', 'full'] + adapter_label: str | None = None + state: TrainingState = 'PENDING' + create_request_id: str + creation_deadline: float + last_activity_at: float + pending_operation_ids: dict[str, dict[str, Any]] = Field(default_factory=dict) + terminal_reason: str | None = None + terminal_at: float | None = None + + +class InstanceRecord(BaseModel): + instance_id: str + controller_epoch: str + base_model: str + mode: Literal['lora', 'full'] + resident: bool = False + state: InstanceState = 'ALLOCATING' + training_ids: list[str] = Field(default_factory=list) + startup_deadline: float + application_name: str | None = None + placement_group_id: str | None = None + node_id: str | None = None + devices: list[str] = Field(default_factory=list) + incarnation: str | None = None + dirty: bool = False + failure: str | None = None + stopped_at: float | None = None + + +class CreateTraining(BaseModel): + base_model: str + mode: Literal['lora', 'full'] + token: str + session_id: str + create_key: str + client_type: Literal['twinkle', 'tinker'] + payload: dict[str, Any] + adapter_label: str | None = None diff --git a/src/twinkle/server/state/training_store.py b/src/twinkle/server/state/training_store.py new file mode 100644 index 000000000..c814e6a73 --- /dev/null +++ b/src/twinkle/server/state/training_store.py @@ -0,0 +1,93 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Controller-owned persistence; other key-spaces are accessed through ServerState.""" +from __future__ import annotations + +import hashlib +import json +import time +from typing import TYPE_CHECKING, Any + +from .backend.base import StateBackend +from .models import FutureFailureRecord +from .training_records import CreateTraining, InstanceRecord, TrainingRecord + +if TYPE_CHECKING: + from .server_state import ServerState + + +class TrainingStore: + + def __init__(self, state: ServerState, namespace: str, backend: StateBackend): + self.state = state + self.backend = backend + self.prefix = f'training::{namespace}::' + + def key(self, kind: str, identity: str) -> str: + return f'{self.prefix}{kind}::{identity}' + + @staticmethod + def create_identity(request: CreateTraining) -> tuple[str, str]: + identity = json.dumps([request.client_type, request.token, request.session_id, request.create_key]) + fingerprint = json.dumps(request.model_dump(mode='json'), sort_keys=True, separators=(',', ':')) + return hashlib.sha256(identity.encode()).hexdigest(), hashlib.sha256(fingerprint.encode()).hexdigest() + + async def reserve(self, record: TrainingRecord) -> None: + await self.state.register_model_record(record.training_id, record) + + async def save_training(self, record: TrainingRecord) -> None: + await self.state.save_model_record(record.training_id, record) + + async def retire(self, record: TrainingRecord) -> None: + await self.backend.set( + self.key('terminal', record.training_id), + record.model_dump(), + ttl=max(1, int(self.state.expiration_timeout))) + await self.state.delete_training_sequences(record.training_id) + # Release the quota row last so a failed earlier step can be idempotently retried. + await self.state.unload_model(record.training_id) + + async def save_instance(self, record: InstanceRecord) -> None: + ttl = max(1, int(self.state.expiration_timeout)) if record.state == 'STOPPED' else None + 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.state.get_training_records(self.prefix) + + async def get_create(self, identity: str) -> dict[str, str] | None: + return await self.backend.get(self.key('create', identity)) + + async def save_create(self, identity: str, fingerprint: str, request_id: str, session_id: str) -> None: + await self.backend.set( + self.key('create', identity), { + 'fingerprint': fingerprint, + 'request_id': request_id, + 'session_id': session_id + }) + + async def cleanup_receipts(self) -> None: + for key in await self.backend.keys(self.key('create', '*')): + receipt = await self.backend.get(key) + if receipt and await self.state.get_session_last_heartbeat(receipt['session_id']) is None: + await self.backend.delete(key) + + async def future(self, + request_id: str, + status: str, + model_id: str | None, + *, + epoch: str | None = None, + result: Any = None, + failure: FutureFailureRecord | None = None, + deadline: float | None = None) -> None: + await self.state.store_future_status( + request_id, status, model_id, result=result, failure=failure, replica_id=epoch, absolute_deadline=deadline) + + async def own_future(self, request_id: str, token: str) -> None: + await self.state.set_future_owner(request_id, token) + + async def begin_epoch(self, epoch: str) -> None: + await self.backend.set(self.key('controller', 'active'), {'epoch': epoch, 'started_at': time.time()}) diff --git a/src/twinkle/server/task_queue/__init__.py b/src/twinkle/server/task_queue/__init__.py index bf5391611..6d7dd81be 100644 --- a/src/twinkle/server/task_queue/__init__.py +++ b/src/twinkle/server/task_queue/__init__.py @@ -5,11 +5,9 @@ Public exports: - TaskStatus - task lifecycle enum - QueueState - queue state enum for tinker client compatibility -- TaskQueueConfig - queue and rate-limit configuration dataclass - TaskQueueMixin - mixin with schedule_task / schedule_background_task - RateLimiter - sliding-window rate limiter """ -from .config import TaskQueueConfig from .mixin import TaskQueueMixin from .rate_limiter import RateLimiter from .types import QueuedTask, QueueState, TaskStatus, UserTaskError @@ -20,7 +18,6 @@ 'QueueState', 'UserTaskError', 'QueuedTask', - 'TaskQueueConfig', 'TaskQueueMixin', 'ComputeWorker', 'RateLimiter', diff --git a/src/twinkle/server/task_queue/mixin.py b/src/twinkle/server/task_queue/mixin.py index f26ec9c5c..b887ed20b 100644 --- a/src/twinkle/server/task_queue/mixin.py +++ b/src/twinkle/server/task_queue/mixin.py @@ -9,12 +9,12 @@ 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.server.config.task_queue import TaskQueueConfig from twinkle.server.exceptions import BatchSizeError, ConfigError, InputTokensExceededError, RateLimitExceededError from twinkle.server.lifecycle.envelope import envelope_from_record from twinkle.server.lifecycle.poll_config import long_poll_window @@ -22,7 +22,6 @@ from twinkle.server.telemetry.metrics import get_task_metrics from twinkle.utils.logger import get_logger from .backend_gate import BackendGate -from .config import TaskQueueConfig from .rate_limiter import RateLimiter from .types import QueuedTask, QueueState, TaskStatus from .worker import ComputeWorker @@ -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 task.on_finished(False, False) + raise + return {'request_id': request_id, 'model_id': model_id} + + async def _prepare_queued_task(self, task: QueuedTask) -> None: + """Optional owner admission, after preflight and before any future write.""" + + async def _enqueue_prepared_task(self, task: QueuedTask) -> None: + request_id, model_id = task.request_id, task.model_id + self._compute_worker.ensure_accepting() await self.state.store_future_status( request_id, TaskStatus.PENDING.value, @@ -239,34 +267,43 @@ async def _schedule_task( absolute_deadline=self._future_deadline(), ) - 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: + if task.on_finished is not None: + await task.on_finished(False, False) + raise + queue_key = self._queue_key(model_id=model_id, token=task.token) self._compute_worker.ensure_queue_registered(queue_key) - await self._compute_worker.ensure_started() - q = self._compute_worker.get_queue(queue_key) - 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: + if task.on_enqueued is not None: + await task.on_enqueued() + await self.state.store_future_status( + request_id, + TaskStatus.QUEUED.value, + model_id, + queue_state=QueueState.ACTIVE.value, + ) + logger.info(f'[TaskQueue] Task {request_id} queued, type={task.task_type or "unknown"}, ' + f'model_id={model_id}, queue_key={queue_key}, ' + f'queue_depth={q.qsize()}, input_tokens={task.input_tokens}') + finally: + # Once q.put succeeds the sleeping worker must always be woken, even + # if an acknowledgement or status write fails afterwards. + self._compute_worker.notify_new_task() async def schedule_task( self, @@ -420,67 +457,12 @@ async def schedule_background_task( model_id: str | None = None, task_type: str | None = None, ) -> dict[str, Any]: - """Schedule a fire-and-forget task outside the serial queue. - - 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(). + """Run pure I/O outside serial compute, with the same supervised lifecycle. - Args: - coro_factory: Zero-argument callable that creates the coroutine. - model_id: Optional model id for result association. - task_type: Label for logging. - - Returns: - {'request_id': str, 'model_id': str | None} + Capacity, timeout, cancellation and shutdown are owned by ComputeWorker; + background jobs must not create a second unsupervised execution path. """ - request_id = f'req_{uuid.uuid4().hex}' - logger.info(f'[TaskQueue] Scheduling background task {request_id}, ' - f'type={task_type or "unknown"}, model_id={model_id}') - - await self.state.store_future_status( - request_id, - TaskStatus.RUNNING.value, - model_id, - queue_state=QueueState.ACTIVE.value, - replica_id=getattr(self, 'replica_id', None), - absolute_deadline=self._future_deadline(), - ) - - async def _run() -> None: - try: - result = await coro_factory() - await self.state.store_future_status( - request_id, - TaskStatus.COMPLETED.value, - model_id, - result=result, - queue_state=QueueState.ACTIVE.value, - ) - logger.info(f'[TaskQueue] Background task {request_id} completed, type={task_type or "unknown"}') - except Exception as exc: - failure = FutureFailureRecord( - reason_code='internal_error', - message=f'{type(exc).__name__}: {exc}'[:1024], - attribution='server', - diagnostic=traceback.format_exc(), - ) - await self.state.store_future_status( - request_id, - TaskStatus.FAILED.value, - model_id, - failure=failure, - queue_state=QueueState.ACTIVE.value, - ) - logger.error(f'[TaskQueue] Background task {request_id} FAILED, type={task_type or "unknown"}:\n' - f'{traceback.format_exc(limit=3)}') - - asyncio.create_task(_run()) - return {'request_id': request_id, 'model_id': model_id} + return await self.schedule_concurrent_task(coro_factory, model_id=model_id, task_type=task_type) async def _fail_queue_tasks_async(self, queue_key: str, reason: str) -> None: await self._compute_worker.fail_queue_tasks(queue_key, reason) @@ -493,13 +475,23 @@ def fail_pending_tasks_for_model(self, model_id: str, reason: str) -> None: return def _schedule() -> None: - asyncio.create_task(self._fail_queue_tasks_async(queue_key, reason)) + task = asyncio.create_task(self._fail_queue_tasks_async(queue_key, reason)) + self._queue_tasks.add(task) + task.add_done_callback(self._queue_task_done) self._event_loop.call_soon_threadsafe(_schedule) + def _queue_task_done(self, task: asyncio.Task) -> None: + self._queue_tasks.discard(task) + if not task.cancelled(): + error = task.exception() + if error is not None: + logger.error('Queue maintenance failed', exc_info=(type(error), error, error.__traceback__)) + async def shutdown_task_queue(self) -> None: """Gracefully shut down the compute queue and release resources.""" await self._rate_limiter.stop_cleanup_task() await self._compute_worker.stop() + await asyncio.gather(*self._queue_tasks, return_exceptions=True) self._backend_gate.shutdown() logger.debug('[TaskQueue] Task queue shutdown complete') diff --git a/src/twinkle/server/task_queue/types.py b/src/twinkle/server/task_queue/types.py index 936a433a3..6e522785e 100644 --- a/src/twinkle/server/task_queue/types.py +++ b/src/twinkle/server/task_queue/types.py @@ -68,3 +68,5 @@ class QueuedTask: input_tokens: int task_type: str | None created_at: float + on_finished: Callable[[bool, bool], Coroutine] | None = None + on_enqueued: Callable[[], Coroutine] | None = None diff --git a/src/twinkle/server/task_queue/worker.py b/src/twinkle/server/task_queue/worker.py index cf684fe79..86aba9e61 100644 --- a/src/twinkle/server/task_queue/worker.py +++ b/src/twinkle/server/task_queue/worker.py @@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, Callable, Deque from twinkle.protocol.types.errors import ErrorCategory +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.exceptions import (BatchSizeError, EndpointUnavailableError, FullModeBusyError, InputTokensExceededError, RateLimitExceededError, RequestRejectedError, ResourceExhaustedError, ResourceNotFoundError, ResourceQuotaExceededError, @@ -26,7 +27,6 @@ from twinkle.server.telemetry.correlation import MODEL_ID, TOKEN_ID from twinkle.server.telemetry.tracing import traced_operation from twinkle.utils.logger import get_logger -from .config import TaskQueueConfig from .types import BACKEND_TIMEOUT_EXCEPTIONS, BackendBusyError, QueuedTask, QueueState, TaskStatus, UserTaskError if TYPE_CHECKING: @@ -107,6 +107,7 @@ def __init__( self._on_backend_timeout = on_backend_timeout self.task_queues: dict[str, asyncio.Queue] = {} + self._request_ids: set[str] = set() self.queue_order: Deque[str] = deque() self.new_task_event: asyncio.Event = asyncio.Event() @@ -119,14 +120,19 @@ def __init__( self._detached_tasks: dict[asyncio.Task, QueuedTask] = {} self._detached_reserved = 0 + def ensure_accepting(self) -> None: + if self._stopping: + raise ResourceExhaustedError('The replica is shutting down.') + async def ensure_started(self) -> None: - """Ensure the background worker coroutine is running.""" + """Ensure the background worker coroutine is running; never reopen a stopped replica.""" + self.ensure_accepting() if self._started and self._worker_task is not None and not self._worker_task.done(): return async with self._start_lock: + self.ensure_accepting() if self._started and self._worker_task is not None and not self._worker_task.done(): return - self._stopping = False self._worker_task = asyncio.create_task(self._worker_loop()) self._started = True @@ -150,10 +156,11 @@ async def stop(self) -> None: 'Sampler task stopped during replica shutdown.', QueueState.ACTIVE.value, reason_code='orphaned_replica') + for queue_key in list(self.task_queues): + await self.fail_queue_tasks(queue_key, 'Replica stopped before the task could execute.') self._worker_task = None self._started = False - self.task_queues.clear() - self.queue_order.clear() + self._request_ids.clear() def ensure_queue_registered(self, queue_key: str) -> None: """Register a new per-key queue if it does not yet exist.""" @@ -171,6 +178,13 @@ def get_queue(self, queue_key: str) -> asyncio.Queue: """ return self.task_queues[queue_key] + def track_task(self, request_id: str) -> None: + self._request_ids.add(request_id) + + def has_task(self, request_id: str) -> bool: + """True while a task is queued or executing, including a terminal-future race.""" + return request_id in self._request_ids + def total_queued(self) -> int: """Total number of pending tasks across all per-key queues.""" return sum(q.qsize() for q in self.task_queues.values()) @@ -268,6 +282,9 @@ async def fail_queue_tasks(self, queue_key: str, reason: str) -> None: for task in drained: await self._store_task_failed(task, reason, QueueState.UNKNOWN.value, queue_state_reason=reason) q.task_done() + self._request_ids.discard(task.request_id) + if task.on_finished is not None: + await task.on_finished(False, False) self.task_queues.pop(queue_key, None) try: @@ -294,6 +311,9 @@ async def _fail_timed_out_task(self, task: QueuedTask, waited: float, q: asyncio await self._store_task_failed(task, error, QueueState.PAUSED_CAPACITY.value, queue_state_reason=error) self._record_tasks_total(task.task_type or 'unknown', 'timeout') q.task_done() + self._request_ids.discard(task.request_id) + if task.on_finished is not None: + await task.on_finished(False, False) async def _execute_task(self, task: QueuedTask, queue_key: str, q: asyncio.Queue) -> None: """Execute a dequeued task, then release its queue slot.""" @@ -329,7 +349,8 @@ async def submit_detached(self, task: QueuedTask, queue_key: str, *, replica_id: def reserve_detached(self) -> None: """Reserve before the future write, so capacity rejection leaves no record.""" - if self._stopping or self._detached_reserved >= self._config.concurrent_task_limit: + self.ensure_accepting() + if self._detached_reserved >= self._config.concurrent_task_limit: raise ResourceExhaustedError('Concurrent compute capacity is unavailable on this replica.') self._detached_reserved += 1 @@ -379,6 +400,10 @@ async def run(): def done(completed: asyncio.Task) -> None: self._detached_tasks.pop(completed, None) self.release_detached() + if not completed.cancelled(): + error = completed.exception() + if error is not None: + logger.error('Detached task failed', exc_info=(type(error), error, error.__traceback__)) running.add_done_callback(done) @@ -391,6 +416,7 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: exec_start = time.monotonic() task_status = 'completed' exec_time = 0.0 + executed = False # One span per queued task execution, tagged `.` # (e.g. `model.forward`, `sampler.sample`). queue_attrs = { @@ -403,24 +429,15 @@ 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. @@ -514,6 +531,9 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: traceback_text=traceback.format_exc(), ) finally: + self._request_ids.discard(task.request_id) + if task.on_finished is not None: + await task.on_finished(executed, executed and task_status in ('timeout', 'cancelled')) self._record_execution_time(task_type, exec_time) self._record_tasks_total(task_type, task_status) @@ -546,13 +566,6 @@ async def _try_run_one(self) -> bool: await self._fail_timed_out_task(task, queue_wait, q) continue # try the next queue - # A record already in a Terminal_State (e.g. written 'failed' by the - # state-hygiene orphan handling) must not be executed again. - if await self._is_record_terminal(task.request_id): - logger.info(f'[ComputeWorker] Task {task.request_id} already terminal on dequeue; skipping.') - q.task_done() - continue - # Execute the task (serial: stops after the first execution) await self._execute_task(task, queue_key, q) return True @@ -561,11 +574,8 @@ async def _try_run_one(self) -> bool: async def _is_record_terminal(self, request_id: str) -> bool: """True if the future record already holds a Terminal_State.""" - try: - record = await self._state.get_future(request_id) - except Exception: - return False - if not record: + record = await self._state.get_future(request_id) + if record is None: return False return record.get('status') in (TaskStatus.COMPLETED.value, TaskStatus.FAILED.value, TaskStatus.CANCELLED.value) diff --git a/src/twinkle/server/telemetry/metrics.py b/src/twinkle/server/telemetry/metrics.py index 586277708..1c287ab9e 100644 --- a/src/twinkle/server/telemetry/metrics.py +++ b/src/twinkle/server/telemetry/metrics.py @@ -111,6 +111,33 @@ class TaskMetrics(BaseModel): rate_limiter_active_tokens: _Gauge +class TrainingMetrics(BaseModel): + """Low-cardinality metrics for the managed training lifecycle.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + instances_active: _Gauge + startup_seconds: _Histogram + shutdown_seconds: _Histogram + allocation_seconds: _Histogram + removal_seconds: _Histogram + quarantined_total: _Counter + sweep_total: _Counter + sweep_seconds: _Histogram + devices_allocated: _Gauge + placement_groups_active: _Gauge + + +class StateBackendMetrics(BaseModel): + """Metrics for persistent state backend availability and contention.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + redis_health: _Gauge + update_atomic_conflicts: _Counter + update_atomic_retries: _Histogram + + class _RequestMetrics(BaseModel): """HTTP request metrics container (internal).""" @@ -177,6 +204,45 @@ def __init__(self) -> None: description='Number of tokens currently tracked by the rate limiter', ) + # === Managed training === + self.training_instances_active = meter.create_up_down_counter( + 'twinkle.training.instances.active', description='Active managed training instances by lifecycle state') + self.training_startup_seconds = meter.create_histogram( + 'twinkle.training.instance.startup.duration_seconds', + description='Training runtime startup duration', + unit='s') + self.training_shutdown_seconds = meter.create_histogram( + 'twinkle.training.instance.shutdown.duration_seconds', + description='Training runtime shutdown duration', + unit='s') + self.training_allocation_seconds = meter.create_histogram( + 'twinkle.training.placement_group.allocation.duration_seconds', + description='Placement group allocation duration', + unit='s') + self.training_removal_seconds = meter.create_histogram( + 'twinkle.training.placement_group.removal.duration_seconds', + description='Placement group removal duration', + unit='s') + self.training_quarantined_total = meter.create_counter( + 'twinkle.training.quarantined.total', description='Training instances quarantined during cleanup') + self.training_sweep_total = meter.create_counter( + 'twinkle.training.sweep.total', description='Training lifecycle sweep runs by status') + self.training_sweep_seconds = meter.create_histogram( + 'twinkle.training.sweep.duration_seconds', description='Training lifecycle sweep duration', unit='s') + self.training_devices_allocated = meter.create_up_down_counter( + 'twinkle.training.devices.allocated', description='Logical devices owned by managed training instances') + self.training_placement_groups_active = meter.create_up_down_counter( + 'twinkle.training.placement_groups.active', + description='Placement groups owned by managed training instances') + + # === Redis state backend === + self.redis_health = meter.create_up_down_counter( + 'twinkle.state.redis.health', description='Last observed Redis health (1 healthy, 0 unhealthy)') + self.redis_update_atomic_conflicts = meter.create_counter( + 'twinkle.state.redis.update_atomic.conflicts.total', description='Redis WATCH conflicts') + self.redis_update_atomic_retries = meter.create_histogram( + 'twinkle.state.redis.update_atomic.retries', description='Retries used by Redis atomic updates') + # === Resources (ObservableGauge backed by the push cache) === # OTEL holds its own references to the gauges via the meter, so we # only need to keep the cache the callbacks read from. @@ -192,6 +258,8 @@ def __init__(self) -> None: # swaps this singleton to rebind instruments) invalidates them. self._task_metrics: dict[str, TaskMetrics] = {} self._request_metrics: dict[str, _RequestMetrics] = {} + self._training_metrics: TrainingMetrics | None = None + self._state_backend_metrics: StateBackendMetrics | None = None def _make_gauge_callback(self, name: str): """Build the sync OTEL callback that reads ``_resource_cache[name]``.""" @@ -227,6 +295,33 @@ def request_metrics(self, deployment: str) -> _RequestMetrics: ) return cached + def training_metrics(self) -> TrainingMetrics: + """Return process-local managed training metric adapters.""" + if self._training_metrics is None: + self._training_metrics = TrainingMetrics( + instances_active=_Gauge(self.training_instances_active), + startup_seconds=_Histogram(self.training_startup_seconds), + shutdown_seconds=_Histogram(self.training_shutdown_seconds), + allocation_seconds=_Histogram(self.training_allocation_seconds), + removal_seconds=_Histogram(self.training_removal_seconds), + quarantined_total=_Counter(self.training_quarantined_total), + sweep_total=_Counter(self.training_sweep_total), + sweep_seconds=_Histogram(self.training_sweep_seconds), + devices_allocated=_Gauge(self.training_devices_allocated), + placement_groups_active=_Gauge(self.training_placement_groups_active), + ) + return self._training_metrics + + def state_backend_metrics(self) -> StateBackendMetrics: + """Return process-local state backend metric adapters.""" + if self._state_backend_metrics is None: + self._state_backend_metrics = StateBackendMetrics( + redis_health=_Gauge(self.redis_health), + update_atomic_conflicts=_Counter(self.redis_update_atomic_conflicts), + update_atomic_retries=_Histogram(self.redis_update_atomic_retries), + ) + return self._state_backend_metrics + # ----- Push API for the cleanup leader ----- def set_resource_count(self, name: str, value: int) -> None: @@ -268,3 +363,13 @@ def get_task_metrics(deployment: str) -> TaskMetrics: def get_request_metrics(deployment: str) -> _RequestMetrics: """Return the per-deployment HTTP request metric adapters.""" return MetricsRegistry.get().request_metrics(deployment) + + +def get_training_metrics() -> TrainingMetrics: + """Return managed training metric adapters.""" + return MetricsRegistry.get().training_metrics() + + +def get_state_backend_metrics() -> StateBackendMetrics: + """Return persistent state backend metric adapters.""" + return MetricsRegistry.get().state_backend_metrics() diff --git a/src/twinkle/server/training/__init__.py b/src/twinkle/server/training/__init__.py new file mode 100644 index 000000000..33dc7caf2 --- /dev/null +++ b/src/twinkle/server/training/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""On-demand training: catalog, lifecycle controller and runtime ownership.""" diff --git a/src/twinkle/server/training/actor.py b/src/twinkle/server/training/actor.py new file mode 100644 index 000000000..6eb1245f6 --- /dev/null +++ b/src/twinkle/server/training/actor.py @@ -0,0 +1,195 @@ +# 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._unreachable_since = {} + self._failures: dict[str, asyncio.Task] = {} + 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.state.touch_replica_last_seen(self.core.epoch) + except Exception: + logger.exception('Training state heartbeat failed; admission remains fenced by the state backend') + await asyncio.sleep(5) + continue + await asyncio.gather(*(self._check_instance_health(instance) + for instance in list(self.core.instances.values()) + if instance.state in ('READY', 'RESETTING'))) + await asyncio.sleep(5) + + async def _check_instance_health(self, instance): + try: + identity = await self.driver.health(instance) + if identity['unhealthy']: + self._fail_unhealthy(instance, 'Training runtime reported a failed worker') + return + await self.store.state.touch_replica_last_seen(instance.instance_id) + self._unreachable_since.pop(instance.instance_id, None) + except Exception: + since = self._unreachable_since.setdefault(instance.instance_id, time.monotonic()) + if time.monotonic() - since >= 30: + self._fail_unhealthy(instance, 'Training runtime became unavailable') + + def _fail_unhealthy(self, instance, message): + identity = instance.instance_id + if identity in self._failures: + return + task = asyncio.create_task( + self.core.fail_instance( + identity, FutureFailureRecord(reason_code='orphaned_replica', attribution='server', message=message))) + self._failures[identity] = task + + def done(completed): + self._failures.pop(identity, None) + if not completed.cancelled(): + error = completed.exception() + if error is not None: + logger.error('Unhealthy runtime cleanup failed', exc_info=(type(error), error, error.__traceback__)) + + task.add_done_callback(done) + + async def _sweep(self): + while True: + started, status = time.monotonic(), 'success' + 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.own_future(request_id, training.token) + await self.store.future(request_id, 'failed', training.training_id, failure=failure) + ray.kill(ray.get_actor(controller_name(self.namespace), namespace=self.namespace), no_restart=False) + + async def register_runtime(self, instance_id, epoch, incarnation): + await self.core.register_runtime(instance_id, epoch, incarnation) + + async def create(self, request): + await self.initialize() + return await self.core.create(CreateTraining.model_validate(request)) + + async def route(self, *args): + await self.initialize() + return await self.core.route(*args) + + async def admit(self, *args, **kwargs): + await self.initialize() + return await self.core.admit(*args, **kwargs) + + async def commit_operation(self, *args): + return await self.core.commit_operation(*args) + + async def finish_operation(self, *args, **kwargs): + return await self.core.finish_operation(*args, **kwargs) + + async def release(self, *args): + return await self.core.release(*args) + + async def fail_instance(self, *args): + return await self.core.fail_instance(*args) + + async def mark_dirty(self, instance_id, epoch): + await self.core.mark_dirty(instance_id, epoch) + + async def diagnostics(self, base_model=None): + await self.initialize() + instances = [ + instance.model_dump() for instance in self.core.instances.values() + if base_model is None or instance.base_model == base_model + ] + return { + 'ready': self.core._ready and self.core.residents_ready, + 'epoch': self.core.epoch, + 'instances': instances + } + + async def shutdown(self): + 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 asyncio.gather(*self._failures.values(), return_exceptions=True) + await self.core.wait_for_jobs() + await self.driver.wait_for_cleanup() diff --git a/src/twinkle/server/training/catalog.py b/src/twinkle/server/training/catalog.py new file mode 100644 index 000000000..1be0548e6 --- /dev/null +++ b/src/twinkle/server/training/catalog.py @@ -0,0 +1,44 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Pure lookup and request validation, separate from runtime allocation.""" +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +from twinkle.server.config.training import LoraTrainingConfig, TrainingConfig, TrainingModeConfig, TrainingModelConfig +from twinkle.server.exceptions import RequestRejectedError + + +@dataclass(frozen=True) +class TrainingSpec: + model: TrainingModelConfig + mode: Literal['lora', 'full'] + config: TrainingModeConfig + + @property + def key(self) -> tuple[str, str]: + return self.model.name, self.mode + + @property + def capacity(self) -> int: + return self.config.max_loras if isinstance(self.config, LoraTrainingConfig) else 1 + + +class TrainingCatalog: + + def __init__(self, config: TrainingConfig): + self._models = {model.name: model for model in config.models} + + def get(self, name: str, mode: Literal['lora', 'full']) -> TrainingSpec: + model = self._models.get(name) + if model is None: + raise RequestRejectedError(f'Unknown training model: {name}') + config = getattr(model.modes, mode) + if config is None: + raise RequestRejectedError(f'Training model {name} does not provide {mode} mode') + return TrainingSpec(model, mode, config) + + def specs(self) -> tuple[TrainingSpec, ...]: + return tuple( + TrainingSpec(model, mode, config) for model in self._models.values() for mode in ('lora', 'full') + if (config := getattr(model.modes, mode)) is not None) diff --git a/src/twinkle/server/training/controller.py b/src/twinkle/server/training/controller.py new file mode 100644 index 000000000..455c04027 --- /dev/null +++ b/src/twinkle/server/training/controller.py @@ -0,0 +1,847 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Single-writer training lifecycle; slow runtime work stays outside the lock.""" +from __future__ import annotations + +import asyncio +import logging +import time +import traceback +import uuid +from collections.abc import Callable +from datetime import datetime +from typing import Any, Protocol + +from twinkle.server.config.training import TrainingConfig +from twinkle.server.exceptions import (RequestRejectedError, ResourceExhaustedError, ResourceNotFoundError, + ResourceQuotaExceededError) +from twinkle.server.resources.allocation import CapacityUnavailableError +from twinkle.server.state.models import FutureFailureRecord +from twinkle.server.state.training_records import (LIVE_INSTANCES, TERMINAL_TRAINING, CreateTraining, InstanceRecord, + TrainingRecord) +from twinkle.server.state.training_store import TrainingStore +from twinkle.server.telemetry.metrics import get_training_metrics +from .catalog import TrainingCatalog, TrainingSpec +from .errors import TrainingExpiredError + +logger = logging.getLogger(__name__) + + +class InstanceDriver(Protocol): + + async def allocate(self, instance: InstanceRecord, spec: TrainingSpec) -> dict[str, Any]: + ... + + async def deploy(self, instance: InstanceRecord, spec: TrainingSpec) -> None: + ... + + async def create_training(self, instance: InstanceRecord, training: TrainingRecord, + request: CreateTraining) -> None: + ... + + async def release_training(self, instance: InstanceRecord, training: TrainingRecord) -> None: + ... + + async def shutdown(self, instance: InstanceRecord) -> None: + ... + + async def verify_operation(self, instance: InstanceRecord, request_id: str) -> dict[str, Any]: + ... + + async def recover_orphans(self) -> None: + ... + + +def failure_for(exc: Exception) -> FutureFailureRecord: + if isinstance(exc, CapacityUnavailableError): + reason, attribution = 'capacity_unavailable', 'user' + elif isinstance(exc, ResourceQuotaExceededError): + reason, attribution = 'resource_quota_exceeded', 'user' + elif isinstance(exc, TimeoutError): + reason, attribution = 'startup_timeout', 'server' + elif isinstance(exc, RequestRejectedError): + reason, attribution = 'invalid_request', 'user' + else: + reason, attribution = 'startup_failed', 'server' + return FutureFailureRecord( + reason_code=reason, + message=str(exc)[:1024], + attribution=attribution, + diagnostic=traceback.format_exc() if attribution == 'server' else None) + + +class TrainingController: + + def __init__(self, + config: TrainingConfig, + store: TrainingStore, + driver: InstanceDriver, + *, + clock: Callable[[], float] = time.monotonic): + self.config = config + self.catalog = TrainingCatalog(config) + self.store = store + self.driver = driver + self.clock = clock + self.epoch = 'controller_' + uuid.uuid4().hex + self.instances: dict[str, InstanceRecord] = {} + self.trainings: dict[str, TrainingRecord] = {} + self._startup: dict[str, asyncio.Task] = {} + self._jobs: set[asyncio.Task] = set() + self._cleaning: set[str] = set() + self._closing: dict[str, asyncio.Task] = {} + self._reconciling: set[str] = set() + self._failed_instances: dict[str, FutureFailureRecord] = {} + self._lock = asyncio.Lock() + self._ready = False + self._metrics = get_training_metrics() + self._reported_metric_keys: set[tuple[str, str, str, str]] = set() + self._reported_resource_metric_keys: set[tuple[str, str, str]] = set() + + def _instance_metric_tags(self, instance: InstanceRecord) -> dict[str, str]: + spec = self.catalog.get(instance.base_model, instance.mode) + return { + 'backend': spec.model.backend, + 'mode': instance.mode, + 'state': instance.state, + 'resident': str(instance.resident).lower(), + } + + def _report_instance_metrics(self) -> None: + counts: dict[tuple[str, str, str, str], int] = {} + resources: dict[tuple[str, str], tuple[int, int]] = {} + for instance in self.instances.values(): + if instance.state == 'STOPPED': + continue + tags = self._instance_metric_tags(instance) + key = tuple(tags[name] for name in ('backend', 'mode', 'state', 'resident')) + counts[key] = counts.get(key, 0) + 1 + resource_key = tags['backend'], tags['mode'], self.catalog.get(instance.base_model, + instance.mode).model.device_type + devices, groups = resources.get(resource_key, (0, 0)) + resources[resource_key] = (devices + len(instance.devices), groups + int(bool(instance.application_name))) + for key in self._reported_metric_keys | counts.keys(): + tags = dict(zip(('backend', 'mode', 'state', 'resident'), key)) + self._metrics.instances_active.set(counts.get(key, 0), tags=tags) + for key in self._reported_resource_metric_keys | resources.keys(): + tags = dict(zip(('backend', 'mode', 'device_type'), key)) + devices, groups = resources.get(key, (0, 0)) + self._metrics.devices_allocated.set(devices, tags=tags) + self._metrics.placement_groups_active.set(groups, tags=tags) + self._reported_metric_keys = set(counts) + self._reported_resource_metric_keys = set(resources) + + @staticmethod + def _quarantine_reason(exc: Exception) -> str: + message = str(exc).lower() + if isinstance(exc, TimeoutError): + return 'timeout' + if 'worker death' in message: + return 'worker_death' + if 'placement group' in message: + return 'pg_removal' + return 'other' + + def _spawn(self, coroutine) -> asyncio.Task: + task = asyncio.create_task(coroutine) + self._jobs.add(task) + + def done(completed): + self._jobs.discard(completed) + if completed.cancelled(): + return + exc = completed.exception() + if exc is not None: + logger.error('Training background job failed', exc_info=(type(exc), exc, exc.__traceback__)) + + task.add_done_callback(done) + return task + + async def initialize(self) -> None: + """Discard old state only after owned compute has been confirmed dead.""" + old_instances = await self.store.instances() + old_trainings = await self.store.trainings() + old_controller = await self.store.backend.get(self.store.key('controller', 'active')) + await self.store.begin_epoch(self.epoch) + failure = FutureFailureRecord( + reason_code='orphaned_replica', + message='Training controller restarted; create a new training handle.', + attribution='server') + for training in old_trainings: + for request_id in (training.create_request_id, *training.pending_operation_ids): + await self.store.future(request_id, 'failed', training.training_id, failure=failure) + await self.driver.recover_orphans() + for instance in old_instances: + if instance.state != 'STOPPED': + await asyncio.wait_for(self.driver.shutdown(instance), self.config.shutdown_timeout_s) + instance.state = 'STOPPED' + instance.stopped_at = time.time() + await self.store.save_instance(instance) + for training in old_trainings: + training.state = 'FAILED' + training.terminal_reason = 'orphaned_replica' + training.terminal_at = time.time() + await self.store.retire(training) + await self.store.state.rebuild_model_indexes() + if old_controller: + await self.store.state.unregister_replica(old_controller['epoch']) + await self.store.state.register_replica(self.epoch, 0) + await self.store.state.touch_replica_last_seen(self.epoch) + try: + for spec in self.catalog.specs(): + for _ in range(spec.config.min_replicas): + async with self._lock: + instance = await self._new_instance(spec, resident=True) + await asyncio.shield(self._startup[instance.instance_id]) + except Exception: + await self.shutdown() + raise + self._ready = True + + @property + def residents_ready(self) -> bool: + return all(self.residents_ready_for(spec) for spec in self.catalog.specs()) + + def residents_ready_for(self, spec: TrainingSpec) -> bool: + return sum( + instance.resident and instance.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: TrainingSpec, *, resident: bool = False) -> InstanceRecord: + instance = InstanceRecord( + instance_id='instance_' + uuid.uuid4().hex, + controller_epoch=self.epoch, + base_model=spec.model.name, + mode=spec.mode, + resident=resident, + startup_deadline=self.clock() + self.config.startup_timeout_s) + await self.store.save_instance(instance) + self.instances[instance.instance_id] = instance + self._report_instance_metrics() + self._startup[instance.instance_id] = self._spawn(self._start_instance(instance)) + return instance + + async def _start_instance(self, instance: InstanceRecord) -> None: + spec = self.catalog.get(instance.base_model, instance.mode) + started, status = time.monotonic(), 'error' + tags = {'backend': spec.model.backend, 'mode': spec.mode} + try: + + async def startup(): + 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: TrainingSpec) -> 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.catalog.get(request.base_model, request.mode) + identity, fingerprint = self.store.create_identity(request) + async with self._lock: + await self._assert_epoch() + if not self._ready: + raise 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) + await self.store.own_future(request_id, request.token) + await self.store.save_create(identity, fingerprint, request_id, request.session_id) + training = None + try: + # Reserve quota before starting an instance, but leave no usable route. + training = TrainingRecord( + training_id=training_id, + token=request.token, + session_id=request.session_id, + base_model=request.base_model, + mode=request.mode, + controller_epoch=self.epoch, + training_namespace=self.store.prefix, + instance_id='', + adapter_label=request.adapter_label, + create_request_id=request_id, + creation_deadline=self.clock() + self.config.startup_timeout_s, + last_activity_at=self.clock(), + lora_config=request.payload.get('lora_config'), + user_metadata=request.payload.get('user_metadata') or {}) + await self.store.reserve(training) + instance = self._select(spec) + if instance is None: + if not self.residents_ready_for(spec): + raise CapacityUnavailableError( + 'Resident training capacity is recovering; retry a new create later.') + count = sum((item.base_model, item.mode) == spec.key and item.state != 'STOPPED' + for item in self.instances.values()) + if spec.config.max_replicas is not None and count >= spec.config.max_replicas: + raise CapacityUnavailableError('Training replica limit reached; retry a new create later.') + instance = await self._new_instance(spec) + training.instance_id = instance.instance_id + training.replica_id = instance.instance_id + instance.training_ids.append(training_id) + self.trainings[training_id] = training + await self.store.save_training(training) + await self.store.save_instance(instance) + self._spawn(self._create_user(training, request)) + except Exception as exc: + if training is not None: + self.trainings.pop(training_id, None) + if training.instance_id in self.instances: + instance = self.instances[training.instance_id] + if training_id in instance.training_ids: + instance.training_ids.remove(training_id) + await self.store.save_instance(instance) + if not instance.training_ids and not instance.resident: + self._spawn(self.fail_instance(instance.instance_id, failure_for(exc))) + training.state = 'FAILED' + training.terminal_reason = failure_for(exc).reason_code + await self.store.retire(training) + await self.store.future(request_id, 'failed', training_id, failure=failure_for(exc)) + return {'request_id': request_id} + + async def _create_user(self, training: TrainingRecord, request: CreateTraining) -> None: + instance = self.instances[training.instance_id] + try: + await asyncio.shield(self._startup[instance.instance_id]) + async with self._lock: + future = await self.store.state.get_future(training.create_request_id) + if future and future['status'] == 'cancelled': + training.state = 'CLOSING' + training.terminal_reason = 'cancelled' + if training.state != 'PENDING' or instance.state != 'READY': + raise RequestRejectedError('Training creation is no longer pending') + training.creation_deadline = self.clock() + self.config.startup_timeout_s + await self.store.save_training(training) + await self.store.future(training.create_request_id, 'running', training.training_id) + remaining = training.creation_deadline - self.clock() + await asyncio.wait_for(self.driver.create_training(instance, training, request), max(0, remaining)) + # Driver includes training-run metadata commit, before public success. + async with self._lock: + await self._assert_epoch() + if self.clock() >= training.creation_deadline: + raise TimeoutError('Training initialization completed after its creation deadline') + if training.state != 'PENDING' or instance.state != 'READY': + raise RuntimeError('Training became unavailable during initialization') + training.state = 'ACTIVE' + training.last_activity_at = self.clock() + await self.store.save_training(training) + result = ({ + 'model_id': training.training_id + } if request.client_type == 'tinker' else { + 'training_id': training.training_id, + 'adapter_name': training.adapter_label + }) + await self.store.future(training.create_request_id, 'completed', training.training_id, result=result) + except Exception as exc: + if training.state not in TERMINAL_TRAINING: + await self.store.future( + training.create_request_id, 'failed', training.training_id, failure=failure_for(exc)) + async with self._lock: + training.state = 'CLOSING' + training.terminal_reason = ('cancelled' if training.terminal_reason == 'cancelled' else + failure_for(exc).reason_code) + await self.store.save_training(training) + await self._cleanup_user(training) + + async def _expired(self, training: TrainingRecord) -> bool: + heartbeat = await self.store.state.get_session_last_heartbeat(training.session_id) + return self._expired_with_heartbeat(training, heartbeat) + + def _expired_with_heartbeat(self, training: TrainingRecord, heartbeat: float | None) -> bool: + spec = self.catalog.get(training.base_model, training.mode) + config, legacy = spec.config, spec.model.legacy_expiration + if heartbeat is None: + return True + if legacy is not None: + now = time.time() + if legacy.session_timeout_s is not None and now - heartbeat >= legacy.session_timeout_s: + return True + if legacy.max_lifetime_s is not None: + created_at = datetime.fromisoformat(training.created_at).timestamp() + if now - created_at >= legacy.max_lifetime_s: + return True + return config.idle_timeout_s is not None and self.clock() - training.last_activity_at >= config.idle_timeout_s + + async def _authorized(self, training_id: str, token: str, session_id: str | None) -> TrainingRecord: + await self._assert_epoch() + training = self.trainings.get(training_id) + if training is None or training.token != token or (session_id is not None + and session_id != training.session_id): + raise ResourceNotFoundError('Unknown training handle') + if training.state == 'EXPIRED': + raise TrainingExpiredError('Training handle expired; create a new handle and explicitly load a checkpoint.') + if training.state != 'ACTIVE' or training.instance_id in self._failed_instances: + raise ResourceNotFoundError('Training handle is no longer active') + if await self.store.state.get_session_last_heartbeat(training.session_id) is None: + raise TrainingExpiredError('Training session expired; create a new training handle') + return training + + async def admit(self, + training_id: str, + token: str, + session_id: str | None, + request_id: str, + *, + renew: bool = True) -> dict[str, str]: + async with self._lock: + training = await self._authorized(training_id, token, session_id) + if not training.pending_operation_ids and await self._expired(training): + training.state = 'CLOSING' + training.terminal_reason = 'expired' + await self.store.save_training(training) + self._spawn(self._cleanup_user(training)) + raise TrainingExpiredError( + 'Training handle expired; create a new handle and explicitly load a checkpoint.') + instance = self.instances[training.instance_id] + if instance.state != 'READY': + raise ResourceNotFoundError('Training runtime is unavailable') + if request_id not in training.pending_operation_ids: + training.pending_operation_ids[request_id] = { + 'accepted_at': self.clock(), + 'committed': False, + 'renew': renew + } + await self.store.save_training(training) + await self.store.own_future(request_id, token) + return {'instance_id': instance.instance_id, 'epoch': self.epoch} + + async def commit_operation(self, training_id: str, request_id: str, epoch: str) -> None: + async with self._lock: + await self._assert_epoch() + training = self.trainings.get(training_id) + if epoch != self.epoch or training is None: + raise ResourceNotFoundError('Operation admission no longer exists') + if request_id not in training.pending_operation_ids: + # A very short task can finish while the enqueue acknowledgment + # is in flight. Its original terminal future is conclusive. + future = await self.store.state.get_future(request_id) + if future and future['status'] in ('completed', 'failed', 'cancelled'): + return + raise ResourceNotFoundError('Operation admission no longer exists') + training.pending_operation_ids[request_id]['committed'] = True + await self.store.save_training(training) + + async def finish_operation(self, + training_id: str, + request_id: str, + epoch: str, + *, + executed: bool, + finished_at: float | None = None) -> None: + async with self._lock: + try: + await self._assert_epoch() + except ResourceNotFoundError: + return + training = self.trainings.get(training_id) + if epoch != self.epoch or training is None: + return + ticket = training.pending_operation_ids.pop(request_id, None) + if ticket is None: + return + if executed and ticket['renew']: + age = max(0, time.time() - finished_at) if finished_at is not None else 0 + training.last_activity_at = max(training.last_activity_at, self.clock() - age) + await self.store.save_training(training) + cleanup = training.state == 'CLOSING' and not training.pending_operation_ids + if cleanup: + self._spawn(self._cleanup_user(training)) + + async def route(self, training_id: str, token: str, session_id: str | None = None) -> dict[str, str]: + async with self._lock: + training = await self._authorized(training_id, token, session_id) + return {'instance_id': training.instance_id, 'base_model': training.base_model, 'epoch': self.epoch} + + async def release(self, training_id: str, token: str, session_id: str | None = None) -> None: + async with self._lock: + training = await self._authorized(training_id, token, session_id) + training.state = 'CLOSING' + training.terminal_reason = 'released' + await self.store.save_training(training) + cleanup = not training.pending_operation_ids + if cleanup: + await self._cleanup_user(training) + config = self.catalog.get(training.base_model, training.mode).config.queue_config + deadline = asyncio.get_running_loop().time() + config.absolute_future_ttl(1) + self.config.shutdown_timeout_s + while training.state == 'CLOSING' or self.instances[training.instance_id].state == 'DRAINING': + if self.instances[training.instance_id].state == 'QUARANTINED': + raise RuntimeError('Runtime cleanup is quarantined; device release is not confirmed') + if asyncio.get_running_loop().time() >= deadline: + raise TimeoutError('Training cleanup is still pending; its slot remains reserved') + await asyncio.sleep(config.inline_fast_path_timeout) + + async def _cleanup_user(self, training: TrainingRecord) -> None: + instance = self.instances[training.instance_id] + 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.catalog.get(training.base_model, training.mode) + await asyncio.wait_for( + self.driver.release_training(instance, training), + spec.config.queue_config.absolute_future_ttl(spec.config.num_devices)) + except Exception as exc: + # A timed-out create/cleanup RPC does not prove that shared workers failed. + # Keep this slot reserved until serial cleanup can be confirmed. + if instance.mode == 'lora' and any( + self.trainings[identity].state in ('ACTIVE', + 'PENDING') or self.trainings[identity].pending_operation_ids + for identity in instance.training_ids if identity != training.training_id): + logger.warning('Training cleanup remains pending for %s: %s', training.training_id, exc) + else: + await self.fail_instance( + instance.instance_id, + FutureFailureRecord( + reason_code='internal_error', + attribution='server', + message=f'Training cleanup failed: {exc}'[:1024], + diagnostic=traceback.format_exc())) + return + async with self._lock: + terminal = training.model_copy( + deep=True, + update={ + 'state': { + 'expired': 'EXPIRED', + 'released': 'RELEASED', + 'cancelled': 'CANCELLED' + }.get(training.terminal_reason, 'FAILED'), + 'terminal_at': time.time() + }) + await self.store.retire(terminal) + remaining = [identity for identity in instance.training_ids if identity != training.training_id] + close = not remaining and (not instance.resident or instance.dirty) + updated_instance = instance.model_copy( + deep=True, + update={ + 'training_ids': remaining, + 'state': 'DRAINING' if close else ('READY' if resetting else instance.state) + }) + await self.store.save_instance(updated_instance) + training.state = terminal.state + training.terminal_at = terminal.terminal_at + instance.training_ids = remaining + instance.state = updated_instance.state + 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: + close_task = self._closing.get(instance.instance_id) + if close_task is None: + close_task = self._spawn(self._close_instance_once(instance)) + self._closing[instance.instance_id] = close_task + + def clear(completed: asyncio.Task) -> None: + if self._closing.get(instance.instance_id) is completed: + self._closing.pop(instance.instance_id, None) + + close_task.add_done_callback(clear) + await asyncio.shield(close_task) + + async def _close_instance_once(self, instance: InstanceRecord) -> None: + started = time.monotonic() + tags = self._instance_metric_tags(instance) + timing_tags = {'backend': tags['backend'], 'mode': tags['mode']} + try: + await asyncio.wait_for(self.driver.shutdown(instance), self.config.shutdown_timeout_s) + except Exception as exc: + async with self._lock: + quarantined = instance.model_copy(update={'state': 'QUARANTINED', 'failure': str(exc)}) + await self.store.save_instance(quarantined) + instance.state = 'QUARANTINED' + instance.failure = str(exc) + self._report_instance_metrics() + tags = self._instance_metric_tags(instance) + self._metrics.quarantined_total.inc(tags={ + 'backend': tags['backend'], + 'mode': tags['mode'], + 'reason': self._quarantine_reason(exc), + }) + self._metrics.shutdown_seconds.observe(time.monotonic() - started, tags={**timing_tags, 'status': 'error'}) + return + self._metrics.shutdown_seconds.observe(time.monotonic() - started, tags={**timing_tags, 'status': 'success'}) + async with self._lock: + terminal_records = [] + terminal_at = time.time() + 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) + startup = self._startup.get(instance_id) + if startup is not None and startup is not asyncio.current_task() and not startup.done(): + startup.cancel() + await asyncio.gather(startup, return_exceptions=True) + finally: + # Once DRAINING is durable, shutdown must be attempted even if a + # later training/future write fails. Sweep will retry if this fails. + await self._close_instance(instance) + + async def sweep(self) -> None: + # Batch I/O outside the decision lock; re-verify state after re-acquiring, ignore handles added since snapshot. + session_ids = list({ + training.session_id + for training in self.trainings.values() if training.state == 'ACTIVE' and not training.pending_operation_ids + }) + heartbeats = await self.store.state.get_session_last_heartbeats(session_ids) + async with self._lock: + await self._assert_epoch() + for training in self.trainings.values(): + if (training.state != 'ACTIVE' or training.pending_operation_ids + 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._spawn(self._cleanup_user(training)) + for training, request_id in uncertain: + self._spawn(self._reconcile_operation(training, request_id)) + for instance in list(self.instances.values()): + failure = self._failed_instances.get(instance.instance_id) + if failure is not None: + self._spawn(self.fail_instance(instance.instance_id, failure)) + elif instance.state in ('DRAINING', 'QUARANTINED'): + self._spawn(self._close_instance(instance)) + # Rebuilding one specification must not stop expiry or other models. + for spec in self.catalog.specs(): + residents = sum( + instance.resident and instance.instance_id not in self._failed_instances and ( + instance.base_model, instance.mode) == spec.key and instance.state in LIVE_INSTANCES + for instance in self.instances.values()) + for _ in range(spec.config.min_replicas - residents): + async with self._lock: + instance = await self._new_instance(spec, resident=True) + await self.store.cleanup_receipts() + # Terminal diagnostics expire like their backend tombstones. Repeated + # model creation must not grow controller memory for the server lifetime. + async with self._lock: + cutoff = time.time() - self.store.state.expiration_timeout + for identity, training in list(self.trainings.items()): + if training.state in TERMINAL_TRAINING and training.terminal_at is not None: + if training.terminal_at <= cutoff: + self.trainings.pop(identity) + for identity, instance in list(self.instances.items()): + if (instance.state == 'STOPPED' and instance.stopped_at is not None and instance.stopped_at <= cutoff + and not instance.training_ids): + self.instances.pop(identity) + self._startup.pop(identity, None) + 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'] == 'pending': + await self.commit_operation(training.training_id, request_id, self.epoch) + elif result['state'] == 'finished': + await self.finish_operation( + training.training_id, + request_id, + self.epoch, + executed=result['executed'], + finished_at=result.get('finished_at')) + elif result['state'] == 'missing': + await self.store.future( + request_id, + 'failed', + training.training_id, + failure=FutureFailureRecord( + reason_code='internal_error', + attribution='server', + message='Operation was not committed to the runtime queue')) + await self.finish_operation(training.training_id, request_id, self.epoch, executed=False) + # An unsafe/unreachable runtime is handled by the worker health monitor. + except Exception: + logger.exception('Operation reconciliation deferred for %s', request_id) + finally: + self._reconciling.discard(request_id) + + async def wait_for_jobs(self) -> None: + while self._jobs: + jobs = tuple(self._jobs) + await asyncio.gather(*jobs, return_exceptions=True) + self._jobs.difference_update(task for task in jobs if task.done()) + + async def shutdown(self) -> None: + self._ready = False + for instance in list(self.instances.values()): + await self.fail_instance( + instance.instance_id, + FutureFailureRecord( + reason_code='orphaned_replica', message='Training service is shutting down', attribution='server')) diff --git a/src/twinkle/server/training/driver.py b/src/twinkle/server/training/driver.py new file mode 100644 index 000000000..955dd83d1 --- /dev/null +++ b/src/twinkle/server/training/driver.py @@ -0,0 +1,229 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Ray/Serve effects behind the controller's small runtime interface.""" +from __future__ import annotations + +import asyncio +import os +import re +import time +from copy import deepcopy + +from twinkle.server.env_propagation import build_propagated_env_vars +from twinkle.server.resources.allocation import DevicePool +from twinkle.server.state.training_records import InstanceRecord +from twinkle.server.telemetry.metrics import get_training_metrics +from twinkle.utils.logger import get_logger +from .catalog import TrainingCatalog + +logger = get_logger() + + +class RayInstanceDriver: + + def __init__(self, config, store, namespace, controller_name): + self.config, self.store = config, store + self.namespace, self.controller_name = namespace, controller_name + self._cleanup_tasks: set[asyncio.Task] = set() + self._deployments: dict[str, asyncio.Task] = {} + self._metrics = get_training_metrics() + self._catalog = TrainingCatalog(config) + self.pool = DevicePool( + namespace, allocation_timeout=config.allocation_timeout_s, shutdown_timeout=config.shutdown_timeout_s) + + def application_name(self, instance): + return self.pool.name(instance.controller_epoch, instance.instance_id) + + async def allocate(self, instance, spec): + started, status = time.monotonic(), 'error' + tags = {'backend': spec.model.backend, 'mode': spec.mode, 'device_type': spec.model.device_type} + try: + allocation = await self.pool.allocate(spec.config.num_devices, spec.model.device_type, + instance.controller_epoch, instance.instance_id) + status = 'success' + return { + 'application_name': allocation.name, + 'placement_group_id': allocation.placement_group_id, + 'node_id': allocation.node_id, + 'devices': list(allocation.devices) + } + finally: + self._metrics.allocation_seconds.observe(time.monotonic() - started, tags={**tags, 'status': status}) + + async def deploy(self, instance, spec): + from ray import serve + from ray.util.placement_group import get_placement_group + + from twinkle.server.model.app import build_model_app + + name = self.application_name(instance) + env = build_propagated_env_vars() + env.update(spec.model.env_vars) + for key in ('PYTHONPATH', 'TWINKLE_DEFAULT_SAVE_DIR', 'TMPDIR', 'OMP_NUM_THREADS'): + if key in os.environ: + env.setdefault(key, os.environ[key]) + env.update({'TWINKLE_RAY_REGISTRY_NAME': name + '-registry', 'TWINKLE_RAY_NAMESPACE': self.namespace}) + context = { + 'instance_id': instance.instance_id, + 'epoch': instance.controller_epoch, + 'namespace': self.namespace, + 'controller_name': self.controller_name, + 'registry_name': name + '-registry' + } + kwargs = spec.config.model_dump( + exclude={'num_devices', 'device_mesh', 'min_replicas', 'max_replicas', 'idle_timeout_s', 'queue_config'}, + exclude_none=True) + if 'hybrid' in kwargs: + kwargs['hybrid'] = spec.config.hybrid + kwargs['train_mode'] = 'hybrid' if kwargs.get('hybrid') else spec.mode + # Each runtime has one identity. Serve autoscaling belongs to the CPU + # router; device-instance replica bounds are enforced by the controller. + options = deepcopy(spec.model.deploy_options) + options.pop('autoscaling_config', None) + options['num_replicas'] = 1 + actor_options = options.setdefault('ray_actor_options', {}) + actor_options.setdefault('num_cpus', 0.1) + actor_options.setdefault('runtime_env', {})['env_vars'] = env + app = build_model_app( + model_id=spec.model.model_id, + backend=spec.model.backend, + nproc_per_node=spec.config.num_devices, + device_group={ + 'name': 'model', + 'ranks': spec.config.num_devices, + 'device_type': spec.model.device_type + }, + device_mesh={ + **spec.config.device_mesh.model_dump(exclude_none=True), 'device_type': spec.model.device_type + }, + deploy_options=options, + queue_config=spec.config.queue_config, + data_plane_url=spec.model.data_plane_url, + training_context=context, + placement_group=get_placement_group(name), + allocation_devices=list(instance.devices), + **kwargs) + # Shield the actual thread so its late return can still be cleaned. + deployment = asyncio.create_task( + asyncio.to_thread( + serve.run, app, name=name, route_prefix=f'/api/v1/training-runtime/{instance.instance_id}')) + self._deployments[instance.instance_id] = deployment + try: + await asyncio.shield(deployment) + except asyncio.CancelledError: + + async def clean_late(): + try: + await deployment + except Exception: + pass + await self.shutdown(instance) + + task = asyncio.create_task(clean_late()) + self._cleanup_tasks.add(task) + + def done(completed): + self._cleanup_tasks.discard(completed) + if not completed.cancelled(): + exc = completed.exception() + if exc is not None: + logger.error('Late deployment cleanup failed', exc_info=(type(exc), exc, exc.__traceback__)) + + task.add_done_callback(done) + raise + finally: + if deployment.done(): + self._deployments.pop(instance.instance_id, None) + + async def wait_for_cleanup(self): + """Join late deployment cleanup before the controller actor is retired.""" + while self._cleanup_tasks: + await asyncio.gather(*tuple(self._cleanup_tasks), return_exceptions=True) + + def handle(self, instance): + from ray import serve + return serve.get_app_handle(instance.application_name or self.application_name(instance)) + + async def create_training(self, instance, training, request): + await self.handle(instance).initialize_training.remote(training.model_dump(), request.model_dump()) + + async def release_training(self, instance, training): + await self.handle(instance).release_training.remote(training.model_dump(), instance.resident + and not instance.dirty) + + async def verify_operation(self, instance, request_id): + return await self.handle(instance).verify_operation.remote(request_id) + + async def shutdown(self, instance: InstanceRecord): + import ray + from ray import serve + from ray.util import list_named_actors + deployment = self._deployments.get(instance.instance_id) + if deployment is not None: + # Cancellation of Serve's caller does not stop its deployment thread. + # Never certify death while that thread can still create resources. + try: + await asyncio.shield(deployment) + except Exception: + pass + self._deployments.pop(instance.instance_id, None) + name = self.application_name(instance) + owned = [ + actor for actor in list_named_actors(all_namespaces=True) + if actor['namespace'] == self.namespace and instance.instance_id in actor['name'] + ] + actors = [] + for actor in owned: + try: + handle = ray.get_actor(actor['name'], namespace=self.namespace) + except ValueError: + continue + actors.append(handle) + ray.kill(handle, no_restart=True) + await asyncio.to_thread(serve.delete, instance.application_name or name) + removal_started, removal_status = time.monotonic(), 'error' + try: + await self.pool.remove(name) + removal_status = 'success' + finally: + try: + spec = self._catalog.get(instance.base_model, instance.mode) + except Exception: + tags = {'backend': 'unknown', 'mode': instance.mode, 'device_type': 'unknown'} + else: + tags = {'backend': spec.model.backend, 'mode': spec.mode, 'device_type': spec.model.device_type} + self._metrics.removal_seconds.observe( + time.monotonic() - removal_started, tags={ + **tags, 'status': removal_status + }) + # The PG kills even unnamed remaining children. Named workers must also + # have failed readiness RPCs before the controller retires quota/tickets. + for handle in actors: + try: + await asyncio.wait_for(handle.__ray_ready__.remote(), 5) + except ray.exceptions.RayActorError: + continue + raise RuntimeError('Owned worker death could not be confirmed') + await self.store.state.unregister_replica_capacity(instance.instance_id) + + async def health(self, instance): + return await asyncio.wait_for(self.handle(instance).runtime_identity.remote(), 5) + + async def recover_orphans(self): + """Ownership names survive loss of the memory state actor.""" + from ray import serve + from ray.util.placement_group import placement_group_table + pattern = re.compile(r'resources-' + re.escape(self.namespace) + + r'-(controller_[0-9a-f]{32})-(instance_[0-9a-f]{32})$') + names = set(serve.status().applications) + names.update(row.get('name', '') for row in placement_group_table().values() if row.get('state') != 'REMOVED') + for name in sorted(names): + match = pattern.fullmatch(name) + if match: + orphan = InstanceRecord( + controller_epoch=match[1], + instance_id=match[2], + application_name=name, + base_model='', + mode='full', + startup_deadline=0) + await asyncio.wait_for(self.shutdown(orphan), self.config.shutdown_timeout_s) diff --git a/src/twinkle/server/training/errors.py b/src/twinkle/server/training/errors.py new file mode 100644 index 000000000..d0d508bc6 --- /dev/null +++ b/src/twinkle/server/training/errors.py @@ -0,0 +1,10 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from twinkle.server.exceptions import RequestRejectedError, ResourceNotFoundError + + +class TrainingExpiredError(ResourceNotFoundError): + pass + + +class ProtocolMismatchError(RequestRejectedError): + pass diff --git a/src/twinkle/server/training/router.py b/src/twinkle/server/training/router.py new file mode 100644 index 000000000..125336da5 --- /dev/null +++ b/src/twinkle/server/training/router.py @@ -0,0 +1,262 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Stable CPU endpoint for one catalog model; requests route by training ID.""" +from __future__ import annotations + +import asyncio +import json +import ray +import uuid +from fastapi import Depends, FastAPI, Request +from fastapi.exceptions import RequestValidationError +from pydantic import ValidationError +from tinker import types as tinker +from types import SimpleNamespace + +from twinkle.hub import HubOperation +from twinkle.protocol.types import model as native +from twinkle.server.checkpoint import _resolve_client_save_dir, create_checkpoint_manager, validate_user_path +from twinkle.server.deployment import LazyCleanupMixin, bind_deployment, build_deployment_app +from twinkle.server.exceptions import RequestRejectedError, ResourceNotFoundError +from twinkle.server.lifecycle.envelope import envelope_from_record +from twinkle.server.middleware.auth import get_session_id_from_request, get_token_from_request +from twinkle.server.service_proxy import ServiceProxy +from twinkle.server.state import get_server_state +from twinkle.server.state.training_records import CreateTraining +from twinkle.server.task_queue import TaskQueueMixin +from twinkle.server.validation import assert_request_supported +from .actor import controller_name +from .catalog import TrainingCatalog +from .errors import ProtocolMismatchError + +NATIVE_REQUESTS = { + 'forward': native.ForwardRequest, + 'forward_only': native.ForwardOnlyRequest, + 'forward_backward': native.ForwardBackwardTaskRequest, + 'calculate_loss': native.AdapterRequest, + 'backward': native.AdapterRequest, + 'forward_from_data_plane': native.DataPlaneForwardRequest, + 'forward_backward_from_data_plane': native.DataPlaneForwardRequest, + 'forward_only_from_data_plane': native.DataPlaneForwardOnlyRequest, + 'step': native.StepRequest, + 'zero_grad': native.AdapterRequest, + 'lr_step': native.LrStepRequest, + 'clip_grad_norm': native.ClipGradNormRequest, + 'clip_grad_and_step': native.ClipGradAndStepRequest, + 'get_train_configs': native.AdapterRequest, + 'set_loss': native.SetLossRequest, + 'set_optimizer': native.SetOptimizerRequest, + 'set_lr_scheduler': native.SetLrSchedulerRequest, + 'set_template': native.SetTemplateRequest, + 'set_processor': native.SetProcessorRequest, + 'add_metric': native.AddMetricRequest, + 'apply_patch': native.ApplyPatchRequest, + 'calculate_metric': native.CalculateMetricRequest, + 'save': native.SaveRequest, + 'load': native.LoadRequest, + 'resume_from_checkpoint': native.ResumeFromCheckpointRequest, + 'remove_adapter': native.AdapterRequest, + 'add_adapter_to_model': native.AddAdapterRequest, + 'upload_to_hub': native.UploadToHubRequest, + 'create': native.CreateRequest, +} +TINKER_REQUESTS = { + 'create_model': tinker.CreateModelRequest, + 'get_info': tinker.GetInfoRequest, + 'unload_model': tinker.UnloadModelRequest, + 'forward': tinker.ForwardRequest, + 'forward_backward': tinker.ForwardBackwardRequest, + 'optim_step': tinker.OptimStepRequest, + 'save_weights': tinker.SaveWeightsRequest, + 'load_weights': tinker.LoadWeightsRequest, + 'save_weights_for_sampler': tinker.SaveWeightsForSamplerRequest, +} + + +class TrainingRouter(TaskQueueMixin, LazyCleanupMixin): + + async def __init__(self, config, model, namespace, http_options): + self.catalog = TrainingCatalog(config) + self.model = model + self.backend = model.backend + self.state = get_server_state() + self.store = 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.catalog.get(self.model.name, mode) + if config is not None: + rank = config.get('r') if isinstance(config, dict) else None + if (not isinstance(config, dict) or config.get('_TWINKLE_TYPE_') != 'LoraConfig' + or type(rank) is not int or not 0 < rank <= spec.config.max_r): + raise RequestRejectedError('LoRA config/rank is not supported by this training specification') + payload = body.model_dump(mode='json') + if body.save_dir: + try: + payload['save_dir'] = _resolve_client_save_dir(body.save_dir).as_posix() + except ValueError as exc: + raise RequestRejectedError(str(exc)) from exc + session = get_session_id_from_request(request) + await self._check_session(session) + ref = await self.controller.create.remote( + CreateTraining( + base_model=self.model.name, + mode=mode, + token=token, + session_id=session, + create_key=body.client_create_id, + client_type='twinkle', + adapter_label=body.adapter_name, + payload=payload).model_dump()) + return envelope_from_record(ref['request_id'], await self.state.get_future(ref['request_id'])) + if endpoint == 'upload_to_hub': + + async def upload(): + checkpoint = body.checkpoint_dir + if checkpoint.startswith('twinkle://'): + checkpoint = create_checkpoint_manager( + token, client_type='twinkle').resolve_checkpoint_dir(checkpoint) + await self.call_backend( + HubOperation.push_to_hub, + repo_id=body.hub_model_id, + folder_path=checkpoint, + token=body.hub_token or token, + private=True) + + ref = await self.schedule_background_task(upload, task_type=endpoint) + await self.store.own_future(ref['request_id'], token) + return envelope_from_record(ref['request_id'], await self.state.get_future(ref['request_id'])) + if not body.training_id: + raise ProtocolMismatchError('Upgrade Twinkle SDK: training_id is required') + session = get_session_id_from_request(request) + await self._check_session(session) + if endpoint == 'remove_adapter': + return await self._release(body.training_id, token, session, 'twinkle') + capability = endpoint if endpoint in ('forward', 'backward', 'calculate_loss') else None + assert_request_supported(SimpleNamespace(backend=self.backend), body, capability=capability) + route = await self.controller.route.remote(body.training_id, token, session) + if route['base_model'] != self.model.name: + raise ResourceNotFoundError('Training handle belongs to another model endpoint') + return await self.proxy.proxy_request(request, f'twinkle/{endpoint}', route['instance_id'], 'training-runtime') + + async def tinker_request(self, request, endpoint): + schema = TINKER_REQUESTS.get(endpoint) + if schema is None: + raise ResourceNotFoundError('Unknown training endpoint') + try: + body = schema.model_validate(await request.json()) + except ValidationError as exc: + raise RequestValidationError(exc.errors()) from exc + token = get_token_from_request(request) + if endpoint == 'create_model': + if body.base_model != self.model.name: + raise RequestRejectedError('Model does not match the catalog endpoint') + mode = 'lora' if body.lora_config is not None else 'full' + spec = self.catalog.get(self.model.name, mode) + if body.lora_config is not None and not 0 < body.lora_config.rank <= spec.config.max_r: + raise RequestRejectedError('LoRA rank exceeds the training specification') + await self._check_session(body.session_id) + return await self.controller.create.remote( + CreateTraining( + base_model=self.model.name, + mode=mode, + token=token, + session_id=body.session_id, + create_key=str(body.model_seq_id), + client_type='tinker', + payload=body.model_dump(mode='json')).model_dump()) + if endpoint == 'unload_model': + return await self._release(body.model_id, token, None, 'tinker') + route = await self.controller.route.remote(body.model_id, token) + if route['base_model'] != self.model.name: + raise ResourceNotFoundError('Training handle belongs to another model endpoint') + return await self.proxy.proxy_request(request, f'tinker/{endpoint}', route['instance_id'], 'training-runtime') + + async def _release(self, training_id, token, session, client_type): + await self.controller.route.remote(training_id, token, session) + + async def release(): + await self.controller.release.remote(training_id, token, session) + return {'model_id': training_id} if client_type == 'tinker' else None + + ref = await self.schedule_background_task(release, task_type='release_training') + await self.store.own_future(ref['request_id'], token) + if client_type == 'tinker': + return {'request_id': ref['request_id']} + return envelope_from_record(ref['request_id'], await self.state.get_future(ref['request_id'])) + + +def build_training_router(config, model, namespace, http_options, deploy_options): + + def routes(app: FastAPI, get_self): + + @app.get('/healthz') + async def health(self=Depends(get_self)): + diagnostic = await self.controller.diagnostics.remote(self.model.name) + active = [item for item in diagnostic['instances'] if item['state'] != 'STOPPED'] + return { + 'healthy': diagnostic['ready'] and not any(item['state'] == 'QUARANTINED' for item in active), + 'status': 'unloaded' if not active else 'loaded', + 'instances': active + } + + @app.post('/twinkle/{endpoint}') + async def native_route(request: Request, endpoint: str, self=Depends(get_self)): + return await self.native_request(request, endpoint) + + @app.post('/tinker/{endpoint}') + async def tinker_route(request: Request, endpoint: str, self=Depends(get_self)): + return await self.tinker_request(request, endpoint) + + async def shutdown(servable): + await servable.shutdown() + + app = build_deployment_app('TrainingRouter', routes, on_shutdown=shutdown) + return bind_deployment( + app, + TrainingRouter, + deploy_options, + deployment_name='TrainingRouter', + bind_args=(config, model, namespace, http_options)) diff --git a/src/twinkle/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..b8f11333f 100644 --- a/src/twinkle_client/auto/agent/server_tools.py +++ b/src/twinkle_client/auto/agent/server_tools.py @@ -373,8 +373,6 @@ def _generate_server_config( def _short(mid: str) -> str: return mid.split('/')[-1] if '/' in mid else mid - model_short = _short(model_id) - # Collect all model IDs for supported_models all_model_ids = [model_id] + [s['model_id'] for s in sampler_list] @@ -409,62 +407,41 @@ def _short(mid: str) -> str: }], }) - # 2. Build GPU-requiring applications (model + samplers), + # 2. Build static GPU-requiring sampler applications, # then sort by GPU count DESCENDING before appending. # Largest PG deploys first → it has the fewest node choices → # avoids GPU scheduling deadlock on single-machine multi-node. gpu_apps: list[tuple[int, dict]] = [] # (gpu_count, app_config) - # 2a. Training model worker (student) - gpu_apps.append(( - train_gpus, - { - 'name': - f'models-{model_short}', - 'route_prefix': - f'/api/v1/model/{model_id}', - 'import_path': - 'model', - 'args': { - 'backend': backend, - 'model_id': f'ms://{model_id}', - 'max_length': 500000, # total tokens per forward pass (must match max_input_tokens) - 'nproc_per_node': train_gpus, - 'device_group': { - 'name': 'model', - 'ranks': train_gpus, - 'device_type': 'cuda', - }, - 'device_mesh': { - 'device_type': 'cuda', - 'dp_size': train_gpus, - }, - 'queue_config': { - 'rps_limit': 100, - 'tps_limit': 100000, - 'max_input_tokens': 500000, - }, - 'adapter_config': { - 'adapter_timeout': 600, - }, + applications.append({ + 'name': f'model-{_short(model_id)}', + 'route_prefix': f'/api/v1/model/{model_id}', + 'import_path': 'model', + 'args': { + 'model_id': f'ms://{model_id}', + 'backend': backend, + 'device_type': 'cuda', + 'env_vars': { + 'TWINKLE_TRUST_REMOTE_CODE': '1' }, - 'deployments': [{ - 'name': 'ModelManagement', - 'autoscaling_config': { - 'min_replicas': 1, - 'max_replicas': 1, - 'target_ongoing_requests': 16, - }, - 'ray_actor_options': { - 'num_cpus': 0.1, - 'runtime_env': { - 'env_vars': { - 'TWINKLE_TRUST_REMOTE_CODE': '1' - }, + 'modes': { + 'lora': { + 'num_devices': train_gpus, + 'device_mesh': { + 'device_type': 'cuda', + 'dp_size': train_gpus }, - }, - }], - })) + 'max_length': 500000, + 'idle_timeout_s': 600, + 'queue_config': { + 'rps_limit': 100, + 'tps_limit': 100000, + 'max_input_tokens': 500000 + }, + } + }, + }, + }) # 2b. Sampler/teacher models sampler_name_count: dict[str, int] = {} diff --git a/src/twinkle_client/model/multi_lora_transformers.py b/src/twinkle_client/model/multi_lora_transformers.py index a9fbfd8ed..b445aca5c 100644 --- a/src/twinkle_client/model/multi_lora_transformers.py +++ b/src/twinkle_client/model/multi_lora_transformers.py @@ -3,6 +3,7 @@ import itertools import logging import threading +import uuid from collections.abc import Mapping from pathlib import Path from typing import TYPE_CHECKING, Any, Dict, Optional @@ -58,6 +59,9 @@ def __init__( self.model_id = model_id self.server_url = f'{self._transport.context.base_url}/model/{model_id}/twinkle' self.adapter_name = None + self._training_ids: dict[str, str] = {} + self._creating: set[str] = set() + self._training_lock = threading.Lock() # Per-client monotonic sequence for idempotent dedup of stateful training ops: # the server dedups on (session_id, seq_id) so a retried grad/step call is # applied at most once. Reserved once per call and reused on retry. @@ -84,12 +88,21 @@ def _ensure_created(self) -> None: with self._create_lock: if self._created: return - self._transport.post(f'{self.server_url}/create') + response = self._transport.post(f'{self.server_url}/create', json_data={'training_handles': True}) + if response.json().get('training_handles') is not True: + raise RuntimeError('Server does not support training handles; upgrade Twinkle server and SDK together.') self._created = True def _submit(self, endpoint: str, model_cls, response_cls, **values): """Build, send, and resolve one twinkle-native request.""" self._ensure_created() + if issubclass(model_cls, model_types.TrainingRequest): + name = values.get('adapter_name') + with self._training_lock: + training_id = values.get('training_id') or self._training_ids.get(name) + if training_id is None: + raise RuntimeError('Create this training with add_adapter_to_model on this object before using it.') + values['training_id'] = training_id body = build_request(model_cls, **values) response = self._transport.post_model(f'{self.server_url}/{endpoint}', body) return self._await_task(response, response_cls) @@ -132,22 +145,38 @@ def add_adapter_to_model( save_dir = kwargs.pop('save_dir', None) if save_dir: save_dir = Path(save_dir).expanduser().resolve().as_posix() - self._submit( - 'add_adapter_to_model', - model_types.AddAdapterRequest, - None, - adapter_name=adapter_name, - config=config, - save_dir=save_dir, - **kwargs) - self.adapter_name = adapter_name + with self._training_lock: + if adapter_name in self._creating: + raise RuntimeError('A training with this adapter_name is already being created on this object.') + self._creating.add(adapter_name) + try: + result = self._submit( + 'add_adapter_to_model', + model_types.AddAdapterRequest, + model_types.TrainingHandleResponse, + client_create_id=uuid.uuid4().hex, + adapter_name=adapter_name, + config=config, + save_dir=save_dir, + **kwargs) + with self._training_lock: + self._training_ids[adapter_name] = result.training_id + self.adapter_name = adapter_name + finally: + with self._training_lock: + self._creating.discard(adapter_name) def remove_adapter(self, adapter_name: str | None = None) -> None: """Release one client-owned adapter from the training component.""" name = adapter_name or self.adapter_name - self._submit('remove_adapter', model_types.AdapterRequest, None, adapter_name=name) - if name == self.adapter_name: - self.adapter_name = None + with self._training_lock: + training_id = self._training_ids.get(name) + self._submit('remove_adapter', model_types.AdapterRequest, None, adapter_name=name, training_id=training_id) + with self._training_lock: + if self._training_ids.get(name) == training_id: + self._training_ids.pop(name, None) + if name == self.adapter_name: + self.adapter_name = None # ------------------------------------------------------------------ # # Inline forward family diff --git a/tests/server/cli/test_cli.py b/tests/server/cli/test_cli.py index a49ae372e..ef859abca 100644 --- a/tests/server/cli/test_cli.py +++ b/tests/server/cli/test_cli.py @@ -89,5 +89,5 @@ def test_env_var_overrides_when_flag_omitted(monkeypatch) -> None: def test_example_config_loads_via_server_config() -> None: cfg = ServerConfig.from_yaml(EXAMPLE) assert isinstance(cfg, ServerConfig) - assert any(a.import_path == 'model' for a in cfg.applications) + assert cfg.training_config.models[0].backend == 'mock' assert any(a.import_path == 'sampler' for a in cfg.applications) diff --git a/tests/server/config/test_server_config.py b/tests/server/config/test_server_config.py index 927089ff2..8e47461b0 100644 --- a/tests/server/config/test_server_config.py +++ b/tests/server/config/test_server_config.py @@ -33,38 +33,25 @@ }), ) -_MODEL_APP = st.fixed_dictionaries({ - 'name': - st.just('m'), - 'route_prefix': - st.just('/api/v1/m'), - 'import_path': - st.just('model'), - 'args': - st.fixed_dictionaries({ +_TRAINING_MODEL = st.fixed_dictionaries({ + 'name': st.just('m'), 'import_path': st.just('model'), 'route_prefix': st.just('/api/v1/model/m'), + 'args': st.fixed_dictionaries({ 'model_id': st.just('model-id'), - 'device_group': st.just({ - 'name': 'g', - 'ranks': 1, - 'device_type': 'CPU' - }), - 'device_mesh': st.just({ - 'device_type': 'CPU', - 'dp_size': 1 - }), 'backend': st.sampled_from(['mock', 'transformers', 'megatron']), + 'device_type': st.just('cuda'), + 'modes': st.just({'lora': {'num_devices': 1, 'device_mesh': {'dp_size': 1}}}), }), }) _VALID_CONFIG = st.fixed_dictionaries({ 'persistence': _PERSISTENCE_VARIANTS, - 'applications': st.lists(_MODEL_APP, min_size=0, max_size=3), + 'applications': st.lists(_TRAINING_MODEL, min_size=0, max_size=1), }) # ---------- valid → fully validated ---------------------------------------- # -@settings(max_examples=100) +@settings(max_examples=100, deadline=None) @given(payload=_VALID_CONFIG) def test_valid_payload_yields_full_instance(payload: dict) -> None: cfg = ServerConfig.model_validate(payload) @@ -141,7 +128,7 @@ def test_sampler_docstring_values_match_selector() -> None: # ---------- round-trip fidelity ----------------------------------------- # -@settings(max_examples=100) +@settings(max_examples=100, deadline=None) @given(payload=_VALID_CONFIG) def test_round_trip_fidelity(payload: dict) -> None: cfg = ServerConfig.model_validate(payload) diff --git a/tests/server/config/test_training_config.py b/tests/server/config/test_training_config.py new file mode 100644 index 000000000..8bb8dfc72 --- /dev/null +++ b/tests/server/config/test_training_config.py @@ -0,0 +1,276 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +import copy +import math + +import pytest +from pydantic import ValidationError + +from twinkle.server.config import ServerConfig +from twinkle.server.config.training import TrainingConfig +from twinkle.server.exceptions import RequestRejectedError +from twinkle.server.training.catalog import TrainingCatalog + + +def model_config(): + return { + 'name': 'Qwen/small', 'model_id': 'ms://Qwen/small', 'backend': 'transformers', 'device_type': 'cuda', + 'modes': { + 'lora': {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}}, + 'full': {'num_devices': 4, 'device_mesh': {'fsdp_size': 4}}, + }, + } + + +def model_application(model=None): + model = dict(model or model_config()) + name = model.pop('name') + return {'name': 'model-' + name, 'route_prefix': '/api/v1/model/' + name, + 'import_path': 'model', 'args': model} + + +def test_catalog_defaults_and_mode_selection(): + config = ServerConfig(applications=[model_application()]) + catalog = TrainingCatalog(config.training_config) + shared = catalog.get('Qwen/small', 'lora') + exclusive = catalog.get('Qwen/small', 'full') + assert shared.capacity == 5 + assert exclusive.capacity == 1 + assert shared.config.min_replicas == exclusive.config.min_replicas == 0 + assert shared.config.idle_timeout_s == 1800 + assert config.training.allocation_timeout_s == 3 + assert config.to_yaml_dict()['applications'][0]['args']['model_id'] == 'ms://Qwen/small' + with pytest.raises(RequestRejectedError, match='Unknown training model'): + catalog.get('not-in-catalog', 'lora') + + +def test_existing_yaml_retains_heartbeat_timeout_and_max_lifetime(): + from pathlib import Path + paths = sorted(Path('cookbook/client/server').glob('**/server_config*.yaml')) + assert paths + for path in paths: + config = ServerConfig.from_yaml(path) + for app in config.applications: + if app.import_path != 'model': + continue + normalized = config.training_config.models[ + [item.model_id for item in config.training_config.models].index(app.args.model_id)] + mode = normalized.modes.full if app.args.train_mode == 'full' else normalized.modes.lora + legacy = app.args.adapter_config or {} + assert mode.idle_timeout_s == app.args.idle_timeout_s + assert normalized.legacy_expiration.session_timeout_s == legacy.get('adapter_timeout', 1800) + assert normalized.legacy_expiration.max_lifetime_s == legacy.get('adapter_max_lifetime', 36000) + + +@pytest.mark.parametrize('field,value', [ + ('num_devices', 0), ('num_devices', True), ('num_devices', 3), + ('min_replicas', -1), ('idle_timeout_s', math.inf), ('idle_timeout_s', math.nan), + ('idle_timeout_s', 0), ('ranks', [0, 1]), ('max_loras', 0), +]) +def test_invalid_mode_cannot_reach_allocation(field, value): + model = model_config() + model['modes']['lora'][field] = value + with pytest.raises(ValidationError): + TrainingConfig(models=[model]) + + +def test_ep_and_vpp_do_not_multiply_physical_device_count(): + model = model_config() + model['backend'] = 'megatron' + model['modes']['lora'] = { + 'num_devices': 8, + 'device_mesh': {'tp_size': 2, 'dp_size': 2, 'pp_size': 2, 'ep_size': 2, 'vpp_size': 2}, + } + assert TrainingConfig(models=[model]).models[0].modes.lora.num_devices == 8 + + +@pytest.mark.parametrize('mutation', ['duplicate', 'no_modes', 'full_lora_field', 'hybrid_megatron', 'mesh_type', 'min_max']) +def test_cross_field_constraints(mutation): + model = model_config() + models = [model] + if mutation == 'duplicate': + models.append(copy.deepcopy(model)) + elif mutation == 'no_modes': + model['modes'] = {} + elif mutation == 'full_lora_field': + model['modes']['full']['max_loras'] = 2 + elif mutation == 'hybrid_megatron': + model['backend'] = 'megatron' + model['modes']['lora']['hybrid'] = {'allocation_path': '/tmp/allocation.json'} + elif mutation == 'mesh_type': + model['modes']['lora']['device_mesh']['device_type'] = 'npu' + elif mutation == 'min_max': + model['modes']['lora'].update(min_replicas=2, max_replicas=1) + with pytest.raises(ValidationError): + TrainingConfig(models=models) + + +def test_full_only_is_explicit_and_lora_is_unavailable(): + model = model_config() + del model['modes']['lora'] + catalog = TrainingCatalog(TrainingConfig(models=[model])) + assert catalog.get('Qwen/small', 'full').capacity == 1 + with pytest.raises(RequestRejectedError, match='does not provide lora'): + catalog.get('Qwen/small', 'lora') + + +def test_training_namespace_conflicts_are_rejected(): + with pytest.raises(ValidationError, match='unique'): + ServerConfig(applications=[model_application(), model_application()]) + with pytest.raises(ValidationError, match='route conflicts'): + ServerConfig(applications=[model_application(), { + 'name': 'processor', 'import_path': 'processor', 'route_prefix': '/api/v1/model/Qwen/small', + }]) + with pytest.raises(ValidationError, match='models'): + ServerConfig(training={'models': [model_config()]}) + + +def test_existing_yaml_examples_keep_their_model_entries(monkeypatch): + from pathlib import Path + monkeypatch.setenv('TWINKLE_LOCAL_MODEL_PATH', '/tmp/model-cache') + paths = list(Path('cookbook/client/server').rglob('server_config*.yaml')) + paths += list(Path('cookbook/client/async_rl').glob('server_config.yaml')) + paths += list(Path('tests/server/config').glob('server_config*e2e*.yaml')) + paths += list(Path('tests/server/fixtures').glob('server_config_mock*.yaml')) + assert len(paths) == 9 + for path in paths: + config = ServerConfig.from_yaml(path) + models = [app for app in config.applications if app.import_path == 'model'] + assert len(models) == len(config.training_config.models) == 1 + model = config.training_config.models[0] + assert model.model_id == models[0].args.model_id + assert model.route_prefix == models[0].route_prefix + mode = model.modes.full or model.modes.lora + assert mode.min_replicas == 1 + assert ServerConfig.model_validate(config.to_yaml_dict()).training_config == config.training_config + + +def test_legacy_lifecycle_overrides_and_full_instances(): + app = {'name': 'model', 'route_prefix': '/api/v1/model/small', 'import_path': 'model', 'args': { + 'model_id': 'weights', 'backend': 'mock', 'train_mode': 'full', + 'device_group': {'ranks': 2, 'device_type': 'CPU'}, 'device_mesh': {'dp_size': 2}, + 'min_replicas': 0, 'max_replicas': 3, 'idle_timeout_s': 17, + 'adapter_config': {'adapter_timeout': 30}, + }, 'deployments': [{'num_replicas': 1}]} + config = ServerConfig(applications=[app]) + mode = config.training_config.models[0].modes.full + assert (mode.num_devices, mode.min_replicas, mode.max_replicas, mode.idle_timeout_s) == (2, 0, 3, 17) + assert config.training_config.models[0].modes.lora is None + + +def test_extended_modes_do_not_accept_ambiguous_single_mode_settings(): + app = model_application() + app['args']['train_mode'] = 'full' + with pytest.raises(ValidationError, match='cannot be mixed'): + ServerConfig(applications=[app]) + + +def test_launcher_keeps_training_and_sampling_capabilities_separate(monkeypatch): + from ray import serve + from twinkle.server.launcher.server_launcher import ServerLauncher + config = ServerConfig(applications=[model_application(), { + 'name': 'gateway', 'import_path': 'server', 'route_prefix': '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/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 legacy_application(): + return {'name': 'model', 'route_prefix': '/api/v1/model/small', 'import_path': 'model', 'args': { + 'model_id': 'weights', 'backend': 'mock', 'train_mode': 'lora', 'max_length': 128, + 'device_group': {'ranks': 2, 'device_type': 'cpu'}, 'device_mesh': {'dp_size': 2}, + 'env_vars': {'USER_OPTION': 'args-wins'}, 'adapter_config': {'adapter_timeout': 30}, + }, 'deployments': [{ + 'name': 'ModelManagement', 'max_ongoing_requests': 17, + 'autoscaling_config': {'min_replicas': 1, 'max_replicas': 2, 'target_ongoing_requests': 16}, + 'ray_actor_options': {'num_cpus': 2, 'runtime_env': { + 'env_vars': {'USER_OPTION': 'deployment', 'OTHER_OPTION': 'retained'}, 'pip': ['example==1.0'] + }}, + }]} + + +def test_legacy_options_reach_cpu_router_without_changing_yaml(monkeypatch): + from types import SimpleNamespace + from twinkle.server.launcher.server_launcher import ServerLauncher + from twinkle.server.training.actor import TrainingControllerActor + import ray + from ray import serve + config = ServerConfig(applications=[legacy_application()]) + remote = SimpleNamespace(initialize=SimpleNamespace(remote=lambda: None)) + monkeypatch.setattr(TrainingControllerActor, 'options', + lambda **kwargs: SimpleNamespace(remote=lambda *args: remote)) + monkeypatch.setattr(ray, 'get', lambda *args, **kwargs: None) + captured = {} + + def router(*args): + captured.update(args[-1]) + return 'app' + + monkeypatch.setattr('twinkle.server.training.router.build_training_router', router) + monkeypatch.setattr(serve, 'run', lambda *args, **kwargs: None) + ServerLauncher(config)._deploy_training() + assert captured['max_ongoing_requests'] == 17 + assert captured['autoscaling_config']['target_ongoing_requests'] == 16 + actor = captured['ray_actor_options'] + assert actor['num_cpus'] == 2 and actor['runtime_env']['pip'] == ['example==1.0'] + assert actor['runtime_env']['env_vars']['USER_OPTION'] == 'args-wins' + assert actor['runtime_env']['env_vars']['OTHER_OPTION'] == 'retained' + assert config.to_yaml_dict()['applications'][0]['deployments'][0] == legacy_application()['deployments'][0] + + +@pytest.mark.asyncio +async def test_legacy_options_reach_runtime_with_one_controller_owned_replica(monkeypatch): + import importlib + from twinkle.server.state.training_records import InstanceRecord + from twinkle.server.training.driver import RayInstanceDriver + from ray import serve + config = ServerConfig(applications=[legacy_application()]) + model = config.training_config.models[0] + driver = RayInstanceDriver(config.training_config, None, 'test', 'controller') + instance = InstanceRecord(instance_id='instance', controller_epoch='epoch', base_model=model.name, + mode='lora', startup_deadline=0, devices=['0', '1']) + captured = {} + + def builder(**kwargs): + captured.update(kwargs) + return 'app' + + monkeypatch.setattr('twinkle.server.model.app.build_model_app', builder) + monkeypatch.setattr(importlib.import_module('ray.util.placement_group'), 'get_placement_group', lambda name: 'pg') + monkeypatch.setattr(serve, 'run', lambda *args, **kwargs: None) + await driver.deploy(instance, driver._catalog.get(model.name, 'lora')) + options = captured['deploy_options'] + assert options['num_replicas'] == 1 and 'autoscaling_config' not in options + assert options['max_ongoing_requests'] == 17 + assert options['ray_actor_options']['num_cpus'] == 2 + assert options['ray_actor_options']['runtime_env']['pip'] == ['example==1.0'] + assert captured['max_length'] == 128 + assert captured['nproc_per_node'] == 2 + assert captured['queue_config'] == model.modes.lora.queue_config + assert not driver._deployments + + +def test_mutated_configuration_does_not_reuse_a_stale_catalog(): + config = ServerConfig(applications=[legacy_application()]) + assert config.training_config.models[0].modes.lora.min_replicas == 1 + config.applications[0].args.min_replicas = 0 + assert config.training_config.models[0].modes.lora.min_replicas == 0 + + +def test_unsupported_adapter_parameters_are_not_silently_ignored(): + app = legacy_application() + app['args']['adapter_config']['misspelled_timeout'] = 12 + with pytest.raises(ValidationError, match='misspelled_timeout'): + ServerConfig(applications=[app]) diff --git a/tests/server/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/gateway/test_proxy.py b/tests/server/gateway/test_proxy.py index 9b03e4a98..caa35deee 100644 --- a/tests/server/gateway/test_proxy.py +++ b/tests/server/gateway/test_proxy.py @@ -7,13 +7,16 @@ """ 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.service_proxy import ServiceProxy from twinkle.protocol.headers import H_MULTIPLEX, H_MULTIPLEX_LEGACY, H_REQUEST_ID @@ -51,6 +54,37 @@ def test_prepare_headers_sets_multiplex_from_request_id(): assert 'host' not in {k.lower() for k in headers} +@pytest.mark.asyncio +async def test_proxy_waits_are_bounded_and_timeout_keeps_existing_wire(): + closed = asyncio.Event() + + async def stall(reader, writer): + try: + await reader.read(4096) + await reader.read() + finally: + writer.close() + await writer.wait_closed() + closed.set() + + server = await asyncio.start_server(stall, '127.0.0.1', 0) + port = server.sockets[0].getsockname()[1] + proxy = ServiceProxy(http_options={'host': '127.0.0.1', 'port': port}) + assert all(value is not None and value > 0 for value in proxy.client.timeout.as_dict().values()) + proxy.client.timeout = httpx.Timeout(0.05) + request = _make_request(headers=[(H_REQUEST_ID.encode(), b'req-timeout')]) + try: + reply = await asyncio.wait_for(proxy.proxy_request(request, 'tinker/forward', 'small', 'model', + body_override=b'{}'), 1) + assert reply.status_code == 502 + assert json.loads(reply.body)['request_id'] == 'req-timeout' + finally: + await proxy.close() + server.close() + await server.wait_closed() + await asyncio.wait_for(closed.wait(), 1) + + @pytest.mark.asyncio async def test_proxy_request_502_fallback_returns_error_payload(): proxy = ServiceProxy(http_options={'host': 'localhost', 'port': 8000}, route_prefix='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/api/v1') diff --git a/tests/server/integration/on_demand_training/__init__.py b/tests/server/integration/on_demand_training/__init__.py new file mode 100644 index 000000000..0bb7387ac --- /dev/null +++ b/tests/server/integration/on_demand_training/__init__.py @@ -0,0 +1 @@ +"""On-demand training integration probes.""" diff --git a/tests/server/integration/on_demand_training/acceptance_support.py b/tests/server/integration/on_demand_training/acceptance_support.py new file mode 100644 index 000000000..b24b5d9f0 --- /dev/null +++ b/tests/server/integration/on_demand_training/acceptance_support.py @@ -0,0 +1,163 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Real SDK operations and bounded resource assertions for acceptance probes.""" +import json +import time + +import requests + +from tests.server.integration.on_demand_training.on_demand_sdk import unload_tinker_model + + +class TrainingHandle: + """A small fixed-token workload with optimizer checkpoint support.""" + + def __init__(self, base_url, model, protocol, label, mode='lora'): + from twinkle_client import init_tinker_client + from twinkle_client.http import ClientTransport + from twinkle_client.http.context import ClientContext + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + + self.base_url, self.model, self.protocol = base_url, model, protocol + self.label, self.mode = label, mode + self.token = 'acceptance-' + protocol + self.transport = None + self.closed = False + self.consumed_samples = 0 + if protocol == 'twinkle': + from peft import LoraConfig + response = requests.post(base_url + '/api/v1/twinkle/create_session', json={}, timeout=30) + response.raise_for_status() + self.transport = ClientTransport(ClientContext( + base_url=base_url, api_key=self.token, session_id=response.json()['session_id'])) + self.client = MultiLoraTransformersModel(model, transport=self.transport) + self.client.add_adapter_to_model( + label, LoraConfig(r=8, target_modules='all-linear') if mode == 'lora' else None, + gradient_accumulation_steps=2) + self.identity = self.client._training_ids[label] + self.configure() + else: + init_tinker_client() + from tinker import ServiceClient + self.service = ServiceClient(base_url=base_url, api_key=self.token) + self.client = (self.service.create_lora_training_client(base_model=model, rank=8) + if mode == 'lora' else self.service.create_full_training_client(base_model=model)) + self.identity = self.client.model_id + self.inputs = [dict(input_ids=[100, 200, 300, 400], labels=[200, 300, 400, 500], + position_ids=list(range(4)), attention_mask=[1] * 4) for _ in range(2)] + + def configure(self): + if self.protocol == 'twinkle': + # Qwen3.5 is multimodal: this hook connects the visual branch even + # for text-only batches, as in the established SFT helpers. + self.client.set_template('Qwen3_5Template', max_length=512) + self.client.set_loss('CrossEntropyLoss') + self.client.set_optimizer('Adam', lr=1e-4) + self.client.set_processor('InputProcessor', padding_side='right') + + def datums(self): + from twinkle.server.model.tinker_datum import input_feature_to_datum + return [input_feature_to_datum(feature) for feature in self.inputs] + + def step(self): + if self.protocol == 'twinkle': + for _ in range(3): + self.client.forward_backward(self.inputs, gradient_accumulation_steps=2) + self.client.clip_grad_and_step(gradient_accumulation_steps=2) + self.consumed_samples += len(self.inputs) + else: + from tinker import types + self.client.forward_backward(self.datums(), 'cross_entropy').result() + self.client.optim_step(types.AdamParams(learning_rate=1e-4)).result() + self.consumed_samples += len(self.inputs) + + def evaluate(self): + if self.protocol == 'twinkle': + result = self.client.forward_only(self.inputs).result + values = result['logps'] + else: + result = self.client.forward(self.datums(), 'cross_entropy').result() + values = [output['logprobs'].tolist() for output in result.loss_fn_outputs] + import numpy as np + values = np.asarray(values, dtype=float).reshape(-1) + assert values.size and np.isfinite(values).all(), 'Non-finite checkpoint probe' + return values + + def save(self, name): + if self.protocol == 'twinkle': + path = self.client.save( + name=name, save_optimizer=True, consumed_train_samples=self.consumed_samples).twinkle_path + else: + path = self.client.save_state(name).result().path + assert path, 'Checkpoint path is empty' + return path + + def resume(self, path): + if self.protocol == 'twinkle': + progress = self.client.resume_from_checkpoint(path) + self.consumed_samples = progress['consumed_train_samples'] + return progress + old = self.identity + self.client = self.service.create_training_client_from_state_with_optimizer(path=path) + self.identity = self.client.model_id + unload_tinker_model(self.service, old) + + def assert_stale(self): + try: + if self.protocol == 'twinkle': + self.client.get_train_configs().result + else: + self.client.get_info() + except Exception as exc: + assert any(term in str(exc).lower() for term in ('404', 'not found', 'handle', 'active', 'unavailable')), exc + else: + raise AssertionError('An old training handle remained usable') + + def close(self, stale=False): + if self.closed: + return + try: + if not stale: + if self.protocol == 'twinkle': + self.client.remove_adapter() + else: + unload_tinker_model(self.service, self.identity) + finally: + if self.transport is not None: + self.transport.close() + elif hasattr(self, 'service'): + self.service.holder.close() + self.closed = True + + +def wait_until(predicate, timeout, label): + deadline = time.monotonic() + timeout + while True: + result = predicate() + if result: + return result + if time.monotonic() >= deadline: + raise TimeoutError(label) + time.sleep(0.5) + + +def assert_released(controller, namespace, redis_url=None, prefix=None): + import ray + from ray import serve + from ray.util.placement_group import placement_group_table + diagnostic = ray.get(controller.diagnostics.remote(), timeout=30) + assert all(row['state'] == 'STOPPED' for row in diagnostic['instances']), diagnostic + owned_prefix = 'resources-' + namespace + '-' + assert not [name for name in serve.status().applications if name.startswith(owned_prefix)] + assert not [row for row in placement_group_table().values() + if row.get('name', '').startswith(owned_prefix) and row.get('state') != 'REMOVED'] + if redis_url: + import redis + client = redis.Redis.from_url(redis_url, decode_responses=True, socket_timeout=5) + try: + assert not list(client.scan_iter(match=prefix + 'model::training_*')) + assert not list(client.scan_iter(match=prefix + 'seq_ledger::training_*')) + for key in client.scan_iter(match=prefix + 'token_count::*'): + assert int(json.loads(client.get(key)) or 0) == 0, 'Quota drift' + finally: + client.close() + return True diff --git a/tests/server/integration/on_demand_training/allocation_probe.py b/tests/server/integration/on_demand_training/allocation_probe.py new file mode 100644 index 000000000..d3cc874ec --- /dev/null +++ b/tests/server/integration/on_demand_training/allocation_probe.py @@ -0,0 +1,97 @@ +# 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 +from twinkle.server.training.catalog import TrainingCatalog + + +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', + 'modes': { + 'full': { + 'num_devices': count, + 'device_mesh': { + 'fsdp_size': count + } + } + }, + }]) + spec = TrainingCatalog(config).get(str(index), 'full') + allocation = await pool.allocate(spec.config.num_devices, spec.model.device_type, 'epoch', str(index)) + allocations.append(allocation) + devices = [device for allocation in allocations for device in allocation.devices] + assert len(set(devices)) == 8, devices + assert len({allocation.node_id for allocation in allocations}) == 1 + try: + await pool.allocate(spec.config.num_devices, spec.model.device_type, 'epoch', 'overflow') + 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 = TrainingCatalog( + TrainingConfig(models=[{ + 'name': 'replacement', + 'model_id': 'replacement', + 'backend': 'mock', + 'device_type': 'cuda', + 'modes': { + 'full': { + 'num_devices': 2, + 'device_mesh': { + 'fsdp_size': 2 + } + } + }, + }])).get('replacement', 'full') + replacement = await pool.allocate(two.config.num_devices, two.model.device_type, 'epoch', 'replacement') + 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/on_demand_training/fragmentation_probe.py b/tests/server/integration/on_demand_training/fragmentation_probe.py new file mode 100644 index 000000000..d715d7e8c --- /dev/null +++ b/tests/server/integration/on_demand_training/fragmentation_probe.py @@ -0,0 +1,88 @@ +# 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 +from twinkle.server.training.catalog import TrainingCatalog + + +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 = TrainingCatalog( + TrainingConfig(models=[{ + 'name': 'two', + 'model_id': 'two', + 'backend': 'mock', + 'device_type': 'cuda', + 'modes': { + 'full': { + 'num_devices': 2, + 'device_mesh': { + 'fsdp_size': 2 + } + } + }, + }])).get('two', 'full') + for marker in ('left', 'right'): + pg = placement_group([{'GPU': 3, marker: 0.01}], strategy='STRICT_PACK') + occupied.append(pg) + await asyncio.wait_for(pg.ready(), 20) + assert ray.available_resources()['GPU'] == 2 + try: + await pool.allocate(spec.config.num_devices, spec.model.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.model.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/on_demand_training/http_probe.py b/tests/server/integration/on_demand_training/http_probe.py new file mode 100644 index 000000000..b469e7593 --- /dev/null +++ b/tests/server/integration/on_demand_training/http_probe.py @@ -0,0 +1,253 @@ +# 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 + }, + 'training': { + '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', + 'device_type': 'cpu', + 'modes': { + 'lora': { + 'num_devices': 1, + 'device_mesh': { + 'dp_size': 1 + }, + 'max_loras': 2, + 'max_replicas': 1 + }, + 'full': { + 'num_devices': 1, + 'device_mesh': { + 'dp_size': 1 + } + } + } + }, + }, { + 'name': 'other', + 'import_path': 'model', + 'route_prefix': '/api/v1/model/other', + 'args': { + 'model_id': 'other', + 'backend': 'mock', + 'train_mode': 'full', + 'device_group': {'ranks': 1, 'device_type': 'cpu'}, + 'device_mesh': {'dp_size': 1}, + 'min_replicas': 0, + 'max_replicas': 1, + }, + 'deployments': [{ + 'max_ongoing_requests': 8, + 'autoscaling_config': {'min_replicas': 1, 'max_replicas': 1, 'target_ongoing_requests': 4}, + 'ray_actor_options': {'num_cpus': 0.1, 'runtime_env': {'env_vars': {'LEGACY_CONFIG_PROBE': '1'}}}, + }], + }, { + 'name': 'gateway', + 'import_path': 'server', + 'route_prefix': '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/api/v1', + 'deployments': [{ + 'ray_actor_options': { + 'num_cpus': 0.1 + } + }] + }], + }) + launcher = ServerLauncher(config) + launcher._start_serve() + base = f'http://127.0.0.1:{port}' + + def wait_epoch(controller, old_epoch): + deadline = time.monotonic() + 90 + last_error = None + while time.monotonic() < deadline: + try: + diagnostic = ray.get(controller.diagnostics.remote(), timeout=2) + if diagnostic['epoch'] != old_epoch: + return diagnostic + except ray.exceptions.RayError as exc: + last_error = exc + time.sleep(0.1) + raise AssertionError(f'Controller did not recover: {last_error}') + + try: + for spec in config.applications: + if spec.import_path != 'model': + launcher._deploy_application(spec) + launcher._deploy_training() + controller = launcher._training_controller + assert ray.get(controller.diagnostics.remote())['instances'] == [] + session = requests.post(base + '/api/v1/twinkle/create_session', json={}, timeout=10).json()['session_id'] + transport = ClientTransport(ClientContext(base_url=base, api_key='probe-owner', session_id=session)) + first = MultiLoraTransformersModel('small', transport=transport) + second = MultiLoraTransformersModel('small', transport=transport) + first.add_adapter_to_model('same', None) + instance = next( + item for item in ray.get(controller.diagnostics.remote())['instances'] + if first._training_ids['same'] in item['training_ids']) + bypass = requests.post( + base + f"/api/v1/training-runtime/{instance['instance_id']}/twinkle/add_adapter_to_model", + json={ + 'adapter_name': 'bypass', + 'config': None + }, + headers={'Authorization': 'Bearer probe-owner'}, + timeout=10) + assert bypass.status_code == 400 and 'catalog endpoint' in bypass.text, bypass.text + second.add_adapter_to_model('same', None) + old_id = first._training_ids['same'] + assert old_id != second._training_ids['same'] + assert json.loads(first.get_train_configs().result) == {} + first.zero_grad() + first.remove_adapter() + first.add_adapter_to_model('same', None) + assert old_id != first._training_ids['same'] + first.remove_adapter() + second.remove_adapter() + from peft import LoraConfig + native_lora = MultiLoraTransformersModel('small', transport=transport) + native_lora.add_adapter_to_model('native-lora', LoraConfig(r=8, target_modules='all-linear')) + native_lora.zero_grad() + native_lora.remove_adapter() + init_tinker_client() + from tinker import ServiceClient + service = ServiceClient(base_url=base, api_key='probe-tinker') + shared1 = service.create_lora_training_client(base_model='small', rank=8) + shared2 = service.create_lora_training_client(base_model='small', rank=8) + assert shared1.model_id != shared2.model_id + assert shared1.get_info().is_lora + assert shared2.get_info().model_name == 'small' + from tinker import types + datum = types.Datum( + model_input=types.ModelInput.from_ints([1, 2, 3]), + loss_fn_inputs={ + 'target_tokens': types.TensorData(data=[2, 3, 4], dtype='int64', shape=[3]), + 'weights': types.TensorData(data=[1.0, 1.0, 1.0], dtype='float32', shape=[3]), + }) + assert shared1.forward_backward([datum], 'cross_entropy').result().metrics + shared1.optim_step(types.AdamParams(learning_rate=1e-4)).result() + exclusive = service.create_full_training_client(base_model='small') + assert not exclusive.get_info().is_lora + assert exclusive.forward_backward([datum], 'cross_entropy').result().metrics + exclusive.optim_step(types.AdamParams(learning_rate=1e-4)).result() + assert unload_tinker_model(service, exclusive.model_id).model_id == exclusive.model_id + assert shared1.get_info().is_lora and shared2.get_info().is_lora + try: + service.create_lora_training_client(base_model='small', rank=8) + raise AssertionError('A third LoRA user bypassed configured capacity') + except Exception as exc: + assert 'replica limit' in str(exc).lower(), exc + assert unload_tinker_model(service, shared2.model_id).model_id == shared2.model_id + replacement = service.create_lora_training_client(base_model='small', rank=8) + assert replacement.model_id != shared2.model_id + assert shared1.get_info().is_lora + old_epoch = ray.get(controller.diagnostics.remote())['epoch'] + ray.kill(controller, no_restart=False) + diagnostic = wait_epoch(controller, old_epoch) + assert not [item for item in diagnostic['instances'] if item['state'] != 'STOPPED'] + stale = requests.post( + base + '/api/v1/get_info', + json={ + 'model_id': shared1.model_id, + 'type': 'get_info' + }, + headers={ + 'Authorization': 'Bearer probe-tinker', + 'x-request-id': uuid.uuid4().hex + }, + timeout=10) + assert stale.status_code == 404, stale.text + before = ray.available_resources().get('CPU', 0) + for index in range(10): + model = MultiLoraTransformersModel('small' if index % 2 == 0 else 'other', transport=transport) + model.add_adapter_to_model('cycle', None) + model.zero_grad() + model.remove_adapter() + print(json.dumps({'phase': 'cycles', 'completed': index + 1}), flush=True) + assert ray.available_resources().get('CPU', 0) >= before - 0.02 + final = ray.get(controller.diagnostics.remote()) + assert all(item['state'] == 'STOPPED' for item in final['instances']) + live = MultiLoraTransformersModel('other', transport=transport) + live.add_adapter_to_model('state-loss', None) + ray.kill(ray.get_actor('twinkle_state_actor', namespace=namespace), no_restart=False) + after_loss = wait_epoch(controller, final['epoch']) + assert not [item for item in after_loss['instances'] if item['state'] != 'STOPPED'] + session = requests.post(base + '/api/v1/twinkle/create_session', json={}, timeout=10).json()['session_id'] + fresh_transport = ClientTransport(ClientContext(base_url=base, api_key='probe-owner', session_id=session)) + fresh = MultiLoraTransformersModel('other', transport=fresh_transport) + fresh.add_adapter_to_model('state-loss', None) + assert fresh._training_ids['state-loss'] != live._training_ids['state-loss'] + fresh.remove_adapter() + fresh_transport.close() + print( + json.dumps({ + 'passed': True, + 'native_recreation': True, + 'native_lora': True, + 'tinker_shared': True, + 'tinker_full': True, + 'restart': True, + 'state_loss': True, + 'runtime_lifecycle_rejected': True, + 'cycles': 10, + 'instances': len(final['instances']) + })) + transport.close() + finally: + if getattr(launcher, '_training_controller', None) is not None: + ray.get(launcher._training_controller.shutdown.remote(), timeout=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..35a30ee6d --- /dev/null +++ b/tests/server/integration/on_demand_training/matrix_validation.py @@ -0,0 +1,441 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Run the full training matrix on an isolated on-demand server.""" +import argparse +import atexit +import json +import math +import os +import ray +import socket +import subprocess +import sys +import time +import uuid +import yaml +from pathlib import Path +from ray import serve +from xml.etree import ElementTree + +_ROOT = Path(__file__).resolve().parents[4] +if str(_ROOT) not in sys.path: + sys.path.insert(0, str(_ROOT)) + +from twinkle.server.config import ServerConfig +from twinkle.server.launcher.server_launcher import ServerLauncher + + +def validate_controller_restart(base_url, controller, model_name): + """Train, restart, restore optimizer checkpoints and recover a failed worker.""" + import numpy as np + from ray.util import list_named_actors + from tests.server.integration.on_demand_training.acceptance_support import ( + TrainingHandle, assert_released, wait_until) + + namespace = ray.get_runtime_context().namespace + old = [TrainingHandle(base_url, model_name, protocol, 'restart-' + protocol) + for protocol in ('twinkle', 'tinker')] + snapshots = [] + try: + for handle in old: + handle.step() + snapshots.append((handle.save('restart-state'), handle.evaluate(), handle.consumed_samples)) + old_epoch = ray.get(controller.diagnostics.remote())['epoch'] + ray.kill(controller, no_restart=False) + + def restarted(): + try: + diagnostic = ray.get(controller.diagnostics.remote(), timeout=10) + return diagnostic if diagnostic['epoch'] != old_epoch else None + except ray.exceptions.RayError: + return None + + diagnostic = wait_until(restarted, 600, 'Controller restart timed out') + assert_released(controller, namespace) + for handle in old: + handle.assert_stale() + restored = [] + try: + # Keep restored peers on one runtime: checkpoint/optimizer isolation + # is checked without reloading the base model for each protocol. + for protocol, (path, values, consumed) in zip(('twinkle', 'tinker'), snapshots): + fresh = TrainingHandle(base_url, model_name, protocol, 'restored-' + protocol) + restored.append(fresh) + progress = fresh.resume(path) + if protocol == 'twinkle': + assert progress['consumed_train_samples'] == consumed + np.testing.assert_allclose(fresh.evaluate(), values, rtol=0, atol=2e-3) + fresh.step() + active = next(row for row in ray.get(controller.diagnostics.remote())['instances'] + if restored[0].identity in row['training_ids']) + assert set(active['training_ids']) == {handle.identity for handle in restored} + owned = [row for row in list_named_actors(all_namespaces=True) + if row['namespace'] == namespace and active['instance_id'] in row['name'] + and row['name'].startswith('model-') and not row['name'].endswith('-registry')] + assert owned, 'No model worker available for fault injection' + ray.kill(ray.get_actor(owned[0]['name'], namespace=namespace), no_restart=True) + wait_until(lambda: all(row['state'] == 'STOPPED' + for row in ray.get(controller.diagnostics.remote())['instances']), + 600, 'Worker failure cleanup timed out') + for handle in restored: + handle.assert_stale() + assert_released(controller, namespace) + finally: + for handle in restored: + handle.close(stale=True) + fresh = TrainingHandle(base_url, model_name, 'twinkle', 'after-worker-fault') + try: + fresh.step() + finally: + fresh.close() + finally: + for handle in old: + handle.close(stale=True) + return {'epoch': diagnostic['epoch'], 'worker_fault_recovered': True, 'checkpoint_restored': True} + + +def validate_idempotent_create(base_url, controller, model_name): + """Concurrent HTTP retries share one receipt, handle and allocation.""" + from concurrent.futures import ThreadPoolExecutor + import requests + from peft import LoraConfig + from twinkle.protocol.serialize import serialize_object + from twinkle.protocol.types.model import TrainingHandleResponse + from twinkle_client._future import resolve_response + from twinkle_client.http import ClientTransport + from twinkle_client.http.context import ClientContext + from tests.server.integration.on_demand_training.acceptance_support import assert_released + + session_reply = requests.post(base_url + '/api/v1/twinkle/create_session', json={}, timeout=30) + session_reply.raise_for_status() + session = session_reply.json()['session_id'] + transport = ClientTransport(ClientContext(base_url=base_url, api_key='idempotency', session_id=session)) + headers = {'Authorization': 'Bearer idempotency', 'X-Twinkle-Session-Id': session} + payload = {'client_create_id': uuid.uuid4().hex, 'adapter_name': 'duplicate', + 'config': serialize_object(LoraConfig(r=8, target_modules='all-linear')), + 'gradient_accumulation_steps': 2} + endpoint = base_url + '/api/v1/model/' + model_name + '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/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 + '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/twinkle/remove_adapter', + json={'training_id': handle.training_id, 'adapter_name': 'duplicate'}, + headers={**headers, 'x-request-id': uuid.uuid4().hex}, timeout=60) + reply.raise_for_status() + resolve_response(reply, None, transport=transport) + assert_released(controller, ray.get_runtime_context().namespace) + finally: + transport.close() + return True + + +def validate_capacity(base_url, controller, model_name, capacity): + """Assert shared LoRA capacity and full-training exclusivity on real devices.""" + from tests.server.integration.on_demand_training.acceptance_support import TrainingHandle, assert_released + handles = [] + try: + for index in range(capacity): + handles.append(TrainingHandle(base_url, model_name, ('twinkle', 'tinker')[index % 2], + 'capacity-' + str(index))) + live = [row for row in ray.get(controller.diagnostics.remote())['instances'] if row['state'] != 'STOPPED'] + assert len(live) == 1 and len(live[0]['training_ids']) == capacity + try: + overflow = TrainingHandle(base_url, model_name, 'tinker', 'overflow') + except Exception as exc: + assert 'replica limit' in str(exc).lower(), exc + else: + overflow.close() + raise AssertionError('LoRA capacity was exceeded') + handles[0].step() + finally: + for handle in reversed(handles): + handle.close() + assert_released(controller, ray.get_runtime_context().namespace) + exclusive = TrainingHandle(base_url, model_name, 'twinkle', 'exclusive', mode='full') + try: + exclusive.step() + live = [row for row in ray.get(controller.diagnostics.remote())['instances'] if row['state'] != 'STOPPED'] + assert len(live) == 1 and live[0]['mode'] == 'full' and len(live[0]['training_ids']) == 1 + try: + overflow = TrainingHandle(base_url, model_name, 'tinker', 'exclusive-overflow', mode='full') + except Exception as exc: + assert 'replica limit' in str(exc).lower(), exc + else: + overflow.close() + raise AssertionError('Full-training instance was shared') + finally: + exclusive.close() + assert_released(controller, ray.get_runtime_context().namespace) + return True + + +def validate_redis_restart(base_url, controller, model_name, managed): + """Only interrupt a Redis process started and owned by this harness.""" + import requests + from tests.server.integration.on_demand_training.acceptance_support import TrainingHandle + handle = TrainingHandle(base_url, model_name, 'twinkle', 'redis-fault') + try: + handle.step() + before = ray.get(controller.diagnostics.remote()) + + def unavailable(): + response = requests.post( + base_url + '/api/v1/model/' + model_name + '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/twinkle/get_train_configs', + json={'training_id': handle.identity, 'adapter_name': handle.label}, + headers={'Authorization': 'Bearer ' + handle.token, 'x-request-id': uuid.uuid4().hex, + 'X-Twinkle-Session-Id': handle.transport.context.session_id}, timeout=15) + assert response.status_code >= 500, response.text + + managed.fault(unavailable) + after = ray.get(controller.diagnostics.remote(), timeout=30) + assert after['epoch'] == before['epoch'] + assert {row['instance_id'] for row in before['instances']} == { + row['instance_id'] for row in after['instances']} + handle.step() + finally: + handle.close() + return True + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--backend', required=True, choices=['transformers', 'megatron']) + parser.add_argument('--model', required=True) + parser.add_argument('--run-dir', required=True) + parser.add_argument('--persistence', choices=['memory', 'redis'], default='memory') + parser.add_argument('--managed-redis', action='store_true', help='Start an owned Redis for restart injection') + parser.add_argument('--redis-url', default=os.environ.get('TWINKLE_TEST_REDIS_URL')) + parser.add_argument('--startup-timeout-s', type=float, default=1800) + parser.add_argument('--shutdown-timeout-s', type=float, default=180) + parser.add_argument('--recovery-rounds', type=int, default=1) + parser.add_argument('--skip-full-cycle', action='store_true') + parser.add_argument('--recovery-only', action='store_true') + parser.add_argument('--full-cycle-only', action='store_true') + parser.add_argument('--stress-iterations', type=int, default=0) + parser.add_argument('--soak-seconds', type=float, default=0) + args = parser.parse_args() + if args.persistence == 'redis' and not args.redis_url and not args.managed_redis: + parser.error('--redis-url or --managed-redis is required for Redis persistence') + if args.managed_redis and args.persistence != 'redis': + parser.error('--managed-redis requires --persistence redis') + if any(not math.isfinite(value) or value <= 0 + for value in (args.startup_timeout_s, args.shutdown_timeout_s)): + parser.error('startup and shutdown timeouts must be finite and positive') + if args.recovery_rounds < 0 or args.stress_iterations < 0 or args.soak_seconds < 0: + parser.error('validation budgets must be non-negative') + root = Path(__file__).resolve().parents[4] + run_dir = Path(args.run_dir) + run_dir.mkdir(parents=True, exist_ok=True) + (run_dir / 'checkpoints').mkdir(parents=True, exist_ok=True) + managed = None + if args.managed_redis: + from tests.server.integration.on_demand_training.managed_redis import ManagedRedis + managed = ManagedRedis(run_dir / 'redis') + atexit.register(managed.stop) + managed.start() + args.redis_url = managed.url + namespace = 'matrix_' + uuid.uuid4().hex[:8] + with socket.socket() as listener: + listener.bind(('127.0.0.1', 0)) + port = listener.getsockname()[1] + source = root / 'tests/server/config' / ('server_config_4b_e2e.yaml' if args.backend == 'transformers' else + 'server_config_4b_e2e_megatron.yaml') + raw = yaml.safe_load(source.read_text()) + raw.setdefault('training', {}).update({ + 'startup_timeout_s': args.startup_timeout_s, + 'shutdown_timeout_s': args.shutdown_timeout_s, + }) + redis_prefix = f'on-demand-test::{namespace}::' + raw['persistence'] = ({ + 'mode': 'redis', + 'redis_url': args.redis_url, + 'key_prefix': redis_prefix, + } if args.persistence == 'redis' else {'mode': 'memory'}) + raw['ray_namespace'] = namespace + raw['http_options'] = {'host': '127.0.0.1', 'port': port} + for app in raw['applications']: + params = app['args'] + if app['import_path'] == 'model': + common = { + 'num_devices': 2, 'device_mesh': {'device_type': 'cuda', 'dp_size': 2}, + 'min_replicas': 0, 'max_replicas': 1, 'idle_timeout_s': 3600, + 'max_length': params.get('max_length'), + 'queue_config': {**params.get('queue_config', {}), 'execution_timeout': 600}, + } + app['args'] = { + 'model_id': args.model, 'backend': args.backend, 'device_type': 'cuda', + 'modes': {'lora': {**common, 'max_loras': 10, 'max_r': 32}, 'full': common}, + 'env_vars': {'TWINKLE_FAIL_FAST': '0', 'TWINKLE_TRUST_REMOTE_CODE': '1'}, + } + app.pop('deployments', None) + elif app['import_path'] == 'sampler': + params['model_id'] = args.model + params['engine_args']['attention_config'] = {'flash_attn_version': 3} + if args.recovery_only: + raw['applications'] = [app for app in raw['applications'] if app['import_path'] in ('server', 'model')] + config = ServerConfig.model_validate(raw) + evidence_config = config.to_yaml_dict() + if args.persistence == 'redis': + evidence_config['persistence']['redis_url'] = '' + (run_dir / 'server.yaml').write_text(yaml.safe_dump(evidence_config, sort_keys=False)) + os.environ.update({ + 'TWINKLE_STATE_RAY_NAMESPACE': namespace, + 'TWINKLE_PERSISTENCE_MODE': args.persistence, + 'TWINKLE_DEFAULT_SAVE_DIR': str(run_dir / 'checkpoints'), + 'TWINKLE_MODEL_ID_ALIASES': json.dumps({'Qwen/Qwen3.5-4B': args.model}), + 'TWINKLE_TRUST_REMOTE_CODE': '1', + 'TWINKLE_FAIL_FAST': '0', + 'TWINKLE_TEST_GPU_E2E': '1', + 'TWINKLE_TEST_BACKEND': args.backend, + 'TWINKLE_SERVER_URL': f'http://127.0.0.1:{port}', + 'TWINKLE_TEST_OPERATION_TIMEOUT': '600', + 'PYTHONPATH': os.pathsep.join((str(root / 'src'), str(root))), + }) + if args.persistence == 'redis': + os.environ.update({ + 'TWINKLE_PERSISTENCE_REDIS_URL': args.redis_url, + 'TWINKLE_PERSISTENCE_KEY_PREFIX': redis_prefix, + }) + else: + os.environ.pop('TWINKLE_PERSISTENCE_REDIS_URL', None) + os.environ.pop('TWINKLE_PERSISTENCE_KEY_PREFIX', None) + evidence = { + 'backend': args.backend, + 'persistence': args.persistence, + 'namespace': namespace, + 'cases': [], + 'lazy_start': False, + } + ray.init( + address='local', + namespace=namespace, + num_cpus=24, + num_gpus=2 if args.recovery_only else 3, + _temp_dir='/tmp/twm-' + namespace[-8:], + include_dashboard=False) + launcher = ServerLauncher(config) + try: + launcher._start_serve() + for app in config.applications: + launcher._deploy_application(app) + launcher._deploy_training() + assert ray.get(launcher._training_controller.diagnostics.remote())['instances'] == [] + evidence['lazy_start'] = True + filenames = ['test_sft_e2e.py', 'test_dpo_e2e.py', 'test_grpo_e2e.py'] + if not args.skip_full_cycle: + filenames.append('test_full_cycle_e2e.py') + if args.recovery_only: + filenames = [] + elif args.full_cycle_only: + filenames = ['test_full_cycle_e2e.py'] + for filename in filenames: + start = time.monotonic() + with (run_dir / (filename + '.log')).open('w') as output: + result = subprocess.run([ + sys.executable, + '-m', + 'pytest', + '--noconftest', # HTTP clients do not need a second local Ray fixture. + str(root / 'tests/server/integration' / filename), + '-vv', + '-s', + '--junitxml=' + str(run_dir / (filename + '.xml')), + ], + stdout=output, + stderr=subprocess.STDOUT, + env=os.environ.copy(), + timeout=7200) + junit = run_dir / (filename + '.xml') + cases = ElementTree.parse(junit).getroot().findall('.//testcase') if junit.exists() else [] + skipped = sum(case.find('skipped') is not None for case in cases) + evidence['cases'].append({ + 'file': filename, + 'returncode': result.returncode, + 'tests': len(cases), + 'skipped': skipped, + 'seconds': time.monotonic() - start + }) + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + assert result.returncode == 0, f'{filename} failed; see {run_dir / (filename + ".log")}' + assert len(cases) == (1 if filename == 'test_full_cycle_e2e.py' else 2) and skipped == 0 + base_url = f'http://127.0.0.1:{port}' + model_name = config.training_config.models[0].name + evidence['recovery'] = [] + for _ in range(args.recovery_rounds): + recovery = validate_controller_restart(base_url, launcher._training_controller, model_name) + evidence['recovery'].append(recovery) + evidence['restart_epoch'] = recovery['epoch'] + evidence['worker_fault_recovered'] = recovery['worker_fault_recovered'] + if managed is not None: + evidence['redis_fault_rounds'] = evidence.get('redis_fault_rounds', 0) + int( + validate_redis_restart(base_url, launcher._training_controller, model_name, managed)) + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + if args.recovery_rounds: + evidence['idempotency_checked'] = validate_idempotent_create( + base_url, launcher._training_controller, model_name) + evidence['capacity_checked'] = validate_capacity( + base_url, launcher._training_controller, model_name, + config.training_config.models[0].modes.lora.max_loras) + if args.stress_iterations or args.soak_seconds: + from tests.server.integration.on_demand_training.acceptance_support import assert_released + controller = launcher._training_controller + assert_released(controller, namespace, + args.redis_url if args.persistence == 'redis' else None, redis_prefix) + from tests.server.integration.on_demand_training.stress_validation import run_stress + evidence['stress'] = [] + budgets = ([('short-' + str(index), args.stress_iterations, 0) for index in range(3)] + if args.stress_iterations else []) + if args.soak_seconds: + budgets.append(('soak', 0, args.soak_seconds)) + for name, iterations, duration in budgets: + evidence['stress'].append(run_stress( + f'http://127.0.0.1:{port}', config.training_config.models[0].name, + concurrency=2, iterations=iterations, duration=duration, + controller=launcher._training_controller, namespace=namespace, + redis_url=args.redis_url if args.persistence == 'redis' else None, prefix=redis_prefix, + progress_path=run_dir / (name + '.json'))) + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + finally: + try: + controller = getattr(launcher, '_training_controller', None) + if controller is not None: + ray.get(controller.shutdown.remote(), timeout=args.shutdown_timeout_s) + diagnostic = ray.get(controller.diagnostics.remote(), timeout=30) + assert all(row['state'] == 'STOPPED' for row in diagnostic['instances']) + from tests.server.integration.on_demand_training.acceptance_support import assert_released + assert_released(controller, namespace, + args.redis_url if args.persistence == 'redis' else None, redis_prefix) + evidence['training_resources_released'] = True + launcher._shutdown() + finally: + serve.shutdown() + ray.shutdown() + if args.persistence == 'redis' and evidence.get('training_resources_released'): + import redis + client = redis.Redis.from_url(args.redis_url) + try: + keys = list(client.scan_iter(match=redis_prefix + '*', count=500)) + if keys: + client.delete(*keys) + finally: + client.close() + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + if managed is not None: + managed.stop() + evidence['passed'] = True + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/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/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_full_cycle_e2e.py b/tests/server/integration/test_full_cycle_e2e.py index 6fa31f006..12d17de9c 100644 --- a/tests/server/integration/test_full_cycle_e2e.py +++ b/tests/server/integration/test_full_cycle_e2e.py @@ -6,8 +6,7 @@ Backend selection via env var TWINKLE_TEST_BACKEND: - "transformers" (default): all 6 phases run strictly - - "megatron": Phase C/D skipped (known multi-LoRA strict-load bug), - Phase E/F best-effort (GPU OOM possible) + - "megatron": all 6 phases run strictly (requires a working optimizer-resume backend) Phase A — initial training (STEPS_PHASE_A steps) Phase B — keep training STEPS_PHASE_B more steps, save again @@ -58,9 +57,9 @@ assert BACKEND in ('transformers', 'megatron'), f'Invalid TWINKLE_TEST_BACKEND={BACKEND!r}' BASE_MODEL = 'Qwen/Qwen3.5-4B' -BASE_URL = 'http://localhost:9000' +BASE_URL = os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:9000') API_KEY = 'EMPTY_API_KEY' -SAVE_DIR = '/mnt/nas2/yunlin.myl/twinkle/output/twinkle_e2e_full_cycle' +SAVE_DIR = os.environ.get('TWINKLE_DEFAULT_SAVE_DIR', 'output/twinkle_e2e_full_cycle') STEPS_PHASE_A = 60 STEPS_PHASE_B = 4 STEPS_PHASE_D = 4 diff --git a/tests/server/integration/test_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/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..cde575825 100644 --- a/tests/server/lifecycle/test_retrieve_endpoint.py +++ b/tests/server/lifecycle/test_retrieve_endpoint.py @@ -17,6 +17,9 @@ class _State: + + async def assert_future_owner(self, request_id, token): + return None """A fake ServerState whose get_future returns a fixed record (or None).""" def __init__(self, record): diff --git a/tests/server/lifecycle/test_run_submit_dedup.py b/tests/server/lifecycle/test_run_submit_dedup.py index 81d9c840b..53d1b6925 100644 --- a/tests/server/lifecycle/test_run_submit_dedup.py +++ b/tests/server/lifecycle/test_run_submit_dedup.py @@ -77,7 +77,7 @@ async def test_release_kept_when_task_already_enqueued(): with pytest.raises(RuntimeError): await run_submit(mgmt, _request(), _body(), task_type='forward_backward', backend_call=_call) assert mgmt.state.released == [], 'claim wrongly released for an already-enqueued task' - assert 'seq::sess-1::sess-1-ad::7' in mgmt.state.claimed + assert 'seq::sess-1::sess-1-ad::forward_backward::7' in mgmt.state.claimed @pytest.mark.asyncio @@ -86,7 +86,7 @@ async def test_release_when_never_enqueued(): mgmt = _FakeManagement(record_after_claim=None) with pytest.raises(RuntimeError): await run_submit(mgmt, _request(), _body(), task_type='forward_backward', backend_call=_call) - assert mgmt.state.released == ['seq::sess-1::sess-1-ad::7'], 'claim should be released when nothing enqueued' + assert mgmt.state.released == ['seq::sess-1::sess-1-ad::forward_backward::7'], 'claim should be released when nothing enqueued' @pytest.mark.asyncio @@ -121,5 +121,5 @@ async def test_dedup_key_is_scoped_per_adapter(): backend_call=_call) claimed = set(mgmt.state.claimed) - assert claimed == {'seq::sess-1::sess-1-lora-A::1', 'seq::sess-1::sess-1-lora-B::1'}, ( + assert claimed == {'seq::sess-1::sess-1-lora-A::forward_backward::1', 'seq::sess-1::sess-1-lora-B::forward_backward::1'}, ( f'adapters collided on one dedup key: {claimed}') diff --git a/tests/server/lifecycle/test_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..d49238fca 100644 --- a/tests/server/model/test_full_mode_exclusivity.py +++ b/tests/server/model/test_full_mode_exclusivity.py @@ -5,120 +5,102 @@ 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.test_controller import request, setup +from twinkle.server.config.training import TrainingConfig +from twinkle.server.training.catalog import TrainingCatalog +from twinkle.server.training.router import TrainingRouter -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'][0]['modes']['full']['max_replicas'] = 1 + controller.config = TrainingConfig.model_validate(raw) + controller.catalog = TrainingCatalog(controller.config) + runtime.start_gate.clear() + return controller, state, runtime @pytest.mark.asyncio async def test_concurrent_full_mode_claims_have_exactly_one_winner(): - 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) == '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/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': '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/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.catalog = controller.catalog + router.model = SimpleNamespace(name='small') + router.controller = SimpleNamespace(create=SimpleNamespace(remote=AsyncMock(side_effect=controller.create))) + + async def create_remote(payload): + from twinkle.server.state.training_records import CreateTraining + return await controller.create(CreateTraining.model_validate(payload)) + router.controller.create.remote.side_effect = create_remote + try: + refs = await asyncio.gather(*(router.native_request(_request({ + 'client_create_id': name, 'adapter_name': name, + }), 'add_adapter_to_model') for name in ('adapter-a', 'adapter-b'))) + assert len(controller.instances) == 1 + assert len(next(iter(controller.instances.values())).training_ids) == 1 + runtime.start_gate.set() + await controller.wait_for_jobs() + records = [await state.get_future(ref.request_id) for ref in refs] + assert sum(record['status'] == 'completed' for record in records) == 1 + assert sum(record['status'] == 'failed' and record['failure']['reason_code'] == 'capacity_unavailable' + for record in records) == 1 + assert await state._backend.get('token_count::owner') == 1 + finally: + runtime.start_gate.set() + await controller.wait_for_jobs() + await controller.shutdown() diff --git a/tests/server/model/test_replica_lifecycle.py b/tests/server/model/test_replica_lifecycle.py index 84b961391..0b1160085 100644 --- a/tests/server/model/test_replica_lifecycle.py +++ b/tests/server/model/test_replica_lifecycle.py @@ -34,6 +34,7 @@ def _make_lifecycle_manager(state: _CapacityState, replica_id: str, max_loras: i manager.state = state manager.replica_id = replica_id manager.max_loras = max_loras + manager.train_mode = 'lora' manager._replica_registered = False manager.data_plane = SimpleNamespace(close=AsyncMock()) manager.shutdown_task_queue = AsyncMock() @@ -54,24 +55,27 @@ async def test_replica_lifecycle_updates_shared_capacity() -> None: assert await state.get_capacity_info() == {'max_loras': 6, 'used_loras': 0, 'free_loras': 6} await second.shutdown() + # Only the owner retires capacity after confirming the GPU workers stopped. + assert await state.get_capacity_info() == {'max_loras': 6, 'used_loras': 0, 'free_loras': 6} + await state.unregister_replica(second.replica_id) assert await state.get_capacity_info() == {'max_loras': 3, 'used_loras': 0, 'free_loras': 3} @pytest.mark.asyncio async def test_async_constructor_registers_replica_before_ready() -> None: state = SimpleNamespace(register_replica=AsyncMock(), touch_replica_last_seen=AsyncMock()) - replica_context = SimpleNamespace(replica_id=SimpleNamespace(unique_id='replica-1')) manager = ModelManagement.__new__(ModelManagement) manager._task_queue_config = SimpleNamespace(effective_execution_timeout=1800.0) with patch('twinkle.server.model.app.DeviceGroup', return_value=SimpleNamespace(name='group')), \ patch('twinkle.server.model.app.init_twinkle_runtime', return_value=None), \ - patch('twinkle.server.model.app.serve.get_replica_context', return_value=replica_context), \ patch.object(MODEL_SELECTOR, 'construct', return_value=MagicMock()), \ patch('twinkle.server.model.app.get_server_state', return_value=state), \ patch('twinkle.server.data_plane.DataPlaneProxy', return_value=MagicMock()), \ patch.object(ModelManagement, '_init_task_queue'), \ - patch.object(ModelManagement, '_init_adapter_manager'): + patch.object(ModelManagement, '_init_adapter_manager'), \ + patch.object(ModelManagement, 'initialize_managed_runtime', new_callable=AsyncMock), \ + patch.object(ModelManagement, 'wait_for_training_workers', new_callable=AsyncMock): await ModelManagement.__init__( manager, model_id='model', @@ -80,6 +84,9 @@ async def test_async_constructor_registers_replica_before_ready() -> None: device_mesh={}, backend='mock', max_loras=3, + training_context={'instance_id': 'replica-1'}, + placement_group=None, + allocation_devices=[], ) state.register_replica.assert_awaited_once_with('replica-1', 3) diff --git a/tests/server/model/test_tinker_handlers.py b/tests/server/model/test_tinker_handlers.py index 6daa6d566..9d116ed51 100644 --- a/tests/server/model/test_tinker_handlers.py +++ b/tests/server/model/test_tinker_handlers.py @@ -9,7 +9,7 @@ from twinkle.server.model.tinker_handlers import _register_model_tinker_routes from twinkle.server.deployment import twinkle_server_error_handler from twinkle.server.exceptions import BatchSizeError, TwinkleServerError -from twinkle.server.task_queue.config import TaskQueueConfig +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.task_queue.mixin import TaskQueueMixin diff --git a/tests/server/sampler/test_sampling_concurrency.py b/tests/server/sampler/test_sampling_concurrency.py index f1697ec9c..88a7e6cb8 100644 --- a/tests/server/sampler/test_sampling_concurrency.py +++ b/tests/server/sampler/test_sampling_concurrency.py @@ -21,7 +21,7 @@ from twinkle.server.sampler.tinker_handlers import _register_tinker_sampler_routes from twinkle.server.sampler.twinkle_handlers import _register_twinkle_sampler_routes from twinkle.server.task_queue.backend_gate import BackendGate -from twinkle.server.task_queue.config import TaskQueueConfig +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.task_queue.types import BackendBusyError from twinkle.server.task_queue.worker import ComputeWorker @@ -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/start_e2e_server.py b/tests/server/start_e2e_server.py index 4af2e2b85..1e98c914a 100644 --- a/tests/server/start_e2e_server.py +++ b/tests/server/start_e2e_server.py @@ -38,6 +38,8 @@ MEGATRON_BASE = "tests/server/config/server_config_4b_e2e_megatron.yaml" _FULL = { "train_mode": "full", + "idle_timeout_s": 3600, + "max_replicas": 1, # Full-parameter training holds the deployment exclusively; use long # timeouts so the tenant is not expired mid-test. "adapter_config.adapter_timeout": 3600, diff --git a/tests/server/state/fake_backend.py b/tests/server/state/fake_backend.py index eaac98833..10f5c9e56 100644 --- a/tests/server/state/fake_backend.py +++ b/tests/server/state/fake_backend.py @@ -67,13 +67,15 @@ async def update_atomic( key: str, transform: Callable[[Any | None], Any | None], ttl: int | None = None, + *, + return_changed: bool = False, ) -> Any | None: current = await self.get(key) updated = transform(current) if updated is None: - return current + return False if return_changed else current await self.set(key, updated, ttl) - return updated + return True if return_changed else updated async def flush_all(self) -> None: self._store.clear() diff --git a/tests/server/state/test_future_lifecycle.py b/tests/server/state/test_future_lifecycle.py index 9728780f3..d9c0e63ff 100644 --- a/tests/server/state/test_future_lifecycle.py +++ b/tests/server/state/test_future_lifecycle.py @@ -7,86 +7,20 @@ """ from __future__ import annotations +import asyncio import time -from collections.abc import Callable -from fnmatch import fnmatch -from typing import Any from unittest import mock import pytest -from twinkle.server.state.backend.base import StateBackend +from tests.server.state.fake_backend import FakeBackend from twinkle.server.state.future_manager import FutureManager from twinkle.server.state.models import FutureFailureRecord -class _FakeBackend(StateBackend): - """Dict-backed StateBackend. Single event loop, so ``update_atomic`` is atomic - simply by not awaiting between read and write -- the real backends' guarantee.""" - - def __init__(self) -> None: - self._store: dict[str, tuple[Any, float | None]] = {} - - def _is_expired(self, key: str) -> bool: - entry = self._store.get(key) - if entry is None: - return True - _, expire_at = entry - if expire_at is not None and time.time() >= expire_at: - del self._store[key] - return True - return False - - async def set(self, key: str, value: Any, ttl: int | None = None) -> None: - self._store[key] = (value, (time.time() + ttl) if ttl is not None else None) - - async def get(self, key: str) -> Any | None: - return None if self._is_expired(key) else self._store[key][0] - - async def mget(self, keys: list[str]) -> list[Any | None]: - return [await self.get(key) for key in keys] - - async def delete(self, key: str) -> None: - self._store.pop(key, None) - - async def exists(self, key: str) -> bool: - return not self._is_expired(key) - - async def keys(self, pattern: str) -> list[str]: - return [k for k in list(self._store) if not self._is_expired(k) and fnmatch(k, pattern)] - - async def count(self, pattern: str) -> int: - return len(await self.keys(pattern)) - - async def set_nx(self, key: str, value: Any, ttl: int | None = None) -> bool: - if not self._is_expired(key): - return False - await self.set(key, value, ttl) - return True - - async def update_atomic( - self, - key: str, - transform: Callable[[Any | None], Any | None], - ttl: int | None = None, - ) -> Any | None: - current = await self.get(key) - updated = transform(current) - if updated is None: - return current - await self.set(key, updated, ttl) - return updated - - async def close(self) -> None: - pass - - async def health_check(self) -> bool: - return True - - @pytest.fixture def manager(): - return FutureManager(_FakeBackend(), expiration_timeout=300.0) + return FutureManager(FakeBackend(), expiration_timeout=300.0) async def _store(manager, request_id, status, *, replica_id=None, absolute_deadline=None): @@ -203,7 +137,7 @@ async def test_stored_timestamps_align_with_wall_clock_regardless_of_host_tz(man @pytest.mark.asyncio async def test_claim_seq_dedups_then_release_readmits(): from twinkle.server.state.server_state import ServerState - state = ServerState(backend=_FakeBackend()) + state = ServerState(backend=FakeBackend()) # First claim of a (session, seq_id) is unseen -> None, caller proceeds to enqueue. assert await state.claim_seq('seq::s1::1', 'reqA', ttl=60) is None # A duplicate claim returns the original request_id -> caller returns its envelope. @@ -218,7 +152,7 @@ async def test_claim_seq_dedups_then_release_readmits(): @pytest.mark.asyncio async def test_cancel_drops_pending_but_never_running(): from twinkle.server.state.server_state import ServerState - state = ServerState(backend=_FakeBackend()) + state = ServerState(backend=FakeBackend()) # pending -> cancel drops it to a terminal domain failure. await state.store_future_status('rp', 'pending', 'm1') assert await state.cancel_future('rp') == {'cancelled': True, 'state': 'cancelled'} @@ -247,3 +181,37 @@ def test_cancelled_record_maps_to_error_envelope(): env = envelope_from_record('rc', rec) assert env.status == 'cancelled' assert env.error is not None and env.error.error_code == 499 + + +@pytest.mark.asyncio +async def test_late_queued_and_repeated_running_cannot_regress_or_renew(manager): + await _store(manager, 'race', 'pending') + assert await manager.try_start('race') + before = (await manager.get('race')).model_dump() + await _store(manager, 'race', 'queued') + await _store(manager, 'race', 'running') + assert (await manager.get('race')).model_dump() == before + assert await manager.cancel_if_pending('race') == 'running' + assert not await manager.try_start('race') + assert not await manager.try_start('missing') + + +@pytest.mark.asyncio +@pytest.mark.parametrize('cancel_first', [False, True]) +async def test_start_and_cancel_have_one_atomic_winner(manager, cancel_first): + await _store(manager, 'race', 'queued') + operations = [manager.try_start('race'), manager.cancel_if_pending('race')] + if cancel_first: + operations.reverse() + results = await asyncio.gather(*operations) + started = results[1] if cancel_first else results[0] + status = (await manager.get('race')).status + assert started == (status == 'running') + assert status in ('running', 'cancelled') + assert not await manager.try_start('race') + + +@pytest.mark.asyncio +async def test_concurrent_starters_only_one_obtains_execution(manager): + await _store(manager, 'race', 'pending') + assert sum(await asyncio.gather(*(manager.try_start('race') for _ in range(20)))) == 1 diff --git a/tests/server/state/test_leader_election.py b/tests/server/state/test_leader_election.py index 792989c30..216105e99 100644 --- a/tests/server/state/test_leader_election.py +++ b/tests/server/state/test_leader_election.py @@ -11,10 +11,14 @@ import asyncio import pytest +from types import SimpleNamespace +from unittest.mock import AsyncMock + +from tests.server.state.fake_backend import FakeBackend from twinkle.server.state import ServerState from twinkle.server.state.backend.memory_backend import RayActorBackend -from twinkle.server.state.cleanup_coordinator import LEADER_KEY, LEASE_RENEW +from twinkle.server.state.cleanup_coordinator import LEADER_KEY, LEASE_RENEW, ResourceCleanupCoordinator from twinkle.server.telemetry import MetricsRegistry @@ -100,13 +104,7 @@ async def test_renew_keeps_leader() -> None: @pytest.mark.asyncio async def test_leader_recovers_after_renewal_failure() -> None: - """Regression: when a leader's renewal raises, it releases - the lease best-effort so the very next election tick re-acquires leadership - without waiting LEASE_TTL. - - On unfixed code ``_is_leader`` flips to False while the lease value lingers - in the backend, so ``set_nx`` keeps returning False for up to LEASE_TTL. - """ + """An unconfirmed renewal fences cleanup until the original lease expires.""" backend = RayActorBackend() state = ServerState(backend=backend, cleanup_interval=600.0) @@ -128,11 +126,13 @@ async def flaky_update_atomic(*args, **kwargs): from unittest import mock with mock.patch.object(backend, 'update_atomic', side_effect=flaky_update_atomic): await state._try_acquire_or_renew() - # Renewal failed → no longer leader, and the stale lease was released. assert state._is_leader is False - # The lease key was deleted best-effort, so the next tick re-acquires - # immediately (no LEASE_TTL wait, no lingering self-owned lease). + assert await backend.get(LEADER_KEY) == state._leader_id + await state._try_acquire_or_renew() + assert state._is_leader is False + await backend.set(LEADER_KEY, state._leader_id, ttl=1) + await asyncio.sleep(1.1) await state._try_acquire_or_renew() assert state._is_leader is True @@ -159,6 +159,75 @@ async def test_renewal_failure_does_not_steal_other_leader_lease() -> None: assert await backend.get(LEADER_KEY) == leader_value, 'follower stole the real leader lease' +def coordinator(backend=None): + models = SimpleNamespace(rebuild_indexes=AsyncMock()) + return ResourceCleanupCoordinator(backend or FakeBackend(), {'models': models}, + cleanup_interval=600, expiration_timeout=300, + sweep_processor_quotas=AsyncMock()) + + +@pytest.mark.asyncio +async def test_stale_leader_renewal_failure_preserves_successor(monkeypatch): + backend = FakeBackend() + cleanup = coordinator(backend) + await cleanup._try_acquire_or_renew() + assert cleanup.is_leader + + async def successor_then_fail(*args, **kwargs): + await backend.set(LEADER_KEY, 'successor', ttl=30) + raise ConnectionError('renewal acknowledgement lost') + + monkeypatch.setattr(backend, 'update_atomic', successor_then_fail) + await cleanup._try_acquire_or_renew() + assert not cleanup.is_leader + assert await backend.get(LEADER_KEY) == 'successor' + + +@pytest.mark.asyncio +async def test_cleanup_concurrent_start_rebuilds_once_and_stop_joins(): + cleanup = coordinator() + entered, proceed = asyncio.Event(), asyncio.Event() + + async def rebuild(): + entered.set() + await proceed.wait() + + cleanup._managers['models'].rebuild_indexes.side_effect = rebuild + first = asyncio.create_task(cleanup.start()) + await entered.wait() + second = asyncio.create_task(cleanup.start()) + proceed.set() + assert await asyncio.gather(first, second) == [True, False] + tasks = (cleanup._cleanup_task, cleanup._leader_task) + assert await cleanup.stop() + assert all(task.done() for task in tasks) + assert not await cleanup.stop() + cleanup._managers['models'].rebuild_indexes.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_failed_cleanup_start_can_retry(): + cleanup = coordinator() + cleanup._managers['models'].rebuild_indexes.side_effect = [RuntimeError('offline'), None] + with pytest.raises(RuntimeError, match='offline'): + await cleanup.start() + assert cleanup._cleanup_task is None and cleanup._leader_task is None + assert await cleanup.start() + await cleanup.stop() + + +@pytest.mark.asyncio +async def test_lazy_cleanup_retries_after_failed_start(): + from twinkle.server.deployment import LazyCleanupMixin + client = LazyCleanupMixin() + client.state = SimpleNamespace(start_cleanup_task=AsyncMock(side_effect=[RuntimeError('offline'), True])) + await client._ensure_state_cleanup_started() + assert not client._state_cleanup_started + await client._ensure_state_cleanup_started() + assert client._state_cleanup_started + assert client.state.start_cleanup_task.await_count == 2 + + # ============================================================ # Metrics ObservableGauge (merged from test_metrics_observable_gauge) # ============================================================ diff --git a/tests/server/state/test_update_atomic.py b/tests/server/state/test_update_atomic.py index d03b3a79e..387f69302 100644 --- a/tests/server/state/test_update_atomic.py +++ b/tests/server/state/test_update_atomic.py @@ -129,6 +129,43 @@ async def test_key_prefix_is_hidden_from_logical_keys(backend_pair) -> None: assert await first.keys('session::*') == ['session::one'] +@pytest.mark.asyncio +async def test_atomic_update_reports_this_call_not_the_existing_value(backend_pair): + first, second = backend_pair + assert await first.update_atomic('claim', functools.partial(_replace_with, value=True), + return_changed=True) is True + assert await second.update_atomic('claim', _no_op, return_changed=True) is False + assert await second.get('claim') is True + assert await second.update_atomic('missing', _no_op, return_changed=True) is False + + +@pytest.mark.asyncio +async def test_future_execution_is_claimed_once_across_backend_handles(backend_pair): + from twinkle.server.state.future_manager import FutureManager + first, second = [FutureManager(backend, 300) for backend in backend_pair] + await first.store_status('start', 'pending', 'm') + starts = await asyncio.gather(*(manager.try_start('start') for manager in [first, second] * 8)) + assert sum(starts) == 1 + await second.store_status('start', 'queued', 'm') + assert await first.cancel_if_pending('start') == 'running' + assert (await first.get('start')).status == 'running' + + +@pytest.mark.asyncio +@pytest.mark.parametrize('cancel_first', [False, True]) +async def test_future_start_cancel_race_across_backend_handles(backend_pair, cancel_first): + from twinkle.server.state.future_manager import FutureManager + first, second = [FutureManager(backend, 300) for backend in backend_pair] + await first.store_status('start', 'queued', 'm') + calls = [first.try_start('start'), second.cancel_if_pending('start')] + if cancel_first: + calls.reverse() + results = await asyncio.gather(*calls) + started = results[1] if cancel_first else results[0] + assert started == ((await first.get('start')).status == 'running') + assert not await second.try_start('start') + + # ---------- Memory backend ----------------------------------------------- # diff --git a/tests/server/telemetry/test_training_metrics.py b/tests/server/telemetry/test_training_metrics.py new file mode 100644 index 000000000..ad446cfeb --- /dev/null +++ b/tests/server/telemetry/test_training_metrics.py @@ -0,0 +1,152 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Managed training metrics stay low-cardinality and survive registry resets.""" +from collections import defaultdict +from types import SimpleNamespace + +import pytest + +from twinkle.server.telemetry import metrics +from twinkle.server.training.controller import TrainingController + +_ALLOWED_ATTRIBUTES = {'backend', 'mode', 'state', 'status', 'reason', 'device_type', 'resident'} + + +@pytest.fixture(autouse=True) +def reset_registry(): + yield + metrics.MetricsRegistry.reset() + + +class FakeInstrument: + + def __init__(self, name): + self.name = name + self.calls = [] + + def add(self, value, attributes=None): + self.calls.append((value, attributes or {})) + + def record(self, value, attributes=None): + self.calls.append((value, attributes or {})) + + +class FakeMeter: + + def __init__(self): + self.instruments = {} + + def _create(self, name, **kwargs): + instrument = FakeInstrument(name) + self.instruments[name] = instrument + return instrument + + create_counter = _create + create_histogram = _create + create_up_down_counter = _create + + def create_observable_gauge(self, name, **kwargs): + return self._create(name, **kwargs) + + +def build_registry(monkeypatch): + meter = FakeMeter() + monkeypatch.setattr(metrics, 'get_meter', lambda name: meter) + metrics.MetricsRegistry.reset() + return metrics.MetricsRegistry.get(), meter + + +def test_training_and_redis_instruments_are_registered(monkeypatch): + registry, meter = build_registry(monkeypatch) + training = registry.training_metrics() + state = registry.state_backend_metrics() + + training.instances_active.set(2, tags={'backend': 'transformers', 'mode': 'lora', 'state': 'READY', + 'resident': 'false'}) + training.startup_seconds.observe(1.5, tags={'backend': 'transformers', 'mode': 'lora', 'status': 'success'}) + training.quarantined_total.inc(tags={'backend': 'megatron', 'mode': 'full', 'reason': 'timeout'}) + state.redis_health.set(1) + state.update_atomic_conflicts.inc() + state.update_atomic_retries.observe(2) + + expected = { + 'twinkle.training.instances.active', + 'twinkle.training.instance.startup.duration_seconds', + 'twinkle.training.instance.shutdown.duration_seconds', + 'twinkle.training.placement_group.allocation.duration_seconds', + 'twinkle.training.placement_group.removal.duration_seconds', + 'twinkle.training.quarantined.total', + 'twinkle.training.sweep.total', + 'twinkle.training.sweep.duration_seconds', + 'twinkle.training.devices.allocated', + 'twinkle.training.placement_groups.active', + 'twinkle.state.redis.health', + 'twinkle.state.redis.update_atomic.conflicts.total', + 'twinkle.state.redis.update_atomic.retries', + } + assert expected <= meter.instruments.keys() + for instrument in meter.instruments.values(): + for _, attributes in instrument.calls: + assert set(attributes) <= _ALLOWED_ATTRIBUTES + assert not {'token', 'session_id', 'training_id', 'instance_id', 'request_id'} & set(attributes) + + +def test_training_metric_adapters_rebind_after_registry_reset(monkeypatch): + first, _ = build_registry(monkeypatch) + first_metrics = first.training_metrics() + metrics.MetricsRegistry.reset() + second_metrics = metrics.MetricsRegistry.get().training_metrics() + assert first_metrics is not second_metrics + assert first_metrics.startup_seconds._instrument is not second_metrics.startup_seconds._instrument + + +@pytest.mark.parametrize('backend', ['transformers', 'megatron']) +def test_controller_metrics_clear_old_states_and_released_resources(monkeypatch, backend): + registry, meter = build_registry(monkeypatch) + spec = SimpleNamespace(model=SimpleNamespace(backend=backend, device_type='cuda')) + controller = TrainingController.__new__(TrainingController) + controller.catalog = SimpleNamespace(get=lambda base_model, mode: spec) + controller._metrics = registry.training_metrics() + controller._reported_metric_keys = set() + controller._reported_resource_metric_keys = set() + controller.instances = { + str(index): SimpleNamespace( + instance_id=f'instance_{index}', base_model='private-model', mode='lora', + state='READY', resident=False, devices=[index * 2, index * 2 + 1], + application_name=f'private-app-{index}', training_ids=[f'training_{index}']) + for index in range(2) + } + + def totals(name): + values = defaultdict(float) + for delta, attributes in meter.instruments[name].calls: + assert set(attributes) <= _ALLOWED_ATTRIBUTES + assert not {'instance_id', 'base_model', 'application_name', 'training_id'} & set(attributes) + values[tuple(sorted(attributes.items()))] += delta + return values + + instances_name = 'twinkle.training.instances.active' + devices_name = 'twinkle.training.devices.allocated' + groups_name = 'twinkle.training.placement_groups.active' + controller._report_instance_metrics() + first_calls = {name: len(meter.instruments[name].calls) + for name in (instances_name, devices_name, groups_name)} + controller._report_instance_metrics() + assert first_calls == {name: len(meter.instruments[name].calls) for name in first_calls} + assert sum(totals(instances_name).values()) == 2 + assert sum(totals(devices_name).values()) == 4 + assert sum(totals(groups_name).values()) == 2 + + controller.instances['0'].state = 'DRAINING' + controller._report_instance_metrics() + state_totals = {dict(key)['state']: value for key, value in totals(instances_name).items()} + assert state_totals == {'READY': 1, 'DRAINING': 1} + assert sum(totals(devices_name).values()) == 4 + assert sum(totals(groups_name).values()) == 2 + + for instance in controller.instances.values(): + instance.state = 'STOPPED' + controller._report_instance_metrics() + for name in (instances_name, devices_name, groups_name): + assert all(value == 0 for value in totals(name).values()) + assert not controller._reported_metric_keys + assert not controller._reported_resource_metric_keys diff --git a/tests/server/test_dependencies.py b/tests/server/test_dependencies.py new file mode 100644 index 000000000..8853f9031 --- /dev/null +++ b/tests/server/test_dependencies.py @@ -0,0 +1,99 @@ +# 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', 'resources', 'runtime', 'task_errors', 'utils', 'validation'}), + frozenset({'telemetry'}), + frozenset({'env_propagation', 'service_proxy', 'state'}), + frozenset({'deployment', 'lifecycle', 'session_resource'}), + frozenset({'task_queue'}), + frozenset({'checkpoint', 'data_plane'}), + frozenset({'gateway', 'model', 'processor'}), + frozenset({'sampler', 'training'}), + frozenset({'launcher'}), + frozenset({'cli'}), + frozenset({'__main__', _ROOT_COMPONENT}), +) + + +def _component(module: str) -> str: + if module == _ROOT: + return _ROOT_COMPONENT + return module.removeprefix(f'{_ROOT}.').split('.', 1)[0] + + +@pytest.mark.parametrize('statement', [ + 'import twinkle.server.exceptions', + 'from twinkle.server.config import ServerConfig', + 'from twinkle.server.state.models import FutureRecord', + 'from twinkle.server.state import SessionRecord', +]) +def test_lightweight_imports_do_not_load_runtime(statement): + code = statement + ''' +import sys +assert not any(name in sys.modules for name in ( + 'torch', 'ray', 'twinkle.server.launcher', 'twinkle.server.training.actor', 'twinkle.hub.hub' +)), sorted(sys.modules) +''' + result = subprocess.run([sys.executable, '-c', code], capture_output=True, text=True, timeout=60) + assert result.returncode == 0, result.stderr + + +def test_hub_first_import_does_not_reenter_plugin_loader(): + code = ''' +from twinkle.hub.hub import HFHub +from twinkle.utils.loader import Plugin +assert HFHub.__module__ == 'twinkle.hub.hub' +assert Plugin.__module__ == 'twinkle.utils.loader' +''' + result = subprocess.run([sys.executable, '-c', code], capture_output=True, text=True, timeout=120) + assert result.returncode == 0, result.stderr + + +def test_public_launcher_is_the_actual_implementation(): + import twinkle.server + from twinkle.server.launcher.server_launcher import ServerLauncher, launch_server + assert twinkle.server.ServerLauncher is ServerLauncher + assert twinkle.server.launch_server is launch_server + assert twinkle.server.__all__ == ['ServerLauncher', 'launch_server'] + + +def test_server_components_form_declared_dag(): + """Reject cycles, upward imports, same-layer coupling, and unclassified components.""" + graph = grimp.build_graph('twinkle', include_external_packages=False) + server_modules = { + module for module in graph.modules if module == _ROOT or module.startswith(f'{_ROOT}.') + } + observed_components = {_component(module) for module in server_modules} + declared_components = set().union(*_LAYERS) + assert observed_components == declared_components, ( + f'Update the server DAG when adding or removing components: ' + f'undeclared={sorted(observed_components - declared_components)}, ' + f'missing={sorted(declared_components - observed_components)}') + + layer_by_component = { + component: layer_index for layer_index, layer in enumerate(_LAYERS) for component in layer + } + component_edges = { + (_component(importer), _component(imported)) + for importer in server_modules + for imported in graph.find_modules_directly_imported_by(importer) + if imported in server_modules and _component(importer) != _component(imported) + } + violations = sorted( + (importer, imported) for importer, imported in component_edges + if layer_by_component[importer] <= layer_by_component[imported]) + assert not violations, ( + 'Server dependencies must point strictly down the declared DAG; ' + f'invalid edges={violations}') diff --git a/tests/server/test_gateway_services.py b/tests/server/test_gateway_services.py index 5a5d032d8..e61dfd631 100644 --- a/tests/server/test_gateway_services.py +++ b/tests/server/test_gateway_services.py @@ -10,6 +10,9 @@ class _State: + async def assert_future_owner(self, request_id, token): + return None + def __init__(self, records): self.records = list(records) @@ -43,7 +46,7 @@ def test_gateway_services_do_not_import_protocol_models(): class _FakeGateway: """Minimal ``GatewayServer`` stand-in; ``poll_future`` is patched so state is unused.""" - state = None + state = _State([]) def _parity_client(monkeypatch, canonical_record): diff --git a/tests/server/training/test_acceptance_harness.py b/tests/server/training/test_acceptance_harness.py new file mode 100644 index 000000000..92b8b972c --- /dev/null +++ b/tests/server/training/test_acceptance_harness.py @@ -0,0 +1,84 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Acceptance harness budgets and Redis ownership are independently testable.""" +import shutil + +import pytest + +from tests.server.integration.on_demand_training import stress_validation +from tests.server.integration.on_demand_training.managed_redis import ManagedRedis + + +@pytest.fixture(scope='session', autouse=True) +def _ray_runtime(): + # These tests deliberately use no Ray runtime. + yield + + +class FakeHandle: + calls = [] + + def __init__(self, base, model, protocol, label): + self.calls.append(('create', protocol)) + + def step(self): + self.calls.append(('step', None)) + + def save(self, name): + return 'checkpoint' + + def resume(self, path): + assert path == 'checkpoint' + + def close(self): + self.calls.append(('close', None)) + + +def test_iteration_budget_does_not_round_up(monkeypatch): + FakeHandle.calls = [] + monkeypatch.setattr(stress_validation, 'TrainingHandle', FakeHandle) + evidence = stress_validation.run_stress('unused', 'unused', concurrency=2, iterations=3) + assert evidence['completed'] == 3 + assert len([call for call in FakeHandle.calls if call[0] == 'close']) == 4 # Includes the warm peer. + + +def test_duration_is_not_truncated_by_iterations(monkeypatch): + now = [0] + + def clock(): + now[0] += 1 + return now[0] + + monkeypatch.setattr(stress_validation.time, 'monotonic', clock) + monkeypatch.setattr(stress_validation, 'TrainingHandle', FakeHandle) + evidence = stress_validation.run_stress('unused', 'unused', concurrency=1, iterations=1, duration=10) + assert evidence['completed'] > 1 and evidence['seconds'] >= 10 + + +def test_failed_training_never_reports_success(monkeypatch): + class Broken(FakeHandle): + def step(self): + raise RuntimeError('injected failure') + + monkeypatch.setattr(stress_validation, 'TrainingHandle', Broken) + with pytest.raises(RuntimeError, match='injected failure'): + stress_validation.run_stress('unused', 'unused', iterations=1) + + +@pytest.mark.skipif(not shutil.which('redis-server'), reason='redis-server executable required') +def test_owned_redis_restart_preserves_state(tmp_path): + redis = pytest.importorskip('redis') + managed = ManagedRedis(tmp_path) + managed.start() + client = redis.Redis.from_url(managed.url, socket_timeout=1) + try: + client.set('checkpoint', 'durable') + + def unavailable(): + with pytest.raises(redis.exceptions.ConnectionError): + client.get('checkpoint') + + managed.fault(unavailable) + assert client.get('checkpoint') == b'durable' + finally: + client.close() + managed.stop() diff --git a/tests/server/training/test_actor_initialization.py b/tests/server/training/test_actor_initialization.py new file mode 100644 index 000000000..10c748341 --- /dev/null +++ b/tests/server/training/test_actor_initialization.py @@ -0,0 +1,127 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Initialization failures and backoff during automatic controller restarts.""" +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + +from twinkle.server.training.actor import TrainingControllerActor + + +@pytest.mark.asyncio +async def test_failed_initialization_restarts_only_after_backoff(monkeypatch): + cls = TrainingControllerActor.__ray_metadata__.modified_class + gate, sleeping = asyncio.Event(), asyncio.Event() + delays = [] + handle = object() + kill = Mock() + get_actor = Mock(return_value=handle) + original_sleep = asyncio.sleep + + async def sleep(delay): + delays.append(delay) + sleeping.set() + await gate.wait() + + monkeypatch.setattr('twinkle.server.training.actor.asyncio.sleep', sleep) + monkeypatch.setattr('twinkle.server.training.actor.ray.get_actor', get_actor) + monkeypatch.setattr('twinkle.server.training.actor.ray.kill', kill) + actor = SimpleNamespace(core=SimpleNamespace(initialize=AsyncMock(side_effect=RuntimeError('allocation failed'))), + config=SimpleNamespace(sweep_interval_s=7), namespace='initialization-test', _monitors=[]) + actor._restart_after_failure = lambda: cls._restart_after_failure(actor) + with pytest.raises(RuntimeError, match='allocation failed'): + await cls._initialize(actor) + await sleeping.wait() + assert delays == [7] and not actor._monitors + kill.assert_not_called() + gate.set() + await actor._restart_task + await original_sleep(0) + get_actor.assert_called_once_with('training-controller-initialization-test', namespace='initialization-test') + kill.assert_called_once_with(handle, no_restart=False) + + +@pytest.mark.asyncio +async def test_shutdown_cancels_pending_restart(): + cls = TrainingControllerActor.__ray_metadata__.modified_class + wait = asyncio.Event() + restart = asyncio.create_task(wait.wait()) + initialization = asyncio.get_running_loop().create_future() + initialization.set_result(None) + actor = SimpleNamespace(driver=SimpleNamespace(wait_for_cleanup=AsyncMock()), + core=SimpleNamespace(shutdown=AsyncMock(), wait_for_jobs=AsyncMock()), + _monitors=[], _restart_task=restart, _initialization=initialization, _failures={}) + await cls.shutdown(actor) + assert restart.cancelled() + actor.core.shutdown.assert_awaited_once() + actor.driver.wait_for_cleanup.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_heartbeat_keeps_healthy_peer_alive_while_probe_and_close_block(): + cls = TrainingControllerActor.__ray_metadata__.modified_class + good, bad = [SimpleNamespace(instance_id=name, state='READY') for name in ('good', 'bad')] + probe_gate, close_gate = asyncio.Event(), asyncio.Event() + healthy_seen, close_started = asyncio.Event(), asyncio.Event() + + async def health(instance): + if instance is bad: + await probe_gate.wait() + return {'unhealthy': instance is bad} + + async def fail(*args): + close_started.set() + await close_gate.wait() + + async def touch(identity): + if identity == 'good': + healthy_seen.set() + + actor = SimpleNamespace( + store=SimpleNamespace(backend=SimpleNamespace(get=AsyncMock(return_value='active')), + key=lambda *args: 'active', state=SimpleNamespace(touch_replica_last_seen=touch)), + core=SimpleNamespace(epoch='controller', instances={'bad': bad, 'good': good}, fail_instance=fail), + driver=SimpleNamespace(health=health), _unreachable_since={}, _failures={}) + actor._check_instance_health = lambda instance: cls._check_instance_health(actor, instance) + actor._fail_unhealthy = lambda instance, message: cls._fail_unhealthy(actor, instance, message) + monitor = asyncio.create_task(cls._heartbeat(actor)) + try: + await asyncio.wait_for(healthy_seen.wait(), 1) + assert not probe_gate.is_set() + probe_gate.set() + await asyncio.wait_for(close_started.wait(), 1) + # Closing a bad peer does not occupy the heartbeat monitor. + healthy_seen.clear() + await cls._check_instance_health(actor, good) + assert healthy_seen.is_set() + cls._fail_unhealthy(actor, bad, 'duplicate') + assert len(actor._failures) == 1 + finally: + probe_gate.set() + close_gate.set() + monitor.cancel() + await asyncio.gather(monitor, *actor._failures.values(), return_exceptions=True) + assert not actor._failures + + +@pytest.mark.asyncio +async def test_shutdown_cancels_and_joins_initialization_before_closing_core(): + cls = TrainingControllerActor.__ray_metadata__.modified_class + started, exited = asyncio.Event(), asyncio.Event() + + async def initialize(): + started.set() + try: + await asyncio.Event().wait() + finally: + exited.set() + + initialization = asyncio.create_task(initialize()) + await started.wait() + actor = SimpleNamespace(driver=SimpleNamespace(wait_for_cleanup=AsyncMock()), + core=SimpleNamespace(shutdown=AsyncMock(), wait_for_jobs=AsyncMock()), + _initialization=initialization, _monitors=[], _failures={}) + await cls.shutdown(actor) + assert initialization.cancelled() and exited.is_set() + actor.core.shutdown.assert_awaited_once() diff --git a/tests/server/training/test_allocation.py b/tests/server/training/test_allocation.py new file mode 100644 index 000000000..519347b8c --- /dev/null +++ b/tests/server/training/test_allocation.py @@ -0,0 +1,23 @@ +import json +import subprocess +import sys +from pathlib import Path + + +def test_production_allocator_uses_disjoint_devices_and_reuses_released_group(): + result = subprocess.run([sys.executable, str(Path(__file__).resolve().parents[3] / 'tests/server/integration/on_demand_training/allocation_probe.py')], + capture_output=True, text=True, timeout=120) + assert result.returncode == 0, result.stdout + result.stderr + evidence = json.loads(result.stdout.strip().splitlines()[-1]) + assert evidence['passed'] and len(set(evidence['devices'])) == 8 + assert len(evidence['reused']) == 2 + + +def test_fragmented_devices_cannot_be_combined_across_nodes(): + result = subprocess.run([sys.executable, str(Path(__file__).resolve().parents[3] / 'tests/server/integration/on_demand_training/fragmentation_probe.py')], + capture_output=True, text=True, timeout=120) + assert result.returncode == 0, result.stdout + result.stderr + evidence = json.loads(result.stdout.strip().splitlines()[-1]) + assert evidence['passed'] and evidence['nodes'] == 2 + assert evidence['fragmented_rejected'] and evidence['pending_removed'] + assert len(evidence['devices']) == 2 diff --git a/tests/server/training/test_controller.py b/tests/server/training/test_controller.py new file mode 100644 index 000000000..358c88e15 --- /dev/null +++ b/tests/server/training/test_controller.py @@ -0,0 +1,679 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Exercise real controller/storage decisions against a controllable runtime.""" +from __future__ import annotations + +import asyncio + +import pytest + +from twinkle.server.config.training import TrainingConfig +from twinkle.server.exceptions import RequestRejectedError, ResourceNotFoundError +from twinkle.server.state import ServerState +from twinkle.server.state.backend.memory_backend import RayActorBackend +from twinkle.server.training.controller import TrainingController +from twinkle.server.training.errors import TrainingExpiredError +from twinkle.server.resources.allocation import CapacityUnavailableError +from twinkle.server.state.training_records import CreateTraining + + +class Clock: + now = 0.0 + + def __call__(self): + return self.now + + +class Runtime: + """Allocates a finite pool and stores actual per-user state for assertions.""" + + def __init__(self, devices=16): + self.devices = devices + self.allocations = {} + self.users = {} + self.starts = 0 + self.stopped = [] + self.start_gate = asyncio.Event() + self.start_gate.set() + self.fail_start = False + self.fail_shutdown = False + + async def allocate(self, instance, spec): + self.starts += 1 + if self.fail_start: + raise RuntimeError('injected startup failure') + used = sum(self.allocations.values()) + if used + spec.config.num_devices > self.devices: + raise CapacityUnavailableError('No devices available') + self.allocations[instance.instance_id] = spec.config.num_devices + return {} + + async def deploy(self, instance, spec): + await self.start_gate.wait() + + async def create_training(self, instance, training, request): + assert instance.instance_id in self.allocations + assert training.training_id not in self.users + self.users[training.training_id] = {'instance': instance.instance_id, 'label': training.adapter_label, 'steps': 0} + + async def release_training(self, instance, training): + self.users.pop(training.training_id, None) + + async def shutdown(self, instance): + if self.fail_shutdown: + raise RuntimeError('cannot confirm process death') + self.allocations.pop(instance.instance_id, None) + self.users = {key: value for key, value in self.users.items() if value['instance'] != instance.instance_id} + self.stopped.append(instance.instance_id) + + async def verify_operation(self, instance, request_id): + return {'state': 'pending'} + + async def recover_orphans(self): + pass + + +@pytest.mark.asyncio +async def test_initialization_timeout_only_retires_the_failing_tenant(): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, _ = await setup(backend=FakeBackend()) + active = await create(controller, state, 'active') + runtime.users[active]['steps'] = 7 + instance = controller.instances[controller.trainings[active].instance_id] + async def timeout(instance, training, request): + raise TimeoutError('User initialization timed out') + runtime.create_training = timeout + ref = await controller.create(request('failing')) + await controller.wait_for_jobs() + assert (await state.get_future(ref['request_id']))['failure']['reason_code'] == 'startup_timeout' + assert controller.trainings[active].state == 'ACTIVE' + assert runtime.users[active]['steps'] == 7 + assert instance.state == 'READY' and instance.training_ids == [active] + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_late_initializer_keeps_its_slot_reserved_until_serial_cleanup(): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, _ = await setup(backend=FakeBackend(), max_replicas=1) + active = await create(controller, state, 'active') + raw = controller.config.model_dump() + raw['startup_timeout_s'] = 0.03 + raw['models'][0]['modes']['lora']['max_loras'] = 2 + controller.config = TrainingConfig.model_validate(raw) + from twinkle.server.training.catalog import TrainingCatalog + controller.catalog = TrainingCatalog(controller.config) + entered, unblock = asyncio.Event(), asyncio.Event() + original_create, original_release = runtime.create_training, runtime.release_training + initializer = None + async def initialize(instance, training, req): + nonlocal initializer + async def late(): + entered.set() + await unblock.wait() + await original_create(instance, training, req) + initializer = asyncio.create_task(late()) + await asyncio.shield(initializer) + async def release(instance, training): + await initializer + await original_release(instance, training) + runtime.create_training, runtime.release_training = initialize, release + ref = await controller.create(request('late')) + await entered.wait() + for _ in range(100): + future = await state.get_future(ref['request_id']) + if future['status'] == 'failed': + break + await asyncio.sleep(0.005) + assert future['failure']['reason_code'] == 'startup_timeout' + instance = controller.instances[controller.trainings[active].instance_id] + assert len(instance.training_ids) == 2 and controller.trainings[active].state == 'ACTIVE' + refused = await controller.create(request('no-free-slot')) + assert (await state.get_future(refused['request_id']))['failure']['reason_code'] == 'capacity_unavailable' + unblock.set() + await controller.wait_for_jobs() + assert instance.training_ids == [active] and set(runtime.users) == {active} + runtime.create_training, runtime.release_training = original_create, original_release + replacement = await create(controller, state, 'replacement') + assert controller.trainings[replacement].instance_id == instance.instance_id + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_user_gets_full_initialization_window_after_base_model_startup(): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, clock = await setup(backend=FakeBackend()) + runtime.start_gate.clear() + ref = await controller.create(request('near-startup-deadline')) + instance = next(iter(controller.instances.values())) + clock.now = instance.startup_deadline - 1 + original = runtime.create_training + async def initialize(instance, training, req): + clock.now += 2 + await original(instance, training, req) + runtime.create_training = initialize + runtime.start_gate.set() + await controller.wait_for_jobs() + assert (await state.get_future(ref['request_id']))['status'] == 'completed' + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_resident_rebuild_does_not_block_expiry_or_another_mode(): + from tests.server.state.fake_backend import FakeBackend + from twinkle.server.training.catalog import TrainingCatalog + controller, state, runtime, clock = await setup(backend=FakeBackend()) + expired = await create(controller, state, 'expired-full', mode='full') + raw = controller.config.model_dump() + raw['models'][0]['modes']['lora']['min_replicas'] = 1 + controller.config = TrainingConfig.model_validate(raw) + controller.catalog = TrainingCatalog(controller.config) + runtime.start_gate.clear() + await asyncio.wait_for(controller.sweep(), 0.1) + for _ in range(20): + await asyncio.sleep(0) + if runtime.starts == 2: + break + clock.now = 11 + await asyncio.wait_for(controller.sweep(), 0.1) + for _ in range(20): + await asyncio.sleep(0) + if controller.trainings[expired].state == 'EXPIRED': + break + assert controller.trainings[expired].state == 'EXPIRED' + ref = await controller.create(request('unrelated-full', mode='full')) + assert (await state.get_future(ref['request_id']))['status'] == 'pending' + resident_ids = [item.instance_id for item in controller.instances.values() if item.resident] + assert len(resident_ids) == 1 + await controller.sweep() + assert [item.instance_id for item in controller.instances.values() if item.resident] == resident_ids + runtime.start_gate.set() + await controller.wait_for_jobs() + assert (await state.get_future(ref['request_id']))['status'] == 'completed' + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_reconciliation_of_committed_completion_does_not_extend_idle_time(): + import time + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, clock = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'active') + await controller.admit(identity, 'owner', 'session', 'lost-finish') + await controller.commit_operation(identity, 'lost-finish', controller.epoch) + async def finished(instance, request_id): + return {'state': 'finished', 'executed': True, 'finished_at': time.time() - 31} + runtime.verify_operation = finished + clock.now = 31 + await controller.sweep() + await controller.wait_for_jobs() + assert not controller.trainings[identity].pending_operation_ids + assert controller.trainings[identity].last_activity_at < 0.01 + await controller.sweep() + await controller.wait_for_jobs() + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_legacy_session_and_max_lifetime_bounds_keep_their_original_meaning(): + import time + from datetime import datetime, timezone + from tests.server.state.fake_backend import FakeBackend + from twinkle.server.training.catalog import TrainingCatalog + controller, state, runtime, clock = await setup(backend=FakeBackend()) + raw = controller.config.model_dump() + raw['models'][0]['modes']['lora']['idle_timeout_s'] = None + raw['models'][0]['legacy_expiration'] = {'session_timeout_s': 30, 'max_lifetime_s': 36000} + controller.config = TrainingConfig.model_validate(raw) + controller.catalog = TrainingCatalog(controller.config) + identity = await create(controller, state, 'legacy') + clock.now = 1801 + await state.touch_session('session') + await controller.sweep() + await controller.wait_for_jobs() + assert controller.trainings[identity].state == 'ACTIVE' + controller.trainings[identity].created_at = datetime.fromtimestamp(time.time() - 36001, timezone.utc).isoformat() + await controller.sweep() + await controller.wait_for_jobs() + assert controller.trainings[identity].state == 'EXPIRED' + second = await create(controller, state, 'stale-heartbeat') + session = await state._backend.get('session::session') + session['last_heartbeat'] = time.time() - 31 + await state._backend.set('session::session', session) + await controller.sweep() + await controller.wait_for_jobs() + assert controller.trainings[second].state == 'EXPIRED' + + +async def setup(*, devices=16, min_replicas=0, max_replicas=None, quota=100, backend=None): + mode = {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}, 'max_loras': 5, + 'idle_timeout_s': 10, 'min_replicas': min_replicas, 'max_replicas': max_replicas} + config = TrainingConfig(models=[{ + 'name': 'small', 'model_id': 'ms://small', 'backend': 'mock', 'device_type': 'cuda', + 'modes': {'lora': mode, 'full': {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}, 'idle_timeout_s': 10}}, + }]) + state = ServerState(backend=backend or RayActorBackend(), per_token_model_limit=quota) + await state.create_session({'session_id': 'session'}) + runtime, clock = Runtime(devices), Clock() + controller = TrainingController(config, state.create_training_store('unit'), runtime, clock=clock) + await controller.initialize() + return controller, state, runtime, clock + + +def request(key, *, mode='lora', label='same', token='owner'): + return CreateTraining(base_model='small', mode=mode, token=token, session_id='session', create_key=str(key), + client_type='twinkle', payload={}, adapter_label=label) + + +async def create(controller, state, key, **kwargs): + ref = await controller.create(request(key, **kwargs)) + await controller.wait_for_jobs() + future = await state.get_future(ref['request_id']) + assert future['status'] == 'completed', future + return future['result']['training_id'] + + +@pytest.mark.asyncio +async def test_lazy_queries_allocate_nothing(): + controller, state, runtime, _ = await setup() + assert runtime.starts == 0 + assert controller.catalog.get('small', 'full').capacity == 1 + with pytest.raises(ResourceNotFoundError): + await controller.route('unknown', 'owner') + assert not runtime.allocations + assert await state.get_capacity_info() == {'max_loras': 0, 'used_loras': 0, 'free_loras': 0} + + +@pytest.mark.asyncio +async def test_32_concurrent_creates_coalesce_even_before_allocation_ready(): + controller, state, runtime, _ = await setup() + runtime.start_gate.clear() + refs = await asyncio.gather(*(controller.create(request(index)) for index in range(32))) + assert len(controller.instances) == 7 + assert sorted(len(instance.training_ids) for instance in controller.instances.values()) == [2, 5, 5, 5, 5, 5, 5] + runtime.start_gate.set() + await controller.wait_for_jobs() + records = [await state.get_future(ref['request_id']) for ref in refs] + assert all(record['status'] == 'completed' for record in records) + assert len({record['result']['training_id'] for record in records}) == 32 + assert runtime.starts == 7 + assert len(runtime.users) == 32 + + +@pytest.mark.asyncio +async def test_full_and_lora_isolation_and_device_reuse(): + controller, state, runtime, _ = await setup(devices=6) + shared = await create(controller, state, 1) + full1 = await create(controller, state, 2, mode='full') + full2 = await create(controller, state, 3, mode='full') + identities = {runtime.users[key]['instance'] for key in (shared, full1, full2)} + assert len(identities) == 3 + old_full = runtime.users[full1]['instance'] + await controller.release(full1, 'owner') + replacement = await create(controller, state, 4, mode='full') + assert runtime.users[replacement]['instance'] != old_full + assert shared in runtime.users and full2 in runtime.users + assert len(runtime.allocations) == 3 + + +@pytest.mark.asyncio +async def test_capacity_failure_is_terminal_and_does_not_hold_quota(): + controller, state, runtime, _ = await setup(devices=2) + await create(controller, state, 1, mode='full') + ref = await controller.create(request(2, mode='full')) + await controller.wait_for_jobs() + record = await state.get_future(ref['request_id']) + assert record['status'] == 'failed' + assert record['failure']['reason_code'] == 'capacity_unavailable' + assert await state._backend.get('token_count::owner') == 1 + assert len(runtime.allocations) == 1 + assert await controller.create(request(2, mode='full')) == ref + + +@pytest.mark.asyncio +async def test_create_idempotency_and_changed_payload_conflict(): + controller, state, runtime, _ = await setup() + refs = await asyncio.gather(*(controller.create(request('one')) for _ in range(16))) + assert len({ref['request_id'] for ref in refs}) == 1 + await controller.wait_for_jobs() + assert len(runtime.users) == 1 + changed = request('one') + changed.payload = {'different': True} + with pytest.raises(RequestRejectedError, match='different parameters'): + await controller.create(changed) + + +@pytest.mark.asyncio +async def test_pending_creates_count_toward_quota(): + controller, state, runtime, _ = await setup(quota=1) + runtime.start_gate.clear() + first = await controller.create(request(1)) + second = await controller.create(request(2)) + failure = await state.get_future(second['request_id']) + assert failure['failure']['reason_code'] == 'resource_quota_exceeded' + runtime.start_gate.set() + await controller.wait_for_jobs() + assert (await state.get_future(first['request_id']))['status'] == 'completed' + assert len(runtime.users) == 1 + + +@pytest.mark.asyncio +async def test_heartbeat_does_not_prevent_idle_expiry_or_affect_busy_user(): + controller, state, runtime, clock = await setup() + idle = await create(controller, state, 1) + busy = await create(controller, state, 2) + await controller.admit(busy, 'owner', 'session', 'op') + await controller.commit_operation(busy, 'op', controller.epoch) + clock.now = 11 + await state.touch_session('session') + await controller.sweep() + await controller.wait_for_jobs() + assert idle not in runtime.users and busy in runtime.users + with pytest.raises(TrainingExpiredError): + await controller.route(idle, 'owner') + await controller.finish_operation(busy, 'op', controller.epoch, executed=True) + assert controller.trainings[busy].last_activity_at == 11 + clock.now = 20 + await controller.sweep() + await controller.wait_for_jobs() + assert busy in runtime.users + clock.now = 22 + await controller.sweep() + await controller.wait_for_jobs() + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_release_drains_accepted_operation_before_cleanup(): + controller, state, runtime, _ = await setup() + identity = await create(controller, state, 1) + await controller.admit(identity, 'owner', 'session', 'op') + await controller.commit_operation(identity, 'op', controller.epoch) + releasing = asyncio.create_task(controller.release(identity, 'owner')) + while controller.trainings[identity].state == 'ACTIVE': + await asyncio.sleep(0) + assert identity in runtime.users + with pytest.raises(ResourceNotFoundError): + await controller.admit(identity, 'owner', 'session', 'next') + await controller.finish_operation(identity, 'op', controller.epoch, executed=True) + await controller.finish_operation(identity, 'op', controller.epoch, executed=True) + await controller.wait_for_jobs() + await releasing + assert not runtime.users and not runtime.allocations + assert len(runtime.stopped) == 1 + + +@pytest.mark.asyncio +async def test_same_label_recreation_and_ownership_do_not_reattach_old_handle(): + controller, state, runtime, _ = await setup() + old = await create(controller, state, 1) + with pytest.raises(ResourceNotFoundError): + await controller.route(old, 'other') + with pytest.raises(ResourceNotFoundError): + await controller.route(old, 'owner', 'other-session') + await controller.release(old, 'owner') + new = await create(controller, state, 2) + assert new != old and runtime.users[new]['label'] == 'same' + with pytest.raises(ResourceNotFoundError): + await controller.route(old, 'owner') + + +@pytest.mark.asyncio +async def test_resident_survives_empty_but_extra_replica_does_not(): + controller, state, runtime, _ = await setup(min_replicas=1) + refs = [await controller.create(request(index)) for index in range(6)] + await controller.wait_for_jobs() + identities = [(await state.get_future(ref['request_id']))['result']['training_id'] for ref in refs] + for identity in identities: + await controller.release(identity, 'owner') + assert len(runtime.allocations) == 1 + resident = next(instance for instance in controller.instances.values() if instance.resident) + assert resident.state == 'READY' + + +@pytest.mark.asyncio +async def test_startup_failure_cleans_all_waiters_and_quota(): + controller, state, runtime, _ = await setup() + runtime.fail_start = True + refs = await asyncio.gather(*(controller.create(request(index)) for index in range(5))) + await controller.wait_for_jobs() + records = [await state.get_future(ref['request_id']) for ref in refs] + assert all(record['status'] == 'failed' for record in records) + assert (await state._backend.get('token_count::owner') or 0) == 0 + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_managed_records_are_not_deleted_by_generic_training_ttl(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 1) + await state._model_mgr.cleanup_expired(10**12, ['session']) + assert await state.get_model_metadata(identity) is not None + assert identity in runtime.users + + +@pytest.mark.asyncio +async def test_restart_invalidates_handles_and_cleans_old_allocation(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 1) + replacement = TrainingController(controller.config, 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') + await controller.commit_operation(identity, 'running', controller.epoch) + runtime.fail_shutdown = True + await controller.fail_instance(instance.instance_id, FutureFailureRecord( + reason_code='execution_timeout', message='RPC timed out', attribution='server')) + assert instance.state == 'QUARANTINED' + assert runtime.allocations and controller.trainings[identity].pending_operation_ids + assert await state.get_model_metadata(identity) is not None + failed = await controller.create(request(2, mode='full')) + await controller.wait_for_jobs() + assert (await state.get_future(failed['request_id']))['failure']['reason_code'] == 'capacity_unavailable' + runtime.fail_shutdown = False + await controller.sweep() + await controller.wait_for_jobs() + assert not runtime.allocations + assert instance.state == 'STOPPED' + assert (await state._backend.get('token_count::owner') or 0) == 0 + + +@pytest.mark.asyncio +async def test_uncertain_enqueue_is_failed_before_ticket_is_released(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 1) + await controller.admit(identity, 'owner', 'session', 'lost-submit') + await controller.store.future('lost-submit', 'pending', identity) + async def missing(instance, request_id): + return {'state': 'missing'} + runtime.verify_operation = missing + clock.now = 31 + await controller.sweep() + await controller.wait_for_jobs() + assert (await state.get_future('lost-submit'))['status'] == 'failed' + assert not controller.trainings[identity].pending_operation_ids + assert controller.trainings[identity].last_activity_at == 0 + await controller.sweep() + await controller.wait_for_jobs() + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_restart_fails_pending_operations_and_fences_old_controller(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 1) + await controller.admit(identity, 'owner', 'session', 'queued') + await controller.store.future('queued', 'queued', identity) + replacement = TrainingController(controller.config, 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'][0]['modes']['full']['min_replicas'] = 1 + original['models'][0]['modes']['full']['max_replicas'] = 1 + controller.config = TrainingConfig.model_validate(original) + from twinkle.server.training.catalog import TrainingCatalog + controller.catalog = TrainingCatalog(controller.config) + await controller.sweep() + await controller.wait_for_jobs() + identity = await create(controller, state, 1, mode='full') + entered, resume = asyncio.Event(), asyncio.Event() + previous = runtime.release_training + async def resetting(instance, training): + entered.set() + await resume.wait() + await previous(instance, training) + runtime.release_training = resetting + releasing = asyncio.create_task(controller.release(identity, 'owner')) + await entered.wait() + instance = controller.instances[controller.trainings[identity].instance_id] + assert instance.state == 'RESETTING' + ref = await controller.create(request(2, mode='full')) + assert (await state.get_future(ref['request_id']))['failure']['reason_code'] == 'capacity_unavailable' + resume.set() + await releasing + assert instance.state == 'READY' + next_id = await create(controller, state, 3, mode='full') + assert controller.trainings[next_id].instance_id == instance.instance_id + + +@pytest.mark.asyncio +async def test_dirty_resident_is_rebuilt_after_last_user_leaves(): + controller, state, runtime, _ = await setup(min_replicas=1) + identity = await create(controller, state, 1) + instance = controller.instances[controller.trainings[identity].instance_id] + instance.dirty = True + await controller.release(identity, 'owner') + assert instance.state == 'STOPPED' + await controller.sweep() + await controller.wait_for_jobs() + ready = [item for item in controller.instances.values() if item.state == 'READY'] + assert len(ready) == 1 and ready[0].instance_id != instance.instance_id + + +@pytest.mark.asyncio +@pytest.mark.parametrize('late_phase', ['runtime', 'user']) +async def test_late_startup_result_cannot_publish_success(late_phase): + controller, state, runtime, clock = await setup() + if late_phase == 'runtime': + original = runtime.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]['modes']['lora']['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._startup + assert (await controller.route(live, 'owner', 'session'))['instance_id'] in runtime.allocations + await controller.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('startup_blocked', [False, True]) +async def test_creation_record_write_failure_closes_new_empty_instance(monkeypatch, startup_blocked): + controller, state, runtime, clock = await setup() + if startup_blocked: + runtime.start_gate.clear() + original = state._backend.set + failed = False + async def failing_set(key, value, *args, **kwargs): + nonlocal failed + if (not failed and key.startswith('model::') and isinstance(value, dict) + and value.get('instance_id') and value.get('state') == 'PENDING'): + failed = True + raise RuntimeError('injected training record write failure') + return await original(key, value, *args, **kwargs) + monkeypatch.setattr(state._backend, 'set', failing_set) + ref = await controller.create(request('write-failure')) + await controller.wait_for_jobs() + assert (await state.get_future(ref['request_id']))['status'] == 'failed' + assert not runtime.allocations + assert await state._backend.get('token_count::owner') == 0 diff --git a/tests/server/training/test_failure_recovery.py b/tests/server/training/test_failure_recovery.py new file mode 100644 index 000000000..d93caf3fe --- /dev/null +++ b/tests/server/training/test_failure_recovery.py @@ -0,0 +1,373 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Fault injection ensures persistence and notification failures do not leak resources.""" +import asyncio +import logging +from unittest.mock import AsyncMock + +import pytest + +from tests.server.state.fake_backend import FakeBackend +from tests.server.training.test_controller import create, request, setup +from tests.server.training.test_runtime_queue import RPC, RuntimeQueue, settle +from twinkle.protocol.types.errors import ErrorCategory +from twinkle.server.exceptions import ResourceExhaustedError, ResourceNotFoundError +from twinkle.server.state.models import FAILURE_REASON_CODES, FutureFailureRecord + + +async def wait_idle(queue): + worker = queue._compute_worker + for _ in range(100): + if not worker._request_ids and not worker.new_task_event.is_set(): + await asyncio.sleep(0) + return + await asyncio.sleep(0.001) + raise AssertionError('Worker did not enter idle wait') + + +@pytest.mark.asyncio +@pytest.mark.parametrize('fault', ['commit', 'queued_status']) +async def test_warm_idle_worker_wakes_after_enqueue_failure(fault, monkeypatch): + controller, state, _, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'queue-failure') + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + executed = [] + + async def operation(): + executed.append(True) + return {'ok': True} + + try: + warmup = await queue.schedule_task(operation, model_id=identity, token='owner', task_type='step') + await settle(state, warmup['request_id']) + await wait_idle(queue) + executed.clear() + if fault == 'commit': + async def broken(*args): + raise RuntimeError('injected commit failure') + queue._training_controller.commit_operation = RPC(broken) + ref = await queue.schedule_task(operation, model_id=identity, token='owner', task_type='step') + else: + previous = state.store_future_status + + async def broken(request_id, status, *args, **kwargs): + if status == 'queued': + raise RuntimeError('injected queued write failure') + return await previous(request_id, status, *args, **kwargs) + + monkeypatch.setattr(state, 'store_future_status', broken) + with pytest.raises(RuntimeError, match='queued write'): + await queue.schedule_task(operation, model_id=identity, token='owner', task_type='step', + request_id='queued-failure') + ref = {'request_id': 'queued-failure'} + assert (await settle(state, ref['request_id']))['status'] == 'completed' + await wait_idle(queue) + assert executed == [True] + assert not controller.trainings[identity].pending_operation_ids + assert queue._compute_worker.total_queued() == 0 + finally: + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('fault,mode', [ + ('retire', 'lora'), ('save_draining', 'full'), ('save_stopped', 'full'), + ('close_retire', 'full'), ('post_retire_save', 'lora') +]) +async def test_cleanup_storage_failure_is_retryable(fault, mode, monkeypatch): + controller, state, runtime, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, fault, mode=mode) + training = controller.trainings[identity] + instance = controller.instances[training.instance_id] + previous_retire = controller.store.retire + previous_save = controller.store.save_instance + failed = False + + async def retire(record): + nonlocal failed + if fault in ('retire', 'close_retire') and not failed: + failed = True + raise RuntimeError('injected retirement failure') + return await previous_retire(record) + + async def save(record): + nonlocal failed + matches = ((fault == 'save_draining' and record.state == 'DRAINING') + or (fault == 'save_stopped' and record.state == 'STOPPED') + or (fault == 'post_retire_save' and not record.training_ids)) + if matches and not failed: + failed = True + raise RuntimeError('injected instance write failure') + return await previous_save(record) + + monkeypatch.setattr(controller.store, 'retire', retire) + monkeypatch.setattr(controller.store, 'save_instance', save) + with pytest.raises(RuntimeError, match='injected'): + await controller.release(identity, 'owner', 'session') + assert failed + assert identity not in controller._cleaning + assert instance.instance_id not in controller._closing + assert training.state == 'CLOSING' + assert instance.state != 'STOPPED' + for _ in range(3): + await controller.sweep() + await controller.wait_for_jobs() + if training.state == 'RELEASED': + break + assert training.state == 'RELEASED' + assert not instance.training_ids + assert instance.state == 'STOPPED' + assert not runtime.allocations + assert await state.get_model_metadata(identity) is None + assert (await state._backend.get('token_count::owner') or 0) == 0 + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_failure_transition_write_is_retried_and_admission_fenced(monkeypatch): + controller, state, runtime, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'failure-transition') + training = controller.trainings[identity] + instance = controller.instances[training.instance_id] + await controller.admit(identity, 'owner', 'session', 'pending-op') + await controller.store.future('pending-op', 'pending', identity) + original = controller.store.save_instance + failed = False + + async def save(record): + nonlocal failed + if record.state == 'DRAINING' and not failed: + failed = True + raise RuntimeError('injected draining write failure') + return await original(record) + + monkeypatch.setattr(controller.store, 'save_instance', save) + failure = FutureFailureRecord(reason_code='orphaned_replica', attribution='server', message='Worker died') + with pytest.raises(RuntimeError, match='draining write'): + await controller.fail_instance(instance.instance_id, failure) + assert instance.state == 'READY' + with pytest.raises(ResourceNotFoundError): + await controller.route(identity, 'owner') + await controller.sweep() + await controller.wait_for_jobs() + assert instance.state == 'STOPPED' and training.state == 'FAILED' + assert not runtime.allocations and not training.pending_operation_ids + assert (await state.get_future('pending-op'))['failure']['reason_code'] == 'orphaned_replica' + assert not controller._failed_instances + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_failed_resident_is_excluded_when_rebuilding_minimum_replicas(monkeypatch): + controller, _, runtime, _ = await setup(backend=FakeBackend(), min_replicas=1) + failed_instance = next(iter(controller.instances.values())) + original = controller.store.save_instance + failed = False + + async def save(record): + nonlocal failed + if record.state == 'DRAINING' and not failed: + failed = True + raise RuntimeError('injected draining write failure') + return await original(record) + + monkeypatch.setattr(controller.store, 'save_instance', save) + failure = FutureFailureRecord(reason_code='orphaned_replica', attribution='server', message='Worker died') + with pytest.raises(RuntimeError, match='draining write'): + await controller.fail_instance(failed_instance.instance_id, failure) + + assert failed_instance.state == 'READY' + assert failed_instance.instance_id in controller._failed_instances + await controller.sweep() + await controller.wait_for_jobs() + + ready_residents = [ + instance for instance in controller.instances.values() if instance.resident and instance.state == 'READY' + ] + assert failed_instance.state == 'STOPPED' + assert len(ready_residents) == 1 + assert ready_residents[0].instance_id != failed_instance.instance_id + assert set(runtime.allocations) == {ready_residents[0].instance_id} + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_shutdown_waits_for_an_in_flight_instance_close(): + controller, _, runtime, _ = await setup(backend=FakeBackend(), min_replicas=1) + instance = next(iter(controller.instances.values())) + original_shutdown = runtime.shutdown + entered, resume = asyncio.Event(), asyncio.Event() + + async def slow_shutdown(target): + entered.set() + await resume.wait() + await original_shutdown(target) + + runtime.shutdown = slow_shutdown + closing = asyncio.create_task(controller._close_instance(instance)) + await entered.wait() + shutdown = asyncio.create_task(controller.shutdown()) + for _ in range(100): + if instance.state == 'DRAINING' or shutdown.done(): + break + await asyncio.sleep(0) + + assert instance.state == 'DRAINING' + assert not shutdown.done() + resume.set() + await asyncio.gather(closing, shutdown) + assert instance.state == 'STOPPED' + assert instance.instance_id not in controller._closing + + +@pytest.mark.asyncio +async def test_cleanup_error_is_not_classified_as_startup_failure(): + controller, state, runtime, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'release-failure') + instance = controller.instances[controller.trainings[identity].instance_id] + runtime.release_training = AsyncMock(side_effect=TimeoutError('release timed out')) + await controller.release(identity, 'owner') + assert instance.failure == 'Training cleanup failed: release timed out' + assert controller.trainings[identity].terminal_reason == 'internal_error' + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_unready_service_is_503_server_and_reason_sets_have_no_dead_codes(): + controller, _, _, _ = await setup(backend=FakeBackend()) + controller._ready = False + with pytest.raises(ResourceExhaustedError) as error: + await controller.create(request('unready')) + assert error.value.error_code == 503 and error.value.category is ErrorCategory.Server + assert not {'protocol_mismatch', 'training_expired'} & FAILURE_REASON_CODES + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_spawn_logs_exception_with_traceback(caplog, monkeypatch): + # The ``twinkle`` logger disables propagation, so attach caplog to the controller logger directly. + target = logging.getLogger('twinkle.server.training.controller') + monkeypatch.setattr(target, 'handlers', [*target.handlers, caplog.handler]) + controller, _, _, _ = await setup(backend=FakeBackend()) + + async def broken(): + raise RuntimeError('background sentinel') + + with caplog.at_level(logging.ERROR): + controller._spawn(broken()) + await controller.wait_for_jobs() + records = [record for record in caplog.records if record.message == 'Training background job failed'] + assert records and records[0].exc_info[1].args == ('background sentinel',) + assert not controller._jobs + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_sweep_heartbeat_io_does_not_hold_route_lock(monkeypatch): + controller, state, _, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'slow-heartbeat') + entered, resume = asyncio.Event(), asyncio.Event() + original = state.get_session_last_heartbeats + + async def slow(identities): + entered.set() + await resume.wait() + return await original(identities) + + monkeypatch.setattr(state, 'get_session_last_heartbeats', slow) + sweeping = asyncio.create_task(controller.sweep()) + try: + await entered.wait() + assert await asyncio.wait_for(controller.route(identity, 'owner'), 0.1) + finally: + resume.set() + await sweeping + await controller.wait_for_jobs() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_explicit_tinker_request_id_is_preserved(): + from twinkle.server.model.managed_runtime import _tinker_operation + controller, state, _, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'explicit-id') + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + context = _tinker_operation.set((identity, 10)) + try: + ref = await queue.schedule_task(AsyncMock(return_value='ok'), model_id=identity, token='owner', + task_type='step', request_id='explicit-id') + assert ref['request_id'] == 'explicit-id' + assert (await settle(state, 'explicit-id'))['status'] == 'completed' + finally: + _tinker_operation.reset(context) + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_shutdown_waits_for_late_serve_deployment_before_confirming_release(monkeypatch): + import importlib + import threading + from types import SimpleNamespace + from twinkle.server.config.training import TrainingConfig + from twinkle.server.state.training_records import InstanceRecord + from twinkle.server.training.driver import RayInstanceDriver + from ray import serve + config = TrainingConfig(models=[{ + 'name': 'small', 'model_id': 'weights', 'backend': 'mock', 'device_type': 'cpu', + 'modes': {'lora': {'num_devices': 1, 'device_mesh': {'dp_size': 1}}} + }]) + store = SimpleNamespace(state=SimpleNamespace(unregister_replica_capacity=AsyncMock())) + driver = RayInstanceDriver(config, store, 'test-late', 'controller') + instance = InstanceRecord(instance_id='instance', controller_epoch='epoch', base_model='small', + mode='lora', startup_deadline=0) + entered, release = threading.Event(), threading.Event() + remove = AsyncMock() + + def deploy(*args, **kwargs): + entered.set() + if not release.wait(10): + raise TimeoutError('test did not release Serve') + + monkeypatch.setattr('twinkle.server.model.app.build_model_app', lambda **kwargs: 'app') + monkeypatch.setattr(importlib.import_module('ray.util.placement_group'), 'get_placement_group', lambda name: 'pg') + monkeypatch.setattr('ray.util.list_named_actors', lambda **kwargs: []) + monkeypatch.setattr(serve, 'run', deploy) + monkeypatch.setattr(serve, 'delete', lambda name: None) + monkeypatch.setattr(driver.pool, 'remove', remove) + deploying = asyncio.create_task(driver.deploy(instance, driver._catalog.get('small', 'lora'))) + closing = None + try: + assert await asyncio.to_thread(entered.wait, 2) + deploying.cancel() + await asyncio.gather(deploying, return_exceptions=True) + closing = asyncio.create_task(driver.shutdown(instance)) + await asyncio.sleep(0.01) + assert not closing.done() + remove.assert_not_awaited() + store.state.unregister_replica_capacity.assert_not_awaited() + finally: + release.set() + await asyncio.gather(deploying, return_exceptions=True) + if closing is not None: + await closing + await driver.wait_for_cleanup() + assert not driver._deployments and not driver._cleanup_tasks + store.state.unregister_replica_capacity.assert_awaited() + + +@pytest.mark.asyncio +async def test_capacity_unregister_does_not_retire_managed_model(): + controller, state, _, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'capacity-only') + instance = controller.instances[controller.trainings[identity].instance_id] + await state.register_replica(instance.instance_id, 5) + await state.unregister_replica_capacity(instance.instance_id) + assert await state.get_model_metadata(identity) is not None + assert await state._backend.get('token_count::owner') == 1 + await controller.release(identity, 'owner') + assert await state.get_model_metadata(identity) is None + assert await state._backend.get('token_count::owner') == 0 + await controller.shutdown() diff --git a/tests/server/training/test_http_integration.py b/tests/server/training/test_http_integration.py new file mode 100644 index 000000000..b20f14b29 --- /dev/null +++ b/tests/server/training/test_http_integration.py @@ -0,0 +1,32 @@ +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: + result = subprocess.run([ + sys.executable, + str(Path(__file__).resolve().parents[3] + / 'tests/server/integration/on_demand_training/http_probe.py') + ], capture_output=True, text=True, timeout=1800) + finally: + ray.init( + ignore_reinit_error=True, + num_cpus=2, + log_to_driver=False, + namespace='twinkle_test', + _temp_dir='/tmp/twinkle-test-' + uuid.uuid4().hex[:8]) + assert result.returncode == 0, result.stdout + result.stderr + evidence = next(json.loads(line) for line in result.stdout.splitlines() if line.startswith('{"passed":')) + assert evidence['passed'] and evidence['native_recreation'] and evidence['tinker_shared'] + assert evidence['native_lora'] + assert evidence['runtime_lifecycle_rejected'] + assert evidence['tinker_full'] + assert evidence['restart'] and evidence['state_loss'] and evidence['cycles'] == 10 diff --git a/tests/server/training/test_launcher.py b/tests/server/training/test_launcher.py new file mode 100644 index 000000000..d212e43e0 --- /dev/null +++ b/tests/server/training/test_launcher.py @@ -0,0 +1,39 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import socket +import uuid + +import pytest +import ray + +from twinkle.server.config import ServerConfig +from twinkle.server.launcher.server_launcher import ServerLauncher +from twinkle.server.training.actor import controller_name + + +def test_failed_initialization_does_not_leave_an_actor_that_blocks_a_retry(): + namespace = 'launcher-retry-' + uuid.uuid4().hex + with socket.socket() as listener: + listener.bind(('127.0.0.1', 0)) + port = listener.getsockname()[1] + devices = int(ray.cluster_resources()['CPU']) + 1 + payload = { + 'ray_namespace': namespace, 'http_options': {'host': '127.0.0.1', 'port': port}, + 'applications': [{'name': 'retry', 'route_prefix': '/api/v1/model/retry', 'import_path': 'model', + 'args': {'model_id': 'mock', 'backend': 'mock', 'device_type': 'cpu', + 'modes': {'full': {'num_devices': devices, 'device_mesh': {'dp_size': devices}, + 'min_replicas': 1}}}}], + } + launcher = ServerLauncher(ServerConfig.model_validate(payload)) + launcher._start_serve() + try: + with pytest.raises(ray.exceptions.RayTaskError, match='No single-node allocation'): + launcher._deploy_training() + with pytest.raises(ValueError): + ray.get_actor(controller_name(namespace), namespace=namespace) + payload['applications'][0]['args']['modes']['full']['min_replicas'] = 0 + launcher.config = ServerConfig.model_validate(payload) + launcher._deploy_training() + diagnostic = ray.get(launcher._training_controller.diagnostics.remote()) + assert diagnostic['ready'] and diagnostic['instances'] == [] + finally: + launcher._shutdown() diff --git a/tests/server/training/test_redis_restart_reads.py b/tests/server/training/test_redis_restart_reads.py new file mode 100644 index 000000000..bfde03295 --- /dev/null +++ b/tests/server/training/test_redis_restart_reads.py @@ -0,0 +1,114 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Read-only Redis reconnects are bounded; uncertain writes are never replayed.""" +import shutil +from unittest.mock import AsyncMock, Mock + +import pytest + +redis = pytest.importorskip('redis') +ConnectionError = redis.exceptions.ConnectionError + +from twinkle.server.state.backend.redis_backend import RedisBackend +from tests.server.integration.on_demand_training.managed_redis import ManagedRedis + + +@pytest.fixture(scope='session', autouse=True) +def _ray_runtime(): + yield # These backend tests need no Ray cluster. + + +@pytest.mark.asyncio +async def test_stale_read_is_retried_but_write_is_not(): + backend = RedisBackend('redis://127.0.0.1:1/0') + try: + backend._client.get = AsyncMock(side_effect=[ConnectionError('stale socket'), '{"value": 1}']) + assert await backend.get('owned') == {'value': 1} + assert backend._client.get.await_count == 2 + backend._client.set = AsyncMock(side_effect=ConnectionError('lost acknowledgement')) + with pytest.raises(ConnectionError): + await backend.set('owned', {'value': 2}) + assert backend._client.set.await_count == 1 + finally: + await backend.close() + + +@pytest.mark.asyncio +async def test_persistent_outage_does_not_look_like_missing_state(): + backend = RedisBackend('redis://127.0.0.1:1/0') + try: + backend._client.get = AsyncMock(side_effect=ConnectionError('offline')) + with pytest.raises(ConnectionError): + await backend.get('controller-active') + assert backend._client.get.await_count == 3 + finally: + await backend.close() + + +@pytest.mark.asyncio +@pytest.mark.skipif(not shutil.which('redis-server'), reason='redis-server executable required') +async def test_real_restart_recovers_old_pooled_readers(tmp_path): + managed = ManagedRedis(tmp_path) + managed.start() + first, stale, atomic = (RedisBackend(managed.url) for _ in range(3)) + try: + await first.set('owner', {'epoch': 'original'}) + assert await stale.get('owner') == {'epoch': 'original'} + assert await atomic.get('owner') == {'epoch': 'original'} + managed.stop() + with pytest.raises(ConnectionError): + await first.get('owner') + managed.start() + assert await stale.get('owner') == {'epoch': 'original'} + assert await stale.exists('owner') + assert await first.mget(['owner', 'missing']) == [{'epoch': 'original'}, None] + assert await stale.health_check() + assert await atomic.update_atomic('counter', lambda value: int(value or 0) + 1) == 1 + assert await atomic.get('counter') == 1 + finally: + await first.close() + await stale.close() + await atomic.close() + managed.stop() + + +class PipelineStub: + def __init__(self): + self.watch = AsyncMock() + self.get = AsyncMock(return_value='0') + self.multi = Mock() + self.set = Mock() + self.execute = AsyncMock(return_value=[]) + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + +@pytest.mark.asyncio +async def test_transport_watch_error_after_exec_is_never_replayed(): + backend = RedisBackend('redis://127.0.0.1:1/0') + pipe = PipelineStub() + pipe.execute.side_effect = redis.exceptions.WatchError( + 'A ConnectionError occurred while watching one or more keys') + backend._client.pipeline = lambda: pipe + try: + with pytest.raises(ConnectionError, match='outcome could not be confirmed'): + await backend.update_atomic('counter', lambda value: value + 1) + assert pipe.execute.await_count == 1 + finally: + await backend.close() + + +@pytest.mark.asyncio +async def test_connection_loss_before_exec_is_safe_to_retry(): + backend = RedisBackend('redis://127.0.0.1:1/0') + pipe = PipelineStub() + pipe.watch.side_effect = [ConnectionError('stale'), None] + backend._client.pipeline = lambda: pipe + try: + assert await backend.update_atomic('counter', lambda value: value + 1) == 1 + assert pipe.watch.await_count == 2 and pipe.execute.await_count == 1 + finally: + await backend.close() diff --git a/tests/server/training/test_router.py b/tests/server/training/test_router.py new file mode 100644 index 000000000..37741d0b7 --- /dev/null +++ b/tests/server/training/test_router.py @@ -0,0 +1,101 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +from unittest.mock import Mock +from types import SimpleNamespace + +import pytest +from starlette.requests import Request + +from tests.server.state.fake_backend import FakeBackend +from tests.server.training.test_runtime_queue import settle +from twinkle.hub import HubOperation +from twinkle.server.state import ServerState +from twinkle.server.training.router import TrainingRouter +from twinkle.server.training.catalog import TrainingCatalog +from twinkle.server.config.training import TrainingConfig + + +def http_request(body): + async def receive(): + return {'type': 'http.request', 'body': json.dumps(body).encode(), 'more_body': False} + request = Request({'type': 'http', 'method': 'POST', 'path': '/upload_to_hub', 'headers': []}, receive) + request.state.token = 'owner' + request.state.session_id = '' + return request + + +@pytest.mark.asyncio +async def test_upload_after_recycle_is_cpu_only_and_preserves_future_ownership(monkeypatch, tmp_path): + router = TrainingRouter.__new__(TrainingRouter) + router.state = ServerState(backend=FakeBackend()) + router.store = 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._backend.get('future_owner::' + response.request_id) == 'owner' + assert not await router.state._backend.keys('model::*') + finally: + await router.shutdown_task_queue() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('rank', [8, 0, 33, True]) +async def test_native_lora_admission_does_not_import_peft(monkeypatch, rank): + import builtins + from unittest.mock import AsyncMock + from twinkle.server.exceptions import RequestRejectedError + router = TrainingRouter.__new__(TrainingRouter) + router.model = SimpleNamespace(name='small') + router.state = ServerState(backend=FakeBackend()) + await router.state.create_session({'session_id': 'session'}) + router.catalog = TrainingCatalog(TrainingConfig(models=[{ + 'name': 'small', 'model_id': 'weights', 'backend': 'transformers', 'device_type': 'cuda', + 'modes': {'lora': {'num_devices': 1, 'device_mesh': {'dp_size': 1}}}, + }])) + async def create(payload): + await router.state.store_future_status('created', 'pending', None) + return {'request_id': 'created'} + router.controller = SimpleNamespace(create=SimpleNamespace(remote=AsyncMock(side_effect=create))) + original_import = builtins.__import__ + def cpu_import(name, *args, **kwargs): + if name == 'peft' or name.startswith('peft.'): + raise AssertionError('CPU admission must not import PEFT') + return original_import(name, *args, **kwargs) + monkeypatch.setattr(builtins, '__import__', cpu_import) + request = http_request({'client_create_id': 'create', 'adapter_name': 'same', + 'config': json.dumps({'_TWINKLE_TYPE_': 'LoraConfig', 'r': rank})}) + request.state.session_id = 'session' + if type(rank) is int and rank == 8: + result = await router.native_request(request, 'add_adapter_to_model') + assert result.request_id == 'created' + router.controller.create.remote.assert_awaited_once() + else: + with pytest.raises(RequestRejectedError, match='rank'): + await router.native_request(request, 'add_adapter_to_model') + router.controller.create.remote.assert_not_awaited() + + +def test_router_schemas_match_model_runtime_routes(): + from fastapi import FastAPI + from twinkle.server.model.tinker_handlers import _register_model_tinker_routes + from twinkle.server.model.twinkle_handlers import _register_model_twinkle_routes + from twinkle.server.training.router import NATIVE_REQUESTS, TINKER_REQUESTS + + app = FastAPI() + _register_model_tinker_routes(app, lambda: None) + _register_model_twinkle_routes(app, lambda: None) + for protocol, schemas in [('twinkle', NATIVE_REQUESTS), ('tinker', TINKER_REQUESTS)]: + endpoints = {route.path.rsplit('/', 1)[-1] for route in app.routes + if route.path.startswith(f'/{protocol}/') and 'POST' in route.methods} + assert endpoints == set(schemas) diff --git a/tests/server/training/test_runtime_initialization.py b/tests/server/training/test_runtime_initialization.py new file mode 100644 index 000000000..27c8ec219 --- /dev/null +++ b/tests/server/training/test_runtime_initialization.py @@ -0,0 +1,49 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""The runtime event loop remains responsive during the first dynamic config import.""" +import asyncio +import threading +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from tests.server.state.fake_backend import FakeBackend +from tests.server.training.test_controller import create, setup +from twinkle.server.model.managed_runtime import ManagedTrainingRuntime +from twinkle.server.state.training_records import CreateTraining + + +@pytest.mark.asyncio +async def test_configuration_deserialization_does_not_block_event_loop(monkeypatch): + controller, _, _, _ = await setup(backend=FakeBackend()) + identity = await create(controller, controller.store.state, 'configuration-import', mode='full') + training = controller.trainings[identity] + entered, resume, finished = threading.Event(), threading.Event(), threading.Event() + + def deserialize(config): + entered.set() + resume.wait(3) + finished.set() + return None + + monkeypatch.setattr('twinkle.server.model.managed_runtime.deserialize_object', deserialize) + manager = SimpleNamespace(save=Mock()) + monkeypatch.setattr('twinkle.server.model.managed_runtime.create_training_run_manager', lambda *args, **kw: manager) + + async def serial(factory, task_type): + return await factory() + + runtime = SimpleNamespace(_managed_serial=serial, resolve_model_adapter_name=lambda _: '', + register_resource=Mock(), set_resource_state=Mock(), is_full_mode=True) + request = CreateTraining(base_model='small', mode='full', token='owner', session_id='session', create_key='key', + client_type='twinkle', payload={'adapter_name': 'same', 'config': None}) + initializing = asyncio.create_task(ManagedTrainingRuntime.initialize_training( + runtime, training.model_dump(), request.model_dump())) + try: + assert await asyncio.to_thread(entered.wait, 2) + assert not finished.is_set(), 'Synchronous import blocked the event loop' + finally: + resume.set() + await initializing + await controller.shutdown() + manager.save.assert_called_once() diff --git a/tests/server/training/test_runtime_queue.py b/tests/server/training/test_runtime_queue.py new file mode 100644 index 000000000..57f02083f --- /dev/null +++ b/tests/server/training/test_runtime_queue.py @@ -0,0 +1,191 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Lifecycle protection through the actual compute worker and backend thread.""" +import asyncio +import threading +from types import SimpleNamespace + +import pytest + +from twinkle.server.config.task_queue import TaskQueueConfig +from twinkle.server.task_queue import TaskQueueMixin +from twinkle.server.model.managed_runtime import ManagedTrainingRuntime, _tinker_operation +from tests.server.training.test_controller import create, setup + + +class RPC: + def __init__(self, method): + self.method = method + + def remote(self, *args, **kwargs): + return asyncio.create_task(self.method(*args, **kwargs)) + + +class RuntimeQueue(ManagedTrainingRuntime, TaskQueueMixin): + def __init__(self, controller, state, instance): + self.state, self.replica_id = state, instance.instance_id + self.training_context = {'instance_id': instance.instance_id, 'epoch': controller.epoch} + self._model_unhealthy = False + self._training_controller = SimpleNamespace(**{ + name: RPC(getattr(controller, name)) + for name in ('admit', 'commit_operation', 'finish_operation', 'fail_instance') + }) + self._init_task_queue(TaskQueueConfig(execution_timeout=0.1, inline_fast_path_timeout=0.01), + enable_admission_gate=True) + + +async def settle(state, request_id): + for _ in range(300): + future = await state.get_future(request_id) + if future and future['status'] in ('completed', 'failed', 'cancelled'): + return future + await asyncio.sleep(0.01) + raise AssertionError('Future did not settle') + + +@pytest.mark.asyncio +async def test_execution_timeout_keeps_live_thread_and_allocation_quarantined(): + controller, state, runtime, _ = await setup(devices=2) + identity = await create(controller, state, 1) + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + entered, unblock = threading.Event(), threading.Event() + runtime.fail_shutdown = True + def blocking(): + entered.set() + unblock.wait(3) + async def operation(): + await queue.call_backend(blocking) + try: + ref = await queue.schedule_task(operation, model_id=identity, token='owner', task_type='forward_backward') + future = await settle(state, ref['request_id']) + assert entered.is_set() and not unblock.is_set() + assert future['failure']['reason_code'] == 'execution_timeout' + for _ in range(100): + if instance.state == 'QUARANTINED':break + await asyncio.sleep(0.01) + assert instance.state == 'QUARANTINED' and runtime.allocations + assert controller.trainings[identity].pending_operation_ids + assert queue._model_unhealthy + finally: + unblock.set() + await queue.shutdown_task_queue() + runtime.fail_shutdown = False + await controller.sweep() + 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} + context = _tinker_operation.set((identity, 1)) + try: + first, retry = await asyncio.gather(*( + queue.schedule_task(operation, model_id=identity, token='owner', task_type='forward_backward') + for _ in range(2))) + assert first == retry + assert (await settle(state, first['request_id']))['status'] == 'completed' + step = await queue.schedule_task(operation, model_id=identity, token='owner', task_type='optim_step') + assert step['request_id'] != first['request_id'] + await settle(state, step['request_id']) + assert len(counter) == 2 + finally: + _tinker_operation.reset(context) + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_cancelled_queued_operation_releases_ticket_without_business_renewal(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 1) + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + queue._task_queue_config = TaskQueueConfig(execution_timeout=2) + queue._compute_worker._config = queue._task_queue_config + entered, unblock = asyncio.Event(), asyncio.Event() + executions = [] + async def busy(): + entered.set() + await unblock.wait() + async def cancelled(): + executions.append('must not run') + try: + first = await queue.schedule_task(busy, model_id=identity, token='owner', task_type='forward_backward') + await entered.wait() + second = await queue.schedule_task(cancelled, model_id=identity, token='owner', task_type='step') + assert (await state.cancel_future(second['request_id']))['cancelled'] + clock.now = 9 + unblock.set() + await settle(state, first['request_id']) + for _ in range(100): + if not controller.trainings[identity].pending_operation_ids:break + await asyncio.sleep(0.01) + assert not executions and not controller.trainings[identity].pending_operation_ids + assert controller.trainings[identity].last_activity_at == 9 + assert (await state.get_future(second['request_id']))['status'] == 'cancelled' + finally: + unblock.set() + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_readiness_waits_for_real_ray_worker_constructor(): + import ray + import time + @ray.remote(num_cpus=0.1) + class Worker: + def __init__(self): + time.sleep(0.3) + self.loaded = True + def is_loaded(self): + return self.loaded + worker = Worker.remote() + runtime = ManagedTrainingRuntime() + runtime.model = SimpleNamespace(_actors=[worker]) + try: + await runtime.wait_for_training_workers() + assert await worker.is_loaded.remote() + assert ray._private.state.actors(worker._actor_id.hex())['State'] == 'ALIVE' + finally: + ray.kill(worker, no_restart=True) + + +@pytest.mark.asyncio +async def test_lost_finish_notification_is_reconciled_without_stopping_compute(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 'lost-notification') + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + async def lost(*args, **kwargs): + raise RuntimeError('Completion RPC lost') + queue._training_controller.finish_operation = RPC(lost) + async def status(instance, request_id): + return await queue.verify_operation(request_id) + runtime.verify_operation = status + async def compute(): + return {'updated': True} + try: + ref = await queue.schedule_task(compute, model_id=identity, token='owner', task_type='step') + assert (await settle(state, ref['request_id']))['status'] == 'completed' + assert controller.trainings[identity].pending_operation_ids + clock.now = 6 + await controller.sweep() + await controller.wait_for_jobs() + assert not controller.trainings[identity].pending_operation_ids + assert controller.trainings[identity].state == 'ACTIVE' and not queue._model_unhealthy + clock.now += 11 + await controller.sweep() + await controller.wait_for_jobs() + assert not runtime.allocations + finally: + await queue.shutdown_task_queue() + await controller.shutdown() diff --git a/tests/server/training/test_sequence.py b/tests/server/training/test_sequence.py new file mode 100644 index 000000000..924258bc6 --- /dev/null +++ b/tests/server/training/test_sequence.py @@ -0,0 +1,70 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +import time +import uuid + +import pytest +import pytest_asyncio + +from twinkle.server.exceptions import RequestRejectedError +from twinkle.server.state import ServerState +from twinkle.server.state.backend.memory_backend import RayActorBackend +from twinkle.server.state.backend.redis_backend import RedisBackend + + +@pytest_asyncio.fixture(params=['memory', 'redis']) +async def state(request): + if request.param == 'redis': + url = os.environ.get('TWINKLE_TEST_REDIS_URL') + if not url: + pytest.skip('TWINKLE_TEST_REDIS_URL is required for the Redis regression') + backend = RedisBackend(url, key_prefix='sequence-test-' + uuid.uuid4().hex + ':') + else: + backend = RayActorBackend() + yield ServerState(backend=backend) + if request.param == 'redis': + for key in await backend.keys('*'): + await backend.delete(key) + await backend.close() + + +@pytest.mark.asyncio +async def test_expired_retry_receipts_cannot_execute_old_sequences_again(state, monkeypatch): + backend = state._backend + now = time.time() + monkeypatch.setattr('twinkle.server.state.server_state.time.time', lambda: now) + assert await state.claim_training_seq('training', 'native::session::step', 1, 'first', 10) is None + assert await state.claim_training_seq('training', 'native::session::step', 1, 'retry', 10) == 'first' + now += 11 + assert await state.claim_training_seq('training', 'native::session::step', 2, 'second', 10) is None + with pytest.raises(RequestRejectedError, match='retry window expired'): + await state.claim_training_seq('training', 'native::session::step', 1, 'late', 10) + ledger = await backend.get('seq_ledger::training::native::session::step') + assert list(ledger['receipts']) == ['2'] + assert len(await backend.keys('seq_ledger::*')) == 1 + + +@pytest.mark.asyncio +async def test_out_of_order_sequences_and_preflight_release_remain_retryable(state): + assert await state.claim_training_seq('training', 'tinker::step', 2, 'second', 10) is None + assert await state.claim_training_seq('training', 'tinker::step', 1, 'first', 10) is None + await state.release_training_seq('training', 'tinker::step', 1, 'another-request') + assert await state.claim_training_seq('training', 'tinker::step', 1, 'retry', 10) == 'first' + await state.release_training_seq('training', 'tinker::step', 1, 'first') + assert await state.claim_training_seq('training', 'tinker::step', 1, 'retry', 10) is None + assert await state.claim_training_seq('other-training', 'tinker::step', 1, 'other', 10) is None + + +@pytest.mark.asyncio +async def test_retiring_a_handle_removes_native_and_tinker_stream_ledgers(state): + from tests.server.training.test_controller import create, setup + controller, state, _, _ = await setup(backend=state._backend) + try: + identity = await create(controller, state, 'dedup-retirement') + for stream in ('native::session::step', 'tinker::optim_step'): + await state.claim_training_seq(identity, stream, 1, stream, 10) + assert len(await state._backend.keys(f'seq_ledger::{identity}::*')) == 2 + await controller.release(identity, 'owner', 'session') + assert await state._backend.keys(f'seq_ledger::{identity}::*') == [] + finally: + await controller.shutdown() diff --git a/tests/server/training/test_storage_recovery.py b/tests/server/training/test_storage_recovery.py new file mode 100644 index 000000000..c2529a6ca --- /dev/null +++ b/tests/server/training/test_storage_recovery.py @@ -0,0 +1,63 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""The lifecycle uses the same recovery semantics on real memory and Redis backends.""" +import os +import uuid + +import pytest + +from twinkle.server.state.backend.memory_backend import RayActorBackend +from twinkle.server.state.backend.redis_backend import RedisBackend +from twinkle.server.training.controller import TrainingController +from tests.server.training.test_controller import create, request, setup + + +@pytest.mark.asyncio +@pytest.mark.parametrize('mode', ['memory', 'redis']) +async def test_pending_operation_recovery_preserves_creation_identity_and_quota(mode): + if mode == 'redis' and not os.environ.get('TWINKLE_TEST_REDIS_URL'): + pytest.skip('Set TWINKLE_TEST_REDIS_URL to a dedicated local Redis') + backend = (RayActorBackend() if mode == 'memory' else + RedisBackend(os.environ['TWINKLE_TEST_REDIS_URL'], key_prefix='training-test-' + uuid.uuid4().hex + ':')) + try: + controller, state, runtime, clock = await setup(backend=backend) + identity = await create(controller, state, 'retry') + original = await controller.create(request('retry')) + await controller.admit(identity, 'owner', 'session', 'queued') + await controller.store.future('queued', 'queued', identity) + replacement = TrainingController(controller.config, 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/utils/task_queue/test_config.py b/tests/server/utils/task_queue/test_config.py index 60787f567..511d6bde5 100644 --- a/tests/server/utils/task_queue/test_config.py +++ b/tests/server/utils/task_queue/test_config.py @@ -13,7 +13,7 @@ from hypothesis import strategies as st from pydantic import ValidationError -from twinkle.server.task_queue.config import TaskQueueConfig +from twinkle.server.config.task_queue import TaskQueueConfig # ---------- defaults snapshot used by the default-value test -------------- # diff --git a/tests/server/utils/test_checkpoint_manager.py b/tests/server/utils/test_checkpoint_manager.py index e4d27eb85..5303a876b 100644 --- a/tests/server/utils/test_checkpoint_manager.py +++ b/tests/server/utils/test_checkpoint_manager.py @@ -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/utils/test_task_queue_mixin.py b/tests/server/utils/test_task_queue_mixin.py index a0316c632..23ad79a80 100644 --- a/tests/server/utils/test_task_queue_mixin.py +++ b/tests/server/utils/test_task_queue_mixin.py @@ -2,7 +2,7 @@ import pytest import pytest_asyncio -from twinkle.server.task_queue.config import TaskQueueConfig +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.task_queue.mixin import TaskQueueMixin from twinkle.server.task_queue.types import UserTaskError from twinkle.server.task_queue.worker import ComputeWorker, _exception_summary @@ -65,6 +65,13 @@ async def store_future_status(self, *args, **kwargs): 'queue_state_reason': kwargs.get('queue_state_reason'), } + async def try_start_future(self, request_id): + record = self._latest.get(request_id) + if record is None or record['status'] not in ('pending', 'queued'): + return False + await self.store_future_status(request_id, 'running', None, queue_state='active') + return self._latest[request_id]['status'] == 'running' + async def get_future(self, request_id): return self._latest.get(request_id) @@ -94,6 +101,7 @@ def enable_compute_worker(self): deployment_name=self._deployment_name, ) self._event_loop = None + self._queue_tasks = set() @pytest.mark.asyncio @@ -140,7 +148,7 @@ async def wait_for_with_completion_race(coro, timeout): @pytest.mark.asyncio -async def test_worker_can_restart_after_stopping_while_idle(): +async def test_idle_worker_stop_is_terminal_for_its_replica(): queue = _DummyQueue() queue.enable_compute_worker() worker = queue._compute_worker @@ -152,9 +160,11 @@ async def work(): await worker.ensure_started() await asyncio.sleep(0) # Let the idle worker wait for new work. await asyncio.wait_for(worker.stop(), timeout=1) - env = await queue.submit_and_peek(work, model_id='model1', token='token1') - assert env.status == 'completed' - assert env.result == {'ok': True} + from twinkle.server.exceptions import ResourceExhaustedError + with pytest.raises(ResourceExhaustedError): + await queue.submit_and_peek(work, model_id='model1', token='token1') + assert not queue.state.records + assert not worker._started finally: await worker.stop() @@ -221,6 +231,7 @@ async def test_preflight_accepts_batch_with_per_dp_multiple(): @pytest.mark.asyncio async def test_background_task_tracks_status(): queue = _DummyQueue() + queue.enable_compute_worker() async def work(): return {'ok': True} @@ -229,9 +240,13 @@ async def work(): work, model_id='model1', ) - await asyncio.sleep(0) + for _ in range(100): + if not queue._compute_worker._detached_tasks: + break + await asyncio.sleep(0.001) + await queue._compute_worker.stop() - assert [args[1] for args, _ in queue.state.records] == ['running', 'completed'] + assert [args[1] for args, _ in queue.state.records] == ['pending', 'running', 'completed'] assert queue.state.records[0][1]['absolute_deadline'] > 0 assert queue.state.records[-1][1]['result'] == {'ok': True} @@ -614,6 +629,108 @@ async def write(*args, **kwargs): assert [args[1] for args, _ in queue.state.records] == ['pending', 'failed'] +@pytest.mark.asyncio +async def test_background_capacity_and_shutdown_are_supervised(make_concurrent_queue): + from twinkle.server.exceptions import ResourceExhaustedError + queue = make_concurrent_queue(TaskQueueConfig(concurrent_task_limit=1)) + entered, exited = asyncio.Event(), asyncio.Event() + + async def work(): + entered.set() + try: + await asyncio.Event().wait() + finally: + exited.set() + + ref = await queue.schedule_background_task(work) + await asyncio.wait_for(entered.wait(), 1) + with pytest.raises(ResourceExhaustedError): + await queue.schedule_background_task(work) + assert len(queue.state._latest) == 1 + await queue._compute_worker.stop() + assert exited.is_set() + assert queue.state._latest[ref['request_id']]['status'] == 'failed' + assert queue._compute_worker._detached_reserved == 0 + + +@pytest.mark.asyncio +async def test_stopped_replica_rejects_serial_submission_without_a_future(): + from twinkle.server.exceptions import ResourceExhaustedError + queue = _DummyQueue() + queue.enable_compute_worker() + await queue._compute_worker.stop() + with pytest.raises(ResourceExhaustedError): + await queue.schedule_task(lambda: asyncio.sleep(0)) + assert not queue.state.records + assert not queue._compute_worker._started + + +@pytest.mark.asyncio +async def test_stop_during_pending_write_cannot_restart_the_serial_worker(): + from twinkle.server.exceptions import ResourceExhaustedError + queue = _DummyQueue() + queue.enable_compute_worker() + original = queue.state.store_future_status + + async def write_then_stop(*args, **kwargs): + await original(*args, **kwargs) + if args[1] == 'pending': + await queue._compute_worker.stop() + + queue.state.store_future_status = write_then_stop + with pytest.raises(ResourceExhaustedError): + await queue.schedule_task(lambda: asyncio.sleep(0), request_id='late-serial') + assert queue.state._latest['late-serial']['status'] == 'failed' + assert not queue._compute_worker._started + assert queue._compute_worker.total_queued() == 0 + + +@pytest.mark.asyncio +async def test_shutdown_drains_serial_tasks_that_never_started(): + queue = _DummyQueue() + queue.enable_compute_worker() + executed = [] + + async def work(): + executed.append(True) + + ref = await queue.schedule_task(work) + await queue._compute_worker.stop() + assert not executed + assert queue.state._latest[ref['request_id']]['status'] == 'failed' + assert queue._compute_worker.total_queued() == 0 + assert not queue._compute_worker.has_task(ref['request_id']) + + +@pytest.mark.asyncio +async def test_unknown_start_ack_never_executes_backend(make_concurrent_queue, monkeypatch): + from tests.server.state.fake_backend import FakeBackend + from twinkle.server.state.server_state import ServerState + queue = make_concurrent_queue() + backend = FakeBackend() + queue.state = queue._compute_worker._state = ServerState(backend=backend) + original = backend.update_atomic + executed = [] + + async def lost_ack(*args, **kwargs): + result = await original(*args, **kwargs) + if kwargs.get('return_changed'): + raise ConnectionError('start acknowledgement lost') + return result + + async def work(): + executed.append(True) + + monkeypatch.setattr(backend, 'update_atomic', lost_ack) + ref = await queue.schedule_concurrent_task(work) + for _ in range(100): + if not queue._compute_worker._detached_tasks: + break + await asyncio.sleep(0.001) + assert not executed + assert (await queue.state.get_future(ref['request_id']))['status'] == 'failed' + + @pytest.mark.asyncio async def test_cancel_winning_running_write_race_skips_execution(make_concurrent_queue): queue = make_concurrent_queue() 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/twinkle_client/test_async_components.py b/tests/twinkle_client/test_async_components.py index 0a092627c..313986191 100644 --- a/tests/twinkle_client/test_async_components.py +++ b/tests/twinkle_client/test_async_components.py @@ -41,6 +41,10 @@ def _recorder(calls, result_factory): def post(url, headers=None, data=None, timeout=None, **kwargs): body = json.loads(data) if data else kwargs.get('json') or {} calls.append((url, body)) + if '/model/' in url and url.endswith('/create'): + return _Response({'training_handles': True}) + if url.endswith('/add_adapter_to_model'): + return _Response(_completed({'training_id': 'training-test', 'adapter_name': body['adapter_name']})) return _Response(result_factory(url)) return post @@ -65,7 +69,7 @@ def test_model_forward_backward_sends_multiple_data_refs(monkeypatch) -> None: }})) model = module.MultiLoraTransformersModel('ms://base') - model.adapter_name = 'adapter' + model.add_adapter_to_model('adapter', None) refs = [ DataRef(ref_id='data-1', size=2, fields=['train_input']), DataRef(ref_id='data-2', size=2, fields=['train_input']), @@ -91,7 +95,7 @@ def test_model_inline_forward_methods_keep_the_original_endpoints(monkeypatch) - _patch_transport(monkeypatch, calls, lambda url: {} if url.endswith('/create') else _completed({'result': {}})) model = module.MultiLoraTransformersModel('ms://base') - model.adapter_name = 'adapter' + model.add_adapter_to_model('adapter', None) inputs = [{'input_ids': [1, 2]}] model.forward(inputs, return_logits=True) model.forward_only(inputs, disable_lora=True) @@ -140,7 +144,7 @@ def test_undeclared_forward_arguments_are_routed_to_loss_kwargs(monkeypatch) -> _patch_transport(monkeypatch, calls, lambda url: {} if url.endswith('/create') else _completed({'result': {}})) model = module.MultiLoraTransformersModel('ms://base') - model.adapter_name = 'adapter' + model.add_adapter_to_model('adapter', None) model.forward_backward([{'input_ids': [1, 2]}], advantages=[0.5], old_logps=[[-1.0, -2.0]]) _, body = calls[-1] @@ -158,7 +162,7 @@ def test_model_data_plane_forward_uses_a_separate_api(monkeypatch) -> None: }})) model = module.MultiLoraTransformersModel('ms://base') - model.adapter_name = 'adapter' + model.add_adapter_to_model('adapter', None) ref = DataRef(ref_id='data-1', size=2, fields=['train_input']) model.forward_from_data_plane(ref, input_field='train_input') @@ -180,7 +184,7 @@ def test_model_data_plane_forward_only_can_append_selected_outputs(monkeypatch) if url.endswith('/create') else _completed({'result': updated_ref.model_dump()})) model = module.MultiLoraTransformersModel('ms://base') - model.adapter_name = 'adapter' + model.add_adapter_to_model('adapter', None) result = model.forward_only_from_data_plane( ref, output_ref=ref, @@ -223,3 +227,46 @@ def test_sampler_async_data_plane_path_returns_reference_without_materializing(m assert url.endswith('/sampler/base/twinkle/sample_to_data_plane') assert body['num_samples'] == 4 assert body['group_ids'] == ['group-1'] + + +def test_old_remove_completion_does_not_clear_new_same_name_binding(): + import threading + from twinkle_client.http import ClientContext, ClientTransport + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + entered, resume = threading.Event(), threading.Event() + class Transport(ClientTransport): + def __init__(self): + super().__init__(ClientContext(base_url='http://server', api_key='owner')) + self.created = 0 + self.operations = [] + def post(self, url, json_data): + return _Response({'training_handles': True}) + def post_model(self, url, body): + self.operations.append((url.rsplit('/', 1)[-1], body.model_dump())) + if url.endswith('/add_adapter_to_model'): + self.created += 1 + return _Response(_completed({'training_id': f'training-{self.created}', 'adapter_name': 'same'})) + if url.endswith('/remove_adapter'): + entered.set() + assert resume.wait(3) + return _Response(_completed(None)) + transport = Transport() + model = MultiLoraTransformersModel('small', transport=transport) + model.add_adapter_to_model('same', None) + errors = [] + def remove(): + try:model.remove_adapter() + except Exception as exc:errors.append(exc) + releasing = threading.Thread(target=remove) + releasing.start() + try: + assert entered.wait(3) + model.add_adapter_to_model('same', None) + finally: + resume.set() + releasing.join(3) + assert not releasing.is_alive() and not errors + model.zero_grad() + assert transport.operations[-1][1]['training_id'] == 'training-2' + removal = next(body for endpoint, body in transport.operations if endpoint == 'remove_adapter') + assert removal['training_id'] == 'training-1'