coupling task type to seperate pool

Author: HansVRP

openeo/extra/job_management/_manager.py

@@ -367,7 +367,7 @@ def run_loop():
                 > 0
 
                 or (self._worker_pool.num_pending_tasks() > 0)
-                
+
                 and not self._stop_thread
             ):
                 self._job_update_loop(job_db=job_db, start_job=start_job, stats=stats)
@@ -644,7 +644,7 @@ def _launch_job(self, start_job, df, i, backend_name, stats: Optional[dict] = No
                                 df_idx=i,
                             )
                             _log.info(f"Submitting task {task} to thread pool")
-                            self._worker_pool.submit_start_task(task)
+                            self._worker_pool.submit_task(task)
 
                             stats["job_queued_for_start"] += 1
                             df.loc[i, "status"] = "queued_for_start"
@@ -690,7 +690,7 @@ def _process_threadworker_updates(
         :param stats:       Dictionary accumulating statistic counters
         """
         # Retrieve completed task results immediately
-        results, start_remaining, download_remaining = worker_pool.process_all_updates(timeout=0)
+        results, _ = worker_pool.process_all_updates(timeout=0)
 
         # Collect update dicts
         updates: List[Dict[str, Any]] = []
@@ -746,18 +746,18 @@ def on_job_done(self, job: BatchJob, row):
             self._refresh_bearer_token(connection=job_con)
 
             task = _JobDownloadTask(
-                root_url=job_con.root_url,
-                bearer_token=job_con.auth.bearer if isinstance(job_con.auth, BearerAuth) else None,
                 job_id=job.job_id,
                 df_idx=row.name,  #this is going to be the index in the not saterted dataframe; should not be an issue as there is no db update for download task
+                root_url=job_con.root_url,
+                bearer_token=job_con.auth.bearer if isinstance(job_con.auth, BearerAuth) else None,
                 download_dir=job_dir,
             )
             _log.info(f"Submitting download task {task} to download thread pool")
 
             if self._worker_pool is None:
                 self._worker_pool = _JobManagerWorkerThreadPool()
 
-            self._worker_pool.submit_download_task(task)
+            self._worker_pool.submit_task(task)
 
     def on_job_error(self, job: BatchJob, row):
         """

openeo/extra/job_management/_thread_worker.py

@@ -3,8 +3,6 @@
 """
 
 import concurrent.futures
-import threading
-import queue
 import logging
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
@@ -103,7 +101,7 @@ def get_connection(self, retry: Union[urllib3.util.Retry, dict, bool, None] = No
             connection.authenticate_bearer_token(self.bearer_token)
         return connection
 
-
+@dataclass(frozen=True)
 class _JobStartTask(ConnectedTask):
     """
     Task for starting an openEO batch job (the `POST /jobs/<job_id>/result` request).
@@ -143,18 +141,16 @@ def execute(self) -> _TaskResult:
                 db_update={"status": "start_failed"},
                 stats_update={"start_job error": 1},
             )
-
+        
+@dataclass(frozen=True)
 class _JobDownloadTask(ConnectedTask):
     """
     Task for downloading job results and metadata.
 
     :param download_dir:
         Root directory where job results and metadata will be downloaded.
-    :param download_throttle:
-        A threading.Semaphore to limit concurrent downloads.
     """
-    download_dir: Path = field(repr=False)
-
+    download_dir: Path = field(default=None, repr=False)
 
     def execute(self) -> _TaskResult:
 
@@ -198,9 +194,10 @@ class _TaskThreadPool:
         Defaults to 2.
     """
 
-    def __init__(self, max_workers: int = 2):
+    def __init__(self, max_workers: int = 2, name: str = None):
         self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=max_workers)
         self._future_task_pairs: List[Tuple[concurrent.futures.Future, Task]] = []
+        self._name = name
 
     def submit_task(self, task: Task) -> None:
         """
@@ -264,36 +261,77 @@ def shutdown(self) -> None:
 
 
 class _JobManagerWorkerThreadPool:
-    """WRAPPER that hides two pools behind one interface"""
-    
-    def __init__(self, max_start_workers=2, max_download_workers=10):
-        # These are the TWO pools with their OWN _future_task_pairs
-        self._start_pool = _TaskThreadPool(max_workers=max_start_workers)
-        self._download_pool = _TaskThreadPool(max_workers=max_download_workers)
+
+    """
+    Generic wrapper that manages multiple thread pools with a dict.
+    Uses task class names as pool names automatically.
+    """
 
-    def submit_start_task(self, task):
-        # Delegate to start pool
-        self._start_pool.submit_task(task)
+    def __init__(self, pool_configs: Optional[Dict[str, int]] = None):
+        """
+        :param pool_configs: Dict of task_class_name -> max_workers
+                            Example: {"_JobStartTask": 1, "_JobDownloadTask": 2}
+        """
+        self._pools: Dict[str, _TaskThreadPool] = {}
+        self._pool_configs = pool_configs or {}
 
-    def submit_download_task(self, task):
-        # Delegate to download pool
-        self._download_pool.submit_task(task)
+    def _get_pool_name_for_task(self, task: Task) -> str:
+        """
+        Get pool name from task class name.
+        """
+        return task.__class__.__name__
 
