[CHORE] Enable rust log service, set timeouts, change tilt. (#4338)

## Description of changes

- Change timeouts in tilt to 60s
- Change deadline for proptest to 90s.
- Reenable use_alt_host_for_everything in tilt.
- Template with the right file consistently.

## Test plan

- [X] Tests pass locally with `pytest` for python, `yarn test` for js,
`cargo test` for rust

## Documentation Changes

N/A

Author: rescrv

Tiltfile

@@ -166,7 +166,7 @@ k8s_yaml(
 # We manually call helm template so we can call set-file
 k8s_yaml(
   local(
-    'helm template --set-file rustFrontendService.configuration=rust/frontend/sample_configs/distributed.yaml,rustLogService.configuration=rust/worker/tilt_config.yaml,compaction_service.configuration=rust/worker/tilt_config.yaml,query_service.configuration=rust/worker/tilt_config.yaml --values k8s/distributed-chroma/values.yaml,k8s/distributed-chroma/values.dev.yaml k8s/distributed-chroma'
+    'helm template --set-file rustFrontendService.configuration=rust/frontend/sample_configs/tilt_config.yaml,rustLogService.configuration=rust/worker/tilt_config.yaml,compaction_service.configuration=rust/worker/tilt_config.yaml,query_service.configuration=rust/worker/tilt_config.yaml --values k8s/distributed-chroma/values.yaml,k8s/distributed-chroma/values.dev.yaml k8s/distributed-chroma'
   ),
 )
 watch_file('rust/frontend/sample_configs/distributed.yaml')

chromadb/test/conftest.py

@@ -56,7 +56,7 @@
 
 hypothesis.settings.register_profile(
     "base",
-    deadline=45000,
+    deadline=90000,
     suppress_health_check=[
         hypothesis.HealthCheck.data_too_large,
         hypothesis.HealthCheck.large_base_example,

chromadb/test/property/invariants.py

@@ -108,6 +108,8 @@ def count(collection: Collection, record_set: RecordSet) -> None:
     """The given collection count is equal to the number of embeddings"""
     count = collection.count()
     normalized_record_set = wrap_all(record_set)
+    if count != len(normalized_record_set["ids"]):
+        print('count mismatch:', count, '=!', len(normalized_record_set["ids"]))
     assert count == len(normalized_record_set["ids"])
 
 

chromadb/test/property/test_embeddings.py

@@ -47,6 +47,9 @@
 traces: DefaultDict[str, int] = defaultdict(lambda: 0)
 
 
+VERSION_INCREASE_WAIT_TIME = 300
+
+
 def trace(key: str) -> None:
     global traces
     traces[key] += 1
@@ -375,7 +378,8 @@ def wait_for_compaction(self) -> None:
                 current_version,
             )
             new_version = wait_for_version_increase(
-                self.client, self.collection.name, current_version, additional_time=240
+                self.client, self.collection.name, current_version,
+                additional_time=VERSION_INCREASE_WAIT_TIME
             )
             # Everything got compacted.
             self.log_operation_count = 0
@@ -1420,7 +1424,7 @@ def test_no_op_compaction(client: ClientAPI) -> None:
         coll.delete(ids=[str(i) for i in range(batch, batch + 100)])
     if not NOT_CLUSTER_ONLY:
         wait_for_version_increase(
-            client, coll.name, get_collection_version(client, coll.name), 240
+            client, coll.name, get_collection_version(client, coll.name), VERSION_INCREASE_WAIT_TIME
         )
 
 
@@ -1439,7 +1443,7 @@ def test_add_then_purge(client: ClientAPI) -> None:
         )
     if not NOT_CLUSTER_ONLY:
         wait_for_version_increase(
-            client, coll.name, get_collection_version(client, coll.name), 240
+            client, coll.name, get_collection_version(client, coll.name), VERSION_INCREASE_WAIT_TIME
         )
 
     # Purge records and wait for compaction
@@ -1449,7 +1453,7 @@ def test_add_then_purge(client: ClientAPI) -> None:
         coll.delete(ids=record_ids)
     if not NOT_CLUSTER_ONLY:
         wait_for_version_increase(
-            client, coll.name, get_collection_version(client, coll.name), 240
+            client, coll.name, get_collection_version(client, coll.name), VERSION_INCREASE_WAIT_TIME
         )
 
     # There should be no records left
@@ -1472,7 +1476,7 @@ def test_encompassing_delete(client: ClientAPI) -> None:
 
     if not NOT_CLUSTER_ONLY:
         wait_for_version_increase(
-            client, col.name, get_collection_version(client, col.name), 240
+            client, col.name, get_collection_version(client, col.name), VERSION_INCREASE_WAIT_TIME
         )
 
     # Add and then delete and then add 16
@@ -1486,7 +1490,7 @@ def test_encompassing_delete(client: ClientAPI) -> None:
 
     if not NOT_CLUSTER_ONLY:
         wait_for_version_increase(
-            client, col.name, get_collection_version(client, col.name), 240
+            client, col.name, get_collection_version(client, col.name), VERSION_INCREASE_WAIT_TIME
         )
 
     # Ensure we can get all

rust/frontend/sample_configs/tilt_config.yaml

@@ -23,16 +23,16 @@ log:
     host: "logservice.chroma"
     port: 50051
     connect_timeout_ms: 5000
-    request_timeout_ms: 5000
+    request_timeout_ms: 60000 # 1 minute
     alt_host: "rust-log-service.chroma"
-    #use_alt_host_for_everything: true
+    use_alt_host_for_everything: true
 
 executor:
   distributed:
     connections_per_node: 5
     replication_factor: 2
     connect_timeout_ms: 5000
-    request_timeout_ms: 5000
+    request_timeout_ms: 60000 # 1 minute
     assignment:
       rendezvous_hashing:
         hasher: Murmur3

rust/frontend/src/config.rs

@@ -160,6 +160,11 @@ impl FrontendServerConfig {
     }
 
     pub fn load_from_path(path: &str) -> Self {
+        // SAFETY(rescrv): If we cannot read the config, we panic anyway.
+        eprintln!(
+            "==========\n{}\n==========\n",
+            std::fs::read_to_string(path).unwrap()
+        );
         // Unfortunately, figment doesn't support environment variables with underscores. So we have to map and replace them.
         // Excluding our own environment variables, which are prefixed with CHROMA_.
         let mut f = figment::Figment::from(

rust/garbage_collector/src/helper.rs

@@ -24,7 +24,7 @@ impl ChromaGrpcClients {
         let sysdb_channel = Channel::from_static("http://localhost:50051")
             .connect()
             .await?;
-        let logservice_channel = Channel::from_static("http://localhost:50052")
+        let logservice_channel = Channel::from_static("http://localhost:50054")
             .connect()
             .await?;
         let queryservice_channel = Channel::from_static("http://localhost:50053")

rust/log-service/src/lib.rs

@@ -709,7 +709,7 @@ impl LogService for LogServer {
             );
             let limits = Limits {
                 max_files: Some(pull_logs.batch_size as u64 + 1),
-                max_bytes: Some(pull_logs.batch_size as u64 * 100_000),
+                max_bytes: None,
             };
             let fragments = match log_reader
                 .scan(

rust/wal3/src/reader.rs

@@ -131,14 +131,23 @@ impl LogReader {
         }
         fragments.retain(|f| f.limit > from);
         fragments.sort_by_key(|f| f.start.offset());
-        fragments.truncate(limits.max_files.unwrap_or(u64::MAX) as usize);
+        if let Some(max_files) = limits.max_files {
+            if fragments.len() as u64 > max_files {
+                tracing::info!("truncating to {} files from {}", max_files, fragments.len());
+                fragments.truncate(max_files as usize);
+            }
+        }
         while fragments.len() > 1
             && fragments
                 .iter()
                 .map(|f| f.num_bytes)
                 .fold(0, u64::saturating_add)
                 > limits.max_bytes.unwrap_or(u64::MAX)
         {
+            tracing::info!(
+                "truncating to {} files because bytes restrictions",
+                fragments.len() - 1
+            );
             fragments.pop();
         }
         Ok(fragments)

rust/worker/src/config.rs

@@ -61,6 +61,11 @@ impl RootConfig {
     pub fn load_from_path(path: &str) -> Self {
         // Unfortunately, figment doesn't support environment variables with underscores. So we have to map and replace them.
         // Excluding our own environment variables, which are prefixed with CHROMA_.
+        eprintln!("loading config from {path}");
+        eprintln!(
+            "{}",
+            std::fs::read_to_string(path).unwrap_or("<ERROR>".to_string())
+        );
         let mut f = figment::Figment::from(Env::prefixed("CHROMA_").map(|k| match k {
             k if k == "my_member_id" => k.into(),
             k => k.as_str().replace("__", ".").into(),

rust/worker/src/execution/operators/count_records.rs

@@ -81,7 +81,7 @@ impl Operator<CountRecordsInput, CountRecordsOutput> for CountRecordsOperator {
             Err(e) => {
                 match *e {
                     RecordSegmentReaderCreationError::UninitializedSegment => {
-                        tracing::info!("[CountQueryOrchestrator] Record segment is uninitialized");
+                        tracing::info!("[CountQueryOrchestrator] Record segment is uninitialized; using {} records from log", input.log_records.len());
                         // This means there no compaction has occured.
                         // So we can just traverse the log records
                         // and count the number of records.

rust/worker/src/lib.rs

@@ -22,8 +22,14 @@ const CONFIG_PATH_ENV_VAR: &str = "CONFIG_PATH";
 pub async fn query_service_entrypoint() {
     // Check if the config path is set in the env var
     let config = match std::env::var(CONFIG_PATH_ENV_VAR) {
-        Ok(config_path) => config::RootConfig::load_from_path(&config_path),
-        Err(_) => config::RootConfig::load(),
+        Ok(config_path) => {
+            eprintln!("loading from {config_path}");
+            config::RootConfig::load_from_path(&config_path)
+        }
+        Err(err) => {
+            eprintln!("loading from default path because {err}");
+            config::RootConfig::load()
+        }
     };
 
     let config = config.query_service;

rust/worker/tilt_config.yaml

@@ -40,9 +40,9 @@ query_service:
             host: "logservice.chroma"
             port: 50051
             connect_timeout_ms: 5000
-            request_timeout_ms: 5000
+            request_timeout_ms: 60000 # 1 minute
             alt_host: "rust-log-service.chroma"
-            #use_alt_host_for_everything: true
+            use_alt_host_for_everything: true
     dispatcher:
         num_worker_threads: 4
         dispatcher_queue_size: 100
@@ -78,7 +78,7 @@ query_service:
             weighted_lru:
                 capacity: 8589934592 # 8GB
         permitted_parallelism: 180
-    fetch_log_batch_size: 4000
+    fetch_log_batch_size: 1000
 
 compaction_service:
     service_name: "compaction-service"
@@ -117,9 +117,9 @@ compaction_service:
             host: "logservice.chroma"
             port: 50051
             connect_timeout_ms: 5000
-            request_timeout_ms: 5000
+            request_timeout_ms: 60000 # 1 minute
             alt_host: "rust-log-service.chroma"
-            #use_alt_host_for_everything: true
+            use_alt_host_for_everything: true
     dispatcher:
         num_worker_threads: 4
         dispatcher_queue_size: 100
@@ -134,7 +134,7 @@ compaction_service:
         max_compaction_size: 10000
         max_partition_size: 5000
         disabled_collections: [] # uuids to disable compaction for
-        fetch_log_batch_size: 4000
+        fetch_log_batch_size: 1000
     blockfile_provider:
         arrow:
             block_manager_config: