BUG: Fix size_per_learner (#257)

* BUG: Fix size_per_learner

* Add tests

Author: basnijholt

adaptive_scheduler/_executor.py

@@ -114,16 +114,25 @@ def _get(self) -> Any | None:  # noqa: PLR0911
         if self.done():
             return super().result(timeout=0)
 
-        idx_learner, idx_data = self.task_id
-        learner, fname = self._learner_and_fname
+        func_id, global_index = self.task_id
+        try:
+            learner_idx, local_index = self.executor._task_mapping[(func_id, global_index)]
+        except KeyError as e:
+            msg = "Task mapping not found; finalize() must be called first."
+            raise RuntimeError(msg) from e
+        # Now retrieve the correct learner and filename:
+        run_manager = self.executor._run_manager
+        assert run_manager is not None, "RunManager not initialized"
+        learner = run_manager.learners[learner_idx]
+        fname = run_manager.fnames[learner_idx]
 
         if learner.done():
-            result = learner.data[idx_data]
+            result = learner.data[local_index]
             self.set_result(result)
             return result
 
         assert self.executor._run_manager is not None
-        last_load_time = self.executor._run_manager._last_load_time.get(idx_learner, 0)
+        last_load_time = self.executor._run_manager._last_load_time.get(learner_idx, 0)
         now = time.monotonic()
         time_since_last_load = now - last_load_time
         if time_since_last_load < self.min_load_interval:
@@ -141,10 +150,10 @@ def _get(self) -> Any | None:  # noqa: PLR0911
         learner.load(fname)
         self._load_time = time.monotonic() - now
         self.min_load_interval = max(1.0, 20.0 * self._load_time)
-        self.executor._run_manager._last_load_time[idx_learner] = now
+        self.executor._run_manager._last_load_time[learner_idx] = now
 
-        if idx_data in learner.data:
-            result = learner.data[idx_data]
+        if local_index in learner.data:
+            result = learner.data[local_index]
             self.set_result(result)
             return result
         return None
@@ -371,6 +380,7 @@ class SlurmExecutor(AdaptiveSchedulerExecutorBase):
     _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
+    _task_mapping: dict[tuple[int, int], tuple[int, int]] = field(default_factory=dict)
 
     def __post_init__(self) -> None:
         if self.folder is None:
@@ -390,32 +400,48 @@ def submit(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> SlurmT
         task_id = TaskID(self._sequence_mapping[fn], i)
         return SlurmTask(self, task_id)
 
-    def _to_learners(self) -> tuple[list[SequenceLearner], list[Path]]:
+    def _to_learners(
+        self,
+    ) -> tuple[
+        list[SequenceLearner],
+        list[Path],
+        dict[tuple[int, int], tuple[int, int]],
+    ]:
         learners = []
         fnames = []
-        for func, args_kwargs_list in self._sequences.items():
-            # Chunk the sequence if size_per_learner is specified
+        task_mapping = {}
+        learner_idx = 0
+        for func, args_list in self._sequences.items():
+            func_id = self._sequence_mapping[func]
+            # Chunk the sequence if size_per_learner is set; otherwise one chunk.
             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)
+                chunked_args = [
+                    args_list[i : i + self.size_per_learner]
+                    for i in range(0, len(args_list), self.size_per_learner)
                 ]
             else:
-                chunked_args_kwargs_list = [args_kwargs_list]
+                chunked_args = [args_list]
+
+            global_index = 0  # global index for tasks of this function
+            for chunk in chunked_args:
+                # Map each task in the chunk: global index -> (current learner, local index)
+                for local_index in range(len(chunk)):
+                    task_mapping[(func_id, global_index)] = (learner_idx, local_index)
+                    global_index += 1
 
-            for i, chunk in enumerate(chunked_args_kwargs_list):
                 learner = SequenceLearner(_SerializableFunctionSplatter(func), chunk)
                 learners.append(learner)
+                name = func.__name__ if hasattr(func, "__name__") else "func"
                 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
+                fnames.append(self.folder / f"{name}-{learner_idx}-{uuid.uuid4().hex}.pickle")
+                learner_idx += 1
+        return learners, fnames, task_mapping
 
     def finalize(self, *, start: bool = True) -> adaptive_scheduler.RunManager | None:
         if self._run_manager is not None:
             msg = "RunManager already initialized. Create a new SlurmExecutor instance."
             raise RuntimeError(msg)
-        learners, fnames = self._to_learners()
+        learners, fnames, self._task_mapping = self._to_learners()
         if not learners:
             return None
         assert self.folder is not None

tests/test_slurm_executor.py

@@ -176,7 +176,7 @@ def test_cleanup(executor: SlurmExecutor) -> None:
 def test_task_get_before_finalize(executor: SlurmExecutor) -> None:
     """Test that _get before finalize returns None."""
     task = executor.submit(example_func, 1.0)
-    with pytest.raises(AssertionError, match="RunManager not initialized"):
+    with pytest.raises(RuntimeError, match="Task mapping not found; finalize()"):
         task._get()
 
 
@@ -272,3 +272,108 @@ async def simulate_result() -> None:
     asyncio.create_task(simulate_result())  # noqa: RUF006
     result = await task
     assert result == 42
+
+
+@pytest.mark.usefixtures("_mock_slurm_partitions")
+@pytest.mark.usefixtures("_mock_slurm_queue")
+def test_to_learners_mapping_single_function(tmp_path: Path) -> None:
+    """Test that _to_learners creates the correct mapping for a single function."""
+    executor = SlurmExecutor(folder=tmp_path, size_per_learner=2)
+    # Submit 5 tasks to example_func so that they are split into chunks of 2.
+    for i in range(5):
+        executor.submit(example_func, i)
+    learners, fnames, mapping = executor._to_learners()
+
+    # We expect ceil(5/2) = 3 learners.
+    assert len(learners) == 3
+
+    func_id = executor._sequence_mapping[example_func]
+    expected_mapping = {
+        (func_id, 0): (0, 0),  # first learner, first task
+        (func_id, 1): (0, 1),  # first learner, second task
+        (func_id, 2): (1, 0),  # second learner, first task
+        (func_id, 3): (1, 1),  # second learner, second task
+        (func_id, 4): (2, 0),  # third learner, first task (only one task in this chunk)
+    }
+    assert mapping == expected_mapping
+
+
+@pytest.mark.usefixtures("_mock_slurm_partitions")
+@pytest.mark.usefixtures("_mock_slurm_queue")
+def test_finalize_mapping_and_learners(tmp_path: Path) -> None:
+    """Test that finalize() sets the task mapping correctly and creates the right number of learners."""
+    executor = SlurmExecutor(folder=tmp_path, size_per_learner=2)
+    # Submit 3 tasks to example_func.
+    for i in range(3):
+        executor.submit(example_func, i)
+
+    rm = executor.finalize(start=False)
+    # For 3 tasks with chunk size 2:
+    #   - The first chunk (learner 0) has tasks 0 and 1.
+    #   - The second chunk (learner 1) has task 2.
+    func_id = executor._sequence_mapping[example_func]
+    expected_mapping = {
+        (func_id, 0): (0, 0),
+        (func_id, 1): (0, 1),
+        (func_id, 2): (1, 0),
+    }
+    assert executor._task_mapping == expected_mapping
+    # Also, the run manager should have 2 learners.
+    assert isinstance(rm, RunManager)
+    assert len(rm.learners) == 2
+
+
+@pytest.mark.usefixtures("_mock_slurm_partitions")
+@pytest.mark.usefixtures("_mock_slurm_queue")
+def test_task_get_with_chunking(tmp_path: Path) -> None:
+    """Test that tasks in different learners retrieve the correct result when using size_per_learner."""
+    executor = SlurmExecutor(folder=tmp_path, size_per_learner=2, save_interval=1)
+    # Submit three tasks; with size_per_learner=2, this will produce 2 learners.
+    task1 = executor.submit(example_func, 42)
+    task2 = executor.submit(example_func, 43)
+    task3 = executor.submit(example_func, 44)
+    rm = executor.finalize(start=False)
+
+    # For learner 0 (tasks 0 and 1)
+    assert isinstance(rm, RunManager)
+    learner0 = rm.learners[0]
+    fname0 = rm.fnames[0]
+    learner0.data[0] = 42
+    learner0.data[1] = 43
+    learner0.save(fname0)
+    # For learner 1 (task 2)
+    learner1 = rm.learners[1]
+    fname1 = rm.fnames[1]
+    learner1.data[0] = 44
+    learner1.save(fname1)
+
+    # _get() should now retrieve the correct values based on the mapping.
+    assert task1._get() == 42
+    assert task2._get() == 43
+    assert task3._get() == 44
+
+
+@pytest.mark.usefixtures("_mock_slurm_partitions")
+@pytest.mark.usefixtures("_mock_slurm_queue")
+def test_mapping_multiple_functions(tmp_path: Path) -> None:
+    """Test that the mapping is correct when tasks are submitted for multiple functions."""
+    executor = SlurmExecutor(folder=tmp_path, size_per_learner=2)
+    # Submit two tasks for example_func and two for another_func.
+    executor.submit(example_func, 10)
+    executor.submit(example_func, 20)
+    executor.submit(another_func, 5)
+    executor.submit(another_func, 6)
+
+    # Directly call _to_learners to examine the mapping.
+    learners, fnames, mapping = executor._to_learners()
+
+    expected_mapping = {
+        # For example_func: two tasks in one learner (since 2 tasks fit in one chunk).
+        (executor._sequence_mapping[example_func], 0): (0, 0),
+        (executor._sequence_mapping[example_func], 1): (0, 1),
+        # For another_func: two tasks in one learner.
+        (executor._sequence_mapping[another_func], 0): (1, 0),
+        (executor._sequence_mapping[another_func], 1): (1, 1),
+    }
+    assert mapping == expected_mapping
+    assert len(learners) == 2