ert 16.0.9__py3-none-any.whl → 19.0.0rc2__py3-none-any.whl

ert/scheduler/lsf_driver.py

@@ -19,10 +19,15 @@ from typing import (
     get_args,
 )
 
-from .driver import SIGNAL_OFFSET, Driver, FailedSubmit, create_submit_script
+from .driver import (
+    _POLL_PERIOD,
+    SIGNAL_OFFSET,
+    Driver,
+    FailedSubmit,
+    create_submit_script,
+)
 from .event import DriverEvent, FinishedEvent, StartedEvent
 
-_POLL_PERIOD = 2.0  # seconds
 LSF_FAILED_JOB = SIGNAL_OFFSET + 65  # first non signal returncode
 """Return code we use when lsf reports failed jobs"""
 
@@ -257,6 +262,7 @@ class LsfDriver(Driver):
         bkill_cmd: str | None = None,
         bhist_cmd: str | None = None,
         activate_script: str = "",
+        poll_period: float = _POLL_PERIOD,
     ) -> None:
         super().__init__(activate_script)
         self._queue_name = queue_name
@@ -277,7 +283,7 @@ class LsfDriver(Driver):
         self._sleep_time_between_cmd_retries = 3
         self._max_bsub_attempts = 10
 
-        self._poll_period = _POLL_PERIOD
+        self._poll_period = poll_period
 
         self._bhist_cmd = Path(bhist_cmd or shutil.which("bhist") or "bhist")
         self._bhist_cache: dict[str, dict[str, int]] | None = None
@@ -438,6 +444,7 @@ class LsfDriver(Driver):
 
     async def poll(self) -> None:
         while True:
+            await self._warn_evaluator_if_polling_has_failed_for_some_time()
             if not self._jobs.keys():
                 await asyncio.sleep(self._poll_period)
                 continue
@@ -455,6 +462,7 @@ class LsfDriver(Driver):
                 )
             except OSError as e:
                 logger.error(str(e))
+                self._last_polling_error_message = str(e)
                 await asyncio.sleep(self._poll_period)
                 continue
 
@@ -462,10 +470,11 @@ class LsfDriver(Driver):
             if process.returncode:
                 # bjobs may give nonzero return code even when it is providing
                 # at least some correct information
+                error_msg = stderr.decode()
                 logger.warning(
-                    f"bjobs gave returncode {process.returncode} "
-                    f"and error {stderr.decode()}"
+                    f"bjobs gave returncode {process.returncode} and error {error_msg}"
                 )
+                self._last_polling_error_message = error_msg
             bjobs_states = _parse_jobs_dict(parse_bjobs(stdout.decode(errors="ignore")))
             self.update_and_log_exec_hosts(
                 parse_bjobs_exec_hosts(stdout.decode(errors="ignore"))
@@ -497,6 +506,7 @@ class LsfDriver(Driver):
                     "bhist did not give status for job_ids "
                     f"{missing_in_bhist_and_bjobs}, giving up for now."
                 )
+            self._last_successful_poll = time.time()
             await asyncio.sleep(self._poll_period)
 
     async def _process_job_update(self, job_id: str, new_state: AnyJob) -> None:

ert/scheduler/openpbs_driver.py

@@ -5,17 +5,17 @@ import json
 import logging
 import shlex
 import shutil
+import time
 from collections.abc import Iterable, Mapping, MutableMapping
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Literal, cast, get_type_hints
 
-from .driver import Driver, FailedSubmit, create_submit_script
+from .driver import _POLL_PERIOD, Driver, FailedSubmit, create_submit_script
 from .event import DriverEvent, FinishedEvent, StartedEvent
 
 logger = logging.getLogger(__name__)
 
-_POLL_PERIOD = 2.0  # seconds
 JOB_STATES = [
     "B",  # Begun
     "E",  # Exiting with or without errors
@@ -130,6 +130,7 @@ class OpenPBSDriver(Driver):
         qstat_cmd: str | None = None,
         qdel_cmd: str | None = None,
         activate_script: str = "",
+        poll_period: float = _POLL_PERIOD,
     ) -> None:
         super().__init__(activate_script)
 
