Fix race between server task setup and cancel

We used to handle this with a timeout, but we had a test to make sure
the timeout never actually had to happen. The timeout could in fact need
to happen, and the test was flaky.

I tried to get Anthropic Claude to solve this, and it noticed the race,
but was unable to come up with a design I liked for fixing it, so I did
it myself. I don't *really* like my design either, but at least it's
mine now.

This moves responsibility for marking a WES workflow as CANCELED from
the Celery task to the Celery task if it can and the ToilWorkflow
get_state() method otherwise. When somebody asks for the state of a
workflow, we ask Celery if it's actually stopped or not. If it has
stopped without error and is supposedly CANCELING, we declare it
CANCELED.

I'm removing the timeout-based way to go from CANCELING to CANCELED,
because if the task *is* still there and doing stuff, it can't really be
canceled yet.

It would still be nicer to have the responsibility in one place, but at
least this way I'm reducing and not increasing the number of weird
methods.

To test this I added a sleep that can make the cancel attempt win the
race, which involved adding an ugly argument to the fake-Celery code,
because we can't monkey-patch at class scope and expect a
Multiprocessing process to see it.

Author: adamnovak

src/toil/server/cli/wes_cwl_runner.py

@@ -419,7 +419,9 @@ def submit_run(
 
 
 def poll_run(client: WESClientWithWorkflowEngineParameters, run_id: str) -> bool:
-    """Return True if the given workflow run is in a finished state."""
+    """
+    Return True if the given workflow run is COMPLETE or never will be.
+    """
     status_result = client.get_run_status(run_id)
     state = status_result.get("state")
 

src/toil/server/utils.py

@@ -507,8 +507,6 @@ class WorkflowStateMachine:
     clients never see e.g. CANCELED -> COMPLETE or COMPLETE -> SYSTEM_ERROR, we
     can implement a real distributed state machine here.
 
-    We do handle making sure that tasks don't get stuck in CANCELING.
-
     State can be:
 
     "UNKNOWN"
@@ -569,22 +567,11 @@ def send_cancel(self) -> None:
         non-terminal state.
         """
 
-        state = self.get_current_state()
-        if state != "CANCELING" and state not in TERMINAL_STATES:
-            # If it's not obvious we shouldn't cancel, cancel.
-
-            # If we end up in CANCELING but the workflow runner task isn't around,
-            # or we signal it at the wrong time, we will stay there forever,
-            # because it's responsible for setting the state to anything else.
-            # So, we save a timestamp, and if we see a CANCELING status and an old
-            # timestamp, we move on.
-            self._store.set("cancel_time", get_iso_time())
-            # Set state after time, because having the state but no time is an error.
-            self._store.set("state", "CANCELING")
+        self._set_state("CANCELING")
 
     def send_canceled(self) -> None:
         """
-        Send a canceled message that would move to CANCELED from CANCELLING.
+        Send a canceled message that would move from CANCELING to CANCELED.
         """
         self._set_state("CANCELED")
 
@@ -621,28 +608,6 @@ def get_current_state(self) -> str:
         # Otherwise do an actual read from backing storage.
         state = self._store.get("state")
 
-        if state == "CANCELING":
-            # Make sure it hasn't been CANCELING for too long.
-            # We can get stuck in CANCELING if the workflow-running task goes
-            # away or is stopped while reporting back, because it is
-            # repsonsible for posting back that it has been successfully
-            # canceled.
-            canceled_at = self._store.get("cancel_time")
-            if canceled_at is None:
-                # If there's no timestamp but it's supposedly canceling, put it
-                # into SYSTEM_ERROR, because we didn't move to CANCELING properly.
-                state = "SYSTEM_ERROR"
-                self._store.set("state", state)
-            else:
-                # See if it has been stuck canceling for too long
-                canceled_at = datetime.fromisoformat(canceled_at)
-                canceling_seconds = (datetime.now() - canceled_at).total_seconds()
-                if canceling_seconds > MAX_CANCELING_SECONDS:
-                    # If it has, go to CANCELED instead, because the task is
-                    # nonresponsive and thus not running.
-                    state = "CANCELED"
-                    self._store.set("state", state)
-
         if state in TERMINAL_STATES:
             # We can cache this state forever
             self._store.write_cache("state", state)

src/toil/server/wes/tasks.py

@@ -19,18 +19,21 @@
 import subprocess
 import sys
 import tempfile
+import time
 import zipfile
 from typing import Any
 from urllib.parse import urldefrag
 
 from celery.exceptions import SoftTimeLimitExceeded  # type: ignore
+import celery.states  # type: ignore
 
 import toil.server.wes.amazon_wes_utils as amazon_wes_utils
 from toil.common import Toil
 from toil.jobStores.utils import generate_locator
 from toil.server.celery_app import celery
 from toil.server.utils import (
     MAX_CANCELING_SECONDS,
+    TERMINAL_STATES,
     WorkflowStateMachine,
     connect_to_workflow_state_store,
     download_file_from_internet,
@@ -44,9 +47,8 @@
 
 # How many seconds should we give a Toil workflow to gracefully shut down
 # before we kill it?
-# Ought to be long enough to let it clean up its job store, but shorter than
-# our patience for CANCELING WES workflows to time out to CANCELED.
-WAIT_FOR_DEATH_TIMEOUT = MAX_CANCELING_SECONDS - 15
+# Ought to be long enough to let it clean up its job store.
+WAIT_FOR_DEATH_TIMEOUT = 20
 
 
 class ToilWorkflowRunner:
@@ -456,9 +458,13 @@ def run_wes_task(
             logger.info(f"Fetching output files.")
             runner.write_output_files()
     except (KeyboardInterrupt, SystemExit, SoftTimeLimitExceeded):
-        # We canceled the workflow run
-        logger.info("Canceling the workflow")
-        runner.state_machine.send_canceled()
+        # We canceled the workflow run after setup.
+        # We can confirm cancellation, but only if we haven't already declared
+        # completion.
+        state = runner.get_state()
+        if state not in TERMINAL_STATES:
+            logger.info("Canceling the workflow")
+            runner.state_machine.send_canceled()
     except Exception:
         # The workflow run broke. We still count as the executor here.
         logger.exception("Running Toil produced an exception.")
@@ -474,8 +480,9 @@ def run_wes_task(
 
 def cancel_run(task_id: str) -> None:
     """
-    Send a SIGTERM signal to the process that is running task_id.
+    Send a signal to the process that is running Celery task task_id.
     """
+    # Celery uses SIGUSR1 for raising SoftTimeLimitExceeded.
     celery.control.terminate(task_id, signal="SIGUSR1")
 
 
@@ -484,6 +491,9 @@ class TaskRunner:
     Abstraction over the Celery API. Runs our run_wes task and allows canceling it.
 
     We can swap this out in the server to allow testing without Celery.
+
+    Note that this is not responsible for acting on or having events for things
+    like failed Celery tasks.
     """
 
     @staticmethod
@@ -505,12 +515,27 @@ def cancel(task_id: str) -> None:
     @staticmethod
     def is_ok(task_id: str) -> bool:
         """
-        Make sure that the task running system is working for the given task.
-        If the task system has detected an internal failure, return False.
+        Returns True if the task has not yet failed, and False otherwise.
+
+        Returns True if the task was successfully canceled.
+
+        If False, the task is also not live.
         """
-        # Nothing to do for Celery
+        # Poll Celery about the task. See <https://stackoverflow.com/a/38287835>
+        result = celery.result.AsyncResult(task_id)
+        if result.status == celery.states.FAILURE:
+            return False
         return True
 
+    @staticmethod
+    def is_live(task_id: str) -> bool:
+        """
+        Returns True if the task has not yet stopped, and False otherwise.
+        """
+        result = celery.result.AsyncResult(task_id)
+        # Celery "ready" means the result is as available as it is getting
+        return result.status not in celery.states.READY_STATES
+
 
 # If Celery can't be set up, we can just use this fake version instead.
 
@@ -520,16 +545,23 @@ class MultiprocessingTaskRunner(TaskRunner):
     Version of TaskRunner that just runs tasks with Multiprocessing.
 
     Can't use threading because there's no way to send a cancel signal or
-    exception to a Python thread, if loops in the task (i.e.
-    ToilWorkflowRunner) don't poll for it.
+    exception to a Python thread, if loops and the task (i.e.
+    ToilWorkflowRunner) doesn't poll for it.
     """
 
     _id_to_process: dict[str, multiprocessing.Process] = {}
     _id_to_log: dict[str, str] = {}
 
+    # For testing, we can delay task setup by this many seconds.
+    # This needs to be smuggled into the multiprocessing child process because
+    # it won't inherit any replacements and has its own globals/class scopes
+    setup_delay = 0
+
     @staticmethod
     def set_up_and_run_task(
-        output_path: str, args: tuple[str, str, str, dict[str, Any], list[str]]
+        output_path: str,
+        args: tuple[str, str, str, dict[str, Any], list[str]],
+        setup_delay: int
     ) -> None:
         """
         Set up logging for the process into the given file and then call
@@ -539,6 +571,8 @@ def set_up_and_run_task(
         the process crashes, the caller must clean up the log.
         """
 
+        time.sleep(setup_delay)
+
         # Multiprocessing and the server manage to hide actual task output from
         # the tests. Logging messages will appear in pytest's "live" log but
         # not in the captured log. And replacing sys.stdout and sys.stderr
@@ -604,7 +638,7 @@ def run(
         )
 
         cls._id_to_process[task_id] = multiprocessing.Process(
-            target=cls.set_up_and_run_task, args=(path, args)
+            target=cls.set_up_and_run_task, args=(path, args, cls.setup_delay)
         )
         cls._id_to_process[task_id].start()
 
@@ -622,8 +656,11 @@ def cancel(cls, task_id: str) -> None:
     @classmethod
     def is_ok(cls, task_id: str) -> bool:
         """
-        Make sure that the task running system is working for the given task.
-        If the task system has detected an internal failure, return False.
+        Return True if the task has not yet failed, and False otherwise.
+
+        Returns True if the task was successfully canceled.
+
+        If False, the task is also not live.
         """
 
         process = cls._id_to_process.get(task_id)
@@ -639,7 +676,7 @@ def is_ok(cls, task_id: str) -> bool:
             process.exitcode is not None
             and process.exitcode not in ACCEPTABLE_EXIT_CODES
         ):
-            # Something went wring in the task and it couldn't handle it.
+            # Something went wrong in the task and it couldn't handle it.
             logger.error(
                 "Process for running %s failed with code %s", task_id, process.exitcode
             )
@@ -655,3 +692,15 @@ def is_ok(cls, task_id: str) -> bool:
             return False
 
         return True
+
+    @classmethod
+    def is_live(cls, task_id: str) -> bool:
+        """
+        Returns True if the task has not yet stopped, and False otherwise.
+        """
+        process = cls._id_to_process.get(task_id)
+        if process is None:
+            # Never heard of this task, so it's probably in the process of
+            # getting made
+            return True
+        return process.exitcode is None