From 7e010241d9a34794e0ce0dc19c1a6f0cf94ba856 Mon Sep 17 00:00:00 2001 From: Lucain Date: Mon, 17 Nov 2025 17:05:57 +0100 Subject: [PATCH] Avoid redundant calls to /whoami-v2 (#344) Co-authored-by: gradio-pr-bot --- .changeset/forty-dolls-train.md | 5 +++++ tests/test_run.py | 5 +++++ trackio/run.py | 4 ++-- trackio/ui/fns.py | 3 +++ trackio/utils.py | 21 ++++++++++++++++++--- 5 files changed, 33 insertions(+), 5 deletions(-) create mode 100644 .changeset/forty-dolls-train.md diff --git a/.changeset/forty-dolls-train.md b/.changeset/forty-dolls-train.md new file mode 100644 index 0000000..d8cd500 --- /dev/null +++ b/.changeset/forty-dolls-train.md @@ -0,0 +1,5 @@ +--- +"trackio": patch +--- + +feat:Avoid redundant calls to /whoami-v2 diff --git a/tests/test_run.py b/tests/test_run.py index 51efd68..af590c7 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -4,6 +4,7 @@ import pytest from trackio import Run, init +from trackio.utils import _cached_whoami class DummyClient: @@ -99,6 +100,8 @@ def test_run_name_generation_with_space_id(mock_time, mock_whoami, temp_dir): ) assert run.name == "testuser-1234567890" + _cached_whoami.cache_clear() + def test_reserved_config_keys_rejected(temp_dir): with pytest.raises(ValueError, match="Config key '_test' is reserved"): @@ -130,6 +133,8 @@ def test_automatic_username_and_timestamp_added(mock_whoami, temp_dir): created_time = datetime.fromisoformat(run.config["_Created"]) assert created_time.tzinfo is not None + _cached_whoami.cache_clear() + def test_run_group_added(temp_dir): run = Run( diff --git a/trackio/run.py b/trackio/run.py index 00bc160..1c40fd4 100644 --- a/trackio/run.py +++ b/trackio/run.py @@ -12,6 +12,7 @@ from trackio.sqlite_storage import SQLiteStorage from trackio.table import Table from trackio.typehints import LogEntry, UploadEntry +from trackio.utils import _get_default_namespace BATCH_SEND_INTERVAL = 0.5 @@ -62,8 +63,7 @@ def __init__( def _get_username(self) -> str | None: """Get the current HuggingFace username if logged in, otherwise None.""" try: - who = huggingface_hub.whoami() - return who["name"] if who else None + return _get_default_namespace() except Exception: return None diff --git a/trackio/ui/fns.py b/trackio/ui/fns.py index 8e40fd7..c3234fa 100644 --- a/trackio/ui/fns.py +++ b/trackio/ui/fns.py @@ -1,6 +1,7 @@ """Shared functions for the Trackio UI.""" import os +from functools import lru_cache import gradio as gr import huggingface_hub as hf @@ -82,6 +83,7 @@ def update_navbar_value(project_dd, request: gr.Request): ) +@lru_cache(maxsize=32) def check_hf_token_has_write_access(hf_token: str | None) -> None: """ Checks to see if the provided hf_token is valid and has write access to the Space @@ -137,6 +139,7 @@ def check_hf_token_has_write_access(hf_token: str | None) -> None: ) +@lru_cache(maxsize=32) def check_oauth_token_has_write_access(oauth_token: str | None) -> None: """ Checks to see if the oauth token provided via Gradio's OAuth is valid and has write access diff --git a/trackio/utils.py b/trackio/utils.py index b4662ea..6a4396f 100644 --- a/trackio/utils.py +++ b/trackio/utils.py @@ -3,6 +3,7 @@ import re import time from datetime import datetime, timezone +from functools import lru_cache from pathlib import Path from typing import TYPE_CHECKING @@ -144,7 +145,7 @@ def generate_readable_name(used_names: list[str], space_id: str | None = None) - If space_id is provided, generates username-timestamp format instead. """ if space_id is not None: - username = huggingface_hub.whoami()["name"] + username = _get_default_namespace() timestamp = int(time.time()) return f"{username}-{timestamp}" adjectives = [ @@ -418,10 +419,10 @@ def preprocess_space_and_dataset_ids( space_id: str | None, dataset_id: str | None ) -> tuple[str | None, str | None]: if space_id is not None and "/" not in space_id: - username = huggingface_hub.whoami()["name"] + username = _get_default_namespace() space_id = f"{username}/{space_id}" if dataset_id is not None and "/" not in dataset_id: - username = huggingface_hub.whoami()["name"] + username = _get_default_namespace() dataset_id = f"{username}/{dataset_id}" if space_id is not None and dataset_id is None: dataset_id = f"{space_id}-dataset" @@ -865,3 +866,17 @@ def get_space() -> str | None: def ordered_subset(items: list[str], subset: list[str] | None) -> list[str]: subset_set = set(subset or []) return [item for item in items if item in subset_set] + + +def _get_default_namespace() -> str: + """Get the default namespace (username). + + This function uses caching to avoid repeated API calls to /whoami-v2. + """ + token = huggingface_hub.get_token() + return _cached_whoami(token)["name"] + + +@lru_cache(maxsize=32) +def _cached_whoami(token: str | None) -> dict: + return huggingface_hub.whoami(token=token)