mindstudio-probe 8.3.2__py3-none-any.whl → 26.0.0a1__py3-none-any.whl

msprobe/core/monitor_v2/trainer.py

@@ -0,0 +1,326 @@
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
+#
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
+#
+#          http://license.coscl.org.cn/MulanPSL2
+#
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+from msprobe.core.common.framework_adapter import FmkAdp
+from msprobe.core.common.file_utils import load_json
+from msprobe.core.common.log import logger
+from msprobe.core.monitor.utils import validate_int_arg, validate_step_count_per_record, validate_ranks
+from msprobe.core.monitor_v2.base import BaseMonitorV2
+from msprobe.core.monitor_v2.factory import MonitorFactory
+from msprobe.core.monitor_v2.writer import CSVWriterV2
+
+
+_FRAMEWORK_ALIASES = {
+    "pytorch": "pytorch",
+    "pt": "pytorch",
+    "torch": "pytorch",
+    "mindspore": "mindspore",
+    "ms": "mindspore",
+}
+
+
+def _resolve_framework(config: Dict[str, Any], fr: Optional[str]) -> str:
+    candidate = fr or config.get("framework", "pytorch")
+    key = str(candidate).strip().lower()
+    if key not in _FRAMEWORK_ALIASES:
+        raise ValueError(f"[monitor_v2] Unsupported framework '{candidate}'")
+    return _FRAMEWORK_ALIASES[key]
+
+
+class TrainerMonitorV2:
+    """
+    Thin orchestrator for monitor_v2.
+
+    Usage (PyTorch):
+        mon = TrainerMonitorV2("config.json", fr="pytorch")
+        mon.start(model=model)
+        for step in range(num_steps):
+            ...
+            mon.step()
+        mon.stop()
+
+    Config (minimal example):
+        {
+          "framework": "pytorch",
+          "output_dir": "./monitor_v2_out",
+          "rank": [0, 1],                # optional, list of ranks to monitor
+          "start_step": 0,               # optional, first step to write
+          "step_interval": 1,            # optional, write every N steps
+          "step_count_per_record": 1,    # optional, steps per output file
+          "monitors": {
+            "module": {
+              "enabled": true,
+              "targets": [],
+              "ops": ["min", "max", "mean", "norm", "nans"]
+            }
+          }
+        }
+    """
+
+    def __init__(self, config_path: str, fr: Optional[str] = None) -> None:
+        self.config_path = config_path
+        self.config: Dict[str, Any] = load_json(config_path)
+        if not isinstance(self.config, dict):
+            raise TypeError("[monitor_v2] Configuration must be a dictionary.")
+
+        self.framework = _resolve_framework(self.config, fr)
+
+        # Ensure framework adapter is initialised so that rank resolution and
+        # tensor ops work correctly for the selected framework.
+        try:
+            FmkAdp.set_fmk(self.framework)
+        except Exception as exc:
+            # Best-effort; fall back to default adapter behaviour.
+            logger.warning(f"[monitor_v2] Failed to set framework adapter to '{self.framework}': {exc}")
+
+        # Current rank is always resolved from framework adapter.
+        self.rank = FmkAdp.get_rank_id()
+
+        # `rank` in config is a list (or single int) of ranks to monitor.
+        self._target_ranks = self._parse_target_ranks(self.config.get("rank", []))
+
+        # Step control: start_step / stop_step define active window,
+        # step_interval controls how often we write.
+        # step_count_per_record controls how many steps go into one output file.
+        self.start_step = validate_int_arg(self.config.get("start_step"), "start_step", 0, 0)
+        self.step_interval = validate_int_arg(self.config.get("step_interval"), "step_interval", 1, 1)
+        try:
+            validate_step_count_per_record(self.config.get("step_count_per_record", 1))
+            self.step_count_per_record = int(self.config.get("step_count_per_record", 1))
+        except Exception as exc:
+            logger.warning(f"[monitor_v2] Validate step_count_per_record failed: {exc}; fallback to 1.")
+            self.step_count_per_record = 1
+        # collect_times: 最大采集次数，达到后停止写出；默认大值表示“几乎一直采集”
+        # stop_step：优先使用显式配置，其次由 collect_times 推导，最后使用一个极大默认值
+        self.collect_times = validate_int_arg(self.config.get("collect_times"), "collect_times", 1, 100_000_000)
+        explicit_stop = validate_int_arg(self.config.get("stop_step"), "stop_step", self.start_step, None)
+        if explicit_stop is not None:
+            self.stop_step = int(explicit_stop)
+        else:
+            # start_step + collect_times * interval 对齐“采集次数”语义
+            self.stop_step = self.start_step + self.collect_times * max(self.step_interval, 1)
+        self.current_step = 0
+        self._collected_steps = 0
+        self.output_dir = self.config.get("output_dir", "./")
+
+        self.monitors = self._build_monitors()
+        fmt = str(self.config.get("format", "csv")).strip().lower()
+        if fmt != "csv":
+            raise ValueError(f"[monitor_v2] Unsupported format '{fmt}', only 'csv' is currently supported.")
+        self.writer = CSVWriterV2(
+            self.output_dir,
+            rank=self.rank,
+            async_write=bool(self.config.get("async_write", False)),
+        )
+        self._writer_closed = False
+        self._optimizer = None
+        self._orig_step = None
+        self._step_patched = False
+        self._patch_optimizer_step = bool(self.config.get("patch_optimizer_step", False))
+        self._warned_manual_step_with_optimizer = False
+
+    def start(self, model: Any = None, optimizer: Any = None, **context: Any) -> None:
+        """
+        Start all configured monitors.
+        """
+        if not self._is_target_rank():
+            return
+        static_ctx = {
+            "rank": self.rank,
+            "output_dir": self.output_dir,
+            "step_provider": lambda: self.current_step,
+        }
+        static_ctx.update(context)
+        for mon in self.monitors:
+            mon.start(model=model, optimizer=optimizer, **static_ctx)
+        if optimizer is not None:
+            self._optimizer = optimizer
+            if "patch_optimizer_step" not in self.config:
+                self._patch_optimizer_step = True
+        if self._optimizer is not None and self._patch_optimizer_step:
+            self._patch_step_via_optimizer()
+
+    def step(self, from_optimizer: bool = False) -> None:
+        """
+        Increment step counter and collect/write data from monitors.
+        """
+        if not self._is_target_rank():
+            return
+        if self._patch_optimizer_step and self._optimizer is not None and not from_optimizer:
+            if not self._warned_manual_step_with_optimizer:
+                logger.warning(
+                    "[monitor_v2] step() called manually while optimizer patch is enabled; "
+                    "skip to avoid double increment. Set patch_optimizer_step=false to "
+                    "manage steps manually."
+                )
+                self._warned_manual_step_with_optimizer = True
+            return
+        step_val = self.current_step
+        self.current_step += 1
+
+        # step_interval 语义：每隔多少个 step 写一次。
+        # step_count_per_record 语义：每多少个 step 写到一个 csv 文件。
+        if not self._should_collect_step(step_val):
+            for mon in self.monitors:
+                collect_fn = getattr(mon, "collect", None)
+                if not callable(collect_fn):
+                    continue
+                try:
+                    collect_fn()
+                except Exception as exc:
+                    # 收集失败不影响训练流程，直接跳过该 monitor
+                    logger.warning(f"[monitor_v2] Failed to collect data '{mon.__class__.__name__}': {exc}")
+            return
+
+        for mon in self.monitors:
+            collect_fn = getattr(mon, "collect", None)
+            if not callable(collect_fn):
+                continue
+            out = collect_fn()
+            if not out:
+                continue
+            # Normalise monitor payload so that individual monitors
+            # only need to care about producing "rows". Trainer is
+            # responsible for adding common metadata.
+            if not isinstance(out, dict):
+                continue
+            if not out.get("rows"):
+                continue
+            if "slug" not in out:
+                slug = getattr(mon, "slug", None)
+                if not isinstance(slug, str):
+                    slug = mon.__class__.__name__.lower()
+                out["slug"] = slug
+            if "monitor" not in out:
+                out["monitor"] = out.get("slug")
+            if "rank" not in out:
+                out["rank"] = self.rank
+            if "step" not in out:
+                out["step"] = step_val
+            if "step_interval" not in out:
+                out["step_interval"] = self.step_interval
+            if "step_count_per_record" not in out:
+                out["step_count_per_record"] = self.step_count_per_record
+            if "start_step" not in out:
+                out["start_step"] = self.start_step
+            self.writer.write_monitor_data(out)
+        # 仅在实际写出数据时增加采集次数
+        self._collected_steps += 1
+
+    def stop(self) -> None:
+        """
+        Stop all monitors and close writer.
+        """
+        if not self._is_target_rank():
+            return
+        for mon in self.monitors:
+            mon.stop()
+        self._restore_optimizer_step()
+        self._close_writer()
+
+    def _build_monitors(self):
+        monitors_cfg = self.config.get("monitors", {})
+        if not isinstance(monitors_cfg, dict):
+            logger.warning("[monitor_v2] Monitors config should be a dict; using empty config.")
+            monitors_cfg = {}
+        monitors = []
+
+        for name, cfg in monitors_cfg.items():
+            if not isinstance(cfg, dict) or not cfg.get("enabled", name == "module"):
+                continue
+            mon = MonitorFactory.create(self.framework, name)
+            if mon is None:
+                continue
+            mon_cfg = {k: v for k, v in cfg.items() if k != "enabled"}
+            setattr(mon, "slug", name)
+            mon.set_config(mon_cfg)
+            monitors.append(mon)
+
+        return monitors
+
+    def _parse_target_ranks(self, rank_cfg: Any) -> set[int]:
+        if isinstance(rank_cfg, int):
+            return {rank_cfg}
+        if isinstance(rank_cfg, list):
+            try:
+                validate_ranks(rank_cfg)
+                return {int(r) for r in rank_cfg}
+            except Exception as exc:
+                logger.warning(f"[monitor_v2] Validate rank list failed: {exc}; fallback to all ranks.")
+                return set()
+        logger.warning("[monitor_v2] Rank config should be int or list; fallback to all ranks.")
+        return set()
+
+    def _should_collect_step(self, step: int) -> bool:
+        """
+        Global step filter for monitor_v2:
+        - Only write when start_step <= step < stop_step
+        - And (step - start_step) % step_interval == 0 when interval > 1
+        - Respect collect_times: stop after N successful writes
+        """
+        if self.collect_times > 0 and self._collected_steps >= self.collect_times:
+            return False
+        if not (self.start_step <= step < self.stop_step):
+            return False
+        if self.step_interval <= 1:
+            return True
+        return (step - self.start_step) % self.step_interval == 0
+
+    def _is_target_rank(self) -> bool:
+        """
+        Whether current rank should participate in monitor_v2 collection.
+        Empty `rank` list in config means all ranks.
+        """
+        return not self._target_ranks or self.rank in self._target_ranks
+
+    def _close_writer(self) -> None:
+        if self._writer_closed:
+            return
+        try:
+            self.writer.close()
+        finally:
+            self._writer_closed = True
+
+    def _patch_step_via_optimizer(self) -> None:
+        if self._step_patched or self._optimizer is None:
+            return
+        step_fn = getattr(self._optimizer, "step", None)
+        if not callable(step_fn):
+            return
+        self._orig_step = step_fn
+
+        def wrapped_step(*args: Any, **kwargs: Any):
+            out = self._orig_step(*args, **kwargs)
+            try:
+                self.step(from_optimizer=True)
+            except Exception as exc:
+                logger.warning(f"[monitor_v2] Failed to run step from optimizer patch: {exc}")
+            return out
+
+        self._optimizer.step = wrapped_step  # type: ignore[assignment]
+        self._step_patched = True
+
+    def _restore_optimizer_step(self) -> None:
+        if not self._step_patched or self._optimizer is None:
+            return
+        if self._orig_step is not None:
+            self._optimizer.step = self._orig_step  # type: ignore[assignment]
+        self._step_patched = False
+        self._orig_step = None

msprobe/core/monitor_v2/utils.py

@@ -0,0 +1,122 @@
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
+#
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
+#
+#          http://license.coscl.org.cn/MulanPSL2
+#
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
+
+from __future__ import annotations
+
+from typing import Any, Dict, Iterable, Optional, Tuple
+
+from msprobe.core.common.const import MonitorConst
+from msprobe.core.common.framework_adapter import FmkAdp
+from msprobe.core.common.log import logger
+
+
+def iter_model_chunks(model: Any) -> Iterable[Tuple[Any, str]]:
+    """
+    Iterate over model or model_list, yielding (model_chunk, prefix).
+
+    - Single model: yields (model, "")
+    - List/tuple: yields (sub_model, "<idx><NAME_SEP>")
+    """
+    if model is None:
+        # 保持与后续分支一致的生成器行为
+        return
+        # 不再显式 return ()
+
+    if isinstance(model, (list, tuple)):
+        for idx, sub in enumerate(model):
+            if sub is None:
+                continue
+            prefix = f"{idx}{MonitorConst.NAME_SEP}"
+            yield sub, prefix
+        return
+        # 同样去掉显式的 return ()
+
+    yield model, ""
+
+
+def iter_model_params(model: Any) -> Iterable[Tuple[str, Any]]:
+    """
+    Iterate (name, param) pairs for a model that may be:
+      - a single nn.Module/nn.Cell
+      - a list/tuple of models
+
+    For list/tuple, each sub-model is prefixed with
+    "<idx><NAME_SEP>" to keep names unique and consistent
+    with other multi-model monitor behaviours.
+    """
+    for chunk, prefix in iter_model_chunks(model):
+        try:
+            for name, param in FmkAdp.named_parameters(chunk):
+                yield prefix + name, param
+        except Exception:
+            # If a sub-element is not a valid module, skip it.
+            continue
+
+
+def iter_model_modules(model: Any) -> Iterable[Tuple[str, Any]]:
+    """
+    Iterate (name, module) pairs for a model or model_list, using
+    framework-specific module traversal, with optional vpp_stage prefix.
+    """
+    for chunk, prefix in iter_model_chunks(model):
+        try:
+            for name, module in FmkAdp.iter_named_modules(chunk):
+                yield prefix + name, module
+        except Exception:
+            logger.warning("Failed to iterate modules for model chunk, skipping.")
+            continue
+
+
+def build_param2name(model: Any = None, optimizer: Any = None) -> Dict[Any, str]:
+    """
+    Build a param->name mapping from model/optimizer.
+
+    - Primary source: model (or list/tuple of models), using framework
+      specific named_parameters().
+    - Fallback: optimizer.param_groups when model mapping is empty.
+    """
+    param2name: Dict[Any, str] = {}
+
+    for name, param in iter_model_params(model):
+        param2name[param] = name
+
+    if param2name or optimizer is None:
+        return param2name
+
+    idx = 0
+    for group in getattr(optimizer, "param_groups", []):
+        for param in group.get("params", []):
+            if param not in param2name:
+                param2name[param] = f"param_{idx}"
+                idx += 1
+
+    return param2name
+
+
+def get_vpp_stage_from_tag(tag: str) -> Optional[int]:
+    """
+    Extract vpp_stage index from a tag or name that may start with
+    "<idx><NAME_SEP>". Returns None if no valid stage prefix exists.
+    """
+    if not tag:
+        return None
+    prefix, sep, _ = tag.partition(MonitorConst.NAME_SEP)
+    if not sep:
+        return None
+    try:
+        return int(prefix)
+    except (TypeError, ValueError):
+        return None