Fix router (#69)

Author: tastelikefeet

ROADMAP.md

@@ -65,6 +65,7 @@
 - [ ] 支持DPO对齐训练
 - [ ] 支持colocate RL训练
 - [ ] Preprocess支持batched
+- [ ] 对多replica的支持和粘滞路由
 
 ### 网络能力
 
@@ -84,5 +85,6 @@
 - [ ] Support for DPO alignment training
 - [ ] Support for colocate RL training
 - [ ] Support for batched preprocessing
+- [ ] Support for multiple replicas and sticky routing
 
 ### Networking Capabilities

cookbook/client/tinker/megatron/server_config.yaml

@@ -33,7 +33,6 @@ applications:
           runtime_env:
             env_vars:
               TWINKLE_TRUST_REMOTE_CODE: "0"
-              DEVICE_COUNT_PER_PHYSICAL_NODE: "8"
 
   # 3. Sampler Service - Runs inference / sampling using vLLM engine
   #    Used for generating text from the model (e.g., evaluating LoRA results).
@@ -52,7 +51,7 @@ applications:
       device_group:                   # Logical device group for the sampler
         name: sampler
         gpus_per_worker: 1
-        ranks: [0,1,2,3]                    # GPU rank indices to use
+        ranks: 4                   # GPU rank indices to use
         device_type: cuda
       device_mesh:
         device_type: cuda
@@ -71,7 +70,6 @@ applications:
           runtime_env:
             env_vars:
               TWINKLE_TRUST_REMOTE_CODE: "0"
-              DEVICE_COUNT_PER_PHYSICAL_NODE: "8"
 
   # 2. Model Service (commented out) - Would host the base model for training.
   #    Uncomment and configure if you need a training model worker.
@@ -86,7 +84,7 @@ applications:
       nproc_per_node: 4                           # Number of GPU processes per node
       device_group:
         name: model
-        ranks: [4,5,6,7]                              # GPU rank indices
+        ranks: 4       # GPU rank indices
         device_type: cuda
       device_mesh:
         device_type: cuda
@@ -111,4 +109,3 @@ applications:
           runtime_env:
             env_vars:
               TWINKLE_TRUST_REMOTE_CODE: "0"
-              DEVICE_COUNT_PER_PHYSICAL_NODE: "8"

cookbook/client/tinker/megatron/server_config_7b.yaml

@@ -67,7 +67,6 @@ applications:
           runtime_env:
             env_vars:
               TWINKLE_TRUST_REMOTE_CODE: "0"
-              DEVICE_COUNT_PER_PHYSICAL_NODE: "8"
 
   # 3. Sampler Service - Runs inference / sampling using vLLM engine
   #    Used for generating text from the model (e.g., evaluating LoRA results).
@@ -104,4 +103,3 @@ applications:
           runtime_env:
             env_vars:
               TWINKLE_TRUST_REMOTE_CODE: "0"
-              DEVICE_COUNT_PER_PHYSICAL_NODE: "8"

cookbook/client/tinker/self_congnition.py

@@ -44,7 +44,7 @@ def train():
 
     # Connect to the Twinkle server running locally
     service_client = init_tinker_compat_client(
-        base_url='http://www.modelscope.cn/twinkle', api_key=os.environ.get('MODELSCOPE_TOKEN'))
+        base_url='localhost:9000', api_key=os.environ.get('MODELSCOPE_TOKEN'))
 
     # Create a LoRA training client for the base model (rank=16 for the LoRA adapter)
     training_client = service_client.create_lora_training_client(base_model=base_model, rank=16)
@@ -68,9 +68,10 @@ def train():
             optim_result = optim_future.result()
 
             # Compute weighted average log-loss per token for monitoring
-            logprobs = np.concatenate([output['logprobs'].tolist() for output in fwdbwd_result.loss_fn_outputs])
-            weights = np.concatenate([example.loss_fn_inputs['weights'].tolist() for example in input_datum])
-            print(f'Loss per token: {-np.dot(logprobs, weights) / weights.sum():.4f}')
+            # logprobs = np.concatenate([output['logprobs'].tolist() for output in fwdbwd_result.loss_fn_outputs])
+            # weights = np.concatenate([example.loss_fn_inputs['weights'].tolist() for example in input_datum])
+            # print(f'Loss per token: {-np.dot(logprobs, weights) / weights.sum():.4f}')
+            print(f'Training Metrics: {optim_result}')
 
         # Save a checkpoint after each epoch
         save_future = training_client.save_state(f'twinkle-lora-{epoch}')
