feat: metadata for tracing integrations (#85)

* metadata for tracing integrations
* tests + mocks

Author: Borda

requirements.txt

@@ -1,5 +1,5 @@
 # NOTE: once we add more dependencies, consider update dependabot to check for updates
 
-lightning-sdk >=0.2.7
+lightning-sdk >=0.2.9
 lightning-utilities
 joblib

src/litmodels/integrations/checkpoints.py

@@ -1,3 +1,4 @@
+import inspect
 import queue
 import threading
 from abc import ABC
@@ -82,9 +83,9 @@ def _worker_loop(self) -> None:
                 break
             action, detail = task
             if action == Action.UPLOAD:
-                registry_name, filepath = detail
+                registry_name, filepath, metadata = detail
                 try:
-                    upload_model(registry_name, filepath)
+                    upload_model(name=registry_name, model=filepath, metadata=metadata)
                     rank_zero_debug(f"Finished uploading: {filepath}")
                 except Exception as ex:
                     rank_zero_warn(f"Upload failed {filepath}: {ex}")
@@ -103,10 +104,10 @@ def _worker_loop(self) -> None:
                 rank_zero_warn(f"Unknown task: {task}")
             self.task_queue.task_done()
 
-    def queue_upload(self, registry_name: str, filepath: str) -> None:
+    def queue_upload(self, registry_name: str, filepath: str, metadata: Optional[dict] = None) -> None:
         """Queue an upload task."""
         self.upload_count += 1
-        self.task_queue.put((Action.UPLOAD, (registry_name, filepath)))
+        self.task_queue.put((Action.UPLOAD, (registry_name, filepath, metadata)))
         rank_zero_debug(f"Queued upload: {filepath} (pending uploads: {self.upload_count})")
 
     def queue_remove(self, trainer: "pl.Trainer", filepath: str) -> None:
@@ -148,15 +149,22 @@ def __init__(self, model_name: Optional[str]) -> None:
         self._model_manager = ModelManager()
 
     @rank_zero_only
-    def _upload_model(self, filepath: str) -> None:
+    def _upload_model(self, filepath: str, metadata: Optional[dict] = None) -> None:
         # todo: use filename as version but need to validate that such version does not exists yet
         if not self.model_registry:
             raise RuntimeError(
                 "Model name is not specified neither updated by `setup` method via Trainer."
                 " Please set the model name before uploading or ensure that `setup` method is called."
             )
+        if not metadata:
+            metadata = {}
+        # Add the integration name to the metadata
+        mro = inspect.getmro(type(self))
+        abc_index = mro.index(LitModelCheckpointMixin)
+        ckpt_class = mro[abc_index - 1]
+        metadata.update({"litModels_integration": ckpt_class.__name__})
         # Add to queue instead of uploading directly
-        get_model_manager().queue_upload(self.model_registry, filepath)
+        get_model_manager().queue_upload(registry_name=self.model_registry, filepath=filepath, metadata=metadata)
 
     @rank_zero_only
     def _remove_model(self, trainer: "pl.Trainer", filepath: str) -> None:

src/litmodels/integrations/duplicate.py

@@ -4,11 +4,22 @@
 from pathlib import Path
 from typing import Optional
 
-from litmodels import upload_model
+from lightning_utilities import module_available
+
+from litmodels.io import upload_model_files
+
+if module_available("huggingface_hub"):
+    from huggingface_hub import snapshot_download
+else:
+    snapshot_download = None
 
 
 def duplicate_hf_model(
-    hf_model: str, lit_model: Optional[str] = None, local_workdir: Optional[str] = None, verbose: int = 1
+    hf_model: str,
+    lit_model: Optional[str] = None,
+    local_workdir: Optional[str] = None,
+    verbose: int = 1,
+    metadata: Optional[dict] = None,
 ) -> str:
     """Downloads the model from Hugging Face and uploads it to Lightning Cloud.
 
@@ -18,13 +29,12 @@ def duplicate_hf_model(
         local_workdir:
             The local working directory to use for the duplication process. If not set a temp folder will be created.
         verbose: Shot a progress bar for the upload.
+        metadata: Optional metadata to attach to the model. If not provided, a default metadata will be used.
 
     Returns:
         The name of the duplicated model in Lightning Cloud.
     """
-    try:
-        from huggingface_hub import snapshot_download
-    except ModuleNotFoundError:
+    if not snapshot_download:
         raise ModuleNotFoundError(
             "Hugging Face Hub is not installed. Please install it with `pip install huggingface_hub`."
         )
@@ -52,5 +62,8 @@ def duplicate_hf_model(
     # Upload the model to Lightning Cloud
     if not lit_model:
         lit_model = model_name
-    model = upload_model(name=lit_model, model=local_workdir / model_name, verbose=verbose)
+    if not metadata:
+        metadata = {}
+    metadata.update({"litModels_integration": "duplicate_hf_model", "hf_model": hf_model})
+    model = upload_model_files(name=lit_model, path=local_workdir / model_name, verbose=verbose, metadata=metadata)
     return model.name

src/litmodels/integrations/mixins.py

@@ -5,7 +5,7 @@
 import warnings
 from abc import ABC
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union
 
 from lightning_utilities.core.rank_zero import rank_zero_warn
 
@@ -55,27 +55,45 @@ def _setup(
             temp_folder = tempfile.mkdtemp()
         return name, model_name, temp_folder
 
+    def _upload_model_files(
+        self, name: str, path: Union[str, Path, List[Union[str, Path]]], metadata: Optional[dict] = None
+    ) -> None:
+        """Upload the model files to the registry."""
+        if not metadata:
+            metadata = {}
+        # Add the integration name to the metadata
+        mro = inspect.getmro(type(self))
+        abc_index = mro.index(ModelRegistryMixin)
+        mixin_class = mro[abc_index - 1]
+        metadata.update({"litModels_integration": mixin_class.__name__})
+        upload_model_files(name=name, path=path, metadata=metadata)
+
 
 class PickleRegistryMixin(ModelRegistryMixin):
     """Mixin for pickle registry integration."""
 
     def upload_model(
-        self, name: Optional[str] = None, version: Optional[str] = None, temp_folder: Union[str, Path, None] = None
+        self,
+        name: Optional[str] = None,
+        version: Optional[str] = None,
+        temp_folder: Union[str, Path, None] = None,
+        metadata: Optional[dict] = None,
     ) -> None:
         """Push the model to the registry.
 
         Args:
             name: The name of the model. If not use the class name.
             version: The version of the model. If None, the latest version is used.
             temp_folder: The temporary folder to save the model. If None, a default temporary folder is used.
+            metadata: Optional metadata to attach to the model. If not provided, a default metadata will be used.
         """
         name, model_name, temp_folder = self._setup(name, temp_folder)
         pickle_path = Path(temp_folder) / f"{model_name}.pkl"
         with open(pickle_path, "wb") as fp:
             pickle.dump(self, fp, protocol=pickle.HIGHEST_PROTOCOL)
         if version:
             name = f"{name}:{version}"
-        upload_model_files(name=name, path=pickle_path)
+        self._upload_model_files(name=name, path=pickle_path, metadata=metadata)
 
     @classmethod
     def download_model(
@@ -128,14 +146,19 @@ def __new__(cls, *args: Any, **kwargs: Any) -> "torch.nn.Module":
         return instance
 
     def upload_model(
-        self, name: Optional[str] = None, version: Optional[str] = None, temp_folder: Union[str, Path, None] = None
+        self,
+        name: Optional[str] = None,
+        version: Optional[str] = None,
+        temp_folder: Union[str, Path, None] = None,
+        metadata: Optional[dict] = None,
     ) -> None:
         """Push the model to the registry.
 
         Args:
             name: The name of the model. If not use the class name.
             version: The version of the model. If None, the latest version is used.
             temp_folder: The temporary folder to save the model. If None, a default temporary folder is used.
+            metadata: Optional metadata to attach to the model. If not provided, a default metadata will be used.
         """
         import torch
 
@@ -145,17 +168,18 @@ def upload_model(
 
         name, model_name, temp_folder = self._setup(name, temp_folder)
 
+        init_kwargs_path = None
         if self.__init_kwargs:
             try:
                 # Save the model arguments to a JSON file
                 init_kwargs_path = Path(temp_folder) / f"{model_name}__init_kwargs.json"
                 with open(init_kwargs_path, "w") as fp:
                     json.dump(self.__init_kwargs, fp)
-            except Exception as e:
+            except Exception as ex:
                 raise RuntimeError(
-                    f"Failed to save model arguments: {e}."
+                    f"Failed to save model arguments: {ex}."
                     " Ensure the model's arguments are JSON serializable or use `PickleRegistryMixin`."
-                ) from e
+                ) from ex
         elif not hasattr(self, "__init_kwargs"):
             rank_zero_warn(
                 "The child class is missing `__init_kwargs`."
@@ -168,7 +192,10 @@ def upload_model(
         model_registry = f"{name}:{version}" if version else name
         # todo: consider creating another temp folder and copying these two files
         # todo: updating SDK to support uploading just specific files
-        upload_model_files(name=model_registry, path=[torch_state_dict_path, init_kwargs_path])
+        uploaded_files = [torch_state_dict_path]
+        if init_kwargs_path:
+            uploaded_files.append(init_kwargs_path)
+        self._upload_model_files(name=model_registry, path=uploaded_files, metadata=metadata)
 
     @classmethod
     def download_model(

src/litmodels/io/cloud.py

@@ -10,6 +10,8 @@
 from lightning_sdk.models import download_model as sdk_download_model
 from lightning_sdk.models import upload_model as sdk_upload_model
 
+import litmodels
+
 if TYPE_CHECKING:
     from lightning_sdk.models import UploadedModelInfo
 
@@ -46,6 +48,7 @@ def upload_model_files(
     progress_bar: bool = True,
     cloud_account: Optional[str] = None,
     verbose: Union[bool, int] = 1,
+    metadata: Optional[Dict[str, str]] = None,
 ) -> "UploadedModelInfo":
     """Upload a local checkpoint file to the model store.
 
@@ -57,13 +60,18 @@ def upload_model_files(
         cloud_account: The name of the cloud account to store the Model in. Only required if it can't be determined
             automatically.
         verbose: Whether to print a link to the uploaded model. If set to 0, no link will be printed.
+        metadata: Optional metadata to attach to the model. If not provided, a default metadata will be used.
 
     """
+    if not metadata:
+        metadata = {}
+    metadata.update({"litModels": litmodels.__version__})
     info = sdk_upload_model(
         name=name,
         path=path,
         progress_bar=progress_bar,
         cloud_account=cloud_account,
+        metadata=metadata,
     )
     if verbose:
         _print_model_link(name, verbose)

src/litmodels/io/gateway.py

@@ -1,7 +1,7 @@
 import os
 import tempfile
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 import joblib
 from lightning_utilities import module_available
@@ -24,6 +24,7 @@ def upload_model(
     cloud_account: Optional[str] = None,
     staging_dir: Optional[str] = None,
     verbose: Union[bool, int] = 1,
+    metadata: Optional[Dict[str, str]] = None,
 ) -> "UploadedModelInfo":
     """Upload a checkpoint to the model store.
 
@@ -37,6 +38,7 @@ def upload_model(
         staging_dir: A directory where the model can be saved temporarily. If not provided, a temporary directory will
             be created and used.
         verbose: Whether to print some additional information about the uploaded model.
+        metadata: Optional metadata to attach to the model. If not provided, a default metadata will be used.
 
     """
     if not staging_dir:
@@ -62,6 +64,7 @@ def upload_model(
         progress_bar=progress_bar,
         cloud_account=cloud_account,
         verbose=verbose,
+        metadata=metadata,
     )
 
 

tests/integrations/test_checkpoints.py

@@ -2,6 +2,7 @@
 import re
 from unittest import mock
 
+import litmodels
 import pytest
 
 from tests.integrations import _SKIP_IF_LIGHTNING_MISSING, _SKIP_IF_PYTORCHLIGHTNING_MISSING
@@ -80,6 +81,7 @@ def test_lightning_checkpoint_callback(mock_auth, mock_upload_model, monkeypatch
         path=mock.ANY,
         progress_bar=True,
         cloud_account=None,
+        metadata={"litModels_integration": LitModelCheckpoint.__name__, "litModels": litmodels.__version__},
     )
     assert mock_upload_model.call_args_list == [expected_call] * 2
 

tests/integrations/test_cloud.py

@@ -10,6 +10,7 @@
 from lightning_sdk.lightning_cloud.rest_client import GridRestClient
 from lightning_sdk.utils.resolve import _resolve_teamspace
 from litmodels import download_model, upload_model
+from litmodels.integrations.duplicate import duplicate_hf_model
 from litmodels.integrations.mixins import PickleRegistryMixin, PyTorchRegistryMixin
 from litmodels.io.cloud import _list_available_teamspaces
 
@@ -288,6 +289,26 @@ def test_pytorch_mixin_push_and_pull():
     _cleanup_model(teamspace, model_name, expected_num_versions=1)
 
 
+@pytest.mark.cloud()
+def test_duplicate_real_hf_model(tmp_path):
+    """Verify that the HF model can be duplicated to the teamspace"""
+
+    # model name with random hash
+    model_name = f"litmodels_hf_model+{os.urandom(8).hex()}"
+    teamspace = _resolve_teamspace(org=LIT_ORG, teamspace=LIT_TEAMSPACE, user=None)
+    org_team = f"{teamspace.owner.name}/{teamspace.name}"
+
+    duplicate_hf_model(hf_model="google/t5-efficient-tiny", lit_model=f"{org_team}/{model_name}")
+
+    client = GridRestClient()
+    model = client.models_store_get_model_by_name(
+        project_owner_name=teamspace.owner.name,
+        project_name=teamspace.name,
+        model_name=model_name,
+    )
+    client.models_store_delete_model(project_id=teamspace.id, model_id=model.id)
+
+
 @pytest.mark.cloud()
 def test_list_available_teamspaces():
     teams = _list_available_teamspaces()

tests/integrations/test_duplicate.py

@@ -1,28 +1,31 @@
 import os
+from unittest import mock
 
-import pytest
-from lightning_sdk.lightning_cloud.rest_client import GridRestClient
-from lightning_sdk.utils.resolve import _resolve_teamspace
 from litmodels.integrations.duplicate import duplicate_hf_model
 
-from tests.integrations import LIT_ORG, LIT_TEAMSPACE
 
-
-@pytest.mark.cloud()
-def test_duplicate_hf_model(tmp_path):
+@mock.patch("litmodels.integrations.duplicate.snapshot_download")
+@mock.patch("litmodels.integrations.duplicate.upload_model_files")
+def test_duplicate_hf_model(mock_upload_model, mock_snapshot_download, tmp_path):
     """Verify that the HF model can be duplicated to the teamspace"""
 
+    hf_model = "google/t5-efficient-tiny"
     # model name with random hash
     model_name = f"litmodels_hf_model+{os.urandom(8).hex()}"
-    teamspace = _resolve_teamspace(org=LIT_ORG, teamspace=LIT_TEAMSPACE, user=None)
-    org_team = f"{teamspace.owner.name}/{teamspace.name}"
-
-    duplicate_hf_model(hf_model="google/t5-efficient-tiny", lit_model=f"{org_team}/{model_name}")
+    duplicate_hf_model(hf_model=hf_model, lit_model=model_name, local_workdir=str(tmp_path))
 
-    client = GridRestClient()
-    model = client.models_store_get_model_by_name(
-        project_owner_name=teamspace.owner.name,
-        project_name=teamspace.name,
-        model_name=model_name,
+    mock_snapshot_download.assert_called_with(
+        repo_id=hf_model,
+        revision="main",
+        repo_type="model",
+        local_dir=tmp_path / hf_model.replace("/", "_"),
+        local_dir_use_symlinks=True,
+        ignore_patterns=[".cache*"],
+        max_workers=os.cpu_count(),
+    )
+    mock_upload_model.assert_called_with(
+        name=f"{model_name}",
+        path=tmp_path / hf_model.replace("/", "_"),
+        metadata={"hf_model": hf_model, "litModels_integration": "duplicate_hf_model"},
+        verbose=1,
     )
-    client.models_store_delete_model(project_id=teamspace.id, model_id=model.id)