diff --git a/pkg/controller/v1beta1/benchmark/controller_test.go b/pkg/controller/v1beta1/benchmark/controller_test.go index 7a7d95bfc..cfb675701 100644 --- a/pkg/controller/v1beta1/benchmark/controller_test.go +++ b/pkg/controller/v1beta1/benchmark/controller_test.go @@ -21,6 +21,7 @@ import ( "sigs.k8s.io/controller-runtime/pkg/log/zap" "sigs.k8s.io/ome/pkg/apis/ome/v1beta1" + "sigs.k8s.io/ome/pkg/constants" "sigs.k8s.io/ome/pkg/controller/v1beta1/controllerconfig" ) @@ -29,6 +30,34 @@ var ( StringPtr = ptr.To[string] ) +// engineTestPod builds an engine pod for isvcName serving under servedName. The benchmark +// controller reads --api-model-name off these pods, so fixtures exercising an +// InferenceService endpoint need one. +func engineTestPod(isvcName, namespace, servedName string) *corev1.Pod { + return &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: isvcName + "-engine-t3st", + Namespace: namespace, + Labels: map[string]string{ + constants.InferenceServicePodLabelKey: isvcName, + constants.OMEComponentLabel: string(v1beta1.EngineComponent), + }, + }, + Spec: corev1.PodSpec{ + Containers: []corev1.Container{ + { + Name: "ome-container", + Image: "vllm-image", + Command: []string{ + "python3", "-m", "vllm.entrypoints.openai.api_server", + "--served-model-name", servedName, + }, + }, + }, + }, + } +} + func TestBenchmarkJobReconciler_Reconcile(t *testing.T) { scheme := runtime.NewScheme() _ = v1beta1.AddToScheme(scheme) @@ -384,6 +413,7 @@ func TestBenchmarkJobReconciler_createPodSpec(t *testing.T) { }, }, }). + WithObjects(engineTestPod("test-isvc", "default", "my-served-model")). Build() r := &BenchmarkJobReconciler{ @@ -484,6 +514,7 @@ func TestBenchmarkJobReconciler_buildBenchmarkCommand(t *testing.T) { WithObjects(tt.benchmarkJob). WithObjects(tt.isvc). WithObjects(baseModel). + WithObjects(engineTestPod("test-isvc", "default", "my-served-model")). Build() r := &BenchmarkJobReconciler{ @@ -503,6 +534,7 @@ func TestBenchmarkJobReconciler_buildBenchmarkCommand(t *testing.T) { if len(args) == 0 { t.Error("buildBenchmarkCommand() args is empty") } + assertFlagValue(t, args, "--api-model-name", "my-served-model") } }, ) @@ -761,7 +793,8 @@ func TestBenchmarkJobReconciler_createPodSpec_NodeAffinity(t *testing.T) { client := cfake.NewClientBuilder(). WithScheme(scheme). - WithObjects(benchmarkJob, inferenceService, baseModel). + WithObjects(benchmarkJob, inferenceService, baseModel, + engineTestPod("test-isvc", "default", "my-served-model")). Build() r := &BenchmarkJobReconciler{ @@ -858,7 +891,8 @@ func TestBenchmarkJobReconciler_createPodSpec_NodeAffinity_WithPodOverride(t *te client := cfake.NewClientBuilder(). WithScheme(scheme). - WithObjects(benchmarkJob, inferenceService, baseModel). + WithObjects(benchmarkJob, inferenceService, baseModel, + engineTestPod("test-isvc", "default", "my-served-model")). Build() r := &BenchmarkJobReconciler{ @@ -1030,3 +1064,18 @@ func TestBenchmarkJobReconciler_updateStatus(t *testing.T) { }) } } + +// assertFlagValue fails unless args contains flag immediately followed by want. +func assertFlagValue(t *testing.T, args []string, flag, want string) { + t.Helper() + for i, arg := range args { + if arg == flag { + if i+1 >= len(args) { + t.Fatalf("%s has no value in %v", flag, args) + } + assert.Equal(t, want, args[i+1]) + return + } + } + t.Fatalf("%s not found in %v", flag, args) +} diff --git a/pkg/controller/v1beta1/benchmark/utils/utils.go b/pkg/controller/v1beta1/benchmark/utils/utils.go index 347450045..d65d6f5e3 100644 --- a/pkg/controller/v1beta1/benchmark/utils/utils.go +++ b/pkg/controller/v1beta1/benchmark/utils/utils.go @@ -3,12 +3,14 @@ package benchmarkutils import ( "context" "fmt" + "strings" v1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/types" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/ome/pkg/apis/ome/v1beta1" + "sigs.k8s.io/ome/pkg/constants" isvcutils "sigs.k8s.io/ome/pkg/controller/v1beta1/inferenceservice/utils" "sigs.k8s.io/ome/pkg/utils/storage" ) @@ -66,9 +68,14 @@ func BuildInferenceServiceArgs(ctx context.Context, c client.Client, endpointSpe return nil, fmt.Errorf("BaseModel %s has missing Storage or Path information", baseModelName) } + modelName, err := resolveServedModelName(ctx, c, inferenceService) + if err != nil { + return nil, err + } + args := map[string]string{ "--api-key": "sample-key", // TODO: Use actual service account key later - "--api-model-name": "vllm-model", + "--api-model-name": modelName, "--model-tokenizer": *baseModel.Storage.Path, } @@ -106,6 +113,93 @@ func BuildInferenceServiceArgs(ctx context.Context, c client.Client, endpointSpe return nil, fmt.Errorf("invalid EndpointSpec: both Endpoint and InferenceService are nil") } +// Flags the supported engines use to name the model they serve. When --served-model-name +// is absent both vLLM and SGLang fall back to the model location, so this code does too. +const ( + servedModelNameFlag = "--served-model-name" + vLLMModelFlag = "--model" + sgLangModelPathFlag = "--model-path" +) + +// resolveServedModelName returns the name the engine actually serves the model under. That +// name goes into the `model` field of every request genai-bench sends, and engines reject +// requests naming a model they do not serve. +// +// The value is read from the running engine pods rather than from the ServingRuntime or the +// InferenceService. --served-model-name can be set on either one: runtimes normally carry it, +// while an InferenceService may override the runner. Only the pod reflects the two merged, so +// reading either spec alone silently misses the other. +func resolveServedModelName(ctx context.Context, c client.Client, isvc *v1beta1.InferenceService) (string, error) { + pods := &v1.PodList{} + if err := c.List(ctx, pods, + client.InNamespace(isvc.Namespace), + client.MatchingLabels{ + constants.InferenceServicePodLabelKey: isvc.Name, + constants.OMEComponentLabel: string(v1beta1.EngineComponent), + }); err != nil { + return "", fmt.Errorf("failed to list engine pods of InferenceService %s/%s: %w", + isvc.Namespace, isvc.Name, err) + } + + var modelLocation string + for _, pod := range pods.Items { + for _, container := range pod.Spec.Containers { + argv := append(append([]string{}, container.Command...), container.Args...) + + if name := flagValue(argv, servedModelNameFlag); name != "" { + return name, nil + } + + // Remember the model location in case no container names the model explicitly. + if modelLocation != "" { + continue + } + for _, flag := range []string{vLLMModelFlag, sgLangModelPathFlag} { + if value := flagValue(argv, flag); value != "" { + modelLocation = expandContainerEnv(value, container.Env) + break + } + } + } + } + + if modelLocation != "" { + return modelLocation, nil + } + + return "", fmt.Errorf("cannot determine the served model name of InferenceService %s/%s: "+ + "no engine pod container specifies %s, %s or %s", + isvc.Namespace, isvc.Name, servedModelNameFlag, vLLMModelFlag, sgLangModelPathFlag) +} + +// flagValue returns the first value given to flag in argv, accepting both "--flag value" and +// "--flag=value". vLLM accepts several names for one model and reports the first back to +// clients, so the first is the one that matters here. +func flagValue(argv []string, flag string) string { + for i, arg := range argv { + if value, found := strings.CutPrefix(arg, flag+"="); found { + return value + } + if arg == flag && i+1 < len(argv) && !strings.HasPrefix(argv[i+1], "-") { + return argv[i+1] + } + } + return "" +} + +// expandContainerEnv resolves Kubernetes $(VAR) references against the container's own +// environment. OME runtime templates point the model flag at $(MODEL_PATH), so the raw +// argument is a placeholder rather than the value the engine sees. +func expandContainerEnv(value string, env []v1.EnvVar) string { + for _, envVar := range env { + if envVar.Value == "" { + continue + } + value = strings.ReplaceAll(value, "$("+envVar.Name+")", envVar.Value) + } + return value +} + // buildArgsFromEndpoint constructs the arguments map when an Endpoint is directly provided. func buildArgsFromEndpoint(endpoint *v1beta1.Endpoint) map[string]string { return map[string]string{ diff --git a/pkg/controller/v1beta1/benchmark/utils/utils_test.go b/pkg/controller/v1beta1/benchmark/utils/utils_test.go index e121ba893..cf722a932 100644 --- a/pkg/controller/v1beta1/benchmark/utils/utils_test.go +++ b/pkg/controller/v1beta1/benchmark/utils/utils_test.go @@ -9,9 +9,11 @@ import ( v1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" + "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/client/fake" "sigs.k8s.io/ome/pkg/apis/ome/v1beta1" + "sigs.k8s.io/ome/pkg/constants" "sigs.k8s.io/ome/pkg/utils/storage" ) @@ -523,3 +525,126 @@ func TestUpdateVolumeMounts(t *testing.T) { }) } } + +// enginePod builds an engine pod of isvcName whose single container runs the given argv. +func enginePod(name, namespace, isvcName string, command []string, env []v1.EnvVar) *v1.Pod { + return &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: name, + Namespace: namespace, + Labels: map[string]string{ + constants.InferenceServicePodLabelKey: isvcName, + constants.OMEComponentLabel: string(v1beta1.EngineComponent), + }, + }, + Spec: v1.PodSpec{ + Containers: []v1.Container{{Name: "ome-container", Command: command, Env: env}}, + }, + } +} + +func TestResolveServedModelName(t *testing.T) { + const ns = "default" + const isvcName = "test-isvc" + + isvc := &v1beta1.InferenceService{ + ObjectMeta: metav1.ObjectMeta{Name: isvcName, Namespace: ns}, + } + + tests := []struct { + name string + objects []client.Object + want string + wantErr bool + }{ + { + name: "first of several served model names wins", + objects: []client.Object{enginePod("p", ns, isvcName, []string{ + "python3", "--served-model-name", "primary", "alias", "--dtype", "float16", + }, nil)}, + want: "primary", + }, + { + name: "equals form is accepted", + objects: []client.Object{enginePod("p", ns, isvcName, []string{ + "python3", "--served-model-name=primary", "--dtype=float16", + }, nil)}, + want: "primary", + }, + { + name: "falls back to --model with $(VAR) expanded", + objects: []client.Object{enginePod("p", ns, isvcName, + []string{"python3", "--model", "$(MODEL_PATH)"}, + []v1.EnvVar{{Name: "MODEL_PATH", Value: "/mnt/models/qwen"}}, + )}, + want: "/mnt/models/qwen", + }, + { + name: "falls back to SGLang --model-path", + objects: []client.Object{enginePod("p", ns, isvcName, []string{ + "python3", "-m", "sglang.launch_server", "--model-path", "/mnt/models/llama", + }, nil)}, + want: "/mnt/models/llama", + }, + { + name: "a flag directly after --served-model-name is not its value", + objects: []client.Object{enginePod("p", ns, isvcName, []string{ + "python3", "--served-model-name", "--dtype", "float16", "--model", "/mnt/models/x", + }, nil)}, + want: "/mnt/models/x", + }, + { + name: "pods of other inference services are ignored", + objects: []client.Object{ + enginePod("other", ns, "another-isvc", []string{"python3", "--served-model-name", "wrong"}, nil), + enginePod("mine", ns, isvcName, []string{"python3", "--served-model-name", "right"}, nil), + }, + want: "right", + }, + { + name: "no engine pods is an error", + objects: nil, + wantErr: true, + }, + { + name: "engine pod naming no model is an error", + objects: []client.Object{enginePod("p", ns, isvcName, []string{ + "python3", "--host", "0.0.0.0", + }, nil)}, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + scheme := runtime.NewScheme() + _ = v1beta1.AddToScheme(scheme) + _ = v1.AddToScheme(scheme) + + c := fake.NewClientBuilder().WithScheme(scheme).WithObjects(tt.objects...).Build() + + got, err := resolveServedModelName(context.TODO(), c, isvc) + if tt.wantErr { + assert.Error(t, err) + return + } + assert.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestExpandContainerEnv(t *testing.T) { + env := []v1.EnvVar{ + {Name: "MODEL_PATH", Value: "/mnt/models/x"}, + {Name: "EMPTY", Value: ""}, + {Name: "FROM_REF", ValueFrom: &v1.EnvVarSource{}}, + } + + assert.Equal(t, "/mnt/models/x", expandContainerEnv("$(MODEL_PATH)", env)) + assert.Equal(t, "/mnt/models/x/sub", expandContainerEnv("$(MODEL_PATH)/sub", env)) + assert.Equal(t, "plain", expandContainerEnv("plain", env)) + // Unresolvable references are left as-is rather than blanked out. + assert.Equal(t, "$(UNKNOWN)", expandContainerEnv("$(UNKNOWN)", env)) + assert.Equal(t, "$(FROM_REF)", expandContainerEnv("$(FROM_REF)", env)) +}