Skip to content

Commit db6bfcf

Browse files
authored
[API] Update user agent for torchtune (#237)
Enable tracking of torchtune models via kagglehub api. This is the change for kagglehub part. PR for midtier change is http://go/kaggle-pr/34437 Fix = [b/381021004](https://buganizer.corp.google.com/issues/381021004) Note: To test this change without midtier changes in production, you need to have a midtier running locally and add breakpoint to review the value for `userAgent` in `ProtoHttpRuleController.ApiV1.cs.` With that: - Creating a python virtual environment - Import this kagglehub version `pip install git+https://github.com/Kaggle/kagglehub.git@nan_torch` - Import torchtune and its required dependencies - set environment variables export KAGGLE_USERNAME=[your_user_name] export KAGGLE_KEY=local_api_token - in virtual environment, run command `tune download "google/bert/tensorFlow2/answer-equivalence-bem" --source "kaggle"` - verify user agent has torchtune and its associated version <img width="650" alt="Screenshot 2025-03-30 at 12 18 49 PM" src="https://github.com/user-attachments/assets/3993667f-489c-415a-9446-073bc926089c" />
1 parent 3247d8a commit db6bfcf

File tree

4 files changed

+27
-6
lines changed

4 files changed

+27
-6
lines changed

CHANGELOG.md

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,10 @@
44

55
*
66

7+
## v0.3.11 (April 1, 2025)
8+
9+
* Add `torchtune` as user-agent ([#237](https://github.com/Kaggle/kagglehub/pull/237))
10+
711
## v0.3.10 (February 25, 2025)
812

913
* Fix Colab environment detection logic ([#230](https://github.com/Kaggle/kagglehub/pull/230))

src/kagglehub/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
__version__ = "0.3.10"
1+
__version__ = "0.3.11"
22

33
import kagglehub.logger # configures the library logger.
44
from kagglehub import colab_cache_resolver, http_resolver, kaggle_cache_resolver, registry

src/kagglehub/clients.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -73,11 +73,10 @@ def get_user_agent() -> str:
7373
str: user agent information.
7474
"""
7575
user_agents = [f"kagglehub/{kagglehub.__version__}"]
76-
77-
for keras_lib in ("keras_hub", "keras_nlp", "keras_cv", "keras"):
78-
keras_info = search_lib_in_call_stack(keras_lib)
79-
if keras_info is not None:
80-
user_agents.append(keras_info)
76+
for lib in ("keras_hub", "keras_nlp", "keras_cv", "keras", "torchtune"):
77+
lib_info = search_lib_in_call_stack(lib)
78+
if lib_info is not None:
79+
user_agents.append(lib_info)
8180
break
8281

8382
# Add an appropriate data loader user agent for kagglehub.dataset_load calls

tests/test_kaggle_api_client.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,3 +138,21 @@ def test_get_user_agent_keras_hub(
138138
mock_is_module.return_value = True
139139
mock_version.return_value = "0.17.0"
140140
self.assertEqual(clients.get_user_agent(), f"kagglehub/{kagglehub.__version__} keras_hub/0.17.0")
141+
142+
@patch("importlib.metadata.version")
143+
@patch("inspect.ismodule")
144+
@patch("inspect.stack")
145+
def test_get_user_agent_torch_tune(
146+
self, mock_stack: MagicMock, mock_is_module: MagicMock, mock_version: MagicMock
147+
) -> None:
148+
# Mock the call stack and version information.
149+
mock_stack.return_value = [
150+
MagicMock(frame=MagicMock(__name__="kagglehub.clients")),
151+
MagicMock(frame=MagicMock(__name__="kagglehub.models_helpers")),
152+
MagicMock(frame=MagicMock(__name__="kagglehub.models")),
153+
MagicMock(frame=MagicMock(__name__="torchtune.src.utils.preset_utils")),
154+
MagicMock(frame=MagicMock(None)),
155+
]
156+
mock_is_module.return_value = True
157+
mock_version.return_value = "0.18.0"
158+
self.assertEqual(clients.get_user_agent(), f"kagglehub/{kagglehub.__version__} torchtune/0.18.0")

0 commit comments

Comments
 (0)