diff --git a/cloudbuild/macrobenchmarks/macrobenchmarks-cloudbuild.yaml b/cloudbuild/macrobenchmarks/macrobenchmarks-cloudbuild.yaml index 13c10781..1e0752ba 100644 --- a/cloudbuild/macrobenchmarks/macrobenchmarks-cloudbuild.yaml +++ b/cloudbuild/macrobenchmarks/macrobenchmarks-cloudbuild.yaml @@ -95,6 +95,7 @@ steps: - "LOCATION=${LOCATION}" - "_BUCKET_TYPE=${_BUCKET_TYPE}" - "_ZONE=${_ZONE}" + - "_DATASET_PATH=${_DATASET_PATH}" args: ["cloudbuild/macrobenchmarks/scripts/create_buckets.sh"] waitFor: ["init-variables"] diff --git a/cloudbuild/macrobenchmarks/macrobenchmarks_schema.json b/cloudbuild/macrobenchmarks/macrobenchmarks_schema.json index 5204f9ed..1ffea17e 100644 --- a/cloudbuild/macrobenchmarks/macrobenchmarks_schema.json +++ b/cloudbuild/macrobenchmarks/macrobenchmarks_schema.json @@ -63,7 +63,26 @@ {"name": "checkpoint_delete_time_p100", "type": "FLOAT", "description": "Maximum (p100) checkpoint-delete duration in seconds."}, {"name": "num_checkpoint_delete_datapoints", "type": "INTEGER", "description": "Count of distinct checkpoint-delete datapoints."}, {"name": "accelerator_blocked_time", "type": "FLOAT", "description": "Seconds the accelerator was blocked waiting on data loading (bottleneck rank, run-wide)."}, - {"name": "accelerator_blocked_percent", "type": "FLOAT", "description": "Percent of wall-clock time the accelerator was blocked on data loading (bottleneck rank, run-wide)."} + {"name": "accelerator_blocked_percent", "type": "FLOAT", "description": "Percent of wall-clock time the accelerator was blocked on data loading (bottleneck rank, run-wide)."}, + {"name": "cpu_usage_peak_cores", "type": "FLOAT", "description": "Peak CPU usage in cores over the run window; max across worker pods (bottleneck pod)."}, + {"name": "cpu_usage_mean_cores", "type": "FLOAT", "description": "Mean CPU usage in cores over the run window; max of per-pod means across worker pods (bottleneck pod)."}, + {"name": "memory_usage_peak_bytes", "type": "INTEGER", "description": "Peak container memory usage in bytes over the run window; max across worker pods (bottleneck pod)."}, + {"name": "network_received_peak_bytes_per_sec", "type": "FLOAT", "description": "Peak pod network receive rate in bytes/s over the run window; max across worker pods (bottleneck pod)."}, + {"name": "network_received_mean_bytes_per_sec", "type": "FLOAT", "description": "Mean pod network receive rate in bytes/s over the run window; max of per-pod means across worker pods (bottleneck pod)."}, + {"name": "network_sent_peak_bytes_per_sec", "type": "FLOAT", "description": "Peak pod network send rate in bytes/s over the run window; max across worker pods (bottleneck pod)."}, + {"name": "network_sent_mean_bytes_per_sec", "type": "FLOAT", "description": "Mean pod network send rate in bytes/s over the run window; max of per-pod means across worker pods (bottleneck pod)."}, + {"name": "checkpoint_read_bytes", "type": "INTEGER", "description": "Total bytes served (egress) from the per-run checkpoint bucket over the run window; downloaded during restore."}, + {"name": "checkpoint_read_request_count", "type": "INTEGER", "description": "Count of object-read requests against the checkpoint bucket over the run window."}, + {"name": "checkpoint_restored_bytes", "type": "INTEGER", "description": "Size in bytes (du) of the restored checkpoint; denominator of the checkpoint read-amplification ratio."}, + {"name": "checkpoint_read_amplification_ratio", "type": "FLOAT", "description": "checkpoint_read_bytes / checkpoint_restored_bytes. DDP ~= world_size (every rank pulls the full checkpoint); FSDP ~= 1 (each pulls its shard)."}, + {"name": "dataset_read_bytes", "type": "INTEGER", "description": "Total bytes served (egress) from the per-run dataset bucket over the run window."}, + {"name": "dataset_read_request_count", "type": "INTEGER", "description": "Count of object-read requests against the dataset bucket over the run window."}, + {"name": "dataset_size_bytes", "type": "INTEGER", "description": "Size in bytes (du) of the per-run dataset bucket (a full unique copy)."}, + {"name": "dataset_read_amplification_ratio", "type": "FLOAT", "description": "dataset_read_bytes / (executed_steps * global_batch_size * dataset_size_bytes / dataset_sample_count). Egress per byte a perfectly sharded single pass over the consumed samples would need: ~1 = each needed byte fetched once, ~world_size = every rank re-read the same data. Normalized by samples consumed, so it is independent of dataset size, step count, and epoch count."}, + {"name": "memory_usage_mean_bytes", "type": "INTEGER", "description": "Mean container memory usage in bytes over the run window; max of per-pod means across worker pods (bottleneck pod)."}, + {"name": "memory_limit_utilization_peak", "type": "FLOAT", "description": "Peak fraction of the container memory limit over the run window (bottleneck pod); how close restore got to OOM."}, + {"name": "cpu_limit_utilization_peak", "type": "FLOAT", "description": "Peak fraction of the container CPU limit over the run window (bottleneck pod); whether the workload is pegged (compute-bound) vs I/O idle."}, + {"name": "dataset_sample_count", "type": "INTEGER", "description": "Estimated total rows in the dataset: dataset_size_bytes / per-sample stored bytes, where per-sample bytes are measured from the largest shard via the datasets library (format-agnostic, streamed). Used to derive per-sample bytes for the read-amplification ratio."} ] } } diff --git a/cloudbuild/macrobenchmarks/metrics/calculate.py b/cloudbuild/macrobenchmarks/metrics/calculate.py index bdefdfb3..e08bf5f5 100644 --- a/cloudbuild/macrobenchmarks/metrics/calculate.py +++ b/cloudbuild/macrobenchmarks/metrics/calculate.py @@ -168,6 +168,127 @@ def calc_data_loading_metrics(dl_rows: list) -> dict: } +# Maps series to schema columns. `None` mean-column means the series has no mean. +_SYSTEM_SERIES_COLUMNS = { + "cpu": ("cpu_usage_peak_cores", "cpu_usage_mean_cores"), + "memory": ("memory_usage_peak_bytes", "memory_usage_mean_bytes"), + "network_received": ( + "network_received_peak_bytes_per_sec", + "network_received_mean_bytes_per_sec", + ), + "network_sent": ( + "network_sent_peak_bytes_per_sec", + "network_sent_mean_bytes_per_sec", + ), + "cpu_limit_utilization": ("cpu_limit_utilization_peak", None), + "memory_limit_utilization": ("memory_limit_utilization_peak", None), + "checkpoint_read_bytes": ("checkpoint_read_bytes", None), + "checkpoint_read_request_count": ("checkpoint_read_request_count", None), + "checkpoint_restored_bytes": ("checkpoint_restored_bytes", None), + "dataset_read_bytes": ("dataset_read_bytes", None), + "dataset_read_request_count": ("dataset_read_request_count", None), + "dataset_size_bytes": ("dataset_size_bytes", None), + "dataset_sample_count": ("dataset_sample_count", None), +} + +# Columns reported as whole numbers (bytes / counts) rather than floats. +_INT_COLUMNS = { + "memory_usage_peak_bytes", + "memory_usage_mean_bytes", + "checkpoint_read_bytes", + "checkpoint_read_request_count", + "checkpoint_restored_bytes", + "dataset_read_bytes", + "dataset_read_request_count", + "dataset_size_bytes", + "dataset_sample_count", +} + + +def _amplification(numerator, denominator): + """numerator / denominator, or None when either is absent/zero.""" + if numerator is not None and denominator: + return numerator / denominator + return None + + +def executed_step_count(step_rows: list) -> int: + """Number of distinct optimizer steps observed (deduped across ranks). + + Each rank emits one row per optimizer step, so the count of unique step + numbers -- not rows -- is how many steps the run actually executed. + """ + return len( + { + r["step"] + for r in step_rows + if r.get("step") is not None and r.get("step_duration") is not None + } + ) + + +def dataset_read_amplification_ratio( + *, + dataset_read_bytes, + dataset_size_bytes, + dataset_sample_count, + executed_steps, + global_batch_size, +): + """dataset_read_bytes / ideal_bytes, or None if any input is absent/zero. + + ``ideal_bytes`` is the egress a perfectly sharded single pass over the + samples actually consumed (``executed_steps * global_batch_size * + dataset_size_bytes / dataset_sample_count``) would incur. Normalizing by + samples consumed, not the full dataset, makes the ratio independent of + dataset size and step count: ~1.0 means each byte was fetched once, + ~world_size means every rank re-read the same data. + """ + if None in ( + dataset_read_bytes, + dataset_size_bytes, + dataset_sample_count, + executed_steps, + global_batch_size, + ): + return None + ideal_bytes = ( + executed_steps * global_batch_size * dataset_size_bytes / dataset_sample_count + if dataset_sample_count + else 0 + ) + if not ideal_bytes: + return None + return dataset_read_bytes / ideal_bytes + + +def calc_system_metrics(system_rows: list) -> dict: + """Reduce per-pod/per-bucket metrics to the bottleneck value and derive ratios.""" + out = {} + by_metric = defaultdict(list) + for r in system_rows: + by_metric[r.get("metric")].append(r) + for series, (peak_col, mean_col) in _SYSTEM_SERIES_COLUMNS.items(): + rows = by_metric.get(series, []) + peaks = [r["peak"] for r in rows if r.get("peak") is not None] + if peaks: + val = max(peaks) + out[peak_col] = int(val) if peak_col in _INT_COLUMNS else val + if mean_col: + means = [r["mean"] for r in rows if r.get("mean") is not None] + if means: + val = max(means) + out[mean_col] = int(val) if mean_col in _INT_COLUMNS else val + ratio = _amplification( + out.get("checkpoint_read_bytes"), out.get("checkpoint_restored_bytes") + ) + if ratio is not None: + out["checkpoint_read_amplification_ratio"] = ratio + # Dataset ratio is derived in build_summary_row; it needs step/batch-size + # inputs this reducer doesn't have. + return out + + def build_summary_row( *, run_id: str, @@ -178,6 +299,7 @@ def build_summary_row( restore_rows: list, delete_rows: list, dl_rows: list, + system_rows: list = None, dimensions: dict = None, ) -> dict: row = { @@ -192,6 +314,18 @@ def build_summary_row( row.update(calc_restore_metrics(restore_rows)) row.update(calc_delete_metrics(delete_rows)) row.update(calc_data_loading_metrics(dl_rows)) + row.update(calc_system_metrics(system_rows or [])) + # Derived here, not in calc_system_metrics, since it needs executed_steps + # and global_batch_size alongside the raw dataset columns just produced. + ratio = dataset_read_amplification_ratio( + dataset_read_bytes=row.get("dataset_read_bytes"), + dataset_size_bytes=row.get("dataset_size_bytes"), + dataset_sample_count=row.get("dataset_sample_count"), + executed_steps=executed_step_count(step_rows), + global_batch_size=row.get("global_batch_size"), + ) + if ratio is not None: + row["dataset_read_amplification_ratio"] = ratio return row @@ -335,6 +469,7 @@ def main(argv=None) -> None: restore_rows = tables.restore_rows delete_rows = tables.delete_rows dl_rows = tables.dl_rows + system_rows = tables.system_rows validate_required_metrics( step_rows=step_rows, @@ -394,6 +529,7 @@ def main(argv=None) -> None: restore_rows=restore_rows, delete_rows=delete_rows, dl_rows=dl_rows, + system_rows=system_rows, dimensions=dimensions, ) diff --git a/cloudbuild/macrobenchmarks/metrics/monitoring.py b/cloudbuild/macrobenchmarks/metrics/monitoring.py new file mode 100644 index 00000000..22e982c9 --- /dev/null +++ b/cloudbuild/macrobenchmarks/metrics/monitoring.py @@ -0,0 +1,315 @@ +"""Fetch per-pod and per-bucket system metrics from Cloud Monitoring. + +Usage: + python3 -m metrics.monitoring --project P --run-id R \ + --start-time RFC3339 --end-time RFC3339 --out-dir DIR \ + [--checkpoint-bucket B] [--dataset-bucket B] +""" + +import argparse +import datetime +import statistics +import traceback +from dataclasses import dataclass + +from metrics import raw_store, schema + + +@dataclass(frozen=True) +class Series: + """One monitoring series mapped to our internal metric name.""" + + name: str # internal series name + metric_type: str # Cloud Monitoring metric.type + resource_type: str # k8s_container | k8s_pod | gcs_bucket + aligner: str # per-series aligner name + filter_kind: str = "pod" # "pod" (pod_name prefix) | "bucket" (bucket_name) + method: str = None # optional metric.labels.method filter (bucket series) + + +# CPU: cores (RATE), Memory: peak bytes (MAX), Network: bytes/s (RATE), +# limit utilizations: fraction of the container limit (MAX). +SERIES = [ + Series( + "cpu", + "kubernetes.io/container/cpu/core_usage_time", + "k8s_container", + "ALIGN_RATE", + ), + Series( + "memory", + "kubernetes.io/container/memory/used_bytes", + "k8s_container", + "ALIGN_MAX", + ), + Series( + "network_received", + "kubernetes.io/pod/network/received_bytes_count", + "k8s_pod", + "ALIGN_RATE", + ), + Series( + "network_sent", + "kubernetes.io/pod/network/sent_bytes_count", + "k8s_pod", + "ALIGN_RATE", + ), + Series( + "cpu_limit_utilization", + "kubernetes.io/container/cpu/limit_utilization", + "k8s_container", + "ALIGN_MAX", + ), + Series( + "memory_limit_utilization", + "kubernetes.io/container/memory/limit_utilization", + "k8s_container", + "ALIGN_MAX", + ), +] + +# Per-bucket totals summed over the window; `name` is prefixed with +# "checkpoint"/"dataset" to form the metric/column name. +GCS_BUCKET_SERIES = [ + Series( + "read_bytes", + "storage.googleapis.com/network/sent_bytes_count", + "gcs_bucket", + "ALIGN_DELTA", + filter_kind="bucket", + ), + Series( + "read_request_count", + "storage.googleapis.com/api/request_count", + "gcs_bucket", + "ALIGN_DELTA", + filter_kind="bucket", + method="ReadObject", + ), +] + + +def _to_epoch(rfc3339: str) -> int: + """Parse RFC3339 to epoch seconds.""" + dt = datetime.datetime.fromisoformat(rfc3339.upper().replace("Z", "+00:00")) + return int(dt.timestamp()) + + +def _point_value(point) -> float: + """Read numeric value from point.""" + v = point.value + if getattr(v, "double_value", 0.0): + return float(v.double_value) + return float(getattr(v, "int64_value", 0)) + + +def reduce_points(values: list) -> tuple: + """Return (peak, mean) of values, or (None, None).""" + if not values: + return None, None + return max(values), statistics.mean(values) + + +def _build_request(project, series, target, start_epoch, end_epoch, period): + """Build a list_time_series request; ``target`` is the run id or bucket name.""" + if series.filter_kind == "bucket": + filter_ = ( + f'metric.type = "{series.metric_type}" ' + f'AND resource.type = "{series.resource_type}" ' + f'AND resource.labels.bucket_name = "{target}"' + ) + if series.method: + filter_ += f' AND metric.labels.method = "{series.method}"' + else: + filter_ = ( + f'metric.type = "{series.metric_type}" ' + f'AND resource.type = "{series.resource_type}" ' + f'AND resource.labels.pod_name = starts_with("{target}-workload-0-")' + ) + return { + "name": f"projects/{project}", + "filter": filter_, + "interval": { + "start_time": {"seconds": int(start_epoch)}, + "end_time": {"seconds": int(end_epoch)}, + }, + "aggregation": { + "alignment_period": {"seconds": period}, + "per_series_aligner": series.aligner, + }, + } + + +def collect(client, *, project, run_id, start_epoch, end_epoch, period=60) -> list: + """Collect SystemMetric rows for all per-pod SERIES; a failed series is skipped.""" + rows = [] + for series in SERIES: + try: + request = _build_request( + project, series, run_id, start_epoch, end_epoch, period + ) + for ts in client.list_time_series(request): + pod_name = ts.resource.labels.get("pod_name", "") + values = [_point_value(p) for p in ts.points] + peak, mean = reduce_points(values) + if peak is None: + continue + rows.append( + schema.SystemMetric( + pod_name=pod_name, metric=series.name, peak=peak, mean=mean + ) + ) + except Exception as e: # best-effort: keep the other series + print( + f"Warning: system series '{series.name}' failed, its columns N/A: {e}" + ) + return rows + + +def collect_bucket_totals( + client, *, project, bucket, prefix, start_epoch, end_epoch, period=60 +) -> list: + """Sum each GCS bucket series over the window into one SystemMetric row each.""" + rows = [] + for series in GCS_BUCKET_SERIES: + try: + request = _build_request( + project, series, bucket, start_epoch, end_epoch, period + ) + total = 0.0 + found = False + for ts in client.list_time_series(request): + for p in ts.points: + total += _point_value(p) + found = True + if found: + rows.append( + schema.SystemMetric( + pod_name=bucket, + metric=f"{prefix}_{series.name}", + peak=total, + mean=None, + ) + ) + except Exception as e: # best-effort: keep the other series + print( + f"Warning: bucket series '{prefix}_{series.name}' failed, " + f"its column N/A: {e}" + ) + return rows + + +def assemble_rows( + client, + storage_client, + *, + project, + run_id, + checkpoint_bucket, + dataset_bucket, + restore_rows, + start_epoch, + end_epoch, + period=60, +) -> list: + """Pod gauges + per-bucket totals + du sizes, as one row list.""" + rows = collect( + client, + project=project, + run_id=run_id, + start_epoch=start_epoch, + end_epoch=end_epoch, + period=period, + ) + if checkpoint_bucket: + rows += collect_bucket_totals( + client, + project=project, + bucket=checkpoint_bucket, + prefix="checkpoint", + start_epoch=start_epoch, + end_epoch=end_epoch, + period=period, + ) + if dataset_bucket: + rows += collect_bucket_totals( + client, + project=project, + bucket=dataset_bucket, + prefix="dataset", + start_epoch=start_epoch, + end_epoch=end_epoch, + period=period, + ) + if storage_client is not None: + from metrics import sizes + + # Best-effort; must not discard the rows already collected above. + try: + loc = sizes.restored_checkpoint_location(restore_rows or []) + rows += sizes.size_rows( + storage_client, dataset_bucket=dataset_bucket, restored_location=loc + ) + except Exception as e: + print(f"Warning: du sizes failed, size columns N/A: {e}") + return rows + + +def main(argv=None) -> None: + parser = argparse.ArgumentParser( + description="Fetch per-pod system metrics from Cloud Monitoring." + ) + parser.add_argument("--project", required=True) + parser.add_argument("--run-id", required=True) + parser.add_argument("--start-time", required=True, help="RFC3339") + parser.add_argument("--end-time", required=True, help="RFC3339") + parser.add_argument("--out-dir", required=True) + parser.add_argument("--checkpoint-bucket") + parser.add_argument("--dataset-bucket") + parser.add_argument("--period", type=int, default=60) + args = parser.parse_args(argv) + + # Import here to handle missing library case separately. + try: + from google.cloud import monitoring_v3 + except ImportError as e: + print(f"Warning: google-cloud-monitoring unavailable, columns will be N/A: {e}") + return + + try: + client = monitoring_v3.MetricServiceClient() + storage_client = None + try: + from google.cloud import storage + + storage_client = storage.Client(project=args.project) + except Exception as e: # sizes become N/A, run continues + print(f"Warning: storage client unavailable, size columns N/A: {e}") + restore_rows = [] + try: + restore_rows = raw_store.read_raw_metrics(args.out_dir).restore_rows + except Exception as e: + print(f"Warning: could not read restore rows for checkpoint du: {e}") + rows = assemble_rows( + client, + storage_client, + project=args.project, + run_id=args.run_id, + checkpoint_bucket=args.checkpoint_bucket, + dataset_bucket=args.dataset_bucket, + restore_rows=restore_rows, + start_epoch=_to_epoch(args.start_time), + end_epoch=_to_epoch(args.end_time), + period=args.period, + ) + raw_store.write_system_metrics(rows, args.out_dir) + print(f"Wrote {len(rows)} system-metric rows to {args.out_dir}") + except Exception as e: # best-effort + print( + f"Warning: system-metrics fetch failed, columns will be N/A: {e}\n" + f"{traceback.format_exc()}" + ) + + +if __name__ == "__main__": + main() diff --git a/cloudbuild/macrobenchmarks/metrics/raw_store.py b/cloudbuild/macrobenchmarks/metrics/raw_store.py index 2af53bbd..f419cf79 100644 --- a/cloudbuild/macrobenchmarks/metrics/raw_store.py +++ b/cloudbuild/macrobenchmarks/metrics/raw_store.py @@ -26,6 +26,7 @@ class RawMetricTables: restore_rows: List[dict] = field(default_factory=list) delete_rows: List[dict] = field(default_factory=list) dl_rows: List[dict] = field(default_factory=list) + system_rows: List[dict] = field(default_factory=list) def write_raw_metrics( @@ -93,6 +94,18 @@ def write_raw_metrics( ) +def write_system_metrics(system_rows, out_dir: str) -> None: + """Write SystemMetric rows to the system-metrics CSV (owned here like the rest).""" + if system_rows: + _write_csv( + os.path.join( + out_dir, schema.SYSTEM_METRICS_DIRECTORY, schema.SYSTEM_METRICS_FILE + ), + schema.SystemMetric, + system_rows, + ) + + def read_raw_metrics( in_dir: str, *, run_type: str = "perf_optimization" ) -> RawMetricTables: @@ -129,6 +142,11 @@ def read_raw_metrics( schema.DATA_LOADING_METRICS_FILE, ) ), + system_rows=_read_csv( + os.path.join( + in_dir, schema.SYSTEM_METRICS_DIRECTORY, schema.SYSTEM_METRICS_FILE + ) + ), ) diff --git a/cloudbuild/macrobenchmarks/metrics/requirements.txt b/cloudbuild/macrobenchmarks/metrics/requirements.txt index c638cbce..c174cffd 100644 --- a/cloudbuild/macrobenchmarks/metrics/requirements.txt +++ b/cloudbuild/macrobenchmarks/metrics/requirements.txt @@ -1,4 +1,6 @@ +datasets google-cloud-logging +google-cloud-monitoring google-cloud-storage numpy pytest diff --git a/cloudbuild/macrobenchmarks/metrics/schema.py b/cloudbuild/macrobenchmarks/metrics/schema.py index da9dc8a0..c6302563 100644 --- a/cloudbuild/macrobenchmarks/metrics/schema.py +++ b/cloudbuild/macrobenchmarks/metrics/schema.py @@ -18,6 +18,8 @@ PER_ACCELERATOR_DIRECTORY = "per_accelerator" CALCULATED_METRICS_DIRECTORY = "calculated_metrics" DATA_LOADING_METRICS_FILE = "data_loading_metrics.csv" +SYSTEM_METRICS_DIRECTORY = "system_metrics" +SYSTEM_METRICS_FILE = "system_metrics.csv" def fieldnames(dataclass_type) -> list: @@ -76,3 +78,11 @@ class DataLoadingMetrics: accelerator_blocked_time: float = None accelerator_blocked_percent: float = None update_timestamp: str = None + + +@dataclass(kw_only=True) +class SystemMetric: + pod_name: str + metric: str + peak: float + mean: float = None diff --git a/cloudbuild/macrobenchmarks/metrics/sizes.py b/cloudbuild/macrobenchmarks/metrics/sizes.py new file mode 100644 index 00000000..220eefe5 --- /dev/null +++ b/cloudbuild/macrobenchmarks/metrics/sizes.py @@ -0,0 +1,149 @@ +"""GCS ``du`` sizes for the read-amplification denominators (best-effort).""" + +import os +import tempfile + +from metrics import schema + + +def _split_gs_path(gs_path: str): + """(bucket, prefix) for ``gs://bucket[/prefix]`` (scheme optional).""" + rest = gs_path[len("gs://") :] if gs_path.startswith("gs://") else gs_path + bucket_name, _, prefix = rest.partition("/") + return bucket_name, prefix + + +def gcs_du(storage_client, gs_path: str): + """Sum blob sizes under ``gs_path`` (gs://bucket[/prefix]); None if empty.""" + bucket_name, prefix = _split_gs_path(gs_path) + total = 0 + found = False + for blob in storage_client.list_blobs(bucket_name, prefix=prefix): + total += blob.size or 0 + found = True + return total if found else None + + +def _hf_num_rows(local_dir: str, filename: str) -> int: + """Row count of one local shard, format inferred from its extension. + + Delegates to ``datasets`` to stay format-agnostic (parquet, arrow, json, + csv, webdataset, ...); streaming avoids caching the decoded shard to disk. + """ + # Force offline: local packaged builders need no Hub access, and a + # sandboxed metrics step must never block on a Hub round-trip. + os.environ.setdefault("HF_HUB_OFFLINE", "1") + os.environ.setdefault("HF_DATASETS_OFFLINE", "1") + import datasets + + ds = datasets.load_dataset( + path=local_dir, + data_files=filename, + split="train", + streaming=True, + cache_dir=local_dir, + ) + return sum(1 for _ in ds) + + +def _download_and_count(storage_client, blob) -> int: + """Download one shard to a temp dir and count its rows via ``datasets``.""" + with tempfile.TemporaryDirectory() as d: + name = blob.name.rsplit("/", 1)[-1] + blob.download_to_filename(os.path.join(d, name)) + return _hf_num_rows(d, name) + + +def _largest_shard(blobs): + """Largest object among ``blobs`` (None if none have a size). + + Used as the representative shard for per-sample bytes: a small remainder + shard would over-weight fixed per-file overhead and skew the ratio. + """ + sized = [b for b in blobs if b.size] + return max(sized, key=lambda b: b.size) if sized else None + + +def dataset_du_and_sample_count( + storage_client, gs_path: str, *, count_shard=_download_and_count +): + """``(du_bytes, estimated_total_rows)`` from a single bucket listing. + + ``du`` is the summed object size; the row count is ``du / per_sample_bytes``, + with per-sample bytes measured from one downloaded/parsed shard (the + largest). Counting is best-effort: on failure ``du`` is still returned and + only the count is dropped. Both are None when the path holds no objects. + """ + bucket_name, prefix = _split_gs_path(gs_path) + blobs = list(storage_client.list_blobs(bucket_name, prefix=prefix)) + if not blobs: + return None, None + du = sum(b.size or 0 for b in blobs) + count = None + shard = _largest_shard(blobs) + if du and shard is not None: + try: + rows = count_shard(storage_client, shard) + except Exception as e: # noqa: BLE001 - best-effort metric + print(f"Warning: dataset sample count failed, ratio N/A: {e}") + rows = None + if rows: + count = round(du / (shard.size / rows)) + return du, count + + +def restored_checkpoint_location(restore_rows: list): + """checkpoint_location of the earliest-ending restore (the resume), or None.""" + candidates = [ + r + for r in restore_rows + if r.get("checkpoint_location") and r.get("end_time") is not None + ] + if not candidates: + return None + return min(candidates, key=lambda r: r["end_time"])["checkpoint_location"] + + +def size_rows( + storage_client, + *, + dataset_bucket, + restored_location, + count_shard=_download_and_count, +) -> list: + """SystemMetric rows for the dataset-bucket size and restored-checkpoint size.""" + rows = [] + if dataset_bucket: + du, count = dataset_du_and_sample_count( + storage_client, f"gs://{dataset_bucket}", count_shard=count_shard + ) + if du is not None: + rows.append( + schema.SystemMetric( + pod_name=dataset_bucket, + metric="dataset_size_bytes", + peak=du, + mean=None, + ) + ) + if count is not None: + rows.append( + schema.SystemMetric( + pod_name=dataset_bucket, + metric="dataset_sample_count", + peak=count, + mean=None, + ) + ) + if restored_location: + size = gcs_du(storage_client, restored_location) + if size is not None: + rows.append( + schema.SystemMetric( + pod_name=restored_location, + metric="checkpoint_restored_bytes", + peak=size, + mean=None, + ) + ) + return rows diff --git a/cloudbuild/macrobenchmarks/metrics/tests/test_calculate_dataset_amplification.py b/cloudbuild/macrobenchmarks/metrics/tests/test_calculate_dataset_amplification.py new file mode 100644 index 00000000..33b1da8e --- /dev/null +++ b/cloudbuild/macrobenchmarks/metrics/tests/test_calculate_dataset_amplification.py @@ -0,0 +1,175 @@ +from metrics import calculate + + +def test_ratio_normalizes_by_samples_consumed_not_full_dataset(): + # size=1000 over 100 samples -> 10 bytes/sample. A run of 5 steps at + # global_batch_size 8 consumes 40 samples, so a perfectly sharded single + # pass would read 40*10 = 400 bytes. Actual read of 800 -> 2x amplification. + ratio = calculate.dataset_read_amplification_ratio( + dataset_read_bytes=800, + dataset_size_bytes=1000, + dataset_sample_count=100, + executed_steps=5, + global_batch_size=8, + ) + assert ratio == 2.0 + + +def test_ratio_independent_of_dataset_size_when_partial_pass(): + # Same run, but the dataset is 10x larger (only a tenth is consumed). The + # naive read/du ratio would collapse to 0.2 and hide duplication; the + # normalized ratio stays 2.0 because per-sample bytes are unchanged. + ratio = calculate.dataset_read_amplification_ratio( + dataset_read_bytes=800, + dataset_size_bytes=10000, + dataset_sample_count=1000, + executed_steps=5, + global_batch_size=8, + ) + assert ratio == 2.0 + + +def test_ratio_none_when_any_input_missing(): + for missing in ( + "dataset_read_bytes", + "dataset_size_bytes", + "dataset_sample_count", + "executed_steps", + "global_batch_size", + ): + kwargs = { + "dataset_read_bytes": 800, + "dataset_size_bytes": 1000, + "dataset_sample_count": 100, + "executed_steps": 5, + "global_batch_size": 8, + } + kwargs[missing] = None + assert calculate.dataset_read_amplification_ratio(**kwargs) is None, missing + + +def test_ratio_none_when_denominator_component_zero(): + # Zero sample count, zero steps, or zero batch size would divide by zero. + for zeroed in ("dataset_sample_count", "executed_steps", "global_batch_size"): + kwargs = { + "dataset_read_bytes": 800, + "dataset_size_bytes": 1000, + "dataset_sample_count": 100, + "executed_steps": 5, + "global_batch_size": 8, + } + kwargs[zeroed] = 0 + assert calculate.dataset_read_amplification_ratio(**kwargs) is None, zeroed + + +def test_ratio_zero_when_no_bytes_read(): + # No egress is a valid measurement (ratio 0.0), not a missing one. + ratio = calculate.dataset_read_amplification_ratio( + dataset_read_bytes=0, + dataset_size_bytes=1000, + dataset_sample_count=100, + executed_steps=5, + global_batch_size=8, + ) + assert ratio == 0.0 + + +def test_executed_step_count_dedupes_steps_across_ranks(): + # Under DDP every rank emits a row per optimizer step; the run executed as + # many distinct steps as there are unique step numbers, not rows. + rows = [ + {"step": 0, "step_duration": 1.0}, + {"step": 1, "step_duration": 1.0}, + {"step": 0, "step_duration": 1.1}, # same step, second rank + {"step": 1, "step_duration": 1.2}, + {"step": 2, "step_duration": 1.0}, + ] + assert calculate.executed_step_count(rows) == 3 + + +def test_executed_step_count_ignores_incomplete_rows(): + rows = [ + {"step": 0, "step_duration": 1.0}, + {"step": None, "step_duration": 1.0}, + {"step": 1, "step_duration": None}, + {"step": 2, "step_duration": 1.0}, + ] + assert calculate.executed_step_count(rows) == 2 + + +def test_executed_step_count_empty_is_zero(): + assert calculate.executed_step_count([]) == 0 + + +def _dataset_system_rows(read_bytes=800, size_bytes=1000, sample_count=100): + return [ + { + "pod_name": "ds", + "metric": "dataset_read_bytes", + "peak": read_bytes, + "mean": None, + }, + { + "pod_name": "ds", + "metric": "dataset_size_bytes", + "peak": size_bytes, + "mean": None, + }, + { + "pod_name": "ds", + "metric": "dataset_sample_count", + "peak": sample_count, + "mean": None, + }, + ] + + +def test_build_summary_row_emits_dataset_read_amplification_ratio(): + # 5 executed steps * gbs 8 = 40 samples consumed; per-sample bytes = 1000/100 + # = 10, so a single sharded pass should read 400 bytes. Actual 800 -> 2.0. + row = calculate.build_summary_row( + run_id="r", + workload_name="w", + requirements="gcsfs==1.0", + step_rows=[{"step": s, "step_duration": 1.0} for s in range(5)], + write_rows=[], + restore_rows=[], + delete_rows=[], + dl_rows=[], + system_rows=_dataset_system_rows(), + dimensions={"global_batch_size": 8}, + ) + assert row["dataset_read_amplification_ratio"] == 2.0 + + +def test_build_summary_row_omits_ratio_without_global_batch_size(): + row = calculate.build_summary_row( + run_id="r", + workload_name="w", + requirements="gcsfs==1.0", + step_rows=[{"step": s, "step_duration": 1.0} for s in range(5)], + write_rows=[], + restore_rows=[], + delete_rows=[], + dl_rows=[], + system_rows=_dataset_system_rows(), + dimensions=None, # no global_batch_size + ) + assert "dataset_read_amplification_ratio" not in row + + +def test_build_summary_row_omits_ratio_without_sample_count(): + rows = [r for r in _dataset_system_rows() if r["metric"] != "dataset_sample_count"] + row = calculate.build_summary_row( + run_id="r", + workload_name="w", + requirements="gcsfs==1.0", + step_rows=[{"step": s, "step_duration": 1.0} for s in range(5)], + write_rows=[], + restore_rows=[], + delete_rows=[], + dl_rows=[], + system_rows=rows, + dimensions={"global_batch_size": 8}, + ) + assert "dataset_read_amplification_ratio" not in row diff --git a/cloudbuild/macrobenchmarks/metrics/tests/test_calculate_summary.py b/cloudbuild/macrobenchmarks/metrics/tests/test_calculate_summary.py index 4263c26c..dc7339a1 100644 --- a/cloudbuild/macrobenchmarks/metrics/tests/test_calculate_summary.py +++ b/cloudbuild/macrobenchmarks/metrics/tests/test_calculate_summary.py @@ -258,6 +258,101 @@ def _write_restore_csv(in_dir): w.writerow([0, "gs://b/ckpt", 10.0, 18.0, 0, ""]) +def _write_system_metrics_csv(in_dir): + (in_dir / "system_metrics").mkdir(parents=True) + path = in_dir / "system_metrics" / "system_metrics.csv" + with open(path, "w", newline="") as f: + w = csv.writer(f) + w.writerow(["pod_name", "metric", "peak", "mean"]) + w.writerow(["p0", "cpu", 3.0, 1.0]) + w.writerow(["p1", "cpu", 5.0, 4.0]) + w.writerow(["p0", "memory", 2048.0, ""]) + w.writerow(["p0", "network_received", 10.0, 2.0]) + + +def test_main_emits_system_metric_columns(tmp_path): + # Verify system metrics are reduced to bottleneck pod and typed correctly. + in_dir = tmp_path / "raw" + _write_step_csv(in_dir) + _write_data_loading_csv(in_dir) + _write_system_metrics_csv(in_dir) + out_file = tmp_path / "summary.csv" + calculate.main( + [ + "--run-id", + "r", + "--workload-name", + "hf-pytorch-lightning-cpu", + "--requirements", + "gcsfs==1.0", + "--in-dir", + str(in_dir), + "--out-file", + str(out_file), + "--require-data-loading-metrics", + ] + ) + with open(out_file) as f: + rows = list(csv.DictReader(f)) + assert rows[0]["cpu_usage_peak_cores"] == "5.0" + assert rows[0]["cpu_usage_mean_cores"] == "4.0" # max of per-pod means + assert rows[0]["memory_usage_peak_bytes"] == "2048" # int-typed + assert rows[0]["network_received_peak_bytes_per_sec"] == "10.0" + assert rows[0]["network_received_mean_bytes_per_sec"] == "2.0" + assert rows[0]["network_sent_peak_bytes_per_sec"] == "N/A" + + +def _write_dataset_system_metrics_csv(in_dir): + (in_dir / "system_metrics").mkdir(parents=True) + path = in_dir / "system_metrics" / "system_metrics.csv" + with open(path, "w", newline="") as f: + w = csv.writer(f) + w.writerow(["pod_name", "metric", "peak", "mean"]) + w.writerow(["ds", "dataset_read_bytes", 80.0, ""]) + w.writerow(["ds", "dataset_size_bytes", 1000.0, ""]) + w.writerow(["ds", "dataset_sample_count", 100.0, ""]) + + +def test_main_emits_dataset_read_amplification_ratio(tmp_path): + # Two executed steps (0, 1) * global_batch_size 2 = 4 samples consumed; + # per-sample bytes = 1000/100 = 10, so an ideal single sharded pass reads + # 40 bytes. Actual egress 80 -> ratio 2.0, end to end through the CSV. + in_dir = tmp_path / "raw" + _write_step_csv(in_dir) # steps 0 and 1 + _write_data_loading_csv(in_dir) + _write_dataset_system_metrics_csv(in_dir) + out_file = tmp_path / "summary.csv" + calculate.main( + [ + "--run-id", + "r", + "--workload-name", + "hf-pytorch-lightning-cpu", + "--requirements", + "gcsfs==1.0", + "--in-dir", + str(in_dir), + "--out-file", + str(out_file), + "--require-data-loading-metrics", + "--per-device-batch", + "2", + "--grad-accum", + "1", + "--nodes", + "1", + "--ranks-per-node", + "1", + ] + ) + with open(out_file) as f: + rows = list(csv.DictReader(f)) + assert rows[0]["global_batch_size"] == "2" + assert rows[0]["dataset_sample_count"] == "100" + assert rows[0]["dataset_read_amplification_ratio"] == "2.0" + assert "dataset_read_amplification_ratio" in calculate.SUMMARY_FIELDNAMES + + def test_main_fails_when_required_data_loading_metrics_are_missing(tmp_path): # step metrics present, but no data_loading_metrics.csv -> must fail when # --require-data-loading-metrics is set (the profiler summary is required). diff --git a/cloudbuild/macrobenchmarks/metrics/tests/test_calculate_system.py b/cloudbuild/macrobenchmarks/metrics/tests/test_calculate_system.py new file mode 100644 index 00000000..5155b205 --- /dev/null +++ b/cloudbuild/macrobenchmarks/metrics/tests/test_calculate_system.py @@ -0,0 +1,125 @@ +from metrics import calculate + + +def test_max_across_pods_per_metric(): + rows = [ + {"pod_name": "p0", "metric": "cpu", "peak": 3.0, "mean": 1.0}, + {"pod_name": "p1", "metric": "cpu", "peak": 5.0, "mean": 4.0}, + {"pod_name": "p0", "metric": "memory", "peak": 1024.0, "mean": None}, + {"pod_name": "p1", "metric": "memory", "peak": 2048.0, "mean": None}, + {"pod_name": "p0", "metric": "network_received", "peak": 10.0, "mean": 2.0}, + {"pod_name": "p1", "metric": "network_received", "peak": 8.0, "mean": 3.0}, + {"pod_name": "p0", "metric": "network_sent", "peak": 7.0, "mean": 1.0}, + {"pod_name": "p1", "metric": "network_sent", "peak": 9.0, "mean": 5.0}, + ] + m = calculate.calc_system_metrics(rows) + assert m["cpu_usage_peak_cores"] == 5.0 + assert m["cpu_usage_mean_cores"] == 4.0 # max of per-pod means + assert m["memory_usage_peak_bytes"] == 2048 # int + assert isinstance(m["memory_usage_peak_bytes"], int) + assert m["network_received_peak_bytes_per_sec"] == 10.0 + assert m["network_received_mean_bytes_per_sec"] == 3.0 + assert m["network_sent_peak_bytes_per_sec"] == 9.0 + assert m["network_sent_mean_bytes_per_sec"] == 5.0 + + +def test_empty_rows_yield_no_keys(): + assert calculate.calc_system_metrics([]) == {} + + +def test_missing_series_omits_its_columns(): + rows = [{"pod_name": "p0", "metric": "cpu", "peak": 2.0, "mean": 1.0}] + m = calculate.calc_system_metrics(rows) + assert m["cpu_usage_peak_cores"] == 2.0 + assert "memory_usage_peak_bytes" not in m + assert "network_received_peak_bytes_per_sec" not in m + + +def test_maps_limit_utilization_and_mean_memory(): + rows = [ + {"pod_name": "p0", "metric": "memory", "peak": 4096.0, "mean": 2048.0}, + {"pod_name": "p1", "metric": "memory", "peak": 8192.0, "mean": 1024.0}, + {"pod_name": "p0", "metric": "cpu_limit_utilization", "peak": 0.7, "mean": 0.3}, + {"pod_name": "p1", "metric": "cpu_limit_utilization", "peak": 0.9, "mean": 0.4}, + { + "pod_name": "p0", + "metric": "memory_limit_utilization", + "peak": 0.5, + "mean": None, + }, + ] + m = calculate.calc_system_metrics(rows) + assert m["memory_usage_mean_bytes"] == 2048 # max of per-pod means, int + assert isinstance(m["memory_usage_mean_bytes"], int) + assert m["cpu_limit_utilization_peak"] == 0.9 + assert m["memory_limit_utilization_peak"] == 0.5 + + +def test_checkpoint_amplification_ratio_and_raw_dataset_columns(): + # The checkpoint ratio is derivable from system rows alone, so it stays in + # calc_system_metrics. The dataset ratio needs run-shape inputs (steps, + # global_batch_size) that live outside system rows, so here we only surface + # the raw dataset columns (read/size/sample_count) that feed it downstream. + rows = [ + { + "pod_name": "ckpt-bkt", + "metric": "checkpoint_read_bytes", + "peak": 800.0, + "mean": None, + }, + { + "pod_name": "gs://ckpt", + "metric": "checkpoint_restored_bytes", + "peak": 100.0, + "mean": None, + }, + { + "pod_name": "ds-bkt", + "metric": "dataset_read_bytes", + "peak": 3000.0, + "mean": None, + }, + { + "pod_name": "ds-bkt", + "metric": "dataset_size_bytes", + "peak": 1000.0, + "mean": None, + }, + { + "pod_name": "ds-bkt", + "metric": "dataset_sample_count", + "peak": 100.0, + "mean": None, + }, + { + "pod_name": "ckpt-bkt", + "metric": "checkpoint_read_request_count", + "peak": 42.0, + "mean": None, + }, + ] + m = calculate.calc_system_metrics(rows) + assert m["checkpoint_read_bytes"] == 800 + assert isinstance(m["checkpoint_read_bytes"], int) + assert m["checkpoint_restored_bytes"] == 100 + assert m["checkpoint_read_amplification_ratio"] == 8.0 + assert m["checkpoint_read_request_count"] == 42 + # Raw dataset columns surfaced (int-typed); the ratio is computed later. + assert m["dataset_read_bytes"] == 3000 + assert m["dataset_size_bytes"] == 1000 + assert m["dataset_sample_count"] == 100 + assert isinstance(m["dataset_sample_count"], int) + assert "dataset_read_amplification_ratio" not in m + + +def test_checkpoint_amplification_ratio_omitted_when_denominator_missing(): + rows = [ + { + "pod_name": "b", + "metric": "checkpoint_read_bytes", + "peak": 500.0, + "mean": None, + }, + ] + m = calculate.calc_system_metrics(rows) + assert "checkpoint_read_amplification_ratio" not in m # no restored_bytes diff --git a/cloudbuild/macrobenchmarks/metrics/tests/test_monitoring.py b/cloudbuild/macrobenchmarks/metrics/tests/test_monitoring.py new file mode 100644 index 00000000..121eaf3e --- /dev/null +++ b/cloudbuild/macrobenchmarks/metrics/tests/test_monitoring.py @@ -0,0 +1,305 @@ +from types import SimpleNamespace + +from metrics import monitoring + + +def _point(value): + return SimpleNamespace(value=SimpleNamespace(double_value=value, int64_value=0)) + + +def _series(pod_name, values): + return SimpleNamespace( + resource=SimpleNamespace(labels={"pod_name": pod_name}), + points=[_point(v) for v in values], + ) + + +class _FakeClient: + """Returns a canned series list keyed by the full metric.type in the filter. + + Matches the exact ``metric.type = ""`` token rather than a loose + substring so a pod series (kubernetes.io/pod/network/sent_bytes_count) does + not accidentally satisfy a bucket series query + (storage.googleapis.com/network/sent_bytes_count) or vice-versa. + """ + + def __init__(self, by_metric_type): + self.by_metric_type = by_metric_type + self.requests = [] + + def list_time_series(self, request): + self.requests.append(request) + for metric_type, series in self.by_metric_type.items(): + if f'metric.type = "{metric_type}"' in request["filter"]: + return iter(series) + return iter(()) + + +def test_reduce_points(): + assert monitoring.reduce_points([1.0, 3.0, 2.0]) == (3.0, 2.0) + assert monitoring.reduce_points([]) == (None, None) + + +def test_point_value_reads_int64_when_double_zero(): + # Test fallback to int64_value when double_value is 0.0. + point = SimpleNamespace( + value=SimpleNamespace(double_value=0.0, int64_value=1073741824) + ) + assert monitoring._point_value(point) == 1073741824.0 + + +def test_build_request_shape(): + # Verify request shape matches MetricServiceClient expectations. + cpu = monitoring.SERIES[0] + req = monitoring._build_request("proj", cpu, "run", 100, 700, 60) + assert req["name"] == "projects/proj" + assert ( + 'metric.type = "kubernetes.io/container/cpu/core_usage_time"' in req["filter"] + ) + assert 'resource.type = "k8s_container"' in req["filter"] + assert 'starts_with("run-workload-0-")' in req["filter"] + assert req["interval"]["start_time"]["seconds"] == 100 + assert req["interval"]["end_time"]["seconds"] == 700 + assert req["aggregation"]["alignment_period"]["seconds"] == 60 + assert req["aggregation"]["per_series_aligner"] == "ALIGN_RATE" + + +def test_build_request_bucket_filter_and_method(): + req_count = [s for s in monitoring.GCS_BUCKET_SERIES if s.method][0] + req = monitoring._build_request("proj", req_count, "my-bucket", 0, 600, 60) + assert 'resource.type = "gcs_bucket"' in req["filter"] + assert 'resource.labels.bucket_name = "my-bucket"' in req["filter"] + assert f'metric.labels.method = "{req_count.method}"' in req["filter"] + assert req["aggregation"]["per_series_aligner"] == "ALIGN_DELTA" + + +def _bucket_series(values): + # A gcs_bucket series carries no pod_name label; only points matter here. + return SimpleNamespace( + resource=SimpleNamespace(labels={}), points=[_point(v) for v in values] + ) + + +def test_collect_bucket_totals_sums_all_points(): + client = _FakeClient( + { + "storage.googleapis.com/network/sent_bytes_count": [ + _bucket_series([100.0, 200.0]), + _bucket_series([50.0]), + ], + "storage.googleapis.com/api/request_count": [_bucket_series([3.0, 4.0])], + } + ) + rows = monitoring.collect_bucket_totals( + client, + project="proj", + bucket="ckpt", + prefix="checkpoint", + start_epoch=0, + end_epoch=600, + ) + by_metric = {r.metric: r for r in rows} + assert by_metric["checkpoint_read_bytes"].peak == 350.0 + assert by_metric["checkpoint_read_bytes"].pod_name == "ckpt" + assert by_metric["checkpoint_read_request_count"].peak == 7.0 + + +def test_collect_bucket_totals_omits_empty_series(): + client = _FakeClient({}) # nothing returned + rows = monitoring.collect_bucket_totals( + client, + project="proj", + bucket="ds", + prefix="dataset", + start_epoch=0, + end_epoch=600, + ) + assert rows == [] + + +def test_to_epoch_handles_zulu(): + + assert monitoring._to_epoch("1970-01-01T00:01:00Z") == 60 + + +def test_collect_emits_one_row_per_pod_and_series(): + client = _FakeClient( + { + "kubernetes.io/container/cpu/core_usage_time": [ + _series("p0", [1.0, 5.0]), + _series("p1", [2.0, 2.0]), + ], + "kubernetes.io/container/memory/used_bytes": [_series("p0", [1024.0])], + "kubernetes.io/pod/network/received_bytes_count": [ + _series("p0", [10.0, 20.0]) + ], + "kubernetes.io/pod/network/sent_bytes_count": [_series("p0", [4.0, 6.0])], + } + ) + rows = monitoring.collect( + client, project="proj", run_id="run", start_epoch=0, end_epoch=600 + ) + by_key = {(r.pod_name, r.metric): r for r in rows} + assert by_key[("p0", "cpu")].peak == 5.0 + assert by_key[("p0", "cpu")].mean == 3.0 + assert by_key[("p1", "cpu")].peak == 2.0 + assert by_key[("p0", "memory")].peak == 1024.0 + assert by_key[("p0", "network_received")].peak == 20.0 + assert by_key[("p0", "network_sent")].peak == 6.0 + + assert 'starts_with("run-workload-0-")' in client.requests[0]["filter"] + + +def test_collect_writes_via_raw_store(tmp_path): + from metrics import raw_store + + client = _FakeClient( + {"kubernetes.io/container/cpu/core_usage_time": [_series("p0", [2.0, 4.0])]} + ) + rows = monitoring.collect( + client, project="proj", run_id="run", start_epoch=0, end_epoch=600 + ) + raw_store.write_system_metrics(rows, str(tmp_path)) + tables = raw_store.read_raw_metrics(str(tmp_path)) + cpu = [r for r in tables.system_rows if r["metric"] == "cpu"] + assert cpu and cpu[0]["peak"] == 4.0 and cpu[0]["mean"] == 3.0 + + +def test_assemble_rows_combines_pod_bucket_and_sizes(): + client = _FakeClient( + { + "kubernetes.io/container/cpu/core_usage_time": [ + _series("run-workload-0-a", [2.0, 4.0]) + ], + "storage.googleapis.com/network/sent_bytes_count": [ + _bucket_series([10.0, 20.0]) + ], + "storage.googleapis.com/api/request_count": [_bucket_series([5.0])], + } + ) + + class _FakeStorage: + def list_blobs(self, bucket_name, prefix=""): + data = {"ds": [("f", 1000)], "ckpt": [("checkpoints/s/x", 400)]} + for name, size in data.get(bucket_name, []): + if name.startswith(prefix): + yield SimpleNamespace(size=size) + + restore_rows = [{"checkpoint_location": "gs://ckpt/checkpoints/s", "end_time": 1.0}] + rows = monitoring.assemble_rows( + client, + _FakeStorage(), + project="p", + run_id="run", + checkpoint_bucket="ckpt", + dataset_bucket="ds", + restore_rows=restore_rows, + start_epoch=0, + end_epoch=600, + ) + metrics = {r.metric for r in rows} + assert "cpu" in metrics + assert "checkpoint_read_bytes" in metrics + assert "dataset_read_bytes" in metrics + assert "dataset_size_bytes" in metrics + assert "checkpoint_restored_bytes" in metrics + + +def test_assemble_rows_without_storage_client_skips_sizes(): + client = _FakeClient( + { + "kubernetes.io/container/cpu/core_usage_time": [ + _series("run-workload-0-a", [1.0]) + ] + } + ) + rows = monitoring.assemble_rows( + client, + None, + project="p", + run_id="run", + checkpoint_bucket=None, + dataset_bucket=None, + restore_rows=[], + start_epoch=0, + end_epoch=600, + ) + assert {r.metric for r in rows} == {"cpu"} + + +class _RaisingClient: + """Raises for one metric.type, returns a canned series for another.""" + + def __init__(self, raise_on, return_for, series): + self.raise_on = raise_on + self.return_for = return_for + self.series = series + + def list_time_series(self, request): + if f'metric.type = "{self.raise_on}"' in request["filter"]: + raise RuntimeError("boom") + if f'metric.type = "{self.return_for}"' in request["filter"]: + return iter(self.series) + return iter(()) + + +def test_collect_isolates_a_failing_series(): + # A failure querying one metric type drops only that series, not the rest. + client = _RaisingClient( + raise_on="kubernetes.io/container/cpu/core_usage_time", + return_for="kubernetes.io/container/memory/used_bytes", + series=[_series("p0", [1024.0])], + ) + rows = monitoring.collect( + client, project="proj", run_id="run", start_epoch=0, end_epoch=600 + ) + metrics = {r.metric for r in rows} + assert "memory" in metrics + assert "cpu" not in metrics + + +def test_collect_bucket_totals_isolates_a_failing_series(): + client = _RaisingClient( + raise_on="storage.googleapis.com/api/request_count", + return_for="storage.googleapis.com/network/sent_bytes_count", + series=[_bucket_series([10.0, 20.0])], + ) + rows = monitoring.collect_bucket_totals( + client, + project="proj", + bucket="ckpt", + prefix="checkpoint", + start_epoch=0, + end_epoch=600, + ) + by_metric = {r.metric: r for r in rows} + assert by_metric["checkpoint_read_bytes"].peak == 30.0 + assert "checkpoint_read_request_count" not in by_metric + + +def test_assemble_rows_isolates_du_failure(): + # A du (list_blobs) failure must not discard the pod metrics already collected. + client = _FakeClient( + { + "kubernetes.io/container/cpu/core_usage_time": [ + _series("run-workload-0-a", [1.0, 2.0]) + ] + } + ) + + class _BoomStorage: + def list_blobs(self, *args, **kwargs): + raise RuntimeError("denied") + + rows = monitoring.assemble_rows( + client, + _BoomStorage(), + project="p", + run_id="run", + checkpoint_bucket=None, + dataset_bucket="ds", + restore_rows=[], + start_epoch=0, + end_epoch=600, + ) + assert {r.metric for r in rows} == {"cpu"} diff --git a/cloudbuild/macrobenchmarks/metrics/tests/test_raw_store_system.py b/cloudbuild/macrobenchmarks/metrics/tests/test_raw_store_system.py new file mode 100644 index 00000000..f6020d69 --- /dev/null +++ b/cloudbuild/macrobenchmarks/metrics/tests/test_raw_store_system.py @@ -0,0 +1,19 @@ +from metrics import raw_store, schema + + +def test_system_metrics_roundtrip(tmp_path): + rows = [ + schema.SystemMetric(pod_name="p0", metric="cpu", peak=3.0, mean=1.5), + schema.SystemMetric(pod_name="p1", metric="memory", peak=1024.0, mean=None), + ] + raw_store.write_system_metrics(rows, str(tmp_path)) + tables = raw_store.read_raw_metrics(str(tmp_path)) + assert tables.system_rows == [ + {"pod_name": "p0", "metric": "cpu", "peak": 3.0, "mean": 1.5}, + {"pod_name": "p1", "metric": "memory", "peak": 1024.0, "mean": None}, + ] + + +def test_absent_system_metrics_read_as_empty(tmp_path): + tables = raw_store.read_raw_metrics(str(tmp_path)) + assert tables.system_rows == [] diff --git a/cloudbuild/macrobenchmarks/metrics/tests/test_sizes.py b/cloudbuild/macrobenchmarks/metrics/tests/test_sizes.py new file mode 100644 index 00000000..596a6105 --- /dev/null +++ b/cloudbuild/macrobenchmarks/metrics/tests/test_sizes.py @@ -0,0 +1,117 @@ +from types import SimpleNamespace + +from metrics import sizes + + +class _FakeStorage: + """list_blobs(bucket, prefix) over a canned {bucket: [(name, size), ...]}.""" + + def __init__(self, by_bucket): + self.by_bucket = by_bucket + + def list_blobs(self, bucket_name, prefix=""): + for name, size in self.by_bucket.get(bucket_name, []): + if name.startswith(prefix): + yield SimpleNamespace(name=name, size=size) + + +def test_gcs_du_sums_under_prefix(): + client = _FakeStorage( + {"b": [("checkpoints/a", 100), ("checkpoints/c", 50), ("other", 7)]} + ) + assert sizes.gcs_du(client, "gs://b/checkpoints") == 150 + assert sizes.gcs_du(client, "gs://b") == 157 + + +def test_gcs_du_none_when_no_match(): + client = _FakeStorage({"b": []}) + assert sizes.gcs_du(client, "gs://b/missing") is None + + +def test_restored_checkpoint_location_picks_earliest_end(): + rows = [ + {"checkpoint_location": "gs://c/step200", "end_time": 200.0}, + {"checkpoint_location": "gs://c/step100", "end_time": 100.0}, + {"checkpoint_location": None, "end_time": 5.0}, + ] + assert sizes.restored_checkpoint_location(rows) == "gs://c/step100" + assert sizes.restored_checkpoint_location([]) is None + + +def test_dataset_du_and_sample_count_uses_largest_shard(): + # du and count come from a single listing. The largest object is the + # representative shard (its bytes/row best matches the du-weighted average); + # a tiny sidecar must not be chosen. + client = _FakeStorage( + { + "ds": [ + ("train/_meta.json", 10), + ("train/a.parquet", 990), + ("train/big.parquet", 2000), # largest + ] + } + ) + + def count_shard(storage_client, blob): + assert blob.name == "train/big.parquet" # largest, not the sidecar + return 200 # 2000 bytes / 200 rows -> 10 bytes/sample + + du, count = sizes.dataset_du_and_sample_count( + client, "gs://ds", count_shard=count_shard + ) + assert du == 3000 # 10 + 990 + 2000 + assert count == 300 # du 3000 / 10 bytes-per-sample + + +def test_dataset_du_and_sample_count_none_when_no_objects(): + client = _FakeStorage({"ds": []}) + assert sizes.dataset_du_and_sample_count( + client, "gs://ds", count_shard=lambda c, b: 5 + ) == (None, None) + + +def test_dataset_du_and_sample_count_du_survives_count_failure(): + # Counting a shard (download + parse) is best-effort: on failure du is still + # returned so dataset_size_bytes is preserved; only the count is dropped. + client = _FakeStorage({"ds": [("a.parquet", 1000)]}) + + def boom(storage_client, blob): + raise RuntimeError("unparseable shard") + + du, count = sizes.dataset_du_and_sample_count(client, "gs://ds", count_shard=boom) + assert du == 1000 + assert count is None + + +def test_dataset_du_and_sample_count_no_count_when_shard_uncountable(): + client = _FakeStorage({"ds": [("a.parquet", 100)]}) + du, count = sizes.dataset_du_and_sample_count( + client, "gs://ds", count_shard=lambda c, b: 0 + ) + assert du == 100 + assert count is None + + +def test_size_rows_emits_dataset_size_sample_count_and_checkpoint(): + client = _FakeStorage( + { + "ds": [("train/0.parquet", 1000)], + "ckpt": [("checkpoints/step100/shard0", 400)], + } + ) + rows = sizes.size_rows( + client, + dataset_bucket="ds", + restored_location="gs://ckpt/checkpoints/step100", + count_shard=lambda c, b: 100, # 1000 bytes / 100 rows -> 10 bytes/sample + ) + by_metric = {r.metric: r for r in rows} + assert by_metric["dataset_size_bytes"].peak == 1000 + assert by_metric["dataset_sample_count"].peak == 100 # du 1000 / 10 per sample + assert by_metric["checkpoint_restored_bytes"].peak == 400 + + +def test_size_rows_skips_missing_inputs(): + client = _FakeStorage({"ds": []}) # dataset du -> None + rows = sizes.size_rows(client, dataset_bucket="ds", restored_location=None) + assert rows == [] diff --git a/cloudbuild/macrobenchmarks/scripts/create_buckets.sh b/cloudbuild/macrobenchmarks/scripts/create_buckets.sh index 7ff0550d..f254bbee 100755 --- a/cloudbuild/macrobenchmarks/scripts/create_buckets.sh +++ b/cloudbuild/macrobenchmarks/scripts/create_buckets.sh @@ -7,13 +7,31 @@ source "$(dirname "$0")/lib.sh" trap 'record_failure create-buckets' ERR skip_if_failed source "${BUILD_VARS_FILE}" -if [[ "${_BUCKET_TYPE}" == "regional" ]]; then - gcloud storage buckets create gs://$CHECKPOINT_BUCKET --project=${PROJECT_ID} --location=$REGION -elif [[ "${_BUCKET_TYPE}" == "zonal" ]]; then - gcloud storage buckets create gs://$CHECKPOINT_BUCKET --project=${PROJECT_ID} --location=$REGION --placement=${_ZONE} --default-storage-class=RAPID --enable-hierarchical-namespace --uniform-bucket-level-access -elif [[ "${_BUCKET_TYPE}" == "hns" ]]; then - gcloud storage buckets create gs://$CHECKPOINT_BUCKET --project=${PROJECT_ID} --location=$REGION --enable-hierarchical-namespace --uniform-bucket-level-access +create_typed_bucket "$CHECKPOINT_BUCKET" + +# Per-run dataset bucket (same config as CHECKPOINT_BUCKET), populated by an +# in-region copy, so its egress is attributable to one run for the dataset +# read-amplification metric. +create_typed_bucket "$DATASET_BUCKET" +SRC_OBJECT_PATH=$(echo "${_DATASET_PATH}" | sed -E 's#^gs://[^/]+/?##') +if [ "${_BUCKET_TYPE}" = "zonal" ]; then + # RAPID (zonal) objects lack the server-side rewrite rsync uses, so daisy-chain + # (download+reupload) + ulimit -n 65536 + DEST_PARENT="${SRC_OBJECT_PATH%/*}" + [ "$DEST_PARENT" = "$SRC_OBJECT_PATH" ] && DEST_PARENT="" + # `cp --recursive` on a directory/bucket source without a trailing wildcard + # copies the source's own name into the destination (e.g. _DATASET_PATH + # "gs://bucket" would land at "gs://DATASET_BUCKET/bucket/..."); "/*" makes + # it copy the source's contents instead. + CLOUDSDK_STORAGE_PROCESS_COUNT=16 CLOUDSDK_STORAGE_THREAD_COUNT=16 \ + CLOUDSDK_STORAGE_ATTEMPT_GRPC_DIRECT_PATH=False \ + gcloud storage cp --recursive --daisy-chain "${_DATASET_PATH%/}/*" "gs://${DATASET_BUCKET}${DEST_PARENT:+/$DEST_PARENT}" +else + # Regional/HNS support server-side copy; rsync mirrors the source into the dest. + gcloud storage rsync --recursive "${_DATASET_PATH}" "gs://${DATASET_BUCKET}/${SRC_OBJECT_PATH}" fi +echo "export RUN_DATASET_PATH=gs://${DATASET_BUCKET}/${SRC_OBJECT_PATH}" >> "${BUILD_VARS_FILE}" if gcloud storage buckets describe gs://$RESULTS_BUCKET --project=${PROJECT_ID} >/dev/null 2>&1; then # Reuse only if co-located with this build's LOCATION. The ingestion pipeline # builds the BigQuery dataset (and external table) in LOCATION; a results diff --git a/cloudbuild/macrobenchmarks/scripts/delete_buckets.sh b/cloudbuild/macrobenchmarks/scripts/delete_buckets.sh index e328b2c6..8a3938be 100755 --- a/cloudbuild/macrobenchmarks/scripts/delete_buckets.sh +++ b/cloudbuild/macrobenchmarks/scripts/delete_buckets.sh @@ -1,6 +1,6 @@ #!/usr/bin/env bash -# delete-buckets: delete the per-run checkpoint bucket (best-effort). The shared -# results bucket is intentionally left in place. +# delete-buckets: delete the per-run checkpoint and dataset buckets (best-effort). +# The shared results bucket is intentionally left in place. if [[ "${_SKIP_CLEANUP}" == "true" ]]; then echo "Skipping delete-buckets as requested." exit 0 @@ -8,3 +8,4 @@ fi source "$(dirname "$0")/lib.sh" source "${BUILD_VARS_FILE}" gcloud storage rm --recursive --project="${PROJECT_ID}" gs://$CHECKPOINT_BUCKET || true +gcloud storage rm --recursive --project="${PROJECT_ID}" gs://$DATASET_BUCKET || true diff --git a/cloudbuild/macrobenchmarks/scripts/init_variables.sh b/cloudbuild/macrobenchmarks/scripts/init_variables.sh index 87e0028c..509b9003 100755 --- a/cloudbuild/macrobenchmarks/scripts/init_variables.sh +++ b/cloudbuild/macrobenchmarks/scripts/init_variables.sh @@ -105,5 +105,6 @@ echo "export CLUSTER_NAME=${_INFRA_PREFIX}-gke-${SHORT_BUILD_ID}" >> "${BUILD_VA echo "export NETWORK_NAME=${_INFRA_PREFIX}-net-${SHORT_BUILD_ID}" >> "${BUILD_VARS_FILE}" echo "export SUBNET_NAME=${_INFRA_PREFIX}-subnet-${SHORT_BUILD_ID}" >> "${BUILD_VARS_FILE}" echo "export CHECKPOINT_BUCKET=${_INFRA_PREFIX}-macrobench-checkpoint-${SHORT_BUILD_ID}" >> "${BUILD_VARS_FILE}" +echo "export DATASET_BUCKET=${_INFRA_PREFIX}-macrobench-dataset-${SHORT_BUILD_ID}" >> "${BUILD_VARS_FILE}" echo "export RESULTS_BUCKET=${_INFRA_PREFIX}-macrobench-results" >> "${BUILD_VARS_FILE}" echo "export REGION=${REGION}" >> "${BUILD_VARS_FILE}" diff --git a/cloudbuild/macrobenchmarks/scripts/lib.sh b/cloudbuild/macrobenchmarks/scripts/lib.sh index d2138b08..880d661a 100755 --- a/cloudbuild/macrobenchmarks/scripts/lib.sh +++ b/cloudbuild/macrobenchmarks/scripts/lib.sh @@ -30,9 +30,25 @@ skip_if_failed() { fi } +# Create a per-run bucket per _BUCKET_TYPE (regional | zonal-RAPID | hns). +create_typed_bucket() { + local bucket="$1" + case "${_BUCKET_TYPE}" in + regional) + gcloud storage buckets create "gs://$bucket" --project="${PROJECT_ID}" --location="$REGION" ;; + zonal) + gcloud storage buckets create "gs://$bucket" --project="${PROJECT_ID}" --location="$REGION" --placement="${_ZONE}" --default-storage-class=RAPID --enable-hierarchical-namespace --uniform-bucket-level-access ;; + hns) + gcloud storage buckets create "gs://$bucket" --project="${PROJECT_ID}" --location="$REGION" --enable-hierarchical-namespace --uniform-bucket-level-access ;; + *) + echo "ERROR: unknown _BUCKET_TYPE='${_BUCKET_TYPE}' (expected regional|zonal|hns)" >&2 + return 1 ;; + esac +} + shared_workload_helm_args() { SHARED_HELM_ARGS=( - --set gcsfs.datasetPath="${_DATASET_PATH}" + --set gcsfs.datasetPath="${RUN_DATASET_PATH:-${_DATASET_PATH}}" --set workload.modelId="${_MODEL_ID}" --set-string workload.image="${_IMAGE}" --set workload.hfToken="${_HF_TOKEN}" diff --git a/cloudbuild/macrobenchmarks/scripts/scrape_metrics.sh b/cloudbuild/macrobenchmarks/scripts/scrape_metrics.sh index bbeb0b6b..7cefde68 100755 --- a/cloudbuild/macrobenchmarks/scripts/scrape_metrics.sh +++ b/cloudbuild/macrobenchmarks/scripts/scrape_metrics.sh @@ -29,15 +29,36 @@ if [ -n "${_CHECKPOINT_LOAD_PATH}" ] || [ "${_SEED_CHECKPOINT}" = "true" ]; then MIN_RESTORE_DATAPOINTS=1 RESUME_ARGS=(--resume-run) fi +# Run the calculator over the current raw metrics into the summary file $1. +run_calculate() { + python3 -m metrics.calculate \ + --run-id "$RUN_ID" --workload-name "${_WORKLOAD}" \ + --requirements "${_REQUIREMENTS}" --in-dir "$RAW_DIR" --out-file "$1" \ + --expected-steps "${_STEPS}" \ + --min-write-datapoints "$MIN_WRITE_DATAPOINTS" \ + --min-restore-datapoints "$MIN_RESTORE_DATAPOINTS" \ + "${RESUME_ARGS[@]}" \ + --require-data-loading-metrics \ + --bucket-type "${_BUCKET_TYPE}" --zone "${_ZONE}" --region "$REGION" \ + --machine-type "${_MACHINE_TYPE}" \ + --nodes "${_NODES}" --ranks-per-node "${_RANKS_PER_NODE}" \ + --steps "${_STEPS}" --checkpoint-interval "${_CHECKPOINT_INTERVAL}" \ + --checkpoints-to-keep "${_CKPT_TO_KEEP}" \ + --dataset-path "${_DATASET_PATH}" --model-id "${_MODEL_ID}" \ + --image "${_IMAGE}" \ + --training-strategy "${_TRAINING_STRATEGY}" \ + --simulated-step-compute-seconds "${_SIMULATED_STEP_COMPUTE_SECONDS}" \ + --per-device-batch "${_PER_DEVICE_BATCH}" --grad-accum "${_GRAD_ACCUM}" \ + --dataloader-workers "${_DATALOADER_WORKERS}" +} # Cloud Logging ingestion lags pod termination by seconds-to-minutes, and the # last logs emitted (the final checkpoint write and the profiler summary that # carries the data-loading metric) are the most likely to still be in flight # when the JobSet reports Completed. Settle once, then re-scrape at a fixed 60s # interval until the required metrics validate (or attempts are exhausted). -# calculate -# exits non-zero when metrics are incomplete; running it as an `if` condition -# keeps `set -e`/the ERR trap from aborting the step on a not-yet-complete -# attempt. +# run_calculate exits non-zero when metrics are incomplete; running it as an +# `if` condition keeps `set -e`/the ERR trap from aborting the step on a +# not-yet-complete attempt. sleep 60 SCRAPE_OK=false for attempt in $(seq 1 5); do @@ -55,25 +76,7 @@ for attempt in $(seq 1 5); do sleep 60 continue fi - if python3 -m metrics.calculate \ - --run-id "$RUN_ID" --workload-name "${_WORKLOAD}" \ - --requirements "${_REQUIREMENTS}" --in-dir "$RAW_DIR" --out-file "$SUMMARY" \ - --expected-steps "${_STEPS}" \ - --min-write-datapoints "$MIN_WRITE_DATAPOINTS" \ - --min-restore-datapoints "$MIN_RESTORE_DATAPOINTS" \ - "${RESUME_ARGS[@]}" \ - --require-data-loading-metrics \ - --bucket-type "${_BUCKET_TYPE}" --zone "${_ZONE}" --region "$REGION" \ - --machine-type "${_MACHINE_TYPE}" \ - --nodes "${_NODES}" --ranks-per-node "${_RANKS_PER_NODE}" \ - --steps "${_STEPS}" --checkpoint-interval "${_CHECKPOINT_INTERVAL}" \ - --checkpoints-to-keep "${_CKPT_TO_KEEP}" \ - --dataset-path "${_DATASET_PATH}" --model-id "${_MODEL_ID}" \ - --image "${_IMAGE}" \ - --training-strategy "${_TRAINING_STRATEGY}" \ - --simulated-step-compute-seconds "${_SIMULATED_STEP_COMPUTE_SECONDS}" \ - --per-device-batch "${_PER_DEVICE_BATCH}" --grad-accum "${_GRAD_ACCUM}" \ - --dataloader-workers "${_DATALOADER_WORKERS}"; then + if run_calculate "$SUMMARY"; then SCRAPE_OK=true break fi @@ -88,4 +91,16 @@ if [ "$SCRAPE_OK" != "true" ]; then record_failure scrape-metrics exit 1 fi +# Best-effort system metrics: fetched once after the required metrics validate +# (not per attempt, to avoid re-du'ing the dataset bucket). Settle for GCS +# metric lag, then fold into the summary; a failure here must not lose the +# metrics-complete summary already written above, so it's `|| true`/warn-only. +sleep "${SYSTEM_METRICS_SETTLE_SECONDS:-180}" +python3 -m metrics.monitoring \ + --project "${PROJECT_ID}" --run-id "$RUN_ID" \ + --start-time "$START_TIME" --end-time "$END_TIME" \ + --checkpoint-bucket "$CHECKPOINT_BUCKET" --dataset-bucket "$DATASET_BUCKET" \ + --out-dir "$RAW_DIR" || true +run_calculate "$SUMMARY" \ + || echo "Warning: recompute with system metrics failed; uploading summary without them." gcloud storage cp "$SUMMARY" "gs://$RESULTS_BUCKET/branch=$BRANCH_NAME/$DATE_DIR/$RUN_ID/${TS_DIR}.csv"