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

ert/storage/local_experiment.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
+import contextlib
 import json
 import shutil
 from collections.abc import Generator
 from datetime import datetime
+from enum import StrEnum, auto
 from functools import cached_property
 from pathlib import Path
 from typing import TYPE_CHECKING, Annotated, Any
@@ -14,14 +16,11 @@ from pydantic import BaseModel, Field, TypeAdapter
 from surfio import IrapSurface
 
 from ert.config import (
-    EverestConstraintsConfig,
-    EverestObjectivesConfig,
-    ExtParamConfig,
-    GenDataConfig,
+    EverestControl,
     GenKwConfig,
+    KnownResponseTypes,
     ParameterConfig,
     ResponseConfig,
-    SummaryConfig,
     SurfaceConfig,
 )
 from ert.config import Field as FieldConfig
@@ -34,6 +33,20 @@ if TYPE_CHECKING:
     from .local_storage import LocalStorage
 
 
+class ExperimentState(StrEnum):
+    pending = auto()
+    running = auto()
+    completed = auto()
+    stopped = auto()
+    failed = auto()
+    never_run = auto()
+
+
+class ExperimentStatus(BaseModel):
+    message: str = Field(default="")
+    status: ExperimentState = Field(default=ExperimentState.pending)
+
+
 class _Index(BaseModel):
     id: UUID
     name: str
@@ -41,21 +54,19 @@ class _Index(BaseModel):
     # from a different experiment. For example, a manual update
     # is a separate experiment from the one that created the prior.
     ensembles: list[UUID]
+    status: ExperimentStatus | None = Field(default=None)
 
 
 _responses_adapter = TypeAdapter(  # type: ignore
     Annotated[
-        GenDataConfig
-        | SummaryConfig
-        | EverestConstraintsConfig
-        | EverestObjectivesConfig,
+        KnownResponseTypes,
         Field(discriminator="type"),
     ]
 )
 
 _parameters_adapter = TypeAdapter(  # type: ignore
     Annotated[
-        (GenKwConfig | SurfaceConfig | FieldConfig | ExtParamConfig),
+        (GenKwConfig | SurfaceConfig | FieldConfig | EverestControl),
         Field(discriminator="type"),
     ]
 )
@@ -73,6 +84,7 @@ class LocalExperiment(BaseMode):
     _responses_file = Path("responses.json")
     _metadata_file = Path("metadata.json")
     _templates_file = Path("templates.json")
