MultiBackendJobManager: Dynamically add all usage metrics in job DB. … (#947)

* MultiBackendJobManager: Dynamically add all usage metrics in job DB. Remove hardcoded usage metrics from required columns

* Update job manager unit tests to no longer require costs and duration metrics

* Inintialize non-existing columns in job DB as None

* Further clean-up of MBJM unit tests

* Ensure backwards compatibility of MBJM for legacy job databases + add unit test for it

Author: VincentVerelst

CHANGELOG.md

@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Make `_DerivedFrom._from_url` more resilient against unresolvable/unparsable `derived_from` links ([#928](https://github.com/Open-EO/openeo-python-client/issues/928), eu-cdse/openeo-cdse-infra#1338)
 - `download_url()`: favor explicit HEAD status check over unhelpful auto-check ([#939](https://github.com/Open-EO/openeo-python-client/issues/939))
 - Modernize Xarray version support: remove `<2025.0.2` upper limit but require at least `2024.07.0` ([#721](https://github.com/Open-EO/openeo-python-client/issues/721))
+- `MultiBackendJobManager`: all fields reported under "usage" in the job metadata are now dynamically included as columns in the job database, instead of only the hardcoded "cpu", "memory" and "duration" fields (which are no longer predefined columns)
 
 ### Removed
 

openeo/extra/job_management/_job_db.py

@@ -87,6 +87,16 @@ def _merge_into_df(self, df: pd.DataFrame):
             unknown_indices = set(df.index).difference(df.index)
             if unknown_indices:
                 _log.warning(f"Merging DataFrame with {unknown_indices=} which will be lost.")
+            # `DataFrame.update` ignores columns that don't exist yet in `self._df`, so add them first
+            # (e.g. dynamically added columns for job usage metadata).
+            new_columns = [c for c in df.columns if c not in self._df.columns]
+            for column in new_columns:
+                self._df[column] = None
+            # Align dtypes to avoid incompatible-dtype warnings/errors on update
+            # (e.g. empty legacy columns loaded as float64, updated with strings).
+            for column in df.columns.intersection(self._df.columns):
+                if df[column].dtype == object and self._df[column].dtype != object:
+                    self._df[column] = self._df[column].astype(object)
             self._df.update(df, overwrite=True)
         else:
             self._df = df

openeo/extra/job_management/_manager.py

@@ -85,6 +85,9 @@ class _ColumnRequirements:
     def __init__(self, requirements: Mapping[str, _ColumnProperties]):
         self._requirements = dict(requirements)
 
+    def __contains__(self, column: str) -> bool:
+        return column in self._requirements
+
     def normalize_df(self, df: pd.DataFrame) -> pd.DataFrame:
         """
         Normalize given pandas dataframe (creating a new one):
@@ -190,12 +193,9 @@ def start_job(
             # TODO: use proper date/time dtype instead of legacy str for start times?
             "start_time": _ColumnProperties(dtype="str"),
             "running_start_time": _ColumnProperties(dtype="str"),
-            # TODO: these columns "cpu", "memory", "duration" are not referenced explicitly from MultiBackendJobManager,
-            #       but are indirectly coupled through handling of VITO-specific "usage" metadata in `_track_statuses`.
-            #       Since bfd99e34 they are not really required to be present anymore, can we make that more explicit?
-            "cpu": _ColumnProperties(dtype="str"),
-            "memory": _ColumnProperties(dtype="str"),
-            "duration": _ColumnProperties(dtype="str"),
+            # Note: columns for job usage metadata (e.g. "cpu", "memory", "duration", ...)
+            # are not predefined here, but are added dynamically in `_track_statuses`
+            # based on the "usage" fields reported by the backend in the job metadata.
             "costs": _ColumnProperties(dtype="float64"),
         }
     )
@@ -903,10 +903,20 @@ def _track_statuses(self, job_db: JobDatabaseInterface, stats: Optional[dict] =
 
                 active.loc[i, "status"] = new_status
 
-                # TODO: there is well hidden coupling here with "cpu", "memory" and "duration" from `_normalize_df`
+                # Dynamically add all "usage" metadata fields as columns.
                 for key in job_metadata.get("usage", {}).keys():
-                    if key in active.columns:
-                        active.loc[i, key] = _format_usage_stat(job_metadata, key)
+                    if key in self._column_requirements:
+                        _log.warning(
+                            f"Skipping usage field {key!r} of job {job_id!r}: collides with existing job db column"
+                        )
+                        continue
+                    if key not in active.columns:
+                        # Initialize new column with None (instead of NaN) to keep it JSON-serializable
+                        active[key] = None
+                    elif active[key].dtype != object:
+                        # Legacy job dbs may have empty usage columns loaded as float64, which can't hold strings
+                        active[key] = active[key].astype(object)
+                    active.loc[i, key] = _format_usage_stat(job_metadata, key)
                 if "costs" in job_metadata.keys():
                     active.loc[i, "costs"] = job_metadata.get("costs")
 

openeo/rest/_testing.py

@@ -53,6 +53,7 @@ class DummyBackend:
         "job_status_updater",
         "job_id_generator",
         "extra_job_metadata_fields",
+        "next_batch_job_usage",
     )
 
     # Default result (can serve both as JSON or binary data)
@@ -77,6 +78,11 @@ def __init__(
         self._forced_job_status: Dict[str, str] = {}
         self._fail_on_job_start = {}
 
+        # "usage" metadata to report on finished batch jobs
+        self.next_batch_job_usage: dict = {
+            "cpu": {"unit": "cpu-seconds", "value": 1234.5},
+        }
+
         # Job status update hook:
         #   callable that is called on starting a job, and getting job metadata
         #   allows to dynamically change how the status of a job evolves
@@ -297,11 +303,7 @@ def _handle_get_job(self, request, context):
         }
         if self.batch_jobs[job_id]["status"] == "finished":  # HACK some realistic values for a small job
             result["costs"] = 123
-            result["usage"] = {
-                "cpu": {"unit": "cpu-seconds", "value": 1234.5},
-                "memory": {"unit": "mb-seconds", "value": 34567.89},
-                "duration": {"unit": "seconds", "value": 2345},
-            }
+            result["usage"] = self.next_batch_job_usage
         return result
 
     def _handle_get_job_results(self, request, context):

tests/extra/job_management/test_job_db.py

@@ -52,9 +52,6 @@ def test_initialize_from_df(self, tmp_path, db_class):
             "id",
             "start_time",
             "running_start_time",
-            "cpu",
-            "memory",
-            "duration",
             "backend_name",
             "costs",
         }
@@ -90,6 +87,21 @@ def test_initialize_from_df_on_exists_skip(self, tmp_path, db_class):
         )
         assert set(db.read()["some_number"]) == {1, 2, 3}
 
+    @pytest.mark.parametrize("db_class", [CsvJobDatabase, ParquetJobDatabase])
+    def test_persist_with_new_columns(self, tmp_path, db_class):
+        """Persisting an update with new columns (e.g. dynamically added usage columns) should preserve them."""
+        path = tmp_path / "jobs.db"
+        db = db_class(path).initialize_from_df(pd.DataFrame({"some_number": [3, 2, 1]}))
+
+        update = db.df.iloc[[0, 1]].copy()
+        update["cpu"] = ["123 cpu-seconds", "456 cpu-seconds"]
+        db.persist(update)
+
+        result = db_class(path).read()
+        assert "cpu" in result.columns
+        assert list(result["cpu"].fillna("")) == ["123 cpu-seconds", "456 cpu-seconds", ""]
+        assert list(result["some_number"]) == [3, 2, 1]
+
     @pytest.mark.parametrize("db_class", [CsvJobDatabase, ParquetJobDatabase])
     def test_count_by_status(self, tmp_path, db_class):
         path = tmp_path / "jobs.db"
@@ -231,9 +243,6 @@ def test_initialize_from_df(self, tmp_path):
             "id",
             "start_time",
             "running_start_time",
-            "cpu",
-            "memory",
-            "duration",
             "backend_name",
             "costs",
         }
@@ -329,9 +338,6 @@ def test_initialize_from_df(self, tmp_path):
             "id",
             "start_time",
             "running_start_time",
-            "cpu",
-            "memory",
-            "duration",
             "backend_name",
             "costs",
         }

tests/extra/job_management/test_manager.py

@@ -178,14 +178,14 @@ def test_basic(self, tmp_path, job_manager, job_manager_root_dir, sleep_mock):
         )
 
         assert [
-            (r.id, r.status, r.backend_name, r.cpu, r.memory, r.duration, r.costs)
+            (r.id, r.status, r.backend_name, r.cpu, r.costs)
             for r in pd.read_csv(job_db_path).itertuples()
         ] == [
-            ("job-2018", "finished", "foo", "1234.5 cpu-seconds", "34567.89 mb-seconds", "2345 seconds", 123),
-            ("job-2019", "finished", "foo", "1234.5 cpu-seconds", "34567.89 mb-seconds", "2345 seconds", 123),
-            ("job-2020", "finished", "bar", "1234.5 cpu-seconds", "34567.89 mb-seconds", "2345 seconds", 123),
-            ("job-2021", "finished", "bar", "1234.5 cpu-seconds", "34567.89 mb-seconds", "2345 seconds", 123),
-            ("job-2022", "finished", "foo", "1234.5 cpu-seconds", "34567.89 mb-seconds", "2345 seconds", 123),
+            ("job-2018", "finished", "foo", "1234.5 cpu-seconds", 123),
+            ("job-2019", "finished", "foo", "1234.5 cpu-seconds", 123),
+            ("job-2020", "finished", "bar", "1234.5 cpu-seconds", 123),
+            ("job-2021", "finished", "bar", "1234.5 cpu-seconds", 123),
+            ("job-2022", "finished", "foo", "1234.5 cpu-seconds", 123),
         ]
 
         # Check downloaded results and metadata.
@@ -219,9 +219,89 @@ def test_db_class(self, tmp_path, job_manager, job_manager_root_dir, sleep_mock,
         assert set(result.status) == {"finished"}
         assert set(result.backend_name) == {"foo", "bar"}
         assert set(result.cpu) == {"1234.5 cpu-seconds"}
+        assert set(result.costs) == {123}
+
+    @pytest.mark.parametrize("db_class", [CsvJobDatabase, ParquetJobDatabase])
+    def test_usage_fields_dynamic_columns(
+        self, tmp_path, job_manager, dummy_backend_foo, dummy_backend_bar, sleep_mock, db_class
+    ):
+        """
+        All fields reported under "usage" in the job metadata
+        should dynamically be included as columns in the job db.
+        """
+        usage = {
+            "cpu": {"unit": "cpu-seconds", "value": 1234.5},
+            "memory": {"unit": "mb-seconds", "value": 34567.89},
+            "duration": {"unit": "seconds", "value": 2345},
+            "network": {"unit": "b", "value": 1000},
+            "sentinelhub": {"unit": "sentinelhub_processing_unit", "value": 7.5},
+        }
+        dummy_backend_foo.next_batch_job_usage = usage
+        dummy_backend_bar.next_batch_job_usage = usage
+
+        df = pd.DataFrame({"year": [2018, 2019, 2020, 2021, 2022]})
+        output_file = tmp_path / "jobs.db"
+        job_db = db_class(output_file).initialize_from_df(df)
+
+        run_stats = job_manager.run_jobs(job_db=job_db, start_job=self._create_year_job)
+        assert run_stats == dirty_equals.IsPartialDict({"job finished": 5})
+
+        result = db_class(output_file).read()
+        assert len(result) == 5
+        assert set(result.status) == {"finished"}
+        assert set(result.cpu) == {"1234.5 cpu-seconds"}
         assert set(result.memory) == {"34567.89 mb-seconds"}
         assert set(result.duration) == {"2345 seconds"}
-        assert set(result.costs) == {123}
+        assert set(result.network) == {"1000 b"}
+        assert set(result.sentinelhub) == {"7.5 sentinelhub_processing_unit"}
+
+    def test_usage_fields_reserved_column_collision(
+        self, tmp_path, job_manager, dummy_backend_foo, dummy_backend_bar, sleep_mock, caplog
+    ):
+        """
+        Usage fields that collide with reserved job db columns (e.g. "status")
+        should be skipped instead of overwriting the column.
+        """
+        usage = {
+            "status": {"unit": "sneaky", "value": 666},
+            "network": {"unit": "b", "value": 1000},
+        }
+        dummy_backend_foo.next_batch_job_usage = usage
+        dummy_backend_bar.next_batch_job_usage = usage
+
+        df = pd.DataFrame({"year": [2018, 2019]})
+        output_file = tmp_path / "jobs.csv"
+        job_db = CsvJobDatabase(output_file).initialize_from_df(df)
+
+        with caplog.at_level(logging.WARNING):
+            job_manager.run_jobs(job_db=job_db, start_job=self._create_year_job)
+
+        result = CsvJobDatabase(output_file).read()
+        assert set(result.status) == {"finished"}
+        assert set(result.network) == {"1000 b"}
+        assert "Skipping usage field 'status'" in caplog.text
+
+    def test_usage_fields_legacy_job_db(self, tmp_path, job_manager, sleep_mock, recwarn):
+        """
+        Resuming from a legacy job db that still has (empty) predefined usage columns
+        (e.g. "cpu", "memory", "duration") should work without warnings/errors.
+        """
+        job_db_path = tmp_path / "jobs.csv"
+        job_db_path.write_text(
+            "year,id,backend_name,status,start_time,running_start_time,cpu,memory,duration,costs\n"
+            "2021,,,not_started,,,,,,\n"
+            "2022,,,not_started,,,,,,\n"
+        )
+
+        job_db = CsvJobDatabase(job_db_path)
+        run_stats = job_manager.run_jobs(job_db=job_db, start_job=self._create_year_job)
+        assert run_stats == dirty_equals.IsPartialDict({"start_job call": 2, "job finished": 2})
+
+        result = pd.read_csv(job_db_path)
+        assert set(result.status) == {"finished"}
+        assert set(result.cpu) == {"1234.5 cpu-seconds"}
+
+        assert [(w.category, w.message, str(w)) for w in recwarn.list] == []
 
     @pytest.mark.parametrize(
         ["filename", "expected_db_class"],
@@ -274,14 +354,14 @@ def test_start_job_thread_basic(self, tmp_path, job_manager, job_manager_root_di
         assert sleep_mock.call_count > 10
 
         assert [
-            (r.id, r.status, r.backend_name, r.cpu, r.memory, r.duration, r.costs)
+            (r.id, r.status, r.backend_name, r.cpu, r.costs)
             for r in pd.read_csv(job_db_path).itertuples()
         ] == [
-            ("job-2018", "finished", "foo", "1234.5 cpu-seconds", "34567.89 mb-seconds", "2345 seconds", 123),
-            ("job-2019", "finished", "foo", "1234.5 cpu-seconds", "34567.89 mb-seconds", "2345 seconds", 123),
-            ("job-2020", "finished", "bar", "1234.5 cpu-seconds", "34567.89 mb-seconds", "2345 seconds", 123),
-            ("job-2021", "finished", "bar", "1234.5 cpu-seconds", "34567.89 mb-seconds", "2345 seconds", 123),
-            ("job-2022", "finished", "foo", "1234.5 cpu-seconds", "34567.89 mb-seconds", "2345 seconds", 123),
+            ("job-2018", "finished", "foo", "1234.5 cpu-seconds", 123),
+            ("job-2019", "finished", "foo", "1234.5 cpu-seconds", 123),
+            ("job-2020", "finished", "bar", "1234.5 cpu-seconds", 123),
+            ("job-2021", "finished", "bar", "1234.5 cpu-seconds", 123),
+            ("job-2022", "finished", "foo", "1234.5 cpu-seconds", 123),
         ]
 
         # Check downloaded results and metadata.
@@ -301,9 +381,6 @@ def test_normalize_df(self):
                 "id",
                 "start_time",
                 "running_start_time",
-                "cpu",
-                "memory",
-                "duration",
                 "backend_name",
                 "costs",
             ]
@@ -358,13 +435,13 @@ def start_worker_thread():
         # Also check that we got sensible end results in the job db.
         results = pd.read_csv(job_db_path).replace({np.nan: None})  # np.nan's are replaced by None for easy comparison
         assert [
-            (r.id, r.status, r.backend_name, r.cpu, r.memory, r.duration, r.costs) for r in results.itertuples()
+            (r.id, r.status, r.backend_name, r.cpu, r.costs) for r in results.itertuples()
         ] == [
-            ("job-2018", "finished", "foo", "1234.5 cpu-seconds", "34567.89 mb-seconds", "2345 seconds", 123),
-            ("job-2019", "finished", "foo", "1234.5 cpu-seconds", "34567.89 mb-seconds", "2345 seconds", 123),
-            ("job-2020", "finished", "bar", "1234.5 cpu-seconds", "34567.89 mb-seconds", "2345 seconds", 123),
-            ("job-2021", "finished", "bar", "1234.5 cpu-seconds", "34567.89 mb-seconds", "2345 seconds", 123),
-            ("job-2022", "error", "foo", None, None, None, None),
+            ("job-2018", "finished", "foo", "1234.5 cpu-seconds", 123),
+            ("job-2019", "finished", "foo", "1234.5 cpu-seconds", 123),
+            ("job-2020", "finished", "bar", "1234.5 cpu-seconds", 123),
+            ("job-2021", "finished", "bar", "1234.5 cpu-seconds", 123),
+            ("job-2022", "error", "foo", None, None),
         ]
 
         # Check downloaded results and metadata.
@@ -627,7 +704,7 @@ def test_empty_csv_handling(self, tmp_path, sleep_mock, recwarn, job_manager):
 
         assert job_db_path.exists()
         # Simple check for empty columns in the CSV file
-        assert ",,,,," in job_db_path.read_text()
+        assert "2021,,,not_started,,,\n" in job_db_path.read_text()
 
         # Start over with existing file
         job_db = CsvJobDatabase(job_db_path)

tests/extra/job_management/test_stac_job_db.py

@@ -64,9 +64,6 @@ def _common_normalized_df_data(rows: int = 1) -> dict:
         "status": ["not_started"] * rows,
         "start_time": None,
         "running_start_time": None,
-        "cpu": None,
-        "memory": None,
-        "duration": None,
         "costs": None,
     }