Repository: DataBiosphere/toil

Issue:
Celery tasks end with an exception when canceled instead of finishing properly
I keep seeing WES server logs like this:
```
Fri, 20 May 2022 16:12:13 -0700	[toil] INFO:werkzeug:10.0.112.123 - - [20/May/2022 23:12:13] "POST /ga4gh/wes/v1/runs/run-e294bea6d7364c6da1b1d68b37574d3b/cancel HTTP/1.1" 200 -
Fri, 20 May 2022 16:12:13 -0700	[celery] [2022-05-20 23:12:13,933: INFO/MainProcess] Terminating run-e294bea6d7364c6da1b1d68b37574d3b (Signals.SIGUSR1)
Fri, 20 May 2022 16:12:13 -0700	[celery] [2022-05-20 23:12:13,934: WARNING/ForkPoolWorker-3] SoftTimeLimitExceeded()
Fri, 20 May 2022 16:12:17 -0700	[toil] INFO:werkzeug:10.0.155.56 - - [20/May/2022 23:12:17] "GET /engine/v1/status HTTP/1.1" 200 -
Fri, 20 May 2022 16:12:18 -0700	[toil] INFO:werkzeug:10.0.174.246 - - [20/May/2022 23:12:18] "GET /engine/v1/status HTTP/1.1" 200 -
Fri, 20 May 2022 16:12:18 -0700	[toil] INFO:werkzeug:10.0.112.123 - - [20/May/2022 23:12:18] "GET /engine/v1/status HTTP/1.1" 200 -
Fri, 20 May 2022 16:12:19 -0700	[celery] [2022-05-20 23:12:19,104: ERROR/ForkPoolWorker-3] Task run_wes[run-e294bea6d7364c6da1b1d68b37574d3b] raised unexpected: TimeoutExpired(['toil-cwl-runner', '--batchSystem=aws_batch', '--awsBatchQueue=arn:aws:batch:us-west-2:318423852362:job-queue/BatchTaskBatchJobQueue15-Ov0dpO450QfKNSng', '--awsBatchRegion=us-west-2', '--disableCaching', '--destBucket=s3://agc-318423852362-us-west-2/project/CWLDemo/userid/anovak4rBHiA/context/myContext/toil-execution/output/run-e294bea6d7364c6da1b1d68b37574d3b', '--clean=always', '--jobStore=aws:us-west-2:toil-ad45452c-c5f4-43a4-a0d1-e645869ad4e9', '/opt/work/workflows/run-e294bea6d7364c6da1b1d68b37574d3b/execution/nontrivial.cwl', '/opt/work/workflows/run-e294bea6d7364c6da1b1d68b37574d3b/execution/wes_inputs.json'], 5)
Fri, 20 May 2022 16:12:19 -0700	[celery] Traceback (most recent call last):
Fri, 20 May 2022 16:12:19 -0700	[celery]   File "/usr/local/lib/python3.7/site-packages/toil/server/wes/tasks.py", line 325, in run
Fri, 20 May 2022 16:12:19 -0700	[celery]     exit_code = process.wait()
Fri, 20 May 2022 16:12:19 -0700	[celery]   File "/usr/lib64/python3.7/subprocess.py", line 1019, in wait
Fri, 20 May 2022 16:12:19 -0700	[celery]     return self._wait(timeout=timeout)
Fri, 20 May 2022 16:12:19 -0700	[celery]   File "/usr/lib64/python3.7/subprocess.py", line 1653, in _wait
Fri, 20 May 2022 16:12:19 -0700	[celery]     (pid, sts) = self._try_wait(0)
Fri, 20 May 2022 16:12:19 -0700	[celery]   File "/usr/lib64/python3.7/subprocess.py", line 1611, in _try_wait
Fri, 20 May 2022 16:12:19 -0700	[celery]     (pid, sts) = os.waitpid(self.pid, wait_flags)
Fri, 20 May 2022 16:12:19 -0700	[celery]   File "/usr/local/lib/python3.7/site-packages/billiard/pool.py", line 229, in soft_timeout_sighandler
Fri, 20 May 2022 16:12:19 -0700	[celery]     raise SoftTimeLimitExceeded()
Fri, 20 May 2022 16:12:19 -0700	[celery] billiard.exceptions.SoftTimeLimitExceeded: SoftTimeLimitExceeded()
Fri, 20 May 2022 16:12:19 -0700	[celery] 
Fri, 20 May 2022 16:12:19 -0700	[celery] During handling of the above exception, another exception occurred:
Fri, 20 May 2022 16:12:19 -0700	[celery] 
Fri, 20 May 2022 16:12:19 -0700	[celery] Traceback (most recent call last):
Fri, 20 May 2022 16:12:19 -0700	[celery]   File "/usr/local/lib/python3.7/site-packages/celery/app/trace.py", line 451, in trace_task
Fri, 20 May 2022 16:12:19 -0700	[celery]     R = retval = fun(*args, **kwargs)
Fri, 20 May 2022 16:12:19 -0700	[celery]   File "/usr/local/lib/python3.7/site-packages/celery/app/trace.py", line 734, in __protected_call__
Fri, 20 May 2022 16:12:19 -0700	[celery]     return self.run(*args, **kwargs)
Fri, 20 May 2022 16:12:19 -0700	[celery]   File "/usr/local/lib/python3.7/site-packages/toil/server/wes/tasks.py", line 405, in run_wes_task
Fri, 20 May 2022 16:12:19 -0700	[celery]     raise e
Fri, 20 May 2022 16:12:19 -0700	[celery]   File "/usr/local/lib/python3.7/site-packages/toil/server/wes/tasks.py", line 391, in run_wes_task
Fri, 20 May 2022 16:12:19 -0700	[celery]     runner.run()
Fri, 20 May 2022 16:12:19 -0700	[celery]   File "/usr/local/lib/python3.7/site-packages/toil/server/wes/tasks.py", line 330, in run
Fri, 20 May 2022 16:12:19 -0700	[celery]     if process.wait(timeout=5) is None:
Fri, 20 May 2022 16:12:19 -0700	[celery]   File "/usr/lib64/python3.7/subprocess.py", line 1019, in wait
Fri, 20 May 2022 16:12:19 -0700	[celery]     return self._wait(timeout=timeout)
Fri, 20 May 2022 16:12:19 -0700	[celery]   File "/usr/lib64/python3.7/subprocess.py", line 1645, in _wait
Fri, 20 May 2022 16:12:19 -0700	[celery]     raise TimeoutExpired(self.args, timeout)
Fri, 20 May 2022 16:12:19 -0700	[celery] subprocess.TimeoutExpired: Command '['toil-cwl-runner', '--batchSystem=aws_batch', '--awsBatchQueue=arn:aws:batch:us-west-2:318423852362:job-queue/BatchTaskBatchJobQueue15-Ov0dpO450QfKNSng', '--awsBatchRegion=us-west-2', '--disableCaching', '--destBucket=s3://agc-318423852362-us-west-2/project/CWLDemo/userid/anovak4rBHiA/context/myContext/toil-execution/output/run-e294bea6d7364c6da1b1d68b37574d3b', '--clean=always', '--jobStore=aws:us-west-2:toil-ad45452c-c5f4-43a4-a0d1-e645869ad4e9', '/opt/work/workflows/run-e294bea6d7364c6da1b1d68b37574d3b/execution/nontrivial.cwl', '/opt/work/workflows/run-e294bea6d7364c6da1b1d68b37574d3b/execution/wes_inputs.json']' timed out after 5 seconds
```

