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

ert/services/ert_server.py

@@ -0,0 +1,317 @@
+from __future__ import annotations
+
+import json
+import logging
+import os
+import sys
+import threading
+from collections.abc import Mapping
+from inspect import Traceback
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+from time import sleep
+from typing import Any, 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
+
+
+class ErtServerContext:
+    def __init__(self, service: ErtServer) -> None:
+        self._service = service
+
+    def __enter__(self) -> ErtServer:
+        return self._service
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException],
+        exc_value: BaseException,
+        traceback: Traceback,
+    ) -> bool:
+        self._service.shutdown()
+        return exc_type is None
+
+
+class ErtServer:
+    _instance: ErtServer | None = None
+
+    def __init__(
+        self,
+        storage_path: str,
+        timeout: int = 120,
+        parent_pid: int | None = None,
+        connection_info: ErtServerConnectionInfo | Exception | None = None,
+        verbose: bool = False,
+        logging_config: str | None = None,  # Only used from everserver
+    ) -> None:
+        if timeout is None:
+            timeout = 120
+
+        self._storage_path = storage_path
+        self._connection_info: ErtServerConnectionInfo | Exception | None = (
+            connection_info
+        )
+        self._on_connection_info_received_event = threading.Event()
+        self._timeout = timeout
+        self._url: str | None = None
+
+        if self._connection_info is not None:
+            # This means that server is already running
+            if isinstance(connection_info, Mapping) and "urls" not in connection_info:
+                raise KeyError("No URLs found in connection info")
+
+            self._on_connection_info_received_event.set()
+            self._thread_that_starts_server_process = None
+            return
+
+        run_storage_main_cmd = [
+            sys.executable,
+            str(Path(__file__).parent / "_storage_main.py"),
+            "--project",
+            storage_path,
+        ]
+
+        if logging_config is not None:
+            run_storage_main_cmd += ["--logging-config", logging_config]
+
+            traceparent = get_traceparent()
+            if traceparent is not None:
+                run_storage_main_cmd += ["--traceparent", traceparent]
+
+        if parent_pid is not None:
+            run_storage_main_cmd += ["--parent_pid", str(parent_pid)]
+
+        if verbose:
+            run_storage_main_cmd.append("--verbose")
+
+        self._thread_that_starts_server_process = _Proc(
+            service_name="storage",
+            exec_args=run_storage_main_cmd,
+            timeout=120,
+            on_connection_info_received=self.on_connection_info_received_from_server_process,
+            project=Path(self._storage_path),
+        )
+
+    def fetch_auth(self) -> tuple[str, Any]:
+        """
+        Returns a tuple of username and password, compatible with requests' `auth`
+        kwarg.
+
+        Blocks while the server is starting.
+        """
+        return (
+            "__token__",
+            cast(dict[str, Any], self.fetch_connection_info())["authtoken"],
+        )
+
+    @classmethod
+    def init_service(
+        cls,
+        project: Path,
+        timeout: int = 0,
+        logging_config: str | None = None,
+    ) -> ErtServerContext:
+        try:
+            service = cls.connect(
+                project=project or Path.cwd(), timeout=0, logging_config=logging_config
+            )
+            # Check the server is up and running
+            _ = service.fetch_url()
+            return ErtServerContext(service)
+        except (TimeoutError, json.JSONDecodeError, KeyError) as e:
+            logging.getLogger(__name__).warning(
+                "Failed locating existing storage service due to "
+                f"{type(e).__name__}: {e}, starting new service"
+            )
+            return cls.start_server(
+                project=project, timeout=timeout, logging_config=logging_config
+            )
+        except PermissionError as pe:
+            logging.getLogger(__name__).error(
+                f"{type(pe).__name__}: {pe}, cannot connect to storage service "
+                f"due to permission issues."
+            )
+            raise pe
+
+    def fetch_url(self) -> str:
+        """Returns the url. Blocks while the server is starting"""
+        if self._url is not None:
+            return self._url
+
+        for url in self.fetch_connection_info()["urls"]:
+            con_info = self.fetch_connection_info()
+            try:
+                resp = requests.get(
+                    f"{url}/healthcheck",
+                    auth=self.fetch_auth(),
+                    verify=con_info["cert"],
+                )
+                logging.getLogger(__name__).info(
+                    f"Connecting to {url} got status: "
+                    f"{resp.status_code}, {resp.headers}, {resp.reason}, {resp.text}"
+                )
+                if resp.status_code == 200:
+                    self._url = url
+                    return str(url)
+
+            except requests.ConnectionError as ce:
+                logging.getLogger(__name__).info(
+                    f"Could not connect to {url}, but will try something else. "
+                    f"Error: {ce}"
+                )
+        raise TimeoutError(
+            "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")
+
+    def shutdown(self) -> int:
+        """Shutdown the server."""
+        if self._thread_that_starts_server_process is None:
+            return -1
+
+        self.__class__._instance = None
+        error_code = self._thread_that_starts_server_process.shutdown()
+        self._thread_that_starts_server_process = None
+
+        return error_code
+
+    @classmethod
+    def connect(
+        cls,
+        *,
+        project: os.PathLike[str],
+        timeout: int | None = None,
+        logging_config: str | None = None,
+    ) -> ErtServer:
+        if cls._instance is not None:
+            cls._instance.wait_until_ready()
+            assert isinstance(cls._instance, cls)
+            return cls._instance
+
+        path = Path(project)
+
+        # Wait for storage_server.json file to appear
+        try:
+            if timeout is None:
+                timeout = 240
+            t = -1
+            while t < timeout:
+                storage_server_path = path / "storage_server.json"
+                if (
+                    storage_server_path.exists()
+                    and storage_server_path.stat().st_size > 0
+                ):
+                    with (path / "storage_server.json").open() as f:
+                        storage_server_content = json.load(f)
+
+                    return ErtServer(
+                        storage_path=str(path),
+                        connection_info=storage_server_content,
+                        logging_config=logging_config,
+                    )
+
+                sleep(1)
+                t += 1
+
+            raise TimeoutError("Server not started")
+        except PermissionError as pe:
+            logging.getLogger(__name__).error(
+                f"{type(pe).__name__}: {pe}, cannot connect to ert server service "
+                f"due to permission issues."
+            )
+            raise pe
+
+    @classmethod
+    def start_server(
+        cls,
+        project: Path,
+        parent_pid: int | None = None,
+        verbose: bool = False,
+        timeout: int | None = None,
+        logging_config: str | None = None,
+    ) -> ErtServerContext:
+        if cls._instance is not None:
+            raise RuntimeError("Server already running")
+        cls._instance = obj = cls(
+            storage_path=str(project),
+            parent_pid=parent_pid,
+            verbose=verbose,
+            timeout=timeout or 120,
+            logging_config=logging_config,
+        )
+        if obj._thread_that_starts_server_process is not None:
+            obj._thread_that_starts_server_process.start()
+        return ErtServerContext(obj)
+
+    def on_connection_info_received_from_server_process(
+        self, info: ErtServerConnectionInfo | Exception | None
+    ) -> None:
+        if self._connection_info is not None:
+            raise ValueError("Connection information already set")
+        if info is None:
+            raise ValueError
+        self._connection_info = info
+
+        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"
+        else:
+            path = "storage_server.json"
+
+        if isinstance(info, Mapping):
+            with NamedTemporaryFile(dir=f"{self._storage_path}", delete=False) as f:
+                f.write(json.dumps(info, indent=4).encode("utf-8"))
+                f.flush()
+                os.rename(f.name, path)
+
+        self._on_connection_info_received_event.set()
+
+    def wait_until_ready(self, timeout: int | None = None) -> bool:
+        if timeout is None:
+            timeout = self._timeout
+
+        if self._on_connection_info_received_event.wait(timeout):
+            return not (
+                self._connection_info is None
+                or isinstance(self._connection_info, Exception)
+            )
+        if isinstance(self._connection_info, TimeoutError):
+            self.logger.critical(f"startup exceeded defined timeout {timeout}s")
+        return False  # Timeout reached
+
+    def fetch_connection_info(self) -> ErtServerConnectionInfo:
+        is_ready = self.wait_until_ready(self._timeout)
+        if isinstance(self._connection_info, Exception):
+            raise self._connection_info
+        if not is_ready:
+            raise TimeoutError
+        if self._connection_info is None:
+            raise ValueError("conn_info is None")
+        return self._connection_info
+
+    def wait(self) -> None:
+        if self._thread_that_starts_server_process is not None:
+            self._thread_that_starts_server_process.join()

ert/shared/_doc_utils/ert_jobs.py

@@ -180,10 +180,7 @@ class _ErtDocumentation(SphinxDirective):
 
 class ErtForwardModelDocumentation(_ErtDocumentation):
     pm = ErtPluginManager()
-    _JOBS: ClassVar[dict[str, Any]] = {
-        **pm.get_documentation_for_jobs(),
-        **pm.get_documentation_for_forward_model_steps(),
-    }
+    _JOBS: ClassVar[dict[str, Any]] = pm.get_documentation_for_forward_model_steps()
 
     def run(self) -> list[nodes.section]:
         return self._generate_job_documentation_without_title(

ert/shared/storage/connection.py

@@ -3,15 +3,15 @@ from __future__ import annotations
 import os
 from typing import Any
 
-from ert.services import StorageService
+from ert.services import ErtServer
 
 
 def get_info(
     project_id: os.PathLike[str],
 ) -> dict[str, str | tuple[str, Any]]:
-    client = StorageService.connect(project=project_id)
+    client = ErtServer.connect(project=project_id)
     return {
         "baseurl": client.fetch_url(),
         "auth": client.fetch_auth(),
-        "cert": client.fetch_conn_info()["cert"],
+        "cert": client.fetch_connection_info()["cert"],
     }

ert/shared/version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '17.1.9'
-__version_tuple__ = version_tuple = (17, 1, 9)
+__version__ = version = '18.0.0'
+__version_tuple__ = version_tuple = (18, 0, 0)
 
-__commit_id__ = commit_id = 'g37d22236c'
+__commit_id__ = commit_id = 'g8fc6e0edc'

ert/storage/local_ensemble.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import logging
 import os
 import time
@@ -660,10 +661,29 @@ class LocalEnsemble(BaseMode):
         iens_active_index: npt.NDArray[np.int_],
     ) -> None:
         config_node = self.experiment.parameter_configuration[param_group]
+        complete_df: pl.DataFrame | None = None
+        with contextlib.suppress(KeyError):
+            complete_df = self._load_parameters_lazy(SCALAR_FILENAME).collect(
+                engine="streaming"
+            )
         for real, ds in config_node.create_storage_datasets(
             parameters, iens_active_index
         ):
-            self.save_parameters(ds, config_node.name, real)
+            if isinstance(ds, pl.DataFrame):
+                if complete_df is None:
+                    complete_df = ds
+                else:
+                    complete_df = (
+                        complete_df.join(ds, on="realization", how="left")
+                        .unique(subset=["realization"], keep="first")
+                        .sort("realization")
+                    )
+            else:
+                self.save_parameters(ds, config_node.name, real)
+
+        group_path = self.mount_point / f"{_escape_filename(SCALAR_FILENAME)}.parquet"
+        if complete_df is not None:
+            self._storage._to_parquet_transaction(group_path, complete_df)
 
     def load_scalars(
         self, group: str | None = None, realizations: npt.NDArray[np.int_] | None = None
@@ -1329,7 +1349,7 @@ async def _write_responses_to_storage(
     for config in response_configs:
         try:
             start_time = time.perf_counter()
-            logger.debug(f"Starting to load response: {config.response_type}")
+            logger.debug(f"Starting to load response: {config.type}")
             try:
                 if isinstance(config, SummaryConfig) and realization == 0:
                     _log_grid_contents(
@@ -1344,14 +1364,14 @@ async def _write_responses_to_storage(
                 continue
             await asyncio.sleep(0)
             logger.debug(
-                f"Loaded {config.response_type}",
+                f"Loaded {config.type}",
                 extra={"Time": f"{(time.perf_counter() - start_time):.4f}s"},
             )
             start_time = time.perf_counter()
-            ensemble.save_response(config.response_type, ds, realization)
+            ensemble.save_response(config.type, ds, realization)
             await asyncio.sleep(0)
             logger.debug(
-                f"Saved {config.response_type} to storage",
+                f"Saved {config.type} to storage",
                 extra={"Time": f"{(time.perf_counter() - start_time):.4f}s"},
             )
         except Exception as err:

ert/storage/local_experiment.py

@@ -16,14 +16,11 @@ from pydantic import BaseModel, Field, TypeAdapter
 from surfio import IrapSurface
 
 from ert.config import (
-    EverestConstraintsConfig,
-    EverestObjectivesConfig,
     ExtParamConfig,
-    GenDataConfig,
     GenKwConfig,
+    KnownResponseTypes,
     ParameterConfig,
     ResponseConfig,
-    SummaryConfig,
     SurfaceConfig,
 )
 from ert.config import Field as FieldConfig
@@ -62,10 +59,7 @@ class _Index(BaseModel):
 
 _responses_adapter = TypeAdapter(  # type: ignore
     Annotated[
-        GenDataConfig
-        | SummaryConfig
-        | EverestConstraintsConfig
-        | EverestObjectivesConfig,
+        KnownResponseTypes,
         Field(discriminator="type"),
     ]
 )
@@ -199,9 +193,7 @@ class LocalExperiment(BaseMode):
 
         response_data = {}
         for response in responses or []:
-            response_data.update(
-                {response.response_type: response.model_dump(mode="json")}
-            )
+            response_data.update({response.type: response.model_dump(mode="json")})
         storage._write_transaction(
             path / cls._responses_file,
             json.dumps(response_data, default=str, indent=2).encode("utf-8"),
@@ -418,7 +410,7 @@ class LocalExperiment(BaseMode):
 
         for data in self.response_info.values():
             response_instance = _responses_adapter.validate_python(data)
-            responses[response_instance.response_type] = response_instance
+            responses[response_instance.type] = response_instance
 
         return responses
 
@@ -451,7 +443,7 @@ class LocalExperiment(BaseMode):
         mapping = {}
         for config in self.response_configuration.values():
             for key in config.keys if config.has_finalized_keys else []:
-                mapping[key] = config.response_type
+                mapping[key] = config.type
 
         return mapping
 
@@ -503,7 +495,7 @@ class LocalExperiment(BaseMode):
             self._path / self._responses_file,
             json.dumps(
                 {
-                    c.response_type: c.model_dump(mode="json")
+                    c.type: c.model_dump(mode="json")
                     for c in responses_configuration.values()
                 },
                 default=str,

ert/storage/local_storage.py

@@ -30,7 +30,7 @@ from .realization_storage_state import RealizationStorageState
 
 logger = logging.getLogger(__name__)
 
-_LOCAL_STORAGE_VERSION = 17
+_LOCAL_STORAGE_VERSION = 18
 
 
 class _Migrations(BaseModel):
@@ -64,8 +64,6 @@ class LocalStorage(BaseMode):
         self,
         path: str | os.PathLike[str],
         mode: Mode,
-        *,
-        ignore_migration_check: bool = False,
     ) -> None:
         """
         Initializes the LocalStorage instance.
@@ -76,19 +74,20 @@ class LocalStorage(BaseMode):
             The file system path to the storage.
         mode : Mode
             The access mode for the storage (read/write).
-        ignore_migration_check : bool
-            If True, skips migration checks during initialization.
         """
 
-        super().__init__(mode)
         self.path = Path(path).absolute()
+        super().__init__(mode)
+
+        if mode.can_write:
+            self._acquire_lock()
 
         self._experiments: dict[UUID, LocalExperiment]
         self._ensembles: dict[UUID, LocalEnsemble]
         self._index: _Index
 
         try:
-            version = _storage_version(self.path)
+            self.version = _storage_version(self.path)
         except FileNotFoundError as err:
             # No index json, will have a problem if other components of storage exists
             errors = []
@@ -100,26 +99,36 @@ class LocalStorage(BaseMode):
                 errors.append(f"ensemble path: {self.path / self.ENSEMBLES_PATH}")
             if errors:
                 raise ValueError(f"No index.json, but found: {errors}") from err
-            version = _LOCAL_STORAGE_VERSION
+            self.version = _LOCAL_STORAGE_VERSION
+
+        if self.check_migration_needed():
+            self.perform_migration()
 
-        if version > _LOCAL_STORAGE_VERSION:
+        self.refresh()
+        if mode.can_write:
+            self._save_index()
+
+    def check_migration_needed(self) -> bool:
+        if self.version > _LOCAL_STORAGE_VERSION:
             raise RuntimeError(
-                f"Cannot open storage '{self.path}': Storage version {version} "
+                f"Cannot open storage '{self.path}': Storage version {self.version} "
                 f"is newer than the current version {_LOCAL_STORAGE_VERSION}, "
                 "upgrade ert to continue, or run with a different ENSPATH"
             )
-        if self.can_write:
-            self._acquire_lock()
-            if version < _LOCAL_STORAGE_VERSION and not ignore_migration_check:
-                self._migrate(version)
-            self._index = self._load_index()
-            self._save_index()
-        elif version < _LOCAL_STORAGE_VERSION:
-            raise RuntimeError(
-                f"Cannot open storage '{self.path}' in read-only mode: "
-                f"Storage version {version} is too old. Run ert to initiate migration."
-            )
-        self.refresh()
+
+        return self.version < _LOCAL_STORAGE_VERSION
+
+    def perform_migration(self) -> None:
+        if self.check_migration_needed():
+            if self.can_write:
+                self._migrate(self.version)
+                self._save_index()
+            else:
+                raise RuntimeError(
+                    f"Cannot open storage '{self.path}' in read-only mode: "
+                    f"Storage version {self.version} is too old. "
+                    f"Run ert to initiate migration."
+                )
 
     def refresh(self) -> None:
         """
@@ -283,22 +292,16 @@ class LocalStorage(BaseMode):
     def close(self) -> None:
         """
         Closes the storage, releasing any acquired locks and saving the index.
-
         This method should be called to cleanly close the storage, especially
         when it was opened in write mode. Failing to call this method may leave
         a lock file behind, which would interfere with subsequent access to
         the storage.
         """
-
         self._ensembles.clear()
         self._experiments.clear()
 
-        if not self.can_write:
-            return
-
-        self._save_index()
-
         if self.can_write:
+            self._save_index()
             self._release_lock()
 
     def _release_lock(self) -> None:
@@ -499,6 +502,7 @@ class LocalStorage(BaseMode):
             to15,
             to16,
             to17,
+            to18,
         )
 
         try:
@@ -545,6 +549,7 @@ class LocalStorage(BaseMode):
                     14: to15,
                     15: to16,
                     16: to17,
+                    17: to18,
                 }
                 for from_version in range(version, _LOCAL_STORAGE_VERSION):
                     migrations[from_version].migrate(self.path)
@@ -562,7 +567,7 @@ class LocalStorage(BaseMode):
         Get a unique experiment name
 
         If an experiment with the given name exists an _0 is appended
-        or _n+1 where n is the the largest postfix found for the given experiment name
+        or _n+1 where n is the largest postfix found for the given experiment name
         """
         if not experiment_name:
             return self.get_unique_experiment_name("default")

ert/storage/migration/to18.py

@@ -0,0 +1,12 @@
+from pathlib import Path
+
+info = (
+    "Added localization attributes to summary observations."
+    "Added RFT observations."
+    "No change to current storage, only additions. "
+    "Bumping to 17 to indicate appended schema in storage."
+)
+
+
+def migrate(path: Path) -> None:
+    pass

ert/storage/migration/to8.py

@@ -138,14 +138,14 @@ def _migrate_observations_to_grouped_parquet(path: Path) -> None:
 
         for response_type in ["gen_data", "summary"]:
             infos = [
-                _info for _info in obs_ds_infos if _info.response_type == response_type
+                info_ for info_ in obs_ds_infos if info_.response_type == response_type
             ]
             if len(infos) > 0:
-                concatd_df = pl.concat([_info.polars_df for _info in infos])
+                concatd_df = pl.concat([info_.polars_df for info_ in infos])
                 concatd_df.write_parquet(experiment / "observations" / response_type)
 
-                for _info in infos:
-                    os.remove(_info.original_ds_path)
+                for info_ in infos:
+                    os.remove(info_.original_ds_path)
 
 
 def migrate(path: Path) -> None: