podstack 1.2.3__tar.gz → 1.3.1__tar.gz

{podstack-1.2.3 → podstack-1.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: podstack
-Version: 1.2.3
+Version: 1.3.1
 Summary: Official Python SDK for Podstack GPU Notebook Platform
 Author-email: Podstack <support@podstack.ai>
 License-Expression: MIT

{podstack-1.2.3 → podstack-1.3.1}/podstack/__init__.py

@@ -54,7 +54,7 @@ Decorators:
         ...
 """
 
-__version__ = "1.2.3"
+__version__ = "1.3.1"
 
 from .client import Client
 from .notebook import Notebook, NotebookStatus
@@ -84,12 +84,15 @@ from .gpu_runner import (
     GPURunner,
     GPUExecutionResult,
     run as run_on_gpu,
+    list_runners,
 )
 
 # Annotations module import
 from . import annotations
 from .annotations import (
     gpu,
+    runner,
+    RunnerConfig,
     environment,
     auto_shutdown,
     experiment,
@@ -207,9 +210,12 @@ __all__ = [
     "GPURunner",
     "GPUExecutionResult",
     "run_on_gpu",
+    "list_runners",
     # Annotations
     "annotations",
     "gpu",
+    "runner",
+    "RunnerConfig",
     "environment",
     "auto_shutdown",
     "experiment",

{podstack-1.2.3 → podstack-1.3.1}/podstack/annotations.py

@@ -51,6 +51,7 @@ from .gpu_runner import (
 
 # Global state for configuration
 _current_gpu_config: Dict[str, Any] = {}
+_current_runner_config: Optional[Dict[str, Any]] = None
 _current_environment: str = "pytorch"
 _auto_shutdown_minutes: int = 60
 _remote_execution_enabled: bool = True
@@ -212,6 +213,23 @@ if __podstack_result__ is not None:
     print(repr(__podstack_result__))
 """
 
+            # Check for active runner config from @podstack.runner() decorator
+            global _current_runner_config
+            runner_name = None
+            effective_pip = self.pip
+            effective_uv = self.uv
+            effective_conda = self.conda
+            if _current_runner_config is not None:
+                runner_name = _current_runner_config.get("image")
+                # Merge package lists from runner config
+                if _current_runner_config.get("pip") and not effective_pip:
+                    effective_pip = _current_runner_config["pip"]
+                if _current_runner_config.get("uv") and not effective_uv:
+                    effective_uv = _current_runner_config["uv"]
+                if _current_runner_config.get("conda") and not effective_conda:
+                    effective_conda = _current_runner_config["conda"]
+                _current_runner_config = None  # Clear after use
+
             # Execute on remote GPU with improved error handling
             # Auto-enable streaming in Jupyter notebooks for real-time output
             try:
@@ -222,11 +240,12 @@ if __podstack_result__ is not None:
                     fraction=self.fraction,
                     timeout=self.timeout,
                     env=self.env,
-                    pip=self.pip,
-                    uv=self.uv,
-                    conda=self.conda,
+                    pip=effective_pip,
+                    uv=effective_uv,
+                    conda=effective_conda,
                     requirements=self.requirements,
                     use_uv=self.use_uv,
+                    runner=runner_name,
                     wait=True,
                     stream=None  # Auto-detect: True in Jupyter, False otherwise
                 )
@@ -402,6 +421,91 @@ def gpu(
     )
 
 