@@ -85,7 +86,7 @@ def eval():
     weight_path = 'twinkle://20260212_174205-Qwen_Qwen2_5-7B-Instruct-51edc9ed/weights/twinkle-lora-2'
 
     # Connect to the server and create a sampling client with the trained weights
-    service_client = init_tinker_compat_client(base_url='http://localhost:8000')
+    service_client = init_tinker_compat_client(base_url='http://localhost:9000')
     sampling_client = service_client.create_sampling_client(model_path=weight_path, base_model=base_model)
 
     # Step 2: Prepare the chat prompt

cookbook/client/tinker/transformer/server_config.yaml

@@ -65,7 +65,6 @@ applications:
           runtime_env:
             env_vars:
               TWINKLE_TRUST_REMOTE_CODE: "0"
-              DEVICE_COUNT_PER_PHYSICAL_NODE: "8"
 
   # 3. Sampler Service - Runs inference / sampling using vLLM engine
   #    Used for generating text from the model (e.g., evaluating LoRA results).
@@ -102,4 +101,3 @@ applications:
           runtime_env:
             env_vars:
               TWINKLE_TRUST_REMOTE_CODE: "0"
-              DEVICE_COUNT_PER_PHYSICAL_NODE: "8"

cookbook/client/twinkle/transformer/server_config.yaml

@@ -61,7 +61,6 @@ applications:
           runtime_env:
             env_vars:
               TWINKLE_TRUST_REMOTE_CODE: "0"
-              DEVICE_COUNT_PER_PHYSICAL_NODE: "8"
 
   # 3. Processor Service - Handles data preprocessing on CPU
   #    Runs tokenization, template application, and other CPU-bound tasks.
@@ -90,7 +89,6 @@ applications:
           runtime_env:
             env_vars:
               TWINKLE_TRUST_REMOTE_CODE: "0"
-              DEVICE_COUNT_PER_PHYSICAL_NODE: "8"
 
   # 4. Sampler Service - Handles text generation inference
   #    Uses vLLM for efficient batched generation with optional LoRA adapters.
@@ -125,4 +123,3 @@ applications:
           runtime_env:
             env_vars:
               TWINKLE_TRUST_REMOTE_CODE: "0"
-              DEVICE_COUNT_PER_PHYSICAL_NODE: "8"

docs/source_en/Usage Guide/Server and Client/Server.md

@@ -55,12 +55,9 @@ This configuration starts 3 nodes:
 Before starting the Server, you need to set the following environment variables:
 
 ```bash
-export DEVICE_COUNT_PER_PHYSICAL_NODE=8  # Specify the total number of GPUs on each physical machine
 export TWINKLE_TRUST_REMOTE_CODE=0       # Whether to trust remote code (security consideration)
 ```
 
-> **Important Note**: `DEVICE_COUNT_PER_PHYSICAL_NODE` must be set to the actual number of physical GPUs on the machine, which is crucial for correctly parsing the `ranks` configuration.
-
 ### Node Rank in YAML Configuration
 
 In the YAML configuration file, **each component needs to occupy a separate Node**.
@@ -117,7 +114,6 @@ applications:
 **Important notes:**
 - The `ranks` configuration uses **physical GPU card numbers**, directly corresponding to the actual GPU devices on the machine
 - The `device_mesh` configuration uses parameters like `dp_size`, `tp_size`, `pp_size`, `ep_size` instead of the original `mesh` and `mesh_dim_names`
-- The environment variable `DEVICE_COUNT_PER_PHYSICAL_NODE` must be set to inform the system of the total number of physical GPUs on each machine
 - Different components will be automatically assigned to different Nodes
 - Ray will automatically schedule to the appropriate Node based on resource requirements (`num_gpus`, `num_cpus` in `ray_actor_options`)
 
@@ -393,7 +389,6 @@ applications:
           num_cpus: 0.1
           runtime_env:
             env_vars:
-              DEVICE_COUNT_PER_PHYSICAL_NODE: "8"  # Total number of physical GPUs on each machine
 
   # 3. Sampler service (optional, for inference sampling)
   - name: sampler-Qwen2.5-0.5B-Instruct
@@ -425,7 +420,6 @@ applications:
           num_gpus: 1                 # Sampler needs independent GPU
           runtime_env:
             env_vars:
-              DEVICE_COUNT_PER_PHYSICAL_NODE: "8"  # Total number of physical GPUs on each machine
 ```
 
 ## Configuration Item Description
@@ -471,6 +465,5 @@ device_mesh:
 **Environment variables:**
 
 ```bash
