podstack 1.3.17__tar.gz → 1.3.20__tar.gz

{podstack-1.3.17 → podstack-1.3.20}/PKG-INFO

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

{podstack-1.3.17 → podstack-1.3.20}/podstack/annotations.py

@@ -90,6 +90,7 @@ class GPUConfig:
         conda: Union[str, list] = None,
         requirements: str = None,
         use_uv: bool = False,
+        env_vars: Optional[Dict[str, str]] = None,
         remote: bool = None
     ):
         """
@@ -121,6 +122,7 @@ class GPUConfig:
         self.conda = conda
         self.requirements = requirements
         self.use_uv = use_uv
+        self.env_vars = env_vars or {}
         self.remote = remote if remote is not None else _remote_execution_enabled
 
         # Store in global config
@@ -138,6 +140,7 @@ class GPUConfig:
             "conda": conda,
             "requirements": requirements,
             "use_uv": use_uv,
+            "env_vars": env_vars or {},
         }
 
     def __call__(self, func: Callable) -> Callable:
@@ -155,9 +158,6 @@ class GPUConfig:
                 print(f"[Podstack] GPU Config (local): {self.type} x{self.count} @ {self.fraction}%")
                 return func(*args, **kwargs)
 
-            # Remote execution on GPU
-            print(f"[Podstack] Provisioning GPU: {self.type} x{self.count} @ {self.fraction}%")
-
             try:
                 runner = get_runner()
             except ValueError as e:
@@ -249,6 +249,7 @@ if __podstack_result__ is not None:
                     conda=effective_conda,
                     requirements=self.requirements,
                     use_uv=self.use_uv,
+                    env_vars=self.env_vars,
                     runner=runner_name,
                     wait=True,
                     stream=None  # Auto-detect: True in Jupyter, False otherwise
@@ -340,6 +341,7 @@ def gpu(
     conda: Union[str, list] = None,
     requirements: str = None,
     use_uv: bool = False,
+    env_vars: Optional[Dict[str, str]] = None,
     remote: bool = None
 ) -> GPUConfig:
     """
@@ -421,6 +423,7 @@ def gpu(
         conda=conda,
         requirements=requirements,
         use_uv=use_uv,
+        env_vars=env_vars,
         remote=remote
     )
 

{podstack-1.3.17 → podstack-1.3.20}/podstack/gpu_runner.py

@@ -18,6 +18,99 @@ import httpx
 # Configure logging
 logger = logging.getLogger("podstack.gpu_runner")
 
