Merge pull request #4302 from ollama/mxyng/forward-env

only forward some env vars
This commit is contained in:
Michael Yang 2024-05-09 16:21:05 -07:00 committed by GitHub
commit 0e1ba65855
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -292,31 +292,25 @@ func NewLlamaServer(gpus gpu.GpuInfoList, model string, ggml *GGML, adapters, pr
done: make(chan error, 1), done: make(chan error, 1),
} }
s.cmd.Env = os.Environ()
s.cmd.Stdout = os.Stdout s.cmd.Stdout = os.Stdout
s.cmd.Stderr = s.status s.cmd.Stderr = s.status
visibleDevicesEnv, visibleDevicesEnvVal := gpu.GpuInfoList(gpus).GetVisibleDevicesEnv() if v := strings.Join(libraryPaths, string(filepath.ListSeparator)); v != "" {
pathEnvVal := strings.Join(libraryPaths, string(filepath.ListSeparator)) s.cmd.Env = append(s.cmd.Env, pathEnv+"="+v)
}
// Update or add the path and visible devices variable with our adjusted version if k, v := gpu.GpuInfoList(gpus).GetVisibleDevicesEnv(); k != "" {
pathNeeded := true s.cmd.Env = append(s.cmd.Env, k+"="+v)
devicesNeeded := visibleDevicesEnv != ""
for i := range s.cmd.Env {
cmp := strings.SplitN(s.cmd.Env[i], "=", 2)
if strings.EqualFold(cmp[0], pathEnv) {
s.cmd.Env[i] = pathEnv + "=" + pathEnvVal
pathNeeded = false
} else if devicesNeeded && strings.EqualFold(cmp[0], visibleDevicesEnv) {
s.cmd.Env[i] = visibleDevicesEnv + "=" + visibleDevicesEnvVal
devicesNeeded = false
} }
for _, ev := range os.Environ() {
if strings.HasPrefix(ev, "CUDA_") ||
strings.HasPrefix(ev, "ROCM_") ||
strings.HasPrefix(ev, "HIP_") ||
strings.HasPrefix(ev, "HSA_") ||
strings.HasPrefix(ev, "GGML_") {
s.cmd.Env = append(s.cmd.Env, ev)
} }
if pathNeeded {
s.cmd.Env = append(s.cmd.Env, pathEnv+"="+pathEnvVal)
}
if devicesNeeded {
s.cmd.Env = append(s.cmd.Env, visibleDevicesEnv+"="+visibleDevicesEnvVal)
} }
slog.Info("starting llama server", "cmd", s.cmd.String()) slog.Info("starting llama server", "cmd", s.cmd.String())