@@ -140,7 +141,7 @@ class OpenPBSDriver(Driver):
         self._job_prefix = job_prefix
         self._max_pbs_cmd_attempts = 10
         self._sleep_time_between_cmd_retries = 2
-        self._poll_period = _POLL_PERIOD
+        self._poll_period = poll_period
 
         self._qsub_cmd = Path(qsub_cmd or shutil.which("qsub") or "qsub")
         self._qstat_cmd = Path(qstat_cmd or shutil.which("qstat") or "qstat")
@@ -264,6 +265,7 @@ class OpenPBSDriver(Driver):
 
     async def poll(self) -> None:
         while True:
+            await self._warn_evaluator_if_polling_has_failed_for_some_time()
             if not self._jobs:
                 await asyncio.sleep(self._poll_period)
                 continue
@@ -280,6 +282,7 @@ class OpenPBSDriver(Driver):
                     )
                 except OSError as e:
                     logger.error(str(e))
+                    self._last_polling_error_message = str(e)
                     await asyncio.sleep(self._poll_period)
                     continue
                 stdout, stderr = await process.communicate()
@@ -289,10 +292,12 @@ class OpenPBSDriver(Driver):
                     await asyncio.sleep(self._poll_period)
                     continue
                 if process.returncode == QSTAT_UNKNOWN_JOB_ID:
+                    error_msg = stderr.decode(errors="ignore")
                     logger.debug(
                         f"qstat gave returncode {QSTAT_UNKNOWN_JOB_ID} "
-                        f"with message {stderr.decode(errors='ignore')}"
+                        f"with message {error_msg}"
                     )
+                    self._last_polling_error_message = error_msg
                 parsed_jobs = _parse_jobs_dict(
                     parse_qstat(stdout.decode(errors="ignore"))
                 )
@@ -330,6 +335,7 @@ class OpenPBSDriver(Driver):
                 for job_id, job in parsed_jobs_dict.items():
                     await self._process_job_update(job_id, job)
 
+            self._last_successful_poll = time.time()
             await asyncio.sleep(self._poll_period)
 
     async def _process_job_update(self, job_id: str, new_state: AnyJob) -> None:

ert/scheduler/scheduler.py

@@ -15,6 +15,7 @@ import orjson
 from pydantic.dataclasses import dataclass
 
 from _ert.events import (
+    EnsembleEvaluationWarning,
     ForwardModelStepChecksum,
     RealizationEvent,
     RealizationFailed,
@@ -343,6 +344,10 @@ class Scheduler:
     async def _process_event_queue(self) -> None:
         while True:
             event = await self.driver.event_queue.get()
+            if isinstance(event, EnsembleEvaluationWarning):
+                if self._ensemble_evaluator_queue:
+                    await self._ensemble_evaluator_queue.put(event)
+                continue
             job = self._jobs[event.iens]
 
             # Any event implies the job has at least started

ert/scheduler/slurm_driver.py

@@ -14,7 +14,13 @@ from enum import Enum, auto
 from pathlib import Path
 from tempfile import NamedTemporaryFile
 
-from .driver import SIGNAL_OFFSET, Driver, FailedSubmit, create_submit_script
+from .driver import (
+    _POLL_PERIOD,
+    SIGNAL_OFFSET,
+    Driver,
+    FailedSubmit,
+    create_submit_script,
+)
 from .event import DriverEvent, FinishedEvent, StartedEvent
 
 SLURM_FAILED_EXIT_CODE_FETCH = SIGNAL_OFFSET + 66
@@ -77,9 +83,10 @@ class SlurmDriver(Driver):
         user: str | None = None,
         queue_name: str | None = None,
         max_runtime: float | None = None,
-        squeue_timeout: float = 2,
+        squeue_timeout: float | None = None,
         project_code: str | None = None,
         activate_script: str = "",
+        poll_period: float = _POLL_PERIOD,
     ) -> None:
         super().__init__(activate_script)
         self._submit_locks: dict[int, asyncio.Lock] = {}
@@ -109,7 +116,11 @@ class SlurmDriver(Driver):
 
         self._sleep_time_between_cmd_retries = 3
         self._sleep_time_between_kills = 30
