move job input file copying logic to executors

Author: andrii-i

jupyter_scheduler/executors.py

@@ -4,7 +4,7 @@
 import tarfile
 import traceback
 from abc import ABC, abstractmethod
-from typing import Dict
+from typing import Dict, Optional
 
 import fsspec
 import nbconvert
@@ -14,7 +14,7 @@
 from jupyter_scheduler.models import DescribeJob, JobFeature, Status
 from jupyter_scheduler.orm import Job, create_session
 from jupyter_scheduler.parameterize import add_parameters
-from jupyter_scheduler.utils import get_utc_timestamp
+from jupyter_scheduler.utils import copy_file, get_utc_timestamp
 
 
 class ExecutionManager(ABC):
@@ -29,11 +29,21 @@ 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],
+        input_uri: Optional[str],
+        package_input_folder: Optional[bool],
+    ):
         self.job_id = job_id
         self.staging_paths = staging_paths
         self.root_dir = root_dir
         self.db_url = db_url
+        self.input_uri = input_uri
+        self.package_input_folder = package_input_folder
 
     @property
     def model(self):
@@ -98,6 +108,17 @@ def before_start(self):
             )
             session.commit()
 
+    def copy_input(self, input_uri: str, copy_to_path: str):
+        if self.package_input_folder:
+            self.copy_input_folder(input_uri, copy_to_path)
+        else:
+            self.copy_input_file(input_uri, copy_to_path)
+
+    def copy_input_file(self, input_uri: str, copy_to_path: str):
+        """Copies the input file to the staging directory in a new process."""
+        input_filepath = os.path.join(self.root_dir, input_uri)
+        copy_file(input_filepath=input_filepath, copy_to_path=copy_to_path)
+
     def on_failure(self, e: Exception):
         """Called after failure of execute"""
         job = self.model

jupyter_scheduler/scheduler.py

@@ -423,11 +423,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 in a new process."""
-        input_filepath = os.path.join(self.root_dir, input_uri)
-        spawn_process(target=copy_file, input_filepath=input_filepath, copy_to_path=copy_to_path)
-
     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))
@@ -468,22 +463,15 @@ def create_job(self, model: CreateJob) -> str:
             session.commit()
 
             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"])
-                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"])
 
             p = spawn_process(
                 target=self.execution_manager_class(
                     job_id=job.job_id,
                     staging_paths=staging_paths,
                     root_dir=self.root_dir,
                     db_url=self.db_url,
+                    input_uri=model.input_uri,
+                    package_input_folder=model.package_input_folder,
                 ).process
             )