adding download task and creating seperate download pool

Author: HansVRP

openeo/extra/job_management/_manager.py

@@ -32,6 +32,7 @@
 from openeo.extra.job_management._thread_worker import (
     _JobManagerWorkerThreadPool,
     _JobStartTask,
+    _JobDownloadTask
 )
 from openeo.rest import OpenEoApiError
 from openeo.rest.auth.auth import BearerAuth
@@ -172,6 +173,7 @@ def start_job(
 
     .. versionchanged:: 0.47.0
         Added ``download_results`` parameter.
+
     """
 
     # Expected columns in the job DB dataframes.
@@ -219,6 +221,7 @@ def __init__(
         )
         self._thread = None
         self._worker_pool = None
+        self._download_pool = None
         # Generic cache
         self._cache = {}
 
@@ -351,6 +354,7 @@ def start_job_thread(self, start_job: Callable[[], BatchJob], job_db: JobDatabas
 
         self._stop_thread = False
         self._worker_pool = _JobManagerWorkerThreadPool()
+        self._download_pool = _JobManagerWorkerThreadPool()
 
         def run_loop():
             # TODO: support user-provided `stats`
@@ -388,7 +392,13 @@ def stop_job_thread(self, timeout_seconds: Optional[float] = _UNSET):
 
         .. versionadded:: 0.32.0
         """
-        self._worker_pool.shutdown()
+        if self._worker_pool is not None:
+            self._worker_pool.shutdown()
+            self._worker_pool = None
+            
+        if self._download_pool is not None:
+            self._download_pool.shutdown()
+            self._download_pool = None
 
         if self._thread is not None:
             self._stop_thread = True
@@ -493,6 +503,8 @@ def run_jobs(
         stats = collections.defaultdict(int)
 
         self._worker_pool = _JobManagerWorkerThreadPool()
+        self._download_pool = _JobManagerWorkerThreadPool()
+
 
         while (
             sum(
@@ -511,7 +523,7 @@ def run_jobs(
             stats["sleep"] += 1
 
         # TODO; run post process after shutdown once more to ensure completion?
-        self._worker_pool.shutdown()
+        self.stop_job_thread()
 
         return stats
 
@@ -553,7 +565,11 @@ def _job_update_loop(
                         stats["job_db persist"] += 1
                         total_added += 1
 
-        self._process_threadworker_updates(self._worker_pool, job_db=job_db, stats=stats)
+        if self._worker_pool is not None:
+            self._process_threadworker_updates(worker_pool=self._worker_pool, job_db=job_db, stats=stats)
+            
+        if self._download_pool is not None:
+            self._process_threadworker_updates(worker_pool=self._download_pool, job_db=job_db, stats=stats)
 
         # TODO: move this back closer to the `_track_statuses` call above, once job done/error handling is also handled in threads?
         for job, row in jobs_done:
@@ -565,6 +581,7 @@ def _job_update_loop(
         for job, row in jobs_cancel:
             self.on_job_cancel(job, row)
 
+
     def _launch_job(self, start_job, df, i, backend_name, stats: Optional[dict] = None):
         """Helper method for launching jobs
 
@@ -657,7 +674,7 @@ def _refresh_bearer_token(self, connection: Connection, *, max_age: float = 60)
             else:
                 _log.warning("Failed to proactively refresh bearer token")
 
-    def _process_threadworker_updates(
+    def _process_task_results(
         self,
         worker_pool: _JobManagerWorkerThreadPool,
         *,
@@ -723,15 +740,23 @@ def on_job_done(self, job: BatchJob, row):
         """
         # TODO: param `row` is never accessed in this method. Remove it? Is this intended for future use?
         if self._download_results:
-            job_metadata = job.describe()
-            job_dir = self.get_job_dir(job.job_id)
-            metadata_path = self.get_job_metadata_path(job.job_id)
 
+            job_dir = self.get_job_dir(job.job_id)
             self.ensure_job_dir_exists(job.job_id)
-            job.get_results().download_files(target=job_dir)
 
-            with metadata_path.open("w", encoding="utf-8") as f:
-                json.dump(job_metadata, f, ensure_ascii=False)
+            # Proactively refresh bearer token (because task in thread will not be able to do that
+            job_con =  job.connection
+            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,  # TODO figure out correct index usage
+                download_dir=job_dir,
+            )
+            _log.info(f"Submitting download task {task} to download thread pool")
+            self._download_pool.submit_task(task)
 
     def on_job_error(self, job: BatchJob, row):
         """

openeo/extra/job_management/_thread_worker.py

@@ -7,7 +7,9 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from typing import Any, Dict, List, Optional, Tuple, Union
+from pathlib import Path
 
+import json
 import urllib3.util
 
 import openeo
@@ -140,6 +142,50 @@ def execute(self) -> _TaskResult:
                 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.
+    """
+    download_dir: Path
+
+    def execute(self) -> _TaskResult:
+        """
+        Download job results and metadata.
+        """
+        try:
+            job = self.get_connection(retry=True).job(self.job_id)
+            
+            # Ensure download directory exists
+            self.download_dir.mkdir(parents=True, exist_ok=True)
+            
+            # Download results
+            job.get_results().download_files(target=self.download_dir)
+            
+            # Download metadata
+            job_metadata = job.describe()
+            metadata_path = self.download_dir / f"job_{self.job_id}.json"
+            with metadata_path.open("w", encoding="utf-8") as f:
+                json.dump(job_metadata, f, ensure_ascii=False)
+            
+            _log.info(f"Job {self.job_id!r} results downloaded successfully")
+            return _TaskResult(
+                job_id=self.job_id,
+                df_idx=self.df_idx,
+                db_update={}, #TODO consider db updates
+                stats_update={"job download": 1},
+            )
+        except Exception as e:
+            _log.error(f"Failed to download results for job {self.job_id!r}: {e!r}")
+            return _TaskResult(
+                job_id=self.job_id,
+                df_idx=self.df_idx,
+                db_update={},
+                stats_update={"job download error": 1},
+            )
 
 class _JobManagerWorkerThreadPool:
     """