Coverage for src/fsl_pipe/job.py: 84%
546 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-22 15:00 +0100
1"""Defines pipeline interface after it has been split into individual jobs based on a FileTree."""
3from file_tree.file_tree import FileTree
4from typing import Optional, Set, Dict, Collection, Any, Callable, Sequence, List, Tuple
5from enum import Enum
6from fsl.utils.fslsub import func_to_cmd
7from loguru import logger
8from pathlib import Path
9from fsl_sub import submit, config
10import re
11import inspect
12from contextlib import contextmanager
13from wcmatch.glob import globmatch, GLOBSTAR
14from warnings import warn
15import numpy as np
18class OutputMissing(IOError):
19 """Raised if a job misses the output files."""
21 pass
24class InputMissingRun(IOError):
25 """Raised if a job misses the input files when it is being launched."""
27 pass
30class InputMissingPipe(IOError):
31 """Raised if a job misses input files while it is being added to the pipeline."""
33 def __init__(self, target, missing, dependency_graph=None):
34 self.target = target
35 self.missing = missing
36 if dependency_graph is None:
37 dependency_graph = []
38 self.dependency_graph = dependency_graph
39 super().__init__("")
41 def __str__(
42 self,
43 ):
44 dependency = " -> ".join([str(d) for d in self.dependency_graph])
45 return f"{self.target} can not be added to pipeline as it is missing required input files: {self.missing} (dependency_graph: {dependency} )"
48class RunMethod(Enum):
49 """How to run the individual jobs."""
51 local = 1
52 submit = 2
53 dask = 3
55 def default():
56 """Return default RunMethod, which is `submit` on a cluster and `local` otherwise."""
57 return RunMethod.submit if config.has_queues() else RunMethod.local
60class JobList:
61 """Pipeline with a concrete set of jobs.
63 Produced from a :class:`fsl_pipe.pipeline.Pipeline` based on the provided `FileTree`.
65 The default `JobList` produced by `Pipeline.generate_jobs` is unsorted and
66 needs to be sorted by running `filter` before running.
67 """
69 def __init__(
70 self,
71 tree: FileTree,
72 jobs: Sequence["SingleJob"],
73 targets: Dict[str, "FileTarget"],
74 ):
75 """Create a new concrete set of jobs."""
76 self.tree = tree
77 self.jobs = jobs
78 self.targets = targets
79 self._sort()
81 def __len__(
82 self,
83 ):
84 return len(self.jobs)
86 def filter(
87 self,
88 templates: Collection[str] = None,
89 overwrite=False,
90 overwrite_dependencies=False,
91 skip_missing=False,
92 ) -> "JobList":
93 """Filter out a subset of jobs.
95 This does a total of three things:
96 - Only include jobs needed to create the templates or file patterns given in `templates`.
97 - Exclude jobs for which output files already exist, unless requested by `overwrite` or `overwrite_dependencies`.
98 - Sort the jobs, so that dependencies are run before the jobs that depend on them.
100 :param templates: target template keys (default: all templates); Any jobs required to produce those files will be kept in the returned pipeline.
101 :param overwrite: if True overwrite the files matching the template even if they already exist, defaults to False
102 :param overwrite_dependencies: if True rerun any dependent jobs even if the output of those jobs already exists, defaults to False
103 :param skip_missing: if True remove any jobs missing input dependencies. If not set, an error is raised whenever inputs are missing.
104 :return: Concrete pipeline with the jobs needed to produces `templates` (with optional overriding)
105 """
106 if isinstance(templates, str):
107 templates = [templates]
109 def add_target(target: "FileTarget", warn_msg=""):
110 if target.exists() and (not overwrite or target.producer is None):
111 return
112 if target.producer is None:
113 warn(f"No job found that produces {target.filename} {warn_msg}")
114 return
115 try:
116 target.producer.add_to_jobs(
117 all_jobs,
118 overwrite=overwrite,
119 overwrite_dependencies=overwrite_dependencies,
120 )
121 except InputMissingPipe as e:
122 if skip_missing:
123 return
124 e.dependency_graph.insert(0, e.target)
125 e.target = target.filename
126 raise
128 all_jobs: Tuple[List[SingleJob], Set[SingleJob]] = ([], set())
129 if templates is None:
130 for target in self.targets.values():
131 if target.producer is None or target.intermediate:
132 continue
133 add_target(target)
134 else:
135 for template in templates:
136 if template in self.tree.template_keys():
137 for filename in self.tree.get_mult(template).data.flatten():
138 target = get_target(
139 filename, self.targets, from_template=template
140 )
141 add_target(target, f"(part of template {template})")
142 else:
143 matches = get_matching_targets(template, self.targets)
144 if len(matches) == 0:
145 warn(
146 f"No files were found that match {template}. It is not a template in the FileTree nor does it match any files that can be produced by this pipeline."
147 )
148 for target in matches:
149 add_target(target, f"(matches {template})")
150 return JobList(self.tree, all_jobs[0], self.targets)
152 def _sort(
153 self,
154 ):
155 """Sorts the job in the `jobs` list in place based on dependencies.
157 A job dependencies should always come before the job itself.
158 This sorting is done purely based on the job inputs/outputs and does not depend on which files actually exist on disk.
159 """
160 all_jobs: Tuple[List[SingleJob], Set[SingleJob]] = ([], set())
161 for job in self.jobs:
162 job.add_to_jobs(all_jobs, only_sort=True)
164 # Remove jobs that were already filtered out
165 sorted_jobs = [j for j in all_jobs[0] if j in self.jobs]
166 assert len(self.jobs) == len(sorted_jobs)
167 self.jobs = sorted_jobs
169 def batch(
170 self,
171 use_label=False,
172 only_connected=False,
173 split_on_ram=False,
174 use_placeholders=True,
175 batch_unlabeled=False,
176 ):
177 """Batches groups of jobs into a single job.
179 Two jobs will be batched if:
181 1. There are no intermediate jobs that need to run between the two and cannot be included in the same batch (because of #3 below).
182 2. One of the jobs depend on the other. This is disabled by default. Enable it by setting `only_connected` to True.
183 3. The jobs are similar. Jobs are similar if they match all of the (enabled) properties listed below:
185 a) They have the same batch label (set during pipeline creating using the `@pipe(batch=<batch label>)`. This can be enabled by setting the `use_label` keyword to True. If enabled, jobs without a batch label will never be merged.
186 b) The jobs have the same submit parameters (e.g., "architecture", "coprocessor"), except for "jobtime", "name", and "export_vars". "jobram" can also be ignored by setting the `split_on_ram` keyword to True.
187 c) The jobs have the same placeholder values. This will prevent jobs from different subjects to be merged with each other. It can be disabled by setting the `use_placeholders` to False. Alternatively, a subset of placeholder values could be considered by passing `use_placeholders` to a sequence of placeholders to consider (e.g., `use_placeholders=["subject", "session"]` will not merge jobs with different "subject" or "session" placeholder values, but will ignore any other placeholders).
189 A new `JobList` with the batched jobs will be returned.
190 """
191 reference = self.copy()
193 all_flags = {
194 key for job in reference.jobs for key in job.submit_params.keys()
195 }.difference(["jobtime", "name", "export_vars"])
196 if not split_on_ram and "jobram" in all_flags:
197 all_flags.remove("jobram")
198 all_flags = sorted(all_flags)
200 job_labels = []
201 for job in reference.jobs:
202 label = []
203 label.append(job.batch if use_label else 1)
204 label.extend(job.submit_params.get(key, None) for key in all_flags)
205 if use_placeholders:
206 if use_placeholders is True:
207 parameters = job.set_parameters
208 else:
209 parameters = {
210 key: value
211 for key, value in job.set_parameters.items()
212 if key in use_placeholders
213 }
214 label.append(frozenset(parameters.items()))
215 job_labels.append(tuple(label))
217 hashable_job_labels = [
218 tuple(label) if isinstance(label, list) else label for label in job_labels
219 ]
221 to_batch = {
222 tb: set(
223 job for (job, label) in zip(reference.jobs, job_labels) if label == tb
224 )
225 for tb in set(hashable_job_labels)
226 }
227 jobs = set(reference.jobs)
228 for to_batch, batching in to_batch.items():
229 if to_batch[0] is None and not batch_unlabeled:
230 # do not batch jobs without a label set (if `use_label` is True)
231 continue
232 other_jobs = jobs.difference(batching)
233 new_jobs = batch_connected_jobs(batching, other_jobs)
234 if not only_connected:
235 new_jobs = batch_unconnected_jobs(new_jobs, other_jobs)
236 jobs = set(new_jobs).union(other_jobs)
237 return JobList(reference.tree, jobs, reference.targets)
239 def split_pipeline(self, use_label=False, split_on_ram=False):
240 """Split the pipeline into multiple stages that require different hardware.
242 This uses the same rules as :meth:`batch`, except that placeholder values are always ignored.
243 """
244 batched_jobs = self.batch(
245 use_label=use_label,
246 split_on_ram=split_on_ram,
247 only_connected=False,
248 use_placeholders=False,
249 batch_unlabeled=True,
250 )
252 return [stage.to_job_list(batched_jobs.tree) for stage in batched_jobs.jobs]
254 def copy(
255 self,
256 ):
257 """Create a new, independent copy of the JobList."""
258 targets = dict()
259 new_jobs = [job.copy(targets) for job in self.jobs]
260 return JobList(self.tree, new_jobs, targets)
262 def report(self, console=None):
263 """Produce tree reports with the relevant input/output templates."""
264 from rich import tree
266 if console is None:
267 from rich.console import Console
269 console = Console()
271 if len(self.jobs) == 0:
272 console.print("No jobs will be run.")
273 return
275 all_inputs = set.union(*[set(j.input_targets) for j in self.jobs])
276 all_outputs = set.union(*[set(j.output_targets) for j in self.jobs])
277 templates = {
278 template
279 for fns in (all_inputs, all_outputs)
280 for target in fns
281 for template in target.from_templates
282 }
284 def proc_line(tree_obj):
285 for t in tree_obj.children:
286 proc_line(t)
287 line = tree_obj.label
288 start = line.index("[cyan]") + 6
289 end = line.index("[/cyan]")
290 key = line[start:end]
291 if key not in templates:
292 return
293 input_count = 0
294 output_count = 0
295 overwrite_count = 0
296 exists_count = 0
298 for fn in self.tree.get_mult(key).data.flatten():
299 target = get_target(fn, self.targets)
300 is_output = False
301 is_input = False
302 if target in all_outputs:
303 is_output = True
304 if target in all_inputs:
305 is_input = True
306 if target.exists() and is_output:
307 overwrite_count += 1
308 elif is_output:
309 output_count += 1
310 elif target.exists():
311 exists_count += 1
312 if is_input:
313 input_count += 1
314 if (input_count + output_count + overwrite_count) == 0:
315 return
316 counter = "/".join(
317 [
318 str(number) if color is None else f"[{color}]{number}[/{color}]"
319 for number, color in [
320 (overwrite_count, "red"),
321 (output_count, "yellow"),
322 (input_count, "blue"),
323 ]
324 ]
325 )
326 tree_obj.label = f"{line} [{counter}]"
328 for rich_obj in (
329 self.tree.filter_templates(templates).fill()._generate_rich_report()
330 ):
331 if isinstance(rich_obj, tree.Tree):
332 proc_line(rich_obj)
333 console.print(rich_obj)
335 def run_datalad(
336 self,
337 ):
338 """Make sure we can run the pipeline.
340 Calls `datalad.get` on all input files and `datalad.unlock` on all output files.
341 """
342 input_targets = set()
343 output_targets = set()
345 for job_group in self.jobs.values():
346 for j in job_group:
347 input_targets.update(j.exists_before)
348 output_targets.update(j.exists_after)
350 input_targets.difference_update(output_targets)
351 input_fns = [t.filename for t in input_targets]
352 output_fns = [t.filename for t in output_targets if t.exists]
354 from .datalad import get_dataset
356 ds = get_dataset()
357 ds.get(input_fns)
358 ds.unlock(output_fns)
360 def run(self, method: RunMethod = None, wait_for=(), clean_script="on_success"):
361 """Run all the jobs that are required to produce the given templates.
363 :param method: defines how to run the job
364 :param wait_for: job IDs to wait for before running pipeline
365 :param clean_script: Sets whether the script produced in the log directory when submitting a job to the cluster should be kept after the job finishes. Only used if `method` is "submit". Options:
366 - "never": Script is kept
367 - "on_success": (default) Only remove if script successfully finishes (i.e., no error is raised)
368 - "always": Always remove the script, even if the script raises an error
369 """
370 if method is None:
371 method = RunMethod.default()
372 elif not isinstance(method, RunMethod):
373 method = RunMethod[method]
374 if len(self.jobs) == 0:
375 logger.info("No new jobs being run/submitted")
376 return
378 prev_count = 0
379 run_jobs: Dict[SingleJob, Any] = {}
380 while len(run_jobs) < len(self.jobs):
381 for job in self.jobs:
382 if job in run_jobs or any(
383 j in self.jobs and j not in run_jobs
384 for _, j in job.dependencies(only_missing=False)
385 ):
386 continue
387 dependencies = [
388 run_jobs[j]
389 for _, j in job.dependencies(only_missing=False)
390 if j in run_jobs
391 ]
392 if len(dependencies) == 0:
393 dependencies = wait_for
394 run_jobs[job] = job(
395 method=method,
396 wait_for=dependencies,
397 clean_script=clean_script,
398 )
399 if len(run_jobs) == prev_count:
400 raise ValueError(
401 "Unable to run/submit all jobs. Are there circular dependencies?"
402 )
403 prev_count = len(run_jobs)
405 if method == RunMethod.dask:
406 import dask
408 def last_dask_job(*all_jobs):
409 if any(a != 0 for a in all_jobs):
410 return 1
411 return 0
413 if (
414 dask.delayed(last_dask_job, name="last_job")(
415 *run_jobs.values()
416 ).compute()
417 == 1
418 ):
419 raise ValueError("One or more of the jobs have failed.")
420 logger.info("Successfully finished running all jobs using Dask.")
422 def scale_jobtime(self, scaling):
423 """
424 Scale the submit job times of all jobs by `scaling`.
426 This will only affect jobs submitted to the cluster.
427 """
428 for job in self.jobs:
429 job.scale_jobtime(scaling)
430 return self
433class JobParent:
434 """Parent for `SingleJob` and `BatchJob`."""
436 input_targets: Set["FileTarget"]
437 output_targets: Set["FileTarget"]
438 optional_targets: Set["FileTarget"]
440 # @lru_cache(None)
441 def dependencies(self, only_missing=True) -> Set[Optional["SingleJob"]]:
442 """Return jobs on which this job depends.
444 By default it only returns those related with missing input files.
446 :param only_missing: set to False to also return dependencies that produce files that already exist on disk
447 """
448 jobs = set()
449 for target in self.input_targets:
450 if not (only_missing and target.exists()):
451 jobs.add((target in self.optional_targets, target.producer))
452 return jobs
454 def missing_output(self, reset_cache=False):
455 """
456 Create a list of filenames that do not exist on disk.
458 Optional outputs are not considered.
460 :param reset_cache: set to True to not rely on an existing cached existence check
461 """
462 missing = set()
463 for to_check in self.output_targets:
464 if to_check in self.optional_targets:
465 continue
466 if reset_cache:
467 to_check.reset_existence()
468 if not to_check.exists():
469 missing.add(to_check)
470 return missing
472 def missing_input(self, reset_cache=False, ignore_expected=False):
473 """
474 Create a list of filenames that do not exist on disk.
476 Optional inputs are not considered.
478 :param reset_cache: set to True to not rely on an existing cached existence check
479 :param ignore_expected: set to True to ignore any missing files that have a job that will produce them in the pipeline
480 """
481 missing = set()
482 for to_check in self.input_targets:
483 if to_check in self.optional_targets:
484 continue
485 if reset_cache:
486 to_check.reset_existence()
487 if ignore_expected and to_check.producer is not None:
488 continue
489 if not to_check.exists():
490 missing.add(to_check)
491 return missing
493 def add_to_jobs(
494 self, all_jobs, overwrite=False, overwrite_dependencies=False, only_sort=False
495 ):
496 """Mark this job and all of its dependencies to run.
498 This job is marked to run, if any of the output does not yet exist or overwrite is True.
499 The dependencies are marked to run, if this job runs and either their output does not exist or overwrite_dependencies is True.
501 :param all_jobs: list and set of individual jobs. This job and all required jobs are added to this list.
502 :param overwrite: if True mark this job even if the output already exists, defaults to False
503 :param overwrite_dependencies: if True mark the dependencies of this job even if their output already exists, defaults to False
504 """
505 if self in all_jobs[1]:
506 return
507 if not only_sort:
508 if not overwrite and len(self.missing_output()) == 0:
509 return
510 missing = self.missing_input(ignore_expected=True)
511 if len(missing) > 0:
512 raise InputMissingPipe(self, {m.filename for m in missing})
513 subjobs = ([], set(all_jobs[1]))
514 for optional, job in self.dependencies(
515 only_missing=not (overwrite_dependencies or only_sort)
516 ):
517 try:
518 if job is not None:
519 job.add_to_jobs(
520 subjobs,
521 overwrite_dependencies,
522 overwrite_dependencies,
523 only_sort,
524 ) and not optional
525 except InputMissingPipe as e:
526 if optional:
527 continue
528 e.dependency_graph.insert(0, e.target)
529 e.target = self
530 raise
531 all_jobs[0].extend(subjobs[0])
532 all_jobs[1].update(subjobs[0])
533 all_jobs[0].append(self)
534 all_jobs[1].add(self)
536 def __call__(self, method: RunMethod, wait_for=(), clean_script="on_success"):
537 """Run the job using the specified `method`."""
538 if method == RunMethod.local:
539 self.prepare_run()
540 missing = self.missing_input()
541 if len(missing) > 0:
542 raise InputMissingRun(
543 f"{self} can not run as it misses required input files: {missing}"
544 )
545 logger.info(f"running {self}")
546 self.job = self.function(**self.kwargs)
547 missing = self.missing_output(reset_cache=True)
548 if len(missing) > 0:
549 raise OutputMissing(
550 f"{self} failed to produce required output files: {missing}"
551 )
552 elif method == RunMethod.submit:
553 self.prepare_run()
554 local_submit = dict(self.submit_params)
555 if "logdir" not in local_submit:
556 local_submit["logdir"] = "log"
557 Path(local_submit["logdir"]).mkdir(exist_ok=True, parents=True)
558 if "name" not in local_submit:
559 local_submit["name"] = self.job_name()
560 cmd = func_to_cmd(
561 self.function,
562 (),
563 self.kwargs,
564 local_submit["logdir"],
565 clean=clean_script,
566 )
567 if len(wait_for) == 0:
568 wait_for = None
569 self.job = submit(cmd, jobhold=wait_for, **local_submit)
570 logger.debug(f"submitted {self} with job ID {self.job}")
571 elif method == RunMethod.dask:
572 import dask
574 def dask_job(*other_jobs):
575 if any(a != 0 for a in other_jobs):
576 logger.debug(f"{self} skipped because dependencies failed")
577 return 1
578 try:
579 logger.debug(f"Running {self} using dask")
580 self(RunMethod.local)
581 except Exception as e:
582 logger.exception(f"{self} failed: {e}")
583 return 1
584 logger.debug(f"Running {self} using dask")
585 return 0
587 self.job = dask.delayed(dask_job, name=str(self))(*wait_for)
588 return self.job
590 def prepare_run(self):
591 """
592 Prepare to run this job.
594 Steps:
595 1. Creates output directory
596 """
597 for target in self.output_targets:
598 target.filename.parent.mkdir(parents=True, exist_ok=True)
600 def expected(
601 self,
602 ):
603 """
604 Return true if this job is expected to be able to run.
605 """
606 for target in self.input_targets:
607 if not target.expected():
608 return False
609 return True
612class SingleJob(JobParent):
613 """A single job within a larger pipeline."""
615 def __init__(
616 self,
617 function: Callable,
618 kwargs,
619 submit_params,
620 input_targets,
621 output_targets,
622 optionals,
623 set_parameters=None,
624 batch=None,
625 ):
626 """
627 Create a single job that can be run locally or submitted.
629 :param function: python function
630 :param kwargs: keyword arguments
631 :param submit_params: instructions to submit job to cluster using `fsl_sub`
632 :param input_targets: set of all filenames expected to exist before this job runs
633 :param output_targets: set of all filenames expected to exist after this job runs
634 :param optionals: set of filenames that are used to generate the dependency graph and yet might not exist
635 :param set_parameters: dictionary with placeholder values used to distinguish this SingleJob with all those produced from the same function
636 :param batch: label used to batch multiple jobs into one when submitting to the cluster
637 """
638 self.function = function
639 if set_parameters is None:
640 set_parameters = {}
641 self.set_parameters = set_parameters
642 self.kwargs = kwargs
643 self.input_targets = input_targets
644 self.output_targets = output_targets
645 self.optional_targets = set(optionals)
646 self.submit_params = dict(submit_params)
647 self.batch = batch
648 for target in self.input_targets:
649 target.required_by.add(self)
650 for target in self.output_targets:
651 target.producer = self
653 def copy(self, targets: Dict[str, "FileTarget"]):
654 """
655 Create a copy of the `SingleJob` to be included in a new :class:`JobList`.
657 `targets` contain the set of FileTargets recognised by this new `JobList`.
658 This will be updated based on the input/output targets of this job.
659 """
661 def copy_targets(job_targets):
662 new_targets = set()
663 for target in job_targets:
664 new_target = get_target(target.filename, targets)
665 new_target.from_templates = target.from_templates
666 new_targets.add(new_target)
667 return new_targets
669 return SingleJob(
670 self.function,
671 self.kwargs,
672 self.submit_params,
673 copy_targets(self.input_targets),
674 copy_targets(self.output_targets),
675 copy_targets(self.optional_targets),
676 self.set_parameters,
677 self.batch,
678 )
680 @property
681 def _function_name(
682 self,
683 ):
684 return self.submit_params.get("name", self.function.__name__)
686 def job_name(
687 self,
688 ):
689 """Return a string representation of this job."""
690 if len(self.set_parameters) > 0:
691 parameter_string = "_".join(
692 [f"{key}-{value}" for key, value in self.set_parameters.items()]
693 )
694 name = f"{self._function_name}_{parameter_string}"
695 else:
696 name = self._function_name
697 value = re.sub(r"[^\w\s-]", "", name).strip().lower()
698 return re.sub(r"[-\s]+", "-", value)
700 def __repr__(
701 self,
702 ):
703 """Print job as a function call."""
704 parameter_string = ", ".join(
705 [f"{key}={value}" for key, value in self.set_parameters.items()]
706 )
707 return f"{self._function_name}({parameter_string})"
709 def to_job_list(self, tree: FileTree):
710 """Convert single job into its own :class:`JobList`."""
711 result = JobList(tree, [self], {}).copy()
712 return self.batch, self.submit_params, result
714 def scale_jobtime(self, scaling):
715 """
716 Scale the submit job time in place by `scaling`.
718 This will only affect jobs submitted to the cluster.
719 """
720 if "jobtime" in self.submit_params.keys():
721 self.submit_params["jobtime"] = int(
722 np.ceil(scaling * self.submit_params["jobtime"])
723 )
726def get_target(filename: Path, all_targets, from_template=None) -> "FileTarget":
727 """
728 Return a :class:`FileTarget` matching the input `filename`.
730 If the `FileTarget` for `filename` is already in `all_targets`, it will be returned.
731 Otherwise a new `FileTarget` will be added to `all_targets` and returned.
733 :param filename: path to the input/intermediate/output filename
734 :param all_targets: dictionary of all FileTargets
735 :param from_template: template key used to obtain the filename
736 """
737 abs_path = Path(filename).absolute()
738 if abs_path not in all_targets:
739 all_targets[abs_path] = FileTarget(filename)
740 if from_template is not None:
741 all_targets[abs_path].from_templates.add(from_template)
742 return all_targets[abs_path]
745def get_matching_targets(pattern: str, all_targets) -> List["FileTarget"]:
746 """
747 Return all :class:`FileTarget` that match the given input pattern.
749 :param pattern: filename definition supporting Unix shell-style wildcards
750 :param all_targets: dictionary of all FileTargets
751 """
752 abs_pattern = str(Path(pattern).absolute())
753 matches = []
754 for path, target in all_targets.items():
755 if globmatch(str(path), abs_pattern, flags=GLOBSTAR):
756 matches.append(target)
757 return matches
760class FileTarget:
761 """Input, intermediate, or output file.
763 See :func:`get_target` for instructions on creating a new `FileTarget`.
765 If a specific :class:`SingleJob` produces a filename, this can be indicated by setting :attr:`producer`:
767 .. code-block:: python
769 get_target(filename, all_targets).producer = job
771 This will raise a `ValueError` if the filename is already produced by another job.
773 If a specific :class:`SingleJob` requires a filename, this can be indicated by adding it to :attr:`required_by`:
775 .. code-block:: python
777 get_target(filename, all_targets).required_by.add(job)
779 Filename existence can be checked using :meth:`exists`.
780 This method uses caching. To reset the cache run :meth:`reset_existence`.
782 To check if the filename can be created by this pipeline (or already exists) run :meth:`expected`.
783 """
785 def __init__(self, filename: Path):
786 """
787 Create a new target based on the provided filename.
789 Do not call this method directly.
790 Instead use :func:`get_target`.
792 :param filename: filename
793 """
794 self.filename = Path(filename)
795 self._producer = None
796 self.required_by: Set[SingleJob] = set()
797 self.from_templates: Set[str] = set()
798 self.intermediate: bool = False
800 def exists(self) -> bool:
801 """
802 Test whether file exists on disk.
804 This function is lazy; once it has been checked once it will keep returning the same result.
806 To reset use :meth:`reset_existence`.
807 """
808 if not hasattr(self, "_exists"):
809 self._exists = self.filename.is_symlink() or self.filename.exists()
810 return self._exists
812 def reset_existence(
813 self,
814 ):
815 """Ensure existence is checked again when running :meth:`exists`."""
816 if hasattr(self, "_exists"):
817 del self._exists
819 def expected(
820 self,
821 ):
822 """
823 Return whether the file can be produced by the pipeline (or already exists).
825 Returns False if the file does not exist and there is no way to produce it. Otherwise, True is returned
826 """
827 if self.exists():
828 return True
829 if self.producer is None:
830 return False
831 return self.producer.expected()
833 def mark_intermediate(self, intermediate=True):
834 """
835 Mark this file as an intermediate file.
837 Intermediate files will not be reproduced if they are missing when checking which jobs need to be run.
838 """
839 self.intermediate = intermediate
841 @property
842 def producer(
843 self,
844 ) -> SingleJob:
845 """Job that can produce this file."""
846 return self._producer
848 @producer.setter
849 def producer(self, new_value):
850 if self._producer is not None:
851 if self._producer is new_value:
852 return
853 raise ValueError(
854 f"{self} can be produced by both {self.producer} and {new_value}"
855 )
856 self._producer = new_value
858 def __repr__(
859 self,
860 ):
861 """Print filename of target."""
862 return f"FileTarget({str(self.filename)})"
865@contextmanager
866def update_closure(*dictionaries, **kwargs):
867 """Add the provided dictionaries to the globals dictionary.
869 Inside the `with` block all the dictionary keys will be available in the local environment.
870 After the `with` block ends the local environment will be cleaned up again.
872 Use like this:
874 .. code-block:: python
876 def my_func()
877 with update_closure({'a': 3}):
878 print(a) # prints 3
879 print(a) # raises a NameError (or prints whatever `a` is set to in the global environment)
880 """
881 func_globals = inspect.stack()[2].frame.f_globals
883 # merge input dictionaries
884 new_kwargs = {}
885 for d in dictionaries:
886 new_kwargs.update(d)
887 new_kwargs.update(kwargs)
889 to_restore = {}
890 to_delete = set()
891 for key in new_kwargs:
892 if key in func_globals:
893 to_restore[key] = func_globals[key]
894 else:
895 to_delete.add(key)
897 func_globals.update(new_kwargs)
898 yield # return to run the code within the with-block
900 # clean up the environment
901 func_globals.update(to_restore)
902 for key in to_delete:
903 del func_globals[key]
904 del func_globals
907def call_batched_jobs(funcs, kwargs):
908 for f, k in zip(funcs, kwargs):
909 f(**k)
912class BatchJob(JobParent):
913 """Batched combination of multiple `SingleJob` instances that will be submitted together."""
915 function = staticmethod(call_batched_jobs)
917 def __init__(self, *sub_jobs):
918 """Creates a new `BatchJob` from sub_jobs.
920 `sub_jobs` can be either `SingleJob` or `BatchJob`. In either case they will be run in the order in which they are supplied.
921 """
922 self.torun: Sequence[SingleJob] = []
923 for job in sub_jobs:
924 if isinstance(job, BatchJob):
925 self.torun.extend(job.torun)
926 else:
927 assert isinstance(job, SingleJob)
928 self.torun.append(job)
929 self.torun_set = set(self.torun)
931 self.output_targets = {t for job in self.torun for t in job.output_targets}
932 for t in self.output_targets:
933 t._producer = self
934 self.input_targets = {
935 t
936 for job in self.torun
937 for t in job.input_targets
938 if t not in self.output_targets
939 }
940 for t in self.input_targets:
941 t.required_by.add(self)
943 self.optional_targets = set.union(
944 {
945 t
946 for t in self.output_targets
947 if not any(
948 t in job.output_targets and t not in job.optional_targets
949 for job in self.torun
950 )
951 },
952 {
953 t
954 for t in self.input_targets
955 if not any(
956 t in job.input_targets and t not in job.optional_targets
957 for job in self.torun
958 )
959 },
960 )
962 as_tuples = set.intersection(
963 *[
964 {(key, value) for key, value in job.set_parameters.items()}
965 for job in self.torun
966 ]
967 )
968 self.set_parameters = {key: value for key, value in as_tuples}
970 def copy(self, targets: Dict[str, "FileTarget"]):
971 """
972 Create a copy of the `BatchJob` to be included in a new :class:`JobList`.
974 `targets` contain the set of FileTargets recognised by this new `JobList`.
975 This will be updated based on the input/output targets of each job within this batch.
976 """
977 return BatchJob(*[job.copy(targets) for job in self.torun])
979 @property
980 def batch(
981 self,
982 ):
983 b = self.torun[0].batch
984 if all(job.batch is not None and job.batch == b for job in self.torun):
985 return b
986 return None
988 @property
989 def kwargs(
990 self,
991 ):
992 return {
993 "funcs": [job.function for job in self.torun],
994 "kwargs": [job.kwargs for job in self.torun],
995 }
997 @property
998 def submit_params(
999 self,
1000 ):
1001 def sum_value(name, *times):
1002 total = sum(t for t in times if t is not None)
1003 return None if total == 0 else total
1005 def max_value(name, *rams):
1006 if any(r is not None for r in rams):
1007 return max(int(r) for r in rams if r is not None)
1008 return None
1010 def extend(name, *vars):
1011 all_vars = set()
1012 for v in vars:
1013 if v is not None:
1014 all_vars.update(v)
1015 return list(all_vars)
1017 def merge_name(_, *names):
1018 if len(names) <= 3:
1019 return "-".join(names)
1020 return "batch_job"
1022 def unique_param(name, *flags):
1023 not_none = {f for f in flags if f is not None}
1024 if len(not_none) == 0:
1025 return None
1026 elif len(not_none) == 1:
1027 return not_none.pop()
1028 else:
1029 raise ValueError(
1030 f"Multiple different values for submit parameter flag '{name}' found. fsl-pipe does not know which one to choose out of: {not_none}"
1031 )
1033 to_combine = {
1034 "jobtime": sum_value,
1035 "jobram": max_value,
1036 "threads": max_value,
1037 "export_vars": extend,
1038 "name": merge_name,
1039 }
1041 all_flags = {flag for job in self.torun for flag in job.submit_params.keys()}
1042 return {
1043 flag: to_combine.get(flag, unique_param)(
1044 flag, *[job.submit_params.get(flag, None) for job in self.torun]
1045 )
1046 for flag in all_flags
1047 }
1049 def job_name(
1050 self,
1051 ):
1052 """Return a string representation of this job."""
1053 base = "" if self.batch is None else self.batch + "_"
1054 func_name = "-".join(sorted({job._function_name for job in self.torun}))
1055 if len(self.set_parameters) > 0:
1056 parameter_string = "_".join(
1057 [f"{key}-{value}" for key, value in self.set_parameters.items()]
1058 )
1059 name = f"{func_name}_{parameter_string}"
1060 else:
1061 name = func_name
1062 value = re.sub(r"[^\w\s-]", "", name).strip().lower()
1063 return base + re.sub(r"[-\s]+", "-", value)
1065 def scale_jobtime(self, scaling):
1066 """
1067 Scale the submit job time in place by `scaling`.
1069 This will only affect jobs submitted to the cluster.
1070 """
1071 for job in self.torun:
1072 job.scale_jobtime(scaling)
1074 def __repr__(
1075 self,
1076 ):
1077 """Print job as a function call."""
1078 return f"Batch({self.torun})"
1080 def to_job_list(self, tree: FileTree):
1081 """Convert batched jobs back into a :class:`JobList`."""
1082 result = JobList(tree, self.torun, {}).copy()
1083 return self.batch, self.submit_params, result
1086def has_dependencies(
1087 possible_parent: JobParent, possible_children: Set[JobParent], all_jobs
1088):
1089 if possible_parent in possible_children:
1090 return True
1091 for _, child in possible_parent.dependencies(only_missing=False):
1092 if child in all_jobs and has_dependencies(child, possible_children, all_jobs):
1093 return True
1094 return False
1097def batch_connected_jobs(
1098 to_batch: Sequence[SingleJob], other_jobs: Collection[SingleJob]
1099) -> List[JobParent]:
1100 """Iteratively combines two jobs
1102 Two jobs will be batched if:
1103 1. They are both in `to_batch`
1104 2. One of the jobs depends on the other
1105 3. There is no other job that needs to run between these two jobs.
1106 Iteration through the jobs will continue until no further batching is possible.
1107 """
1108 other_jobs = set(other_jobs)
1109 batches = set(to_batch)
1110 nbatches_prev = len(batches) * 2
1111 while len(batches) != nbatches_prev:
1112 nbatches_prev = len(batches)
1113 for job in set(batches):
1114 if job not in batches:
1115 continue
1116 for _, dependency in job.dependencies(only_missing=False):
1117 if job in batches and dependency in batches:
1118 # properties 1 and 2 are met; now checking property 3
1119 all_jobs = set.union(batches, other_jobs)
1120 indirect_dependency = any(
1121 has_dependencies(d, {dependency}, all_jobs)
1122 for _, d in job.dependencies(only_missing=False)
1123 if d in all_jobs and d != dependency
1124 )
1125 if not indirect_dependency:
1126 batches.remove(job)
1127 batches.remove(dependency)
1128 batches.add(BatchJob(dependency, job))
1129 break
1130 return list(batches)
1133def batch_unconnected_jobs(
1134 to_batch: Sequence[JobParent], other_jobs: Collection[JobParent]
1135) -> List[BatchJob]:
1136 """Batch jobs into generational sets
1138 All jobs that do not depend on any other job in `to_batch` will be added to the first generational batch.
1139 All jobs that only depend on the first generational batch will be added to the second.
1140 And so forth, until all jobs have been added to a generation.
1142 Only dependencies through `other_jobs` will be considered.
1143 """
1144 to_batch = set(to_batch)
1145 all_jobs = set.union(to_batch, other_jobs)
1147 generations = []
1149 while len(to_batch) > 0:
1150 generations.append(set())
1151 for job in to_batch:
1152 if not any(
1153 has_dependencies(dependency, to_batch, all_jobs)
1154 for _, dependency in job.dependencies(only_missing=False)
1155 if dependency in all_jobs
1156 ):
1157 generations[-1].add(job)
1158 to_batch.difference_update(generations[-1])
1159 return [jobs.pop() if len(jobs) == 1 else BatchJob(*jobs) for jobs in generations]