Update

basetenlabs · Nov 5, 2024 · 0145dc7 · 0145dc7
1 parent 5788801
commit 0145dc7
Showing 1 changed file with 6 additions and 7 deletions.
diff --git a/truss/tests/test_model_inference.py b/truss/tests/test_model_inference.py
@@ -786,7 +786,7 @@ def predict(self, model_input):
             return self.environment_name
     """
     config = "model_name: init-environment-truss"
-    with ensure_kill_all(), temp_truss(model, config) as tr:
+    with ensure_kill_all(), _temp_truss(model, config) as tr:
         # Mimic environment changing to staging
         staging_env = {"name": "staging"}
         staging_env_str = json.dumps(staging_env)
@@ -810,7 +810,7 @@ def predict(self, model_input):
 
     # Test a truss deployment with no associated environment
     config = "model_name: init-no-environment-truss"
-    with ensure_kill_all(), temp_truss(model, config) as tr:
+    with ensure_kill_all(), _temp_truss(model, config) as tr:
         container = tr.docker_run(
             local_port=8090,
             detach=True,
@@ -836,8 +836,7 @@ def setup_environment(self, environment: Optional[dict]):
         def predict(self, model_input):
             return model_input
     """
-    config = "model_name: setup-environment-truss"
-    with ensure_kill_all(), temp_truss(model, config) as tr:
+    with ensure_kill_all(), _temp_truss(model, "") as tr:
         container = tr.docker_run(
             local_port=8090,
             detach=True,
@@ -887,8 +886,7 @@ def load(self):
         def predict(self, model_input):
             return model_input
     """
-    config = "model_name: setup-environment-and-load-truss"
-    with ensure_kill_all(), temp_truss(model, config) as tr:
+    with ensure_kill_all(), _temp_truss(model, "") as tr:
         # Mimic environment changing to staging
         staging_env = {"name": "staging"}
         staging_env_str = json.dumps(staging_env)
@@ -934,7 +932,8 @@ def predict(self, model_input):
             ]
         )
 
-def _patch_termination_timeout(container: Container, truss_container_fs):
+
+def _patch_termination_timeout(container: Container, seconds: int, truss_container_fs):
     app_path = truss_container_fs / "app"
     sys.path.append(str(app_path))
     import truss_server