Repository: AgnostiqHQ/covalent
Pull request #1711 (2023-06-26 15:22:15)

graceful shutdown with cancellation
<!--
⚠️ Your pull request title should be short detailed and understandable for all.
⚠️ Add a note to /CHANGELOG.md summarizing the changes.
⚠️ If your pull request fixes an open issue, please link to the issue.
⚠️ Ensure that your branch is up-to-date with the base branch.
-->
This is an update to an out-of-date PR #1258 which cancels pending or running workflows when the server is stopped.  This was not previously possible because of how we stopped the Uvicorn server using SIGKILL instead of SIGINT (gracefully).  When SIGINT is used instead, the second half of the FastAPI lifespan is able to execute before the server process terminates.

- [ ] I have added the tests to cover my changes.
- [x] I have updated the documentation and CHANGELOG accordingly.
- [x] I have read the CONTRIBUTING document.


To test this, consider the following workflow:

```
import covalent as ct
from time import sleep

@ct.electron
def task(x, y):
    sleep(10)
    return x + y

@ct.lattice
def workflow(x, y):
    for _ in range(10):
        x = task(x, y)
    return x
```

Execute this workflow, and then restart the server while it's running.  The workflow and task should be marked cancelled.

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`CHANGELOG.md`](/AgnostiqHQ/covalent/blob/master/CHANGELOG.md) with a one-to-two sentence description of the change in the UNRELEASED section. You may include a small working example for new features.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/AgnostiqHQ/covalent/pull/1711?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) Report
> Merging [#1711](https://app.codecov.io/gh/AgnostiqHQ/covalent/pull/1711?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) (2fc2304) into [develop](https://app.codecov.io/gh/AgnostiqHQ/covalent/commit/7e7bf51c5c319833e1dc0bc371a1ed86df12a6e3?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) (7e7bf51) will **increase** coverage by `2.07%`.
> The diff coverage is `100.00%`.

<details><summary>Additional details and impacted files</summary>


[![Impacted file tree graph](https://app.codecov.io/gh/AgnostiqHQ/covalent/pull/1711/graphs/tree.svg?width=650&height=150&src=pr&token=YGHCB3DE4P&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ)](https://app.codecov.io/gh/AgnostiqHQ/covalent/pull/1711?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ)

```diff
@@             Coverage Diff             @@
##           develop    #1711      +/-   ##
===========================================
+ Coverage    88.10%   90.18%   +2.07%     
===========================================
  Files          161       85      -76     
  Lines         6962     5248    -1714     
  Branches       151        0     -151     
===========================================
- Hits          6134     4733    -1401     
+ Misses         723      515     -208     
+ Partials       105        0     -105     
```

| Flag | Coverage Δ | | *Carryforward flag |
|---|---|---|---|
| Dispatcher | `87.89% <100.00%> (+0.18%)` | :arrow_up: | |
| Functional_Tests | `?` | | |
| SDK | `91.49% <ø> (ø)` | | Carriedforward from [7e7bf51](https://app.codecov.io/gh/AgnostiqHQ/covalent/commit/7e7bf51c5c319833e1dc0bc371a1ed86df12a6e3?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) |
| UI_Frontend | `?` | | |

*This pull request uses carry forward flags. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) to find out more.

</details>

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8afeafb5..764eb3fa 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,14 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [UNRELEASED]
 
-### Docs
-
-- Fix autodoc for SSH, Slurm, AWS Braket, AWS Lambda, AWS EC2, AWS Batch, Google Batch 
-
 ### Operations
 
 - Bumped versions in pre-commit config
 - Added prettier for markdown files.
+- Reduce the number of pinned version numbers in the `setup.py`, `requirements.txt`, and `requirements-client.txt`
+- Updated the `wci.yml` file with new features
 
 ### Added
 
@@ -29,23 +27,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Enhanced the quickstart with a set of commonly used features
 - Removed duplicate "stop server" warning in the First Experimemnt page
 - Fixed typo in quickstart
+- Fix autodoc for SSH, Slurm, AWS Braket, AWS Lambda, AWS EC2, AWS Batch, Google Batch 
 
 ### Changed
 
 - Removed the upper limit from `dask` and `distributed` packages' versions until we find a version which is incompatible with Covalent.
+- When the server is stopped, any workflows in a non-terminal state are first cancelled
 - Pinned sqlalchemy version with upper limit <2.0.0.
 
 ### Tests
 
 - Skipping functional tests for azure blob storage and gcp storage how to guides since they require credentials to run.
 
-### Operations
-
-- Reduce the number of pinned version numbers in the `setup.py`, `requirements.txt`, and `requirements-client.txt`
-- Updated the `wci.yml` file with new features
-
 ### Fixed
 
+- When stopping the server, we send the proper `SIGINT` signal to uvicorn instead of `SIGKILL` which allows the second part of the FastAPI `lifespan` to execute properly.
 - Fixed the outstanding incompatibities between front-end data layer and a postgres database
 - Reverted file-lock changes
 
diff --git a/covalent_dispatcher/_cli/service.py b/covalent_dispatcher/_cli/service.py
index ebf345c5..95939cec 100644
--- a/covalent_dispatcher/_cli/service.py
+++ b/covalent_dispatcher/_cli/service.py
@@ -26,6 +26,7 @@ import contextlib
 import json
 import os
 import shutil
+import signal
 import socket
 import sys
 import time
@@ -42,7 +43,6 @@ from distributed.comm import unparse_address
 from distributed.core import connect, rpc
 
 from covalent._shared_files.config import ConfigManager, get_config, set_config
-from covalent_ui.heartbeat import Heartbeat
 
 from .._db.datastore import DataStore
 from .migrate import migrate_pickled_result_object
@@ -243,10 +243,22 @@ def _terminate_child_processes(pid: int) -> None:
     Returns:
         None
     """
-    for child_proc in psutil.Process(pid).children(recursive=True):
+
+    # Uvicorn
+    leader = psutil.Process(pid).children()[0]
+
+    # Dask
+    children = psutil.Process(leader.pid).children(recursive=True)
+
+    with contextlib.suppress(psutil.NoSuchProcess):
+        leader.send_signal(signal.SIGINT)
+
+    for child_proc in children:
         with contextlib.suppress(psutil.NoSuchProcess):
             child_proc.kill()
-            child_proc.wait()
+
+    psutil.wait_procs(children)
+    leader.wait()
 
 
 def _graceful_shutdown(pidfile: str) -> None:
@@ -268,7 +280,6 @@ def _graceful_shutdown(pidfile: str) -> None:
         with contextlib.suppress(psutil.NoSuchProcess):
             proc.terminate()
             proc.wait()
-        Heartbeat.stop()
         click.echo("Covalent server has stopped.")
 
     else:
diff --git a/covalent_ui/heartbeat.py b/covalent_ui/heartbeat.py
index 6a40b93f..22d058e3 100644
--- a/covalent_ui/heartbeat.py
+++ b/covalent_ui/heartbeat.py
@@ -21,11 +21,14 @@
 import asyncio
 from contextlib import asynccontextmanager
 from datetime import datetime, timezone
+from typing import List
 
 import aiofiles
 from fastapi import FastAPI
 
 from covalent._shared_files.config import get_config
+from covalent_ui.api.v1.routes.end_points.summary_routes import get_all_dispatches
+from covalent_ui.api.v1.utils.status import Status
 
 
 class Heartbeat:
@@ -58,9 +61,46 @@ class Heartbeat:
             )
 
 
+async def cancel_all_with_status(status: Status) -> List[str]:
+    from covalent_dispatcher._core.dispatcher import cancel_dispatch
+
+    dispatch_ids = []
+    page = 0
+    count = 100
+
+    while True:
+        dispatches = get_all_dispatches(
+            count=count,
+            offset=page * count,
+            status_filter=status,
+        )
+
+        dispatch_ids += [dispatch.dispatch_id for dispatch in dispatches.items]
+
+        if dispatches.total_count == page * count + len(dispatches.items):
+            break
+
+        page += 1
+
+    for dispatch in dispatch_ids:
+        await cancel_dispatch(dispatch.dispatch_id)
+
+    return dispatch_ids
+
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     heartbeat = Heartbeat()
     asyncio.create_task(heartbeat.start())
 
     yield
+
+    for status in [
+        Status.NEW_OBJECT,
+        Status.POSTPROCESSING,
+        Status.PENDING_POSTPROCESSING,
+        Status.RUNNING,
+    ]:
+        await cancel_all_with_status(status)
+
+    Heartbeat.stop()


Test patch:
diff --git a/tests/covalent_dispatcher_tests/_cli/service_test.py b/tests/covalent_dispatcher_tests/_cli/service_test.py
index f0e452c8..8b18491a 100644
--- a/tests/covalent_dispatcher_tests/_cli/service_test.py
+++ b/tests/covalent_dispatcher_tests/_cli/service_test.py
@@ -21,12 +21,13 @@
 """Tests for Covalent command line interface (CLI) Tool."""
 
 import os
+import signal
 import subprocess
 import sys
 import tempfile
 import venv
 from unittest import mock
-from unittest.mock import MagicMock, Mock
+from unittest.mock import MagicMock, Mock, call
 
 import psutil
 import pytest
@@ -1140,10 +1141,28 @@ def test_terminate_child_processes(mocker):
     psutil_process_mock = mocker.patch(
         "covalent_dispatcher._cli.service.psutil.Process", return_value=MagicMock()
     )
+    children_mock = MagicMock()
+    psutil_process_mock.return_value.children.return_value = [children_mock]
+    wait_procs_mock = mocker.patch("covalent_dispatcher._cli.service.psutil.wait_procs")
 
     _terminate_child_processes(1)
 
-    psutil_process_mock.assert_called_with(1)
+    psutil_process_mock.assert_has_calls(
+        [
+            call(1),
+            call(children_mock.pid),
+        ]
+    )
+    psutil_process_mock.return_value.children.assert_has_calls(
+        [
+            call(),
+            call(recursive=True),
+        ]
+    )
+    children_mock.send_signal.assert_called_once_with(signal.SIGINT)
+    children_mock.kill.assert_called_once_with()
+    wait_procs_mock.assert_called_once_with([children_mock])
+    children_mock.wait.assert_called_once_with()
 
 
 def test_graceful_start_permission_exception(mocker):