diff --git a/modules/ollama/src/main/java/org/testcontainers/ollama/OllamaContainer.java b/modules/ollama/src/main/java/org/testcontainers/ollama/OllamaContainer.java index db611a34a39..3b14e05803d 100644 --- a/modules/ollama/src/main/java/org/testcontainers/ollama/OllamaContainer.java +++ b/modules/ollama/src/main/java/org/testcontainers/ollama/OllamaContainer.java @@ -3,15 +3,12 @@ import com.github.dockerjava.api.DockerClient; import com.github.dockerjava.api.model.DeviceRequest; import com.github.dockerjava.api.model.Image; -import com.github.dockerjava.api.model.Info; -import com.github.dockerjava.api.model.RuntimeInfo; import org.testcontainers.DockerClientFactory; import org.testcontainers.containers.GenericContainer; import org.testcontainers.utility.DockerImageName; import java.util.Collections; import java.util.List; -import java.util.Map; /** * Testcontainers implementation for Ollama. @@ -19,6 +16,8 @@ * Supported image: {@code ollama/ollama} *
* Exposed ports: 11434 + *
+ * GPU access is opt-in: use {@link #withGpuAccess()} to request access to the host GPU(s).
*/
public class OllamaContainer extends GenericContainer
+ * GPU access is opt-in and not enabled automatically: the presence of an {@code nvidia} runtime does not
+ * guarantee that a usable GPU is present. Notably, Docker Desktop (WSL2) registers such a runtime even on
+ * machines without NVIDIA GPUs, which would make the container fail to start if a GPU device request was
+ * applied automatically.
+ *
+ * @return this
+ */
+ public OllamaContainer withGpuAccess() {
+ withCreateContainerCmdModifier(cmd -> {
+ cmd
+ .getHostConfig()
+ .withDeviceRequests(
+ Collections.singletonList(
+ new DeviceRequest()
+ .withCapabilities(Collections.singletonList(Collections.singletonList("gpu")))
+ .withCount(-1)
+ )
+ );
+ });
+ return self();
+ }
+
/**
* Commits the current file system changes in the container into a new image.
* Should be used for creating an image that contains a loaded model.
diff --git a/modules/ollama/src/test/java/org/testcontainers/ollama/OllamaContainerGpuTest.java b/modules/ollama/src/test/java/org/testcontainers/ollama/OllamaContainerGpuTest.java
new file mode 100644
index 00000000000..528866b18cd
--- /dev/null
+++ b/modules/ollama/src/test/java/org/testcontainers/ollama/OllamaContainerGpuTest.java
@@ -0,0 +1,65 @@
+package org.testcontainers.ollama;
+
+import com.github.dockerjava.api.command.CreateContainerCmd;
+import com.github.dockerjava.api.model.DeviceRequest;
+import com.github.dockerjava.core.command.CreateContainerCmdImpl;
+import org.junit.jupiter.api.Test;
+import org.testcontainers.containers.GenericContainer;
+import org.testcontainers.core.CreateContainerCmdModifier;
+import org.testcontainers.utility.DockerImageName;
+
+import java.lang.reflect.Field;
+import java.util.Collections;
+import java.util.Set;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Unit tests for the opt-in GPU access of {@link OllamaContainer}. The tests do not start any container: they
+ * apply the registered {@link CreateContainerCmdModifier}s to a {@link CreateContainerCmd} instance and assert
+ * on the resulting HostConfig.
+ */
+class OllamaContainerGpuTest {
+
+ @Test
+ void shouldNotRequestGpuAccessByDefault() throws Exception {
+ try (OllamaContainer container = new OllamaContainer(DockerImageName.parse("ollama/ollama"))) {
+ CreateContainerCmd createCommand = applyCreateContainerCmdModifiers(container);
+
+ assertThat(createCommand.getHostConfig().getDeviceRequests()).isNull();
+ }
+ }
+
+ @Test
+ void shouldRequestGpuAccessWhenOptedIn() throws Exception {
+ try (OllamaContainer container = new OllamaContainer(DockerImageName.parse("ollama/ollama"))) {
+ OllamaContainer result = container.withGpuAccess();
+
+ assertThat(result).isSameAs(container);
+
+ CreateContainerCmd createCommand = applyCreateContainerCmdModifiers(container);
+ DeviceRequest deviceRequest = createCommand.getHostConfig().getDeviceRequests().get(0);
+
+ assertThat(createCommand.getHostConfig().getDeviceRequests()).hasSize(1);
+ assertThat(deviceRequest.getCapabilities()).containsExactly(Collections.singletonList("gpu"));
+ assertThat(deviceRequest.getCount()).isEqualTo(-1);
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ private static CreateContainerCmd applyCreateContainerCmdModifiers(OllamaContainer container) throws Exception {
+ CreateContainerCmd createCommand = new CreateContainerCmdImpl(
+ command -> null, // the command is never executed in this test
+ null,
+ "ollama/ollama"
+ );
+
+ Field modifiersField = GenericContainer.class.getDeclaredField("createContainerCmdModifiers");
+ modifiersField.setAccessible(true);
+ Set