python asyncio - how to wait for a cancelled shielded task? - python-3.x

If I have a coroutine which runs a task which should not be cancelled, I will wrap that task in asyncio.shield().
It seems the behavior of cancel and shield is not what I would expect. If I have a task wrapped in shield and I cancel it, the await-ing coroutine returns from that await statement immediately rather than awaiting for the task to finish as shield would suggest. Additionally, the task that was run with shield continues to run but its future is now cancelled an not await-able.
From the docs:
except that if the coroutine containing it is cancelled, the Task running in something() is not cancelled. From the point of view of something(), the cancellation did not happen. Although its caller is still cancelled, so the “await” expression still raises a CancelledError.
These docs do not imply strongly that the caller is cancelled potentially before the callee finishes, which is the heart of my issue.
What is the proper method to shield a task from cancellation and then wait for it to complete before returning.
It would make more sense if asyncio.shield() raised the asyncio.CancelledError after the await-ed task has completed, but obviously there is some other idea going on here that I don't understand.
Here is a simple example:
import asyncio
async def count(n):
for i in range(n):
print(i)
await asyncio.sleep(1)
async def t():
try:
await asyncio.shield(count(5))
except asyncio.CancelledError:
print('This gets called at 3, not 5')
return 42
async def c(ft):
await asyncio.sleep(3)
ft.cancel()
async def m():
ft = asyncio.ensure_future(t())
ct = asyncio.ensure_future(c(ft))
r = await ft
print(r)
loop = asyncio.get_event_loop()
loop.run_until_complete(m())
# Running loop forever continues to run shielded task
# but I'd rather not do that
#loop.run_forever()

