Add capacity limiter to ConcurrentTaskRunner #7013

jashbycu · 2022-09-29T14:37:06Z

Adds a CapacityLimiter to the ConcurrentTaskRunner to limit the number of concurrent threads. This is helpful for flows that would otherwise spawn hundreds of threads.

Example

from prefect.task_runners import ConcurrentTaskRunner

@flow(task_runner=ConcurrentTaskRunner(max_workers=10))
def example_flow():
    # Flow that runs with at most 10 threads ...

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

netlify · 2022-09-29T14:37:10Z

✅ Deploy Preview for prefect-orion ready!

Built without sensitive environment variables

Name	Link
🔨 Latest commit	`8073c6a`
🔍 Latest deploy log	https://app.netlify.com/sites/prefect-orion/deploys/6335ccd609874e00089791eb
😎 Deploy Preview	https://deploy-preview-7013--prefect-orion.netlify.app
📱 Preview on mobile	Toggle QR Code... Use your smartphone camera to open QR code link.

To edit notification comments on pull requests, go to your Netlify site settings.

zanieb · 2022-09-29T15:37:45Z

src/prefect/task_runners.py

+        if self._max_workers is not None:
+            self._limiter = CapacityLimiter(self._max_workers)


Could we set this to an asyncnullcontext when _max_workers is null so that the code in _run_and_store_result doesn't need a branch?

zanieb · 2022-09-29T15:38:28Z

src/prefect/task_runners.py

 import anyio
+from anyio import CapacityLimiter


Can you use anyio.CapacityLimiter instead of importing the name? It's nice to have the clarity of where the implementation is.

zanieb

Thanks for contributing! I've been meaning to add something like this. A couple comments on the implementation. We'll also need some sort of test before we can accept it.

jashbycu · 2022-09-29T15:49:40Z

Looking at test_task_runners, I'm not sure the best way to test this. Do we just need an additional test suite that yields (for example) ConcurrentTaskRunner(max_workers=10)?

zanieb · 2022-09-29T15:53:28Z

Looking at test_task_runners, I'm not sure the best way to test this. Do we just need an additional test suite that yields (for example) ConcurrentTaskRunner(max_workers=10)?

That'd be a good start! It seems quite challenging to "prove" that it's being limited to some number of workers. Maybe would be useful to look at the AnyIO capacity limiter tests and see how they check it.

zanieb · 2022-09-29T15:54:55Z

src/prefect/task_runners.py

        # Runtime attributes
        self._task_group: anyio.abc.TaskGroup = None
        self._result_events: Dict[UUID, anyio.abc.Event] = {}
        self._results: Dict[UUID, Any] = {}
        self._keys: Set[UUID] = set()
+        self._max_workers = max_workers
+        self._limiter = asyncnullcontext


Suggested change

self._limiter = asyncnullcontext

self._limiter = asyncnullcontext()

I'm also curious, is there a reason we need to instantiate the limiter in _start instead of __init__? Does it need the event loop?

zanieb · 2022-09-29T15:56:20Z

We might be able to set max_workers=1 then assert that the task runner behaves sequentially instead of concurrently?

jashbycu · 2022-09-29T16:28:32Z

Thanks Michael. I've pushed some test ideas, but I can't get the tests to run with pytest tests/test_task_runners.py. I get pytest: error: unrecognized arguments: --cov=src/ --cov=tests/ --no-cov-on-fail --cov-report=term. I don't see any detailed testing instructions in this repo. Can you help me?

One thing I'm not sure of is if max_workers=1 will result in truly sequential task running, or if the ConcurrentTask Runner might run nearly-simultaneously-submitted tasks in an arbitrary order. So if the tests fail due to this, I have a workaround idea. But I have to be able to run the tests to see.

zanieb · 2022-09-29T16:31:46Z

tests/test_task_runners.py

+class TestConcurrentTaskRunnerSingleThreaded(TaskRunnerStandardTestSuite):
+    @pytest.fixture
+    def task_runner(self):
+        yield ConcurrentTaskRunner(max_workers=1)


You may need to / be able to change the the concurrency_type on this instance of the task runner from TaskConcurrencyType.CONCURRENT to TaskConcurrencyType.SEQUENTIAL so it goes through the sequential standard tests

zanieb · 2022-09-29T16:33:29Z

I've pushed some test ideas, but I can't get the tests to run with pytest tests/test_task_runners.py. I get pytest: error: unrecognized arguments: --cov=src/ --cov=tests/ --no-cov-on-fail --cov-report=term. I don't see any detailed testing instructions in this repo. Can you help me?

Sounds like you're missing some test requirements. Did you install with the dev extra? https://docs.prefect.io/contributing/overview/#setting-up-a-development-environment

