Add ollama module (#778)

packages/modules/ollama/jest.config.ts

@@ -0,0 +1,11 @@
+import type { Config } from "jest";
+import * as path from "path";
+
+const config: Config = {
+  preset: "ts-jest",
+  moduleNameMapper: {
+    "^testcontainers$": path.resolve(__dirname, "../../testcontainers/src"),
+  },
+};
+
+export default config;

packages/modules/ollama/package.json

@@ -0,0 +1,34 @@
+{
+  "name": "@testcontainers/ollama",
+  "version": "10.7.2",
+  "license": "MIT",
+  "keywords": [
+    "ollama",
+    "testing",
+    "docker",
+    "testcontainers"
+  ],
+  "description": "Ollama module for Testcontainers",
+  "homepage": "https://github.com/testcontainers/testcontainers-node#readme",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/testcontainers/testcontainers-node"
+  },
+  "bugs": {
+    "url": "https://github.com/testcontainers/testcontainers-node/issues"
+  },
+  "main": "build/index.js",
+  "files": [
+    "build"
+  ],
+  "publishConfig": {
+    "access": "public"
+  },
+  "scripts": {
+    "prepack": "shx cp ../../../README.md . && shx cp ../../../LICENSE .",
+    "build": "tsc --project tsconfig.build.json"
+  },
+  "dependencies": {
+    "testcontainers": "^10.7.2"
+  }
+}

packages/modules/ollama/src/index.ts

@@ -0,0 +1 @@
+export { OllamaContainer, StartedOllamaContainer } from "./ollama-container";

packages/modules/ollama/src/ollama-container.test.ts

@@ -0,0 +1,35 @@
+import { OllamaContainer } from "./ollama-container";
+
+describe("OllamaContainer", () => {
+  jest.setTimeout(180_000);
+
+  it("should run ollama with default config", async () => {
+    const container = await new OllamaContainer("ollama/ollama:0.1.44").start();
+    const response = await fetch(`${container.getEndpoint()}/api/version`);
+    expect(response.status).toEqual(200);
+    const body = await response.json();
+    expect(body.version).toEqual("0.1.44");
+    await container.stop();
+  });
+
+  it("download model and commit to image", async () => {
+    const container = await new OllamaContainer("ollama/ollama:0.1.44").start();
+    const execResult = await container.exec(["ollama", "pull", "all-minilm"]);
+    console.log(execResult.output);
+    const response = await fetch(`${container.getEndpoint()}/api/tags`);
+    expect(response.status).toEqual(200);
+    const body = await response.json();
+    expect(body.models[0].name).toContain("all-minilm");
+
+    const newImageName: string = "tc-ollama-allminilm-" + (Math.random() + 1).toString(36).substring(4).toLowerCase();
+    await container.commitToImage(newImageName);
+
+    const newContainer = await new OllamaContainer(newImageName).start();
+    const response2 = await fetch(`${newContainer.getEndpoint()}/api/tags`);
+    expect(response2.status).toEqual(200);
+    const body2 = await response2.json();
+    expect(body2.models[0].name).toContain("all-minilm");
+    await container.stop();
+    await newContainer.stop();
+  });
+});

packages/modules/ollama/src/ollama-container.ts

@@ -0,0 +1,54 @@
+import {
+  AbstractStartedContainer,
+  GenericContainer,
+  getContainerRuntimeClient,
+  StartedTestContainer,
+  Wait,
+} from "testcontainers";
+
+export const OLLAMA_PORT = 11434;
+
+export class OllamaContainer extends GenericContainer {
+  constructor(image = "ollama/ollama") {
+    super(image);
+    this.withExposedPorts(OLLAMA_PORT)
+      .withWaitStrategy(Wait.forLogMessage("Listening on "))
+      .withStartupTimeout(120_000);
+  }
+
+  public override async start(): Promise<StartedOllamaContainer> {
+    const containerRuntimeClient = await getContainerRuntimeClient();
+    const runtimes = containerRuntimeClient.info.containerRuntime.runtimes;
+    if (runtimes.includes("nvidia")) {
+      this.hostConfig.DeviceRequests = [
+        {
+          Driver: "nvidia",
+          Count: -1,
+          Capabilities: [["gpu"]],
+        },
+      ];
+    }
+
+    return new StartedOllamaContainer(await super.start());
+  }
+}
+
+export class StartedOllamaContainer extends AbstractStartedContainer {
+  constructor(startedTestContainer: StartedTestContainer) {
+    super(startedTestContainer);
+  }
+
+  public getPort(): number {
+    return this.startedTestContainer.getMappedPort(OLLAMA_PORT);
+  }
+
+  public getEndpoint(): string {
+    return `http://${this.getHost()}:${this.getPort().toString()}`;
+  }
+
+  public async commitToImage(imageName: string): Promise<NodeJS.ReadableStream> {
+    const client = await getContainerRuntimeClient();
+    const dockerode = client.container.dockerode;
+    return dockerode.getContainer(this.getId()).commit({ repo: imageName });
+  }
+}

packages/modules/ollama/tsconfig.build.json

@@ -0,0 +1,13 @@
+{
+  "extends": "./tsconfig.json",
+  "exclude": [
+    "build",
+    "jest.config.ts",
+    "src/**/*.test.ts"
+  ],
+  "references": [
+    {
+      "path": "../../testcontainers"
+    }
+  ]
+}

packages/modules/ollama/tsconfig.json

@@ -0,0 +1,21 @@
+{
+  "extends": "../../../tsconfig.base.json",
+  "compilerOptions": {
+    "rootDir": "src",
+    "outDir": "build",
+    "paths": {
+      "testcontainers": [
+        "../../testcontainers/src"
+      ]
+    }
+  },
+  "exclude": [
+    "build",
+    "jest.config.ts"
+  ],
+  "references": [
+    {
+      "path": "../../testcontainers"
+    }
+  ]
+}

packages/testcontainers/src/container-runtime/clients/client.ts

@@ -106,6 +106,7 @@ async function initStrategy(strategy: ContainerRuntimeClientStrategy): Promise<C
     architecture: dockerodeInfo.Architecture,
     cpus: dockerodeInfo.NCPU,
     memory: dockerodeInfo.MemTotal,
+    runtimes: Object.keys(dockerodeInfo.Runtimes),
   };
 
   const composeInfo: ComposeInfo = composeClient.info;

packages/testcontainers/src/container-runtime/clients/types.ts

@@ -21,6 +21,7 @@ export type ContainerRuntimeInfo = {
   architecture: string;
   cpus: number;
   memory: number;
+  runtimes: string[];
 };
 
 export type ComposeInfo =