Redundant logic after using aiologic

[ultmaster]

agent-lightning/agentlightning/store/memory.py

Line 188 in 9f8a25f

result = await asyncio.to_thread(evt.wait, timeout)

Can be replaced with (however, note for the future: it is better to rely on asyncio.timeout here):

try:
    result = await asyncio.wait_for(evt, timeout)
except asyncio.TimeoutError:
    result = False

And the previous branch can be deleted accordingly, since there will be no more threads here.

Originally posted by @x42005e1f in #356 (comment)

[x42005e1f]

And _ThreadSafeAsyncLock can also be removed and completely replaced with aiologic.Lock, since the latter has the same methods as the former.

Originally posted by @x42005e1f in #356 (comment)

[x42005e1f]

agent-lightning/agentlightning/utils/server_launcher.py

Lines 851 to 852 in 9f8a25f

	self._thread_event_queue = queue.Queue()
	self._thread_stop_event = threading.Event()

The queue can be replaced with aiologic.SimpleQueue, and the event can be replaced with aiologic.Event (or, for a more lightweight option, aiologic.lowlevel.create_async_event(), since there is only one waiter). This will eliminate both polling and unnecessary threads (the corresponding changes will need to be made). But note that all these types override __bool__(), so constructs like if self._thread_stop_event will have to be replaced with if self._thread_stop_event is not None.

[x42005e1f]

agent-lightning/agentlightning/execution/shared_memory.py

Lines 91 to 99 in 9f8a25f

	async def watcher() -> None:
	# Poll the threading event without blocking the event loop. Using a
	# background thread via ``asyncio.to_thread`` makes cancellation
	# difficult because ``ThreadingEvent.wait`` is not interruptible.
	# Instead we cooperatively check the flag from the loop so the
	# watcher task stays cancellable and tests don't hang when the
	# bundle finishes naturally before the stop event is set.
	while not stop_evt.is_set():
	await asyncio.sleep(self.poll_interval)

It seems that this could be optimized using aiologic.REvent (it basically corresponds to the ExecutionEvent protocol, but it is more heavyweight than aiologic.Event):

async def watcher() -> None: 
    await stop_evt  # for `aiologic.REvent`
    ...

[ultmaster]

What I'm interested in is a universal interface for asyncio/thread/multiprocessing queue/event/locks that all support asyncio. A part of utils in this repo is actually doing that. Do you have any suggestions for this?

[x42005e1f]

Support for multiprocessing will be somewhat non-trivial, since it is... quite problematic. For everything else, there is aiologic (it supports much, much more than you listed). In any case, you will have to work with processes in a special way, especially because of the different nature of memory, so the best thing here is to try to isolate this specificity somehow. For example, GuideLLM does something similar (see src/guidellm/utils/messaging.py), relying on Culsans queues - my second library, derived from aiologic.

As for the method of waiting for processes (and the corresponding primitives), in one comment I provided a useful helper function that can do this most efficiently via long polling. You can use it as a basis for implementing asynchronous interfaces.

[x42005e1f]

However, as I see it, the easiest way to work with processes is to design the architecture so that you do not have to work with them. Because they are quite painful: we have different start methods, copying everything during communication, binding primitives to these methods, signals, zombies, daemons, pentagrams, exorcists, and so on. I have seen that sometimes this is solved via a thread-supervisor, but any way that only requires you to know about threads will do.

[ultmaster]

Thanks for the comment. How about forking processes? Will carry aiologic objects into a forked process cause issues? Do I need to register atexit callbacks to close them when the processes exit? I'm not expecting one aiologic object to be reusable or sharable across multiple processes; I'm just wondering about the friendliness.

Also I suppose aiologic lock is somewhat less efficient that asyncio.Lock? I think it would be nice if I can replace that as well because asyncio.Lock is binded to some event loop and unfriendly to pytest asyncio, which creates a new event loop for every test (?).

[x42005e1f]

Will carry aiologic objects into a forked process cause issues?

If you fork when primitives are already in use - yes, as in the case of asyncio/threading/queue/... primitives, because the internal wait queue will remain non-empty. You will have to recreate them. Otherwise, there should be no issues, and they can also be safely serialized and thus transferred between processes (returning to their original state, just like your _LoopAwareAsyncLock, but with the exception of queues and aiologic.Flag - they save items and markers, respectively, but also clear the wait queues).

Do I need to register atexit callbacks to close them when the processes exit?

No, because aiologic does not create tasks or threads anywhere (except in its own tests, via aiologic._testing). It relies on existing internal mechanisms of concurrency libraries (such as loop.call_soon_threadsafe()), and thus has no state that needs to be cleaned up in any way.

Also I suppose aiologic lock is somewhat less efficient that asyncio.Lock?

Yes, that is right. asyncio lives in the cooperative multitasking world, which is why it has a very simple implementation of its primitives. However, aiologic primitives are not much slower: as far as I remember, about twice as slow for pure async code (which is actually a relatively small performance degradation). Starting with 0.15.0, memory usage has been significantly improved, which also makes memory usage by unused primitives much lower.

some old memory usage measurements

from pympler.asizeof import asizeof, flatsize


def sizeof(factory, /, *args, limit=2**20):
    number = 1
    
    while True:
        array = [factory(*args) for _ in range(number)]
        
        total_bytes = asizeof(array)
        
        if total_bytes >= limit:
            break
        
        delta = int((limit - total_bytes) / (total_bytes / number))
        
        if not delta:
            break
        
        number += delta
    
    return round((total_bytes - flatsize(array)) / number)

>>> sizeof(collections.deque)
760
>>> sizeof(queue.SimpleQueue)
72  # see https://github.com/python/cpython/issues/140025
>>> sizeof(queue.Queue)
3730
>>> sizeof(asyncio.Queue)
3346
>>> sizeof(janus.Queue)
7765
>>> sizeof(culsans.Queue)
2152
>>> sizeof(aiologic.Queue)
680
>>> sizeof(aiologic.SimpleQueue)
448
>>> sizeof(aiologic.SimpleLifoQueue)
376
>>> sizeof(aiologic.lowlevel.lazydeque)
128

>>> sizeof(threading.Event)
1128
>>> sizeof(asyncio.Event)
896
>>> sizeof(anyio.Event)
1024
>>> sizeof(aiologic.Event)
192
>>> sizeof(threading.Lock)
56
>>> sizeof(asyncio.Lock)
136
>>> sizeof(anyio.Lock)
904
>>> sizeof(aiologic.Lock)
280
>>> sizeof(threading.Semaphore, 1)
1128
>>> sizeof(asyncio.Semaphore, 1)
136
>>> sizeof(anyio.Semaphore, 1)
912
>>> sizeof(aiologic.Semaphore, 1)
264
>>> sizeof(threading.Condition)
1000
>>> sizeof(asyncio.Condition)
1056
>>> sizeof(anyio.Condition)
1800
>>> sizeof(aiologic.Condition)
616

I think it would be nice if I can replace that as well because asyncio.Lock is binded to some event loop and unfriendly to pytest asyncio, which creates a new event loop for every test (?).

Yes, of course.

[x42005e1f]

I would also add that if you care about performance, try aiologic.BinarySemaphore instead of aiologic.Lock if you always ensure release for each acquire. The reason for this is described in one discussion.

[ultmaster]

Thanks for the explanation. Really learned a lot!

I think I'll migrate one step at a time and have each migration fully tested.