Spaces:
Sleeping
Sleeping
| import hashlib | |
| import json | |
| import logging | |
| import os | |
| import warnings | |
| import webbrowser | |
| from pathlib import Path | |
| from typing import Any | |
| from gradio.blocks import BUILT_IN_THEMES | |
| from gradio.themes import Default as DefaultTheme | |
| from gradio.themes import ThemeClass | |
| from gradio_client import Client | |
| from huggingface_hub import SpaceStorage | |
| from trackio import context_vars, deploy, utils | |
| from trackio.histogram import Histogram | |
| from trackio.imports import import_csv, import_tf_events | |
| from trackio.media import TrackioImage, TrackioVideo | |
| from trackio.run import Run | |
| from trackio.sqlite_storage import SQLiteStorage | |
| from trackio.table import Table | |
| from trackio.ui.main import demo | |
| from trackio.utils import TRACKIO_DIR, TRACKIO_LOGO_DIR | |
| logging.getLogger("httpx").setLevel(logging.WARNING) | |
| warnings.filterwarnings( | |
| "ignore", | |
| message="Empty session being created. Install gradio\\[oauth\\]", | |
| category=UserWarning, | |
| module="gradio.helpers", | |
| ) | |
| __version__ = json.loads(Path(__file__).parent.joinpath("package.json").read_text())[ | |
| "version" | |
| ] | |
| __all__ = [ | |
| "init", | |
| "log", | |
| "finish", | |
| "show", | |
| "import_csv", | |
| "import_tf_events", | |
| "Image", | |
| "Video", | |
| "Table", | |
| "Histogram", | |
| ] | |
| Image = TrackioImage | |
| Video = TrackioVideo | |
| config = {} | |
| DEFAULT_THEME = "default" | |
| def init( | |
| project: str, | |
| name: str | None = None, | |
| group: str | None = None, | |
| space_id: str | None = None, | |
| space_storage: SpaceStorage | None = None, | |
| dataset_id: str | None = None, | |
| config: dict | None = None, | |
| resume: str = "never", | |
| settings: Any = None, | |
| private: bool | None = None, | |
| embed: bool = True, | |
| ) -> Run: | |
| """ | |
| Creates a new Trackio project and returns a [`Run`] object. | |
| Args: | |
| project (`str`): | |
| The name of the project (can be an existing project to continue tracking or | |
| a new project to start tracking from scratch). | |
| name (`str`, *optional*): | |
| The name of the run (if not provided, a default name will be generated). | |
| group (`str`, *optional*): | |
| The name of the group which this run belongs to in order to help organize | |
| related runs together. You can toggle the entire group's visibilitiy in the | |
| dashboard. | |
| space_id (`str`, *optional*): | |
| If provided, the project will be logged to a Hugging Face Space instead of | |
| a local directory. Should be a complete Space name like | |
| `"username/reponame"` or `"orgname/reponame"`, or just `"reponame"` in which | |
| case the Space will be created in the currently-logged-in Hugging Face | |
| user's namespace. If the Space does not exist, it will be created. If the | |
| Space already exists, the project will be logged to it. | |
| space_storage ([`~huggingface_hub.SpaceStorage`], *optional*): | |
| Choice of persistent storage tier. | |
| dataset_id (`str`, *optional*): | |
| If a `space_id` is provided, a persistent Hugging Face Dataset will be | |
| created and the metrics will be synced to it every 5 minutes. Specify a | |
| Dataset with name like `"username/datasetname"` or `"orgname/datasetname"`, | |
| or `"datasetname"` (uses currently-logged-in Hugging Face user's namespace), | |
| or `None` (uses the same name as the Space but with the `"_dataset"` | |
| suffix). If the Dataset does not exist, it will be created. If the Dataset | |
| already exists, the project will be appended to it. | |
| config (`dict`, *optional*): | |
| A dictionary of configuration options. Provided for compatibility with | |
| `wandb.init()`. | |
| resume (`str`, *optional*, defaults to `"never"`): | |
| Controls how to handle resuming a run. Can be one of: | |
| - `"must"`: Must resume the run with the given name, raises error if run | |
| doesn't exist | |
| - `"allow"`: Resume the run if it exists, otherwise create a new run | |
| - `"never"`: Never resume a run, always create a new one | |
| private (`bool`, *optional*): | |
| Whether to make the Space private. If None (default), the repo will be | |
| public unless the organization's default is private. This value is ignored | |
| if the repo already exists. | |
| settings (`Any`, *optional*): | |
| Not used. Provided for compatibility with `wandb.init()`. | |
| embed (`bool`, *optional*, defaults to `True`): | |
| If running inside a jupyter/Colab notebook, whether the dashboard should | |
| automatically be embedded in the cell when trackio.init() is called. | |
| Returns: | |
| `Run`: A [`Run`] object that can be used to log metrics and finish the run. | |
| """ | |
| if settings is not None: | |
| warnings.warn( | |
| "* Warning: settings is not used. Provided for compatibility with wandb.init(). Please create an issue at: https://github.com/gradio-app/trackio/issues if you need a specific feature implemented." | |
| ) | |
| if space_id is None and dataset_id is not None: | |
| raise ValueError("Must provide a `space_id` when `dataset_id` is provided.") | |
| space_id, dataset_id = utils.preprocess_space_and_dataset_ids(space_id, dataset_id) | |
| url = context_vars.current_server.get() | |
| share_url = context_vars.current_share_server.get() | |
| if url is None: | |
| if space_id is None: | |
| _, url, share_url = demo.launch( | |
| show_api=False, | |
| inline=False, | |
| quiet=True, | |
| prevent_thread_lock=True, | |
| show_error=True, | |
| favicon_path=TRACKIO_LOGO_DIR / "trackio_logo_light.png", | |
| allowed_paths=[TRACKIO_LOGO_DIR, TRACKIO_DIR], | |
| ) | |
| else: | |
| url = space_id | |
| share_url = None | |
| context_vars.current_server.set(url) | |
| context_vars.current_share_server.set(share_url) | |
| if ( | |
| context_vars.current_project.get() is None | |
| or context_vars.current_project.get() != project | |
| ): | |
| print(f"* Trackio project initialized: {project}") | |
| if dataset_id is not None: | |
| os.environ["TRACKIO_DATASET_ID"] = dataset_id | |
| print( | |
| f"* Trackio metrics will be synced to Hugging Face Dataset: {dataset_id}" | |
| ) | |
| if space_id is None: | |
| print(f"* Trackio metrics logged to: {TRACKIO_DIR}") | |
| if utils.is_in_notebook() and embed: | |
| base_url = share_url + "/" if share_url else url | |
| full_url = utils.get_full_url( | |
| base_url, project=project, write_token=demo.write_token | |
| ) | |
| utils.embed_url_in_notebook(full_url) | |
| else: | |
| utils.print_dashboard_instructions(project) | |
| else: | |
| deploy.create_space_if_not_exists( | |
| space_id, space_storage, dataset_id, private | |
| ) | |
| user_name, space_name = space_id.split("/") | |
| space_url = deploy.SPACE_HOST_URL.format( | |
| user_name=user_name, space_name=space_name | |
| ) | |
| print(f"* View dashboard by going to: {space_url}") | |
| if utils.is_in_notebook() and embed: | |
| utils.embed_url_in_notebook(space_url) | |
| context_vars.current_project.set(project) | |
| client = None | |
| if not space_id: | |
| client = Client(url, verbose=False) | |
| if resume == "must": | |
| if name is None: | |
| raise ValueError("Must provide a run name when resume='must'") | |
| if name not in SQLiteStorage.get_runs(project): | |
| raise ValueError(f"Run '{name}' does not exist in project '{project}'") | |
| resumed = True | |
| elif resume == "allow": | |
| resumed = name is not None and name in SQLiteStorage.get_runs(project) | |
| elif resume == "never": | |
| if name is not None and name in SQLiteStorage.get_runs(project): | |
| warnings.warn( | |
| f"* Warning: resume='never' but a run '{name}' already exists in " | |
| f"project '{project}'. Generating a new name and instead. If you want " | |
| "to resume this run, call init() with resume='must' or resume='allow'." | |
| ) | |
| name = None | |
| resumed = False | |
| else: | |
| raise ValueError("resume must be one of: 'must', 'allow', or 'never'") | |
| run = Run( | |
| url=url, | |
| project=project, | |
| client=client, | |
| name=name, | |
| group=group, | |
| config=config, | |
| space_id=space_id, | |
| ) | |
| if resumed: | |
| print(f"* Resumed existing run: {run.name}") | |
| else: | |
| print(f"* Created new run: {run.name}") | |
| context_vars.current_run.set(run) | |
| globals()["config"] = run.config | |
| return run | |
| def log(metrics: dict, step: int | None = None) -> None: | |
| """ | |
| Logs metrics to the current run. | |
| Args: | |
| metrics (`dict`): | |
| A dictionary of metrics to log. | |
| step (`int`, *optional*): | |
| The step number. If not provided, the step will be incremented | |
| automatically. | |
| """ | |
| run = context_vars.current_run.get() | |
| if run is None: | |
| raise RuntimeError("Call trackio.init() before trackio.log().") | |
| run.log( | |
| metrics=metrics, | |
| step=step, | |
| ) | |
| def finish(): | |
| """ | |
| Finishes the current run. | |
| """ | |
| run = context_vars.current_run.get() | |
| if run is None: | |
| raise RuntimeError("Call trackio.init() before trackio.finish().") | |
| run.finish() | |
| def show( | |
| project: str | None = None, | |
| theme: str | ThemeClass | None = None, | |
| mcp_server: bool | None = None, | |
| ): | |
| """ | |
| Launches the Trackio dashboard. | |
| Args: | |
| project (`str`, *optional*): | |
| The name of the project whose runs to show. If not provided, all projects | |
| will be shown and the user can select one. | |
| theme (`str` or `ThemeClass`, *optional*): | |
| A Gradio Theme to use for the dashboard instead of the default Gradio theme, | |
| can be a built-in theme (e.g. `'soft'`, `'citrus'`), a theme from the Hub | |
| (e.g. `"gstaff/xkcd"`), or a custom Theme class. If not provided, the | |
| `TRACKIO_THEME` environment variable will be used, or if that is not set, the | |
| default Gradio theme will be used. | |
| mcp_server (`bool`, *optional*): | |
| If `True`, the Trackio dashboard will be set up as an MCP server and certain | |
| 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). | |
| """ | |
| theme = theme or os.environ.get("TRACKIO_THEME", DEFAULT_THEME) | |
| if theme != DEFAULT_THEME: | |
| # TODO: It's a little hacky to reproduce this theme-setting logic from Gradio Blocks, | |
| # but in Gradio 6.0, the theme will be set in `launch()` instead, which means that we | |
| # will be able to remove this code. | |
| if isinstance(theme, str): | |
| if theme.lower() in BUILT_IN_THEMES: | |
| theme = BUILT_IN_THEMES[theme.lower()] | |
| else: | |
| try: | |
| theme = ThemeClass.from_hub(theme) | |
| except Exception as e: | |
| warnings.warn(f"Cannot load {theme}. Caught Exception: {str(e)}") | |
| theme = DefaultTheme() | |
| if not isinstance(theme, ThemeClass): | |
| warnings.warn("Theme should be a class loaded from gradio.themes") | |
| theme = DefaultTheme() | |
| demo.theme: ThemeClass = theme | |
| demo.theme_css = theme._get_theme_css() | |
| demo.stylesheets = theme._stylesheets | |
| theme_hasher = hashlib.sha256() | |
| theme_hasher.update(demo.theme_css.encode("utf-8")) | |
| demo.theme_hash = theme_hasher.hexdigest() | |
| _mcp_server = ( | |
| mcp_server | |
| if mcp_server is not None | |
| else os.environ.get("GRADIO_MCP_SERVER", "False") == "True" | |
| ) | |
| _, url, share_url = demo.launch( | |
| show_api=_mcp_server, | |
| quiet=True, | |
| inline=False, | |
| prevent_thread_lock=True, | |
| favicon_path=TRACKIO_LOGO_DIR / "trackio_logo_light.png", | |
| allowed_paths=[TRACKIO_LOGO_DIR, TRACKIO_DIR], | |
| mcp_server=_mcp_server, | |
| ) | |
| base_url = share_url + "/" if share_url else url | |
| full_url = utils.get_full_url( | |
| base_url, project=project, write_token=demo.write_token | |
| ) | |
| if not utils.is_in_notebook(): | |
| print(f"* Trackio UI launched at: {full_url}") | |
| webbrowser.open(full_url) | |
| utils.block_main_thread_until_keyboard_interrupt() | |
| else: | |
| utils.embed_url_in_notebook(full_url) | |