luminarycloud 0.22.1__py3-none-any.whl → 0.22.3__py3-none-any.whl

luminarycloud/pipelines/__init__.py

@@ -2,8 +2,12 @@
 from .core import (
     Pipeline as Pipeline,
     PipelineParameter as PipelineParameter,
+    stage as stage,
     # Stage base class, mainly exported for testing
     Stage as Stage,
+    # RunScript stage lives in core because it's a special snowflake
+    RunScript as RunScript,
+    StopRun as StopRun,
 )
 
 from .parameters import (
@@ -36,6 +40,11 @@ from .arguments import (
     ArgNamedVariableSet as ArgNamedVariableSet,
 )
 
+from .flowables import (
+    FlowableType as FlowableType,
+    FlowableIOSchema as FlowableIOSchema,
+)
+
 from .api import (
     create_pipeline as create_pipeline,
     list_pipelines as list_pipelines,
@@ -46,5 +55,7 @@ from .api import (
     PipelineJobRecord as PipelineJobRecord,
     PipelineRecord as PipelineRecord,
     PipelineJobRunRecord as PipelineJobRunRecord,
+    PipelineTaskRecord as PipelineTaskRecord,
+    StageDefinition as StageDefinition,
     LogLine as LogLine,
 )

luminarycloud/pipelines/api.py

@@ -10,6 +10,7 @@ from .arguments import PipelineArgValueType
 from .core import Stage
 from ..pipelines import Pipeline, PipelineArgs
 from .._client import get_default_client
+from .._helpers import parse_iso_datetime
 
 logger = logging.getLogger(__name__)
 
@@ -23,12 +24,56 @@ class LogLine:
     @classmethod
     def from_json(cls, json: dict) -> "LogLine":
         return cls(
-            timestamp=datetime.fromisoformat(json["timestamp"]),
+            timestamp=parse_iso_datetime(json["timestamp"]),
             level=json["level"],
             message=json["message"],
         )
 
 
+@dataclass
+class StageDefinition:
+    """
+    Represents a stage definition from a pipeline.
+    """
+
+    id: str
+    name: str
+    stage_type: str
+
+    @classmethod
+    def from_json(cls, json: dict) -> "StageDefinition":
+        return cls(
+            id=json["id"],
+            name=json["name"],
+            stage_type=json["stage_type"],
+        )
+
+
+@dataclass
+class PipelineTaskRecord:
+    """
+    A PipelineTaskRecord represents a task within a pipeline job run.
+    """
+
+    status: Literal["pending", "running", "completed", "failed", "upstream_failed", "cancelled"]
+    artifacts: dict[str, dict]
+    stage: StageDefinition | None
+    created_at: datetime
+    updated_at: datetime
+    error_messages: list[str] | None
+
+    @classmethod
+    def from_json(cls, json: dict) -> "PipelineTaskRecord":
+        return cls(
+            status=json["status"],
+            artifacts=json["artifacts"],
+            stage=StageDefinition.from_json(json["stage"]) if json.get("stage") else None,
+            created_at=parse_iso_datetime(json["created_at"]),
+            updated_at=parse_iso_datetime(json["updated_at"]),
+            error_messages=json.get("error_messages"),
+        )
+
+
 @dataclass
 class PipelineRecord:
     """
@@ -54,8 +99,8 @@ class PipelineRecord:
             name=json["name"],
             description=json["description"],
             definition_yaml=json["definition_yaml"],
-            created_at=datetime.fromisoformat(json["created_at"]),
-            updated_at=datetime.fromisoformat(json["updated_at"]),
+            created_at=parse_iso_datetime(json["created_at"]),
+            updated_at=parse_iso_datetime(json["updated_at"]),
         )
 
     def pipeline_jobs(self) -> "list[PipelineJobRecord]":
@@ -100,11 +145,12 @@ class PipelineJobRecord:
     pipeline_id: str
     name: str
     description: str | None
-    status: Literal["pending", "running", "completed", "failed", "cancelled"]
+    status: Literal["pending", "running", "completed", "failed", "cancelled", "paused"]
     created_at: datetime
     updated_at: datetime
     started_at: datetime | None
     completed_at: datetime | None
+    paused_at: datetime | None
 
     @classmethod
     def from_json(cls, json: dict) -> "PipelineJobRecord":
@@ -112,14 +158,15 @@ class PipelineJobRecord:
             id=json["id"],
             pipeline_id=json["pipeline_id"],
             name=json["name"],
-            description=json["description"],
+            description=json.get("description"),
             status=json["status"],
-            created_at=datetime.fromisoformat(json["created_at"]),
-            updated_at=datetime.fromisoformat(json["updated_at"]),
-            started_at=datetime.fromisoformat(json["started_at"]) if json["started_at"] else None,
+            created_at=parse_iso_datetime(json["created_at"]),
+            updated_at=parse_iso_datetime(json["updated_at"]),
+            started_at=(parse_iso_datetime(json["started_at"]) if json.get("started_at") else None),
             completed_at=(
-                datetime.fromisoformat(json["completed_at"]) if json["completed_at"] else None
+                parse_iso_datetime(json["completed_at"]) if json.get("completed_at") else None
             ),
+            paused_at=(parse_iso_datetime(json["paused_at"]) if json.get("paused_at") else None),
         )
 
     def pipeline(self) -> PipelineRecord:
@@ -315,6 +362,39 @@ class PipelineJobRecord:
         get_default_client().http.post(f"/rest/v0/pipeline_jobs/{self.id}/cancel", {})
         logger.info(f"Cancelled pipeline job {self.id}")
 
+    def pause(self) -> None:
+        """Pause this running pipeline job.
+
+        This will prevent new tasks from being scheduled while allowing
+        in-progress tasks to complete. The job status will be set to PAUSED
+        and all stage concurrency limits will be temporarily set to 0.
+
+        Call resume() to continue execution.
+
+        Raises
+        ------
+        HTTPError
+            If the pipeline job cannot be paused (e.g., not found or not in
+            RUNNING state).
+        """
+        get_default_client().http.post(f"/rest/v0/pipeline_jobs/{self.id}/pause", {})
+        logger.info(f"Paused pipeline job {self.id}")
+
+    def resume(self) -> None:
+        """Resume this paused pipeline job.
+
+        This will restore the job status to RUNNING and restore the original
+        concurrency limits, allowing new tasks to be scheduled again.
+
+        Raises
+        ------
+        HTTPError
+            If the pipeline job cannot be resumed (e.g., not found or not in
+            PAUSED state).
+        """
+        get_default_client().http.post(f"/rest/v0/pipeline_jobs/{self.id}/resume", {})
+        logger.info(f"Resumed pipeline job {self.id}")
+
 
 @dataclass
 class PipelineJobRunRecord:
@@ -378,6 +458,23 @@ class PipelineJobRunRecord:
         )
         return res["data"]
 
+    def tasks(self) -> list[PipelineTaskRecord]:
+        """
+        Returns a list of tasks for this pipeline job run.
+
+        Each task represents an execution of a stage of the pipeline, with its own
+        status, artifacts, and stage information.
+
+        Returns
+        -------
+        list[PipelineTaskRecord]
+            The tasks associated with this pipeline job run.
+        """
+        res = get_default_client().http.get(
+            f"/rest/v0/pipeline_jobs/{self.pipeline_job_id}/runs/{self.idx}/tasks"
+        )
+        return [PipelineTaskRecord.from_json(t) for t in res["data"]]
+
 
 def create_pipeline(
     name: str, pipeline: Pipeline | str, description: str | None = None

luminarycloud/pipelines/core.py

@@ -1,12 +1,26 @@
 # Copyright 2025 Luminary Cloud, Inc. All Rights Reserved.
+from __future__ import annotations
 from abc import ABC, abstractmethod
 from dataclasses import is_dataclass, fields
-from typing import Any, Type, TypeVar, Generic
+from typing import Any, Callable, Mapping, Type, TypeVar, Generic, TYPE_CHECKING
 from typing_extensions import Self
+import inspect
 import re
+import textwrap
 import yaml
 
 from ..pipeline_util.yaml import ensure_yamlizable
+from .flowables import (
+    PipelineOutput,
+    PipelineInput,
+    FlowableType,
+    flowable_class_to_name,
+    flowable_name_to_class,
+    FlowableIOSchema,
+)
+
+if TYPE_CHECKING:
+    from .arguments import PipelineArgValueType
 
 
 class PipelineParameterRegistry:
@@ -86,40 +100,12 @@ class PipelineParameter(ABC):
         return self.__hash__() == other.__hash__()
 
 
-class PipelineInput:
-    """
-    A named input for a Stage. Explicitly connected to a PipelineOutput.
-    """
-
-    def __init__(self, upstream_output: "PipelineOutput", owner: "Stage", name: str):
-        self.upstream_output = upstream_output
-        self.owner = owner
-        self.name = name
-
-    def _to_dict(self, id_for_stage: dict) -> dict:
-        if self.upstream_output.owner not in id_for_stage:
-            raise ValueError(
-                f"Stage {self.owner} depends on a stage, {self.upstream_output.owner}, that isn't in the Pipeline. Did you forget to add it?"
-            )
-        upstream_stage_id = id_for_stage[self.upstream_output.owner]
-        upstream_output_name = self.upstream_output.name
-        return {self.name: f"{upstream_stage_id}.{upstream_output_name}"}
-
-
-class PipelineOutput(ABC):
+class StopRun(RuntimeError):
     """
-    A named output for a Stage. Can be used to spawn any number of connected PipelineInputs.
+    Raised by RunScript code to indicate that the pipeline run should stop intentionally.
     """
 
-    def __init__(self, owner: "Stage", name: str):
-        self.owner = owner
-        self.name = name
-        self.downstream_inputs: list[PipelineInput] = []
-
-    def _spawn_input(self, owner: "Stage", name: str) -> PipelineInput:
-        input = PipelineInput(self, owner, name)
-        self.downstream_inputs.append(input)
-        return input
+    pass
 
 
 class StageInputs:
@@ -187,11 +173,52 @@ class StageOutputs(ABC):
         return inputs
 
 
+class DynamicStageOutputs(StageOutputs):
+    def __init__(self, owner: "RunScript", output_types: dict[str, FlowableType]):
+        self.owner = owner
+        self._order = list(output_types.keys())
+        self.outputs: dict[str, PipelineOutput] = {}
+        for name in self._order:
+            output_type = output_types[name]
+            output_cls = flowable_name_to_class(output_type)
+            self.outputs[name] = output_cls(owner, name)
+
+    def downstream_inputs(self) -> list[PipelineInput]:
+        inputs = []
+        for output in self.outputs.values():
+            inputs.extend(output.downstream_inputs)
+        return inputs
+
+    def __getattr__(self, name: str) -> PipelineOutput:
+        return self.outputs[name]
+
+    def __getitem__(self, key: int | str) -> PipelineOutput:
+        if isinstance(key, int):
+            name = self._order[key]
+            return self.outputs[name]
+        return self.outputs[key]
+
+    def __iter__(self):
+        return iter(self._order)
+
+    def __len__(self) -> int:
+        return len(self.outputs)
+
+    def keys(self):
+        return self.outputs.keys()
+
+    def values(self):
+        return self.outputs.values()
+
+    def items(self):
+        return self.outputs.items()
+
+
 class StageRegistry:
     def __init__(self):
         self.stages = {}
 
-    def register(self, stage_class: Type["Stage"]) -> None:
+    def register(self, stage_class: Type["StandardStage"] | Type["RunScript"]) -> None:
         self.stages[stage_class.__name__] = stage_class
 
     def get(self, stage_type_name: str) -> Type["Stage"]:
@@ -203,7 +230,7 @@ class StageRegistry:
 TOutputs = TypeVar("TOutputs", bound=StageOutputs)
 
 
-class Stage(Generic[TOutputs], ABC):
+class StandardStage(Generic[TOutputs], ABC):
     def __init__(
         self,
         stage_name: str | None,
@@ -268,7 +295,7 @@ class Stage(Generic[TOutputs], ABC):
 
     def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
-        Stage._registry.register(cls)
+        StandardStage._registry.register(cls)
 
     @classmethod
     def _get_subclass(cls, stage_type_name: str) -> Type["Stage"]:
@@ -281,6 +308,276 @@ class Stage(Generic[TOutputs], ABC):
         return params
 
 
+class RunScript:
+    """
+    RunScript is a stage that runs a user-provided Python function.
+
+    While you can instantiate a RunScript stage directly, the usual way to construct one is to
+    decorate a function with the `@stage` decorator.
+
+    Examples
+    --------
+    >>> @pipelines.stage(
+    ...     inputs={"geometry": read_geo.outputs.geometry},
+    ...     outputs={"geometry": pipelines.PipelineOutputGeometry},
+    ... )
+    ... def ensure_single_volume(geometry: lc.Geometry):
+    ...     _, volumes = geometry.list_entities()
+    ...     if len(volumes) != 1:
+    ...         raise pipelines.StopRun("expected exactly one volume")
+    ...     return {"geometry": geometry}
+    """
+
+    def __init__(
+        self,
+        script: Callable[..., dict[str, Any]] | str,
+        *,
+        stage_name: str | None = None,
+        inputs: dict[str, PipelineOutput] | None = None,
+        outputs: Mapping[str, type[PipelineOutput] | str] | None = None,
+        entrypoint: str | None = None,
+        params: dict[str, Any] | None = None,
+    ):
+        inputs = inputs or {}
+        params = params or {}
+        outputs = outputs or {}
+        overlapping = set(inputs.keys()).intersection(params.keys())
+        if overlapping:
+            overlap = ", ".join(sorted(overlapping))
+            raise ValueError(f"RunScript params and inputs cannot share names: {overlap}")
+
+        inputs_and_params = set(inputs.keys()).union(params.keys())
+        script_source, callable_entrypoint = self._get_script_source(script, inputs_and_params)
+        self._stage_type_name = "RunScript"
+        self._entrypoint = (
+            entrypoint or callable_entrypoint or self._infer_entrypoint(script_source)
+        )
+        self._name = (
+            stage_name if stage_name is not None else self._default_stage_name(self._entrypoint)
+        )
+
+        for input_name, upstream_output in inputs.items():
+            if not isinstance(upstream_output, PipelineOutput):
+                raise TypeError(
+                    f"Input '{input_name}' must be a PipelineOutput, got {type(upstream_output).__name__}"
+                )
+
+        stage_inputs_kwargs = {
+            input_name: (PipelineOutput, upstream_output)
+            for input_name, upstream_output in inputs.items()
+        }
+        self._inputs = StageInputs(self, **stage_inputs_kwargs)
+
+        input_types = {
+            input_name: flowable_class_to_name(type(upstream_output))
+            for input_name, upstream_output in inputs.items()
+        }
+        output_flowable_types = self._normalize_output_types(outputs)
+        self._io_schema = FlowableIOSchema(
+            inputs=input_types,
+            outputs=output_flowable_types,
+        )
+
+        self.outputs = DynamicStageOutputs(self, output_flowable_types)
+
+        reserved_params = {
+            "$script": script_source,
+            "$output_types": {name: ft.value for name, ft in output_flowable_types.items()},
+            "$entrypoint": self._entrypoint,
+        }
+        user_params = dict(params or {})
+        invalid_param_names = ({"context"} | reserved_params.keys()).intersection(
+            user_params.keys()
+        )
+        if invalid_param_names:
+            invalid = ", ".join(sorted(invalid_param_names))
+            raise ValueError(f"RunScript params cannot use reserved names: {invalid}")
+        overlapping_input_names = set(inputs.keys()).intersection(user_params.keys())
+        if overlapping_input_names:
+            overlap = ", ".join(sorted(overlapping_input_names))
+            raise ValueError(f"RunScript params and inputs cannot share names: {overlap}")
+        if "context" in inputs.keys():
+            raise ValueError("RunScript inputs cannot include reserved name 'context'")
+
+        self._params = reserved_params | user_params
+        ensure_yamlizable(self._params_dict()[0], "RunScript parameters")
+
+    @staticmethod
+    def _default_stage_name(entrypoint: str) -> str:
+        words = entrypoint.replace("_", " ").split()
+        if not words:
+            return "RunScript"
+        return " ".join(word.capitalize() for word in words)
+
+    @staticmethod
+    def _normalize_output_types(
+        output_types: Mapping[str, type[PipelineOutput] | str | FlowableType],
+    ) -> dict[str, FlowableType]:
+        normalized: dict[str, FlowableType] = {}
+        if not output_types:
+            raise ValueError("RunScript stages must declare at least one output")
+        for name, value in output_types.items():
+            if isinstance(value, FlowableType):
+                normalized[name] = value
+            elif isinstance(value, str):
+                normalized[name] = FlowableType(value)
+            elif isinstance(value, type) and issubclass(value, PipelineOutput):
+                normalized[name] = flowable_class_to_name(value)
+            else:
+                raise TypeError(
+                    f"Output '{name}' must be a PipelineOutput subclass or flowable type string, got {value}"
+                )
+        return normalized
+
+    @staticmethod
+    def _validate_script(
+        script: Callable[..., dict[str, Any]], inputs_and_params: set[str]
+    ) -> None:
+        closurevars = inspect.getclosurevars(script)
+        if closurevars.nonlocals:
+            raise ValueError(
+                f"RunScript functions must not close over non-local variables. Found these non-local variables: {', '.join(closurevars.nonlocals.keys())}"
+            )
+        globals_except_lc = {
+            k for k in closurevars.globals.keys() if k != "lc" and k != "luminarycloud"
+        }
+        if globals_except_lc:
+            raise ValueError(
+                f"RunScript functions must not rely on global variables, including imports. All modules your script needs (except `luminarycloud` or `lc`) must be imported in the function body. Found globals: {', '.join(globals_except_lc)}"
+            )
+        script_params = set(inspect.signature(script).parameters.keys())
+        if script_params != inputs_and_params and script_params != inputs_and_params | {"context"}:
+            raise ValueError(
+                f"RunScript function must take exactly the same parameters as the inputs and params (and optionally `context`): {script_params} != {inputs_and_params}"
+            )
+
+    @staticmethod
+    def _get_script_source(
+        script: Callable[..., dict[str, Any]] | str,
+        inputs_and_params: set[str],
+    ) -> tuple[str, str | None]:
+        if callable(script):
+            RunScript._validate_script(script, inputs_and_params)
+            try:
+                source_lines, _ = inspect.getsourcelines(script)  # type: ignore[arg-type]
+            except (OSError, IOError, TypeError) as exc:
+                raise ValueError(f"Unable to retrieve source for {script.__name__}: {exc}") from exc
+            # Drop decorator lines (everything before the `def`)
+            for i, line in enumerate(source_lines):
+                if line.lstrip().startswith("def "):
+                    source_lines = source_lines[i:]
+                    break
+            source = "".join(source_lines)
+            entrypoint = script.__name__
+        else:
+            source = script
+            entrypoint = None
+        dedented = textwrap.dedent(source).strip()
+        if not dedented:
+            raise ValueError("RunScript code cannot be empty")
+        return dedented + "\n", entrypoint
+
+    @staticmethod
+    def _infer_entrypoint(script_source: str) -> str:
+        matches = re.findall(r"^def\s+([A-Za-z_][\w]*)\s*\(", script_source, re.MULTILINE)
+        if not matches:
+            raise ValueError(
+                "Could not determine the entrypoint for the RunScript code. Please set the `entrypoint` argument."
+            )
+        unique_matches = [match for match in matches if match]
+        if len(unique_matches) > 1:
+            raise ValueError(
+                "Multiple top-level functions were found in the RunScript code. Please specify the `entrypoint` argument."
+            )
+        return unique_matches[0]
+
+    def is_source(self) -> bool:
+        return len(self._inputs.inputs) == 0
+
+    def inputs_dict(self) -> dict[str, tuple["Stage", str]]:
+        inputs: dict[str, tuple["Stage", str]] = {}
+        for pipeline_input in self._inputs.inputs:
+            inputs[pipeline_input.name] = (
+                pipeline_input.upstream_output.owner,
+                pipeline_input.upstream_output.name,
+            )
+        return inputs
+
+    def downstream_stages(self) -> list["Stage"]:
+        return [inp.owner for inp in self.outputs.downstream_inputs()]
+
+    def _params_dict(self) -> tuple[dict, set[PipelineParameter]]:
+        d: dict[str, Any] = {}
+        pipeline_params = set()
+        for name, value in self._params.items():
+            if hasattr(value, "_to_pipeline_dict"):
+                d[name], downstream_params = value._to_pipeline_dict()
+                for param in downstream_params:
+                    if not isinstance(param, PipelineParameter):
+                        raise ValueError(
+                            f"Expected `_to_pipeline_dict()` to only return PipelineParameters, but got {type(param)}"
+                        )
+                pipeline_params.update(downstream_params)
+            else:
+                d[name] = value
+        d = {k: v for k, v in d.items() if v is not None}
+        return d, pipeline_params
+
+    def _to_dict(self, id_for_task: dict) -> tuple[dict, set[PipelineParameter]]:
+        params, pipeline_params = self._params_dict()
+        d = {
+            "name": self._name,
+            "operator": self._stage_type_name,
+            "params": params,
+            "inputs": self._inputs._to_dict(id_for_task),
+        }
+        return d, pipeline_params
+
+    @classmethod
+    def _parse_params(cls, params: dict) -> dict:
+        return params
+
+
+def stage(
+    *,
+    inputs: dict[str, PipelineOutput] | None = None,
+    outputs: dict[str, type[PipelineOutput]] | None = None,
+    stage_name: str | None = None,
+    params: dict[str, PipelineParameter | PipelineArgValueType] | None = None,
+) -> Callable[[Callable[..., dict[str, Any]]], RunScript]:
+    """
+    Decorator for building a RunScript stage from a Python function.
+
+    Examples
+    --------
+    >>> @pipelines.stage(
+    ...     inputs={"geometry": read_geo.outputs.geometry},
+    ...     outputs={"geometry": pipelines.PipelineOutputGeometry},
+    ... )
+    ... def ensure_single_volume(geometry: lc.Geometry):
+    ...     _, volumes = geometry.list_entities()
+    ...     if len(volumes) != 1:
+    ...         raise pipelines.StopRun("expected exactly one volume")
+    ...     return {"geometry": geometry}
+    """
+
+    def decorator(fn: Callable[..., dict[str, Any]]) -> RunScript:
+        return RunScript(
+            script=fn,
+            stage_name=stage_name,
+            inputs=inputs,
+            outputs=outputs,
+            params=params,
+        )
+
+    return decorator
+
+
+StandardStage._registry.register(RunScript)
+
+Stage = StandardStage | RunScript
+
+
 class Pipeline:
     def __init__(self, stages: list[Stage]):
         self.stages = stages
@@ -350,7 +647,7 @@ class Pipeline:
 
         # first, parse the pipeline parameters...
         parsed_params = {}
-        for param_name, param_metadata in d.get("params", {}).items():
+        for param_name, param_metadata in (d.get("params") or {}).items():
             parsed_params[param_name] = PipelineParameter._get_subclass(param_metadata["type"])(
                 param_name
             )
@@ -405,7 +702,7 @@ def _parse_stage(pipeline_dict: dict, stage_id: str, all_stages: dict[str, Stage
     stage_type_name = stage_dict[
         "operator"
     ]  # TODO: change key to "stage_type" when we're ready to bump the yaml schema version
-    stage_class = Stage._get_subclass(stage_type_name)
+    stage_class = StandardStage._get_subclass(stage_type_name)
 
     parsed_inputs = {}
     for input_name, input_value in stage_dict["inputs"].items():
@@ -414,13 +711,29 @@ def _parse_stage(pipeline_dict: dict, stage_id: str, all_stages: dict[str, Stage
         source_output = getattr(source_stage.outputs, source_output_name)
         parsed_inputs[input_name] = source_output
 
-    parsed_params = stage_class._parse_params(stage_dict["params"])
-
-    stage_params = {
-        "stage_name": stage_dict["name"],
-        **parsed_params,
-        **parsed_inputs,
-    }
-    stage = stage_class(**stage_params)
+    parsed_params = stage_class._parse_params(stage_dict.get("params"))
+
+    if stage_class == RunScript:
+        user_params = parsed_params.copy()
+        script = user_params.pop("$script", None)
+        output_types = user_params.pop("$output_types", None)
+        entrypoint = user_params.pop("$entrypoint", None)
+        if script is None or output_types is None:
+            raise ValueError("RunScript stages must define both `$script` and `$output_types`")
+        stage = RunScript(
+            stage_name=stage_dict["name"],
+            script=script,
+            inputs=parsed_inputs,
+            outputs=output_types,
+            entrypoint=entrypoint,
+            params=user_params,
+        )
+    else:
+        stage_params = {
+            "stage_name": stage_dict["name"],
+            **parsed_params,
+            **parsed_inputs,
+        }
+        stage = stage_class(**stage_params)
     all_stages[stage_id] = stage
     return stage