Store optimization data using blobs (#13948)

Author: frode-aarstad

src/ert/storage/blob_data.py

@@ -17,6 +17,7 @@ class BlobType(StrEnum):
     MATRIX = "matrix"
     SCALING_FACTORS = "scaling_factors"
     RHO_MATRIX = "rho_matrix"
+    EVEREST_BATCH_DATA = "everest_batch_data"
 
 
 class ObservationReportData(BaseModel):
@@ -49,8 +50,17 @@ class RhoStorageData(_MatrixBase):
     observation_keys: list[str] = []
 
 
+class EverestBatchData(BaseModel):
+    blob_type: Literal[BlobType.EVEREST_BATCH_DATA] = BlobType.EVEREST_BATCH_DATA
+    dataframe_name: str
+
+
 BlobInfo = (
-    MatrixStorageData | ObservationReportData | ScalingFactorsData | RhoStorageData
+    MatrixStorageData
+    | ObservationReportData
+    | ScalingFactorsData
+    | RhoStorageData
+    | EverestBatchData
 )
 
 
@@ -70,7 +80,11 @@ class BlobStorageData(BaseModel):
     file_type: str
     name: str
     blob_info: Annotated[
-        MatrixStorageData | ObservationReportData | ScalingFactorsData | RhoStorageData,
+        MatrixStorageData
+        | ObservationReportData
+        | ScalingFactorsData
+        | RhoStorageData
+        | EverestBatchData,
         Discriminator("blob_type"),
     ]
 

src/ert/storage/local_ensemble.py

@@ -42,6 +42,7 @@
 from .blob_data import (
     BlobStorageData,
     BlobType,
+    EverestBatchData,
     MatrixStorageData,
     ObservationReportData,
     ScalingFactorsData,
@@ -1416,13 +1417,31 @@ def save_blob(
 
     @require_write
     def save_batch_dataframes(self, dataframes: BatchDataframes) -> None:
+        blob_dir = self._path / BLOB_DATA_DIR
         for df_name, df in dataframes.items():
-            if isinstance(df, pl.DataFrame):
-                df.write_parquet(self._path / f"{df_name}.parquet")
+            if not isinstance(df, pl.DataFrame):
+                continue
+            buf = io.BytesIO()
+            df.write_parquet(buf)
+            data = buf.getvalue()
+            BlobStorageData.save_blob(
+                name=df_name,
+                data=data,
+                blob_info=EverestBatchData(dataframe_name=df_name),
+                file_type="application/parquet",
+                storage=self._storage,
+                blob_dir=blob_dir,
+            )
 
     @property
     def has_function_results(self) -> bool:
-        return (self._path / "batch_objectives.parquet").exists()
+        for meta in self.load_blobs(BlobType.EVEREST_BATCH_DATA):
+            if (
+                isinstance(meta.blob_info, EverestBatchData)
+                and meta.blob_info.dataframe_name == "batch_objectives"
+            ):
+                return meta.file_size > 0
+        return False
 
     @property
     def has_gradient_results(self) -> bool:
@@ -1433,10 +1452,13 @@ def has_gradient_results(self) -> bool:
             info["perturbation"] != -1 for _, info in self.simulations_with_responses
         )
 
-    @staticmethod
-    def _read_df_if_exists(path: Path) -> pl.DataFrame | None:
-        if path.exists():
-            return pl.read_parquet(path)
+    def _read_batch_dataframe(self, dataframe_name: str) -> pl.DataFrame | None:
+        for meta in self.load_blobs(BlobType.EVEREST_BATCH_DATA):
+            if (
+                isinstance(meta.blob_info, EverestBatchData)
+                and meta.blob_info.dataframe_name == dataframe_name
+            ):
+                return pl.read_parquet(io.BytesIO(self.load_blob(meta.uri)))
         return None
 
     @property
@@ -1533,7 +1555,7 @@ def perturbation_controls(self) -> pl.DataFrame | None:
 
     @property
     def batch_objectives(self) -> pl.DataFrame | None:
-        return self._read_df_if_exists(self._path / "batch_objectives.parquet")
+        return self._read_batch_dataframe("batch_objectives")
 
     @property
     def realization_objectives(self) -> pl.DataFrame | None:
@@ -1570,7 +1592,7 @@ def realization_objectives(self) -> pl.DataFrame | None:
 
     @property
     def batch_constraints(self) -> pl.DataFrame | None:
-        return self._read_df_if_exists(self._path / "batch_constraints.parquet")
+        return self._read_batch_dataframe("batch_constraints")
 
     @property
     def realization_constraints(self) -> pl.DataFrame | None:
@@ -1610,25 +1632,19 @@ def realization_constraints(self) -> pl.DataFrame | None:
 
     @property
     def batch_bound_constraint_violations(self) -> pl.DataFrame | None:
-        return self._read_df_if_exists(
-            self._path / "batch_bound_constraint_violations.parquet"
-        )
+        return self._read_batch_dataframe("batch_bound_constraint_violations")
 
     @property
     def batch_input_constraint_violations(self) -> pl.DataFrame | None:
-        return self._read_df_if_exists(
-            self._path / "batch_input_constraint_violations.parquet"
-        )
+        return self._read_batch_dataframe("batch_input_constraint_violations")
 
     @property
     def batch_output_constraint_violations(self) -> pl.DataFrame | None:
-        return self._read_df_if_exists(
-            self._path / "batch_output_constraint_violations.parquet"
-        )
+        return self._read_batch_dataframe("batch_output_constraint_violations")
 
     @property
     def batch_objective_gradient(self) -> pl.DataFrame | None:
-        return self._read_df_if_exists(self._path / "batch_objective_gradient.parquet")
+        return self._read_batch_dataframe("batch_objective_gradient")
 
     @property
     def simulations(self) -> list[tuple[int, EverestRealizationInfo]]:
@@ -1686,7 +1702,7 @@ def perturbation_objectives(self) -> pl.DataFrame | None:
 
     @property
     def batch_constraint_gradient(self) -> pl.DataFrame | None:
-        return self._read_df_if_exists(self._path / "batch_constraint_gradient.parquet")
+        return self._read_batch_dataframe("batch_constraint_gradient")
 
     @property
     def perturbation_constraints(self) -> pl.DataFrame | None:

src/ert/storage/local_storage.py

@@ -31,7 +31,8 @@
 
 logger = logging.getLogger(__name__)
 
-_LOCAL_STORAGE_VERSION = 37
+
+_LOCAL_STORAGE_VERSION = 38
 
 
 def open_storage(
@@ -642,6 +643,7 @@ def _migrate(self, version: int) -> None:
             to35,
             to36,
             to37,
+            to38,
         )
 
         try:  # ruff: ignore[too-many-statements-in-try-clause]
@@ -708,6 +710,7 @@ def _migrate(self, version: int) -> None:
                     34: to35,
                     35: to36,
                     36: to37,
+                    37: to38,
                 }
                 for from_version in range(version, _LOCAL_STORAGE_VERSION):
                     migrations[from_version].migrate(self.path)

src/ert/storage/migration/to38.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+import json
+import logging
+import uuid as _uuid
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+info = "Move everest batch dataframes into ensemble blobs"
+
+_BATCH_DATAFRAME_NAMES = (
+    "batch_objectives",
+    "batch_constraints",
+    "batch_bound_constraint_violations",
+    "batch_input_constraint_violations",
+    "batch_output_constraint_violations",
+    "batch_objective_gradient",
+    "batch_constraint_gradient",
+)
+
+
+def _move_batch_dataframes_into_blobs(path: Path) -> None:
+    ensembles_dir = path / "ensembles"
+    if not ensembles_dir.exists():
+        return
+
+    for ens_dir in ensembles_dir.iterdir():
+        if not ens_dir.is_dir():
+            continue
+
+        for dataframe_name in _BATCH_DATAFRAME_NAMES:
+            parquet_file = ens_dir / f"{dataframe_name}.parquet"
+            if not parquet_file.exists():
+                continue
+
+            blob_dir = ens_dir / "blobs"
+            blob_dir.mkdir(parents=True, exist_ok=True)
+
+            data = parquet_file.read_bytes()
+            uri = f"{_uuid.uuid4().hex[:8]}.blob"
+            blob_data = {
+                "uri": uri,
+                "file_size": len(data),
+                "file_type": "application/parquet",
+                "name": dataframe_name,
+                "blob_info": {
+                    "blob_type": "everest_batch_data",
+                    "dataframe_name": dataframe_name,
+                },
+            }
+
+            (blob_dir / uri).write_bytes(data)
+            (blob_dir / f"{uri}.json").write_text(
+                json.dumps(blob_data, indent=2), encoding="utf-8"
+            )
+            parquet_file.unlink()
+            logger.info("Moved %s into blob %s", parquet_file, uri)
+
+
+def migrate(path: Path) -> None:
+    _move_batch_dataframes_into_blobs(path)

tests/ert/unit_tests/dark_storage/test_http_endpoints.py

@@ -4,6 +4,7 @@
 import re
 
 import pandas as pd
+import polars as pl
 import pytest
 from requests import Response
 from starlette.testclient import TestClient
@@ -379,6 +380,64 @@ def test_that_blob_endpoint_returns_blob_bytes(tmp_path, monkeypatch, dark_stora
     assert resp.content == blob_bytes
 
 
+def test_that_blobs_endpoint_lists_everest_batch_dataframes(
+    tmp_path, monkeypatch, dark_storage_app
+):
+    storage_path = tmp_path / "storage"
+    with open_storage(storage_path, mode="w") as storage:
+        experiment = storage.create_experiment(name="test-experiment")
+        ensemble = storage.create_ensemble(
+            experiment, ensemble_size=1, iteration=0, name="batch_0"
+        )
+        ensemble.save_batch_dataframes(
+            {
+                "batch_objectives": pl.DataFrame(
+                    {"batch_id": [0], "total_objective_value": [1.5]}
+                ),
+                "batch_objective_gradient": pl.DataFrame(
+                    {"batch_id": [0], "control_name": ["x"], "distance": [2.0]}
+                ),
+            }
+        )
+        ensemble_id = ensemble.id
+
+    monkeypatch.setenv("ERT_STORAGE_ENS_PATH", str(storage_path))
+    with TestClient(dark_storage_app) as client:
+        resp = client.get(f"/ensembles/{ensemble_id}/blobs")
+
+    assert resp.status_code == 200
+    blobs = resp.json()
+    by_name = {blob["name"]: blob for blob in blobs}
+    assert set(by_name) == {"batch_objectives", "batch_objective_gradient"}
+    for name, blob in by_name.items():
+        assert blob["file_type"] == "application/parquet"
+        assert blob["blob_info"]["blob_type"] == "everest_batch_data"
+        assert blob["blob_info"]["dataframe_name"] == name
+
+
+def test_that_blob_endpoint_returns_everest_batch_dataframe_parquet(
+    tmp_path, monkeypatch, dark_storage_app
+):
+    storage_path = tmp_path / "storage"
+    objectives = pl.DataFrame({"batch_id": [0], "total_objective_value": [1.5]})
+    with open_storage(storage_path, mode="w") as storage:
+        experiment = storage.create_experiment(name="test-experiment")
+        ensemble = storage.create_ensemble(
+            experiment, ensemble_size=1, iteration=0, name="batch_0"
+        )
+        ensemble.save_batch_dataframes({"batch_objectives": objectives})
+        [blob] = ensemble.load_blobs()
+        ensemble_id = ensemble.id
+
+    monkeypatch.setenv("ERT_STORAGE_ENS_PATH", str(storage_path))
+    with TestClient(dark_storage_app) as client:
+        resp = client.get(f"/ensembles/{ensemble_id}/blobs/{blob.uri}")
+
+    assert resp.status_code == 200
+    assert resp.headers["content-type"] == "application/octet-stream"
+    assert pl.read_parquet(io.BytesIO(resp.content)).equals(objectives)
+
+
 @pytest.mark.slow
 def test_get_record_observations(poly_example_tmp_dir, dark_storage_client):
     resp: Response = dark_storage_client.get("/experiments")