use download queue

andrii-i · andrii-i · commit 9636c7feab99 · 2024-04-10T11:03:06.000-07:00
diff --git a/jupyter_scheduler/download_manager.py b/jupyter_scheduler/download_manager.py
@@ -0,0 +1,129 @@
+from dataclasses import dataclass
+from datetime import datetime
+from multiprocessing import Queue
+from typing import List, Optional
+
+from jupyter_scheduler.orm import DownloadCacheRecord, create_session, generate_uuid
+from jupyter_scheduler.utils import get_utc_timestamp
+from jupyter_scheduler.pydantic_v1 import BaseModel
+
+
+class DescribeDownloadCache(BaseModel):
+    job_id: str
+    download_id: str
+    download_initiated_time: int
+
+    class Config:
+        orm_mode = True
+
+
+@dataclass
+class DownloadTask:
+    job_id: str
+    download_id: str
+    download_initiated_time: int
+
+    def __lt__(self, other):
+        return self.download_initiated_time < other.download_initiated_time
+
+    def __str__(self):
+        download_initiated_time = datetime.fromtimestamp(self.download_initiated_time / 1e3)
+        return f"Id: {self.job_id}, Download initiated: {download_initiated_time}"
+
+
+class MultiprocessQueue:
+    """A multiprocess-safe queue using multiprocessing.Queue()"""
+
+    def __init__(self):
+        self.queue = Queue()
+
+    def put(self, download: DownloadTask):
+        self.queue.put(download)
+
+    def get(self) -> Optional[DownloadTask]:
+        return self.queue.get() if not self.queue.empty() else None
+
+    def isempty(self) -> bool:
+        return self.queue.empty()
+
+
+class DownloadCache:
+    def __init__(self, db_url):
+        self.session = create_session(db_url)
+
+    def put(self, download: DescribeDownloadCache):
+        with self.session() as session:
+            new_download = DownloadCacheRecord(**download.dict())
+            session.add(new_download)
+            session.commit()
+
+    def get(self, job_id: str) -> Optional[DescribeDownloadCache]:
+        with self.session() as session:
+            download = (
+                session.query(DownloadCacheRecord)
+                .filter(DownloadCacheRecord.job_id == job_id)
+                .first()
+            )
+
+        if download:
+            return DescribeDownloadCache.from_orm(download)
+        else:
+            return None
+
+    def get_tasks(self) -> List[DescribeDownloadCache]:
+        with self.session() as session:
+            return (
+                session.query(DownloadCacheRecord)
+                .order_by(DownloadCacheRecord.download_initiated_time)
+                .all()
+            )
+
+    def delete_download(self, download_id: str):
+        with self.session() as session:
+            session.query(DownloadCacheRecord).filter(
+                DownloadCacheRecord.download_id == download_id
+            ).delete()
+            session.commit()
+
+    def delete_job_downloads(self, job_id: str):
+        with self.session() as session:
+            session.query(DownloadCacheRecord).filter(DownloadCacheRecord.job_id == job_id).delete()
+            session.commit()
+
+
+class DownloadManager:
+    def __init__(self, db_url: str):
+        self.cache = DownloadCache(db_url=db_url)
+        self.queue = MultiprocessQueue()
+
+    def download_from_staging(self, job_id: str):
+        download_initiated_time = get_utc_timestamp()
+        download_id = generate_uuid()
+        download_cache = DescribeDownloadCache(
+            job_id=job_id,
+            download_id=download_id,
+            download_initiated_time=download_initiated_time,
+        )
+        self.cache.put(download_cache)
+        download_task = DownloadTask(
+            job_id=job_id,
+            download_id=download_id,
+            download_initiated_time=download_initiated_time,
+        )
+        self.queue.put(download_task)
+
+    def delete_download(self, download_id: str):
+        self.cache.delete_download(download_id)
+
+    def delete_job_downloads(self, job_id: str):
+        self.cache.delete_job_downloads(job_id)
+
+    def populate_queue(self):
+        tasks = self.cache.get_tasks()
+        for task in tasks:
+            download_task = DownloadTask(
+                job_id=task.job_id,
+                download_id=task.download_id,
+                download_initiated_time=task.download_initiated_time,
+            )
+            self.queue.put(download_task)
diff --git a/jupyter_scheduler/download_runner.py b/jupyter_scheduler/download_runner.py
@@ -0,0 +1,72 @@
+import asyncio
+
+from jupyter_scheduler.download_manager import DownloadManager
+from jupyter_scheduler.job_files_manager import JobFilesManager
+
+
+class BaseDownloadRunner:
+    def start(self):
+        raise NotImplementedError("Must be implemented by subclass")
+
+
+class DownloadRunner(BaseDownloadRunner):
+    def __init__(
+        self,
+        download_manager: DownloadManager,
+        job_files_manager: JobFilesManager,
+        poll_interval: int = 10,
+    ):
+        self.download_manager = download_manager
+        self.job_files_manager = job_files_manager
+        self.poll_interval = poll_interval
+
+    # def add_download(self, job_id: str):
+    #     download_initiated_time = get_utc_timestamp()
+    #     download_id = generate_uuid()
+    #     download_cache = DescribeDownloadCache(
+    #         job_id=job_id,
+    #         download_id=download_id,
+    #         download_initiated_time=download_initiated_time,
+    #     )
+    #     self.download_cache.put(download_cache)
+    #     download_task = DownloadTask(
+    #         job_id=job_id,
+    #         download_id=download_id,
+    #         download_initiated_time=download_initiated_time,
+    #     )
+    #     self.download_queue.put(download_task)
+
+    # def delete_download(self, download_id: str):
+    #     self.download_cache.delete_download(download_id)
+
+    # def delete_job_downloads(self, job_id: str):
+    #     self.download_cache.delete_job_downloads(job_id)
+
+    async def process_download_queue(self):
+        print("\n\n***\nDownloadRunner.process_download_queue isempty")
+        print(self.download_manager.queue.isempty())
+        while not self.download_manager.queue.isempty():
+            download = self.download_manager.queue.get()
+            print(download)
+            cache = self.download_manager.cache.get(download.job_id)
+            print(cache)
+            if not cache or not download:
+                continue
+            await self.job_files_manager.copy_from_staging(cache.job_id)
+            self.download_manager.cache.delete_download(cache.download_id)
+
+    # def populate_queue(self):
+    #     tasks = self.download_manager.cache.get_tasks()
+    #     for task in tasks:
+    #         download_task = DownloadTask(
+    #             job_id=task.job_id,
+    #             download_id=task.download_id,
+    #             download_initiated_time=task.download_initiated_time,
+    #         )
+    #         self.download_manager.queue.put(download_task)
+
+    async def start(self):
+        self.download_manager.populate_queue()
+        while True:
+            await self.process_download_queue()
+            await asyncio.sleep(self.poll_interval)
diff --git a/jupyter_scheduler/executors.py b/jupyter_scheduler/executors.py
@@ -1,4 +1,5 @@
 import io