It seems the behavior of cancel and shield is not what I would expect. If I have a task wrapped in shield and I cancel it, the await-ing coroutine returns from that await statement immediately rather than awaiting for the task to finish as shield would suggest. Additionally, the task that was run with shield continues to run but its future is now cancelled an not await-able.
Conceptually shield is like a bullet-proof vest that absorbs the bullet and protects the wearer, but is itself destroyed by the impact. shield absorbs the cancellation, and reports itself as canceled, raising a CancelledError when asked for result, but allows the protected task to continue running. (Artemiy's answer explains the implementation.)
Cancellation of the future returned by shield could have been implemented differently, e.g. by completely ignoring the cancel request. The current approach ensures that the cancellation "succeeds", i.e. that the canceller can't tell that the cancellation was in fact circumvented. This is by design, and it makes the cancellation mechanism more consistent on the whole.
What is the proper method to shield a task from cancellation and then wait for it to complete before returning
By keeping two objects: the original task, and the shielded task. You pass the shielded task to whatever function it is that might end up canceling it, and you await the original one. For example:
async def coro():
print('starting')
await asyncio.sleep(2)
print('done sleep')
async def cancel_it(some_task):
await asyncio.sleep(0.5)
some_task.cancel()
print('cancellation effected')
async def main():
loop = asyncio.get_event_loop()
real_task = loop.create_task(coro())
shield = asyncio.shield(real_task)
# cancel the shield in the background while we're waiting
loop.create_task(cancel_it(shield))
await real_task
assert not real_task.cancelled()
assert shield.cancelled()
asyncio.get_event_loop().run_until_complete(main())
The code waits for the task to fully complete, despite its shield getting cancelled.

It would make more sense if asyncio.shield() raised the asyncio.CancelledError after the await-ed task has completed, but obviously there is some other idea going on here that I don't understand.
asyncio.shield
creates a dummy future, that may be cancelled
executes the wrapped coroutine as future and bind to it a callback on done to setting a result for the dummy future from the completed wrapped coroutine
returns the dummy future
You can see the implementation here
What is the proper method to shield a task from cancellation and then wait for it to complete before returning
You should shield count(5) future
async def t():
c_ft = asyncio.ensure_future(count(5))
try:
await asyncio.shield(c_ft)
except asyncio.CancelledError:
print('This gets called at 3, not 5')
await c_ft
return 42
or t() future
async def t():
await count(5)
return 42
async def m():
ft = asyncio.ensure_future(t())
shielded_ft = asyncio.shield(ft)
ct = asyncio.ensure_future(c(shielded_ft))
try:
r = await shielded_ft
except asyncio.CancelledError:
print('Shield cancelled')
r = await ft

Shielding a coroutine
we can also shield a task (but this code is about shilding a coroutine)
import asyncio
async def task1():
print("Starting task1")
await asyncio.sleep(1)
print("Ending task1")
print("SUCCESS !!")
async def task2(some_task):
print("Starting task2")
await asyncio.sleep(2)
print("Cancelling task1")
some_task.cancel()
print("Ending task2")
async def main():
# coroutines
co_task1 = task1()
# creating task from coroutines
task1_shielded = asyncio.shield(co_task1) # Create a shielded task1
task2_obj = asyncio.create_task(coro=task2(task1_shielded))
await task2_obj
await task1_shielded
asyncio.run(main())
out put:
Starting task1
Starting task2
Ending task1
SUCCESS !!
Cancelling task1
Ending task2

Related

How can I force asyncio task to run?

I'm noticing that when I spawn an asyncio task using create_task, it's first completing the rest of the logic rather than starting that task. I'm forced to add an await asyncio.sleep(0) to get the task started, which seems a bit hacky and unclean to me.
Here is some example code:
async def make_rpc_calls(...some args...)
val_1, val_2 = await asyncio.gather(rpc_call_1(...), rpc_call_2(...))
return process(val_1, val_2)
def some_very_cpu_intensive_function(...some args...):
// Does a lot of computation, can take 20 seconds to run
task_1 = asyncio.get_running_loop().create_task(make_rpc_calls(...))
intensive_result = some_very_cpu_intensive_function(...)
await task_1
process(intensive_result, task_1.result())
Anytime I run the above, it runs the some_very_cpu_intensive_function function before the kicking off the expensive RPCs. The only way I've gotten this to work is to do:
async def make_rpc_calls(...some args...)
val_1, val_2 = await asyncio.gather(rpc_call_1(...), rpc_call_2(...))
return process(val_1, val_2)
def some_very_cpu_intensive_function(...some args...):
// Does a lot of computation, can take 20 seconds to run
task_1 = asyncio.get_running_loop().create_task(make_rpc_calls(...))
await asyncio.sleep(0)
intensive_result = some_very_cpu_intensive_function(...)
await task_1
process(intensive_result, task_1.result())
This feels like a hack to me - I'm forcing the event loop to context switch, and doesn't feel like I'm using the asyncio framework correctly. Is there another way I should be approaching this?
sleep() always suspends the current task, allowing other tasks to run.
Setting the delay to 0 provides an optimized path to allow other tasks to run. This can be used by long-running functions to avoid blocking the event loop for the full duration of the function call.
Source: https://docs.python.org/3/library/asyncio-task.html

Is there a way to await on an asyncio.Task.result from sync code

I've a situation like below,
event_loop = asyncio.new_event_loop()
async def second_async():
# some async job
print("I'm here")
return 123
def sync():
return asyncio.run_coroutine_threadsafe(second_async(), loop=event_loop).result()
async def first_async():
sync()
event_loop.run_until_complete(first_async())
I call the sync function from a different thread(where the event_loop is not running), it works fine. The problem is if I run the event_loop.run_complete... line, the .result() call on the Task returned by run_coroutine_threadsafe blocks the execution of the loop, which makes sense. To avoid this, I tried changing this as follows,
event_loop = asyncio.new_event_loop()
async def second_async():
# some async job
print("I'm here")
return 123
def sync():
# if event_loop is running on current thread
res = loop.create_task(second_async()).result()
# else
res = asyncio.run_coroutine_threadsafe(second_async(), loop=event_loop).result()
# Additional processing on res
# Need to evaluate the result of task right here in sync.
return res
async def first_async():
sync()
event_loop.run_until_complete(first_async())
This works fine, but the .result() call on the Task object returned by create_task always raises an InvalidStateError. The set_result is never called on the Task object.
Basically, I want the flow to be as such
(async code) -> sync code (a non blocking call ->) async code
I know this is a bad way of doing things, but I'm integrating stuff, so I don't really have an option.
Here is a little single-threaded program that illustrates the problem.
If you un-comment the line asyncio.run(first_async1()), you see the same error as you're seeing, and for the same reason. You're trying to access the result of a task without awaiting it first.
import asyncio
event_loop = asyncio.new_event_loop()
async def second_async():
# some async job
print("I'm here")
return 123
def sync1():
return asyncio.create_task(second_async()).result()
async def first_async1():
print(sync1())
def sync2():
return asyncio.create_task(second_async())
async def first_async2():
print(await sync2())
# This prints I'm here,
# the raises invalid state error:
# asyncio.run(first_async1())
# This works, prints "I'm here" and "123"
asyncio.run(first_async2())
With that line commented out again, the second version of the program (first_async2) runs just fine. The only difference is that the ordinary function, sync2, returns an awaitable instead of a result. The await is done in the async function that called it.
I don't see why this is a bad practice. To me, it seems like there are situations where it's absolutely necessary.
Another approach is to create a second daemon thread and set up an event loop there. Coroutines can be executed in this second thread with asyncio.run_coroutine_threadsafe, which returns a concurrent.futures.Future. Its result method will block until the Future's value is set by the other thread.
#! python3.8
import asyncio
import threading
def a_second_thread(loop):
asyncio.set_event_loop(loop)
loop.run_forever()
loop2 = asyncio.new_event_loop()
threading.Thread(target=a_second_thread, args=(loop2,), daemon=True).start()
async def second_async():
# some async job
print("I'm here")
for _ in range(4):
await asyncio.sleep(0.25)
print("I'm done")
return 123
def sync1():
# Run the coroutine in the second thread -> get a concurrent.futures.Future
fut = asyncio.run_coroutine_threadsafe(second_async(), loop2)
return fut.result()
async def first_async1():
print(sync1())
def sync2():
return asyncio.create_task(second_async())
async def first_async2():
print(await sync2())
# This works, prints "I'm here", "I'm done", and "123"
asyncio.run(first_async1())
# This works, prints "I'm here", "I'm done", and "123"
asyncio.run(first_async2())
Of course this will still block the event loop in the main thread until fut.result() returns. There is no avoiding that. But the program runs.

Starvation in `asyncio` loop

I have a system where two "processes" A and B run on the same asyncio event loop.
I notice that the order of the initiation of processes matters - i.e. if I start process B first then process B runs all the time, while it seems that A is being "starved" of resources vise-a-versa.
In my experience, the only reason this might happen is due to a mutex which is not being released by B, but in the following toy example it happens without any mutexs being used:
import asyncio
async def A():
while True:
print('A')
await asyncio.sleep(2)
async def B():
while True:
print('B')
await asyncio.sleep(8)
async def main():
await B()
await A()
if __name__ == '__main__':
loop = asyncio.get_event_loop()
loop.run_until_complete(main())
Is in python the processes do not perform context-switch automatically? If not - how can I make both processes participate, each one in the time the other one is idle (i.e., sleeping)?
TLDR: Coroutines merely enable concurrency, they do not automatically trigger concurrency. Explicitly launch separate tasks, e.g. via create_task or gather, to run the coroutines concurrently.
async def main():
await asyncio.gather(B(), A())
Concurrency in asyncio is handled via Tasks – a close equivalent to Threads – which merely consist of coroutines/awaitables – like Threads consist of functions/callables. In general, a coroutine/awaitable itself does not equate to a separate task.
Using await X() means "start X and wait for it to complete". When using several such constructs in sequence:
async def main():
await B()
await A()
this means launching B first, and only launching A after B has completed: while async def and await allows for concurrency towards other tasks, B and A are run sequentially with respect to each other in a single task.
The simplest means to add concurrency is to explicitly create a task:
async def main():
# execute B in a new task
b_task = asyncio.create_task(B())
# execute A in the current task
await A()
await b_task
Note how B is offloaded to a new task, while one can still do a final await A() to re-use the current task.
Most async frameworks ship with high-level helpers for common concurrency scenarios. In this case, asyncio.gather is appropriate to launch several tasks at once:
async def main():
# execute B and A in new tasks
await asyncio.gather(B(), A())

Why are all the tasks completed in asyncio.wait() when I clearly indicate that I want only the first one completed?

I'm trying to fetch some data from OpenSubtitles using asyncio and then download a file who's information is contained in that data. I want to fetch that data and download the file at the same time using asyncio.
The problem is that I want to wait for 1 task from the list tasks to finish before commencing with the rest of the tasks in the list or the download_tasks. The reason for this is that in self._perform_query() I am writing information to a file and in self._download_and_save_file() I am reading that same information from that file. So in other words, the download_tasks need to wait for at least one task in tasks to finish before starting.
I found out I can do that with asyncio.wait(return_when=FIRST_COMPLETED) but for some reason it is not working properly:
payloads = [create_payloads(entry) for entry in retreive(table_in_database)]
tasks = [asyncio.create_task(self._perform_query(payload, proxy)) for payload in payloads]
download_tasks = [asyncio.create_task(self._download_and_save_file(url, proxy) for url in url_list]
done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
print(done)
print(len(done))
print(pending)
print(len(pending))
await asyncio.wait(download_tasks)
The output is completely different than expected. It seems that out of 3 tasks in the list tasks all 3 of them are being completed despite me passing asyncio.FIRST_COMPLETED. Why is this happening?
{<Task finished coro=<SubtitleDownloader._perform_query() done, defined at C:\Users\...\subtitles.py:71> result=None>, <Task finished coro=<SubtitleDownloader._perform_query() done, defined at C:\Users\...\subtitles.py:71> result=None>, <Task finished coro=<SubtitleDownloader._perform_query() done, defined at C:\Users\...\subtitles.py:71> result=None>}
3
set()
0
Exiting
As far as I can tell, the code in self._perform_query() shouldn't affect this problem. Here it is anyway just to make sure:
async def _perform_query(self, payload, proxy):
try:
query_result = proxy.SearchSubtitles(self.opensubs_token, [payload], {"limit": 25})
except Fault as e:
raise "A fault has occurred:\n{}".format(e)
except ProtocolError as e:
raise "A ProtocolError has occurred:\n{}".format(e)
else:
if query_result["status"] == "200 OK":
with open("dl_links.json", "w") as dl_links_json:
result = query_result["data"][0]
subtitle_name = result["SubFileName"]
download_link = result["SubDownloadLink"]
download_data = {"download link": download_link,
"file name": subtitle_name}
json.dump(download_data, dl_links_json)
else:
print("Wrong status code: {}".format(query_result["status"]))
For now, I've been testing this without running download_tasks but I have included it here for context. Maybe I am going about this problem in a completely wrong manner. If so, I would much appreciate your input!
Edit:
The problem was very simple as answered below. _perform_query wasn't an awaitable function, instead it ran synchronously. I changed that by editing the file writing part of _perform_query to be asynchronous with aiofiles:
def _perform_query(self, payload, proxy):
query_result = proxy.SearchSubtitles(self.opensubs_token, [payload], {"limit": 25})
if query_result["status"] == "200 OK":
async with aiofiles.open("dl_links.json", mode="w") as dl_links_json:
result = query_result["data"][0]
download_link = result["SubDownloadLink"]
await dl_links_json.write(download_link)
return_when=FIRST_COMPLETED doesn't guarantee that only a single task will complete. It guarantees that the wait will complete as soon as a task completes, but it is perfectly possible that other tasks complete "at the same time", which for asyncio means in the same iteration of the event loop. Consider, for example, the following code:
async def noop():
pass
async def main():
done, pending = await asyncio.wait(
[noop(), noop(), noop()], return_when=asyncio.FIRST_COMPLETED)
print(len(done), len(pending))
asyncio.run(main())
This prints 3 0, just like your code. Why?
asyncio.wait does two things: it submits the coroutines to the event loop, and it sets up callbacks to notify it when any of them is complete. However, the noop coroutine doesn't contain an await, so none of the calls to noop() suspends, each just does its thing and immediately returns. As a result, all three coroutine instances finish within the same pass of the event loop. wait is then informed that all three coroutines have finished, a fact it dutifully reports.
If you change noop to await a random sleep, e.g. change pass to await asyncio.sleep(0.1 * random.random()), you get the expected behavior. With an await the coroutines no longer complete at the same time, and wait will report the first one as soon as it detects it.
This reveals the true underlying issue with your code: _perform_query doesn't await. This indicates that you are not using an async underlying library, or that you are using it incorrectly. The call to SearchSubtitles likely simply blocks the event loop, which appears to work in trivial tests, but breaks essential asyncio features such as concurrent execution of tasks.

Sharing a dynamically started worker among several consumers

I am building a worker class that connects to an external event stream using asyncio. It is a single stream, but several consumer may enable it. The goal is to only maintain the connection while one or more consumer requires it.
My requirements are as follow:
The worker instance is created dynamically first time a consumer requires it.
When other consumers then require it, they re-use the same worker instance.
When the last consumer closes the stream, it cleans up its resources.
This sounds easy enough. However, the startup sequence is causing me issues, because it is itself asynchronous. Thus, assuming this interface:
class Stream:
async def start(self, *, timeout=DEFAULT_TIMEOUT):
pass
async def stop(self):
pass
I have the following scenarios:
Scenario 1 - exception at startup
Consumer 1 requests worker to start.
Worker startup sequence begins
Consumer 2 requests worker to start.
Worker startup sequence raises an exception.
Both consumer should see the exception as the result of their call to start().
Scenario 2 - partial asynchronous cancellation
Consumer 1 requests worker to start.
Worker startup sequence begins
Consumer 2 requests worker to start.
Consumer 1 gets cancelled.
Worker startup sequence completes.
Consumer 2 should see a successful start.
Scenario 3 - complete asynchronous cancellation
Consumer 1 requests worker to start.
Worker startup sequence begins
Consumer 2 requests worker to start.
Consumer 1 gets cancelled.
Consumer 2 gets cancelled.
Worker startup sequence must be cancelled as a result.
I struggle to cover all scenarios without getting any race condition and a spaghetti mess of either bare Future or Event objects.
Here is an attempt at writing start(). It relies on _worker() setting an asyncio.Event named self._worker_ready when it completes the startup sequence:
async def start(self, timeout=None):
assert not self.closing
if not self._task:
self._task = asyncio.ensure_future(self._worker())
# Wait until worker is ready, has failed, or timeout triggers
try:
self._waiting_start += 1
wait_ready = asyncio.ensure_future(self._worker_ready.wait())
done, pending = await asyncio.wait(
[self._task, wait_ready],
return_when=asyncio.FIRST_COMPLETED, timeout=timeout
)
except asyncio.CancelledError:
wait_ready.cancel()
if self._waiting_start == 1:
self.closing = True
self._task.cancel()
with suppress(asyncio.CancelledError):
await self._task # let worker shutdown
raise
finally:
self._waiting_start -= 1
# worker failed to start - either throwing or timeout triggering
if not self._worker_ready.is_set():
self.closing = True
self._task.cancel()
wait_ready.cancel()
try:
await self._task # let worker shutdown
except asyncio.CancelledError:
raise FeedTimeoutError('stream failed to start within %ss' % timeout)
else:
assert False, 'worker must propagate the exception'
That seems to work, but it seems too complex, and is really hard to test: the worker has many await points, leading to combinatoric explosion if I am to try all possible cancellation points and execution orders.
I need a better way. I am thus wondering:
Are my requirements reasonable?
Is there a common pattern to do this?
Does my question raise some code smell?
Your requirements sound reasonable. I would try to simplify start by replacing Event with a future (in this case a task), using it to both wait for the startup to finish and to propagate exceptions that occur during its course, if any. Something like:
class Stream:
async def start(self, *, timeout=DEFAULT_TIMEOUT):
loop = asyncio.get_event_loop()
if self._worker_startup_task is None:
self._worker_startup_task = \
loop.create_task(self._worker_startup())
self._add_user()
try:
await asyncio.shield(asyncio.wait_for(
self._worker_startup_task, timeout))
except:
self._rm_user()
raise
async def _worker_startup(self):
loop = asyncio.get_event_loop()
await asyncio.sleep(1) # ...
self._worker_task = loop.create_task(self._worker())
In this code the worker startup is separated from the worker coroutine, and is also moved to a separate task. This separate task can be awaited and removes the need for a dedicated Event, but more importantly, it allows scenarios 1 and 2 to be handled by the same code. Even if someone cancels the first consumer, the worker startup task will not be canceled - the cancellation just means that there is one less consumer waiting for it.
Thus in case of consumer cancellation, await self._worker_startup_task will work just fine for other consumers, whereas in case of an actual exception in worker startup, all other waiters will see the same exception because the task will have completed.
Scenario 3 should work automatically because we always cancel the startup that can no longer be observed by a consumer, regardless of the reason. If the consumers are gone because the startup itself has failed, then self._worker_startup_task will have completed (with an exception) and its cancellation will be a no-op. If it is because all consumers have been themselves canceled while awaiting the startup, then self._worker_startup_task.cancel() will cancel the startup sequence, as required by scenario 3.
The rest of the code would look like this (untested):
def __init__(self):
self._users = 0
self._worker_startup = None
def _add_user(self):
self._users += 1
def _rm_user(self):
self._users -= 1
if self._users:
return
self._worker_startup_task.cancel()
self._worker_startup_task = None
if self._worker_task is not None:
self._worker_task.cancel()
self._worker_task = None
async def stop(self):
self._rm_user()
async def _worker(self):
# actual worker...
while True:
await asyncio.sleep(1)
With my previous tests and integrating suggestions from #user4815162342 I came up with a re-usable solution:
st = SharedTask(test())
task1 = asyncio.ensure_future(st.wait())
task2 = asyncio.ensure_future(st.wait(timeout=15))
task3 = asyncio.ensure_future(st.wait())
This does the right thing: task2 cancels itself after 15s. Cancelling tasks has no effect on test() unless they all get cancelled. In that case, the last task to get cancelled will manually cancel test() and wait for cancellation handling to complete.
If passed a coroutine, it is only scheduled when first task starts waiting.
Lastly, awaiting the shared task after it has completed simply yields its result immediately (seems obvious, but initial version did not).
import asyncio
from contextlib import suppress
class SharedTask:
__slots__ = ('_clients', '_task')
def __init__(self, task):
if not (asyncio.isfuture(task) or asyncio.iscoroutine(task)):
raise TypeError('task must be either a Future or a coroutine object')
self._clients = 0
self._task = task
#property
def started(self):
return asyncio.isfuture(self._task)
async def wait(self, *, timeout=None):
self._task = asyncio.ensure_future(self._task)
self._clients += 1
try:
return await asyncio.wait_for(asyncio.shield(self._task), timeout=timeout)
except:
self._clients -= 1
if self._clients == 0 and not self._task.done():
self._task.cancel()
with suppress(asyncio.CancelledError):
await self._task
raise
def cancel(self):
if asyncio.iscoroutine(self._task):
self._task.close()
elif asyncio.isfuture(self._task):
self._task.cancel()
The re-raising of task exception cancellation (mentioned in comments) is intentional. It allows this pattern:
async def my_task():
try:
await do_stuff()
except asyncio.CancelledError as exc:
await flush_some_stuff() # might raise an exception
raise exc
The clients can cancel the shared task and handle an exception that might arise as a result, it will work the same whether my_task is wrapped in a SharedTask or not.

Resources