Skip to content

Commit a0a1616

Browse files
shukladivyanshcopybara-github
shukladivyansh
authored and
copybara-github
committed
feat: Add ContainerCodeExecutor
PiperOrigin-RevId: 795529787
1 parent e0602b3 commit a0a1616

File tree

2 files changed

+216
-0
lines changed

2 files changed

+216
-0
lines changed

‎core/pom.xml

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
<java-websocket.version>1.6.0</java-websocket.version>
4242
<jackson.version>2.19.0</jackson.version>
4343
<okhttp.version>4.12.0</okhttp.version>
44+
<docker-java.version>3.3.6</docker-java.version>
4445
</properties>
4546

4647
<dependencies>
@@ -62,6 +63,16 @@
6263
<groupId>com.google.cloud</groupId>
6364
<artifactId>google-cloud-aiplatform</artifactId>
6465
</dependency>
66+
<dependency>
67+
<groupId>com.github.docker-java</groupId>
68+
<artifactId>docker-java</artifactId>
69+
<version>${docker-java.version}</version>
70+
</dependency>
71+
<dependency>
72+
<groupId>com.github.docker-java</groupId>
73+
<artifactId>docker-java-transport-httpclient5</artifactId>
74+
<version>${docker-java.version}</version>
75+
</dependency>
6576
<dependency>
6677
<groupId> io.modelcontextprotocol.sdk</groupId>
6778
<artifactId>mcp</artifactId>

‎core/src/main/java/com/google/adk/codeexecutors/ContainerCodeExecutor.java

