diff --git a/openhcs/core/orchestrator/execution_result.py b/openhcs/core/orchestrator/execution_result.py index c516da01b6..b5c4af3471 100644 --- a/openhcs/core/orchestrator/execution_result.py +++ b/openhcs/core/orchestrator/execution_result.py @@ -8,6 +8,7 @@ import copyreg from dataclasses import dataclass, field from enum import Enum +from pathlib import Path from types import MappingProxyType from typing import Mapping, Optional @@ -56,6 +57,7 @@ class RuntimeContextObservation: context_key: str records: tuple[StoredRuntimeValue, ...] + runtime_export_paths: tuple[Path, ...] = field(default_factory=tuple) @dataclass(frozen=True) diff --git a/openhcs/core/orchestrator/worker_execution.py b/openhcs/core/orchestrator/worker_execution.py index f1a171a088..2c2beaff03 100644 --- a/openhcs/core/orchestrator/worker_execution.py +++ b/openhcs/core/orchestrator/worker_execution.py @@ -52,6 +52,7 @@ from openhcs.core.steps.abstract import AbstractStep, StepExecutionObservation from openhcs.core.steps.function_artifact_materialization import ( observed_materialized_artifact_locations_by_address, + observed_runtime_export_artifact_output_paths, ) from openhcs.utils.environment import OpenHCSProcessEnvironment @@ -796,22 +797,39 @@ def _execute_axis_with_sequential_combinations( lane_context, cancellation=cancellation, ) + observed_records = runtime_store.observed_values_after( + execution_observation_cursor + ) + runtime_export_paths = tuple( + dict.fromkeys( + path + for step_plan in frozen_context.step_plans.values() + if step_plan.owns_runtime_outputs + for path in observed_runtime_export_artifact_output_paths( + step_plan, + frozen_context, + observed_records, + ) + ) + ) finally: # This cache is context-local even when lanes share a process. # Runtime observations remain owned by the value store below. frozen_context.release_execution_image_cache() if release_axis_resources: _release_runtime_resources((frozen_context,), owner=f"axis {axis_id}") - if runtime_observation_mode.collects_records: + retained_records = ( + observed_records if runtime_observation_mode.collects_records else () + ) + if retained_records or runtime_export_paths: runtime_observations.append( RuntimeContextObservation( context_key=context_key, - records=runtime_store.observed_values_after( - execution_observation_cursor - ), + records=retained_records, + runtime_export_paths=runtime_export_paths, ) ) - elif runtime_observation_mode.releases_worker_records: + if runtime_observation_mode.releases_worker_records: frozen_context.runtime_value_store.clear() if not result.is_success(): diff --git a/openhcs/core/runtime_exports.py b/openhcs/core/runtime_exports.py index 98d703391b..1cc9d50317 100644 --- a/openhcs/core/runtime_exports.py +++ b/openhcs/core/runtime_exports.py @@ -29,6 +29,7 @@ if TYPE_CHECKING: from openhcs.core.context.processing_context import ProcessingContext + from openhcs.core.orchestrator.execution_result import RuntimeExecutionObservation @dataclass(frozen=True, slots=True) @@ -132,6 +133,22 @@ def from_execution_contexts( ) ) + @classmethod + def from_runtime_observations( + cls, + observations: Sequence[RuntimeExecutionObservation], + ) -> RuntimeExportObservation: + """Build exports from worker-projected paths without retaining values.""" + + return cls.from_output_paths( + tuple( + path + for observation in observations + for context in observation.contexts + for path in context.runtime_export_paths + ) + ) + @classmethod def from_output_root( cls, diff --git a/openhcs/core/steps/function_artifact_materialization.py b/openhcs/core/steps/function_artifact_materialization.py index 6bfc6e5b63..71e1ddf9b0 100644 --- a/openhcs/core/steps/function_artifact_materialization.py +++ b/openhcs/core/steps/function_artifact_materialization.py @@ -1322,6 +1322,27 @@ def observed_materialized_artifact_output_paths( ) +def observed_runtime_export_artifact_output_paths( + plan: CompiledStepPlan, + context: "ProcessingContext", + records: tuple[StoredRuntimeValue, ...], +) -> tuple[Path, ...]: + """Derive exact pipeline-declared exports from worker-observed records.""" + + if not plan.runtime_artifact_materialization.has_persistent_target: + return () + return tuple( + Path(output.path) + for materialization in runtime_artifact_materializations_from_records( + plan, + context, + records, + ) + if materialization.spec.participates_in_runtime_export_observation() + for output in materialization.outputs(plan, context) + ) + + def observed_materialized_artifact_locations_by_address( plan: CompiledStepPlan, context: "ProcessingContext", diff --git a/openhcs/runtime/zmq_execution_server.py b/openhcs/runtime/zmq_execution_server.py index c1b80db0a0..5b8e05d3da 100644 --- a/openhcs/runtime/zmq_execution_server.py +++ b/openhcs/runtime/zmq_execution_server.py @@ -766,15 +766,25 @@ def _export_runtime_observation( request_context.auxiliary_params.runtime_observation_export_scope is ZMQRuntimeObservationExportScope.OUTCOMES ): + observed_exports = RuntimeExportObservation.from_runtime_observations( + tuple( + result.runtime_observation for result in execution_results.values() + ) + ) + if execution_bundle.requires_parent_runtime_observation: + parent_exports = RuntimeExportObservation.from_execution_contexts( + execution_bundle.runtime_contexts + ) + observed_exports = RuntimeExportObservation.from_output_paths( + (*observed_exports.output_files, *parent_exports.output_files) + ) export = ZMQRuntimeExecutionOutcomeExport.from_execution( compiled_axis_ids=execution_bundle.runtime_contexts, execution_results=execution_results, output_roots=output_roots, server_environment=self._server_environment, execution_id=request_context.execution_id, - exports=RuntimeExportObservation.from_execution_contexts( - execution_bundle.runtime_contexts - ), + exports=observed_exports, ) else: export = ZMQRuntimeExecutionObservationExport.from_execution( diff --git a/tests/unit/test_compiled_execution.py b/tests/unit/test_compiled_execution.py index f856a6a244..e795860a76 100644 --- a/tests/unit/test_compiled_execution.py +++ b/tests/unit/test_compiled_execution.py @@ -356,6 +356,7 @@ def execute_current_run( 0: SimpleNamespace( execution_scope=FunctionStepExecutionScope.AXIS, device_assignment=FrameworkDeviceAssignment(), + owns_runtime_outputs=False, ) }, ) diff --git a/tests/unit/test_zmq_execution_server.py b/tests/unit/test_zmq_execution_server.py index 080ec4c739..8551094da0 100644 --- a/tests/unit/test_zmq_execution_server.py +++ b/tests/unit/test_zmq_execution_server.py @@ -17,6 +17,8 @@ from openhcs.core.execution_state import ExecutionOutputPlateSummary from openhcs.core.orchestrator.execution_result import ( ExecutionResult, + RuntimeContextObservation, + RuntimeExecutionObservation, RuntimeObservationMode, ) from openhcs.core.progress import ( @@ -160,7 +162,10 @@ def test_server_exports_outcomes_without_projecting_compiled_values( ), ) compilation = SimpleNamespace( - execution_bundle=SimpleNamespace(runtime_contexts={}), + execution_bundle=SimpleNamespace( + runtime_contexts={}, + requires_parent_runtime_observation=False, + ), output_plate=SimpleNamespace(output_plate_root=tmp_path), ) monkeypatch.setattr( @@ -168,17 +173,33 @@ def test_server_exports_outcomes_without_projecting_compiled_values( lambda contexts, root: (root,), ) + declared_output = tmp_path / "exports" / "Saved.tiff" + declared_output.parent.mkdir() + declared_output.write_bytes(b"image evidence") server._export_runtime_observation( request_context=request_context, compilation=compilation, - execution_results={"A01": ExecutionResult.success("A01")}, + execution_results={ + "A01": ExecutionResult.success( + "A01", + runtime_observation=RuntimeExecutionObservation( + contexts=( + RuntimeContextObservation( + "context", + (), + runtime_export_paths=(declared_output,), + ), + ) + ), + ) + }, ) export = ZMQRuntimeExecutionOutcomeExport.read(export_path) assert export.successful_axis_count == 1 assert export.execution_id == record.execution_id assert export.exports is not None - assert export.exports.output_files == () + assert export.exports.output_files == (declared_output,) assert record.get_extra("runtime_observation_export_path") == str(export_path) assert record.get_extra("runtime_observation_export_scope") == "outcomes" diff --git a/tests/unit/test_zmq_runtime_outcome_export.py b/tests/unit/test_zmq_runtime_outcome_export.py index 5ea213a0a3..f358c89ca1 100644 --- a/tests/unit/test_zmq_runtime_outcome_export.py +++ b/tests/unit/test_zmq_runtime_outcome_export.py @@ -28,7 +28,13 @@ def test_outcome_export_round_trip_excludes_runtime_values(tmp_path: Path) -> No "A01": ExecutionResult.success( "A01", runtime_observation=RuntimeExecutionObservation( - contexts=(RuntimeContextObservation("context", (runtime_value,)),) + contexts=( + RuntimeContextObservation( + "context", + (runtime_value,), + runtime_export_paths=(declared_output,), + ), + ) ), ), "B01": ExecutionResult.error(