From 3fedeb56c832b9e241e687ebfd41be939f5f08e5 Mon Sep 17 00:00:00 2001 From: magnesj <1793152+magnesj@users.noreply.github.com> Date: Sat, 24 Oct 2026 06:26:53 +0000 Subject: [PATCH] Update rips module and proto files --- .../alter_wbs_plot.py | 0 .../create_wbs_plot.py | 0 .../geomech/read_geomech_results.py | 56 +++ .../select_geomech_cell_result.py} | 0 docs/rips/case.py | 71 ++++ docs/rips/generated/generated_classes.py | 133 +++--- docs/rips/taskmaestro_helper/__init__.py | 10 +- docs/rips/taskmaestro_helper/__main__.py | 124 +++++- docs/rips/taskmaestro_helper/catalog.py | 132 ++++++ docs/rips/taskmaestro_helper/definition.py | 117 ++++++ docs/rips/taskmaestro_helper/refs.py | 28 ++ docs/rips/taskmaestro_helper/run.py | 73 +++- docs/rips/taskmaestro_helper/schema.py | 56 +++ docs/rips/tests/test_cases.py | 99 +++++ .../tm_helper_pipeline.py | 79 ++++ docs/rips/tests/test_taskmaestro_helper.py | 80 +++- .../test_taskmaestro_helper_definition.py | 385 ++++++++++++++++++ 17 files changed, 1355 insertions(+), 88 deletions(-) rename docs/rips/PythonExamples/{export_and_plotting => geomech}/alter_wbs_plot.py (100%) rename docs/rips/PythonExamples/{export_and_plotting => geomech}/create_wbs_plot.py (100%) create mode 100644 docs/rips/PythonExamples/geomech/read_geomech_results.py rename docs/rips/PythonExamples/{case_and_grid_operations/set_geomech_cell_result.py => geomech/select_geomech_cell_result.py} (100%) create mode 100644 docs/rips/taskmaestro_helper/catalog.py create mode 100644 docs/rips/taskmaestro_helper/definition.py create mode 100644 docs/rips/taskmaestro_helper/schema.py create mode 100644 docs/rips/tests/test_data/taskmaestro_workflows/tm_helper_pipeline.py create mode 100644 docs/rips/tests/test_taskmaestro_helper_definition.py diff --git a/docs/rips/PythonExamples/export_and_plotting/alter_wbs_plot.py b/docs/rips/PythonExamples/geomech/alter_wbs_plot.py similarity index 100% rename from docs/rips/PythonExamples/export_and_plotting/alter_wbs_plot.py rename to docs/rips/PythonExamples/geomech/alter_wbs_plot.py diff --git a/docs/rips/PythonExamples/export_and_plotting/create_wbs_plot.py b/docs/rips/PythonExamples/geomech/create_wbs_plot.py similarity index 100% rename from docs/rips/PythonExamples/export_and_plotting/create_wbs_plot.py rename to docs/rips/PythonExamples/geomech/create_wbs_plot.py diff --git a/docs/rips/PythonExamples/geomech/read_geomech_results.py b/docs/rips/PythonExamples/geomech/read_geomech_results.py new file mode 100644 index 0000000000..9afb15252a --- /dev/null +++ b/docs/rips/PythonExamples/geomech/read_geomech_results.py @@ -0,0 +1,56 @@ +###################################################################### +# This script loads a geomechanical case (.odb, .inp or VTK .pvd) and +# reads time steps, result names and result values into Python lists. +# +# Usage: python read_geomech_results.py +###################################################################### +import math +import sys + +import rips + +if len(sys.argv) < 2: + print("Usage: read_geomech_results.py ") + sys.exit(1) + +try: + resinsight = rips.Instance.find() + case = resinsight.project.load_case(sys.argv[1]) +except Exception as exc: + print(f"Failed to load geomechanical case: {exc}") + sys.exit(1) + +if not isinstance(case, rips.GeoMechCase): + print("The file is not a geomechanical case") + sys.exit(1) + +time_steps = case.time_step_strings() +print("Time steps:", time_steps) + +# Result names are "FIELD:COMPONENT", or "FIELD" for results without components +for position in rips.GeoMechResultPosition: + if position in ( + rips.GeoMechResultPosition.WELLPATH_DERIVED, + rips.GeoMechResultPosition.DIFFERENTIALS, + ): + continue + print(position, case.result_names(position)) + +# Print statistics for every result with defined values at one position. +# A negative frame_index (the default) selects the last frame of the time step. +position = rips.GeoMechResultPosition.NODAL +for name in case.result_names(position): + field_name, _, component_name = name.partition(":") + for time_step, label in enumerate(time_steps): + values = case.result_values( + field_name, + component_name, + time_step=time_step, + result_position=position, + ) + finite = [v for v in values if math.isfinite(v)] + if finite: + print( + f"{name}, {label}: {len(values)} values, " + f"min={min(finite):.4g}, max={max(finite):.4g}" + ) diff --git a/docs/rips/PythonExamples/case_and_grid_operations/set_geomech_cell_result.py b/docs/rips/PythonExamples/geomech/select_geomech_cell_result.py similarity index 100% rename from docs/rips/PythonExamples/case_and_grid_operations/set_geomech_cell_result.py rename to docs/rips/PythonExamples/geomech/select_geomech_cell_result.py diff --git a/docs/rips/case.py b/docs/rips/case.py index cfa9340262..909b0c6b31 100644 --- a/docs/rips/case.py +++ b/docs/rips/case.py @@ -54,6 +54,7 @@ Case as Case, EclipseCase as EclipseCase, GeoMechCase as GeoMechCase, + GeoMechResultPosition as GeoMechResultPosition, PorosityModelType as PorosityModelType, PropertyDataType as PropertyDataType, PropertyType as PropertyType, @@ -303,6 +304,76 @@ def time_steps(self) -> List[Case_pb2.TimeStepDate]: return self.__case_stub.GetTimeSteps(self.__request()).dates +@add_method(GeoMechCase) +def time_step_strings(self: GeoMechCase) -> List[str]: + """Get the time step strings of this geomechanical case + + Returns: + A list of time step strings, one per time step. The list index + matches the time_step argument of result_values(). + """ + return self.time_step_strings_internal().values + + +@add_method(GeoMechCase) +def result_names( + self: GeoMechCase, + result_position: GeoMechResultPosition = GeoMechResultPosition.ELEMENT_NODAL, +) -> List[str]: + """Get the result names of this geomechanical case + + Arguments: + result_position(GeoMechResultPosition): Result position, for example + GeoMechResultPosition.NODAL or GeoMechResultPosition.ELEMENT_NODAL. + + Returns: + A list of names on the form "FIELD:COMPONENT", or "FIELD" for results + without components. + """ + return self.result_names_internal(result_position=result_position).values + + +@add_method(GeoMechCase) +def result_values( + self: GeoMechCase, + field_name: str, + component_name: str = "", + time_step: int = 0, + result_position: GeoMechResultPosition = GeoMechResultPosition.ELEMENT_NODAL, + frame_index: int = -1, +) -> List[float]: + """Read the values of a geomechanical result + + Arguments: + field_name(str): Result field name, e.g. "ST" or "POR-Bar". + component_name(str): Result component name, e.g. "S1". Empty for results without components. + time_step(int): Zero-based time step index. + result_position(GeoMechResultPosition): Result position of the result. + frame_index(int): Zero-based frame index. Negative selects the last frame of the time step. + + Returns: + The values of all parts, in the order used by ResInsight for the result position. + Values are NaN where the result is not defined. + """ + project = self.ancestor(Project) + if project is None: + raise RipsError("Could not find parent project") + + key = f"{uuid.uuid4()}_geomech_result" + try: + self.read_result_values_internal( + result_position=result_position, + field_name=field_name, + component_name=component_name, + time_step=time_step, + frame_index=frame_index, + key=key, + ) + return project.key_values(key) # type: ignore[attr-defined] + finally: + project.remove_key_values(key) # type: ignore[attr-defined] + + @add_method(Case) def reservoir_boundingbox(self) -> Case_pb2.BoundingBox: """Get the reservoir bounding box diff --git a/docs/rips/generated/generated_classes.py b/docs/rips/generated/generated_classes.py index d4845fbd92..e81a8bacfc 100644 --- a/docs/rips/generated/generated_classes.py +++ b/docs/rips/generated/generated_classes.py @@ -50,6 +50,13 @@ class ControlMode(StrEnum): BHP = "BHP" THP = "THP" +class CrossSectionType(StrEnum): + CS_WELL_PATH = "CS_WELL_PATH" + CS_SIMULATION_WELL = "CS_SIMULATION_WELL" + CS_POLYLINE = "CS_POLYLINE" + CS_AZIMUTHLINE = "CS_AZIMUTHLINE" + CS_POLYGON = "CS_POLYGON" + class DepthType(StrEnum): MEASURED_DEPTH = "MEASURED_DEPTH" TRUE_VERTICAL_DEPTH = "TRUE_VERTICAL_DEPTH" @@ -97,6 +104,16 @@ class GasInflowEq(StrEnum): P_P = "P-P" GPP = "GPP" +class GeoMechResultPosition(StrEnum): + NODAL = "NODAL" + ELEMENT_NODAL = "ELEMENT_NODAL" + INTEGRATION_POINT = "INTEGRATION_POINT" + ELEMENT_NODAL_FACE = "ELEMENT_NODAL_FACE" + FORMATION_NAMES = "FORMATION_NAMES" + ELEMENT = "ELEMENT" + WELLPATH_DERIVED = "WELLPATH_DERIVED" + DIFFERENTIALS = "DIFFERENTIALS" + class GeometryType(StrEnum): FULL_3D = "FULL_3D" PROJECTED_TO_PLANE = "PROJECTED_TO_PLANE" @@ -109,10 +126,6 @@ class LengthAndDepth(StrEnum): INC = "INC" ABS = "ABS" -class LengthAndDepth2(StrEnum): - INC = "INC" - ABS = "ABS" - class LgrSplitType(StrEnum): LGR_PER_CELL = "LGR_PER_CELL" LGR_PER_COMPLETION = "LGR_PER_COMPLETION" @@ -137,11 +150,6 @@ class NonDarcyFlowType(StrEnum): Computed = "Computed" UserDefined = "UserDefined" -class NonDarcyFlowType2(StrEnum): - None_ = "None" - Computed = "Computed" - UserDefined = "UserDefined" - class NumberOfColumns(StrEnum): _1 = "1" _2 = "2" @@ -193,11 +201,6 @@ class PressureDrop(StrEnum): HF_ = "HF-" HFA = "HFA" -class PressureDrop2(StrEnum): - H__ = "H--" - HF_ = "HF-" - HFA = "HFA" - class Property(StrEnum): UNDEFINED = "UNDEFINED" FACIES = "FACIES" @@ -250,20 +253,6 @@ class ReferenceMdType(StrEnum): GridEntryPoint = "GridEntryPoint" UserDefined = "UserDefined" -class ReferenceMdType2(StrEnum): - GridEntryPoint = "GridEntryPoint" - UserDefined = "UserDefined" - -class ResultPositionType(StrEnum): - NODAL = "NODAL" - ELEMENT_NODAL = "ELEMENT_NODAL" - INTEGRATION_POINT = "INTEGRATION_POINT" - ELEMENT_NODAL_FACE = "ELEMENT_NODAL_FACE" - FORMATION_NAMES = "FORMATION_NAMES" - ELEMENT = "ELEMENT" - WELLPATH_DERIVED = "WELLPATH_DERIVED" - DIFFERENTIALS = "DIFFERENTIALS" - class RowsPerPage(StrEnum): _1 = "1" _2 = "2" @@ -288,10 +277,6 @@ class State(StrEnum): OPEN = "OPEN" SHUT = "SHUT" -class State2(StrEnum): - OPEN = "OPEN" - SHUT = "SHUT" - class SubTitleFontSize(StrEnum): XX_Small = "XX_Small" X_Small = "X_Small" @@ -309,13 +294,6 @@ class TransScalingWbhpSource(StrEnum): WBHP_SUMMARY = "WBHP_SUMMARY" WBHP_USER_DEFINED = "WBHP_USER_DEFINED" -class Type2(StrEnum): - CS_WELL_PATH = "CS_WELL_PATH" - CS_SIMULATION_WELL = "CS_SIMULATION_WELL" - CS_POLYLINE = "CS_POLYLINE" - CS_AZIMUTHLINE = "CS_AZIMUTHLINE" - CS_POLYGON = "CS_POLYGON" - class ValveType(StrEnum): ICV = "ICV" ICD = "ICD" @@ -984,7 +962,7 @@ class CurveIntersection(PdmObjectBase): name (str): Name points (List[List[float]]): Points simulation_well (Optional[SimulationWell]): Simulation Well - type (Type2): One of [CS_WELL_PATH, CS_SIMULATION_WELL, CS_POLYLINE, CS_AZIMUTHLINE, CS_POLYGON] + type (CrossSectionType): One of [CS_WELL_PATH, CS_SIMULATION_WELL, CS_POLYLINE, CS_AZIMUTHLINE, CS_POLYGON] well_path (Optional[WellPath]): Well Path """ __custom_init__ = None #: Assign a custom init routine to be run at __init__ @@ -993,7 +971,7 @@ def __init__(self, pb2_object: Optional[PdmObject_pb2.PdmObject]=None, channel: self.name: str = "Intersection Name" self.points: List[List[float]] = [] self.simulation_well: Optional[SimulationWell] = None - self.type: Type2 = Type2.CS_POLYLINE + self.type: CrossSectionType = CrossSectionType.CS_POLYLINE self.well_path: Optional[WellPath] = None PdmObjectBase.__init__(self, pb2_object, channel) if CurveIntersection.__custom_init__ is not None: @@ -1809,14 +1787,14 @@ class GeoMechResult(PdmObjectBase): Attributes: result_component_name (str): Component result_field_name (str): Field Name - result_position_type (ResultPositionType): One of [NODAL, ELEMENT_NODAL, INTEGRATION_POINT, ELEMENT_NODAL_FACE, FORMATION_NAMES, ELEMENT, WELLPATH_DERIVED, DIFFERENTIALS] + result_position_type (GeoMechResultPosition): One of [NODAL, ELEMENT_NODAL, INTEGRATION_POINT, ELEMENT_NODAL_FACE, FORMATION_NAMES, ELEMENT, WELLPATH_DERIVED, DIFFERENTIALS] """ __custom_init__ = None #: Assign a custom init routine to be run at __init__ def __init__(self, pb2_object: Optional[PdmObject_pb2.PdmObject]=None, channel: Optional[grpc.Channel]=None) -> None: self.result_component_name: str = "" self.result_field_name: str = "" - self.result_position_type: ResultPositionType = ResultPositionType.NODAL + self.result_position_type: GeoMechResultPosition = GeoMechResultPosition.NODAL PdmObjectBase.__init__(self, pb2_object, channel) if GeoMechResult.__custom_init__ is not None: GeoMechResult.__custom_init__(self, pb2_object=pb2_object, channel=channel) @@ -2930,6 +2908,47 @@ def create_well_bore_stability_plot(self, well_path: Optional[WellPath]=None, ti return self._call_pdm_method_return_value("createWellBoreStabilityPlot", WellBoreStabilityPlot, well_path=well_path, time_step=time_step) + def read_result_values_internal(self, result_position: GeoMechResultPosition=GeoMechResultPosition.ELEMENT_NODAL, field_name: str="", component_name: str="", time_step: int=0, frame_index: int=-1, key: str="") -> None: + """ + Read GeoMech result values + + Arguments: + result_position (GeoMechResultPosition): One of [NODAL, ELEMENT_NODAL, INTEGRATION_POINT, ELEMENT_NODAL_FACE, FORMATION_NAMES, ELEMENT, WELLPATH_DERIVED, DIFFERENTIALS] + field_name (str): + component_name (str): + time_step (int): Zero-based time step index + frame_index (int): Zero-based frame index. Negative selects the last frame + key (str): + Returns: + + """ + self._call_pdm_method_void("ReadResultValuesInternal", result_position=result_position, field_name=field_name, component_name=component_name, time_step=time_step, frame_index=frame_index, key=key) + + + def result_names_internal(self, result_position: GeoMechResultPosition=GeoMechResultPosition.ELEMENT_NODAL) -> DataContainerString: + """ + Result names for this case. + + Arguments: + result_position (GeoMechResultPosition): One of [NODAL, ELEMENT_NODAL, INTEGRATION_POINT, ELEMENT_NODAL_FACE, FORMATION_NAMES, ELEMENT, WELLPATH_DERIVED, DIFFERENTIALS] + Returns: + DataContainerString + """ + return self._call_pdm_method_return_value("ResultNamesInternal", DataContainerString, result_position=result_position) + + + def time_step_strings_internal(self, ) -> DataContainerString: + """ + Time step strings for this case. + + Arguments: + + Returns: + DataContainerString + """ + return self._call_pdm_method_return_value("TimeStepStringsInternal", DataContainerString) + + def views(self) -> List[GeoMechView]: """All GeoMech Views in the Case @@ -3726,7 +3745,7 @@ class SegmentCollection(PdmObjectBase): compsegs_segment_number (CompsegsSegmentNumber): One of [AssignedBySimulator, Exported] custom_values_for_lateral (bool): Custom Values for Lateral liner_diameter (float): Liner Inner Diameter - reference_md_type (ReferenceMdType2): One of [GridEntryPoint, UserDefined] + reference_md_type (ReferenceMdType): One of [GridEntryPoint, UserDefined] roughness_factor (float): Roughness Factor single_segment_after_last_perforation (bool): Single Segment After Last Perforation single_segment_before_first_perforation (bool): Single Segment Before First Perforation @@ -3738,7 +3757,7 @@ def __init__(self, pb2_object: Optional[PdmObject_pb2.PdmObject]=None, channel: self.compsegs_segment_number: CompsegsSegmentNumber = CompsegsSegmentNumber.AssignedBySimulator self.custom_values_for_lateral: bool = False self.liner_diameter: float = 1.520000000000000e-01 - self.reference_md_type: ReferenceMdType2 = ReferenceMdType2.GridEntryPoint + self.reference_md_type: ReferenceMdType = ReferenceMdType.GridEntryPoint self.roughness_factor: float = 1.000000000000000e-05 self.single_segment_after_last_perforation: bool = False self.single_segment_before_first_perforation: bool = False @@ -3852,7 +3871,7 @@ class FractureTemplate(PdmObjectBase): height_scale_factor (float): Height id (int): ID inertial_coefficient (float): Inertial Coefficient (β) [Forch. unit] - non_darcy_flow_type (NonDarcyFlowType2): One of [None, Computed, UserDefined] + non_darcy_flow_type (NonDarcyFlowType): One of [None, Computed, UserDefined] orientation (Orientation): One of [Azimuth, Longitudinal, Transverse] perforation_length (float): Perforation Length permeability_type (PermeabilityType): One of [UserDefinedPermeability, FractureConductivity] @@ -3878,7 +3897,7 @@ def __init__(self, pb2_object: Optional[PdmObject_pb2.PdmObject]=None, channel: self.height_scale_factor: float = 1.000000000000000e+00 self.id: int = -1 self.inertial_coefficient: float = 6.083236000000000e-03 - self.non_darcy_flow_type: NonDarcyFlowType2 = NonDarcyFlowType2.None_ + self.non_darcy_flow_type: NonDarcyFlowType = NonDarcyFlowType.None_ self.orientation: Orientation = Orientation.Transverse self.perforation_length: float = 1.000000000000000e+00 self.permeability_type: PermeabilityType = PermeabilityType.FractureConductivity @@ -4994,7 +5013,7 @@ def add_state_event(self, event_date: str="2024-01-01", well_path: Optional[Well return self._call_pdm_method_return_value("AddStateEvent", WellEventState, event_date=event_date, well_path=well_path, well_state=well_state) - def add_valve_event(self, event_date: str="2024-01-01", well_path: Optional[WellPath]=None, measured_depth: float=0.000000000000000e+00, valve_type: ValveType=ValveType.ICV, state: State2=State2.OPEN, flow_coefficient: float=7.000000000000000e-01, area: float=1.000000000000000e-04, aicd_strength: float=2.100000000000000e-04, aicd_density_calib_fluid: float=1.000000000000000e+03, aicd_viscosity_calib_fluid: float=1.000000000000000e+00, aicd_vol_flow_exp: float=2.100000000000000e+00, aicd_visc_func_exp: float=5.000000000000000e-01) -> WellEventValve: + def add_valve_event(self, event_date: str="2024-01-01", well_path: Optional[WellPath]=None, measured_depth: float=0.000000000000000e+00, valve_type: ValveType=ValveType.ICV, state: State=State.OPEN, flow_coefficient: float=7.000000000000000e-01, area: float=1.000000000000000e-04, aicd_strength: float=2.100000000000000e-04, aicd_density_calib_fluid: float=1.000000000000000e+03, aicd_viscosity_calib_fluid: float=1.000000000000000e+00, aicd_vol_flow_exp: float=2.100000000000000e+00, aicd_visc_func_exp: float=5.000000000000000e-01) -> WellEventValve: """ Add a valve event to the timeline @@ -5003,7 +5022,7 @@ def add_valve_event(self, event_date: str="2024-01-01", well_path: Optional[Well well_path (Optional[WellPath]): Well Path measured_depth (float): Measured Depth valve_type (ValveType): One of [ICV, ICD, AICD] - state (State2): One of [OPEN, SHUT] + state (State): One of [OPEN, SHUT] flow_coefficient (float): Flow Coefficient area (float): Area [m2] aicd_strength (float): AICD Strength @@ -5145,7 +5164,7 @@ class WellEventValve(WellEvent): area (float): Area [m2] flow_coefficient (float): Flow Coefficient measured_depth (float): Measured Depth - state (State2): One of [OPEN, SHUT] + state (State): One of [OPEN, SHUT] valve_template (Optional[ValveTemplate]): Valve Template valve_type (ValveType): One of [ICV, ICD, AICD] """ @@ -5160,7 +5179,7 @@ def __init__(self, pb2_object: Optional[PdmObject_pb2.PdmObject]=None, channel: self.area: float = 1.000000000000000e-04 self.flow_coefficient: float = 7.000000000000000e-01 self.measured_depth: float = 0.000000000000000e+00 - self.state: State2 = State2.OPEN + self.state: State = State.OPEN self.valve_template: Optional[ValveTemplate] = None self.valve_type: ValveType = ValveType.ICV WellEvent.__init__(self, pb2_object, channel) @@ -5174,7 +5193,7 @@ class WellEventWellSpec(WellEvent): Attributes: allow_cross_flow (bool): Allow Cross-Flow group_name (str): Group Name - pressure_drop (PressureDrop2): One of [H--, HF-, HFA] + pressure_drop (PressureDrop): One of [H--, HF-, HFA] reference_depth (Optional[float]): Reference Depth well_type (WellTypeForExport): One of [OIL, GAS, WATER, LIQUID] """ @@ -5183,7 +5202,7 @@ class WellEventWellSpec(WellEvent): def __init__(self, pb2_object: Optional[PdmObject_pb2.PdmObject]=None, channel: Optional[grpc.Channel]=None) -> None: self.allow_cross_flow: bool = True self.group_name: str = "" - self.pressure_drop: PressureDrop2 = PressureDrop2.HF_ + self.pressure_drop: PressureDrop = PressureDrop.HF_ self.reference_depth: Optional[float] = None self.well_type: WellTypeForExport = WellTypeForExport.OIL WellEvent.__init__(self, pb2_object, channel) @@ -5364,8 +5383,8 @@ class WellPathCompletionSettings(PdmObjectBase): gas_inflow_eq (GasInflowEq): One of [STD, R-G, P-P, GPP] group_name_for_export (str): Group Name hydrostatic_density (HydrostaticDensity): One of [SEG, AVG] - length_and_depth (LengthAndDepth2): One of [INC, ABS] - pressure_drop (PressureDrop2): One of [H--, HF-, HFA] + length_and_depth (LengthAndDepth): One of [INC, ABS] + pressure_drop (PressureDrop): One of [H--, HF-, HFA] reference_depth_for_export (Optional[float]): BHP Reference Depth well_bore_fluid_pvt_table (int): Wellbore Fluid PVT table well_name_for_export (str): Well Name @@ -5381,8 +5400,8 @@ def __init__(self, pb2_object: Optional[PdmObject_pb2.PdmObject]=None, channel: self.gas_inflow_eq: GasInflowEq = GasInflowEq.STD self.group_name_for_export: str = "" self.hydrostatic_density: HydrostaticDensity = HydrostaticDensity.SEG - self.length_and_depth: LengthAndDepth2 = LengthAndDepth2.ABS - self.pressure_drop: PressureDrop2 = PressureDrop2.HF_ + self.length_and_depth: LengthAndDepth = LengthAndDepth.ABS + self.pressure_drop: PressureDrop = PressureDrop.HF_ self.reference_depth_for_export: Optional[float] = None self.well_bore_fluid_pvt_table: int = 0 self.well_name_for_export: str = "" diff --git a/docs/rips/taskmaestro_helper/__init__.py b/docs/rips/taskmaestro_helper/__init__.py index cc4a3947c6..d961e18f0f 100644 --- a/docs/rips/taskmaestro_helper/__init__.py +++ b/docs/rips/taskmaestro_helper/__init__.py @@ -1,6 +1,10 @@ """Bridge between ResInsight and the taskmaestro workflow library. -ResInsight invokes this package as a subprocess to execute workflows against -a running ResInsight. Workflow schemas are read with the taskmaestro CLI -(`python -m taskmaestro workflow describe`). +ResInsight invokes this package as a subprocess (``python -m +rips.taskmaestro_helper ``) to: + +- list installed tasks and workflows (``catalog``), +- convert workflow folders and registered workflows to editable workflow + definitions (``load``, ``export-registered``) and back to YAML (``save``), +- execute workflows against a running ResInsight (``run``). """ diff --git a/docs/rips/taskmaestro_helper/__main__.py b/docs/rips/taskmaestro_helper/__main__.py index 5396a160bd..c63161f6e8 100644 --- a/docs/rips/taskmaestro_helper/__main__.py +++ b/docs/rips/taskmaestro_helper/__main__.py @@ -1,28 +1,134 @@ -"""Command-line entry point: `python -m rips.taskmaestro_helper `.""" +"""Command-line entry point: `python -m rips.taskmaestro_helper `. + +Subcommands other than ``run`` print exactly one JSON document to stdout. Any +output written while importing task modules is redirected to stderr, including +output written by extension modules and child processes straight to file +descriptor 1. On +failure they exit with code 2 and print +``{"status": "invalid", "error": {"message", "task", "field"}}``. +""" from __future__ import annotations +import json +import os import sys +from collections.abc import Callable, Iterator +from contextlib import contextmanager, redirect_stdout +from typing import Any + +USAGE = ( + "usage: python -m rips.taskmaestro_helper " + "{catalog,load,save,export-registered,run} ..." +) + + +def _catalog(argv: list[str]) -> dict[str, Any]: + from .catalog import catalog_main + + return catalog_main(argv) + + +def _load(argv: list[str]) -> dict[str, Any]: + from .definition import load_main + + return load_main(argv) + + +def _save(argv: list[str]) -> dict[str, Any]: + from .definition import save_main + + return save_main(argv, sys.stdin.read()) + + +def _export_registered(argv: list[str]) -> dict[str, Any]: + from .catalog import export_registered_main + return export_registered_main(argv) -def main() -> int: - if len(sys.argv) < 2: + +JSON_COMMANDS: dict[str, Callable[[list[str]], dict[str, Any]]] = { + "catalog": _catalog, + "load": _load, + "save": _save, + "export-registered": _export_registered, +} + + +def _invalid(error: dict[str, Any]) -> str: + return json.dumps({"status": "invalid", "error": error}) + + +def _error_json(exc: Exception) -> dict[str, Any]: + to_json = getattr(exc, "to_json", None) # taskmaestro's DefinitionError + if callable(to_json): + return dict(to_json()) + if isinstance(exc, ImportError) and (exc.name or "").startswith("taskmaestro"): + message = f"taskmaestro>=0.4,<0.5 is required to edit workflows: {exc}" + else: + message = f"{type(exc).__name__}: {exc}" + return {"message": message, "task": None, "field": None} + + +@contextmanager +def _stdout_to_stderr() -> Iterator[None]: + """Send everything written to stdout to stderr, at the file descriptor level.""" + sys.stdout.flush() + try: + saved_fd = os.dup(1) + except OSError: # no usable stdout file descriptor; redirect Python output only + saved_fd = None + try: + if saved_fd is not None: + os.dup2(2, 1) + with redirect_stdout(sys.stderr): + yield + finally: + sys.stderr.flush() + if saved_fd is not None: + os.dup2(saved_fd, 1) + os.close(saved_fd) + + +def run_json_command( + command: Callable[[list[str]], dict[str, Any]], argv: list[str] +) -> int: + stdout = sys.stdout + try: + with _stdout_to_stderr(): + result = command(argv) + except SystemExit as exc: # argparse errors print( - "usage: python -m rips.taskmaestro_helper run ...", - file=sys.stderr, + _invalid( + {"message": f"Invalid arguments: {exc}", "task": None, "field": None} + ), + file=stdout, ) return 2 + except Exception as exc: # noqa: BLE001 - always answer ResInsight with JSON + print(_invalid(_error_json(exc)), file=stdout) + return 2 + print(json.dumps(result), file=stdout) + return 0 - subcommand = sys.argv[1] - rest = sys.argv[2:] +def main(argv: list[str] | None = None) -> int: + args = sys.argv[1:] if argv is None else argv + if not args: + print(USAGE, file=sys.stderr) + return 2 + + subcommand, rest = args[0], args[1:] if subcommand == "run": from . import run return run.main(rest) - print(f"unknown subcommand: {subcommand}", file=sys.stderr) - return 2 + command = JSON_COMMANDS.get(subcommand) + if command is None: + print(f"unknown subcommand: {subcommand}\n{USAGE}", file=sys.stderr) + return 2 + return run_json_command(command, rest) if __name__ == "__main__": diff --git a/docs/rips/taskmaestro_helper/catalog.py b/docs/rips/taskmaestro_helper/catalog.py new file mode 100644 index 0000000000..f5accaf28e --- /dev/null +++ b/docs/rips/taskmaestro_helper/catalog.py @@ -0,0 +1,132 @@ +"""Catalog of installed taskmaestro tasks and workflows. + +Every entry point is loaded on its own, so one broken plugin is reported in +``errors`` while the rest of the catalog still loads. +""" + +from __future__ import annotations + +from typing import Any + +import taskmaestro +from taskmaestro.exceptions import DefinitionError + +from .definition import definition_from_workflow +from .schema import task_descriptor + +CATALOG_FORMAT = 1 + + +def _error(problem: Any) -> dict[str, Any]: + return { + "kind": problem.kind, + "id": problem.name, + "entry_point": problem.entry_point, + "message": problem.message, + } + + +def class_ids(tasks: dict[str, type[Any]]) -> dict[type[Any], str]: + """Map task classes back to their (first) entry-point id.""" + result: dict[type[Any], str] = {} + for task_id, cls in sorted(tasks.items()): + result.setdefault(cls, task_id) + return result + + +def _workflow_entry( + workflow_id: str, workflow: Any, ids: dict[type[Any], str] +) -> dict[str, Any]: + definition = definition_from_workflow(workflow, workflow_id, ids, describe=True) + return { + "id": workflow_id, + "name": workflow.name, + "definition": definition, + "describe": definition["describe"], + "editable": definition["editable"], + "readonly_reasons": definition["readonly_reasons"], + } + + +def build_catalog(*, include_workflows: bool = True) -> dict[str, Any]: + from taskmaestro.discovery import load_registry + + registry = load_registry(include_workflows=include_workflows) + errors = [_error(problem) for problem in registry.problems] + descriptors: list[dict[str, Any]] = [] + for task_id, cls in sorted(registry.tasks.items()): + try: + descriptors.append(task_descriptor(task_id, cls)) + except Exception as exc: # noqa: BLE001 - report plugin errors, keep the rest + errors.append( + { + "kind": "task", + "id": task_id, + "entry_point": "", + "message": f"Cannot describe task: {exc}", + } + ) + + workflows: list[dict[str, Any]] = [] + ids = class_ids(registry.tasks) + for workflow_id, workflow in sorted(registry.workflows.items()): + try: + workflows.append(_workflow_entry(workflow_id, workflow, ids)) + except Exception as exc: # noqa: BLE001 - report plugin errors, keep the rest + errors.append( + { + "kind": "workflow", + "id": workflow_id, + "entry_point": "", + "message": str(exc), + } + ) + + return { + "format": CATALOG_FORMAT, + "status": "ok", + "taskmaestro_version": taskmaestro.__version__, + "tasks": descriptors, + "workflows": workflows, + "errors": errors, + } + + +def export_registered(workflow_id: str, *, describe: bool = False) -> dict[str, Any]: + from taskmaestro.discovery import get_registered_workflow, load_registry + from taskmaestro.exceptions import PluginLoadError + + try: + workflow = get_registered_workflow(workflow_id) + except PluginLoadError as exc: + raise DefinitionError(str(exc)) from exc + tasks = load_registry(include_workflows=False).tasks + return definition_from_workflow( + workflow, workflow_id, class_ids(tasks), describe=describe + ) + + +def catalog_main(argv: list[str]) -> dict[str, Any]: + import argparse + + parser = argparse.ArgumentParser( + prog="catalog", description="List installed taskmaestro tasks and workflows." + ) + parser.add_argument("--no-workflows", action="store_true") + args = parser.parse_args(argv) + return build_catalog(include_workflows=not args.no_workflows) + + +def export_registered_main(argv: list[str]) -> dict[str, Any]: + import argparse + + parser = argparse.ArgumentParser( + prog="export-registered", + description="Convert a registered workflow to a workflow definition.", + ) + parser.add_argument("workflow_id") + parser.add_argument("--describe", action="store_true") + args = parser.parse_args(argv) + result = export_registered(args.workflow_id, describe=args.describe) + result["status"] = "ok" + return result diff --git a/docs/rips/taskmaestro_helper/definition.py b/docs/rips/taskmaestro_helper/definition.py new file mode 100644 index 0000000000..6c0fd4c232 --- /dev/null +++ b/docs/rips/taskmaestro_helper/definition.py @@ -0,0 +1,117 @@ +"""Workflow definition subcommands (``load`` and ``save``). + +The conversion between workflow folders and editable definitions is done by +``taskmaestro.definition``. This module adds ResInsight's schema annotations and +makes sure object references, which are only valid for one run, are never saved. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +from taskmaestro import definition as tm_definition +from taskmaestro.exceptions import DefinitionError + +from .refs import strip_refs +from .schema import ResInsightSchemaGenerator + +DEFINITION_FORMAT = tm_definition.DEFINITION_FORMAT +WORKFLOW_FILE = tm_definition.WORKFLOW_FILE +INPUT_FILE = tm_definition.INPUT_FILE + + +def load_definition(workflow_dir: Path, *, describe: bool = False) -> dict[str, Any]: + """Load ``workflow.yaml`` (and ``input.yaml``) from a workflow folder.""" + return tm_definition.load_definition( + workflow_dir, describe=describe, schema_generator=ResInsightSchemaGenerator + ) + + +def definition_from_workflow( + workflow: Any, + registered_id: str | None, + class_ids: dict[type[Any], str], + *, + describe: bool = False, +) -> dict[str, Any]: + """Convert an in-memory workflow to a definition.""" + return tm_definition.definition_from_workflow( + workflow, + registered_id, + class_ids, + describe=describe, + schema_generator=ResInsightSchemaGenerator, + ) + + +def describe_yaml( + workflow_yaml: Path, input_yaml: Path | None = None +) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: + """Describe a YAML workflow. Returns ``(description, error)``.""" + return tm_definition.describe_yaml( + workflow_yaml, input_yaml, ResInsightSchemaGenerator + ) + + +def write_definition( + definition: dict[str, Any], + out_dir: Path, + *, + workflow_only: bool = False, + backup: bool = False, +) -> dict[str, Any]: + """Write ``workflow.yaml`` (and ``input.yaml``) for a definition.""" + return tm_definition.write_definition( + definition, + out_dir, + workflow_only=workflow_only, + backup=backup, + value_filter=strip_refs, + ) + + +# --- Command-line entry points --- + + +def load_main(argv: list[str]) -> dict[str, Any]: + parser = argparse.ArgumentParser(prog="load", description="Load a workflow folder.") + parser.add_argument("workflow_dir", type=Path) + parser.add_argument("--describe", action="store_true") + args = parser.parse_args(argv) + result = load_definition(args.workflow_dir, describe=args.describe) + result["status"] = "ok" + return result + + +def save_main(argv: list[str], stdin_text: str) -> dict[str, Any]: + parser = argparse.ArgumentParser( + prog="save", + description="Write a workflow definition (read from stdin) as YAML.", + ) + parser.add_argument("--out-dir", type=Path, required=True) + parser.add_argument("--workflow-only", action="store_true") + parser.add_argument("--describe", action="store_true") + parser.add_argument("--backup", action="store_true") + args = parser.parse_args(argv) + try: + definition = json.loads(stdin_text) + except json.JSONDecodeError as exc: + raise DefinitionError(f"Invalid definition JSON: {exc}") from exc + if not isinstance(definition, dict): + raise DefinitionError("The definition must be a JSON object") + + result = write_definition( + definition, + args.out_dir, + workflow_only=args.workflow_only, + backup=args.backup, + ) + if args.describe: + input_path = Path(result["input_yaml"]) if result["input_yaml"] else None + result["describe"], result["describe_error"] = describe_yaml( + Path(result["workflow_yaml"]), input_path + ) + return result diff --git a/docs/rips/taskmaestro_helper/refs.py b/docs/rips/taskmaestro_helper/refs.py index afa2d8f030..830115919c 100644 --- a/docs/rips/taskmaestro_helper/refs.py +++ b/docs/rips/taskmaestro_helper/refs.py @@ -8,5 +8,33 @@ from __future__ import annotations +from typing import Any + # Marker used inside input.yaml to denote a ResInsight object reference. REF_MARKER = "__resinsight_ref__" + + +def _is_ref(value: Any) -> bool: + return isinstance(value, dict) and REF_MARKER in value + + +def _holds_ref(value: Any) -> bool: + """A reference, or a list with a reference among its items.""" + return _is_ref(value) or ( + isinstance(value, list) and any(_holds_ref(item) for item in value) + ) + + +def strip_refs(value: Any) -> Any: + """Remove ResInsight object references, which are only valid for one run. + + A list that holds a reference is removed as a whole: dropping only the + reference would shorten the list and change what the remaining items mean. + """ + if isinstance(value, dict): + return { + key: strip_refs(item) for key, item in value.items() if not _holds_ref(item) + } + if isinstance(value, list): + return [strip_refs(item) for item in value if not _holds_ref(item)] + return value diff --git a/docs/rips/taskmaestro_helper/run.py b/docs/rips/taskmaestro_helper/run.py index 8ed05d409e..dd1ed6d73d 100644 --- a/docs/rips/taskmaestro_helper/run.py +++ b/docs/rips/taskmaestro_helper/run.py @@ -15,19 +15,25 @@ import json import sys +from collections.abc import Callable from pathlib import Path -from typing import Any, Callable +from typing import Any import yaml from .refs import REF_MARKER - PROGRESS_PREFIX = "@@RI_WORKFLOW_EVENT@@" -def _emit_progress(run_id: str, task: str, state: str, error: str = "") -> None: +def _emit_progress( + run_id: str, task: str, state: str, error: str = "", item: str | None = None +) -> None: payload = {"event": "task_state", "run_id": run_id, "task": task, "state": state} + if item is not None: + # One item of a mapped task (`map:`); the task itself reports its own states + payload["event"] = "map_item_state" + payload["item"] = item if error: payload["error"] = error sys.stdout.write(PROGRESS_PREFIX + json.dumps(payload) + "\n") @@ -49,6 +55,15 @@ def on_task_complete(self, job: Any, task: Any, output: Any) -> None: def on_task_fail(self, job: Any, task: Any, error: Exception) -> None: _emit_progress(self.run_id, task.name, "failed", str(error)) + def on_map_item_start(self, job: Any, task: Any, key: str) -> None: + _emit_progress(self.run_id, task.name, "running", item=key) + + def on_map_item_complete(self, job: Any, task: Any, key: str, output: Any) -> None: + _emit_progress(self.run_id, task.name, "completed", item=key) + + def on_map_item_fail(self, job: Any, task: Any, key: str, error: Exception) -> None: + _emit_progress(self.run_id, task.name, "failed", str(error), item=key) + def _emit(event: str, **fields: Any) -> None: sys.stdout.write(json.dumps({"event": event, **fields}) + "\n") @@ -91,12 +106,31 @@ def resolve(ri_type: str, extras: dict[str, Any]) -> Any: return resolve +def _load_workflow(workflow_dir: Path | None, registered_id: str | None) -> Any: + if registered_id is not None: + from taskmaestro.discovery import get_registered_workflow + + return get_registered_workflow(registered_id) + + from taskmaestro.yaml_config import load_workflow_only + + assert workflow_dir is not None + workflow, _ = load_workflow_only(workflow_dir / "workflow.yaml") + return workflow + + def run_workflow( - workflow_dir: Path, input_path: Path, grpc_port: int, run_id: str = "" + workflow_dir: Path | None, + input_path: Path, + grpc_port: int, + run_id: str = "", + registered_id: str | None = None, ) -> int: - workflow_dir_str = str(workflow_dir) - if workflow_dir_str not in sys.path: - sys.path.insert(0, workflow_dir_str) + """Run a YAML workflow folder, or a registered workflow by its id.""" + if workflow_dir is not None: + workflow_dir_str = str(workflow_dir) + if workflow_dir_str not in sys.path: + sys.path.insert(0, workflow_dir_str) import rips @@ -114,15 +148,17 @@ def run_workflow( resolver = make_rips_resolver(instance) try: resolved = resolve_refs(raw_input, resolver) - except Exception as exc: + except Exception as exc: # noqa: BLE001 - reported as an error event _emit("error", message=f"Failed to resolve references: {exc}") return 1 from taskmaestro import EmptyConfig, ExecutionContext, Job, JobConfiguration, Runner - from taskmaestro.yaml_config import _load_workflow_only - workflow_yaml = workflow_dir / "workflow.yaml" - workflow, _ = _load_workflow_only(workflow_yaml) + try: + workflow = _load_workflow(workflow_dir, registered_id) + except Exception as exc: # noqa: BLE001 - reported as an error event + _emit("error", message=f"Failed to load workflow: {exc}") + return 1 job_config = JobConfiguration(resolved) job: Job[Any] = Job( @@ -132,7 +168,7 @@ def run_workflow( _emit("starting", workflow=workflow.name) try: result = Runner(hooks=[ProgressHook(run_id)]).run(job, ctx=ExecutionContext()) - except Exception as exc: + except Exception as exc: # noqa: BLE001 - reported as an error event _emit("error", message=str(exc)) return 1 @@ -146,11 +182,20 @@ def main(argv: list[str]) -> int: parser = argparse.ArgumentParser( prog="run", description="Run a taskmaestro workflow." ) - parser.add_argument("workflow_dir", type=Path) + parser.add_argument("workflow_dir", type=Path, nargs="?") + parser.add_argument( + "--registered", metavar="ID", help="Run a registered workflow by its id" + ) parser.add_argument("--input", type=Path, required=True) parser.add_argument("--grpc-port", type=int, required=True) parser.add_argument("--run-id", default="") args = parser.parse_args(argv) + if (args.workflow_dir is None) == (args.registered is None): + parser.error("specify either a workflow folder or --registered ") return run_workflow( - args.workflow_dir.resolve(), args.input.resolve(), args.grpc_port, args.run_id + args.workflow_dir.resolve() if args.workflow_dir is not None else None, + args.input.resolve(), + args.grpc_port, + args.run_id, + registered_id=args.registered, ) diff --git a/docs/rips/taskmaestro_helper/schema.py b/docs/rips/taskmaestro_helper/schema.py new file mode 100644 index 0000000000..eafb6bfd52 --- /dev/null +++ b/docs/rips/taskmaestro_helper/schema.py @@ -0,0 +1,56 @@ +"""JSON schemas for task input/output types, as consumed by ResInsight. + +taskmaestro's ``TaskSchemaGenerator`` marks opaque Python objects as +``{"not": {}, "x-taskmaestro-opaque": true, ...}``. This module adds +``x-ri-python-type``/``x-ri-python-bases`` to models and opaque types. The bases +list the class MRO so ResInsight can check port compatibility the same way +taskmaestro does (``issubclass``) without importing anything. +""" + +from __future__ import annotations + +import inspect +from typing import Any + +from pydantic import BaseModel +from taskmaestro.introspection import TaskSchemaGenerator, python_type_name +from taskmaestro.introspection import model_json_schema as tm_model_json_schema +from taskmaestro.introspection import task_descriptor as tm_task_descriptor + +_IGNORED_BASES = {object, BaseModel} + + +def python_bases(cls: type[Any]) -> list[str]: + """Return the dotted names of ``cls`` and its bases, most derived first.""" + names: list[str] = [] + for base in inspect.getmro(cls): + if base in _IGNORED_BASES or base.__module__ == "typing": + continue + name = python_type_name(base) + if name not in names: + names.append(name) + return names + + +def _type_annotations(cls: type[Any]) -> dict[str, Any]: + return { + "x-ri-python-type": python_type_name(cls), + "x-ri-python-bases": python_bases(cls), + } + + +class ResInsightSchemaGenerator(TaskSchemaGenerator): + """Schema generator that records Python class bases for port compatibility.""" + + @classmethod + def type_annotations(cls, python_type: type[Any]) -> dict[str, Any]: + return _type_annotations(python_type) + + +def model_json_schema(model: type[Any]) -> dict[str, Any]: + return tm_model_json_schema(model, ResInsightSchemaGenerator) + + +def task_descriptor(task_id: str, cls: type[Any]) -> dict[str, Any]: + """Describe a task class for the ResInsight task palette and port typing.""" + return tm_task_descriptor(task_id, cls, ResInsightSchemaGenerator) diff --git a/docs/rips/tests/test_cases.py b/docs/rips/tests/test_cases.py index e367d79996..a77a185ffe 100644 --- a/docs/rips/tests/test_cases.py +++ b/docs/rips/tests/test_cases.py @@ -470,3 +470,102 @@ def test_create_well_bore_stability_plot_is_geomech_only( case = rips_instance.project.load_case(case_path) assert not hasattr(case, "create_well_bore_stability_plot") assert hasattr(rips.GeoMechCase, "create_well_bore_stability_plot") + + +def _write_small_vtk_case(folder): + # Two hexahedron cells sharing a face, with two time steps and one cell property + points = [(x, y, z) for x in range(3) for y in range(2) for z in range(2)] + point_index = {p: i for i, p in enumerate(points)} + + def hex_connectivity(x0): + corners = [ + (x0, 0, 0), + (x0 + 1, 0, 0), + (x0 + 1, 1, 0), + (x0, 1, 0), + (x0, 0, 1), + (x0 + 1, 0, 1), + (x0 + 1, 1, 1), + (x0, 1, 1), + ] + return [point_index[c] for c in corners] + + connectivity = hex_connectivity(0) + hex_connectivity(1) + coordinates = "\n".join(f"{x} {y} {z}" for x, y, z in points) + displacements = "\n".join("0 0 0" for _ in points) + offsets = " ".join(str(8 * (i + 1)) for i in range(2)) + types = " ".join("12" for _ in range(2)) + + property_values = [[0.1, 0.2], [0.3, 0.4]] + for step, values in enumerate(property_values): + vtu = f""" + + + + + +{coordinates} + + + + {" ".join(map(str, connectivity))} + {offsets} + {types} + + + +{displacements} + + + + {" ".join(map(str, values))} + + + + +""" + with open(os.path.join(folder, f"step_{step}.vtu"), "w") as f: + f.write(vtu) + + pvd = """ + + + + + + +""" + pvd_path = os.path.join(folder, "small.pvd") + with open(pvd_path, "w") as f: + f.write(pvd) + return pvd_path, property_values + + +def test_geomech_case_results(rips_instance, initialize_test, tmp_path): + pvd_path, property_values = _write_small_vtk_case(str(tmp_path)) + case = rips_instance.project.load_case(pvd_path) + assert isinstance(case, rips.GeoMechCase) + + assert len(case.time_step_strings()) == 2 + + positions = rips.GeoMechResultPosition + names = case.result_names(positions.ELEMENT) + assert "porosity" in names + + for time_step, expected in enumerate(property_values): + values = case.result_values( + "porosity", time_step=time_step, result_position=positions.ELEMENT + ) + assert values == pytest.approx(expected) + + # The last frame is the default, and is the only frame in this case + assert case.result_values( + "porosity", time_step=1, result_position=positions.ELEMENT, frame_index=0 + ) == pytest.approx(property_values[1]) + + with pytest.raises(rips.RipsError): + case.result_values("porosity", time_step=2, result_position=positions.ELEMENT) + with pytest.raises(rips.RipsError): + case.result_values( + "porosity", time_step=0, result_position=positions.ELEMENT, frame_index=1 + ) diff --git a/docs/rips/tests/test_data/taskmaestro_workflows/tm_helper_pipeline.py b/docs/rips/tests/test_data/taskmaestro_workflows/tm_helper_pipeline.py new file mode 100644 index 0000000000..4bf325c88b --- /dev/null +++ b/docs/rips/tests/test_data/taskmaestro_workflows/tm_helper_pipeline.py @@ -0,0 +1,79 @@ +"""Small taskmaestro tasks used by the taskmaestro helper tests.""" + +from __future__ import annotations + +from pydantic import BaseModel +from taskmaestro import EmptyConfig, ExecutionContext, Task, Workflow + + +class Number(BaseModel): + """A single number.""" + + value: float + + +class StartInput(BaseModel): + """Configured start value.""" + + value: float + + +class SpecialNumber(Number): + """A number subclass, compatible with Number.""" + + +class Start(Task[StartInput, Number]): + """Produce the configured value.""" + + name = "start" + + def run(self, input: StartInput, ctx: ExecutionContext) -> Number: + return Number(value=input.value) + + +class Constant(Task[EmptyConfig, Number]): + """Produce the number one.""" + + name = "constant" + + def run(self, input: EmptyConfig, ctx: ExecutionContext) -> Number: + return Number(value=1.0) + + +class Double(Task[Number, Number]): + """Double a number.""" + + name = "double" + + def run(self, input: Number, ctx: ExecutionContext) -> Number: + return Number(value=2 * input.value) + + +class Special(Task[Number, SpecialNumber]): + """Produce a SpecialNumber.""" + + name = "special" + + def run(self, input: Number, ctx: ExecutionContext) -> SpecialNumber: + return SpecialNumber(value=input.value) + + +registered_workflow = ( + Workflow.builder("registered_flow") + .add_task(Start, config_fields=["value"]) + .add_task(Double, depends_on=Start) + .build() +) + + +def _local_workflow() -> Workflow: + class LocalTask(Task[EmptyConfig, Number]): + name = "local" + + def run(self, input: EmptyConfig, ctx: ExecutionContext) -> Number: + return Number(value=0.0) + + return Workflow.builder("local_flow").add_task(LocalTask).build() + + +local_workflow = _local_workflow() diff --git a/docs/rips/tests/test_taskmaestro_helper.py b/docs/rips/tests/test_taskmaestro_helper.py index 81b2be4f11..4252074631 100644 --- a/docs/rips/tests/test_taskmaestro_helper.py +++ b/docs/rips/tests/test_taskmaestro_helper.py @@ -14,6 +14,22 @@ pytest.importorskip("taskmaestro") +def test_strip_refs_drops_references_and_lists_holding_them() -> None: + from rips.taskmaestro_helper.refs import strip_refs + + ref = {"__resinsight_ref__": "EclipseCase", "case_id": 0} + inputs = { + "start": {"case": ref, "count": 1}, + "sum": {"cases": [ref, ref], "mixed": [1.0, ref], "values": [1.0, 2.0]}, + "only_refs": {"case": ref}, + } + assert strip_refs(inputs) == { + "start": {"count": 1}, + "sum": {"values": [1.0, 2.0]}, + "only_refs": {}, + } + + def test_resolve_refs_substitutes_object_model_value() -> None: from rips.taskmaestro_helper.run import resolve_refs @@ -73,9 +89,8 @@ def test_progress_hook_reports_named_task_instances( def test_progress_hook_follows_runner_instance_names( capsys: pytest.CaptureFixture[str], ) -> None: - from taskmaestro import EmptyConfig, ExecutionContext, Job, Runner, Task, Workflow - from rips.taskmaestro_helper.run import PROGRESS_PREFIX, ProgressHook + from taskmaestro import EmptyConfig, ExecutionContext, Job, Runner, Task, Workflow class Step(Task[EmptyConfig, EmptyConfig]): name = "step" @@ -108,9 +123,8 @@ def run(self, input: EmptyConfig, ctx: ExecutionContext) -> EmptyConfig: def test_progress_hook_reports_runner_failure( capsys: pytest.CaptureFixture[str], ) -> None: - from taskmaestro import EmptyConfig, ExecutionContext, Job, Runner, Task, Workflow - from rips.taskmaestro_helper.run import PROGRESS_PREFIX, ProgressHook + from taskmaestro import EmptyConfig, ExecutionContext, Job, Runner, Task, Workflow class Fails(Task[EmptyConfig, EmptyConfig]): name = "fails" @@ -135,11 +149,67 @@ def run(self, input: EmptyConfig, ctx: ExecutionContext) -> EmptyConfig: assert events[-1]["error"] == "bad range" +def test_progress_hook_reports_mapped_items( + capsys: pytest.CaptureFixture[str], +) -> None: + from pydantic import BaseModel + from rips.taskmaestro_helper.run import PROGRESS_PREFIX, ProgressHook + from taskmaestro import EmptyConfig, ExecutionContext, Job, Runner, Task, Workflow + from taskmaestro.job import JobConfiguration + from taskmaestro.mapping import TaskMap + + class Item(BaseModel): + key: str + value: int + + class Out(BaseModel): + value: int + + class Check(Task[Item, Out]): + name = "check" + + def run(self, input: Item, ctx: ExecutionContext) -> Out: + if input.value < 0: + raise ValueError("negative") + return Out(value=input.value) + + workflow = ( + Workflow.builder("mapped") + .add_task( + Check, + name="each", + mapped_over=TaskMap( + over="items", key_as="key", value_as="value", error_mode="collect_all" + ), + ) + .build() + ) + config = JobConfiguration({"each": {"items": {"a": 1, "b": -1}}}) + Runner(hooks=[ProgressHook("run-map")]).run( + Job(workflow, EmptyConfig(), job_configuration=config), ctx=ExecutionContext() + ) + events = [ + json.loads(line.removeprefix(PROGRESS_PREFIX)) + for line in capsys.readouterr().out.splitlines() + if line.startswith(PROGRESS_PREFIX) + ] + assert [ + (event["event"], event.get("item"), event["state"]) for event in events + ] == [ + ("task_state", None, "running"), + ("map_item_state", "a", "running"), + ("map_item_state", "a", "completed"), + ("map_item_state", "b", "running"), + ("map_item_state", "b", "failed"), + ("task_state", None, "failed"), + ] + assert events[4]["error"] == "negative" + + def test_run_helper_emits_progress_without_grpc( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] ) -> None: import rips - from rips.taskmaestro_helper.run import PROGRESS_PREFIX, run_workflow (tmp_path / "progress_pipeline.py").write_text( diff --git a/docs/rips/tests/test_taskmaestro_helper_definition.py b/docs/rips/tests/test_taskmaestro_helper_definition.py new file mode 100644 index 0000000000..b6363f2f15 --- /dev/null +++ b/docs/rips/tests/test_taskmaestro_helper_definition.py @@ -0,0 +1,385 @@ +"""Tests for the workflow definition subcommands of rips.taskmaestro_helper. + +Run with ``pytest --noconftest``: these tests do not need a running ResInsight. +""" + +from __future__ import annotations + +import json +import shutil +import textwrap +from importlib.metadata import EntryPoint +from pathlib import Path +from typing import Any + +import pytest + +pytest.importorskip("taskmaestro") + +from rips.taskmaestro_helper import __main__ as helper_main +from rips.taskmaestro_helper.catalog import build_catalog, export_registered +from rips.taskmaestro_helper.definition import ( + describe_yaml, + load_definition, + write_definition, +) + +TEST_DATA = Path(__file__).parent / "test_data" / "taskmaestro_workflows" +PIPELINE = "tm_helper_pipeline" + +STRUCTURAL_KEYS = ( + "name", + "result_task", + "nodes", + "edges", + "inputs", + "editable", + "header_comment", + "passthrough", + "task_types", +) + + +@pytest.fixture(autouse=True) +def pipeline_on_path(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.syspath_prepend(str(TEST_DATA)) + + +def _fake_entry_points( + monkeypatch: pytest.MonkeyPatch, + tasks: dict[str, str], + workflows: dict[str, str] | None = None, +) -> None: + from taskmaestro import discovery + + groups = { + discovery.TASK_ENTRY_POINT_GROUP: [ + EntryPoint(name, value, discovery.TASK_ENTRY_POINT_GROUP) + for name, value in tasks.items() + ], + discovery.WORKFLOW_ENTRY_POINT_GROUP: [ + EntryPoint(name, value, discovery.WORKFLOW_ENTRY_POINT_GROUP) + for name, value in (workflows or {}).items() + ], + } + + def entry_points(*, group: str) -> list[EntryPoint]: + return groups.get(group, []) + + monkeypatch.setattr(discovery, "entry_points", entry_points) + + +@pytest.fixture +def plugins(monkeypatch: pytest.MonkeyPatch) -> None: + _fake_entry_points( + monkeypatch, + tasks={ + "test.start": f"{PIPELINE}:Start", + "test.double": f"{PIPELINE}:Double", + "test.constant": f"{PIPELINE}:Constant", + "test.broken": "tm_helper_missing_module:Broken", + }, + workflows={ + "test.flow": f"{PIPELINE}:registered_workflow", + "test.local": f"{PIPELINE}:local_workflow", + }, + ) + + +def _write_workflow(directory: Path, workflow: str, inputs: str | None = None) -> Path: + directory.mkdir(parents=True, exist_ok=True) + (directory / "workflow.yaml").write_text(textwrap.dedent(workflow)) + if inputs is not None: + (directory / "input.yaml").write_text(textwrap.dedent(inputs)) + return directory + + +def _structure(definition: dict[str, Any]) -> dict[str, Any]: + return {key: definition[key] for key in STRUCTURAL_KEYS} + + +def _describe_summary(description: dict[str, Any]) -> dict[str, Any]: + return { + "result_task": description["result_task"], + "tasks": [ + (task["name"], task["depends_on"], task["config_fields"]) + for task in description["tasks"] + ], + } + + +def _round_trip(source: Path, out_dir: Path) -> tuple[dict[str, Any], dict[str, Any]]: + first = load_definition(source, describe=True) + assert first["editable"], first["readonly_reasons"] + assert first["describe_error"] is None, first["describe_error"] + write_definition(first, out_dir) + second = load_definition(out_dir, describe=True) + assert second["describe_error"] is None, second["describe_error"] + assert _structure(second) == _structure(first) + assert _describe_summary(second["describe"]) == _describe_summary(first["describe"]) + third_dir = out_dir.parent / (out_dir.name + "_again") + write_definition(second, third_dir) + assert (third_dir / "workflow.yaml").read_text() == ( + out_dir / "workflow.yaml" + ).read_text() + return first, second + + +# --- Catalog --- + + +def test_catalog_reports_broken_plugin_and_loads_the_rest(plugins: None) -> None: + catalog = build_catalog() + ids = [task["id"] for task in catalog["tasks"]] + assert ids == ["test.constant", "test.double", "test.start"] + assert [error["id"] for error in catalog["errors"]] == ["test.broken"] + + start = next(task for task in catalog["tasks"] if task["id"] == "test.start") + assert start["name"] == "start" + assert start["description"] == "Produce the configured value." + assert start["python_type"] == f"{PIPELINE}.Start" + assert start["input_schema"]["properties"]["value"]["type"] == "number" + + workflows = {workflow["id"]: workflow for workflow in catalog["workflows"]} + assert set(workflows) == {"test.flow", "test.local"} + assert workflows["test.flow"]["editable"] + assert not workflows["test.local"]["editable"] + assert workflows["test.flow"]["describe"]["result_task"] == "double" + + +def test_catalog_without_workflows(plugins: None) -> None: + assert build_catalog(include_workflows=False)["workflows"] == [] + + +def test_catalog_cli_prints_one_json_document( + plugins: None, capsys: pytest.CaptureFixture[str] +) -> None: + assert helper_main.main(["catalog", "--no-workflows"]) == 0 + out = capsys.readouterr().out + document = json.loads(out) + assert document["status"] == "ok" + assert len(document["tasks"]) == 3 + + +def test_cli_reports_missing_taskmaestro( + monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +) -> None: + def missing(argv: list[str]) -> dict[str, Any]: + raise ModuleNotFoundError("No module named 'taskmaestro'", name="taskmaestro") + + monkeypatch.setitem(helper_main.JSON_COMMANDS, "catalog", missing) + assert helper_main.main(["catalog"]) == 2 + error = json.loads(capsys.readouterr().out)["error"] + assert error["message"].startswith("taskmaestro>=0.4,<0.5 is required") + + +def test_cli_failure_prints_invalid_document( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + assert helper_main.main(["load", str(tmp_path / "missing")]) == 2 + document = json.loads(capsys.readouterr().out) + assert document["status"] == "invalid" + assert "Cannot read" in document["error"]["message"] + + +# --- Schema --- + + +def test_schema_records_python_bases() -> None: + from rips.taskmaestro_helper.schema import task_descriptor + from tm_helper_pipeline import Special + + descriptor = task_descriptor("test.special", Special) + output = descriptor["output_schema"] + assert output["x-ri-python-type"] == f"{PIPELINE}.SpecialNumber" + assert output["x-ri-python-bases"] == [ + f"{PIPELINE}.SpecialNumber", + f"{PIPELINE}.Number", + ] + + +def test_schema_perforation_output_is_not_a_well_path() -> None: + pytest.importorskip("taskmaestro_resinsight") + from rips.taskmaestro_helper.schema import task_descriptor + from taskmaestro import get_registered_task + + add_perforation = task_descriptor( + "resinsight.add_perforation", get_registered_task("resinsight.add_perforation") + ) + output_bases = add_perforation["output_schema"]["x-ri-python-bases"] + well_path = add_perforation["input_schema"]["$defs"]["WellPath"] + assert well_path["x-ri-python-type"] == "taskmaestro_resinsight.models.WellPath" + assert well_path["x-ri-python-type"] not in output_bases + opaque = well_path["properties"]["value"] + assert opaque["x-taskmaestro-opaque"] is True + assert "rips.generated.generated_classes.WellPath" in opaque["x-ri-python-bases"] + + +# --- Round trips --- + + +def test_round_trip_whole_output_dependency(tmp_path: Path) -> None: + source = _write_workflow( + tmp_path / "src", + f"""\ + workflow: + name: whole + tasks: + - task: {PIPELINE}.Start + config_fields: [value] + - task: {PIPELINE}.Double + depends_on: {PIPELINE}.Start + """, + "start:\n value: 2.5\n", + ) + first, _ = _round_trip(source, tmp_path / "out") + assert first["edges"] == [ + {"from": "start", "output": None, "to": "double", "input": None} + ] + assert first["inputs"] == {"start": {"value": 2.5}, "double": {}} + text = (tmp_path / "out" / "workflow.yaml").read_text() + assert "name: start" in text and "depends_on: start" in text + assert "result_task" not in text + + +# --- References --- + + +def test_saved_input_drops_resinsight_refs(tmp_path: Path) -> None: + source = _write_workflow( + tmp_path / "src", + f"""\ + workflow: + name: with_ref + tasks: + - task: {PIPELINE}.Start + config_fields: [value] + """, + "start:\n value: {__resinsight_ref__: EclipseCase, case_id: 0}\n", + ) + write_definition(load_definition(source), tmp_path / "out") + assert (tmp_path / "out" / "input.yaml").read_text() == "start: {}\n" + + +def test_save_cli_reads_definition_from_stdin( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +) -> None: + import io + + definition = load_definition( + _write_workflow( + tmp_path / "src", + f"""\ + workflow: + name: via_cli + tasks: + - task: {PIPELINE}.Start + config_fields: [value] + """, + ) + ) + monkeypatch.setattr("sys.stdin", io.StringIO(json.dumps(definition))) + out_dir = tmp_path / "out" + assert helper_main.main(["save", "--out-dir", str(out_dir), "--describe"]) == 0 + result = json.loads(capsys.readouterr().out) + assert result["status"] == "ok" + assert result["describe"]["tasks"][0]["missing_config_fields"] == ["value"] + assert (out_dir / "input.yaml").read_text() == "start: {}\n" + + +# --- Registered workflows --- + + +def test_export_registered(plugins: None) -> None: + definition = export_registered("test.flow", describe=True) + assert definition["editable"] + assert definition["source"]["registered_id"] == "test.flow" + assert definition["nodes"] == [ + {"name": "start", "task": "test.start", "config_fields": ["value"]}, + {"name": "double", "task": "test.double", "config_fields": []}, + ] + assert definition["edges"] == [ + {"from": "start", "output": None, "to": "double", "input": None} + ] + assert definition["result_task"] is None + assert definition["describe"]["result_task"] == "double" + + +def test_export_registered_local_class_is_read_only(plugins: None) -> None: + definition = export_registered("test.local") + assert not definition["editable"] + assert "cannot be referenced" in definition["readonly_reasons"][0] + + +def test_exported_registered_workflow_saves_and_validates( + plugins: None, tmp_path: Path +) -> None: + definition = export_registered("test.flow") + definition["inputs"] = {"start": {"value": 3}} + write_definition(definition, tmp_path / "out") + description, error = describe_yaml( + tmp_path / "out" / "workflow.yaml", tmp_path / "out" / "input.yaml" + ) + assert error is None + assert description is not None and description["result_task"] == "double" + + +def test_run_registered_workflow( + plugins: None, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + import rips + from rips.taskmaestro_helper.run import PROGRESS_PREFIX, main + + input_path = tmp_path / "input.yaml" + input_path.write_text("start:\n value: 2\n") + monkeypatch.setattr(rips.Instance, "find", lambda **kwargs: object()) + + assert ( + main( + [ + "--registered", + "test.flow", + "--input", + str(input_path), + "--grpc-port", + "1234", + "--run-id", + "rid", + ] + ) + == 0 + ) + events = [ + json.loads(line.removeprefix(PROGRESS_PREFIX)) + for line in capsys.readouterr().out.splitlines() + if line.startswith(PROGRESS_PREFIX) + ] + assert [(event["task"], event["state"]) for event in events] == [ + ("start", "running"), + ("start", "completed"), + ("double", "running"), + ("double", "completed"), + ] + + +def test_run_requires_folder_or_registered_id(tmp_path: Path) -> None: + from rips.taskmaestro_helper.run import main + + with pytest.raises(SystemExit): + main(["--input", str(tmp_path / "input.yaml"), "--grpc-port", "1"]) + + +# --- The taskmaestro-resinsight examples --- + + +@pytest.mark.parametrize("name", ["resinsight_completions", "resinsight_load_data"]) +def test_example_workflows_round_trip(tmp_path: Path, name: str) -> None: + pytest.importorskip("taskmaestro_resinsight") + source = tmp_path / name + shutil.copytree(TEST_DATA / name, source) + first, _ = _round_trip(source, tmp_path / "out") + assert first["warnings"] == [] + assert all(node["task"].startswith("resinsight.") for node in first["nodes"])