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

ert/services/ert_server.py

@@ -1,23 +1,187 @@
 from __future__ import annotations
 
+import contextlib
+import io
 import json
 import logging
 import os
+import signal
 import sys
 import threading
 import types
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping, Sequence
 from pathlib import Path
+from select import PIPE_BUF, select
+from subprocess import Popen, TimeoutExpired
 from tempfile import NamedTemporaryFile
 from time import sleep
-from typing import Any, cast
+from typing import Any, TypedDict, cast
 
 import requests
 
 from ert.dark_storage.client import Client, ErtClientConnectionInfo
-from ert.services._base_service import ErtServerConnectionInfo, _Proc
 from ert.trace import get_traceparent
 
+SERVICE_CONF_PATHS: set[str] = set()
+
+
+class ErtServerConnectionInfo(TypedDict):
+    urls: list[str]
+    authtoken: str
+    host: str
+    port: str
+    cert: str
+    auth: str
+
+
+class ErtServerExit(OSError):
+    pass
+
+
+def cleanup_service_files(signum: int, frame: types.FrameType | None) -> None:
+    for file_path in SERVICE_CONF_PATHS:
+        file = Path(file_path)
+        if file.exists():
+            file.unlink()
+    raise ErtServerExit(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)
+
+
+class ServerBootFail(RuntimeError):
+    pass
+
+
+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) -> logging.Logger:
+        return logging.getLogger("ert.shared.storage")
+
+
+_ERT_SERVER_CONNECTION_INFO_FILE = "storage_server.json"
+_ERT_SERVER_EXECUTABLE_FILE = str(Path(__file__).parent / "_storage_main.py")
+
 
 class ErtServerContext:
     def __init__(self, service: ErtServer) -> None:
@@ -70,7 +234,7 @@ class ErtServer:
 
         run_storage_main_cmd = [
             sys.executable,
-            str(Path(__file__).parent / "_storage_main.py"),
+            _ERT_SERVER_EXECUTABLE_FILE,
             "--project",
             storage_path,
         ]
@@ -91,7 +255,7 @@ class ErtServer:
         self._thread_that_starts_server_process = _Proc(
             service_name="storage",
             exec_args=run_storage_main_cmd,
-            timeout=120,
+            timeout=timeout,
             on_connection_info_received=self.on_connection_info_received_from_server_process,
             project=Path(self._storage_path),
         )
@@ -167,21 +331,6 @@ class ErtServer:
             "None of the URLs provided for the ert storage server worked."
         )
 
-    @classmethod
-    def session(cls, project: os.PathLike[str], timeout: int | None = None) -> Client:
-        """
-        Start a HTTP transaction with the server
-        """
-        inst = cls.connect(timeout=timeout, project=project)
-        info = inst.fetch_connection_info()
-        return Client(
-            conn_info=ErtClientConnectionInfo(
-                base_url=inst.fetch_url(),
-                auth_token=inst.fetch_auth()[1],
-                cert=info["cert"],
-            )
-        )
-
     @property
     def logger(self) -> logging.Logger:
         return logging.getLogger("ert.shared.storage")
@@ -218,12 +367,12 @@ class ErtServer:
                 timeout = 240
             t = -1
             while t < timeout:
-                storage_server_path = path / "storage_server.json"
+                storage_server_path = path / _ERT_SERVER_CONNECTION_INFO_FILE
                 if (
                     storage_server_path.exists()
                     and storage_server_path.stat().st_size > 0
                 ):
