Add cancellation loop to agent #7671

zanieb · 2022-11-28T22:03:23Z

Example

Flow:

import prefect
import time

@prefect.flow
def foo():
    print("hi!")
    time.sleep(1000)

if __name__ == "__main__":
    from prefect.deployments import Deployment

    d = Deployment.build_from_flow(foo, name="test")
    d.apply()

Trigger:

❯ prefect deployment run foo/test
❯ prefect flow-run cancel 9344f984-f5b9-46c6-936c-2f6e44d2d291
16:18:49.350 | DEBUG   | prefect.client - Connecting to API at https://api.prefect.cloud/api/accounts/f6f508a4-f8a7-49a1-bb4c-4970a9d8662e/workspaces/897bd37e-a9aa-49ea-9f5e-75c6f6b2622d/
Flow run '9344f984-f5b9-46c6-936c-2f6e44d2d291' was succcessfully scheduled for cancellation.

Agent:

16:18:51.147 | DEBUG   | prefect.agent - Checking for cancelled flow runs...
16:18:51.330 | INFO    | prefect.agent - Found 1 flow runs awaiting cancellation.
16:18:51.597 | INFO    | prefect.agent - Killing process Michaels-MacBook-Pro.local:45276 for flow run '9344f984-f5b9-46c6-936c-2f6e44d2d291'...
16:18:53.086 | ERROR   | prefect.infrastructure.process - Process 'hospitable-coucal' exited with status code: -15; This indicates that the process exited due to a SIGTERM signal. Typically, this is caused by manual cancellation.

Checklist

This pull request references any related issue by including "closes <link to issue>"
- If no issue exists and your change is not a small fix, please create an issue first.
This pull request includes tests or only affects documentation.
This pull request includes a label categorizing the change e.g. fix, feature, enhancement

src/prefect/agent.py

bunchesofdonald · 2022-11-29T16:15:19Z

src/prefect/agent.py

+            f"Killing {infrastructure.type} {flow_run.infrastructure_pid} for flow run '{flow_run.id}'..."
+        )
+        try:
+            await infrastructure.kill(flow_run.infrastructure_pid)


Should we handle the case that infrastructure_pid is null?

Yeah let's handle that and do a hasattr check for kill.

On second thought, we should maybe handle this during the proposal of a cancelled state? We can add a kill-infrastructure capability to infrastructure blocks that support killing and block cancellation when it is not supported?

That works for me. We do still have the, very small, edge case that infrastructure_pid is null because the flow_run update failed. Checking that here probably still makes sense.

I added these in 52f5565 and we mark the run as cancelled so we don't pick it up forever. Kind of an awkward situation.

peytonrunyan

Left a couple small comments, otherwise lgtm!

