fix: park a migration whose chain task is killed

pilot/core/bench/migration/operation.py

@@ -246,6 +246,30 @@ def migrate_site(
         if self.next_migrate_site() is None:
             self._complete(on_step)
 
+    def strand(self, reason: str) -> None:
+        """Park a chain whose task died without unwinding: a kill, an OOM, a reboot.
+
+        Every other write to this record happens inside the task process, so a
+        task that never raises leaves the operation spinning in a working state
+        forever - no retry, no restore, and no new update allowed past it.
+        """
+        target = self.state.failure_target
+        if target is None:
+            return
+        self.diagnosis = {"phase": self.state.name, "message": reason, "output_excerpt": ""}
+        if target == "revert_failed":
+            self._transition(target)
+            return
+        self._enter_needs_attention(self.state.name, self._interrupted_site())
+
+    def _interrupted_site(self) -> str | None:
+        """The site whose per-site work was in flight when the chain died."""
+        if self.state == "backing_up":
+            return next((site.name for site in self.sites if site.backup_status == "backing_up"), None)
+        if self.state == "migrating":
+            return next((site.name for site in self.sites if site.migration_status == "running"), None)
+        return None
+
     def retry(self) -> None:
         """Resume the chain from needs_attention so the failed unit runs again."""
         if self.state != "needs_attention":

pilot/core/bench/migration/state.py

@@ -20,6 +20,7 @@ class MigrationState:
     is_terminal: ClassVar[bool] = False
     is_failure: ClassVar[bool] = False  # paused on a failure, waiting for the user
     starts_work: ClassVar[bool] = False
+    failure_target: ClassVar[str | None] = None  # where work parks when its task dies
     allowed: ClassVar[frozenset[str]] = frozenset()
 
     def next_step(self, operation: "MigrationOperation") -> ChainStep | None:
@@ -51,6 +52,7 @@ class BackingUp(MigrationState):
     name = "backing_up"
     label = "Backing up"
     starts_work = True
+    failure_target = "needs_attention"
     allowed = frozenset({"updating", "migrating", "needs_attention"})
 
     def next_step(self, operation: "MigrationOperation") -> ChainStep | None:
@@ -62,6 +64,7 @@ class Updating(MigrationState):
     name = "updating"
     label = "Updating apps"
     starts_work = True
+    failure_target = "needs_attention"
     allowed = frozenset({"migrating", "needs_attention"})
 
     def next_step(self, operation: "MigrationOperation") -> ChainStep | None:
@@ -72,6 +75,7 @@ class Migrating(MigrationState):
     name = "migrating"
     label = "Migrating"
     starts_work = True
+    failure_target = "needs_attention"
     allowed = frozenset({"completed", "needs_attention"})
 
     def next_step(self, operation: "MigrationOperation") -> ChainStep | None:
@@ -96,6 +100,7 @@ class RevertingApps(MigrationState):
     name = "reverting_apps"
     label = "Reverting app revisions"
     starts_work = True
+    failure_target = "revert_failed"
     allowed = frozenset({"reverting_sites", "restarting", "revert_failed"})
 
     def next_step(self, operation: "MigrationOperation") -> ChainStep | None:
@@ -106,6 +111,7 @@ class RevertingSites(MigrationState):
     name = "reverting_sites"
     label = "Recovering sites"
     starts_work = True
+    failure_target = "revert_failed"
     allowed = frozenset({"restarting", "revert_failed"})
 
     def next_step(self, operation: "MigrationOperation") -> ChainStep | None:
@@ -117,6 +123,7 @@ class Restarting(MigrationState):
     name = "restarting"
     label = "Restarting services"
     starts_work = True
+    failure_target = "revert_failed"
     allowed = frozenset({"reverted", "revert_failed"})
 
     def next_step(self, operation: "MigrationOperation") -> ChainStep | None:

pilot/core/bench/migration/store.py

@@ -7,8 +7,14 @@
 
 from pilot.core.bench.migration.operation import AppRevision, MigrationOperation, SiteProgress
 from pilot.core.bench.migration.state import get_state
-from pilot.exceptions import BenchError, MigrationConflictError, MigrationNotFoundError
+from pilot.exceptions import (
+    BenchError,
+    MigrationConflictError,
+    MigrationNotFoundError,
+    TaskNotFoundError,
+)
 from pilot.internal.atomic_file import atomic_write_private_text
+from pilot.internal.tasks.store import TaskStore
 from pilot.utils import make_private_directory
 
 if TYPE_CHECKING:
@@ -74,9 +80,30 @@ def get(self, operation_id: str) -> MigrationOperation:
             raise MigrationNotFoundError(f"Migration operation not found: {operation_id}")
         try:
             data = json.loads(path.read_text(encoding="utf-8"))
-            return MigrationOperation.from_dict(data, self.bench, self)
+            operation = MigrationOperation.from_dict(data, self.bench, self)
         except (OSError, ValueError, KeyError, TypeError) as error:
             raise BenchError(f"Could not load migration operation {path.name}: {error}") from error
+        self._strand_if_abandoned(operation)
+        return operation
+
+    def _strand_if_abandoned(self, operation: MigrationOperation) -> None:
+        """Park an operation whose chain ended without anything advancing the record.
+
+        The task callback covers a killed link. This covers the runs where the
+        callback never got to fire either - a host reboot, or the task wrapper
+        itself being killed - so the operation still heals on the next read.
+        """
+        if operation.state.failure_target is None or not operation.chain:
+            return
+        task_id = operation.chain[-1].get("task_id")
+        if not isinstance(task_id, str):
+            return
+        try:
+            status = TaskStore(self.bench.path).read_status(task_id)
+        except (OSError, ValueError, TaskNotFoundError):
+            return
+        if status.is_terminal:
+            operation.strand(f"The {operation.state.label.lower()} step stopped without reporting back.")
 
     def get_all(self) -> list[MigrationOperation]:
         if not self.root.exists():

pilot/internal/tasks/callbacks.py

@@ -93,11 +93,25 @@ def _reload_workers(meta: dict, args: dict) -> None:
     Bench(Path(meta["bench_root"])).reload_workers(web_only=args.get("web_only", False))
 
 
+def _strand_migration(meta: dict, args: dict) -> None:
+    """Park a migration whose chain link died. Runs for every failed link, so it is
+    a no-op once the link already reported the failure from inside its own process."""
+    from pilot.core.bench import Bench
+    from pilot.exceptions import MigrationNotFoundError
+
+    try:
+        operation = Bench(Path(meta["bench_root"])).migrations.get(args["operation_id"])
+    except MigrationNotFoundError:
+        return
+    operation.strand(f"The {operation.state.label.lower()} step stopped before it could report back.")
+
+
 _OPERATIONS: dict[str, CallbackOperation] = {
     "cleanup-site-restore": _cleanup_site_restore,
     "remove-failed-site": _remove_failed_site,
     "disable-site-ssl": _disable_site_ssl,
     "reload-workers": _reload_workers,
+    "strand-migration": _strand_migration,
 }
 
 

pilot/tasks/migrate.py

@@ -1,27 +1,19 @@
-import sys
 from dataclasses import dataclass
 from typing import ClassVar
 
-from pilot.tasks import Task
+from pilot.tasks.migration_chain import MigrationChainTask
 
 
 @dataclass(kw_only=True)
-class MigrateTask(Task):
+class MigrateTask(MigrationChainTask):
     """Chain link: migrate one site, then queue the next site (or complete)."""
 
     command: ClassVar[str] = "migrate"
 
-    operation_id: str
     site: str
 
-    def run(self) -> None:
-        operation = self.bench.migrations.get(self.operation_id)
-        try:
-            operation.migrate_site(self.site, on_step=self.step, on_progress=self.report)
-        except Exception:
-            self.step_failed()
-            sys.exit(1)
-        operation.enqueue_next(handoff_from=operation.chain[-1]["task_id"])
+    def run_step(self, operation) -> None:
+        operation.migrate_site(self.site, on_step=self.step, on_progress=self.report)
 
 
 if __name__ == "__main__":

pilot/tasks/migration_backup.py

@@ -1,29 +1,20 @@
-import sys
 from dataclasses import dataclass
 from typing import ClassVar
 
-from pilot.tasks import Task
+from pilot.tasks.migration_chain import MigrationChainTask
 
 
 @dataclass(kw_only=True)
-class MigrationBackupTask(Task):
+class MigrationBackupTask(MigrationChainTask):
     """Chain link: back up one site's tables before a migration, then queue the next step."""
 
     command: ClassVar[str] = "migration-backup"
 
-    operation_id: str
     site: str
 
-    def run(self) -> None:
-        operation = self.bench.migrations.get(self.operation_id)
-        try:
-            operation.back_up_site(self.site, on_step=self.step, on_progress=self.report)
-        except Exception:
-            self.step_failed()
-            sys.exit(1)
-        operation.enqueue_next(handoff_from=operation.chain[-1]["task_id"])
+    def run_step(self, operation) -> None:
+        operation.back_up_site(self.site, on_step=self.step, on_progress=self.report)
 
 
 if __name__ == "__main__":
     MigrationBackupTask.main()
-

pilot/tasks/migration_chain.py

@@ -0,0 +1,34 @@
+import sys
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from pilot.tasks import Task, on_cancel, on_failure
+
+if TYPE_CHECKING:
+    from pilot.core.bench.migration.operation import MigrationOperation
+
+
+@dataclass(kw_only=True)
+class MigrationChainTask(Task):
+    """One link of a migration chain: run its step, then queue whatever the operation wants next."""
+
+    operation_id: str
+
+    def run(self) -> None:
+        operation = self.bench.migrations.get(self.operation_id)
+        try:
+            self.run_step(operation)
+        except Exception:
+            self.step_failed()
+            sys.exit(1)
+        operation.enqueue_next(handoff_from=operation.chain[-1]["task_id"])
+
+    def run_step(self, operation: "MigrationOperation") -> None:
+        raise NotImplementedError
+
+    @on_failure
+    @on_cancel
+    def strand_migration(self) -> dict:
+        """A killed link never reaches its own error handling, so the operation is
+        parked from outside instead - by the wrapper, or by task reconciliation."""
+        return {"operation_id": self.operation_id}

pilot/tasks/restart_services.py

@@ -1,26 +1,17 @@
-import sys
 from dataclasses import dataclass
 from typing import ClassVar
 
-from pilot.tasks import Task
+from pilot.tasks.migration_chain import MigrationChainTask
 
 
 @dataclass(kw_only=True)
-class RestartServicesTask(Task):
+class RestartServicesTask(MigrationChainTask):
     """Chain link: restart services to finish a restore, then mark the operation reverted."""
 
     command: ClassVar[str] = "restart-services"
 
-    operation_id: str
-
-    def run(self) -> None:
-        operation = self.bench.migrations.get(self.operation_id)
-        try:
-            operation.restart(on_step=self.step)
-        except Exception:
-            self.step_failed()
-            sys.exit(1)
-        operation.enqueue_next(handoff_from=operation.chain[-1]["task_id"])
+    def run_step(self, operation) -> None:
+        operation.restart(on_step=self.step)
 
 
 if __name__ == "__main__":

pilot/tasks/revert_apps.py

@@ -1,26 +1,17 @@
-import sys
 from dataclasses import dataclass
 from typing import ClassVar
 
-from pilot.tasks import Task
+from pilot.tasks.migration_chain import MigrationChainTask
 
 
 @dataclass(kw_only=True)
-class RevertAppsTask(Task):
+class RevertAppsTask(MigrationChainTask):
     """Chain link: roll app revisions back and rebuild, then queue the next revert step."""
 
     command: ClassVar[str] = "revert-apps"
 
-    operation_id: str
-
-    def run(self) -> None:
-        operation = self.bench.migrations.get(self.operation_id)
-        try:
-            operation.revert_apps(on_step=self.step, on_progress=self.report)
-        except Exception:
-            self.step_failed()
-            sys.exit(1)
-        operation.enqueue_next(handoff_from=operation.chain[-1]["task_id"])
+    def run_step(self, operation) -> None:
+        operation.revert_apps(on_step=self.step, on_progress=self.report)
 
 
 if __name__ == "__main__":

pilot/tasks/revert_site.py

@@ -1,27 +1,19 @@
-import sys
 from dataclasses import dataclass
 from typing import ClassVar
 
-from pilot.tasks import Task
+from pilot.tasks.migration_chain import MigrationChainTask
 
 
 @dataclass(kw_only=True)
-class RevertSiteTask(Task):
+class RevertSiteTask(MigrationChainTask):
     """Chain link: restore one site's database and clear its cache, then queue the next site."""
 
     command: ClassVar[str] = "revert-site"
 
-    operation_id: str
     site: str
 
-    def run(self) -> None:
-        operation = self.bench.migrations.get(self.operation_id)
-        try:
-            operation.revert_site(self.site, on_step=self.step, on_progress=self.report)
-        except Exception:
-            self.step_failed()
-            sys.exit(1)
-        operation.enqueue_next(handoff_from=operation.chain[-1]["task_id"])
+    def run_step(self, operation) -> None:
+        operation.revert_site(self.site, on_step=self.step, on_progress=self.report)
 
 
 if __name__ == "__main__":

pilot/tasks/update.py

@@ -1,26 +1,17 @@
-import sys
 from dataclasses import dataclass
 from typing import ClassVar
 
-from pilot.tasks import Task
+from pilot.tasks.migration_chain import MigrationChainTask
 
 
 @dataclass(kw_only=True)
-class UpdateTask(Task):
+class UpdateTask(MigrationChainTask):
     """Chain link: update/reinstall/rebuild apps, then queue the first site migration."""
 
     command: ClassVar[str] = "update"
 
-    operation_id: str
-
-    def run(self) -> None:
-        operation = self.bench.migrations.get(self.operation_id)
-        try:
-            operation.update_apps(on_step=self.step, on_progress=self.report)
-        except Exception:
-            self.step_failed()
-            sys.exit(1)
-        operation.enqueue_next(handoff_from=operation.chain[-1]["task_id"])
+    def run_step(self, operation) -> None:
+        operation.update_apps(on_step=self.step, on_progress=self.report)
 
 
 if __name__ == "__main__":