+class RunnerConfig:
+    """
+    Runner configuration decorator.
+
+    Used alongside @podstack.gpu() to specify which runner image and
+    additional packages to use.
+
+    Usage:
+        @podstack.gpu(type="L40S", count=1)
+        @podstack.runner(image="pytorch-2.3.0-cuda-12.1-py3.11")
+        def my_func():
+            ...
+    """
+
+    def __init__(
+        self,
+        image: str = None,
+        pip: Union[str, list] = None,
+        uv: Union[str, list] = None,
+        conda: Union[str, list] = None,
+    ):
+        self.image = image
+        self.pip = pip
+        self.uv = uv
+        self.conda = conda
+
+        # Store in global state
+        global _current_runner_config
+        _current_runner_config = {
+            "image": image,
+            "pip": pip,
+            "uv": uv,
+            "conda": conda,
+        }
+
+    def __call__(self, func: Callable) -> Callable:
+        """Decorator usage - stores runner config for the next gpu() call."""
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            global _current_runner_config
+            _current_runner_config = {
+                "image": self.image,
+                "pip": self.pip,
+                "uv": self.uv,
+                "conda": self.conda,
+            }
+            return func(*args, **kwargs)
+        return wrapper
+
+
+def runner(
+    image: str = None,
+    pip: Union[str, list] = None,
+    uv: Union[str, list] = None,
+    conda: Union[str, list] = None,
+) -> RunnerConfig:
+    """
+    Configure which runner image and packages to use for GPU execution.
+
+    Used alongside @podstack.gpu() to select a pre-built runner image.
+
+    Examples:
+        @podstack.gpu(type="L40S")
+        @podstack.runner(image="pytorch-2.3.0-cuda-12.1-py3.11")
+        def train():
+            import torch
+            ...
+
+        @podstack.gpu(type="L40S")
+        @podstack.runner(image="pytorch-2.3.0-cuda-12.1-py3.11", pip=["transformers"])
+        def train_llm():
+            ...
+
+    Args:
+        image: Runner name from podstack runners (e.g. "pytorch-2.3.0-cuda-12.1-py3.11")
+        pip: Additional pip packages - string "pkg1,pkg2" or list ["pkg1", "pkg2"]
+        uv: Additional UV packages (faster than pip) - string or list
+        conda: Additional conda packages - string or list
+
+    Returns:
+        RunnerConfig instance
+    """
+    return RunnerConfig(image=image, pip=pip, uv=uv, conda=conda)
+
+
 def environment(env: str = "pytorch") -> str:
     """
     Set the runtime environment.
@@ -708,6 +812,8 @@ def get_auto_shutdown_minutes() -> int:
 
 __all__ = [
     "gpu",
+    "runner",
+    "RunnerConfig",
     "environment",
     "auto_shutdown",
     "experiment",

{podstack-1.2.3 → podstack-1.3.1}/podstack/gpu_runner.py

@@ -268,7 +268,8 @@ class GPURunner:
         uv: Union[str, list] = None,
         conda: Union[str, list] = None,
         requirements: str = None,
-        use_uv: bool = False
+        use_uv: bool = False,
+        runner: str = None
     ) -> str:
         """Build the @podstack annotation string."""
         parts = [f"#@podstack gpu={gpu}"]
@@ -307,6 +308,10 @@ class GPURunner:
         if use_uv:
             parts.append("use_uv=true")
 
+        # Runner name (resolved by notebook-service via platform-service)
+        if runner:
+            parts.append(f"runner={runner}")
+
         return " ".join(parts)
 
     def submit(
@@ -322,7 +327,8 @@ class GPURunner:
         conda: Union[str, list] = None,
         requirements: str = None,
         use_uv: bool = False,
-        add_annotation: bool = True
+        add_annotation: bool = True,
+        runner: str = None
     ) -> Dict[str, Any]:
         """
         Submit code for GPU execution (non-blocking).
@@ -352,7 +358,7 @@ class GPURunner:
 
         # Add annotation if not present
         if add_annotation and not code.strip().startswith("#@podstack"):
-            annotation = self._build_annotation(gpu, count, fraction, timeout, env, pip, uv, conda, requirements, use_uv)
+            annotation = self._build_annotation(gpu, count, fraction, timeout, env, pip, uv, conda, requirements, use_uv, runner)
             code = f"{annotation}\n\n{code}"
 
         # Build installation code for packages
