ionworks-api 0.1.0__py3-none-any.whl → 0.1.3__py3-none-any.whl

ionworks/pipeline.py

@@ -1,5 +1,14 @@
+"""
+Pipeline client for running parameterization workflows.
+
+This module provides the :class:`PipelineClient` for creating and managing
+pipelines that combine data fitting, calculations, and validation steps
+for battery model parameterization.
+"""
+
 import os
 import re
+import time
 from typing import Any
 
 from pydantic import (
@@ -20,6 +29,8 @@ def _prepare_payload(data: Any) -> Any:
 
 
 class DataFitConfig(BaseModel):
+    """Configuration for data fitting step in a pipeline."""
+
     objectives: dict[str, Any]
     parameters: dict[str, Any]
     cost: dict[str, Any] | None = None
@@ -28,14 +39,20 @@ class DataFitConfig(BaseModel):
 
 
 class EntryConfig(BaseModel):
+    """Configuration for entry point in a pipeline."""
+
     values: dict[str, Any]
 
 
 class BuiltInEntryConfig(BaseModel):
+    """Configuration for built-in entry point in a pipeline."""
+
     name: str
 
 
 class CalculationConfig(BaseModel):
+    """Configuration for calculation step in a pipeline."""
+
     calculation: str
     electrode: str | None = None
     method: str | None = None
@@ -43,12 +60,16 @@ class CalculationConfig(BaseModel):
 
 
 class ValidationConfig(BaseModel):
+    """Configuration for validation step in a pipeline."""
+
     objectives: dict[str, Any]
     summary_stats: list[Any]
     existing_parameters: dict[str, Any] | None = None
 
 
 class PipelineConfig(BaseModel):
+    """Configuration for a complete pipeline workflow."""
+
     project_id: str | None = Field(
         default=None,
         description="The project id to submit the pipeline to. "
@@ -108,23 +129,33 @@ class PipelineConfig(BaseModel):
 
 
 class DataFitResponse(BaseModel):
+    """Response from a data fitting step containing fitted parameters."""
+
     parameter_values: dict[str, Any]
 
 
 class CalculationResponse(BaseModel):
+    """Response from a calculation step containing calculated parameters."""
+
     parameter_values: dict[str, Any]
 
 
 class ValidationResponse(BaseModel):
+    """Response from a validation step containing validation results."""
+
     validation_results: dict[str, Any]
     summary_stats: dict[str, list[Any]]
 
 
 class EntryResponse(BaseModel):
+    """Response from an entry point containing parameter values."""
+
     parameter_values: dict[str, Any]
 
 
 class PipelineSubmissionResponse(BaseModel):
+    """Response from submitting a pipeline to the API."""
+
     id: str
     name: str
     description: str | None = None
@@ -133,12 +164,23 @@ class PipelineSubmissionResponse(BaseModel):
 
 
 class PipelineResponse(BaseModel):
+    """Complete response from retrieving pipeline results."""
+
     result: dict[str, Any]
     element_results: dict[str, Any]
 
 
 class PipelineClient:
-    def __init__(self, client: Any):
+    """Client for creating and managing pipeline workflows."""
+
+    def __init__(self, client: Any) -> None:
+        """Initialize the pipeline client.
+
+        Parameters
+        ----------
+        client : Any
+            The HTTP client to use for API requests.
+        """
         self.client = client
 
     def create(self, config: dict[str, Any]) -> PipelineSubmissionResponse:
@@ -189,19 +231,29 @@ class PipelineClient:
             # Re-raise original error for other cases
             raise
 
-    def list(self, project_id: str | None = None) -> list[PipelineSubmissionResponse]:
+    def list(
+        self, project_id: str | None = None, limit: int | None = None
+    ) -> list[PipelineSubmissionResponse]:
         """List all pipelines.
 
         Parameters
         ----------
-        project_id : str, optional
-            The project id to filter pipelines. If not provided, will use
+        project_id : str | None
+            The project id to filter pipelines. If not provided, uses
             PROJECT_ID environment variable.
+        limit : int | None
+            Maximum number of pipelines to return. If not provided, returns
+            all pipelines (up to the API's default limit).
 
         Returns
         -------
         list[PipelineSubmissionResponse]
             List of pipeline submission responses.
+
+        Raises
+        ------
+        ValueError
+            If response data is not a list or project_id is missing.
         """
         if project_id is None:
             project_id = os.getenv("PROJECT_ID")
@@ -212,6 +264,8 @@ class PipelineClient:
                 )
 
         endpoint = f"/pipelines?project_id={project_id}"
+        if limit is not None:
+            endpoint += f"&limit={limit}"
         try:
             response_data = self.client.get(endpoint)
             if not isinstance(response_data, list):
@@ -254,3 +308,58 @@ class PipelineClient:
         """
         response_data = self.client.get(f"/pipelines/{job_id}/result")
         return PipelineResponse(**response_data)
+
+    def wait_for_completion(
+        self,
+        pipeline_id: str,
+        timeout: int = 600,
+        poll_interval: int = 2,
+        verbose: bool = True,
+    ) -> PipelineSubmissionResponse:
+        """Wait for a pipeline to complete by polling until done or timeout.
+
+        Parameters
+        ----------
+        pipeline_id : str
+            The pipeline ID to wait for.
+        timeout : int, optional
+            Maximum time to wait in seconds (default: 600).
+        poll_interval : int, optional
+            Time between polls in seconds (default: 2).
+        verbose : bool, optional
+            Whether to print status updates (default: True).
+
+        Returns
+        -------
+        PipelineSubmissionResponse
+            The completed or failed pipeline response.
+
+        Raises
+        ------
+        TimeoutError
+            If timeout is reached before the pipeline completes.
+        """
+        deadline = time.time() + timeout
+        pipeline = self.get(pipeline_id)
+
+        if verbose:
+            print(f"Polling pipeline {pipeline_id} for completion...")
+
+        while pipeline.status not in ("completed", "failed"):
+            if time.time() >= deadline:
+                raise TimeoutError(
+                    f"Pipeline {pipeline_id} did not complete within "
+                    f"{timeout} seconds (status: {pipeline.status})"
+                )
+            time.sleep(poll_interval)
+            pipeline = self.get(pipeline_id)
+            if verbose:
+                elapsed = int(timeout - (deadline - time.time()))
+                print(f"  Status: {pipeline.status} (elapsed: {elapsed}s)")
+
+        if verbose:
+            print(f"Pipeline finished with status: {pipeline.status}")
+            if pipeline.status == "failed" and pipeline.error:
+                print(f"  Error: {pipeline.error}")
+
+        return pipeline

ionworks/simulation.py

@@ -1,10 +1,22 @@
+"""
+Simulation client for running battery simulations.
+
+This module provides the :class:`SimulationClient` for running battery
+simulations using the Universal Cycler Protocol (UCP) format. It supports
+single simulations, batch simulations with design of experiments (DOE),
+and PyBaMM-based modeling.
+"""
+
 from __future__ import annotations
 
 from datetime import datetime, timedelta
 import time
-from typing import Any, List, cast
+from typing import Any, cast
 
 from pydantic import BaseModel, Field, ValidationError
+import requests
+
+from .errors import IonworksError
 
 
 class QuickModelConfig(BaseModel):
@@ -129,7 +141,14 @@ class SimulationResponse(BaseModel):
 class SimulationClient:
     """Client for running simulations."""
 
-    def __init__(self, client: Any):
+    def __init__(self, client: Any) -> None:
+        """Initialize the SimulationClient.
+
+        Parameters
+        ----------
+        client : Any
+            The HTTP client instance for making API requests.
+        """
         self.client = client
 
     def protocol(self, config: dict[str, Any]) -> SimulationResponse:
@@ -139,16 +158,18 @@ class SimulationClient:
         ----------
         config : dict[str, Any]
             Configuration dictionary containing:
-            - parameterized_model: quick_model dict, full model dict, or model ID string
-            - protocol_experiment: ProtocolExperimentConfig dict with protocol
-              and name fields
-            - experiment_parameters: Optional ProtocolExperimentParameters dict
-              with initial_soc and initial_temperature
+
+            - parameterized_model: quick_model dict, full model dict,
+              or model ID string
+            - protocol_experiment: ProtocolExperimentConfig dict with
+              protocol and name fields
+            - experiment_parameters: Optional dict with initial_soc and
+              initial_temperature
             - design_parameters: Optional dict[str, float]
             - max_backward_jumps: Optional int
             - study_id: Optional str
-            - extra_variables: Optional list[str] - Extra variables to include
-              in simulation output
+            - extra_variables: Optional list[str] — extra variables to
+              include in simulation output
 
         Returns
         -------
@@ -177,15 +198,17 @@ class SimulationClient:
         ----------
         config : dict[str, Any]
             Configuration dictionary containing:
-            - parameterized_model: quick_model dict, full model dict, or model ID string
-            - protocol_experiment: ProtocolExperimentConfig dict with protocol
-              and name fields
+
+            - parameterized_model: quick_model dict, full model dict,
+              or model ID string
+            - protocol_experiment: ProtocolExperimentConfig dict with
+              protocol and name fields
             - design_parameters_doe: DesignParametersDOE dict
-            - experiment_parameters: Optional ProtocolExperimentParameters dict
+            - experiment_parameters: Optional dict
             - max_backward_jumps: Optional int
             - study_id: Optional str
-            - extra_variables: Optional list[str] - Extra variables to include
-              in simulation output
+            - extra_variables: Optional list[str] — extra variables to
+              include in simulation output
 
         Returns
         -------
@@ -204,10 +227,11 @@ class SimulationClient:
                 endpoint, json_payload=validated_config.model_dump(exclude_none=True)
             )
             if not isinstance(response_data, list):
-                raise ValueError(
+                msg = (
                     f"Unexpected response format from {endpoint}: expected a "
                     f"list, got {type(response_data).__name__}"
                 )
+                raise ValueError(msg)
             return [SimulationResponse(**item) for item in response_data]
         except ValidationError as e:
             raise ValueError(
@@ -225,10 +249,11 @@ class SimulationClient:
         endpoint = "/simulations"
         response_data = self.client.get(endpoint)
         if not isinstance(response_data, list):
-            raise ValueError(
+            msg = (
                 f"Unexpected response format from {endpoint}: expected a list, "
                 f"got {type(response_data).__name__}"
             )
+            raise ValueError(msg)
         return response_data
 
     def get(self, simulation_id: str) -> dict[str, Any]:
@@ -273,99 +298,127 @@ class SimulationClient:
         response_data = self.client.get(endpoint)
         return cast(dict[str, Any], response_data)
 
-    def wait_for_completion(
+    def _poll_simulations(
         self,
-        simulation_id: str | List[str],  # List to avoid conflict with list() method
-        timeout: int = 60,
-        poll_interval: int = 2,
-        verbose: bool = True,
-    ) -> dict[str, Any] | List[dict[str, Any]]:
-        """Wait for simulation(s) to complete by polling until done or timeout.
+        simulation_ids: list[str],
+        timeout: int,
+        poll_interval: int,
+        verbose: bool,
+    ) -> dict[str, dict[str, Any]]:
+        """Poll simulations until all complete or timeout is reached.
 
         Parameters
         ----------
-        simulation_id : str | list[str]
-            Single simulation ID or list of simulation IDs to wait for.
-        timeout : int, optional
-            Maximum time to wait in seconds (default: 60).
-        poll_interval : int, optional
-            Time between polls in seconds (default: 2).
-        verbose : bool, optional
-            Whether to print status updates (default: True).
+        simulation_ids : list[str]
+            List of simulation IDs to poll.
+        timeout : int
+            Maximum time to wait in seconds.
+        poll_interval : int
+            Time between polls in seconds.
+        verbose : bool
+            Whether to print status updates.
 
         Returns
         -------
-        dict[str, Any] | list[dict[str, Any]]
-            Completed simulation(s). Returns single dict if single ID provided,
-            list of dicts if list of IDs provided. Only returns completed
-            simulations if timeout is reached.
+        dict[str, dict[str, Any]]
+            Dict mapping simulation IDs to their completed result dicts.
 
         Raises
         ------
         TimeoutError
-            If timeout is reached before all simulations complete.
+            If no simulations complete within the timeout.
         """
