Add container memory limit

lzy/execution-env/src/main/java/ai/lzy/env/Environment.java

@@ -53,8 +53,15 @@ public void signal(int sigValue) {}
 
         InputStream err();
 
-        int waitFor() throws InterruptedException;
+        int waitFor() throws InterruptedException, OomKilledException;
 
         void signal(int sigValue);
     }
+
+
+    class OomKilledException extends RuntimeException {
+        public OomKilledException(String msg) {
+            super(msg);
+        }
+    }
 }

lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvDescription.java

@@ -18,7 +18,9 @@ public record DockerEnvDescription(
     String networkMode,
     DockerClientConfig dockerClientConfig,
     String user,
-    Set<String> allowedPlatforms // In format os/arch like "linux/amd64". Empty means all are allowed
+    Set<String> allowedPlatforms, // In format os/arch like "linux/amd64". Empty means all are allowed
+    @Nullable
+    Long memLimitMb
 ) {
 
     public static Builder newBuilder() {
@@ -57,6 +59,7 @@ public static class Builder {
         DockerClientConfig dockerClientConfig;
         String user = ROOT_USER_UID;
         Set<String> allowedPlatforms = new HashSet<>();
+        Long memLimitMb = null;
 
         public Builder withName(String name) {
             this.name = name;
@@ -108,12 +111,17 @@ public Builder withAllowedPlatforms(Collection<String> allowedPlatforms) {
             return this;
         }
 
+        public Builder setMemLimitMb(Long memLimitMb) {
+            this.memLimitMb = memLimitMb;
+            return this;
+        }
+
         public DockerEnvDescription build() {
             if (StringUtils.isBlank(name)) {
                 name = "job-" + RandomStringUtils.randomAlphanumeric(5);
             }
             return new DockerEnvDescription(name, image, mounts, gpu, envVars, networkMode, dockerClientConfig, user,
-                allowedPlatforms);
+                allowedPlatforms, memLimitMb);
         }
     }
 

lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java

@@ -28,6 +28,7 @@
 import java.io.OutputStream;
 import java.io.PipedInputStream;
 import java.io.PipedOutputStream;
+import java.lang.management.ManagementFactory;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
@@ -114,6 +115,10 @@ public void install(LogStream outStream, LogStream errStream) throws Environment
             hostConfig.withNetworkMode(config.networkMode());
         }
 
+        if (config.memLimitMb() != null) {
+            hostConfig.withMemory(config.memLimitMb() * 1024 * 1024);
+        }
+
         AtomicInteger containerCreatingAttempt = new AtomicInteger(0);
         final var container = retry.executeSupplier(() -> {
             LOG.info("Creating container {}... (attempt {}); image: {}, config: {}",
@@ -243,11 +248,21 @@ public InputStream err() {
             }
 
             @Override
-            public int waitFor() throws InterruptedException {
+            public int waitFor() throws InterruptedException, OomKilledException {
                 try {
                     feature.get();
-                    return Math.toIntExact(retry.executeSupplier(() -> client.inspectExecCmd(exec.getId()).exec())
+                    var rc = Math.toIntExact(retry.executeSupplier(() -> client.inspectExecCmd(exec.getId()).exec())
                         .getExitCodeLong());
+
+                    if (rc == 0) {
+                        return 0;
+                    }
+
+                    if (isOomKilled()) {
+                        throw new OomKilledException("Process exited with rc %s, and it was killed by OOM killer".formatted(rc));
+                    }
+
+                    return rc;
                 } catch (InterruptedException e) {
                     try {
                         startCmd.close();
@@ -361,4 +376,20 @@ private void checkPlatform(InspectImageResponse inspectImageResponse, LogStream
             throw new RuntimeException(msg);
         }
     }
+
+    private boolean isOomKilled() {
+        if (containerId == null) {
+            return false;
+        }
+
+        var killed = client.inspectContainerCmd(containerId)
+            .exec()
+            .getState().getOOMKilled();
+
+        if (killed != null) {
+            return killed;
+        }
+
+        return false;
+    }
 }