@@ -635,7 +641,8 @@ _stream_install(
         max_retries: int = 3,
         provisioning_timeout: int = 300,
         cancel_on_timeout: bool = True,
-        stream: bool = None
+        stream: bool = None,
+        runner: str = None
     ) -> GPUExecutionResult:
         """
         Execute code on GPU and optionally wait for completion.
@@ -670,7 +677,7 @@ _stream_install(
             ValueError: If parameters are invalid
         """
         # Submit the code
-        submission = self.submit(code, gpu, count, fraction, timeout, env, pip, uv, conda, requirements, use_uv)
+        submission = self.submit(code, gpu, count, fraction, timeout, env, pip, uv, conda, requirements, use_uv, runner=runner)
         execution_id = submission.get("execution_id")
 
         if not execution_id:
@@ -958,6 +965,49 @@ _stream_install(
             except:
                 return {"status": "cancelled"}
 
+    def _get_platform_url(self) -> str:
+        """Derive the platform service URL from the notebook service URL."""
+        if os.environ.get("PODSTACK_PLATFORM_URL"):
+            return os.environ.get("PODSTACK_PLATFORM_URL")
+        # External: cloud.podstack.ai/notebooks -> cloud.podstack.ai/platform
+        if "/notebooks" in self.api_url:
+            return self.api_url.replace("/notebooks", "/platform")
+        # Internal K8s: notebook service URL -> platform service URL
+        if "notebook" in self.api_url and ".svc.cluster.local" in self.api_url:
+            return self.api_url.replace("notebook", "platform").replace(":8084", ":8081")
+        # Fallback: use the same base URL
+        return self.api_url
+
+    def list_runners(self) -> list:
+        """
+        List available GPU/CPU runners.
+
+        Returns:
+            List of runner dictionaries with name, type, description, image, libraries, scenarios.
+        """
+        platform_url = self._get_platform_url()
+        url = f"{platform_url}/api/v1/runners"
+
+        with httpx.Client(timeout=self.timeout) as client:
+            try:
+                response = client.get(url, headers=self._get_headers())
+            except httpx.ConnectError as e:
+                raise ConnectionError(f"Failed to connect to {url}: {e}")
+            except httpx.TimeoutException:
+                raise TimeoutError(f"Request to {url} timed out")
+
+            if response.status_code >= 400:
+                try:
+                    error_msg = response.json().get("error", response.text)
+                except Exception:
+                    error_msg = response.text[:500] if response.text else f"HTTP {response.status_code}"
+                raise RuntimeError(f"Failed to list runners: {error_msg}")
+
+            try:
+                return response.json()
+            except Exception:
+                raise RuntimeError(f"Invalid JSON response: {response.text[:200]}")
+
     def run_function(
         self,
         func: Callable,
@@ -1131,6 +1181,21 @@ def run(
     )
 
 
+def list_runners() -> list:
+    """
+    List available GPU/CPU runners.
+
+    Returns:
+        List of runner dictionaries with name, type, description, image, libraries, scenarios.
+
+    Example:
+        runners = podstack.gpu_runner.list_runners()
+        for r in runners:
+            print(f"{r['name']} ({r['type']}): {r['description']}")
+    """
+    return get_runner().list_runners()
+
+
 def stream_output(execution_id: str, show_output: bool = True) -> Iterator[Dict[str, Any]]:
     """
     Stream real-time output from a running execution.

{podstack-1.2.3 → podstack-1.3.1}/podstack.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: podstack
-Version: 1.2.3
+Version: 1.3.1
 Summary: Official Python SDK for Podstack GPU Notebook Platform
 Author-email: Podstack <support@podstack.ai>
 License-Expression: MIT

{podstack-1.2.3 → podstack-1.3.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "podstack"
-version = "1.2.3"
+version = "1.3.1"
 description = "Official Python SDK for Podstack GPU Notebook Platform"
 readme = "README.md"
 license = "MIT"