Repository navigation
fix(operator): preserve argv in multinode vLLM and SGLang shell wrappers #17274
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,97 @@ | ||
| /* | ||
| * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
|
|
||
| package dynamo | ||
|
|
||
| import ( | ||
| "os" | ||
| "os/exec" | ||
| "path/filepath" | ||
| "strings" | ||
| "testing" | ||
|
|
||
| "github.com/ai-dynamo/dynamo/deploy/operator/api/v1alpha1" | ||
| commonconsts "github.com/ai-dynamo/dynamo/deploy/operator/internal/consts" | ||
| "github.com/stretchr/testify/require" | ||
| corev1 "k8s.io/api/core/v1" | ||
| ) | ||
|
|
||
| // TestMultinodeShellWrappersPreserveArgv runs each single-layer `sh -c` wrapper | ||
| // the operator renders for multinode vLLM and SGLang through a real shell and | ||
| // checks that the user's argv reaches the engine unchanged. | ||
| func TestMultinodeShellWrappersPreserveArgv(t *testing.T) { | ||
| t.Log("Stub python3 to print its argv NUL-delimited and ray to succeed, so the rendered wrapper can run locally") | ||
| binDir := t.TempDir() | ||
| require.NoError(t, os.WriteFile(filepath.Join(binDir, "python3"), []byte("#!/bin/sh\nprintf '%s\\000' \"$@\"\n"), 0o755)) | ||
| require.NoError(t, os.WriteFile(filepath.Join(binDir, "ray"), []byte("#!/bin/sh\nexit 0\n"), 0o755)) | ||
| t.Setenv("PATH", binDir+string(os.PathListSeparator)+os.Getenv("PATH")) | ||
| t.Setenv("GROVE_PCLQ_POD_INDEX", "0") | ||
|
|
||
| // Every value here is legal in a container args list and must survive the | ||
| // shell verbatim: a single quote, an empty string, whitespace, JSON, and | ||
| // characters the shell would otherwise expand. | ||
| userArgs := []string{ | ||
| "--model", "models/it's a model", tensorParallelSizeFlag, "8", | ||
| "--served-model-name", "", | ||
| "--chat-template", `{{ 'hello' }}`, | ||
| "--kv-transfer-config", `{"kv_connector": "NixlConnector", "kv_role": "$ROLE"}`, | ||
| "--literal", `glob*?[x]`, | ||
| "--whitespace", "tab\tand\nnewline", | ||
| } | ||
|
|
||
| tests := []struct { | ||
| name string | ||
| backend Backend | ||
| module string | ||
| role Role | ||
| executor string | ||
| }{ | ||
| {name: "vllm ray leader", backend: &VLLMBackend{}, module: "dynamo.vllm", role: RoleLeader, executor: "ray"}, | ||
| {name: "vllm mp grove worker", backend: &VLLMBackend{}, module: "dynamo.vllm", role: RoleWorker, executor: "mp"}, | ||
| {name: "sglang grove worker", backend: &SGLangBackend{}, module: "dynamo.sglang", role: RoleWorker}, | ||
| } | ||
| for _, tt := range tests { | ||
| t.Run(tt.name, func(t *testing.T) { | ||
| t.Log("Render the multinode wrapper for an argv-form python command") | ||
| container := &corev1.Container{ | ||
| Command: []string{"python3", "-m", tt.module}, | ||
| Args: append([]string(nil), userArgs...), | ||
| } | ||
| spec := &v1alpha1.DynamoComponentDeploymentSharedSpec{} | ||
| if tt.executor != "" { | ||
| spec.Annotations = map[string]string{commonconsts.KubeAnnotationVLLMDistributedExecutorBackend: tt.executor} | ||
| } | ||
| require.NoError(t, tt.backend.UpdateContainer(container, 2, tt.role, betaComponent(t, spec), "worker", &GroveMultinodeDeployer{}, staticContainerGPUCount(4))) | ||
| require.Equal(t, "-c", container.Command[len(container.Command)-1], "expected a sh -c wrapper, got %v", container.Command) | ||
|
|
||
| t.Log("Stand in for kubelet $(VAR) expansion, then execute the wrapper through the shell it targets") | ||
| kubelet := strings.NewReplacer("$(GROVE_PCSG_NAME)", "pcs", "$(GROVE_PCSG_INDEX)", "0", "$(GROVE_HEADLESS_SERVICE)", "svc") | ||
| script := kubelet.Replace(container.Args[0]) | ||
| out, err := exec.Command(container.Command[0], append(container.Command[1:], script)...).CombinedOutput() | ||
| require.NoError(t, err, "shell wrapper failed:\n%s\n%s", script, out) | ||
| got := strings.Split(strings.TrimSuffix(string(out), "\x00"), "\x00") | ||
|
|
||
| t.Log("The user's argv must lead the engine command unchanged, followed by the injected multinode flags") | ||
| want := append([]string{"-m", tt.module}, userArgs...) | ||
| require.GreaterOrEqual(t, len(got), len(want), "engine argv: %q", got) | ||
| require.Equal(t, want, got[:len(want)]) | ||
| if tt.role == RoleWorker { | ||
| require.Contains(t, got[len(want):], "1", "injected node rank should evaluate to 1, got %q", got[len(want):]) | ||
| } | ||
| }) | ||
| } | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -38,9 +38,12 @@ import ( | |
| * - Insert flags after the Python command but before any shell operators (|, &, ;) | ||
| */ | ||
|
|
||
| // shellQuoteForBashC quotes a string so it survives shell interpretation inside sh -c. | ||
| // Simple args (flags, paths) pass through unchanged; args containing special characters | ||
| // (JSON, env vars, spaces, quotes) are wrapped in double quotes with inner escaping. | ||
| // shellQuoteForBashC quotes a string for the nested `bash -c '...'` command the | ||
| // TensorRT-LLM mpirun launcher builds. Simple args (flags, paths) pass through | ||
| // unchanged; args containing special characters (JSON, env vars, spaces, quotes) | ||
| // are wrapped in double quotes with inner escaping, and single quotes are closed | ||
| // and reopened so they survive the outer single-quoted wrapper. It is not | ||
| // argv-preserving through a single sh -c layer: use shellQuotePOSIX there. | ||
| func shellQuoteForBashC(s string) string { | ||
| if strings.ContainsAny(s, " \t\n'\"\\{}[]$`!") { | ||
| escaped := s | ||
|
|
@@ -123,16 +126,16 @@ func injectFlagsIntoContainerCommand(container *corev1.Container, flags string, | |
| // Direct python command case | ||
| if needsShell { | ||
| // Transform to shell wrapper for env var interpretation. | ||
| // Quote each token individually so paths with spaces or special | ||
| // characters survive shell interpretation. | ||
| // Quote each token for a single sh -c layer so every argv value, | ||
| // including empty strings and embedded quotes, reaches the engine intact. | ||
| quotedCmd := make([]string, len(container.Command)) | ||
| for i, tok := range container.Command { | ||
| quotedCmd[i] = shellQuoteForBashC(tok) | ||
| quotedCmd[i] = shellQuotePOSIX(tok) | ||
| } | ||
| fullCommand := strings.Join(quotedCmd, " ") | ||
| quotedArgs := make([]string, len(container.Args)) | ||
| for i, arg := range container.Args { | ||
| quotedArgs[i] = shellQuoteForBashC(arg) | ||
| quotedArgs[i] = shellQuotePOSIX(arg) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🔎 Supported by static analysis
Keep kubelet-expanded arguments out of the shell script. Kubelet expands This is not a new root cause: the previous helper already corrupted this case by leaving a leading backslash in the substituted value. However, the new code still does not preserve kubelet-expanded arguments. Pass the original command and arguments as positional
|
||
| } | ||
| originalArgs := strings.Join(quotedArgs, " ") | ||
| var shellCommand 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.
This comment only expands the test name and narrates behavior already visible in the table and
exec.Commandcall; it preserves no non-obvious constraint.🤖 AI Fix
Remove this three-line comment.