ml-dash 0.0.17__py3-none-any.whl → 0.4.0__py3-none-any.whl

ml_dash/components/files.py

@@ -0,0 +1,246 @@
+"""File management component for ML-Logger."""
+
+import json
+import pickle
+from typing import Any, Optional
+from pathlib import Path
+
+from ..backends.base import StorageBackend
+
+
+class FileManager:
+    """Manages file storage and retrieval.
+
+    Files are stored in the files/ subdirectory.
+
+    Args:
+        backend: Storage backend
+        prefix: Experiment prefix path
+        namespace: Optional namespace for files (e.g., "checkpoints")
+    """
+
+    def __init__(
+        self,
+        backend: StorageBackend,
+        prefix: str,
+        namespace: str = ""
+    ):
+        """Initialize file manager.
+
+        Args:
+            backend: Storage backend
+            prefix: Experiment prefix path
+            namespace: Optional namespace subdirectory
+        """
+        self.backend = backend
+        self.prefix = prefix
+        self.namespace = namespace
+
+    def _get_file_path(self, filename: str) -> str:
+        """Get full file path with namespace.
+
+        Args:
+            filename: File name
+
+        Returns:
+            Full path including prefix, files/, and namespace
+        """
+        parts = [self.prefix, "files"]
+        if self.namespace:
+            parts.append(self.namespace)
+        parts.append(filename)
+        return "/".join(parts)
+
+    def save(self, data: Any, filename: str) -> None:
+        """Save data to a file (auto-detects format).
+
+        Supports: JSON (.json), pickle (.pkl, .pickle), PyTorch (.pt, .pth),
+        NumPy (.npy, .npz), and raw bytes.
+
+        Args:
+            data: Data to save
+            filename: File name
+        """
+        file_path = self._get_file_path(filename)
+        suffix = Path(filename).suffix.lower()
+
+        if suffix == ".json":
+            # Save as JSON
+            json_str = json.dumps(data, indent=2)
+            self.backend.write_text(file_path, json_str)
+
+        elif suffix in [".pkl", ".pickle"]:
+            # Save as pickle
+            pickled = pickle.dumps(data)
+            self.backend.write_bytes(file_path, pickled)
+
+        elif suffix in [".pt", ".pth"]:
+            # Save PyTorch tensor/model
+            try:
+                import torch
+                import io
+                buffer = io.BytesIO()
+                torch.save(data, buffer)
+                self.backend.write_bytes(file_path, buffer.getvalue())
+            except ImportError:
+                raise ImportError("PyTorch is required to save .pt/.pth files")
+
+        elif suffix in [".npy", ".npz"]:
+            # Save NumPy array
+            try:
+                import numpy as np
+                import io
+                buffer = io.BytesIO()
+                if suffix == ".npy":
+                    np.save(buffer, data)
+                else:
+                    np.savez(buffer, data)
+                self.backend.write_bytes(file_path, buffer.getvalue())
+            except ImportError:
+                raise ImportError("NumPy is required to save .npy/.npz files")
+
+        else:
+            # Save as raw bytes
+            if isinstance(data, bytes):
+                self.backend.write_bytes(file_path, data)
+            elif isinstance(data, str):
+                self.backend.write_text(file_path, data)
+            else:
+                # Fallback to pickle
+                pickled = pickle.dumps(data)
+                self.backend.write_bytes(file_path, pickled)
+
+    def save_pkl(self, data: Any, filename: str) -> None:
+        """Save data as pickle file.
+
+        Args:
+            data: Data to save
+            filename: File name (will add .pkl if missing)
+        """
+        if not filename.endswith((".pkl", ".pickle")):
+            filename = f"{filename}.pkl"
+        self.save(data, filename)
+
+    def load(self, filename: str) -> Any:
+        """Load data from a file (auto-detects format).
+
+        Args:
+            filename: File name
+
+        Returns:
+            Loaded data
+        """
+        file_path = self._get_file_path(filename)
+        suffix = Path(filename).suffix.lower()
+
+        if suffix == ".json":
+            # Load JSON
+            json_str = self.backend.read_text(file_path)
+            return json.loads(json_str)
+
+        elif suffix in [".pkl", ".pickle"]:
+            # Load pickle
+            pickled = self.backend.read_bytes(file_path)
+            return pickle.loads(pickled)
+
+        elif suffix in [".pt", ".pth"]:
+            # Load PyTorch
+            try:
+                import torch
+                import io
+                data = self.backend.read_bytes(file_path)
+                buffer = io.BytesIO(data)
+                return torch.load(buffer)
+            except ImportError:
+                raise ImportError("PyTorch is required to load .pt/.pth files")
+
+        elif suffix in [".npy", ".npz"]:
+            # Load NumPy
+            try:
+                import numpy as np
+                import io
+                data = self.backend.read_bytes(file_path)
+                buffer = io.BytesIO(data)
+                if suffix == ".npy":
+                    return np.load(buffer)
+                else:
+                    return np.load(buffer, allow_pickle=True)
+            except ImportError:
+                raise ImportError("NumPy is required to load .npy/.npz files")
+
+        else:
+            # For unknown extensions, try different strategies
+            data = self.backend.read_bytes(file_path)
+
+            # If it looks like a binary extension, return bytes directly
+            if suffix in [".bin", ".dat", ".raw"]:
+                return data
+
+            # Try to unpickle first (handles custom extensions from save())
+            try:
+                return pickle.loads(data)
+            except (pickle.UnpicklingError, EOFError, AttributeError):
+                pass
+
+            # Try to decode as text
+            try:
+                return data.decode('utf-8')
+            except UnicodeDecodeError:
+                # Return raw bytes as fallback
+                return data
+
+    def load_torch(self, filename: str) -> Any:
+        """Load PyTorch checkpoint.
+
+        Args:
+            filename: File name
+
+        Returns:
+            Loaded PyTorch data
+        """
+        if not filename.endswith((".pt", ".pth")):
+            filename = f"{filename}.pt"
+        return self.load(filename)
+
+    def __call__(self, namespace: str) -> "FileManager":
+        """Create a namespaced file manager.
+
+        Args:
+            namespace: Namespace subdirectory (e.g., "checkpoints")
+
+        Returns:
+            New FileManager with the namespace
+        """
+        new_namespace = f"{self.namespace}/{namespace}" if self.namespace else namespace
+        return FileManager(
+            backend=self.backend,
+            prefix=self.prefix,
+            namespace=new_namespace
+        )
+
+    def exists(self, filename: str) -> bool:
+        """Check if a file exists.
+
+        Args:
+            filename: File name
+
+        Returns:
+            True if file exists
+        """
+        file_path = self._get_file_path(filename)
+        return self.backend.exists(file_path)
+
+    def list(self) -> list:
+        """List files in the current namespace.
+
+        Returns:
+            List of file names
+        """
+        dir_path = f"{self.prefix}/files"
+        if self.namespace:
+            dir_path = f"{dir_path}/{self.namespace}"
+
+        try:
+            return self.backend.list_dir(dir_path)
+        except Exception:
+            return []

ml_dash/components/logs.py

@@ -0,0 +1,104 @@
+"""Text logging component for ML-Logger."""
+
+import json
+import time
+from typing import Any, Dict, Optional, List
+
+from ..backends.base import StorageBackend
+
+
+class LogManager:
+    """Manages structured text logging.
+
+    Logs are stored in a JSONL file (logs.jsonl).
+
+    Args:
+        backend: Storage backend
+        prefix: Experiment prefix path
+    """
+
+    def __init__(self, backend: StorageBackend, prefix: str):
+        """Initialize log manager.
+
+        Args:
+            backend: Storage backend
+            prefix: Experiment prefix path
+        """
+        self.backend = backend
+        self.prefix = prefix
+        self.logs_file = f"{prefix}/logs.jsonl"
+
+    def log(self, message: str, level: str = "INFO", **context) -> None:
+        """Log a message with context.
+
+        Args:
+            message: Log message
+            level: Log level (DEBUG, INFO, WARNING, ERROR)
+            **context: Additional context fields
+        """
+        entry = {
+            "timestamp": time.time(),
+            "level": level.upper(),
+            "message": message,
+        }
+
+        if context:
+            entry["context"] = context
+
+        line = json.dumps(entry) + "\n"
+        self.backend.append_text(self.logs_file, line)
+
+    def info(self, message: str, **context) -> None:
+        """Log an info message.
+
+        Args:
+            message: Log message
+            **context: Additional context fields
+        """
+        self.log(message, level="INFO", **context)
+
+    def warning(self, message: str, **context) -> None:
+        """Log a warning message.
+
+        Args:
+            message: Log message
+            **context: Additional context fields
+        """
+        self.log(message, level="WARNING", **context)
+
+    def error(self, message: str, **context) -> None:
+        """Log an error message.
+
+        Args:
+            message: Log message
+            **context: Additional context fields
+        """
+        self.log(message, level="ERROR", **context)
+
+    def debug(self, message: str, **context) -> None:
+        """Log a debug message.
+
+        Args:
+            message: Log message
+            **context: Additional context fields
+        """
+        self.log(message, level="DEBUG", **context)
+
+    def read(self) -> List[Dict[str, Any]]:
+        """Read all logs from file.
+
+        Returns:
+            List of log entries
+        """
+        if not self.backend.exists(self.logs_file):
+            return []
+
+        content = self.backend.read_text(self.logs_file)
+        logs = []
+
+        for line in content.strip().split("\n"):
+            if not line:
+                continue
+            logs.append(json.loads(line))
+
+        return logs

