Skip to content

Commit 6b40d2c

Browse files
eculvereculver
authored
[CHORE] Add pprof server to compaction service (#5092)
## Description of changes Similar to what we did in #4845 and #5072, this adds the jemalloc pprof server to the compaction-service server so that we can do memory profiling for the service using Polar Signals. - New functionality - pprof server started when config given ## Test plan Build should succeed and updating chroma_config.yaml with local tilt environment to test that pprof endpoint is up.
1 parent b56819c commit 6b40d2c

File tree

11 files changed

+60
-4
lines changed

11 files changed

+60
-4
lines changed

‎k8s/distributed-chroma/Chart.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ apiVersion: v2
1616
name: distributed-chroma
1717
description: A helm chart for distributed Chroma
1818
type: application
19-
version: 0.1.50
19+
version: 0.1.51
2020
appVersion: "0.4.24"
2121
keywords:
2222
- chroma

‎k8s/distributed-chroma/templates/compaction-service.yaml

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,9 @@ spec:
5959
{{ end }}
6060
ports:
6161
- containerPort: 50051
62+
- containerPort: 6060
63+
protocol: TCP
64+
name: pprof
6265
env:
6366
{{if .Values.compactionService.configuration}}
6467
- name: CONFIG_PATH
@@ -69,6 +72,10 @@ spec:
6972
# TODO properly use flow control here to check which type of value we need.
7073
{{ .value | nindent 14 }}
7174
{{ end }}
75+
{{ if .Values.compactionService.jemallocConfig }}
76+
- name: _RJEM_MALLOC_CONF
77+
value: {{ .Values.compactionService.jemallocConfig }}
78+
{{ end }}
7279
- name: CHROMA_COMPACTION_SERVICE__MY_MEMBER_ID
7380
valueFrom:
7481
fieldRef:

‎k8s/distributed-chroma/templates/query-service.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,7 @@ spec:
9494
{{ end }}
9595
{{ if .Values.queryService.jemallocConfig }}
9696
- name: _RJEM_MALLOC_CONF
97-
value: {{ .Values.queryService.jemallocConfig}}
97+
value: {{ .Values.queryService.jemallocConfig }}
9898
{{ end }}
9999
- name: CHROMA_QUERY_SERVICE__MY_MEMBER_ID
100100
valueFrom:

‎k8s/distributed-chroma/values.dev.yaml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ compactionService:
3232
value: 'value: "1"'
3333
- name: CONFIG_PATH
3434
value: 'value: "/tilt_config.yaml"'
35+
jemallocConfig: "prof:true,prof_active:true,lg_prof_sample:19"
3536

3637
rustLogService:
3738
replicaCount: 1

‎rust/worker/src/compactor/compaction_server.rs

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
use async_trait::async_trait;
2+
use chroma_jemalloc_pprof_server::spawn_pprof_server;
23
use chroma_system::ComponentHandle;
34
use chroma_types::chroma_proto::{
45
compactor_server::{Compactor, CompactorServer},
@@ -18,6 +19,7 @@ use super::{CompactionManager, RebuildMessage};
1819
pub struct CompactionServer {
1920
pub manager: ComponentHandle<CompactionManager>,
2021
pub port: u16,
22+
pub jemalloc_pprof_server_port: Option<u16>,
2123
}
2224

2325
impl CompactionServer {
@@ -42,9 +44,19 @@ impl CompactionServer {
4244
.send(RegisterOnReadySignal { on_ready_tx }, None)
4345
.await?;
4446

47+
// Start pprof server
48+
let mut pprof_shutdown_tx = None;
49+
if let Some(port) = self.jemalloc_pprof_server_port {
50+
tracing::info!("Starting jemalloc pprof server on port {}", port);
51+
let shutdown_channel = tokio::sync::oneshot::channel();
52+
pprof_shutdown_tx = Some(shutdown_channel.0);
53+
spawn_pprof_server(port, shutdown_channel.1).await;
54+
}
55+
4556
let server = Server::builder()
4657
.add_service(health_service)
4758
.add_service(CompactorServer::new(self));
59+
4860
server
4961
.serve_with_shutdown(addr, async {
5062
match signal(SignalKind::terminate()) {
@@ -58,6 +70,12 @@ impl CompactionServer {
5870
}
5971
})
6072
.await?;
73+
74+
// Shutdown pprof server after server is finished shutting down
75+
if let Some(shutdown_tx) = pprof_shutdown_tx {
76+
let _ = shutdown_tx.send(());
77+
}
78+
6179
Ok(())
6280
}
6381
}

‎rust/worker/src/config.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,8 @@ pub struct CompactionServiceConfig {
202202
pub hnsw_provider: chroma_index::config::HnswProviderConfig,
203203
#[serde(default)]
204204
pub spann_provider: chroma_index::config::SpannProviderConfig,
205+
#[serde(default)]
206+
pub jemalloc_pprof_server_port: Option<u16>,
205207
}
206208

207209
impl CompactionServiceConfig {

‎rust/worker/src/lib.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,7 @@ pub async fn compaction_service_entrypoint() {
148148
let compaction_server = CompactionServer {
149149
manager: compaction_manager_handle.clone(),
150150
port: config.my_port,
151+
jemalloc_pprof_server_port: config.jemalloc_pprof_server_port,
151152
};
152153

153154
let server_join_handle = tokio::spawn(async move {

‎rust/worker/tests/config_from_default_path.rs

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ fn test_config_from_default_path() {
1616
otel_endpoint: "http://jaeger:4317"
1717
my_member_id: "query-service-0"
1818
my_port: 50051
19+
jemalloc_pprof_server_port: 6060
1920
assignment_policy:
2021
rendezvous_hashing:
2122
hasher: Murmur3
@@ -78,6 +79,7 @@ fn test_config_from_default_path() {
7879
otel_endpoint: "http://jaeger:4317"
7980
my_member_id: "compaction-service-0"
8081
my_port: 50051
82+
jemalloc_pprof_server_port: 6060
8183
assignment_policy:
8284
rendezvous_hashing:
8385
hasher: Murmur3
@@ -155,12 +157,16 @@ fn test_config_from_default_path() {
155157
let config = RootConfig::load();
156158
assert_eq!(config.query_service.my_member_id, "query-service-0");
157159
assert_eq!(config.query_service.my_port, 50051);
158-
160+
assert_eq!(config.query_service.jemalloc_pprof_server_port, Some(6060));
159161
assert_eq!(
160162
config.compaction_service.my_member_id,
161163
"compaction-service-0"
162164
);
163165
assert_eq!(config.compaction_service.my_port, 50051);
166+
assert_eq!(
167+
config.compaction_service.jemalloc_pprof_server_port,
168+
Some(6060)
169+
);
164170
assert_eq!(
165171
config
166172
.compaction_service

‎rust/worker/tests/config_from_specific_path.rs

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ fn test_config_from_specific_path() {
1515
otel_endpoint: "http://jaeger:4317"
1616
my_member_id: "query-service-0"
1717
my_port: 50051
18+
jemalloc_pprof_server_port: 6060
1819
assignment_policy:
1920
rendezvous_hashing:
2021
hasher: Murmur3
@@ -76,6 +77,7 @@ fn test_config_from_specific_path() {
7677
otel_endpoint: "http://jaeger:4317"
7778
my_member_id: "compaction-service-0"
7879
my_port: 50051
80+
jemalloc_pprof_server_port: 6060
7981
assignment_policy:
8082
rendezvous_hashing:
8183
hasher: Murmur3
@@ -155,12 +157,16 @@ fn test_config_from_specific_path() {
155157
let config = RootConfig::load_from_path("random_path.yaml");
156158
assert_eq!(config.query_service.my_member_id, "query-service-0");
157159
assert_eq!(config.query_service.my_port, 50051);
158-
160+
assert_eq!(config.query_service.jemalloc_pprof_server_port, Some(6060));
159161
assert_eq!(
160162
config.compaction_service.my_member_id,
161163
"compaction-service-0"
162164
);
163165
assert_eq!(config.compaction_service.my_port, 50051);
166+
assert_eq!(
167+
config.compaction_service.jemalloc_pprof_server_port,
168+
Some(6060)
169+
);
164170
assert!(
165171
config
166172
.compaction_service

‎rust/worker/tests/config_with_env_override.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,16 @@ fn test_config_with_env_override() {
99
Jail::expect_with(|jail| {
1010
jail.set_env("CHROMA_QUERY_SERVICE__MY_MEMBER_ID", "query-service-0");
1111
jail.set_env("CHROMA_QUERY_SERVICE__MY_PORT", 50051);
12+
jail.set_env("CHROMA_QUERY_SERVICE__JEMALLOC_PPROF_SERVER_PORT", 6060);
1213
jail.set_env(
1314
"CHROMA_COMPACTION_SERVICE__MY_MEMBER_ID",
1415
"compaction-service-0",
1516
);
1617
jail.set_env("CHROMA_COMPACTION_SERVICE__MY_PORT", 50051);
18+
jail.set_env(
19+
"CHROMA_COMPACTION_SERVICE__JEMALLOC_PPROF_SERVER_PORT",
20+
6060,
21+
);
1722
jail.set_env("CHROMA_COMPACTION_SERVICE__STORAGE__S3__BUCKET", "buckets!");
1823
jail.set_env("CHROMA_COMPACTION_SERVICE__STORAGE__S3__CREDENTIALS", "AWS");
1924
jail.set_env(
@@ -156,11 +161,16 @@ fn test_config_with_env_override() {
156161
let config = RootConfig::load();
157162
assert_eq!(config.query_service.my_member_id, "query-service-0");
158163
assert_eq!(config.query_service.my_port, 50051);
164+
assert_eq!(config.query_service.jemalloc_pprof_server_port, Some(6060));
159165
assert_eq!(
160166
config.compaction_service.my_member_id,
161167
"compaction-service-0"
162168
);
163169
assert_eq!(config.compaction_service.my_port, 50051);
170+
assert_eq!(
171+
config.compaction_service.jemalloc_pprof_server_port,
172+
Some(6060)
173+
);
164174
match &config.compaction_service.storage {
165175
chroma_storage::config::StorageConfig::S3(s) => {
166176
assert_eq!(s.bucket, "buckets!");

0 commit comments

Comments
 (0)