+    _index_file = Path("index.json")
 
     def __init__(
         self,
@@ -97,7 +109,7 @@ class LocalExperiment(BaseMode):
         self._storage = storage
         self._path = path
         self._index = _Index.model_validate_json(
-            (path / "index.json").read_text(encoding="utf-8")
+            (path / self._index_file).read_text(encoding="utf-8")
         )
 
     @classmethod
@@ -147,9 +159,9 @@ class LocalExperiment(BaseMode):
             name = datetime.today().isoformat()
 
         storage._write_transaction(
-            path / "index.json",
+            path / cls._index_file,
             _Index(id=uuid, name=name, ensembles=[])
-            .model_dump_json(indent=2)
+            .model_dump_json(indent=2, exclude_none=True)
             .encode("utf-8"),
         )
 
@@ -181,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"),
@@ -280,8 +290,7 @@ class LocalExperiment(BaseMode):
         path = self.mount_point / self._metadata_file
         if not path.exists():
             raise ValueError(f"{self._metadata_file!s} does not exist")
-        with open(path, encoding="utf-8") as f:
-            return json.load(f)
+        return json.loads(path.read_text(encoding="utf-8"))
 
     @property
     def relative_weights(self) -> str:
@@ -295,25 +304,37 @@ class LocalExperiment(BaseMode):
     def id(self) -> UUID:
         return self._index.id
 
+    @property
+    def status(self) -> ExperimentStatus | None:
+        return self._index.status
+
+    @status.setter
+    @require_write
+    def status(self, status: ExperimentStatus) -> None:
+        if status != self._index.status:
+            self._index.status = status
+            self._storage._write_transaction(
+                self.mount_point / self._index_file,
+                self._index.model_dump_json(indent=2).encode("utf-8"),
+            )
+
     @property
     def mount_point(self) -> Path:
         return self._path
 
     @property
     def parameter_info(self) -> dict[str, Any]:
-        info: dict[str, Any]
-        with open(self.mount_point / self._parameter_file, encoding="utf-8") as f:
-            info = json.load(f)
-        return info
+        return json.loads(
+            (self.mount_point / self._parameter_file).read_text(encoding="utf-8")
+        )
 
     @property
     def templates_configuration(self) -> list[tuple[str, str]]:
         templates: list[tuple[str, str]] = []
-        try:
-            with open(self.mount_point / self._templates_file, encoding="utf-8") as f:
-                templates = json.load(f)
-        except (FileNotFoundError, json.JSONDecodeError):
-            pass
+        with contextlib.suppress(FileNotFoundError, json.JSONDecodeError):
+            templates = json.loads(
+                (self.mount_point / self._templates_file).read_text(encoding="utf-8")
+            )
         templates_with_content: list[tuple[str, str]] = []
         for source_file, target_file in templates:
             try:
@@ -327,10 +348,9 @@ class LocalExperiment(BaseMode):
 
     @property
     def response_info(self) -> dict[str, Any]:
-        info: dict[str, Any]
-        with open(self.mount_point / self._responses_file, encoding="utf-8") as f:
-            info = json.load(f)
-        return info
+        return json.loads(
+            (self.mount_point / self._responses_file).read_text(encoding="utf-8")
+        )
 
     def get_surface(self, name: str) -> IrapSurface:
         """
@@ -390,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
 
@@ -423,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
 
@@ -475,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

@@ -1,19 +1,18 @@
 from __future__ import annotations
 
-import contextlib
 import json
 import logging
 import os
 import re
 import shutil
+import types
 from collections.abc import Generator, MutableSequence
 from datetime import datetime
 from functools import cached_property
 from pathlib import Path
 from tempfile import NamedTemporaryFile
 from textwrap import dedent
-from types import TracebackType
-from typing import Any
+from typing import Any, Self
 from uuid import UUID, uuid4
 
 import polars as pl
@@ -21,6 +20,7 @@ import xarray as xr
 from filelock import FileLock, Timeout
 from pydantic import BaseModel, Field
 
+import ert.storage
 from ert.config import ErtConfig, ParameterConfig, ResponseConfig
 from ert.shared import __version__
 
@@ -31,7 +31,7 @@ from .realization_storage_state import RealizationStorageState
 
 logger = logging.getLogger(__name__)
 
-_LOCAL_STORAGE_VERSION = 14
+_LOCAL_STORAGE_VERSION = 21
 
 
 class _Migrations(BaseModel):
@@ -65,8 +65,7 @@ class LocalStorage(BaseMode):
         self,
         path: str | os.PathLike[str],
         mode: Mode,
-        *,
-        ignore_migration_check: bool = False,
+        stage_for_migration: bool = False,
     ) -> None:
         """
         Initializes the LocalStorage instance.
@@ -77,19 +76,22 @@ 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.
+        stage_for_migration : bool
+            Whether to avoid reloading storage to allow migration
         """
 
-        super().__init__(mode)
         self.path = Path(path).absolute()
+        super().__init__(mode)
 
-        self._experiments: dict[UUID, LocalExperiment]
-        self._ensembles: dict[UUID, LocalEnsemble]
-        self._index: _Index
+        if mode.can_write:
+            self._acquire_lock()
+
+        self._experiments: dict[UUID, LocalExperiment] = {}
+        self._ensembles: dict[UUID, LocalEnsemble] = {}
+        self._index: _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 = []
@@ -101,29 +103,45 @@ 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 version > _LOCAL_STORAGE_VERSION:
-            raise RuntimeError(
-                f"Cannot open storage '{self.path}': Storage version {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._ensure_fs_version_exists()
-            self._save_index()
-        elif version < _LOCAL_STORAGE_VERSION:
+        if self.check_migration_needed(Path(self.path)) and not self.can_write:
             raise RuntimeError(
                 f"Cannot open storage '{self.path}' in read-only mode: "
-                f"Storage version {version} is too old. Run ert to initiate migration."
+                f"Storage version {self.version} is too old. "
+                f"Run ert to initiate migration."
+            )
+
+        if not stage_for_migration:
+            self.reload()
+
+            if mode.can_write:
+                self._save_index()
+
+    @staticmethod
+    def check_migration_needed(storage_dir: Path) -> bool:
+        try:
+            version = _storage_version(storage_dir)
+        except FileNotFoundError:
+            version = _LOCAL_STORAGE_VERSION
+
+        if version > _LOCAL_STORAGE_VERSION:
+            raise ert.storage.ErtStorageException(
+                f"Cannot open storage '{storage_dir.absolute()}': Storage version "
+                f"{version} is newer than the current version {_LOCAL_STORAGE_VERSION}"
+                f", upgrade ert to continue, or run with a different ENSPATH"
             )
-        self.refresh()
 
-    def refresh(self) -> None:
+        return version < _LOCAL_STORAGE_VERSION
+
+    @staticmethod
+    def perform_migration(path: Path) -> None:
+        if LocalStorage.check_migration_needed(path):
+            with LocalStorage(path, Mode("w"), True) as storage:
+                storage._migrate(storage.version)
+                storage.reload()
+
+    def reload(self) -> None:
         """
         Reloads the index, experiments, and ensembles from the storage.
 
@@ -207,6 +225,12 @@ class LocalStorage(BaseMode):
             return _Index.model_validate_json(
                 (self.path / "index.json").read_text(encoding="utf-8")
             )
+        except PermissionError as e:
+            logger.error(
+                "Permission error when loading index from path: "
+                f"{self.path / 'index.json'}. Error: {e}",
+            )
+            raise e
         except FileNotFoundError:
             return _Index()
 
@@ -218,6 +242,12 @@ class LocalStorage(BaseMode):
             try:
                 ensemble = LocalEnsemble(self, ensemble_path, self.mode)
                 ensembles.append(ensemble)
+            except PermissionError as e:
+                logger.error(
+                    "Permission error when loading ensemble from path: "
+                    f"{ensemble_path}. Error: {e}",
+                )
+                raise e
             except FileNotFoundError:
                 logger.exception(
                     "Failed to load an ensemble from path: %s", ensemble_path
@@ -247,24 +277,17 @@ class LocalStorage(BaseMode):
     def _swap_path(self) -> Path:
         return self.path / self.SWAP_PATH
 
-    def __enter__(self) -> LocalStorage:
+    def __enter__(self) -> Self:
         return self
 
     def __exit__(
         self,
-        exception: Exception,
-        exception_type: type[Exception],
-        traceback: TracebackType,
+        exception: type[BaseException] | None,
+        exception_type: BaseException | None,
+        traceback: types.TracebackType | None,
     ) -> None:
         self.close()
 
-    @require_write
-    def _ensure_fs_version_exists(self) -> None:
-        # ERT 4 checks that this file exists and if it exists tells the user
-        # that their ERT storage is incompatible
-        with contextlib.suppress(FileExistsError):
-            (self.path / ".fs_version").symlink_to("index.json")
-
     @require_write
     def _acquire_lock(self) -> None:
         self._lock = FileLock(self.path / "storage.lock")
@@ -280,22 +303,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:
@@ -493,6 +510,13 @@ class LocalStorage(BaseMode):
             to12,
             to13,
             to14,
+            to15,
+            to16,
+            to17,
+            to18,
+            to19,
+            to20,
+            to21,
         )
 
         try:
@@ -522,7 +546,7 @@ class LocalStorage(BaseMode):
 
                 self._index = self._load_index()
 
-                logger.info("storage backed up for version less than 6")
+                logger.info("Storage backed up for version less than 5")
                 print(self._legacy_storage_migration_message(bkup_path, "14.6.*"))
                 return None
             elif version < _LOCAL_STORAGE_VERSION:
@@ -536,6 +560,13 @@ class LocalStorage(BaseMode):
                     11: to12,
                     12: to13,
                     13: to14,
+                    14: to15,
+                    15: to16,
+                    16: to17,
+                    17: to18,
+                    18: to19,
+                    19: to20,
+                    20: to21,
                 }
                 for from_version in range(version, _LOCAL_STORAGE_VERSION):
                     migrations[from_version].migrate(self.path)
