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

tracdap/rt/_exec/dev_mode.py

@@ -12,8 +12,6 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from __future__ import annotations
-
 import re
 import typing as tp
 import copy
@@ -31,12 +29,12 @@ import tracdap.rt._impl.util as _util  # noqa
 
 
 DEV_MODE_JOB_CONFIG = [
-    re.compile(r"job\.run(Model|Flow)\.parameters\.\w+"),
-    re.compile(r"job\.run(Model|Flow)\.inputs\.\w+"),
-    re.compile(r"job\.run(Model|Flow)\.outputs\.\w+"),
-    re.compile(r"job\.runModel\.model"),
-    re.compile(r"job\.runFlow\.flow"),
-    re.compile(r"job\.runFlow\.models\.\w+")]
+    re.compile(r"job\.\w+\.parameters\.\w+"),
+    re.compile(r"job\.\w+\.inputs\.\w+"),
+    re.compile(r"job\.\w+\.outputs\.\w+"),
+    re.compile(r"job\.\w+\.models\.\w+"),
+    re.compile(r"job\.\w+\.model"),
+    re.compile(r"job\.\w+\.flow")]
 
 DEV_MODE_SYS_CONFIG = []
 
@@ -56,7 +54,7 @@ class DevModeTranslator:
             sys_config.storage = _cfg.StorageConfig()
 
         sys_config = cls._add_integrated_repo(sys_config)
-        sys_config = cls._resolve_relative_storage_root(sys_config, config_mgr)
+        sys_config = cls._process_storage(sys_config, config_mgr)
 
         return sys_config
 
@@ -72,24 +70,23 @@ class DevModeTranslator:
 
         cls._log.info(f"Applying dev mode config translation to job config")
 
-        if not job_config.jobId:
+        # Protobuf semantics for a blank jobId should be an object, but objectId will be an empty string
+        if not job_config.jobId or not job_config.jobId.objectId:
             job_config = cls._process_job_id(job_config)
 
         if job_config.job.jobType is None or job_config.job.jobType == _meta.JobType.JOB_TYPE_NOT_SET:
             job_config = cls._process_job_type(job_config)
 
         # Load and populate any models provided as a Python class or class name
-        if job_config.job.jobType in [_meta.JobType.RUN_MODEL, _meta.JobType.RUN_FLOW]:
-            job_config = cls._process_models(sys_config, job_config, scratch_dir, model_class)
+        job_config = cls._process_models(sys_config, job_config, scratch_dir, model_class)
 
         # Fow flows, load external flow definitions then perform auto-wiring and type inference
         if job_config.job.jobType == _meta.JobType.RUN_FLOW:
             job_config = cls._process_flow_definition(job_config, config_mgr)
 
-        # For run (model|flow) jobs, apply processing to the parameters, inputs and outputs
-        if job_config.job.jobType in [_meta.JobType.RUN_MODEL, _meta.JobType.RUN_FLOW]:
-            job_config = cls._process_parameters(job_config)
-            job_config = cls._process_inputs_and_outputs(sys_config, job_config)
+        # Apply processing to the parameters, inputs and outputs
+        job_config = cls._process_parameters(job_config)
+        job_config = cls._process_inputs_and_outputs(sys_config, job_config)
 
         return job_config
 
@@ -107,51 +104,60 @@ class DevModeTranslator:
         return sys_config
 
     @classmethod
