ert 19.0.0rc4__py3-none-any.whl → 20.0.0b0__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.0rc4'
-__version_tuple__ = version_tuple = (19, 0, 0, 'rc4')
+__version__ = version = '20.0.0b0'
+__version_tuple__ = version_tuple = (20, 0, 0, 'b0')
 
-__commit_id__ = commit_id = 'g2518c4485'
+__commit_id__ = commit_id = 'g23469e640'

ert/storage/local_ensemble.py

@@ -15,15 +15,18 @@ 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
 from pydantic import BaseModel
 from typing_extensions import TypedDict
 
-from ert.config import ParameterCardinality, ParameterConfig, SummaryConfig
-from ert.config.response_config import InvalidResponseFile
+from ert.config import (
+    InvalidResponseFile,
+    ParameterCardinality,
+    ParameterConfig,
+    SummaryConfig,
+)
 from ert.substitutions import substitute_runpath_name
 
 from .load_status import LoadResult
@@ -570,30 +573,51 @@ class LocalEnsemble(BaseMode):
     ) -> pl.DataFrame:
         if keys is None:
             keys = self.experiment.parameter_keys
-        elif set(keys) - set(self.experiment.parameter_keys):
-            missing = set(keys) - set(self.experiment.parameter_keys)
-            raise KeyError(f"Parameters not registered to the experiment: {missing}")
 
         df_lazy = self._load_parameters_lazy(SCALAR_FILENAME)
-        df_lazy = df_lazy.select(["realization", *keys])
+        names = df_lazy.collect_schema().names()
+        matches = [key for key in keys if any(key in item for item in names)]
+        if len(matches) != len(keys):
+            missing = set(keys) - set(matches)
+            raise KeyError(f"Parameters not registered to the experiment: {missing}")
+
+        parameter_keys = [
+            key
+            for e in keys
+            for key in self.experiment.parameter_configuration[e].parameter_keys
+        ]
+        df_lazy_filtered = df_lazy.select(["realization", *parameter_keys])
+
         if realizations is not None:
             if isinstance(realizations, int):
                 realizations = np.array([realizations])
-            df_lazy = df_lazy.filter(pl.col("realization").is_in(realizations))
-        df = df_lazy.collect(engine="streaming")
+            df_lazy_filtered = df_lazy_filtered.filter(
+                pl.col("realization").is_in(realizations)
+            )
+        df = df_lazy_filtered.collect(engine="streaming")
         if df.is_empty():
             raise IndexError(
                 f"No matching realizations {realizations} found for {keys}"
             )
 
         if transformed:
