[ENH]: publish image with Rust CLI, update sample configs (#3843)

codetheweb · web-flow · commit 5e7f70ed0e43 · 2025-02-26T14:07:40.000-08:00
## Description of changes

- The Rust CLI is now published in our container image instead of the
Python server.
- Removed Python from frontend/CLI &amp; worker Dockerfiles.
- Put all sample configs under `rust/frontend/sample_configs`.

## Test plan
*How are these changes tested?*

Built the new image with 
```
docker build -f rust/cli/Dockerfile -t rust-chromadb .
```
and verified that it exposes a working API.
diff --git a/.github/workflows/release-chromadb.yml b/.github/workflows/release-chromadb.yml
@@ -98,17 +98,23 @@ jobs:
         uses: depot/build-push-action@v1
         with:
           context: .
+          file: rust/cli/Dockerfile
           platforms: ${{ env.PLATFORMS }}
           push: true
+          build-args: |
+            RELEASE_MODE=1
           tags: "${{ env.GHCR_IMAGE_NAME }}:${{ needs.get-version.outputs.version}},${{ env.DOCKERHUB_IMAGE_NAME }}:${{ needs.get-version.outputs.version }}"
           project: ${{ vars.DEPOT_PROJECT_ID }}
       - name: Build and push release Docker image
         if: ${{ needs.check-tag.outputs.tag_matches == 'true' }}
         uses: depot/build-push-action@v1
         with:
           context: .
+          file: rust/cli/Dockerfile
           platforms: ${{ env.PLATFORMS }}
           push: true
+          build-args: |
+            RELEASE_MODE=1
           tags: "${{ env.GHCR_IMAGE_NAME }}:${{ needs.get-version.outputs.version }},${{ env.DOCKERHUB_IMAGE_NAME }}:${{ needs.get-version.outputs.version }},${{ env.GHCR_IMAGE_NAME }}:latest,${{ env.DOCKERHUB_IMAGE_NAME }}:latest"
           project: ${{ vars.DEPOT_PROJECT_ID }}
 
@@ -198,7 +204,7 @@ jobs:
             PIP Package: `chroma-${{needs.get-version.outputs.version}}.tar.gz`
             Github Container Registry Image: `${{ env.GHCR_IMAGE_NAME }}:${{ needs.get-version.outputs.version }}`
             DockerHub Image: `${{ env.DOCKERHUB_IMAGE_NAME }}:${{ needs.get-version.outputs.version }}`
-          artifacts: "dist/chromadb-${{needs.get-version.outputs.version}}.tar.gz"
+          artifacts: "dist/*"
           prerelease: false
           makeLatest: true
           generateReleaseNotes: true
@@ -222,7 +228,7 @@ jobs:
             PIP Package: `chroma-${{needs.get-version.outputs.version}}.tar.gz`
             Github Container Registry Image: `${{ env.GHCR_IMAGE_NAME }}:${{ needs.get-version.outputs.version }}`
             DockerHub Image: `${{ env.DOCKERHUB_IMAGE_NAME }}:${{ needs.get-version.outputs.version }}`
-          artifacts: "dist/chromadb-${{needs.get-version.outputs.version}}.tar.gz"
+          artifacts: "dist/*"
           allowUpdates: true
           removeArtifacts: true
           prerelease: true
diff --git a/Cargo.toml b/Cargo.toml
@@ -48,7 +48,7 @@ pyo3 = { version = "0.23.3", features = ["abi3-py39"] }
 tower-http = { version = "0.6.2", features = ["trace"] }
 bytemuck = "1.21.0"
 validator = { version = "0.19", features = ["derive"] }
-rust-embed = { version = "8.5.0", features = ["include-exclude"] }
+rust-embed = { version = "8.5.0", features = ["include-exclude", "debug-embed"] }
 hnswlib = { version = "0.8.0", git = "https://github.com/chroma-core/hnswlib.git" }
 
 chroma-benchmark = { path = "rust/benchmark" }
diff --git a/Tiltfile b/Tiltfile
@@ -50,8 +50,7 @@ docker_build(
   'local:rust-frontend-service',
   '.',
   only=["rust/", "idl/", "Cargo.toml", "Cargo.lock"],
-  dockerfile='./rust/frontend/Dockerfile',
-  target='frontend_service'
+  dockerfile='./rust/cli/Dockerfile',
 )
 
 
diff --git a/k8s/distributed-chroma/Chart.yaml b/k8s/distributed-chroma/Chart.yaml
@@ -16,7 +16,7 @@ apiVersion: v2
 name: distributed-chroma
 description: A helm chart for distributed Chroma
 type: application
-version: 0.1.18
+version: 0.1.19
 appVersion: "0.4.24"
 keywords:
   - chroma
diff --git a/k8s/distributed-chroma/templates/rust-frontend-service.yaml b/k8s/distributed-chroma/templates/rust-frontend-service.yaml
@@ -15,6 +15,7 @@ spec:
     spec:
       containers:
         - name: rust-frontend-service
+          command: ["chroma", "run", "$(CONFIG_PATH)"]
           image: "{{ .Values.rustFrontendService.image.repository }}:{{ .Values.rustFrontendService.image.tag }}"
           imagePullPolicy: IfNotPresent
           readinessProbe:
@@ -32,10 +33,16 @@ spec:
               cpu: {{ .Values.rustFrontendService.resources.requests.cpu }}
               memory: {{ .Values.rustFrontendService.resources.requests.memory }}
           {{ end }}
-{{ if .Values.frontendService.otherEnvConfig }}
           env:
-{{ .Values.frontendService.otherEnvConfig | nindent 12 }}
-{{ end }}
+            - name: CONFIG_PATH
+              value: "/sample_configs/distributed.yaml"
+            {{ if .Values.frontendService.otherEnvConfig }}
+              {{ .Values.frontendService.otherEnvConfig | nindent 12 }}
+            {{ end }}
+            {{ if .Values.rustFrontendService.otherEnvConfig }}
+              {{ .Values.rustFrontendService.otherEnvConfig | nindent 12 }}
+            {{ end }}
+
       {{if .Values.rustFrontendService.tolerations}}
       tolerations:
         {{ toYaml .Values.rustFrontendService.tolerations | nindent 8 }}
diff --git a/rust/cli/Dockerfile b/rust/cli/Dockerfile
@@ -0,0 +1,37 @@
+FROM rust:1.81.0-alpine AS builder
+
+ARG RELEASE_MODE=
+
+WORKDIR /chroma
+
+ENV PROTOC_ZIP=protoc-25.1-linux-x86_64.zip
+RUN apk add curl musl-dev alpine-sdk \
+  && curl -OL https://github.com/protocolbuffers/protobuf/releases/download/v25.1/$PROTOC_ZIP \
+  && unzip -o $PROTOC_ZIP -d /usr/local bin/protoc \
+  && unzip -o $PROTOC_ZIP -d /usr/local 'include/*' \
+  && rm -f $PROTOC_ZIP
+
+COPY idl/ idl/
+COPY Cargo.toml Cargo.toml
+COPY Cargo.lock Cargo.lock
+COPY rust/ rust/
+
+# sharing=locked is necessary to prevent cargo build from running concurrently on the same mounted directory
+RUN --mount=type=cache,sharing=locked,target=/chroma/target/ \
+  --mount=type=cache,sharing=locked,target=/usr/local/cargo/registry/ \
+  cd rust/cli && \
+  if [ "$RELEASE_MODE" = "1" ]; then cargo build --bin chroma --release; else cargo build --bin chroma; fi && \
+  cd ../.. && \
+  if [ "$RELEASE_MODE" = "1" ]; then mv target/release/chroma ./chroma; else mv target/debug/chroma ./chroma; fi
+
+
+FROM scratch AS runner
+
+COPY --from=builder /chroma/rust/frontend/sample_configs /sample_configs
+COPY --from=builder /chroma/rust/frontend/sample_configs/docker_single_node.yaml /config.yaml
+COPY --from=builder /chroma/chroma /usr/local/bin/chroma
+
+EXPOSE 8000
+
+ENTRYPOINT [ "chroma" ]
+CMD [ "run", "/config.yaml" ]
diff --git a/rust/frontend/Dockerfile b/rust/frontend/Dockerfile
diff --git a/rust/frontend/sample_configs/distributed.yaml b/rust/frontend/sample_configs/distributed.yaml
diff --git a/rust/frontend/sample_configs/docker_single_node.yaml b/rust/frontend/sample_configs/docker_single_node.yaml
@@ -0,0 +1 @@
+persist_path: "/data"
diff --git a/rust/frontend/sample_configs/single_node.yaml b/rust/frontend/sample_configs/single_node.yaml
diff --git a/rust/frontend/src/config.rs b/rust/frontend/src/config.rs
@@ -102,8 +102,8 @@ pub struct FrontendServerConfig {
     pub persist_path: Option<String>,
 }
 
-const DEFAULT_CONFIG_PATH: &str = "./frontend_config.yaml";
-const DEFAULT_SINGLE_NODE_CONFIG_FILENAME: &str = "single_node_frontend_config.yaml";
+const DEFAULT_CONFIG_PATH: &str = "./sample_configs/distributed.yaml";
+const DEFAULT_SINGLE_NODE_CONFIG_FILENAME: &str = "./sample_configs/single_node.yaml";
 
 #[derive(Embed)]
 #[folder = "./"]
diff --git a/rust/worker/Dockerfile b/rust/worker/Dockerfile
@@ -10,9 +10,6 @@ RUN curl -OL https://github.com/protocolbuffers/protobuf/releases/download/v25.1
     && unzip -o $PROTOC_ZIP -d /usr/local 'include/*' \
     && rm -f $PROTOC_ZIP
 
-# Install python for pyo3
-RUN apt-get update && apt-get install -y python3 python3-pip && rm -rf /var/lib/apt/lists/*
-
 COPY Cargo.toml Cargo.toml
 COPY Cargo.lock Cargo.lock
 COPY idl/ idl/
@@ -22,28 +19,28 @@ FROM builder AS query_service_builder
 # sharing=locked is necessary to prevent cargo build from running concurrently on the same mounted directory
 RUN --mount=type=cache,sharing=locked,target=/chroma/target/ \
     --mount=type=cache,sharing=locked,target=/usr/local/cargo/registry/ \
+    cd rust/worker && \
     if [ "$RELEASE_MODE" = "1" ]; then cargo build --bin query_service --release; else cargo build --bin query_service; fi && \
+    cd ../.. && \
     if [ "$RELEASE_MODE" = "1" ]; then mv target/release/query_service ./query_service; else mv target/debug/query_service ./query_service; fi
 
 FROM builder AS compaction_service_builder
 RUN --mount=type=cache,sharing=locked,target=/chroma/target/ \
     --mount=type=cache,sharing=locked,target=/usr/local/cargo/registry/ \
+    cd rust/worker && \
     if [ "$RELEASE_MODE" = "1" ]; then cargo build --bin compaction_service --release; else cargo build --bin compaction_service; fi && \
+    cd ../.. && \
     if [ "$RELEASE_MODE" = "1" ]; then mv target/release/compaction_service ./compaction_service; else mv target/debug/compaction_service ./compaction_service; fi
 
 
 FROM debian:bookworm-slim AS runner
 RUN apt-get update && apt-get install -y libssl-dev ca-certificates && rm -rf /var/lib/apt/lists/*
 COPY --from=builder /chroma/rust/worker/chroma_config.yaml .
 
-# Install python for pyo3
-RUN apt-get update && apt-get install -y python3 python3-pip && rm -rf /var/lib/apt/lists/*
-
-
 FROM runner AS query_service
 COPY --from=query_service_builder /chroma/query_service .
 ENTRYPOINT [ "./query_service" ]
 
 FROM runner AS compaction_service
 COPY --from=compaction_service_builder /chroma/compaction_service .
-ENTRYPOINT [ "./compaction_service" ]
+ENTRYPOINT [ "./compaction_service" ]

Original file line number	Diff line number	Diff line change
`@@ -50,8 +50,7 @@ docker_build(`
`50`	`50`	`'local:rust-frontend-service',`
`51`	`51`	`'.',`
`52`	`52`	`only=["rust/", "idl/", "Cargo.toml", "Cargo.lock"],`
`53`		`- dockerfile='./rust/frontend/Dockerfile',`
`54`		`- target='frontend_service'`
	`53`	`+ dockerfile='./rust/cli/Dockerfile',`
`55`	`54`	`)`
`56`	`55`
`57`	`56`