-    def _resolve_relative_storage_root(
+    def _process_storage(
             cls, sys_config: _cfg.RuntimeConfig,
             config_mgr: _cfg_p.ConfigManager):
 
         storage_config = copy.deepcopy(sys_config.storage)
 
         for bucket_key, bucket_config in storage_config.buckets.items():
+            storage_config.buckets[bucket_key] = cls._resolve_storage_location(
+                bucket_key, bucket_config, config_mgr)
+
+        for bucket_key, bucket_config in storage_config.external.items():
+            storage_config.external[bucket_key] = cls._resolve_storage_location(
+                bucket_key, bucket_config, config_mgr)
 
-            if bucket_config.protocol != "LOCAL":
-                continue
+        sys_config = copy.copy(sys_config)
+        sys_config.storage = storage_config
 
-            if "rootPath" not in bucket_config.properties:
-                continue
+        return sys_config
 
-            root_path = pathlib.Path(bucket_config.properties["rootPath"])
+    @classmethod
+    def _resolve_storage_location(cls, bucket_key, bucket_config, config_mgr: _cfg_p.ConfigManager):
 
-            if root_path.is_absolute():
-                continue
+        if bucket_config.protocol != "LOCAL":
+            return bucket_config
 
-            cls._log.info(f"Resolving relative path for [{bucket_key}] local storage...")
+        if "rootPath" not in bucket_config.properties:
+            return bucket_config
 
-            sys_config_path = config_mgr.config_dir_path()
-            if sys_config_path is not None:
-                absolute_path = sys_config_path.joinpath(root_path).resolve()
-                if absolute_path.exists():
-                    cls._log.info(f"Resolved [{root_path}] -> [{absolute_path}]")
-                    bucket_config.properties["rootPath"] = str(absolute_path)
-                    continue
+        root_path = pathlib.Path(bucket_config.properties["rootPath"])
 
-            cwd = pathlib.Path.cwd()
-            absolute_path = cwd.joinpath(root_path).resolve()
+        if root_path.is_absolute():
+            return bucket_config
 
+        cls._log.info(f"Resolving relative path for [{bucket_key}] local storage...")
+
+        sys_config_path = config_mgr.config_dir_path()
+        if sys_config_path is not None:
+            absolute_path = sys_config_path.joinpath(root_path).resolve()
             if absolute_path.exists():
                 cls._log.info(f"Resolved [{root_path}] -> [{absolute_path}]")
                 bucket_config.properties["rootPath"] = str(absolute_path)
-                continue
+                return bucket_config
 
-            msg = f"Failed to resolve relative storage path [{root_path}]"
-            cls._log.error(msg)
-            raise _ex.EConfigParse(msg)
+        cwd = pathlib.Path.cwd()
+        absolute_path = cwd.joinpath(root_path).resolve()
 
-        sys_config = copy.copy(sys_config)
-        sys_config.storage = storage_config
+        if absolute_path.exists():
+            cls._log.info(f"Resolved [{root_path}] -> [{absolute_path}]")
+            bucket_config.properties["rootPath"] = str(absolute_path)
+            return bucket_config
 
-        return sys_config
+        msg = f"Failed to resolve relative storage path [{root_path}]"
+        cls._log.error(msg)
+        raise _ex.EConfigParse(msg)
 
     @classmethod
     def _add_job_resource(
@@ -188,6 +194,12 @@ class DevModeTranslator:
         elif job_config.job.importModel is not None:
             job_type = _meta.JobType.IMPORT_MODEL
 
+        elif job_config.job.importData is not None:
+            job_type = _meta.JobType.IMPORT_DATA
+
+        elif job_config.job.exportData is not None:
+            job_type = _meta.JobType.EXPORT_DATA
+
         else:
             cls._log.error("Could not infer job type")
             raise _ex.EConfigParse("Could not infer job type")
@@ -202,6 +214,26 @@ class DevModeTranslator:
 
         return job_config
 
+    @classmethod
+    def _get_job_detail(cls, job_config: _cfg.JobConfig):
+
+        if job_config.job.jobType == _meta.JobType.RUN_MODEL:
+            return job_config.job.runModel
+
+        if job_config.job.jobType == _meta.JobType.RUN_FLOW:
+            return job_config.job.runFlow
+
+        if job_config.job.jobType == _meta.JobType.IMPORT_MODEL:
+            return job_config.job.importModel
+
+        if job_config.job.jobType == _meta.JobType.IMPORT_DATA:
+            return job_config.job.importData
+
+        if job_config.job.jobType == _meta.JobType.EXPORT_DATA:
+            return job_config.job.exportData
+
+        raise _ex.EConfigParse(f"Could not get job details for job type [{job_config.job.jobType}]")
+
     @classmethod
     def _process_models(
             cls,
@@ -214,41 +246,39 @@ class DevModeTranslator:
         model_loader = _models.ModelLoader(sys_config, scratch_dir)
         model_loader.create_scope("DEV_MODE_TRANSLATION")
 
-        original_config = job_config
+        # This processing works on the assumption that job details follow a convention for addressing models
+        # Jobs requiring a single model have a field called "model"
+        # Jobs requiring multiple models have a field called "models@, which is a dict
 
-        job_config = copy.copy(job_config)
-        job_config.job = copy.copy(job_config.job)
-        job_config.resources = copy.copy(job_config.resources)
+        job_detail = cls._get_job_detail(job_config)
 
-        if job_config.job.jobType == _meta.JobType.RUN_MODEL:
+        # If a model class is supplied in code, use that to generate the model def
+        if model_class is not None:
 
-            job_config.job.runModel = copy.copy(job_config.job.runModel)
+            # Passing a model class via launch_model() is only supported for job types with a single model
+            if not hasattr(job_detail, "model"):
+                raise _ex.EJobValidation(f"Job type [{job_config.job.jobType}] cannot be launched using launch_model()")
 
-            # If a model class is supplied in code, use that to generate the model def
-            if model_class is not None:
-                model_id, model_obj = cls._generate_model_for_class(model_loader, model_class)
-                job_config = cls._add_job_resource(job_config, model_id, model_obj)
-                job_config.job.runModel.model = _util.selector_for(model_id)
-
-            # Otherwise if model specified as a string instead of a selector, apply the translation
-            elif isinstance(original_config.job.runModel.model, str):
-                model_detail = original_config.job.runModel.model
-                model_id, model_obj = cls._generate_model_for_entry_point(model_loader, model_detail)  # noqa
-                job_config = cls._add_job_resource(job_config, model_id, model_obj)
-                job_config.job.runModel.model = _util.selector_for(model_id)
-
-        if job_config.job.jobType == _meta.JobType.RUN_FLOW:
+            model_id, model_obj = cls._generate_model_for_class(model_loader, model_class)
+            job_detail.model = _util.selector_for(model_id)
+            job_config = cls._add_job_resource(job_config, model_id, model_obj)
 
-            job_config.job.runFlow = copy.copy(job_config.job.runFlow)
-            job_config.job.runFlow.models = copy.copy(job_config.job.runFlow.models)
+        # Otherwise look for models specified as a single string, and take that as the entry point
+        else:
 
-            for model_key, model_detail in original_config.job.runFlow.models.items():
+            # Jobs with a single model
+            if hasattr(job_detail, "model") and isinstance(job_detail.model, str):
+                model_id, model_obj = cls._generate_model_for_entry_point(model_loader, job_detail.model)  # noqa
+                job_detail.model = _util.selector_for(model_id)
+                job_config = cls._add_job_resource(job_config, model_id, model_obj)
 
-                # Only apply translation if the model is specified as a string instead of a selector
-                if isinstance(model_detail, str):
-                    model_id, model_obj = cls._generate_model_for_entry_point(model_loader, model_detail)
-                    job_config = cls._add_job_resource(job_config, model_id, model_obj)
-                    job_config.job.runFlow.models[model_key] = _util.selector_for(model_id)
+            # Jobs with multiple modlels
+            elif hasattr(job_detail, "models") and isinstance(job_detail.models, dict):
+                for model_key, model_detail in job_detail.models.items():
+                    if isinstance(model_detail, str):
+                        model_id, model_obj = cls._generate_model_for_entry_point(model_loader, model_detail)
+                        job_detail.models[model_key] = _util.selector_for(model_id)
+                        job_config = cls._add_job_resource(job_config, model_id, model_obj)
 
         model_loader.destroy_scope("DEV_MODE_TRANSLATION")
 
@@ -596,35 +626,34 @@ class DevModeTranslator:
     @classmethod
     def _process_parameters(cls, job_config: _cfg.JobConfig) -> _cfg.JobConfig:
 
-        if job_config.job.jobType == _meta.JobType.RUN_MODEL:
-
-            job_details = job_config.job.runModel
-            model_key = _util.object_key(job_details.model)
-            model_or_flow = job_config.resources[model_key].model
+        # This relies on convention for naming properties across similar job types
 
-        elif job_config.job.jobType == _meta.JobType.RUN_FLOW:
+        job_detail = cls._get_job_detail(job_config)
 
-            job_details = job_config.job.runFlow
-            flow_key = _util.object_key(job_details.flow)
+        if hasattr(job_detail, "model"):
+            model_key = _util.object_key(job_detail.model)
+            model_or_flow = job_config.resources[model_key].model
+        elif hasattr(job_detail, "flow"):
+            flow_key = _util.object_key(job_detail.flow)
             model_or_flow = job_config.resources[flow_key].flow
-
         else:
-            raise _ex.EUnexpected()
+            model_or_flow = None
+
+        if model_or_flow is not None:
 
-        param_specs = model_or_flow.parameters
-        param_values = job_details.parameters
+            param_specs = model_or_flow.parameters
+            raw_values = job_detail.parameters
 
-        # Set encoded params on runModel or runFlow depending on the job type
-        job_details.parameters = cls._process_parameters_dict(param_specs, param_values)
+            job_detail.parameters = cls._process_parameters_dict(param_specs, raw_values)
 
         return job_config
 
     @classmethod
     def _process_parameters_dict(
             cls, param_specs: tp.Dict[str, _meta.ModelParameter],
-            param_values: tp.Dict[str, _meta.Value]) -> tp.Dict[str, _meta.Value]:
+            raw_values: tp.Dict[str, _meta.Value]) -> tp.Dict[str, _meta.Value]:
 
-        unknown_params = list(filter(lambda p: p not in param_specs, param_values))
+        unknown_params = list(filter(lambda p: p not in param_specs, raw_values))
 
         if any(unknown_params):
             msg = f"Unknown parameters cannot be translated: [{', '.join(unknown_params)}]"
@@ -633,7 +662,7 @@ class DevModeTranslator:
 
         encoded_values = dict()
 
-        for p_name, p_value in param_values.items():
+        for p_name, p_value in raw_values.items():
 
             if isinstance(p_value, _meta.Value):
                 encoded_values[p_name] = p_value
@@ -651,23 +680,23 @@ class DevModeTranslator:
     @classmethod
     def _process_inputs_and_outputs(cls, sys_config: _cfg.RuntimeConfig, job_config: _cfg.JobConfig) -> _cfg.JobConfig:
 
-        if job_config.job.jobType == _meta.JobType.RUN_MODEL:
-            job_details = job_config.job.runModel
-            model_obj = _util.get_job_resource(job_details.model, job_config)
+        job_detail = cls._get_job_detail(job_config)
+
+        if hasattr(job_detail, "model"):
+            model_obj = _util.get_job_resource(job_detail.model, job_config)
             required_inputs = model_obj.model.inputs
             required_outputs = model_obj.model.outputs
 
-        elif job_config.job.jobType == _meta.JobType.RUN_FLOW:
-            job_details = job_config.job.runFlow
-            flow_obj = _util.get_job_resource(job_details.flow, job_config)
+        elif hasattr(job_detail, "flow"):
+            flow_obj = _util.get_job_resource(job_detail.flow, job_config)
             required_inputs = flow_obj.flow.inputs
             required_outputs = flow_obj.flow.outputs
 
         else:
             return job_config
 
-        job_inputs = job_details.inputs
-        job_outputs = job_details.outputs
+        job_inputs = job_detail.inputs
+        job_outputs = job_detail.outputs
         job_resources = job_config.resources
 
         for input_key, input_value in job_inputs.items():

tracdap/rt/_exec/engine.py

@@ -12,8 +12,6 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from __future__ import annotations
-
 import copy as cp
 import dataclasses as dc
 import enum
@@ -272,12 +270,13 @@ class JobProcessor(_actors.Actor):
         self.result_spec = result_spec
         self._models = models
         self._storage = storage
+        self._resolver = _func.FunctionResolver(models, storage)
         self._log = _util.logger_for_object(self)
 
     def on_start(self):
         self._log.info(f"Starting job [{self.job_key}]")
         self._models.create_scope(self.job_key)
-        self.actors().spawn(GraphBuilder(self.job_config, self.result_spec, self._models, self._storage))
+        self.actors().spawn(GraphBuilder(self.job_config, self.result_spec, self._resolver))
 
     def on_stop(self):
         self._log.info(f"Cleaning up job [{self.job_key}]")
@@ -305,7 +304,7 @@ class JobProcessor(_actors.Actor):
 
     @_actors.Message
     def job_graph(self, graph: _EngineContext, root_id: NodeId):
-        self.actors().spawn(GraphProcessor(graph, root_id))
+        self.actors().spawn(GraphProcessor(graph, root_id, self._resolver))
         self.actors().stop(self.actors().sender)
 
     @_actors.Message
@@ -331,15 +330,14 @@ class GraphBuilder(_actors.Actor):
     def __init__(
             self, job_config: _cfg.JobConfig,
             result_spec: _graph.JobResultSpec,
-            models: _models.ModelLoader,
-            storage: _storage.StorageManager):
+            resolver: _func.FunctionResolver):
 
         super().__init__()
         self.job_config = job_config
         self.result_spec = result_spec
         self.graph: tp.Optional[_EngineContext] = None
 
-        self._resolver = _func.FunctionResolver(models, storage)
+        self._resolver = resolver
         self._log = _util.logger_for_object(self)
 
     def on_start(self):
@@ -378,11 +376,12 @@ class GraphProcessor(_actors.Actor):
     Once all running nodes are stopped, an error is reported to the parent
     """
 
-    def __init__(self, graph: _EngineContext, root_id: NodeId):
+    def __init__(self, graph: _EngineContext, root_id: NodeId, resolver: _func.FunctionResolver):
         super().__init__()
         self.graph = graph
         self.root_id = root_id
         self.processors: tp.Dict[NodeId, _actors.ActorId] = dict()
+        self._resolver = resolver
         self._log = _util.logger_for_object(self)
 
     def on_start(self):
@@ -463,6 +462,62 @@ class GraphProcessor(_actors.Actor):
         # Job may have completed due to error propagation
         self.check_job_status(do_submit=False)
 
+    @_actors.Message
+    def update_graph(
+            self, requestor_id: NodeId,
+            new_nodes: tp.Dict[NodeId, _graph.Node],
+            new_deps: tp.Dict[NodeId, tp.List[_graph.Dependency]]):
+
+        new_graph = cp.copy(self.graph)
+        new_graph.nodes = cp.copy(new_graph.nodes)
+
+        # Attempt to insert a duplicate node is always an error
+        node_collision = list(filter(lambda nid: nid in self.graph.nodes, new_nodes))
+
+        # Only allow adding deps to pending nodes for now (adding deps to active nodes will require more work)
+        dep_collision = list(filter(lambda nid: nid not in self.graph.pending_nodes, new_deps))
+
+        dep_invalid = list(filter(
+            lambda dds: any(filter(lambda dd: dd.node_id not in new_nodes, dds)),
+            new_deps.values()))
+
+        if any(node_collision) or any(dep_collision) or any(dep_invalid):
+
+            self._log.error(f"Node collision during graph update (requested by {requestor_id})")
+            self._log.error(f"Duplicate node IDs: {node_collision or 'None'}")
+            self._log.error(f"Dependency updates for dead nodes: {dep_collision or 'None'}")
+            self._log.error(f"Dependencies added for existing nodes: {dep_invalid or 'None'}")
+
+            # Set an error on the node, and wait for it to complete normally
+            # The error will be picked up when the result is recorded
+            # If dependencies are added for an active node, more signalling will be needed
+            requestor = cp.copy(new_graph.nodes[requestor_id])
+            requestor.error = _ex.ETracInternal("Node collision during graph update")
+            new_graph.nodes[requestor_id] = requestor
+
+            return
+
+        new_graph.pending_nodes = cp.copy(new_graph.pending_nodes)
+
+        for node_id, node in new_nodes.items():
+            GraphLogger.log_node_add(node)
+            node_func = self._resolver.resolve_node(node)
+            new_node = _EngineNode(node, {}, function=node_func)
+            new_graph.nodes[node_id] = new_node
+            new_graph.pending_nodes.add(node_id)
+
+        for node_id, deps in new_deps.items():
+            engine_node = cp.copy(new_graph.nodes[node_id])
+            engine_node.dependencies = cp.copy(engine_node.dependencies)
+            for dep in deps:
+                GraphLogger.log_dependency_add(node_id, dep.node_id)
+                engine_node.dependencies[dep.node_id] = dep.dependency_type
+            new_graph.nodes[node_id] = engine_node
+
+        self.graph = new_graph
+
+        self.actors().send(self.actors().id, "submit_viable_nodes")
+
     @classmethod
     def _is_required_node(cls, node: _EngineNode, graph: _EngineContext):
 
@@ -654,8 +709,15 @@ class NodeProcessor(_actors.Actor):
 
             NodeLogger.log_node_start(self.node)
 
+            # Context contains only node states available when the context is set up
             ctx = NodeContextImpl(self.graph.nodes)
-            result = self.node.function(ctx)
+
+            # Callback remains valid because it only lives inside the call stack for this message
+            callback = NodeCallbackImpl(self.actors(), self.node_id)
+
+            # Execute the node function
+            result = self.node.function(ctx, callback)
+
             self._check_result_type(result)
 
             NodeLogger.log_node_succeeded(self.node)
@@ -730,6 +792,37 @@ class DataNodeProcessor(NodeProcessor):
         super().__init__(graph, node_id, node)
 
 
+class GraphLogger:
+
+    """
+    Log the activity of the GraphProcessor
+    """
+
+    _log = _util.logger_for_class(GraphProcessor)
+
+    @classmethod
+    def log_node_add(cls, node: _graph.Node):
+
+        node_name = node.id.name
+        namespace = node.id.namespace
+
+        cls._log.info(f"ADD {cls._func_type(node)} [{node_name}] / {namespace}")
+
+    @classmethod
+    def log_dependency_add(cls, node_id: NodeId, dep_id: NodeId):
+
+        if node_id.namespace == dep_id.namespace:
+            cls._log.info(f"ADD DEPENDENCY [{node_id.name}] -> [{dep_id.name}] / {node_id.namespace}")
+        else:
+            cls._log.info(f"ADD DEPENDENCY [{node_id.name}] / {node_id.namespace} -> [{dep_id.name}] / {dep_id.namespace}")
+
+    @classmethod
+    def _func_type(cls, node: _graph.Node):
+
+        func_type = type(node)
+        return func_type.__name__[:-4]
+
+
 class NodeLogger:
 
     """
@@ -912,3 +1005,21 @@ class NodeContextImpl(_func.NodeContext):
         for node_id, node in self.__nodes.items():
             if node.complete and not node.error:
                 yield node_id, node.result
+
+
+class NodeCallbackImpl(_func.NodeCallback):
+
+    """
+    Callback impl is passed to node functions so they can call into the engine
+    It is only valid as long as the node function runs inside the call stack of a single message
+    """
+
+    def __init__(self, actor_ctx: _actors.ActorContext, node_id: NodeId):
+        self.__actor_ctx = actor_ctx
+        self.__node_id = node_id
+
+    def send_graph_updates(
+            self, new_nodes: tp.Dict[NodeId, _graph.Node],
+            new_deps: tp.Dict[NodeId, tp.List[_graph.Dependency]]):
+
+        self.__actor_ctx.send_parent("update_graph", self.__node_id, new_nodes, new_deps)