Add uid to docker config (#1167)

lambdazy · Mar 18, 2024 · 2ccb2b9 · 2ccb2b9
1 parent d6376d1
commit 2ccb2b9
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
@@ -17,6 +17,7 @@ public record DockerEnvDescription(
     @Nullable
     String networkMode,
     DockerClientConfig dockerClientConfig,
+    String user,
     Set<String> allowedPlatforms // In format os/arch like "linux/amd64". Empty means all are allowed
 ) {
 
@@ -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 user = ROOT_USER_UID;
         Set<String> allowedPlatforms = new HashSet<>();
 
         public Builder withName(String name) {
@@ -94,6 +98,11 @@ public Builder withDockerClientConfig(DockerClientConfig dockerClientConfig) {
             return this;
         }
 
+        public Builder withUser(String user) {
+            this.user = user;
+            return this;
+        }
+
         public Builder withAllowedPlatforms(Collection<String> allowedPlatforms) {
             this.allowedPlatforms.addAll(allowedPlatforms);
             return this;
@@ -103,8 +112,8 @@ public DockerEnvDescription build() {
             if (StringUtils.isBlank(name)) {
                 name = "job-" + RandomStringUtils.randomAlphanumeric(5);
             }
-            return new DockerEnvDescription(name, image, mounts, gpu, envVars, networkMode, dockerClientConfig,
-                    allowedPlatforms);
+            return new DockerEnvDescription(name, image, mounts, gpu, envVars, networkMode, dockerClientConfig, user,
+                allowedPlatforms);
         }
     }
 

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
@@ -40,7 +40,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";
     private static final String NO_MATCHING_MANIFEST_ERROR = "no matching manifest";
     private static final String NOT_MATCH_PLATFORM_ERROR = "was found but does not match the specified platform";
 
@@ -125,7 +124,7 @@ public void install(LogStream outStream, LogStream errStream) throws Environment
                 .withAttachStdout(true)
                 .withAttachStderr(true)
                 .withTty(true)
-                .withUser(ROOT_USER_UID)
+                .withUser(config.user())
                 .withEnv(config.envVars())
                 .withEntrypoint("/bin/sh")
                 .exec();
@@ -167,7 +166,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.user())
             .withAttachStdout(true)
             .withAttachStderr(true);