+import multiprocessing
 import os
 import shutil
 import tarfile
@@ -11,8 +12,14 @@
 import nbformat
 from nbconvert.preprocessors import CellExecutionError, ExecutePreprocessor
 
+from jupyter_scheduler.download_manager import (
+    DescribeDownloadCache,
+    DownloadCacheRecord,
+    DownloadTask,
+    MultiprocessQueue,
+)
 from jupyter_scheduler.models import DescribeJob, JobFeature, JobFile, Status
-from jupyter_scheduler.orm import Job, create_session
+from jupyter_scheduler.orm import Job, create_session, generate_uuid
 from jupyter_scheduler.parameterize import add_parameters
 from jupyter_scheduler.utils import get_utc_timestamp
 
@@ -29,11 +36,19 @@ class ExecutionManager(ABC):
     _model = None
     _db_session = None
 
-    def __init__(self, job_id: str, root_dir: str, db_url: str, staging_paths: Dict[str, str]):
+    def __init__(
+        self,
+        job_id: str,
+        root_dir: str,
+        db_url: str,
+        staging_paths: Dict[str, str],
+        download_queue,
+    ):
         self.job_id = job_id
         self.staging_paths = staging_paths
         self.root_dir = root_dir
         self.db_url = db_url
+        self.download_queue = download_queue
 
     @property
     def model(self):
