From d59a0f9b63774034923ec0532b81098b1ccb500d Mon Sep 17 00:00:00 2001 From: Justus Schock <12886177+justusschock@users.noreply.github.com> Date: Tue, 31 Mar 2026 16:48:24 +0200 Subject: [PATCH 1/7] update --- requirements.txt | 2 +- src/litmodels/__about__.py | 2 +- src/litmodels/io/cloud.py | 4 ++++ src/litmodels/io/gateway.py | 7 +++++++ 4 files changed, 13 insertions(+), 2 deletions(-) diff --git a/requirements.txt b/requirements.txt index ddc5da4..f90a76c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,2 @@ -lightning-sdk >=0.2.11 +lightning-sdk >=2026.03.31 lightning-utilities diff --git a/src/litmodels/__about__.py b/src/litmodels/__about__.py index c71714e..e4bd307 100644 --- a/src/litmodels/__about__.py +++ b/src/litmodels/__about__.py @@ -1,4 +1,4 @@ -__version__ = "0.1.8" +__version__ = "0.2.0" __author__ = "Lightning-AI et al." __author_email__ = "community@lightning.ai" __license__ = "Apache-2.0" diff --git a/src/litmodels/io/cloud.py b/src/litmodels/io/cloud.py index 4061c5f..65183a6 100644 --- a/src/litmodels/io/cloud.py +++ b/src/litmodels/io/cloud.py @@ -10,6 +10,7 @@ from lightning_sdk.models import delete_model as sdk_delete_model from lightning_sdk.models import download_model as sdk_download_model from lightning_sdk.models import upload_model as sdk_upload_model +from lightning_sdk.api.utils import Experiment import litmodels @@ -49,6 +50,7 @@ def upload_model_files( cloud_account: Optional[str] = None, verbose: Union[bool, int] = 1, metadata: Optional[dict[str, str]] = None, + experiment: Optional[Experiment] = None, ) -> "UploadedModelInfo": """Upload local artifact(s) to Lightning Cloud using the SDK. @@ -59,6 +61,7 @@ def upload_model_files( cloud_account: Optional cloud account to store the model in, when it cannot be auto-resolved. verbose: Verbosity for printing the model link (0 = no output, 1 = print once, 2 = print always). metadata: Optional metadata to attach to the model/version. The package version is added automatically. + experiment: Optional experiment producing this model Returns: UploadedModelInfo describing the created or updated model version. @@ -72,6 +75,7 @@ def upload_model_files( progress_bar=progress_bar, cloud_account=cloud_account, metadata=metadata, + experiment=experiment ) if verbose: _print_model_link(name, verbose) diff --git a/src/litmodels/io/gateway.py b/src/litmodels/io/gateway.py index b0d7d2d..8eacf44 100644 --- a/src/litmodels/io/gateway.py +++ b/src/litmodels/io/gateway.py @@ -5,6 +5,7 @@ from litmodels.io.cloud import download_model_files, upload_model_files from litmodels.io.utils import _KERAS_AVAILABLE, _PYTORCH_AVAILABLE, dump_pickle, load_pickle +from lightning_sdk.api.utils import Experiment if _PYTORCH_AVAILABLE: import torch @@ -23,6 +24,7 @@ def upload_model( cloud_account: Optional[str] = None, verbose: Union[bool, int] = 1, metadata: Optional[dict[str, str]] = None, + experiment: Optional[Experiment] = None, ) -> "UploadedModelInfo": """Upload a local artifact (file or directory) to Lightning Cloud Models. @@ -34,6 +36,7 @@ def upload_model( cloud_account: Optional cloud account to store the model in, when it cannot be auto-resolved. verbose: Verbosity of informational output (0 = silent, 1 = print link once, 2 = print link always). metadata: Optional metadata key/value pairs to attach to the uploaded model/version. + experiment: Optional experiment producing this model. Returns: UploadedModelInfo describing the created or updated model version. @@ -54,6 +57,7 @@ def upload_model( cloud_account=cloud_account, verbose=verbose, metadata=metadata, + experiment=experiment, ) @@ -65,6 +69,7 @@ def save_model( staging_dir: Optional[str] = None, verbose: Union[bool, int] = 1, metadata: Optional[dict[str, str]] = None, + experiment: Optional[Experiment] = None, ) -> "UploadedModelInfo": """Serialize an in-memory model and upload it to Lightning Cloud Models. @@ -83,6 +88,7 @@ def save_model( verbose: Verbosity of informational output (0 = silent, 1 = print link once, 2 = print link always). metadata: Optional metadata key/value pairs to attach to the uploaded model/version. Integration markers are added automatically. + experiment: Optional experiment producing this model. Returns: UploadedModelInfo describing the created or updated model version. @@ -125,6 +131,7 @@ def save_model( cloud_account=cloud_account, verbose=verbose, metadata=metadata, + experiment=experiment, ) From 35a814967cb298980b9055f5eee435c95eaceea2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 31 Mar 2026 14:50:08 +0000 Subject: [PATCH 2/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/litmodels/io/cloud.py | 4 ++-- src/litmodels/io/gateway.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/litmodels/io/cloud.py b/src/litmodels/io/cloud.py index 65183a6..02d9616 100644 --- a/src/litmodels/io/cloud.py +++ b/src/litmodels/io/cloud.py @@ -5,12 +5,12 @@ from pathlib import Path from typing import TYPE_CHECKING, Optional, Union +from lightning_sdk.api.utils import Experiment from lightning_sdk.lightning_cloud.env import LIGHTNING_CLOUD_URL from lightning_sdk.models import _extend_model_name_with_teamspace, _parse_org_teamspace_model_version from lightning_sdk.models import delete_model as sdk_delete_model from lightning_sdk.models import download_model as sdk_download_model from lightning_sdk.models import upload_model as sdk_upload_model -from lightning_sdk.api.utils import Experiment import litmodels @@ -75,7 +75,7 @@ def upload_model_files( progress_bar=progress_bar, cloud_account=cloud_account, metadata=metadata, - experiment=experiment + experiment=experiment, ) if verbose: _print_model_link(name, verbose) diff --git a/src/litmodels/io/gateway.py b/src/litmodels/io/gateway.py index 8eacf44..a67449b 100644 --- a/src/litmodels/io/gateway.py +++ b/src/litmodels/io/gateway.py @@ -3,9 +3,10 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Optional, Union +from lightning_sdk.api.utils import Experiment + from litmodels.io.cloud import download_model_files, upload_model_files from litmodels.io.utils import _KERAS_AVAILABLE, _PYTORCH_AVAILABLE, dump_pickle, load_pickle -from lightning_sdk.api.utils import Experiment if _PYTORCH_AVAILABLE: import torch From 0be921985fd7b95c58b54aa29a53ea383aa19a46 Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Wed, 1 Apr 2026 13:03:08 +0000 Subject: [PATCH 3/7] update --- requirements.txt | 3 +- setup.py | 2 +- src/litmodels/__init__.py | 3 +- src/litmodels/io/__init__.py | 5 +- src/litmodels/io/cloud.py | 141 +---------------- src/litmodels/io/gateway.py | 192 +----------------------- src/litmodels/io/utils.py | 87 +++-------- tests/integrations/test_checkpoints.py | 17 +-- tests/test_io_cloud.py | 200 +++++-------------------- 9 files changed, 73 insertions(+), 577 deletions(-) diff --git a/requirements.txt b/requirements.txt index f90a76c..56ce879 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ +litlogger >=2026.03.17 lightning-sdk >=2026.03.31 -lightning-utilities +lightning-utilities<=0.15.3 diff --git a/setup.py b/setup.py index 053a597..dce9bc3 100755 --- a/setup.py +++ b/setup.py @@ -68,7 +68,7 @@ def _prepare_extras(requirements_dir: str = _PATH_REQUIRES, skip_files: tuple = include_package_data=True, zip_safe=False, keywords=["deep learning", "pytorch", "AI"], - python_requires=">=3.8", + python_requires=">=3.10", setup_requires=["wheel"], install_requires=_load_requirements(), extras_require=_prepare_extras(), diff --git a/src/litmodels/__init__.py b/src/litmodels/__init__.py index 479a80b..7d9a6db 100644 --- a/src/litmodels/__init__.py +++ b/src/litmodels/__init__.py @@ -3,10 +3,9 @@ import os from litmodels.__about__ import * # noqa: F401, F403 +from litlogger.models import download_model, load_model, save_model, upload_model, upload_model_files _PACKAGE_ROOT = os.path.dirname(__file__) _PROJECT_ROOT = os.path.dirname(_PACKAGE_ROOT) -from litmodels.io import download_model, load_model, save_model, upload_model, upload_model_files # noqa: F401 - __all__ = ["download_model", "upload_model", "load_model", "save_model"] diff --git a/src/litmodels/io/__init__.py b/src/litmodels/io/__init__.py index 363f207..2ccad08 100644 --- a/src/litmodels/io/__init__.py +++ b/src/litmodels/io/__init__.py @@ -1,6 +1,5 @@ """Root package for Input/output.""" -from litmodels.io.cloud import download_model_files, upload_model_files # noqa: F401 -from litmodels.io.gateway import download_model, load_model, save_model, upload_model +from litlogger.models import download_model, download_model_files, load_model, save_model, upload_model, upload_model_files -__all__ = ["download_model", "upload_model", "upload_model_files", "load_model", "save_model"] +__all__ = ["download_model", "download_model_files", "upload_model", "upload_model_files", "load_model", "save_model"] diff --git a/src/litmodels/io/cloud.py b/src/litmodels/io/cloud.py index 02d9616..d2952af 100644 --- a/src/litmodels/io/cloud.py +++ b/src/litmodels/io/cloud.py @@ -1,140 +1,5 @@ -# Copyright The Lightning AI team. -# Licensed under the Apache License, Version 2.0 (the "License"); -# http://www.apache.org/licenses/LICENSE-2.0 -# -from pathlib import Path -from typing import TYPE_CHECKING, Optional, Union +"""Compatibility exports for the vendored litlogger model helpers.""" -from lightning_sdk.api.utils import Experiment -from lightning_sdk.lightning_cloud.env import LIGHTNING_CLOUD_URL -from lightning_sdk.models import _extend_model_name_with_teamspace, _parse_org_teamspace_model_version -from lightning_sdk.models import delete_model as sdk_delete_model -from lightning_sdk.models import download_model as sdk_download_model -from lightning_sdk.models import upload_model as sdk_upload_model +from litlogger.models.cloud import _list_available_teamspaces, delete_model_version, download_model_files, upload_model_files -import litmodels - -if TYPE_CHECKING: - from lightning_sdk.models import UploadedModelInfo - - -_SHOWED_MODEL_LINKS = [] - - -def _print_model_link(name: str, verbose: Union[bool, int]) -> None: - """Print a stable URL to the uploaded model. - - Args: - name: Model registry name. Teamspace defaults may be applied before URL construction. - verbose: Controls printing behavior: - - 0: do not print - - 1: print the link only once for a given model - - 2: always print the link - """ - name = _extend_model_name_with_teamspace(name) - org_name, teamspace_name, model_name, _ = _parse_org_teamspace_model_version(name) - - url = f"{LIGHTNING_CLOUD_URL}/{org_name}/{teamspace_name}/models/{model_name}" - msg = f"Model uploaded successfully. Link to the model: '{url}'" - if int(verbose) > 1: - print(msg) - elif url not in _SHOWED_MODEL_LINKS: - print(msg) - _SHOWED_MODEL_LINKS.append(url) - - -def upload_model_files( - name: str, - path: Union[str, Path, list[Union[str, Path]]], - progress_bar: bool = True, - cloud_account: Optional[str] = None, - verbose: Union[bool, int] = 1, - metadata: Optional[dict[str, str]] = None, - experiment: Optional[Experiment] = None, -) -> "UploadedModelInfo": - """Upload local artifact(s) to Lightning Cloud using the SDK. - - Args: - name: Model registry name in the form 'organization/teamspace/modelname[:version]'. - path: File path, directory path, or list of paths to upload. - progress_bar: Whether to show a progress bar during upload. - cloud_account: Optional cloud account to store the model in, when it cannot be auto-resolved. - verbose: Verbosity for printing the model link (0 = no output, 1 = print once, 2 = print always). - metadata: Optional metadata to attach to the model/version. The package version is added automatically. - experiment: Optional experiment producing this model - - Returns: - UploadedModelInfo describing the created or updated model version. - """ - 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, - experiment=experiment, - ) - if verbose: - _print_model_link(name, verbose) - return info - - -def download_model_files( - name: str, - download_dir: Union[str, Path] = ".", - progress_bar: bool = True, -) -> Union[str, list[str]]: - """Download artifact(s) for a model version using the SDK. - - Args: - name: Model registry name in the form 'organization/teamspace/modelname[:version]'. - download_dir: Directory where downloaded artifact(s) will be stored. Defaults to the current directory. - progress_bar: Whether to show a progress bar during download. - - Returns: - str | list[str]: Absolute path(s) to the downloaded artifact(s). - """ - return sdk_download_model( - name=name, - download_dir=download_dir, - progress_bar=progress_bar, - ) - - -def _list_available_teamspaces() -> dict[str, dict]: - """List teamspaces available to the authenticated user. - - Returns: - dict[str, dict]: Mapping of 'org/teamspace' to a metadata dictionary with details. - """ - from lightning_sdk.api import OrgApi, UserApi - from lightning_sdk.utils import resolve as sdk_resolvers - - org_api = OrgApi() - user = sdk_resolvers._get_authed_user() - teamspaces = {} - for ts in UserApi()._get_all_teamspace_memberships(""): - if ts.owner_type == "organization": - org = org_api._get_org_by_id(ts.owner_id) - teamspaces[f"{org.name}/{ts.name}"] = {"name": ts.name, "org": org.name} - elif ts.owner_type == "user": # todo: check also the name - teamspaces[f"{user.name}/{ts.name}"] = {"name": ts.name, "user": user} - else: - raise RuntimeError(f"Unknown organization type {ts.organization_type}") - return teamspaces - - -def delete_model_version( - name: str, - version: str, -) -> None: - """Delete a specific model version from the model store. - - Args: - name: Base model registry name in the form 'organization/teamspace/modelname'. - version: Identifier of the version to delete. This argument is required. - """ - sdk_delete_model(name=f"{name}:{version}") +__all__ = ["_list_available_teamspaces", "delete_model_version", "download_model_files", "upload_model_files"] diff --git a/src/litmodels/io/gateway.py b/src/litmodels/io/gateway.py index a67449b..5b00b32 100644 --- a/src/litmodels/io/gateway.py +++ b/src/litmodels/io/gateway.py @@ -1,191 +1,5 @@ -import os -import tempfile -from pathlib import Path -from typing import TYPE_CHECKING, Any, Optional, Union +"""Compatibility exports for the vendored litlogger model helpers.""" -from lightning_sdk.api.utils import Experiment +from litlogger.models import download_model, load_model, save_model, upload_model -from litmodels.io.cloud import download_model_files, upload_model_files -from litmodels.io.utils import _KERAS_AVAILABLE, _PYTORCH_AVAILABLE, dump_pickle, load_pickle - -if _PYTORCH_AVAILABLE: - import torch - -if _KERAS_AVAILABLE: - from tensorflow import keras - -if TYPE_CHECKING: - from lightning_sdk.models import UploadedModelInfo - - -def upload_model( - name: str, - model: Union[str, Path], - progress_bar: bool = True, - cloud_account: Optional[str] = None, - verbose: Union[bool, int] = 1, - metadata: Optional[dict[str, str]] = None, - experiment: Optional[Experiment] = None, -) -> "UploadedModelInfo": - """Upload a local artifact (file or directory) to Lightning Cloud Models. - - Args: - name: Model registry name in the form 'organization/teamspace/modelname[:version]'. - If the version is omitted, one may be assigned automatically by the service. - model: Path to a checkpoint file or a directory containing model artifacts. - progress_bar: Whether to show a progress bar during the upload. - cloud_account: Optional cloud account to store the model in, when it cannot be auto-resolved. - verbose: Verbosity of informational output (0 = silent, 1 = print link once, 2 = print link always). - metadata: Optional metadata key/value pairs to attach to the uploaded model/version. - experiment: Optional experiment producing this model. - - Returns: - UploadedModelInfo describing the created or updated model version. - - Raises: - ValueError: If `model` is not a filesystem path. For in-memory objects, use `save_model()` instead. - """ - if not isinstance(model, (str, Path)): - raise ValueError( - "The `model` argument should be a path to a file or folder, not an python object." - " For smooth integrations with PyTorch model, Lightning model and many more, use `save_model` instead." - ) - - return upload_model_files( - path=model, - name=name, - progress_bar=progress_bar, - cloud_account=cloud_account, - verbose=verbose, - metadata=metadata, - experiment=experiment, - ) - - -def save_model( - name: str, - model: Union["torch.nn.Module", Any], - progress_bar: bool = True, - cloud_account: Optional[str] = None, - staging_dir: Optional[str] = None, - verbose: Union[bool, int] = 1, - metadata: Optional[dict[str, str]] = None, - experiment: Optional[Experiment] = None, -) -> "UploadedModelInfo": - """Serialize an in-memory model and upload it to Lightning Cloud Models. - - Supported models: - - TorchScript (torch.jit.ScriptModule) → saved as .ts via model.save() - - PyTorch nn.Module → saved as .pth (state_dict via torch.save) - - Keras (tf.keras.Model) → saved as .keras via model.save() - - Any other Python object → saved as .pkl via pickle or joblib - - Args: - name: Model registry name in the form 'organization/teamspace/modelname[:version]'. - model: The in-memory model instance to serialize and upload. - progress_bar: Whether to show a progress bar during the upload. - cloud_account: Optional cloud account to store the model in, when it cannot be auto-resolved. - staging_dir: Optional temporary directory used for serialization. A new temp directory is created if omitted. - verbose: Verbosity of informational output (0 = silent, 1 = print link once, 2 = print link always). - metadata: Optional metadata key/value pairs to attach to the uploaded model/version. Integration markers are - added automatically. - experiment: Optional experiment producing this model. - - Returns: - UploadedModelInfo describing the created or updated model version. - - Raises: - ValueError: If `model` is a path. For file/folder uploads use `upload_model()` instead. - """ - if isinstance(model, (str, Path)): - raise ValueError( - "The `model` argument should be a PyTorch model or a Lightning model, not a path to a file." - " With file or folder path use `upload_model` instead." - ) - - if not staging_dir: - staging_dir = tempfile.mkdtemp() - # if LightningModule and isinstance(model, LightningModule): - # path = os.path.join(staging_dir, f"{model.__class__.__name__}.ckpt") - # model.save_checkpoint(path) - if _PYTORCH_AVAILABLE and isinstance(model, torch.jit.ScriptModule): - path = os.path.join(staging_dir, f"{model.__class__.__name__}.ts") - model.save(path) - elif _PYTORCH_AVAILABLE and isinstance(model, torch.nn.Module): - path = os.path.join(staging_dir, f"{model.__class__.__name__}.pth") - torch.save(model.state_dict(), path) - elif _KERAS_AVAILABLE and isinstance(model, keras.models.Model): - path = os.path.join(staging_dir, f"{model.__class__.__name__}.keras") - model.save(path) - else: - path = os.path.join(staging_dir, f"{model.__class__.__name__}.pkl") - dump_pickle(model=model, path=path) - - if not metadata: - metadata = {} - metadata.update({"litModels.integration": "save_model"}) - - return upload_model( - model=path, - name=name, - progress_bar=progress_bar, - cloud_account=cloud_account, - verbose=verbose, - metadata=metadata, - experiment=experiment, - ) - - -def download_model( - name: str, - download_dir: Union[str, Path] = ".", - progress_bar: bool = True, -) -> Union[str, list[str]]: - """Download a model version from Lightning Cloud Models to a local directory. - - Args: - name: Model registry name in the form 'organization/teamspace/modelname[:version]'. - download_dir: Directory where the artifact(s) will be stored. Defaults to the current working directory. - progress_bar: Whether to show a progress bar during the download. - - Returns: - str | list[str]: Absolute path(s) to the downloaded file(s) or directory content. - """ - return download_model_files( - name=name, - download_dir=download_dir, - progress_bar=progress_bar, - ) - - -def load_model(name: str, download_dir: str = ".") -> Any: - """Download a model and load it into memory based on its file extension. - - Supported formats: - - .ts → torch.jit.load - - .keras → keras.models.load_model - - .pkl → pickle/joblib via load_pickle - - Args: - name: Model registry name in the form 'organization/teamspace/modelname[:version]'. - download_dir: Directory to store the downloaded artifact(s) before loading. Defaults to the current directory. - - Returns: - Any: The loaded model object. - - Raises: - NotImplementedError: If multiple files are downloaded or the file extension is not supported. - """ - download_paths = download_model(name=name, download_dir=download_dir) - # filter out all Markdown, TXT and RST files - download_paths = [p for p in download_paths if Path(p).suffix.lower() not in {".md", ".txt", ".rst"}] - if len(download_paths) > 1: - raise NotImplementedError("Downloaded model with multiple files is not supported yet.") - model_path = Path(download_dir) / download_paths[0] - if model_path.suffix.lower() == ".ts": - return torch.jit.load(model_path) - if model_path.suffix.lower() == ".keras": - return keras.models.load_model(model_path) - if model_path.suffix.lower() == ".pkl": - return load_pickle(path=model_path) - raise NotImplementedError(f"Loading model from {model_path.suffix} is not supported yet.") +__all__ = ["download_model", "load_model", "save_model", "upload_model"] diff --git a/src/litmodels/io/utils.py b/src/litmodels/io/utils.py index 32f490c..fb25ca3 100644 --- a/src/litmodels/io/utils.py +++ b/src/litmodels/io/utils.py @@ -1,68 +1,19 @@ -import os -import pickle -from collections.abc import Iterator -from contextlib import contextmanager -from pathlib import Path -from typing import Any, Union - -from lightning_utilities import module_available -from lightning_utilities.core.imports import RequirementCache - - -@contextmanager -def _suppress_os_stderr() -> Iterator[None]: - devnull_fd = os.open(os.devnull, os.O_WRONLY) - old_stderr_fd = os.dup(2) - os.dup2(devnull_fd, 2) # redirect stderr (fd 2) to /dev/null - os.close(devnull_fd) - try: - yield - finally: - os.dup2(old_stderr_fd, 2) # restore stderr - os.close(old_stderr_fd) - - -_JOBLIB_AVAILABLE = module_available("joblib") -_PYTORCH_AVAILABLE = module_available("torch") -with _suppress_os_stderr(): - _TENSORFLOW_AVAILABLE = module_available("tensorflow") - _KERAS_AVAILABLE = RequirementCache("tensorflow >=2.0.0") - -if _JOBLIB_AVAILABLE: - import joblib - - -def dump_pickle(model: Any, path: Union[str, Path]) -> None: - """Serialize a Python object to disk using joblib (if available) or pickle. - - Args: - model: The object to serialize. - path: Destination file path. - - Notes: - - Uses joblib with compression (level 7) when available for smaller artifacts. - - Falls back to pickle with the highest protocol otherwise. - """ - if _JOBLIB_AVAILABLE: - joblib.dump(model, filename=path, compress=7) - else: - with open(path, "wb") as fp: - pickle.dump(model, fp, protocol=pickle.HIGHEST_PROTOCOL) - - -def load_pickle(path: Union[str, Path]) -> Any: - """Load a Python object from a joblib/pickle file. - - Args: - path: Path to the serialized artifact. - - Returns: - Any: The deserialized object. - - Warning: - Loading pickle/joblib files can execute arbitrary code. Only open files from trusted sources. - """ - if _JOBLIB_AVAILABLE: - return joblib.load(path) - with open(path, "rb") as fp: - return pickle.load(fp) +"""Compatibility exports for the vendored litlogger model helpers.""" + +from litlogger.models.serialization import ( + _JOBLIB_AVAILABLE, + _KERAS_AVAILABLE, + _PYTORCH_AVAILABLE, + _TENSORFLOW_AVAILABLE, + dump_pickle, + load_pickle, +) + +__all__ = [ + "_JOBLIB_AVAILABLE", + "_KERAS_AVAILABLE", + "_PYTORCH_AVAILABLE", + "_TENSORFLOW_AVAILABLE", + "dump_pickle", + "load_pickle", +] diff --git a/tests/integrations/test_checkpoints.py b/tests/integrations/test_checkpoints.py index a3aec32..2efbf83 100644 --- a/tests/integrations/test_checkpoints.py +++ b/tests/integrations/test_checkpoints.py @@ -4,7 +4,6 @@ import pytest -import litmodels from tests.integrations import _SKIP_IF_LIGHTNING_MISSING, _SKIP_IF_PYTORCHLIGHTNING_MISSING @@ -20,8 +19,8 @@ ) @pytest.mark.parametrize("clear_all_local", [True, False]) @pytest.mark.parametrize("keep_all_uploaded", [True, False]) -@mock.patch("litmodels.io.cloud.sdk_delete_model") -@mock.patch("litmodels.io.cloud.sdk_upload_model") +@mock.patch("litmodels.integrations.checkpoints.delete_model_version") +@mock.patch("litmodels.integrations.checkpoints.upload_model") @mock.patch("litmodels.integrations.checkpoints.Auth") def test_lightning_checkpoint_callback( mock_auth, @@ -67,7 +66,6 @@ def test_lightning_checkpoint_callback( expected_org = expected_model_registry["org"] expected_teamspace = expected_model_registry["teamspace"] expected_model = expected_model_registry["model"] - mock_upload_model.return_value.name = f"{expected_org}/{expected_teamspace}/{expected_model}" monkeypatch.setattr( "litmodels.integrations.checkpoints.LitModelCheckpointMixin.default_model_name", mock.MagicMock(return_value=expected_boring_model), @@ -101,10 +99,8 @@ def test_lightning_checkpoint_callback( assert mock_upload_model.call_args_list == [ mock.call( name=f"{expected_org}/{expected_teamspace}/{expected_model}:{v}", - path=mock.ANY, - progress_bar=True, - cloud_account=None, - metadata={"litModels.integration": LitModelCheckpoint.__name__, "litModels": litmodels.__version__}, + model=mock.ANY, + metadata={"litModels.integration": LitModelCheckpoint.__name__}, ) for v in ("epoch=0-step=64", "epoch=1-step=128") ] @@ -115,12 +111,13 @@ def test_lightning_checkpoint_callback( assert mock_delete_model.call_count == expected_cloud_removals if expected_cloud_removals: mock_delete_model.assert_called_once_with( - name=f"{expected_org}/{expected_teamspace}/{expected_model}:epoch=0-step=64" + name=f"{expected_org}/{expected_teamspace}/{expected_model}", + version="epoch=0-step=64", ) # Verify paths match the expected pattern for call_args in mock_upload_model.call_args_list: - path = call_args[1]["path"] + path = call_args[1]["model"] assert re.match(r".*[/\\]lightning_logs[/\\]version_\d+[/\\]checkpoints[/\\]epoch=\d+-step=\d+\.ckpt$", path) diff --git a/tests/test_io_cloud.py b/tests/test_io_cloud.py index 09ed4b0..0a859de 100644 --- a/tests/test_io_cloud.py +++ b/tests/test_io_cloud.py @@ -1,167 +1,37 @@ -import os -from contextlib import nullcontext -from unittest import mock - -import joblib -import pytest -import torch -import torch.jit as torch_jit -from sklearn import svm -from torch.nn import Module +from litlogger import models as litlogger_models import litmodels -from litmodels import download_model, load_model, save_model -from litmodels.io import upload_model_files -from litmodels.io.utils import _KERAS_AVAILABLE -from tests.integrations import LIT_ORG, LIT_TEAMSPACE - - -@pytest.mark.parametrize("name", ["/too/many/slashes", "org/model", "model-name"]) -@pytest.mark.parametrize("in_studio", [True, False]) -@mock.patch("litmodels.io.cloud.sdk_upload_model") -def test_upload_wrong_model_name(mock_sdk_upload, name, in_studio, monkeypatch): - if in_studio: - # mock env variables as it would run in studio - monkeypatch.setenv("LIGHTNING_ORG", LIT_ORG) - monkeypatch.setenv("LIGHTNING_TEAMSPACE", LIT_TEAMSPACE) - monkeypatch.setattr("lightning_sdk.organization.Organization", mock.MagicMock) - monkeypatch.setattr("lightning_sdk.teamspace.Teamspace", mock.MagicMock) - monkeypatch.setattr("lightning_sdk.teamspace.TeamspaceApi", mock.MagicMock) - monkeypatch.setattr("lightning_sdk.models._get_teamspace", mock.MagicMock) - - in_studio_only_name = in_studio and name == "model-name" - with ( - pytest.raises(ValueError, match=r".*organization/teamspace/model.*") - if not in_studio_only_name - else nullcontext() - ): - upload_model_files(path="path/to/checkpoint", name=name) - - -@pytest.mark.parametrize("name", ["/too/many/slashes", "org/model", "model-name"]) -@pytest.mark.parametrize("in_studio", [True, False]) -def test_download_wrong_model_name(name, in_studio, monkeypatch): - if in_studio: - # mock env variables as it would run in studio - monkeypatch.setenv("LIGHTNING_ORG", LIT_ORG) - monkeypatch.setenv("LIGHTNING_TEAMSPACE", LIT_TEAMSPACE) - monkeypatch.setattr("lightning_sdk.organization.Organization", mock.MagicMock) - monkeypatch.setattr("lightning_sdk.teamspace.Teamspace", mock.MagicMock) - monkeypatch.setattr("lightning_sdk.models.TeamspaceApi", mock.MagicMock) - in_studio_only_name = in_studio and name == "model-name" - with ( - pytest.raises(ValueError, match=r".*organization/teamspace/model.*") - if not in_studio_only_name - else nullcontext() - ): - download_model(name=name) - - -@pytest.mark.parametrize( - ("model", "model_path", "verbose"), - [ - # ("path/to/checkpoint", "path/to/checkpoint", False), - # (BoringModel(), "%s/BoringModel.ckpt"), - (torch_jit.script(Module()), f"%s{os.path.sep}RecursiveScriptModule.ts", True), - (Module(), f"%s{os.path.sep}Module.pth", True), - (svm.SVC(), f"%s{os.path.sep}SVC.pkl", 1), - ], -) -@mock.patch("litmodels.io.cloud.sdk_upload_model") -def test_upload_model(mock_upload_model, tmp_path, model, model_path, verbose): - mock_upload_model.return_value.name = "org-name/teamspace/model-name" - - # The lit-logger function is just a wrapper around the SDK function - save_model( - model=model, - name="org-name/teamspace/model-name", - cloud_account="cluster_id", - staging_dir=str(tmp_path), - verbose=verbose, - ) - expected_path = model_path % str(tmp_path) if "%" in model_path else model_path - mock_upload_model.assert_called_once_with( - path=expected_path, - name="org-name/teamspace/model-name", - cloud_account="cluster_id", - progress_bar=True, - metadata={"litModels": litmodels.__version__, "litModels.integration": "save_model"}, - ) - - -@mock.patch("litmodels.io.cloud.sdk_download_model") -def test_download_model(mock_download_model): - # The lit-logger function is just a wrapper around the SDK function - download_model( - name="org-name/teamspace/model-name", - download_dir="where/to/download", - ) - mock_download_model.assert_called_once_with( - name="org-name/teamspace/model-name", download_dir="where/to/download", progress_bar=True - ) - - -@mock.patch("litmodels.io.cloud.sdk_download_model") -def test_load_model_pickle(mock_download_model, tmp_path): - # create a dummy model file - model_file = tmp_path / "dummy_model.pkl" - test_data = svm.SVC() - joblib.dump(test_data, model_file) - mock_download_model.return_value = [str(model_file.name)] - - # The lit-logger function is just a wrapper around the SDK function - model = load_model( - name="org-name/teamspace/model-name", - download_dir=str(tmp_path), - ) - mock_download_model.assert_called_once_with( - name="org-name/teamspace/model-name", download_dir=str(tmp_path), progress_bar=True - ) - assert isinstance(model, svm.SVC) - - -@mock.patch("litmodels.io.cloud.sdk_download_model") -def test_load_model_torch_jit(mock_download_model, tmp_path): - # create a dummy model file - model_file = tmp_path / "dummy_model.ts" - test_data = torch_jit.script(Module()) - test_data.save(model_file) - mock_download_model.return_value = [str(model_file.name)] - - # The lit-logger function is just a wrapper around the SDK function - model = load_model( - name="org-name/teamspace/model-name", - download_dir=str(tmp_path), - ) - mock_download_model.assert_called_once_with( - name="org-name/teamspace/model-name", download_dir=str(tmp_path), progress_bar=True - ) - assert isinstance(model, torch.jit.ScriptModule) - - -@pytest.mark.skipif(not _KERAS_AVAILABLE, reason="TensorFlow/Keras is not available") -@mock.patch("litmodels.io.cloud.sdk_download_model") -def test_load_model_tf_keras(mock_download_model, tmp_path): - from tensorflow import keras - - # create a dummy model file - model_file = tmp_path / "dummy_model.keras" - # Define the model - model = keras.Sequential([ - keras.layers.Dense(10, input_shape=(784,), name="dense_1"), - keras.layers.Dense(10, name="dense_2"), - ]) - model.compile(optimizer="adam", loss="categorical_crossentropy") - model.save(model_file) - # prepare mocked SDK download function - mock_download_model.return_value = [str(model_file.name)] - - # The lit-logger function is just a wrapper around the SDK function - model = load_model( - name="org-name/teamspace/model-name", - download_dir=str(tmp_path), - ) - mock_download_model.assert_called_once_with( - name="org-name/teamspace/model-name", download_dir=str(tmp_path), progress_bar=True - ) - assert isinstance(model, keras.models.Model) +from litmodels import io as litmodels_io +from litmodels.io import cloud as litmodels_cloud +from litmodels.io import gateway as litmodels_gateway +from litmodels.io import utils as litmodels_utils + + +def test_top_level_exports_are_reexported_from_litlogger(): + assert litmodels.download_model is litlogger_models.download_model + assert litmodels.load_model is litlogger_models.load_model + assert litmodels.save_model is litlogger_models.save_model + assert litmodels.upload_model is litlogger_models.upload_model + assert litmodels.upload_model_files is litlogger_models.upload_model_files + + +def test_io_exports_are_reexported_from_litlogger(): + assert litmodels_io.download_model is litlogger_models.download_model + assert litmodels_io.download_model_files is litlogger_models.download_model_files + assert litmodels_io.load_model is litlogger_models.load_model + assert litmodels_io.save_model is litlogger_models.save_model + assert litmodels_io.upload_model is litlogger_models.upload_model + assert litmodels_io.upload_model_files is litlogger_models.upload_model_files + assert litmodels_gateway.download_model is litlogger_models.download_model + assert litmodels_gateway.load_model is litlogger_models.load_model + assert litmodels_gateway.save_model is litlogger_models.save_model + assert litmodels_gateway.upload_model is litlogger_models.upload_model + + +def test_cloud_and_utils_dependencies_are_reexported(): + assert litmodels_cloud.download_model_files is litlogger_models.download_model_files + assert litmodels_cloud.upload_model_files is litlogger_models.upload_model_files + assert callable(litmodels_cloud.delete_model_version) + assert callable(litmodels_cloud._list_available_teamspaces) + assert callable(litmodels_utils.dump_pickle) + assert callable(litmodels_utils.load_pickle) From cbf2a12e0407c476bc1e3b716f5be896366d2388 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 1 Apr 2026 13:03:31 +0000 Subject: [PATCH 4/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/litmodels/__init__.py | 3 ++- src/litmodels/io/__init__.py | 9 ++++++++- src/litmodels/io/cloud.py | 7 ++++++- 3 files changed, 16 insertions(+), 3 deletions(-) diff --git a/src/litmodels/__init__.py b/src/litmodels/__init__.py index 7d9a6db..e407229 100644 --- a/src/litmodels/__init__.py +++ b/src/litmodels/__init__.py @@ -2,9 +2,10 @@ import os -from litmodels.__about__ import * # noqa: F401, F403 from litlogger.models import download_model, load_model, save_model, upload_model, upload_model_files +from litmodels.__about__ import * # noqa: F401, F403 + _PACKAGE_ROOT = os.path.dirname(__file__) _PROJECT_ROOT = os.path.dirname(_PACKAGE_ROOT) diff --git a/src/litmodels/io/__init__.py b/src/litmodels/io/__init__.py index 2ccad08..4b749c3 100644 --- a/src/litmodels/io/__init__.py +++ b/src/litmodels/io/__init__.py @@ -1,5 +1,12 @@ """Root package for Input/output.""" -from litlogger.models import download_model, download_model_files, load_model, save_model, upload_model, upload_model_files +from litlogger.models import ( + download_model, + download_model_files, + load_model, + save_model, + upload_model, + upload_model_files, +) __all__ = ["download_model", "download_model_files", "upload_model", "upload_model_files", "load_model", "save_model"] diff --git a/src/litmodels/io/cloud.py b/src/litmodels/io/cloud.py index d2952af..8c36ddb 100644 --- a/src/litmodels/io/cloud.py +++ b/src/litmodels/io/cloud.py @@ -1,5 +1,10 @@ """Compatibility exports for the vendored litlogger model helpers.""" -from litlogger.models.cloud import _list_available_teamspaces, delete_model_version, download_model_files, upload_model_files +from litlogger.models.cloud import ( + _list_available_teamspaces, + delete_model_version, + download_model_files, + upload_model_files, +) __all__ = ["_list_available_teamspaces", "delete_model_version", "download_model_files", "upload_model_files"] From 4598d50329eef99e66e25f6e4902b6067d81287a Mon Sep 17 00:00:00 2001 From: Justus Perillieux Date: Fri, 25 Sep 2026 16:14:04 +0200 Subject: [PATCH 5/7] Modernize packaging with setuptools-scm --- .github/workflows/ci-cloud.yml | 4 +++- .github/workflows/ci-testing.yml | 4 +++- .github/workflows/docs-build.yml | 2 ++ .github/workflows/release-pypi.yml | 2 ++ MANIFEST.in | 3 --- docs/source/conf.py | 19 ++++++++++--------- pyproject.toml | 15 ++++++++------- requirements.txt | 3 --- src/litmodels/__about__.py | 4 +++- 9 files changed, 31 insertions(+), 25 deletions(-) delete mode 100644 requirements.txt diff --git a/.github/workflows/ci-cloud.yml b/.github/workflows/ci-cloud.yml index ce8cccd..2a059d0 100644 --- a/.github/workflows/ci-cloud.yml +++ b/.github/workflows/ci-cloud.yml @@ -34,6 +34,8 @@ jobs: steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + fetch-depth: 0 - name: Install uv and set Python ${{ matrix.python-version }} uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 @@ -46,7 +48,7 @@ jobs: if: matrix.requires == 'oldest' run: | uv pip install 'lightning-utilities[cli]' - python -m lightning_utilities.cli requirements set-oldest --req_files='["requirements.txt", "pyproject.toml"]' + python -m lightning_utilities.cli requirements set-oldest --req_files='["pyproject.toml"]' - name: Install package & dependencies run: | diff --git a/.github/workflows/ci-testing.yml b/.github/workflows/ci-testing.yml index 376cfa1..523dc43 100644 --- a/.github/workflows/ci-testing.yml +++ b/.github/workflows/ci-testing.yml @@ -67,6 +67,8 @@ jobs: steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + fetch-depth: 0 - name: Install uv and set Python ${{ matrix.python-version }} uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 @@ -79,7 +81,7 @@ jobs: if: matrix.requires == 'oldest' run: | uv pip install 'lightning-utilities[cli]' - python -m lightning_utilities.cli requirements set-oldest --req_files='["requirements.txt", "pyproject.toml"]' + python -m lightning_utilities.cli requirements set-oldest --req_files='["pyproject.toml"]' - name: Adjust requirements run: | diff --git a/.github/workflows/docs-build.yml b/.github/workflows/docs-build.yml index ed94c13..60421e8 100644 --- a/.github/workflows/docs-build.yml +++ b/.github/workflows/docs-build.yml @@ -11,6 +11,8 @@ jobs: UV_EXCLUDE_NEWER: "2 days" steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + fetch-depth: 0 - name: Install uv and set Python 3.10 uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 diff --git a/.github/workflows/release-pypi.yml b/.github/workflows/release-pypi.yml index f4aa05c..51f2230 100644 --- a/.github/workflows/release-pypi.yml +++ b/.github/workflows/release-pypi.yml @@ -21,6 +21,8 @@ jobs: steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + fetch-depth: 0 - name: Install uv and set Python 3.10 uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 diff --git a/MANIFEST.in b/MANIFEST.in index de2f42d..43059f0 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -26,9 +26,6 @@ exclude tests recursive-exclude docs * exclude docs -# Include the Requirements -include requirements.txt - # Exclude Makefile exclude Makefile diff --git a/docs/source/conf.py b/docs/source/conf.py index 08d31d8..347034f 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -9,6 +9,7 @@ import os import re import sys +from importlib.metadata import requires from importlib.util import module_from_spec, spec_from_file_location import pt_lightning_sphinx_theme @@ -288,13 +289,13 @@ def setup(app): # Ignoring Third-party packages # https://stackoverflow.com/questions/15889621/sphinx-how-to-exclude-imports-in-automodule -def _package_list_from_file(file): +def _package_list_from_metadata(package): list_pkgs = [] - with open(file) as fp: - lines = fp.readlines() - for ln in lines: - found = [ln.index(ch) for ch in list(",=<>#") if ch in ln] - pkg = ln[: min(found)] if found else ln + for requirement in requires(package) or []: + if "extra ==" in requirement: + continue + found = [requirement.index(ch) for ch in list(",=<>!~;[") if ch in requirement] + pkg = requirement[: min(found)] if found else requirement if pkg.rstrip(): list_pkgs.append(pkg.rstrip()) return list_pkgs @@ -303,13 +304,13 @@ def _package_list_from_file(file): # define mapping from PyPI names to python imports PACKAGE_MAPPING = { "PyYAML": "yaml", + "lightning-sdk": "lightning_sdk", + "lightning-utilities": "lightning_utilities", } MOCK_PACKAGES = [] if SPHINX_MOCK_REQUIREMENTS: # mock also base packages when we are on RTD since we don't install them there - MOCK_PACKAGES += _package_list_from_file( - os.path.join(_PATH_ROOT, "requirements.txt") - ) + MOCK_PACKAGES += _package_list_from_metadata("litmodels") MOCK_PACKAGES = [PACKAGE_MAPPING.get(pkg, pkg) for pkg in MOCK_PACKAGES] autodoc_mock_imports = MOCK_PACKAGES diff --git a/pyproject.toml b/pyproject.toml index 9bc5015..0c81341 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,7 @@ [build-system] requires = [ - "setuptools>=69", + "setuptools>=80", + "setuptools-scm[simple]>=9.2", "wheel", ] build-backend = "setuptools.build_meta" @@ -30,7 +31,12 @@ classifiers = [ "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", ] -dynamic = ["version", "dependencies"] +dynamic = ["version"] +dependencies = [ + "litlogger >=2026.04.16post0", + "lightning-sdk >=2026.04.16", + "lightning-utilities <=0.15.3", +] [project.urls] "Bug Tracker" = "https://github.com/Lightning-AI/litModels/issues" @@ -86,11 +92,6 @@ zip-safe = false [tool.setuptools.packages.find] where = ["src"] -[tool.setuptools.dynamic] -version = { attr = "litmodels.__about__.__version__" } -dependencies = { file = ["requirements.txt"] } - - [tool.check-manifest] ignore = [ "*.yml", diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 69d5629..0000000 --- a/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -litlogger >=2026.04.16post0 -lightning-sdk >=2026.04.16 -lightning-utilities<=0.15.3 diff --git a/src/litmodels/__about__.py b/src/litmodels/__about__.py index e4bd307..58dc8ca 100644 --- a/src/litmodels/__about__.py +++ b/src/litmodels/__about__.py @@ -1,4 +1,6 @@ -__version__ = "0.2.0" +from importlib.metadata import version + +__version__ = version("litmodels") __author__ = "Lightning-AI et al." __author_email__ = "community@lightning.ai" __license__ = "Apache-2.0" From ceec9a967f0c09b3965753271a5dc1961da5945e Mon Sep 17 00:00:00 2001 From: Justus Perillieux Date: Fri, 25 Sep 2026 16:27:12 +0200 Subject: [PATCH 6/7] Align dependencies with latest LitLogger --- .github/workflows/ci-checks.yml | 2 +- .github/workflows/ci-cloud.yml | 4 +-- .github/workflows/ci-testing.yml | 4 +-- .github/workflows/docs-build.yml | 4 +-- .github/workflows/release-pypi.yml | 4 +-- pyproject.toml | 7 +++-- src/litmodels/integrations/checkpoints.py | 16 +----------- src/litmodels/io/cloud.py | 3 +-- tests/integrations/test_checkpoints.py | 32 +++-------------------- tests/integrations/test_real_cloud.py | 8 ------ tests/test_io_cloud.py | 1 - 11 files changed, 17 insertions(+), 68 deletions(-) diff --git a/.github/workflows/ci-checks.yml b/.github/workflows/ci-checks.yml index 0f48a9b..bbf06e8 100644 --- a/.github/workflows/ci-checks.yml +++ b/.github/workflows/ci-checks.yml @@ -34,7 +34,7 @@ jobs: testing-matrix: | { "os": ["ubuntu-latest", "macos-latest", "windows-latest"], - "python-version": ["3.10", "3.12"] + "python-version": ["3.11", "3.12"] } # check-docs: diff --git a/.github/workflows/ci-cloud.yml b/.github/workflows/ci-cloud.yml index 2a059d0..48d22bd 100644 --- a/.github/workflows/ci-cloud.yml +++ b/.github/workflows/ci-cloud.yml @@ -18,10 +18,10 @@ jobs: fail-fast: false matrix: os: ["ubuntu-22.04", "macOS-14", "windows-2022"] - python-version: ["3.10"] + python-version: ["3.11"] requires: ["latest"] include: - - { os: "ubuntu-22.04", python-version: "3.10", requires: "oldest" } + - { os: "ubuntu-22.04", python-version: "3.11", requires: "oldest" } # Timeout: https://stackoverflow.com/a/59076067/4521646 timeout-minutes: 25 diff --git a/.github/workflows/ci-testing.yml b/.github/workflows/ci-testing.yml index 523dc43..99515fe 100644 --- a/.github/workflows/ci-testing.yml +++ b/.github/workflows/ci-testing.yml @@ -23,14 +23,14 @@ jobs: fail-fast: false matrix: os: ["ubuntu-24.04", "macOS-14", "windows-2022"] - python-version: ["3.10", "3.12"] + python-version: ["3.11", "3.12"] requires: ["latest"] dependency: ["lightning"] dependency_package: ["lightning"] include: - { os: "ubuntu-22.04", - python-version: "3.10", + python-version: "3.11", requires: "oldest", dependency: "lightning", dependency_package: "lightning", diff --git a/.github/workflows/docs-build.yml b/.github/workflows/docs-build.yml index 60421e8..0545392 100644 --- a/.github/workflows/docs-build.yml +++ b/.github/workflows/docs-build.yml @@ -14,10 +14,10 @@ jobs: with: fetch-depth: 0 - - name: Install uv and set Python 3.10 + - name: Install uv and set Python 3.11 uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 with: - python-version: "3.10" + python-version: "3.11" activate-environment: true enable-cache: true diff --git a/.github/workflows/release-pypi.yml b/.github/workflows/release-pypi.yml index 51f2230..9cebd67 100644 --- a/.github/workflows/release-pypi.yml +++ b/.github/workflows/release-pypi.yml @@ -24,10 +24,10 @@ jobs: with: fetch-depth: 0 - - name: Install uv and set Python 3.10 + - name: Install uv and set Python 3.11 uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 with: - python-version: "3.10" + python-version: "3.11" activate-environment: true enable-cache: true diff --git a/pyproject.toml b/pyproject.toml index 0c81341..c18e7c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,7 @@ license = { file = "LICENSE" } authors = [ { name = "Lightning-AI et al.", email = "community@lightning.ai" }, ] -requires-python = ">=3.10" +requires-python = ">=3.11" keywords = ["deep learning", "pytorch", "AI"] classifiers = [ "Environment :: Console", @@ -26,15 +26,14 @@ classifiers = [ "License :: OSI Approved :: Apache Software License", "Operating System :: OS Independent", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", ] dynamic = ["version"] dependencies = [ - "litlogger >=2026.04.16post0", - "lightning-sdk >=2026.04.16", + "litlogger >=2026.8.28", + "lightning-sdk >=2026.8.5", "lightning-utilities <=0.15.3", ] diff --git a/src/litmodels/integrations/checkpoints.py b/src/litmodels/integrations/checkpoints.py index d2c3e7d..68b2cf4 100644 --- a/src/litmodels/integrations/checkpoints.py +++ b/src/litmodels/integrations/checkpoints.py @@ -9,13 +9,12 @@ from typing import TYPE_CHECKING, Any, Optional from lightning_sdk.lightning_cloud.login import Auth -from lightning_sdk.utils.resolve import _resolve_teamspace from lightning_utilities import StrEnum from lightning_utilities.core.rank_zero import rank_zero_debug, rank_zero_only, rank_zero_warn from litmodels import upload_model from litmodels.integrations.imports import _LIGHTNING_AVAILABLE, _PYTORCHLIGHTNING_AVAILABLE -from litmodels.io.cloud import _list_available_teamspaces, delete_model_version +from litmodels.io.cloud import delete_model_version if _LIGHTNING_AVAILABLE: from lightning.pytorch.callbacks import ModelCheckpoint as _LightningModelCheckpoint @@ -243,19 +242,6 @@ def _update_model_name(self, pl_model: "pl.LightningModule") -> None: elif count_slashes_in_name == 0: if not self.model_registry: self.model_registry = default_model_name - teamspace = _resolve_teamspace(None, None, None) - if teamspace: - # case you use default model name and teamspace determined from env. variables aka running in studio - self.model_registry = f"{teamspace.owner.name}/{teamspace.name}/{self.model_registry}" - else: # try to load default users teamspace - ts_names = list(_list_available_teamspaces().keys()) - if len(ts_names) == 1: - self.model_registry = f"{ts_names[0]}/{self.model_registry}" - else: - options = "\n\t".join(ts_names) - raise RuntimeError( - f"Teamspace is not defined and there are multiple teamspaces available:\n{options}" - ) else: raise RuntimeError(f"Invalid model name: '{self.model_registry}'") diff --git a/src/litmodels/io/cloud.py b/src/litmodels/io/cloud.py index 46e03f3..215bcb1 100644 --- a/src/litmodels/io/cloud.py +++ b/src/litmodels/io/cloud.py @@ -1,10 +1,9 @@ """Compatibility exports for LitLogger's cloud model helpers.""" from litlogger.models.cloud import ( - _list_available_teamspaces, delete_model_version, download_model_files, upload_model_files, ) -__all__ = ["_list_available_teamspaces", "delete_model_version", "download_model_files", "upload_model_files"] +__all__ = ["delete_model_version", "download_model_files", "upload_model_files"] diff --git a/tests/integrations/test_checkpoints.py b/tests/integrations/test_checkpoints.py index 2efbf83..bab36fc 100644 --- a/tests/integrations/test_checkpoints.py +++ b/tests/integrations/test_checkpoints.py @@ -53,38 +53,12 @@ def test_lightning_checkpoint_callback( if model_name: ckpt_args.update({"model_registry": model_name}) - all_model_registry = { - "org-name/teamspace/model-name": {"org": "org-name", "teamspace": "teamspace", "model": "model-name"}, - "model-in-studio": {"org": "my-org", "teamspace": "dream-team", "model": "model-in-studio"}, - "model-user-only-project": {"org": "my-org", "teamspace": "default-ts", "model": "model-user-only-project"}, - } expected_boring_model = "BoringModel_20250102-1213" - expected_model_registry = all_model_registry.get( - model_name, - {"org": "org-name", "teamspace": "teamspace", "model": expected_boring_model}, - ) - expected_org = expected_model_registry["org"] - expected_teamspace = expected_model_registry["teamspace"] - expected_model = expected_model_registry["model"] + expected_model_registry = model_name or expected_boring_model monkeypatch.setattr( "litmodels.integrations.checkpoints.LitModelCheckpointMixin.default_model_name", mock.MagicMock(return_value=expected_boring_model), ) - if model_name is None or model_name == "model-in-studio": - mock_teamspace = mock.Mock(owner=mock.Mock()) - mock_teamspace.owner.name = expected_org - mock_teamspace.name = expected_teamspace - - monkeypatch.setattr( - "litmodels.integrations.checkpoints._resolve_teamspace", mock.MagicMock(return_value=mock_teamspace) - ) - elif model_name == "model-user-only-project": - monkeypatch.setattr("litmodels.integrations.checkpoints._resolve_teamspace", mock.MagicMock(return_value=None)) - monkeypatch.setattr( - "litmodels.integrations.checkpoints._list_available_teamspaces", - mock.MagicMock(return_value={f"{expected_org}/{expected_teamspace}": {}}), - ) - # mocking the trainer delete checkpoint removal mock_remove_ckpt = mock.Mock() # setting the Trainer and custom checkpointing @@ -98,7 +72,7 @@ def test_lightning_checkpoint_callback( assert mock_auth.call_count == 1 assert mock_upload_model.call_args_list == [ mock.call( - name=f"{expected_org}/{expected_teamspace}/{expected_model}:{v}", + name=f"{expected_model_registry}:{v}", model=mock.ANY, metadata={"litModels.integration": LitModelCheckpoint.__name__}, ) @@ -111,7 +85,7 @@ def test_lightning_checkpoint_callback( assert mock_delete_model.call_count == expected_cloud_removals if expected_cloud_removals: mock_delete_model.assert_called_once_with( - name=f"{expected_org}/{expected_teamspace}/{expected_model}", + name=expected_model_registry, version="epoch=0-step=64", ) diff --git a/tests/integrations/test_real_cloud.py b/tests/integrations/test_real_cloud.py index 315af37..1a1dc80 100644 --- a/tests/integrations/test_real_cloud.py +++ b/tests/integrations/test_real_cloud.py @@ -13,7 +13,6 @@ from litmodels import download_model, load_model, save_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 from litmodels.io.utils import _KERAS_AVAILABLE from tests.integrations import ( _SKIP_IF_LIGHTNING_BELLOW_2_5_1, @@ -323,13 +322,6 @@ def test_duplicate_real_hf_model(tmp_path): 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() - assert len(teams) > 0 - assert f"{LIT_USER}/{LIT_TEAMSPACE}" in teams - - @pytest.mark.cloud @pytest.mark.skipif( not _KERAS_AVAILABLE, diff --git a/tests/test_io_cloud.py b/tests/test_io_cloud.py index 93c8674..dffb598 100644 --- a/tests/test_io_cloud.py +++ b/tests/test_io_cloud.py @@ -31,7 +31,6 @@ def test_io_exports_are_reexported_from_litlogger(): def test_cloud_exports_are_reexported_from_litlogger(): - assert litmodels_cloud._list_available_teamspaces is litlogger_cloud._list_available_teamspaces assert litmodels_cloud.delete_model_version is litlogger_cloud.delete_model_version assert litmodels_cloud.download_model_files is litlogger_cloud.download_model_files assert litmodels_cloud.upload_model_files is litlogger_cloud.upload_model_files From 21a89c7ffbcf890e654e82b05e8036046108e518 Mon Sep 17 00:00:00 2001 From: Justus Perillieux Date: Fri, 25 Sep 2026 16:28:59 +0200 Subject: [PATCH 7/7] Run typing checks on Python 3.11 --- .github/workflows/ci-checks.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/ci-checks.yml b/.github/workflows/ci-checks.yml index bbf06e8..cd22688 100644 --- a/.github/workflows/ci-checks.yml +++ b/.github/workflows/ci-checks.yml @@ -19,6 +19,7 @@ jobs: with: actions-ref: 86fe1b20b4609835ba9e8c8739cd39707ba76868 extra-typing: "typing" + python-version: "3.11" check-schema: uses: Lightning-AI/utilities/.github/workflows/check-schema.yml@86fe1b20b4609835ba9e8c8739cd39707ba76868 # v0.15.3