diff --git a/swanlab/data/run/main.py b/swanlab/data/run/main.py index c7d118b9f..519559bad 100644 --- a/swanlab/data/run/main.py +++ b/swanlab/data/run/main.py @@ -8,6 +8,7 @@ 在此处定义SwanLabRun类并导出 """ import os +import sys from typing import Any, Dict, Optional, List, Tuple from swanlab.core_python import timer @@ -24,9 +25,25 @@ from .public import SwanLabPublicConfig from ..store import get_run_store, reset_run_store +# Windows 不支持 fcntl，使用 msvcrt +if sys.platform == 'win32': + import msvcrt +else: + import fcntl + MAX_LIST_LENGTH = 108 +def _get_monitor_lock_path(run_dir: str, swanlog_dir: Optional[str], run_id: Optional[str], parallel: str) -> str: + """ + 计算硬件监控锁文件路径。 + shared 模式下需要使用稳定路径，确保多个进程竞争同一把锁。 + """ + if parallel == "shared" and swanlog_dir is not None and run_id is not None: + return os.path.join(swanlog_dir, f".hardware_monitor.{run_id}.lock") + return os.path.join(run_dir, ".hardware_monitor.lock") + + class SwanLabRun: """ The SwanLabRun class is used for logging during a single experiment. @@ -64,12 +81,16 @@ def __init__( self.__operator = operator self.__state = SwanLabRunState.RUNNING self.__monitor_timer: Optional[timer.Timer] = None + self.__monitor_lock_path: Optional[str] = None + self.__monitor_lock_file = None self.__config: Optional[SwanLabConfig] = None # 1. 设置常规参数 self.__mode = get_mode() self.__public = SwanLabPublicConfig() self.__operator.before_run(None) # 2. 初始化配置 + # 全局 config 对象会跨 run 复用，这里先清理运行态，避免旧 setter 在 on_run 之前被误触发。 + config.clean() if run_store.config is not None: config.update(run_store.config) config.update(run_config) @@ -96,37 +117,66 @@ def __init__( operator.on_runtime_info_update(RuntimeInfo(requirements=requirements, conda=conda, metadata=metadata)) # 定时采集系统信息 # 测试时不开启此功能 - # resume时不开启此功能 - if "PYTEST_VERSION" not in os.environ and run_store.resume == 'never': + # 使用文件锁确保多进程环境下只有一个进程采集硬件数据 + if "PYTEST_VERSION" not in os.environ: if monitor_funcs is not None and len(monitor_funcs) != 0: - swanlog.debug("Monitor on.") - - # 定义定时任务函数 - def monitor_func(): - monitor_info_list = [f() for f in monitor_funcs] - # 剔除其中为None的数据 - for monitor_info in monitor_info_list: - if monitor_info is None: - swanlog.debug("Hardware info is empty. Skip it.") - continue - for info in monitor_info: - key, name, value, cfg = ( - info['key'], - info['name'], - info['value'], - info['config'], - ) - v = DataWrapper(key, [Line(value)], reference="TIME") - self.__exp.add( - data=v, - key=key, - name=name, - column_config=cfg, - column_class="SYSTEM", - section_type="SYSTEM", - ) - - self.__monitor_timer = timer.Timer(monitor_func, interval=monitor_interval, immediate=True).run() + # 尝试获取硬件监控锁 + lock_path = _get_monitor_lock_path( + run_dir=run_store.run_dir, + swanlog_dir=run_store.swanlog_dir, + run_id=run_store.run_id, + parallel=run_store.parallel, + ) + lock_file = None + try: + lock_file = open(lock_path, "w") + if sys.platform == 'win32': + # Windows: 使用 msvcrt.locking + msvcrt.locking(lock_file.fileno(), msvcrt.LK_NBLCK, 1) + else: + # Unix/Linux/macOS: 使用 fcntl.flock + fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + self.__monitor_lock_file = lock_file + self.__monitor_lock_path = lock_path + swanlog.debug("Hardware monitor lock acquired, starting monitor.") + except (IOError, OSError): + # 其他进程已持有锁，跳过硬件监控 + swanlog.debug("Hardware monitor already running in another process, skipping.") + try: + lock_file is not None and lock_file.close() + except Exception: # noqa + pass + monitor_funcs = None + + if self.__monitor_lock_file is not None and monitor_funcs is not None and len(monitor_funcs) != 0: + swanlog.debug("Monitor on.") + + # 定义定时任务函数 + def monitor_func(): + monitor_info_list = [f() for f in monitor_funcs] + # 剔除其中为None的数据 + for monitor_info in monitor_info_list: + if monitor_info is None: + swanlog.debug("Hardware info is empty. Skip it.") + continue + for info in monitor_info: + key, name, value, cfg = ( + info['key'], + info['name'], + info['value'], + info['config'], + ) + v = DataWrapper(key, [Line(value)], reference="TIME") + self.__exp.add( + data=v, + key=key, + name=name, + column_config=cfg, + column_class="SYSTEM", + section_type="SYSTEM", + ) + + self.__monitor_timer = timer.Timer(monitor_func, interval=monitor_interval, immediate=True).run() def __cleanup(self, error: str = None, interrupt: bool = False): """ @@ -136,7 +186,19 @@ def __cleanup(self, error: str = None, interrupt: bool = False): if self.__monitor_timer is not None: self.__monitor_timer.cancel() self.__monitor_timer.join() - # 2. 更新状态 + # 2. 释放硬件监控锁 + if self.__monitor_lock_file is not None: + try: + if sys.platform == 'win32': + # Windows: 使用 msvcrt.locking 解锁 + msvcrt.locking(self.__monitor_lock_file.fileno(), msvcrt.LK_UNLCK, 1) + else: + # Unix/Linux/macOS: 使用 fcntl.flock 解锁 + fcntl.flock(self.__monitor_lock_file.fileno(), fcntl.LOCK_UN) + self.__monitor_lock_file.close() + except Exception: # noqa + pass + # 3. 更新状态 self.__state = SwanLabRunState.SUCCESS if error is None else SwanLabRunState.CRASHED # 3. 触发回调 if get_settings().log_proxy_type not in ['stderr', 'all']: diff --git a/swanlab/data/run/metadata/hardware/gpu/nvidia.py b/swanlab/data/run/metadata/hardware/gpu/nvidia.py index 3a41113c1..0fb75cfc4 100644 --- a/swanlab/data/run/metadata/hardware/gpu/nvidia.py +++ b/swanlab/data/run/metadata/hardware/gpu/nvidia.py @@ -75,7 +75,10 @@ def get_nvidia_gpu_info() -> HardwareFuncResult: except pynvml.NVMLError: pass finally: - pynvml.nvmlShutdown() + try: + pynvml.nvmlShutdown() + except Exception: + pass count = info["cores"] return info, None if not count else GpuCollector(count=count, max_mem_mb=max_gpu_mem_mb) @@ -264,7 +267,10 @@ def collect(self) -> HardwareInfoList: return result def __del__(self): - pynvml.nvmlShutdown() + try: + pynvml.nvmlShutdown() + except Exception: + pass def before_collect_impl(self): # 低频采集下（30s以下），应该每次采集时都执行pynvml.nvmlInit() diff --git a/swanlab/data/sdk.py b/swanlab/data/sdk.py index 6b65de4c2..9ff5c89ad 100644 --- a/swanlab/data/sdk.py +++ b/swanlab/data/sdk.py @@ -327,7 +327,8 @@ def init( callbacks = check_callback_format(self.cbs + callbacks) self.cbs = [] # 4. 校验并行模式 - if str(parallel).lower() in ["shared", "true", "yes"]: + parallel_mode = "shared" if str(parallel).lower() in ["shared", "true", "yes"] else "none" + if parallel_mode == "shared": resume = "allow" mode = "cloud" id = id or secrets.token_hex(4) @@ -375,6 +376,7 @@ def init( # ---------------------------------- 设置运行时配置 ---------------------------------- # 1. 写入运行时配置 run_store.resume = resume + run_store.parallel = parallel_mode run_store.run_id = id run_store.project = project run_store.workspace = workspace @@ -413,7 +415,7 @@ def init( run_name = "run-{}-{}".format(timestamp, run_id) run_dir = os.path.join(logdir, run_name) try: - os.mkdir(run_dir) + os.makedirs(run_dir, exist_ok=False) break except FileExistsError: pass @@ -427,9 +429,9 @@ def init( # ---------------------------------- 初始化运行实例 ---------------------------------- # 系统信息检测 meta, monitor_funcs = None, None - # 新实验开启系统信息检测，旧实验暂时不开启 + # 新实验或共享模式(parallel="shared")开启系统信息检测 # 并且只在用户设置开启了元数据收集或硬件监控时才开启 - if run_store.new is True and user_settings.metadata_collect: + if (run_store.new is True or run_store.parallel == "shared") and user_settings.metadata_collect: meta, monitor_funcs = get_metadata(run_store.run_dir) run = SwanLabRun(run_config=config, operator=operator, metadata=meta, monitor_funcs=monitor_funcs) return run diff --git a/swanlab/data/store.py b/swanlab/data/store.py index 62426dc42..5aca0649a 100644 --- a/swanlab/data/store.py +++ b/swanlab/data/store.py @@ -29,6 +29,8 @@ class RunStore(BaseModel): # ---------------------------------- 实验 ---------------------------------- # 实验模式 resume: Literal['must', 'allow', 'never'] = 'never' + # 并行模式，用于区分普通 resume 和 parallel="shared" + parallel: Literal['none', 'shared'] = 'none' # 实验名称 run_name: Optional[str] = None # 实验颜色 diff --git a/test/unit/data/run/test_main.py b/test/unit/data/run/test_main.py index ff9195da1..777fafe96 100644 --- a/test/unit/data/run/test_main.py +++ b/test/unit/data/run/test_main.py @@ -11,6 +11,7 @@ import json import math import os +import builtins import nanoid import numpy as np @@ -23,7 +24,7 @@ import tutils as T from swanlab import Image, Audio, Text from swanlab.data.modules import Line -from swanlab.data.run.main import SwanLabRun, get_run, SwanLabRunState, swanlog, get_url, get_project_url +from swanlab.data.run.main import SwanLabRun, get_run, SwanLabRunState, swanlog, get_url, get_project_url, get_config from swanlab.env import SwanLabEnv from tutils import TEMP_PATH from tutils.setup import UseMockRunState @@ -72,6 +73,24 @@ def test_duplicate_init(self): _run = run.finish() assert swanlog.proxied is False + def test_init_does_not_use_stale_config_setter(self): + os.environ[SwanLabEnv.MODE.value] = "disabled" + with UseMockRunState(): + stale_called = False + + def stale_setter(_): + nonlocal stale_called + stale_called = True + raise AssertionError("stale setter should not be called during init") + + config = get_config() + setattr(config, "_SwanLabConfig__on_setter", stale_setter) + + run = SwanLabRun(run_config={"a": 1}) + + assert stale_called is False + assert run.config["a"] == 1 + class TestSwanLabRunState: """ @@ -435,3 +454,57 @@ def test_run_id_none(): state.store.run_id = None with pytest.raises(AssertionError): SwanLabRun() + + +class TestSwanLabRunMonitor: + @staticmethod + def setup_method(): + os.environ[SwanLabEnv.MODE.value] = "disabled" + + def test_shared_parallel_uses_run_id_scoped_monitor_lock(self, monkeypatch): + monkeypatch.setattr("swanlab.data.run.main._get_runtime_info", lambda metadata: (metadata, None, None)) + + class DummyTimer: + def __init__(self, func, interval=None, immediate=False): + self.func = func + self.interval = interval + self.immediate = immediate + + def run(self): + return self + + def cancel(self): + return None + + def join(self): + return None + + lock_paths = [] + real_open = builtins.open + + def spy_open(path, *args, **kwargs): + path_str = os.fspath(path) + if path_str.endswith(".lock"): + lock_paths.append(path_str) + return real_open(path, *args, **kwargs) + + monkeypatch.setattr("swanlab.data.run.main.timer.Timer", DummyTimer) + monkeypatch.setattr(builtins, "open", spy_open) + + monkeypatch.setenv("PYTEST_VERSION", "1") + with UseMockRunState(run_id="abcdefghijklmnopqrstu") as run_state: + try: + monkeypatch.delenv("PYTEST_VERSION", raising=False) + run_state.store.parallel = "shared" + expected_lock_path = os.path.join( + run_state.store.swanlog_dir, f".hardware_monitor.{run_state.store.run_id}.lock" + ) + + run = SwanLabRun(metadata={"system": "metadata"}, monitor_funcs=[lambda: []]) + + assert lock_paths[-1] == expected_lock_path + assert os.path.exists(expected_lock_path) + + run.finish() + finally: + os.environ["PYTEST_VERSION"] = "1" diff --git a/test/unit/data/test_sdk.py b/test/unit/data/test_sdk.py index 8e974a3c3..a5a3af64a 100644 --- a/test/unit/data/test_sdk.py +++ b/test/unit/data/test_sdk.py @@ -335,6 +335,65 @@ def test_init_priority_3(self): run.log({"TestInitMode": 1}) +class TestParallelMode: + @staticmethod + def _mock_cloud_init(monkeypatch, *, is_new: bool): + class DummyOperator: + def on_init(self, *args, **kwargs): + run_store = get_run_store() + run_store.run_name = run_store.run_name or "Mock Run" + run_store.run_colors = run_store.run_colors or ("#111111", "#222222") + run_store.run_id = run_store.run_id or "abcdefghijklmnopqrstu" + run_store.new = is_new + if is_new: + run_store.config = None + run_store.metrics = None + else: + run_store.config = {} + run_store.metrics = {} + + class DummyRun: + @staticmethod + def is_started(): + return False + + def __init__(self, run_config=None, operator=None, metadata=None, monitor_funcs=None): + self.run_config = run_config + self.operator = operator + self.metadata = metadata + self.monitor_funcs = monitor_funcs + + monkeypatch.setattr(S, "_init_mode", lambda mode, default_mode=None: ("cloud", object())) + monkeypatch.setattr(S, "_create_operator", lambda *args, **kwargs: DummyOperator()) + monkeypatch.setattr(S, "SwanLabRun", DummyRun) + + def test_resume_allow_without_parallel_does_not_collect_metadata(self, monkeypatch): + self._mock_cloud_init(monkeypatch, is_new=False) + get_metadata_mock = Mock(return_value=({"system": "metadata"}, [lambda: []])) + monkeypatch.setattr(S, "get_metadata", get_metadata_mock) + + S.init(mode="cloud", resume="allow", id="abcdefghijklmnopqrstu", logdir=T.TEMP_PATH) + + run_store = get_run_store() + assert run_store.resume == "allow" + assert run_store.parallel == "none" + get_metadata_mock.assert_not_called() + + def test_parallel_shared_marks_parallel_and_collects_metadata(self, monkeypatch): + self._mock_cloud_init(monkeypatch, is_new=False) + get_metadata_mock = Mock(return_value=({"system": "metadata"}, [lambda: []])) + monkeypatch.setattr(S, "get_metadata", get_metadata_mock) + + run = S.init(parallel="shared", id="abcdefghijklmnopqrstu", logdir=T.TEMP_PATH) + + run_store = get_run_store() + assert run_store.resume == "allow" + assert run_store.parallel == "shared" + get_metadata_mock.assert_called_once_with(run_store.run_dir) + assert run.metadata == {"system": "metadata"} + assert len(run.monitor_funcs) == 1 + + class TestInitProject: """ 测试init时函数的project参数设置行为