Better exception catching (#69)

* Catch server exceptions
* Catch status and in-thread exceptions
* Add unit tests

Author: destefy

centml/compiler/backend.py

@@ -24,16 +24,16 @@ def __init__(self, module: GraphModule, inputs: List[torch.Tensor]):
         self._inputs: List[torch.Tensor] = inputs
         self.compiled_forward_function: Optional[Callable[[torch.Tensor], tuple]] = None
         self.lock = th.Lock()
-        self.child_thread = th.Thread(target=self.remote_compilation)
+        self.child_thread = th.Thread(target=self.remote_compilation_starter)
 
         self.serialized_model_dir: Optional[TemporaryDirectory] = None
         self.serialized_model_path: Optional[str] = None
         self.serialized_input_path: Optional[str] = None
 
         try:
             self.child_thread.start()
-        except Exception:
-            logging.getLogger(__name__).exception("Remote compilation failed with the following exception: \n")
+        except Exception as e:
+            logging.getLogger(__name__).exception(f"Failed to start compilation thread\n{e}")
 
     @property
     def module(self) -> Optional[GraphModule]:
@@ -108,7 +108,6 @@ def _compile_model(self, model_id: str):
                 files={"model": model_file, "inputs": input_file},
                 timeout=settings.TIMEOUT,
             )
