from __future__ import annotations import argparse import json import math import os import sys import urllib.error import urllib.parse import urllib.request from dataclasses import dataclass from datetime import UTC from datetime import datetime from datetime import timedelta from pathlib import Path from typing import Any from zoneinfo import ZoneInfo DEFAULT_TIMEOUT_S = 60 DEFAULT_O11Y_AUTH_URL = "https://www.ds.us-west-2.aws.observability.tidbcloud.com/client/auth" BEIJING_TZ = ZoneInfo("Asia/Shanghai") @dataclass(frozen=True) class TiKVNodeInventory: node: str tikv_node: str class MetricsQueryError(RuntimeError): pass class O11yMetricsClient: def __init__( self, *, provider: str, region: str, infra_uuid: str, github_pat: str, timeout_s: int = DEFAULT_TIMEOUT_S, ) -> None: self._provider = provider self._region = region self._infra_uuid = infra_uuid self._timeout_s = timeout_s self._auth0_token = self._authenticate(github_pat=github_pat) self._metrics_base = ( "https://www.ds." f"{self._region}.{self._provider}.observability.tidbcloud.com/internal/v2/metrics/" f"{self._infra_uuid}/api/v1" ) def query_range( self, *, promql: str, start: int, end: int, step_seconds: int, ) -> dict[str, Any]: return self._get( "/query_range", params={ "query": promql, "start": str(start), "end": str(end), "step": str(step_seconds), }, ) def query_instant( self, *, promql: str, query_time: int, ) -> dict[str, Any]: return self._get( "/query", params={ "query": promql, "time": str(query_time), }, ) def _authenticate(self, *, github_pat: str) -> str: body = json.dumps( { "github_token": github_pat, "provider": self._provider, } ).encode("utf-8") request = urllib.request.Request( DEFAULT_O11Y_AUTH_URL, headers={"Content-Type": "application/json"}, data=body, method="POST", ) try: with urllib.request.urlopen(request, timeout=self._timeout_s) as response: payload = response.read().decode("utf-8") except urllib.error.HTTPError as exc: body = exc.read().decode("utf-8", errors="replace") raise MetricsQueryError(f"HTTP {exc.code} during O11Y auth: {body}") from exc except urllib.error.URLError as exc: raise MetricsQueryError(f"O11Y auth request failed: {exc}") from exc try: data = json.loads(payload) except json.JSONDecodeError as exc: raise MetricsQueryError(f"invalid O11Y auth JSON response: {exc}") from exc auth0_token = str(data.get("auth0_token") or "").strip() if not auth0_token: raise MetricsQueryError("missing auth0_token in O11Y auth response") return auth0_token def _get(self, path: str, *, params: dict[str, str]) -> dict[str, Any]: url = f"{self._metrics_base}{path}?{urllib.parse.urlencode(params)}" request = urllib.request.Request( url, headers={"Authorization": f"Bearer {self._auth0_token}"}, method="GET", ) try: with urllib.request.urlopen(request, timeout=self._timeout_s) as response: payload = response.read().decode("utf-8") except urllib.error.HTTPError as exc: body = exc.read().decode("utf-8", errors="replace") raise MetricsQueryError(f"HTTP {exc.code} for {path}: {body}") from exc except urllib.error.URLError as exc: raise MetricsQueryError(f"request failed for {path}: {exc}") from exc try: data = json.loads(payload) except json.JSONDecodeError as exc: raise MetricsQueryError(f"invalid JSON response for {path}: {exc}") from exc status = str(data.get("status") or "") if status and status != "success": raise MetricsQueryError(f"unexpected response status for {path}: {status}") return data def _env_github_token() -> str: for key in ( "SKILL_GITHUB_TOKEN", "INVESTIBOOK_GITHUB_TOKEN", "GITHUB_TOKEN", "GH_TOKEN", "GITHUB_PAT", ): value = os.environ.get(key, "").strip() if value: return value raise MetricsQueryError( "missing GitHub token in " "SKILL_GITHUB_TOKEN/INVESTIBOOK_GITHUB_TOKEN/GITHUB_TOKEN/GH_TOKEN/GITHUB_PAT" ) def _inventory_query(namespace: str) -> str: return f'kube_pod_info{{namespace="{namespace}",created_by_kind="TiKV"}}' def _disk_read_query(instance_regex: str, *, rate_window_seconds: int) -> str: selector = f'node_disk_reads_completed_total{{instance=~"{instance_regex}",device!~"dm.*"}}' return f"sum by (instance) (rate({selector}[{rate_window_seconds}s]))" def _safe_float(value: Any) -> float | None: try: numeric = float(value) except (TypeError, ValueError): return None if not math.isfinite(numeric): return None return numeric def _isoformat_utc(ts: int) -> str: return datetime.fromtimestamp(ts, tz=UTC).isoformat().replace("+00:00", "Z") def _parse_iso8601(value: str) -> int: normalized = value.strip() if normalized.endswith("Z"): normalized = normalized[:-1] + "+00:00" return int(datetime.fromisoformat(normalized).timestamp()) def _result_series(payload: dict[str, Any]) -> list[dict[str, Any]]: data = payload.get("data") if not isinstance(data, dict): return [] result = data.get("result") return [item for item in result if isinstance(item, dict)] if isinstance(result, list) else [] def _inventory_pairs(payload: dict[str, Any]) -> list[TiKVNodeInventory]: resolved: dict[str, str] = {} for item in _result_series(payload): metric = item.get("metric") if not isinstance(metric, dict): continue node = str(metric.get("node") or "").strip() tikv_node = str(metric.get("pod") or "").strip() if not node or not tikv_node: continue previous = resolved.get(node) if previous is None: resolved[node] = tikv_node continue if previous != tikv_node: raise MetricsQueryError( f"node '{node}' maps to multiple TiKV pods: '{previous}' and '{tikv_node}'" ) return [ TiKVNodeInventory(node=node, tikv_node=tikv_node) for node, tikv_node in sorted(resolved.items()) ] def _load_json_object_from_path(path: Path) -> dict[str, Any]: try: raw = path.read_text(encoding="utf-8") except OSError as exc: raise MetricsQueryError(f"failed to read JSON file '{path}': {exc}") from exc try: payload = json.loads(raw) except json.JSONDecodeError as exc: raise MetricsQueryError(f"invalid JSON file '{path}': {exc}") from exc if not isinstance(payload, dict): raise MetricsQueryError(f"JSON file '{path}' must contain an object") return payload def _resolve_case_relative_path(raw_path: str) -> Path: path = Path(raw_path).expanduser() if path.is_absolute(): return path return Path(__file__).resolve().parent / path def _prefetched_cluster_rows(cluster_payload: dict[str, Any], *, cluster_id: str) -> list[dict[str, Any]]: direct_rows = cluster_payload.get("nodes") if isinstance(direct_rows, list): rows: list[dict[str, Any]] = [] for item in direct_rows: if not isinstance(item, list) or len(item) != 2: raise MetricsQueryError("prefetched cluster nodes entries must be [tikv_node, max_1m_read_iops]") tikv_node = str(item[0]).strip() max_value = _safe_float(item[1]) if not tikv_node or max_value is None: raise MetricsQueryError("prefetched cluster nodes entries require tikv_node and max_1m_read_iops") rows.append( { "cluster_id": cluster_id, "tikv_node": tikv_node, "max_1m_read_iops": round(max_value, 3), } ) rows.sort(key=lambda item: (-float(item["max_1m_read_iops"]), str(item["tikv_node"]))) return rows inventory_raw = cluster_payload.get("inventory") metrics_raw = cluster_payload.get("per_instance_read_iops") if not isinstance(inventory_raw, list) or not isinstance(metrics_raw, list): raise MetricsQueryError( "prefetched cluster input requires either nodes or both inventory and per_instance_read_iops" ) instance_to_tikv_node: dict[str, str] = {} for item in inventory_raw: if not isinstance(item, dict): raise MetricsQueryError("prefetched inventory entries must be objects") instance = str(item.get("instance") or item.get("node") or "").strip() tikv_node = str(item.get("tikv_node") or item.get("pod") or "").strip() if not instance or not tikv_node: raise MetricsQueryError("prefetched inventory entries require instance and tikv_node") previous = instance_to_tikv_node.get(instance) if previous and previous != tikv_node: raise MetricsQueryError( f"prefetched inventory maps instance '{instance}' to multiple TiKV nodes for cluster {cluster_id}" ) instance_to_tikv_node[instance] = tikv_node rows: list[dict[str, Any]] = [] for item in metrics_raw: if not isinstance(item, dict): raise MetricsQueryError("prefetched per_instance_read_iops entries must be objects") instance = str(item.get("instance") or "").strip() max_value = _safe_float(item.get("max_1m_read_iops")) if not instance or max_value is None: raise MetricsQueryError( "prefetched per_instance_read_iops entries require instance and max_1m_read_iops" ) tikv_node = instance_to_tikv_node.get(instance) if not tikv_node: raise MetricsQueryError( f"prefetched metric instance '{instance}' was not found in prefetched inventory for cluster {cluster_id}" ) rows.append( { "cluster_id": cluster_id, "tikv_node": tikv_node, "max_1m_read_iops": round(max_value, 3), } ) rows.sort(key=lambda item: (-float(item["max_1m_read_iops"]), str(item["tikv_node"]))) return rows def _workspace_dir() -> Path: raw = os.environ.get("INVESTIBOOK_EVAL_WORKSPACE_DIR", "").strip() if not raw: raise MetricsQueryError("missing INVESTIBOOK_EVAL_WORKSPACE_DIR") return Path(raw) def _runtime_context_path() -> Path: raw = os.environ.get("INVESTIBOOK_EVAL_RUNTIME_CONTEXT_PATH", "").strip() if not raw: raise MetricsQueryError("missing INVESTIBOOK_EVAL_RUNTIME_CONTEXT_PATH") return Path(raw) def _write_workspace_prefetched_artifact(*, payload: dict[str, Any], filename: str) -> str: workspace_dir = _workspace_dir() evidence_dir = workspace_dir / "evidences" evidence_dir.mkdir(parents=True, exist_ok=True) artifact_path = evidence_dir / filename artifact_path.write_text( json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8", ) return str(artifact_path.relative_to(workspace_dir)) def _workspace_path_from_relative(relative_path: str) -> Path: path = Path(relative_path) if path.is_absolute(): return path return _workspace_dir() / path def _load_runtime_output_data(output_id: str) -> dict[str, Any]: runtime_context = _load_json_object_from_path(_runtime_context_path()) outputs = runtime_context.get("outputs") if not isinstance(outputs, dict): raise MetricsQueryError("runtime context is missing outputs") item = outputs.get(output_id) if not isinstance(item, dict): raise MetricsQueryError(f"runtime context is missing output '{output_id}'") data = item.get("data") if not isinstance(data, dict): raise MetricsQueryError(f"runtime output '{output_id}' does not contain an object payload") return data def _cluster_payloads_from_prefetched_input(payload: dict[str, Any]) -> dict[str, dict[str, Any]]: if "nodes_by_cluster" in payload: raw_groups = payload.get("nodes_by_cluster") if not isinstance(raw_groups, list): raise MetricsQueryError("prefetched nodes_by_cluster must be a list") else: raw_groups = payload.get("clusters") if not isinstance(raw_groups, list): raise MetricsQueryError("prefetched input requires nodes_by_cluster or clusters") cluster_payloads: dict[str, dict[str, Any]] = {} for item in raw_groups: if not isinstance(item, dict): continue cluster_id = str(item.get("cluster_id") or "").strip() if not cluster_id: continue cluster_payloads[cluster_id] = item return cluster_payloads def _validate_prefetched_window( payload: dict[str, Any], *, expected_start: str, expected_end: str, ) -> None: window = payload.get("window") if not isinstance(window, dict): return start = str(window.get("start") or "").strip() end = str(window.get("end") or "").strip() if start and start != expected_start: raise MetricsQueryError( f"prefetched input window.start '{start}' does not match expected '{expected_start}'" ) if end and end != expected_end: raise MetricsQueryError( f"prefetched input window.end '{end}' does not match expected '{expected_end}'" ) def _prepare_prefetched_case_input( *, prefetched_input_path: str, cluster_ids: list[str], artifact_filename: str, expected_start: str, expected_end: str, ) -> tuple[str, str]: source_path = _resolve_case_relative_path(prefetched_input_path) payload = _load_json_object_from_path(source_path) _validate_prefetched_window( payload, expected_start=expected_start, expected_end=expected_end, ) cluster_payloads = _cluster_payloads_from_prefetched_input(payload) for cluster_id in cluster_ids: cluster_payload = cluster_payloads.get(cluster_id) if cluster_payload is None: raise MetricsQueryError(f"prefetched input is missing cluster_id {cluster_id}") _prefetched_cluster_rows(cluster_payload, cluster_id=cluster_id) relative_path = _write_workspace_prefetched_artifact( payload=payload, filename=artifact_filename, ) prompt = ( f"workspace 中已提供预取好的基础设施证据文件 `{relative_path}`。" "请优先使用该文件完成排序，不要重复查询同类 `node_*` / `kube_*` / `aws_ebs_*` 指标。" ) return relative_path, prompt def _instance_regex(instances: list[str]) -> str: unique_instances = [instance for instance in sorted(set(instances)) if instance] if not unique_instances: raise MetricsQueryError("resolved empty TiKV instance inventory") return "^(?:" + "|".join(instance.replace(".", "[.]") for instance in unique_instances) + ")$" def _series_max_value(item: dict[str, Any]) -> float | None: values = item.get("values") if not isinstance(values, list): return None observed: list[float] = [] for row in values: if not isinstance(row, list) or len(row) < 2: continue parsed = _safe_float(row[1]) if parsed is not None: observed.append(parsed) return max(observed) if observed else None def _build_cluster_rows( *, cluster_id: str, parent_id: str, provider: str, region: str, infra_uuid: str, github_pat: str, start_ts: int, end_ts: int, step_seconds: int, ) -> tuple[list[dict[str, Any]], dict[str, Any]]: namespace = f"tidb{parent_id}" client = O11yMetricsClient( provider=provider, region=region, infra_uuid=infra_uuid, github_pat=github_pat, ) inventory_payload = client.query_instant( promql=_inventory_query(namespace), query_time=end_ts, ) inventory = _inventory_pairs(inventory_payload) if not inventory: return [], { "inventory_query": _inventory_query(namespace), "metrics_scope": { "provider": provider, "region": region, "infra_uuid": infra_uuid, "namespace": namespace, }, "inventory_count": 0, "ranked_count": 0, } node_to_tikv_node = {item.node: item.tikv_node for item in inventory} sample_start_ts = start_ts + 60 if sample_start_ts > end_ts: raise MetricsQueryError("window is too short to compute max 1-minute read IOPS") range_query = _disk_read_query(_instance_regex(list(node_to_tikv_node)), rate_window_seconds=60) range_payload = client.query_range( promql=range_query, start=sample_start_ts, end=end_ts, step_seconds=step_seconds, ) rows: list[dict[str, Any]] = [] observed_instances: set[str] = set() for item in _result_series(range_payload): metric = item.get("metric") if not isinstance(metric, dict): continue instance = str(metric.get("instance") or "").strip() if not instance: continue tikv_node = node_to_tikv_node.get(instance) if not tikv_node: raise MetricsQueryError( f"metric instance '{instance}' was not found in resolved TiKV node inventory for cluster {cluster_id}" ) max_value = _series_max_value(item) if max_value is None: continue observed_instances.add(instance) rows.append( { "cluster_id": cluster_id, "tikv_node": tikv_node, "max_1m_read_iops": round(max_value, 3), } ) missing_instances = sorted(set(node_to_tikv_node) - observed_instances) rows.sort(key=lambda item: (-float(item["max_1m_read_iops"]), str(item["tikv_node"]))) return rows, { "inventory_query": _inventory_query(namespace), "metrics_scope": { "provider": provider, "region": region, "infra_uuid": infra_uuid, "namespace": namespace, }, "inventory_count": len(inventory), "ranked_count": len(rows), "range_query": range_query, "sample_start": _isoformat_utc(sample_start_ts), "sample_end": _isoformat_utc(end_ts), "missing_instances": missing_instances, } def _latest_completed_beijing_window( *, local_hour: int, local_minute: int, duration_minutes: int, ) -> tuple[int, int, str]: now_local = datetime.now(tz=BEIJING_TZ) candidate_start = now_local.replace( hour=local_hour, minute=local_minute, second=0, microsecond=0, ) candidate_end = candidate_start + timedelta(minutes=duration_minutes) if now_local < candidate_end: candidate_start = candidate_start - timedelta(days=1) candidate_end = candidate_start + timedelta(minutes=duration_minutes) return ( int(candidate_start.astimezone(UTC).timestamp()), int(candidate_end.astimezone(UTC).timestamp()), candidate_start.date().isoformat(), ) def _prepare_case_input(args: argparse.Namespace) -> dict[str, Any]: window_start, window_end, local_date = _latest_completed_beijing_window( local_hour=args.local_hour, local_minute=args.local_minute, duration_minutes=args.duration_minutes, ) cluster_ids = sorted([args.cluster_id_a, args.cluster_id_b]) expected_start = _isoformat_utc(window_start) expected_end = _isoformat_utc(window_end) prefetched_relative_path = "" prefetched_prompt = "" evidence_mode = "live_query" prefetched_input_path = str(args.prefetched_input_path or "").strip() if prefetched_input_path: prefetched_relative_path, prefetched_prompt = _prepare_prefetched_case_input( prefetched_input_path=prefetched_input_path, cluster_ids=cluster_ids, artifact_filename="prefetched_read_iops_input.json", expected_start=expected_start, expected_end=expected_end, ) evidence_mode = "prefetched_input" return { "id": f"bjt-{local_date}-{args.local_hour:02d}{args.local_minute:02d}-{args.duration_minutes}m", "cluster_ids": cluster_ids, "start": expected_start, "end": expected_end, "local_timezone": "Asia/Shanghai", "local_date": local_date, "local_hour": args.local_hour, "local_minute": args.local_minute, "duration_minutes": args.duration_minutes, "step_seconds": args.step_seconds, "mapping_strategy": "node_filter", "evidence_mode": evidence_mode, "prefetched_evidence_path": prefetched_relative_path, "prefetched_infra_prompt": prefetched_prompt, } def _build_reference(args: argparse.Namespace) -> dict[str, Any]: start_ts = _parse_iso8601(args.start) end_ts = _parse_iso8601(args.end) if end_ts <= start_ts: raise MetricsQueryError("reference window end must be greater than start") cluster_specs = sorted( [ { "cluster_id": args.cluster_id_a, "parent_id": args.parent_id_a, "provider": args.provider_a, "region": args.region_a, "infra_uuid": args.infra_uuid_a, }, { "cluster_id": args.cluster_id_b, "parent_id": args.parent_id_b, "provider": args.provider_b, "region": args.region_b, "infra_uuid": args.infra_uuid_b, }, ], key=lambda item: str(item["cluster_id"]), ) cluster_ids = [str(item["cluster_id"]) for item in cluster_specs] cluster_queries: dict[str, Any] = {} cluster_node_counts: dict[str, int] = {} grouped_rows: list[dict[str, Any]] = [] case_input = _load_runtime_output_data(args.case_input_output_id) evidence_mode = str(case_input.get("evidence_mode") or "live_query").strip() or "live_query" if evidence_mode == "prefetched_input": relative_path = str(case_input.get("prefetched_evidence_path") or "").strip() if not relative_path: raise MetricsQueryError("case_input is missing prefetched_evidence_path") prefetched_input = _load_json_object_from_path(_workspace_path_from_relative(relative_path)) cluster_payloads = _cluster_payloads_from_prefetched_input(prefetched_input) for spec in cluster_specs: cluster_id = str(spec["cluster_id"]) cluster_payload = cluster_payloads.get(cluster_id) if cluster_payload is None: raise MetricsQueryError(f"prefetched input is missing cluster_id {cluster_id}") cluster_rows = _prefetched_cluster_rows(cluster_payload, cluster_id=cluster_id) cluster_queries[cluster_id] = { "inventory_query": "prefetched_input.inventory", "metrics_scope": { "source": "prefetched_workspace_artifact", "path": relative_path, }, "ranked_count": len(cluster_rows), } cluster_node_counts[cluster_id] = len(cluster_rows) grouped_rows.append( { "cluster_id": cluster_id, "nodes": [ [str(item["tikv_node"]), float(item["max_1m_read_iops"])] for item in cluster_rows ], } ) else: github_pat = _env_github_token() for spec in cluster_specs: cluster_rows, cluster_debug = _build_cluster_rows( cluster_id=str(spec["cluster_id"]), parent_id=str(spec["parent_id"]), provider=str(spec["provider"]), region=str(spec["region"]), infra_uuid=str(spec["infra_uuid"]), github_pat=github_pat, start_ts=start_ts, end_ts=end_ts, step_seconds=args.step_seconds, ) cluster_id = str(spec["cluster_id"]) cluster_queries[cluster_id] = cluster_debug cluster_node_counts[cluster_id] = len(cluster_rows) grouped_rows.append( { "cluster_id": cluster_id, "nodes": [ [str(item["tikv_node"]), float(item["max_1m_read_iops"])] for item in cluster_rows ], } ) grouped_rows.sort(key=lambda item: str(item["cluster_id"])) evidence_sufficient = all(item["nodes"] for item in grouped_rows) return { "window": { "start": _isoformat_utc(start_ts), "end": _isoformat_utc(end_ts), }, "mapping_strategy": "node_filter", "evidence_sufficient": evidence_sufficient, "cluster_ids": cluster_ids, "cluster_node_counts": cluster_node_counts, "nodes_by_cluster": grouped_rows if evidence_sufficient else [], "queries": cluster_queries, } def _build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="Runtime hook for premium full TiKV read IOPS ranking") subparsers = parser.add_subparsers(dest="command", required=True) prepare_parser = subparsers.add_parser("prepare-case-input") prepare_parser.add_argument("--cluster-id-a", required=True) prepare_parser.add_argument("--cluster-id-b", required=True) prepare_parser.add_argument("--local-hour", type=int, default=11) prepare_parser.add_argument("--local-minute", type=int, default=0) prepare_parser.add_argument("--duration-minutes", type=int, default=5) prepare_parser.add_argument("--step-seconds", type=int, default=60) prepare_parser.add_argument("--prefetched-input-path", default="") reference_parser = subparsers.add_parser("build-reference") reference_parser.add_argument("--cluster-id-a", required=True) reference_parser.add_argument("--cluster-id-b", required=True) reference_parser.add_argument("--parent-id-a", required=True) reference_parser.add_argument("--parent-id-b", required=True) reference_parser.add_argument("--provider-a", required=True) reference_parser.add_argument("--provider-b", required=True) reference_parser.add_argument("--region-a", required=True) reference_parser.add_argument("--region-b", required=True) reference_parser.add_argument("--infra-uuid-a", required=True) reference_parser.add_argument("--infra-uuid-b", required=True) reference_parser.add_argument("--start", required=True) reference_parser.add_argument("--end", required=True) reference_parser.add_argument("--step-seconds", type=int, default=60) reference_parser.add_argument("--case-input-output-id", default="case_input") return parser def main(argv: list[str]) -> int: parser = _build_parser() args = parser.parse_args(argv) if args.command == "prepare-case-input": payload = _prepare_case_input(args) elif args.command == "build-reference": payload = _build_reference(args) else: raise MetricsQueryError(f"unsupported command: {args.command}") sys.stdout.write(json.dumps(payload, ensure_ascii=False)) sys.stdout.write("\n") return 0 if __name__ == "__main__": raise SystemExit(main(sys.argv[1:]))