def initiate_download_standalonefunction and use it in DownloadManager and ExecutionManager

Author: andrii-i

jupyter_scheduler/download_manager.py

@@ -6,6 +6,24 @@
 from jupyter_scheduler.utils import get_utc_timestamp
 
 
+def initiate_download_standalone(job_id: str, queue: Queue, db_session, redownload: bool = False):
+    """
+    This static method initiates a download in a standalone manner independent of the DownloadManager instance. It is suitable for use in multiprocessing environment where a direct reference to DownloadManager instance is not feasible.
+    """
+    download_initiated_time = get_utc_timestamp()
+    download_id = generate_uuid()
+    download = DescribeDownload(
+        job_id=job_id,
+        download_id=download_id,
+        download_initiated_time=download_initiated_time,
+        redownload=redownload,
+    )
+    download_record = Download(**download.dict())
+    db_session.add(download_record)
+    db_session.commit()
+    queue.put(download)
+
+
 class DownloadRecordManager:
     def __init__(self, db_url):
         self.session = create_session(db_url)
@@ -45,17 +63,11 @@ def __init__(self, db_url: str):
         self.record_manager = DownloadRecordManager(db_url=db_url)
         self.queue = Queue()
 
-    def download_from_staging(self, job_id: str, redownload: bool):
-        download_initiated_time = get_utc_timestamp()
-        download_id = generate_uuid()
-        download = DescribeDownload(
-            job_id=job_id,
-            download_id=download_id,
-            download_initiated_time=download_initiated_time,
-            redownload=redownload,
-        )
-        self.record_manager.put(download)
-        self.queue.put(download)
+    def initiate_download(self, job_id: str, redownload: bool):
+        with self.record_manager.session() as session:
+            initiate_download_standalone(
+                job_id=job_id, queue=self.queue, db_session=session, redownload=redownload
+            )
 
     def delete_download(self, download_id: str):
         self.record_manager.delete_download(download_id)

jupyter_scheduler/executors.py

@@ -11,7 +11,11 @@
 import nbformat
 from nbconvert.preprocessors import CellExecutionError, ExecutePreprocessor
 
-from jupyter_scheduler.download_manager import DescribeDownload, Download
+from jupyter_scheduler.download_manager import (
+    DescribeDownload,
+    Download,
+    initiate_download_standalone,
+)
 from jupyter_scheduler.models import DescribeJob, JobFeature, Status
 from jupyter_scheduler.orm import Job, create_session, generate_uuid
 from jupyter_scheduler.parameterize import add_parameters
@@ -156,22 +160,13 @@ def execute(self):
                 output, _ = cls().from_notebook_node(nb)
                 with fsspec.open(self.staging_paths[output_format], "w", encoding="utf-8") as f:
                     f.write(output)
-            self._download_from_staging(job.job_id)
-
-    def _download_from_staging(self, job_id: str):
-        download_initiated_time = get_utc_timestamp()
-        download_id = generate_uuid()
-        download = DescribeDownload(
-            job_id=job_id,
-            download_id=download_id,
-            download_initiated_time=download_initiated_time,
-            redownload=True,
-        )
-        with self.db_session() as session:
-            download_record = Download(**download.dict())
-            session.add(download_record)
-            session.commit()
-        self.download_queue.put(download)
+            with self.db_session() as session:
+                initiate_download_standalone(
+                    job_id=job.job_id,
+                    queue=self.download_queue,
+                    db_session=session,
+                    redownload=True,
+                )
 
     def add_side_effects_files(self, staging_dir):
         """Scan for side effect files potentially created after input file execution and update the job's packaged_files with these files"""

jupyter_scheduler/extension.py

@@ -81,8 +81,8 @@ def initialize_settings(self):
             root_dir=self.serverapp.root_dir,
             environments_manager=environments_manager,
             db_url=self.db_url,
-            config=self.config,
             download_queue=download_manager.queue,
+            config=self.config,
         )
 
         job_files_manager = self.job_files_manager_class(scheduler=scheduler)
@@ -95,7 +95,7 @@ def initialize_settings(self):
             environments_manager=environments_manager,
             scheduler=scheduler,
             job_files_manager=job_files_manager,
-            download_from_staging=download_manager.download_from_staging,
+            initiate_download=download_manager.initiate_download,
         )
 
         if scheduler.task_runner:

jupyter_scheduler/handlers.py

@@ -395,20 +395,20 @@ def get(self):
 
 
 class FilesDownloadHandler(ExtensionHandlerMixin, APIHandler):
-    _download_from_staging = None
+    _initiate_download = None
 
     @property
-    def download_from_staging(self):
-        if not self._download_from_staging:
-            self._download_from_staging = self.settings.get("download_from_staging", None)
+    def initiate_download(self):
+        if not self._initiate_download:
+            self._initiate_download = self.settings.get("initiate_download", None)
 
-        return self._download_from_staging
+        return self._initiate_download
 
     @authenticated
     async def get(self, job_id):
         redownload = self.get_query_argument("redownload", False)
         try:
-            self.download_from_staging(job_id, redownload)
+            self.initiate_download(job_id, redownload)
         except Exception as e:
             self.log.exception(e)
             raise HTTPError(500, str(e)) from e