tracdap-runtime 0.6.1.dev3__py3-none-any.whl → 0.6.3__py3-none-any.whl

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,22 +87,35 @@ 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
 
+        # 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()
 
@@ -127,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")
 
@@ -149,8 +173,17 @@ 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
 
         except Exception as e:
             self._handle_startup_error(e)
@@ -159,6 +192,10 @@ class TracRuntime:
 
         try:
 
+            # Ensure pre-start has been run
+            if not self._pre_start_complete:
+                self.pre_start()
+
             self._log.info("Starting the engine...")
 
             self._models = _models.ModelLoader(self._sys_config, self._scratch_dir)
@@ -175,11 +212,26 @@ class TracRuntime:
 
             self._system.start(wait=wait)
 
+            # If the runtime server has been enabled, start it up
+            if self._server_enabled:
+
+                self._log.info("Starting the runtime API server...")
+
+                # 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._system, self._server_port)
+                self._server.start()
+
         except Exception as e:
             self._handle_startup_error(e)
 
     def stop(self, due_to_error=False):
 
+        if self._server is not None:
+            self._log.info("Stopping the runtime API server...")
+            self._server.stop()
+
         if due_to_error:
             self._log.info("Shutting down the engine in response to an error")
         else:
@@ -209,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:
@@ -246,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
@@ -269,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 "")
@@ -281,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

tracdap/rt/_exec/server.py