-        self._poll_period = squeue_timeout
+
+        self._poll_period = (
+            squeue_timeout if squeue_timeout is not None else poll_period
+        )
+
         self._project_code = project_code
 
     def _submit_cmd(
@@ -251,6 +262,7 @@ class SlurmDriver(Driver):
 
     async def poll(self) -> None:
         while True:
+            await self._warn_evaluator_if_polling_has_failed_for_some_time()
             if not self._jobs.keys():
                 await asyncio.sleep(self._poll_period)
                 continue
@@ -266,14 +278,16 @@ class SlurmDriver(Driver):
                 )
             except OSError as e:
                 logger.error(str(e))
+                self._last_polling_error_message = str(e)
                 await asyncio.sleep(self._poll_period)
                 continue
             stdout, stderr = await process.communicate()
             if process.returncode:
+                error_msg = stderr.decode()
                 logger.warning(
-                    f"squeue gave returncode {process.returncode} "
-                    f"and error {stderr.decode()}"
+                    f"squeue gave returncode {process.returncode} and error {error_msg}"
                 )
+                self._last_polling_error_message = error_msg
             squeue_states = dict(_parse_squeue_output(stdout.decode(errors="ignore")))
 
             job_ids_found_in_squeue_output = set(squeue_states.keys())
@@ -306,6 +320,7 @@ class SlurmDriver(Driver):
                     "scontrol did not give status for job_ids "
                     f"{missing_in_squeue_and_scontrol}, giving up for now."
                 )
+            self._last_successful_poll = time.time()
             await asyncio.sleep(self._poll_period)
 
     async def _process_job_update(self, job_id: str, new_info: JobInfo) -> None:

ert/services/__init__.py

@@ -1,4 +1,4 @@
-from .storage_service import StorageService
+from .ert_server import ErtServer
 from .webviz_ert_service import WebvizErt
 
-__all__ = ["StorageService", "WebvizErt"]
+__all__ = ["ErtServer", "WebvizErt"]

ert/services/_base_service.py

