luminarycloud 0.21.1__py3-none-any.whl → 0.21.2__py3-none-any.whl

luminarycloud/_proto/quantity/quantity_pb2.pyi

@@ -1,7 +1,7 @@
 """
 @generated by mypy-protobuf.  Do not edit manually!
 isort:skip_file
-Copyright 2020-2025 Luminary Cloud, Inc. All Rights Reserved.
+Copyright 2025 Luminary Cloud, Inc. All Rights Reserved.
 Generated by quantities.py. DO NOT EDIT
 """
 import builtins

luminarycloud/enum/gpu_type.py

@@ -10,3 +10,5 @@ class GPUType(IntEnum):
     UNSPECIFIED = simulationpb.SimulationOptions.GPU_TYPE_UNSPECIFIED
     V100 = simulationpb.SimulationOptions.GPU_TYPE_V100
     A100 = simulationpb.SimulationOptions.GPU_TYPE_A100
+    T4 = simulationpb.SimulationOptions.GPU_TYPE_T4
+    H100 = simulationpb.SimulationOptions.GPU_TYPE_H100

luminarycloud/feature_modification.py

@@ -1,12 +1,12 @@
 # Copyright 2025 Luminary Cloud, Inc. All Rights Reserved.
 from enum import Enum, auto
-from typing import Dict, Iterable, List, Optional, cast
+from typing import Dict, Iterable, List, Optional
 from copy import deepcopy
 
-from luminarycloud.types.adfloat import _to_ad_proto, _from_ad_proto
+from luminarycloud.types.adfloat import _to_ad_proto
 from ._proto.geometry import geometry_pb2 as gpb
 from .types import Vector3Like
-from .types.vector3 import _to_vector3
+from .types.vector3 import _to_vector3_ad_proto
 from .params.geometry import Shape, Sphere, Cube, Cylinder, Torus, Cone, HalfSphere, Volume
 from google.protobuf.internal.containers import RepeatedScalarFieldContainer
 
