ert 17.1.9__py3-none-any.whl → 18.0.0__py3-none-any.whl

ert/sample_prior.py

@@ -33,6 +33,7 @@ def sample_prior(
     parameter_configs = ensemble.experiment.parameter_configuration
     if parameters is None:
         parameters = list(parameter_configs.keys())
+    complete_dataset: pl.DataFrame | None = None
     for parameter in parameters:
         config_node = parameter_configs[parameter]
         if config_node.forward_init:
@@ -75,13 +76,17 @@ def sample_prior(
                 if datasets:
                     dataset = pl.concat(datasets, how="vertical")
 
-            if dataset is not None:
-                ensemble.save_parameters(
-                    dataset=dataset,
-                )
+            if complete_dataset is None:
+                complete_dataset = dataset
+            elif dataset is not None:
+                complete_dataset = complete_dataset.join(dataset, on="realization")
         else:
             for realization_nr in active_realizations:
                 ds = config_node.read_from_runpath(Path(), realization_nr, 0)
                 ensemble.save_parameters(ds, parameter, realization_nr)
 
+    if complete_dataset is not None:
+        ensemble.save_parameters(
+            dataset=complete_dataset,
+        )
     ensemble.refresh_ensemble_state()

ert/scheduler/driver.py

@@ -3,10 +3,13 @@ from __future__ import annotations
 import asyncio
 import logging
 import shlex
+import time
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
 from pathlib import Path
 
+from _ert.events import EnsembleEvaluationWarning
+
 from .event import DriverEvent
 
 SIGNAL_OFFSET = 128
@@ -34,12 +37,19 @@ class FailedSubmit(RuntimeError):
 class Driver(ABC):
     """Adapter for the HPC cluster."""
 
+    POLLING_TIMEOUT_PERIOD = 600
+
     def __init__(self, activate_script: str = "") -> None:
         self._event_queue: asyncio.Queue[DriverEvent] | None = None
         self._job_error_message_by_iens: dict[int, str] = {}
         self.activate_script = activate_script
         self._poll_period = _POLL_PERIOD
 
+        self._polling_timeout_period = Driver.POLLING_TIMEOUT_PERIOD
+        self._last_successful_poll = time.time()
+        self._last_polling_error_message: str | None = None
+        self._has_warned_evaluator_of_polling_error = False
+
     @property
     def event_queue(self) -> asyncio.Queue[DriverEvent]:
         if self._event_queue is None:
@@ -178,3 +188,30 @@ class Driver(ABC):
         )
         logger.error(error_message)
         return False, error_message
+
+    async def _warn_evaluator_if_polling_has_failed_for_some_time(self) -> None:
+        if (
+            (self._last_successful_poll < time.time() - self._polling_timeout_period)
+            and self._last_polling_error_message
+            and not self._has_warned_evaluator_of_polling_error
+        ):
+            await self._warn_evaluator_about_polling_difficulties()
+            self._has_warned_evaluator_of_polling_error = True
+
+    async def _warn_evaluator_about_polling_difficulties(self) -> None:
+        last_polling_error_message = self._last_polling_error_message
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            "Driver has not successfully polled statuses for "
+            f"{self._polling_timeout_period}s. The previous error "
+            f"was due to '{last_polling_error_message}'"
+        )
+        formatted_msg = (
+            "ert has not been able to update the job status for some time. This might "
+            "be resolved by itself, and it does not mean that the run has crashed.\n"
+            "Please check the runpath if it seems to still be running.\n"
+            f"The last error message was '{last_polling_error_message}'"
+        )
+        await self.event_queue.put(
+            EnsembleEvaluationWarning(warning_message=formatted_msg)
+        )

ert/scheduler/event.py

@@ -2,6 +2,8 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 
+from _ert.events import EnsembleEvaluationWarning
+
 
 @dataclass
 class StartedEvent:
@@ -16,4 +18,4 @@ class FinishedEvent:
     exec_hosts: str = "-"
 
 
-DriverEvent = StartedEvent | FinishedEvent
+DriverEvent = StartedEvent | FinishedEvent | EnsembleEvaluationWarning

ert/scheduler/job.py

@@ -12,6 +12,7 @@ from enum import StrEnum
 from pathlib import Path
 from typing import TYPE_CHECKING, assert_never
 
