From afd33e0f43779d8ae39e27f33ca3a09500833679 Mon Sep 17 00:00:00 2001 From: Akanksha Akkihal Date: Tue, 28 Jul 2026 11:46:44 -0700 Subject: [PATCH 1/4] fix(tools): batch ClickHouse JSON-lines bulk loads Bulk-loading a JSON-lines dataset piped the entire file through a single INSERT ... FORMAT JSONEachRow, which is unreliable for the multi-GB datasets the ClickHouse benchmark now uses. Load in bounded batches via a dedicated loader script instead, poll ClickHouse HTTP until it is actually reachable before loading, and validate each line so malformed input fails with the offending line number rather than an opaque client error. Also let init_sql_file own DROP/CREATE for its own objects, so a schema that defines dependent materialized views is not broken by a standalone DROP TABLE. Co-authored-by: Cursor --- .../services/clickhouse_service.py | 169 ++++++++++-------- .../experiment_utils/services/json/loader.py | 151 ++++++++++++++++ 2 files changed, 250 insertions(+), 70 deletions(-) create mode 100644 asap-tools/experiments/experiment_utils/services/json/loader.py diff --git a/asap-tools/experiments/experiment_utils/services/clickhouse_service.py b/asap-tools/experiments/experiment_utils/services/clickhouse_service.py index a423c94..6e9c6eb 100644 --- a/asap-tools/experiments/experiment_utils/services/clickhouse_service.py +++ b/asap-tools/experiments/experiment_utils/services/clickhouse_service.py @@ -5,6 +5,7 @@ import os import shlex import subprocess +import time from typing import Optional from jinja2 import Template @@ -192,8 +193,10 @@ class ClickHouseDataLoaderService(BaseService): # H2O loader script (relative to _ASSETS_DIR). H2O_LOADER_SCRIPT = "h2o/loader.py" + JSON_LOADER_SCRIPT = "json/loader.py" H2O_BATCH_SIZE = 50_000 + JSON_BATCH_SIZE = 100_000 DEFAULT_TABLES = { "clickbench": "hits", @@ -287,13 +290,15 @@ def start( url = f"http://localhost:{self.clickhouse_http_port}/" - print(f"Dropping table {table!r}...") - self._exec_sql(f"DROP TABLE IF EXISTS {table}", url) - if init_sql_file is not None: + # init SQL owns DROP/CREATE for the raw table and any MVs / agg + # tables (e.g. netflow_init.sql). Skipping the standalone DROP + # avoids failing on dependents that still reference ``table``. print(f"Running init SQL from {init_sql_file!r}...") self._exec_sql_file(init_sql_file, url) elif dataset_name in self.BUILTIN_DDL_FILES: + print(f"Dropping table {table!r}...") + self._exec_sql(f"DROP TABLE IF EXISTS {table}", url) local_ddl = os.path.join( self._ASSETS_DIR, self.BUILTIN_DDL_FILES[dataset_name] ) @@ -309,6 +314,8 @@ def start( f"No built-in DDL for dataset_name={dataset_name!r}; pass init_sql_file" ) + self._ensure_clickhouse_http_ready(url) + if dataset_name == "clickbench": self._load_clickbench(remote_data_file, url, table, max_rows) elif dataset_name == "h2o": @@ -403,6 +410,84 @@ def _exec_sql_file(self, remote_sql_file: str, url: str) -> None: for stmt in (s.strip() for s in result.stdout.split(";") if s.strip()): self._exec_sql(stmt, url) + def _ensure_clickhouse_http_ready(self, url: str, max_retries: int = 30) -> None: + """Wait until ClickHouse HTTP /ping returns Ok. before bulk load.""" + ping_url = url.rstrip("/") + "/ping" + for attempt in range(max_retries): + result = self.provider.execute_command( + node_idx=self.node_offset, + cmd=f"curl -sS {shlex.quote(ping_url)}", + cmd_dir=None, + nohup=False, + popen=False, + ignore_errors=True, + ) + if ( + isinstance(result, subprocess.CompletedProcess) + and result.returncode == 0 + and result.stdout.strip() == "Ok." + ): + return + print( + f"Waiting for ClickHouse HTTP before data load... " + f"({attempt + 1}/{max_retries})" + ) + time.sleep(2) + + logs = self.provider.execute_command( + node_idx=self.node_offset, + cmd="docker logs clickhouse-server --tail 30 2>&1", + cmd_dir=None, + nohup=False, + popen=False, + ignore_errors=True, + ) + log_tail = "" + if isinstance(logs, subprocess.CompletedProcess): + log_tail = (logs.stdout or logs.stderr or "").strip()[-500:] + raise RuntimeError( + "ClickHouse HTTP is not ready before data load. " + f"Check clickhouse_logs/ on the experiment node. Recent logs: {log_tail}" + ) + + def _load_json_batched( + self, + remote_data_file: str, + url: str, + table: str, + max_rows: int, + dataset_label: str, + batch_size: int = JSON_BATCH_SIZE, + ) -> None: + """Load JSON-lines via batched HTTP INSERTs (safe for multi-GB files).""" + local_script = os.path.join(self._ASSETS_DIR, self.JSON_LOADER_SCRIPT) + remote_script = f"/tmp/json_loader_{os.getpid()}.py" + self._rsync_to_remote(local_script, remote_script) + try: + cmd = ( + "python3 {} --data-file {} --table {} " + "--batch-size {} --max-rows {} --container {}" + ).format( + shlex.quote(remote_script), + shlex.quote(remote_data_file), + shlex.quote(table), + batch_size, + max_rows, + shlex.quote(ClickHouseService.CONTAINER_NAME), + ) + result = self.provider.execute_command( + node_idx=self.node_offset, + cmd=cmd, + cmd_dir=None, + nohup=False, + popen=False, + ) + if isinstance(result, subprocess.CompletedProcess) and result.returncode != 0: + detail = (result.stderr or result.stdout or "").strip()[:300] + raise RuntimeError(f"{dataset_label} data load failed: {detail}") + finally: + self._remote_rm(remote_script) + def _check_row_count(self, table: str, url: str) -> int: """Return the row count for a table on the remote node, or 0 on error.""" cmd = "curl -sS {} --data {}".format( @@ -429,40 +514,9 @@ def _load_clickbench( ) -> None: """Stream a JSON-lines file (optionally gzipped) into ClickHouse.""" print(f"Loading ClickBench data from {remote_data_file!r}...") - file_lower = remote_data_file.lower() - is_gz = file_lower.endswith(".json.gz") or file_lower.endswith(".jsonl.gz") - insert_sql = shlex.quote(f"INSERT INTO {table} FORMAT JSONEachRow") - - if is_gz: - if max_rows > 0: - reader = "zcat {} | head -n {}".format( - shlex.quote(remote_data_file), max_rows - ) - else: - reader = "zcat {}".format(shlex.quote(remote_data_file)) - else: - if max_rows > 0: - reader = "head -n {} {}".format(max_rows, shlex.quote(remote_data_file)) - else: - reader = "cat {}".format(shlex.quote(remote_data_file)) - - cmd = ( - "{} | docker exec -i clickhouse-server clickhouse-client --query {}".format( - reader, insert_sql - ) - ) - - result = self.provider.execute_command( - node_idx=self.node_offset, - cmd=cmd, - cmd_dir=None, - nohup=False, - popen=False, + self._load_json_batched( + remote_data_file, url, table, max_rows, "ClickBench" ) - if isinstance(result, subprocess.CompletedProcess) and result.returncode != 0: - raise RuntimeError( - f"ClickBench data load failed: {result.stderr.strip()[:200]}" - ) def _load_h2o( self, remote_data_file: str, url: str, batch_size: int, max_rows: int @@ -503,42 +557,17 @@ def _load_custom( """Stream a custom JSON-lines file (plain or gzipped) into ClickHouse.""" print(f"Loading custom data from {remote_data_file!r} into {table!r}...") file_lower = remote_data_file.lower() - is_gz = file_lower.endswith(".json.gz") or file_lower.endswith(".jsonl.gz") - is_json = file_lower.endswith(".json") or file_lower.endswith(".jsonl") - insert_sql = shlex.quote(f"INSERT INTO {table} FORMAT JSONEachRow") - - if is_gz: - if max_rows > 0: - reader = "zcat {} | head -n {}".format( - shlex.quote(remote_data_file), max_rows - ) - else: - reader = "zcat {}".format(shlex.quote(remote_data_file)) - elif is_json: - if max_rows > 0: - reader = "head -n {} {}".format(max_rows, shlex.quote(remote_data_file)) - else: - reader = "cat {}".format(shlex.quote(remote_data_file)) - else: + is_json = ( + file_lower.endswith(".json") + or file_lower.endswith(".jsonl") + or file_lower.endswith(".json.gz") + or file_lower.endswith(".jsonl.gz") + ) + if not is_json: raise ValueError( f"Unsupported file format for {remote_data_file!r}. " "Use dataset_name='h2o' for CSV files." ) - - cmd = ( - "{} | docker exec -i clickhouse-server clickhouse-client --query {}".format( - reader, insert_sql - ) + self._load_json_batched( + remote_data_file, url, table, max_rows, "Custom" ) - - result = self.provider.execute_command( - node_idx=self.node_offset, - cmd=cmd, - cmd_dir=None, - nohup=False, - popen=False, - ) - if isinstance(result, subprocess.CompletedProcess) and result.returncode != 0: - raise RuntimeError( - f"Custom data load failed: {result.stderr.strip()[:200]}" - ) diff --git a/asap-tools/experiments/experiment_utils/services/json/loader.py b/asap-tools/experiments/experiment_utils/services/json/loader.py new file mode 100644 index 0000000..fb9dac5 --- /dev/null +++ b/asap-tools/experiments/experiment_utils/services/json/loader.py @@ -0,0 +1,151 @@ +#!/usr/bin/env python3 +""" +Load a JSON-lines file into ClickHouse in bounded-size batches. + +Uses ``docker exec … clickhouse-client`` on the experiment node so multi-GB +files stream without loading the whole payload into curl/HTTP memory. + +Usage: + python3 loader.py \\ + --data-file /path/to/netflow.jsonl \\ + --table netflow_table \\ + --batch-size 100000 \\ + --max-rows 0 +""" + +import argparse +import gzip +import json +import subprocess +import sys + + +DEFAULT_CONTAINER = "clickhouse-server" +PROGRESS_ROW_INTERVAL = 500_000 + + +def _open_data_file(path: str): + lower = path.lower() + if lower.endswith(".gz"): + return gzip.open(path, "rt", encoding="utf-8") + return open(path, "r", encoding="utf-8") + + +def _validate_line(line: str, line_no: int) -> str: + stripped = line.strip() + if not stripped: + return "" + if "\0" in stripped: + preview = stripped[:80].replace("\0", "\\0") + raise RuntimeError( + f"Null byte in JSON line {line_no} (file may be corrupt — " + f"regenerate on a native Linux filesystem, not WSL /mnt/c): {preview!r}" + ) + try: + json.loads(stripped) + except json.JSONDecodeError as exc: + preview = stripped[:120] + raise RuntimeError( + f"Invalid JSON on line {line_no}: {exc}. Preview: {preview!r}" + ) from exc + return stripped + + +def flush_batch( + table: str, + lines: list[str], + container: str, +) -> None: + body = "\n".join(lines).encode("utf-8") + result = subprocess.run( + [ + "docker", + "exec", + "-i", + container, + "clickhouse-client", + "--query", + f"INSERT INTO {table} FORMAT JSONEachRow", + ], + input=body, + capture_output=True, + ) + if result.returncode != 0: + detail = (result.stderr or result.stdout or b"").decode( + "utf-8", errors="replace" + )[:500] + raise RuntimeError(f"ClickHouse insert failed: {detail}") + + +def load( + data_file: str, + table: str, + batch_size: int, + max_rows: int, + container: str = DEFAULT_CONTAINER, +) -> None: + batch: list[str] = [] + total = 0 + file_line_no = 0 + next_progress_report = PROGRESS_ROW_INTERVAL + + def flush(lines: list[str]) -> None: + nonlocal total + if not lines: + return + if max_rows > 0: + lines = lines[: max_rows - total] + if not lines: + return + flush_batch(table, lines, container) + total += len(lines) + + with _open_data_file(data_file) as fin: + for raw_line in fin: + file_line_no += 1 + if max_rows > 0 and total >= max_rows: + break + stripped = _validate_line(raw_line, file_line_no) + if not stripped: + continue + batch.append(stripped) + if len(batch) >= batch_size: + flush(batch) + batch = [] + if total >= next_progress_report: + print(f" Inserted {total:,} rows...", flush=True) + next_progress_report = total + PROGRESS_ROW_INTERVAL + + flush(batch) + print(f"JSON load complete: {total:,} rows into {table!r}") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--data-file", required=True, help="Path to JSON-lines file") + parser.add_argument("--table", required=True, help="Target table name") + parser.add_argument( + "--batch-size", type=int, default=100_000, help="Rows per INSERT batch" + ) + parser.add_argument( + "--max-rows", type=int, default=0, help="Max rows to load (0 = all)" + ) + parser.add_argument( + "--container", + default=DEFAULT_CONTAINER, + help=f"ClickHouse docker container name (default: {DEFAULT_CONTAINER})", + ) + args = parser.parse_args() + if args.batch_size < 1: + parser.error("--batch-size must be >= 1") + try: + load(args.data_file, args.table, args.batch_size, args.max_rows, args.container) + except (RuntimeError, OSError) as exc: + # The caller surfaces only the head of stderr, so print the message + # itself rather than letting a traceback bury it. + print(f"ERROR: {exc}", file=sys.stderr) + sys.exit(1) + + +if __name__ == "__main__": + main() From 089de18428003350552b1ca77810bb1177df2597 Mon Sep 17 00:00:00 2001 From: Akanksha Akkihal Date: Tue, 28 Jul 2026 12:11:54 -0700 Subject: [PATCH 2/4] docs(tools): correct _load_json_batched docstring --- .../experiments/experiment_utils/services/clickhouse_service.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/asap-tools/experiments/experiment_utils/services/clickhouse_service.py b/asap-tools/experiments/experiment_utils/services/clickhouse_service.py index 6e9c6eb..95014e7 100644 --- a/asap-tools/experiments/experiment_utils/services/clickhouse_service.py +++ b/asap-tools/experiments/experiment_utils/services/clickhouse_service.py @@ -459,7 +459,7 @@ def _load_json_batched( dataset_label: str, batch_size: int = JSON_BATCH_SIZE, ) -> None: - """Load JSON-lines via batched HTTP INSERTs (safe for multi-GB files).""" + """Load JSON-lines in bounded batches via clickhouse-client.""" local_script = os.path.join(self._ASSETS_DIR, self.JSON_LOADER_SCRIPT) remote_script = f"/tmp/json_loader_{os.getpid()}.py" self._rsync_to_remote(local_script, remote_script) From 726640aadd446b3d7ce73a2b6daed86fabbb6b8d Mon Sep 17 00:00:00 2001 From: Akanksha Akkihal Date: Tue, 28 Jul 2026 12:21:20 -0700 Subject: [PATCH 3/4] docs(tools): correct _load_json_batched docstring and apply black formatting --- .../experiment_utils/services/clickhouse_service.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/asap-tools/experiments/experiment_utils/services/clickhouse_service.py b/asap-tools/experiments/experiment_utils/services/clickhouse_service.py index 95014e7..b3e4f6a 100644 --- a/asap-tools/experiments/experiment_utils/services/clickhouse_service.py +++ b/asap-tools/experiments/experiment_utils/services/clickhouse_service.py @@ -482,7 +482,10 @@ def _load_json_batched( nohup=False, popen=False, ) - if isinstance(result, subprocess.CompletedProcess) and result.returncode != 0: + if ( + isinstance(result, subprocess.CompletedProcess) + and result.returncode != 0 + ): detail = (result.stderr or result.stdout or "").strip()[:300] raise RuntimeError(f"{dataset_label} data load failed: {detail}") finally: @@ -514,9 +517,7 @@ def _load_clickbench( ) -> None: """Stream a JSON-lines file (optionally gzipped) into ClickHouse.""" print(f"Loading ClickBench data from {remote_data_file!r}...") - self._load_json_batched( - remote_data_file, url, table, max_rows, "ClickBench" - ) + self._load_json_batched(remote_data_file, url, table, max_rows, "ClickBench") def _load_h2o( self, remote_data_file: str, url: str, batch_size: int, max_rows: int @@ -568,6 +569,4 @@ def _load_custom( f"Unsupported file format for {remote_data_file!r}. " "Use dataset_name='h2o' for CSV files." ) - self._load_json_batched( - remote_data_file, url, table, max_rows, "Custom" - ) + self._load_json_batched(remote_data_file, url, table, max_rows, "Custom") From 3039055b214e0eeee49fb4cf8108f4d9b1796ed9 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Mon, 17 Aug 2026 17:48:54 -0400 Subject: [PATCH 4/4] fix(asap-tools): address PR #542 review findings in ClickHouse loader - Run the ClickHouse HTTP readiness wait before DROP/init-SQL execution instead of after, so it actually guards those calls. - Drop the target table for dataset_name="custom" with no init_sql_file, restoring clean-table-per-run behavior for that path. - Bound json/loader.py max_rows against a live total+len(batch) count instead of the post-flush total, so it stops reading at the requested row count instead of validating up to a full batch_size of lines first. - Dedupe the ClickHouse /ping readiness check into a shared _clickhouse_ping_ok() helper used by both ClickHouseService.is_healthy() and the new ClickHouseDataLoaderService.is_healthy(), and collapse _ensure_clickhouse_http_ready() onto the inherited wait_until_ready() while keeping its docker-logs-on-timeout diagnostic. - Use ClickHouseService.CONTAINER_NAME instead of a hardcoded container name in the readiness-timeout log fetch. Co-Authored-By: Claude Sonnet 5 --- .../services/clickhouse_service.py | 91 +++++++++---------- .../experiment_utils/services/json/loader.py | 2 +- 2 files changed, 45 insertions(+), 48 deletions(-) diff --git a/asap-tools/experiments/experiment_utils/services/clickhouse_service.py b/asap-tools/experiments/experiment_utils/services/clickhouse_service.py index b3e4f6a..246401e 100644 --- a/asap-tools/experiments/experiment_utils/services/clickhouse_service.py +++ b/asap-tools/experiments/experiment_utils/services/clickhouse_service.py @@ -5,7 +5,6 @@ import os import shlex import subprocess -import time from typing import Optional from jinja2 import Template @@ -15,6 +14,25 @@ import utils +def _clickhouse_ping_ok( + provider: InfrastructureProvider, node_offset: int, http_port: int +) -> bool: + """True iff ClickHouse's HTTP /ping on ``http_port`` returns exactly 'Ok.'.""" + result = provider.execute_command( + node_idx=node_offset, + cmd=f"curl -sS http://localhost:{http_port}/ping", + cmd_dir=None, + nohup=False, + popen=False, + ignore_errors=True, + ) + return ( + isinstance(result, subprocess.CompletedProcess) + and result.returncode == 0 + and result.stdout.strip() == "Ok." + ) + + class ClickHouseService(DockerServiceBase): """Manages a ClickHouse Docker container on a remote CloudLab node.""" @@ -46,17 +64,7 @@ def get_http_port(self) -> int: def is_healthy(self) -> bool: """ClickHouse is ready only when /ping returns exactly 'Ok.'""" - result = self.provider.execute_command( - node_idx=self.node_offset, - cmd=f"curl -s http://localhost:{self._http_port}/ping", - cmd_dir=None, - nohup=False, - popen=False, - ignore_errors=True, - ) - if not isinstance(result, subprocess.CompletedProcess): - return False - return result.returncode == 0 and result.stdout.strip() == "Ok." + return _clickhouse_ping_ok(self.provider, self.node_offset, self._http_port) def start( self, @@ -216,6 +224,12 @@ def __init__( self.clickhouse_http_port = clickhouse_http_port self.remote_data_file: Optional[str] = None + def is_healthy(self) -> bool: + """ClickHouse is ready only when /ping returns exactly 'Ok.'""" + return _clickhouse_ping_ok( + self.provider, self.node_offset, self.clickhouse_http_port + ) + def prepare(self, local_data_file: str, remote_dir: str) -> str: """Rsync a local data file to the remote node. @@ -289,6 +303,7 @@ def start( ) url = f"http://localhost:{self.clickhouse_http_port}/" + self._ensure_clickhouse_http_ready() if init_sql_file is not None: # init SQL owns DROP/CREATE for the raw table and any MVs / agg @@ -309,13 +324,14 @@ def start( self._exec_sql_file(remote_ddl, url) finally: self._remote_rm(remote_ddl) - elif dataset_name != "custom": + elif dataset_name == "custom": + print(f"Dropping table {table!r}...") + self._exec_sql(f"DROP TABLE IF EXISTS {table}", url) + else: raise ValueError( f"No built-in DDL for dataset_name={dataset_name!r}; pass init_sql_file" ) - self._ensure_clickhouse_http_ready(url) - if dataset_name == "clickbench": self._load_clickbench(remote_data_file, url, table, max_rows) elif dataset_name == "h2o": @@ -410,45 +426,26 @@ def _exec_sql_file(self, remote_sql_file: str, url: str) -> None: for stmt in (s.strip() for s in result.stdout.split(";") if s.strip()): self._exec_sql(stmt, url) - def _ensure_clickhouse_http_ready(self, url: str, max_retries: int = 30) -> None: + def _ensure_clickhouse_http_ready(self, max_retries: int = 30) -> None: """Wait until ClickHouse HTTP /ping returns Ok. before bulk load.""" - ping_url = url.rstrip("/") + "/ping" - for attempt in range(max_retries): - result = self.provider.execute_command( + try: + self.wait_until_ready(timeout=max_retries * 2) + except RuntimeError as exc: + logs = self.provider.execute_command( node_idx=self.node_offset, - cmd=f"curl -sS {shlex.quote(ping_url)}", + cmd=f"docker logs {shlex.quote(ClickHouseService.CONTAINER_NAME)} --tail 30 2>&1", cmd_dir=None, nohup=False, popen=False, ignore_errors=True, ) - if ( - isinstance(result, subprocess.CompletedProcess) - and result.returncode == 0 - and result.stdout.strip() == "Ok." - ): - return - print( - f"Waiting for ClickHouse HTTP before data load... " - f"({attempt + 1}/{max_retries})" - ) - time.sleep(2) - - logs = self.provider.execute_command( - node_idx=self.node_offset, - cmd="docker logs clickhouse-server --tail 30 2>&1", - cmd_dir=None, - nohup=False, - popen=False, - ignore_errors=True, - ) - log_tail = "" - if isinstance(logs, subprocess.CompletedProcess): - log_tail = (logs.stdout or logs.stderr or "").strip()[-500:] - raise RuntimeError( - "ClickHouse HTTP is not ready before data load. " - f"Check clickhouse_logs/ on the experiment node. Recent logs: {log_tail}" - ) + log_tail = "" + if isinstance(logs, subprocess.CompletedProcess): + log_tail = (logs.stdout or logs.stderr or "").strip()[-500:] + raise RuntimeError( + "ClickHouse HTTP is not ready before data load. " + f"Check clickhouse_logs/ on the experiment node. Recent logs: {log_tail}" + ) from exc def _load_json_batched( self, diff --git a/asap-tools/experiments/experiment_utils/services/json/loader.py b/asap-tools/experiments/experiment_utils/services/json/loader.py index fb9dac5..3fb526b 100644 --- a/asap-tools/experiments/experiment_utils/services/json/loader.py +++ b/asap-tools/experiments/experiment_utils/services/json/loader.py @@ -103,7 +103,7 @@ def flush(lines: list[str]) -> None: with _open_data_file(data_file) as fin: for raw_line in fin: file_line_no += 1 - if max_rows > 0 and total >= max_rows: + if max_rows > 0 and total + len(batch) >= max_rows: break stripped = _validate_line(raw_line, file_line_no) if not stripped: