Use job scheduling information to allow more flexible job submission

ParProcCo/job_scheduler.py

@@ -212,12 +212,13 @@ def fetch_and_update_state(
         return slurm_state
 
     def get_output_paths(
-        self, job_scheduling_info_list: list[JobSchedulingInformation] | ValuesView[JobSchedulingInformation]
+        self,
+        job_scheduling_info_list: list[JobSchedulingInformation]
+        | ValuesView[JobSchedulingInformation],
     ) -> tuple[Path, ...]:
-        return tuple(p for p in (
-            jsi.get_output_path()
-            for jsi in job_scheduling_info_list
-        ) if p)
+        return tuple(
+            p for p in (jsi.get_output_path() for jsi in job_scheduling_info_list) if p
+        )
 
     def get_success(
         self, job_scheduling_info_list: list[JobSchedulingInformation]
@@ -346,7 +347,9 @@ def make_job_submission(
 
     def wait_all_jobs(
         self,
-        job_scheduling_info_list: Sequence[JobSchedulingInformation] | ValuesView[JobSchedulingInformation],
+        job_scheduling_info_list: Sequence[JobSchedulingInformation]
+        | ValuesView[JobSchedulingInformation],
+        in_group: bool,
         state_group: STATEGROUP,
         deadline: datetime,
         sleep_time: int,
@@ -355,7 +358,7 @@ def wait_all_jobs(
         while len(remaining_jobs) > 0 and datetime.now() <= deadline:
             for jsi in list(remaining_jobs):
                 current_state = self.fetch_and_update_state(jsi)
-                if current_state in state_group:
+                if (current_state in state_group) == in_group:
                     remaining_jobs.remove(jsi)
             if len(remaining_jobs) > 0:
                 time.sleep(sleep_time)
@@ -375,51 +378,60 @@ def get_deadline(
             allow_from_submission: bool = False,
         ) -> datetime | None:
             # Timeout shouldn't include queue time
-            if job_scheduling_info.status_info is None:
+            status_info = job_scheduling_info.status_info
+            if status_info is None:
                 return None
-            elif job_scheduling_info.status_info.start_time is None:
+            elif status_info.start_time is None:
                 if allow_from_submission:
-                    return (
-                        job_scheduling_info.status_info.submit_time
-                        + job_scheduling_info.timeout
-                    )
+                    return status_info.submit_time + job_scheduling_info.timeout
                 return None
-            return (
-                job_scheduling_info.status_info.start_time + job_scheduling_info.timeout
-            )
+            return status_info.start_time + job_scheduling_info.timeout
 
         def handle_not_started(
-            job_scheduling_info_list: Sequence[JobSchedulingInformation] | ValuesView[JobSchedulingInformation],
+            job_scheduling_info_list: Sequence[JobSchedulingInformation]
+            | ValuesView[JobSchedulingInformation],
             check_time: timedelta,
         ) -> list[JobSchedulingInformation]:
             # Wait for jobs to start (timeout shouldn't include queue time)
             starting_jobs = list(job_scheduling_info_list)
             timeout = datetime.now() + check_time
-            while len(job_scheduling_info_list) > 0 and datetime.now() < timeout:
-                for jsi in self.wait_all_jobs(
+            logging.debug(
+                "Wait for jobs (%d) to start up to %s", len(starting_jobs), timeout
+            )
+            while len(starting_jobs) > 0 and datetime.now() < timeout:
+                starting_jobs = self.wait_all_jobs(
                     starting_jobs,
+                    False,
                     STATEGROUP.STARTING,
                     timeout,
-                    0,
-                ):
-                    starting_jobs.remove(jsi)
+                    5,
+                )
                 if len(starting_jobs) > 0:
-                    # We want ot sleep only if there are jobs waiting to start
+                    # We want to sleep only if there are jobs waiting to start
                     time.sleep(5)
                     logging.info("Jobs left to start: %d", len(starting_jobs))
             return starting_jobs
 
         def wait_for_ended(
-            job_scheduling_info_list: Sequence[JobSchedulingInformation] | ValuesView[JobSchedulingInformation],
+            job_scheduling_info_list: Sequence[JobSchedulingInformation]
+            | ValuesView[JobSchedulingInformation],
             deadline: datetime,
             check_time: timedelta,
         ) -> list[JobSchedulingInformation]:
+            sleep_time = int(round(check_time.total_seconds()))
+            logging.debug(
+                "Wait for ending in %d jobs up to %s with sleeps of %ss",
+                len(job_scheduling_info_list),
+                deadline,
+                sleep_time,
+            )
             # Wait for jobs to complete
             self.wait_all_jobs(
                 job_scheduling_info_list,
+                True,
                 STATEGROUP.ENDED,
                 deadline,
-                int(round(check_time.total_seconds())),
+                sleep_time,
             )
             ended_jobs = handle_ended_jobs(job_scheduling_info_list)
             logging.info(
@@ -430,7 +442,8 @@ def wait_for_ended(
             return ended_jobs
 
         def handle_ended_jobs(
-            job_scheduling_info_list: Sequence[JobSchedulingInformation] | ValuesView[JobSchedulingInformation],
+            job_scheduling_info_list: Sequence[JobSchedulingInformation]
+            | ValuesView[JobSchedulingInformation],
         ) -> list[JobSchedulingInformation]:
             ended_jobs = []
             for job_scheduling_info in job_scheduling_info_list:
@@ -442,7 +455,8 @@ def handle_ended_jobs(
             return ended_jobs
 
         def handle_timeouts(
-            job_scheduling_info_list: Sequence[JobSchedulingInformation] | ValuesView[JobSchedulingInformation],
+            job_scheduling_info_list: Sequence[JobSchedulingInformation]
+            | ValuesView[JobSchedulingInformation],
         ) -> list[JobSchedulingInformation]:
             deadlines = (
                 (jsi, get_deadline(jsi, allow_from_submission=False))
@@ -492,8 +506,8 @@ def handle_timeouts(
 
         if not running_jobs:
             logging.warning("All jobs ended before wait began")
-            handle_ended_jobs(ended_jobs.values())
             return
+
         try:
             while datetime.now() < wait_deadline and len(running_jobs) > 0:
                 # Handle none started (empty deadline list)
@@ -503,7 +517,8 @@ def handle_timeouts(
                         for deadline in (
                             get_deadline(jsi, allow_from_submission=True)
                             for jsi in running_jobs.values()
-                        ) if deadline is not None
+                        )
+                        if deadline is not None
                     ]
                 )
                 check_time = min(
@@ -662,7 +677,8 @@ def resubmit_killed_jobs(
             failed_jobs = [
                 jsi
                 for jsi in job_scheduling_info_dict.values()
-                if jsi.status_info and jsi.status_info.final_state != SLURMSTATE.COMPLETED
+                if jsi.status_info
+                and jsi.status_info.final_state != SLURMSTATE.COMPLETED
             ]
             killed_jobs = self.filter_killed_jobs(failed_jobs)
             logging.info(

ParProcCo/passthru_wrapper.py

@@ -4,8 +4,12 @@
 
 from .job_slicer_interface import JobSlicerInterface
 from .program_wrapper import ProgramWrapper
-from .utils import (check_jobscript_is_readable, check_location,
-                    format_timestamp, get_absolute_path)
+from .utils import (
+    check_jobscript_is_readable,
+    check_location,
+    format_timestamp,
+    get_absolute_path,
+)
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -29,7 +33,7 @@ def create_slice_jobs(
         timestamp = format_timestamp(job_scheduling_information.timestamp)
         job_scheduling_information.stdout_filename = f"out_{timestamp}"
         job_scheduling_information.stderr_filename = f"err_{timestamp}"
-        old_args= job_scheduling_information.job_script_arguments
+        old_args = job_scheduling_information.job_script_arguments
         job_script = str(
             check_jobscript_is_readable(check_location(get_absolute_path(old_args[0])))
         )
@@ -42,7 +46,13 @@ def create_slice_jobs(
             str(job_scheduling_information.job_resources.cpu_cores),
         )
         if job_scheduling_information.output_filename:
-            args += ("--output", job_scheduling_information.output_filename)
+            args += (
+                "--output",
+                str(
+                    job_scheduling_information.output_dir
+                    / job_scheduling_information.output_filename
+                ),
+            )
         if len(old_args) > 1:
             args += old_args[1:]
         job_scheduling_information.job_script_arguments = args

scripts/ppc_cluster_submit

@@ -149,10 +149,12 @@ def run_ppc(args: argparse.Namespace, script_args: list[str]) -> None:
     jc.run(args.jobs, wrapper_args, "PPC-" + program,
            JobResources(memory=memory, cpu_cores=args.cores,
                         extra_properties=extra_properties))
-    print("complete")
+    print("Jobs completed")
 
 
 if __name__ == "__main__":
     args, script_args = create_parser().parse_known_args()
     logging.getLogger().setLevel(logging.DEBUG if args.debug else logging.INFO)
+    if args.debug:
+        logging.basicConfig(format="%(asctime)s:%(levelname)s:%(name)s:%(module)s:%(funcName)s:%(lineno)d:%(message)s")
     run_ppc(args, script_args)