tracdap-runtime 0.6.4__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:
@@ -267,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:
@@ -290,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:
@@ -324,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
 
@@ -434,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)
@@ -463,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)
@@ -518,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
 
@@ -535,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]:
 
@@ -550,23 +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 = {}
+        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
+        # Set up access to external storage if required
 
-            if node_id.name in model_def.parameters:
-                param_name = node_id.name
-                local_ctx[param_name] = node_result
+        storage_map = {}
 
-            if node_id.name in model_def.inputs:
-                input_name = node_id.name
-                local_ctx[input_name] = node_result
+        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, self.checkout_directory)
+        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()
@@ -580,7 +649,10 @@ 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():
 
@@ -589,7 +661,6 @@ class RunModelFunc(NodeFunction[Bundle[_data.DataView]]):
             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}]")
 
                 # Create a placeholder for optional outputs that were not emitted
@@ -598,6 +669,30 @@ class RunModelFunc(NodeFunction[Bundle[_data.DataView]]):
 
             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
 
 
@@ -621,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):
@@ -655,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__] = {
 
@@ -674,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}
 
 

tracdap/rt/_exec/graph_builder.py

@@ -12,8 +12,6 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from __future__ import annotations
-
 import tracdap.rt.config as config
 import tracdap.rt.exceptions as _ex
 import tracdap.rt._impl.data as _data  # noqa
@@ -42,6 +40,9 @@ class GraphBuilder:
         if job_config.job.jobType == meta.JobType.RUN_FLOW:
             return cls.build_standard_job(job_config, result_spec, cls.build_run_flow_job)
 
+        if job_config.job.jobType in [meta.JobType.IMPORT_DATA, meta.JobType.EXPORT_DATA]:
+            return cls.build_standard_job(job_config, result_spec, cls.build_import_export_data_job)
+
         raise _ex.EConfigParse(f"Job type [{job_config.job.jobType}] is not supported yet")
 
     @classmethod
@@ -114,6 +115,28 @@ class GraphBuilder:
 
         return cls._join_sections(main_section, result_section)
 
+    @classmethod
+    def build_import_export_data_job(
+            cls, job_config: config.JobConfig, result_spec: JobResultSpec,
+            job_namespace: NodeNamespace, job_push_id: NodeId) \
+            -> GraphSection:
+
+        # TODO: These are processed as regular calculation jobs for now
+        # That might be ok, but is worth reviewing
+
+        if job_config.job.jobType == meta.JobType.IMPORT_DATA:
+            job_def = job_config.job.importData
+        else:
+            job_def = job_config.job.exportData
+
+        target_selector = job_def.model
+        target_obj = _util.get_job_resource(target_selector, job_config)
+        target_def = target_obj.model
+
+        return cls.build_calculation_job(
+            job_config, result_spec, job_namespace, job_push_id,
+            target_selector, target_def, job_def)
+
     @classmethod
     def build_run_model_job(
             cls, job_config: config.JobConfig, result_spec: JobResultSpec,
@@ -380,6 +403,65 @@ class GraphBuilder:
 
         return GraphSection(nodes, inputs=inputs)
 
+    @classmethod
+    def build_runtime_outputs(cls, output_names: tp.List[str], job_namespace: NodeNamespace):
+
+        # TODO: Factor out common logic with regular job outputs (including static / dynamic)
+
+        nodes = {}
+        inputs = set()
+        outputs = list()
+
+        for output_name in output_names:
+
+            # Output data view must already exist in the namespace
+            data_view_id = NodeId.of(output_name, job_namespace, _data.DataView)
+            data_spec_id = NodeId.of(f"{output_name}:SPEC", job_namespace, _data.DataSpec)
+
+            data_key = output_name + ":DATA"
+            data_id = _util.new_object_id(meta.ObjectType.DATA)
+            storage_key = output_name + ":STORAGE"
+            storage_id = _util.new_object_id(meta.ObjectType.STORAGE)
+
+            data_spec_node = DynamicDataSpecNode(
+                data_spec_id, data_view_id,
+                data_id, storage_id,
+                prior_data_spec=None)
+
+            output_data_key = _util.object_key(data_id)
+            output_storage_key = _util.object_key(storage_id)
+
+            # Map one data item from each view, since outputs are single part/delta
+            data_item_id = NodeId(f"{output_name}:ITEM", job_namespace, _data.DataItem)
+            data_item_node = DataItemNode(data_item_id, data_view_id)
+
+            # Create a physical save operation for the data item
+            data_save_id = NodeId.of(f"{output_name}:SAVE", job_namespace, None)
+            data_save_node = SaveDataNode(data_save_id, data_spec_id, data_item_id)
+
+            data_result_id = NodeId.of(f"{output_name}:RESULT", job_namespace, ObjectBundle)
+            data_result_node = DataResultNode(
+                data_result_id, output_name,
+                data_item_id, data_spec_id, data_save_id,
+                output_data_key, output_storage_key)
+
+            nodes[data_spec_id] = data_spec_node
+            nodes[data_item_id] = data_item_node
+            nodes[data_save_id] = data_save_node
+            nodes[data_result_id] = data_result_node
+
+            # Job-level data view is an input to the save operation
+            inputs.add(data_view_id)
+            outputs.append(data_result_id)
+
+        runtime_outputs = JobOutputs(bundles=outputs)
+        runtime_outputs_id = NodeId.of("trac_runtime_outputs", job_namespace, JobOutputs)
+        runtime_outputs_node = RuntimeOutputsNode(runtime_outputs_id, runtime_outputs)
+
+        nodes[runtime_outputs_id] = runtime_outputs_node
+
+        return GraphSection(nodes, inputs=inputs, outputs={runtime_outputs_id})
+
     @classmethod
     def build_job_results(
             cls, job_config: cfg.JobConfig, job_namespace: NodeNamespace, result_spec: JobResultSpec,
@@ -396,7 +478,8 @@ class GraphBuilder:
 
             build_result_node = BuildJobResultNode(
                 build_result_id, job_config.jobId,
-                objects=objects, explicit_deps=explicit_deps)
+                outputs = JobOutputs(objects=objects),
+                explicit_deps=explicit_deps)
 
         elif bundles is not None:
 
@@ -404,7 +487,8 @@ class GraphBuilder:
 
             build_result_node = BuildJobResultNode(
                 build_result_id, job_config.jobId,
-                bundles=bundles, explicit_deps=explicit_deps)
+                outputs = JobOutputs(bundles=bundles),
+                explicit_deps=explicit_deps)
 
         else:
             raise _ex.EUnexpected()