@@ -0,0 +1,345 @@
+#  Copyright 2024 Accenture Global Solutions Limited
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import asyncio
+import threading
+import typing as tp
+
+import tracdap.rt.config as config
+import tracdap.rt.exceptions as ex
+import tracdap.rt._exec.actors as actors
+import tracdap.rt._impl.grpc.codec as codec  # noqa
+import tracdap.rt._impl.util as util  # noqa
+
+# Check whether gRPC is installed before trying to load any of the generated modules
+try:
+    import grpc.aio  # noqa
+    import google.protobuf.message as _msg  # noqa
+except ImportError:
+    raise ex.EStartup("The runtime API server cannot be enabled because gRPC libraries are not installed")
+
+# Imports for gRPC generated code, these are managed by build_runtime.py for distribution
+import tracdap.rt._impl.grpc.tracdap.api.internal.runtime_pb2 as runtime_pb2
+import tracdap.rt._impl.grpc.tracdap.api.internal.runtime_pb2_grpc as runtime_grpc
+
+
+class RuntimeApiServer(runtime_grpc.TracRuntimeApiServicer):
+
+    # Default timeout values in seconds
+    __DEFAULT_STARTUP_TIMEOUT = 5.0
+    __DEFAULT_SHUTDOWN_TIMEOUT = 10.0
+    __DEFAULT_REQUEST_TIMEOUT = 10.0
+
+    def __init__(self, system: actors.ActorSystem, port: int):
+
+        self.__log = util.logger_for_object(self)
+
+        self.__system = system
+        self.__engine_id = system.main_id()
+        self.__agent: tp.Optional[ApiAgent] = None
+
+        self.__port = port
+        self.__request_timeout = self.__DEFAULT_REQUEST_TIMEOUT  # Not configurable atm
+        self.__server: tp.Optional[grpc.aio.Server] = None
+        self.__server_thread: tp.Optional[threading.Thread] = None
+        self.__event_loop: tp.Optional[asyncio.AbstractEventLoop] = None
+
+        self.__start_signal: tp.Optional[threading.Event] = None
+        self.__stop_signal: tp.Optional[asyncio.Event] = None
+
+    def start(self, startup_timeout: float = None):
+
+        if self.__start_signal is not None:
+            return
+
+        timeout = startup_timeout or self.__DEFAULT_SHUTDOWN_TIMEOUT
+
+        self.__start_signal = threading.Event()
+        self.__server_thread = threading.Thread(target=self.__server_main, name="api_server", daemon=True)
+        self.__server_thread.start()
+
+        try:
+            self.__start_signal.wait(timeout)
+        except TimeoutError as e:
+            raise ex.EStartup("Runtime API failed to start") from e
+
+    def stop(self, shutdown_timeout: float = None):
+
+        if self.__server is None:
+            return
+
+        timeout = shutdown_timeout or self.__DEFAULT_SHUTDOWN_TIMEOUT
+
+        self.__event_loop.call_soon_threadsafe(lambda: self.__stop_signal.set())
+        self.__server_thread.join(timeout)
+
+        if self.__server_thread.is_alive():
+            self.__log.warning("Runtime API server did not go down cleanly")
+
+    def __server_main(self):
+
+        self.__event_loop = asyncio.new_event_loop()
+        self.__event_loop.run_until_complete(self.__server_main_async())
+        self.__event_loop.close()
+
+    async def __server_main_async(self):
+
+        server_address = f"[::]:{self.__port}"
+
+        # Asyncio events must be created inside the event loop for Python <= 3.9
+        self.__stop_signal = asyncio.Event()
+
+        # Agent using asyncio, so must be created inside the event loop
+        self.__agent = ApiAgent()
+        self.__system.spawn_agent(self.__agent)
+
+        self.__server = grpc.aio.server()
+        self.__server.add_insecure_port(server_address)
+        runtime_grpc.add_TracRuntimeApiServicer_to_server(self, self.__server)
+
+        await self.__server.start()
+        await self.__agent.started()
+
+        self.__start_signal.set()
+
+        self.__log.info(f"Runtime API server is up and listening on port [{self.__port}]")
+
+        await asyncio.create_task(self.__stop_signal.wait())
+
+        self.__log.info(f"Shutdown signal received, runtime API server is going down...")
+
+        await self.__server.stop(self.__DEFAULT_SHUTDOWN_TIMEOUT)
+        self.__server = None
+
+        self.__log.info("Runtime API server has gone down cleanly")
+
+    async def listJobs(self, request: runtime_pb2.RuntimeListJobsRequest, context: grpc.aio.ServicerContext):
+
+        request_task = ListJobsRequest(self.__engine_id, request, context)
+        self.__agent.threadsafe().spawn(request_task)
+
+        return await request_task.complete(self.__request_timeout)
+
+    async def getJobStatus(self, request: runtime_pb2.RuntimeJobInfoRequest, context: grpc.ServicerContext):
+
+        request_task = GetJobStatusRequest(self.__engine_id, request, context)
+        self.__agent.threadsafe().spawn(request_task)
+
+        return await request_task.complete(self.__request_timeout)
+
+    async def getJobResult(self, request: runtime_pb2.RuntimeJobInfoRequest, context: grpc.ServicerContext):
+
+        request_task = GetJobResultRequest(self.__engine_id, request, context)
+        self.__agent.threadsafe().spawn(request_task)
+
+        return await request_task.complete(self.__request_timeout)
+
+
+_T_REQUEST = tp.TypeVar("_T_REQUEST", bound=_msg.Message)
+_T_RESPONSE = tp.TypeVar("_T_RESPONSE", bound=_msg.Message)
+
+
+class ApiAgent(actors.ThreadsafeActor):
+
+    # API Agent is the parent actor that will be used to spawn API requests
+    # It must be created inside the asyncio event loop
+
+    def __init__(self):
+        super().__init__()
+        self._log = util.logger_for_object(self)
+        self._event_loop = asyncio.get_event_loop()
+        self.__start_signal = asyncio.Event()
+
+    def on_start(self):
+        self._event_loop.call_soon_threadsafe(lambda: self.__start_signal.set())
+
+    def on_signal(self, signal: actors.Signal) -> tp.Optional[bool]:
+
+        # Do not allow a failed request to bring down the API server
+        if signal.message == actors.SignalNames.FAILED:
+            error = signal.error if isinstance(signal, actors.ErrorSignal) else None
+            self._log.warning("Unhandled error during API request: " + str(error))
+            self._log.warning("The API agent will continue running")
+            return True
+
+        return False
+
+    async def started(self):
+        await self.__start_signal.wait()
+
+
+class ApiRequest(actors.ThreadsafeActor, tp.Generic[_T_REQUEST, _T_RESPONSE]):
+
+    # API request is the bridge between asyncio events (gRPC) and actor messages (TRAC runtime engine)
+    # Requests objects must be created inside the asyncio event loop
+
+    _log = None
+
+    def __init__(
+            self, engine_id, method: str, request: _T_REQUEST,
+            context: grpc.aio.ServicerContext):
+
+        super().__init__()
+
+        self._engine_id = engine_id
+        self._method = method
+        self._request = request
+        self._response: tp.Optional[_T_RESPONSE] = None
+        self._error: tp.Optional[Exception] = None
+        self._grpc_code = grpc.StatusCode.OK
+        self._grpc_message = ""
+
+        self._context = context
+        self._event_loop = asyncio.get_event_loop()
+        self._completion = asyncio.Event()
+
+        self._log.info("API call start: %s()", self._method)
+
+    def _mark_complete(self):
+
+        self._event_loop.call_soon_threadsafe(lambda: self._completion.set())
+
+    def on_stop(self):
+
+        if self.state() == actors.ActorState.ERROR:
+            self._error = self.error()
+
+        self._mark_complete()
+
+    async def complete(self, request_timeout: float) -> _T_RESPONSE:
+
+        try:
+
+            completion_task = asyncio.create_task(self._completion.wait())
+            await asyncio.wait_for(completion_task, request_timeout)
+
+            if self._error:
+                raise self._error
+
+            elif self._grpc_code != grpc.StatusCode.OK:
+                self._log.info("API call failed: %s() %s %s", self._method, self._grpc_code.name, self._grpc_message)
+                self._context.set_code(self._grpc_code)
+                self._context.set_details(self._grpc_message)
+
+            elif self._response is not None:
+                self._log.info("API call succeeded: %s()", self._method)
+                return self._response
+
+            else:
+                raise ex.EUnexpected()
+
+        except TimeoutError:
+            self._completion.set()
+            self._context.set_code(grpc.StatusCode.DEADLINE_EXCEEDED)
+            self._context.set_details("The TRAC runtime engine did not respond")
+            self._log.error("API call failed: %s() %s", self._method, "The TRAC runtime engine did not respond")
+            raise
+
+        except Exception as e:
+            self._context.set_code(grpc.StatusCode.INTERNAL)
+            self._context.set_details("Internal server error")
+            self._log.error("API call failed: %s() %s", self._method, str(e))
+            self._log.exception(e)
+            raise
+
+        finally:
+            self.threadsafe().stop()
+
+
+ApiRequest._log = util.logger_for_class(ApiRequest)
+
+
+class ListJobsRequest(ApiRequest[runtime_pb2.RuntimeListJobsRequest, runtime_pb2.RuntimeListJobsResponse]):
+
+    def __init__(self, engine_id, request, context):
+        super().__init__(engine_id, "get_job_list", request, context)
+
+    def on_start(self):
+        self.actors().send(self._engine_id, "get_job_list")
+
+    @actors.Message
+    def job_list(self, job_list):
+
+        self._response = runtime_pb2.RuntimeListJobsResponse(
+            jobs=codec.encode(job_list))
+
+        self._mark_complete()
+
+
+class GetJobStatusRequest(ApiRequest[runtime_pb2.RuntimeJobInfoRequest, runtime_pb2.RuntimeJobStatus]):
+
+    def __init__(self, engine_id, request, context):
+
+        super().__init__(engine_id, "get_job_status", request, context)
+
+        if request.HasField("jobKey"):
+            self._job_key = self._request.jobKey
+        elif request.HasField("jobSelector"):
+            self._job_key = util.object_key(self._request.jobSelector)
+        else:
+            raise ex.EValidation("Bad request: Neither jobKey nor jobSelector is specified")
+
+    def on_start(self):
+        self.actors().send(self._engine_id, "get_job_details", self._job_key, details=False)
+
+    @actors.Message
+    def job_details(self, job_details: tp.Optional[config.JobResult]):
+
+        if job_details is None:
+            self._grpc_code = grpc.StatusCode.NOT_FOUND
+            self._grpc_message = f"Job not found: [{self._job_key}]"
+
+        else:
+            self._response = runtime_pb2.RuntimeJobStatus(
+                jobId=codec.encode(job_details.jobId),
+                statusCode=codec.encode(job_details.statusCode),
+                statusMessage=codec.encode(job_details.statusMessage))
+
+        self._mark_complete()
+
+
+class GetJobResultRequest(ApiRequest[runtime_pb2.RuntimeJobInfoRequest, runtime_pb2.RuntimeJobResult]):
+
+    def __init__(self, engine_id, request, context):
+
+        super().__init__(engine_id, "get_job_result", request, context)
+
+        if request.HasField("jobKey"):
+            self._job_key = self._request.jobKey
+        elif request.HasField("jobSelector"):
+            self._job_key = util.object_key(self._request.jobSelector)
+        else:
+            raise ex.EValidation("Bad request: Neither jobKey nor jobSelector is specified")
+
+    def on_start(self):
+        self.actors().send(self._engine_id, "get_job_details", self._job_key, details=True)
+
+    @actors.Message
+    def job_details(self, job_details: tp.Optional[config.JobResult]):
+
+        if job_details is None:
+            self._grpc_code = grpc.StatusCode.NOT_FOUND
+            self._grpc_message = f"Job not found: [{self._job_key}]"
+
+        else:
+
+            encoded_results = dict((k, codec.encode(v)) for k, v in job_details.results.items())
+
+            self._response = runtime_pb2.RuntimeJobResult(
+                jobId=codec.encode(job_details.jobId),
+                statusCode=codec.encode(job_details.statusCode),
+                statusMessage=codec.encode(job_details.statusMessage),
+                results=encoded_results)
+
+        self._mark_complete()