-                    with (path / "storage_server.json").open() as f:
+                    with (path / _ERT_SERVER_CONNECTION_INFO_FILE).open() as f:
                         storage_server_content = json.load(f)
 
                     return ErtServer(
@@ -277,9 +426,9 @@ class ErtServer:
         if self._storage_path is not None:
             if not Path(self._storage_path).exists():
                 raise RuntimeError(f"No storage exists at : {self._storage_path}")
-            path = f"{self._storage_path}/storage_server.json"
+            path = f"{self._storage_path}/{_ERT_SERVER_CONNECTION_INFO_FILE}"
         else:
-            path = "storage_server.json"
+            path = _ERT_SERVER_CONNECTION_INFO_FILE
 
         if isinstance(info, Mapping):
             with NamedTemporaryFile(dir=f"{self._storage_path}", delete=False) as f:
@@ -315,3 +464,16 @@ class ErtServer:
     def wait(self) -> None:
         if self._thread_that_starts_server_process is not None:
             self._thread_that_starts_server_process.join()
+
+
+def create_ertserver_client(project: Path, timeout: int | None = None) -> Client:
+    """Read connection info from file in path and create HTTP client."""
+    connection = ErtServer.connect(timeout=timeout, project=project)
+    info = connection.fetch_connection_info()
+    return Client(
+        conn_info=ErtClientConnectionInfo(
+            base_url=connection.fetch_url(),
+            auth_token=connection.fetch_auth()[1],
+            cert=info["cert"],
+        )
+    )

ert/shared/version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '19.0.0'
-__version_tuple__ = version_tuple = (19, 0, 0)
+__version__ = version = '19.0.0rc1'
+__version_tuple__ = version_tuple = (19, 0, 0, 'rc1')
 
-__commit_id__ = commit_id = 'g2518c4485'
+__commit_id__ = commit_id = 'g6995ffc2f'

ert/storage/local_ensemble.py

@@ -15,7 +15,6 @@ from typing import TYPE_CHECKING
 from uuid import UUID
 
 import numpy as np
-import pandas as pd
 import polars as pl
 import resfo
 import xarray as xr
@@ -592,8 +591,8 @@ class LocalEnsemble(BaseMode):
                     pl.col(col)
                     .map_elements(
                         self.experiment.parameter_configuration[col].transform_data(),
-                        return_dtype=df[col].dtype,
                     )
+                    .cast(df[col].dtype)
                     .alias(col)
                     for col in df.columns
                     if col != "realization"
@@ -679,9 +678,6 @@ class LocalEnsemble(BaseMode):
                 if complete_df is None:
                     complete_df = ds
                 else:
-                    complete_df = complete_df.drop(
-                        [c for c in ds.columns if c != "realization"], strict=False
-                    )
                     complete_df = (
                         complete_df.join(ds, on="realization", how="left")
                         .unique(subset=["realization"], keep="first")
@@ -1057,7 +1053,7 @@ class LocalEnsemble(BaseMode):
                                 pl.col(col).is_in(observed_values.implode())
                             )
 
-                    pivoted = responses.collect(engine="streaming").pivot(  # noqa: PD010
+                    pivoted = responses.collect(engine="streaming").pivot(
                         on="realization",
                         index=["response_key", *response_cls.primary_key],
                         values="values",
@@ -1099,11 +1095,6 @@ class LocalEnsemble(BaseMode):
                             on=["response_key", *response_cls.primary_key],
                         )
 
-                    # Do not drop primary keys which
-                    # overlap with localization attributes
-                    primary_keys_to_drop = set(response_cls.primary_key).difference(
-                        {"north", "east", "radius"}
-                    )
                     joined = (
                         joined.with_columns(
                             pl.concat_str(
@@ -1113,7 +1104,7 @@ class LocalEnsemble(BaseMode):
                                 # Avoid potential collisions w/ primary key
                             )
                         )
-                        .drop(primary_keys_to_drop)
+                        .drop(response_cls.primary_key)
                         .rename({"__tmp_index_key__": "index"})
                     )
 
@@ -1129,9 +1120,6 @@ class LocalEnsemble(BaseMode):
                                 "observation_key",
                                 "observations",
                                 "std",
-                                "east",
-                                "north",
-                                "radius",
                             ]
                         )
 
@@ -1190,98 +1178,6 @@ class LocalEnsemble(BaseMode):
             self._path / "index.json", self._index.model_dump_json().encode("utf-8")
         )
 
