diff --git a/src/litmodels/integrations/checkpoints.py b/src/litmodels/integrations/checkpoints.py index 7d23dfc..ddd4154 100644 --- a/src/litmodels/integrations/checkpoints.py +++ b/src/litmodels/integrations/checkpoints.py @@ -193,7 +193,7 @@ def _upload_model(self, trainer: "pl.Trainer", filepath: Union[str, Path], metad mro = inspect.getmro(type(self)) abc_index = mro.index(LitModelCheckpointMixin) ckpt_class = mro[abc_index - 1] - metadata.update({"litModels_integration": ckpt_class.__name__}) + metadata.update({"litModels.integration": ckpt_class.__name__}) # Add to queue instead of uploading directly get_model_manager().queue_upload(registry_name=model_registry, filepath=filepath, metadata=metadata) if self._clear_all_local: diff --git a/src/litmodels/integrations/duplicate.py b/src/litmodels/integrations/duplicate.py index 90eb74f..80123c4 100644 --- a/src/litmodels/integrations/duplicate.py +++ b/src/litmodels/integrations/duplicate.py @@ -64,6 +64,6 @@ def duplicate_hf_model( lit_model = model_name if not metadata: metadata = {} - metadata.update({"litModels_integration": "duplicate_hf_model", "hf_model": hf_model}) + 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 diff --git a/src/litmodels/integrations/mixins.py b/src/litmodels/integrations/mixins.py index ff28ed8..64e0cba 100644 --- a/src/litmodels/integrations/mixins.py +++ b/src/litmodels/integrations/mixins.py @@ -65,7 +65,7 @@ def _upload_model_files( mro = inspect.getmro(type(self)) abc_index = mro.index(ModelRegistryMixin) mixin_class = mro[abc_index - 1] - metadata.update({"litModels_integration": mixin_class.__name__}) + metadata.update({"litModels.integration": mixin_class.__name__}) upload_model_files(name=name, path=path, metadata=metadata) diff --git a/src/litmodels/io/gateway.py b/src/litmodels/io/gateway.py index ea909eb..97626d7 100644 --- a/src/litmodels/io/gateway.py +++ b/src/litmodels/io/gateway.py @@ -103,7 +103,7 @@ def save_model( if not metadata: metadata = {} - metadata.update({"litModels_integration": "save_model"}) + metadata.update({"litModels.integration": "save_model"}) return upload_model( model=path, diff --git a/tests/integrations/test_checkpoints.py b/tests/integrations/test_checkpoints.py index da65529..a3aec32 100644 --- a/tests/integrations/test_checkpoints.py +++ b/tests/integrations/test_checkpoints.py @@ -104,7 +104,7 @@ def test_lightning_checkpoint_callback( path=mock.ANY, progress_bar=True, cloud_account=None, - metadata={"litModels_integration": LitModelCheckpoint.__name__, "litModels": litmodels.__version__}, + metadata={"litModels.integration": LitModelCheckpoint.__name__, "litModels": litmodels.__version__}, ) for v in ("epoch=0-step=64", "epoch=1-step=128") ] diff --git a/tests/integrations/test_duplicate.py b/tests/integrations/test_duplicate.py index f0d8e1f..0f517e9 100644 --- a/tests/integrations/test_duplicate.py +++ b/tests/integrations/test_duplicate.py @@ -26,6 +26,6 @@ def test_duplicate_hf_model(mock_upload_model, mock_snapshot_download, tmp_path) 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"}, + metadata={"hf_model": hf_model, "litModels.integration": "duplicate_hf_model"}, verbose=1, ) diff --git a/tests/integrations/test_mixins.py b/tests/integrations/test_mixins.py index b48517f..33592dd 100644 --- a/tests/integrations/test_mixins.py +++ b/tests/integrations/test_mixins.py @@ -24,7 +24,7 @@ def test_pickle_push_and_pull(mock_download_model, mock_upload_model, tmp_path): # The expected registry name is "dummy_model:v1" and the file should be placed in the temp folder. expected_path = tmp_path / "DummyModel.pkl" mock_upload_model.assert_called_once_with( - name="DummyModel:v1", path=expected_path, metadata={"litModels_integration": "PickleRegistryMixin"} + name="DummyModel:v1", path=expected_path, metadata={"litModels.integration": "PickleRegistryMixin"} ) # Set the mock to return the full path to the pickle file. @@ -78,7 +78,7 @@ def test_pytorch_push_and_pull(mock_download_model, mock_upload_model, torch_cla mock_upload_model.assert_called_once_with( name=torch_class.__name__, path=[tmp_path / f"{torch_class.__name__}.pth", json_path], - metadata={"litModels_integration": "PyTorchRegistryMixin"}, + metadata={"litModels.integration": "PyTorchRegistryMixin"}, ) # Prepare mocking for pull_from_registry. diff --git a/tests/test_io_cloud.py b/tests/test_io_cloud.py index 5222641..62c3779 100644 --- a/tests/test_io_cloud.py +++ b/tests/test_io_cloud.py @@ -84,7 +84,7 @@ def test_upload_model(mock_upload_model, tmp_path, model, model_path, verbose): name="org-name/teamspace/model-name", cloud_account="cluster_id", progress_bar=True, - metadata={"litModels": litmodels.__version__, "litModels_integration": "save_model"}, + metadata={"litModels": litmodels.__version__, "litModels.integration": "save_model"}, )