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

tracdap/rt/_exec/dev_mode.py

@@ -46,7 +46,7 @@ class DevModeTranslator:
     _log: tp.Optional[_util.logging.Logger] = None
 
     @classmethod
-    def translate_sys_config(cls, sys_config: _cfg.RuntimeConfig, config_dir: tp.Optional[pathlib.Path]):
+    def translate_sys_config(cls, sys_config: _cfg.RuntimeConfig, config_mgr: _cfg_p.ConfigManager):
 
         cls._log.info(f"Applying dev mode config translation to system config")
 
@@ -56,7 +56,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_dir)
+        sys_config = cls._resolve_relative_storage_root(sys_config, config_mgr)
 
         return sys_config
 
@@ -66,7 +66,7 @@ class DevModeTranslator:
             sys_config: _cfg.RuntimeConfig,
             job_config: _cfg.JobConfig,
             scratch_dir: pathlib.Path,
-            config_dir: tp.Optional[pathlib.Path],
+            config_mgr: _cfg_p.ConfigManager,
             model_class: tp.Optional[_api.TracModel.__class__]) \
             -> _cfg.JobConfig:
 
@@ -84,7 +84,7 @@ class DevModeTranslator:
 
         # 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_dir)
+            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]:
@@ -109,7 +109,7 @@ class DevModeTranslator:
     @classmethod
     def _resolve_relative_storage_root(
             cls, sys_config: _cfg.RuntimeConfig,
-            sys_config_path: tp.Optional[pathlib.Path]):
+            config_mgr: _cfg_p.ConfigManager):
 
         storage_config = copy.deepcopy(sys_config.storage)
 
@@ -128,6 +128,7 @@ class DevModeTranslator:
 
             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():
@@ -291,7 +292,7 @@ class DevModeTranslator:
         return model_id, model_object
 
     @classmethod
-    def _process_flow_definition(cls, job_config: _cfg.JobConfig, config_dir: pathlib.Path) -> _cfg.JobConfig:
+    def _process_flow_definition(cls, job_config: _cfg.JobConfig, config_mgr: _cfg_p.ConfigManager) -> _cfg.JobConfig:
 
         flow_details = job_config.job.runFlow.flow
 
@@ -305,21 +306,15 @@ class DevModeTranslator:
             cls._log.error(err)
             raise _ex.EConfigParse(err)
 
-        flow_path = config_dir.joinpath(flow_details) if config_dir is not None else pathlib.Path(flow_details)
-
-        if not flow_path.exists():
-            err = f"Flow definition not available for [{flow_details}]: File not found ({flow_path})"
-            cls._log.error(err)
-            raise _ex.EConfigParse(err)
-
         flow_id = _util.new_object_id(_meta.ObjectType.FLOW)
         flow_key = _util.object_key(flow_id)
 
-        cls._log.info(f"Generating flow definition for [{flow_details}] with ID = [{flow_key}]")
+        cls._log.info(f"Generating flow definition from [{flow_details}] with ID = [{flow_key}]")
+
+        flow_def = config_mgr.load_config_object(flow_details, _meta.FlowDefinition)
 
-        flow_parser = _cfg_p.ConfigParser(_meta.FlowDefinition)
-        flow_raw_data = flow_parser.load_raw_config(flow_path, flow_path.name)
-        flow_def = flow_parser.parse(flow_raw_data, flow_path.name)
+        # Validate models against the flow (this could move to _impl.validation and check prod jobs as well)
+        cls._check_models_for_flow(flow_def, job_config)
 
         # Auto-wiring and inference only applied to externally loaded flows for now
         flow_def = cls._autowire_flow(flow_def, job_config)
@@ -339,6 +334,37 @@ class DevModeTranslator:
 
         return job_config
 
