This commit improves the synchronization of multiple recovery workflows

Author: LanderOtto

streamflow/core/recovery.py

@@ -113,14 +113,22 @@ async def recover(self, failed_job: Job, failed_step: Step) -> None: ...
 
 
 class RetryRequest:
-    __slots__ = ("job_token", "lock", "output_tokens", "version", "workflow")
+    __slots__ = (
+        "job_token",
+        "lock",
+        "output_tokens",
+        "version",
+        "workflow",
+        "workflow_ready",
+    )
 
     def __init__(self) -> None:
         self.job_token: JobToken | None = None
         self.lock: asyncio.Lock = asyncio.Lock()
         self.output_tokens: MutableMapping[str, Token] = {}
         self.version: int = 1
         self.workflow: Workflow | None = None
+        self.workflow_ready: asyncio.Event = asyncio.Event()
 
 
 class TokenAvailability(IntEnum):

streamflow/core/utils.py

@@ -129,7 +129,7 @@ def create_command(
 
 
 def get_job_step_name(job_name: str) -> str:
-    return PurePosixPath(job_name).parent.name
+    return PurePosixPath(job_name).parent.as_posix()
 
 
 def get_job_tag(job_name: str) -> str:

streamflow/persistence/sqlite.py

@@ -470,7 +470,9 @@ async def get_token(self, token_id: int) -> MutableMapping[str, Any]:
                 "WHERE id =:id",
                 {"id": token_id},
             ) as cursor:
-                return _load_keys(dict(await cursor.fetchone()), keys=["value"])
+                row = _load_keys(dict(await cursor.fetchone()), keys=["value"])
+                row["recoverable"] = bool(row["recoverable"])
+                return row
 
     async def get_workflow(self, workflow_id: int) -> MutableMapping[str, Any]:
         async with self.connection as db:

streamflow/recovery/policy/recovery.py

@@ -8,8 +8,8 @@
 
 from streamflow.core.exception import FailureHandlingException
 from streamflow.core.recovery import RecoveryPolicy
-from streamflow.core.utils import get_tag
-from streamflow.core.workflow import Job, Step, Token, Workflow
+from streamflow.core.utils import get_job_tag, get_tag
+from streamflow.core.workflow import Job, Status, Step, Token, Workflow
 from streamflow.log_handler import logger
 from streamflow.persistence.loading_context import WorkflowBuilder
 from streamflow.recovery.utils import (
@@ -93,7 +93,7 @@ async def _inject_tokens(mapper: GraphMapper, new_workflow: Workflow) -> None:
         ):
             if logger.isEnabledFor(logging.DEBUG):
                 logger.debug(f"Injecting termination token on port {port.name}")
-            port.put(TerminationToken())
+            port.put(TerminationToken(Status.SKIPPED))
 
 
 async def _populate_workflow(
@@ -111,14 +111,16 @@ async def _populate_workflow(
             for step_id in step_ids
         )
     )
-    # Add failed step into new_workflow
+    # Add the failed step to the new workflow
     await workflow_builder.load_step(
         new_workflow.context,
         failed_step.persistent_id,
     )
-    # Instantiate ports capable of moving tokens across workflows
+    # Instantiate ports that can transfer tokens between workflows
     for port in new_workflow.ports.values():
-        if not isinstance(port, ConnectorPort):
+        if not isinstance(
+            port, (ConnectorPort, InterWorkflowJobPort, InterWorkflowPort)
+        ):
             new_workflow.create_port(
                 (
                     InterWorkflowJobPort
@@ -129,7 +131,7 @@ async def _populate_workflow(
             )
     for port in failed_step.get_output_ports().values():
         cast(InterWorkflowPort, new_workflow.ports[port.name]).add_inter_port(
-            port, border_tag=get_tag(failed_job.inputs.values())
+            port, boundary_tag=get_tag(failed_job.inputs.values()), terminate=False
         )
 
 
@@ -175,21 +177,23 @@ async def _recover_workflow(self, failed_job: Job, failed_step: Step) -> Workflo
             ]
         )
         mapper = await create_graph_mapper(self.context, provenance)
-        # Synchronize across multiple recovery workflows
+        # Synchronize between multiple recovery workflows
         job_tokens = list(
             filter(lambda t: isinstance(t, JobToken), mapper.token_instances.values())
         )
-        await self._sync_workflows(
-            {*(t.value.name for t in job_tokens), failed_job.name},
-            job_tokens,
-            mapper,
-            new_workflow,
+        job_names = 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,
         )
         # Populate new workflow
         steps = await mapper.get_port_and_step_ids(failed_step.output_ports.values())
         await _populate_workflow(
             steps, failed_step, new_workflow, workflow_builder, failed_job
         )
+        for job_name in job_names:
+            self.context.failure_manager.get_request(job_name).workflow_ready.set()
         await _inject_tokens(mapper, new_workflow)
         await _set_step_states(mapper, new_workflow)
         return new_workflow
