tracdap-runtime 0.6.3__py3-none-any.whl → 0.6.5__py3-none-any.whl

tracdap/rt/_exec/functions.py

@@ -23,6 +23,7 @@ import tracdap.rt.api as _api
 import tracdap.rt.config as _config
 import tracdap.rt.exceptions as _ex
 import tracdap.rt._exec.context as _ctx
+import tracdap.rt._exec.graph_builder as _graph
 import tracdap.rt._impl.config_parser as _cfg_p  # noqa
 import tracdap.rt._impl.type_system as _types  # noqa
 import tracdap.rt._impl.data as _data  # noqa
@@ -59,6 +60,13 @@ class NodeContext:
         pass
 
 
+class NodeCallback:
+
+    @abc.abstractmethod
+    def send_graph_updates(self, new_nodes: tp.Dict[NodeId, Node], new_deps: tp.Dict[NodeId, tp.List[Dependency]]):
+        pass
+
+
 # Helper functions to access the node context (in case the NodeContext interface needs to change)
 
 def _ctx_lookup(node_id: NodeId[_T], ctx: NodeContext) -> _T:
@@ -89,8 +97,15 @@ class NodeFunction(tp.Generic[_T]):
         :py:class:`NodeContext <NodeContext>`
     """
 
-    def __call__(self, ctx: NodeContext) -> _T:
-        return self._execute(ctx)
+    def __init__(self):
+        self.node_callback: tp.Optional[NodeCallback] = None
+
+    def __call__(self, ctx: NodeContext, callback: NodeCallback = None) -> _T:
+        try:
+            self.node_callback = callback
+            return self._execute(ctx)
+        finally:
+            self.node_callback = None
 
     @abc.abstractmethod
     def _execute(self, ctx: NodeContext) -> _T:
@@ -105,6 +120,7 @@ class NodeFunction(tp.Generic[_T]):
 class NoopFunc(NodeFunction[None]):
 
     def __init__(self, node: NoopNode):
+        super().__init__()
         self.node = node
 
     def _execute(self, _: NodeContext) -> None:
@@ -114,6 +130,7 @@ class NoopFunc(NodeFunction[None]):
 class StaticValueFunc(NodeFunction[_T]):
 
     def __init__(self, node: StaticValueNode[_T]):
+        super().__init__()
         self.node = node
 
     def _execute(self, ctx: NodeContext) -> _T:
@@ -123,6 +140,7 @@ class StaticValueFunc(NodeFunction[_T]):
 class IdentityFunc(NodeFunction[_T]):
 
     def __init__(self, node: IdentityNode[_T]):
+        super().__init__()
         self.node = node
 
     def _execute(self, ctx: NodeContext) -> _T:
@@ -138,6 +156,7 @@ class _ContextPushPopFunc(NodeFunction[Bundle[tp.Any]], abc.ABC):
     _POP = False
 
     def __init__(self, node: tp.Union[ContextPushNode, ContextPopNode], direction: bool):
+        super().__init__()
         self.node = node
         self.direction = direction
 
@@ -176,6 +195,7 @@ class ContextPopFunc(_ContextPushPopFunc):
 class KeyedItemFunc(NodeFunction[_T]):
 
     def __init__(self, node: KeyedItemNode[_T]):
+        super().__init__()
         self.node = node
 
     def _execute(self, ctx: NodeContext) -> _T:
@@ -184,9 +204,20 @@ class KeyedItemFunc(NodeFunction[_T]):
         return src_item
 
 
+class RuntimeOutputsFunc(NodeFunction[JobOutputs]):
+
+    def __init__(self, node: RuntimeOutputsNode):
+        super().__init__()
+        self.node = node
+
+    def _execute(self, ctx: NodeContext) -> JobOutputs:
+        return self.node.outputs
+
+
 class BuildJobResultFunc(NodeFunction[_config.JobResult]):
 
     def __init__(self, node: BuildJobResultNode):
+        super().__init__()
         self.node = node
 
     def _execute(self, ctx: NodeContext) -> _config.JobResult:
@@ -197,20 +228,33 @@ class BuildJobResultFunc(NodeFunction[_config.JobResult]):
 
         # TODO: Handle individual failed results
 
-        for obj_id, node_id in self.node.objects.items():
+        for obj_id, node_id in self.node.outputs.objects.items():
             obj_def = _ctx_lookup(node_id, ctx)
             job_result.results[obj_id] = obj_def
 
-        for bundle_id in self.node.bundles:
+        for bundle_id in self.node.outputs.bundles:
             bundle = _ctx_lookup(bundle_id, ctx)
             job_result.results.update(bundle.items())
 
+        if self.node.runtime_outputs is not None:
+
+            runtime_outputs = _ctx_lookup(self.node.runtime_outputs, ctx)
+
+            for obj_id, node_id in runtime_outputs.objects.items():
+                obj_def = _ctx_lookup(node_id, ctx)
+                job_result.results[obj_id] = obj_def
+
+            for bundle_id in runtime_outputs.bundles:
+                bundle = _ctx_lookup(bundle_id, ctx)
+                job_result.results.update(bundle.items())
+
         return job_result
 
 
 class SaveJobResultFunc(NodeFunction[None]):
 
     def __init__(self, node: SaveJobResultNode):
+        super().__init__()
         self.node = node
 
     def _execute(self, ctx: NodeContext) -> None:
@@ -241,6 +285,7 @@ class SaveJobResultFunc(NodeFunction[None]):
 class DataViewFunc(NodeFunction[_data.DataView]):
 
     def __init__(self, node: DataViewNode):
+        super().__init__()
         self.node = node
 
     def _execute(self, ctx: NodeContext) -> _data.DataView:
@@ -252,7 +297,13 @@ class DataViewFunc(NodeFunction[_data.DataView]):
         if root_item.is_empty():
             return _data.DataView.create_empty()
 
-        data_view = _data.DataView.for_trac_schema(self.node.schema)
+        if self.node.schema is not None and len(self.node.schema.table.fields) > 0:
+            trac_schema = self.node.schema
+        else:
+            arrow_schema = root_item.schema
+            trac_schema = _data.DataMapping.arrow_to_trac_schema(arrow_schema)
+
+        data_view = _data.DataView.for_trac_schema(trac_schema)
         data_view = _data.DataMapping.add_item_to_view(data_view, root_part_key, root_item)
 
         return data_view
@@ -261,6 +312,7 @@ class DataViewFunc(NodeFunction[_data.DataView]):
 class DataItemFunc(NodeFunction[_data.DataItem]):
 
     def __init__(self, node: DataItemNode):
+        super().__init__()
         self.node = node
 
     def _execute(self, ctx: NodeContext) -> _data.DataItem:
@@ -284,6 +336,7 @@ class DataItemFunc(NodeFunction[_data.DataItem]):
 class DataResultFunc(NodeFunction[ObjectBundle]):
 
     def __init__(self, node: DataResultNode):
+        super().__init__()
         self.node = node
 
     def _execute(self, ctx: NodeContext) -> ObjectBundle:
@@ -318,6 +371,7 @@ class DynamicDataSpecFunc(NodeFunction[_data.DataSpec]):
     RANDOM.seed()
 
     def __init__(self, node: DynamicDataSpecNode, storage: _storage.StorageManager):
+        super().__init__()
         self.node = node
         self.storage = storage
 
@@ -428,7 +482,7 @@ class _LoadSaveDataFunc(abc.ABC):
         return copy_
 
 
-class LoadDataFunc(NodeFunction[_data.DataItem], _LoadSaveDataFunc):
+class LoadDataFunc( _LoadSaveDataFunc, NodeFunction[_data.DataItem],):
 
     def __init__(self, node: LoadDataNode, storage: _storage.StorageManager):
         super().__init__(storage)
@@ -457,7 +511,7 @@ class LoadDataFunc(NodeFunction[_data.DataItem], _LoadSaveDataFunc):
         return _data.DataItem(table.schema, table)
 
 
-class SaveDataFunc(NodeFunction[None], _LoadSaveDataFunc):
+class SaveDataFunc(_LoadSaveDataFunc, NodeFunction[None]):
 
     def __init__(self, node: SaveDataNode, storage: _storage.StorageManager):
         super().__init__(storage)
@@ -512,6 +566,7 @@ def _model_def_for_import(import_details: meta.ImportModelJob):
 class ImportModelFunc(NodeFunction[meta.ObjectDefinition]):
 
     def __init__(self, node: ImportModelNode, models: _models.ModelLoader):
+        super().__init__()
         self.node = node
         self._models = models
 
@@ -529,11 +584,17 @@ class ImportModelFunc(NodeFunction[meta.ObjectDefinition]):
 
 class RunModelFunc(NodeFunction[Bundle[_data.DataView]]):
 
-    def __init__(self, node: RunModelNode, model_class: _api.TracModel.__class__, checkout_directory: pathlib.Path):
+    def __init__(
+            self, node: RunModelNode,
+            model_class: _api.TracModel.__class__,
+            checkout_directory: pathlib.Path,
+            storage_manager: _storage.StorageManager):
+
         super().__init__()
         self.node = node
         self.model_class = model_class
         self.checkout_directory = checkout_directory
+        self.storage_manager = storage_manager
 
     def _execute(self, ctx: NodeContext) -> Bundle[_data.DataView]:
 
@@ -544,36 +605,37 @@ class RunModelFunc(NodeFunction[Bundle[_data.DataView]]):
         # Still, if any nodes are missing or have the wrong type TracContextImpl will raise ERuntimeValidation
 
         local_ctx = {}
-        static_schemas = {}
+        dynamic_outputs = []
 
         for node_id, node_result in _ctx_iter_items(ctx):
+            if node_id.namespace == self.node.id.namespace:
+                if node_id.name in model_def.parameters or node_id.name in model_def.inputs:
+                    local_ctx[node_id.name] = node_result
 
-            if node_id.namespace != self.node.id.namespace:
-                continue
-
-            if node_id.name in model_def.parameters:
-                param_name = node_id.name
-                local_ctx[param_name] = node_result
-
-            if node_id.name in model_def.inputs:
-                input_name = node_id.name
-                local_ctx[input_name] = node_result
-                # At the moment, all model inputs have static schemas
-                static_schemas[input_name] = model_def.inputs[input_name].schema
+        # Set up access to external storage if required
 
-        # Add empty data views to the local context to hold model outputs
-        # Assuming outputs are all defined with static schemas
+        storage_map = {}
 
-        for output_name in model_def.outputs:
-            output_schema = self.node.model_def.outputs[output_name].schema
-            empty_data_view = _data.DataView.for_trac_schema(output_schema)
-            local_ctx[output_name] = empty_data_view
-            # At the moment, all model outputs have static schemas
-            static_schemas[output_name] = output_schema
+        if self.node.storage_access:
+            write_access = True if self.node.model_def.modelType == meta.ModelType.DATA_EXPORT_MODEL else False
+            for storage_key in self.node.storage_access:
+                if self.storage_manager.has_file_storage(storage_key, external=True):
+                    storage_impl = self.storage_manager.get_file_storage(storage_key, external=True)
+                    storage = _ctx.TracFileStorageImpl(storage_key, storage_impl, write_access, self.checkout_directory)
+                    storage_map[storage_key] = storage
 
         # Run the model against the mapped local context
 
-        trac_ctx = _ctx.TracContextImpl(self.node.model_def, self.model_class, local_ctx, static_schemas)
+        if model_def.modelType in [meta.ModelType.DATA_IMPORT_MODEL, meta.ModelType.DATA_EXPORT_MODEL]:
+            trac_ctx = _ctx.TracDataContextImpl(
+                self.node.model_def, self.model_class,
+                local_ctx, dynamic_outputs, storage_map,
+                self.checkout_directory)
+        else:
+            trac_ctx = _ctx.TracContextImpl(
+                self.node.model_def, self.model_class,
+                local_ctx, dynamic_outputs,
+                self.checkout_directory)
 
         try:
             model = self.model_class()
@@ -587,20 +649,50 @@ class RunModelFunc(NodeFunction[Bundle[_data.DataView]]):
 
         # Check required outputs are present and build the results bundle
 
+        model_name = self.model_class.__name__
         results: Bundle[_data.DataView] = dict()
+        new_nodes = dict()
+        new_deps = dict()
 
         for output_name, output_schema in model_def.outputs.items():
 
             result: _data.DataView = local_ctx.get(output_name)
 
             if result is None or result.is_empty():
+
                 if not output_schema.optional:
-                    model_name = self.model_class.__name__
                     raise _ex.ERuntimeValidation(f"Missing required output [{output_name}] from model [{model_name}]")
 
-            if result is not None:
+                # Create a placeholder for optional outputs that were not emitted
+                elif result is None:
+                    result = _data.DataView.create_empty()
+
+            results[output_name] = result
+
+        if dynamic_outputs:
+
+            for output_name in dynamic_outputs:
+
+                result: _data.DataView = local_ctx.get(output_name)
+
+                if result is None or result.is_empty():
+                    raise _ex.ERuntimeValidation(f"No data provided for [{output_name}] from model [{model_name}]")
+
                 results[output_name] = result
 
+                result_node_id = NodeId.of(output_name, self.node.id.namespace, _data.DataView)
+                result_node = BundleItemNode(result_node_id, self.node.id, output_name)
+
+                new_nodes[result_node_id] = result_node
+
+            output_section = _graph.GraphBuilder.build_runtime_outputs(dynamic_outputs, self.node.id.namespace)
+            new_nodes.update(output_section.nodes)
+
+            ctx_id = NodeId.of("trac_build_result", self.node.id.namespace, result_type=None)
+            new_deps[ctx_id] = list(_graph.Dependency(nid, _graph.DependencyType.HARD) for nid in output_section.outputs)
+
+            self.node_callback.send_graph_updates(new_nodes, new_deps)
+
         return results
 
 
@@ -624,6 +716,14 @@ class FunctionResolver:
         :py:class:`NodeFunction <NodeFunction>`
     """
 
