From b366079ab80e10fca490689811ecd712aa23e438 Mon Sep 17 00:00:00 2001
From: Tianshu Cheng <tianshuc@Tianshus-MBP.localdomain>
Date: Thu, 5 Dec 2024 15:11:29 -0800
Subject: [PATCH] update

---
 truss/base/truss_config.py                                | 8 ++++----
 truss/contexts/image_builder/serving_image_builder.py     | 6 ++++++
 truss/templates/docker_server/proxy.conf.jinja            | 4 ----
 .../tests/test_data/test_custom_server_truss/config.yaml  | 2 ++
 4 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/truss/base/truss_config.py b/truss/base/truss_config.py
index 59a428e32..8809599d3 100644
--- a/truss/base/truss_config.py
+++ b/truss/base/truss_config.py
@@ -421,8 +421,8 @@ class DockerServer:
     start_command: str
     server_port: int
     predict_endpoint: str
-    readiness_endpoint: Optional[str] = None
-    liveness_endpoint: Optional[str] = None
+    readiness_endpoint: str
+    liveness_endpoint: str
 
     @staticmethod
     def from_dict(d) -> "DockerServer":
@@ -430,8 +430,8 @@ def from_dict(d) -> "DockerServer":
             start_command=d.get("start_command"),
             server_port=d.get("server_port"),
             predict_endpoint=d.get("predict_endpoint"),
-            readiness_endpoint=d.get("readiness_endpoint", None),
-            liveness_endpoint=d.get("liveness_endpoint", None),
+            readiness_endpoint=d.get("readiness_endpoint"),
+            liveness_endpoint=d.get("liveness_endpoint"),
         )
 
     def to_dict(self):
diff --git a/truss/contexts/image_builder/serving_image_builder.py b/truss/contexts/image_builder/serving_image_builder.py
index cf1f6879e..70d7ebb46 100644
--- a/truss/contexts/image_builder/serving_image_builder.py
+++ b/truss/contexts/image_builder/serving_image_builder.py
@@ -312,6 +312,12 @@ def generate_docker_server_nginx_config(build_dir, config):
     assert (
         config.docker_server.server_port is not None
     ), "docker_server.server_port is required to use custom server"
+    assert (
+        config.docker_server.readiness_endpoint is not None
+    ), "docker_server.readiness_endpoint is required to use custom server"
+    assert (
+        config.docker_server.liveness_endpoint is not None
+    ), "docker_server.liveness_endpoint is required to use custom server"
 
     nginx_content = nginx_template.render(
         server_endpoint=config.docker_server.predict_endpoint,
diff --git a/truss/templates/docker_server/proxy.conf.jinja b/truss/templates/docker_server/proxy.conf.jinja
index ecd059d26..cb210635b 100644
--- a/truss/templates/docker_server/proxy.conf.jinja
+++ b/truss/templates/docker_server/proxy.conf.jinja
@@ -3,7 +3,6 @@ server {
     listen 8080;
     client_max_body_size {{client_max_body_size}};
 
-{%- if liveness_endpoint %}
     # Liveness endpoint override
     location = / {
         proxy_redirect off;
@@ -13,8 +12,6 @@ server {
 
         proxy_pass http://127.0.0.1:{{server_port}};
     }
-{%- endif %}
-{%- if readiness_endpoint %}
     # Readiness endpoint override
     location ~ ^/v1/models/model$ {
         proxy_redirect off;
@@ -24,7 +21,6 @@ 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
index e9112f04e..95c1aa15e 100644
--- a/truss/tests/test_data/test_custom_server_truss/config.yaml
+++ b/truss/tests/test_data/test_custom_server_truss/config.yaml
@@ -4,6 +4,8 @@ docker_server:
   start_command: python main.py
   predict_endpoint: /predict
   server_port: 8000
+  readiness_endpoint: /
+  liveness_endpoint: /
 resources:
   accelerator: null
   cpu: '1'