+SPINNER_FRAMES = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
+
+
+class LiveDisplay:
+    """Animated phase display for GPU runner lifecycle in Jupyter + terminal."""
+
+    PHASES = {
+        "pending":      ("🔍", "Searching for GPU..."),
+        "queued":       ("📋", "Queued — waiting for available GPU..."),
+        "provisioning": ("🚀", "Allocating GPU pod..."),
+        "running":      None,   # No spinner — logs stream directly
+    }
+    CHECKMARKS = {
+        "pending":      "🔍 Submitted",
+        "queued":       "📋 In queue",
+        "provisioning": "✓  GPU pod provisioning...",
+        "running":      "✓  Pod ready — logging live output:",
+    }
+
+    def __init__(self, gpu_type: str, gpu_count: int, fraction: int):
+        self._gpu_type = gpu_type
+        self._gpu_count = gpu_count
+        self._fraction = fraction
+        self._phase = None
+        self._spinner_thread: Optional[threading.Thread] = None
+        self._stop_evt = threading.Event()
+        self._logs_started = False
+
+    def set_phase(self, status: str, extra: str = ""):
+        """Transition to a new lifecycle phase."""
+        self._stop_spinner()
+        if self._phase and self._phase != status:
+            label = self.CHECKMARKS.get(self._phase, f"✓ {self._phase}")
+            self._print(f"\r[Podstack] {label}{' ' * 30}\n")
+        self._phase = status
+        if status == "running":
+            if not self._logs_started:
+                self._logs_started = True
+                self._print(f"[Podstack] ─── Live Logs ({self._gpu_type} x{self._gpu_count}) ───\n\n")
+        else:
+            phase_info = self.PHASES.get(status)
+            if phase_info:
+                _, label = phase_info
+                if extra:
+                    label = f"{label} {extra}"
+                self._start_spinner(f"[Podstack] {{spinner}} {label}")
+
+    def log(self, content: str, output_type: str = "stdout"):
+        """Write a log line from the pod (indented, real-time)."""
+        if output_type == "stderr":
+            sys.stderr.write(f"  {content}" if not content.startswith("  ") else content)
+            sys.stderr.flush()
+        else:
+            sys.stdout.write(f"  {content}" if not content.startswith("  ") else content)
+            sys.stdout.flush()
+
+    def complete(self, success: bool, gpu_seconds: float, cost_paise: int, error: str = None):
+        """Print final summary line."""
+        self._stop_spinner()
+        if self._logs_started:
+            self._print(f"\n[Podstack] ─────────────────────────────────────\n")
+        if success:
+            cost_str = f"₹{cost_paise/100:.2f}" if cost_paise else ""
+            self._print(f"[Podstack] ✓ Completed in {gpu_seconds:.1f}s  |  {self._gpu_type} x{self._gpu_count}  |  {cost_str}\n")
+        else:
+            self._print(f"[Podstack] ✗ Failed: {error}\n")
+
+    def _start_spinner(self, template: str):
+        self._stop_evt.clear()
+
+        def _spin():
+            i = 0
+            while not self._stop_evt.is_set():
+                frame = SPINNER_FRAMES[i % len(SPINNER_FRAMES)]
+                sys.stdout.write(f"\r{template.format(spinner=frame)}   ")
+                sys.stdout.flush()
+                self._stop_evt.wait(0.1)
+                i += 1
+
+        self._spinner_thread = threading.Thread(target=_spin, daemon=True)
+        self._spinner_thread.start()
+
+    def _stop_spinner(self):
+        if self._spinner_thread and self._spinner_thread.is_alive():
+            self._stop_evt.set()
+            self._spinner_thread.join(timeout=0.5)
+        self._spinner_thread = None
+
+    @staticmethod
+    def _print(msg: str):
+        sys.stdout.write(msg)
+        sys.stdout.flush()
+
 
 def is_jupyter() -> bool:
     """Check if running in a Jupyter notebook."""
@@ -361,6 +454,7 @@ class GPURunner:
         conda: Union[str, list] = None,
         requirements: str = None,
         use_uv: bool = False,
+        env_vars: Dict[str, str] = None,
         add_annotation: bool = True,
         runner: str = None
     ) -> Dict[str, Any]:
@@ -395,6 +489,13 @@ class GPURunner:
             annotation = self._build_annotation(gpu, count, fraction, timeout, env, pip, uv, conda, requirements, use_uv, runner)
             code = f"{annotation}\n\n{code}"
 
+        # Inject environment variables
+        if env_vars:
+            env_lines = ["import os"]
+            for k, v in env_vars.items():
+                env_lines.append(f"os.environ[{repr(k)}] = {repr(str(v))}")
+            code = "\n".join(env_lines) + "\n\n" + code
+
         # Build installation code for packages
         install_parts = []
 
@@ -672,6 +773,7 @@ _stream_install(
         conda: Union[str, list] = None,
         requirements: str = None,
         use_uv: bool = False,
+        env_vars: Dict[str, str] = None,
         wait: bool = True,
         poll_interval: float = 2.0,
         max_retries: int = 3,
@@ -713,14 +815,12 @@ _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, runner=runner)
+        submission = self.submit(code, gpu, count, fraction, timeout, env, pip, uv, conda, requirements, use_uv, env_vars=env_vars, runner=runner)
         execution_id = submission.get("execution_id")
 
         if not execution_id:
             raise RuntimeError(f"No execution_id in response: {submission}")
 