+            tmp_configuration: dict[str, ParameterConfig] = {}
+            for key in keys:
+                for col in df.columns:
+                    if col == "realization":
+                        continue
+                    if col.startswith(key):
+                        tmp_configuration[col] = (
+                            self.experiment.parameter_configuration[key]
+                        )
+
             df = df.with_columns(
                 [
                     pl.col(col)
                     .map_elements(
-                        self.experiment.parameter_configuration[col].transform_data(),
-                        return_dtype=df[col].dtype,
+                        tmp_configuration[col].transform_data(),
                     )
+                    .cast(df[col].dtype)
                     .alias(col)
                     for col in df.columns
                     if col != "realization"
@@ -618,13 +642,11 @@ class LocalEnsemble(BaseMode):
             for p in self.experiment.parameter_configuration.values()
             if group in {p.name, p.group_name}
         ]
-
         if not cfgs:
             raise KeyError(f"{group} is not registered to the experiment.")
 
         # if group refers to a group name, we expect the same cardinality
         cardinality = next(cfg.cardinality for cfg in cfgs)
-
         if cardinality == ParameterCardinality.multiple_configs_per_ensemble_dataset:
             return self.load_scalar_keys(
                 [cfg.name for cfg in cfgs], realizations, transformed
@@ -741,20 +763,21 @@ class LocalEnsemble(BaseMode):
     @staticmethod
     def sample_parameter(
         parameter: ParameterConfig,
-        real_nr: int,
+        active_realizations: list[int],
         random_seed: int,
+        num_realizations: int,
     ) -> pl.DataFrame:
-        parameter_value = parameter.sample_value(
-            str(random_seed),
-            real_nr,
+        parameter_values = parameter.sample_values(
+            str(random_seed), active_realizations, num_realizations=num_realizations
         )
 
-        parameter_dict = {parameter.name: parameter_value[0]}
-        parameter_dict["realization"] = real_nr
-        return pl.DataFrame(
-            parameter_dict,
-            schema={parameter.name: pl.Float64, "realization": pl.Int64},
+        parameters = pl.DataFrame(
+            {parameter.name: parameter_values},
+            schema={parameter.name: pl.Float64},
         )
+        realizations_series = pl.Series("realization", list(active_realizations))
+
+        return parameters.with_columns(realizations_series)
 
     def load_responses(self, key: str, realizations: tuple[int, ...]) -> pl.DataFrame:
         """Load responses for key and realizations into xarray Dataset.
@@ -1057,7 +1080,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",
@@ -1190,98 +1213,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/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/warnings/specific_warning_handler.py

@@ -7,7 +7,7 @@ from typing import TextIO
 @contextmanager
 def capture_specific_warning(
     warning_class_to_capture: type[Warning],
-    propagate_warning: Callable[[Warning | str], None],
+    propagate_warning: Callable[[Warning | str], None] | None = None,
 ) -> Generator[None, None, None]:
     original_warning_handler = warnings.showwarning
 
@@ -20,7 +20,8 @@ def capture_specific_warning(
         line: str | None = None,
     ) -> None:
         if issubclass(category, warning_class_to_capture):
-            propagate_warning(message)
+            if propagate_warning is not None:
+                propagate_warning(message)
         else:
             original_warning_handler(message, category, filename, lineno, file, line)
 

{ert-19.0.0rc4.dist-info → ert-20.0.0b0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ert
-Version: 19.0.0rc4
+Version: 20.0.0b0
 Summary: Ensemble based Reservoir Tool (ERT)
 Author-email: Equinor ASA <fg_sib-scout@equinor.com>
 License-Expression: GPL-3.0-only
@@ -45,9 +45,9 @@ Requires-Dist: opentelemetry-instrumentation-threading
 Requires-Dist: opentelemetry-sdk
 Requires-Dist: orjson
 Requires-Dist: packaging
-Requires-Dist: pandas
+Requires-Dist: pandas<3
 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
@@ -74,53 +74,6 @@ Requires-Dist: websockets
 Requires-Dist: xarray
 Requires-Dist: xtgeo>=3.3.0
 Requires-Dist: resfo-utilities>=0.5.0
-Provides-Extra: dev
-Requires-Dist: furo; extra == "dev"
-Requires-Dist: hypothesis!=6.102.0,!=6.112.3,>=6.85; extra == "dev"
-Requires-Dist: jsonpath_ng; extra == "dev"
-Requires-Dist: jupyter; extra == "dev"
-Requires-Dist: jupytext; extra == "dev"
-Requires-Dist: nbsphinx; extra == "dev"
-Requires-Dist: oil_reservoir_synthesizer; extra == "dev"
-Requires-Dist: pytest-asyncio; extra == "dev"
-Requires-Dist: pytest-benchmark; extra == "dev"
-Requires-Dist: pytest-cov; extra == "dev"
-Requires-Dist: pytest-memray; extra == "dev"
-Requires-Dist: pytest-mock; extra == "dev"
-Requires-Dist: pytest-mpl; extra == "dev"
-Requires-Dist: pytest-qt; extra == "dev"
-Requires-Dist: pytest-raises; extra == "dev"
-Requires-Dist: pytest-rerunfailures!=16.0; extra == "dev"
-Requires-Dist: pytest-snapshot; extra == "dev"
-Requires-Dist: pytest-timeout; extra == "dev"
-Requires-Dist: pytest-xdist; extra == "dev"
-Requires-Dist: pytest-durations; extra == "dev"
-Requires-Dist: pytest>6; extra == "dev"
-Requires-Dist: resdata; extra == "dev"
-Requires-Dist: rust-just; extra == "dev"
-Requires-Dist: scikit-image; extra == "dev"
-Requires-Dist: sphinx; extra == "dev"
-Requires-Dist: sphinx-argparse; extra == "dev"
-Requires-Dist: sphinx-autoapi; extra == "dev"
-Requires-Dist: sphinx-copybutton; extra == "dev"
-Requires-Dist: sphinxcontrib.datatemplates; extra == "dev"
-Requires-Dist: json-schema-for-humans; extra == "dev"
-Requires-Dist: xlsxwriter>=3.2.3; extra == "dev"
-Requires-Dist: resfo-utilities[testing]>=0.5.0; extra == "dev"
-Provides-Extra: style
-Requires-Dist: pre-commit; extra == "style"
-Provides-Extra: types
-Requires-Dist: mypy; extra == "types"
-Requires-Dist: types-lxml; extra == "types"
-Requires-Dist: types-requests; extra == "types"
-Requires-Dist: types-PyYAML; extra == "types"
-Requires-Dist: types-python-dateutil; extra == "types"
-Requires-Dist: types-decorator; extra == "types"
-Requires-Dist: types-docutils; extra == "types"
-Requires-Dist: types-tqdm; extra == "types"
-Requires-Dist: types-psutil; extra == "types"
-Requires-Dist: types-setuptools; extra == "types"
-Requires-Dist: types-networkx; extra == "types"
 Dynamic: license-file
 
 <h1 align="center">
@@ -184,7 +137,7 @@ Once uv is installed, you can get a development environment by running:
 ```sh
 git clone https://github.com/equinor/ert
 cd ert
-uv sync --all-extras
+uv sync --all-groups
 ```
 
 ### Test setup