diff --git a/requirements/pytorch/loggers.info b/requirements/pytorch/loggers.info index 35471ad2ae086..1c01893cda9d9 100644 --- a/requirements/pytorch/loggers.info +++ b/requirements/pytorch/loggers.info @@ -1,6 +1,6 @@ # all supported loggers. this list is here as a reference, but they are not installed in CI -litlogger >= 0.1.7 +litlogger >= 2026.03.17 neptune >=1.0.0 comet-ml >=3.31.0 mlflow >=1.0.0 diff --git a/src/lightning/pytorch/loggers/litlogger.py b/src/lightning/pytorch/loggers/litlogger.py index c4cb45861dad1..a734e6639970e 100644 --- a/src/lightning/pytorch/loggers/litlogger.py +++ b/src/lightning/pytorch/loggers/litlogger.py @@ -20,7 +20,7 @@ from argparse import Namespace from collections.abc import Mapping from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, Optional, Union +from typing import TYPE_CHECKING, Any, Optional, Union, cast from lightning_utilities.core.imports import RequirementCache from torch import Tensor @@ -108,6 +108,7 @@ def training_step(self, batch, batch_idx: int): self._sub_dir = None self._prefix = "" self._fs = get_filesystem(self._root_dir) + self._experiment: Optional[Experiment] = None self._step = -1 self._metadata = metadata or {} self._is_ready = False @@ -163,14 +164,40 @@ def sub_dir(self) -> Optional[str]: """Gets the sub directory where the TensorBoard experiments are saved.""" return self._sub_dir + @property + def _experiment_name(self) -> str: + if self.version is None: + return self.name + return f"{self.name}-{self.version}" + + @staticmethod + def _default_artifact_key(path: str) -> str: + try: + rel = os.path.relpath(path) + except ValueError: + rel = None + key = rel if rel is not None and not rel.startswith("..") else os.path.basename(path) + return key.replace("\\", "/") + + def _model_key(self) -> str: + return self._experiment_name + + @staticmethod + def _model_version(version: Optional[str], step: Optional[int]) -> Optional[str]: + if version is not None: + return version + if step is not None and step >= 0: + return str(step) + return None + @property @rank_zero_experiment def experiment(self) -> Optional["Experiment"]: """Returns the underlying litlogger Experiment object.""" import litlogger - if litlogger.experiment is not None: - return litlogger.experiment + if self._experiment is not None: + return self._experiment if not self._is_ready: self._is_ready = True @@ -182,24 +209,31 @@ def experiment(self) -> Optional["Experiment"]: if self.version is None: # Generate version as proper RFC 3339 timestamp with Z suffix (required by protobuf) timestamp = datetime.now(timezone.utc).isoformat(timespec="milliseconds") - self._version = timestamp.replace("+00:00", "Z") + self._version = timestamp.replace(":", "-").replace("+00:00", "Z") - litlogger.init( - name=self._name, - root_dir=self._root_dir, + self._experiment = litlogger.Experiment( + name=self._experiment_name, teamspace=self._teamspace, metadata={k: str(v) for k, v in self._metadata.items()}, store_step=True, store_created_at=True, + log_dir=self.log_dir, save_logs=self._save_logs, ) + self._experiment.print_url() - return litlogger.experiment + return self._experiment + + def _require_experiment(self) -> "Experiment": + experiment = self.experiment + if experiment is None: + raise RuntimeError("Experiment is not initialized") + return experiment @property @rank_zero_only def url(self) -> str: - return self.experiment.url + return self._require_experiment().url # ────────────────────────────────────────────────────────────────────────────── # Override methods from Logger @@ -208,18 +242,17 @@ def url(self) -> str: @override @rank_zero_only def log_metrics(self, metrics: Mapping[str, float], step: Optional[int] = None) -> None: - import litlogger - assert rank_zero_only.rank == 0, "experiment tried to log from global_rank != 0" # Ensure experiment is initialized - _ = self.experiment + experiment = self._require_experiment() self._step = self._step + 1 if step is None else step metrics = _add_prefix(metrics, self._prefix, self.LOGGER_JOIN_CHAR) metrics = {k: v.item() if isinstance(v, Tensor) else v for k, v in metrics.items()} - litlogger.log_metrics(metrics, step=self._step) + for key, value in metrics.items(): + experiment[key].append(value, step=self._step) @override @rank_zero_only @@ -231,8 +264,9 @@ def log_hyperparams( """Log hyperparams.""" if isinstance(params, Namespace): params = params.__dict__ - params.update(self._metadata or {}) - self._metadata = params + experiment = self._require_experiment() + for key, value in params.items(): + experiment[key] = str(value) @override @rank_zero_only @@ -247,13 +281,11 @@ def save(self) -> None: @override @rank_zero_only def finalize(self, status: Optional[str] = None) -> None: - import litlogger - - if litlogger.experiment is not None: + if self._experiment is not None: # log checkpoints as artifacts before finalizing if self._checkpoint_callback: self._scan_and_log_checkpoints(self._checkpoint_callback) - litlogger.finalize(status) + self._experiment.finalize(status) # ────────────────────────────────────────────────────────────────────────────── # Public methods @@ -267,8 +299,9 @@ def log_metadata( """Log hyperparams.""" if isinstance(params, Namespace): params = params.__dict__ - params.update(self._metadata or {}) - self._metadata = params + experiment = self._require_experiment() + for key, value in params.items(): + experiment[key] = str(value) @rank_zero_only def log_model( @@ -289,10 +322,14 @@ def log_model( metadata: Optional metadata dictionary to store with the model. """ - import litlogger + from litlogger import Model - _ = self.experiment - litlogger.log_model(model, staging_dir, verbose, version, metadata) + self._require_experiment()[self._model_key()] = Model( + model, + version=self._model_version(version, self._step), + metadata=cast(Optional[dict[str, str]], metadata), + staging_dir=staging_dir, + ) @rank_zero_only def log_model_artifact( @@ -309,10 +346,9 @@ def log_model_artifact( version: Optional version string for the model. Defaults to the experiment version. """ - import litlogger + from litlogger import Model - _ = self.experiment - litlogger.log_model_artifact(path, verbose, version) + self._require_experiment()[self._model_key()] = Model(path, version=self._model_version(version, self._step)) @rank_zero_only def get_file(self, path: str, verbose: bool = True) -> str: @@ -326,46 +362,8 @@ def get_file(self, path: str, verbose: bool = True) -> str: str: The local path where the file was saved. """ - import litlogger - - _ = self.experiment - return litlogger.get_file(path, verbose=verbose) - - @rank_zero_only - def get_model(self, staging_dir: Optional[str] = None, verbose: bool = False, version: Optional[str] = None) -> Any: - """Download and load a model object using litmodels. - - Args: - staging_dir: Optional directory where the model will be downloaded. - verbose: Whether to show progress bar. - version: Optional version string for the model. - - Returns: - The loaded model object. - - """ - import litlogger - - _ = self.experiment - return litlogger.get_model(staging_dir, verbose, version) - - @rank_zero_only - def get_model_artifact(self, path: str, verbose: bool = False, version: Optional[str] = None) -> str: - """Download a model artifact file or directory from cloud storage using litmodels. - - Args: - path: Path where the model should be saved locally. - verbose: Whether to show progress bar during download. - version: Optional version string for the model. - - Returns: - str: The local path where the model was saved. - - """ - import litlogger - - _ = self.experiment - return litlogger.get_model_artifact(path, verbose, version) + file = cast(Any, self._require_experiment()[self._default_artifact_key(path)]) + return file.save(path) @rank_zero_only def log_file(self, path: str) -> None: @@ -382,10 +380,9 @@ def log_file(self, path: str) -> None: logger.log_file('config.yaml') """ - import litlogger + from litlogger import File - _ = self.experiment - litlogger.log_file(path) + self._require_experiment()[self._default_artifact_key(path)] = File(path) # ────────────────────────────────────────────────────────────────────────────── # Callback methods @@ -412,11 +409,12 @@ def _scan_and_log_checkpoints(self, checkpoint_callback: ModelCheckpoint) -> Non """Find new checkpoints from the callback and log them as model artifacts.""" checkpoints = _scan_checkpoints(checkpoint_callback, self._logged_model_time) - for timestamp, path_ckpt, _score, tag in checkpoints: - if not self._checkpoint_name: - self._checkpoint_name = self.experiment.name - # Ensure the version tag is unique by appending a timestamp - unique_tag = f"{tag}-{int(datetime.now(timezone.utc).timestamp())}" - self.log_model_artifact(path_ckpt, verbose=True, version=unique_tag) + for timestamp, path_ckpt, _score, _tag in checkpoints: + experiment = self._require_experiment() + checkpoint_key = self._checkpoint_name or experiment.name + checkpoint_step = getattr(checkpoint_callback, "_last_global_step_saved", None) + from litlogger import Model + + experiment[checkpoint_key] = Model(path_ckpt, version=self._model_version(None, checkpoint_step)) # remember logged models - timestamp needed in case filename didn't change self._logged_model_time[path_ckpt] = timestamp diff --git a/tests/tests_pytorch/loggers/conftest.py b/tests/tests_pytorch/loggers/conftest.py index bc72493b02fee..241374e619220 100644 --- a/tests/tests_pytorch/loggers/conftest.py +++ b/tests/tests_pytorch/loggers/conftest.py @@ -164,24 +164,23 @@ def litlogger_mock(monkeypatch): experiment_mock.url = "https://lightning.ai/test/experiments/test-experiment" experiment_mock.name = "test-experiment" experiment_mock.version = "2024-01-01T00:00:00.000Z" + experiment_mock.get_file.return_value = "/path/to/file" + experiment_mock.get_model.return_value = MagicMock() + experiment_mock.get_model_artifact.return_value = "/path/to/artifact" + experiment_mock.series_mocks = {} + + def get_series(key): + if key not in experiment_mock.series_mocks: + experiment_mock.series_mocks[key] = MagicMock() + return experiment_mock.series_mocks[key] + + experiment_mock.__getitem__.side_effect = get_series litlogger = ModuleType("litlogger") litlogger.experiment = None - litlogger.Experiment = MagicMock - - def mock_init(**kwargs): - litlogger.experiment = experiment_mock - return experiment_mock - - litlogger.init = Mock(side_effect=mock_init) - litlogger.log_metrics = Mock() - litlogger.log_file = Mock() - litlogger.get_file = Mock(return_value="/path/to/file") - litlogger.log_model = Mock() - litlogger.get_model = Mock(return_value=MagicMock()) - litlogger.log_model_artifact = Mock() - litlogger.get_model_artifact = Mock(return_value="/path/to/artifact") - litlogger.finalize = Mock() + litlogger.Experiment = Mock(return_value=experiment_mock) + litlogger.File = Mock() + litlogger.Model = Mock() monkeypatch.setitem(sys.modules, "litlogger", litlogger) # Create generator submodule diff --git a/tests/tests_pytorch/loggers/test_litlogger.py b/tests/tests_pytorch/loggers/test_litlogger.py index abbe2f899ea25..6ffbf510fdcb2 100644 --- a/tests/tests_pytorch/loggers/test_litlogger.py +++ b/tests/tests_pytorch/loggers/test_litlogger.py @@ -13,7 +13,7 @@ # limitations under the License. import os from argparse import Namespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call import pytest import torch @@ -71,15 +71,18 @@ def test_litlogger_experiment_property(litlogger_mock, tmp_path): experiment = logger.experiment assert experiment is not None - litlogger_mock.init.assert_called_once() - - # Check init was called with correct arguments - call_kwargs = litlogger_mock.init.call_args[1] - assert call_kwargs["name"] == "test" - assert call_kwargs["root_dir"] == str(tmp_path) + litlogger_mock.Experiment.assert_called_once() + + # Check Experiment was called with correct arguments + call_kwargs = litlogger_mock.Experiment.call_args[1] + assert call_kwargs["name"].startswith("test-") + assert call_kwargs["name"].endswith(logger.version) + assert ":" not in logger.version + assert call_kwargs["log_dir"] == os.path.join(str(tmp_path), "test") assert call_kwargs["teamspace"] == "my-teamspace" assert call_kwargs["store_step"] is True assert call_kwargs["store_created_at"] is True + experiment.print_url.assert_called_once_with() def test_litlogger_experiment_reuses_existing(litlogger_mock, tmp_path): @@ -90,8 +93,8 @@ def test_litlogger_experiment_reuses_existing(litlogger_mock, tmp_path): _ = logger.experiment _ = logger.experiment - # init should only be called once - assert litlogger_mock.init.call_count == 1 + # Experiment should only be created once + assert litlogger_mock.Experiment.call_count == 1 @pytest.mark.parametrize("step_idx", [10, None]) @@ -102,13 +105,18 @@ def test_litlogger_log_metrics(litlogger_mock, tmp_path, step_idx): metrics = {"float": 0.3, "int": 1, "FloatTensor": torch.tensor(0.1), "IntTensor": torch.tensor(1)} logger.log_metrics(metrics, step_idx) - litlogger_mock.log_metrics.assert_called_once() - call_args = litlogger_mock.log_metrics.call_args - logged_metrics = call_args[0][0] + experiment = litlogger_mock.Experiment.return_value + expected_step = 0 if step_idx is None else step_idx # Verify tensors are converted to Python scalars - assert isinstance(logged_metrics["FloatTensor"], float) - assert isinstance(logged_metrics["IntTensor"], int) + experiment.series_mocks["float"].append.assert_called_once_with(0.3, step=expected_step) + experiment.series_mocks["int"].append.assert_called_once_with(1, step=expected_step) + float_tensor_call = experiment.series_mocks["FloatTensor"].append.call_args + int_tensor_call = experiment.series_mocks["IntTensor"].append.call_args + assert isinstance(float_tensor_call.args[0], float) + assert isinstance(int_tensor_call.args[0], int) + assert float_tensor_call.kwargs == {"step": expected_step} + assert int_tensor_call.kwargs == {"step": expected_step} def test_litlogger_log_metrics_with_prefix(litlogger_mock, tmp_path): @@ -118,11 +126,7 @@ def test_litlogger_log_metrics_with_prefix(litlogger_mock, tmp_path): logger.log_metrics({"loss": 0.5}, step=1) - litlogger_mock.log_metrics.assert_called_once() - call_args = litlogger_mock.log_metrics.call_args - logged_metrics = call_args[0][0] - - assert "train-loss" in logged_metrics + litlogger_mock.Experiment.return_value.series_mocks["train-loss"].append.assert_called_once_with(0.5, step=1) def test_litlogger_log_hyperparams_dict(litlogger_mock, tmp_path): @@ -131,8 +135,11 @@ def test_litlogger_log_hyperparams_dict(litlogger_mock, tmp_path): hparams = {"learning_rate": 0.001, "batch_size": 32} logger.log_hyperparams(hparams) - assert logger._metadata["learning_rate"] == 0.001 - assert logger._metadata["batch_size"] == 32 + experiment = litlogger_mock.Experiment.return_value + assert experiment.__setitem__.mock_calls == [ + call("learning_rate", "0.001"), + call("batch_size", "32"), + ] def test_litlogger_log_hyperparams_namespace(litlogger_mock, tmp_path): @@ -141,8 +148,11 @@ def test_litlogger_log_hyperparams_namespace(litlogger_mock, tmp_path): hparams = Namespace(learning_rate=0.001, batch_size=32) logger.log_hyperparams(hparams) - assert logger._metadata["learning_rate"] == 0.001 - assert logger._metadata["batch_size"] == 32 + experiment = litlogger_mock.Experiment.return_value + assert experiment.__setitem__.mock_calls == [ + call("learning_rate", "0.001"), + call("batch_size", "32"), + ] def test_litlogger_log_graph_warning(litlogger_mock, tmp_path): @@ -163,7 +173,7 @@ def test_litlogger_finalize(litlogger_mock, tmp_path): logger.finalize("success") - litlogger_mock.finalize.assert_called_once_with("success") + litlogger_mock.Experiment.return_value.finalize.assert_called_once_with("success") def test_litlogger_finalize_no_experiment(litlogger_mock, tmp_path): @@ -174,7 +184,7 @@ def test_litlogger_finalize_no_experiment(litlogger_mock, tmp_path): logger.finalize("success") # finalize should not be called since experiment is None - litlogger_mock.finalize.assert_not_called() + litlogger_mock.Experiment.return_value.finalize.assert_not_called() def test_litlogger_log_file(litlogger_mock, tmp_path): @@ -182,7 +192,8 @@ def test_litlogger_log_file(litlogger_mock, tmp_path): logger = LitLogger(name="test", root_dir=tmp_path) logger.log_file("config.yaml") - litlogger_mock.log_file.assert_called_once_with("config.yaml") + litlogger_mock.File.assert_called_once_with("config.yaml") + litlogger_mock.Experiment.return_value.__setitem__.assert_any_call("config.yaml", litlogger_mock.File.return_value) def test_litlogger_get_file(litlogger_mock, tmp_path): @@ -190,8 +201,9 @@ def test_litlogger_get_file(litlogger_mock, tmp_path): logger = LitLogger(name="test", root_dir=tmp_path) result = logger.get_file("config.yaml", verbose=True) - litlogger_mock.get_file.assert_called_once_with("config.yaml", verbose=True) - assert result == "/path/to/file" + litlogger_mock.Experiment.return_value.__getitem__.assert_any_call("config.yaml") + litlogger_mock.Experiment.return_value.series_mocks["config.yaml"].save.assert_called_once_with("config.yaml") + assert result == litlogger_mock.Experiment.return_value.series_mocks["config.yaml"].save.return_value def test_litlogger_log_model(litlogger_mock, tmp_path): @@ -200,16 +212,21 @@ def test_litlogger_log_model(litlogger_mock, tmp_path): model = torch.nn.Linear(10, 10) logger.log_model(model, staging_dir="/tmp", verbose=True, version="v1", metadata={"epoch": 10}) - litlogger_mock.log_model.assert_called_once_with(model, "/tmp", True, "v1", {"epoch": 10}) + litlogger_mock.Model.assert_called_once_with(model, version="v1", metadata={"epoch": 10}, staging_dir="/tmp") + litlogger_mock.Experiment.return_value.__setitem__.assert_any_call( + logger._experiment_name, litlogger_mock.Model.return_value + ) -def test_litlogger_get_model(litlogger_mock, tmp_path): - """Test get_model method.""" +def test_litlogger_log_model_uses_step_as_default_version(litlogger_mock, tmp_path): + """Test log_model defaults the model version to the current step.""" logger = LitLogger(name="test", root_dir=tmp_path) - result = logger.get_model(staging_dir="/tmp", verbose=True, version="v1") + logger._step = 7 + model = torch.nn.Linear(10, 10) - litlogger_mock.get_model.assert_called_once_with("/tmp", True, "v1") - assert result is not None + logger.log_model(model) + + litlogger_mock.Model.assert_called_once_with(model, version="7", metadata=None, staging_dir=None) def test_litlogger_log_model_artifact(litlogger_mock, tmp_path): @@ -217,16 +234,20 @@ def test_litlogger_log_model_artifact(litlogger_mock, tmp_path): logger = LitLogger(name="test", root_dir=tmp_path) logger.log_model_artifact("/path/to/model.ckpt", verbose=True, version="v1") - litlogger_mock.log_model_artifact.assert_called_once_with("/path/to/model.ckpt", True, "v1") + litlogger_mock.Model.assert_called_once_with("/path/to/model.ckpt", version="v1") + litlogger_mock.Experiment.return_value.__setitem__.assert_any_call( + logger._experiment_name, litlogger_mock.Model.return_value + ) -def test_litlogger_get_model_artifact(litlogger_mock, tmp_path): - """Test get_model_artifact method.""" +def test_litlogger_log_model_artifact_uses_step_as_default_version(litlogger_mock, tmp_path): + """Test log_model_artifact defaults the model version to the current step.""" logger = LitLogger(name="test", root_dir=tmp_path) - result = logger.get_model_artifact("/path/to/model", verbose=True, version="v1") + logger._step = 11 + + logger.log_model_artifact("/path/to/model.ckpt") - litlogger_mock.get_model_artifact.assert_called_once_with("/path/to/model", True, "v1") - assert result == "/path/to/artifact" + litlogger_mock.Model.assert_called_once_with("/path/to/model.ckpt", version="11") def test_litlogger_url_property(litlogger_mock, tmp_path): @@ -247,6 +268,7 @@ def test_litlogger_version_property(litlogger_mock, tmp_path): # After accessing experiment, version is set _ = logger.experiment assert logger.version is not None + assert ":" not in logger.version def test_litlogger_with_trainer(litlogger_mock, tmp_path): @@ -274,11 +296,11 @@ def training_step(self, batch, batch_idx): trainer.fit(model) # Verify metrics were logged - assert litlogger_mock.log_metrics.called + assert any(series.append.called for series in litlogger_mock.Experiment.return_value.series_mocks.values()) def test_litlogger_metadata_in_init(litlogger_mock, tmp_path): - """Test metadata is passed to litlogger.init.""" + """Test metadata is passed to litlogger.Experiment.""" logger = LitLogger( name="test", root_dir=tmp_path, @@ -287,7 +309,10 @@ def test_litlogger_metadata_in_init(litlogger_mock, tmp_path): _ = logger.experiment - call_kwargs = litlogger_mock.init.call_args[1] + call_kwargs = litlogger_mock.Experiment.call_args[1] + assert call_kwargs["name"].startswith("test-") + assert call_kwargs["name"].endswith(logger.version) + assert ":" not in logger.version assert call_kwargs["metadata"] == {"experiment_type": "test", "version": "1.0"} @@ -328,10 +353,10 @@ def test_litlogger_after_save_checkpoint_enabled(litlogger_mock, tmp_path): def test_litlogger_save_logs_option(litlogger_mock, tmp_path): - """Test save_logs option is passed to init.""" + """Test save_logs option is passed to Experiment.""" logger = LitLogger(name="test", root_dir=tmp_path, save_logs=True) _ = logger.experiment - call_kwargs = litlogger_mock.init.call_args[1] + call_kwargs = litlogger_mock.Experiment.call_args[1] assert call_kwargs["save_logs"] is True