luminarycloud 0.19.0__py3-none-any.whl → 0.22.0__py3-none-any.whl

luminarycloud/pipelines/api.py

@@ -1,15 +1,39 @@
 # Copyright 2023-2024 Luminary Cloud, Inc. All Rights Reserved.
-from typing import Literal
+from typing import Any, Literal
 from dataclasses import dataclass
 
 from datetime import datetime
+from time import time, sleep
+import logging
 
+from .arguments import PipelineArgValueType
 from ..pipelines import Pipeline, PipelineArgs
 from .._client import get_default_client
 
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class LogLine:
+    timestamp: datetime
+    level: int
+    message: str
+
+    @classmethod
+    def from_json(cls, json: dict) -> "LogLine":
+        return cls(
+            timestamp=datetime.fromisoformat(json["timestamp"]),
+            level=json["level"],
+            message=json["message"],
+        )
+
 
 @dataclass
 class PipelineRecord:
+    """
+    A PipelineRecord represents a persisted pipeline.
+    """
+
     id: str
     name: str
     description: str | None
@@ -17,8 +41,10 @@ class PipelineRecord:
     created_at: datetime
     updated_at: datetime
 
-    def pipeline(self) -> Pipeline:
-        return Pipeline._from_yaml(self.definition_yaml)
+    # I don't think users need to get the Pipeline object from a PipelineRecord, but if they did,
+    # it would be done like this.
+    # def pipeline(self) -> Pipeline:
+    #     return Pipeline._from_yaml(self.definition_yaml)
 
     @classmethod
     def from_json(cls, json: dict) -> "PipelineRecord":
@@ -31,15 +57,30 @@ class PipelineRecord:
             updated_at=datetime.fromisoformat(json["updated_at"]),
         )
 
+    def pipeline_jobs(self) -> "list[PipelineJobRecord]":
+        """
+        Returns a list of pipeline jobs that were created from this pipeline.
+
+        Returns
+        -------
+        list[PipelineJobRecord]
+            A list of PipelineJobRecord objects.
+        """
+        res = get_default_client().http.get(f"/rest/v0/pipelines/{self.id}/pipeline_jobs")
+        return [PipelineJobRecord.from_json(p) for p in res["data"]]
+
 
 @dataclass
 class PipelineJobRecord:
+    """
+    A PipelineJobRecord represents a persisted pipeline job.
+    """
+
     id: str
     pipeline_id: str
-    project_id: str
     name: str
     description: str | None
-    status: Literal["pending", "running", "completed", "failed", "cancelled"]
+    status: Literal["pending", "running", "completed", "failed"]
     created_at: datetime
     updated_at: datetime
     started_at: datetime | None
@@ -50,7 +91,6 @@ class PipelineJobRecord:
         return cls(
             id=json["id"],
             pipeline_id=json["pipeline_id"],
-            project_id=json["project_id"],
             name=json["name"],
             description=json["description"],
             status=json["status"],
@@ -62,6 +102,200 @@ class PipelineJobRecord:
             ),
         )
 
