diff --git a/backend/backend.proto b/backend/backend.proto index 01c5b63a7b5c..abbbdafd0137 100644 --- a/backend/backend.proto +++ b/backend/backend.proto @@ -445,6 +445,9 @@ message ModelOptions { // Unknown keys produce an error at LoadModel time. string EngineArgs = 73; + // EnvVars carries environment variables to be passed to the backend process. + map EnvVars = 75; + // Proxy carries the cloud-proxy backend's per-model configuration. // Empty for non-proxy backends. ProxyOptions Proxy = 74; diff --git a/core/backend/options.go b/core/backend/options.go index 8fc47f231e16..877d41a3fb42 100644 --- a/core/backend/options.go +++ b/core/backend/options.go @@ -187,6 +187,10 @@ func ModelOptions(c config.ModelConfig, so *config.ApplicationConfig, opts ...mo defOpts = append(defOpts, model.WithModelSizeBytes(sizeBytes)) } + if c.Environment != nil && len(c.Environment) > 0 { + defOpts = append(defOpts, model.WithEnvVars(c.Environment)) + } + return append(defOpts, opts...) } @@ -421,6 +425,14 @@ func grpcModelOpts(c config.ModelConfig, modelPath string) *pb.ModelOptions { opts.DraftModel = filepath.Join(modelPath, c.DraftModel) } + // Add environment variables from model configuration + if c.Environment != nil && len(c.Environment) > 0 { + opts.EnvVars = make(map[string]string) + for k, v := range c.Environment { + opts.EnvVars[k] = v + } + } + return opts } diff --git a/core/config/meta/registry.go b/core/config/meta/registry.go index 4fa555d65641..61df9914e51a 100644 --- a/core/config/meta/registry.go +++ b/core/config/meta/registry.go @@ -66,6 +66,13 @@ func DefaultRegistry() map[string]FieldMetaOverride { Options: UsecaseOptions, Order: 6, }, + "env": { + Section: "general", + Label: "Environment Variables", + Description: "Environment variables to be applied to the backend process", + Component: "map-editor", + Order: 7, + }, // --- LLM --- "context_size": { diff --git a/core/config/model_config.go b/core/config/model_config.go index 0038f4f8d118..b8b7fc573494 100644 --- a/core/config/model_config.go +++ b/core/config/model_config.go @@ -134,6 +134,9 @@ type ModelConfig struct { Proxy ProxyConfig `yaml:"proxy,omitempty" json:"proxy,omitempty"` MITM MITMModelConfig `yaml:"mitm,omitempty" json:"mitm,omitempty"` Limits LimitsConfig `yaml:"limits,omitempty" json:"limits,omitempty"` + + // Environment variables to set when starting the backend process + Environment map[string]string `yaml:"env,omitempty" json:"env,omitempty"` } // @Description Admission-control limits applied per request. The diff --git a/core/config/model_config_test.go b/core/config/model_config_test.go index 2f2f3fd82b42..d6e27bc5f73e 100644 --- a/core/config/model_config_test.go +++ b/core/config/model_config_test.go @@ -98,6 +98,12 @@ parameters: Expect(valid).To(BeTrue()) Expect(err).NotTo(HaveOccurred()) + // Test environment variables configuration parsing from YAML + envYAML := "name: env-test-model\nbackend: test-backend\nenv:\n TEST_ENV_VAR: test_value\n ANOTHER_VAR: another_value\n" + var envConfig ModelConfig + Expect(yaml.Unmarshal([]byte(envYAML), &envConfig)).To(Succeed()) + Expect(envConfig.Environment).To(HaveKeyWithValue("TEST_ENV_VAR", "test_value")) + Expect(envConfig.Environment).To(HaveKeyWithValue("ANOTHER_VAR", "another_value")) tcAndChat := FLAG_TOKEN_CLASSIFY | FLAG_CHAT tcCombined := ModelConfig{ Name: "ner-and-chat", diff --git a/core/services/worker/supervisor.go b/core/services/worker/supervisor.go index 5fda9bae8af2..dc3591dc09e1 100644 --- a/core/services/worker/supervisor.go +++ b/core/services/worker/supervisor.go @@ -81,7 +81,7 @@ func (s *backendSupervisor) startBackend(backend, backendPath string) (string, e bindAddr := fmt.Sprintf("0.0.0.0:%d", port) clientAddr := fmt.Sprintf("127.0.0.1:%d", port) - proc, err := s.ml.StartProcess(backendPath, backend, bindAddr) + proc, err := s.ml.StartProcess(backendPath, backend, bindAddr, nil) if err != nil { s.mu.Unlock() return "", fmt.Errorf("starting backend process: %w", err) diff --git a/docs/content/advanced/model-configuration.md b/docs/content/advanced/model-configuration.md index 8092c162ad98..a70c157429ce 100644 --- a/docs/content/advanced/model-configuration.md +++ b/docs/content/advanced/model-configuration.md @@ -909,6 +909,24 @@ Multiple detectors union their detections; overlapping spans resolve to the stro > The earlier regex pattern tier (`pii.patterns`, the global pattern catalogue, `--pii-config`, and the `/api/pii/patterns` admin endpoints) has been removed, along with response/streaming-side redaction. Those keys now no-op with a startup warning; migrate to `pii.detectors` + a detector's `pii_detection` block. +## Environment Variables Configuration + +Model configurations can specify environment variables passed to the backend process: + +```yaml +name: vllm-model +backend: vllm +parameters: + model: my-vllm-model + +env: + VLLM_WORKER_MULTIPROC_METHOD: "spawn" + VLLM_CACHE_DIR: "/tmp/vllm_cache" + CUDA_VISIBLE_DEVICES: "0,1" +``` + +Environment variables are appended to the system environment variables and will override any conflicting system variables with the same name. + ## Complete Example Here's a comprehensive example combining many options: diff --git a/pkg/model/initializers.go b/pkg/model/initializers.go index 43273a662946..d0ff76aeb4b4 100644 --- a/pkg/model/initializers.go +++ b/pkg/model/initializers.go @@ -119,7 +119,7 @@ func (ml *ModelLoader) spawnGRPCModel(backend, uri string, o *Options, modelID, return nil, fmt.Errorf("failed allocating free ports: %s", err.Error()) } // Make sure the process is executable - process, err := ml.startProcess(uri, modelID, serverAddress) + process, err := ml.StartProcess(uri, modelID, serverAddress, o.envVars) if err != nil { xlog.Error("failed to launch", "error", err, "path", uri) return nil, err diff --git a/pkg/model/loader_options.go b/pkg/model/loader_options.go index d247a00c8d3c..90f14ef13e7f 100644 --- a/pkg/model/loader_options.go +++ b/pkg/model/loader_options.go @@ -24,6 +24,9 @@ type Options struct { // by the caller using the vram estimation scaffolding. When non-zero it is // registered with the watchdog so size-aware eviction can rank models. modelSizeBytes int64 + + // envVars contains model-specific environment variables to pass to the backend process + envVars map[string]string } type Option func(*Options) @@ -97,6 +100,12 @@ func WithModelSizeBytes(bytes int64) Option { } } +func WithEnvVars(envVars map[string]string) Option { + return func(o *Options) { + o.envVars = envVars + } +} + func NewOptions(opts ...Option) *Options { o := &Options{ gRPCOptions: &pb.ModelOptions{}, diff --git a/pkg/model/process.go b/pkg/model/process.go index 8a796a4f5325..7896a2b4ea51 100644 --- a/pkg/model/process.go +++ b/pkg/model/process.go @@ -141,11 +141,11 @@ func (ml *ModelLoader) GetGRPCPID(id string) (int, error) { // StartProcess starts a gRPC backend process and returns its process handle. // This is the public wrapper for the internal startProcess method, used by // the serve-backend CLI subcommand to start a backend on a specified address. -func (ml *ModelLoader) StartProcess(grpcProcess, id string, serverAddress string, args ...string) (*process.Process, error) { - return ml.startProcess(grpcProcess, id, serverAddress, args...) +func (ml *ModelLoader) StartProcess(grpcProcess, id string, serverAddress string, envVars map[string]string, args ...string) (*process.Process, error) { + return ml.startProcess(grpcProcess, id, serverAddress, envVars, args...) } -func (ml *ModelLoader) startProcess(grpcProcess, id string, serverAddress string, args ...string) (*process.Process, error) { +func (ml *ModelLoader) startProcess(grpcProcess, id string, serverAddress string, envVars map[string]string, args ...string) (*process.Process, error) { // Make sure the process is executable // Check first if it has executable permissions if fi, err := os.Stat(grpcProcess); err == nil { @@ -175,6 +175,13 @@ func (ml *ModelLoader) startProcess(grpcProcess, id string, serverAddress string // and the GPU would silently fall back to CPU). No-op for other backends. env = append(env, vulkanICDEnv(workDir)...) + // Add model-specific environment variables + if envVars != nil { + for key, value := range envVars { + env = append(env, fmt.Sprintf("%s=%s", key, value)) + } + } + grpcControlProcess := process.New( process.WithTemporaryStateDir(), process.WithName(filepath.Base(grpcProcess)),