Serve exit fixes & debug logs (#141)

Co-authored-by: Luke Lombardi <[email protected]>

internal/abstractions/endpoint/buffer.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"log"
 	"net/http"
 	"sort"
 	"sync"
@@ -168,6 +169,8 @@ func (rb *RequestBuffer) discoverContainers() {
 						return
 					}
 
+					log.Printf("<%s> Retrieved container address: %s\n", cs.ContainerId, containerAddress)
+
 					inFlight, err := rb.requestsInFlight(cs.ContainerId)
 					if err != nil {
 						return
@@ -179,7 +182,11 @@ func (rb *RequestBuffer) discoverContainers() {
 							address:  containerAddress,
 							inFlight: inFlight,
 						}
+						return
 					}
+
+					log.Printf("<%s> Container not ready\n", cs.ContainerId)
+
 				}(containerState)
 			}
 

sdk/src/beta9/runner/serve.py

@@ -1,4 +1,3 @@
-import atexit
 import os
 import signal
 import subprocess
@@ -47,6 +46,9 @@ def __init__(self) -> None:
         self.restart_event = Event()
         self.exit_event = Event()
 
+        # Register signal handlers
+        signal.signal(signal.SIGTERM, self.shutdown)
+
         # Set up the file change event handler & observer
         self.event_handler = SyncEventHandler(self.trigger_restart)
         self.observer = PollingObserver()
@@ -56,7 +58,7 @@ def __init__(self) -> None:
     def shutdown(self, signum=None, frame=None) -> None:
         self.kill_subprocess()
         self.observer.stop()
-        self.observer.join()
+        self.observer.join(timeout=0.1)
         self.exit_event.set()
         self.restart_event.set()
 
@@ -117,7 +119,6 @@ def _command() -> List[str]:
 
 if __name__ == "__main__":
     sg = ServeGateway()
-    atexit.register(sg.shutdown)
 
     try:
         sg.run(command=_command())