Merge pull request #1 from GaoHuaZhang/mf

refactor backend cmd logits

Author: muqing-li

ais_bench/benchmark/models/local_models/base.py

@@ -35,6 +35,7 @@ class BaseModel:
     """
 
     is_api: bool = False
+    launcher: str = "torchrun"
 
     def __init__(self,
                  path: str,

ais_bench/benchmark/models/local_models/mindformers_model.py

@@ -55,6 +55,8 @@ def drop_error_generation_kwargs(generation_kwargs: dict) -> dict:
 @MODELS.register_module()
 class MindFormerModel(BaseModel):
 
+    launcher: str = "msrun"
+
     def __init__(self,
                  path: str,
                  checkpoint: Optional[str] = None,

ais_bench/benchmark/tasks/openicl_infer.py

@@ -5,14 +5,14 @@
 import threading
 import sys
 import time
-from typing import Any
+from typing import Any, Callable, Dict
 
 from mmengine.config import Config, ConfigDict
 from mmengine.utils import mkdir_or_exist
 
 from ais_bench.benchmark.registry import (ICL_INFERENCERS, ICL_RETRIEVERS, TASKS)
 from ais_bench.benchmark.tasks.base import BaseTask
-from ais_bench.benchmark.utils.config import build_dataset_from_cfg
+from ais_bench.benchmark.utils.config import build_dataset_from_cfg, get_model_cls_from_cfg
 from ais_bench.benchmark.utils.core.abbr import task_abbr_from_cfg, model_abbr_from_cfg
 from ais_bench.benchmark.tasks.base import TaskStateManager
 from ais_bench.benchmark.utils.logging import AISLogger
@@ -22,6 +22,51 @@
 from ais_bench.benchmark.openicl.icl_inferencer.icl_base_local_inferencer import BaseLocalInferencer
 
 
+def _build_msrun_cmd(ctx: Dict[str, Any]) -> str:
+    """Build msrun command from context dict."""
+    parts = [
+        "msrun",
+        f"--worker_num={ctx['worker_num']}",
+        f"--local_worker_num={ctx['local_worker_num']}",
+        f"--master_port={ctx['port']}",
+    ]
+    if ctx.get("master_addr") is not None and ctx.get("node_rank") is not None:
+        parts.extend([
+            f"--master_addr={ctx['master_addr']}",
+            f"--node_rank={ctx['node_rank']}",
+        ])
+    parts.extend([
+        "--log_dir='output/msrun_log'",
+        "--join=True",
+        "--cluster_time_out=7200",
+        ctx["script_path"],
+        ctx["cfg_path"],
+    ])
+    return " ".join(parts)
+
+
+def _build_torchrun_cmd(ctx: Dict[str, Any]) -> str:
+    """Build torchrun command from context dict."""
+    parts = [
+        "torchrun",
+        f"--master_port={ctx['port']}",
+        f"--nproc_per_node {ctx['nproc_per_node']}",
+    ]
+    if ctx.get("nnodes", 1) > 1:
+        parts.extend([
+            f"--nnodes {ctx['nnodes']}",
+            f"--node_rank {ctx['node_rank']}",
+            f"--master_addr {ctx['master_addr']}",
+        ])
+    parts.extend([ctx["script_path"], ctx["cfg_path"]])
+    return " ".join(parts)
+
+
+LAUNCHER_CMD_BUILDERS: Dict[str, Callable[[Dict[str, Any]], str]] = {
+    "msrun": _build_msrun_cmd,
+    "torchrun": _build_torchrun_cmd,
+}
+
 
 @TASKS.register_module()
 class OpenICLInferTask(BaseTask):
@@ -43,7 +88,6 @@ def __init__(self, cfg: ConfigDict):
         self.node_rank = run_cfg.get('node_rank', 0)
         self.master_addr = run_cfg.get('master_addr', "localhost")
         self.logger.debug(f"Local infer task config: {run_cfg}")
-        self.abbr = self.model_cfg.get('abbr', '')
 
     def get_command(self, cfg_path, template):
         """Get the command template for the task.
@@ -60,45 +104,36 @@ def get_command(self, cfg_path, template):
             key in str(self.model_cfg.get('type', ''))
             or key in str(self.model_cfg.get('llm', {}).get('type', ''))
             for key in backend_keys)
+        model_cls = get_model_cls_from_cfg(self.model_cfg)
+        launcher = getattr(model_cls, 'launcher', 'torchrun') if model_cls else 'torchrun'
         if self.num_gpus > 1 and not use_backend and self.nnodes == 1:
-            port = random.randint(12000, 32000)
-            if self.abbr == 'mindformer-model':
-                command = (
-                    f"msrun "
-                    f"--worker_num={self.num_gpus} "
-                    f"--local_worker_num={self.num_gpus} "
-                    f"--master_port={port} "
-                    f"--log_dir='output/msrun_log' "
-                    f"--join=True "
-                    f"--cluster_time_out=7200 "
-                    f'{script_path} {cfg_path}'
-                )
-            else :
-                command = (f'torchrun --master_port={port} '
-                           f'--nproc_per_node {self.num_procs} '
-                           f'{script_path} {cfg_path}')
+            ctx = dict(
+                port=random.randint(12000, 32000),
+                script_path=script_path,
+                cfg_path=cfg_path,
+                worker_num=self.num_gpus,
+                local_worker_num=self.num_gpus,
+                nproc_per_node=self.num_procs,
+                nnodes=1,
+                master_addr=None,
+                node_rank=None,
+            )
+            builder = LAUNCHER_CMD_BUILDERS.get(launcher, LAUNCHER_CMD_BUILDERS["torchrun"])
+            command = builder(ctx)
         elif self.nnodes > 1:
-            port = 12345
-            if self.abbr == "mindformer-model" :
-                command = (
-                    f"msrun "
-                    f"--worker_num={self.nnodes*self.num_procs} "
-                    f"--local_worker_num={self.num_procs} "
-                    f"--master_port={port} "
-                    f"--master_addr={self.master_addr} "
-                    f"--node_rank={self.node_rank} "
-                    f"--log_dir='output/msrun_log' "
-                    f"--join=True "
-                    f"--cluster_time_out=7200 "
-                    f'{script_path} {cfg_path}'
-                )
-            else :
-                command = (f'torchrun --master_port={port} '
-                       f'--nproc_per_node {self.num_procs} '
-                       f'--nnodes {self.nnodes} '
-                       f'--node_rank {self.node_rank} '
-                       f'--master_addr {self.master_addr} '
-                       f'{script_path} {cfg_path}')
+            ctx = dict(
+                port=12345,
+                script_path=script_path,
+                cfg_path=cfg_path,
+                worker_num=self.nnodes * self.num_procs,
+                local_worker_num=self.num_procs,
+                nproc_per_node=self.num_procs,
+                nnodes=self.nnodes,
+                master_addr=self.master_addr,
+                node_rank=self.node_rank,
+            )
+            builder = LAUNCHER_CMD_BUILDERS.get(launcher, LAUNCHER_CMD_BUILDERS["torchrun"])
+            command = builder(ctx)
         else:
             python = sys.executable
             command = f'{python} {script_path} {cfg_path}'

ais_bench/benchmark/utils/config/build.py

@@ -16,6 +16,33 @@
 
 logger = AISLogger()
 
+
+def get_model_type_name(model_cfg: ConfigDict) -> str:
+    """从 model_cfg 得到稳定的模型类型名（用于分支判断）。兼容 type 为类或字符串。"""
+    t = model_cfg.get("type", "")
+    if t is None or t == "":
+        return ""
+    if hasattr(t, "__name__"):
+        return t.__name__
+    if isinstance(t, str):
+        return t.split(".")[-1] if "." in t else t
+    return ""
+
+
+def get_model_cls_from_cfg(model_cfg: ConfigDict):
+    """从 model_cfg 解析出模型类，用于读取 launcher 等类属性。"""
+    t = model_cfg.get("type", None)
+    if t is None:
+        return None
+    if isinstance(t, type) and hasattr(t, "__name__"):
+        return t
+    if isinstance(t, str):
+        name = t.split(".")[-1] if "." in t else t
+        if name in MODELS.module_dict:
+            return MODELS.get(name)
+    return None
+
+
 def _validate_model_cfg(model_cfg: ConfigDict) -> dict:
     errors = {}
 
@@ -110,7 +137,7 @@ def check(condition, key, message):
                     if not validator(traffic_cfg[field]):
                         errors[f"traffic_cfg.{field}"] = error_msg
     return errors
-    
+
 
 def build_dataset_from_cfg(dataset_cfg: ConfigDict):
     logger.debug(f"Building dataset from config: type={dataset_cfg.get('type')} abbr={dataset_cfg.get('abbr')}")
@@ -123,7 +150,7 @@ def build_dataset_from_cfg(dataset_cfg: ConfigDict):
 def build_model_from_cfg(model_cfg: ConfigDict):
     logger.debug(f"Building model from config: type={model_cfg.get('type')} abbr={model_cfg.get('abbr')}")
     model_cfg = copy.deepcopy(model_cfg)
-    model_name = model_cfg.get("type", "").split(".")[-1]
+    model_name = get_model_type_name(model_cfg)
     errors = _validate_model_cfg(model_cfg)
     if errors:
         logger.warning(f"Model config validation failed for {model_name}: {errors}")
@@ -134,17 +161,14 @@ def build_model_from_cfg(model_cfg: ConfigDict):
     model_cfg.pop("run_cfg", None)
     model_cfg.pop("request_rate", None)
     batch_size = model_cfg.pop("batch_size", None)
-    abbr = model_cfg.pop("abbr", None)
-    attr = model_cfg.pop("attr", None)
+    model_cfg.pop("abbr", None)
+    model_cfg.pop("attr", None)
     model_cfg.pop("summarizer_abbr", None)
     model_cfg.pop("pred_postprocessor", None)
     model_cfg.pop("min_out_len", None)
     model_cfg.pop("returns_tool_calls", None)
     model_cfg.pop("traffic_cfg", None)
-    if attr == "local" and abbr == "mindformer-model" :
-        return MODELS.build(model_cfg, batch_size = batch_size)
-    else :
-        return MODELS.build(model_cfg)
+    return MODELS.build(model_cfg)
 
 def build_perf_metric_calculator_from_cfg(metric_cfg: ConfigDict):
     logger.debug(f"Building perf metric calculator config: type={metric_cfg.get('type')}")