-        is_single = isinstance(simulation_id, str)
-        if is_single:
-            simulation_ids: List[str] = [simulation_id]  # type: ignore[list-item]
-        else:
-            simulation_ids = simulation_id  # type: ignore[assignment]
         timeout_delta = timedelta(seconds=timeout)
         start_time = datetime.now()
-        completed_simulations: dict[str, dict[str, Any]] = {}
+        completed: dict[str, dict[str, Any]] = {}
 
         if verbose:
             print(f"Polling for {len(simulation_ids)} simulation(s) completion...")
 
         while datetime.now() - start_time < timeout_delta:
-            completed_count = len(completed_simulations)
             for sim_id in simulation_ids:
-                if sim_id in completed_simulations:
+                if sim_id in completed:
                     continue
                 try:
                     simulation = self.get(sim_id)
-                    if simulation.get("simulation_data") is not None:
-                        completed_simulations[sim_id] = simulation
-                        completed_count += 1
-                except Exception:
-                    # Continue polling if there's an error
-                    pass
+                    if (
+                        simulation.get("storage_folder")
+                        or simulation.get("simulation_data")  # Legacy fallback
+                    ):
+                        completed[sim_id] = simulation
+                except (IonworksError, requests.exceptions.RequestException):
+                    pass  # Continue polling on transient errors
 
             elapsed = (datetime.now() - start_time).seconds
             if verbose:
                 print(
-                    f"  Status: {completed_count}/{len(simulation_ids)} completed "
+                    f"  Status: {len(completed)}/{len(simulation_ids)} completed "
                     f"(elapsed: {elapsed}s)"
                 )
 
