update documentation (#123)

* Apply suggestions from code review

Author: Borda

src/litmodels/integrations/checkpoints.py

@@ -35,7 +35,11 @@
 # Create a singleton upload manager
 @cache
 def get_model_manager() -> "ModelManager":
-    """Get or create the singleton upload manager."""
+    """Get or create the singleton background manager for uploads/removals.
+
+    Returns:
+        ModelManager: A process-wide, cached instance managing asynchronous tasks.
+    """
     return ModelManager()
 
 
@@ -55,12 +59,16 @@ class RemoveType(StrEnum):
 
 
 class ModelManager:
-    """Manages uploads and removals with a single queue but separate counters."""
+    """Manage asynchronous uploads and removals via a single worker queue.
+
+    This manager runs a daemon worker thread that processes queued upload and removal tasks.
+    It maintains separate counters for pending uploads and removals and supports graceful shutdown.
+    """
 
     task_queue: queue.Queue
 
     def __init__(self) -> None:
-        """Initialize the ModelManager with a task queue and counters."""
+        """Initialize the manager with a task queue, counters, and a daemon worker thread."""
         self.task_queue = queue.Queue()
         self.upload_count = 0
         self.remove_count = 0
@@ -142,7 +150,11 @@ def shutdown(self) -> None:
 
 # Base class to be inherited
 class LitModelCheckpointMixin(ABC):
-    """Mixin class for LitModel checkpoint functionality."""
+    """Mixin adding upload/remove behavior for Lightning checkpoint callbacks.
+
+    This mixin queues uploads to Lightning Cloud upon checkpoint save and can optionally
+    remove local checkpoints or skip cloud-side pruning based on configuration.
+    """
 
     _datetime_stamp: str
     model_registry: Optional[str] = None
@@ -151,12 +163,12 @@ class LitModelCheckpointMixin(ABC):
     def __init__(
         self, model_registry: Optional[str], keep_all_uploaded: bool = False, clear_all_local: bool = False
     ) -> None:
-        """Initialize with model name.
+        """Configure model registry and pruning behavior.
 
         Args:
-            model_registry: Name of the model to upload in format 'organization/teamspace/modelname'.
-            keep_all_uploaded: Whether prevent deleting models from cloud if the checkpointing logic asks to do so.
-            clear_all_local: Whether to clear local models after uploading to the cloud.
+            model_registry: Target model registry in the form 'organization/teamspace/modelname'.
+            keep_all_uploaded: If True, never delete uploaded cloud versions even if local pruning occurs.
+            clear_all_local: If True, remove local checkpoint files after they are uploaded to cloud.
         """
         if not model_registry:
             rank_zero_warn(
@@ -201,7 +213,7 @@ def _upload_model(self, trainer: "pl.Trainer", filepath: Union[str, Path], metad
 
     @rank_zero_only
     def _remove_model(self, trainer: "pl.Trainer", filepath: Union[str, Path]) -> None:
-        """Remove the local version of the model if requested."""
+        """Queue removal of local and/or cloud artifacts according to configuration."""
         get_model_manager().queue_remove(
             filepath=filepath,
             # skip the local removal we put it in the queue right after the upload
@@ -211,7 +223,7 @@ def _remove_model(self, trainer: "pl.Trainer", filepath: Union[str, Path]) -> No
         )
 
     def default_model_name(self, pl_model: "pl.LightningModule") -> str:
-        """Generate a default model name based on the class name and timestamp."""
+        """Generate a default model name using the LightningModule class name and a timestamp."""
         return pl_model.__class__.__name__ + f"_{self._datetime_stamp}"
 
     def _update_model_name(self, pl_model: "pl.LightningModule") -> None:
@@ -252,14 +264,14 @@ def _update_model_name(self, pl_model: "pl.LightningModule") -> None:
 if _LIGHTNING_AVAILABLE:
 
     class LightningModelCheckpoint(LitModelCheckpointMixin, _LightningModelCheckpoint):
-        """Lightning ModelCheckpoint with LitModel support.
+        """Drop-in ModelCheckpoint that uploads saved checkpoints to Lightning Cloud.
 
         Args:
-            model_registry: Name of the model to upload in format 'organization/teamspace/modelname'.
-            keep_all_uploaded: Whether prevent deleting models from cloud if the checkpointing logic asks to do so.
-            clear_all_local: Whether to clear local models after uploading to the cloud.
-            *args: Additional arguments to pass to the parent class.
-            **kwargs: Additional keyword arguments to pass to the parent class.
+            model_registry: Target model registry in the form 'organization/teamspace/modelname'.
+            keep_all_uploaded: If True, does not remove cloud versions when local pruning occurs.
+            clear_all_local: If True, removes local checkpoint files after successful upload.
+            *args: Additional positional arguments forwarded to the base ModelCheckpoint.
+            **kwargs: Additional keyword arguments forwarded to the base ModelCheckpoint.
         """
 
         def __init__(
@@ -291,7 +303,7 @@ def setup(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", stage: s
             self._update_model_name(pl_module)
 
         def _save_checkpoint(self, trainer: "pl.Trainer", filepath: str) -> None:
-            """Extend the save checkpoint method to upload the model."""
+            """Save the checkpoint and queue an upload from the global-zero process."""
             _LightningModelCheckpoint._save_checkpoint(self, trainer, filepath)
             if trainer.is_global_zero:  # Only upload from the main process
                 self._upload_model(trainer=trainer, filepath=filepath)
@@ -311,14 +323,14 @@ def _remove_checkpoint(self, trainer: "pl.Trainer", filepath: str) -> None:
 if _PYTORCHLIGHTNING_AVAILABLE:
 
     class PytorchLightningModelCheckpoint(LitModelCheckpointMixin, _PytorchLightningModelCheckpoint):
-        """PyTorch Lightning ModelCheckpoint with LitModel support.
+        """Drop-in ModelCheckpoint for PyTorch Lightning that uploads to Lightning Cloud.
 
         Args:
             model_registry: Name of the model to upload in format 'organization/teamspace/modelname'.
-            keep_all_uploaded: Whether prevent deleting models from cloud if the checkpointing logic asks to do so.
-            clear_all_local: Whether to clear local models after uploading to the cloud.
-            args: Additional arguments to pass to the parent class.
-            kwargs: Additional keyword arguments to pass to the parent class.
+            keep_all_uploaded: If True, does not remove cloud versions when local pruning occurs.
+            clear_all_local: If True, removes local checkpoint files after successful upload.
+            args: Additional positional arguments forwarded to the base ModelCheckpoint.
+            kwargs: Additional keyword arguments forwarded to the base ModelCheckpoint.
         """
 
         def __init__(

src/litmodels/integrations/duplicate.py

@@ -21,15 +21,14 @@ def duplicate_hf_model(
     verbose: int = 1,
     metadata: Optional[dict] = None,
 ) -> str:
-    """Downloads the model from Hugging Face and uploads it to Lightning Cloud.
+    """Download a model from Hugging Face and upload it to Lightning Cloud as a new model.
 
     Args:
-        hf_model: The name of the Hugging Face model to duplicate.
-        lit_model: The name of the Lightning Cloud model to create.
-        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.
+        hf_model: Hugging Face model identifier, for example 'org/name' or 'user/name'.
+        lit_model: Target Lightning Cloud model name. If omitted, derived from `hf_model` by replacing '/' with '_'.
+        local_workdir: Working directory used for download and staging. A temporary directory is created if omitted.
+        verbose: Verbosity for upload progress (0 = silent, 1 = print link once, 2 = print link always).
+        metadata: Optional metadata to attach to the uploaded model. Integration markers are added automatically.
 
     Returns:
         The name of the duplicated model in Lightning Cloud.

src/litmodels/io/cloud.py

@@ -21,15 +21,14 @@
 
 
 def _print_model_link(name: str, verbose: Union[bool, int]) -> None:
-    """Print a link to the uploaded model.
+    """Print a stable URL to the uploaded model.
 
     Args:
-        name: Name of the model.
-        verbose: Whether to print the link:
-
-            - If set to 0, no link will be printed.
-            - If set to 1, the link will be printed only once.
-            - If set to 2, the link will be printed every time.
+        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)
@@ -51,18 +50,18 @@ def upload_model_files(
     verbose: Union[bool, int] = 1,
     metadata: Optional[dict[str, str]] = None,
 ) -> "UploadedModelInfo":
-    """Upload a local checkpoint file to the model store.
+    """Upload local artifact(s) to Lightning Cloud using the SDK.
 
     Args:
-        name: Name of the model to upload. Must be in the format 'organization/teamspace/modelname'
-            where entity is either your username or the name of an organization you are part of.
-        path: Path to the model file to upload.
-        progress_bar: Whether to show a progress bar for the upload.
-        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.
+        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.
 
+    Returns:
+        UploadedModelInfo describing the created or updated model version.
     """
     if not metadata:
         metadata = {}
@@ -84,17 +83,15 @@ def download_model_files(
     download_dir: Union[str, Path] = ".",
     progress_bar: bool = True,
 ) -> Union[str, list[str]]:
-    """Download a checkpoint from the model store.
+    """Download artifact(s) for a model version using the SDK.
 
     Args:
-        name: Name of the model to download. Must be in the format 'organization/teamspace/modelname'
-            where entity is either your username or the name of an organization you are part of.
-        download_dir: A path to directory where the model should be downloaded. Defaults
-            to the current working directory.
-        progress_bar: Whether to show a progress bar for the download.
+        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:
-        The absolute path to the downloaded model file or folder.
+        str | list[str]: Absolute path(s) to the downloaded artifact(s).
     """
     return sdk_download_model(
         name=name,
@@ -104,10 +101,10 @@ def download_model_files(
 
 
 def _list_available_teamspaces() -> dict[str, dict]:
-    """List available teamspaces for the authenticated user.
+    """List teamspaces available to the authenticated user.
 
     Returns:
-        Dict with teamspace names as keys and their details as values.
+        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
@@ -128,13 +125,12 @@ def _list_available_teamspaces() -> dict[str, dict]:
 
 def delete_model_version(
     name: str,
-    version: Optional[str] = None,
+    version: str,
 ) -> None:
-    """Delete a model version from the model store.
+    """Delete a specific model version from the model store.
 
     Args:
-        name: Name of the model to delete. Must be in the format 'organization/teamspace/modelname'
-            where entity is either your username or the name of an organization you are part of.
-        version: Version of the model to delete. If not provided, all versions will be deleted.
+        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}")

src/litmodels/io/gateway.py

@@ -24,18 +24,22 @@ def upload_model(
     verbose: Union[bool, int] = 1,
     metadata: Optional[dict[str, str]] = None,
 ) -> "UploadedModelInfo":
-    """Upload a checkpoint to the model store.
+    """Upload a local artifact (file or directory) to Lightning Cloud Models.
 
     Args:
-        name: Name of the model to upload. Must be in the format 'organization/teamspace/modelname'
-            where entity is either your username or the name of an organization you are part of.
-        model: The model to upload. Can be a path to a checkpoint file or a folder.
-        progress_bar: Whether to show a progress bar for the upload.
-        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 some additional information about the uploaded model.
-        metadata: Optional metadata to attach to the model. If not provided, a default metadata will be used.
+        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.
 
+    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(
@@ -62,20 +66,29 @@ def save_model(
     verbose: Union[bool, int] = 1,
     metadata: Optional[dict[str, str]] = None,
 ) -> "UploadedModelInfo":
-    """Upload a checkpoint to the model store.
+    """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: Name of the model to upload. Must be in the format 'organization/teamspace/modelname'
-            where entity is either your username or the name of an organization you are part of.
-        model: The model to upload. Can be a PyTorch model, or a Lightning model a.
-        progress_bar: Whether to show a progress bar for the upload.
-        cloud_account: The name of the cloud account to store the Model in. Only required if it can't be determined
-            automatically.
-        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.
+        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.
 
+    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(
@@ -120,17 +133,15 @@ def download_model(
     download_dir: Union[str, Path] = ".",
     progress_bar: bool = True,
 ) -> Union[str, list[str]]:
-    """Download a checkpoint from the model store.
+    """Download a model version from Lightning Cloud Models to a local directory.
 
     Args:
-        name: Name of the model to download. Must be in the format 'organization/teamspace/modelname'
-            where entity is either your username or the name of an organization you are part of.
-        download_dir: A path to directory where the model should be downloaded. Defaults
-            to the current working directory.
-        progress_bar: Whether to show a progress bar for the download.
+        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:
-        The absolute path to the downloaded model file or folder.
+        str | list[str]: Absolute path(s) to the downloaded file(s) or directory content.
     """
     return download_model_files(
         name=name,
@@ -140,16 +151,22 @@ def download_model(
 
 
 def load_model(name: str, download_dir: str = ".") -> Any:
-    """Download a model from the model store and load it into memory.
+    """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: Name of the model to download. Must be in the format 'organization/teamspace/modelname'
-            where entity is either your username or the name of an organization you are part of.
-        download_dir: A path to directory where the model should be downloaded. Defaults
-            to the current working directory.
+        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:
-        The loaded model.
+        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