Task in tags context doesn't reuse thread

[nkanazawa1989]

Bug summary

I'm writing a scientific workflow which includes optimization loops. Multiple optimization loops and cost function evaluations run in parallel with the ThreadPoolTaskRunner.

After certain loop number, I encounter

Task run failed with exception: OSError(24, 'Too many open files') - Retries are exhausted

which indicates a soft limit of file descriptor open per process.

The simplified structure of my program looks like:

from prefect import flow, task, tags

@task
def inner_task():
    # do something
    pass

@task
def outer_task():
    for i in range(5):
        with tags(f"loop {i}"):
            futures = []
            for j in range(4):
                with tags(f"instance {j}"):
                    fut = inner_task.submit()
                futures.append(fut)
            for fut in futures:
                fut.result()

@flow
def main():
    fut1 = outer_task.submit()
    ...
    fut1.wait()

if __name__ == "__main__":
    main()

Here I assume the inner_task is a time consuming IO bound subroutine of the cost function. For better traceability and visibility of my cost function execution I want to tag them recursively -- this is why tags context is helpful here. As a side effect, each cost function spawns a new thread and the thread is pooled by the thread pool executor even after execution.

This accumulates the file descriptor associated with threads during the loop, and at certain point the workflow crashes due to their soft limit. Of course I can change the soft limit, but this is just putting off the problem.

To confirm above description of the problem, you can try the following complete script.

import os
import psutil

from prefect import flow, task, tags
from prefect.logging import get_run_logger

@task
def inner_task():
    # do something
    pass

@task
def outer_task():
    logger = get_run_logger()
    
    for i in range(5):
        with tags(f"loop {i}"):  # loop tag
            num_fd = psutil.Process(os.getpid()).num_fds()
            num_th = psutil.Process(os.getpid()).num_threads()
            logger.info(f"Open FDs at loop {i} begin: {num_fd}")
            logger.info(f"Threads at loop {i} begin: {num_th}")
        
            futures = []
            for j in range(4):
                with tags(f"instance {j}"):  # instance tag
                    fut = inner_task.submit()
                futures.append(fut)
                
            for fut in futures:
                fut.result()
                
            num_fd_after = psutil.Process(os.getpid()).num_fds()
            num_th_after = psutil.Process(os.getpid()).num_threads()
            logger.info(f"Open FDs at loop {i} end: {num_fd_after} (+ {num_fd_after - num_fd})")
            logger.info(f"Threads at loop {i} end: {num_th_after} (+ {num_th_after - num_th})")

@flow
def main():
    fut = outer_task.submit()
    fut.wait()

if __name__ == "__main__":
    main()

Version info

Version:             3.4.2
API version:         0.8.4
Python version:      3.11.11
Git commit:          c3c1c119
Built:               Mon, May 19, 2025 04:04 PM
OS/Arch:             darwin/x86_64
Profile:             local
Server type:         server
Pydantic version:    2.11.4

Additional context

I did a small experiment with different settings.

As-is (every function call spawns new thread / 4 instances per loop):

...
Threads at loop 4 end: 32 (+ 4)
Open FDs at loop 4 end: 33 (+ 5)
...

I remove instance tag (some threads are reused -- behavior is a bit stochastic):

...
Threads at loop 4 end: 17 (+ 1)
Open FDs at loop 4 end: 18 (+ 2)
...

I remove loop tag (no thread is accumulated, 8 threads are added at loop 0):

...
Threads at loop 4 end: 16 (+ 0)
Open FDs at loop 4 end: 17 (+ 1)
...

I remove both tags (no thread is accumulated, 4 threads are added at loop 0):

...
Threads at loop 4 end: 12 (+ 0)
Open FDs at loop 4 end: 12 (+ 0)
...

This result indicates tags context doesn't work well with the thread pool executor.

[nkanazawa1989]

Seems like this is a problem of contextvar (i.e. context.run in multithreaded environment).

prefect/src/prefect/task_runners.py

Lines 326 to 330 in 7d4d416

	future = self._executor.submit(
	context.run,
	run_task_sync,
	**submit_kwargs,
	)

If I understand correctly context is thread local and this is why always new thread is created? -- so this is expected behavior, not a bug. It would be great if this file descriptor accumulation with Prefect context is documented as a note in ThreadPoolTaskRunner API docs.

[nkanazawa1989]

@zzstoatzz Thanks for the quick response, but I found another edge case. Can you reopen this or give me your thoughts?

import asyncio
import psutil

from prefect import flow, task
from prefect.logging import get_run_logger

@task
async def my_task():
    await asyncio.sleep(1)

@task
async def outer():
    logger = get_run_logger()

    for i in range(3):
        num_fd = psutil.Process().num_fds()
        num_th = psutil.Process().num_threads()
        logger.info(f"Start of the loop {i}: FD={num_fd}, Thread={num_th}")

        coros = []
        for j in range(3):
            tagged = my_task.with_options(tags=[f"loop{i}", f"instance{j}"])
            coro = tagged()
            coros.append(coro)
        await asyncio.gather(*coros)
                    
        num_fd_after = psutil.Process().num_fds()
        num_th_after = psutil.Process().num_threads()
        logger.info(f"End of the loop {i}: FD={num_fd_after} (+{num_fd_after-num_fd}), Thread={num_th_after} (+{num_th_after-num_th})")


@flow
def main():
    outer.submit().wait()

if __name__ == "__main__":
    main()

In this example I don't use the task runner. Instead, I run multiple async tasks with asyncio.gather. Unlike the multithreading, this workflow is expected to run on a single thread with context switching. However, the log indicates each loop accumulates 3 threads with file descriptors -- this is totally unexpected and the behavior doesn't look async. Indeed this behavior is a bit surprising.

I scrutinized the problem and found how this happened:

The async task is called inside AsyncTaskRunEngine context:

prefect/src/prefect/task_engine.py

Line 1444 in 857102f

async with engine.start(task_run_id=task_run_id, dependencies=dependencies):

This checks concurrency slot with tags:

prefect/src/prefect/task_engine.py

Lines 774 to 776 in 857102f

	# Acquire a concurrency slot for each tag, but only if a limit
	# matching the tag already exists.
	with concurrency(list(self.task_run.tags), self.task_run.id):

This slot check invokes ConcurrencySlotAcquisitionService with tags:

prefect/src/prefect/concurrency/_asyncio.py

Line 29 in 857102f

service = ConcurrencySlotAcquisitionService.instance(frozenset(names))

Usually the service instance is cached by the QueueService, but now tags are considered in the cache key evaluation:

prefect/src/prefect/_internal/concurrency/services.py

Lines 260 to 263 in 857102f

	with cls._instance_lock:
	key = hash((cls, *args))
	if key not in cls._instances:
	cls._instances[key] = cls._new_instance(*args)

This results in creation of new service per tag, and new thread is created here with a file descriptor.

prefect/src/prefect/_internal/concurrency/services.py

Lines 281 to 283 in 857102f

	# Otherwise, bind the service to the global loop
	else:
	from_sync.call_soon_in_loop_thread(create_call(instance.start)).result()

This means Prefect always creates unnecessary threads for concurrency check when tags are assigned dynamically. The task runner issue above might be also due to this concurrency slot check -- not context copy but I didn't confirm.

It would be nice if Prefect can support extra tags that are not considered in the concurrency control, i.e. pure metadata. The only workaround I can come up with is to embed such information in the task_run_name, but this looks a bit awkward.