feat(core): Add support for ollama module

- Added a new class OllamaContainer with few methods to handle the Ollama container.

- The `_check_and_add_gpu_capabilities` method checks if the host has GPUs and adds the necessary capabilities to the container.

- The `commit_to_image` allows to save somehow the state of a container into an image so that we can reuse it, especially for the ones having some models pulled.
- Added tests to check the functionality of the new class.

Author: bricefotzo

modules/ollama/README.rst

@@ -0,0 +1,2 @@
+.. autoclass:: testcontainers.ollama.OllamaContainer
+.. title:: testcontainers.ollama.OllamaContainer

modules/ollama/testcontainers/ollama/__init__.py

@@ -0,0 +1,75 @@
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+import docker
+from docker.types.containers import DeviceRequest
+
+from testcontainers.core.generic import ServerContainer
+from testcontainers.core.waiting_utils import wait_container_is_ready
+
+
+class OllamaContainer(ServerContainer):
+    def __init__(self, image="ollama/ollama:latest"):
+        self.port = 11434
+        super().__init__(port=self.port, image=image)
+        self.with_exposed_ports(self.port)
+        self._check_and_add_gpu_capabilities()
+
+    def _check_and_add_gpu_capabilities(self):
+        info = docker.APIClient().info()
+        if "nvidia" in info["Runtimes"]:
+            self.with_kwargs(device_requests=DeviceRequest(count=-1, capabilities=[["gpu"]]))
+
+    def start(self) -> "OllamaContainer":
+        """
+        Start the Ollama server
+        """
+        super().start()
+        wait_container_is_ready(self, "Ollama started successfully")
+        self._connect()
+
+        return self
+
+    def get_endpoint(self):
+        """
+        Return the endpoint of the Ollama server
+        """
+        return self._create_connection_url()
+
+    @property
+    def id(self) -> str:
+        """
+        Return the container object
+        """
+        return self._container.id
+
+    def pull_model(self, model_name: str) -> None:
+        """
+        Pull a model from the Ollama server
+
+        Args:
+            model_name (str): Name of the model
+        """
+        self.exec(f"ollama pull {model_name}")
+
+    def commit_to_image(self, image_name: str) -> None:
+        """
+        Commit the current container to a new image
+
+        Args:
+            image_name (str): Name of the new image
+        """
+        docker_client = self.get_docker_client()
+        existing_images = docker_client.client.images.list(name=image_name)
+        if not existing_images and self.id:
+            docker_client.client.containers.get(self.id).commit(repository=image_name, conf={"Labels": {"org.testcontainers.session-id": ""}})

modules/ollama/tests/test_ollama.py

@@ -0,0 +1,46 @@
+import requests
+from testcontainers.ollama import OllamaContainer
+import random
+import string
+
+
+def random_string(length=6):
+    return "".join(random.choices(string.ascii_lowercase, k=length))
+
+
+def test_ollama_container():
+    with OllamaContainer() as ollama:
+        url = ollama.get_endpoint()
+        response = requests.get(url)
+        assert response.status_code == 200
+        assert response.text == "Ollama is running"
+
+
+def test_with_default_config():
+    with OllamaContainer("ollama/ollama:0.1.26") as ollama:
+        ollama.start()
+        response = requests.get(f"{ollama.get_endpoint()}/api/version")
+        version = response.json().get("version")
+        assert version == "0.1.26"
+
+
+def test_download_model_and_commit_to_image():
+    new_image_name = f"tc-ollama-allminilm-{random_string(length=4).lower()}"
+    with OllamaContainer("ollama/ollama:0.1.26") as ollama:
+        ollama.start()
+        # Pull the model
+        ollama.pull_model("all-minilm")
+
+        response = requests.get(f"{ollama.get_endpoint()}/api/tags")
+        model_name = response.json().get("models", [])[0].get("name")
+        assert "all-minilm" in model_name
+
+        # Commit the container state to a new image
+        ollama.commit_to_image(new_image_name)
+
+    # Verify the new image
+    with OllamaContainer(new_image_name) as ollama:
+        ollama.start()
+        response = requests.get(f"{ollama.get_endpoint()}/api/tags")
+        model_name = response.json().get("models", [])[0].get("name")
+        assert "all-minilm" in model_name

poetry.lock

@@ -51,6 +51,7 @@ packages = [
     { include = "testcontainers", from = "modules/nats" },
     { include = "testcontainers", from = "modules/neo4j" },
     { include = "testcontainers", from = "modules/nginx" },
+    { include = "testcontainers", from = "modules/ollama" },
     { include = "testcontainers", from = "modules/opensearch" },
     { include = "testcontainers", from = "modules/oracle-free" },
     { include = "testcontainers", from = "modules/postgres" },
@@ -127,6 +128,7 @@ nats = ["nats-py"]
 neo4j = ["neo4j"]
 nginx = []
 opensearch = ["opensearch-py"]
+ollama = []
 oracle = ["sqlalchemy", "oracledb"]
 oracle-free = ["sqlalchemy", "oracledb"]
 postgres = []
@@ -272,6 +274,7 @@ mypy_path = [
 #    "modules/mysql",
 #    "modules/neo4j",
 #    "modules/nginx",
+#    "modules/ollama",
 #    "modules/opensearch",
 #    "modules/oracle",
 #    "modules/postgres",