@@ -553,7 +584,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")
@@ -623,16 +654,17 @@ class LocalStorage(BaseMode):
 def _storage_version(path: Path) -> int:
     if not path.exists():
         return _LOCAL_STORAGE_VERSION
-    try:
-        with open(path / "index.json", encoding="utf-8") as f:
-            return int(json.load(f)["version"])
-    except KeyError as exc:
-        raise NotImplementedError("Incompatible ERT Local Storage") from exc
-    except FileNotFoundError:
+    if not (path / "index.json").exists():
         if _is_block_storage(path):
             return 0
         else:
-            raise
+            raise FileNotFoundError(path / "index.json")
+    try:
+        return int(
+            json.loads((path / "index.json").read_text(encoding="utf-8"))["version"]
+        )
+    except KeyError as exc:
+        raise NotImplementedError("Incompatible ERT Local Storage") from exc
 
 
 _migration_ert_config: ErtConfig | None = None

ert/storage/migration/to10.py

@@ -22,8 +22,9 @@ def migrate(path: Path) -> None:
             shutil.copyfile(incoming_template, template_file_path)
             templates_abs.append((str(template_file_path.relative_to(experiment)), dst))
 
-        with open(experiment / "parameter.json", encoding="utf-8") as fin:
-            parameters_json = json.load(fin)
+        parameters_json = json.loads(
+            (experiment / "parameter.json").read_text(encoding="utf-8")
+        )
 
         for param in parameters_json.values():
             if param["_ert_kind"] == "GenKwConfig":

