[BUG] Load Fixes (#1333)

## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
- Introduced "chroma_server_nofile" which allows for configuring the
nofile limit on unix systems. This is needed since with many concurrent
requests, we easily exceed platform defaults like 256 on osx (sockets
alone for 100 rps eat a large % of the budget).
- Added cursor closing() in sqlite pool for extra safety (otherwise it
relies on __del__)
	 - Increased the keep alive timeout for uvicorn.

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

- [x] Tests pass locally with `pytest` for python, `yarn test` for js

## Documentation Changes
*Are all docstrings for user-facing APIs updated if required? Do we need
to make documentation changes in the [docs
repository](https://github.com/chroma-core/docs)?*
None required

bin/docker_entrypoint.sh

@@ -3,4 +3,5 @@
 echo "Rebuilding hnsw to ensure architecture compatibility"
 pip install --force-reinstall --no-cache-dir chroma-hnswlib
 export IS_PERSISTENT=1
-uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000 --proxy-headers --log-config chromadb/log_config.yml
+export CHROMA_SERVER_NOFILE=65535
+uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000 --proxy-headers --log-config chromadb/log_config.yml --timeout-keep-alive 30

bin/integration-test

@@ -51,6 +51,7 @@ export CHROMA_INTEGRATION_TEST_ONLY=1
 export CHROMA_API_IMPL=chromadb.api.fastapi.FastAPI
 export CHROMA_SERVER_HOST=localhost
 export CHROMA_SERVER_HTTP_PORT=8000
+export CHROMA_SERVER_NOFILE=65535
 
 echo testing: python -m pytest "$@"
 python -m pytest "$@"

chromadb/cli/cli.py

@@ -28,11 +28,11 @@ def run(
     path: str = typer.Option(
         "./chroma_data", help="The path to the file or directory."
     ),
-    host: Annotated[Optional[str], typer.Option(
-        help="The host to listen to. Default: localhost")] = "localhost",
+    host: Annotated[
+        Optional[str], typer.Option(help="The host to listen to. Default: localhost")
+    ] = "localhost",
     port: int = typer.Option(8000, help="The port to run the server on."),
-    test: bool = typer.Option(False, help="Test mode.",
-                              show_envvar=False, hidden=True),
+    test: bool = typer.Option(False, help="Test mode.", show_envvar=False, hidden=True),
 ) -> None:
     """Run a chroma server"""
 
@@ -53,6 +53,7 @@ def run(
     # set ENV variable for PERSIST_DIRECTORY to path
     os.environ["IS_PERSISTENT"] = "True"
     os.environ["PERSIST_DIRECTORY"] = path
+    os.environ["CHROMA_SERVER_NOFILE"] = "65535"
 
     # get the path where chromadb is installed
     chromadb_path = os.path.dirname(os.path.realpath(__file__))
@@ -66,6 +67,7 @@ def run(
         "port": port,
         "workers": 1,
         "log_config": f"{chromadb_path}/log_config.yml",
+        "timeout_keep_alive": 30,
     }
 
     if test:

chromadb/config.py

@@ -10,6 +10,7 @@
 from overrides import EnforceOverrides
 from overrides import override
 from typing_extensions import Literal
+import platform
 
 
 in_pydantic_v2 = False
@@ -123,6 +124,7 @@ class Settings(BaseSettings):  # type: ignore
     chroma_server_grpc_port: Optional[str] = None
     # eg ["http://localhost:3000"]
     chroma_server_cors_allow_origins: List[str] = []
+    chroma_server_nofile: Optional[int] = None
 
     pulsar_broker_url: Optional[str] = None
     pulsar_admin_port: Optional[str] = "8080"
@@ -195,7 +197,7 @@ def chroma_server_auth_credentials_file_non_empty_file_exists(
         "chroma_server_authz_config_file", pre=True, always=True, allow_reuse=True
     )
     def chroma_server_authz_config_file_non_empty_file_exists(
-        cls: Type["Settings"], v: str  # type: ignore
+        cls: Type["Settings"], v: str
     ) -> Optional[str]:
         if v and not v.strip():
             raise ValueError(
@@ -298,6 +300,39 @@ def __init__(self, settings: Settings):
             if settings[key] is not None:
                 raise ValueError(LEGACY_ERROR)
 
+        # Apply the nofile limit if set
+        if settings["chroma_server_nofile"] is not None:
+            if platform.system() != "Windows":
+                import resource
+
+                curr_soft, curr_hard = resource.getrlimit(resource.RLIMIT_NOFILE)
+                desired_soft = settings["chroma_server_nofile"]
+                # Validate
+                if desired_soft > curr_hard:
+                    raise ValueError(
+                        f"chroma_server_nofile cannot be set to a value greater than the current hard limit of {curr_hard}"
+                    )
+                # Apply
+                if desired_soft > curr_soft:
+                    try:
+                        resource.setrlimit(
+                            resource.RLIMIT_NOFILE, (desired_soft, curr_hard)
+                        )
+                        logger.info(f"Set chroma_server_nofile to {desired_soft}")
+                    except Exception as e:
+                        logger.error(
+                            f"Failed to set chroma_server_nofile to {desired_soft}: {e} nofile soft limit will remain at {curr_soft}"
+                        )
+                # Don't apply if reducing the limit
+                elif desired_soft < curr_soft:
+                    logger.warning(
+                        f"chroma_server_nofile is set to {desired_soft}, but this is less than current soft limit of {curr_soft}. chroma_server_nofile will not be set."
+                    )
+            else:
+                logger.warning(
+                    "chroma_server_nofile is not supported on Windows. chroma_server_nofile will not be set."
+                )
+
         self.settings = settings
         self._instances = {}
         super().__init__(self)

chromadb/db/impl/sqlite.py

@@ -51,6 +51,7 @@ def __exit__(
                 self._conn.commit()
             else:
                 self._conn.rollback()
+        self._conn.cursor().close()
         self._pool.return_to_pool(self._conn)
         return False
 

chromadb/test/conftest.py

@@ -117,7 +117,13 @@ def _run_server(
             chroma_server_authz_config=chroma_server_authz_config,
         )
     server = chromadb.server.fastapi.FastAPI(settings)
-    uvicorn.run(server.app(), host="0.0.0.0", port=port, log_level="error")
+    uvicorn.run(
+        server.app(),
+        host="0.0.0.0",
+        port=port,
+        log_level="error",
+        timeout_keep_alive=30,
+    )
 
 
 def _await_server(api: ServerAPI, attempts: int = 0) -> None:

clients/python/integration-test.sh

@@ -21,6 +21,7 @@ export CHROMA_INTEGRATION_TEST_ONLY=1
 export CHROMA_API_IMPL=chromadb.api.fastapi.FastAPI
 export CHROMA_SERVER_HOST=localhost
 export CHROMA_SERVER_HTTP_PORT=8000
+export CHROMA_SERVER_NOFILE=65535
 
 echo testing: python -m pytest "$@"
 

docker-compose.test-auth.yml

@@ -12,7 +12,7 @@ services:
     volumes:
       - ./:/chroma
       - test_index_data:/index_data
-    command: uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000 --log-config chromadb/log_config.yml
+    command: uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000 --log-config chromadb/log_config.yml --timeout-keep-alive 30
     environment:
       - ANONYMIZED_TELEMETRY=False
       - ALLOW_RESET=True

docker-compose.test.yml

@@ -12,7 +12,7 @@ services:
     volumes:
       - ./:/chroma
       - test_index_data:/index_data
-    command: uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000 --log-config chromadb/log_config.yml
+    command: uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000 --log-config chromadb/log_config.yml --timeout-keep-alive 30
     environment:
       - ANONYMIZED_TELEMETRY=False
       - ALLOW_RESET=True

docker-compose.yml

@@ -14,7 +14,7 @@ services:
       - ./:/chroma
       # Be aware that indexed data are located in "/chroma/chroma/"
       # Default configuration for persist_directory in chromadb/config.py
-    command: uvicorn chromadb.app:app --reload --workers 1 --host 0.0.0.0 --port 8000 --log-config chromadb/log_config.yml
+    command: uvicorn chromadb.app:app --reload --workers 1 --host 0.0.0.0 --port 8000 --log-config chromadb/log_config.yml --timeout-keep-alive 30
     environment:
       - IS_PERSISTENT=TRUE
       - CHROMA_SERVER_AUTH_PROVIDER=${CHROMA_SERVER_AUTH_PROVIDER}
@@ -26,6 +26,7 @@ services:
       - CHROMA_OTEL_EXPORTER_HEADERS=${CHROMA_OTEL_EXPORTER_HEADERS}
       - CHROMA_OTEL_SERVICE_NAME=${CHROMA_OTEL_SERVICE_NAME}
       - CHROMA_OTEL_GRANULARITY=${CHROMA_OTEL_GRANULARITY}
+      - CHROMA_SERVER_NOFILE=${CHROMA_SERVER_NOFILE}
     ports:
       - 8000:8000
     networks:

examples/basic_functionality/authz/README.md

@@ -134,7 +134,7 @@ CHROMA_SERVER_AUTH_CREDENTIALS_FILE=examples/basic_functionality/authz/authz.yam
 CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER="user_token_config" \
 CHROMA_SERVER_AUTH_PROVIDER="chromadb.auth.token.TokenAuthServerProvider" \
 CHROMA_SERVER_AUTHZ_CONFIG_FILE=examples/basic_functionality/authz/authz.yaml \
-uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000 --proxy-headers --log-config chromadb/log_config.yml --reload
+uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000 --proxy-headers --log-config chromadb/log_config.yml --reload --timeout-keep-alive 30
 ```
 
 ## Testing the authorization

examples/deployments/render-terraform/chroma.tf

@@ -1,18 +1,18 @@
 terraform {
   required_providers {
     render = {
-      source = "jackall3n/render"
+      source  = "jackall3n/render"
       version = "~> 1.3.0"
     }
   }
 }
 
 variable "render_api_token" {
-  sensitive   = true
+  sensitive = true
 }
 
 variable "render_user_email" {
-  sensitive   = true
+  sensitive = true
 }
 
 provider "render" {
@@ -24,52 +24,52 @@ data "render_owner" "render_owner" {
 }
 
 resource "render_service" "chroma" {
-  name = "chroma"
-  owner = data.render_owner.render_owner.id
-  type = "web_service"
+  name        = "chroma"
+  owner       = data.render_owner.render_owner.id
+  type        = "web_service"
   auto_deploy = true
 
   env_vars = concat([{
-    key = "IS_PERSISTENT"
+    key   = "IS_PERSISTENT"
     value = "1"
-  },
+    },
     {
-      key = "PERSIST_DIRECTORY"
+      key   = "PERSIST_DIRECTORY"
       value = var.chroma_data_volume_mount_path
     },
-  ],
-    var.enable_auth? [{
-      key = "CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER"
+    ],
+    var.enable_auth ? [{
+      key   = "CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER"
       value = "chromadb.auth.token.TokenConfigServerAuthCredentialsProvider"
-    },
+      },
       {
-        key = "CHROMA_SERVER_AUTH_CREDENTIALS"
+        key   = "CHROMA_SERVER_AUTH_CREDENTIALS"
         value = "${local.token_auth_credentials.token}"
       },
       {
-        key = "CHROMA_SERVER_AUTH_PROVIDER"
+        key   = "CHROMA_SERVER_AUTH_PROVIDER"
         value = var.auth_type
-      }] : []
+    }] : []
   )
 
   image = {
-    owner_id = data.render_owner.render_owner.id
+    owner_id   = data.render_owner.render_owner.id
     image_path = "${var.chroma_image_reg_url}:${var.chroma_release}"
   }
 
   web_service_details = {
-    env = "image"
-    plan = var.render_plan
-    region = var.region
+    env               = "image"
+    plan              = var.render_plan
+    region            = var.region
     health_check_path = "/api/v1/heartbeat"
     disk = {
-      name = var.chroma_data_volume_device_name
+      name       = var.chroma_data_volume_device_name
       mount_path = var.chroma_data_volume_mount_path
-      size_gb = var.chroma_data_volume_size
+      size_gb    = var.chroma_data_volume_size
     }
     docker = {
-      command = "uvicorn chromadb.app:app --reload --workers 1 --host 0.0.0.0 --port 80 --log-config chromadb/log_config.yml"
-      path = "./Dockerfile"
+      command = "uvicorn chromadb.app:app --reload --workers 1 --host 0.0.0.0 --port 80 --log-config chromadb/log_config.yml --timeout-keep-alive 30"
+      path    = "./Dockerfile"
     }
   }
 }