[Bugfix] benchmark: read served model name from engine pods #786
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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": 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 | ||
|
Collaborator
There was a problem hiding this comment. nit. Three bundled runtimes pass it through an env var — - --served-model-name
- $(SERVED_MODEL_NAME)
env:
- name: SERVED_MODEL_NAME
value: "vllm-model"
if name := flagValue(argv, servedModelNameFlag); name != "" {
return expandContainerEnv(name, container.Env), nil
}Worth a test too — the existing |
||
| } | ||
|
|
||
| // 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) | ||
| } | ||
|
Comment on lines
+193
to
+199
There was a problem hiding this comment. 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win 🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -eu
printf '%s\n' '--- file outline ---'
ast-grep outline pkg/controller/v1beta1/benchmark/utils/utils.go --view expanded
printf '%s\n' '--- target implementation ---'
sed -n '100,215p' pkg/controller/v1beta1/benchmark/utils/utils.go
printf '%s\n' '--- related symbols and tests ---'
rg -n --glob '*.go' 'expandContainerEnv|resolveServedModelName|MODEL_ROOT|MODEL_PATH|ContainerCommand|ContainerArgs|Expand' .
printf '%s\n' '--- changed-file summary ---'
git diff --statRepository: ome-projects/ome Length of output: 11602 🏁 Script executed: #!/bin/bash
set -eu
python3 - <<'PY'
from pathlib import Path
p = Path("pkg/controller/v1beta1/benchmark/utils/utils.go")
lines = p.read_text().splitlines()
for start, end in [(1, 105), (175, 215)]:
print(f"--- {p}:{start}-{end} ---")
for n in range(start, min(end, len(lines)) + 1):
print(f"{n:4}: {lines[n-1]}")
PY
printf '%s\n' '--- repository tests near the utility package ---'
find pkg/controller/v1beta1/benchmark/utils -maxdepth 2 -type f -print
rg -n --glob '*.go' 'expandContainerEnv|Resolve.*Model|model.*flag|env expansion|command.*args' pkg/controller/v1beta1/benchmark pkg/controllerRepository: ome-projects/ome Length of output: 9924 🏁 Script executed: #!/bin/bash
set -eu
printf '%s\n' '--- Kubernetes expansion references in repository ---'
rg -n --glob '*.{go,yaml,yml,md}' '\$\(|\$\$|ExpandContainer|envFrom|EnvVar' pkg config deploy charts 2>/dev/null || true
printf '%s\n' '--- exact diff for target file ---'
git diff -- pkg/controller/v1beta1/benchmark/utils/utils.go
printf '%s\n' '--- all utility-package test files and relevant test bodies ---'
for f in $(find pkg/controller/v1beta1/benchmark/utils -maxdepth 1 -type f -name '*_test.go' -print); do
echo "--- $f ---"
sed -n '1,280p' "$f"
doneRepository: ome-projects/ome Length of output: 50373 🏁 Script executed: #!/bin/bash
set -eu
sed -n '1,240p' pkg/controller/v1beta1/benchmark/utils/utils.go
printf '\n--- tests ---\n'
find pkg/controller/v1beta1/benchmark/utils -type f -name '*_test.go' -print -exec sed -n '1,320p' {} \;
printf '\n--- usages ---\n'
rg -n --glob '*.go' 'expandContainerEnv|resolveServedModelName' .Repository: ome-projects/ome Length of output: 18454 🏁 Script executed: #!/bin/bash
set -eu
printf '%s\n' '--- focused resolver tests ---'
sed -n '480,670p' pkg/controller/v1beta1/benchmark/utils/utils_test.go
printf '%s\n' '--- representative runtime env and command sections ---'
for f in config/runtimes/srt/mistralai/mistral-7b-instruct-v0-3-rt.yaml config/runtimes/vllm/llama-3-2-90b-vision-instruct-rt.yaml charts/ome-serving/templates/clusterservingruntime.yaml; do
if [ -f "$f" ]; then
echo "--- $f ---"
rg -n -C 8 'MODEL_PATH|command:|args:|env:' "$f" | head -220
fi
done
printf '%s\n' '--- Kubernetes API dependency/version ---'
rg -n 'k8s.io/api' go.mod go.sum | head -20Repository: ome-projects/ome Length of output: 12241 🏁 Script executed: #!/bin/bash
set -eu
python3 - <<'PY'
import re
def current(value, env):
for name, raw in env:
if raw:
value = value.replace(f"$({name})", raw)
return value
def kube_command(value, env):
# Minimal read-only model of the documented command/args expansion:
# process env declarations in order, then expand command/args with $$ escapes.
values = {}
for name, raw in env:
values[name] = expand(raw, values)
return expand(value, values)
def expand(value, values):
sentinel = "\0"
value = value.replace("$$", sentinel)
value = re.sub(r"\$\(([A-Za-z_][A-Za-z0-9_]*)\)",
lambda m: values.get(m.group(1), m.group(0)), value)
return value.replace(sentinel, "$")
env = [
("MODEL_ROOT", "/mnt"),
("MODEL_PATH", "$(MODEL_ROOT)/model"),
]
cases = [
("$(MODEL_PATH)", "/mnt/model"),
("$$(MODEL_PATH)", "$(MODEL_PATH)"),
]
for value, expected in cases:
got_current = current(value, env)
got_kube = kube_command(value, env)
print(f"{value!r}: current={got_current!r}, documented={got_kube!r}, expected={expected!r}")
assert got_kube == expected
assert current("$(MODEL_PATH)", env) != "/mnt/model"
assert current("$$(MODEL_PATH)", env) != "$(MODEL_PATH)"
PYRepository: ome-projects/ome Length of output: 352 🏁 Script executed: #!/bin/bash
set -eu
printf '%s\n' '--- focused test range ---'
sed -n '530,665p' pkg/controller/v1beta1/benchmark/utils/utils_test.go
printf '%s\n' '--- concrete runtime snippets ---'
for f in config/runtimes/srt/mistralai/mistral-7b-instruct-v0-3-rt.yaml config/runtimes/vllm/llama-3-2-90b-vision-instruct-rt.yaml charts/ome-serving/templates/clusterservingruntime.yaml; do
[ -f "$f" ] || continue
echo "--- $f ---"
rg -n -C 5 'MODEL_PATH|command:|args:|env:' "$f" | head -160
doneRepository: ome-projects/ome Length of output: 8169 🌐 Web query:
💡 Result: In Kubernetes, the command and args fields support variable expansion using the Citations:
Match Kubernetes environment expansion semantics.
Resolve direct environment values in declaration order, then apply Kubernetes-compatible 🤖 Prompt for AI Agents |
||
| 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{ | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
Repository: ome-projects/ome
Length of output: 50372
🏁 Script executed:
Repository: ome-projects/ome
Length of output: 17152
🏁 Script executed:
Repository: ome-projects/ome
Length of output: 50372
🏁 Script executed:
Repository: ome-projects/ome
Length of output: 50373
🏁 Script executed:
Repository: ome-projects/ome
Length of output: 14701
🏁 Script executed:
Repository: ome-projects/ome
Length of output: 50372
🏁 Script executed:
Repository: ome-projects/ome
Length of output: 50372
🏁 Script executed:
Repository: ome-projects/ome
Length of output: 242
🏁 Script executed:
Repository: ome-projects/ome
Length of output: 244
Select a serving engine pod.
Skip pods with
DeletionTimestampset,Status.Phase != v1.PodRunning, orPodReady != Truebefore scanning containers. The current selector includes old and replacement pods, so the first match can return a stale model name. Add a regression test with a stale pod and a current ready pod.🤖 Prompt for AI Agents