Add experimental log_media method (#16)

owbone · web-flow · commit a77d6cdead9f · 2026-02-17T16:27:11.000Z
diff --git a/src/litlogger/experiment.py b/src/litlogger/experiment.py
@@ -15,6 +15,7 @@
 
 import atexit
 import contextlib
+import mimetypes
 import os
 import signal
 import sys
@@ -26,15 +27,18 @@
 from types import FrameType
 from typing import TYPE_CHECKING, Any, Dict, List, Union
 
+from lightning_sdk.lightning_cloud.openapi import V1MediaType
+
 from litlogger.api.artifacts_api import ArtifactsApi
 from litlogger.api.auth_api import AuthApi
+from litlogger.api.media_api import MediaApi
 from litlogger.api.metrics_api import MetricsApi
 from litlogger.api.utils import _resolve_teamspace, build_experiment_url, get_accessible_url, get_guest_url
 from litlogger.artifacts import Artifact, Model, ModelArtifact
 from litlogger.background import _BackgroundThread
 from litlogger.capture import rerun_and_record
 from litlogger.printer import Printer, RunStats
-from litlogger.types import Metrics, MetricValue
+from litlogger.types import MediaType, Metrics, MetricValue
 
 if TYPE_CHECKING:
     from lightning_sdk import Teamspace
@@ -111,6 +115,7 @@ def __init__(
             teamspace = None
 
         self._metrics_api = MetricsApi()
+        self._media_api = MediaApi(client=self._metrics_api.client)
         self._artifacts_api = ArtifactsApi()
         self._teamspace = _resolve_teamspace(teamspace)
 
@@ -561,6 +566,67 @@ def get_model(self, staging_dir: str | None = None, verbose: bool = False, versi
             self._printer.print_success("Retrieved model object")
         return result
 
+    def log_media(
+        self,
+        name: str,
+        path: str,
+        kind: MediaType | None = None,
+        step: int | None = None,
+        epoch: int | None = None,
+        caption: str | None = None,
+        verbose: bool = False,
+    ) -> None:
+        """Upload a media file (image, text, etc.) to the experiment.
+
+        Args:
+            name: Name of the media.
+            path: Local path to the media file.
+            kind: Type of media (MediaType.IMAGE or MediaType.TEXT).
+                  If None, attempts to guess from file extension or mime type.
+            step: Optional training step.
+            epoch: Optional training epoch.
+            caption: Optional caption for the media.
+            verbose: Whether to print a confirmation message after upload.
+
+        Raises:
+            ValueError: If the file type cannot be determined or is not supported.
+            FileNotFoundError: If the file does not exist.
+        """
+        if not os.path.exists(path):
+            raise FileNotFoundError(f"Media file not found: {path}")
+
+        media_type = V1MediaType.UNSPECIFIED
+
+        if kind is not None:
+            if kind == MediaType.IMAGE:
+                media_type = V1MediaType.IMAGE
+            elif kind == MediaType.TEXT:
+                media_type = V1MediaType.TEXT
+        else:
+            mime_type, _ = mimetypes.guess_type(path)
+            if mime_type:
+                if mime_type.startswith("image/"):
+                    media_type = V1MediaType.IMAGE
+                elif mime_type.startswith("text/"):
+                    media_type = V1MediaType.TEXT
+
+        if media_type == V1MediaType.UNSPECIFIED:
+            raise ValueError(f"Unsupported media type for file: {path}")
+
+        self._media_api.upload_media(
+            experiment_id=self._metrics_store.id,
+            teamspace=self._teamspace,
+            file_path=path,
+            name=name,
+            media_type=media_type,
+            step=step,
+            epoch=epoch,
+            caption=caption,
+        )
+        self._stats.media_logged += 1
+        if verbose:
+            self._printer.media_logged(path, step)
+
     def print_url(self) -> None:
         """Print the experiment URL and initialization info with styled output."""
         self._printer.experiment_start(
diff --git a/src/litlogger/logger.py b/src/litlogger/logger.py
@@ -28,6 +28,7 @@
 
 from litlogger.experiment import Experiment
 from litlogger.generator import _create_name
+from litlogger.types import MediaType
 
 _base_classes = []
 
@@ -418,3 +419,29 @@ def log_file(self, path: str) -> None:
             self._is_ready = True
             self._store_step = True
             self.experiment.log_file(path)
+
+        def log_media(
+            self,
+            name: str,
+            path: str,
+            kind: MediaType | None = None,
+            step: int | None = None,
+            epoch: int | None = None,
+            caption: str | None = None,
+            verbose: bool = False,
+        ) -> None:
+            """Log a media file to the experiment.
+
+            Args:
+                name: Name of the media.
+                path: Local path to the media file.
+                kind: Kind of media (MediaType.IMAGE or MediaType.TEXT).
+                      If None, attempts to guess from file extension or mime type.
+                step: Optional training step.
+                epoch: Optional training epoch.
+                caption: Optional caption.
+                verbose: Whether to print a confirmation message after upload.
+            """
+            self._is_ready = True
+            self._store_step = True
+            self.experiment.log_media(name, path, kind, step, epoch, caption, verbose)
diff --git a/src/litlogger/printer.py b/src/litlogger/printer.py
@@ -121,6 +121,7 @@ class RunStats:
 
     metrics_logged: int = 0
     artifacts_logged: int = 0
+    media_logged: int = 0
     models_logged: int = 0
     # Store recent values for sparklines (metric_name -> list of values)
     metrics_history: Dict[str, List[float]] = field(default_factory=dict)
@@ -406,6 +407,8 @@ def _print_run_stats(self, stats: RunStats) -> None:
             rows.append(("Metrics logged", f"{stats.metrics_logged:,}"))
         if stats.artifacts_logged > 0:
             rows.append(("Artifacts", f"{stats.artifacts_logged:,}"))
+        if stats.media_logged > 0:
+            rows.append(("Media logged", f"{stats.media_logged:,}"))
         if stats.models_logged > 0:
             rows.append(("Models", f"{stats.models_logged:,}"))
 
@@ -477,6 +480,19 @@ def artifact_logged(self, path: str, remote_path: str | None = None) -> None:
         display_path = remote_path or path
         self._echo(f"{self.success(check)} Logged {self.files(display_path)}")
 
+    def media_logged(self, path: str, step: int | None = None) -> None:
+        """Print media upload confirmation.
+
+        Example:
+            litlogger: ✓ Logged output.png (step 100)
+        """
+        if not self.verbose:
+            return
+
+        check = self.emoji("check") or "✓"
+        step_str = f" (step {step})" if step is not None else ""
+        self._echo(f"{self.success(check)} Logged {self.files(path)}{step_str}")
+
     def artifact_retrieved(self, path: str) -> None:
         """Print artifact download confirmation.
 
diff --git a/src/litlogger/types.py b/src/litlogger/types.py
@@ -33,6 +33,13 @@ class PhaseType(str, Enum):
     STOPPED = "stopped"
 
 
+class MediaType(str, Enum):
+    """Type of media to upload."""
+
+    IMAGE = "image"
+    TEXT = "text"
+
+
 @dataclass
 class MetricValue:
     """A single metric value with optional step and timestamp.
diff --git a/tests/unittests/test_experiment.py b/tests/unittests/test_experiment.py
@@ -676,4 +676,106 @@ def test_atexit_handler_registered(self):
     def test_signal_handlers_exist(self):
         """Test that signal handler method exists."""
         assert hasattr(experiment_module.Experiment, "_signal_handler")
-        assert callable(experiment_module.Experiment._signal_handler)
+
+
+class TestExperimentLogMedia:
+    """Test log_media method."""
+
+    def test_log_media_basic(self):
+        """Test basic media upload with explicit type."""
+        exp = MagicMock()
+        exp.name = "test_exp"
+        exp._manager = MagicMock()
+        exp._media_api = MagicMock()
+        exp._printer = MagicMock()
+        exp._stats = MagicMock()
+        exp._stats.media_logged = 0
+
+        # Import MediaType from types
+        from unittest.mock import patch
+
+        from lightning_sdk.lightning_cloud.openapi import V1MediaType
+        from litlogger.experiment import Experiment
+        from litlogger.types import MediaType
+
+        with patch("os.path.exists", return_value=True):
+            Experiment.log_media(exp, "image", "/path/to/image.png", kind=MediaType.IMAGE)
+
+        # Verify upload_media called with correct args
+        exp._media_api.upload_media.assert_called_once()
+        _, kwargs = exp._media_api.upload_media.call_args
+        assert kwargs["file_path"] == "/path/to/image.png"
+        assert kwargs["media_type"] == V1MediaType.IMAGE
+        assert exp._stats.media_logged == 1
+
+    def test_log_media_guess_type_image(self):
+        """Test media upload with guessed image type."""
+        exp = MagicMock()
+        exp._media_api = MagicMock()
+        exp._printer = MagicMock()
+        exp._stats = MagicMock()
+        exp._stats.media_logged = 0
+
+        from unittest.mock import patch
+
+        from lightning_sdk.lightning_cloud.openapi import V1MediaType
+        from litlogger.experiment import Experiment
+
+        with patch("os.path.exists", return_value=True):
+            Experiment.log_media(exp, "image", "/path/to/image.jpg")
+
+        _, kwargs = exp._media_api.upload_media.call_args
+        assert kwargs["media_type"] == V1MediaType.IMAGE
+        assert exp._stats.media_logged == 1
+
+    def test_log_media_guess_type_text(self):
+        """Test media upload with guessed text type."""
+        exp = MagicMock()
+        exp._media_api = MagicMock()
+        exp._printer = MagicMock()
+        exp._stats = MagicMock()
+        exp._stats.media_logged = 0
+
+        from unittest.mock import patch
+
+        from lightning_sdk.lightning_cloud.openapi import V1MediaType
+        from litlogger.experiment import Experiment
+
+        with patch("os.path.exists", return_value=True):
+            Experiment.log_media(exp, "file", "/path/to/file.txt")
+
+        _, kwargs = exp._media_api.upload_media.call_args
+        assert kwargs["media_type"] == V1MediaType.TEXT
+        assert exp._stats.media_logged == 1
+
+    def test_log_media_unsupported_type(self):
+        """Test log_media raises ValueError for guessed unsupported media type."""
+        exp = MagicMock()
+        exp._media_api = MagicMock()
+        exp._printer = MagicMock()
+        exp._stats = MagicMock()
+        exp._stats.media_logged = 0
+
+        from unittest.mock import patch
+
+        import pytest
+        from litlogger.experiment import Experiment
+
+        with (
+            patch("os.path.exists", return_value=True),
+            patch("mimetypes.guess_type", return_value=("application/zip", None)),
+            pytest.raises(ValueError, match="Unsupported media type for file: /path/to/file.txt"),
+        ):
+            Experiment.log_media(exp, "file", "/path/to/file.txt")
+
+    def test_log_media_raises_file_not_found(self):
+        """Test log_media raises FileNotFoundError."""
+        exp = MagicMock()
+
+        from unittest.mock import patch
+
+        import pytest
+        from litlogger.experiment import Experiment
+
+        with patch("os.path.exists", return_value=False), pytest.raises(FileNotFoundError):
+            Experiment.log_media(exp, "file", "/non/existent/file.png")
diff --git a/tests/unittests/test_logger.py b/tests/unittests/test_logger.py
@@ -47,3 +47,16 @@ def test_log_file():
     # Call log_file and verify it delegates to the experiment
     logger.log_file("test.txt")
     logger._experiment.log_file.assert_called_once_with("test.txt")
+
+
+def test_log_media():
+    """Test that LightningLogger has a log_media method."""
+    from litlogger import LightningLogger
+    from litlogger.types import MediaType
+
+    logger = object.__new__(LightningLogger)
+    logger._experiment = MagicMock()
+
+    logger.log_media("image", "test.png", kind=MediaType.IMAGE, step=10, caption="Test caption")
+
+    logger._experiment.log_media.assert_called_once_with("image", "test.png", MediaType.IMAGE, 10, None, "Test caption")