-        print(f"[Podstack] Execution submitted: {execution_id}")
-
         if not wait:
             return GPUExecutionResult(
                 execution_id=execution_id,
@@ -733,9 +833,9 @@ _stream_install(
         should_stream = stream if stream is not None else is_jupyter()
 
         if should_stream:
-            return self._run_with_streaming(execution_id, gpu, count, timeout, max_retries, cancel_on_timeout)
+            return self._run_with_streaming(execution_id, gpu, count, timeout, max_retries, cancel_on_timeout, fraction)
         else:
-            return self._run_with_polling(execution_id, gpu, count, timeout, poll_interval, max_retries, provisioning_timeout, cancel_on_timeout)
+            return self._run_with_polling(execution_id, gpu, count, timeout, poll_interval, max_retries, provisioning_timeout, cancel_on_timeout, fraction)
 
     def _run_with_streaming(
         self,
@@ -744,10 +844,12 @@ _stream_install(
         count: int,
         timeout: int,
         max_retries: int,
-        cancel_on_timeout: bool
+        cancel_on_timeout: bool,
+        fraction: int = 100
     ) -> GPUExecutionResult:
         """Run execution with real-time output streaming."""
-        print(f"[Podstack] Waiting for GPU runner ({gpu} x{count})...")
+        display = LiveDisplay(gpu, count, fraction)
+        display.set_phase("pending")
 
         start_time = time.time()
         output_buffer = []
@@ -755,13 +857,14 @@ _stream_install(
         final_event = {}
 
         try:
-            for event in self.stream_output(execution_id, show_output=True):
+            for event in self.stream_output(execution_id, show_output=False):
                 elapsed = time.time() - start_time
                 if elapsed > timeout:
                     if cancel_on_timeout:
                         try:
                             self.cancel(execution_id)
-                            print(f"\n[Podstack] Execution cancelled due to timeout")
+                            display._stop_spinner()
+                            display._print(f"\r[Podstack] Execution cancelled due to timeout{' ' * 30}\n")
                         except Exception as e:
                             logger.warning(f"Failed to cancel execution: {e}")
 
@@ -778,37 +881,41 @@ _stream_install(
                         )
                     )
 
-                # Track output
-                if event.get("type") in ("stdout", "stderr", "output"):
-                    content = event.get("content", "")
-                    if content:
-                        output_buffer.append(content)
-
-                # Track status
+                # Track status transitions
                 if "status" in event:
                     new_status = event["status"]
                     if new_status != final_status:
                         final_status = new_status
-                        if final_status == "provisioning":
-                            print(f"\n[Podstack] Provisioning GPU runner...")
-                        elif final_status == "running":
-                            print(f"\n[Podstack] Running on GPU...")
+                        extra = ""
+                        if new_status == "queued":
+                            pos = event.get("queue_position", "?")
+                            extra = f"(position: {pos})"
+                        display.set_phase(new_status, extra)
 
                     # Check for terminal status
                     if final_status in ("completed", "failed", "timeout", "cancelled"):
                         final_event = event
                         break
 
+                # Stream output lines
+                if event.get("type") in ("stdout", "stderr", "output"):
+                    content = event.get("content", "")
+                    if content:
+                        output_buffer.append(content)
+                        display.log(content, event.get("type", "stdout"))
+
         except RuntimeError as e:
             if "HTTP 401" in str(e):
                 # Auth failed on stream — fall back to polling
-                print(f"\n[Podstack] Streaming auth failed, falling back to polling...")
-                return self._run_with_polling(execution_id, gpu, count, timeout, 2.0, max_retries, 300, cancel_on_timeout)
+                display._stop_spinner()
+                display._print(f"\r[Podstack] Streaming auth failed, falling back to polling...{' ' * 10}\n")
+                return self._run_with_polling(execution_id, gpu, count, timeout, 2.0, max_retries, 300, cancel_on_timeout, fraction)
             raise
         except (ConnectionError, httpx.ConnectError) as e:
             # Try to recover and get the result
             logger.warning(f"Stream connection lost: {e}")
-            print(f"\n[Podstack] Stream connection lost, fetching final result...")
+            display._stop_spinner()
+            display._print(f"\r[Podstack] Stream connection lost, fetching final result...{' ' * 10}\n")
 
         # Get final result
         result = None
@@ -828,11 +935,7 @@ _stream_install(
         if "__PODSTACK_RESULT__" not in result.output and "__PODSTACK_RESULT__" in streamed:
             result.output = streamed
 
-        if result.success:
-            print(f"\n[Podstack] Completed in {result.gpu_seconds:.1f}s (cost: ₹{result.cost_paise/100:.2f})")
-        else:
-            error_msg = result.error or 'Unknown error'
-            print(f"\n[Podstack] Failed: {error_msg}")
+        display.complete(result.success, result.gpu_seconds, result.cost_paise, result.error or "Unknown error")
 
         return result
 
@@ -845,10 +948,13 @@ _stream_install(
         poll_interval: float,
         max_retries: int,
         provisioning_timeout: int,
-        cancel_on_timeout: bool
+        cancel_on_timeout: bool,
+        fraction: int = 100
     ) -> GPUExecutionResult:
         """Run execution with polling (non-streaming mode)."""
-        print(f"[Podstack] Waiting for GPU runner ({gpu} x{count})...")
+        display = LiveDisplay(gpu, count, fraction)
+        display.set_phase("pending")
+
         start_time = time.time()
         provisioning_start = None
         last_status = ""
@@ -863,7 +969,8 @@ _stream_install(
                 if cancel_on_timeout:
                     try:
                         self.cancel(execution_id)
-                        print(f"[Podstack] Execution cancelled due to timeout")
+                        display._stop_spinner()
+                        display._print(f"\r[Podstack] Execution cancelled due to timeout{' ' * 30}\n")
                     except Exception as e:
                         logger.warning(f"Failed to cancel execution: {e}")
 
@@ -905,24 +1012,17 @@ _stream_install(
 
             if current_status != last_status:
                 last_status = current_status
-                if current_status == "pending":
-                    print(f"[Podstack] Pending...")
-                elif current_status == "queued":
-                    pos = status_data.get("queue_position", "?")
-                    print(f"[Podstack] Queued (position: {pos})")
-                elif current_status == "provisioning":
+                if current_status == "provisioning":
                     provisioning_start = time.time()
-                    print(f"[Podstack] Provisioning GPU runner...")
-                elif current_status == "running":
-                    print(f"[Podstack] Running on GPU...")
-                elif current_status == "streaming":
-                    print(f"[Podstack] Streaming output...")
-                elif current_status in ("completed", "failed", "timeout", "cancelled"):
-                    pass  # Terminal states handled below
-                else:
+                if current_status in ("pending", "queued", "provisioning", "running"):
+                    extra = ""
+                    if current_status == "queued":
+                        pos = status_data.get("queue_position", "?")
+                        extra = f"(position: {pos})"
+                    display.set_phase(current_status, extra)
+                elif current_status not in ("completed", "failed", "timeout", "cancelled"):
                     # Unknown status - log but continue
                     logger.warning(f"Unknown status: {current_status}")
-                    print(f"[Podstack] Status: {current_status}")
 
             # Check for provisioning timeout
             if provisioning_start and current_status == "provisioning":
@@ -965,14 +1065,12 @@ _stream_install(
                     raise ConnectionError(f"Failed to get result after {max_retries} attempts: {e}")
                 time.sleep(poll_interval * (attempt + 1))
 
-        if result.success:
-            print(f"[Podstack] Completed in {result.gpu_seconds:.1f}s (cost: ₹{result.cost_paise/100:.2f})")
-        else:
-            error_msg = result.error or 'Unknown error'
-            print(f"[Podstack] Failed: {error_msg}")
-            # Include partial output in the error for debugging
-            if result.output:
-                print(f"[Podstack] Output (last 500 chars):\n{result.output[-500:]}")
+        display.complete(result.success, result.gpu_seconds, result.cost_paise, result.error or "Unknown error")
+
+        # In polling mode, output isn't streamed — print it after completion
+        if result.output and not result.success:
+            sys.stdout.write(f"[Podstack] Output (last 500 chars):\n{result.output[-500:]}\n")
+            sys.stdout.flush()
 
         return result
 

{podstack-1.3.17 → podstack-1.3.20}/podstack/registry/__init__.py

@@ -176,7 +176,13 @@ def list_experiments(limit: int = 20, offset: int = 0) -> list:
     return _get_client().list_experiments(limit, offset)
 
 
-def start_run(name: str = None, tags: dict = None) -> Run:
+def start_run(
+    name: str = None,
+    tags: dict = None,
+    capture_env: bool = True,
+    system_metrics: bool = True,
+    system_metrics_interval: float = 10.0,
+) -> Run:
     """
     Start a new run in the active experiment.
 
@@ -188,11 +194,14 @@ def start_run(name: str = None, tags: dict = None) -> Run:
     Args:
         name: Optional run name
         tags: Optional tags dict
+        capture_env: Auto-capture Python/pip/git/CUDA as _env.* params
+        system_metrics: Log CPU/RAM/GPU metrics every system_metrics_interval seconds
+        system_metrics_interval: Seconds between metric samples (default 10)
 
     Returns:
         Run object (context manager)
     """
-    return _get_client().start_run(name, tags)
+    return _get_client().start_run(name, tags, capture_env, system_metrics, system_metrics_interval)
 
 
 def end_run(status: str = "completed"):
@@ -558,9 +567,23 @@ def create_trial(sweep_id: str, run_id: str, params: dict) -> dict:
     return _get_client().create_trial(sweep_id, run_id, params)
 
 
-def complete_trial(sweep_id: str, trial_id: str, value: float) -> None:
-    """Mark a trial as completed with its objective metric value."""
-    _get_client().complete_trial(sweep_id, trial_id, value)
+def complete_trial(
+    sweep_id: str,
+    trial_id: str,
+    value: float = None,
+    metrics: dict = None,
+    run_id: str = None,
+) -> None:
+    """Mark a trial as completed with its objective metric value.
+
+    Args:
+        sweep_id: Sweep ID.
+        trial_id: Trial ID (from suggest_trial_params).
+        value: Scalar objective value.
+        metrics: Dict of metric name → value; first entry used if value not set.
+        run_id: Optional run ID to link to this trial.
+    """
+    _get_client().complete_trial(sweep_id, trial_id, value, metrics, run_id)
 
 
 def list_trials(sweep_id: str) -> list:

{podstack-1.3.17 → podstack-1.3.20}/podstack/registry/client.py

@@ -1372,9 +1372,12 @@ class RegistryClient:
         return data.get("sweeps", [])
 
     def suggest_trial_params(self, sweep_id: str) -> Dict[str, Any]:
-        """Get suggested hyperparameter values for the next trial."""
+        """Get suggested hyperparameter values for the next trial.
+
+        Returns a dict with ``id`` (trial ID) and ``params`` (hyperparameter map).
+        """
         data = self._request("GET", f"/sweeps/{sweep_id}/suggest")
-        return data.get("params", data)
+        return data.get("trial", data)
 
     def create_trial(
         self,
@@ -1389,11 +1392,31 @@ class RegistryClient:
         })
         return data.get("trial", data)
 
-    def complete_trial(self, sweep_id: str, trial_id: str, value: float) -> None:
-        """Mark a trial as completed with its objective metric value."""
-        self._request("POST", f"/sweeps/{sweep_id}/trials/{trial_id}/complete", json={
-            "value": value
-        })
+    def complete_trial(
+        self,
+        sweep_id: str,
+        trial_id: str,
+        value: float = None,
+        metrics: Dict[str, float] = None,
+        run_id: str = None,
+    ) -> None:
+        """Mark a trial as completed with its objective metric value.
+
+        Args:
+            sweep_id: Sweep ID.
+            trial_id: Trial ID (from suggest_trial_params).
+            value: Scalar objective value.
+            metrics: Dict of metric name → value; first entry used if value not set.
+            run_id: Optional run ID to link to this trial.
+        """
+        body: Dict[str, Any] = {}
+        if value is not None:
+            body["value"] = value
+        if metrics:
+            body["metrics"] = metrics
+        if run_id:
+            body["run_id"] = run_id
+        self._request("POST", f"/sweeps/{sweep_id}/trials/{trial_id}/complete", json=body)
 
     def list_trials(self, sweep_id: str) -> List[Dict[str, Any]]:
         """List all trials for a sweep."""

{podstack-1.3.17 → podstack-1.3.20}/podstack.egg-info/PKG-INFO

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

{podstack-1.3.17 → podstack-1.3.20}/pyproject.toml

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