Add uid to docker config

lambdazy · Mar 18, 2024 · 02c608d · 02c608d
1 parent 50e1686
commit 02c608d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvDescription.java b/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvDescription.java
@@ -18,7 +18,8 @@ public record DockerEnvDescription(
     List<String> envVars,  // In format <NAME>=<value>
     @Nullable
     String networkMode,
-    DockerClientConfig dockerClientConfig
+    DockerClientConfig dockerClientConfig,
+    String uid
 ) {
 
     public static Builder newBuilder() {
@@ -45,13 +46,16 @@ public record MountDescription(
     ) {}
 
     public static class Builder {
+        private static final String ROOT_USER_UID = "0";
+
         String name = null;
         String image = null;
         boolean gpu = false;
         List<MountDescription> mounts = new ArrayList<>();
         List<String> envVars = new ArrayList<>();
         String networkMode = null;
         DockerClientConfig dockerClientConfig;
+        String uid = ROOT_USER_UID;
 
         public Builder withName(String name) {
             this.name = name;
@@ -93,11 +97,16 @@ public Builder withDockerClientConfig(DockerClientConfig dockerClientConfig) {
             return this;
         }
 
+        public Builder withUID(String uid) {
+            this.uid = uid;
+            return this;
+        }
+
         public DockerEnvDescription build() {
             if (StringUtils.isBlank(name)) {
                 name = "job-" + RandomStringUtils.randomAlphanumeric(5);
             }
-            return new DockerEnvDescription(name, image, mounts, gpu, envVars, networkMode, dockerClientConfig);
+            return new DockerEnvDescription(name, image, mounts, gpu, envVars, networkMode, dockerClientConfig, uid);
         }
 
     }

diff --git a/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java b/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java
@@ -36,7 +36,6 @@ public class DockerEnvironment extends BaseEnvironment {
 
     private static final Logger LOG = LogManager.getLogger(DockerEnvironment.class);
     private static final long GB_AS_BYTES = 1073741824;
-    private static final String ROOT_USER_UID = "0";
 
     @Nullable
     public String containerId = null;
@@ -119,7 +118,7 @@ public void install(LogStream outStream, LogStream errStream) throws Environment
                 .withAttachStdout(true)
                 .withAttachStderr(true)
                 .withTty(true)
-                .withUser(ROOT_USER_UID)
+                .withUser(config.uid())
                 .withEnv(config.envVars())
                 .withEntrypoint("/bin/sh")
                 .exec();
@@ -161,7 +160,7 @@ public LzyProcess runProcess(String[] command, @Nullable String[] envp, @Nullabl
         LOG.info("Creating cmd {}", String.join(" ", command));
         final ExecCreateCmd execCmd = client.execCreateCmd(containerId)
             .withCmd(command)
-            .withUser(ROOT_USER_UID)
+            .withUser(config.uid())
             .withAttachStdout(true)
             .withAttachStderr(true);