baalajimaestro/ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

only forward some env vars

Browse source
This commit is contained in:
Michael Yang 2024-05-09 15:11:43 -07:00
parent 58876091f7
commit ce3b212d12
1 changed files with 14 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
llm/server.go
View file

@ -292,32 +292,26 @@ 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) for _, ev := range os.Environ() {
if strings.EqualFold(cmp[0], pathEnv) { if strings.HasPrefix(ev, "CUDA_") ||
s.cmd.Env[i] = pathEnv + "=" + pathEnvVal strings.HasPrefix(ev, "ROCM_") ||
pathNeeded = false strings.HasPrefix(ev, "HIP_") ||
} else if devicesNeeded && strings.EqualFold(cmp[0], visibleDevicesEnv) { strings.HasPrefix(ev, "HSA_") ||
s.cmd.Env[i] = visibleDevicesEnv + "=" + visibleDevicesEnvVal strings.HasPrefix(ev, "GGML_") {
devicesNeeded = false 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())
// Log at debug as the environment is inherited and might contain sensitive information // Log at debug as the environment is inherited and might contain sensitive information