ert 19.0.0rc1__py3-none-any.whl → 19.0.0rc3__py3-none-any.whl

ert/services/_base_service.py

@@ -0,0 +1,387 @@
+"""
+This file contains a more generic version of "ert services", and
+is scheduled for removal when WebvizErt is removed.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import io
+import json
+import os
+import signal
+import sys
+import threading
+import types
+from collections.abc import Callable, Mapping, Sequence
+from logging import Logger, getLogger
+from pathlib import Path
+from select import PIPE_BUF, select
+from subprocess import Popen, TimeoutExpired
+from tempfile import NamedTemporaryFile
+from time import sleep
+from types import FrameType
+from typing import TYPE_CHECKING, Any, Generic, Self, TypedDict, TypeVar
+
+if TYPE_CHECKING:
+    pass
+
+T = TypeVar("T", bound="BaseService")
+
+
+class ErtServerConnectionInfo(TypedDict):
+    urls: list[str]
+    authtoken: str
+    host: str
+    port: str
+    cert: str
+    auth: str
+
+
+SERVICE_CONF_PATHS: set[str] = set()
+
+
+class BaseServiceExit(OSError):
+    pass
+
+
+def cleanup_service_files(signum: int, frame: FrameType | None) -> None:
+    for file_path in SERVICE_CONF_PATHS:
+        file = Path(file_path)
+        if file.exists():
+            file.unlink()
+    raise BaseServiceExit(f"Signal {signum} received.")
+
+
+if threading.current_thread() is threading.main_thread():
+    signal.signal(signal.SIGTERM, cleanup_service_files)
+    signal.signal(signal.SIGINT, cleanup_service_files)
+
+
+def local_exec_args(script_args: str | list[str]) -> list[str]:
+    """
+    Convenience function that returns the exec_args for executing a Python
+    script in the directory of '_base_service.py'.
+
+    This is done instead of using 'python -m [module path]' due to the '-m' flag
+    adding the user's current working directory to sys.path. Executing a Python
+    script by itself will add the directory of the script rather than the
+    current working directory, thus we avoid accidentally importing user's
+    directories that just happen to have the same names as the ones we use.
+    """
+    if isinstance(script_args, str):
+        script = script_args
+        rest: list[str] = []
+    else:
+        script = script_args[0]
+        rest = script_args[1:]
+    script = f"_{script}_main.py"
+    return [sys.executable, str(Path(__file__).parent / script), *rest]
+
+
+class _Context(Generic[T]):
+    def __init__(self, service: T) -> None:
+        self._service = service
+
+    def __enter__(self) -> T:
+        return self._service
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: types.TracebackType | None,
+    ) -> bool:
+        self._service.shutdown()
+        return exc_type is None
+
+
+class _Proc(threading.Thread):
+    def __init__(
+        self,
+        service_name: str,
+        exec_args: Sequence[str],
+        timeout: int,
+        on_connection_info_received: Callable[
+            [ErtServerConnectionInfo | Exception | None], None
+        ],
+        project: Path,
+    ) -> None:
+        super().__init__()
+
+        self._shutdown = threading.Event()
+
+        self._service_name = service_name
+        self._exec_args = exec_args
+        self._timeout = timeout
+        self._propagate_connection_info_from_childproc = on_connection_info_received
+        self._service_config_path = project / f"{self._service_name}_server.json"
+
+        fd_read, fd_write = os.pipe()
+        self._comm_pipe = os.fdopen(fd_read)
+
+        env = os.environ.copy()
+        env["ERT_COMM_FD"] = str(fd_write)
+
+        SERVICE_CONF_PATHS.add(str(self._service_config_path))
+
+        # The process is waited for in _do_shutdown()
+        self._childproc = Popen(
+            self._exec_args,
+            pass_fds=(fd_write,),
+            env=env,
+            close_fds=True,
+        )
+        os.close(fd_write)
+
+    def run(self) -> None:
+        comm = self._read_connection_info_from_process(self._childproc)
+
+        if comm is None:
+            self._propagate_connection_info_from_childproc(TimeoutError())
+            return  # _read_conn_info() has already cleaned up in this case
+
+        conn_info: ErtServerConnectionInfo | Exception | None = None
+        try:
+            conn_info = json.loads(comm)
+        except json.JSONDecodeError:
+            conn_info = ServerBootFail()
+        except Exception as exc:
+            conn_info = exc
+
+        try:
+            self._propagate_connection_info_from_childproc(conn_info)
+
+            while True:
+                if self._childproc.poll() is not None:
+                    break
+                if self._shutdown.wait(1):
+                    self._do_shutdown()
+                    break
+
+        except Exception as e:
+            print(str(e))
+            self.logger.exception(e)
+
+        finally:
+            self._ensure_connection_info_file_is_deleted()
+
+    def shutdown(self) -> int:
+        """Shutdown the server."""
+        self._shutdown.set()
+        self.join()
+
+        return self._childproc.returncode
+
+    def _read_connection_info_from_process(self, proc: Popen[bytes]) -> str | None:
+        comm_buf = io.StringIO()
+        first_iter = True
+        while first_iter or proc.poll() is None:
+            first_iter = False
+            ready = select([self._comm_pipe], [], [], self._timeout)
+
+            # Timeout reached, exit with a failure
+            if ready == ([], [], []):
+                self._do_shutdown()
+                self._ensure_connection_info_file_is_deleted()
+                return None
+
+            x = self._comm_pipe.read(PIPE_BUF)
+            if not x:  # EOF
+                break
+            comm_buf.write(x)
+        return comm_buf.getvalue()
+
+    def _do_shutdown(self) -> None:
+        if self._childproc is None:
+            return
+        try:
+            self._childproc.terminate()
+            self._childproc.wait(10)  # Give it 10s to shut down cleanly..
+        except TimeoutExpired:
+            try:
+                self._childproc.kill()  # ... then kick it harder...
+                self._childproc.wait(self._timeout)  # ... and wait again
+            except TimeoutExpired:
+                self.logger.error(
+                    f"waiting for child-process exceeded timeout {self._timeout}s"
+                )
+
+    def _ensure_connection_info_file_is_deleted(self) -> None:
+        """
+        Ensure that the JSON connection information file is deleted
+        """
+        with contextlib.suppress(OSError):
+            if self._service_config_path.exists():
+                self._service_config_path.unlink()
+
+    @property
+    def logger(self) -> Logger:
+        return getLogger(f"ert.shared.{self._service_name}")
+
+
+class ServerBootFail(RuntimeError):
+    pass
+
+
+class BaseService:
+    """
+    BaseService provides a block-only-when-needed mechanism for starting and
+    maintaining services as subprocesses.
+
+    This is achieved by using a POSIX communication pipe, over which the service
+    can communicate that it has started. The contents of the communication is
+    also written to a file inside of the ERT storage directory.
+
+    The service itself can implement the other side of the pipe as such::
+
+        import os
+
+        # ... perform initialisation ...
+
+        # BaseService provides this environment variable with the pipe's FD
+        comm_fd = os.environ["ERT_COMM_FD"]
+
+        # Open the pipe with Python's IO classes for ease of use
+        with os.fdopen(comm_fd, "wb") as comm:
+            # Write JSON over the pipe, which will be interpreted by a subclass
+            # of BaseService on ERT's side
+            comm.write('{"some": "json"}')
+
+        # The pipe is flushed and closed here. This tells BaseService that
+        # initialisation is finished and it will try to read the JSON data.
+    """
+
+    _instance: BaseService | None = None
+
+    def __init__(
+        self,
+        exec_args: Sequence[str] = (),
+        timeout: int = 120,
+        conn_info: ErtServerConnectionInfo | Exception | None = None,
+        project: str | None = None,
+    ) -> None:
+        self._exec_args = exec_args
+        self._timeout = timeout
+
+        self._proc: _Proc | None = None
+        self._conn_info: ErtServerConnectionInfo | Exception | None = conn_info
+        self._conn_info_event = threading.Event()
+        self._project = Path(project) if project is not None else Path.cwd()
+
+        # Flag that we have connection information
+        if self._conn_info:
+            self._conn_info_event.set()
+        else:
+            self._proc = _Proc(
+                self.service_name, exec_args, timeout, self.set_conn_info, self._project
+            )
+
+    @classmethod
+    def start_server(cls, *args: Any, **kwargs: Any) -> _Context[Self]:
+        if cls._instance is not None:
+            raise RuntimeError("Server already running")
+        cls._instance = obj = cls(*args, **kwargs)
+        if obj._proc is not None:
+            obj._proc.start()
+        return _Context(obj)
+
+    @classmethod
+    def connect(
+        cls,
+        *,
+        project: os.PathLike[str],
+        timeout: int | None = None,
+    ) -> Self:
+        if cls._instance is not None:
+            cls._instance.wait_until_ready()
+            assert isinstance(cls._instance, cls)
+            return cls._instance
+
+        path = Path(project)
+        name = f"{cls.service_name}_server.json"
+        # Note: If the caller actually pass None, we override that here...
+        if timeout is None:
+            timeout = 240
+        t = -1
+        while t < timeout:
+            if (path / name).exists():
+                with (path / name).open() as f:
+                    return cls((), conn_info=json.load(f), project=str(path))
+
+            sleep(1)
+            t += 1
+
+        raise TimeoutError("Server not started")
+
+    def wait_until_ready(self, timeout: int | None = None) -> bool:
+        if timeout is None:
+            timeout = self._timeout
+
+        if self._conn_info_event.wait(timeout):
+            return not (
+                self._conn_info is None or isinstance(self._conn_info, Exception)
+            )
+        if isinstance(self._conn_info, TimeoutError):
+            self.logger.critical(f"startup exceeded defined timeout {timeout}s")
+        return False  # Timeout reached
+
+    def wait(self) -> None:
+        if self._proc is not None:
+            self._proc.join()
+
+    def set_conn_info(self, info: ErtServerConnectionInfo | Exception | None) -> None:
+        if self._conn_info is not None:
+            raise ValueError("Connection information already set")
+        if info is None:
+            raise ValueError
+        self._conn_info = info
+
+        if self._project is not None:
+            if not Path(self._project).exists():
+                raise RuntimeError(f"No storage exists at : {self._project}")
+            path = f"{self._project}/{self.service_name}_server.json"
+        else:
+            path = f"{self.service_name}_server.json"
+
+        if isinstance(info, Mapping):
+            with NamedTemporaryFile(dir=f"{self._project}", delete=False) as f:
+                f.write(json.dumps(info, indent=4).encode("utf-8"))
+                f.flush()
+                os.rename(f.name, path)
+
+        self._conn_info_event.set()
+
+    def fetch_conn_info(self) -> Mapping[str, Any]:
+        is_ready = self.wait_until_ready(self._timeout)
+        if isinstance(self._conn_info, Exception):
+            raise self._conn_info
+        if not is_ready:
+            raise TimeoutError
+        if self._conn_info is None:
+            raise ValueError("conn_info is None")
+        return self._conn_info
+
+    def shutdown(self) -> int:
+        """Shutdown the server."""
+        if self._proc is None:
+            return -1
+        self.__class__._instance = None
+        proc, self._proc = self._proc, None
+        return proc.shutdown()
+
+    @property
+    def service_name(self) -> str:
+        """
+        Subclass should return the name of the service, eg 'storage' for ERT Storage.
+        Used for identifying the server information JSON file.
+        """
+        raise NotImplementedError
+
+    @property
+    def logger(self) -> Logger:
+        return getLogger(f"ert.shared.{self.service_name}")
+
+    @property
+    def _service_file(self) -> str:
+        return f"{self.service_name}_server.json"

ert/services/_storage_main.py

@@ -13,7 +13,6 @@ import sys
 import threading
 import time
 import warnings
-from argparse import ArgumentParser
 from base64 import b64encode
 from pathlib import Path
 from typing import Any
@@ -30,11 +29,12 @@ from uvicorn.supervisors import ChangeReload
 
 from ert.logging import STORAGE_LOG_CONFIG
 from ert.plugins import setup_site_logging
-from ert.services import ErtServerExit
+from ert.services._base_service import BaseServiceExit
 from ert.shared import __file__ as ert_shared_path
 from ert.shared import find_available_socket, get_machine_name
+from ert.shared.storage.command import add_parser_options
 from ert.trace import tracer
-from ert.utils import makedirs_if_needed
+from everest.util import makedirs_if_needed
 
 DARK_STORAGE_APP = "ert.dark_storage.app:app"
 
@@ -82,7 +82,7 @@ def _get_host_list() -> list[str]:
 
 
 def _create_connection_info(
-    sock: socket.socket, authtoken: str, cert: str | os.PathLike[str] | Path
+    sock: socket.socket, authtoken: str, cert: str | os.PathLike[str]
 ) -> dict[str, Any]:
     connection_info = {
         "urls": [
@@ -91,7 +91,7 @@ def _create_connection_info(
         "authtoken": authtoken,
         "host": get_machine_name(),
         "port": sock.getsockname()[1],
-        "cert": str(cert),
+        "cert": cert,
         "auth": authtoken,
     }
 
@@ -102,17 +102,14 @@ def _create_connection_info(
     return connection_info
 
 
-def _generate_certificate(cert_folder: Path) -> tuple[Path, Path, bytes]:
+def _generate_certificate(cert_folder: str) -> tuple[str, str, bytes]:
     """Generate a private key and a certificate signed with it
 
     Both the certificate and the key are written to files in the folder given
     by `get_certificate_dir(config)`. The key is encrypted before being
     stored.
