#!/usr/bin/env python3 """Task-driven S3 size scan for 0/, diagnosis/data/, deltalake/; list1-style walk for 1/. Cluster source: diagnosis/tasks/processor-dedicated/{clusterID}.yaml (+ hardcoded seed/o11y for diagnosis/data and deltalake, matching diagnosis-compute). Scan kinds: 0 - task-driven cluster roots: logs/, profiles/, regions/, ... regions - task-driven regions/ only (filenames YYYYMMDDHHmm.json.log.gz) 1 - full recursive walk under 1/ (same as list1.py; no tenant task mapping) diagnosis/data - diagnosis/data/{storageID}/... filtered by time where possible deltalake - parquet + log-stats paths for task clusters / dates in range """ from collections import defaultdict from datetime import datetime, timedelta, timezone import boto3 import yaml # === User Configuration === BUCKET = 'o11y-prod-shared-ap-southeast-1' TASK_PREFIX = 'diagnosis/tasks/processor-dedicated/' START_TIMESTAMP = 1785196800 END_TIMESTAMP = 1785283200 # Which roots to scan. Use ['regions'] to only size keyviz regions/. SCAN_KINDS = ['0'] # Env used for seed/o11y synthetic clusters (and fallback when task meta.env is empty) DEFAULT_ENV = 'staging' if 'staging' in BUCKET else ('dev' if 'dev' in BUCKET else 'prod') s3 = boto3.client('s3') stats = defaultdict(lambda: {'count': 0, 'total_size': 0}) def stats_filename(kind): safe = kind.replace('/', '_') return f"scan_stats_{BUCKET}_{safe}_tasks_{START_TIMESTAMP}_{END_TIMESTAMP}.txt" def write_stats_to_file(kind, prefix, count, size_mb): with open(stats_filename(kind), 'a') as f: f.write(f"{prefix},{count},{size_mb}\n") def list_prefixes(bucket, prefix): paginator = s3.get_paginator('list_objects_v2') pages = paginator.paginate(Bucket=bucket, Prefix=prefix, Delimiter='/') for page in pages: for cp in page.get('CommonPrefixes', []): yield cp['Prefix'] def list_task_keys(bucket, prefix): paginator = s3.get_paginator('list_objects_v2') pages = paginator.paginate(Bucket=bucket, Prefix=prefix) for page in pages: for obj in page.get('Contents', []): key = obj['Key'] if key.endswith('.yaml') and key.count('/') == prefix.count('/'): yield key def load_task_meta(bucket, key): obj = s3.get_object(Bucket=bucket, Key=key) data = yaml.safe_load(obj['Body'].read()) or {} meta = data.get('meta') or {} cluster_id = meta.get('clusterID') or '' cluster_type = meta.get('clusterType') or '' if not cluster_id or cluster_type not in ('dedicated', 'nextgen'): return None return { 'tenantID': meta.get('tenantID') or '', 'projectID': meta.get('projectID') or '', 'clusterID': cluster_id, 'storageID': cluster_id, 'clusterType': cluster_type, 'env': meta.get('env') or DEFAULT_ENV, 'task_key': key, 'source': 'task', } def infra_metas(): """Hardcoded seed/o11y clusters (Type=1), same as diagnosis-compute main.go.""" return [ { 'tenantID': 'infra-prod', 'projectID': 'seed', 'clusterID': 'k8s-infra', 'storageID': 'seed', 'clusterType': 'dedicated', 'env': DEFAULT_ENV, 'task_key': 'infra:seed', 'source': 'infra', 'raw_type': 1, }, { 'tenantID': 'infra-prod', 'projectID': 'o11y', 'clusterID': 'k8s-infra', 'storageID': 'o11y', 'clusterType': 'dedicated', 'env': DEFAULT_ENV, 'task_key': 'infra:o11y', 'source': 'infra', 'raw_type': 1, }, ] def deploy_type(meta): # Match config.DeployType: nextgen -> premium, else dedicated if meta.get('clusterType') == 'nextgen': return 'premium' return 'dedicated' def resolve_cluster_bases(bucket, meta): """Cluster root prefixes under 0/ (sibling dirs: logs/, profiles/, regions/, ...).""" cluster_id = meta['clusterID'] if meta.get('raw_type') == 1 or meta.get('source') == 'infra': return [f"1/{meta['tenantID']}/{meta['projectID']}/{cluster_id}/"] if meta['clusterType'] == 'dedicated': return [f"0/{meta['tenantID']}/{meta['projectID']}/{cluster_id}/"] # nextgen: detect physical vs logical by whether the cluster root exists physical = f"0/shared_pool/{cluster_id}/" logical = f"0/next-gen/{cluster_id}/" roots = [] if any(True for _ in list_prefixes(bucket, physical)): roots.append(physical) if any(True for _ in list_prefixes(bucket, logical)): roots.append(logical) return roots def iter_dates_utc(start_ts, end_ts): """Yield YYYY-MM-DD for each UTC calendar day overlapping [start, end).""" start = datetime.fromtimestamp(start_ts, tz=timezone.utc).date() end = datetime.fromtimestamp(end_ts - 1, tz=timezone.utc).date() d = start while d <= end: yield d.isoformat() d += timedelta(days=1) def iter_hours_utc(start_ts, end_ts): """Yield YYYYMMDDHH for each UTC hour in [start, end).""" t = start_ts - (start_ts % 3600) while t < end_ts: yield datetime.fromtimestamp(t, tz=timezone.utc).strftime('%Y%m%d%H') t += 3600 def finish_scan(kind, prefix, count): size_mb = stats[prefix]['total_size'] / (1024 * 1024) print(f"scan done [{kind}]: {prefix} count: {count} size: {size_mb:.2f} MB") write_stats_to_file(kind, prefix, count, f"{size_mb:.2f}") def ts_to_regions_name(ts): """regions object names use UTC YYYYMMDDHHmm (e.g. 202606270000.json.log.gz).""" return datetime.fromtimestamp(ts, tz=timezone.utc).strftime('%Y%m%d%H%M') def scan_raw_component(kind, bucket, prefix): """Count objects in [START_TIMESTAMP, END_TIMESTAMP] under a prefix (unix-ts filenames).""" paginator = s3.get_paginator('list_objects_v2') start_prefix = f"{prefix}{START_TIMESTAMP}" end_prefix = f"{prefix}{END_TIMESTAMP + 1}" pages = paginator.paginate( Bucket=bucket, Prefix=prefix, StartAfter=start_prefix, PaginationConfig={'PageSize': 1000}, ) count = 0 for page in pages: for obj in page.get('Contents', []): key = obj['Key'] if key >= end_prefix: finish_scan(kind, prefix, count) return count += 1 if count % 10000 == 0: print(f"scan [{kind}]: {count} key: {key}") stats[prefix]['count'] += 1 stats[prefix]['total_size'] += obj['Size'] finish_scan(kind, prefix, count) def scan_regions_prefix(kind, bucket, prefix): """Count regions/ objects named YYYYMMDDHHmm*.json.log.gz in [START, END).""" paginator = s3.get_paginator('list_objects_v2') start_name = ts_to_regions_name(START_TIMESTAMP) end_name = ts_to_regions_name(END_TIMESTAMP) # StartAfter is exclusive; cursor just before the first minute in range start_after = f"{prefix}{start_name}" end_prefix = f"{prefix}{end_name}" print(f"regions window: [{start_name}, {end_name}) prefix={prefix}") pages = paginator.paginate( Bucket=bucket, Prefix=prefix, StartAfter=start_after, PaginationConfig={'PageSize': 1000}, ) count = 0 for page in pages: for obj in page.get('Contents', []): key = obj['Key'] if key >= end_prefix: finish_scan(kind, prefix, count) return count += 1 if count % 1000 == 0: print(f"scan [{kind}]: {count} key: {key}") stats[prefix]['count'] += 1 stats[prefix]['total_size'] += obj['Size'] finish_scan(kind, prefix, count) def scan_cluster_base(kind, bucket, cluster_base, only_regions=False): """Scan siblings under a cluster root, matching list0.py layout. 0/{tenant}/{project}/{cluster}/logs/{component}/... (unix-ts names) 0/{tenant}/{project}/{cluster}/profiles/... (unix-ts names) 0/{tenant}/{project}/{cluster}/regions/... (YYYYMMDDHHmm names) """ children = list(list_prefixes(bucket, cluster_base)) if not children: print(f"no children under: {cluster_base}") return for child in children: if only_regions: if child.endswith('/regions/'): scan_regions_prefix(kind, bucket, child) continue if child.endswith('/logs/'): components = list(list_prefixes(bucket, child)) if not components: print(f"no components under: {child}") continue for component_prefix in components: scan_raw_component(kind, bucket, component_prefix) elif child.endswith('/regions/'): scan_regions_prefix(kind, bucket, child) else: # profiles/ and any other top-level data dirs (unix-ts filenames) scan_raw_component(kind, bucket, child) def scan_kind_0(bucket, metas, only_regions=False): label = 'regions only' if only_regions else 'logs/profiles/regions/...' kind = 'regions' if only_regions else '0' print(f"\n=== scan kind: 0/ task-driven ({label}) ===") print( f"time window unix=[{START_TIMESTAMP}, {END_TIMESTAMP}) " f"regions_name=[{ts_to_regions_name(START_TIMESTAMP)}, {ts_to_regions_name(END_TIMESTAMP)})" ) scanned = set() for i, meta in enumerate(metas, 1): if meta.get('source') == 'infra': continue bases = resolve_cluster_bases(bucket, meta) if not bases: print(f"[{i}] no cluster base cluster={meta['clusterID']} type={meta['clusterType']}") continue for cluster_base in bases: if cluster_base in scanned: continue scanned.add(cluster_base) print(f"[{i}] cluster={meta['clusterID']} base={cluster_base}") scan_cluster_base(kind, bucket, cluster_base, only_regions=only_regions) print(f"kind {kind} scanned cluster bases: {len(scanned)}") def scan_kind_1(bucket): """Same recursive walk as list1.py — 1/ has no tenant layout from tasks.""" print("\n=== scan kind: 1/ (list1 full walk) ===") prefix_root = '1/' for p1 in list_prefixes(bucket, prefix_root): for p2 in list_prefixes(bucket, p1): for p3 in list_prefixes(bucket, p2): for p4 in list_prefixes(bucket, p3): if p4.endswith('/logs/'): for p5 in list_prefixes(bucket, p4): scan_raw_component('1', bucket, p5) elif p4.endswith('/regions/'): scan_regions_prefix('1', bucket, p4) else: scan_raw_component('1', bucket, p4) def scan_prefix_all(kind, bucket, prefix): """Count all objects under prefix (no StartAfter).""" paginator = s3.get_paginator('list_objects_v2') pages = paginator.paginate(Bucket=bucket, Prefix=prefix, PaginationConfig={'PageSize': 1000}) count = 0 for page in pages: for obj in page.get('Contents', []): count += 1 if count % 10000 == 0: print(f"scan [{kind}]: {count} key: {obj['Key']}") stats[prefix]['count'] += 1 stats[prefix]['total_size'] += obj['Size'] finish_scan(kind, prefix, count) def scan_kind_diagnosis_data(bucket, metas): """diagnosis/data/{storageID}/ — prefer log-stats/1h hour files in range; also count other subtrees.""" print("\n=== scan kind: diagnosis/data (task-driven) ===") hours = set(iter_hours_utc(START_TIMESTAMP, END_TIMESTAMP)) hour_start = min(hours) if hours else '' hour_end = max(hours) if hours else '' seen = set() for i, meta in enumerate(metas, 1): storage_id = meta['storageID'] if storage_id in seen: continue seen.add(storage_id) root = f"diagnosis/data/{storage_id}/" print(f"[{i}] storageID={storage_id} root={root}") # log-stats/1h/{YYYYMMDDHH}.db — filter by hour name in range stats_prefix = f"{root}log-stats/1h/" paginator = s3.get_paginator('list_objects_v2') kwargs = {} if hour_start: # StartAfter is exclusive; cursor just before first hour in range kwargs['StartAfter'] = f"{stats_prefix}{int(hour_start) - 1}" count = 0 for page in paginator.paginate( Bucket=bucket, Prefix=stats_prefix, PaginationConfig={'PageSize': 1000}, **kwargs, ): stop = False for obj in page.get('Contents', []): key = obj['Key'] name = key.rsplit('/', 1)[-1] hour = name[:-3] if name.endswith('.db') else name if hour_end and hour > hour_end: stop = True break if hour not in hours: continue count += 1 stats[stats_prefix]['count'] += 1 stats[stats_prefix]['total_size'] += obj['Size'] if stop: break finish_scan('diagnosis/data', stats_prefix, count) # Other top-level dirs under this storage (e.g. merged-logs) — full count for child in list_prefixes(bucket, root): if child == stats_prefix or child.endswith('/log-stats/'): continue scan_prefix_all('diagnosis/data', bucket, child) def parse_parquet_cluster_id(filename): """Parse HH_clusterID_NN.parquet -> clusterID.""" if not filename.endswith('.parquet'): return None base = filename[:-8] parts = base.split('_') if len(parts) < 3: return None return '_'.join(parts[1:-1]) def scan_kind_deltalake(bucket, metas): """Scan deltalake parquet + log-stats for task clusters and dates in range.""" print("\n=== scan kind: deltalake (task-driven) ===") dates = list(iter_dates_utc(START_TIMESTAMP, END_TIMESTAMP)) dates_compact = [d.replace('-', '') for d in dates] # Group clusterIDs by (env, deployType) groups = defaultdict(set) # (env, deploy) -> {clusterID} org_cluster = [] # for log-stats path for meta in metas: env = meta['env'] or DEFAULT_ENV dt = deploy_type(meta) groups[(env, dt)].add(meta['clusterID']) org_cluster.append((meta['tenantID'], meta['clusterID'], meta['storageID'])) # Parquet detail: list once per (env, deployType, date), attribute by clusterID in filename for (env, dt), cluster_ids in sorted(groups.items()): for date in dates: date_prefix = f"deltalake/global/logs/env={env}/deployType={dt}/date={date}/" print(f"list parquet: {date_prefix} clusters={len(cluster_ids)}") per_cluster = defaultdict(lambda: {'count': 0, 'total_size': 0}) paginator = s3.get_paginator('list_objects_v2') total = 0 matched = 0 for page in paginator.paginate( Bucket=bucket, Prefix=date_prefix, PaginationConfig={'PageSize': 1000}, ): for obj in page.get('Contents', []): total += 1 filename = obj['Key'].rsplit('/', 1)[-1] cid = parse_parquet_cluster_id(filename) if cid not in cluster_ids: continue matched += 1 # Aggregate by component prefix for CSV readability # .../component=tidb/14_xxx_01.parquet -> component dir key = obj['Key'] comp_idx = key.find('/component=') if comp_idx >= 0: # prefix up to and including component=xxx/ rest = key[comp_idx + 1:] # component=tidb/... comp_prefix = date_prefix + rest.split('/', 1)[0] + '/' else: comp_prefix = date_prefix stats_key = f"{comp_prefix}#cluster={cid}" stats[stats_key]['count'] += 1 stats[stats_key]['total_size'] += obj['Size'] per_cluster[cid]['count'] += 1 per_cluster[cid]['total_size'] += obj['Size'] print(f" objects={total} matched_task_clusters={matched}") for cid, data in sorted(per_cluster.items()): size_mb = data['total_size'] / (1024 * 1024) prefix = f"{date_prefix}#cluster={cid}" write_stats_to_file('deltalake', prefix, data['count'], f"{size_mb:.2f}") # log-stats deltalake: deltalake/global/log-stats/date={YYYYMMDD}/org={org}/cluster={cluster}/ seen_oc = set() for tenant_id, cluster_id, storage_id in org_cluster: org = tenant_id or storage_id key = (org, cluster_id) if key in seen_oc: continue seen_oc.add(key) for d in dates_compact: prefix = f"deltalake/global/log-stats/date={d}/org={org}/cluster={cluster_id}/" scan_prefix_all('deltalake', bucket, prefix) def load_all_metas(bucket): task_keys = list(list_task_keys(bucket, TASK_PREFIX)) print(f"found {len(task_keys)} task files under s3://{bucket}/{TASK_PREFIX}") metas = [] for key in task_keys: meta = load_task_meta(bucket, key) if meta is None: print(f"skip invalid task: {key}") continue metas.append(meta) metas.extend(infra_metas()) print(f"active clusters: {len(metas)} (tasks + seed/o11y)") return metas def print_summary(): print(f"\n{'Prefix':80} | Files | Total Size (MB)") print("-" * 100) for prefix, data in sorted(stats.items()): size_mb = data['total_size'] / 1024 / 1024 print(f"{prefix:80} | {data['count']:5d} | {size_mb:12.2f}") def main(): metas = load_all_metas(BUCKET) if '0' in SCAN_KINDS: scan_kind_0(BUCKET, metas, only_regions=False) if 'regions' in SCAN_KINDS: scan_kind_0(BUCKET, metas, only_regions=True) if '1' in SCAN_KINDS: scan_kind_1(BUCKET) if 'diagnosis/data' in SCAN_KINDS: scan_kind_diagnosis_data(BUCKET, metas) if 'deltalake' in SCAN_KINDS: scan_kind_deltalake(BUCKET, metas) print_summary() print("\nCSV outputs:") for kind in SCAN_KINDS: print(f" {stats_filename(kind)}") if __name__ == '__main__': main()