…cancelled_flow_runs`

# Conflicts: # src/prefect/agent.py

zanieb force-pushed the cancel/agent-loop branch from d189892 to c2a1221 Compare November 28, 2022 22:07

zanieb force-pushed the feature/flow-run-cancellation branch from e2bbe58 to 886f129 Compare November 28, 2022 22:08

zanieb force-pushed the cancel/agent-loop branch 2 times, most recently from 85000e5 to b6728e9 Compare November 28, 2022 22:36

zanieb mentioned this pull request Nov 28, 2022

Tracking: Flow run cancellation #7666

Merged

peytonrunyan reviewed Nov 29, 2022

View reviewed changes

src/prefect/agent.py Show resolved Hide resolved

peytonrunyan reviewed Nov 29, 2022

View reviewed changes

src/prefect/agent.py Outdated Show resolved Hide resolved

peytonrunyan reviewed Nov 29, 2022

View reviewed changes

src/prefect/agent.py Show resolved Hide resolved

bunchesofdonald approved these changes Nov 29, 2022

View reviewed changes

peytonrunyan self-requested a review November 29, 2022 16:24

peytonrunyan approved these changes Nov 29, 2022

View reviewed changes

zanieb marked this pull request as ready for review November 30, 2022 18:39

zanieb force-pushed the feature/flow-run-cancellation branch from 9daea9a to 11fe0a0 Compare November 30, 2022 19:33

zanieb force-pushed the cancel/agent-loop branch from 3cf384b to bf46e32 Compare November 30, 2022 19:46

zanieb force-pushed the feature/flow-run-cancellation branch from 11fe0a0 to 39e2bc7 Compare November 30, 2022 19:54

zanieb requested a review from pleek91 as a code owner November 30, 2022 19:54

zanieb force-pushed the cancel/agent-loop branch from bf46e32 to 9a128a2 Compare November 30, 2022 19:55

zanieb force-pushed the feature/flow-run-cancellation branch from 39e2bc7 to e3caa3f Compare November 30, 2022 19:58

zanieb force-pushed the cancel/agent-loop branch from 9a128a2 to 95e2bdc Compare November 30, 2022 19:58

zanieb force-pushed the feature/flow-run-cancellation branch from e3caa3f to d5e6209 Compare November 30, 2022 19:59

zanieb added 10 commits November 30, 2022 13:59

Add cancellation loop to agent

3208add

Add cancelling set to avoid duplicate cancels

e6f051e

Add default grace-seconds to process infrastructure

c8e62ad

Update process infrastructure errors and add agent handling

b5ee31c

Update process infra messaging

3ef8033

Mark flow runs as cancelled after kill completes

1d4a740

Use agent task group for cancel calls instead of blocking `check_for_…

56cb802

…cancelled_flow_runs`

Filter flow run ids in the API query to reduce data exchanged

7d2f7df

Fix set update

bac4519

Add workaround for flow runs cached by cloud

746bb0e

# Conflicts: # src/prefect/agent.py

zanieb added 5 commits November 30, 2022 13:59

Add guards against edge-cases

5227aa0

Update process tests to use new exception types

df4b58e

Move agent tests to dedicated directory

e833c14

Add tests for agent run cancellation

c3b313b

# Conflicts: # src/prefect/agent.py

Move agent tests to new file

bd98425

zanieb force-pushed the cancel/agent-loop branch from 95e2bdc to bd98425 Compare November 30, 2022 19:59

zanieb merged commit e6c4297 into feature/flow-run-cancellation Nov 30, 2022

zanieb deleted the cancel/agent-loop branch November 30, 2022 22:21

zanieb added a commit that referenced this pull request Dec 1, 2022

Add cancellation loop to agent (#7671)

fbac6de

zanieb added a commit that referenced this pull request Dec 1, 2022

Add cancellation loop to agent (#7671)

314cb89

github-actions bot pushed a commit that referenced this pull request Dec 1, 2022

Add cancellation loop to agent (#7671)

c9603a0

github-actions bot pushed a commit to ddelange/prefect that referenced this pull request Dec 15, 2022

Add cancellation loop to agent (PrefectHQ#7671)

81ec3ba

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add cancellation loop to agent #7671

Add cancellation loop to agent #7671

Uh oh!

zanieb commented Nov 28, 2022 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

bunchesofdonald Nov 29, 2022

Uh oh!

zanieb Nov 29, 2022

Uh oh!

zanieb Nov 30, 2022

Uh oh!

bunchesofdonald Nov 30, 2022

Uh oh!

zanieb Nov 30, 2022

Uh oh!

peytonrunyan left a comment

Uh oh!

Uh oh!

Add cancellation loop to agent #7671

Add cancellation loop to agent #7671

Uh oh!

Conversation

zanieb commented Nov 28, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Example

Checklist

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

bunchesofdonald Nov 29, 2022

Choose a reason for hiding this comment

Uh oh!

zanieb Nov 29, 2022

Choose a reason for hiding this comment

Uh oh!

zanieb Nov 30, 2022

Choose a reason for hiding this comment

Uh oh!

bunchesofdonald Nov 30, 2022

Choose a reason for hiding this comment

Uh oh!

zanieb Nov 30, 2022

Choose a reason for hiding this comment

Uh oh!

peytonrunyan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

zanieb commented Nov 28, 2022 •

edited

Loading