+    def pipeline(self) -> PipelineRecord:
+        """
+        Returns the pipeline that this pipeline job was created from.
+
+        Returns
+        -------
+        PipelineRecord
+            The PipelineRecord for the pipeline that this pipeline job was created from.
+        """
+        return get_pipeline(self.pipeline_id)
+
+    def runs(self) -> "list[PipelineJobRunRecord]":
+        """
+        Returns a list of runs for this pipeline job.
+
+        Returns
+        -------
+        list[PipelineJobRunRecord]
+            A list of PipelineJobRunRecord objects.
+        """
+        res = get_default_client().http.get(f"/rest/v0/pipeline_jobs/{self.id}/runs")
+        return [PipelineJobRunRecord.from_json(r) for r in res["data"]]
+
+    def logs(self) -> list[LogLine]:
+        """
+        Returns a list of log lines for this pipeline job.
+
+        Each log line is a LogLine object, which has a timestamp, level, and message.
+
+        Returns
+        -------
+        list[LogLine]
+            A list of LogLine objects.
+        """
+        res = get_default_client().http.get(f"/rest/v0/pipeline_jobs/{self.id}/logs")
+        return [LogLine.from_json(l) for l in res["data"]]
+
+    def artifacts(self) -> list[dict]:
+        """
+        Returns a list of artifacts that were produced by this pipeline job.
+
+        Artifacts are things like Geometries, Meshes, and Simulations. Each artifact is a dictionary
+        with an "id" key, which is an identifier for the artifact.
+
+        .. warning:: This feature is experimental and may change or be removed in the future.
+
+        Returns
+        -------
+        list[dict]
+            A list of artifact dictionaries.
+        """
+        res = get_default_client().http.get(f"/rest/v0/pipeline_jobs/{self.id}/artifacts")
+        return res["data"]
+
+    def wait(
+        self,
+        *,
+        interval_seconds: float = 5,
+        timeout_seconds: float = float("inf"),
+        print_logs: bool = False,
+    ) -> Literal["completed", "failed"]:
+        """
+        Wait for the pipeline job to complete or fail.
+
+        This method polls the pipeline job status at regular intervals until it reaches
+        a terminal state (completed or failed).
+
+        Parameters
+        ----------
+        interval_seconds : float
+            Number of seconds between status polls. Default is 5 seconds.
+        timeout_seconds : float
+            Number of seconds before the operation times out. Default is infinity.
+        print_logs : bool
+            If True, prints new log lines as they become available. Default is False.
+
+        Returns
+        -------
+        Literal["completed", "failed"]
+            The final status of the pipeline job.
+
+        Raises
+        ------
+        TimeoutError
+            If the pipeline job does not complete within the specified timeout.
+
+        Examples
+        --------
+        >>> pipeline_job = pipelines.create_pipeline_job(pipeline.id, args, "My Job")
+        >>> final_status = pipeline_job.wait(timeout_seconds=3600)
+        >>> print(f"Pipeline job finished with status: {final_status}")
+        """
+        deadline = time() + timeout_seconds
+        last_log_count = 0
+
+        while True:
+            # Refresh the pipeline job status
+            updated_job = get_pipeline_job(self.id)
+
+            # Print new logs if requested
+            if print_logs:
+                logs = updated_job.logs()
+                if len(logs) > last_log_count:
+                    for log_line in logs[last_log_count:]:
+                        print(f"[{log_line.timestamp}] {log_line.message}")
+                    last_log_count = len(logs)
+
+            # Check if we've reached a terminal state
+            if updated_job.status == "completed":
+                logger.info(f"Pipeline job {self.id} completed successfully")
+                return "completed"
+            elif updated_job.status == "failed":
+                logger.warning(f"Pipeline job {self.id} failed")
+                return "failed"
+
+            # Check timeout
+            if time() >= deadline:
+                raise TimeoutError(
+                    f"Timed out waiting for pipeline job {self.id} to complete. "
+                    f"Current status: {updated_job.status}"
+                )
+
+            # Wait before next poll
+            sleep(max(0, min(interval_seconds, deadline - time())))
+
+            # Update self with the latest status
+            self.status = updated_job.status
+            self.updated_at = updated_job.updated_at
+            self.started_at = updated_job.started_at
+            self.completed_at = updated_job.completed_at
+
+
+@dataclass
+class PipelineJobRunRecord:
+    pipeline_job_id: str
+    idx: int
+    arguments: list[PipelineArgValueType]
+    status: Literal["pending", "running", "completed", "failed"]
+
+    @classmethod
+    def from_json(cls, json: dict) -> "PipelineJobRunRecord":
+        return cls(
+            pipeline_job_id=json["pipeline_job_id"],
+            idx=json["idx"],
+            arguments=json["arguments"],
+            status=json["status"],
+        )
+
+    def pipeline_job(self) -> PipelineJobRecord:
+        """
+        Returns the pipeline job that this pipeline job run was created from.
+
+        Returns
+        -------
+        PipelineJobRecord
+            The PipelineJobRecord for the pipeline job that this pipeline job run was created from.
+        """
+        return get_pipeline_job(self.pipeline_job_id)
+
+    def logs(self) -> list[LogLine]:
+        """
+        Returns a list of log lines for this pipeline job run.
+
+        Each log line is a LogLine object, which has a timestamp, level, and message.
+
+        Returns
+        -------
+        list[LogLine]
+            A list of LogLine objects.
+        """
+        res = get_default_client().http.get(
+            f"/rest/v0/pipeline_jobs/{self.pipeline_job_id}/runs/{self.idx}/logs"
+        )
+        return [LogLine.from_json(l) for l in res["data"]]
+
+    def artifacts(self) -> list[dict]:
+        """
+        Returns a list of artifacts that were produced by this pipeline job run.
+
+        Artifacts are things like Geometries, Meshes, and Simulations. Each artifact is a dictionary
+        with an "id" key, which is an identifier for the artifact.
+
+        .. warning:: This feature is experimental and may change or be removed in the future.
+
+        Returns
+        -------
+        list[dict]
+            A list of artifact dictionaries.
+        """
+        res = get_default_client().http.get(
+            f"/rest/v0/pipeline_jobs/{self.pipeline_job_id}/runs/{self.idx}/artifacts"
+        )
+        return res["data"]
+
 
 def create_pipeline(
     name: str, pipeline: Pipeline | str, description: str | None = None
@@ -88,7 +322,7 @@ def create_pipeline(
         "description": description,
     }
     res = get_default_client().http.post("/rest/v0/pipelines", body)
-    return PipelineRecord.from_json(res)
+    return PipelineRecord.from_json(res["data"])
 
 
 def list_pipelines() -> list[PipelineRecord]:
@@ -96,7 +330,7 @@ def list_pipelines() -> list[PipelineRecord]:
     List all pipelines.
     """
     res = get_default_client().http.get("/rest/v0/pipelines")
-    return [PipelineRecord.from_json(p) for p in res]
+    return [PipelineRecord.from_json(p) for p in res["data"]]
 
 
 def get_pipeline(id: str) -> PipelineRecord:
@@ -109,11 +343,11 @@ def get_pipeline(id: str) -> PipelineRecord:
         ID of the pipeline to fetch.
     """
     res = get_default_client().http.get(f"/rest/v0/pipelines/{id}")
-    return PipelineRecord.from_json(res)
+    return PipelineRecord.from_json(res["data"])
 
 
 def create_pipeline_job(
-    pipeline_id: str, args: PipelineArgs, project_id: str, name: str, description: str | None = None
+    pipeline_id: str, args: PipelineArgs, name: str, description: str | None = None
 ) -> PipelineJobRecord:
     """
     Create a new pipeline job.
@@ -124,8 +358,6 @@ def create_pipeline_job(
         ID of the pipeline to invoke.
     args : PipelineArgs
         Arguments to pass to the pipeline.
-    project_id : str
-        ID of the project to run the pipeline job in.
     name : str
         Name of the pipeline job.
     description : str, optional
@@ -135,13 +367,13 @@ def create_pipeline_job(
     arg_rows = [row.row_values for row in args.rows]
     body = {
         "name": name,
-        "project_id": project_id,
+        "description": description,
         "argument_names": [p.name for p in args.params],
         "argument_rows": arg_rows,
     }
 
     res = get_default_client().http.post(f"/rest/v0/pipelines/{pipeline_id}/pipeline_jobs", body)
-    return PipelineJobRecord.from_json(res)
+    return PipelineJobRecord.from_json(res["data"])
 
 
 def get_pipeline_job(id: str) -> PipelineJobRecord:
@@ -149,7 +381,7 @@ def get_pipeline_job(id: str) -> PipelineJobRecord:
     Get a pipeline job by ID.
     """
     res = get_default_client().http.get(f"/rest/v0/pipeline_jobs/{id}")
-    return PipelineJobRecord.from_json(res)
+    return PipelineJobRecord.from_json(res["data"])
 
 
 def list_pipeline_jobs() -> list[PipelineJobRecord]:
@@ -157,4 +389,4 @@ def list_pipeline_jobs() -> list[PipelineJobRecord]:
     List all pipeline jobs.
     """
     res = get_default_client().http.get("/rest/v0/pipeline_jobs")
-    return [PipelineJobRecord.from_json(p) for p in res]
+    return [PipelineJobRecord.from_json(p) for p in res["data"]]

luminarycloud/pipelines/arguments.py

@@ -103,3 +103,18 @@ class PipelineArgs:
         return (
             f"PipelineArgs(param_names={[p.name for p in self.params]}, row_count={len(self.rows)})"
         )
+
+    def print_as_table(self) -> None:
+        headers = [p.name for p in self.params]
+        row_strs = [[str(v) for v in row.row_values] for row in self.rows]
+        col_widths = [
+            max(len(headers[i]), *(len(r[i]) for r in row_strs)) for i in range(len(headers))
+        ]
+
+        def format_row(values: list[str]) -> str:
+            return " | ".join(val.ljust(col_widths[i]) for i, val in enumerate(values))
+
+        print(format_row(headers))
+        print("-+-".join("-" * w for w in col_widths))
+        for r in row_strs:
+            print(format_row(r))