ChromaDB on GCP (using Pulumi)

index.ts

import * as pulumi from "@pulumi/pulumi";
import * as gcp from "@pulumi/gcp";

const chromaApiKey = new pulumi.Config("chroma").requireSecret("apiKey");
const zone = gcp.config.region || <your preferred region>;

// Create a new secret for Chroma API key.
const chromaApiKeySecret = new gcp.secretmanager.Secret(
  "chroma-api-key-secret",
  {
    secretId: "CHROMA_API_KEY",
    replication: {
      auto: {},
    },
  },
);

// Create a new version of the Chroma API key secret.
const chromaApiKeySecretVersion = new gcp.secretmanager.SecretVersion(
  "chroma-api-key-secret-version",
  {
    secret: chromaApiKeySecret.id,
    secretData: chromaApiKey,
  },
  {
    parent: chromaApiKeySecret,
  },
);

//enable firewall rules to allow access. We will use port 80 so it's easier to connect across js and python
const computeFirewall = new gcp.compute.Firewall("chroma-firewall-rule", {
  network: "default",
  allows: [
    {
      protocol: "tcp",
      ports: ["22", "80"],
    },
  ],
  // Allow access from internet
  sourceRanges: ["0.0.0.0/0"],
  targetTags: ["chroma"],
});

// Create a persistent disk
const disk = new gcp.compute.Disk("chroma-data-persistent-disk", {
  size: 256, // Size in GB
  zone,
  type: "pd-ssd", // Persistent SSD for better performance
});

// Create a Google Cloud VM instance
const instance = new gcp.compute.Instance("chroma-vm", {
  machineType: "n2d-highmem-2", //Prefer highmem intances. You can choose as you want tho :)
  zone,
  serviceAccount: {
    email: "default",
    scopes: ["https://www.googleapis.com/auth/cloud-platform"],
  },
  allowStoppingForUpdate: true,
  bootDisk: {
    initializeParams: {
      image: "cos-cloud/cos-109-lts", // Keep it at 109 for inbuilt GCP logging
    },
  },
  networkInterfaces: [
    {
      network: "default",
      accessConfigs: [{}], // Allow external access
    },
  ],
  tags: ["chroma"],
  attachedDisks: [
    {
      source: disk.id,
      mode: "READ_WRITE",
      deviceName: "chroma-disk",
    },
  ],
  metadata: {
    "google-logging-enabled": "true", // Logging to GCP logging
  },
  metadataStartupScript: pulumi.interpolate`#!/bin/bash
        # Check if the disk is already formatted
        if ! blkid /dev/disk/by-id/google-chroma-disk; then
            # Format the persistent disk if not formatted
            mkfs.ext4 -m 0 -E lazy_itable_init=0,lazy_journal_init=0,discard /dev/disk/by-id/google-chroma-disk
        fi

        # Mount the persistent disk
        mkdir -p /mnt/disks/chroma/chroma
        mount -t ext4 -o discard,defaults /dev/disk/by-id/google-chroma-disk /mnt/disks/chroma/chroma

        # Ensure the disk mounts on reboot
        echo '/dev/disk/by-id/google-chroma-disk /mnt/disks/chroma/chroma ext4 defaults 0 0' >> /etc/fstab

        # Make directory readable and writable
        chmod -R a+rw /mnt/disks/chroma/chroma

        # Run the chromadb/chroma Docker container
        docker run -d --name chroma\
        -p 80:8000\ //We will use port 80 so it's easier to connect across js and python
        -v /mnt/disks/chroma/chroma:/chroma/chroma\
        -e CHROMA_SERVER_AUTHN_CREDENTIALS="${chromaApiKey}"\
        -e CHROMA_SERVER_AUTHN_PROVIDER="chromadb.auth.token_authn.TokenAuthenticationServerProvider"\
        chromadb/chroma:0.5.11`,
});

export const instanceName = instance.name;
export const instanceIP = instance.networkInterfaces.apply(
  (ni) => ni[0].accessConfigs![0].natIp!,
);
export const diskId = disk.id;