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 { @@ -34,26 +33,34 @@ public OllamaContainer(DockerImageName dockerImageName) { super(dockerImageName); dockerImageName.assertCompatibleWith(DOCKER_IMAGE_NAME); - Info info = this.dockerClient.infoCmd().exec(); - Map runtimes = info.getRuntimes(); - if (runtimes != null) { - if (runtimes.containsKey("nvidia")) { - withCreateContainerCmdModifier(cmd -> { - cmd - .getHostConfig() - .withDeviceRequests( - Collections.singletonList( - new DeviceRequest() - .withCapabilities(Collections.singletonList(Collections.singletonList("gpu"))) - .withCount(-1) - ) - ); - }); - } - } withExposedPorts(OLLAMA_PORT); } + /** + * Requests access to the host GPU(s) by adding a {@link DeviceRequest} to the container's create command. + *

+ * 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 modifiers = (Set) modifiersField.get(container); + for (CreateContainerCmdModifier modifier : modifiers) { + createCommand = modifier.modify(createCommand); + } + return createCommand; + } +}