use dask for input file / folder copying, make copy functions standalone

Author: andrii-i

jupyter_scheduler/scheduler.py

@@ -40,7 +40,8 @@
 )
 from jupyter_scheduler.orm import Job, JobDefinition, create_session
 from jupyter_scheduler.utils import (
-    copy_directory,
+    copy_input_file,
+    copy_input_folder,
     create_output_directory,
     create_output_filename,
 )
@@ -427,22 +428,6 @@ def db_session(self):
 
         return self._db_session
 
-    def copy_input_file(self, input_uri: str, copy_to_path: str):
-        """Copies the input file to the staging directory"""
-        input_filepath = os.path.join(self.root_dir, input_uri)
-        with fsspec.open(input_filepath) as input_file:
-            with fsspec.open(copy_to_path, "wb") as output_file:
-                output_file.write(input_file.read())
-
-    def copy_input_folder(self, input_uri: str, nb_copy_to_path: str) -> List[str]:
-        """Copies the input file along with the input directory to the staging directory, returns the list of copied files relative to the staging directory"""
-        input_dir_path = os.path.dirname(os.path.join(self.root_dir, input_uri))
-        staging_dir = os.path.dirname(nb_copy_to_path)
-        return copy_directory(
-            source_dir=input_dir_path,
-            destination_dir=staging_dir,
-        )
-
     async def create_job(self, model: CreateJob) -> str:
         if not model.job_definition_id and not self.file_exists(model.input_uri):
             raise InputUriError(model.input_uri)
@@ -473,19 +458,26 @@ async def create_job(self, model: CreateJob) -> str:
             session.add(job)
             session.commit()
 
+            dask: DaskClient = await self.dask_client_future
+
             staging_paths = self.get_staging_paths(DescribeJob.from_orm(job))
             if model.package_input_folder:
-                copied_files = self.copy_input_folder(model.input_uri, staging_paths["input"])
+                copy_future = dask.submit(
+                    copy_input_folder, self.root_dir, model.input_uri, staging_paths["input"]
+                )
+                copied_files = await dask.gather(copy_future)
                 input_notebook_filename = os.path.basename(model.input_uri)
                 job.packaged_files = [
                     file for file in copied_files if file != input_notebook_filename
                 ]
                 session.commit()
             else:
-                self.copy_input_file(model.input_uri, staging_paths["input"])
+                copy_future = dask.submit(
+                    copy_input_file, self.root_dir, model.input_uri, staging_paths["input"]
+                )
+                await dask.gather(copy_future)
 
-            dask_client: DaskClient = await self.dask_client_future
-            future = dask_client.submit(
+            process_future = dask.submit(
                 self.execution_manager_class(
                     job_id=job.job_id,
                     staging_paths=staging_paths,
@@ -494,7 +486,7 @@ async def create_job(self, model: CreateJob) -> str:
                 ).process
             )
 
-            job.pid = future.key
+            job.pid = process_future.key
             session.commit()
 
             job_id = job.job_id

jupyter_scheduler/utils.py

@@ -113,3 +113,21 @@ def copy_directory(
             copied_files.append(rel_path)
 
     return copied_files
+
+
+def copy_input_file(root_dir: str, input_uri: str, copy_to_path: str):
+    """Copies the input file to the staging directory"""
+    input_filepath = os.path.join(root_dir, input_uri)
+    with fsspec.open(input_filepath) as input_file:
+        with fsspec.open(copy_to_path, "wb") as output_file:
+            output_file.write(input_file.read())
+
+
+def copy_input_folder(root_dir: str, input_uri: str, nb_copy_to_path: str) -> List[str]:
+    """Copies the input file along with the input directory to the staging directory, returns the list of copied files relative to the staging directory"""
+    input_dir_path = os.path.dirname(os.path.join(root_dir, input_uri))
+    staging_dir = os.path.dirname(nb_copy_to_path)
+    return copy_directory(
+        source_dir=input_dir_path,
+        destination_dir=staging_dir,
+    )