[ENH] Add rust frontend stub (#3549)

## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
   - /
 - New functionality
   - Adds a rust-based frontend stub
   - Adds the frontend to the k8s tiltfile

## Test plan
*How are these changes tested?*
Confirmed I can tilt up and curl the service at `:3000`
- [x] Tests pass locally with `pytest` for python, `yarn test` for js,
`cargo test` for rust

## Documentation Changes
None required.

Cargo.toml

@@ -1,11 +1,12 @@
 [workspace]
 resolver = "2"
 
-members = ["rust/benchmark", "rust/blockstore", "rust/cache", "rust/chroma", "rust/config", "rust/distance", "rust/error", "rust/garbage_collector", "rust/index", "rust/load", "rust/storage", "rust/system", "rust/sysdb", "rust/types", "rust/worker", "rust/segment"]
+members = ["rust/benchmark", "rust/blockstore", "rust/cache", "rust/chroma", "rust/config", "rust/distance", "rust/error", "rust/garbage_collector", "rust/index", "rust/load", "rust/storage", "rust/system", "rust/sysdb", "rust/types", "rust/worker", "rust/segment", "rust/frontend"]
 
 [workspace.dependencies]
 arrow = "52.2.0"
 async-trait = "0.1"
+axum = "0.7"
 chrono = { version = "0.4", features = ["serde"] }
 clap = { version = "4", features = ["derive"] }
 criterion = { version = "0.5", features = ["async_tokio"] }
@@ -46,6 +47,7 @@ chroma-distance = { path = "rust/distance" }
 chroma-segment = { path="rust/segment" }
 chroma-system = { path = "rust/system" }
 chroma-sysdb = { path = "rust/sysdb" }
+chroma-frontend = { path = "rust/frontend" }
 worker = { path = "rust/worker" }
 
 # Dev dependencies

Tiltfile

@@ -46,6 +46,15 @@ docker_build(
   ignore=['**/*.pyc', 'chromadb/test/'],
 )
 
+docker_build(
+  'local:rust-frontend-service',
+  '.',
+  only=["rust/", "idl/", "Cargo.toml", "Cargo.lock"],
+  dockerfile='./rust/frontend/Dockerfile',
+  target='frontend_service'
+)
+
+
 docker_build(
   'local:query-service',
   '.',
@@ -135,6 +144,7 @@ k8s_resource('logservice-migration-logservice-migration', resource_deps=['postgr
 k8s_resource('logservice', resource_deps=['sysdb-migration-sysdb-migration'], labels=["chroma"], port_forwards='50052:50051')
 k8s_resource('sysdb', resource_deps=['sysdb-migration-sysdb-migration'], labels=["chroma"], port_forwards='50051:50051')
 k8s_resource('frontend-service', resource_deps=['sysdb', 'logservice'],labels=["chroma"], port_forwards='8000:8000')
+k8s_resource('rust-frontend-service', resource_deps=['sysdb', 'logservice'], labels=["chroma"], port_forwards='3000:3000')
 k8s_resource('query-service', resource_deps=['sysdb'], labels=["chroma"], port_forwards='50053:50051')
 k8s_resource('compaction-service', resource_deps=['sysdb'], labels=["chroma"])
 

k8s/distributed-chroma/templates/rust-frontend-service.yaml

@@ -0,0 +1,54 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: rust-frontend-service
+  namespace: {{ .Values.namespace }}
+spec:
+  replicas: {{ .Values.rustFrontendService.replicaCount }}
+  selector:
+    matchLabels:
+      app: rust-frontend-service
+  template:
+    metadata:
+      labels:
+        app: rust-frontend-service
+    spec:
+      containers:
+        - name: rust-frontend-service
+          image: "{{ .Values.rustFrontendService.image.repository }}:{{ .Values.rustFrontendService.image.tag }}"
+          imagePullPolicy: IfNotPresent
+          ports:
+            - containerPort: 3000
+          {{ if .Values.rustFrontendService.resources }}
+          resources:
+            limits:
+              cpu: {{ .Values.rustFrontendService.resources.limits.cpu }}
+              memory: {{ .Values.rustFrontendService.resources.limits.memory }}
+            requests:
+              cpu: {{ .Values.rustFrontendService.resources.requests.cpu }}
+              memory: {{ .Values.rustFrontendService.resources.requests.memory }}
+          {{ end }}
+      {{if .Values.rustFrontendService.tolerations}}
+      tolerations:
+        {{ toYaml .Values.rustFrontendService.tolerations | nindent 8 }}
+      {{ end }}
+      {{if .Values.rustFrontendService.nodeSelector}}
+      nodeSelector:
+        {{ toYaml .Values.rustFrontendService.nodeSelector | nindent 8 }}
+      {{ end }}
+
+---
+
+apiVersion: v1
+kind: Service
+metadata:
+  name: rust-frontend-service
+  namespace: {{ .Values.namespace }}
+spec:
+  ports:
+    - name: server-port
+      port: 3000
+      targetPort: 3000
+  selector:
+    app: rust-frontend-service
+  type: ClusterIP

k8s/distributed-chroma/values.yaml

@@ -34,6 +34,21 @@ frontendService:
     - name: CHROMA_OTEL_GRANULARITY
       value: all
   replicaCount: 2
+
+
+rustFrontendService:
+  image:
+    repository: 'local'
+    tag: 'rust-frontend-service'
+  replicaCount: 1
+  resources:
+    limits:
+      cpu: '2000m'
+      memory: '1Gi'
+    requests:
+      cpu: '1000m'
+      memory: '512Mi'
+
 sysdb:
   image:
     repository: 'local'

rust/frontend/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "frontend"
+version = "0.1.0"
+edition = "2021"
+
+[[bin]]
+name = "frontend_service"
+path = "src/bin/frontend_service.rs"
+
+[dependencies]
+tokio = { workspace = true }
+axum = { workspace = true }

rust/frontend/Dockerfile

@@ -0,0 +1,35 @@
+FROM rust:1.81.0 AS builder
+
+ARG RELEASE_MODE=
+
+WORKDIR /
+
+WORKDIR /chroma/
+
+ENV PROTOC_ZIP=protoc-25.1-linux-x86_64.zip
+RUN 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 Cargo.toml Cargo.toml
+COPY Cargo.lock Cargo.lock
+COPY idl/ idl/
+COPY rust/ rust/
+
+
+FROM builder AS frontend_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/ \
+    if [ "$RELEASE_MODE" = "1" ]; then cargo build --bin frontend_service --release; else cargo build --bin frontend_service; fi && \
+    if [ "$RELEASE_MODE" = "1" ]; then mv target/release/frontend_service ./frontend_service; else mv target/debug/frontend_service ./frontend_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 .
+
+FROM runner AS frontend_service
+COPY --from=frontend_service_builder /chroma/frontend_service .
+ENTRYPOINT [ "./frontend_service" ]

rust/frontend/src/bin/frontend_service.rs

@@ -0,0 +1,7 @@
+use frontend::frontend_service_entrypoint;
+
+#[tokio::main]
+async fn main() {
+    // TODO(hammadb): Implement the frontend service entrypoint
+    frontend_service_entrypoint().await;
+}

rust/frontend/src/lib.rs

@@ -0,0 +1,7 @@
+mod server;
+use server::FrontendServer;
+
+pub async fn frontend_service_entrypoint() {
+    let server = FrontendServer::new();
+    FrontendServer::run(server).await;
+}

rust/frontend/src/server.rs

@@ -0,0 +1,46 @@
+use std::sync::Arc;
+
+use axum::{extract::State, routing::get, Router};
+
+struct FrontendServerInner {}
+
+#[derive(Clone)]
+pub(crate) struct FrontendServer {
+    _inner: Arc<FrontendServerInner>,
+}
+
+impl FrontendServer {
+    pub fn new() -> FrontendServer {
+        FrontendServer {
+            _inner: Arc::new(FrontendServerInner {}),
+        }
+    }
+
+    #[allow(dead_code)]
+    pub async fn run(server: FrontendServer) {
+        let app = Router::new()
+            // `GET /` goes to `root`
+            .route("/", get(root))
+            .with_state(server);
+
+        // TODO: configuration for this
+        // TODO: tracing
+        let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
+        axum::serve(listener, app).await.unwrap();
+    }
+
+    ////////////////////////// Method Implementations //////////////////////
+
+    fn root(&self) -> &'static str {
+        "Hello, World!"
+    }
+}
+
+////////////////////////// Method Handlers //////////////////////////
+// These handlers simply proxy the call and the relevant inputs into
+// the appropriate method on the `FrontendServer` struct.
+
+// Dummy implementation for now
+async fn root(State(server): State<FrontendServer>) -> &'static str {
+    server.root()
+}

rust/load/Cargo.toml

@@ -12,6 +12,7 @@ serde = { workspace = true }
 serde_json = { workspace = true }
 tokio = { workspace = true }
 uuid = { workspace = true }
+axum = { workspace = true }
 
 tracing = { workspace = true }
 tracing-bunyan-formatter = { workspace = true }
@@ -23,7 +24,6 @@ opentelemetry-http = { workspace = true }
 opentelemetry_sdk = { workspace = true }
 
 # Unlikely to be used in the workspace.
-axum = "0.7"
 chromadb = { git = "https://github.com/rescrv/chromadb-rs", rev = "7b46aab7366f009388f014929a92785c94134d61" }
 guacamole = { version = "0.9", default-features = false }
 tower-http = { version = "0.6.2", features = ["trace"] }