Merge pull request #595 from Open-EO/issue571-job-manager-db

Start supporting custom job databases in MultiBackendJobManager

Author: soxofaan

CHANGELOG.md

@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add experimental `openeo.testing.results` subpackage with reusable test utilities for comparing batch job results with reference data
 - `MultiBackendJobManager`: add initial support for storing job metadata in Parquet file (instead of CSV) ([#571](https://github.com/Open-EO/openeo-python-client/issues/571))
 - Add `Connection.authenticate_oidc_access_token()` to set up authorization headers with an access token that is obtained "out-of-band" ([#598](https://github.com/Open-EO/openeo-python-client/issues/598))
+- Add `JobDatabaseInterface` to allow custom job metadata storage with `MultiBackendJobManager` ([#571](https://github.com/Open-EO/openeo-python-client/issues/571))
 
 ### Changed
 

docs/cookbook/job_manager.rst

@@ -5,5 +5,12 @@ Multi Backend Job Manager
 .. warning::
     This is a new experimental API, subject to change.
 
-.. automodule:: openeo.extra.job_management
-    :members: MultiBackendJobManager, ignore_connection_errors
+.. autoclass:: openeo.extra.job_management.MultiBackendJobManager
+    :members:
+
+.. autoclass:: openeo.extra.job_management.JobDatabaseInterface
+    :members:
+
+.. autoclass:: openeo.extra.job_management.CsvJobDatabase
+
+.. autoclass:: openeo.extra.job_management.ParquetJobDatabase

openeo/extra/job_management.py

@@ -1,8 +1,10 @@
+import abc
 import contextlib
 import datetime
 import json
 import logging
 import time
+import warnings
 from pathlib import Path
 from typing import Callable, Dict, NamedTuple, Optional, Union
 
@@ -30,6 +32,40 @@ class _Backend(NamedTuple):
 
 MAX_RETRIES = 5
 
+
+class JobDatabaseInterface(metaclass=abc.ABCMeta):
+    """
+    Interface for a database of job metadata to use with the :py:class:`MultiBackendJobManager`,
+    allowing to regularly persist the job metadata while polling the job statuses
+    and resume/restart the job tracking after it was interrupted.
+
+    .. versionadded:: 0.31.0
+    """
+
+    @abc.abstractmethod
+    def exists(self) -> bool:
+        """Does the job database already exist, to read job data from?"""
+        ...
+
+    @abc.abstractmethod
+    def read(self) -> pd.DataFrame:
+        """
+        Read job data from the database as pandas DataFrame.
+
+        :return: loaded job data.
+        """
+        ...
+
+    @abc.abstractmethod
+    def persist(self, df: pd.DataFrame):
+        """
+        Store job data to the database.
+
+        :param df: job data to store.
+        """
+        ...
+
+
 class MultiBackendJobManager:
     """
     Tracker for multiple jobs on multiple backends.
@@ -207,7 +243,8 @@ def run_jobs(
         self,
         df: pd.DataFrame,
         start_job: Callable[[], BatchJob],
-        output_file: Union[str, Path],
+        job_db: Union[str, Path, JobDatabaseInterface, None] = None,
+        **kwargs,
     ):
         """Runs jobs, specified in a dataframe, and tracks parameters.
 
@@ -241,30 +278,52 @@ def run_jobs(
             any of them out, then remember to include the ``*args`` and ``**kwargs`` parameters.
             Otherwise you will have an exception because :py:meth:`run_jobs` passes unknown parameters to ``start_job``.
 
-        :param output_file:
-            Path to output file (CSV or Parquet) containing the status and metadata of the jobs.
+        :param job_db:
+            Job database to load/store existing job status data and other metadata from/to.
+            Can be specified as a path to CSV or Parquet file,
+            or as a custom database object following the :py:class:`JobDatabaseInterface` interface.
 
             .. note::
                 Support for Parquet files depends on the ``pyarrow`` package
                 as :ref:`optional dependency <installation-optional-dependencies>`.
 
         .. versionchanged:: 0.31.0
-            Added support for providing a Parquet ``output_file``
+            Added support for persisting the job metadata in Parquet format.
+
+        .. versionchanged:: 0.31.0
+            Replace ``output_file`` argument with ``job_db`` argument,
+            which can be a path to a CSV or Parquet file,
+            or a user-defined :py:class:`JobDatabaseInterface` object.
+            The deprecated ``output_file`` argument is still supported for now.
         """
         # TODO: Defining start_jobs as a Protocol might make its usage more clear, and avoid complicated doctrings,
         #   but Protocols are only supported in Python 3.8 and higher.
-        output_file = Path(output_file)
 
-        if output_file.suffix.lower() == ".csv":
-            job_db = _CsvJobDatabase(output_file)
-        elif output_file.suffix.lower() == ".parquet":
-            job_db = _ParquetJobDatabase(output_file)
-        else:
-            raise ValueError(f"Unsupported output file type: {output_file}")
+        # Backwards compatibility for deprecated `output_file` argument
+        if "output_file" in kwargs:
+            if job_db is not None:
+                raise ValueError("Only one of `output_file` and `job_db` should be provided")
+            warnings.warn(
+                "The `output_file` argument is deprecated. Use `job_db` instead.", DeprecationWarning, stacklevel=2
+            )
+            job_db = kwargs.pop("output_file")
+        assert not kwargs, f"Unexpected keyword arguments: {kwargs!r}"
+
+        if isinstance(job_db, (str, Path)):
+            job_db_path = Path(job_db)
+            if job_db_path.suffix.lower() == ".csv":
+                job_db = CsvJobDatabase(path=job_db_path)
+            elif job_db_path.suffix.lower() == ".parquet":
+                job_db = ParquetJobDatabase(path=job_db_path)
+            else:
+                raise ValueError(f"Unsupported job database file type {job_db_path!r}")
+
+        if not isinstance(job_db, JobDatabaseInterface):
+            raise ValueError(f"Unsupported job_db {job_db!r}")
 
-        if output_file.exists() and output_file.is_file():
+        if job_db.exists():
             # Resume from existing db
-            _log.info(f"Resuming `run_jobs` from {output_file.absolute()}")
+            _log.info(f"Resuming `run_jobs` from existing {job_db}")
             df = job_db.read()
             status_histogram = df.groupby("status").size().to_dict()
             _log.info(f"Status histogram: {status_histogram}")
@@ -462,28 +521,32 @@ def _format_usage_stat(job_metadata: dict, field: str) -> str:
 
 
 @contextlib.contextmanager
-def ignore_connection_errors(context: Optional[str] = None):
+def ignore_connection_errors(context: Optional[str] = None, sleep: int = 5):
     """Context manager to ignore connection errors."""
+    # TODO: move this out of this module and make it a more public utility?
     try:
         yield
     except requests.exceptions.ConnectionError as e:
         _log.warning(f"Ignoring connection error (context {context or 'n/a'}): {e}")
         # Back off a bit
-        time.sleep(5)
+        time.sleep(sleep)
 
 
-class _JobDatabaseInterface:
-    def read(self) -> pd.DataFrame:
-        raise NotImplementedError()
-
-    def persist(self, df: pd.DataFrame):
-        raise NotImplementedError()
+class CsvJobDatabase(JobDatabaseInterface):
+    """
+    Persist/load job metadata with a CSV file.
 
+    :implements: :py:class:`JobDatabaseInterface`
+    :param path: Path to local CSV file.
 
-class _CsvJobDatabase(_JobDatabaseInterface):
+    .. versionadded:: 0.31.0
+    """
     def __init__(self, path: Union[str, Path]):
         self.path = Path(path)
 
+    def exists(self) -> bool:
+        return self.path.exists()
+
     def _is_valid_wkt(self, wkt: str) -> bool:
         try:
             shapely.wkt.loads(wkt)
@@ -493,23 +556,39 @@ def _is_valid_wkt(self, wkt: str) -> bool:
 
     def read(self) -> pd.DataFrame:
         df = pd.read_csv(self.path)
-        # Workaround for loading of geopandas "geometry" column.
+        # `df.to_csv` in `persist()` will encode geometries as WKT, so we decode that here.
         if (
             "geometry" in df.columns
             and df["geometry"].dtype.name != "geometry"
             and self._is_valid_wkt(df["geometry"].iloc[0])
         ):
             df["geometry"] = df["geometry"].apply(shapely.wkt.loads)
         return df
+
     def persist(self, df: pd.DataFrame):
         self.path.parent.mkdir(parents=True, exist_ok=True)
         df.to_csv(self.path, index=False)
 
 
-class _ParquetJobDatabase(_JobDatabaseInterface):
+class ParquetJobDatabase(JobDatabaseInterface):
+    """
+    Persist/load job metadata with a Parquet file.
+
+    :implements: :py:class:`JobDatabaseInterface`
+    :param path: Path to the Parquet file.
+
+    .. versionadded:: 0.31.0
+
+    .. note::
+        Support for Parquet files depends on the ``pyarrow`` package
+        as :ref:`optional dependency <installation-optional-dependencies>`.
+    """
     def __init__(self, path: Union[str, Path]):
         self.path = Path(path)
 
+    def exists(self) -> bool:
+        return self.path.exists()
+
     def read(self) -> pd.DataFrame:
         return pd.read_parquet(self.path)
 

tests/extra/test_job_management.py

@@ -18,9 +18,9 @@
 from openeo import BatchJob
 from openeo.extra.job_management import (
     MAX_RETRIES,
+    CsvJobDatabase,
     MultiBackendJobManager,
-    _CsvJobDatabase,
-    _ParquetJobDatabase,
+    ParquetJobDatabase,
 )
 
 
@@ -466,7 +466,7 @@ def test_read_wkt(self, tmp_path):
         )
         path = tmp_path / "jobs.csv"
         wkt_df.to_csv(path, index=False)
-        df = _CsvJobDatabase(path).read()
+        df = CsvJobDatabase(path).read()
         assert isinstance(df.geometry[0], shpt.Point)
 
     def test_read_non_wkt(self, tmp_path):
@@ -478,29 +478,37 @@ def test_read_non_wkt(self, tmp_path):
         )
         path = tmp_path / "jobs.csv"
         non_wkt_df.to_csv(path, index=False)
-        df = _CsvJobDatabase(path).read()
+        df = CsvJobDatabase(path).read()
         assert isinstance(df.geometry[0], str)
 
-    def test_persist(self, tmp_path):
-        df = pd.DataFrame(
+    def test_persist_and_read(self, tmp_path):
+        orig = pd.DataFrame(
             {
-                "some_number": [3, 2, 1],
+                "numbers": [3, 2, 1],
+                "names": ["apple", "banana", "coconut"],
             }
         )
-
         path = tmp_path / "jobs.csv"
-        _CsvJobDatabase(path).persist(df)
-        assert _CsvJobDatabase(path).read().equals(df)
+        CsvJobDatabase(path).persist(orig)
+        assert path.exists()
+
+        loaded = CsvJobDatabase(path).read()
+        assert list(loaded.dtypes) == list(orig.dtypes)
+        assert loaded.equals(orig)
 
 
 class TestParquetJobDatabase:
-    def test_read_persist(self, tmp_path):
-        df = pd.DataFrame(
+    def test_persist_and_read(self, tmp_path):
+        orig = pd.DataFrame(
             {
-                "some_number": [3, 2, 1],
+                "numbers": [3, 2, 1],
+                "names": ["apple", "banana", "coconut"],
             }
         )
-
         path = tmp_path / "jobs.parquet"
-        _ParquetJobDatabase(path).persist(df)
-        assert _ParquetJobDatabase(path).read().equals(df)
+        ParquetJobDatabase(path).persist(orig)
+        assert path.exists()
+
+        loaded = ParquetJobDatabase(path).read()
+        assert list(loaded.dtypes) == list(orig.dtypes)
+        assert loaded.equals(orig)