-
-    Returns a 3-tuple with
-        * Certificate file path
-        * Key file path
-        * Password used for encrypting the key
+    Returns the path to the certificate file, the path to the key file, and
+    the password used for encrypting the key
     """
     # Generate private key
     key = rsa.generate_private_key(
@@ -153,11 +150,11 @@ def _generate_certificate(cert_folder: Path) -> tuple[Path, Path, bytes]:
 
     # Write certificate and key to disk
     makedirs_if_needed(cert_folder)
-    cert_path = cert_folder / f"{dns_name}.crt"
-    cert_path.write_bytes(cert.public_bytes(serialization.Encoding.PEM))
-    key_path = cert_folder / f"{dns_name}.key"
+    cert_path = os.path.join(cert_folder, dns_name + ".crt")
+    Path(cert_path).write_bytes(cert.public_bytes(serialization.Encoding.PEM))
+    key_path = os.path.join(cert_folder, dns_name + ".key")
     pw = bytes(os.urandom(28))
-    key_path.write_bytes(
+    Path(key_path).write_bytes(
         key.private_bytes(
             encoding=serialization.Encoding.PEM,
             format=serialization.PrivateFormat.TraditionalOpenSSL,
@@ -187,16 +184,16 @@ def run_server(
 
     config_args: dict[str, Any] = {}
     if args.debug or debug:
-        config_args.update(reload=True, reload_dirs=[Path(ert_shared_path).parent])
+        config_args.update(reload=True, reload_dirs=[os.path.dirname(ert_shared_path)])
         os.environ["ERT_STORAGE_DEBUG"] = "1"
 
-    sock: socket.socket = find_available_socket(
+    sock = find_available_socket(
         host=get_machine_name(), port_range=range(51850, 51870 + 1)
     )
 
     # Appropriated from uvicorn.main:run
     os.environ["ERT_STORAGE_NO_TOKEN"] = "1"
-    os.environ["ERT_STORAGE_ENS_PATH"] = str(args.project.absolute())
+    os.environ["ERT_STORAGE_ENS_PATH"] = os.path.abspath(args.project)
     config = (
         # uvicorn.Config() resets the logging config (overriding additional
         # handlers added to loggers like e.g. the ert_azurelogger handler
@@ -256,11 +253,11 @@ def _join_terminate_thread(terminate_on_parent_death_thread: threading.Thread) -
     """Join the terminate thread, handling BaseServiceExit (which is used by Everest)"""
     try:
         terminate_on_parent_death_thread.join()
-    except ErtServerExit:
+    except BaseServiceExit:
         logger = logging.getLogger("ert.shared.storage.info")
         logger.info(
             "Got BaseServiceExit while joining terminate thread, "
-            "as expected from ert_server.py"
+            "as expected from _base_service.py"
         )
 
 
@@ -268,7 +265,9 @@ def main() -> None:
     args = parse_args()
     authentication = _generate_authentication()
     os.environ["ERT_STORAGE_TOKEN"] = authentication
-    cert_path, key_path, key_pw = _generate_certificate(args.project / "cert")
+    cert_path, key_path, key_pw = _generate_certificate(
+        os.path.join(args.project, "cert")
+    )
     config_args: dict[str, Any] = {
         "ssl_keyfile": key_path,
         "ssl_certfile": cert_path,
@@ -284,7 +283,7 @@ def main() -> None:
     warnings.filterwarnings("ignore", category=DeprecationWarning)
 
     if args.debug:
-        config_args.update(reload=True, reload_dirs=[Path(ert_shared_path).parent])
+        config_args.update(reload=True, reload_dirs=[os.path.dirname(ert_shared_path)])
 
     # Need to run uvicorn.Config before entering the ErtPluginContext because
     # uvicorn.Config overrides the configuration of existing loggers, thus removing
@@ -311,48 +310,12 @@ def main() -> None:
             logger.info("Starting dark storage")
             logger.info(f"Started dark storage with parent {args.parent_pid}")
             run_server(args, debug=False, uvicorn_config=uvicorn_config)
-        except (SystemExit, ErtServerExit):
+        except (SystemExit, BaseServiceExit):
             logger.info("Stopping dark storage")
         finally:
             stopped.set()
             _join_terminate_thread(terminate_on_parent_death_thread)
 
 
-def add_parser_options(ap: ArgumentParser) -> None:
-    ap.add_argument(
-        "config",
-        type=str,
-        help=("ERT config file to start the server from "),
-        nargs="?",  # optional
-    )
-    ap.add_argument(
-        "--project",
-        "-p",
-        type=Path,
-        help="Path to directory in which to create storage_server.json",
-        default=Path.cwd(),
-    )
-    ap.add_argument(
-        "--traceparent",
-        type=str,
-        help="Trace parent id to be used by the storage root span",
-        default=None,
-    )
-    ap.add_argument(
-        "--parent_pid",
-        type=int,
-        help="The parent process id",
-        default=os.getppid(),
-    )
-    ap.add_argument(
-        "--host", type=str, default=os.environ.get("ERT_STORAGE_HOST", "127.0.0.1")
-    )
-    ap.add_argument("--logging-config", type=str, default=None)
-    ap.add_argument(
-        "--verbose", action="store_true", help="Show verbose output.", default=False
-    )
-    ap.add_argument("--debug", action="store_true", default=False)
-
-
 if __name__ == "__main__":
     main()