@@ -200,7 +204,8 @@ async def _sync_workflows(
         job_tokens: MutableSequence[Token],
         mapper: GraphMapper,
         workflow: Workflow,
-    ) -> None:
+    ) -> MutableSequence[str]:
+        new_job_names = []
         for job_name in job_names:
             retry_request = self.context.failure_manager.get_request(job_name)
             if (
@@ -209,20 +214,35 @@ async def _sync_workflows(
                 )
             ) == TokenAvailability.FutureAvailable:
                 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
+                # `retry_request` represents the currently running job.
+                # `job_token` refers to the token that needs to be removed from the graph,
+                # as the workflow depends 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
+                    if not (is_wf_ready := retry_request.workflow_ready.is_set()):
+                        logger.debug(
+                            f"Synchronizing rollbacks: Job {job_name} is waiting for the rollback workflow to be ready."
+                        )
+                    else:
+                        logger.debug(
+                            f"Synchronizing rollbacks: Job {job_name} is currently executing."
+                        )
+                else:
+                    is_wf_ready = True
+                await retry_request.workflow_ready.wait()
+                if logger.isEnabledFor(logging.DEBUG) and not is_wf_ready:
+                    logger.debug(
+                        f"Synchronizing rollbacks: Job {job_name} has resumed after the rollback workflow is ready."
+                    )
                 for port_name in await mapper.get_output_ports(job_token):
                     if port_name in retry_request.workflow.ports.keys():
                         cast(
                             InterWorkflowPort, retry_request.workflow.ports[port_name]
                         ).add_inter_port(
-                            workflow.create_port(cls=InterWorkflowPort, name=port_name)
+                            workflow.create_port(cls=InterWorkflowPort, name=port_name),
+                            boundary_tag=get_job_tag(job_token.value.name),
+                            terminate=True,
                         )
-                # Remove tokens recovered in other workflows
+                # Remove tokens that will be recovered in other workflows
                 for token_id in await mapper.get_output_tokens(job_token.persistent_id):
                     mapper.remove_token(token_id, preserve_token=True)
             elif is_available == TokenAvailability.Available:
@@ -247,6 +267,9 @@ async def _sync_workflows(
             else:
                 await self.context.failure_manager.update_request(job_name)
                 retry_request.workflow = workflow
+                retry_request.workflow_ready.clear()
+                new_job_names.append(job_name)
+        return new_job_names
 
     async def recover(self, failed_job: Job, failed_step: Step) -> None:
         # Create recover workflow

streamflow/recovery/utils.py

@@ -68,7 +68,7 @@ def prev(self, vertex: Any) -> MutableSet[Any]:
     def remove(self, vertex: Any) -> MutableSequence[Any]:
         self.graph.pop(vertex, None)
         removed = [vertex]
-        # Delete nodes which are not connected to the leaves nodes
+        # Delete the nodes that are not connected to any leaf nodes
         dead_end_nodes = set()
         for node, values in self.graph.items():
             if vertex in values:
@@ -78,7 +78,7 @@ def remove(self, vertex: Any) -> MutableSequence[Any]:
         for node in dead_end_nodes:
             removed.extend(self.remove(node))
 
-        # Assign the root node to vertices without parent
+        # Assign the root node to the vertices that do not have a parent
         orphan_nodes = set()
         for node in self.keys():
             if node != DirectGraph.ROOT and not self.prev(node):
@@ -266,8 +266,8 @@ async def get_port_and_step_ids(
             )
             for dependency_row in dependency_rows
         }
-        # Remove steps with some missing input ports
-        # A port can have multiple input steps. It is necessary to load only the needed steps
+        # Remove steps with missing input ports
+        # A port may have multiple input steps, so it is important to load only the necessary steps.
         step_to_remove = set()
         for step_id, dependency_rows in zip(
             step_ids,
@@ -306,7 +306,7 @@ def remove_port(self, port_name: str) -> None:
         for token_id in orphan_tokens:
             self.remove_token(token_id)
 
-    def remove_token(self, token_id: int, preserve_token: bool = True):
+    def remove_token(self, token_id: int, preserve_token: bool = True) -> None:
         if logger.isEnabledFor(logging.INFO):
             logger.info(f"Remove token id {token_id}")
         if token_id == DirectGraph.ROOT:
@@ -325,13 +325,13 @@ def remove_token(self, token_id: int, preserve_token: bool = True):
                 token_leaves.add(prev_token_id)
         # Delete end-road branches
         for leaf_id in token_leaves:
-            self.remove_token(leaf_id)
+            self.remove_token(leaf_id, preserve_token=False)
         # Delete token (if needed)
         if not preserve_token:
             self.token_available.pop(token_id, None)
             self.token_instances.pop(token_id, None)
             self.dag_tokens.remove(token_id)
-        if not preserve_token:
+            # Remove ports
             empty_ports = set()
             for port_name, token_list in self.port_tokens.items():
                 if token_id in token_list:
@@ -342,10 +342,22 @@ def remove_token(self, token_id: int, preserve_token: bool = True):
                 self.remove_port(port_name)
 
     def replace_token(self, port_name: str, token: Token, is_available: bool) -> None:
-        old_token_id = self.get_equal_token(port_name, token)
-        if old_token_id is None:
-            raise FailureHandlingException("Impossible replace token")
-        if logger.isEnabledFor(logging.INFO):
+        if (old_token_id := self.get_equal_token(port_name, token)) is None:
+            raise FailureHandlingException(
+                f"Unable to find a token for replacement with {token.persistent_id}."
+            )
+        if old_token_id == token.persistent_id:
+            if self.token_available[old_token_id] != is_available:
+                raise FailureHandlingException(
+                    f"Availability mismatch for token {old_token_id}. "
+                    f"Expected: {self.token_available[old_token_id]}, Got: {is_available}."
+                )
+            elif logger.isEnabledFor(logging.INFO):
+                logger.info(
+                    f"Token {old_token_id} is already in desired state. Skipping replacement."
+                )
+            return
+        elif logger.isEnabledFor(logging.INFO):
             logger.info(f"Replacing {old_token_id} with {token.persistent_id}")
         # Replace
         self.dag_tokens.replace(old_token_id, token.persistent_id)
@@ -373,13 +385,13 @@ def add(self, src_token: Token | None, dst_token: Token | None) -> None:
             dst_token.persistent_id if dst_token is not None else dst_token,
         )
 