@@ -459,7 +543,7 @@ class GraphBuilder:
             -> GraphSection:
 
         if model_or_flow.objectType == meta.ObjectType.MODEL:
-            return cls.build_model(namespace, model_or_flow.model, explicit_deps)
+            return cls.build_model(job_config, namespace, model_or_flow.model, explicit_deps)
 
         elif model_or_flow.objectType == meta.ObjectType.FLOW:
             return cls.build_flow(job_config, namespace, model_or_flow.flow)
@@ -469,11 +553,13 @@ class GraphBuilder:
 
     @classmethod
     def build_model(
-            cls, namespace: NodeNamespace,
+            cls, job_config: config.JobConfig, namespace: NodeNamespace,
             model_def: meta.ModelDefinition,
             explicit_deps: tp.Optional[tp.List[NodeId]] = None) \
             -> GraphSection:
 
+        cls.check_model_type(job_config, model_def)
+
         def param_id(node_name):
             return NodeId(node_name, namespace, meta.Value)
 
@@ -485,6 +571,14 @@ class GraphBuilder:
         input_ids = set(map(data_id, model_def.inputs))
         output_ids = set(map(data_id, model_def.outputs))
 
+        # Set up storage access for import / export data jobs
+        if job_config.job.jobType == meta.JobType.IMPORT_DATA:
+            storage_access = job_config.job.importData.storageAccess
+        elif job_config.job.jobType == meta.JobType.EXPORT_DATA:
+            storage_access = job_config.job.exportData.storageAccess
+        else:
+            storage_access = None
+
         # Create the model node
         # Always add the prior graph root ID as a dependency
         # This is to ensure dependencies are still pulled in for models with no inputs!
@@ -500,7 +594,8 @@ class GraphBuilder:
         model_node = RunModelNode(
             model_id, model_scope, model_def,
             frozenset(parameter_ids), frozenset(input_ids),
-            explicit_deps=explicit_deps, bundle=model_id.namespace)
+            explicit_deps=explicit_deps, bundle=model_id.namespace,
+            storage_access=storage_access)
 
         model_result_id = NodeId(f"{model_name}:RESULT", namespace)
         model_result_node = RunModelResultNode(model_result_id, model_id)
@@ -637,6 +732,7 @@ class GraphBuilder:
 
             # Explicit check for model compatibility - report an error now, do not try build_model()
             cls.check_model_compatibility(model_selector, model_obj.model, node_name, node)
+            cls.check_model_type(job_config, model_obj.model)
 
             return cls.build_model_or_flow_with_context(
                 job_config, namespace, node_name, model_obj,
@@ -647,8 +743,8 @@ class GraphBuilder:
 
     @classmethod
     def check_model_compatibility(
-            cls, model_selector: meta.TagSelector, model_def: meta.ModelDefinition,
-            node_name: str, flow_node: meta.FlowNode):
+            cls, model_selector: meta.TagSelector,
+            model_def: meta.ModelDefinition, node_name: str, flow_node: meta.FlowNode):
 
         model_params = list(sorted(model_def.parameters.keys()))
         model_inputs = list(sorted(model_def.inputs.keys()))
@@ -662,6 +758,21 @@ class GraphBuilder:
             model_key = _util.object_key(model_selector)
             raise _ex.EJobValidation(f"Incompatible model for flow node [{node_name}] (Model: [{model_key}])")
 
+    @classmethod
+    def check_model_type(cls, job_config: config.JobConfig, model_def: meta.ModelDefinition):
+
+        if job_config.job.jobType == meta.JobType.IMPORT_DATA:
+            allowed_model_types = [meta.ModelType.DATA_IMPORT_MODEL]
+        elif job_config.job.jobType == meta.JobType.EXPORT_DATA:
+            allowed_model_types = [meta.ModelType.DATA_EXPORT_MODEL]
+        else:
+            allowed_model_types = [meta.ModelType.STANDARD_MODEL]
+
+        if model_def.modelType not in allowed_model_types:
+            job_type = job_config.job.jobType.name
+            model_type = model_def.modelType.name
+            raise _ex.EJobValidation(f"Job type [{job_type}] cannot use model type [{model_type}]")
+
     @staticmethod
     def build_context_push(
             namespace: NodeNamespace, input_mapping: tp.Dict[str, NodeId],