ml-dash 0.6.5__py3-none-any.whl → 0.6.7__py3-none-any.whl

ml_dash/files.py

@@ -56,6 +56,7 @@ class FileBuilder:
         experiment.files.save_blob(b"xxx", to="data.bin")
         experiment.files.save_torch(model, to="model.pt")
         experiment.files.save_pkl(data, to="data.pkl")
+        experiment.files.save_image(array, to="frame.png")
         experiment.files.save_fig(fig, to="plot.png")
         experiment.files.save_video(frames, to="video.mp4")
     """
@@ -184,8 +185,9 @@ class FileBuilder:
                 - str (file path): Uploads existing file
                 - bytes: Saves as binary blob (requires 'to' parameter)
                 - dict/list: Saves as JSON (requires 'to' parameter)
+                - numpy.ndarray: Saves as image (requires 'to' parameter with image extension)
                 - None: Uses file_path from constructor (backwards compatibility)
-            to: Target filename (required for bytes/dict/list, optional for file paths)
+            to: Target filename (required for bytes/dict/list/arrays, optional for file paths)
             description: Optional description
             tags: Optional list of tags
             metadata: Optional metadata dict
@@ -224,7 +226,17 @@ class FileBuilder:
                 raise ValueError("'to' parameter is required when saving dict/list")
             return self.save_json(content, to=to)
 
-        raise ValueError(f"Unsupported content type: {type(content)}. Expected str (file path), bytes, dict, or list.")
+        # Check if content is a numpy array (save as image)
+        try:
+            import numpy as np
+            if isinstance(content, np.ndarray):
+                if not to:
+                    raise ValueError("'to' parameter is required when saving numpy arrays")
+                return self.save_image(content, to=to)
+        except ImportError:
+            pass  # numpy not available
+
+        raise ValueError(f"Unsupported content type: {type(content)}. Expected str (file path), bytes, dict, list, or numpy.ndarray.")
 
     def _save_file(
         self,
@@ -288,23 +300,27 @@ class FileBuilder:
         temp_path = os.path.join(temp_dir, filename)
         # Create parent directories if filename contains path
         os.makedirs(os.path.dirname(temp_path), exist_ok=True)
-        try:
-            with open(temp_path, 'wb') as f:
-                f.write(data)
-            return self._save_file(
-                fpath=temp_path,
-                prefix=prefix,
-                description=description,
-                tags=tags,
-                metadata=metadata
-            )
-        finally:
+
+        with open(temp_path, 'wb') as f:
+            f.write(data)
+        result = self._save_file(
+            fpath=temp_path,
+            prefix=prefix,
+            description=description,
+            tags=tags,
+            metadata=metadata
+        )
+
+        # Only clean up if NOT queued for buffered upload
+        if result.get("status") != "queued":
             try:
                 os.unlink(temp_path)
                 os.rmdir(temp_dir)
             except Exception:
                 pass
 
+        return result
+
     def _save_json(
         self,
         content: Any,
@@ -323,23 +339,27 @@ class FileBuilder:
         temp_path = os.path.join(temp_dir, filename)
         # Create parent directories if filename contains path
         os.makedirs(os.path.dirname(temp_path), exist_ok=True)
-        try:
-            with open(temp_path, 'w') as f:
-                json.dump(content, f, indent=2)
-            return self._save_file(
-                fpath=temp_path,
-                prefix=prefix,
-                description=description,
-                tags=tags,
-                metadata=metadata
-            )
-        finally:
+
+        with open(temp_path, 'w') as f:
+            json.dump(content, f, indent=2)
+        result = self._save_file(
+            fpath=temp_path,
+            prefix=prefix,
+            description=description,
+            tags=tags,
+            metadata=metadata
+        )
+
+        # Only clean up if NOT queued for buffered upload
+        if result.get("status") != "queued":
             try:
                 os.unlink(temp_path)
                 os.rmdir(temp_dir)
             except Exception:
                 pass
 
+        return result
+
     def _save_torch(
         self,
         model: Any,
@@ -358,22 +378,26 @@ class FileBuilder:
         temp_path = os.path.join(temp_dir, filename)
         # Create parent directories if filename contains path
         os.makedirs(os.path.dirname(temp_path), exist_ok=True)
-        try:
-            torch.save(model, temp_path)
-            return self._save_file(
-                fpath=temp_path,
-                prefix=prefix,
-                description=description,
-                tags=tags,
-                metadata=metadata
-            )
-        finally:
+
+        torch.save(model, temp_path)
+        result = self._save_file(
+            fpath=temp_path,
+            prefix=prefix,
+            description=description,
+            tags=tags,
+            metadata=metadata
+        )
+
+        # Only clean up if NOT queued for buffered upload
+        if result.get("status") != "queued":
             try:
                 os.unlink(temp_path)
                 os.rmdir(temp_dir)
             except Exception:
                 pass
 
+        return result
+
     def _save_fig(
         self,
         fig: Any,
@@ -393,23 +417,27 @@ class FileBuilder:
         temp_path = os.path.join(temp_dir, filename)
         # Create parent directories if filename contains path
         os.makedirs(os.path.dirname(temp_path), exist_ok=True)
-        try:
-            fig.savefig(temp_path, **kwargs)
-            plt.close(fig)
-            return self._save_file(
-                fpath=temp_path,
-                prefix=prefix,
-                description=description,
-                tags=tags,
-                metadata=metadata
-            )
-        finally:
+
+        fig.savefig(temp_path, **kwargs)
+        plt.close(fig)
+        result = self._save_file(
+            fpath=temp_path,
+            prefix=prefix,
+            description=description,
+            tags=tags,
+            metadata=metadata
+        )
+
+        # Only clean up if NOT queued for buffered upload
+        if result.get("status") != "queued":
             try:
                 os.unlink(temp_path)
                 os.rmdir(temp_dir)
             except Exception:
                 pass
 
+        return result
+
     def _save_pickle(
         self,
         content: Any,
@@ -428,23 +456,107 @@ class FileBuilder:
         temp_path = os.path.join(temp_dir, filename)
         # Create parent directories if filename contains path
         os.makedirs(os.path.dirname(temp_path), exist_ok=True)
+
+        with open(temp_path, 'wb') as f:
+            pickle.dump(content, f)
+        result = self._save_file(
+            fpath=temp_path,
+            prefix=prefix,
+            description=description,
+            tags=tags,
+            metadata=metadata
+        )
+
+        # Only clean up if NOT queued for buffered upload
+        if result.get("status") != "queued":
+            try:
+                os.unlink(temp_path)
+                os.rmdir(temp_dir)
+            except Exception:
+                pass
+
+        return result
+
+    def _save_image(
+        self,
+        array: Any,
+        filename: str,
+        prefix: str,
+        description: Optional[str],
+        tags: Optional[List[str]],
+        metadata: Optional[Dict[str, Any]],
+        quality: int = 95
+    ) -> Dict[str, Any]:
+        """Save numpy array as image file."""
+        import tempfile
+        import os
+
         try:
-            with open(temp_path, 'wb') as f:
-                pickle.dump(content, f)
-            return self._save_file(
-                fpath=temp_path,
-                prefix=prefix,
-                description=description,
-                tags=tags,
-                metadata=metadata
-            )
-        finally:
+            from PIL import Image
+            import numpy as np
+        except ImportError:
+            raise ImportError("PIL/Pillow is required for saving images. Install it with: pip install Pillow")
+
+        temp_dir = tempfile.mkdtemp()
+        temp_path = os.path.join(temp_dir, filename)
+        # Create parent directories if filename contains path
+        os.makedirs(os.path.dirname(temp_path), exist_ok=True)
+
+        # Convert numpy array to PIL Image
+        # Handle different array shapes and dtypes
+        if array.dtype == np.uint8:
+            img = Image.fromarray(array)
+        else:
+            # Normalize to 0-255 range for non-uint8 arrays
+            if array.max() <= 1.0:
+                # Assume normalized float in [0, 1]
+                array_uint8 = (array * 255).astype(np.uint8)
+            else:
+                # Scale to 0-255
+                array_uint8 = ((array - array.min()) / (array.max() - array.min()) * 255).astype(np.uint8)
+            img = Image.fromarray(array_uint8)
+
+        # Handle JPEG-specific requirements
+        file_ext = os.path.splitext(filename)[1].lower()
+        save_kwargs = {}
+        if file_ext in ['.jpg', '.jpeg']:
+            # JPEG doesn't support alpha channel, convert RGBA to RGB
+            if img.mode in ('RGBA', 'LA', 'P'):
+                # Create white background
+                background = Image.new('RGB', img.size, (255, 255, 255))
+                if img.mode == 'P':
+                    img = img.convert('RGBA')
+                background.paste(img, mask=img.split()[-1] if img.mode in ('RGBA', 'LA') else None)
+                img = background
+            elif img.mode not in ('RGB', 'L'):
+                # Convert other modes to RGB
+                img = img.convert('RGB')
+            # Set JPEG quality
+            save_kwargs['quality'] = quality
+            save_kwargs['optimize'] = True
+
+        # Save image to temp file
+        img.save(temp_path, **save_kwargs)
+
+        result = self._save_file(
+            fpath=temp_path,
+            prefix=prefix,
+            description=description,
+            tags=tags,
+            metadata=metadata
+        )
+
+        # Only clean up if NOT queued for buffered upload
+        # If queued, the buffer manager will clean up after upload
+        if result.get("status") != "queued":
             try:
                 os.unlink(temp_path)
                 os.rmdir(temp_dir)
             except Exception:
                 pass
 
+        return result
+
     def list(self, pattern: Optional[str] = None) -> List[Dict[str, Any]]:
         """
         List files with optional glob pattern filtering.