-    def process_all_updates(self, timeout=0):
-        # Get results from BOTH pools
-        start_results, start_remaining = self._start_pool.process_futures(timeout)
-        download_results, download_remaining = self._download_pool.process_futures(timeout)
+    def submit_task(self, task: Task) -> None:
+        """
+        Submit a task to a pool named after its class.
+        Creates pool dynamically if it doesn't exist.
+        """
+        pool_name = self._get_pool_name_for_task(task)
 
-        # Combine and return
-        all_results = start_results + download_results
-        return all_results, start_remaining, download_remaining
+        if pool_name not in self._pools:
+            # Create pool on-demand
+            max_workers = self._pool_configs.get(pool_name, 1)  # Default 1 worker
+            self._pools[pool_name] = _TaskThreadPool(max_workers=max_workers, name=pool_name)
+            _log.info(f"Created pool '{pool_name}' with {max_workers} workers")
+        
+        self._pools[pool_name].submit_task(task)
+
+    def process_all_updates(self, timeout: Union[float, None] = 0) -> Tuple[List[_TaskResult], Dict[str, int]]:
+        """
+        Process updates from ALL pools.
+        Returns: (all_results, dict of remaining tasks per pool)
+        """
+        all_results = []
+        remaining_by_pool = {}
+        
+        for pool_name, pool in self._pools.items():
+            results, remaining = pool.process_futures(timeout)
+            all_results.extend(results)
+            remaining_by_pool[pool_name] = remaining
+            
+        return all_results, remaining_by_pool
 
-    def num_pending_tasks(self):
-        # Sum of BOTH pools
-        return (self._start_pool.num_pending_tasks() + 
-                self._download_pool.num_pending_tasks())
+    def num_pending_tasks(self, pool_name: Optional[str] = None) -> int:
+        if pool_name:
+            pool = self._pools.get(pool_name)
+            return pool.num_pending_tasks() if pool else 0
+        else:
+            return sum(pool.num_pending_tasks() for pool in self._pools.values())
+    
+    def shutdown(self, pool_name: Optional[str] = None) -> None:
+        """
+        Shutdown pools.
+        If pool_name is None, shuts down all pools.
+        """
+        if pool_name:
+            if pool_name in self._pools:
+                self._pools[pool_name].shutdown()
+                del self._pools[pool_name]
+        else:
+            for pool_name, pool in list(self._pools.items()):
+                pool.shutdown()
+                del self._pools[pool_name]
 
-    def shutdown(self):
-        # Shutdown BOTH pools
-        self._start_pool.shutdown()
-        self._download_pool.shutdown()
+    def list_pools(self) -> List[str]:
+        """List all active pool names."""
+        return list(self._pools.keys())

tests/extra/job_management/test_manager.py

@@ -729,7 +729,7 @@ def get_status(job_id, current_status):
         assert isinstance(rfc3339.parse_datetime(filled_running_start_time), datetime.datetime)
 
     def test_process_threadworker_updates(self, tmp_path, caplog):
-        pool = _JobManagerWorkerThreadPool(max_workers=2)
+        pool = _JobManagerWorkerThreadPool()
         stats = collections.defaultdict(int)
 
         # Submit tasks covering all cases
@@ -769,7 +769,7 @@ def test_process_threadworker_updates(self, tmp_path, caplog):
         assert caplog.messages == []
 
     def test_process_threadworker_updates_unknown(self, tmp_path, caplog):
-        pool = _JobManagerWorkerThreadPool(max_workers=2)
+        pool = _JobManagerWorkerThreadPool()
         stats = collections.defaultdict(int)
 
         pool.submit_task(DummyResultTask("j-123", df_idx=0, db_update={"status": "queued"}, stats_update={"queued": 1}))
@@ -806,7 +806,7 @@ def test_process_threadworker_updates_unknown(self, tmp_path, caplog):
         assert caplog.messages == [dirty_equals.IsStr(regex=".*Ignoring unknown.*indices.*4.*")]
 
     def test_no_results_leaves_db_and_stats_untouched(self, tmp_path, caplog):
-        pool = _JobManagerWorkerThreadPool(max_workers=2)
+        pool = _JobManagerWorkerThreadPool()
         stats = collections.defaultdict(int)
 
         df_initial = pd.DataFrame({"id": ["j-0"], "status": ["created"]})
@@ -820,7 +820,7 @@ def test_no_results_leaves_db_and_stats_untouched(self, tmp_path, caplog):
         assert stats == {}
 
     def test_logs_on_invalid_update(self, tmp_path, caplog):
-        pool = _JobManagerWorkerThreadPool(max_workers=2)
+        pool = _JobManagerWorkerThreadPool()
         stats = collections.defaultdict(int)
 
         # Malformed db_update (not a dict unpackable via **)

tests/extra/job_management/test_thread_worker.py

@@ -221,7 +221,7 @@ class TestJobManagerWorkerThreadPool:
     @pytest.fixture
     def worker_pool(self) -> Iterator[_JobManagerWorkerThreadPool]:
         """Fixture for creating and cleaning up a worker thread pool."""
-        pool = _JobManagerWorkerThreadPool(max_workers=2)
+        pool = _JobManagerWorkerThreadPool()
         yield pool
         pool.shutdown()