[shim] Pass proxy variables to the container

runner/cmd/shim/main.go

@@ -27,6 +27,13 @@ import (
 // Version is a build-time variable. The value is overridden by ldflags.
 var Version string
 
+// https://everything.curl.dev/usingcurl/proxies/env.html
+// https://cs.opensource.google/go/x/net/+/657eb1317b5dd33038d683297c6be9cae05fa97d:http/httpproxy/proxy.go
+// We accept HTTP_PROXY in upper case without additional checks as it's unlikely that
+// the shim is running in the CGI context
+// The lower case form should be used as some applications ignore the upper case form, e.g., curl, apt
+const defaultPassEnv = "http_proxy,https_proxy,no_proxy,HTTP_PROXY,HTTPS_PROXY,NO_PROXY"
+
 func main() {
 	os.Exit(mainInner())
 }
@@ -147,6 +154,13 @@ func mainInner() int {
 				Sources:     cli.EnvVars("DSTACK_DCGM_ADDRESS"),
 			},
 			/* Docker Parameters */
+			&cli.StringFlag{
+				Name:        "pass-env",
+				Usage:       "Environment variables to pass on to the container, a comma-separated list of names",
+				Value:       defaultPassEnv,
+				Destination: &args.Docker.PassEnv,
+				Sources:     cli.EnvVars("DSTACK_DOCKER_PASS_ENV"),
+			},
 			&cli.BoolFlag{
 				Name:        "privileged",
 				Usage:       "Give extended privileges to the container",

runner/internal/shim/docker.go

@@ -135,6 +135,7 @@ type DockerRunner struct {
 	client       *docker.Client
 	dockerParams DockerParameters
 	dockerInfo   dockersystem.Info
+	baseEnv      []string
 	gpus         []host.GpuInfo
 	gpuVendor    gpu.GpuVendor
 	gpuLock      *GpuLock
@@ -151,6 +152,15 @@ func NewDockerRunner(ctx context.Context, dockerParams DockerParameters) (*Docke
 		return nil, fmt.Errorf("get docker info: %w", err)
 	}
 
+	// Copy variables once rather than on a per-task basis
+	// We don't expect variables to change during the shim's lifetime
+	baseEnv := []string{}
+	for _, name := range dockerParams.DockerPassEnv() {
+		if value, ok := os.LookupEnv(name); ok {
+			baseEnv = append(baseEnv, fmt.Sprintf("%s=%s", name, value))
+		}
+	}
+
 	var gpuVendor gpu.GpuVendor
 	gpus := host.GetGpuInfo(ctx)
 	if len(gpus) > 0 {
@@ -167,6 +177,7 @@ func NewDockerRunner(ctx context.Context, dockerParams DockerParameters) (*Docke
 		client:       client,
 		dockerParams: dockerParams,
 		dockerInfo:   dockerInfo,
+		baseEnv:      baseEnv,
 		gpus:         gpus,
 		gpuVendor:    gpuVendor,
 		gpuLock:      gpuLock,
@@ -861,8 +872,9 @@ func (d *DockerRunner) createContainer(ctx context.Context, task *Task) error {
 
 	// Set the environment variables
 	envVars := []string{}
-	if d.dockerParams.DockerPJRTDevice() != "" {
-		envVars = append(envVars, fmt.Sprintf("PJRT_DEVICE=%s", d.dockerParams.DockerPJRTDevice()))
+	envVars = append(envVars, d.baseEnv...)
+	if pjrtDevice := d.dockerParams.DockerPJRTDevice(); pjrtDevice != "" {
+		envVars = append(envVars, fmt.Sprintf("PJRT_DEVICE=%s", pjrtDevice))
 	}
 
 	// Override /dev/shm with tmpfs mount with `exec` option (the default is `noexec`)
@@ -1235,6 +1247,16 @@ func getContainerLastLogs(ctx context.Context, client docker.APIClient, containe
 
 /* DockerParameters interface implementation for CLIArgs */
 
+func (c *CLIArgs) DockerPassEnv() []string {
+	names := []string{}
+	for _, name := range strings.Split(c.Docker.PassEnv, ",") {
+		if name = strings.TrimSpace(name); name != "" {
+			names = append(names, name)
+		}
+	}
+	return names
+}
+
 func (c *CLIArgs) DockerPrivileged() bool {
 	return c.Docker.Privileged
 }

runner/internal/shim/docker_test.go

@@ -119,6 +119,10 @@ func (c *dockerParametersMock) DockerShellCommands(authorizedKeys []string, runn
 	return commands
 }
 
+func (c *dockerParametersMock) DockerPassEnv() []string {
+	return []string{}
+}
+
 func (c *dockerParametersMock) DockerPorts() []int {
 	return []int{}
 }

runner/internal/shim/models.go

@@ -5,6 +5,7 @@ import (
 )
 
 type DockerParameters interface {
+	DockerPassEnv() []string
 	DockerPrivileged() bool
 	DockerShellCommands(authorizedKeys []string, runnerHttpAddress string) []string
 	DockerMounts(string) ([]mount.Mount, error)
@@ -40,6 +41,7 @@ type CLIArgs struct {
 	}
 
 	Docker struct {
+		PassEnv    string
 		Privileged bool
 		PJRTDevice string
 	}