copy staging folder and side effects to output after job runs, track and redownload files

Author: andrii-i

jupyter_scheduler/executors.py

@@ -149,27 +149,41 @@ def execute(self):
         except CellExecutionError as e:
             raise e
         finally:
+            self.add_side_effects_files(staging_dir)
             for output_format in job.output_formats:
                 cls = nbconvert.get_exporter(output_format)
                 output, _ = cls().from_notebook_node(nb)
                 with fsspec.open(self.staging_paths[output_format], "w", encoding="utf-8") as f:
                     f.write(output)
-
-        self.copy_staged_files_to_output()
-
+            self.copy_staged_files_to_output()
+
+    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"""
+        input_notebook = os.path.relpath(self.staging_paths["input"])
+        new_files_set = set()
+        for root, _, files in os.walk(staging_dir):
+            for file in files:
+                file_rel_path = os.path.relpath(os.path.join(root, file), staging_dir)
+                if file_rel_path != input_notebook:
+                    new_files_set.add(file_rel_path)
+
+        if new_files_set:
+            with self.db_session() as session:
+                current_packaged_files_set = set(
+                    session.query(Job.packaged_files).filter(Job.job_id == self.job_id).scalar()
+                    or []
+                )
+                updated_packaged_files = list(current_packaged_files_set.union(new_files_set))
+                session.query(Job).filter(Job.job_id == self.job_id).update(
+                    {"packaged_files": updated_packaged_files}
+                )
+                session.commit()
+
+    # TODO: copy via downloader and remove this function or use this function and utilize return for putting side efects into packaged_files
     def copy_staged_files_to_output(self):
         """Copies snapshot of the original notebook and staged input files from the staging directory to the output directory and includes them into job_files."""
         staging_dir = os.path.dirname(self.staging_paths["input"])
-        copied_files = copy_directory(
-            source_dir=staging_dir, destination_dir=self.output_dir, base_dir=self.root_dir
-        )
-
-        if copied_files:
-            for rel_path in copied_files:
-                if not any(job_file.file_path == rel_path for job_file in self.model.job_files):
-                    self.model.job_files.append(
-                        JobFile(display_name="File", file_format="File", file_path=rel_path)
-                    )
+        copy_directory(source_dir=staging_dir, destination_dir=self.output_dir)
 
     def supported_features(cls) -> Dict[JobFeature, bool]:
         return {

jupyter_scheduler/job_files_manager.py

@@ -57,21 +57,17 @@ def __init__(
 
     def generate_filepaths(self):
         """A generator that produces filepaths"""
+        output_formats = self.output_formats + ["input"]
+        for output_format in output_formats:
+            input_filepath = self.staging_paths[output_format]
+            output_filepath = os.path.join(self.output_dir, self.output_filenames[output_format])
+            if not os.path.exists(output_filepath) or self.redownload:
+                yield input_filepath, output_filepath
         if self.include_staging_files:
             staging_dir = os.path.dirname(self.staging_paths["input"])
-            for root, _, files in os.walk(staging_dir):
-                for file in files:
-                    input_filepath = os.path.join(root, file)
-                    relative_path = os.path.relpath(input_filepath, staging_dir)
-                    output_filepath = os.path.join(self.output_dir, relative_path)
-                    yield input_filepath, output_filepath
-        else:
-            output_formats = self.output_formats + ["input"]
-            for output_format in output_formats:
-                input_filepath = self.staging_paths[output_format]
-                output_filepath = os.path.join(
-                    self.output_dir, self.output_filenames[output_format]
-                )
+            for file_relative_path in self.output_filenames["files"]:
+                input_filepath = os.path.join(staging_dir, file_relative_path)
+                output_filepath = os.path.join(self.output_dir, file_relative_path)
                 if not os.path.exists(output_filepath) or self.redownload:
                     yield input_filepath, output_filepath
 

jupyter_scheduler/models.py

@@ -147,7 +147,7 @@ class DescribeJob(BaseModel):
     status_message: Optional[str] = None
     downloaded: bool = False
     package_input_folder: Optional[bool] = None
-    output_folder: Optional[str] = None
+    packaged_files: Optional[List[str]] = []
 
     class Config:
         orm_mode = True

jupyter_scheduler/orm.py

@@ -86,6 +86,7 @@ class CommonColumns:
     update_time = Column(Integer, default=get_utc_timestamp, onupdate=get_utc_timestamp)
     create_time = Column(Integer, default=get_utc_timestamp)
     package_input_folder = Column(Boolean)
+    packaged_files = Column(JsonType, default=[])
 
 
 class Job(CommonColumns, Base):

jupyter_scheduler/scheduler.py

@@ -2,7 +2,7 @@
 import os
 import random
 import shutil
-from typing import Dict, Optional, Type, Union
+from typing import Dict, List, Optional, Type, Union
 
 import fsspec
 import psutil
@@ -252,7 +252,29 @@ def file_exists(self, path: str):
         else:
             return os.path.isfile(os_path)
 
-    def get_job_filenames(self, model: DescribeJob) -> Dict[str, str]:
+    def dir_exists(self, path: str):
+        """Returns True if the directory exists, else returns False.
+
+        API-style wrapper for os.path.isdir
+
+        Parameters
+        ----------
+        path : string
+            The relative path to the directory (with '/' as separator)
+
+        Returns
+        -------
+        exists : bool
+            Whether the directory exists.
+        """
+        root = os.path.abspath(self.root_dir)
+        os_path = to_os_path(path, root)
+        if not (os.path.abspath(os_path) + os.path.sep).startswith(root):
+            return False
+        else:
+            return os.path.isdir(os_path)
+
+    def get_job_filenames(self, model: DescribeJob) -> Dict[str, Union[str, List[str]]]:
         """Returns dictionary mapping output formats to
         the job filenames in the JupyterLab workspace.
 
