implement exit_code hooks

wandb/internal/internal.py

@@ -145,6 +145,8 @@ def __init__(self, settings, q, resp_q):
         # TODO(jhr): do something better, why do we need to send full lines?
         self._partial_output = dict()
 
+        self._exit_code = 0
+
     def _flatten(self, dictionary):
         if type(dictionary) == dict:
             for k, v in list(dictionary.items()):
@@ -155,6 +157,9 @@ def _flatten(self, dictionary):
                         dictionary[k + "." + k2] = v2
 
     def handle_exit(self, data):
+        exit = data.exit
+        self._exit_code = exit.exit_code
+
         # Ensure we've at least noticed every file in the run directory. Sometimes
         # we miss things because asynchronously watching filesystems isn't reliable.
         run_dir = self._settings.files_dir
@@ -289,7 +294,7 @@ def finish(self):
             self._pusher.finish()
         if self._fs:
             # FIXME(jhr): now is a good time to output pending output lines
-            self._fs.finish(0)
+            self._fs.finish(self._exit_code)
         if self._pusher:
             self._pusher.update_all_files()
             files = self._pusher.files()

wandb/sdk/wandb_init.py

@@ -11,6 +11,7 @@
 from wandb.backend.backend import Backend
 from wandb.stuff import util2
 from wandb.util import reporting
+from wandb.errors import Error
 
 import time
 import json
@@ -22,6 +23,7 @@
 import sys
 import os
 from wandb.util import redirect
+import traceback
 
 from .wandb_settings import Settings
 
@@ -35,6 +37,40 @@ def online_status(*args, **kwargs):
     pass
 
 
+class ExitHooks(object):
+    def __init__(self):
+        self.exit_code = 0
+        self.exception = None
+
+    def hook(self):
+        self._orig_exit = sys.exit
+        sys.exit = self.exit
+        sys.excepthook = self.exc_handler
+
+    def exit(self, code=0):
+        orig_code = code
+        if code is None:
+            code = 0
+        elif not isinstance(code, int):
+            code = 1
+        self.exit_code = code
+        self._orig_exit(orig_code)
+
+    def was_ctrl_c(self):
+        return isinstance(self.exception, KeyboardInterrupt)
+
+    def exc_handler(self, exc_type, exc, *tb):
+        self.exit_code = 1
+        self.exception = exc
+        if issubclass(exc_type, Error):
+            termerror(str(exc))
+
+        if self.was_ctrl_c():
+            self.exit_code = 255
+
+        traceback.print_exception(exc_type, exc, *tb)
+
+
 def win32_redirect(stdout_slave_fd, stderr_slave_fd):
     # import win32api
 
@@ -92,6 +128,7 @@ def __init__(self):
         self._save_stdout = None
         self._save_stderr = None
 
+        self._hooks = None
         self._atexit_cleanup_called = None
 
     def setup(self, kwargs):
@@ -141,7 +178,9 @@ def _atexit_cleanup(self):
             return
         self._atexit_cleanup_called = True
 
-        ret = self.backend.interface.send_exit_sync(0, timeout=60)
+        exit_code = self._hooks.exit_code if self._hooks else 0
+        logger.info("got exitcode: %d", exit_code)
+        ret = self.backend.interface.send_exit_sync(exit_code, timeout=60)
         logger.info("got exit ret: %s", ret)
 
         self._restore()
@@ -305,6 +344,8 @@ def init(self):
         run.on_start()
 
         logger.info("atexit reg")
+        self._hooks = ExitHooks()
+        self._hooks.hook()
         atexit.register(lambda: self._atexit_cleanup())
 
         if self._use_redirect:

wandb/sdk_py27/wandb_init.py

@@ -35,6 +35,40 @@ def online_status(*args, **kwargs):
     pass
 
 
+class ExitHooks(object):
+    def __init__(self):
+        self.exit_code = 0
+        self.exception = None
+
+    def hook(self):
+        self._orig_exit = sys.exit
+        sys.exit = self.exit
+        sys.excepthook = self.exc_handler
+
+    def exit(self, code=0):
+        orig_code = code
+        if code is None:
+            code = 0
+        elif not isinstance(code, int):
+            code = 1
+        self.exit_code = code
+        self._orig_exit(orig_code)
+
+    def was_ctrl_c(self):
+        return isinstance(self.exception, KeyboardInterrupt)
+
+    def exc_handler(self, exc_type, exc, *tb):
+        self.exit_code = 1
+        self.exception = exc
+        if issubclass(exc_type, Error):
+            termerror(str(exc))
+
+        if self.was_ctrl_c():
+            self.exit_code = 255
+
+        traceback.print_exception(exc_type, exc, *tb)
+
+
 def win32_redirect(stdout_slave_fd, stderr_slave_fd):
     # import win32api
 
@@ -92,6 +126,7 @@ def __init__(self):
         self._save_stdout = None
         self._save_stderr = None
 
+        self._hooks = None
         self._atexit_cleanup_called = None
 
     def setup(self, kwargs):
@@ -141,7 +176,9 @@ def _atexit_cleanup(self):
             return
         self._atexit_cleanup_called = True
 
-        ret = self.backend.interface.send_exit_sync(0, timeout=60)
+        exit_code = self._hooks.exit_code if self._hooks else 0
+        logger.info("got exitcode: %d", exit_code)
+        ret = self.backend.interface.send_exit_sync(exit_code, timeout=60)
         logger.info("got exit ret: %s", ret)
 
         self._restore()
@@ -305,6 +342,8 @@ def init(self):
         run.on_start()
 
         logger.info("atexit reg")
+        self._hooks = ExitHooks()
+        self._hooks.hook()
         atexit.register(lambda: self._atexit_cleanup())
 
         if self._use_redirect: