feat: add on-demand models with shared LoRA and exclusive training

.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

.gitignore

@@ -77,6 +77,10 @@ instance/
 # Sphinx documentation
 docs/_build/
 
+# Internal design notes and developer tools are local-only.
+/docs/design/
+/tools/
+
 # PyBuilder
 target/
 

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

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.

docs/source_zh/使用指引/服务端和客户端/服务端.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`。

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",

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__)

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,14 +425,16 @@ 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,
                 }
 
                 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

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)

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()
 

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."""

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