MP checkpointing fix

Author: americalexander

activitysim/core/calibration.py

@@ -290,6 +290,99 @@ def run_calibration_loop(
         state.filesystem.pipeline_file_name = original_pipeline_name
 
 
+def _get_pipeline_checkpoint_names(
+    state: workflow.State,
+    pipeline_path,
+) -> list[str]:
+    """Read checkpoint names from a pipeline store file without fully restoring."""
+    from activitysim.core.workflow.checkpoint import HdfStore, ParquetStore
+
+    pipeline_path = Path(pipeline_path)
+    if state.settings.checkpoint_format == "hdf":
+        store = HdfStore(pipeline_path, mode="r")
+    else:
+        store = ParquetStore(pipeline_path, mode="r")
+    try:
+        return store.list_checkpoint_names()
+    finally:
+        store.close()
+
+
+def _resolve_checkpoint_name(
+    state: workflow.State,
+    resume_after: str,
+    checkpoint_names: list[str],
+) -> str:
+    """Resolve a model-name ``resume_after`` to an existing pipeline checkpoint.
+
+    In single-process pipelines model names are used directly as checkpoint
+    names.  In multiprocess pipelines only step-level checkpoint names (e.g.
+    ``mp_households``) exist in the main pipeline store.  This maps a model
+    name to the enclosing multiprocess step checkpoint so the calibration
+    module can restore the correct pipeline state.
+    """
+    from activitysim.core.exceptions import CheckpointNameNotFoundError
+
+    if resume_after in checkpoint_names:
+        return resume_after
+
+    all_models = state.settings.models
+    mp_steps = state.settings.multiprocess_steps
+
+    if not mp_steps or resume_after not in all_models:
+        raise CheckpointNameNotFoundError(
+            f"Checkpoint '{resume_after}' not found in pipeline and cannot "
+            f"be resolved. Available checkpoints: {checkpoint_names}"
+        )
+
+    resume_idx = all_models.index(resume_after)
+    step_boundaries = [all_models.index(s.begin) for s in mp_steps]
+    step_boundaries.append(len(all_models))
+
+    for i, step in enumerate(mp_steps):
+        if step_boundaries[i] <= resume_idx < step_boundaries[i + 1]:
+            if step.name in checkpoint_names:
+                return step.name
+            # Also check for calibration-prefixed step names from a prior
+            # calibration run's multiprocess steps.
+            for cp in checkpoint_names:
+                if cp.startswith(f"calibration_{step.name}"):
+                    return cp
+            break
+
+    raise CheckpointNameNotFoundError(
+        f"Cannot resolve '{resume_after}' to any available checkpoint. "
+        f"Available: {checkpoint_names}"
+    )
+
+
+def _effective_resume_index(
+    state: workflow.State,
+    checkpoint_name: str,
+) -> int:
+    """Return the index of the last model covered by a checkpoint.
+
+    For model-level checkpoints this is the model's own index in
+    ``settings.models``.  For step-level checkpoints this is the index of
+    the last model in the enclosing multiprocess step.
+    """
+    all_models = state.settings.models
+    if checkpoint_name in all_models:
+        return all_models.index(checkpoint_name)
+
+    mp_steps = state.settings.multiprocess_steps or []
+    step_boundaries = [all_models.index(s.begin) for s in mp_steps]
+    step_boundaries.append(len(all_models))
+
+    for i, step in enumerate(mp_steps):
+        if step.name == checkpoint_name or (
+            checkpoint_name.startswith(f"calibration_{step.name}")
+        ):
+            return step_boundaries[i + 1] - 1
+
+    return -1
+
+
 def _run_precursor_components(
     state: workflow.State,
     models: list[str],
@@ -306,7 +399,11 @@ def _run_precursor_components(
         prior_pipeline = state.checkpoint.store.filename
         state.checkpoint.close_store()
         state.filesystem.pipeline_file_name = f"pipeline_calibration_iter_{global_iter}"
-        state.checkpoint.restore_from(prior_pipeline, checkpoint_name=resume_after)
+        # In multiprocess pipelines the model-level checkpoint may not exist;
+        # resolve to the enclosing step checkpoint.
+        checkpoint_names = _get_pipeline_checkpoint_names(state, prior_pipeline)
+        resolved = _resolve_checkpoint_name(state, resume_after, checkpoint_names)
+        state.checkpoint.restore_from(prior_pipeline, checkpoint_name=resolved)
     else:
         _run_in_configured_mode(
             state,
@@ -1394,7 +1491,29 @@ def _run_in_configured_mode(
     if not models:
         return
 
+    # Remember the original last model — downstream code (e.g.
+    # _calibrate_component) references it as a checkpoint name.
+    original_last_model = models[-1]
+
     if state.settings.multiprocess:
+        # In multiprocess mode the checkpoint system uses step-level names,
+        # so model-level resume_after cannot be passed through directly.
+        # Instead, trim the models list to exclude those already completed.
+        if resume_after:
+            all_models = state.settings.models
+            if resume_after in all_models:
+                resume_idx = all_models.index(resume_after)
+                models = [m for m in models if all_models.index(m) > resume_idx]
+
+        if not models:
+            # All models are covered by the resume point — restore the
+            # pipeline state and create a model-level checkpoint so that
+            # downstream code can reference it by name.
+            if not state.checkpoint.store_is_open():
+                state.checkpoint.restore(resume_after="_")
+            state.checkpoint.add(original_last_model)
+            return
+
         _run_multiprocess_with_overrides(
             state,
             models=models,
@@ -1405,13 +1524,43 @@ def _run_in_configured_mode(
         # Restore it into the parent process state so tables are accessible
         # for calibration expression evaluation.
         _restore_parent_state_from_pipeline(state)
-        # Add a checkpoint named after the last model so that model-name
-        # references (e.g. _prior_step_name, resume_after on global_iter > 1)
-        # resolve correctly. Without this, only the step-level coalesce name
-        # exists in the pipeline.
-        state.checkpoint.add(models[-1])
+        # Add a checkpoint named after the original last model so that
+        # model-name references (e.g. _prior_step_name, resume_after on
+        # global_iter > 1) resolve correctly.  Without this, only the
+        # step-level coalesce name exists in the pipeline.
+        state.checkpoint.add(original_last_model)
         return
 
+    # Single-process mode: resolve resume_after to an existing checkpoint
+    # when the pipeline has step-level checkpoints from a prior multiprocess
+    # run (model-level checkpoints won't exist in that pipeline).
+    if resume_after:
+        pipeline_path = state.checkpoint.default_pipeline_file_path()
+        try:
+            checkpoint_names = _get_pipeline_checkpoint_names(
+                state, pipeline_path
+            )
+        except Exception:
+            checkpoint_names = []
+
+        if checkpoint_names and resume_after not in checkpoint_names:
+            resolved = _resolve_checkpoint_name(
+                state, resume_after, checkpoint_names
+            )
+            # The resolved checkpoint may cover more models than
+            # resume_after — trim the models list accordingly.
+            end_idx = _effective_resume_index(state, resolved)
+            if end_idx >= 0:
+                all_models = state.settings.models
+                models = [m for m in models if all_models.index(m) > end_idx]
+
+            if not models:
+                state.checkpoint.restore(resolved)
+                state.checkpoint.add(original_last_model)
+                return
+
+            resume_after = resolved
+
     state.run(
         models=models,
         resume_after=resume_after,