drop debug

Author: technillogue

python/cog/server/connection.py

@@ -18,9 +18,6 @@ class AsyncConnection(Generic[X]):
     def __init__(self, conn: Connection) -> None:
         self.wrapped_conn = conn
         self.started = False
-        print("conn __init__")
-        # perhaps the lock should be here
-        # self.sync_lock = threading.Lock()
 
     async def async_init(self) -> None:
         if self.started:
@@ -29,10 +26,10 @@ async def async_init(self) -> None:
         # mp may have handled something already but let's dup so exit is clean
         dup_fd = os.dup(fd)
         sock = socket.socket(fileno=dup_fd)
-        # sock = socket.socket(fileno=fd)
         # we don't want to see EAGAIN, we'd rather wait
         # however, perhaps this is wrong and in some cases this could still block terribly
         # sock.setblocking(False)
+        sock.setblocking(True)
         # TODO: use /proc/sys/net/core/rmem_max, but special-case language models
         sz = 65536
         sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, sz)
@@ -94,11 +91,6 @@ def _send_bytes(self, buf: bytes) -> None:
     def send(self, obj: Any) -> None:
         self._send_bytes(_ForkingPickler.dumps(obj, protocol=5))
 
-    # # perhaps we could do it like this
-    # def send_sync(self, obj: Any) -> None:
-    #     with self.sync_lock:
-    #         self.wrapped_conn.send(obj)
-
     # we could implement async def drain() but it's not really necessary for our purposes
 
     def close(self) -> None:

python/cog/server/helpers.py

@@ -7,10 +7,6 @@
 from typing import Callable, Optional, Sequence, TextIO
 
 
-def debug(*args: str, f: io.IOBase = open("/tmp/debug", "a")) -> None:  # noqa
-    print(*args, file=f, flush=True)
-
-
 async def async_fdopen(fd: int) -> asyncio.StreamReader:
     loop = asyncio.get_running_loop()
     reader = asyncio.StreamReader()
