podstack 1.2.0__py3-none-any.whl

podstack_gpu/image.py

@@ -0,0 +1,325 @@
+"""Podstack Image - Fluent container image builder."""
+
+from __future__ import annotations
+from typing import List, Dict, Optional, Union
+from dataclasses import dataclass, field
+import json
+
+
+@dataclass
+class ImageDefinition:
+    """Serializable image definition."""
+    base: str = "python:3.11-slim"
+    python_version: str = "3.11"
+    pip_packages: List[str] = field(default_factory=list)
+    requirements_file: Optional[str] = None
+    apt_packages: List[str] = field(default_factory=list)
+    conda_packages: List[str] = field(default_factory=list)
+    env_vars: Dict[str, str] = field(default_factory=dict)
+    workdir: str = "/app"
+    run_commands: List[str] = field(default_factory=list)
+    copy_files: List[tuple] = field(default_factory=list)
+    env_preset: Optional[str] = None  # ml, cv, nlp, audio, tabular, rl, scientific
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for API serialization."""
+        return {
+            "base": self.base,
+            "python_version": self.python_version,
+            "pip_packages": self.pip_packages,
+            "requirements_file": self.requirements_file,
+            "apt_packages": self.apt_packages,
+            "conda_packages": self.conda_packages,
+            "env_vars": self.env_vars,
+            "workdir": self.workdir,
+            "run_commands": self.run_commands,
+            "env_preset": self.env_preset,
+        }
+
+    def to_annotation(self) -> str:
+        """Convert to annotation string for legacy compatibility."""
+        parts = []
+
+        if self.env_preset:
+            parts.append(f"env={self.env_preset}")
+
+        if self.pip_packages:
+            parts.append(f"pip=[{','.join(self.pip_packages)}]")
+
+        if self.requirements_file:
+            parts.append(f"requirements={self.requirements_file}")
+
+        if self.conda_packages:
+            parts.append(f"conda=[{','.join(self.conda_packages)}]")
+
+        if self.python_version != "3.11":
+            parts.append(f"python={self.python_version}")
+
+        return " ".join(parts)
+
+
+class Image:
+    """
+    Fluent container image builder for Podstack GPU.
+
+    Example:
+        image = (
+            podstack.Image.debian_slim()
+            .pip_install("torch", "transformers")
+            .pip_install_from_requirements("requirements.txt")
+            .apt_install("ffmpeg", "libsndfile1")
+            .env({"CUDA_VISIBLE_DEVICES": "0"})
+        )
+
+        # Full GPU for training
+        @app.function(gpu="H100", image=image)
+        def train():
+            import torch
+            ...
+
+        # Fractional GPU for inference (25%, 50%, 75%, or 100%)
+        @app.function(gpu="L40S", fraction=25, image=image)
+        def inference(data):
+            return model.predict(data)
+    """
+
+    def __init__(self, definition: Optional[ImageDefinition] = None):
+        self._definition = definition or ImageDefinition()
+
+    def _copy(self) -> Image:
+        """Create a copy of this image for chaining."""
+        import copy
+        return Image(copy.deepcopy(self._definition))
+
+    # ===== Base Image Factory Methods =====
+
+    @classmethod
+    def debian_slim(cls, python_version: str = "3.11") -> Image:
+        """Create image based on Debian slim with Python."""
+        return cls(ImageDefinition(
+            base=f"python:{python_version}-slim",
+            python_version=python_version,
+        ))
+
+    @classmethod
+    def from_registry(cls, image_name: str) -> Image:
+        """Create image from a Docker registry."""
+        return cls(ImageDefinition(base=image_name))
+
+    @classmethod
+    def from_dockerfile(cls, path: str, context_dir: str = ".") -> Image:
+        """Create image from a Dockerfile."""
+        img = cls(ImageDefinition())
+        img._definition.base = f"dockerfile:{path}"
+        return img
+
+    @classmethod
+    def micromamba(cls, python_version: str = "3.11") -> Image:
+        """Create image with micromamba for conda packages."""
+        return cls(ImageDefinition(
+            base="mambaorg/micromamba:latest",
+            python_version=python_version,
+        ))
+
+    # ===== Environment Preset Factory Methods =====
+
+    @classmethod
+    def ml(cls) -> Image:
+        """Pre-configured image for machine learning (PyTorch, scikit-learn, etc.)."""
+        img = cls.debian_slim()
+        img._definition.env_preset = "ml"
+        return img
+
+    @classmethod
+    def cv(cls) -> Image:
+        """Pre-configured image for computer vision (PyTorch, torchvision, opencv, etc.)."""
+        img = cls.debian_slim()
+        img._definition.env_preset = "cv"
+        return img
+
+    @classmethod
+    def nlp(cls) -> Image:
+        """Pre-configured image for NLP (transformers, tokenizers, etc.)."""
+        img = cls.debian_slim()
+        img._definition.env_preset = "nlp"
+        return img
+
+    @classmethod
+    def audio(cls) -> Image:
+        """Pre-configured image for audio processing (torchaudio, librosa, etc.)."""
+        img = cls.debian_slim()
+        img._definition.env_preset = "audio"
+        return img
+
+    @classmethod
+    def tabular(cls) -> Image:
+        """Pre-configured image for tabular data (pandas, xgboost, etc.)."""
+        img = cls.debian_slim()
+        img._definition.env_preset = "tabular"
+        return img
+
+    @classmethod
+    def rl(cls) -> Image:
+        """Pre-configured image for reinforcement learning (gymnasium, stable-baselines3, etc.)."""
+        img = cls.debian_slim()
+        img._definition.env_preset = "rl"
+        return img
+
+    @classmethod
+    def scientific(cls) -> Image:
+        """Pre-configured image for scientific computing (numpy, scipy, matplotlib, etc.)."""
+        img = cls.debian_slim()
+        img._definition.env_preset = "scientific"
+        return img
+
+    # ===== Fluent Builder Methods =====
+
+    def pip_install(self, *packages: str, find_links: str = None,
+                    index_url: str = None, extra_index_url: str = None) -> Image:
+        """
+        Install pip packages.
+
+        Args:
+            *packages: Package names (e.g., "torch", "transformers>=4.30.0")
+            find_links: URL to find packages
+            index_url: Custom PyPI index URL
+            extra_index_url: Additional PyPI index URL
+
+        Example:
+            image.pip_install("torch", "transformers>=4.30.0")
+            image.pip_install("torch", index_url="https://download.pytorch.org/whl/cu118")
+        """
+        new_img = self._copy()
+        new_img._definition.pip_packages.extend(packages)
+        return new_img
+
+    def pip_install_from_requirements(self, path: str) -> Image:
+        """
+        Install packages from a requirements.txt file.
+
+        Args:
+            path: Path to requirements.txt file
+
+        Example:
+            image.pip_install_from_requirements("requirements.txt")
+        """
+        new_img = self._copy()
+        new_img._definition.requirements_file = path
+        return new_img
+
+    def apt_install(self, *packages: str) -> Image:
+        """
+        Install apt packages (Debian/Ubuntu).
+
+        Args:
+            *packages: Package names (e.g., "ffmpeg", "libsndfile1")
+
+        Example:
+            image.apt_install("ffmpeg", "libsndfile1")
+        """
+        new_img = self._copy()
+        new_img._definition.apt_packages.extend(packages)
+        return new_img
+
+    def conda_install(self, *packages: str, channels: List[str] = None) -> Image:
+        """
+        Install conda packages.
+
+        Args:
+            *packages: Package names
+            channels: Conda channels to use
+
+        Example:
+            image.conda_install("pytorch", "cudatoolkit=11.8", channels=["pytorch"])
+        """
+        new_img = self._copy()
+        new_img._definition.conda_packages.extend(packages)
+        return new_img
+
+    def env(self, env_vars: Dict[str, str]) -> Image:
+        """
+        Set environment variables.
+
+        Args:
+            env_vars: Dictionary of environment variables
+
+        Example:
+            image.env({"CUDA_VISIBLE_DEVICES": "0", "HF_HOME": "/cache/huggingface"})
+        """
+        new_img = self._copy()
+        new_img._definition.env_vars.update(env_vars)
+        return new_img
+
+    def workdir(self, path: str) -> Image:
+        """
+        Set the working directory.
+
+        Args:
+            path: Working directory path
+
+        Example:
+            image.workdir("/app/src")
+        """
+        new_img = self._copy()
+        new_img._definition.workdir = path
+        return new_img
+
+    def run_commands(self, *commands: str) -> Image:
+        """
+        Run shell commands during image build.
+
+        Args:
+            *commands: Shell commands to run
+
+        Example:
+            image.run_commands("apt-get update", "apt-get install -y curl")
+        """
+        new_img = self._copy()
+        new_img._definition.run_commands.extend(commands)
+        return new_img
+
+    def copy_local_file(self, local_path: str, remote_path: str) -> Image:
+        """
+        Copy a local file into the image.
+
+        Args:
+            local_path: Local file path
+            remote_path: Path in the container
+
+        Example:
+            image.copy_local_file("./model.py", "/app/model.py")
+        """
+        new_img = self._copy()
+        new_img._definition.copy_files.append((local_path, remote_path))
+        return new_img
+
+    def add_local_dir(self, local_path: str, remote_path: str) -> Image:
+        """
+        Add a local directory to the image.
+
+        Args:
+            local_path: Local directory path
+            remote_path: Path in the container
+        """
+        return self.copy_local_file(local_path, remote_path)
+
+    # ===== Properties and Serialization =====
+
+    @property
+    def definition(self) -> ImageDefinition:
+        """Get the image definition."""
+        return self._definition
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for API serialization."""
+        return self._definition.to_dict()
+
+    def __repr__(self) -> str:
+        parts = [f"Image(base={self._definition.base!r}"]
+        if self._definition.env_preset:
+            parts.append(f", env={self._definition.env_preset!r}")
+        if self._definition.pip_packages:
+            parts.append(f", pip={self._definition.pip_packages}")
+        if self._definition.apt_packages:
+            parts.append(f", apt={self._definition.apt_packages}")
+        return "".join(parts) + ")"