+    # TODO: Validate consistency for resource keys
+    # Storage key should be validated for load data, save data and run model with storage access
+    # Repository key should be validated for import model (and explicitly for run model)
+
+    # Currently jobs with missing resources will fail at runtime, with a suitable error
+    # The resolver is called during graph building
+    # Putting the check here will raise a consistency error before the job starts processing
+
     __ResolveFunc = tp.Callable[['FunctionResolver', Node[_T]], NodeFunction[_T]]
 
     def __init__(self, models: _models.ModelLoader, storage: _storage.StorageManager):
@@ -658,12 +758,13 @@ class FunctionResolver:
 
     def resolve_run_model_node(self, node: RunModelNode) -> NodeFunction:
 
+        # TODO: Verify model_class against model_def
+
         model_class = self._models.load_model_class(node.model_scope, node.model_def)
         checkout_directory = self._models.model_load_checkout_directory(node.model_scope, node.model_def)
+        storage_manager = self._storage if node.storage_access else None
 
-        # TODO: Verify model_class against model_def
-
-        return RunModelFunc(node, model_class, checkout_directory)
+        return RunModelFunc(node, model_class, checkout_directory, storage_manager)
 
     __basic_node_mapping: tp.Dict[Node.__class__, NodeFunction.__class__] = {
 
@@ -677,6 +778,7 @@ class FunctionResolver:
         SaveJobResultNode: SaveJobResultFunc,
         DataResultNode: DataResultFunc,
         StaticValueNode: StaticValueFunc,
+        RuntimeOutputsNode: RuntimeOutputsFunc,
         BundleItemNode: NoopFunc,
         NoopNode: NoopFunc,
         RunModelResultNode: NoopFunc

tracdap/rt/_exec/graph.py

@@ -12,8 +12,6 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from __future__ import annotations
-
 import pathlib
 import typing as tp
 import dataclasses as dc
@@ -38,7 +36,7 @@ class NodeNamespace:
         return cls.__ROOT
 
     name: str
-    parent: tp.Optional[NodeNamespace] = dc.field(default_factory=lambda: NodeNamespace.root())
+    parent: "tp.Optional[NodeNamespace]" = dc.field(default_factory=lambda: NodeNamespace.root())
 
     def __str__(self):
         if self is self.__ROOT:
@@ -62,7 +60,7 @@ class NodeNamespace:
 class NodeId(tp.Generic[_T]):
 
     @staticmethod
-    def of(name: str, namespace: NodeNamespace, result_type: tp.Type[_T]) -> NodeId[_T]:
+    def of(name: str, namespace: NodeNamespace, result_type: tp.Type[_T]) -> "NodeId[_T]":
         return NodeId(name, namespace, result_type)
 
     name: str
@@ -83,8 +81,8 @@ class DependencyType:
     immediate: bool = True
     tolerant: bool = False
 
-    HARD: tp.ClassVar[DependencyType]
-    TOLERANT: tp.ClassVar[DependencyType]
+    HARD: "tp.ClassVar[DependencyType]"
+    TOLERANT: "tp.ClassVar[DependencyType]"
 
 
 DependencyType.HARD = DependencyType(immediate=True, tolerant=False)
@@ -93,6 +91,13 @@ DependencyType.TOLERANT = DependencyType(immediate=True, tolerant=True)
 DependencyType.DELAYED = DependencyType(immediate=False, tolerant=False)
 
 
+@dc.dataclass(frozen=True)
+class Dependency:
+
+    node_id: NodeId
+    dependency_type: DependencyType
+
+
 @dc.dataclass(frozen=True)
 class Node(tp.Generic[_T]):
 
@@ -165,6 +170,17 @@ class GraphSection:
     must_run: tp.List[NodeId] = dc.field(default_factory=list)
 
 
+Bundle: tp.Generic[_T] = tp.Dict[str, _T]
+ObjectBundle = Bundle[meta.ObjectDefinition]
+
+
+@dc.dataclass(frozen=True)
+class JobOutputs:
+
+    objects: tp.Dict[str, NodeId[meta.ObjectDefinition]] = dc.field(default_factory=dict)
+    bundles: tp.List[NodeId[ObjectBundle]] = dc.field(default_factory=list)
+
+
 # TODO: Where does this go?
 @dc.dataclass(frozen=True)
 class JobResultSpec:
@@ -179,10 +195,6 @@ class JobResultSpec:
 # ----------------------------------------------------------------------------------------------------------------------
 
 
-Bundle: tp.Generic[_T] = tp.Dict[str, _T]
-ObjectBundle = Bundle[meta.ObjectDefinition]
-
-
 @_node_type
 class NoopNode(Node):
     pass
@@ -354,6 +366,7 @@ class RunModelNode(Node[Bundle[_data.DataView]]):
     model_def: meta.ModelDefinition
     parameter_ids: tp.FrozenSet[NodeId]
     input_ids: tp.FrozenSet[NodeId]
+    storage_access: tp.Optional[tp.List[str]] = None
 
     def _node_dependencies(self) -> tp.Dict[NodeId, DependencyType]:
         return {dep_id: DependencyType.HARD for dep_id in [*self.parameter_ids, *self.input_ids]}
@@ -368,16 +381,28 @@ class RunModelResultNode(Node[None]):
         return {self.model_id: DependencyType.HARD}
 
 
+@_node_type
+class RuntimeOutputsNode(Node[JobOutputs]):
+
+    outputs: JobOutputs
+
+    def _node_dependencies(self) -> tp.Dict[NodeId, DependencyType]:
+        dep_ids = [*self.outputs.bundles, *self.outputs.objects.values()]
+        return {node_id: DependencyType.HARD for node_id in dep_ids}
+
+
 @_node_type
 class BuildJobResultNode(Node[cfg.JobResult]):
 
     job_id: meta.TagHeader
 
-    objects: tp.Dict[str, NodeId[meta.ObjectDefinition]] = dc.field(default_factory=dict)
-    bundles: tp.List[NodeId[ObjectBundle]] = dc.field(default_factory=list)
+    outputs: JobOutputs
+    runtime_outputs: tp.Optional[NodeId[JobOutputs]] = None
 
     def _node_dependencies(self) -> tp.Dict[NodeId, DependencyType]:
-        dep_ids = [*self.bundles, *self.objects.values()]
+        dep_ids = [*self.outputs.bundles, *self.outputs.objects.values()]
+        if self.runtime_outputs is not None:
+            dep_ids.append(self.runtime_outputs)
         return {node_id: DependencyType.HARD for node_id in dep_ids}