@@ -102,31 +98,23 @@ def __init__(
 
     def drain(self) -> None:
         if self.is_async:
-            debug("ignoring drain")
             # if we're async, we assume that logs will be processed promptly,
             # and we don't want to block the event loop
             return
         self.drain_event.clear()
         for stream in self._streams:
-            debug(repr(stream), stream.name)
             stream.write(self.drain_token + "\n")
-            debug(repr(stream), "flush")
             stream.flush()
-        debug("wait drain")
         if not self.drain_event.wait(timeout=1):
-            debug("drain timed out")
             raise RuntimeError("output streams failed to drain")
-        debug("drain done")
 
     def shutdown(self) -> None:
         if not self.is_alive():
-            debug("skipping shutdown because not alive")
             return
         for stream in self._streams:
             stream.write(self.terminate_token + "\n")
             stream.flush()
             break  # only need to write one terminate token
-        debug("joining")
         self.join()
 
     async def shutdown_async(self) -> None:
@@ -149,13 +137,10 @@ def run(self) -> None:
             drain_tokens_needed += 1
 
         while not should_exit:
-            debug("selector.select")
             for key, _ in selector.select():
-                debug("selector key")
                 stream = key.data
 
                 for line in stream.wrapped:
-                    debug("redirector saw", line)
                     if not line.endswith("\n"):
                         # TODO: limit how much we're prepared to buffer on a
                         # single line
@@ -180,12 +165,9 @@ def run(self) -> None:
                     # thing in the line was a drain token (or a terminate
                     # token).
                     if full_line:
-                        debug("write hook")
                         self._write_hook(stream.name, stream.original, full_line + "\n")
-                        debug("write hook done")
 
                     if drain_tokens_seen >= drain_tokens_needed:
-                        debug("drain event set")
                         self.drain_event.set()
                         drain_tokens_seen = 0
 
@@ -203,22 +185,18 @@ async def switch_to_async(self) -> None:
 
         We must not call write_hook twice for the same data during the switch.
         """
-        debug("switch async, drain")
         # Drain the streams to ensure all buffered data is processed
         try:
             self.drain()
         except RuntimeError:
-            debug("drain failed")
             raise
-        debug("drain done, shutdown")
 
         # Shut down the thread
         # we do this before starting a coroutine that will also read from the same fd
         # so that shutdown can find the terminate tokens correctly
         self.shutdown()
         self.stream_tasks = []
         self.is_async = True
-        debug("set is async")
 
         for stream in self._streams:
             # Open each stream as a StreamReader
@@ -235,7 +213,6 @@ async def switch_to_async(self) -> None:
     async def process_stream(
         self, stream: WrappedStream, reader: asyncio.StreamReader
     ) -> None:
-        debug("process_stream", stream.name)
         buffer = io.StringIO()
         drain_tokens_seen = 0
         should_exit = False
@@ -246,7 +223,6 @@ async def process_stream(
                 break
 
             line = line.decode()
-            debug("redirector saw", line)
 
             if not line.endswith("\n"):
                 buffer.write(line)
@@ -270,7 +246,6 @@ async def process_stream(
                 self._write_hook(stream.name, stream.original, full_line + "\n")
 
             if drain_tokens_seen >= len(self._streams):
-                debug("drain event set")
                 self.drain_event.set()
                 drain_tokens_seen = 0
             if should_exit:

python/cog/server/worker.py

@@ -119,7 +119,6 @@ def __init__(
         super().__init__()
 
     def run(self) -> None:
-        debug("run")
         self._sync_events_lock = threading.Lock()
         # If we're running at a shell, SIGINT will be sent to every process in
         # the process group. We ignore it in the child process and require that
@@ -146,33 +145,24 @@ def run(self) -> None:
         self._stream_redirector.start()
         # </could be moved into StreamRedirector>
 
-        debug("setup")
         self._setup()
-        debug("loop")
         self._loop()  # shuts down stream redirector the correct way
-        debug("loop done")
         self._events.close()
 
     async def _async_init(self) -> None:
-        debug("async_init start")
         if self._events_async:
-            debug("async_init finished")
             return
         # if AsyncConnection is created before switch_to_async, a race condition can cause drain to fail
         # and write, seemingly, to block
         # maybe because we're trying to call StreamWriter.write when no event loop is running?
         await self._stream_redirector.switch_to_async()
         self._events_async = AsyncConnection(self._events)
         await self._events_async.async_init()
-        debug("async_init done")
 
     def _setup(self) -> None:
-        debug("_setup start")
         with self._handle_setup_error():
             # we need to load the predictor to know if setup is async
-            debug("'about to load")
             self._predictor = load_predictor_from_ref(self._predictor_ref)
-            debug("loaded ref")
             self._predictor.log = self._log
             # if users want to access the same event loop from setup and predict,
             # both have to be async. if setup isn't async, it doesn't matter if we
@@ -181,36 +171,24 @@ def _setup(self) -> None:
             # otherwise, if setup is sync and the user does new_event_loop to use a ClientSession,
             # then tries to use the same session from async predict, they would get an error.
             # that's significant if connections are open and would need to be discarded
-            debug("async predictor")
             if is_async_predictor(self._predictor):
-                debug("getting loop")
                 self.loop = get_loop()
-                debug("got loop")
-            debug("getattr")
             # Could be a function or a class
             if hasattr(self._predictor, "setup"):
-                debug("inspect")
                 if inspect.iscoroutinefunction(self._predictor.setup):
                     # we should probably handle Shutdown during this process?
                     # possibly we prefer to not stop-start the event loop
                     # between these calls
                     self.loop.run_until_complete(self._async_init())
                     self.loop.run_until_complete(run_setup_async(self._predictor))
-                    debug("run_setup_async done")
                 else:
-                    debug("sync setup")
                     run_setup(self._predictor)
-            debug("_setup done inside ctx mgr")
-        debug("_setup done")
 
     @contextlib.contextmanager
     def _handle_setup_error(self) -> Iterator[None]:
         done = Done()
-        debug("done")
         try:
-            debug("yield")
             yield
-            debug("yield done")
         except Exception as e:
             traceback.print_exc()
             done.error = True
@@ -226,15 +204,11 @@ def _handle_setup_error(self) -> Iterator[None]:
             # we can arrive here if there was an error setting up stream_redirector
             # for example, because drain failed
             # in this case this drain could block or fail
-            debug("setup done, calling drain")
             try:
                 self._stream_redirector.drain()
             except Exception as e:
-                debug("exc", str(e))
                 raise
-            debug("sending setup done")
             self.send(("SETUP", done))
-            debug("sent setup done")
 
     def _loop_sync(self) -> None:
         while True:
@@ -253,7 +227,6 @@ def _loop_sync(self) -> None:
         self._stream_redirector.shutdown()
 
     async def _loop_async(self) -> None:
-        debug("loop async")
         await self._async_init()
         assert self._events_async
         tasks: dict[str, asyncio.Task[None]] = {}
@@ -277,14 +250,11 @@ async def _loop_async(self) -> None:
                     print(f"Got unexpected cancellation: {ev}", file=sys.stderr)
             else:
                 print(f"Got unexpected event: {ev}", file=sys.stderr)
-        debug("shutdown_async")
         await self._stream_redirector.shutdown_async()
         self._events_async.close()
 
     def _loop(self) -> None:
-        debug("in loop")
         if is_async(get_predict(self._predictor)):
-            debug("async loop")
             self.loop.run_until_complete(self._loop_async())
         else:
             self._loop_sync()
@@ -320,14 +290,10 @@ def _emit_metric(self, name: str, value: "int | float") -> None:
 
     def send(self, obj: Any) -> None:
         if self._events_async:
-            debug("sending on async")
             self._events_async.send(obj)
-            debug("sent on async")
         else:
-            debug("send lock")
             with self._sync_events_lock:
                 self._events.send(obj)
-                debug("finished sync send")
 
     def _mk_send(self, id: str) -> Callable[[PublicEventType], None]:
         def send(event: PublicEventType) -> None: