Skip to content

Commit

Permalink
Revert "only forward some env vars"
Browse files Browse the repository at this point in the history
This reverts commit ce3b212.
  • Loading branch information
jmorganca committed May 11, 2024
1 parent 1e1634d commit 92ca2cc
Showing 1 changed file with 20 additions and 14 deletions.
34 changes: 20 additions & 14 deletions llm/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -288,25 +288,31 @@ func NewLlamaServer(gpus gpu.GpuInfoList, model string, ggml *GGML, adapters, pr
done: make(chan error, 1),
}

s.cmd.Env = os.Environ()
s.cmd.Stdout = os.Stdout
s.cmd.Stderr = s.status

if v := strings.Join(libraryPaths, string(filepath.ListSeparator)); v != "" {
s.cmd.Env = append(s.cmd.Env, pathEnv+"="+v)
visibleDevicesEnv, visibleDevicesEnvVal := gpu.GpuInfoList(gpus).GetVisibleDevicesEnv()
pathEnvVal := strings.Join(libraryPaths, string(filepath.ListSeparator))

// Update or add the path and visible devices variable with our adjusted version
pathNeeded := true
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
}
}

if k, v := gpu.GpuInfoList(gpus).GetVisibleDevicesEnv(); k != "" {
s.cmd.Env = append(s.cmd.Env, k+"="+v)
if pathNeeded {
s.cmd.Env = append(s.cmd.Env, pathEnv+"="+pathEnvVal)
}

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 devicesNeeded {
s.cmd.Env = append(s.cmd.Env, visibleDevicesEnv+"="+visibleDevicesEnvVal)
}

slog.Info("starting llama server", "cmd", s.cmd.String())
Expand Down

0 comments on commit 92ca2cc

Please sign in to comment.