Try fix tests

basetenlabs · Sep 19, 2024 · 5882451 · 5882451
1 parent 54eaa90
commit 5882451
Show file tree

Hide file tree

Showing 8 changed files with 444 additions and 426 deletions.
diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml
@@ -33,7 +33,7 @@ jobs:
       - uses: ./.github/actions/setup-python/
       - run: poetry install
       - name: run tests
-        run: poetry run pytest -v --cov=truss -m 'not integration' --junitxml=report.xml
+        run: poetry run pytest --durations=0  -m 'not integration' --junitxml=report.xml
       - name: Publish Test Report  # Not sure how to display this in the UI for non PRs.
         uses: mikepenz/action-junit-report@v4
         if: always()

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.9.35"
+version = "0.9.36rc01"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"

diff --git a/truss/templates/control/control/endpoints.py b/truss/templates/control/control/endpoints.py
@@ -57,6 +57,7 @@ async def proxy(request: Request):
         ),
         stop=stop_after_attempt(INFERENCE_SERVER_START_WAIT_SECS),
         wait=wait_fixed(1),
+        reraise=True,
     ):
         with attempt:
             try:

diff --git a/truss/templates/control/control/helpers/inference_server_starter.py b/truss/templates/control/control/helpers/inference_server_starter.py
@@ -41,6 +41,7 @@ def inference_server_startup_flow(
     for attempt in Retrying(
         stop=stop_after_attempt(15),
         wait=wait_exponential(multiplier=2, min=1, max=4),
+        reraise=True,
     ):
         with attempt:
             try:

diff --git a/truss/templates/server/model_wrapper.py b/truss/templates/server/model_wrapper.py
@@ -462,7 +462,7 @@ async def _write_response_to_queue(
                     await queue.put(chunk)
             except Exception as e:
                 self._logger.exception(
-                    "Exception while reading stream response: " + str(e)
+                    "Exception while generating streamed response: " + str(e)
                 )
             finally:
                 await queue.put(SENTINEL)

diff --git a/truss/tests/templates/server/test_model_wrapper.py b/truss/tests/templates/server/test_model_wrapper.py
@@ -54,7 +54,7 @@ async def test_model_wrapper_load_error_once(app_path):
     model_wrapper = model_wrapper_class(config, sdk_trace.NoOpTracer())
     model_wrapper.load()
     # Allow load thread to execute
-    time.sleep(1)
+    time.sleep(2)
     output = await model_wrapper.predict({}, MagicMock(spec=Request))
     assert output == {}
     assert model_wrapper._model.load_count == 2

diff --git a/truss/tests/test_model_inference.py b/truss/tests/test_model_inference.py
@@ -334,7 +334,7 @@ def test_streaming_with_error_and_stacktrace():
         assert_logs_contain_error(
             container.logs(),
             error=expected_stack_trace,
-            message="Exception while reading stream response: Crashed in `bar`.",
+            message="Exception while generating streamed response: Crashed in `bar`.",
         )