Add DockerModelRunnerContainer to core

Author: kiview

core/build.gradle

@@ -123,6 +123,9 @@ dependencies {
     testImplementation 'org.assertj:assertj-core:3.26.3'
     testImplementation 'io.rest-assured:rest-assured:5.5.0'
 
+    // DockerModelRunnerContainer tests
+    testImplementation("com.openai:openai-java:1.3.0")
+
     jarFileTestCompileOnly "org.projectlombok:lombok:${lombok.version}"
     jarFileTestAnnotationProcessor "org.projectlombok:lombok:${lombok.version}"
     jarFileTestImplementation 'junit:junit:4.13.2'

core/src/main/java/org/testcontainers/containers/DockerModelRunnerContainer.java

@@ -0,0 +1,71 @@
+package org.testcontainers.containers;
+
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStreamReader;
+import java.io.OutputStream;
+import java.net.HttpURLConnection;
+import java.net.URL;
+import java.nio.charset.StandardCharsets;
+
+public class DockerModelRunnerContainer extends GenericContainer {
+
+    public static final String MODEL_RUNNER_ENDPOINT = "model-runner.docker.internal";
+    private SocatContainer socat;
+    private String model;
+
+    @Override
+    public void start() {
+        socat = new SocatContainer()
+            .withTarget(80, MODEL_RUNNER_ENDPOINT, 80);
+        socat.start();
+        pullModel();
+    }
+
+    private void pullModel() {
+        logger().info("Pulling model: {}. Please be patient, no progress bar yet!", model);
+        try {
+            // Construct JSON payload
+            String json = String.format("{\"from\":\"%s\"}", model);
+            String endpoint = "http://" + socat.getHost() + ":" + socat.getMappedPort(80) + "/models/create";
+
+            URL url = new URL(endpoint);
+            HttpURLConnection connection = (HttpURLConnection) url.openConnection();
+            connection.setRequestMethod("POST");
+            connection.setRequestProperty("Content-Type", "application/json");
+            connection.setDoOutput(true);
+
+            try (OutputStream os = connection.getOutputStream()) {
+                byte[] input = json.getBytes(StandardCharsets.UTF_8);
+                os.write(input, 0, input.length);
+            }
+
+            try (BufferedReader br = new BufferedReader(new InputStreamReader(connection.getInputStream(), StandardCharsets.UTF_8))) {
+                StringBuilder response = new StringBuilder();
+                String responseLine;
+                while ((responseLine = br.readLine()) != null) {
+                    response.append(responseLine.trim());
+                }
+                logger().info(response.toString());
+            }
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+        logger().info("Finished pulling model: {}", model);
+    }
+
+    @Override
+    public void stop() {
+        socat.stop();
+    }
+
+    public String getOpenAIEndpoint() {
+        return "http://" + socat.getHost() + ":" + socat.getMappedPort(80) + "/engines";
+    }
+
+    public DockerModelRunnerContainer withModel(String modelName) {
+        this.model = modelName;
+        return this;
+    }
+
+}

core/src/test/java/org/testcontainers/containers/DockerModelRunnerContainerTest.java

@@ -0,0 +1,35 @@
+package org.testcontainers.containers;
+
+import com.openai.client.OpenAIClient;
+import com.openai.client.okhttp.OpenAIOkHttpClient;
+import com.openai.models.chat.completions.ChatCompletion;
+import com.openai.models.chat.completions.ChatCompletionCreateParams;
+import org.junit.Test;
+
+public class DockerModelRunnerContainerTest {
+
+    @Test
+    public void pullsModelAndExposesInference() {
+        String modelName = "ai/smollm2:360M-Q4_K_M";
+
+        try (DockerModelRunnerContainer dmr = new DockerModelRunnerContainer()
+            .withModel(modelName);) {
+            dmr.start();
+
+            OpenAIClient client = OpenAIOkHttpClient.builder()
+                .baseUrl(dmr.getOpenAIEndpoint())
+                .build();
+
+            ChatCompletionCreateParams params = ChatCompletionCreateParams.builder()
+                .addUserMessage("Say this is a test")
+                .model(modelName)
+                .build();
+            ChatCompletion chatCompletion = client.chat().completions().create(params);
+
+            String answer = chatCompletion.toString();
+            System.out.println(answer);
+
+        }
+    }
+
+}