Add budget for asyncio Task

[hidva]

Feature or enhancement

Proposal:

Should we add a mechanism similar to Per-task operation budget:

Even though Tokio is not able to preempt, there is still an opportunity to nudge a task to yield back to the scheduler. As of 0.2.14, each Tokio task has an operation budget. This budget is reset when the scheduler switches to the task. Each Tokio resource (socket, timer, channel, ...) is aware of this budget. As long as the task has budget remaining, the resource operates as it did previously. Each asynchronous operation (actions that users must .await on) decrements the task's budget. Once the task is out of budget, all Tokio resources will perpetually return "not ready" until the task yields back to the scheduler. At that point, the budget is reset, and future .awaits on Tokio resources will again function normally.

The possible modifications are as follows:

diff --git a/Lib/asyncio/futures.py b/Lib/asyncio/futures.py
index 5f6fa234872..6356d9c2f6d 100644
--- a/Lib/asyncio/futures.py
+++ b/Lib/asyncio/futures.py
@@ -285,6 +285,8 @@ def __await__(self):
         if not self.done():
             self._asyncio_future_blocking = True
             yield self  # This tells Task to wait for completion.
+        if current_task_has_no_budget:
+            yield self  # We don't have a budget! Let others run.
         if not self.done():
             raise RuntimeError("await wasn't used with future")
         return self.result()  # May raise too.

Has this already been discussed elsewhere?

This is a minor feature, which does not need previous discussion elsewhere

Links to previous discussion of this feature:

No response

[asvetlov]

Could you please provide more comprehensive example?
In particular, I'm curious where asyncio should reset and decrease the budget.

A working (but maybe incomplete) example would be awesome!

[hidva]

Could you please provide more comprehensive example?

Our online LLM service currently has two Coroutine Tasks that communicate through an asyncio.Queue. Here is a code example:

q = asyncio.Queue()

async def task1():
  while True:
    req = await get_req_from_connection()
    do_something(req)
    await q.put(req)

async def task2():
  req = await q.get()
  do_something2(req)

Since there are always requests in the connection during load testing, task1 keeps executing, which blocks task2 from running. Our current workaround is using q = asyncio.Queue(maxsize=128). However, we still want to know if it's possible to implement a mechanism similar to the Tokio Task budget.

where asyncio should reset budget

When a task switches from non-runnable to the runnable state.

where asyncio should decrease budget

Currently, I only have a preliminary idea that hasn't been fully validated, as demonstrated in the following code:

diff --git a/Lib/asyncio/futures.py b/Lib/asyncio/futures.py
index 5f6fa234872..e4839e51a7d 100644
--- a/Lib/asyncio/futures.py
+++ b/Lib/asyncio/futures.py
@@ -285,6 +285,10 @@ def __await__(self):
         if not self.done():
             self._asyncio_future_blocking = True
             yield self  # This tells Task to wait for completion.
+        # Returning True from acquire means that the current task has a budget available
+        # and we have deducted from it.
+        if not current_task_budget.acquire():
+            yield self
         if not self.done():
             raise RuntimeError("await wasn't used with future")
         return self.result()  # May raise too.

[asvetlov]

I see your request.

Unfortunately, asyncio works in a different way.
It has no callback or something that is called on every await thing.
Say, if q.put(req) has enough capacity in the queue it returns without creating a future and awaiting for it.
Thus, I don't see a single place where we can put the budget decreasing.

Modifying all async functions in asyncio and third-party libraries to decrease a budget doesn't sound like a right way to do things, sorry.
Also, budget management can have a visible performance degradation.

P.S.
Your example could be modified by adding explicit task switch to avoid starvation problem:

async def task1():
  while True:
    req = await get_req_from_connection()
    do_something(req)
    await asyncio.sleep(0)  # enforce task switch
    await q.put(req)

[hidva]

Thanks

[sanjeevagasta]

hello****
bitween me ### **
to much **_
Hello my_

dear****