-
         if compile_response.status_code != HTTPStatus.OK:
             raise Exception(
                 f"Compile model: request failed, exception from server:\n{compile_response.json().get('detail')}\n"
@@ -118,21 +117,30 @@ def _wait_for_status(self, model_id: str) -> bool:
         tries = 0
         while True:
             # get server compilation status
-            status_response = requests.get(f"{settings.CENTML_SERVER_URL}/status/{model_id}", timeout=settings.TIMEOUT)
-            if status_response.status_code != HTTPStatus.OK:
-                raise Exception(
-                    f"Status check: request failed, exception from server:\n{status_response.json().get('detail')}"
+            status = None
+            try:
+                status_response = requests.get(
+                    f"{settings.CENTML_SERVER_URL}/status/{model_id}", timeout=settings.TIMEOUT
                 )
-            status = status_response.json().get("status")
+                if status_response.status_code != HTTPStatus.OK:
+                    raise Exception(
+                        f"Status check: request failed, exception from server:\n{status_response.json().get('detail')}"
+                    )
+                status = status_response.json().get("status")
+            except Exception as e:
+                logging.getLogger(__name__).exception(f"Status check failed:\n{e}")
 
             if status == CompilationStatus.DONE.value:
                 return True
             elif status == CompilationStatus.COMPILING.value:
                 pass
             elif status == CompilationStatus.NOT_FOUND.value:
-                tries += 1
                 logging.info("Submitting model to server for compilation.")
-                self._compile_model(model_id)
+                try:
+                    self._compile_model(model_id)
+                except Exception as e:
+                    logging.getLogger(__name__).exception(f"Submitting compilation failed:\n{e}")
+                tries += 1
             else:
                 tries += 1
 
@@ -141,6 +149,12 @@ def _wait_for_status(self, model_id: str) -> bool:
 
             time.sleep(settings.COMPILING_SLEEP_TIME)
 
+    def remote_compilation_starter(self):
+        try:
+            self.remote_compilation()
+        except Exception as e:
+            logging.getLogger(__name__).exception(f"Compilation thread failed:\n{e}")
+
     def remote_compilation(self):
         self._serialize_model_and_inputs()
 

tests/test_backend.py

@@ -119,8 +119,10 @@ def test_successful_download(self, mock_requests, mock_load, mock_open, mock_mak
 
 
 class TestWaitForStatus(SetUpGraphModule):
+    @patch("centml.compiler.config.settings.COMPILING_SLEEP_TIME", new=0)
     @patch("centml.compiler.backend.requests")
-    def test_invalid_status(self, mock_requests):
+    @patch("logging.Logger.exception")
+    def test_invalid_status(self, mock_logger, mock_requests):
         mock_response = MagicMock()
         mock_response.status_code = HTTPStatus.BAD_REQUEST
         mock_requests.get.return_value = mock_response
@@ -129,8 +131,28 @@ def test_invalid_status(self, mock_requests):
         with self.assertRaises(Exception) as context:
             self.runner._wait_for_status(model_id)
 
-        mock_requests.get.assert_called_once()
-        self.assertIn("Status check: request failed, exception from server", str(context.exception))
+        mock_requests.get.assert_called()
+        assert mock_requests.get.call_count == settings.MAX_RETRIES + 1
+        assert len(mock_logger.call_args_list) == settings.MAX_RETRIES + 1
+        print(mock_logger.call_args_list)
+        assert mock_logger.call_args_list[0].startswith("Status check failed:")
+        assert "Waiting for status: compilation failed too many times.\n" == str(context.exception)
+
+    @patch("centml.compiler.config.settings.COMPILING_SLEEP_TIME", new=0)
+    @patch("centml.compiler.backend.requests")
+    @patch("logging.Logger.exception")
+    def test_exception_in_status(self, mock_logger, mock_requests):
+        exception_message = "Exiting early"
+        mock_requests.get.side_effect = Exception(exception_message)
+
+        model_id = "exception_in_status"
+        with self.assertRaises(Exception) as context:
+            self.runner._wait_for_status(model_id)
+
+        mock_requests.get.assert_called()
+        assert mock_requests.get.call_count == settings.MAX_RETRIES + 1
+        mock_logger.assert_called_with(f"Status check failed:\n{exception_message}")
+        assert str(context.exception) == "Waiting for status: compilation failed too many times.\n"
 
     @patch("centml.compiler.config.settings.COMPILING_SLEEP_TIME", new=0)
     @patch("centml.compiler.backend.Runner._compile_model")
@@ -151,6 +173,7 @@ def test_max_tries(self, mock_requests, mock_compile):
     @patch("centml.compiler.config.settings.COMPILING_SLEEP_TIME", new=0)
     @patch("centml.compiler.backend.requests")
     def test_wait_on_compilation(self, mock_requests):
+        # Mock the status check
         COMPILATION_STEPS = 10
         mock_response = MagicMock()
         mock_response.status_code = HTTPStatus.OK
@@ -163,6 +186,34 @@ def test_wait_on_compilation(self, mock_requests):
         # _wait_for_status should return True when compilation DONE
         assert self.runner._wait_for_status(model_id)
 
+    @patch("centml.compiler.config.settings.COMPILING_SLEEP_TIME", new=0)
+    @patch("centml.compiler.backend.requests")
+    @patch("centml.compiler.backend.Runner._compile_model")
+    @patch("logging.Logger.exception")
+    def test_exception_in_compilation(self, mock_logger, mock_compile, mock_requests):
+        # Mock the status check
+        mock_response = MagicMock()
+        mock_response.status_code = HTTPStatus.OK
+        mock_response.json.return_value = {"status": CompilationStatus.NOT_FOUND.value}
+        mock_requests.get.return_value = mock_response
+
+        # Mock the compile model function
+        exception_message = "Exiting early"
+        mock_compile.side_effect = Exception(exception_message)
+
+        model_id = "exception_in_compilation"
+        with self.assertRaises(Exception) as context:
+            self.runner._wait_for_status(model_id)
+
+        mock_requests.get.assert_called()
+        assert mock_requests.get.call_count == settings.MAX_RETRIES + 1
+
+        mock_compile.assert_called()
+        assert mock_compile.call_count == settings.MAX_RETRIES + 1
+
+        mock_logger.assert_called_with(f"Submitting compilation failed:\n{exception_message}")
+        assert str(context.exception) == "Waiting for status: compilation failed too many times.\n"
+
     @patch("centml.compiler.backend.requests")
     def test_compilation_done(self, mock_requests):
         mock_response = MagicMock()