+import anyio
 from lxml import etree
 from opentelemetry.trace import Status, StatusCode
 
@@ -31,7 +32,7 @@ from ert.storage import (
     RealizationStorageState,
     load_realization_parameters_and_responses,
 )
-from ert.trace import trace, tracer
+from ert.trace import trace
 from ert.warnings import PostSimulationWarning
 
 from .driver import Driver, FailedSubmit
@@ -238,7 +239,6 @@ class Job:
                 f"{method_name} spent {elapsed_time} seconds waiting for files"
             )
 
-    @tracer.start_as_current_span(f"{__name__}.run")
     async def run(
         self,
         sem: asyncio.BoundedSemaphore,
@@ -335,8 +335,12 @@ class Job:
 
         valid_checksums = [info for info in checksum.values() if "error" not in info]
 
-        # Wait for files in checksum
-        while not all(Path(info["path"]).exists() for info in valid_checksums):
+        async def all_paths_exist(paths: list[Path]) -> bool:
+            return all(
+                await asyncio.gather(*[anyio.Path(path).exists() for path in paths])
+            )
+
+        while not await all_paths_exist([info["path"] for info in valid_checksums]):
             if timeout <= 0:
                 break
             timeout -= DISK_SYNCHRONIZATION_POLLING_INTERVAL
@@ -344,17 +348,19 @@ class Job:
             await asyncio.sleep(DISK_SYNCHRONIZATION_POLLING_INTERVAL)
         async with checksum_lock:
             for info in valid_checksums:
-                file_path = Path(info["path"])
+                file_path = anyio.Path(info["path"])
                 expected_md5sum = info.get("md5sum")
-                if file_path.exists() and expected_md5sum:
-                    actual_md5sum = hashlib.md5(file_path.read_bytes()).hexdigest()
+                file_path_exists = await file_path.exists()
+                if file_path_exists and expected_md5sum:
+                    file_bytes = await file_path.read_bytes()
+                    actual_md5sum = hashlib.md5(file_bytes).hexdigest()
                     if expected_md5sum == actual_md5sum:
                         logger.debug(f"File {file_path} checksum successful.")
                     else:
                         logger.warning(
                             f"File {file_path} checksum verification failed."
                         )
-                elif file_path.exists() and expected_md5sum is None:
+                elif file_path_exists and expected_md5sum is None:
                     logger.warning(f"Checksum not received for file {file_path}")
                 else:
                     logger.error(f"Disk synchronization failed for {file_path}")
@@ -506,11 +512,12 @@ async def log_warnings_from_forward_model(
             or "- ERROR - " in line
         )
 
-    async def log_warnings_from_file(  # noqa
+    async def log_warnings_from_file(
         file: Path, iens: int, step: ForwardModelStep, step_idx: int, filetype: str
     ) -> None:
         captured: list[str] = []
-        for line in file.read_text(encoding="utf-8").splitlines():
+        file_text = await anyio.Path(file).read_text(encoding="utf-8")
+        for line in file_text.splitlines():
             if line_contains_warning(line):
                 captured.append(line[:max_length])
 
@@ -527,9 +534,12 @@ async def log_warnings_from_forward_model(
             return 0
         remaining_timeout = _timeout
         for _ in range(_timeout):
-            if not (
-                file_path.exists() and file_path.stat().st_mtime >= job_submission_time
-            ):
+            file_path_exists = await anyio.Path(file_path).exists()
+            if file_path_exists:
+                st_mtime = (await anyio.Path(file_path).stat()).st_mtime
+            else:
+                st_mtime = 0
+            if not (file_path_exists and st_mtime >= job_submission_time):
                 remaining_timeout -= 1
                 await asyncio.sleep(1)
             else:

ert/scheduler/lsf_driver.py

@@ -444,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
@@ -461,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
 
@@ -468,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"))
@@ -503,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,6 +5,7 @@ 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
@@ -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

@@ -262,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
@@ -277,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())
@@ -317,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
@@ -15,13 +20,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
 
 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()
@@ -88,7 +101,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 +113,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 +134,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 +149,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 +163,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 +182,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 +206,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 +257,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()
 
@@ -313,7 +329,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: