feat: add color parameter to swanlab.init for experiment color control (#1481)

* feat: add color parameter to swanlab.init for experiment color control

- Add color parameter supporting preset colors, RGB, and hex formats
- Add SWANLAB_EXP_COLOR environment variable
- Support color formats: preset names (green, blue, etc), rgb(r,g,b), #hex
- Update callbackers to respect user-defined colors

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* refactor: improve color format validation code quality

- Merge RGB regex patterns using optional group (?:rgb\s*)?
- Use re.fullmatch for hex color validation

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* improve

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: Zeyi-Lin

swanlab/data/callbacker/disabled.py

@@ -13,7 +13,8 @@
 class DisabledCallback(SwanLabRunCallback):
     def on_init(self, proj_name: str, workspace: str, public: bool = None, logdir: str = None, *args, **kwargs):
         self.run_store.run_name = "run-disabled"
-        self.run_store.run_colors = N.generate_colors(0)
+        if self.run_store.run_colors is None:
+            self.run_store.run_colors = N.generate_colors(0)
         self.run_store.run_id = N.generate_run_id()
         self.run_store.new = True
 

swanlab/data/callbacker/local.py

@@ -78,7 +78,8 @@ def on_init(self, proj_name: str, workspace: str, public: bool = None, logdir: s
         run_store.tags = [] if run_store.tags is None else run_store.tags
         exp_count = random.randint(0, 20)
         run_store.run_name = N.generate_name(exp_count) if run_store.run_name is None else run_store.run_name
-        run_store.run_colors = generate_colors(random.randint(0, 20))
+        if run_store.run_colors is None:
+            run_store.run_colors = generate_colors(exp_count)
         run_store.run_id = N.generate_run_id()
         run_store.new = True
 

swanlab/data/callbacker/offline.py

@@ -30,12 +30,11 @@ def on_init(self, proj_name: str, workspace: str, public: bool = None, logdir: s
         run_store.workspace = workspace
         run_store.visibility = public
         run_store.tags = [] if run_store.tags is None else run_store.tags
-        # 设置颜色，随机生成一个
+        # 设置颜色和名称
         exp_count = random.randint(0, 20)
-        run_store.run_colors = N.generate_colors(exp_count)
-        # 设置名称，随机生成
+        if run_store.run_colors is None:
+            run_store.run_colors = N.generate_colors(exp_count)
         run_store.run_name = N.generate_name(exp_count) if run_store.run_name is None else run_store.run_name
-        run_store.run_colors = N.generate_colors(exp_count)
         run_store.run_id = N.generate_run_id()
         run_store.new = True
 

swanlab/data/sdk.py

@@ -23,6 +23,7 @@
     check_desc_format,
     check_tags_format,
     check_run_id_format,
+    check_color_format,
 )
 from swanlab.log import swanlog
 from swanlab.swanlab_settings import Settings, get_settings, set_settings, read_folder_settings
@@ -108,6 +109,7 @@ def init(
         id: str = None,
         resume: Union[Literal['must', 'allow', 'never'], bool] = None,
         reinit: bool = None,
+        color: str = None,
         **kwargs,
     ) -> SwanLabRun:
         """
@@ -188,6 +190,10 @@ def init(
             The run ID of the previous run, which is used to resume the previous run.
         reinit : bool, optional
             Whether to reinitialize SwanLabRun, the default is False.
+        color : str, optional
+            The experiment color displayed in the web interface.
+            Supports preset colors (e.g., "green"), RGB format (e.g., "rgb(82,141,89)"),
+            or hex format (e.g., "#528d59", "528d59").
         """
         # 一个进程同时只能有一个实验在运行
         if SwanLabRun.is_started():
@@ -256,6 +262,7 @@ def init(
         resume = _load_from_env(SwanLabEnv.RESUME.value, resume)
         id = _load_from_env(SwanLabEnv.RUN_ID.value, id)
         logdir = _load_from_env(SwanLabEnv.SWANLOG_FOLDER.value, logdir)
+        color = _load_from_env(SwanLabEnv.EXP_COLOR.value, color)
         # 2. 部分格式校验
         # 2.1 校验项目名称，默认实验名称为当前目录名
         project = project if project else os.path.basename(os.getcwd())
@@ -298,6 +305,13 @@ def init(
                 if tags[i] != new_tags[i]:
                     swanlog.warning("The tag has been truncated automatically.")
                     tags[i] = new_tags[i]
+        # 2.7 校验颜色格式
+        if color:
+            try:
+                color = check_color_format(color)
+            except ValueError as e:
+                swanlog.warning(f"Invalid color format: {e}, will use random color")
+                color = None
         # 3. 校验回调函数
         callbacks = check_callback_format(self.cbs + callbacks)
         self.cbs = []
@@ -352,6 +366,7 @@ def init(
         run_store.tags = tags
         run_store.description = description
         run_store.run_name = experiment_name
+        run_store.run_colors = (color, color) if color else None
         run_store.swanlog_dir = logdir
         # 2. 启动操作员，注册运行实例
         operator = _create_operator(mode, login_info, callbacks)

swanlab/env.py

@@ -133,6 +133,10 @@ class SwanLabEnv(enum.Enum):
     """
     实验标签，用于标注当前实验，多个标签用逗号分隔
     """
+    EXP_COLOR = "SWANLAB_EXP_COLOR"
+    """
+    实验颜色，用于控制实验在网页端的显示颜色
+    """
     DISABLE_GIT = "SWANLAB_DISABLE_GIT"
     """
     禁用Git功能，设置为true时不会采集Git信息

swanlab/formatter.py

@@ -256,3 +256,52 @@ def check_callback_format(callback: Optional[Union[SwanKitCallback, List[SwanKit
     if not isinstance(callback, list):
         raise TypeError(f"Only support SwanKitCallback or List[SwanKitCallback], but got {type(callback)}")
     return callback
+
+
+PRESET_COLORS = {
+    "green": "#528d59",
+    "blue": "#587ad2",
+    "red": "#c24d46",
+    "cyan": "#6ebad3",
+    "orange": "#dfb142",
+    "purple": "#6d4ba4",
+    "pink": "#d47694",
+    "brown": "#905f4a",
+    "gray": "#989fa3",
+}
+
+
+def check_color_format(color: str) -> Optional[str]:
+    """
+    验证并转换颜色格式为十六进制
+    支持：预置颜色名、RGB格式、十六进制
+    :param color: 颜色字符串
+    :return: 十六进制颜色字符串
+    :raises ValueError: 颜色格式不正确
+    """
+    if not color:
+        return None
+
+    color = color.strip().lower()
+
+    # 预置颜色
+    if color in PRESET_COLORS:
+        return PRESET_COLORS[color]
+
+    # RGB格式
+    rgb_match = re.match(r'(?:rgb\s*)?\(\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\)', color)
+    if rgb_match:
+        r, g, b = map(int, rgb_match.groups())
+        if all(0 <= v <= 255 for v in [r, g, b]):
+            return f"#{r:02x}{g:02x}{b:02x}"
+        raise ValueError(f"RGB values must be 0-255: ({r}, {g}, {b})")
+
+    # 十六进制格式
+    hex_color = color.lstrip("#")
+    if len(hex_color) == 3:
+        hex_color = "".join(c * 2 for c in hex_color)
+
+    if re.fullmatch(r'[0-9a-f]{6}', hex_color):
+        return f"#{hex_color}"
+
+    raise ValueError(f"Invalid color format: {color}")