Refactor/data-code (#1126)

* refactor: backup and proto

* refactor: transfer

* refactor: swanlab settings

* Refactor backup and transfer modules structure

Moved backup-related modules from swanlab.data.backup to swanlab.log.backup and transfer modules from swanlab.data.transfers to swanlab.transfers. Updated all relevant imports and usages throughout the codebase. Refactored Transfer and ProtoV0Transfer to accept media_dir and file_dir as constructor arguments. Improved error log uploading in CloudPyCallback and unified backup handler usage.

* Refactor backup handler and callbackers for improved modularity

Refactored BackupHandler to accept run parameters directly instead of accessing run_store internally, improving modularity and testability. Updated cloud, local, and offline callbackers to pass run parameters explicitly and to use run_store consistently. Added access control to get_run_store to restrict usage to swanlab.data module. Moved uploader imports to explicit usage and centralized error log uploading in ProtoV0Transfer.

* Refactor callbacker module and unify callback logic

Moved and refactored callback logic from swanlab/data/run/callback.py into swanlab/data/callbacker/callback.py, and updated cloud, local, and offline callback implementations to inherit from the new base class. Centralized utility functions in swanlab/data/callbacker/utils.py and updated all callbackers to use these shared utilities for printing and path formatting. Removed redundant code and improved maintainability by reducing duplication and clarifying callback registration and cleanup logic.

* Refactor data transfer and backup architecture

Removed legacy transfer and backup modules, consolidating data transfer logic into a new ProtoTransfer singleton in swanlab/data/transfer.py. Updated callbackers to use ProtoTransfer for logging, metric, and runtime info handling. Migrated DataStore to swanlab/data, removed async_io utility, and cleaned up related imports and usages. Adjusted tests and internal references to reflect new module structure and APIs.

* Refactor DataPorter for unified data sync and trace modes

Replaces the previous ProtoTransfer and ModelsParser logic with a new DataPorter class that supports both experiment trace and sync upload modes, centralizing backup file parsing, data publishing, and resource management. Updates all callbackers and sync logic to use DataPorter, removes ModelsParser from proto/v0.py, and adds platformdirs to requirements. This refactor improves maintainability and consistency for experiment data handling and synchronization.

* Add synced decorator and improve backup handling

Applied the @synced decorator to DataPorter methods to ensure thread safety. Improved backup handling in SwanLabRun by deleting the run directory when backup is disabled. Updated SwanLabInitializer to use a system runtime directory for logs when backup is off, ensuring proper log storage.

* Update __init__.py

* Refactor test setup and store access control

Introduced an @inside decorator to restrict access to RunStore functions to the swanlab.data module or test runtime. Refactored test utilities to use a new UseMockRunState context manager for consistent client and store state management in tests. Updated imports and test code to align with these changes, improving test isolation and reliability.

* Refactor local callback and improve config error handling

Refactored LocalRunCallback to set logdir via environment variable and updated tests to use UseMockRunState for better isolation. Changed SwanLabRun to default operator to None. Cleared callbacks after validation in SwanLabInitializer. Improved and relocated error handling for invalid config parameters from test_config.py to test_sdk.py, adding more comprehensive tests for config input validation.

* Add disabled callback mode and refactor logdir handling

Introduced DisabledCallback for 'disabled' mode, ensuring proper handling when runs are not to be saved or uploaded. Refactored logdir initialization logic in SwanLabInitializer to unify and simplify directory setup for all modes, and moved run_dir cleanup logic from main.py to the cloud callback. Updated operator creation to use DisabledCallback in disabled mode.

* chore: fix test

* fix: test

* chore: fix typing

* chore: change cache dir

* chore: some misc

* fix: code

Author: SAKURA-CAT

requirements.txt

@@ -11,6 +11,7 @@ botocore
 pydantic>=2.9.0
 pyecharts>=2.0.0
 wrapt>=1.17.0
+platformdirs>=4.2.0
 typing_extensions; python_version < '3.9'
 protobuf>=3.12.0,!=4.21.0,!=5.28.0,<7; python_version < '3.9' and sys_platform == 'linux'
 protobuf>=3.15.0,!=4.21.0,!=5.28.0,<7; python_version == '3.9' and sys_platform == 'linux'

swanlab/__init__.py

@@ -1,5 +1,13 @@
 # 导出初始化函数和log函数，以及一些数据处理模块
 from .data import *
+from .data.modules import (
+    Audio,
+    Image,
+    Object3D,
+    Molecule,
+    Text,
+    echarts,
+)
 from .env import SwanLabEnv
 from .package import get_package_version
 from .swanlab_settings import Settings

swanlab/cli/commands/sync/__init__.py

@@ -67,9 +67,9 @@ def sync(path, api_key, workspace, project, host):
         api_key = get_key() if api_key is None else api_key
     except KeyFileError:
         pass
-    for path in path:
+    for p in path:
         # 1.3 登录，创建 http 对象
         log_info = auth.terminal_login(api_key=api_key, save_key=False)
         create_client(log_info)
         # 2. 同步日志
-        sync_logs(path, workspace=workspace, project_name=project, login_required=False, raise_error=len(path) == 1)
+        sync_logs(p, workspace=workspace, project_name=project, login_required=False, raise_error=len(path) == 1)

swanlab/core_python/__init__.py

@@ -8,5 +8,3 @@
 """
 
 from .client import *
-from .uploader import *
-from .uploader import thread

swanlab/core_python/auth/providers/api_key.py

@@ -13,10 +13,11 @@
 from rich.status import Status
 from rich.text import Text
 
-from swanlab.env import is_windows
-from swanlab.error import ValidationError, APIKeyFormatError
+from swanlab.core_python import auth
+from swanlab.env import is_windows, is_interactive
+from swanlab.error import ValidationError, APIKeyFormatError, KeyFileError
 from swanlab.log import swanlog
-from swanlab.package import get_setting_url, get_host_api, get_host_web, fmt_web_host, save_key as sk
+from swanlab.package import get_setting_url, get_host_api, get_host_web, fmt_web_host, save_key as sk, get_key
 
 
 class LoginInfo:
@@ -206,6 +207,22 @@ def login_again(error: Exception):
             api_key = login_again(e)
 
 
+def create_login_info(save: bool = True):
+    """
+    在代码运行时发起登录，获取登录信息，执行此方法会覆盖原有的login_info
+    """
+    key = None
+    try:
+        key = get_key()
+    except KeyFileError:
+        pass
+    if key is None and not is_interactive():
+        raise KeyFileError(
+            "api key not configured (no-tty), call `swanlab.login(api_key=[your_api_key])` or set `swanlab.init(mode=\"local\")`."
+        )
+    return auth.terminal_login(key, save)
+
+
 def _abort_tip(tp, _, __):
     """处理用户在input_api_key输入时按下CTRL+C的情况"""
     if tp == KeyboardInterrupt:
@@ -219,4 +236,5 @@ def _abort_tip(tp, _, __):
     "terminal_login",
     "code_login",
     "LoginInfo",
+    "create_login_info",
 ]

swanlab/core_python/client/__init__.py

@@ -11,7 +11,6 @@
 
 import requests
 from requests.adapters import HTTPAdapter
-from rich.status import Status
 from urllib3.util.retry import Retry
 
 from swanlab.error import NetworkError, ApiError
@@ -255,46 +254,44 @@ def mount_project(self, name: str, username: str = None, public: bool = None):
         :param public: 项目是否公开
         :return: 项目信息
         """
-        with Status("Getting project...", spinner="dots"):
-            try:
-                data = {"name": name}
-                if username is not None:
-                    data["username"] = username
-                if public is not None:
-                    data["visibility"] = "PUBLIC" if public else "PRIVATE"
-                resp = self.post(f"/project", data=data)
-            except ApiError as e:
-                if e.resp.status_code == 409:
-                    # 项目已经存在，从对象中解析信息
-                    resp = decode_response(e.resp)
-                elif e.resp.status_code == 404 and e.resp.reason == "Not Found":
-                    # WARNING: 早期 （私有化） swanlab 后端没有 /project 接口，需要使用 /project/{username} 接口，此时没有默认空间的特性
-                    self.__groupname = self.__groupname if username is None else username
-                    try:
-                        visibility = "PUBLIC" if public else "PRIVATE"
-                        resp = self.post(f"/project/{self.groupname}", data={"name": name, "visibility": visibility})
-                    except ApiError as e:
-                        # 如果为409，表示已经存在，获取项目信息
-                        if e.resp.status_code == 409:
-                            resp = self.get(f"/project/{self.groupname}/{name}")
-                        elif e.resp.status_code == 404:
-                            # 组织/用户不存在
-                            raise ValueError(f"Space `{self.groupname}` not found")
-                        elif e.resp.status_code == 403:
-                            # 权限不足
-                            raise ValueError(f"Space permission denied: " + self.groupname)
-                        else:
-                            raise e
-                    return ProjectInfo(resp)
-                else:
-                    # 此接口为后端处理，sdk 在理论上不会出现其他错误，因此不需要处理其他错误
-                    raise e
-            # 设置当前项目所属的用户名
-            self.__groupname = resp['username']
-            # 获取详细信息
-            resp = self.get(f"/project/{self.groupname}/{name}")
-            project = ProjectInfo(resp)
-        self.__proj = project
+        try:
+            data = {"name": name}
+            if username is not None:
+                data["username"] = username
+            if public is not None:
+                data["visibility"] = "PUBLIC" if public else "PRIVATE"
+            resp = self.post(f"/project", data=data)
+        except ApiError as e:
+            if e.resp.status_code == 409:
+                # 项目已经存在，从对象中解析信息
+                resp = decode_response(e.resp)
+            elif e.resp.status_code == 404 and e.resp.reason == "Not Found":
+                # WARNING: 早期 （私有化） swanlab 后端没有 /project 接口，需要使用 /project/{username} 接口，此时没有默认空间的特性
+                self.__groupname = self.__groupname if username is None else username
+                try:
+                    visibility = "PUBLIC" if public else "PRIVATE"
+                    resp = self.post(f"/project/{self.groupname}", data={"name": name, "visibility": visibility})
+                except ApiError as e:
+                    # 如果为409，表示已经存在，获取项目信息
+                    if e.resp.status_code == 409:
+                        resp = self.get(f"/project/{self.groupname}/{name}")
+                    elif e.resp.status_code == 404:
+                        # 组织/用户不存在
+                        raise ValueError(f"Space `{self.groupname}` not found")
+                    elif e.resp.status_code == 403:
+                        # 权限不足
+                        raise ValueError(f"Space permission denied: " + self.groupname)
+                    else:
+                        raise e
+                return ProjectInfo(resp)
+            else:
+                # 此接口为后端处理，sdk 在理论上不会出现其他错误，因此不需要处理其他错误
+                raise e
+        # 设置当前项目所属的用户名
+        self.__groupname = resp['username']
+        # 获取详细信息
+        resp = self.get(f"/project/{self.groupname}/{name}")
+        self.__proj = ProjectInfo(resp)
 
     def mount_exp(self, exp_name, colors: Tuple[str, str], description: str = None, tags: List[str] = None):
         """
@@ -304,20 +301,19 @@ def mount_exp(self, exp_name, colors: Tuple[str, str], description: str = None,
         :param description: 实验描述
         :param tags: 实验标签
         """
-        with Status("Getting experiment...", spinner="dots"):
-            post_data = {
-                "name": exp_name,
-                "colors": list(colors),
-            }
-            if description is not None:
-                post_data["description"] = description
-            if tags is not None:
-                post_data["labels"] = [{"name": tag} for tag in tags]
-
-            data = self.post(f"/project/{self.groupname}/{self.__proj.name}/runs", post_data)
-            self.__exp = ExperimentInfo(data)
-            # 获取cos信息
-            self.__get_cos()
+        post_data = {
+            "name": exp_name,
+            "colors": list(colors),
+        }
+        if description is not None:
+            post_data["description"] = description
+        if tags is not None:
+            post_data["labels"] = [{"name": tag} for tag in tags]
+
+        data = self.post(f"/project/{self.groupname}/{self.__proj.name}/runs", post_data)
+        self.__exp = ExperimentInfo(data)
+        # 获取cos信息
+        self.__get_cos()
 
     def update_state(self, success: bool):
         """

swanlab/data/__init__.py

@@ -7,7 +7,6 @@
 @Description:
     在此处完成回调注册、swanlog注册，并为外界提供api，提供运行时生成的配置
 """
-from .modules import Audio, Image, Text, Object3D, Molecule, echarts
 from .run import (
     SwanLabRun as Run,
     SwanLabRunState as State,

swanlab/data/callbacker/__init__.py

@@ -0,0 +1,15 @@
+"""
+@author: cunyue
+@file: __init__.py
+@time: 2025/6/21 17:10
+@description: 回调器模块，四大回调器对应 swanlab 的四种运行模式。
+1. local： 本地模式回调器
+2. cloud： 云端模式回调器
+3. offline： 离线模式回调器
+4. disabled： 禁用回调器
+"""
+
+from .cloud import CloudPyCallback
+from .disabled import DisabledCallback
+from .local import LocalRunCallback
+from .offline import OfflineCallback

swanlab/data/callbacker/callback.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+r"""
+@DATE: 2024/6/19 16:46
+@File: callback.py
+@IDE: pycharm
+@Description:
+    回调函数注册抽象模块
+"""
+import atexit
+import sys
+import traceback
+
+from swanlab.data.run import SwanLabRunState, get_run
+from swanlab.log import swanlog
+from swanlab.swanlab_settings import get_settings
+from swanlab.toolkit import SwanKitCallback
+from . import utils
+from ..porter import DataPorter
+from ..store import get_run_store
+
+
+class SwanLabRunCallback(SwanKitCallback):
+    """
+    SwanLabRunCallback，回调函数注册类，所有以`on_`和`before_`开头的函数都会在对应的时机被调用
+    为了方便管理：
+    1. `_`开头的函数为内部函数，不会被调用，且写在最开头
+    2. 所有回调按照逻辑上的触发顺序排列
+    3. 带有from_*后缀的回调函数代表调用者来自其他地方，比如config、operator等，这将通过settings对象传递
+    4. 所有回调不要求全部实现，只需实现需要的回调即可
+    """
+
+    def __init__(self):
+        self.run_store = get_run_store()
+        self.porter = DataPorter()
+        self.user_settings = get_settings()
+
+    def _register_sys_callback(self):
+        """
+        注册系统回调，内部使用
+        """
+        sys.excepthook = self._except_handler
+        atexit.register(self._clean_handler)
+
+    def _unregister_sys_callback(self):
+        """
+        注销系统回调，内部使用
+        """
+        sys.excepthook = sys.__excepthook__
+        atexit.unregister(self._clean_handler)
+
+    def _clean_handler(self):
+        """
+        正常退出清理函数，此函数调用`run.finish`
+        """
+        run = get_run()
+        if run is None:
+            return swanlog.debug("SwanLab Runtime has been cleaned manually.")
+        # 打印训练结束信息
+        utils.print_train_finish(self.run_store.run_name)
+        # 如果正在运行
+        run.finish() if run.running else swanlog.debug("Duplicate finish, ignore it.")
+
+    @staticmethod
+    def _except_handler(tp, val, tb):
+        """
+        异常退出清理函数
+        """
+        # 1. 如果是KeyboardInterrupt异常，特殊显示
+        if tp == KeyboardInterrupt:
+            swanlog.info("KeyboardInterrupt by user")
+        else:
+            swanlog.info("Error happened while training")
+        # 2. 生成错误堆栈
+        trace_list = traceback.format_tb(tb)
+        error = ""
+        for line in trace_list:
+            error += line
+        error += str(val)
+        # 3. 结束运行，注意此时终端错误还没打印
+        get_run().finish(SwanLabRunState.CRASHED, error=error)
+        assert swanlog.proxied is False, "except_handler should be called after swanlog.stop_proxy()"
+        # 4. 打印终端错误，此时终端代理已经停止，不必担心此副作用
+        print(error, file=sys.stderr)
+
+    def _start_terminal_proxy(self, handler=None):
+        """
+        启动终端代理
+        """
+        if handler is None:
+            handler = lambda data: self.porter.trace_log(data)
+        swanlog.start_proxy(
+            proxy_type=self.user_settings.log_proxy_type,
+            max_log_length=self.user_settings.max_log_length,
+            handler=handler,
+        )
+
+    def __str__(self):
+        raise NotImplementedError("Please implement this method")