@@ -483,10 +483,7 @@ def modify_translate(
         _update_repeated_field(transform_op.body, vol_ids)
 
     if displacement is not None:
-        vec = _to_vector3(displacement)
-        transform_op.translation.vector.x = vec.x
-        transform_op.translation.vector.y = vec.y
-        transform_op.translation.vector.z = vec.z
+        transform_op.translation.vector.CopyFrom(_to_vector3_ad_proto(displacement))
 
     if keep is not None:
         transform_op.keep = keep
@@ -532,19 +529,13 @@ def modify_rotate(
 
     # Update existing rotation
     if angle is not None:
-        transform_op.rotation.angle = angle
+        transform_op.rotation.angle.CopyFrom(_to_ad_proto(angle))
 
     if axis is not None:
-        axis_vec = _to_vector3(axis)
-        transform_op.rotation.arbitrary.direction.x = axis_vec.x
-        transform_op.rotation.arbitrary.direction.y = axis_vec.y
-        transform_op.rotation.arbitrary.direction.z = axis_vec.z
+        transform_op.rotation.arbitrary.direction.CopyFrom(_to_vector3_ad_proto(axis))
 
     if origin is not None:
-        origin_vec = _to_vector3(origin)
-        transform_op.rotation.arbitrary.origin.x = origin_vec.x
-        transform_op.rotation.arbitrary.origin.y = origin_vec.y
-        transform_op.rotation.arbitrary.origin.z = origin_vec.z
+        transform_op.rotation.arbitrary.origin.CopyFrom(_to_vector3_ad_proto(origin))
 
     if keep is not None:
         transform_op.keep = keep
@@ -587,13 +578,10 @@ def modify_scale(
         _update_repeated_field(transform_op.body, vol_ids)
 
     if scale_factor is not None:
-        transform_op.scaling.isotropic = scale_factor
+        transform_op.scaling.isotropic.CopyFrom(_to_ad_proto(scale_factor))
 
     if origin is not None:
-        origin_vec = _to_vector3(origin)
-        transform_op.scaling.arbitrary.x = origin_vec.x
-        transform_op.scaling.arbitrary.y = origin_vec.y
-        transform_op.scaling.arbitrary.z = origin_vec.z
+        transform_op.scaling.arbitrary.CopyFrom(_to_vector3_ad_proto(origin))
 
     if keep is not None:
         transform_op.keep = keep
@@ -769,10 +757,7 @@ def modify_linear_pattern(
 
     if direction is not None:
         # Update existing linear pattern direction
-        dir_vec = _to_vector3(direction)
-        pattern_op.direction.linear_spacing.vector.x = dir_vec.x
-        pattern_op.direction.linear_spacing.vector.y = dir_vec.y
-        pattern_op.direction.linear_spacing.vector.z = dir_vec.z
+        pattern_op.direction.linear_spacing.vector.CopyFrom(_to_vector3_ad_proto(direction))
 
     if quantity is not None:
         pattern_op.direction.quantity = quantity
@@ -832,19 +817,13 @@ def modify_circular_pattern(
     circular = pattern_op.direction.circular_distribution
 
     if angle is not None:
-        circular.rotation.angle = angle
+        circular.rotation.angle.CopyFrom(_to_ad_proto(angle))
 
     if axis is not None:
-        axis_vec = _to_vector3(axis)
-        circular.rotation.arbitrary.direction.x = axis_vec.x
-        circular.rotation.arbitrary.direction.y = axis_vec.y
-        circular.rotation.arbitrary.direction.z = axis_vec.z
+        circular.rotation.arbitrary.direction.CopyFrom(_to_vector3_ad_proto(axis))
 
     if origin is not None:
-        origin_vec = _to_vector3(origin)
-        circular.rotation.arbitrary.origin.x = origin_vec.x
-        circular.rotation.arbitrary.origin.y = origin_vec.y
-        circular.rotation.arbitrary.origin.z = origin_vec.z
+        circular.rotation.arbitrary.origin.CopyFrom(_to_vector3_ad_proto(origin))
 
     if full_rotation is not None:
         circular.full = full_rotation

luminarycloud/physics_ai/solution.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import tarfile
-from typing import List, Optional, BinaryIO, cast
+from typing import List, Optional, BinaryIO, cast, Dict
 
 from .._client import get_default_client
 from .._helpers.download import download_solution_physics_ai as _download_solution_physics_ai
@@ -16,6 +16,7 @@ def _download_processed_solution_physics_ai(  # noqa: F841
     volume_fields_to_keep: Optional[List[QuantityType]] = None,
     process_volume: bool = False,
     single_precision: bool = True,
+    internal_options: Optional[Dict[str, str]] = None,
 ) -> tarfile.TarFile:
     """
     Download solution data with physics AI processing applied.
@@ -50,6 +51,7 @@ def _download_processed_solution_physics_ai(  # noqa: F841
         volume_fields_to_keep=volume_fields_to_keep,
         process_volume=process_volume,
         single_precision=single_precision,
+        internal_options=internal_options,
     )
 
     assert stream is not None, "Failed to download solution data"

luminarycloud/pipelines/__init__.py

@@ -2,6 +2,8 @@
 from .core import (
     Pipeline as Pipeline,
     PipelineParameter as PipelineParameter,
+    # Stage base class, mainly exported for testing
+    Stage as Stage,
 )
 
 from .parameters import (
@@ -11,9 +13,7 @@ from .parameters import (
     BoolPipelineParameter as BoolPipelineParameter,
 )
 
-from .operators import (
-    # Operator base class, mainly exported for testing
-    Operator as Operator,
+from .stages import (
     # PipelineOutputs, i.e. things that "flow" in a Pipeline
     PipelineOutputGeometry as PipelineOutputGeometry,
     PipelineOutputMesh as PipelineOutputMesh,

luminarycloud/pipelines/api.py

@@ -3,11 +3,15 @@ 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:
@@ -152,6 +156,83 @@ class PipelineJobRecord:
         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:

luminarycloud/pipelines/core.py

@@ -88,48 +88,47 @@ class PipelineParameter(ABC):
 
 class PipelineInput:
     """
-    A named input for an Operator instance (i.e. a Task). Explicitly connected to a PipelineOutput.
+    A named input for a Stage. Explicitly connected to a PipelineOutput.
     """
 
-    def __init__(self, upstream_output: "PipelineOutput", owner: "Operator", name: str):
+    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_task: dict) -> dict:
-        if self.upstream_output.owner not in id_for_task:
+    def _to_dict(self, id_for_stage: dict) -> dict:
+        if self.upstream_output.owner not in id_for_stage:
             raise ValueError(
-                f"Task {self.owner} depends on a task, {self.upstream_output.owner}, that isn't in the Pipeline. Did you forget to add it?"
+                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_task_id = id_for_task[self.upstream_output.owner]
+        upstream_stage_id = id_for_stage[self.upstream_output.owner]
         upstream_output_name = self.upstream_output.name
-        return {self.name: f"{upstream_task_id}.{upstream_output_name}"}
+        return {self.name: f"{upstream_stage_id}.{upstream_output_name}"}
 
 
 class PipelineOutput(ABC):
     """
-    A named output for an Operator instance (i.e. a Task). Can be used to spawn any number of
-    connected PipelineInputs.
+    A named output for a Stage. Can be used to spawn any number of connected PipelineInputs.
     """
 
-    def __init__(self, owner: "Operator", name: str):
+    def __init__(self, owner: "Stage", name: str):
         self.owner = owner
         self.name = name
         self.downstream_inputs: list[PipelineInput] = []
 
-    def _spawn_input(self, owner: "Operator", name: str) -> PipelineInput:
+    def _spawn_input(self, owner: "Stage", name: str) -> PipelineInput:
         input = PipelineInput(self, owner, name)
         self.downstream_inputs.append(input)
         return input
 
 
-class OperatorInputs:
+class StageInputs:
     """
-    A collection of all PipelineInputs for an Operator instance (i.e. a Task).
+    A collection of all PipelineInputs for a Stage.
     """
 
     def __init__(
-        self, owner: "Operator", **input_descriptors: tuple[Type[PipelineOutput], PipelineOutput]
+        self, owner: "Stage", **input_descriptors: tuple[Type[PipelineOutput], PipelineOutput]
     ):
         """
         input_descriptors is a dict of input name -> (required_upstream_output_type, upstream_output)
@@ -144,26 +143,26 @@ class OperatorInputs:
                 )
             self.inputs.add(upstream_output._spawn_input(owner, name))
 
-    def _to_dict(self, id_for_task: dict) -> dict[str, str]:
+    def _to_dict(self, id_for_stage: dict) -> dict[str, str]:
         d: dict[str, str] = {}
         for input in self.inputs:
-            d |= input._to_dict(id_for_task)
+            d |= input._to_dict(id_for_stage)
         return d
 
 
-T = TypeVar("T", bound="OperatorOutputs")
+T = TypeVar("T", bound="StageOutputs")
 
 
-class OperatorOutputs(ABC):
+class StageOutputs(ABC):
     """
-    A collection of all PipelineOutputs for an Operator instance (i.e. a Task). Must be subclassed,
-    and the subclass must also be a dataclass whose fields are all PipelineOutput subclasses. Then
-    that subclass should be instantiated with `_instantiate_for`. Sounds a little complicated,
-    perhaps, but it's not bad. See the existing subclasses in `./operators.py` for examples.
+    A collection of all PipelineOutputs for a Stage. Must be subclassed, and the subclass must also
+    be a dataclass whose fields are all PipelineOutput subclasses. Then that subclass should be
+    instantiated with `_instantiate_for`. Sounds a little complicated, perhaps, but it's not bad.
+    See the existing subclasses in `./stages.py` for examples.
     """
 
     @classmethod
-    def _instantiate_for(cls: type[T], owner: "Operator") -> T:
+    def _instantiate_for(cls: type[T], owner: "Stage") -> T:
         # create an instance with all fields instantiated with the given owner, and named by the
         # field name.
         # Also validate here that we are a dataclass, and all our fields are PipelineOutput types.
@@ -188,41 +187,41 @@ class OperatorOutputs(ABC):
         return inputs
 
 
-class OperatorRegistry:
+class StageRegistry:
     def __init__(self):
-        self.operators = {}
+        self.stages = {}
 
-    def register(self, operator_class: Type["Operator"]) -> None:
-        self.operators[operator_class.__name__] = operator_class
+    def register(self, stage_class: Type["Stage"]) -> None:
+        self.stages[stage_class.__name__] = stage_class
 
-    def get(self, operator_name: str) -> Type["Operator"]:
-        if operator_name not in self.operators:
-            raise ValueError(f"Unknown operator: {operator_name}")
-        return self.operators[operator_name]
+    def get(self, stage_type_name: str) -> Type["Stage"]:
+        if stage_type_name not in self.stages:
+            raise ValueError(f"Unknown stage type: {stage_type_name}")
+        return self.stages[stage_type_name]
 
 
-TOutputs = TypeVar("TOutputs", bound=OperatorOutputs)
+TOutputs = TypeVar("TOutputs", bound=StageOutputs)
 
 
-class Operator(Generic[TOutputs], ABC):
+class Stage(Generic[TOutputs], ABC):
     def __init__(
         self,
-        task_name: str | None,
+        stage_name: str | None,
         params: dict,
-        inputs: OperatorInputs,
+        inputs: StageInputs,
         outputs: TOutputs,
     ):
-        self._operator_name = self.__class__.__name__
-        self._task_name = task_name if task_name is not None else self._operator_name
+        self._stage_type_name = self.__class__.__name__
+        self._name = stage_name if stage_name is not None else self._stage_type_name
         self._params = params
         self._inputs = inputs
         self.outputs = outputs
-        ensure_yamlizable(self._params_dict()[0], "Operator parameters")
+        ensure_yamlizable(self._params_dict()[0], "Stage parameters")
 
     def is_source(self) -> bool:
         return len(self._inputs.inputs) == 0
 
-    def inputs_dict(self) -> dict[str, tuple["Operator", str]]:
+    def inputs_dict(self) -> dict[str, tuple["Stage", str]]:
         inputs = {}
         for pipeline_input in self._inputs.inputs:
             inputs[pipeline_input.name] = (
@@ -231,16 +230,16 @@ class Operator(Generic[TOutputs], ABC):
             )
         return inputs
 
-    def downstream_tasks(self) -> list["Operator"]:
+    def downstream_stages(self) -> list["Stage"]:
         return [input.owner for input in self.outputs.downstream_inputs()]
 
-    def _to_dict(self, id_for_task: dict) -> tuple[dict, set[PipelineParameter]]:
+    def _to_dict(self, id_for_stage: dict) -> tuple[dict, set[PipelineParameter]]:
         params, pipeline_params_set = self._params_dict()
         d = {
-            "name": self._task_name,
-            "operator": self._operator_name,
+            "name": self._name,
+            "operator": self._stage_type_name,  # TODO: change key to "stage_type" when we're ready to bump the yaml schema version
             "params": params,
-            "inputs": self._inputs._to_dict(id_for_task),
+            "inputs": self._inputs._to_dict(id_for_stage),
         }
         return d, pipeline_params_set
 
@@ -256,64 +255,64 @@ class Operator(Generic[TOutputs], ABC):
         return d, pipeline_params
 
     def __str__(self) -> str:
-        return f'{self._operator_name}(name="{self._task_name}")'
+        return f'{self._stage_type_name}(name="{self._name}")'
 
-    _registry = OperatorRegistry()
+    _registry = StageRegistry()
 
     def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
-        Operator._registry.register(cls)
+        Stage._registry.register(cls)
 
     @classmethod
-    def _get_subclass(cls, operator_name: str) -> Type["Operator"]:
-        return cls._registry.get(operator_name)
+    def _get_subclass(cls, stage_type_name: str) -> Type["Stage"]:
+        return cls._registry.get(stage_type_name)
 
     @classmethod
     def _parse_params(cls, params: dict) -> dict:
-        # Operators with params that are just primitives or PipelineParams have no parsing to do.
-        # Operators with more complicated params should override this method.
+        # Stages with params that are just primitives or PipelineParams have no parsing to do.
+        # Stages with more complicated params should override this method.
         return params
 
 
 class Pipeline:
-    def __init__(self, tasks: list[Operator]):
-        self.tasks = tasks
-        self._task_ids = self._assign_ids_to_tasks()
+    def __init__(self, stages: list[Stage]):
+        self.stages = stages
+        self._stage_ids = self._assign_ids_to_stages()
 
     def to_yaml(self) -> str:
         return yaml.safe_dump(self._to_dict())
 
     def pipeline_params(self) -> set[PipelineParameter]:
-        return self._tasks_dict_and_params()[1]
+        return self._stages_dict_and_params()[1]
 
-    def _get_task_id(self, task: Operator) -> str:
-        return self._task_ids[task]
+    def _get_stage_id(self, stage: Stage) -> str:
+        return self._stage_ids[stage]
 
-    def _tasks_dict_and_params(self) -> tuple[dict, set[PipelineParameter]]:
-        id_for_task = self._task_ids
-        tasks = {}
+    def _stages_dict_and_params(self) -> tuple[dict, set[PipelineParameter]]:
+        id_for_stage = self._stage_ids
+        stages = {}
         params = set()
-        for task in id_for_task.keys():
-            task_dict, referenced_params = task._to_dict(id_for_task)
-            tasks[id_for_task[task]] = task_dict
+        for stage in id_for_stage.keys():
+            stage_dict, referenced_params = stage._to_dict(id_for_stage)
+            stages[id_for_stage[stage]] = stage_dict
             params.update(referenced_params)
-        return tasks, params
+        return stages, params
 
     def _to_dict(self) -> dict:
-        tasks, params = self._tasks_dict_and_params()
+        stages, params = self._stages_dict_and_params()
 
         d = {
             "lc_pipeline": {
                 "schema_version": 1,
                 "params": self._pipeline_params_dict(params),
-                "tasks": tasks,
+                "tasks": stages,  # TODO: change key to "stages" when we're ready to bump the yaml schema version
             }
         }
         ensure_yamlizable(d, "Pipeline")
         return d
 
-    def _assign_ids_to_tasks(self) -> dict[Operator, str]:
-        return {task: f"t{i + 1}-{task._operator_name}" for i, task in enumerate(self.tasks)}
+    def _assign_ids_to_stages(self) -> dict[Stage, str]:
+        return {stage: f"s{i + 1}-{stage._stage_type_name}" for i, stage in enumerate(self.stages)}
 
     def _pipeline_params_dict(self, params: set[PipelineParameter]) -> dict:
         d: dict[str, dict] = {}
@@ -334,7 +333,9 @@ class Pipeline:
         d = d["lc_pipeline"]
         if "schema_version" not in d:
             raise ValueError("Invalid pipeline YAML: missing 'schema_version' key")
-        if "tasks" not in d:
+        if (
+            "tasks" not in d
+        ):  # TODO: change key to "stages" when we're ready to bump the yaml schema version
             raise ValueError("Invalid pipeline YAML: missing 'tasks' key")
 
         if d["schema_version"] != 1:
@@ -347,18 +348,20 @@ class Pipeline:
                 param_name
             )
 
-        # ...and use them as replacements for any references in the tasks' parameters
-        for task_dict in d["tasks"].values():
-            task_dict["params"] = _recursive_replace_pipeline_params(
-                task_dict["params"], parsed_params
+        # ...and use them as replacements for any references in the stages' parameters
+        for stage_dict in d[
+            "tasks"
+        ].values():  # TODO: change key to "stages" when we're ready to bump the yaml schema version
+            stage_dict["params"] = _recursive_replace_pipeline_params(
+                stage_dict["params"], parsed_params
             )
 
-        # then, finish parsing the tasks
-        parsed_tasks = {}
-        for task_id in d["tasks"]:
-            _parse_task(d, task_id, parsed_tasks)
+        # then, finish parsing the stages
+        parsed_stages = {}
+        for stage_id in d["tasks"]:
+            _parse_stage(d, stage_id, parsed_stages)
 
-        return cls(list(parsed_tasks.values()))
+        return cls(list(parsed_stages.values()))
 
 
 def _recursive_replace_pipeline_params(d: Any, parsed_params: dict) -> Any:
@@ -368,7 +371,7 @@ def _recursive_replace_pipeline_params(d: Any, parsed_params: dict) -> Any:
             pp_name = d["$pipeline_param"]
             if pp_name not in parsed_params:
                 raise ValueError(
-                    f'Pipeline parameter "{pp_name}" referenced in a pipeline task, but not found in pipeline\'s declared parameters'
+                    f'Pipeline parameter "{pp_name}" referenced in a pipeline stage, but not found in pipeline\'s declared parameters'
                 )
             return parsed_params[pp_name]
         else:
@@ -382,28 +385,32 @@ def _recursive_replace_pipeline_params(d: Any, parsed_params: dict) -> Any:
         return d
 
 
-def _parse_task(pipeline_dict: dict, task_id: str, all_tasks: dict[str, Operator]) -> Operator:
-    all_tasks_dict = pipeline_dict["tasks"]
-    if task_id in all_tasks:
-        return all_tasks[task_id]
-    task_dict = all_tasks_dict[task_id]
-    operator_name = task_dict["operator"]
-    operator_class = Operator._get_subclass(operator_name)
+def _parse_stage(pipeline_dict: dict, stage_id: str, all_stages: dict[str, Stage]) -> Stage:
+    all_stages_dict = pipeline_dict[
+        "tasks"
+    ]  # TODO: change key to "stages" when we're ready to bump the yaml schema version
+    if stage_id in all_stages:
+        return all_stages[stage_id]
+    stage_dict = all_stages_dict[stage_id]
+    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)
 
     parsed_inputs = {}
-    for input_name, input_value in task_dict["inputs"].items():
-        source_task_id, source_output_name = input_value.split(".")
-        source_task = _parse_task(pipeline_dict, source_task_id, all_tasks)
-        source_output = getattr(source_task.outputs, source_output_name)
+    for input_name, input_value in stage_dict["inputs"].items():
+        source_stage_id, source_output_name = input_value.split(".")
+        source_stage = _parse_stage(pipeline_dict, source_stage_id, all_stages)
+        source_output = getattr(source_stage.outputs, source_output_name)
         parsed_inputs[input_name] = source_output
 
-    parsed_params = operator_class._parse_params(task_dict["params"])
+    parsed_params = stage_class._parse_params(stage_dict["params"])
 
-    op_params = {
-        "task_name": task_dict["name"],
+    stage_params = {
+        "stage_name": stage_dict["name"],
         **parsed_params,
         **parsed_inputs,
     }
-    operator = operator_class(**op_params)
-    all_tasks[task_id] = operator
-    return operator
+    stage = stage_class(**stage_params)
+    all_stages[stage_id] = stage
+    return stage