-export DEVICE_COUNT_PER_PHYSICAL_NODE=8  # Total number of GPUs on each physical machine (must be set)
 export TWINKLE_TRUST_REMOTE_CODE=0       # Whether to trust remote code
 ```

docs/source_en/Usage Guide/Server and Client/Tinker-Compatible-Client.md

@@ -25,7 +25,7 @@ for item in service_client.get_server_capabilities().supported_models:
 When calling `init_tinker_compat_client`, the following operations are automatically executed:
 
 1. **Patch Tinker SDK**: Bypass Tinker's `tinker://` prefix validation, allowing it to connect to standard HTTP addresses
-2. **Set Request Headers**: Inject necessary authentication headers such as `X-Ray-Serve-Request-Id` and `Authorization`
+2. **Set Request Headers**: Inject necessary authentication headers such as `serve_multiplexed_model_id` and `Authorization`
 3. **Return `ServiceClient`**: Returns a standard Tinker `ServiceClient` object, subsequent operations are completely identical to native Tinker
 
 This means that after initialization, **all existing Tinker training code can be used directly** without any modifications.

docs/source_zh/使用指引/服务端和客户端/Tinker兼容客户端.md

@@ -25,7 +25,7 @@ for item in service_client.get_server_capabilities().supported_models:
 调用 `init_tinker_compat_client` 时，会自动执行以下操作：
 
 1. **Patch Tinker SDK**：绕过 Tinker 的 `tinker://` 前缀校验，使其可以连接到标准 HTTP 地址
-2. **设置请求头**：注入 `X-Ray-Serve-Request-Id` 和 `Authorization` 等必要的认证头
+2. **设置请求头**：注入 `serve_multiplexed_model_id` 和 `Authorization` 等必要的认证头
 3. **返回 `ServiceClient`**：返回一个标准的 Tinker `ServiceClient` 对象，后续操作与原生 Tinker 完全一致
 
 这意味着在初始化之后，**所有已有的 Tinker 训练代码都可以直接使用**，无需任何修改。

docs/source_zh/使用指引/服务端和客户端/服务端.md

@@ -55,12 +55,9 @@ ray start --address=10.28.252.9:6379 --num-gpus=0
 在启动 Server 之前，需要设置以下环境变量：
 
 ```bash
-export DEVICE_COUNT_PER_PHYSICAL_NODE=8  # 指定每台物理机上的 GPU 总数
 export TWINKLE_TRUST_REMOTE_CODE=0       # 是否信任远程代码（安全考虑）
 ```
 
-> **重要提示**：`DEVICE_COUNT_PER_PHYSICAL_NODE` 必须设置为机器上实际的物理 GPU 数量，这对于正确解析 `ranks` 配置至关重要。
-
 ### YAML 配置中的 Node Rank
 
 在 YAML 配置文件中，**每个组件需要占用一个独立的 Node**。
@@ -117,7 +114,6 @@ applications:
 **重要提示：**
 - `ranks` 配置使用**物理 GPU 卡号**，直接对应机器上的实际 GPU 设备
 - `device_mesh` 配置使用 `dp_size`、`tp_size`、`pp_size`、`ep_size` 等参数替代原来的 `mesh` 和 `mesh_dim_names`
-- 必须设置环境变量 `DEVICE_COUNT_PER_PHYSICAL_NODE` 来告知系统每台机器的物理 GPU 总数
 - 不同组件会自动分配到不同的 Node 上
 - Ray 会根据资源需求（`ray_actor_options` 中的 `num_gpus`、`num_cpus`）自动调度到合适的 Node
 
@@ -336,7 +332,6 @@ applications:
           num_cpus: 0.1
           runtime_env:
             env_vars:
-              DEVICE_COUNT_PER_PHYSICAL_NODE: "8"  # 每台机器的物理 GPU 总数
 
   # 3. Sampler 服务（可选，用于推理采样）
   - name: sampler-Qwen2.5-0.5B-Instruct
@@ -368,7 +363,6 @@ applications:
           num_gpus: 1                 # Sampler 需要独立 GPU
           runtime_env:
             env_vars:
-              DEVICE_COUNT_PER_PHYSICAL_NODE: "8"  # 每台机器的物理 GPU 总数
 ```
 
 ## 配置项说明
@@ -414,6 +408,5 @@ device_mesh:
 **环境变量：**
 
 ```bash
-export DEVICE_COUNT_PER_PHYSICAL_NODE=8  # 每台物理机上的 GPU 总数（必须设置）
 export TWINKLE_TRUST_REMOTE_CODE=0       # 是否信任远程代码
 ```