@@ -147,6 +162,31 @@ 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_cache = DescribeDownloadCache(
+            job_id=job_id,
+            download_id=download_id,
+            download_initiated_time=download_initiated_time,
+        )
+        with self.db_session() as session:
+            new_download = DownloadCacheRecord(**download_cache.dict())
+            session.add(new_download)
+            session.commit()
+        download_task = DownloadTask(
+            job_id=job_id,
+            download_id=download_id,
+            download_initiated_time=download_initiated_time,
+        )
+        self.download_queue.put(download_task)
+        # print(
+        #     "\n\n***\n ExecutionManager.download_from_staging uuid and task being put on a qeueue"
+        # )
+        # print(download_id)
+        # print(download_task)
 
     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"""
diff --git a/jupyter_scheduler/extension.py b/jupyter_scheduler/extension.py
@@ -5,6 +5,8 @@
 from jupyter_server.transutils import _i18n
 from traitlets import Bool, Type, Unicode, default
 
+from jupyter_scheduler.download_manager import DownloadManager
+from jupyter_scheduler.download_runner import DownloadRunner
 from jupyter_scheduler.orm import create_tables
 
 from .handlers import (
@@ -73,21 +75,33 @@ def initialize_settings(self):
 
         environments_manager = self.environment_manager_class()
 
+        download_manager = DownloadManager(db_url=self.db_url)
+
         scheduler = self.scheduler_class(
             root_dir=self.serverapp.root_dir,
             environments_manager=environments_manager,
             db_url=self.db_url,
             config=self.config,
+            download_queue=download_manager.queue,
         )
 
         job_files_manager = self.job_files_manager_class(scheduler=scheduler)
 
+        download_runner = DownloadRunner(
+            download_manager=download_manager, job_files_manager=job_files_manager
+        )
+
         self.settings.update(
             environments_manager=environments_manager,
             scheduler=scheduler,
             job_files_manager=job_files_manager,
+            download_from_staging=download_manager.download_from_staging,
         )
 
         if scheduler.task_runner:
             loop = asyncio.get_event_loop()
             loop.create_task(scheduler.task_runner.start())
+
+        if download_runner:
+            loop = asyncio.get_event_loop()
+            loop.create_task(download_runner.start())
diff --git a/jupyter_scheduler/handlers.py b/jupyter_scheduler/handlers.py
@@ -394,20 +394,29 @@ def get(self):
 
 
 class FilesDownloadHandler(ExtensionHandlerMixin, APIHandler):
-    _job_files_manager = None
+    # _job_files_manager = None
+    _download_from_staging = None
+
+    # @property
+    # def job_files_manager(self):
+    #     if not self._job_files_manager:
+    #         self._job_files_manager = self.settings.get("job_files_manager", None)
+
+    #     return self._job_files_manager
 
     @property
-    def job_files_manager(self):
-        if not self._job_files_manager:
-            self._job_files_manager = self.settings.get("job_files_manager", None)
+    def download_from_staging(self):
+        if not self._download_from_staging:
+            self._download_from_staging = self.settings.get("download_from_staging", None)
 
-        return self._job_files_manager
+        return self._download_from_staging
 
     @authenticated
     async def get(self, job_id):
-        redownload = self.get_query_argument("redownload", False)
+        # redownload = self.get_query_argument("redownload", False)
         try:
-            await self.job_files_manager.copy_from_staging(job_id=job_id, redownload=redownload)
+            # await self.job_files_manager.copy_from_staging(job_id=job_id, redownload=redownload)
+            self.download_from_staging(job_id)
         except Exception as e:
             self.log.exception(e)
             raise HTTPError(500, str(e)) from e
diff --git a/jupyter_scheduler/orm.py b/jupyter_scheduler/orm.py
@@ -1,5 +1,4 @@
 import json
-import os
 from sqlite3 import OperationalError
 from uuid import uuid4
 
@@ -112,6 +111,13 @@ class JobDefinition(CommonColumns, Base):
     active = Column(Boolean, default=True)
 
 
+class DownloadCacheRecord(Base):
+    __tablename__ = "download_cache"
+    job_id = Column(String(36), primary_key=True)
+    download_id = Column(String(36), primary_key=True)
+    download_initiated_time = Column(Integer)
+
+
 def create_tables(db_url, drop_tables=False):
     engine = create_engine(db_url)
     try:
diff --git a/jupyter_scheduler/scheduler.py b/jupyter_scheduler/scheduler.py