Lines changed: 205 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,205 @@
1+
/*
2+
* Copyright 2025 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may not in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
package com.google.adk.codeexecutors;
19+
20+
import com.github.dockerjava.api.DockerClient;
21+
import com.github.dockerjava.api.command.ExecCreateCmdResponse;
22+
import com.github.dockerjava.api.model.Container;
23+
import com.github.dockerjava.core.DefaultDockerClientConfig;
24+
import com.github.dockerjava.core.DockerClientBuilder;
25+
import com.github.dockerjava.core.command.ExecStartResultCallback;
26+
import com.google.adk.agents.InvocationContext;
27+
import com.google.adk.codeexecutors.CodeExecutionUtils.CodeExecutionInput;
28+
import com.google.adk.codeexecutors.CodeExecutionUtils.CodeExecutionResult;
29+
import java.io.ByteArrayOutputStream;
30+
import java.io.File;
31+
import java.io.IOException;
32+
import java.io.UncheckedIOException;
33+
import java.nio.charset.StandardCharsets;
34+
import java.nio.file.Paths;
35+
import java.util.Optional;
36+
import org.slf4j.Logger;
37+
import org.slf4j.LoggerFactory;
38+
39+
/** A code executor that uses a custom container to execute code. */
40+
public class ContainerCodeExecutor extends BaseCodeExecutor {
41+
private static final Logger logger = LoggerFactory.getLogger(ContainerCodeExecutor.class);
42+
private static final String DEFAULT_IMAGE_TAG = "adk-code-executor:latest";
43+
44+
private final Optional<String> baseUrl;
45+
private final String image;
46+
private final Optional<String> dockerPath;
47+
private final DockerClient dockerClient;
48+
private Container container;
49+
50+
/**
51+
* Initializes the ContainerCodeExecutor.
52+
*
53+
* @param baseUrl Optional. The base url of the user hosted Docker client.
54+
* @param image The tag of the predefined image or custom image to run on the container. Either
55+
* dockerPath or image must be set.
56+
* @param dockerPath The path to the directory containing the Dockerfile. If set, build the image
57+
* from the dockerfile path instead of using the predefined image. Either dockerPath or image
58+
* must be set.
59+
*/
60+
public ContainerCodeExecutor(
61+
Optional<String> baseUrl, Optional<String> image, Optional<String> dockerPath) {
62+
if (image.isEmpty() && dockerPath.isEmpty()) {
63+
throw new IllegalArgumentException(
64+
"Either image or dockerPath must be set for ContainerCodeExecutor.");
65+
}
66+
this.baseUrl = baseUrl;
67+
this.image = image.orElse(DEFAULT_IMAGE_TAG);
68+
this.dockerPath = dockerPath.map(p -> Paths.get(p).toAbsolutePath().toString());
69+
70+
if (baseUrl.isPresent()) {
71+
var config =
72+
DefaultDockerClientConfig.createDefaultConfigBuilder()
73+
.withDockerHost(baseUrl.get())
74+
.build();
75+
this.dockerClient = DockerClientBuilder.getInstance(config).build();
76+
} else {
77+
this.dockerClient = DockerClientBuilder.getInstance().build();
78+
}
79+
80+
initContainer();
81+
Runtime.getRuntime().addShutdownHook(new Thread(this::cleanupContainer));
82+
}
83+
84+
@Override
85+
public boolean stateful() {
86+
return false;
87+
}
88+
89+
@Override
90+
public boolean optimizeDataFile() {
91+
return false;
92+
}
93+
94+
@Override
95+
public CodeExecutionResult executeCode(
96+
InvocationContext invocationContext, CodeExecutionInput codeExecutionInput) {
97+
ByteArrayOutputStream stdout = new ByteArrayOutputStream();
98+
ByteArrayOutputStream stderr = new ByteArrayOutputStream();
99+
100+
ExecCreateCmdResponse execCreateCmdResponse =
101+
dockerClient
102+
.execCreateCmd(container.getId())
103+
.withAttachStdout(true)
104+
.withAttachStderr(true)
105+
.withCmd("python3", "-c", codeExecutionInput.code())
106+
.exec();
107+
try {
108+
dockerClient
109+
.execStartCmd(execCreateCmdResponse.getId())
110+
.exec(new ExecStartResultCallback(stdout, stderr))
111+
.awaitCompletion();
112+
} catch (InterruptedException e) {
113+
Thread.currentThread().interrupt();
114+
throw new RuntimeException("Code execution was interrupted.", e);
115+
}
116+
117+
return CodeExecutionResult.builder()
118+
.stdout(stdout.toString(StandardCharsets.UTF_8))
119+
.stderr(stderr.toString(StandardCharsets.UTF_8))
120+
.build();
121+
}
122+
123+
private void buildDockerImage() {
124+
if (dockerPath.isEmpty()) {
125+
throw new IllegalStateException("Docker path is not set.");
126+
}
127+
File dockerfile = new File(dockerPath.get());
128+
if (!dockerfile.exists()) {
129+
throw new UncheckedIOException(new IOException("Invalid Docker path: " + dockerPath.get()));
130+
}
131+
132+
logger.info("Building Docker image...");
133+
try {
134+
dockerClient.buildImageCmd(dockerfile).withTag(image).start().awaitCompletion();
135+
} catch (InterruptedException e) {
136+
Thread.currentThread().interrupt();
137+
throw new RuntimeException("Docker image build was interrupted.", e);
138+
}
139+
logger.info("Docker image: {} built.", image);
140+
}
141+
142+
private void verifyPythonInstallation() {
143+
ExecCreateCmdResponse execCreateCmdResponse =
144+
dockerClient.execCreateCmd(container.getId()).withCmd("which", "python3").exec();
145+
ByteArrayOutputStream stdout = new ByteArrayOutputStream();
146+
ByteArrayOutputStream stderr = new ByteArrayOutputStream();
147+
try (ExecStartResultCallback callback = new ExecStartResultCallback(stdout, stderr)) {
148+
dockerClient.execStartCmd(execCreateCmdResponse.getId()).exec(callback).awaitCompletion();
149+
} catch (InterruptedException e) {
150+
Thread.currentThread().interrupt();
151+
throw new RuntimeException("Python verification was interrupted.", e);
152+
} catch (IOException e) {
153+
throw new UncheckedIOException(e);
154+
}
155+
}
156+
157+
private void initContainer() {
158+
if (dockerClient == null) {
159+
throw new IllegalStateException("Docker client is not initialized.");
160+
}
161+
if (dockerPath.isPresent()) {
162+
buildDockerImage();
163+
} else {
164+
// If a dockerPath is not provided, always pull the image to ensure it's up-to-date.
165+
// If the image already exists locally, this will be a quick no-op.
166+
logger.info("Ensuring image {} is available locally...", image);
167+
try {
168+
dockerClient.pullImageCmd(image).start().awaitCompletion();
169+
} catch (InterruptedException e) {
170+
Thread.currentThread().interrupt();
171+
throw new RuntimeException("Docker image pull was interrupted.", e);
172+
}
173+
logger.info("Image {} is available.", image);
174+
}
175+
logger.info("Starting container for ContainerCodeExecutor...");
176+
var createContainerResponse =
177+
dockerClient.createContainerCmd(image).withTty(true).withAttachStdin(true).exec();
178+
dockerClient.startContainerCmd(createContainerResponse.getId()).exec();
179+
180+
var containers = dockerClient.listContainersCmd().withShowAll(true).exec();
181+
this.container =
182+
containers.stream()
183+
.filter(c -> c.getId().equals(createContainerResponse.getId()))
184+
.findFirst()
185+
.orElseThrow(() -> new IllegalStateException("Failed to find the created container."));
186+
187+
logger.info("Container {} started.", container.getId());
188+
verifyPythonInstallation();
189+
}
190+
191+
private void cleanupContainer() {
192+
if (container == null) {
193+
return;
194+
}
195+
logger.info("[Cleanup] Stopping the container...");
196+
dockerClient.stopContainerCmd(container.getId()).exec();
197+
dockerClient.removeContainerCmd(container.getId()).exec();
198+
logger.info("Container {} stopped and removed.", container.getId());
199+
try {
200+
dockerClient.close();
201+
} catch (IOException e) {
202+
logger.warn("Failed to close docker client", e);
203+
}
204+
}
205+
}

0 commit comments

Comments
 (0)