ml_dash/components/metrics.py

@@ -0,0 +1,169 @@
+"""Metrics logging component for ML-Logger."""
+
+import json
+import time
+from typing import Any, Dict, Optional, List
+from collections import defaultdict
+
+from ..backends.base import StorageBackend
+
+
+class MetricsLogger:
+    """Logs metrics with support for namespacing and aggregation.
+
+    Metrics are stored in a single JSONL file (metrics.jsonl).
+
+    Args:
+        backend: Storage backend
+        prefix: Experiment prefix path
+        namespace: Optional namespace for metrics (e.g., "train", "val")
+    """
+
+    def __init__(
+        self,
+        backend: StorageBackend,
+        prefix: str,
+        namespace: str = ""
+    ):
+        """Initialize metrics logger.
+
+        Args:
+            backend: Storage backend
+            prefix: Experiment prefix path
+            namespace: Optional namespace prefix
+        """
+        self.backend = backend
+        self.prefix = prefix
+        self.namespace = namespace
+        self.metrics_file = f"{prefix}/metrics.jsonl"
+        self._collect_buffer: Dict[str, List[float]] = defaultdict(list)
+
+    def log(self, step: Optional[int] = None, **metrics) -> None:
+        """Log metrics immediately.
+
+        Args:
+            step: Step number (epoch, iteration, etc.)
+            **metrics: Metric name-value pairs
+        """
+        # Apply namespace to metric names
+        namespaced_metrics = {}
+        for key, value in metrics.items():
+            if self.namespace:
+                key = f"{self.namespace}.{key}"
+            namespaced_metrics[key] = value
+
+        entry = {
+            "timestamp": time.time(),
+            "metrics": namespaced_metrics
+        }
+
+        if step is not None:
+            entry["step"] = step
+
+        line = json.dumps(entry) + "\n"
+        self.backend.append_text(self.metrics_file, line)
+
+    def collect(self, step: Optional[int] = None, **metrics) -> None:
+        """Collect metrics for later aggregation.
+
+        Args:
+            step: Step number (optional, used by flush)
+            **metrics: Metric name-value pairs
+        """
+        for key, value in metrics.items():
+            if self.namespace:
+                key = f"{self.namespace}.{key}"
+            self._collect_buffer[key].append(float(value))
+
+    def flush(
+        self,
+        _aggregation: str = "mean",
+        step: Optional[int] = None,
+        **additional_metrics
+    ) -> None:
+        """Flush collected metrics with aggregation.
+
+        Args:
+            _aggregation: Aggregation method ("mean", "sum", "min", "max", "last")
+            step: Step number for logged metrics
+            **additional_metrics: Additional metrics to log (not aggregated)
+        """
+        if not self._collect_buffer and not additional_metrics:
+            return
+
+        aggregated = {}
+
+        # Aggregate collected metrics
+        for key, values in self._collect_buffer.items():
+            if not values:
+                continue
+
+            if _aggregation == "mean":
+                aggregated[key] = sum(values) / len(values)
+            elif _aggregation == "sum":
+                aggregated[key] = sum(values)
+            elif _aggregation == "min":
+                aggregated[key] = min(values)
+            elif _aggregation == "max":
+                aggregated[key] = max(values)
+            elif _aggregation == "last":
+                aggregated[key] = values[-1]
+            else:
+                raise ValueError(f"Unknown aggregation method: {_aggregation}")
+
+        # Add non-aggregated metrics
+        for key, value in additional_metrics.items():
+            if self.namespace:
+                key = f"{self.namespace}.{key}"
+            aggregated[key] = value
+
+        # Log aggregated metrics
+        if aggregated:
+            entry = {
+                "timestamp": time.time(),
+                "metrics": aggregated
+            }
+
+            if step is not None:
+                entry["step"] = step
+
+            line = json.dumps(entry) + "\n"
+            self.backend.append_text(self.metrics_file, line)
+
+        # Clear buffer
+        self._collect_buffer.clear()
+
+    def __call__(self, namespace: str) -> "MetricsLogger":
+        """Create a namespaced metrics logger.
+
+        Args:
+            namespace: Namespace name (e.g., "train", "val")
+
+        Returns:
+            New MetricsLogger with the namespace
+        """
+        new_namespace = f"{self.namespace}.{namespace}" if self.namespace else namespace
+        return MetricsLogger(
+            backend=self.backend,
+            prefix=self.prefix,
+            namespace=new_namespace
+        )
+
+    def read(self) -> List[Dict[str, Any]]:
+        """Read all metrics from file.
+
+        Returns:
+            List of metric entries
+        """
+        if not self.backend.exists(self.metrics_file):
+            return []
+
+        content = self.backend.read_text(self.metrics_file)
+        metrics = []
+
+        for line in content.strip().split("\n"):
+            if not line:
+                continue
+            metrics.append(json.loads(line))
+
+        return metrics