@@ -269,7 +291,8 @@ def get_job_filenames(self, model: DescribeJob) -> Dict[str, str]:
         {
             'ipynb': 'helloworld-2022-10-10.ipynb',
             'html': 'helloworld-2022-10-10.html',
-            'input': 'helloworld.ipynb'
+            'input': 'helloworld.ipynb',
+            'files': ['data/helloworld.csv', 'images/helloworld.png']
         }
 
         """
@@ -282,6 +305,9 @@ def get_job_filenames(self, model: DescribeJob) -> Dict[str, str]:
 
         filenames["input"] = model.input_filename
 
+        if model.package_input_folder and model.packaged_files:
+            filenames["files"] = [relative_path for relative_path in model.packaged_files]
+
         return filenames
 
     def add_job_files(self, model: DescribeJob):
@@ -320,8 +346,27 @@ def add_job_files(self, model: DescribeJob):
             )
         )
 
+        # Add link to output folder with packaged input files and side effects
+        if model.package_input_folder and model.packaged_files:
+            job_files.append(
+                JobFile(
+                    display_name="Files",
+                    file_format="files",
+                    file_path=output_dir if self.dir_exists(output_dir) else None,
+                )
+            )
+
         model.job_files = job_files
-        model.downloaded = all(job_file.file_path for job_file in job_files)
+
+        packaged_files = []
+        if model.package_input_folder and model.packaged_files:
+            packaged_files = [
+                os.path.join(output_dir, packaged_file_rel_path)
+                for packaged_file_rel_path in model.packaged_files
+            ]
+        model.downloaded = all(job_file.file_path for job_file in job_files) and all(
+            self.file_exists(file_path) for file_path in packaged_files
+        )
 
     def get_local_output_path(
         self, input_filename: str, job_id: str, root_dir_relative: Optional[bool] = False
@@ -385,11 +430,11 @@ def copy_input_file(self, input_uri: str, copy_to_path: str):
             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):
-        """Copies the input file along with the input directory to the staging directory"""
+    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)
-        copy_directory(
+        return copy_directory(
             source_dir=input_dir_path,
             destination_dir=staging_dir,
         )
@@ -420,12 +465,18 @@ def create_job(self, model: CreateJob) -> str:
                 model.output_formats = []
 
             job = Job(**model.dict(exclude_none=True, exclude={"input_uri"}))
+
             session.add(job)
             session.commit()
 
             staging_paths = self.get_staging_paths(DescribeJob.from_orm(job))
             if model.package_input_folder:
-                self.copy_input_folder(model.input_uri, staging_paths["input"])
+                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"])
 
@@ -497,10 +548,6 @@ def list_jobs(self, query: ListJobsQuery) -> ListJobsResponse:
         for job in jobs:
             model = DescribeJob.from_orm(job)
             self.add_job_files(model=model)
-            if model.package_input_folder:
-                model.output_folder = self.get_local_output_path(
-                    input_filename=model.input_filename, job_id=model.job_id, root_dir_relative=True
-                )
             jobs_list.append(model)
 
         list_jobs_response = ListJobsResponse(

jupyter_scheduler/utils.py

@@ -90,11 +90,10 @@ def get_localized_timestamp(timezone) -> int:
 def copy_directory(
     source_dir: str,
     destination_dir: str,
-    base_dir: Optional[str] = None,
     exclude_files: Optional[List[str]] = [],
 ) -> List[str]:
     """Copies content of source_dir to destination_dir excluding exclude_files.
