Validate config.yaml fields depending on transport kind, fix model wrapper test

truss/base/truss_config.py

@@ -542,10 +542,10 @@ def _validate_path(cls, v: str) -> str:
 
 class DockerServer(custom_types.ConfigModel):
     start_command: str
-    server_port: int
-    predict_endpoint: str
-    readiness_endpoint: str
-    liveness_endpoint: str
+    server_port: Optional[int] = None
+    predict_endpoint: Optional[str] = None
+    readiness_endpoint: Optional[str] = None
+    liveness_endpoint: Optional[str] = None
 
 
 class Checkpoint(custom_types.ConfigModel):
@@ -616,6 +616,13 @@ class TrussConfig(custom_types.ConfigModel):
     apply_library_patches: bool = True
     spec_version: str = "2.0"
 
+    DOCKER_SERVER_OPTIONAL_FIELDS: ClassVar[list[str]] = [
+        "server_port",
+        "predict_endpoint",
+        "readiness_endpoint",
+        "liveness_endpoint",
+    ]
+
     class Config:
         protected_namespaces = ()  # Silence warnings about fields starting with `model_`.
 
@@ -720,6 +727,34 @@ def _serialize_trt_llm(
         exclude_unset = bool(info.context and "verbose" in info.context)
         return trt_llm.model_dump(exclude_unset=exclude_unset)
 
+    @pydantic.model_validator(mode="after")
+    def _validate_docker_server(self) -> "TrussConfig":
+        is_grpc = self.runtime.transport.kind == TransportKind.GRPC
+        has_docker_server = self.docker_server is not None
+
+        if is_grpc:
+            if not has_docker_server:
+                raise ValueError(
+                    "docker_server is required when transport kind is gRPC"
+                )
+            if any(
+                getattr(self.docker_server, field) is not None
+                for field in TrussConfig.DOCKER_SERVER_OPTIONAL_FIELDS
+            ):
+                raise ValueError(
+                    "When transport kind is gRPC, docker_server should only have start_command defined"
+                )
+        elif has_docker_server:
+            if any(
+                getattr(self.docker_server, field) is None
+                for field in TrussConfig.DOCKER_SERVER_OPTIONAL_FIELDS
+            ):
+                raise ValueError(
+                    "Please define server_port, predict_endpoint, readiness_endpoint, and liveness_endpoint for docker_server"
+                )
+
+        return self
+
 
 def _map_to_supported_python_version(python_version: str) -> str:
     """Map python version to truss supported python version.

truss/tests/test_config.py

@@ -977,3 +977,90 @@ def test_supported_versions_are_sorted():
     assert semvers == semvers_sorted, (
         f"{constants.SUPPORTED_PYTHON_VERSIONS} must be sorted ascendingly"
     )
+
+
+@pytest.mark.parametrize(
+    "transport_config",
+    [
+        pytest.param(
+            {
+                "runtime": {"transport": {"kind": "grpc"}},
+                "docker_server": {"start_command": "python main.py"},
+            },
+            id="valid-grpc-minimal",
+        ),
+        pytest.param(
+            {
+                "runtime": {"transport": {"kind": "http"}},
+                "docker_server": {
+                    "start_command": "./start.sh",
+                    "server_port": 8080,
+                    "predict_endpoint": "/predict",
+                    "readiness_endpoint": "/ready",
+                    "liveness_endpoint": "/health",
+                },
+            },
+            id="valid-http-full",
+        ),
+        pytest.param(
+            {
+                "runtime": {"transport": {"kind": "websocket"}},
+                "docker_server": {
+                    "start_command": "./start.sh",
+                    "server_port": 8080,
+                    "predict_endpoint": "/predict",
+                    "readiness_endpoint": "/ready",
+                    "liveness_endpoint": "/health",
+                },
+            },
+            id="valid-websocket-full",
+        ),
+        pytest.param(
+            {"runtime": {"transport": {"kind": "http"}}}, id="valid-http-no-docker"
+        ),
+    ],
+)
+def test_valid_transport_configurations(transport_config, tmp_path):
+    config_path = tmp_path / "config.yaml"
+    config_path.write_text(yaml.dump(transport_config))
+    config = TrussConfig.from_yaml(config_path)
+    assert config.runtime.transport.kind == TransportKind(
+        transport_config["runtime"]["transport"]["kind"]
+    )
+
+
+@pytest.mark.parametrize(
+    "invalid_config,expected_error",
+    [
+        pytest.param(
+            {
+                "runtime": {"transport": {"kind": "grpc"}},
+                "docker_server": {
+                    "start_command": "./start.sh",
+                    "server_port": 8080,
+                    "predict_endpoint": "/predict",
+                },
+            },
+            "When transport kind is gRPC, docker_server should only have start_command defined",
+            id="invalid-grpc-extra-fields",
+        ),
+        pytest.param(
+            {
+                "runtime": {"transport": {"kind": "http"}},
+                "docker_server": {"start_command": "./start.sh", "server_port": 8080},
+            },
+            "Please define server_port, predict_endpoint, readiness_endpoint, and liveness_endpoint for docker_server",
+            id="invalid-http-missing-fields",
+        ),
+        pytest.param(
+            {"runtime": {"transport": {"kind": "grpc"}}},
+            "docker_server is required when transport kind is gRPC",
+            id="invalid-grpc-missing-docker",
+        ),
+    ],
+)
+def test_invalid_transport_configurations(invalid_config, expected_error, tmp_path):
+    config_path = tmp_path / "config.yaml"
+    config_path.write_text(yaml.dump(invalid_config))
+    with pytest.raises(ValueError, match=expected_error):
+        TrussConfig.from_yaml(config_path)