@@ -1,3 +1,8 @@
+"""
+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
@@ -7,6 +12,7 @@ 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
@@ -15,13 +21,21 @@ 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, TypeVar
+from typing import TYPE_CHECKING, Any, Generic, Self, TypedDict, TypeVar
 
 if TYPE_CHECKING:
-    from inspect import Traceback
+    pass
 
 T = TypeVar("T", bound="BaseService")
-ConnInfo = Mapping[str, Any] | Exception | None
+
+
+class ErtServerConnectionInfo(TypedDict):
+    urls: list[str]
+    authtoken: str
+    host: str
+    port: str
+    cert: str
+    auth: str
 
 
 SERVICE_CONF_PATHS: set[str] = set()
@@ -74,9 +88,9 @@ class _Context(Generic[T]):
 
     def __exit__(
         self,
-        exc_type: type[BaseException],
-        exc_value: BaseException,
-        traceback: Traceback,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: types.TracebackType | None,
     ) -> bool:
         self._service.shutdown()
         return exc_type is None
@@ -88,7 +102,9 @@ class _Proc(threading.Thread):
         service_name: str,
         exec_args: Sequence[str],
         timeout: int,
-        set_conn_info: Callable[[ConnInfo], None],
+        on_connection_info_received: Callable[
+            [ErtServerConnectionInfo | Exception | None], None
+        ],
         project: Path,
     ) -> None:
         super().__init__()
@@ -98,7 +114,7 @@ class _Proc(threading.Thread):
         self._service_name = service_name
         self._exec_args = exec_args
         self._timeout = timeout
-        self._set_conn_info = set_conn_info
+        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()
@@ -119,13 +135,13 @@ class _Proc(threading.Thread):
         os.close(fd_write)
 
     def run(self) -> None:
-        comm = self._read_conn_info(self._childproc)
+        comm = self._read_connection_info_from_process(self._childproc)
 
         if comm is None:
-            self._set_conn_info(TimeoutError())
+            self._propagate_connection_info_from_childproc(TimeoutError())
             return  # _read_conn_info() has already cleaned up in this case
 
-        conn_info: ConnInfo = None
+        conn_info: ErtServerConnectionInfo | Exception | None = None
         try:
             conn_info = json.loads(comm)
         except json.JSONDecodeError:
@@ -134,7 +150,7 @@ class _Proc(threading.Thread):
             conn_info = exc
 
         try:
-            self._set_conn_info(conn_info)
+            self._propagate_connection_info_from_childproc(conn_info)
 
             while True:
                 if self._childproc.poll() is not None:
@@ -148,15 +164,16 @@ class _Proc(threading.Thread):
             self.logger.exception(e)
 
         finally:
-            self._ensure_delete_conn_info()
+            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_conn_info(self, proc: Popen[bytes]) -> str | None:
+    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:
@@ -166,7 +183,7 @@ class _Proc(threading.Thread):
             # Timeout reached, exit with a failure
             if ready == ([], [], []):
                 self._do_shutdown()
-                self._ensure_delete_conn_info()
+                self._ensure_connection_info_file_is_deleted()
                 return None
 
             x = self._comm_pipe.read(PIPE_BUF)
@@ -190,7 +207,7 @@ class _Proc(threading.Thread):
                     f"waiting for child-process exceeded timeout {self._timeout}s"
                 )
 
-    def _ensure_delete_conn_info(self) -> None:
+    def _ensure_connection_info_file_is_deleted(self) -> None:
         """
         Ensure that the JSON connection information file is deleted
         """
@@ -241,14 +258,14 @@ class BaseService:
         self,
         exec_args: Sequence[str] = (),
         timeout: int = 120,
-        conn_info: ConnInfo = None,
+        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: ConnInfo = conn_info
+        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()
 
@@ -261,7 +278,7 @@ class BaseService:
             )
 
     @classmethod
-    def start_server(cls: type[T], *args: Any, **kwargs: Any) -> _Context[T]:
+    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)
@@ -313,7 +330,7 @@ class BaseService:
         if self._proc is not None:
             self._proc.join()
 
-    def set_conn_info(self, info: ConnInfo) -> None:
+    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:

ert/services/_storage_main.py

@@ -28,7 +28,7 @@ from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapProp
 from uvicorn.supervisors import ChangeReload
 
 from ert.logging import STORAGE_LOG_CONFIG
-from ert.plugins import ErtPluginContext
+from ert.plugins import setup_site_logging
 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
@@ -274,10 +274,12 @@ def main() -> None:
         "ssl_keyfile_password": key_pw,
         "ssl_version": ssl.PROTOCOL_TLS_SERVER,
     }
-    with open(args.logging_config or STORAGE_LOG_CONFIG, encoding="utf-8") as conf_file:
-        logging_conf = yaml.safe_load(conf_file)
-        logging.config.dictConfig(logging_conf)
-        config_args.update(log_config=logging_conf)
+
+    logging_conf = yaml.safe_load(
+        Path(args.logging_config or STORAGE_LOG_CONFIG).read_text(encoding="utf-8")
+    )
+    logging.config.dictConfig(logging_conf)
+    config_args.update(log_config=logging_conf)
     warnings.filterwarnings("ignore", category=DeprecationWarning)
 
     if args.debug:
@@ -300,19 +302,19 @@ def main() -> None:
     terminate_on_parent_death_thread = threading.Thread(
         target=terminate_on_parent_death, args=[stopped, args.parent_pid, 1.0]
     )
-    with ErtPluginContext(logger=logging.getLogger()):
-        terminate_on_parent_death_thread.start()
-        with tracer.start_as_current_span("run_storage_server", ctx):
-            logger = logging.getLogger("ert.shared.storage.info")
-            try:
-                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, BaseServiceExit):
-                logger.info("Stopping dark storage")
-            finally:
-                stopped.set()
-                _join_terminate_thread(terminate_on_parent_death_thread)
+    setup_site_logging(logging.getLogger())
+    terminate_on_parent_death_thread.start()
+    with tracer.start_as_current_span("run_storage_server", ctx):
+        logger = logging.getLogger("ert.shared.storage.info")
+        try:
+            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, BaseServiceExit):
+            logger.info("Stopping dark storage")
+        finally:
+            stopped.set()
+            _join_terminate_thread(terminate_on_parent_death_thread)
 
 
 if __name__ == "__main__":