-            if len(completed_simulations) == len(simulation_ids):
+            if len(completed) == len(simulation_ids):
                 if verbose:
                     print("All simulations completed!")
-                break
+                return completed
 
             time.sleep(poll_interval)
-        else:
-            # Timeout reached
-            if verbose:
-                print(
-                    f"Timeout: Only {len(completed_simulations)}/{len(simulation_ids)} "
-                    f"simulations completed within {timeout} seconds"
-                )
-            if not completed_simulations:
-                raise TimeoutError(f"No simulations completed within {timeout} seconds")
 
-        # Return results in the same format as input
+        # Timeout reached
+        if verbose:
+            print(
+                f"Timeout: Only {len(completed)}/{len(simulation_ids)} "
+                f"simulations completed within {timeout} seconds"
+            )
+        if not completed:
+            msg = f"No simulations completed within {timeout} seconds"
+            raise TimeoutError(msg)
+        return completed
+
+    def wait_for_completion(
+        self,
+        simulation_id: str | list[str],
+        timeout: int = 60,
+        poll_interval: int = 2,
+        verbose: bool = True,
+    ) -> dict[str, Any] | list[dict[str, Any]]:
+        """Wait for simulation(s) to complete by polling until done or timeout.
+
+        Parameters
+        ----------
+        simulation_id : str | list[str]
+            Single simulation ID or list of simulation IDs to wait for.
+        timeout : int
+            Maximum time to wait in seconds (default: 60).
+        poll_interval : int
+            Time between polls in seconds (default: 2).
+        verbose : bool
+            Whether to print status updates (default: True).
+
+        Returns
+        -------
+        dict[str, Any] | list[dict[str, Any]]
+            Completed simulation(s). Returns single dict if single ID
+            provided, list of dicts if list of IDs provided. Only returns
+            completed simulations if timeout is reached.
+
+        Raises
+        ------
+        TimeoutError
+            If timeout is reached before all simulations complete.
+        """
+        is_single = isinstance(simulation_id, str)
+        simulation_ids = [simulation_id] if is_single else simulation_id  # type: ignore[list-item]
+
+        completed = self._poll_simulations(
+            simulation_ids, timeout, poll_interval, verbose
+        )
+
         if is_single:
-            if simulation_ids[0] not in completed_simulations:
-                raise TimeoutError(
+            if simulation_ids[0] not in completed:
+                msg = (
                     f"Simulation {simulation_ids[0]} did not complete within "
                     f"{timeout} seconds"
                 )
-            return completed_simulations[simulation_ids[0]]
-        else:
-            # Return list of completed simulations in order
-            return [
-                completed_simulations[sim_id]
-                for sim_id in simulation_ids
-                if sim_id in completed_simulations
-            ]
+                raise TimeoutError(msg)
+            return completed[simulation_ids[0]]
+        return [completed[sim_id] for sim_id in simulation_ids if sim_id in completed]