Fix send race [debug] #1786

technillogue · 2024-07-03T23:13:16Z

This solves a similar problem as #1758: in some cases, currently emitting a metric or yielding an output can result in an EAGAIN error, and we believe Done events are sometimes dropped. In async land, we shouldn't need a lock; multiple writers should be able to write to a StreamWriter safely. We also keep a sync Lock when the predictor is not async.

However, I'm still not sure if StreamWriter is thread-safe.

After some research, StreamWriter.write calls self._transport.write and when StreamWriter is wrapping a unix socket and using the default selector, the _transport is a _SelectorSocketTransport, which calls [self._buffer.extend](https://github.com/python/cpython/blob/3.11/Lib/asyncio/selector_events.py#L1077), and self._buffer is a [bytearray](https://github.com/python/cpython/blob/3.11/Lib/asyncio/selector_events.py#L761). However, on 3.12, [append is called](https://github.com/python/cpython/blob/3.12/Lib/asyncio/selector_events.py#L1091) instead of extend.

After more soul-searching, I think write should be basically thread-safe: the two critical calls, socket.send() and bytearray.append/extend should be represented as a single bytecode instruction and some native code which works on python objects and does not release the GIL. As far as I understand, python threads are switched "between bytecodes". Even though the individual calls are thread-safe, the overall write method is not quite thread safe: a relevant thread switch could occur between the if not self._buffer, _sock.send, and _buffer.extend lines. However, in that case the race condition would result data could being sent out of order or incorrectly delayed/sent early, but it I think it shouldn't be corrupted, and we don't really care about the ordering of log lines vs outputs.

I can see two solutions:

Ideally, we would probably add an alternate implementation of StreamWriter that uses the main event loop and wraps stderr/stdout in StreamReaders, start a separate task for each stream, and use await wrapped_stream.read() so that threads are not necessary. there's a tricky moment where we need to use the threaded StreamRedirector to capture logs while the predictor is being imported, since we don't know if we're going to be async or not.
Alternatively, we could try to use a queue or deque to communicate between threads, so instead of calling stream_write_hook StreamRedirector would do queue.put, the main event loop would do queue.get and not worry about thread safety for the pipe. The problem with this is that asyncio.Queue is not thread-safe and queue.Queue could block the event loop. You could use a busy wait or similar and get_nowait, but that has other downsides.

…cellation to succeed while I struggle with #1786, we see queue spikes because not responding to cancellation promptly causes the pod to get restarted. this is a dirty hack to pretend like cancellation works immediately. as soon as we fix the race condition (and possibly any issues with task.cancel() behaving differently from signal handlers), we can drop this.

…cellation to succeed while I struggle with #1786, we see queue spikes because not responding to cancellation promptly causes the pod to get restarted. this is a dirty hack to pretend like cancellation works immediately. as soon as we fix the race condition (and possibly any issues with task.cancel() behaving differently from signal handlers), we can drop this. Signed-off-by: technillogue <technillogue@gmail.com>

mattt

This is looking good! I think if we can get tests passing (integration got stuck on the last commit, and I cancelled after 38 minutes) and clean up the debug statements / comments, this should be ready to go.

python/cog/server/helpers.py

mattt · 2024-07-17T19:00:30Z

python/cog/server/helpers.py

+            if full_line.endswith(self.terminate_token):
+                full_line = full_line[: -len(self.terminate_token)]
+                should_exit = True
+
+            if full_line.endswith(self.drain_token):
+                drain_tokens_seen += 1
+                full_line = full_line[: -len(self.drain_token)]
+
+            if full_line:
+                # Call write_hook from the main thread
+                self._write_hook(stream.name, stream.original, full_line + "\n")


Worth profiling, but precomputed / fixed lengths instead of a string suffix and len could be noticeably faster if this is a hot code path.

technillogue · 2024-07-17T22:53:04Z

python/cog/server/worker.py

+                    async def setup_async() -> None:
+                        # we prefer to not stop-start the event loop between these calls
+                        await self._async_init()  # this creates tasks
+                        debug("async_init done")
+                        await run_setup_async(self._predictor)
+                        debug("ran setup async")


not sure if this is still needed

Repeating back to confirm my own understanding: We no longer need setup_async, because our new approach is to always run (sync) setup and then transition to async at the end after we know the signature of predict. Right?

no, setup may be async -- what I mean here is that

Suggested change

async def setup_async() -> None:

# we prefer to not stop-start the event loop between these calls

await self._async_init() # this creates tasks

debug("async_init done")

await run_setup_async(self._predictor)

debug("ran setup async")

self.loop.run_until_complete(self._async_init())

self.loop.run_until_complete(run_setup_async(self._predictor))

I added the setup_async function there so that the event loop kept going between the two coros because I thought that was causing problems, but I'm not sure that's necessary

…cellation to succeed while I struggle with #1786, we see queue spikes because not responding to cancellation promptly causes the pod to get restarted. this is a dirty hack to pretend like cancellation works immediately. as soon as we fix the race condition (and possibly any issues with task.cancel() behaving differently from signal handlers), we can drop this. Signed-off-by: technillogue <technillogue@gmail.com>

* unconditionally mark predictions as cancelled without waiting for cancellation to succeed while I struggle with #1786, we see queue spikes because not responding to cancellation promptly causes the pod to get restarted. this is a dirty hack to pretend like cancellation works immediately. as soon as we fix the race condition (and possibly any issues with task.cancel() behaving differently from signal handlers), we can drop this. Signed-off-by: technillogue <technillogue@gmail.com> * make Mux.write a sync method this makes the cancel patch cleaner, but might be a small speedup for high-throughput outputs Signed-off-by: technillogue <technillogue@gmail.com> * gate this behind a feature flag --------- Signed-off-by: technillogue <technillogue@gmail.com>

technillogue · 2024-07-25T16:23:08Z

https://replicatehq.slack.com/archives/C0617RF4HHP/p1719896679090459?thread_ts=1719867269.654439&cid=C0617RF4HHP

previously, AsyncConnection would only be used in _loop_async, and _events.send would always be used, which usually immediately makes a write(2) syscall. in contrast, AsyncConnection has a StreamWriter, which should be safe to call from different coroutines.

…r loop

…_queue

mattt · 2024-07-31T16:38:42Z

Closing in favor of #1831

technillogue requested a review from a team July 3, 2024 23:13

technillogue force-pushed the syl/fix-send-race-async branch from b3b7df9 to d42afa5 Compare July 11, 2024 21:35

technillogue mentioned this pull request Jul 11, 2024

patch cancel bug: immediately mark predictions as cancelled #1798

Merged

technillogue force-pushed the syl/fix-send-race-async branch 2 times, most recently from e7a5ff4 to 4982b97 Compare July 17, 2024 18:33

mattt reviewed Jul 17, 2024

View reviewed changes

technillogue commented Jul 17, 2024

View reviewed changes

technillogue force-pushed the syl/fix-send-race-async branch 3 times, most recently from 96f3da1 to 1bff4ee Compare July 19, 2024 00:05

technillogue added 8 commits July 26, 2024 00:01

use a lock to protect sync writes

1cd6c81

process_log_queue + a lot of debugging

593b3f6

don't start asyncconn or async queue processing during setup, only fo…

a59f334

…r loop

sketch

d9816c1

add a switch_to_async method to StreamRedirector and drop process_log…

2d6c0d8

…_queue

drain before starting new readers and don't strip newlines

affb023

maybe fix the race condition

90f3f07

technillogue mentioned this pull request Jul 26, 2024

fix send race by handling logs in the main thread when async #1831

Closed

technillogue force-pushed the syl/fix-send-race-async branch from 1bff4ee to 90f3f07 Compare July 26, 2024 20:17

finally a test

a068801

technillogue force-pushed the syl/fix-send-race-async branch from 00c530a to a068801 Compare July 27, 2024 18:23

idk test in ci

0cc54c8

mattt closed this Jul 31, 2024

technillogue changed the title ~~Fix send race~~ Fix send race [debug] Oct 22, 2024

nickstenning deleted the syl/fix-send-race-async branch February 11, 2025 10:09

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Fix send race [debug] #1786

Fix send race [debug] #1786

Uh oh!

technillogue commented Jul 3, 2024 •

edited

Loading

Uh oh!

mattt left a comment

Uh oh!

Uh oh!

mattt Jul 17, 2024

Uh oh!

technillogue Jul 17, 2024

Uh oh!

mattt Jul 18, 2024

Uh oh!

technillogue Jul 18, 2024

Uh oh!

technillogue commented Jul 25, 2024

Uh oh!

mattt commented Jul 31, 2024

Uh oh!

Uh oh!

Fix send race [debug] #1786

Fix send race [debug] #1786

Uh oh!

Conversation

technillogue commented Jul 3, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mattt left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mattt Jul 17, 2024

Choose a reason for hiding this comment

Uh oh!

technillogue Jul 17, 2024

Choose a reason for hiding this comment

Uh oh!

mattt Jul 18, 2024

Choose a reason for hiding this comment

Uh oh!

technillogue Jul 18, 2024

Choose a reason for hiding this comment

Uh oh!

technillogue commented Jul 25, 2024

Uh oh!

mattt commented Jul 31, 2024

Uh oh!

Uh oh!

technillogue commented Jul 3, 2024 •

edited

Loading