If we cancel a WES workflow, and the actual workflow process doesn't stop within 5 seconds of getting a SIGTERM, the Celery task will throw an exception like this instead of actually killing the Toil workflow. This puts the workflow into EXECUTOR_ERROR instead of CANCELED.

┆Issue is synchronized with this [Jira Story](https://ucsc-cgl.atlassian.net/browse/TOIL-1175)
┆friendlyId: TOIL-1175



Discussion:

Resolving patch:
diff --git a/src/toil/lib/aws/utils.py b/src/toil/lib/aws/utils.py
--- a/src/toil/lib/aws/utils.py
+++ b/src/toil/lib/aws/utils.py
@@ -208,7 +208,7 @@ def create_s3_bucket(
 def get_bucket_region(bucket_name: str, endpoint_url: Optional[str] = None) -> str:
     """
     Get the AWS region name associated with the given S3 bucket.
-    
+
     Takes an optional S3 API URL override.
     """
     s3_client = cast(S3Client, session.client('s3', endpoint_url=endpoint_url))
@@ -216,7 +216,7 @@ def get_bucket_region(bucket_name: str, endpoint_url: Optional[str] = None) -> s
         with attempt:
             loc = s3_client.get_bucket_location(Bucket=bucket_name)
             return bucket_location_to_region(loc.get('LocationConstraint', None))
-            
+
 def region_to_bucket_location(region: str) -> str:
     return '' if region == 'us-east-1' else region
 
@@ -234,7 +234,7 @@ def get_object_for_url(url: ParseResult, existing: Optional[bool] = None) -> Obj
 
         keyName = url.path[1:]
         bucketName = url.netloc
-        
+
         # Decide if we need to override Boto's built-in URL here.
         endpoint_url: Optional[str] = None
         host = os.environ.get('TOIL_S3_HOST', None)
diff --git a/src/toil/server/utils.py b/src/toil/server/utils.py
--- a/src/toil/server/utils.py
+++ b/src/toil/server/utils.py
@@ -332,7 +332,10 @@ def __init__(self, url: str) -> None:
                 raise RuntimeError(f"{url} doesn't look like an S3 URL")
 
             self._bucket = parse.netloc
-            self._base_path = parse.path
+            # urlparse keeps the leading '/', but here we want a path in the
+            # bucket without a leading '/'. We also need to support an empty
+            # path.
+            self._base_path = parse.path[1:] if parse.path.startswith('/') else parse.path
             self._client = client('s3', region_name=get_current_aws_region())
 
             logger.debug("Connected to S3StateStore at %s", url)
@@ -465,7 +468,7 @@ def connect_to_workflow_state_store(url: str, workflow_id: str) -> WorkflowState
 
 # How long can a workflow be in CANCELING state before we conclude that the
 # workflow running task is gone and move it to CANCELED?
-MAX_CANCELING_SECONDS = 600
+MAX_CANCELING_SECONDS = 30
 
 class WorkflowStateMachine:
     """
diff --git a/src/toil/server/wes/tasks.py b/src/toil/server/wes/tasks.py
--- a/src/toil/server/wes/tasks.py
+++ b/src/toil/server/wes/tasks.py
@@ -36,6 +36,11 @@
 
 logger = logging.getLogger(__name__)
 
+# 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 = 20
 
 class ToilWorkflowRunner:
     """
@@ -327,7 +332,10 @@ def run(self) -> None:
             logger.warning(str(e))
             process.terminate()
 
-            if process.wait(timeout=5) is None:
+            try:
+                process.wait(timeout=WAIT_FOR_DEATH_TIMEOUT)
+            except subprocess.TimeoutExpired:
+                # We need to actually stop now.
                 process.kill()
 
             logger.info("Child process terminated by interruption.")
@@ -383,7 +391,7 @@ def run_wes_task(base_scratch_dir: str, state_store_url: str, workflow_id: str,
 
     :param workflow_id: ID of the workflow run.
     """
-    
+
     runner = ToilWorkflowRunner(base_scratch_dir, state_store_url, workflow_id,
                                 request=request, engine_options=engine_options)
 
@@ -418,10 +426,10 @@ def cancel_run(task_id: str) -> None:
 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.
     """
-    
+
     @staticmethod
     def run(args: Tuple[str, str, str, Dict[str, Any], List[str]], task_id: str) -> None:
         """
@@ -430,14 +438,14 @@ def run(args: Tuple[str, str, str, Dict[str, Any], List[str]], task_id: str) ->
         run_wes.apply_async(args=args,
                             task_id=task_id,
                             ignore_result=True)
-                            
+
     @staticmethod
     def cancel(task_id: str) -> None:
         """
         Cancel the task with the given ID on Celery.
         """
         cancel_run(task_id)
-        
+
 
 # If Celery can't be set up, we can just use this fake version instead.
 
@@ -446,7 +454,7 @@ class MultiprocessingTaskRunner(TaskRunner):
     """
     Version of TaskRunner that just runs tasks with Multiprocessing.
     """
-    
+
     @staticmethod
     def run(args: Tuple[str, str, str, Dict[str, Any], List[str]], task_id: str) -> None:
         """
@@ -455,7 +463,7 @@ def run(args: Tuple[str, str, str, Dict[str, Any], List[str]], task_id: str) ->
         logger.info("Starting task %s in a process", task_id)
         _id_to_process[task_id] = multiprocessing.Process(target=run_wes_task, args=args)
         _id_to_process[task_id].start()
-        
+
     @staticmethod
     def cancel(task_id: str) -> None:
         """


Test patch:
diff --git a/src/toil/test/server/serverTest.py b/src/toil/test/server/serverTest.py
--- a/src/toil/test/server/serverTest.py
+++ b/src/toil/test/server/serverTest.py
@@ -22,6 +22,7 @@
 from abc import abstractmethod
 from io import BytesIO
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+from urllib.parse import urlparse
 
 try:
     from flask import Flask
@@ -175,12 +176,11 @@ def get_state_store(self) -> AbstractStateStore:
         return FileStateStore(self.state_store_dir)
 
 @needs_aws_s3
-class AWSStateStoreTest(hidden.AbstractStateStoreTest):
+class BucketUsingTest(ToilTest):
     """
-    Test AWS-based state storage.
+    Base class for tests that need a bucket.
     """
 
-    from toil.server.utils import AbstractStateStore
     from mypy_boto3_s3 import S3ServiceResource
     from mypy_boto3_s3.service_resource import Bucket
 
@@ -213,6 +213,15 @@ def tearDownClass(cls) -> None:
             delete_s3_bucket(cls.s3_resource, cls.bucket_name, cls.region)
         super().tearDownClass()
 
+class AWSStateStoreTest(hidden.AbstractStateStoreTest, BucketUsingTest):
+    """
+    Test AWS-based state storage.
+    """
+
+    from toil.server.utils import AbstractStateStore
+
+    bucket_path = "prefix/of/keys"
+
     def get_state_store(self) -> AbstractStateStore:
         """
         Make a state store to test, on a single fixed URL.
@@ -220,7 +229,31 @@ def get_state_store(self) -> AbstractStateStore:
 
         from toil.server.utils import S3StateStore
 
-        return S3StateStore('s3://' + self.bucket_name)
+        return S3StateStore('s3://' + self.bucket_name + '/' + self.bucket_path)
+
+    def test_state_store_paths(self) -> None:
+        """
+        Make sure that the S3 state store puts things in the right places.
+
+        We don't *really* care about the exact internal structure, but we do
+        care about actually being under the path we are supposed to use.
+        """
+
+        from toil.lib.aws.utils import get_object_for_url
+
+        store = self.get_state_store()
+
+        # Should hold a value
+        store.set('testid', 'testkey', 'testvalue')
+        self.assertEqual(store.get('testid', 'testkey'), 'testvalue')
+
+        expected_url = urlparse('s3://' + self.bucket_name + '/' +
+            os.path.join(self.bucket_path, 'testid', 'testkey'))
+
+        obj = get_object_for_url(expected_url, True)
+        self.assertEqual(obj.content_length, len('testvalue'))
+
+
 
 
 @needs_server
@@ -269,6 +302,21 @@ def setUp(self) -> None:
                 type: stdout
             """)
 
+        self.slow_cwl = textwrap.dedent("""
+            cwlVersion: v1.0
+            class: CommandLineTool
+            baseCommand: sleep
+            stdout: output.txt
+            inputs:
+              delay:
+                type: string
+                inputBinding:
+                  position: 1
+            outputs:
+              output:
+                type: stdout
+            """)
+
     def tearDown(self) -> None:
         super().tearDown()
 
@@ -307,33 +355,73 @@ def _report_absolute_url(self, client: "FlaskClient", url: str):
         self.assertEqual(rv.status_code, 200)
         logger.info("Got %s:\n%s", url, rv.data.decode('utf-8'))
 
-    def _wait_for_success(self, client: "FlaskClient", run_id: str) -> None:
+    def _start_slow_workflow(self, client: "FlaskClient") -> str:
         """
-        Wait for the given workflow run to succeed. If it fails, raise an exception.
+        Start a slow workflow and return its ID.
+        """
+        rv = client.post("/ga4gh/wes/v1/runs", data={
+            "workflow_url": "slow.cwl",
+            "workflow_type": "CWL",
+            "workflow_type_version": "v1.0",
+            "workflow_params": json.dumps({"delay": "5"}),
+            "workflow_attachment": [
+                (BytesIO(self.slow_cwl.encode()), "slow.cwl"),
+            ],
+        })
+        # workflow is submitted successfully
+        self.assertEqual(rv.status_code, 200)
+        self.assertTrue(rv.is_json)
+        run_id = rv.json.get("run_id")
+        self.assertIsNotNone(run_id)
+
+        return run_id
+
+    def _poll_status(self, client: "FlaskClient", run_id: str) -> str:
+        """
+        Get the status of the given workflow.
+        """
+
+        rv = client.get(f"/ga4gh/wes/v1/runs/{run_id}/status")
+        self.assertEqual(rv.status_code, 200)
+        self.assertTrue(rv.is_json)
+        self.assertIn("run_id", rv.json)
+        self.assertEqual(rv.json.get("run_id"), run_id)
+        self.assertIn("state", rv.json)
+        state = rv.json.get("state")
+        self.assertIn(state, ["UNKNOWN", "QUEUED", "INITIALIZING", "RUNNING",
+                              "PAUSED", "COMPLETE", "EXECUTOR_ERROR", "SYSTEM_ERROR",
+                              "CANCELED", "CANCELING"])
+        return state
+
+    def _cancel_workflow(self, client: "FlaskClient", run_id: str) -> None:
+        rv = client.post(f"/ga4gh/wes/v1/runs/{run_id}/cancel")
+        self.assertEqual(rv.status_code, 200)
+
+    def _wait_for_status(self, client: "FlaskClient", run_id: str, target_status: str) -> None:
+        """
+        Wait for the given workflow run to reach the given state. If it reaches
+        a different terminal state, raise an exception.
         """
 
         while True:
-            rv = client.get(f"/ga4gh/wes/v1/runs/{run_id}/status")
-            self.assertEqual(rv.status_code, 200)
-            self.assertTrue(rv.is_json)
-            self.assertIn("run_id", rv.json)
-            self.assertEqual(rv.json.get("run_id"), run_id)
-            self.assertIn("state", rv.json)
-            state = rv.json.get("state")
-            self.assertIn(state, ["UNKNOWN", "QUEUED", "INITIALIZING", "RUNNING",
-                                  "PAUSED", "COMPLETE", "EXECUTOR_ERROR", "SYSTEM_ERROR",
-                                  "CANCELED", "CANCELING"])
-            if state == "COMPLETE":
+            state = self._poll_status(client, run_id)
+            if state == target_status:
                 # We're done!
                 logger.info("Workflow reached state %s", state)
                 return
-            if state in ["EXECUTOR_ERROR", "SYSTEM_ERROR", "CANCELED"]:
-                logger.error("Workflow run failed")
+            if state in ["EXECUTOR_ERROR", "SYSTEM_ERROR", "CANCELED", "COMPLETE"]:
+                logger.error("Workflow run reached unexpected terminal state %s", state)
                 self._report_log(client, run_id)
-                raise RuntimeError("Workflow is in fail state " + state)
+                raise RuntimeError("Workflow is in unexpected state " + state)
             logger.info("Waiting on workflow in state %s", state)
             time.sleep(2)
 
+    def _wait_for_success(self, client: "FlaskClient", run_id: str) -> None:
+        """
+        Wait for the given workflow run to succeed. If it fails, raise an exception.
+        """
+        self._wait_for_status(client, run_id, "COMPLETE")
+
 
 class ToilWESServerBenchTest(AbstractToilWESServerTest):
     """
@@ -499,8 +587,36 @@ def test_run_workflow_manifest_and_inputs_zip(self) -> None:
             zip_file.writestr('MANIFEST.json', json.dumps({"mainWorkflowURL": "actual.cwl", "inputFileURLs": ["data.json"]}))
         self.run_zip_workflow(zip_path, include_message=False)
 
-    # TODO: When we can check the output value, add tests for overriding
-    # packaged inputs.
+    def test_run_and_cancel_workflows(self) -> None:
+        """
+        Run two workflows, cancel one of them, and make sure they all exist.
+        """
+        with self.app.test_client() as client:
+            # Start 1 workflow
+            run1 = self._start_slow_workflow(client)
+            time.sleep(2)
+            status1 = self._poll_status(client, run1)
+            self.assertIn(status1, ["QUEUED", "INITIALIZING", "RUNNING"])
+
+            # Start another workflow
+            run2 = self._start_slow_workflow(client)
+            self.assertNotEqual(run1, run2)
+            time.sleep(2)
+            status2 = self._poll_status(client, run2)
+            self.assertIn(status2, ["QUEUED", "INITIALIZING", "RUNNING"])
+
+            # Cancel the second one
+            self._cancel_workflow(client, run2)
+            time.sleep(1)
+            status2 = self._poll_status(client, run2)
+            self.assertIn(status2, ["CANCELING", "CANCELED"])
+
+            # Make sure the first one still exists too
+            status1 = self._poll_status(client, run1)
+            self.assertIn(status1, ["QUEUED", "INITIALIZING", "RUNNING"])
+
+            self._wait_for_status(client, run2, "CANCELED")
+            self._wait_for_success(client, run1)
 
 @needs_celery_broker
 class ToilWESServerCeleryWorkflowTest(ToilWESServerWorkflowTest):
@@ -515,7 +631,16 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._server_args = []
 
+@needs_celery_broker
+class ToilWESServerCeleryS3StateWorkflowTest(ToilWESServerWorkflowTest, BucketUsingTest):
+    """
+    Test the server with Celery and state stored in S3.
+    """
 
+    def setUp(self) -> None:
+        # Overwrite server args from __init__. The bucket name isn't available when __init__ runs.
+        self._server_args = ["--state_store", "s3://" + self.bucket_name + "/state"]
+        super().setUp()
 
 if __name__ == "__main__":
     unittest.main()