Fixed synchronization of multiple recovery workflows (#822)

This commit improves the synchronization of concurrent recovery
workflows. Previously, an error occurred when a dependent workflow
attempted to attach to an `InterWorkflowPort` of the dependee before
it was created. This race condition happened because port creation
occurred in a successor phase of job synchronization. This is now
resolved by acquiring a lock to check the status and holding it until
the recovery workflow responsible for the rollback is ready.

Additionally, this commit fixes the `move_token_to_root` and
`remove_port` methods. Token movement now correctly drives
port deletion; previously, port deletion drove token deletion, which
occasionally resulted in an empty token graph.

Finally, this commit corrects the status precedence in the
`_reduce_statuses` function. The `FAILED` and `CANCELLED`
statuses now take priority over `RECOVERED`.

Author: LanderOtto

streamflow/core/recovery.py

@@ -3,8 +3,6 @@
 import asyncio
 import functools
 from abc import ABC, abstractmethod
-from collections.abc import MutableMapping
-from enum import IntEnum
 from typing import TYPE_CHECKING
 
 from streamflow.core.context import SchemaEntity
@@ -86,10 +84,7 @@ async def close(self) -> None: ...
     def get_request(self, job_name: str) -> RetryRequest: ...
 
     @abstractmethod
-    async def recover(self, job: Job, step: Step, exception: BaseException) -> None: ...
-
-    @abstractmethod
-    async def is_recovered(self, job_name: str) -> TokenAvailability: ...
+    async def is_recovering(self, job_name: str) -> bool: ...
 
     @abstractmethod
     async def notify(
@@ -99,6 +94,9 @@ async def notify(
         job_token: JobToken | None = None,
     ) -> None: ...
 
+    @abstractmethod
+    async def recover(self, job: Job, step: Step, exception: BaseException) -> None: ...
+
     @abstractmethod
     async def update_request(self, job_name: str) -> None: ...
 
@@ -112,17 +110,10 @@ async def recover(self, failed_job: Job, failed_step: Step) -> None: ...
 
 
 class RetryRequest:
-    __slots__ = ("job_token", "lock", "output_tokens", "version", "workflow")
+    __slots__ = ("lock", "name", "version", "workflow")
 
-    def __init__(self) -> None:
-        self.job_token: JobToken | None = None
+    def __init__(self, name: str) -> None:
         self.lock: asyncio.Lock = asyncio.Lock()
-        self.output_tokens: MutableMapping[str, Token] = {}
+        self.name: str = name
         self.version: int = 1
         self.workflow: Workflow | None = None
-
-
-class TokenAvailability(IntEnum):
-    Unavailable = 0
-    Available = 1
-    FutureAvailable = 2

streamflow/recovery/failure_manager.py

@@ -7,12 +7,7 @@
 
 from streamflow.core.context import StreamFlowContext
 from streamflow.core.exception import FailureHandlingException
-from streamflow.core.recovery import (
-    FailureManager,
-    RetryRequest,
-    TokenAvailability,
-    recoverable,
-)
+from streamflow.core.recovery import FailureManager, RetryRequest, recoverable
 from streamflow.core.workflow import Job, Status, Step, Token
 from streamflow.log_handler import logger
 from streamflow.recovery.policy.recovery import RollbackRecoveryPolicy
@@ -36,9 +31,14 @@ async def _do_handle_failure(self, job: Job, step: Step) -> None:
         # Delay rescheduling to manage temporary failures (e.g. connection lost)
         if self.retry_delay is not None:
             await asyncio.sleep(self.retry_delay)
-        await RollbackRecoveryPolicy(self.context).recover(job, step)
-        if logger.isEnabledFor(logging.INFO):
-            logger.info(f"COMPLETED Recovery execution of failed job {job.name}")
+        try:
+            await RollbackRecoveryPolicy(self.context).recover(job, step)
+            if logger.isEnabledFor(logging.INFO):
+                logger.info(f"COMPLETED Recovery execution of failed job {job.name}")
+        except FailureHandlingException as e:
+            if logger.isEnabledFor(logging.INFO):
+                logger.info(f"FAILED Recovery execution of failed job {job.name}")
+            raise e
 
     async def close(self) -> None:
         pass
@@ -47,7 +47,7 @@ def get_request(self, job_name: str) -> RetryRequest:
         if job_name in self._retry_requests.keys():
             return self._retry_requests[job_name]
         else:
-            return self._retry_requests.setdefault(job_name, RetryRequest())
+            return self._retry_requests.setdefault(job_name, RetryRequest(job_name))
 
     @classmethod
     def get_schema(cls) -> str:
@@ -58,6 +58,13 @@ def get_schema(cls) -> str:
             .read_text("utf-8")
         )
 
+    async def is_recovering(self, job_name: str) -> bool:
+        return self.context.scheduler.get_allocation(job_name).status in (
+            Status.ROLLBACK,
+            Status.RUNNING,
+            Status.FIREABLE,
+        )
+
     async def recover(self, job: Job, step: Step, exception: BaseException) -> None:
         if logger.isEnabledFor(logging.INFO):
             logger.info(
@@ -66,46 +73,16 @@ async def recover(self, job: Job, step: Step, exception: BaseException) -> None:
         await self.context.scheduler.notify_status(job.name, Status.RECOVERY)
         await self._do_handle_failure(job, step)
 
-    async def is_recovered(self, job_name: str) -> TokenAvailability:
-        if request := self._retry_requests.get(job_name):
-            async with request.lock:
-                if self.context.scheduler.get_allocation(job_name).status in (
-                    Status.ROLLBACK,
-                    Status.RUNNING,
-                    Status.FIREABLE,
-                ):
-                    return TokenAvailability.FutureAvailable
-                elif len(request.output_tokens) > 0 and all(
-                    await asyncio.gather(
-                        *(
-                            asyncio.create_task(t.is_available(self.context))
-                            for t in request.output_tokens.values()
-                        )
-                    )
-                ):
-                    return TokenAvailability.Available
-        return TokenAvailability.Unavailable
-
     async def notify(
         self,
         output_port: str,
         output_token: Token,
         job_token: JobToken | None = None,
     ) -> None:
-        if job_token is not None:
-            job_name = job_token.value.name
-            if job_name in self._retry_requests.keys():
-                async with self._retry_requests[job_name].lock:
-                    self._retry_requests[job_name].job_token = job_token
-                    self._retry_requests[job_name].output_tokens.setdefault(
-                        output_port, output_token
-                    )
+        pass
 
     async def update_request(self, job_name: str) -> None:
         retry_request = self._retry_requests[job_name]
-        async with retry_request.lock:
-            retry_request.job_token = None
-            retry_request.output_tokens = {}
         if self.max_retries is None or retry_request.version < self.max_retries:
             retry_request.version += 1
             if logger.isEnabledFor(logging.DEBUG):
@@ -138,16 +115,16 @@ def get_schema(cls) -> str:
     def get_request(self, job_name: str) -> RetryRequest:
         pass
 
+    async def is_recovering(self, job_name: str) -> bool:
+        return False
+
     async def recover(self, job: Job, step: Step, exception: BaseException) -> None:
         if logger.isEnabledFor(logging.WARNING):
             logger.warning(
                 f"Job {job.name} failure can not be recovered. Failure manager is not enabled."
             )
         raise exception
 
-    async def is_recovered(self, job_name: str) -> TokenAvailability:
-        return TokenAvailability.Unavailable
-
     async def notify(
         self,
         output_port: str,

streamflow/recovery/policy/recovery.py

@@ -1,22 +1,18 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import logging
 from collections.abc import MutableSequence, MutableSet
 from typing import cast
 
 from streamflow.core.exception import FailureHandlingException
-from streamflow.core.recovery import RecoveryPolicy
+from streamflow.core.recovery import RecoveryPolicy, RetryRequest
 from streamflow.core.utils import get_tag
-from streamflow.core.workflow import Job, Step, Token, Workflow
+from streamflow.core.workflow import Job, Port, Step, Token, Workflow
 from streamflow.log_handler import logger
 from streamflow.persistence.loading_context import WorkflowBuilder
-from streamflow.recovery.utils import (
-    GraphMapper,
-    ProvenanceGraph,
-    TokenAvailability,
-    create_graph_mapper,
-)
+from streamflow.recovery.utils import GraphMapper, ProvenanceGraph, create_graph_mapper
 from streamflow.workflow.executor import StreamFlowExecutor
 from streamflow.workflow.port import (
     BoundaryAction,
@@ -29,6 +25,28 @@
 from streamflow.workflow.utils import get_job_token
 
 
+def _get_recovery_port(
+    token_id: int,
+    mapper: GraphMapper,
+    original_workflow: Workflow,
+    recovery_workflow: Workflow,
+) -> Port:
+    port_name = next(
+        curr_port
+        for curr_port, curr_tokens in mapper.port_tokens.items()
+        if token_id in curr_tokens
+    )
+    if port_name not in recovery_workflow.ports.keys() or not isinstance(
+        recovery_workflow.ports[port_name], InterWorkflowPort
+    ):
+        return recovery_workflow.create_port(
+            cls=type(original_workflow.ports[port_name]),
+            name=port_name,
+        )
+    else:
+        return recovery_workflow.ports[port_name]
+
+
 async def _inject_tokens(
     failed_job: Job,
     failed_step: Step,
@@ -124,62 +142,48 @@ async def _populate_workflow(
 
 
 class RollbackRecoveryPolicy(RecoveryPolicy):
-    async def _sync_workflows(
+    async def _synchronize_workflows(
         self,
-        job_names: MutableSet[str],
+        failed_job: str,
         job_tokens: MutableSequence[Token],
         mapper: GraphMapper,
+        retry_requests: MutableSequence[RetryRequest],
         workflow: Workflow,
     ) -> None:
-        for job_name in job_names:
-            retry_request = self.context.failure_manager.get_request(job_name)
-            if (
-                is_available := await self.context.failure_manager.is_recovered(
-                    job_name
-                )
-            ) == TokenAvailability.FutureAvailable:
+        for retry_request in retry_requests:
+            job_name = retry_request.name
+            if await self.context.failure_manager.is_recovering(job_name):
                 job_token = get_job_token(job_name, job_tokens)
-                # The `retry_request` is the current job running, instead
-                # the `job_token` is the token to remove in the graph because
-                # the workflow will depend on the already running job
                 if logger.isEnabledFor(logging.DEBUG):
-                    logger.debug(f"Synchronize rollbacks: job {job_name} is running")
-                # todo: create a unit test for this case
-                for port_name in await mapper.get_output_ports(job_token):
-                    if port_name in retry_request.workflow.ports.keys():
-                        cast(
-                            InterWorkflowJobPort,
-                            retry_request.workflow.ports[port_name],
-                        ).add_inter_port(
-                            workflow.create_port(
-                                cls=InterWorkflowJobPort, name=port_name
-                            ),
-                            boundary_tags=[job_token.tag],
-                            boundary_action=(
-                                BoundaryAction.PROPAGATE | BoundaryAction.TERMINATE
-                            ),
-                        )
-                # Remove tokens recovered in other workflows
-                for token_id in await mapper.get_output_tokens(job_token.persistent_id):
+                    logger.debug(
+                        f"Synchronizing rollbacks for failed job {failed_job}: Job {job_name} is currently executing."
+                    )
+                available_tokens = set()
+                for token_id in (
+                    mapper.dag_tokens.successors(job_token.persistent_id)
+                    if mapper.dag_tokens.contains(job_token.persistent_id)
+                    else []
+                ):
                     mapper.move_token_to_root(token_id)
-            elif is_available == TokenAvailability.Available:
-                job_token = get_job_token(job_name, job_tokens)
+                    available_tokens.add(token_id)
+                # Some tokens could be discarded by the `move_token_to_root` method
+                for token_id in available_tokens & mapper.token_instances.keys():
+                    new_port = _get_recovery_port(
+                        token_id, mapper, retry_request.workflow, workflow
+                    )
+                    cast(
+                        InterWorkflowPort,
+                        retry_request.workflow.ports[new_port.name],
+                    ).add_inter_port(
+                        port=new_port,
+                        boundary_tags=[job_token.tag],
+                        boundary_action=BoundaryAction.PROPAGATE,
+                    )
+            else:
                 if logger.isEnabledFor(logging.DEBUG):
                     logger.debug(
-                        f"Synchronize rollbacks: job {job_token.value.name} output available"
+                        f"Synchronizing rollbacks for failed job {failed_job}: Job {job_name} rollback"
                     )
-                # Search execute token after job token, replace this token with job_req token.
-                # Then remove all the prev tokens
-                for port_name in await mapper.get_output_ports(job_token):
-                    if port_name in retry_request.output_tokens.keys():
-                        new_token = retry_request.output_tokens[port_name]
-                        mapper.replace_token(
-                            port_name,
-                            new_token,
-                            True,
-                        )
-                        mapper.move_token_to_root(new_token.persistent_id)
-            else:
                 await self.context.failure_manager.update_request(job_name)
                 retry_request.workflow = workflow
 
@@ -211,24 +215,35 @@ async def recover(self, failed_job: Job, failed_step: Step) -> None:
         job_tokens = list(
             filter(lambda t: isinstance(t, JobToken), mapper.token_instances.values())
         )
-        await self._sync_workflows(
-            job_names={*(t.value.name for t in job_tokens), failed_job.name},
-            job_tokens=job_tokens,
-            mapper=mapper,
-            workflow=new_workflow,
-        )
-        if mapper.dag_tokens.empty():
-            raise FailureHandlingException(
-                f"Impossible to recover {failed_job.name}: empty token graph"
+        retry_requests = [
+            self.context.failure_manager.get_request(job_name)
+            for job_name in {*(t.value.name for t in job_tokens), failed_job.name}
+        ]
+        async with contextlib.AsyncExitStack() as exit_stack:
+            await asyncio.gather(
+                *(
+                    asyncio.create_task(exit_stack.enter_async_context(request.lock))
+                    for request in retry_requests
+                )
+            )
+            await self._synchronize_workflows(
+                failed_job=failed_job.name,
+                job_tokens=job_tokens,
+                mapper=mapper,
+                retry_requests=retry_requests,
+                workflow=new_workflow,
+            )
+            if mapper.dag_tokens.empty():
+                raise FailureHandlingException(
+                    f"Impossible to recover {failed_job.name}: empty token graph"
+                )
+            # Populate new workflow
+            await _populate_workflow(
+                failed_step=failed_step,
+                step_ids=await mapper.get_step_ids(failed_step.output_ports.values()),
+                workflow=new_workflow,
+                workflow_builder=workflow_builder,
             )
-        # Populate new workflow
-        step_ids = await mapper.get_step_ids(failed_step.output_ports.values())
-        await _populate_workflow(
-            failed_step=failed_step,
-            step_ids=step_ids,
-            workflow=new_workflow,
-            workflow_builder=workflow_builder,
-        )
         await _inject_tokens(
             failed_job=failed_job,
             failed_step=failed_step,