+    @classmethod
+    def _check_models_for_flow(cls, flow: _meta.FlowDefinition, job_config: _cfg.JobConfig):
+
+        model_nodes = dict(filter(lambda n: n[1].nodeType == _meta.FlowNodeType.MODEL_NODE, flow.nodes.items()))
+
+        missing_models = list(filter(lambda m: m not in job_config.job.runFlow.models, model_nodes.keys()))
+        extra_models = list(filter(lambda m: m not in model_nodes, job_config.job.runFlow.models.keys()))
+
+        if any(missing_models):
+            error = f"Missing models in job definition: {', '.join(missing_models)}"
+            cls._log.error(error)
+            raise _ex.EJobValidation(error)
+
+        if any (extra_models):
+            error = f"Extra models in job definition: {', '.join(extra_models)}"
+            cls._log.error(error)
+            raise _ex.EJobValidation(error)
+
+        for model_name, model_node in model_nodes.items():
+
+            model_selector = job_config.job.runFlow.models[model_name]
+            model_obj = _util.get_job_resource(model_selector, job_config)
+
+            model_inputs = set(model_obj.model.inputs.keys())
+            model_outputs = set(model_obj.model.outputs.keys())
+
+            if model_inputs != set(model_node.inputs) or model_outputs != set(model_node.outputs):
+                error = f"The model supplied for [{model_name}] does not match the flow definition"
+                cls._log.error(error)
+                raise _ex.EJobValidation(error)
+
     @classmethod
     def _autowire_flow(cls, flow: _meta.FlowDefinition, job_config: _cfg.JobConfig):
 
@@ -629,11 +655,13 @@ class DevModeTranslator:
             job_details = job_config.job.runModel
             model_obj = _util.get_job_resource(job_details.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)
             required_inputs = flow_obj.flow.inputs
+            required_outputs = flow_obj.flow.outputs
 
         else:
             return job_config
@@ -645,7 +673,8 @@ class DevModeTranslator:
         for input_key, input_value in job_inputs.items():
             if not (isinstance(input_value, str) and input_value in job_resources):
 
-                input_schema = required_inputs[input_key].schema
+                model_input = required_inputs[input_key]
+                input_schema = model_input.schema if model_input and not model_input.dynamic else None
 
                 input_id = cls._process_input_or_output(
                     sys_config, input_key, input_value, job_resources,
@@ -656,9 +685,12 @@ class DevModeTranslator:
         for output_key, output_value in job_outputs.items():
             if not (isinstance(output_value, str) and output_value in job_resources):
 
+                model_output= required_outputs[output_key]
+                output_schema = model_output.schema if model_output and not model_output.dynamic else None
+
                 output_id = cls._process_input_or_output(
                     sys_config, output_key, output_value, job_resources,
-                    new_unique_file=True, schema=None)
+                    new_unique_file=True, schema=output_schema)
 
                 job_outputs[output_key] = _util.selector_for(output_id)
 
@@ -776,7 +808,7 @@ class DevModeTranslator:
         if schema is not None:
             data_def.schema = schema
         else:
