Skip to content

update metadata notation #101

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Borda merged 1 commit into from
Apr 22, 2025
Merged

update metadata notation #101

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/litmodels/integrations/checkpoints.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion src/litmodels/integrations/duplicate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
2 changes: 1 addition & 1 deletion src/litmodels/integrations/mixins.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)


Expand Down
2 changes: 1 addition & 1 deletion src/litmodels/io/gateway.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion tests/integrations/test_checkpoints.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
]
Expand Down
2 changes: 1 addition & 1 deletion tests/integrations/test_duplicate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
4 changes: 2 additions & 2 deletions tests/integrations/test_mixins.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down
2 changes: 1 addition & 1 deletion tests/test_io_cloud.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"},
)


Expand Down
Loading