refactor: tidy swankit (#1111)

Author: SAKURA-CAT

docs/插件化设计.md

@@ -25,7 +25,7 @@ swanlab一切的起点都是`swanlab.init`，此函数允许传入`callbacks`参
 `SwanKitCallback`，类似下面这样：
 
 ```python
-from swankit.callback import SwanKitCallback
+from swanlab.toolkit import SwanKitCallback
 
 
 class CustomCallback(SwanKitCallback):
@@ -41,7 +41,7 @@ class CustomCallback(SwanKitCallback):
 
 ```python
 import swanlab
-from swankit.callback import SwanKitCallback
+from swanlab.toolkit import SwanKitCallback
 
 
 class CustomCallback(SwanKitCallback):

requirements.txt

@@ -15,4 +15,5 @@ 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'
 protobuf>=3.19.0,!=4.21.0,!=5.28.0,<7; python_version > '3.9' and sys_platform == 'linux'
-protobuf>=3.19.0,!=4.21.0,!=5.28.0,<7; sys_platform != 'linux'
+protobuf>=3.19.0,!=4.21.0,!=5.28.0,<7; sys_platform != 'linux'
+rich>=13.6.0, <14.0.0

swanlab/cli/commands/auth/login.py

@@ -10,10 +10,11 @@
 import os
 
 import click
-from swankit.log import FONT
+from rich.text import Text
 
 from swanlab.core_python import auth
 from swanlab.env import SwanLabEnv
+from swanlab.log import swanlog
 from swanlab.package import has_api_key, HostFormatter
 
 
@@ -51,9 +52,12 @@ def login(api_key: str, relogin: bool, host: str, web_host: str):
     """Login to the swanlab cloud."""
     if not relogin and has_api_key():
         # 此时代表token已经获取，需要打印一条信息：已经登录
-        command = FONT.bold("swanlab login --relogin")
-        tip = FONT.swanlab("You are already logged in. Use `" + command + "` to force relogin.")
-        return print(tip)
+        return swanlog.info(
+            "You are already logged in. Use `",
+            Text("swanlab login --relogin", "bold"),
+            "` to force relogin.",
+            sep='',
+        )
     # 清除环境变量
     if relogin:
         del os.environ[SwanLabEnv.API_HOST.value]
@@ -64,4 +68,4 @@ def login(api_key: str, relogin: bool, host: str, web_host: str):
         raise click.BadParameter(str(e))
     # 进行登录，此时将直接覆盖本地token文件
     login_info = auth.terminal_login(api_key)
-    print(FONT.swanlab("Login successfully. Hi, " + FONT.bold(FONT.default(login_info.username))) + "!")
+    swanlog.info("Login successfully. Hi, ", Text(login_info.username, "bold"), "!", sep='')

swanlab/cli/commands/auth/logout.py

@@ -11,28 +11,30 @@
 import sys
 
 import click
-from swankit.log import FONT
+from rich.text import Text
 
 from swanlab.env import get_save_dir
+from swanlab.log import swanlog
 from swanlab.package import has_api_key
 
 
 @click.command()
 def logout():
     """Logout to the swanlab cloud."""
-    command = FONT.bold("swanlab login")
+    command = Text("swanlab login", "bold")
     if has_api_key():
         # 如果已经是登录状态，那么则询问用户是否确认，如果确认则删除token文件夹
-        confirm = input(FONT.swanlab("Are you sure you want to logout? (y/N): "))
+
+        swanlog.info("Are you sure you want to logout? (y/N): ", end='')
+        confirm = input()
         if confirm.lower() == "y":
             try:
                 shutil.rmtree(get_save_dir())
-                return print(FONT.swanlab("Logout successfully. You can use `" + command + "` to login again."))
+                return swanlog.info(" Logout successfully. You can use `", command, "` to login again.", sep="")
             except Exception as e:
-                return print(FONT.swanlab("Logout failed. Reason:" + str(e)))
+                return swanlog.info("Logout failed. Reason:" + str(e))
         else:
-            return print(FONT.swanlab("Logout canceled."))
+            return swanlog.info("Logout canceled.")
     # 如果还未登录，则不做任何处理，并告知用户如何登录
-    tip = FONT.swanlab("You are not logged in. If you want to login in, please use `" + command + "` to login.")
-    print(tip)
+    swanlog.info("You are not logged in. If you want to login in, please use `", command, "` to login.", sep="")
     return sys.exit(1)

swanlab/core_python/auth/providers/api_key.py

@@ -10,10 +10,10 @@
 from typing import Union
 
 import requests
-from swankit.env import is_windows
-from swankit.log import FONT
+from rich.status import Status
+from rich.text import Text
 
-from swanlab.env import in_jupyter
+from swanlab.env import is_windows
 from swanlab.error import ValidationError, APIKeyFormatError
 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
@@ -143,19 +143,21 @@ def input_api_key(
     _t = sys.excepthook
     sys.excepthook = _abort_tip
     if not again:
-        swanlog.info("You can find your API key at: " + FONT.yellow(get_setting_url()))
+        swanlog.info("You can find your API key at:", Text(get_setting_url(), "yellow"))
+    swanlog.info(tip, end='')
     # windows 额外打印提示信息
     if is_windows():
-        tip += (
-            '\nOn Windows, '
-            f'use {FONT.yellow("Ctrl + Shift + V")} or {FONT.yellow("right-click")} '
-            f'to paste the API key'
+        swanlog.console.print(
+            '\nOn Windows, ',
+            'use',
+            Text("Ctrl + Shift + V", 'yellow'),
+            'or',
+            Text("right-click", 'yellow'),
+            'to paste the API key',
+            end='',
         )
-    tip += ': '
-    tip = FONT.swanlab(tip)
-    ij = in_jupyter()
-    ij and print(tip)
-    key = getpass.getpass("" if ij else tip)
+    swanlog.console.print(': ', end='')
+    key = getpass.getpass("")
     sys.excepthook = _t
     return key
 
@@ -170,8 +172,8 @@ def code_login(api_key: str, save_key: bool = True) -> LoginInfo:
     :raise APIKeyFormatError: api_key格式错误
     """
     APIKeyFormatError.check(api_key)
-    tip = "Waiting for the swanlab cloud response."
-    login_info: LoginInfo = FONT.loading(tip, login_by_key, args=(api_key, 20, save_key), interval=0.5)
+    with Status("Waiting for the swanlab cloud response.", spinner="dots"):
+        login_info = login_by_key(api_key, 20, save_key)
     if login_info.is_fail:
         raise ValidationError("Login failed: " + str(login_info))
     return login_info

swanlab/core_python/client/__init__.py

@@ -11,13 +11,13 @@
 
 import requests
 from requests.adapters import HTTPAdapter
-from swankit.core import MediaBuffer
-from swankit.log import FONT
+from rich.status import Status
 from urllib3.util.retry import Retry
 
 from swanlab.error import NetworkError, ApiError
 from swanlab.log import swanlog
 from swanlab.package import get_package_version
+from swanlab.toolkit import MediaBuffer
 from .cos import CosClient
 from .model import ProjectInfo, ExperimentInfo
 from .. import auth
@@ -255,8 +255,7 @@ def mount_project(self, name: str, username: str = None, public: bool = None):
         :param public: 项目是否公开
         :return: 项目信息
         """
-
-        def _():
+        with Status("Getting project...", spinner="dots"):
             try:
                 data = {"name": name}
                 if username is not None:
@@ -294,9 +293,7 @@ def _():
             self.__groupname = resp['username']
             # 获取详细信息
             resp = self.get(f"/project/{self.groupname}/{name}")
-            return ProjectInfo(resp)
-
-        project: ProjectInfo = FONT.loading("Getting project...", _)
+            project = ProjectInfo(resp)
         self.__proj = project
 
     def mount_exp(self, exp_name, colors: Tuple[str, str], description: str = None, tags: List[str] = None):
@@ -307,11 +304,7 @@ def mount_exp(self, exp_name, colors: Tuple[str, str], description: str = None,
         :param description: 实验描述
         :param tags: 实验标签
         """
-
-        def _():
-            """
-            先创建实验，后生成cos凭证
-            """
+        with Status("Getting experiment...", spinner="dots"):
             post_data = {
                 "name": exp_name,
                 "colors": list(colors),
@@ -326,21 +319,15 @@ def _():
             # 获取cos信息
             self.__get_cos()
 
-        FONT.loading("Creating experiment...", _)
-
     def update_state(self, success: bool):
         """
         更新实验状态
         :param success: 实验是否成功
         """
-
-        def _():
-            self.put(
-                f"/project/{self.groupname}/{self.projname}/runs/{self.exp_id}/state",
-                {"state": "FINISHED" if success else "CRASHED", "from": "sdk"},
-            )
-
-        FONT.loading("Updating experiment status...", _)
+        self.put(
+            f"/project/{self.groupname}/{self.projname}/runs/{self.exp_id}/state",
+            {"state": "FINISHED" if success else "CRASHED", "from": "sdk"},
+        )
 
 
 client: Optional["Client"] = None

swanlab/core_python/uploader/model.py

@@ -10,9 +10,8 @@
 from enum import Enum
 from typing import List, Optional, TypedDict, Literal
 
-from swankit.callback.models import ColumnClass, ColumnConfig
-
 from swanlab.data.modules import MediaBuffer
+from swanlab.toolkit import ColumnClass, ColumnConfig
 
 
 class ColumnModel:

swanlab/data/callbacker/cloud.py

@@ -9,9 +9,7 @@
 import sys
 from typing import Literal
 
-from swankit.callback.models import RuntimeInfo, MetricInfo, ColumnInfo
-from swankit.env import create_time
-from swankit.log import FONT
+from rich.text import Text
 
 from swanlab.env import in_jupyter, is_interactive
 from swanlab.error import KeyFileError
@@ -21,6 +19,12 @@
     get_package_latest_version,
     get_key,
 )
+from swanlab.toolkit import (
+    RuntimeInfo,
+    MetricInfo,
+    ColumnInfo,
+    create_time,
+)
 from ..run import get_run, SwanLabRunState
 from ..run.callback import SwanLabRunCallback
 from ...core_python import *
@@ -71,8 +75,8 @@ def _get_package_latest_version():
     def _view_web_print():
         http = get_client()
         proj_url, exp_url = http.web_proj_url, http.web_exp_url
-        swanlog.info("🏠 View project at " + FONT.blue(FONT.underline(proj_url)))
-        swanlog.info("🚀 View run at " + FONT.blue(FONT.underline(exp_url)))
+        swanlog.info("🏠 View project at", Text(proj_url, "blue underline"))
+        swanlog.info("🚀 View run at", Text(exp_url, "blue underline"))
         return exp_url
 
     def _clean_handler(self):
@@ -129,8 +133,9 @@ def on_run(self, *args, **kwargs):
         # 打印实验开始信息，在 cloud 模式下如果没有开启 backup 的话不打印“数据保存在 xxx”的信息
         swanlab_settings = get_settings()
         self._train_begin_print(save_dir=self.settings.run_dir if swanlab_settings.backup else None)
-        swanlog.info("👋 Hi " + FONT.bold(FONT.default(http.username)) + ", welcome to swanlab!")
-        swanlog.info("Syncing run " + FONT.yellow(self.settings.exp_name) + " to the cloud")
+
+        swanlog.info(" 👋 Hi ", Text(http.username, "bold default"), ",welcome to swanlab!", sep="")
+        swanlog.info("Syncing run", Text(self.settings.exp_name, "yellow"), "to the cloud")
         experiment_url = self._view_web_print()
         # 在Jupyter Notebook环境下，显示按钮
         if in_jupyter():

swanlab/data/callbacker/local.py

@@ -6,12 +6,12 @@
 local模式（目前）将自动调用swanboard，如果不存在则报错
 """
 
-from swankit.callback import ColumnInfo
-from swankit.log import FONT
+from rich.text import Text
 
 from swanlab.log.backup import backup
 from swanlab.log.backup.writer import write_media_buffer, write_runtime_info
 from swanlab.log.type import LogData
+from swanlab.toolkit import ColumnInfo
 from ..run import get_run
 
 try:
@@ -34,9 +34,7 @@
 from datetime import datetime
 from typing import Tuple, Optional, TextIO
 
-from swankit.callback.models import RuntimeInfo, MetricInfo
-from swankit.core import SwanLabSharedSettings
-
+from swanlab.toolkit import RuntimeInfo, MetricInfo, SwanLabSharedSettings
 from swanlab.data.run.callback import SwanLabRunCallback
 from swanlab.log import swanlog
 
@@ -57,9 +55,10 @@ def _watch_tip_print(self):
         watch命令提示打印
         """
         swanlog.info(
-            "🌟 Run `"
-            + FONT.bold("swanlab watch {}".format(self.fmt_windows_path(self.settings.swanlog_dir)))
-            + "` to view SwanLab Experiment Dashboard locally"
+            " 🌟 Run `",
+            Text("swanlab watch {}".format(self.fmt_windows_path(self.settings.swanlog_dir)), "bold"),
+            "` to view SwanLab Experiment Dashboard locally",
+            sep="",
         )
 
     @backup("terminal")

swanlab/data/callbacker/offline.py

@@ -5,13 +5,13 @@
 @description: 日志备份回调
 """
 
-from swankit.callback import ColumnInfo, MetricInfo, RuntimeInfo
-from swankit.log import FONT
+from rich.text import Text
 
 from swanlab.data.run.callback import SwanLabRunCallback
 from swanlab.log import swanlog
 from swanlab.log.backup import backup
 from swanlab.log.type import LogData
+from swanlab.toolkit import ColumnInfo, MetricInfo, RuntimeInfo
 from ..run import get_run
 
 
@@ -28,9 +28,10 @@ def _sync_tip_print(self):
         提示用户可以通过命令上传日志到远程服务器
         """
         swanlog.info(
-            "☁️ Run `"
-            + FONT.bold("swanlab sync {}".format(self.fmt_windows_path(self.settings.run_dir)))
-            + "` to sync logs to remote server"
+            " ☁️ Run `",
+            Text("swanlab sync {}".format(self.fmt_windows_path(self.settings.run_dir))),
+            "` to sync logs to remote server",
+            sep="",
         )
 
     @backup("terminal")
@@ -51,7 +52,7 @@ def on_init(self, proj_name: str, workspace: str, public: bool = None, logdir: s
     def on_run(self, *args, **kwargs):
         self.handle_run()
         self._train_begin_print(self.settings.run_dir)
-        swanlog.info("Backing up run " + FONT.yellow(self.settings.exp_name) + " locally")
+        swanlog.info("Backing up run", Text(self.settings.exp_name, "yellow"), "locally")
         self._sync_tip_print()
 
     @backup('runtime')