make liveness_endpoint and readiness_endpoint optional

basetenlabs · Nov 28, 2024 · 5063ef5 · 5063ef5
1 parent 45158fc
commit 5063ef5
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 7 deletions.
diff --git a/truss/templates/docker_server/proxy.conf.jinja b/truss/templates/docker_server/proxy.conf.jinja
@@ -2,7 +2,9 @@ server {
     # We use the proxy_read_timeout directive here (instead of proxy_send_timeout) as it sets the timeout for reading a response from the proxied server vs. setting a timeout for sending a request to the proxied server.
     listen 8080;
     client_max_body_size {{client_max_body_size}};
-    # Liveness
+
+{%- if liveness_endpoint %}
+    # Liveness endpoint override
     location = / {
         proxy_redirect off;
         proxy_read_timeout 300s;
@@ -11,8 +13,9 @@ server {
 
         proxy_pass http://127.0.0.1:{{server_port}};
     }
-
-    # Readiness
+{%- endif %}
+{%- if readiness_endpoint %}
+    # Readiness endpoint override
     location ~ ^/v1/models/model$ {
         proxy_redirect off;
         proxy_read_timeout 300s;
@@ -21,7 +24,7 @@ server {
 
         proxy_pass http://127.0.0.1:{{server_port}};
     }
-
+{%- endif %}
     # Predict
     location ~ ^/v1/models/model:predict$ {
         proxy_redirect off;

diff --git a/truss/tests/test_data/test_custom_server_truss/config.yaml b/truss/tests/test_data/test_custom_server_truss/config.yaml
@@ -2,8 +2,6 @@ base_image:
   image: baseten/fastapi-test:0.1.2
 docker_server:
   start_command: python main.py
-  readiness_endpoint: /health
-  liveness_endpoint: /health
   predict_endpoint: /predict
   server_port: 8000
 resources:

diff --git a/truss/truss_handle/truss_handle.py b/truss/truss_handle/truss_handle.py
@@ -1114,7 +1114,7 @@ def wait_for_truss(
     try:
         _wait_for_docker_build(container)
         if wait_for_server_ready:
-            if model_server_stop_retry_override:
+            if model_server_stop_retry_override is not None:
                 _wait_for_model_server(url, stop=model_server_stop_retry_override)
             else:
                 _wait_for_model_server(url)