-    async def build_graph(self, inputs: Iterable[Token]):
+    async def build_graph(self, inputs: Iterable[Token]) -> None:
         """
-        The provenance graph represent the execution, and is always a DAG.
-        Visit the provenance graph with a breadth-first search and is done
-        backward starting from the input tokens. At the end of the search,
-        we have a tree where at root there are token which data are available
-        in some location and leaves will be the input tokens.
+        The provenance graph represents the execution and is always a DAG.
+        To traverse the provenance graph, a breadth-first search is performed
+        starting from the input tokens and moving backward. At the end of the search,
+        we obtain a tree where the root node represents the tokens whose data are available
+        in a specific location, and the leaves correspond to the input tokens.
         """
         token_frontier = deque(inputs)
         loading_context = DefaultDatabaseLoadingContext()

streamflow/workflow/port.py

@@ -5,7 +5,7 @@
 from typing import Callable
 
 from streamflow.core.deployment import Connector
-from streamflow.core.workflow import Job, Port, Token, Workflow
+from streamflow.core.workflow import Job, Port, Status, Token, Workflow
 from streamflow.log_handler import logger
 from streamflow.workflow.token import TerminationToken
 
@@ -53,16 +53,23 @@ def put(self, token: Token) -> None:
 class InterWorkflowPort(Port):
     def __init__(self, workflow: Workflow, name: str):
         super().__init__(workflow, name)
-        self.inter_ports: MutableSequence[tuple[Port, str | None]] = []
+        self.inter_ports: MutableSequence[tuple[Port, str, bool]] = []
 
-    def add_inter_port(self, port: Port, border_tag: str | None = None) -> None:
-        self.inter_ports.append((port, border_tag))
+    def add_inter_port(self, port: Port, boundary_tag: str, terminate: bool) -> None:
+        self.inter_ports.append((port, boundary_tag, terminate))
+        for token in self.token_list:
+            if boundary_tag == token.tag:
+                port.put(token)
+                if terminate:
+                    port.put(TerminationToken(Status.SKIPPED))
 
     def put(self, token: Token) -> None:
         if not isinstance(token, TerminationToken):
-            for port, border_tag in self.inter_ports:
-                if border_tag is None or border_tag == token.tag:
+            for port, boundary_tag, terminate in self.inter_ports:
+                if boundary_tag == token.tag:
                     port.put(token)
+                    if terminate:
+                        port.put(TerminationToken(Status.SKIPPED))
         super().put(token)
 
 

streamflow/workflow/step.py

@@ -1044,7 +1044,9 @@ async def run(self) -> None:
                     logger.debug(f"Step {self.name} forces gather on key {key}")
 
                 # Update size_map with the current size
-                self.size_map[key] = Token(value=len(self.token_map[key]), tag=key)
+                self.size_map[key] = Token(
+                    value=len(self.token_map[key]), tag=key, recoverable=True
+                )
                 await self.size_map[key].save(
                     self.workflow.context, size_port.persistent_id
                 )
@@ -1218,7 +1220,9 @@ async def run(self) -> None:
                             )
 
                         async for schema in self.combinator.combine(task_name, token):
-                            ins = [id for t in schema.values() for id in t["input_ids"]]
+                            ins = [
+                                id_ for t in schema.values() for id_ in t["input_ids"]
+                            ]
                             for port_name, token in schema.items():
                                 self.get_output_port(port_name).put(
                                     await self._persist_token(
@@ -1652,7 +1656,7 @@ async def _save_additional_params(
             "size_port": self.get_size_port().persistent_id
         }
 
-    async def _scatter(self, token: Token) -> Token:
+    async def _scatter(self, token: Token) -> None:
         if isinstance(token, ListToken):
             output_port = self.get_output_port()
             for i, t in enumerate(token.value):