Allow customzing the trackio color palette (#339)

Co-authored-by: gradio-pr-bot <gradio-pr-bot@users.noreply.github.com>

Author: abidlabs

.changeset/thin-ducks-grow.md

@@ -0,0 +1,5 @@
+---
+"trackio": minor
+---
+
+feat:Allow customzing the trackio color palette

README.md

@@ -22,7 +22,8 @@
 
 `trackio` is a lightweight, free experiment tracking Python library built by Hugging Face 🤗.
 
-![Screen Recording 2025-07-28 at 5 26 32 PM](https://github.com/user-attachments/assets/f3eac49e-d8ee-4fc0-b1ca-aedfc6d6fae1)
+![Screen Recording 2025-11-06 at 5 34 50 PM](https://github.com/user-attachments/assets/8c9c1b96-f17a-401c-83a4-26ac754f89c7)
+
 
 - **API compatible** with `wandb.init`, `wandb.log`, and `wandb.finish`. Drop-in replacement: just 
 

docs/source/environment_variables.md

@@ -59,6 +59,17 @@ export TRACKIO_THEME="gstaff/xkcd"
 export TRACKIO_THEME="ParityError/Anime"
 ```
 
+### `TRACKIO_COLOR_PALETTE`
+
+Customizes the color palette used for plot lines in the Trackio dashboard. The value should be a comma-separated list of hex color codes. These colors will be cycled through when plotting multiple runs.
+
+```bash
+export TRACKIO_COLOR_PALETTE="#FF0000,#00FF00,#0000FF,#FFFF00,#FF00FF,#00FFFF"
+```
+
+**Default palette:**
+`#A8769B, #E89957, #3B82F6, #10B981, #EF4444, #8B5CF6, #14B8A6, #F59E0B, #EC4899, #06B6D4`
+
 ### `TRACKIO_TABLE_TRUNCATE_LENGTH`
 
 Controls the maximum length of string values displayed in table cells before they are truncated. When a cell value exceeds this length, it will be truncated with an expandable element that allows viewing the full text. Defaults to `250` characters.

docs/source/launch.md

@@ -71,6 +71,31 @@ trackio.show(theme="soft")
 
 To see the available themes, check out the [themes gallery](https://huggingface.co/spaces/gradio/theme-gallery).
 
+## Customizing Plot Colors
+
+You can customize the color palette used for plot lines by providing a `color_palette` argument. This is useful if you want to match your organization's branding or have specific color preferences.
+
+<hfoptions id="language">
+<hfoption id="Shell">
+
+```sh
+trackio show --color-palette "#FF0000,#00FF00,#0000FF"
+```
+
+</hfoption>
+<hfoption id="Python">
+
+```py
+import trackio 
+
+trackio.show(color_palette=["#FF0000", "#00FF00", "#0000FF"])
+```
+
+</hfoption>
+</hfoptions>
+
+The colors will be cycled through when displaying multiple runs. You can provide as many or as few colors as you like.
+
 ## Launching a Dashboard in Jupyter Notebooks
 
 You can also launch the dashboard directly within a Jupyter Notebook. Just use the same command as above:

trackio/__init__.py

@@ -311,6 +311,7 @@ def show(
     project: str | None = None,
     theme: str | ThemeClass | None = None,
     mcp_server: bool | None = None,
+    color_palette: list[str] | None = None,
     *,
     open_browser: bool = True,
     block_thread: bool | None = None,
@@ -333,6 +334,11 @@ def show(
             functions will be added as MCP tools. If `None` (default behavior), then the
             `GRADIO_MCP_SERVER` environment variable will be used to determine if the
             MCP server should be enabled (which is `"True"` on Hugging Face Spaces).
+        color_palette (`list[str]`, *optional*):
+            A list of hex color codes to use for plot lines. If not provided, the
+            `TRACKIO_COLOR_PALETTE` environment variable will be used (comma-separated
+            hex codes), or if that is not set, the default color palette will be used.
+            Example: `['#FF0000', '#00FF00', '#0000FF']`
         open_browser (`bool`, *optional*, defaults to `True`):
             If `True` and not in a notebook, a new browser tab will be opened with the dashboard.
             If `False`, the browser will not be opened.
@@ -347,6 +353,9 @@ def show(
             `share_url`: The public share URL of the dashboard.
             `full_url`: The full URL of the dashboard including the write token (will use the public share URL if launched publicly, otherwise the local URL).
     """
+    if color_palette is not None:
+        os.environ["TRACKIO_COLOR_PALETTE"] = ",".join(color_palette)
+
     theme = theme or os.environ.get("TRACKIO_THEME", DEFAULT_THEME)
 
     if theme != DEFAULT_THEME:

trackio/cli.py

@@ -24,11 +24,19 @@ def main():
         action="store_true",
         help="Enable MCP server functionality. The Trackio dashboard will be set up as an MCP server and certain functions will be exposed as MCP tools.",
     )
+    ui_parser.add_argument(
+        "--color-palette",
+        required=False,
+        help="Comma-separated list of hex color codes for plot lines (e.g. '#FF0000,#00FF00,#0000FF'). If not provided, the TRACKIO_COLOR_PALETTE environment variable will be used, or the default palette if not set.",
+    )
 
     args = parser.parse_args()
 
     if args.command == "show":
-        show(args.project, args.theme, args.mcp_server)
+        color_palette = None
+        if args.color_palette:
+            color_palette = [color.strip() for color in args.color_palette.split(",")]
+        show(args.project, args.theme, args.mcp_server, color_palette)
     else:
         parser.print_help()
 

trackio/utils.py

@@ -465,26 +465,39 @@ def format_timestamp(timestamp_str):
         return "Unknown"
 
 
-COLOR_PALETTE = [
+DEFAULT_COLOR_PALETTE = [
+    "#A8769B",
+    "#E89957",
     "#3B82F6",
-    "#EF4444",
     "#10B981",
-    "#F59E0B",
+    "#EF4444",
     "#8B5CF6",
+    "#14B8A6",
+    "#F59E0B",
     "#EC4899",
     "#06B6D4",
-    "#84CC16",
-    "#F97316",
-    "#6366F1",
 ]
 
 
-def get_color_mapping(runs: list[str], smoothing: bool) -> dict[str, str]:
+def get_color_palette() -> list[str]:
+    """Get the color palette from environment variable or use default."""
+    env_palette = os.environ.get("TRACKIO_COLOR_PALETTE")
+    if env_palette:
+        return [color.strip() for color in env_palette.split(",")]
+    return DEFAULT_COLOR_PALETTE
+
+
+def get_color_mapping(
+    runs: list[str], smoothing: bool, color_palette: list[str] | None = None
+) -> dict[str, str]:
     """Generate color mapping for runs, with transparency for original data when smoothing is enabled."""
+    if color_palette is None:
+        color_palette = get_color_palette()
+
     color_map = {}
 
     for i, run in enumerate(runs):
-        base_color = COLOR_PALETTE[i % len(COLOR_PALETTE)]
+        base_color = color_palette[i % len(color_palette)]
 
         if smoothing:
             color_map[run] = base_color + "4D"