benchmark/benchmark_cloud_load.py
Rama main · réplica de lectura del repositorio original.
#!/usr/bin/env python3
"""Databend Cloud load benchmark runner backed by bendsql."""
import json
import logging
import os
import shutil
import subprocess
import sys
import time
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Optional
logging.basicConfig(level=logging.INFO, format="%(message)s")
logger = logging.getLogger(__name__)
@dataclass
class BenchmarkConfig:
benchmark_id: str
dataset: str
size: str
cache_size: str
version: str
database: str
tries: int
user: str
password: str
gateway: str
warehouse: str
source: str
source_id: str
sha: str
@dataclass
class ResultRecord:
date: str
dataset: str
database: str
version: str
warehouse: str
machine: str
tags: List[str]
result: List[List[float]]
values: Dict[str, List[float]]
run_id: str
size: str
tries: int
storage: str
cache_size: str
load_time: Optional[float] = None
data_size: Optional[int] = None
system: Optional[str] = None
comment: Optional[str] = None
class BendSQLRunner:
def __init__(self) -> None:
self._env = os.environ.copy()
def set_dsn(self, dsn: str) -> None:
self._env["BENDSQL_DSN"] = dsn
logger.debug("Using DSN: %s", dsn)
def run(
self,
args: Optional[List[str]] = None,
*,
sql: Optional[str] = None,
capture_output: bool = False,
) -> subprocess.CompletedProcess:
command = ["bendsql"] + (args or [])
logger.debug("Running command: %s", " ".join(command))
try:
result = subprocess.run(
command,
input=sql,
text=True,
env=self._env,
capture_output=True, # Always capture output
check=True,
)
# If not capturing output, print it to stdout/stderr
if not capture_output:
if result.stdout:
print(result.stdout, end="", file=sys.stdout, flush=True)
if result.stderr:
print(result.stderr, end="", file=sys.stderr, flush=True)
return result
except subprocess.CalledProcessError as exc: # pragma: no cover - passthrough
stdout = exc.stdout.strip() if exc.stdout else ""
stderr = exc.stderr.strip() if exc.stderr else ""
logger.error("bendsql failed: %s", stderr or stdout)
raise
def load_config() -> BenchmarkConfig:
benchmark_id = os.environ.get("BENCHMARK_ID", str(int(time.time())))
dataset = os.environ.get("BENCHMARK_DATASET", "load")
size = os.environ.get("BENCHMARK_SIZE", "Small")
raw_cache_size = os.environ.get("BENCHMARK_CACHE_SIZE", "")
cache_size = raw_cache_size.strip() or "0"
version = os.environ.get("BENCHMARK_VERSION", "")
database = os.environ.get("BENCHMARK_DATABASE", "default")
tries_raw = os.environ.get("BENCHMARK_TRIES", "3")
source = os.environ.get("BENCHMARK_SOURCE", "")
source_id = os.environ.get("BENCHMARK_SOURCE_ID", "")
sha = os.environ.get("BENCHMARK_SHA", "")
if dataset != "load":
logger.error("benchmark_cloud_load.py only supports BENCHMARK_DATASET=load")
sys.exit(1)
if not version:
logger.error("Please set BENCHMARK_VERSION to run the benchmark.")
sys.exit(1)
try:
tries = int(tries_raw)
except ValueError:
logger.error("BENCHMARK_TRIES must be an integer, got %s", tries_raw)
sys.exit(1)
if not 1 <= tries <= 3:
logger.error("BENCHMARK_TRIES must be between 1 and 3, got %s", tries)
sys.exit(1)
user = os.environ.get("CLOUD_USER", "")
password = os.environ.get("CLOUD_PASSWORD", "")
gateway = os.environ.get("CLOUD_GATEWAY", "")
warehouse = os.environ.get("CLOUD_WAREHOUSE", f"benchmark-{benchmark_id}")
if not user or not password or not gateway:
logger.error(
"Please set CLOUD_USER, CLOUD_PASSWORD and CLOUD_GATEWAY to run the benchmark.",
)
sys.exit(1)
return BenchmarkConfig(
benchmark_id=benchmark_id,
dataset=dataset,
size=size,
cache_size=cache_size,
version=version,
database=database,
tries=tries,
user=user,
password=password,
gateway=gateway,
warehouse=warehouse,
source=source,
source_id=source_id,
sha=sha,
)
def ensure_dependencies() -> None:
if not shutil.which("bendsql"):
logger.error("bendsql is required but was not found in PATH.")
sys.exit(1)
logger.info("Checking script dependencies...")
logger.info("bendsql version: %s", subprocess.check_output(["bendsql", "--version"]).decode().strip())
SIZE_MAPPING: Dict[str, str] = {
"Small": "Small",
"Large": "Large",
}
def build_dsn(
config: BenchmarkConfig,
*,
database: Optional[str] = None,
warehouse: Optional[str] = None,
login_disable: bool = False,
) -> str:
params = []
if login_disable:
params.append("login=disable")
if warehouse:
params.append(f"warehouse={warehouse}")
query = f"?{'&'.join(params)}" if params else ""
db_path = f"/{database}" if database else ""
return f"databend://{config.user}:{config.password}@{config.gateway}:443{db_path}{query}"
def quote_literal(value: str) -> str:
return value.replace("'", "''")
def wait_for_warehouse(runner: BendSQLRunner, warehouse: str, retries: int = 20, delay: int = 10) -> None:
logger.info("Waiting for warehouse %s to be ready...", warehouse)
for attempt in range(retries + 1):
completed = runner.run(
["--output", "csv"],
sql=f"SHOW WAREHOUSES LIKE '{quote_literal(warehouse)}'",
capture_output=True,
)
output = completed.stdout.strip()
if output and "Running" in output:
logger.info("Warehouse %s is running.", warehouse)
return
logger.info("Warehouse not ready yet. Sleeping %s seconds...", delay)
time.sleep(delay)
logger.error("Failed to start warehouse %s in time.", warehouse)
sys.exit(1)
def execute_sql_file(runner: BendSQLRunner, path: Path) -> None:
sql = path.read_text()
if sql.strip():
runner.run(sql=sql)
def parse_time_output(raw_value: str) -> Optional[float]:
# --time=server outputs only the time value (e.g., "0.014")
value = raw_value.strip()
if not value:
return None
try:
return float(value)
except ValueError:
return None
def run_timed_query(runner: BendSQLRunner, sql: str) -> Optional[float]:
completed = runner.run(["--time=server"], sql=sql, capture_output=True)
output = completed.stdout.strip()
return parse_time_output(output)
def pad_attempts(values: List[float], target_size: int) -> None:
if not values:
return
while len(values) < target_size:
values.append(values[-1])
def write_result_files(script_dir: Path, record: ResultRecord) -> None:
result_path = script_dir / "result.json"
cache_suffix = f"-cache-{record.cache_size}" if record.cache_size else ""
final_result_path = script_dir / f"result-{record.dataset}-cloud-{record.size}{cache_suffix}.json"
ndjson_name = (
f"result-{record.dataset}-cloud-{record.size}{cache_suffix}"
f"-{record.run_id}.ndjson"
)
ndjson_path = script_dir / ndjson_name
payload = {key: value for key, value in asdict(record).items() if value is not None}
result_path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
final_result_path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
ndjson_path.write_text(json.dumps(payload) + "\n", encoding="utf-8")
logger.info("Wrote JSON results to %s and %s", result_path, final_result_path)
logger.info("Wrote NDJSON results to %s", ndjson_path)
def main() -> None:
config = load_config()
ensure_dependencies()
logger.info("#######################################################")
logger.info("Running load benchmark for Databend Cloud with %s storage...", "S3")
script_dir = Path(__file__).resolve().parent
dataset_dir = script_dir / "load"
if not dataset_dir.exists():
logger.error("Dataset directory %s does not exist", dataset_dir)
sys.exit(1)
if config.size not in SIZE_MAPPING:
logger.error("Unsupported benchmark size: %s", config.size)
sys.exit(1)
run_date = datetime.now(timezone.utc).strftime("%Y-%m-%d")
tags = ["s3", f"cache-{config.cache_size}"]
machine = SIZE_MAPPING[config.size]
system: Optional[str] = None
comment: Optional[str] = None
if config.source and config.source_id:
if config.source == "pr":
system = f"Databend(PR#{config.source_id})"
elif config.source == "release":
system = f"Databend(Release@{config.source_id})"
else:
logger.error("Unsupported benchmark source: %s", config.source)
sys.exit(1)
elif config.source or config.source_id:
logger.error("Both BENCHMARK_SOURCE and BENCHMARK_SOURCE_ID must be provided together.")
sys.exit(1)
if config.sha:
comment = f"commit: {config.sha}"
record = ResultRecord(
date=run_date,
dataset=config.dataset,
database=config.database,
version=config.version,
warehouse=config.warehouse,
machine=machine,
tags=tags,
result=[],
values={},
run_id=config.benchmark_id,
size=config.size,
tries=config.tries,
storage="s3",
cache_size=config.cache_size,
system=system,
comment=comment,
)
runner = BendSQLRunner()
runner.set_dsn(build_dsn(config, warehouse="default", login_disable=True))
warehouse_literal = quote_literal(config.warehouse)
logger.info("Creating warehouse %s...", config.warehouse)
runner.run(sql=f"DROP WAREHOUSE IF EXISTS '{warehouse_literal}';")
runner.run(
sql=(
f"CREATE WAREHOUSE '{warehouse_literal}' "
f"WITH version='{config.version}' "
f"warehouse_size='{config.size}' "
f"cache_size={config.cache_size};"
),
)
runner.run(["--output", "table"], sql="SHOW WAREHOUSES;")
wait_for_warehouse(runner, config.warehouse)
runner.set_dsn(build_dsn(config, warehouse=config.warehouse))
logger.info("Creating database %s for load dataset...", config.database)
runner.run(["--database", "default"], sql=f"CREATE DATABASE {config.database};")
runner.set_dsn(build_dsn(config, database=config.database, warehouse=config.warehouse))
logger.info("Checking session settings...")
runner.run(
["--output", "table"],
sql="select * from system.settings where value != default;",
)
analyze_sql = dataset_dir / "analyze.sql"
if analyze_sql.exists():
logger.info("Analyze tables...")
execute_sql_file(runner, analyze_sql)
logger.info("Running queries...")
queries = sorted(dataset_dir.glob("queries/*.sql"))
if not queries:
logger.error("No queries found under %s", dataset_dir)
sys.exit(1)
for query_num, query_file in enumerate(queries):
logger.info("==> Running Q%s: %s", query_num, query_file)
query_sql = query_file.read_text()
record.result.append([])
record.values[f"Q{query_num}"] = []
for attempt in range(1, config.tries + 1):
try:
q_time = run_timed_query(runner, query_sql)
except subprocess.CalledProcessError:
logger.error("Q%s[%s] failed", query_num, attempt)
continue
if q_time is None:
logger.error("Q%s[%s] returned no timing info", query_num, attempt)
continue
logger.info("Q%s[%s] succeeded in %.3f seconds", query_num, attempt, q_time)
record.result[query_num].append(round(q_time, 3))
record.values[f"Q{query_num}"].append(round(q_time, 3))
if config.tries < 3:
pad_attempts(record.result[query_num], 3)
pad_attempts(record.values[f"Q{query_num}"], 3)
cleanup_runner = BendSQLRunner()
cleanup_runner.set_dsn(build_dsn(config, warehouse="default", login_disable=True))
try:
logger.info("Dropping database %s...", config.database)
cleanup_runner.run(sql=f"DROP DATABASE IF EXISTS {config.database};")
logger.info("Dropping warehouse %s...", config.warehouse)
cleanup_runner.run(sql=f"DROP WAREHOUSE IF EXISTS '{warehouse_literal}';")
finally:
write_result_files(script_dir, record)
if __name__ == "__main__":
main()