-            data_def.schema = _meta.SchemaDefinition(schemaType=_meta.SchemaType.TABLE, table=_meta.TableSchema())
+            data_def.schema = None
 
         data_def.storageId = _meta.TagSelector(
             _meta.ObjectType.STORAGE, storage_id.objectId,

tracdap/rt/_exec/engine.py

@@ -19,6 +19,7 @@ import dataclasses as dc
 import enum
 import typing as tp
 
+import tracdap.rt.metadata as _meta
 import tracdap.rt.config as _cfg
 import tracdap.rt.exceptions as _ex
 import tracdap.rt._exec.actors as _actors
@@ -28,7 +29,6 @@ import tracdap.rt._impl.models as _models  # noqa
 import tracdap.rt._impl.data as _data  # noqa
 import tracdap.rt._impl.storage as _storage  # noqa
 import tracdap.rt._impl.util as _util  # noqa
-from .actors import Signal
 
 from .graph import NodeId
 
@@ -66,6 +66,18 @@ class _EngineContext:
     failed_nodes: tp.Set[NodeId] = dc.field(default_factory=set)
 
 
+@dc.dataclass
+class _JobState:
+
+    job_id: _meta.TagHeader
+    job_config: _cfg.JobConfig
+
+    actor_id: _actors.ActorId = None
+
+    job_result: _cfg.JobResult = None
+    job_error: Exception = None
+
+
 class TracEngine(_actors.Actor):
 
     """
@@ -88,7 +100,7 @@ class TracEngine(_actors.Actor):
         self._storage = storage
         self._notify_callback = notify_callback
 
-        self._job_actors = dict()
+        self._jobs: tp.Dict[str, _JobState] = dict()
 
     def on_start(self):
 
@@ -98,7 +110,7 @@ class TracEngine(_actors.Actor):
 
         self._log.info("Engine shutdown complete")
 
-    def on_signal(self, signal: Signal) -> tp.Optional[bool]:
+    def on_signal(self, signal: _actors.Signal) -> tp.Optional[bool]:
 
         # Failed signals can propagate from leaf nodes up the actor tree for a job
         # If the failure goes all the way up the tree without being handled, it will reach the engine node
@@ -110,8 +122,8 @@ class TracEngine(_actors.Actor):
             failed_job_key = None
 
             # Look for the job key corresponding to the failed actor
-            for job_key, job_actor in self._job_actors.items():
-                if job_actor == signal.sender:
+            for job_key, job_state in self._jobs.items():
+                if job_state.actor_id == signal.sender:
                     failed_job_key = job_key
 
             # If the job is still live, call job_failed explicitly
@@ -147,19 +159,34 @@ class TracEngine(_actors.Actor):
         job_processor = JobProcessor(job_key, job_config, result_spec,self._models, self._storage)
         job_actor_id = self.actors().spawn(job_processor)
 
-        job_actors = {**self._job_actors, job_key: job_actor_id}
-        self._job_actors = job_actors
+        job_state = _JobState(job_config.jobId, job_config)
+        job_state.actor_id = job_actor_id
+
+        self._jobs[job_key] = job_state
+
+    @_actors.Message
+    def get_job_list(self):
+
+        job_list = list(map(self._get_job_info, self._jobs.keys()))
+        self.actors().reply("job_list", job_list)
+
+    @_actors.Message
+    def get_job_details(self, job_key: str, details: bool):
+
+        details = self._get_job_info(job_key, details)
+        self.actors().reply("job_details", details)
 
     @_actors.Message
     def job_succeeded(self, job_key: str, job_result: _cfg.JobResult):
 
         # Ignore duplicate messages from the job processor (can happen in unusual error cases)
-        if job_key not in self._job_actors:
+        if job_key not in self._jobs:
             self._log.warning(f"Ignoring [job_succeeded] message, job [{job_key}] has already completed")
             return
 
         self._log.info(f"Recording job as successful: {job_key}")
 
+        self._jobs[job_key].job_result = job_result
         self._finalize_job(job_key)
 
         if self._notify_callback is not None:
@@ -169,12 +196,13 @@ class TracEngine(_actors.Actor):
     def job_failed(self, job_key: str, error: Exception):
 
         # Ignore duplicate messages from the job processor (can happen in unusual error cases)
-        if job_key not in self._job_actors:
+        if job_key not in self._jobs:
             self._log.warning(f"Ignoring [job_failed] message, job [{job_key}] has already completed")
             return
 
         self._log.error(f"Recording job as failed: {job_key}")
 
+        self._jobs[job_key].job_error = error
         self._finalize_job(job_key)
 
         if self._notify_callback is not None:
@@ -182,10 +210,47 @@ class TracEngine(_actors.Actor):
 
     def _finalize_job(self, job_key: str):
 
-        job_actors = self._job_actors
-        job_actor_id = job_actors.pop(job_key)
-        self.actors().stop(job_actor_id)
-        self._job_actors = job_actors
+        # Stop the actor but keep the job state available for status / results queries
+
+        # In the future, job state will need to be expunged after some period of time
+        # For now each instance of the runtime only processes one job so no need to worry
+
+        job_state = self._jobs.get(job_key)
+        job_actor_id = job_state.actor_id if job_state is not None else None
+
+        if job_actor_id is not None:
+            self.actors().stop(job_actor_id)
+            job_state.actor_id = None
+
+    def _get_job_info(self, job_key: str, details: bool = False) -> tp.Optional[_cfg.JobResult]:
+
+        job_state = self._jobs.get(job_key)
+
+        if job_state is None:
+            return None
+
+        job_result = _cfg.JobResult()
+        job_result.jobId = job_state.job_id
+
+        if job_state.actor_id is not None:
+            job_result.statusCode = _meta.JobStatusCode.RUNNING
+
+        elif job_state.job_result is not None:
+            job_result.statusCode = job_state.job_result.statusCode
+            job_result.statusMessage = job_state.job_result.statusMessage
+            if details:
+                job_result.results = job_state.job_result.results or dict()
+
+        elif job_state.job_error is not None:
+            job_result.statusCode = _meta.JobStatusCode.FAILED
+            job_result.statusMessage = str(job_state.job_error.args[0])
+
+        else:
+            # Alternatively return UNKNOWN status or throw an error here
+            job_result.statusCode = _meta.JobStatusCode.FAILED
+            job_result.statusMessage = "No details available"
+
+        return job_result
 
 
 class JobProcessor(_actors.Actor):
@@ -218,7 +283,7 @@ class JobProcessor(_actors.Actor):
         self._log.info(f"Cleaning up job [{self.job_key}]")
         self._models.destroy_scope(self.job_key)
 
-    def on_signal(self, signal: Signal) -> tp.Optional[bool]:
+    def on_signal(self, signal: _actors.Signal) -> tp.Optional[bool]:
 
         if signal.message == _actors.SignalNames.FAILED and isinstance(signal, _actors.ErrorSignal):
 

tracdap/rt/_exec/functions.py

@@ -252,7 +252,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
@@ -544,7 +550,6 @@ 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 = {}
 
         for node_id, node_result in _ctx_iter_items(ctx):
 
@@ -558,22 +563,10 @@ class RunModelFunc(NodeFunction[Bundle[_data.DataView]]):
             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
-
-        # Add empty data views to the local context to hold model outputs
-        # Assuming outputs are all defined with static schemas
-
-        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
 
         # Run the model against the mapped local context
 
-        trac_ctx = _ctx.TracContextImpl(self.node.model_def, self.model_class, local_ctx, static_schemas)
+        trac_ctx = _ctx.TracContextImpl(self.node.model_def, self.model_class, local_ctx, self.checkout_directory)
 
         try:
             model = self.model_class()
@@ -594,12 +587,16 @@ class RunModelFunc(NodeFunction[Bundle[_data.DataView]]):
             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:
-                results[output_name] = result
+                # Create a placeholder for optional outputs that were not emitted
+                elif result is None:
+                    result = _data.DataView.create_empty()
+
+            results[output_name] = result
 
         return results
 

tracdap/rt/_exec/runtime.py

@@ -16,6 +16,7 @@ from __future__ import annotations
 
 import dataclasses as dc
 import datetime as dt
+import signal
 import threading
 
 import sys
@@ -54,6 +55,8 @@ class TracRuntime:
         _engine.ModelNodeProcessor: "model",
         _engine.DataNodeProcessor: "data"}
 
+    __DEFAULT_API_PORT = 9000
+
     def __init__(
             self,
             sys_config: tp.Union[str, pathlib.Path, _cfg.RuntimeConfig],
@@ -61,6 +64,7 @@ class TracRuntime:
             job_result_format: tp.Optional[str] = None,
             scratch_dir: tp.Union[str, pathlib.Path, None] = None,
             scratch_dir_persist: bool = False,
+            plugin_packages: tp.List[str] = None,
             dev_mode: bool = False):
 
         trac_version = _version.__version__
@@ -83,28 +87,34 @@ class TracRuntime:
         self._log.info(f"TRAC D.A.P. Python Runtime {trac_version}")
 
         self._sys_config = sys_config if isinstance(sys_config, _cfg.RuntimeConfig) else None
-        self._sys_config_path = pathlib.Path(sys_config) if not self._sys_config else None
+        self._sys_config_path = sys_config if not self._sys_config else None
         self._job_result_dir = job_result_dir
         self._job_result_format = job_result_format
         self._scratch_dir = scratch_dir
         self._scratch_dir_provided = True if scratch_dir is not None else False
         self._scratch_dir_persist = scratch_dir_persist
+        self._plugin_packages = plugin_packages or []
         self._dev_mode = dev_mode
-        self._server_enabled = False
-        self._server_port = 0
 
+        # Runtime control
+        self._runtime_lock = threading.Lock()
+        self._runtime_event = threading.Condition(self._runtime_lock)
         self._pre_start_complete = False
+        self._shutdown_requested = False
+        self._oneshot_job = None
 
         # Top level resources
+        self._config_mgr: tp.Optional[_cparse.ConfigManager] = None
         self._models: tp.Optional[_models.ModelLoader] = None
         self._storage: tp.Optional[_storage.StorageManager] = None
 
         # The execution engine
         self._system: tp.Optional[_actors.ActorSystem] = None
         self._engine: tp.Optional[_engine.TracEngine] = None
-        self._engine_event = threading.Condition()
 
         # Runtime API server
+        self._server_enabled = False
+        self._server_port = 0
         self._server = None
 
         self._jobs: tp.Dict[str, _RuntimeJobInfo] = dict()
@@ -134,21 +144,28 @@ class TracRuntime:
 
             self._prepare_scratch_dir()
 
-            # Plugin manager and static API impl are singletons
-            # If these methods are called multiple times, the second and subsequent calls are ignored
+            # Plugin manager, static API and guard rails are singletons
+            # Calling these methods multiple times is safe (e.g. for embedded or testing scenarios)
+            # However, plugins are never un-registered for the lifetime of the processes
 
             _plugins.PluginManager.register_core_plugins()
+
+            for plugin_package in self._plugin_packages:
+                _plugins.PluginManager.register_plugin_package(plugin_package)
+
             _static_api.StaticApiImpl.register_impl()
             _guard.PythonGuardRails.protect_dangerous_functions()
 
             # Load sys config (or use embedded), config errors are detected before start()
             # Job config can also be checked before start() by using load_job_config()
 
+            self._config_mgr = _cparse.ConfigManager.for_root_config(self._sys_config_path)
+
             if self._sys_config is None:
                 sys_config_dev_mode = _dev_mode.DEV_MODE_SYS_CONFIG if self._dev_mode else None
-                sys_config_parser = _cparse.ConfigParser(_cfg.RuntimeConfig, sys_config_dev_mode)
-                sys_config_raw = sys_config_parser.load_raw_config(self._sys_config_path, config_file_name="system")
-                self._sys_config = sys_config_parser.parse(sys_config_raw, self._sys_config_path)
+                self._sys_config = self._config_mgr.load_root_object(
+                    _cfg.RuntimeConfig, sys_config_dev_mode,
+                    config_file_name="system")
             else:
                 self._log.info("Using embedded system config")
 
@@ -156,8 +173,15 @@ class TracRuntime:
             # I.e. it can be applied to embedded configs
 
             if self._dev_mode:
-                config_dir = self._sys_config_path.parent if self._sys_config_path is not None else None
-                self._sys_config = _dev_mode.DevModeTranslator.translate_sys_config(self._sys_config, config_dir)
+                self._sys_config = _dev_mode.DevModeTranslator.translate_sys_config(self._sys_config, self._config_mgr)
+
+            # Runtime API server is controlled by the sys config
+
+            if self._sys_config.runtimeApi is not None:
+                api_config = self._sys_config.runtimeApi
+                if api_config.enabled:
+                    self._server_enabled = True
+                    self._server_port = api_config.port or self.__DEFAULT_API_PORT
 
             self._pre_start_complete = True
 
@@ -196,7 +220,7 @@ class TracRuntime:
                 # The server module pulls in all the gRPC dependencies, don't import it unless we have to
                 import tracdap.rt._exec.server as _server
 
-                self._server = _server.RuntimeApiServer(self._server_port)
+                self._server = _server.RuntimeApiServer(self._system, self._server_port)
                 self._server.start()
 
         except Exception as e:
@@ -237,6 +261,28 @@ class TracRuntime:
         else:
             self._log.info("TRAC runtime has gone down cleanly")
 
+    def is_oneshot(self):
+        return not self._server_enabled
+
+    def run_until_done(self):
+
+        if self._server_enabled == False and len(self._jobs) == 0:
+            self._log.error("No job config supplied, TRAC runtime will not run")
+            raise _ex.EStartup("No job config supplied")
+
+        signal.signal(signal.SIGTERM, self._request_shutdown)
+        signal.signal(signal.SIGINT, self._request_shutdown)
+
+        with self._runtime_lock:
+            while not self._shutdown_requested:
+                self._runtime_event.wait()
+
+    def _request_shutdown(self, _signum = None, _frame = None):
+
+        with self._runtime_lock:
+            self._shutdown_requested = True
+            self._runtime_event.notify()
+
     def _prepare_scratch_dir(self):
 
         if not self._scratch_dir_provided:
@@ -274,20 +320,18 @@ class TracRuntime:
 
         if isinstance(job_config, _cfg.JobConfig):
             self._log.info("Using embedded job config")
-            job_config_path = None
 
         else:
-            job_config_path = job_config
             job_config_dev_mode = _dev_mode.DEV_MODE_JOB_CONFIG if self._dev_mode else None
-            job_config_parser = _cparse.ConfigParser(_cfg.JobConfig, job_config_dev_mode)
-            job_config_raw = job_config_parser.load_raw_config(job_config_path, config_file_name="job")
-            job_config = job_config_parser.parse(job_config_raw, job_config_path)
+            job_config = self._config_mgr.load_config_object(
+                job_config, _cfg.JobConfig,
+                job_config_dev_mode,
+                config_file_name="job")
 
         if self._dev_mode:
-            config_dir = job_config_path.parent if job_config_path is not None else None
             job_config = _dev_mode.DevModeTranslator.translate_job_config(
                 self._sys_config, job_config,
-                self._scratch_dir, config_dir,
+                self._scratch_dir, self._config_mgr,
                 model_class)
 
         return job_config
@@ -297,7 +341,7 @@ class TracRuntime:
         job_key = _util.object_key(job_config.jobId)
         self._jobs[job_key] = _RuntimeJobInfo()
 
-        self._system.send(
+        self._system.send_main(
             "submit_job", job_config,
             str(self._job_result_dir) if self._job_result_dir else "",
             self._job_result_format if self._job_result_format else "")
@@ -309,35 +353,34 @@ class TracRuntime:
         if job_key not in self._jobs:
             raise _ex.ETracInternal(f"Attempt to wait for a job that was never started")
 
-        with self._engine_event:
-            while True:
+        self._oneshot_job = job_key
 
-                job_info = self._jobs[job_key]
+        self.run_until_done()
 
-                if job_info.error is not None:
-                    raise job_info.error
+        job_info = self._jobs[job_key]
 
-                if job_info.result is not None:
-                    return job_info.result
+        if job_info.error is not None:
+            raise job_info.error
 
-                # TODO: Timeout / heartbeat
+        elif job_info.result is not None:
+            return job_info.result
 
-                self._engine_event.wait(1)
+        else:
+            err = f"No result or error information is available for job [{job_key}]"
+            self._log.error(err)
+            raise _ex.ETracInternal(err)
 
     def _engine_callback(self, job_key, job_result, job_error):
 
-        with self._engine_event:
-
-            if job_result is not None:
-                self._jobs[job_key].done = True
-                self._jobs[job_key].result = job_result
-            elif job_error is not None:
-                self._jobs[job_key].done = True
-                self._jobs[job_key].error = job_error
-            else:
-                pass
+        if job_result is not None:
+            self._jobs[job_key].done = True
+            self._jobs[job_key].result = job_result
+        elif job_error is not None:
+            self._jobs[job_key].done = True
+            self._jobs[job_key].error = job_error
 
-            self._engine_event.notify()
+        if self._oneshot_job == job_key:
+            self._request_shutdown()
 
     # ------------------------------------------------------------------------------------------------------------------
     # Error handling