diff --git a/docs/credentials.md b/docs/credentials.md index 98391391b00..a3ffb989945 100644 --- a/docs/credentials.md +++ b/docs/credentials.md @@ -23,6 +23,31 @@ Here are the keys that must be set for to access these platforms: - Perspective: `perspectiveApiKey` - Writer: `writerApiKey` +## MongoDB cache credentials + +If you use MongoDB for caching, you can avoid putting credentials in the +`--mongo-uri` command-line argument by storing the URI and credentials in +`credentials.conf`: + +``` +mongoUri: "mongodb://host:27017/cache" +mongoUsername: "username" +mongoPassword: "password" +``` + +You can also provide the same values with environment variables: + +``` +HELM_MONGODB_URI="mongodb://host:27017/cache" +HELM_MONGODB_USERNAME="username" +HELM_MONGODB_PASSWORD="password" +``` + +For backwards compatibility, `--mongo-uri` still accepts a full MongoDB URI. +However, using `credentials.conf` or environment variables avoids exposing +secrets in shell history. HELM redacts MongoDB credentials when logging cache +configuration. + ## Additional setup ### Google @@ -44,4 +69,4 @@ If you are attempting to access models that are private, restricted, or require huggingface-cli login ``` -Refer to [Hugging Face's documentation](https://huggingface.co/docs/huggingface_hub/en/quick-start#login-command) for more information. \ No newline at end of file +Refer to [Hugging Face's documentation](https://huggingface.co/docs/huggingface_hub/en/quick-start#login-command) for more information. diff --git a/src/helm/benchmark/run.py b/src/helm/benchmark/run.py index 6c5218a7cc4..b595fa31159 100644 --- a/src/helm/benchmark/run.py +++ b/src/helm/benchmark/run.py @@ -9,6 +9,7 @@ from helm.benchmark import model_metadata_registry from helm.benchmark.presentation.run_entry import RunEntry, read_run_entries from helm.common.cache_backend_config import MongoCacheBackendConfig, SqliteCacheBackendConfig +from helm.common.cache_backend_config import get_mongo_cache_backend_config from helm.common.general import ensure_directory_exists from helm.common.hierarchical_logger import hlog, htrack, htrack_block, setup_default_logging, hwarn, hdebug from helm.common.authentication import Authentication @@ -173,8 +174,9 @@ def run_benchmarking( mongo_cache_backend_config: Optional[MongoCacheBackendConfig] = None if not disable_cache: - if mongo_uri: - mongo_cache_backend_config = MongoCacheBackendConfig(mongo_uri) + resolved_mongo_cache_backend_config = get_mongo_cache_backend_config(mongo_uri, local_path) + if resolved_mongo_cache_backend_config: + mongo_cache_backend_config = resolved_mongo_cache_backend_config else: sqlite_cache_path = os.path.join(local_path, CACHE_DIR) ensure_directory_exists(sqlite_cache_path) diff --git a/src/helm/common/cache.py b/src/helm/common/cache.py index a2686bd13ed..dd512a7a4c8 100644 --- a/src/helm/common/cache.py +++ b/src/helm/common/cache.py @@ -1,6 +1,6 @@ from collections import defaultdict from dataclasses import dataclass -from typing import Dict, Callable, Generator, Mapping, Tuple +from typing import Dict, Callable, Generator, Mapping, Optional, Tuple import json import threading @@ -79,9 +79,26 @@ class MongoCacheConfig(KeyValueStoreCacheConfig): # Name of the MongoDB collection. collection_name: str + # Optional credentials supplied separately from the URI. + username: Optional[str] = None + password: Optional[str] = None + + def __repr__(self) -> str: + from helm.common.cache_backend_config import redact_mongo_uri + + username = "" if self.username else None + password = "" if self.password else None + return ( + f"MongoCacheConfig(uri={redact_mongo_uri(self.uri)!r}, " + f"collection_name={self.collection_name!r}, " + f"username={username!r}, password={password!r})" + ) + @property def cache_stats_key(self) -> str: - return f"{self.uri}/{self.collection_name}" + from helm.common.cache_backend_config import redact_mongo_uri + + return f"{redact_mongo_uri(self.uri)}/{self.collection_name}" def get_all_from_sqlite(path: str) -> Generator[Tuple[Dict, Dict], None, None]: @@ -108,7 +125,7 @@ def create_key_value_store(config: KeyValueStoreCacheConfig) -> KeyValueStore: if isinstance(config, MongoCacheConfig): from helm.common.mongo_key_value_store import MongoKeyValueStore - return MongoKeyValueStore(config.uri, config.collection_name) + return MongoKeyValueStore(config.uri, config.collection_name, config.username, config.password) elif isinstance(config, SqliteCacheConfig): return SqliteKeyValueStore(config.path) elif isinstance(config, BlackHoleCacheConfig): diff --git a/src/helm/common/cache_backend_config.py b/src/helm/common/cache_backend_config.py index 4253224ec43..57b4a8bdf2c 100644 --- a/src/helm/common/cache_backend_config.py +++ b/src/helm/common/cache_backend_config.py @@ -1,8 +1,33 @@ from abc import ABC, abstractmethod from dataclasses import dataclass import os +from typing import Optional +from urllib.parse import SplitResult, urlsplit, urlunsplit from helm.common.cache import CacheConfig, MongoCacheConfig, BlackHoleCacheConfig, SqliteCacheConfig +from helm.common.general import get_credentials + + +MONGO_URI_ENV_NAME = "HELM_MONGODB_URI" +MONGO_USERNAME_ENV_NAME = "HELM_MONGODB_USERNAME" +MONGO_PASSWORD_ENV_NAME = "HELM_MONGODB_PASSWORD" + +MONGO_URI_CREDENTIALS_KEY = "mongoUri" +MONGO_USERNAME_CREDENTIALS_KEY = "mongoUsername" +MONGO_PASSWORD_CREDENTIALS_KEY = "mongoPassword" + + +def redact_mongo_uri(uri: str) -> str: + """Return a MongoDB URI with credentials removed from the authority.""" + parsed: SplitResult = urlsplit(uri) + if not parsed.username and not parsed.password: + return uri + + hostname = parsed.hostname or "" + netloc = f"[{hostname}]" if ":" in hostname else hostname + if parsed.port is not None: + netloc = f"{netloc}:{parsed.port}" + return urlunsplit((parsed.scheme, netloc, parsed.path, parsed.query, parsed.fragment)) class CacheBackendConfig(ABC): @@ -24,8 +49,49 @@ class MongoCacheBackendConfig(CacheBackendConfig): Example format: mongodb://[username:password@]host1[:port1]/[dbname] For full format, see: https://www.mongodb.com/docs/manual/reference/connection-string/""" + username: Optional[str] = None + """Optional MongoDB username passed separately from the URI.""" + + password: Optional[str] = None + """Optional MongoDB password passed separately from the URI.""" + + def __repr__(self) -> str: + username = "" if self.username else None + password = "" if self.password else None + return ( + f"MongoCacheBackendConfig(uri={redact_mongo_uri(self.uri)!r}, " + f"username={username!r}, password={password!r})" + ) + def get_cache_config(self, shard_name: str) -> CacheConfig: - return MongoCacheConfig(uri=self.uri, collection_name=shard_name) + return MongoCacheConfig( + uri=self.uri, + collection_name=shard_name, + username=self.username, + password=self.password, + ) + + +def get_mongo_cache_backend_config( + mongo_uri: Optional[str], + base_path: str, +) -> Optional[MongoCacheBackendConfig]: + """Resolve MongoDB cache settings from CLI args, env vars, and credentials.conf.""" + credentials = get_credentials(base_path) + resolved_uri = mongo_uri or os.getenv(MONGO_URI_ENV_NAME) or credentials.get(MONGO_URI_CREDENTIALS_KEY, None) + if not resolved_uri: + return None + + username = os.getenv(MONGO_USERNAME_ENV_NAME) or credentials.get(MONGO_USERNAME_CREDENTIALS_KEY, None) + password = os.getenv(MONGO_PASSWORD_ENV_NAME) or credentials.get(MONGO_PASSWORD_CREDENTIALS_KEY, None) + if bool(username) != bool(password): + raise ValueError( + f"Set both {MONGO_USERNAME_ENV_NAME} and {MONGO_PASSWORD_ENV_NAME}, " + f"or both {MONGO_USERNAME_CREDENTIALS_KEY} and {MONGO_PASSWORD_CREDENTIALS_KEY} " + "in credentials.conf." + ) + + return MongoCacheBackendConfig(uri=resolved_uri, username=username, password=password) @dataclass(frozen=True) diff --git a/src/helm/common/mongo_key_value_store.py b/src/helm/common/mongo_key_value_store.py index 43a16d493fc..90f8ad8d9c1 100644 --- a/src/helm/common/mongo_key_value_store.py +++ b/src/helm/common/mongo_key_value_store.py @@ -21,9 +21,15 @@ class MongoKeyValueStore(KeyValueStore): _REQUEST_KEY = "request" _RESPONSE_KEY = "response" - def __init__(self, uri: str, collection_name: str): + def __init__( + self, + uri: str, + collection_name: str, + username: Optional[str] = None, + password: Optional[str] = None, + ): # TODO: Create client in __enter__ and clean up client in __exit__ - self._mongodb_client: MongoClient = MongoClient(uri) + self._mongodb_client: MongoClient = MongoClient(uri, username=username, password=password) self._database = self._mongodb_client.get_default_database() self._collection = self._database.get_collection(collection_name) self._collection.create_index(self._REQUEST_KEY, unique=True) diff --git a/src/helm/common/test_cache_backend_config.py b/src/helm/common/test_cache_backend_config.py new file mode 100644 index 00000000000..9a387b6215d --- /dev/null +++ b/src/helm/common/test_cache_backend_config.py @@ -0,0 +1,73 @@ +import os +import tempfile + +import pytest + +from helm.common.cache import MongoCacheConfig +from helm.common.cache_backend_config import ( + MongoCacheBackendConfig, + get_mongo_cache_backend_config, + redact_mongo_uri, +) + + +def test_redact_mongo_uri_removes_password(): + uri = "mongodb://user:password@example.com:27017/cache?retryWrites=true" + + assert redact_mongo_uri(uri) == "mongodb://example.com:27017/cache?retryWrites=true" + + +def test_mongo_cache_backend_config_repr_redacts_password(): + config = MongoCacheBackendConfig( + uri="mongodb://user:password@example.com/cache", + username="user", + password="password", + ) + + assert "mongodb://user:password@" not in repr(config) + assert "password@example.com" not in repr(config) + assert "example.com" in repr(config) + + +def test_mongo_cache_config_redacts_password(): + config = MongoCacheConfig( + uri="mongodb://user:password@example.com/cache", + collection_name="requests", + username="user", + password="password", + ) + + assert "user:password@" not in repr(config) + assert config.cache_stats_key == "mongodb://example.com/cache/requests" + + +def test_get_mongo_cache_backend_config_from_credentials_file(monkeypatch): + monkeypatch.delenv("HELM_MONGODB_URI", raising=False) + monkeypatch.delenv("HELM_MONGODB_USERNAME", raising=False) + monkeypatch.delenv("HELM_MONGODB_PASSWORD", raising=False) + with tempfile.TemporaryDirectory() as base_path: + with open(os.path.join(base_path, "credentials.conf"), "w") as f: + f.write( + """ + mongoUri: "mongodb://example.com/cache" + mongoUsername: "user" + mongoPassword: "password" + """ + ) + + config = get_mongo_cache_backend_config(None, base_path) + + assert config == MongoCacheBackendConfig( + uri="mongodb://example.com/cache", + username="user", + password="password", + ) + + +def test_get_mongo_cache_backend_config_requires_username_and_password(monkeypatch): + monkeypatch.delenv("HELM_MONGODB_PASSWORD", raising=False) + monkeypatch.setenv("HELM_MONGODB_URI", "mongodb://example.com/cache") + monkeypatch.setenv("HELM_MONGODB_USERNAME", "user") + + with pytest.raises(ValueError, match="Set both HELM_MONGODB_USERNAME and HELM_MONGODB_PASSWORD"): + get_mongo_cache_backend_config(None, "prod_env") diff --git a/src/helm/proxy/server.py b/src/helm/proxy/server.py index b1a6ce554ac..31228fa0e27 100644 --- a/src/helm/proxy/server.py +++ b/src/helm/proxy/server.py @@ -21,7 +21,8 @@ ) from helm.benchmark.model_deployment_registry import get_default_model_deployment_for_model from helm.common.authentication import Authentication -from helm.common.cache_backend_config import CacheBackendConfig, MongoCacheBackendConfig, SqliteCacheBackendConfig +from helm.common.cache_backend_config import CacheBackendConfig, SqliteCacheBackendConfig +from helm.common.cache_backend_config import get_mongo_cache_backend_config from helm.common.general import ensure_directory_exists from helm.common.hierarchical_logger import hlog, setup_default_logging from helm.common.optional_dependencies import handle_module_not_found_error @@ -279,8 +280,9 @@ def main(): register_configs_from_directory(args.base_path) cache_backend_config: CacheBackendConfig - if args.mongo_uri: - cache_backend_config = MongoCacheBackendConfig(args.mongo_uri) + resolved_mongo_cache_backend_config = get_mongo_cache_backend_config(args.mongo_uri, args.base_path) + if resolved_mongo_cache_backend_config: + cache_backend_config = resolved_mongo_cache_backend_config else: sqlite_cache_path = os.path.join(args.base_path, CACHE_DIR) ensure_directory_exists(sqlite_cache_path)