Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Return early from concurrency() and rate_limit() without limit names #14789

Merged
merged 2 commits into from
Jul 30, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions src/prefect/concurrency/asyncio.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,12 @@ async def main():
await resource_heavy()
```
"""
if not names:
yield
return

names = names if isinstance(names, list) else [names]

limits = await _acquire_concurrency_slots(
names,
occupy,
Expand Down Expand Up @@ -99,7 +104,11 @@ async def rate_limit(
raising a `TimeoutError`. A timeout of `None` will wait indefinitely.
create_if_missing: Whether to create the concurrency limits if they do not exist.
"""
if not names:
return

names = names if isinstance(names, list) else [names]

limits = await _acquire_concurrency_slots(
names,
occupy,
Expand Down
8 changes: 8 additions & 0 deletions src/prefect/concurrency/sync.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,10 @@ def main():
resource_heavy()
```
"""
if not names:
yield
return

names = names if isinstance(names, list) else [names]

limits: List[MinimalConcurrencyLimitResponse] = _call_async_function_from_sync(
Expand Down Expand Up @@ -110,7 +114,11 @@ def rate_limit(
raising a `TimeoutError`. A timeout of `None` will wait indefinitely.
create_if_missing: Whether to create the concurrency limits if they do not exist.
"""
if not names:
return

names = names if isinstance(names, list) else [names]

limits = _call_async_function_from_sync(
_acquire_concurrency_slots,
names,
Expand Down
54 changes: 54 additions & 0 deletions tests/concurrency/test_concurrency_asyncio.py
Original file line number Diff line number Diff line change
Expand Up @@ -365,6 +365,33 @@ async def resource_heavy():
}


@pytest.mark.parametrize("names", [[], None])
async def test_rate_limit_without_limit_names(names):
executed = False

async def resource_heavy():
nonlocal executed
await rate_limit(names=names, occupy=1)
executed = True

assert not executed

with mock.patch(
"prefect.concurrency.asyncio._acquire_concurrency_slots",
wraps=lambda *args, **kwargs: None,
) as acquire_spy:
with mock.patch(
"prefect.concurrency.asyncio._release_concurrency_slots",
wraps=lambda *args, **kwargs: None,
) as release_spy:
await resource_heavy()

acquire_spy.assert_not_called()
release_spy.assert_not_called()

assert executed


async def test_concurrency_creates_new_limits_if_requested(
concurrency_limit: ConcurrencyLimitV2,
):
Expand Down Expand Up @@ -401,3 +428,30 @@ async def resource_heavy():
assert occupy_seconds > 0

assert executed


@pytest.mark.parametrize("names", [[], None])
async def test_concurrency_without_limit_names(names):
executed = False

async def resource_heavy():
nonlocal executed
async with concurrency(names=names, occupy=1):
executed = True

assert not executed

with mock.patch(
"prefect.concurrency.asyncio._acquire_concurrency_slots",
wraps=lambda *args, **kwargs: None,
) as acquire_spy:
with mock.patch(
"prefect.concurrency.asyncio._release_concurrency_slots",
wraps=lambda *args, **kwargs: None,
) as release_spy:
await resource_heavy()

acquire_spy.assert_not_called()
release_spy.assert_not_called()

assert executed
54 changes: 54 additions & 0 deletions tests/concurrency/test_concurrency_sync.py
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,33 @@ def my_flow():
assert executed


@pytest.mark.parametrize("names", [[], None])
def test_rate_limit_without_limit_names_sync(names):
executed = False

def resource_heavy():
nonlocal executed
rate_limit(names=names, occupy=1)
executed = True

assert not executed

with mock.patch(
"prefect.concurrency.sync._acquire_concurrency_slots",
wraps=lambda *args, **kwargs: None,
) as acquire_spy:
with mock.patch(
"prefect.concurrency.sync._release_concurrency_slots",
wraps=lambda *args, **kwargs: None,
) as release_spy:
resource_heavy()

acquire_spy.assert_not_called()
release_spy.assert_not_called()

assert executed


async def test_concurrency_can_be_used_while_event_loop_is_running(
concurrency_limit: ConcurrencyLimitV2,
):
Expand Down Expand Up @@ -350,3 +377,30 @@ def resource_heavy():
),
"prefect.resource.role": "concurrency-limit",
}


@pytest.mark.parametrize("names", [[], None])
def test_concurrency_without_limit_names_sync(names):
executed = False

def resource_heavy():
nonlocal executed
with concurrency(names=names, occupy=1):
executed = True

assert not executed

with mock.patch(
"prefect.concurrency.sync._acquire_concurrency_slots",
wraps=lambda *args, **kwargs: None,
) as acquire_spy:
with mock.patch(
"prefect.concurrency.sync._release_concurrency_slots",
wraps=lambda *args, **kwargs: None,
) as release_spy:
resource_heavy()

acquire_spy.assert_not_called()
release_spy.assert_not_called()

assert executed
Loading