@@ -881,6 +993,49 @@ class FileBuilder:
             metadata=self._metadata
         )
 
+    def save_image(
+        self,
+        array: Any,
+        *,
+        to: str,
+        quality: int = 95
+    ) -> Dict[str, Any]:
+        """
+        Save numpy array as an image file.
+
+        Args:
+            array: Numpy array representing the image (HxW or HxWxC)
+            to: Target filename (must have image extension like .png, .jpg, .jpeg)
+            quality: JPEG quality (1-100, default: 95). Only used for JPEG files.
+
+        Returns:
+            File metadata dict with id, path, filename, checksum, etc.
+
+        Examples:
+            # Save rendered frame from MuJoCo as PNG
+            pixels = renderer.render()
+            result = dxp.files("frames").save_image(pixels, to="frame_001.png")
+
+            # Save as JPEG with custom quality
+            result = dxp.files("frames").save_image(pixels, to="frame_001.jpg", quality=85)
+
+            # Save numpy array as image
+            import numpy as np
+            img_array = np.random.rand(480, 640, 3) * 255
+            result = dxp.files("images").save_image(img_array, to="random.png")
+        """
+        prefix = '/' + self._path.lstrip('/') if self._path else self._prefix
+
+        return self._save_image(
+            array=array,
+            filename=to,
+            prefix=prefix,
+            description=self._description,
+            tags=self._tags,
+            metadata=self._metadata,
+            quality=quality
+        )
+
     def save_fig(
         self,
         fig: Optional[Any] = None,
@@ -974,29 +1129,32 @@ class FileBuilder:
         # Create parent directories if filename contains path
         os.makedirs(os.path.dirname(temp_path), exist_ok=True)
 
+        frames_ubyte = img_as_ubyte(frame_stack)
         try:
-            frames_ubyte = img_as_ubyte(frame_stack)
-            try:
-                iio.imwrite(temp_path, frames_ubyte, fps=fps, **imageio_kwargs)
-            except iio.core.NeedDownloadError:
-                import imageio.plugins.ffmpeg
-                imageio.plugins.ffmpeg.download()
-                iio.imwrite(temp_path, frames_ubyte, fps=fps, **imageio_kwargs)
+            iio.imwrite(temp_path, frames_ubyte, fps=fps, **imageio_kwargs)
+        except iio.core.NeedDownloadError:
+            import imageio.plugins.ffmpeg
+            imageio.plugins.ffmpeg.download()
+            iio.imwrite(temp_path, frames_ubyte, fps=fps, **imageio_kwargs)
+
+        result = self._save_file(
+            fpath=temp_path,
+            prefix=prefix,
+            description=self._description,
+            tags=self._tags,
+            metadata=self._metadata
+        )
 
-            return self._save_file(
-                fpath=temp_path,
-                prefix=prefix,
-                description=self._description,
-                tags=self._tags,
-                metadata=self._metadata
-            )
-        finally:
+        # Only clean up if NOT queued for buffered upload
+        if result.get("status") != "queued":
             try:
                 os.unlink(temp_path)
                 os.rmdir(temp_dir)
             except Exception:
                 pass
 
+        return result
+
     def duplicate(self, source: Union[str, Dict[str, Any]], to: str) -> Dict[str, Any]:
         """
         Duplicate an existing file to a new path within the same experiment.

ml_dash/run.py

@@ -22,14 +22,43 @@ Usage:
         exp.logs.info(f"Running {RUN.name}")
 """
 
+import functools
 import os
 import sys
+import typing
 from datetime import datetime
 from pathlib import Path
-from typing import Union
+from typing import Optional, Union
 
 from params_proto import EnvVar, proto
 
+
+def requires_open(func):
+  """
+  Decorator that ensures the experiment is open before executing a method.
+
+  Raises:
+      RuntimeError: If experiment is not started
+  """
+
+  @functools.wraps(func)
+  def wrapper(self, *args, **kwargs):
+    if not self._is_open:
+      raise RuntimeError(
+        "Experiment not started. Use 'with experiment.run:' or call experiment.run.start() first.\n"
+        "Example:\n"
+        "  with dxp.run:\n"
+        "      dxp.params.set(lr=0.001)"
+      )
+    return func(self, *args, **kwargs)
+
+  return wrapper
+
+if typing.TYPE_CHECKING:
+  from .client import RemoteClient
+  from .experiment import Experiment
+  from .storage import LocalStorage
+
 PROJECT_ROOT_FILES = ("pyproject.toml", "requirements.txt", "setup.py", "setup.cfg")
 
 
@@ -77,12 +106,19 @@ def find_project_root(
 @proto.prefix
 class RUN:
   """
-  Global Experiment Run Configuration.
+  Global Experiment Run Configuration and
+  Lifecycle manager for experiments
 
   This class is the single source of truth for experiment metadata.
   Configure it before starting an experiment, or through the Experiment
   constructor.
 
+  Supports three usage patterns:
+  1. Method calls: experiment.run.start(), experiment.run.complete()
+  2. Context manager: with Experiment(...).run as exp:
+      3. Decorator: @exp.run or @Experiment(...).run
+
+
   Default prefix template:
       {project}/{now:%Y/%m-%d}/{path_stem}/{job_name}
 
@@ -156,10 +192,21 @@ class RUN:
   debug = "pydevd" in sys.modules
   "set to True automatically for pyCharm"
 
+  _experiment: "Experiment" = None
+  _client: Optional["RemoteClient"] = None
+  _storage: Optional["LocalStorage"] = None
+
+  # Prefix components (parsed from prefix)
+  owner: Optional[str] = None
+  name: Optional[str] = None
+  _folder_path: Optional[str] = None
+
   def __post_init__(self):
     """
+
     Initialize RUN with auto-detected prefix from entry path.
 
+
     Args:
         entry: Path to entry file/directory (e.g., __file__ or directory
                containing sweep.jsonl). If not provided, uses caller's
@@ -179,7 +226,6 @@ class RUN:
         RUN.__post_init__(entry=__file__)
         # Result: RUN.prefix = "vision/resnet", RUN.name = "resnet"
     """
-
     # Use provided entry or try to auto-detect from caller
     if self.entry is None:
       import inspect
@@ -221,6 +267,49 @@ class RUN:
     # for k, v in data.items():
     #   print(f"> {k:>30}: {v}")
 
+    # Parse prefix into components: {owner}/{project}/path.../[name]
+    if self.prefix:
+      self._folder_path = self.prefix
+      parts = self.prefix.strip("/").split("/")
+      if len(parts) >= 2:
+        self.owner = parts[0]
+        self.project = parts[1]
+        # self.name is the last segment
+        self.name = parts[-1] if len(parts) > 2 else parts[1]
+
+  def start(self) -> "Experiment":
+    """
+    Start the experiment (sets status to RUNNING).
+
+    Returns:
+        The experiment instance for chaining
+    """
+    return self._experiment._open()
+
+  def complete(self) -> None:
+    """Mark experiment as completed (status: COMPLETED)."""
+    self._experiment._close(status="COMPLETED")
+
+  def fail(self) -> None:
+    """Mark experiment as failed (status: FAILED)."""
+    self._experiment._close(status="FAILED")
+
+  def cancel(self) -> None:
+    """Mark experiment as cancelled (status: CANCELLED)."""
+    self._experiment._close(status="CANCELLED")
+
+  def __enter__(self) -> "Experiment":
+    """Context manager entry - starts the experiment."""
+    return self.start()
+
+  def __exit__(self, exc_type, exc_val, exc_tb):
+    """Context manager exit - completes or fails the experiment."""
+    if exc_type is not None:
+      self.fail()
+    else:
+      self.complete()
+    return False
+
 
 if __name__ == "__main__":
   RUN.description = ""