One thing I'm not sure of is if max_workers=1 will result in truly sequential task running, or if the ConcurrentTask Runner might run nearly-simultaneously-submitted tasks in an arbitrary order. So if the tests fail due to this, I have a workaround idea. But I have to be able to run the tests to see

Also not sure about what will happen there :) that seems like a possibility. If it's really a pain, we can explore a different test pattern.

jashbycu · 2022-09-29T16:57:18Z

Ok I see that about the tests. I adjusted it to simply override the concurrency type, which hopefully will let it pass the tests. But I'm having trouble getting a dev env working on my computer, and unfortunately I have to focus on some other high priority work for me now. Short of tinkering with the tests, is there anything else I can do to help? I've added you as a collaborator to my fork.

zanieb · 2022-09-29T18:32:26Z

I've run the tests in CI and it looks like there are some issues. We don't have the bandwidth to prioritize this right now, but if it languishes we'll eventually get an engineer to get it across the line. Futzing with the tests is definitely the biggest part of this change :)

jashbycu · 2022-09-29T21:58:40Z

Some more testing on my end shows that with more complex flows, this method clogs up and results in flows indefinitely stalling. I think it's because pending tasks count toward the thread limit. For my use case, I ended up using task tag concurrency limits, where each flow run I generate a unique tag to limit all tasks of that flow run, and then delete it when I'm done. I'm not sure what the Prefect solution should be, but maybe it needs to follow a similar pattern to that used by the task tags. In any case, I'm going to close this. Thank you!

zanieb · 2022-09-30T15:20:56Z

Thanks for investigating! It seems like we'd need to take the concurrency limit right before calling the user's function which is indeed trickier.

Add capacity limiter to ConcurrentTaskRunner

1fd0841

jashbycu requested a review from cicdw as a code owner September 29, 2022 14:37

Set CapacityLimiter context correctly

fdc0eee

zanieb reviewed Sep 29, 2022

View reviewed changes

zanieb suggested changes Sep 29, 2022

View reviewed changes

zanieb added the enhancement An improvement of an existing feature label Sep 29, 2022

Use asyncnullcontext for cleaner code

5a3e192

gabcoyne added the from:sales Submitted by a sales engineer label Sep 29, 2022

zanieb reviewed Sep 29, 2022

View reviewed changes

add tests

3a90e80

zanieb reviewed Sep 29, 2022

View reviewed changes

Run standard tasks with sequential type

8073c6a

jashbycu closed this Sep 29, 2022

zanieb mentioned this pull request Apr 18, 2023

[Feature] Add max_workers optional parameter to ConcurrentTaskRunner #9244

Closed

3 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add capacity limiter to ConcurrentTaskRunner #7013

Add capacity limiter to ConcurrentTaskRunner #7013

jashbycu commented Sep 29, 2022 •

edited by zanieb

netlify bot commented Sep 29, 2022 •

edited

zanieb Sep 29, 2022

zanieb Sep 29, 2022

zanieb left a comment

jashbycu commented Sep 29, 2022

zanieb commented Sep 29, 2022

zanieb Sep 29, 2022

zanieb commented Sep 29, 2022

jashbycu commented Sep 29, 2022

zanieb Sep 29, 2022

zanieb commented Sep 29, 2022

jashbycu commented Sep 29, 2022

zanieb commented Sep 29, 2022

jashbycu commented Sep 29, 2022

zanieb commented Sep 30, 2022

		if self._max_workers is not None:
		self._limiter = CapacityLimiter(self._max_workers)

	self._limiter = asyncnullcontext
	self._limiter = asyncnullcontext()

Add capacity limiter to ConcurrentTaskRunner #7013

Add capacity limiter to ConcurrentTaskRunner #7013

Conversation

jashbycu commented Sep 29, 2022 • edited by zanieb

Example

Checklist

netlify bot commented Sep 29, 2022 • edited

✅ Deploy Preview for prefect-orion ready!

zanieb Sep 29, 2022

Choose a reason for hiding this comment

zanieb Sep 29, 2022

Choose a reason for hiding this comment

zanieb left a comment

Choose a reason for hiding this comment

jashbycu commented Sep 29, 2022

zanieb commented Sep 29, 2022

zanieb Sep 29, 2022

Choose a reason for hiding this comment

zanieb commented Sep 29, 2022

jashbycu commented Sep 29, 2022

zanieb Sep 29, 2022

Choose a reason for hiding this comment

zanieb commented Sep 29, 2022

jashbycu commented Sep 29, 2022

zanieb commented Sep 29, 2022

jashbycu commented Sep 29, 2022

zanieb commented Sep 30, 2022

jashbycu commented Sep 29, 2022 •

edited by zanieb

netlify bot commented Sep 29, 2022 •

edited