Set i_batch for SubmitterLocal when submitting (#164)

XENONnT · Jun 7, 2024 · cce3ea0 · cce3ea0
1 parent 41119c7
commit cce3ea0
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/alea/submitter.py b/alea/submitter.py
@@ -62,6 +62,7 @@ class Submitter:
     config_file_path: str
     template_path: str
     combine_n_jobs: int = 1
+    first_i_batch: int = 0
     allowed_special_args: List[str] = []
     logging = logging.getLogger("submitter_logger")
 
@@ -336,7 +337,7 @@ def computation_tickets_generator(self):
         _, _, annotations = Runner.runner_arguments()
 
         for runner_args in self.merged_arguments_generator():
-            for i_batch in range(runner_args.get("n_batch", 1)):
+            for i_batch in range(self.first_i_batch, runner_args.get("n_batch", 1)):
                 i_args = deepcopy(runner_args)
                 i_args["i_batch"] = i_batch
 

diff --git a/alea/submitters/local.py b/alea/submitters/local.py
@@ -33,6 +33,7 @@ def __init__(self, *args, **kwargs):
         self.local_configurations = kwargs.get("local_configurations", {})
         self.template_path = self.local_configurations.pop("template_path", None)
         self.combine_n_jobs = self.local_configurations.pop("combine_n_jobs", 1)
+        self.first_i_batch = kwargs.pop("i_batch", 1)
         super().__init__(*args, **kwargs)
 
     @staticmethod