ml_dash/components/parameters.py

@@ -0,0 +1,144 @@
+"""Parameter management component for ML-Logger."""
+
+import json
+import time
+from typing import Any, Dict, Optional
+
+from ..backends.base import StorageBackend
+
+
+def deep_merge(base: Dict, updates: Dict) -> Dict:
+    """Deep merge two dictionaries.
+
+    Args:
+        base: Base dictionary
+        updates: Updates to merge in
+
+    Returns:
+        Merged dictionary
+    """
+    result = base.copy()
+    for key, value in updates.items():
+        if key in result and isinstance(result[key], dict) and isinstance(value, dict):
+            result[key] = deep_merge(result[key], value)
+        else:
+            result[key] = value
+    return result
+
+
+class ParameterManager:
+    """Manages experiment parameters.
+
+    Parameters are stored in an append-only JSONL file (parameters.jsonl)
+    with operations: set, extend, update.
+
+    Args:
+        backend: Storage backend
+        prefix: Experiment prefix path
+    """
+
+    def __init__(self, backend: StorageBackend, prefix: str):
+        """Initialize parameter manager.
+
+        Args:
+            backend: Storage backend
+            prefix: Experiment prefix path
+        """
+        self.backend = backend
+        self.prefix = prefix
+        self.params_file = f"{prefix}/parameters.jsonl"
+        self._cache: Optional[Dict[str, Any]] = None
+
+    def set(self, **kwargs) -> None:
+        """Set parameters (replaces existing).
+
+        Args:
+            **kwargs: Parameter key-value pairs
+        """
+        self._append_operation("set", data=kwargs)
+        self._cache = None  # Invalidate cache
+
+    def extend(self, **kwargs) -> None:
+        """Extend parameters (deep merge with existing).
+
+        Args:
+            **kwargs: Parameter key-value pairs to merge
+        """
+        self._append_operation("extend", data=kwargs)
+        self._cache = None  # Invalidate cache
+
+    def update(self, key: str, value: Any) -> None:
+        """Update a single parameter.
+
+        Args:
+            key: Parameter key (can be dot-separated like "model.layers")
+            value: New value
+        """
+        self._append_operation("update", key=key, value=value)
+        self._cache = None  # Invalidate cache
+
+    def read(self) -> Dict[str, Any]:
+        """Read current parameters by replaying operations.
+
+        Returns:
+            Current parameter dictionary
+        """
+        if self._cache is not None:
+            return self._cache.copy()
+
+        params = {}
+
+        if not self.backend.exists(self.params_file):
+            return params
+
+        # Read and replay all operations
+        content = self.backend.read_text(self.params_file)
+        for line in content.strip().split("\n"):
+            if not line:
+                continue
+
+            operation = json.loads(line)
+            op_type = operation.get("operation")
+
+            if op_type == "set":
+                params = operation.get("data", {})
+            elif op_type == "extend":
+                params = deep_merge(params, operation.get("data", {}))
+            elif op_type == "update":
+                key = operation.get("key")
+                value = operation.get("value")
+                if key:
+                    # Support dot-separated keys
+                    keys = key.split(".")
+                    current = params
+                    for k in keys[:-1]:
+                        if k not in current:
+                            current[k] = {}
+                        current = current[k]
+                    current[keys[-1]] = value
+
+        self._cache = params
+        return params.copy()
+
+    def log(self, **kwargs) -> None:
+        """Alias for set() to match API documentation.
+
+        Args:
+            **kwargs: Parameter key-value pairs
+        """
+        self.set(**kwargs)
+
+    def _append_operation(self, operation: str, **kwargs) -> None:
+        """Append an operation to the parameters file.
+
+        Args:
+            operation: Operation type (set, extend, update)
+            **kwargs: Operation-specific data
+        """
+        entry = {
+            "timestamp": time.time(),
+            "operation": operation,
+            **kwargs
+        }
+        line = json.dumps(entry) + "\n"
+        self.backend.append_text(self.params_file, line)

ml_dash/job_logger.py

@@ -0,0 +1,42 @@
+"""JobLogger - Job-based logging wrapper around ML_Logger.
+
+This class provides a simple wrapper around ML_Logger for job-based logging.
+"""
+
+from typing import Optional
+
+from .ml_logger import ML_Logger
+from .backends.base import StorageBackend
+
+
+class JobLogger(ML_Logger):
+    """Job-based logger wrapper.
+
+    This is a simple wrapper around ML_Logger that can be extended
+    with job-specific functionality in the future.
+
+    Args:
+        prefix: Directory prefix for logging
+        backend: Storage backend (optional)
+        job_id: Optional job identifier
+    """
+
+    def __init__(
+        self,
+        prefix: str,
+        backend: Optional[StorageBackend] = None,
+        job_id: Optional[str] = None,
+    ):
+        """Initialize JobLogger.
+
+        Args:
+            prefix: Directory prefix for logging
+            backend: Storage backend (optional)
+            job_id: Optional job identifier
+        """
+        super().__init__(prefix, backend)
+        self.job_id = job_id
+
+    def __repr__(self) -> str:
+        """String representation."""
+        return f"JobLogger(prefix='{self.prefix}', job_id='{self.job_id}', entries={len(self.buffer)})"