Merge branch 'master' into patch-2

Author: Borda

src/lightning/pytorch/loggers/wandb.py

@@ -278,6 +278,7 @@ def any_lightning_module_function_or_hook(self):
         prefix: A string to put at the beginning of metric keys.
         experiment: WandB experiment object. Automatically set when creating a run.
         checkpoint_name: Name of the model checkpoint artifact being logged.
+        add_file_policy: If "mutable", copies file to tempdirectory before upload.
         \**kwargs: Arguments passed to :func:`wandb.init` like `entity`, `group`, `tags`, etc.
 
     Raises:
@@ -304,6 +305,7 @@ def __init__(
         experiment: Union["Run", "RunDisabled", None] = None,
         prefix: str = "",
         checkpoint_name: Optional[str] = None,
+        add_file_policy: Literal["mutable", "immutable"] = "mutable",
         **kwargs: Any,
     ) -> None:
         if not _WANDB_AVAILABLE:
@@ -323,6 +325,7 @@ def __init__(
         self._experiment = experiment
         self._logged_model_time: dict[str, float] = {}
         self._checkpoint_callbacks: dict[int, ModelCheckpoint] = {}
+        self.add_file_policy = add_file_policy
 
         # paths are processed as strings
         if save_dir is not None:
@@ -676,7 +679,7 @@ def _scan_and_log_checkpoints(self, checkpoint_callback: ModelCheckpoint) -> Non
             if not self._checkpoint_name:
                 self._checkpoint_name = f"model-{self.experiment.id}"
             artifact = wandb.Artifact(name=self._checkpoint_name, type="model", metadata=metadata)
-            artifact.add_file(p, name="model.ckpt")
+            artifact.add_file(p, name="model.ckpt", policy=self.add_file_policy)
             aliases = ["latest", "best"] if p == checkpoint_callback.best_model_path else ["latest"]
             self.experiment.log_artifact(artifact, aliases=aliases)
             # remember logged models - timestamp needed in case filename didn't change (lastkckpt or custom name)

src/lightning/pytorch/trainer/connectors/callback_connector.py

@@ -18,7 +18,7 @@
 from datetime import timedelta
 from typing import Optional, Union
 
-from lightning_utilities import module_available
+from lightning_utilities.core.imports import RequirementCache
 
 import lightning.pytorch as pl
 from lightning.fabric.utilities.registry import _load_external_callbacks
@@ -93,7 +93,7 @@ def _configure_checkpoint_callbacks(self, enable_checkpointing: bool) -> None:
                     " but found `ModelCheckpoint` in callbacks list."
                 )
         elif enable_checkpointing:
-            if module_available("litmodels") and self.trainer._model_registry:
+            if RequirementCache("litmodels >=0.1.7") and self.trainer._model_registry:
                 trainer_source = inspect.getmodule(self.trainer)
                 if trainer_source is None or not isinstance(trainer_source.__package__, str):
                     raise RuntimeError("Unable to determine the source of the trainer.")
@@ -103,11 +103,11 @@ def _configure_checkpoint_callbacks(self, enable_checkpointing: bool) -> None:
                 else:
                     from litmodels.integrations.checkpoints import LightningModelCheckpoint as LitModelCheckpoint
 
-                model_checkpoint = LitModelCheckpoint(model_name=self.trainer._model_registry)
+                model_checkpoint = LitModelCheckpoint(model_registry=self.trainer._model_registry)
             else:
                 rank_zero_info(
                     "You are using the default ModelCheckpoint callback."
-                    " Install `litmodels` package to use the `LitModelCheckpoint` instead"
+                    " Install `pip install litmodels` package to use the `LitModelCheckpoint` instead"
                     " for seamless uploading to the Lightning model registry."
                 )
                 model_checkpoint = ModelCheckpoint()