Hi there I m running into an issue where tasks in v2 don t s Prefect Community #ask-community

Hi there! I'm running into an issue where tasks in...

Nick Batchelder

06/01/2023, 12:42 AM

Hi there! I'm running into an issue where tasks in v2 don't seem to clean up pycuda context correctly. This issue doesn't appear in prefect v1. Here's a minimal example:

Copy code

from prefect import task, flow

@task
def initialize_pycuda():
    import pycuda.autoinit

    print("Initialized...")

@flow(log_prints=True)
def my_flow():
    f = initialize_pycuda()

if __name__ == "__main__":
    my_flow()

When I run, I get:

Copy code

00:38:51.036 | INFO    | prefect.engine - Created flow run 'smiling-hog' for flow 'my-flow'
00:38:51.039 | INFO    | Flow run 'smiling-hog' - View at <https://app.prefect.cloud/account/>...
00:38:51.600 | INFO    | Flow run 'smiling-hog' - Created task run 'initialize_pycuda-0' for task 'initialize_pycuda'
00:38:51.601 | INFO    | Flow run 'smiling-hog' - Executing 'initialize_pycuda-0' immediately...
00:38:52.864 | INFO    | Task run 'initialize_pycuda-0' - Initialized...
-------------------------------------------------------------------
PyCUDA ERROR: The context stack was not empty upon module cleanup.
-------------------------------------------------------------------
A context was still active when the context stack was being
cleaned up. At this point in our execution, CUDA may already
have been deinitialized, so there is no way we can finish
cleanly. The program will be aborted now.
Use Context.pop() to avoid this problem.
-------------------------------------------------------------------
Aborted (core dumped)

When I remove the @task decorator, this flow works perfectly.

Nick Batchelder

06/01/2023, 12:42 AM

And the equivalent flow in prefect v1 works fine:

Copy code

import prefect
from prefect import task, Flow


@task
def initialize_pycuda():
    import pycuda.autoinit

    logger = prefect.context.get("logger")
    <http://logger.info|logger.info>("Initialized...")


with Flow("Initialize PyCUDA") as flow:
    f = initialize_pycuda()

When I run

prefect run --path v1_pycuda_flow.py

, I get:

Copy code

Retrieving local flow... Done
Running flow locally...
└── 00:28:47 | INFO    | Beginning Flow run for 'Initialize PyCUDA'
└── 00:28:47 | INFO    | Task 'initialize_pycuda': Starting task run...
└── 00:28:48 | INFO    | Initialized...
└── 00:28:48 | INFO    | Task 'initialize_pycuda': Finished task run for task with final state: 'Success'
└── 00:28:48 | INFO    | Flow run SUCCESS: all reference tasks succeeded
Flow run succeeded!

Open in Slack

Previous Next