rename Download tables, use mp.quque directly without a wrapper

Author: andrii-i

jupyter_scheduler/download_manager.py

@@ -3,12 +3,12 @@
 from multiprocessing import Queue
 from typing import List, Optional
 
-from jupyter_scheduler.orm import DownloadCacheRecord, create_session, generate_uuid
+from jupyter_scheduler.orm import Downloads, create_session, generate_uuid
 from jupyter_scheduler.pydantic_v1 import BaseModel
 from jupyter_scheduler.utils import get_utc_timestamp
 
 
-class DescribeDownloadCache(BaseModel):
+class DescribeDownload(BaseModel):
     job_id: str
     download_id: str
     download_initiated_time: int
@@ -31,80 +31,54 @@ def __str__(self):
         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:
+class DownloadRecordManager:
     def __init__(self, db_url):
         self.session = create_session(db_url)
 
-    def put(self, download: DescribeDownloadCache):
+    def put(self, download: DescribeDownload):
         with self.session() as session:
-            new_download = DownloadCacheRecord(**download.dict())
+            new_download = Downloads(**download.dict())
             session.add(new_download)
             session.commit()
 
-    def get(self, job_id: str) -> Optional[DescribeDownloadCache]:
+    def get(self, job_id: str) -> Optional[DescribeDownload]:
         with self.session() as session:
-            download = (
-                session.query(DownloadCacheRecord)
-                .filter(DownloadCacheRecord.job_id == job_id)
-                .first()
-            )
+            download = session.query(Downloads).filter(Downloads.job_id == job_id).first()
 
         if download:
-            return DescribeDownloadCache.from_orm(download)
+            return DescribeDownload.from_orm(download)
         else:
             return None
 
-    def get_tasks(self) -> List[DescribeDownloadCache]:
+    def get_tasks(self) -> List[DescribeDownload]:
         with self.session() as session:
-            return (
-                session.query(DownloadCacheRecord)
-                .order_by(DownloadCacheRecord.download_initiated_time)
-                .all()
-            )
+            return session.query(Downloads).order_by(Downloads.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.query(Downloads).filter(Downloads.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.query(Downloads).filter(Downloads.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()
+        self.record_manager = DownloadRecordManager(db_url=db_url)
+        self.queue = Queue()
 
     def download_from_staging(self, job_id: str):
         download_initiated_time = get_utc_timestamp()
         download_id = generate_uuid()
-        download_cache = DescribeDownloadCache(
+        download_cache = DescribeDownload(
             job_id=job_id,
             download_id=download_id,
             download_initiated_time=download_initiated_time,
         )
-        self.cache.put(download_cache)
+        self.record_manager.put(download_cache)
         download_task = DownloadTask(
             job_id=job_id,
             download_id=download_id,
@@ -113,13 +87,13 @@ def download_from_staging(self, job_id: str):
         self.queue.put(download_task)
 
     def delete_download(self, download_id: str):
-        self.cache.delete_download(download_id)
+        self.record_manager.delete_download(download_id)
 
     def delete_job_downloads(self, job_id: str):
-        self.cache.delete_job_downloads(job_id)
+        self.record_manager.delete_job_downloads(job_id)
 
     def populate_queue(self):
-        tasks = self.cache.get_tasks()
+        tasks = self.record_manager.get_tasks()
         for task in tasks:
             download_task = DownloadTask(
                 job_id=task.job_id,

jupyter_scheduler/download_runner.py

@@ -14,56 +14,20 @@ def __init__(
         self,
         download_manager: DownloadManager,
         job_files_manager: JobFilesManager,
-        poll_interval: int = 10,
+        poll_interval: int = 5,
     ):
         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():
+        while not self.download_manager.queue.empty():
             download = self.download_manager.queue.get()
-            print(download)
-            cache = self.download_manager.cache.get(download.job_id)
-            print(cache)
+            cache = self.download_manager.record_manager.get(download.job_id)
             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)
+            self.download_manager.record_manager.delete_download(cache.download_id)
 
     async def start(self):
         self.download_manager.populate_queue()

jupyter_scheduler/executors.py

@@ -13,10 +13,9 @@
 from nbconvert.preprocessors import CellExecutionError, ExecutePreprocessor
 
 from jupyter_scheduler.download_manager import (
-    DescribeDownloadCache,
-    DownloadCacheRecord,
+    DescribeDownload,
+    Downloads,
     DownloadTask,
-    MultiprocessQueue,
 )
 from jupyter_scheduler.models import DescribeJob, JobFeature, JobFile, Status
 from jupyter_scheduler.orm import Job, create_session, generate_uuid
@@ -167,13 +166,13 @@ def execute(self):
     def download_from_staging(self, job_id: str):
         download_initiated_time = get_utc_timestamp()
         download_id = generate_uuid()
-        download_cache = DescribeDownloadCache(
+        download_cache = DescribeDownload(
             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())
+            new_download = Downloads(**download_cache.dict())
             session.add(new_download)
             session.commit()
         download_task = DownloadTask(
@@ -182,11 +181,6 @@ def download_from_staging(self, job_id: str):
             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"""

jupyter_scheduler/orm.py

@@ -111,8 +111,8 @@ class JobDefinition(CommonColumns, Base):
     active = Column(Boolean, default=True)
 
 
-class DownloadCacheRecord(Base):
-    __tablename__ = "download_cache"
+class Downloads(Base):
+    __tablename__ = "downloads"
     job_id = Column(String(36), primary_key=True)
     download_id = Column(String(36), primary_key=True)
     download_initiated_time = Column(Integer)

jupyter_scheduler/scheduler.py

@@ -15,7 +15,6 @@
 from traitlets import Unicode, default
 from traitlets.config import LoggingConfigurable
 
-from jupyter_scheduler.download_manager import MultiprocessQueue
 from jupyter_scheduler.environments import EnvironmentManager
 from jupyter_scheduler.exceptions import (
     IdempotencyTokenError,
@@ -405,7 +404,7 @@ def __init__(
         root_dir: str,
         environments_manager: Type[EnvironmentManager],
         db_url: str,
-        download_queue: MultiprocessQueue,
+        download_queue: mp.Queue,
         config=None,
         **kwargs,
     ):
@@ -495,7 +494,7 @@ def create_job(self, model: CreateJob) -> str:
                     staging_paths=staging_paths,
                     root_dir=self.root_dir,
                     db_url=self.db_url,
-                    download_queue=self.download_queue.queue,
+                    download_queue=self.download_queue,
                 ).process
             )
             p.start()