-    @property
-    def all_parameters_and_gen_data(self) -> pl.DataFrame | None:
-        """
-        Only for Everest wrt objectives/constraints,
-        disregards summary data and primary key values
-        """
-        param_dfs = []
-        for param_group in self.experiment.parameter_configuration:
-            params_pd = self.load_parameters(param_group)["values"].to_pandas()
-
-            assert isinstance(params_pd, pd.DataFrame)
-            params_pd = params_pd.reset_index()
-            param_df = pl.from_pandas(params_pd)
-
-            param_columns = [c for c in param_df.columns if c != "realizations"]
-            param_df = param_df.rename(
-                {
-                    **{
-                        c: param_group + "." + c.replace("\0", ".")
-                        for c in param_columns
-                    },
-                    "realizations": "realization",
-                }
-            )
-            param_df = param_df.cast(
-                {
-                    "realization": pl.UInt16,
-                    **{c: pl.Float64 for c in param_df.columns if c != "realization"},
-                }
-            )
-            param_dfs.append(param_df)
-
-        responses = self.load_responses(
-            "gen_data", tuple(self.get_realization_list_with_responses())
-        )
-
-        if responses is None:
-            return pl.concat(param_dfs)
-
-        params_wide = pl.concat(
-            [
-                (
-                    pdf.sort("realization").drop("realization")
-                    if i > 0
-                    else pdf.sort("realization")
-                )
-                for i, pdf in enumerate(param_dfs)
-            ],
-            how="horizontal",
-        )
-
-        responses_wide = responses["realization", "response_key", "values"].pivot(  # noqa: PD010
-            on="response_key", values="values"
-        )
-
-        # If responses are missing for some realizations, this _left_ join will
-        # put null (polars) which maps to nan when doing .to_numpy() into the
-        # response columns for those realizations
-        params_and_responses = params_wide.join(
-            responses_wide, on="realization", how="left"
-        ).with_columns(pl.lit(self.iteration).alias("batch"))
-
-        assert self.everest_realization_info is not None
-
-        model_realization_mapping = {
-            k: v["model_realization"] for k, v in self.everest_realization_info.items()
-        }
-        perturbation_mapping = {
-            k: v["perturbation"] for k, v in self.everest_realization_info.items()
-        }
-
-        params_and_responses = params_and_responses.with_columns(
-            pl.col("realization")
-            .replace(model_realization_mapping)
-            .alias("model_realization"),
-            pl.col("realization")
-            .cast(pl.Int32)
-            .replace(perturbation_mapping)
-            .alias("perturbation"),
-        )
-
-        column_order = [
-            "batch",
-            "model_realization",
-            "perturbation",
-            "realization",
-            *[c for c in responses_wide.columns if c != "realization"],
-            *[c for c in params_wide.columns if c != "realization"],
-        ]
-
-        return params_and_responses[column_order]
-
 
 async def _read_parameters(
     run_path: str,

ert/storage/local_experiment.py

@@ -508,19 +508,3 @@ class LocalExperiment(BaseMode):
 
         if self.response_type_to_response_keys is not None:
             del self.response_type_to_response_keys
-
-    @property
-    def all_parameters_and_gen_data(self) -> pl.DataFrame | None:
-        if not self.ensembles:
-            return None
-
-        ensemble_dfs = [
-            e.all_parameters_and_gen_data
-            for e in self.ensembles
-            if e.all_parameters_and_gen_data is not None
-        ]
-
-        if not ensemble_dfs:
-            return None
-
-        return pl.concat(ensemble_dfs)

ert/storage/local_storage.py

@@ -31,7 +31,7 @@ from .realization_storage_state import RealizationStorageState
 
 logger = logging.getLogger(__name__)
 
-_LOCAL_STORAGE_VERSION = 23
+_LOCAL_STORAGE_VERSION = 22
 
 
 class _Migrations(BaseModel):
@@ -518,7 +518,6 @@ class LocalStorage(BaseMode):
             to20,
             to21,
             to22,
-            to23,
         )
 
         try:
@@ -570,7 +569,6 @@ class LocalStorage(BaseMode):
                     19: to20,
                     20: to21,
                     21: to22,
-                    22: to23,
                 }
                 for from_version in range(version, _LOCAL_STORAGE_VERSION):
                     migrations[from_version].migrate(self.path)

ert/utils/__init__.py

@@ -1,9 +1,13 @@
 import logging
 import time
 from collections.abc import Callable
+from datetime import UTC, datetime
 from functools import wraps
+from pathlib import Path
 from typing import ParamSpec, TypeVar
 
+from _ert.utils import file_safe_timestamp
+
 P = ParamSpec("P")
 R = TypeVar("R")
 
@@ -30,3 +34,19 @@ def log_duration(
         return wrapper
 
     return decorator
+
+
+def makedirs_if_needed(path: Path, roll_if_exists: bool = False) -> None:
+    if path.is_dir():
+        if not roll_if_exists:
+            return
+        _roll_dir(path)  # exists and should be rolled
+    path.mkdir(parents=True, exist_ok=False)
+
+
+def _roll_dir(old_name: Path) -> None:
+    old_name = old_name.resolve()
+    timestamp = file_safe_timestamp(datetime.now(UTC).isoformat())
+    new_name = f"{old_name}__{timestamp}"
+    old_name.rename(new_name)
+    logging.getLogger().info(f"renamed {old_name} to {new_name}")

{ert-19.0.0.dist-info → ert-19.0.0rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ert
-Version: 19.0.0
+Version: 19.0.0rc1
 Summary: Ensemble based Reservoir Tool (ERT)
 Author-email: Equinor ASA <fg_sib-scout@equinor.com>
 License-Expression: GPL-3.0-only
@@ -47,7 +47,7 @@ Requires-Dist: orjson
 Requires-Dist: packaging
 Requires-Dist: pandas
 Requires-Dist: pluggy>=1.3.0
-Requires-Dist: polars<1.35,>=1.32.3
+Requires-Dist: polars!=1.35,>=1.32.3
 Requires-Dist: progressbar2
 Requires-Dist: psutil
 Requires-Dist: pyarrow