ert/storage/migration/to11.py

@@ -15,18 +15,17 @@ def migrate(path: Path) -> None:
         ensembles = path.glob("ensembles/*")
 
         experiment_id = None
-        with open(experiment / "index.json", encoding="utf-8") as f:
-            exp_index = json.load(f)
-            experiment_id = exp_index["id"]
+        exp_index = json.loads((experiment / "index.json").read_text(encoding="utf-8"))
+        experiment_id = exp_index["id"]
 
-        with open(experiment / "parameter.json", encoding="utf-8") as fin:
-            parameters_json = json.load(fin)
+        parameters_json = json.loads(
+            (experiment / "parameter.json").read_text(encoding="utf-8")
+        )
 
         for ens in ensembles:
-            with open(ens / "index.json", encoding="utf-8") as f:
-                ens_file = json.load(f)
-                if ens_file["experiment_id"] != experiment_id:
-                    continue
+            ens_file = json.loads((ens / "index.json").read_text(encoding="utf-8"))
+            if ens_file["experiment_id"] != experiment_id:
+                continue
 
             real_dirs = [*ens.glob("realization-*")]
             for param_config in parameters_json.values():
@@ -45,7 +44,7 @@ def migrate(path: Path) -> None:
                         array = ds.isel(realizations=0, drop=True)["values"]
                         realization = int(real_dir.name.split("-")[1])
 
-                        def parse_value(value: float | int | str) -> float | int | str:
+                        def parse_value(value: float | str) -> float | int | str:
                             if isinstance(value, float | int):
                                 return value
                             try:

ert/storage/migration/to12.py

@@ -25,22 +25,20 @@ def migrate_everest_param(config: dict[str, Any]) -> dict[str, Any]:
 
 def migrate(path: Path) -> None:
     def _replace_ert_kind(file: Path, kind_to_type: dict[str, str]) -> None:
-        with open(file, encoding="utf-8") as fin:
-            old_json = json.load(fin)
-            new_json = {}
+        old_json = json.loads(file.read_text(encoding="utf-8"))
+        new_json = {}
 
-            for key, config in old_json.items():
-                ert_kind = config.pop("_ert_kind")
+        for key, config in old_json.items():
+            ert_kind = config.pop("_ert_kind")
 
-                if ert_kind == "ExtParamConfig":
-                    new_json[key] = migrate_everest_param(config) | {
-                        "type": "everest_parameters"
-                    }
-                else:
-                    new_json[key] = config | {"type": kind_to_type[ert_kind]}
+            if ert_kind == "ExtParamConfig":
+                new_json[key] = migrate_everest_param(config) | {
+                    "type": "everest_parameters"
+                }
+            else:
+                new_json[key] = config | {"type": kind_to_type[ert_kind]}
 
-        with open(file, "w", encoding="utf-8") as fout:
-            json.dump(new_json, fout, indent=2)
+        file.write_text(json.dumps(new_json, indent=2), encoding="utf-8")
 
     for experiment in path.glob("experiments/*"):
         _replace_ert_kind(
@@ -63,12 +61,13 @@ def migrate(path: Path) -> None:
             },
         )
 
-        with open(experiment / "responses.json", encoding="utf-8") as fin:
-            old_json = json.load(fin)
-            new_json = {}
+        old_json = json.loads(
+            (experiment / "responses.json").read_text(encoding="utf-8")
+        )
+        new_json = {}
 
-            for key, config in old_json.items():
-                if config["type"] == "summary" and "refcase" in config:
-                    config.pop("refcase")
+        for key, config in old_json.items():
+            if config["type"] == "summary" and "refcase" in config:
+                config.pop("refcase")
 
-                new_json[key] = config
+            new_json[key] = config