diff --git a/src/compute-plane-services/nvca/vendor/github.com/NVIDIA/nvcf/src/libraries/go/lib/pkg/icms-translate/translate/function/llm.go b/src/compute-plane-services/nvca/vendor/github.com/NVIDIA/nvcf/src/libraries/go/lib/pkg/icms-translate/translate/function/llm.go index 21ec72965..3369cae12 100644 --- a/src/compute-plane-services/nvca/vendor/github.com/NVIDIA/nvcf/src/libraries/go/lib/pkg/icms-translate/translate/function/llm.go +++ b/src/compute-plane-services/nvca/vendor/github.com/NVIDIA/nvcf/src/libraries/go/lib/pkg/icms-translate/translate/function/llm.go @@ -19,6 +19,8 @@ package function import ( "fmt" + "strconv" + "strings" corev1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/resource" @@ -58,6 +60,24 @@ func normalizeLLMRequestRouterAddressEnvAliases(envSet map[string]string) { } } +// pylon probes the upstream over HTTP on the inference port, so the function's +// declared health endpoint only carries over when the function keeps both +// aligned. Otherwise pylon falls back to its own candidate paths. +func upstreamHealthPath(allEnvSet map[string]string) string { + path := strings.TrimSpace(allEnvSet["INFERENCE_HEALTH_ENDPOINT"]) + if path == "" { + return "" + } + if strings.EqualFold(strings.TrimSpace(allEnvSet["INFERENCE_HEALTH_PROTOCOL"]), "grpc") { + return "" + } + healthPort, err := strconv.Atoi(strings.TrimSpace(allEnvSet["INFERENCE_HEALTH_PORT"])) + if err == nil && healthPort > 0 && strconv.Itoa(healthPort) != strings.TrimSpace(allEnvSet["INFERENCE_PORT"]) { + return "" + } + return path +} + func newLLMRouterClientContainer( ls *LaunchSpecification, allEnvSet map[string]string, @@ -128,6 +148,9 @@ func newLLMRouterClientContainer( "--backend-connectivity=reverse", "--initial-input-tps=100", } + if healthPath := upstreamHealthPath(allEnvSet); healthPath != "" { + args = append(args, fmt.Sprintf("--upstream-health-path=%s", healthPath)) + } if tcfg.StargateQUICInsecure { args = append(args, "--quic-insecure") } diff --git a/src/libraries/go/lib/pkg/icms-translate/translate/function/llm.go b/src/libraries/go/lib/pkg/icms-translate/translate/function/llm.go index 21ec72965..3369cae12 100644 --- a/src/libraries/go/lib/pkg/icms-translate/translate/function/llm.go +++ b/src/libraries/go/lib/pkg/icms-translate/translate/function/llm.go @@ -19,6 +19,8 @@ package function import ( "fmt" + "strconv" + "strings" corev1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/resource" @@ -58,6 +60,24 @@ func normalizeLLMRequestRouterAddressEnvAliases(envSet map[string]string) { } } +// pylon probes the upstream over HTTP on the inference port, so the function's +// declared health endpoint only carries over when the function keeps both +// aligned. Otherwise pylon falls back to its own candidate paths. +func upstreamHealthPath(allEnvSet map[string]string) string { + path := strings.TrimSpace(allEnvSet["INFERENCE_HEALTH_ENDPOINT"]) + if path == "" { + return "" + } + if strings.EqualFold(strings.TrimSpace(allEnvSet["INFERENCE_HEALTH_PROTOCOL"]), "grpc") { + return "" + } + healthPort, err := strconv.Atoi(strings.TrimSpace(allEnvSet["INFERENCE_HEALTH_PORT"])) + if err == nil && healthPort > 0 && strconv.Itoa(healthPort) != strings.TrimSpace(allEnvSet["INFERENCE_PORT"]) { + return "" + } + return path +} + func newLLMRouterClientContainer( ls *LaunchSpecification, allEnvSet map[string]string, @@ -128,6 +148,9 @@ func newLLMRouterClientContainer( "--backend-connectivity=reverse", "--initial-input-tps=100", } + if healthPath := upstreamHealthPath(allEnvSet); healthPath != "" { + args = append(args, fmt.Sprintf("--upstream-health-path=%s", healthPath)) + } if tcfg.StargateQUICInsecure { args = append(args, "--quic-insecure") } diff --git a/src/libraries/go/lib/pkg/icms-translate/translate/function/llm_test.go b/src/libraries/go/lib/pkg/icms-translate/translate/function/llm_test.go index 557e701e7..8157e6b0c 100644 --- a/src/libraries/go/lib/pkg/icms-translate/translate/function/llm_test.go +++ b/src/libraries/go/lib/pkg/icms-translate/translate/function/llm_test.go @@ -47,6 +47,16 @@ func assertCanonicalPylonBootstrapArgs(t *testing.T, args []string) { assert.Equal(t, []string{"--initial-input-tps=100"}, initialInputTPSArgs) } +func healthPathArgs(args []string) []string { + var healthPaths []string + for _, arg := range args { + if strings.HasPrefix(arg, "--upstream-health-path=") { + healthPaths = append(healthPaths, arg) + } + } + return healthPaths +} + func TestNewLLMRouterClientContainer(t *testing.T) { type spec struct { name string @@ -298,6 +308,69 @@ func TestNewLLMRouterClientContainer(t *testing.T) { assert.Equal(t, "stargate.example.com:443", envMap["STARGATE_ADDRESS"]) }, }, + { + name: "declared health endpoint is passed to pylon", + ls: &LaunchSpecification{}, + allEnvSet: map[string]string{ + "STARGATE_ADDRESS": "stargate.example.com:443", + "INFERENCE_PORT": "8080", + "INFERENCE_HEALTH_ENDPOINT": "/v1/health/ready", + "INFERENCE_HEALTH_PORT": "8080", + "INFERENCE_HEALTH_PROTOCOL": "http", + }, + tcfg: TranslateConfig{}, + instanceID: "inst-health", + isHelm: false, + validate: func(t *testing.T, c corev1.Container) { + assert.Equal(t, []string{"--upstream-health-path=/v1/health/ready"}, healthPathArgs(c.Args)) + }, + }, + { + name: "health endpoint on a port other than the inference port is skipped", + ls: &LaunchSpecification{}, + allEnvSet: map[string]string{ + "STARGATE_ADDRESS": "stargate.example.com:443", + "INFERENCE_PORT": "8080", + "INFERENCE_HEALTH_ENDPOINT": "/custom/ready", + "INFERENCE_HEALTH_PORT": "9090", + }, + tcfg: TranslateConfig{}, + instanceID: "inst-health-port", + isHelm: false, + validate: func(t *testing.T, c corev1.Container) { + assert.Empty(t, healthPathArgs(c.Args)) + }, + }, + { + name: "gRPC health protocol is skipped", + ls: &LaunchSpecification{}, + allEnvSet: map[string]string{ + "STARGATE_ADDRESS": "stargate.example.com:443", + "INFERENCE_PORT": "8080", + "INFERENCE_HEALTH_ENDPOINT": "/grpc.health.v1.Health/Check", + "INFERENCE_HEALTH_PROTOCOL": "gRPC", + }, + tcfg: TranslateConfig{}, + instanceID: "inst-health-grpc", + isHelm: false, + validate: func(t *testing.T, c corev1.Container) { + assert.Empty(t, healthPathArgs(c.Args)) + }, + }, + { + name: "no declared health endpoint leaves the pylon defaults in place", + ls: &LaunchSpecification{}, + allEnvSet: map[string]string{ + "STARGATE_ADDRESS": "stargate.example.com:443", + "INFERENCE_PORT": "8080", + }, + tcfg: TranslateConfig{}, + instanceID: "inst-health-absent", + isHelm: false, + validate: func(t *testing.T, c corev1.Container) { + assert.Empty(t, healthPathArgs(c.Args)) + }, + }, } for _, tt := range cases { diff --git a/src/libraries/go/lib/testdata/icms-translate/function/container/llm/exp.yaml b/src/libraries/go/lib/testdata/icms-translate/function/container/llm/exp.yaml index 9c416a60f..e421e46d7 100644 --- a/src/libraries/go/lib/testdata/icms-translate/function/container/llm/exp.yaml +++ b/src/libraries/go/lib/testdata/icms-translate/function/container/llm/exp.yaml @@ -187,6 +187,7 @@ spec: - --auth-token-file=/var/run/llm/worker-token - --backend-connectivity=reverse - --initial-input-tps=100 + - --upstream-health-path=/v1/health/ready - --quic-insecure - --model-name=model-gamma - --model-name=model-epsilon diff --git a/src/libraries/go/lib/testdata/icms-translate/function/helmchart/llm/exp.yaml b/src/libraries/go/lib/testdata/icms-translate/function/helmchart/llm/exp.yaml index 41484e5c6..8c6bc6e55 100644 --- a/src/libraries/go/lib/testdata/icms-translate/function/helmchart/llm/exp.yaml +++ b/src/libraries/go/lib/testdata/icms-translate/function/helmchart/llm/exp.yaml @@ -150,6 +150,7 @@ spec: - --auth-token-file=/var/run/llm/worker-token - --backend-connectivity=reverse - --initial-input-tps=100 + - --upstream-health-path=/v2/health/ready - --quic-insecure - --model-name=model-gamma - --model-name=model-epsilon