-    Returns a list of relative paths to copied files, relative to base_dir if provided.
+    Returns a list of relative paths to copied files from destination_dir.
     """
     copied_files = []
     for item in os.listdir(source_dir):
@@ -104,15 +103,13 @@ def copy_directory(
             shutil.copytree(source, destination, ignore=shutil.ignore_patterns(*exclude_files))
             for dirpath, _, filenames in os.walk(destination):
                 for filename in filenames:
-                    rel_path = os.path.relpath(
-                        os.path.join(dirpath, filename), base_dir if base_dir else destination_dir
-                    )
+                    rel_path = os.path.relpath(os.path.join(dirpath, filename), destination_dir)
                     copied_files.append(rel_path)
         elif os.path.isfile(source) and item not in exclude_files:
             with fsspec.open(source, "rb") as source_file:
                 with fsspec.open(destination, "wb") as output_file:
                     output_file.write(source_file.read())
-            rel_path = os.path.relpath(destination, base_dir if base_dir else destination_dir)
+            rel_path = os.path.relpath(destination, destination_dir)
             copied_files.append(rel_path)
 
     return copied_files

src/components/files-directory-link.tsx

@@ -0,0 +1,26 @@
+import React from 'react';
+import { Scheduler } from '../handler';
+import { JupyterFrontEnd } from '@jupyterlab/application';
+import { useTranslator } from '../hooks';
+import { Link } from '@mui/material';
+
+export function FilesDirectoryLink(props: {
+  jobFile: Scheduler.IJobFile;
+  app: JupyterFrontEnd;
+}): JSX.Element | null {
+  const trans = useTranslator('jupyterlab');
+  return (
+    <Link
+      href={`/lab/tree/${props.jobFile.file_path}`}
+      title={trans.__('Open output directory with files')}
+      onClick={e => {
+        e.preventDefault();
+        props.app.commands.execute('filebrowser:open-path', {
+          path: props.jobFile.file_path
+        });
+      }}
+    >
+      {trans.__('Files')}
+    </Link>
+  );
+}

src/components/job-file-link.tsx

@@ -6,6 +6,7 @@ import { useEventLogger, useTranslator } from '../hooks';
 import { JupyterFrontEnd } from '@jupyterlab/application';
 
 import { Link } from '@mui/material';
+import { FilesDirectoryLink } from './files-directory-link';
 
 export interface IJobFileLinkProps {
   jobFile: Scheduler.IJobFile;
@@ -30,6 +31,9 @@ export function JobFileLink(props: IJobFileLinkProps): JSX.Element | null {
       ? trans.__('Open input file "%1"', fileBaseName)
       : trans.__('Open output file "%1"', fileBaseName);
 
+  if (props.jobFile.file_format === 'files') {
+    return <FilesDirectoryLink jobFile={props.jobFile} app={props.app} />;
+  }
   return (
     <Link
       key={props.jobFile.file_format}

src/components/job-row.tsx

@@ -75,33 +75,6 @@ function JobFiles(props: {
   );
 }
 
-function FilesDirectoryLink(props: {
-  job: Scheduler.IDescribeJob;
-  app: JupyterFrontEnd;
-}): JSX.Element | null {
-  if (!props.job.package_input_folder || !props.job.output_folder) {
-    return null;
-  }
-  const trans = useTranslator('jupyterlab');
-  return (
-    <Link
-      href={`/lab/tree/${props.job.output_folder}`}
-      title={trans.__(
-        'Open output directory with files for "%1"',
-        props.job.name
-      )}
-      onClick={e => {
-        e.preventDefault();
-        props.app.commands.execute('filebrowser:open-path', {
-          path: props.job.output_folder
-        });
-      }}
-    >
-      {trans.__('Files')}
-    </Link>
-  );
-}
-
 type DownloadFilesButtonProps = {
   app: JupyterFrontEnd;
   job: Scheduler.IDescribeJob;
@@ -194,9 +167,6 @@ export function buildJobRow(
           />
         )}
       <JobFiles job={job} app={app} />
-      {(job.status === 'COMPLETED' || job.status === 'FAILED') && (
-        <FilesDirectoryLink job={job} app={app} />
-      )}
     </>,
     <Timestamp job={job} />,
     translateStatus(job.status),

src/handler.ts

@@ -467,7 +467,6 @@ export namespace Scheduler {
     end_time?: number;
     downloaded: boolean;
     package_input_folder?: boolean;
-    output_folder?: string;
   }
 
   export interface ICreateJobResponse {