ENH: Implement size_per_learner in SlurmExecutor

adaptive_scheduler/_executor.py

@@ -317,6 +317,8 @@ class SlurmExecutor(AdaptiveSchedulerExecutorBase):
         Extra keyword arguments to pass to the `RunManager`.
     extra_scheduler_kwargs
         Extra keyword arguments to pass to the `adaptive_scheduler.scheduler.SLURM`.
+    size_per_learner
+        The size of each learner. If None, the whole sequence is passed to the learner.
 
     """
 
@@ -362,6 +364,7 @@ class SlurmExecutor(AdaptiveSchedulerExecutorBase):
     extra_run_manager_kwargs: dict[str, Any] | None = None
     extra_scheduler_kwargs: dict[str, Any] | None = None
     # Internal
+    size_per_learner: int | None = None
     _sequences: dict[Callable[..., Any], list[Any]] = field(default_factory=dict)
     _sequence_mapping: dict[Callable[..., Any], int] = field(default_factory=dict)
     _run_manager: adaptive_scheduler.RunManager | None = None
@@ -388,11 +391,21 @@ def _to_learners(self) -> tuple[list[SequenceLearner], list[Path]]:
         learners = []
         fnames = []
         for func, args_kwargs_list in self._sequences.items():
-            learner = SequenceLearner(_SerializableFunctionSplatter(func), args_kwargs_list)
-            learners.append(learner)
-            assert isinstance(self.folder, Path)
-            name = func.__name__ if hasattr(func, "__name__") else ""
-            fnames.append(self.folder / f"{name}-{uuid.uuid4().hex}.pickle")
+            # Chunk the sequence if size_per_learner is specified
+            if self.size_per_learner is not None:
+                chunked_args_kwargs_list = [
+                    args_kwargs_list[i : i + self.size_per_learner]
+                    for i in range(0, len(args_kwargs_list), self.size_per_learner)
+                ]
+            else:
+                chunked_args_kwargs_list = [args_kwargs_list]
+
+            for i, chunk in enumerate(chunked_args_kwargs_list):
+                learner = SequenceLearner(_SerializableFunctionSplatter(func), chunk)
+                learners.append(learner)
+                assert isinstance(self.folder, Path)
+                name = func.__name__ if hasattr(func, "__name__") else ""
+                fnames.append(self.folder / f"{name}-{i}-{uuid.uuid4().hex}.pickle")
         return learners, fnames
 
     def finalize(self, *, start: bool = True) -> adaptive_scheduler.RunManager:

adaptive_scheduler/widgets.py

@@ -106,7 +106,7 @@ def _sort_key(value: tuple[float | str, str]) -> tuple[float | int, str]:
 
     def _vec_timedelta(ts: pd.Timestamp) -> str:
         now = np.datetime64(datetime.now())  # noqa: DTZ005
-        dt = np.timedelta64(now - ts, "s")  # type: ignore[operator]
+        dt = np.timedelta64(now - ts, "s")  # type: ignore[operator, call-overload]
         return f"{dt} ago"
 
     mapping = {