diff --git a/README.md b/README.md index 674dc7d..2a9f7ef 100644 --- a/README.md +++ b/README.md @@ -56,6 +56,33 @@ sparkparse analyze ./logs sparkparse analyze ./logs --compact --top-n 25 --redact ``` +### Databricks runs and experiments + +Collect an existing job/run (it never launches or waits for a workload), save a +portable snapshot, and compare explicit trials offline. See +[docs/databricks-runs.md](docs/databricks-runs.md) for the tested collection +contract and measure semantics. + +```bash +# report only (JSON to stdout, progress to stderr) +sparkparse databricks analyze --run-id 123 --profile free + +# save a snapshot, or record a trial in an explicit experiment +sparkparse databricks analyze --run-id 123 --profile free --out runs/123 +sparkparse databricks analyze --run-id 124 --profile free \ + --experiment ./experiments/join-tuning --variant broadcast --revision abc123 + +# compare and view saved trials entirely offline +sparkparse experiments compare ./experiments/join-tuning --baseline 123 --candidate 124 +sparkparse experiments viz ./experiments/join-tuning +``` + +Workflow elapsed time is kept separate from summed task durations, physical plan +changes never exclude a comparison, and a Query History subtotal over only the +queries that reported a metric is labelled an observed subtotal. Failed, active +and warmup trials stay visible but are excluded from default aggregates. + + ### context manager ```python diff --git a/docs/databricks-runs.md b/docs/databricks-runs.md new file mode 100644 index 0000000..da5c154 --- /dev/null +++ b/docs/databricks-runs.md @@ -0,0 +1,121 @@ +# Databricks run collection + +`sparkparse databricks analyze` collects an existing Databricks job/run into a +portable, versioned report. It never launches, waits for, repairs or cancels a +workload. `sparkparse experiments` compares saved reports offline. + +## Tested collection contract + +| Item | Value | +|---|---| +| Databricks CLI | v1.15.0 | +| Profile used | `free` (serverless; Databricks Free Edition) | +| Validated | 2026-09-13 against retained serverless runs of the bundled bundle jobs | +| Auth | Existing CLI profile; no SDK transport | + +Required access: Jobs API read access for the run, and the `query-history` scope +to read Query History with metrics. Git snapshot fields appear only when the job +was configured with a Git source. + +### Jobs + +- `databricks api get /api/2.1/jobs/runs/get?run_id=&include_resolved_values=true` + returns identity, `state`/`status`, `run_duration`/`setup_duration`/ + `execution_duration`/`cleanup_duration`, `tasks[]`, `environments[]`, and + `git_source` when present. Array properties are paged with `next_page_token`. +- `databricks api get /api/2.1/jobs/runs/list?job_id=&limit=` returns + `runs[]`, `has_more` and `next_page_token`. The generated `jobs list-runs` + command exposes `--offset` but not the page token, so the raw API is used. +- Task output: `databricks jobs get-run-output ` returns + `metadata`, `notebook_output{result,truncated}` and optionally `error`/`logs`. +- The generated `jobs get-run` command omits `environments`; the raw API does not. + +### Query History + +Query History filters use dot notation, for example: + +``` +/api/2.0/sql/history/queries?include_metrics=true&max_results=100 + &filter_by.query_start_time_range.start_time_ms= + &filter_by.query_start_time_range.end_time_ms= +``` + +The generated `query-history list` command in v1.15.0 has no filter flags, so +the raw API path is used. The response is +`{res: [QueryInfo...], has_next_page, next_page_token}`. Queries are matched to +collected task runs through `query_source.job_info.job_task_run_id` (a parent +`job_run_id` may be absent). Attribution is never made from time overlap alone. + +Metrics observed per query: `read_bytes`, `read_remote_bytes`, +`read_cache_bytes`, `write_remote_bytes`, `spill_to_disk_bytes`, +`task_total_time_ms`, `execution_time_ms`, `total_time_ms`, `rows_read_count`, +`rows_produced_count`, `network_sent_bytes`, `result_from_cache`, +`read_files_count`, `read_partitions_count`. + +### Not available from this collection path + +- `git_source.git_snapshot.used_commit` on workloads without a configured Git + source (the bundled bundle jobs did not report one; fixtures cover it). +- Classic cluster details for serverless runs, and cluster snapshots taken at + run time when only an `existing_cluster_id` is reported (lookups are labelled + `current_lookup`). +- Memory spill and per-executor distributions; only disk spill is reported. +- Query profiles and full query text (query text is never persisted). + +## Report measure semantics + +| Measure | Definition | +|---|---| +| Workflow elapsed | `run_duration`; never the sum of task durations (tasks may overlap) | +| Task execution | Per-task `execution_duration`, plus `setup`/`cleanup` | +| Query execution | Query History `execution_time_ms`; not workflow elapsed | +| Read bytes | Query History `read_bytes`; cache reads kept separately | +| Written bytes | Remote persistent bytes (`write_remote_bytes`), labelled as such | +| Disk spill | Query History `spill_to_disk_bytes` | +| Aggregate task time | Sum of query `task_total_time_ms`, not elapsed | + +Every aggregate records its source/version, unit, aggregation basis, observed +query count, count with the metric, and a completeness status. A subtotal over +only the queries that reported a metric is `observed_subtotal`, never a complete +total. Missing metrics are `null`; a measured zero stays `0`. + +## Usage + +```bash +# report only (JSON to stdout, progress to stderr) +sparkparse databricks analyze --run-id 123 --profile free + +# save a portable snapshot +sparkparse databricks analyze --run-id 123 --profile free --out runs/123 + +# resolve the latest terminal run for a job +sparkparse databricks analyze --job-id 456 --profile free --out runs/latest + +# record trials in an explicit experiment +sparkparse databricks analyze --run-id 123 --profile free \ + --experiment ./experiments/join-tuning --variant baseline +sparkparse databricks analyze --run-id 124 --profile free \ + --experiment ./experiments/join-tuning --variant broadcast --revision abc123 + +# compare and view saved trials entirely offline +sparkparse experiments compare ./experiments/join-tuning --baseline 123 --candidate 124 +sparkparse experiments viz ./experiments/join-tuning +``` + +An experiment directory holds an editable `experiment.json` plus immutable +snapshots: + +```text +experiments/join-tuning/ + experiment.json + runs/ + 123/.json + 124/.json +``` + +Recollecting a run appends a snapshot and updates its trial reference without +adding a trial. Storage is local only; cloud URIs are rejected. Manifest updates +use atomic local file replacement (one writer per experiment). + +Exit codes: `0` for a valid report (including partial/active), `1` for a +collection or experiment failure, `2` for invalid arguments. diff --git a/sparkparse/app.py b/sparkparse/app.py index dc11e08..5688c78 100644 --- a/sparkparse/app.py +++ b/sparkparse/app.py @@ -11,9 +11,21 @@ from sparkparse import alerts, history from sparkparse.analyze import to_analysis_export, to_plan_summary from sparkparse.artifact import save_capture_artifact +from sparkparse.databricks import ( + DEFAULT_LIMITS, + DatabricksCLI, + DatabricksError, + collect_run, +) from sparkparse.eventlog import discover_sources -from sparkparse.models import OutputFormat, ParsedLogDataFrames +from sparkparse.models import ( + OutputFormat, + ParsedLogDataFrames, + TrialMetadata, + WorkflowCollectionLimits, +) from sparkparse.parse import get_all_parsed_metrics, get_parsed_metrics +from sparkparse.runreport import build_report, format_report_text from sparkparse.storage import ( get_path_name, get_path_stem, @@ -63,6 +75,12 @@ class HistoryFormat(StrEnum): json = "json" +class OutputsOption(StrEnum): + failed = "failed" + all = "all" + none = "none" + + def _version_callback(value: bool) -> None: if value: typer.echo(f"sparkparse {__version__}") @@ -530,5 +548,218 @@ def compare_cmd( typer.echo(f" excluded: {reason}") +databricks_app = typer.Typer(help="Collect Databricks job/run reports.") +experiments_app = typer.Typer(help="Compare saved performance experiment trials.") + + +@databricks_app.command("analyze") +def databricks_analyze( + run_id: Annotated[ + str | None, typer.Option("--run-id", help="Job run ID to collect.") + ] = None, + job_id: Annotated[ + str | None, + typer.Option("--job-id", help="Resolve the latest terminal run for this job."), + ] = None, + profile: Annotated[ + str | None, typer.Option("--profile", help="Databricks CLI profile to use.") + ] = None, + out: Annotated[ + str | None, + typer.Option("--out", help="Local directory to save a portable snapshot."), + ] = None, + experiment: Annotated[ + str | None, + typer.Option( + "--experiment", help="Local experiment directory to record a trial." + ), + ] = None, + variant: Annotated[ + str, typer.Option("--variant", help="Variant label for the recorded trial.") + ] = "default", + revision: Annotated[ + str | None, + typer.Option("--revision", help="Explicitly asserted code revision fallback."), + ] = None, + outputs: Annotated[ + OutputsOption, + typer.Option( + "--outputs", help="Task output retrieval: failed (default), all, none." + ), + ] = OutputsOption.failed, + format: Annotated[ + AnalysisFormat, + typer.Option(help="Output format: 'json' (default) or human-readable 'text'."), + ] = AnalysisFormat.json, + deadline_seconds: Annotated[ + float, typer.Option(help="Wall-clock collection deadline in seconds.") + ] = DEFAULT_LIMITS.deadline_seconds, + max_pages: Annotated[ + int, typer.Option(help="Maximum pages to follow per paged collection.") + ] = DEFAULT_LIMITS.max_pages, + max_requests: Annotated[ + int, typer.Option(help="Maximum CLI requests per collection.") + ] = DEFAULT_LIMITS.max_requests, + max_retries: Annotated[ + int, typer.Option(help="Maximum rate-limit retries per request.") + ] = DEFAULT_LIMITS.max_retries, +) -> None: + """Collect a Databricks run report without launching or waiting for a workload.""" + if (run_id is None) == (job_id is None): + raise typer.BadParameter("pass exactly one of --run-id or --job-id") + if out is not None and experiment is not None: + raise typer.BadParameter("pass at most one of --out or --experiment") + + limits = WorkflowCollectionLimits( + deadline_seconds=deadline_seconds, + max_pages=max_pages, + max_requests=max_requests, + max_retries=max_retries, + retry_backoff_seconds=DEFAULT_LIMITS.retry_backoff_seconds, + outputs=outputs.value, + ) + client = DatabricksCLI(profile=profile, limits=limits) + target = run_id if run_id is not None else f"job {job_id}" + typer.echo(f"Collecting {target}...", err=True) + try: + raw = collect_run( + client, + run_id=run_id, + job_id=job_id, + outputs=outputs.value, + ) + report = build_report(raw, asserted_revision=revision) + except DatabricksError as exc: + typer.echo(f"Collection failed: {exc}", err=True) + raise typer.Exit(1) from exc + + if experiment is not None: + from sparkparse import experiments as experiments_mod + + try: + _ref, path = experiments_mod.record_trial( + experiment, report, variant=variant, revision=revision + ) + except experiments_mod.ExperimentError as exc: + typer.echo(str(exc), err=True) + raise typer.Exit(1) from exc + typer.echo(f"Trial recorded at {path}", err=True) + elif out is not None: + from sparkparse import experiments as experiments_mod + + snapshot = experiments_mod.build_snapshot( + report, trial=TrialMetadata(variant=variant, revision=revision) + ) + path = experiments_mod.save_snapshot(out, snapshot) + typer.echo(f"Snapshot written to {path}", err=True) + + if format == AnalysisFormat.json: + sys.stdout.write(report.model_dump_json(indent=2) + "\n") + else: + sys.stdout.write(format_report_text(report) + "\n") + + +@experiments_app.command("compare") +def experiments_compare( + experiment_dir: Annotated[str, typer.Argument(help="Local experiment directory.")], + baseline: Annotated[ + str | None, typer.Option("--baseline", help="Baseline run ID.") + ] = None, + candidate: Annotated[ + str | None, typer.Option("--candidate", help="Candidate run ID.") + ] = None, + baseline_variant: Annotated[ + str | None, + typer.Option("--baseline-variant", help="Baseline variant label group."), + ] = None, + candidate_variant: Annotated[ + str | None, + typer.Option("--candidate-variant", help="Candidate variant label group."), + ] = None, + format: Annotated[ + AnalysisFormat, + typer.Option(help="Output format: 'json' (default) or human-readable 'text'."), + ] = AnalysisFormat.json, +) -> None: + """Compare explicit trials in a saved experiment, entirely offline.""" + from sparkparse import experiments as experiments_mod + + if (baseline is None) == (baseline_variant is None): + raise typer.BadParameter("pass exactly one of --baseline or --baseline-variant") + if (candidate is None) == (candidate_variant is None): + raise typer.BadParameter( + "pass exactly one of --candidate or --candidate-variant" + ) + + try: + comparison = experiments_mod.compare_experiment( + experiment_dir, + baseline_run_id=baseline, + candidate_run_id=candidate, + baseline_variant=baseline_variant, + candidate_variant=candidate_variant, + ) + except experiments_mod.ExperimentError as exc: + typer.echo(str(exc), err=True) + raise typer.Exit(1) from exc + + if format == AnalysisFormat.json: + typer.echo(comparison.model_dump_json(indent=2)) + return + typer.echo( + f"Experiment {comparison.experiment_id} " + f"(baseline n={comparison.baseline.sample_count}, " + f"candidate n={comparison.candidate.sample_count})" + ) + rows = pl.DataFrame( + [ + { + "metric": m.metric, + "baseline": m.baseline_median, + "candidate": m.candidate_median, + "delta": m.delta, + "pct_change": m.pct_change, + "caveat": m.caveat, + } + for m in comparison.metrics + ] + ) + typer.echo(str(rows) if comparison.metrics else "No comparable metrics.") + for note in comparison.notes: + typer.echo(f" note: {note}") + + +@experiments_app.command("viz") +def experiments_viz( + experiment_dir: Annotated[str, typer.Argument(help="Local experiment directory.")], + force_port: Annotated[ + bool, + typer.Option(help="Force kill any process using port 8050 before starting."), + ] = False, +) -> None: + """Launch the local experiment dashboard (read-only, loopback only).""" + try: + from sparkparse.dashboard import run_app + from sparkparse.experiments_dashboard import init_experiments_dashboard + except ImportError as exc: # pragma: no cover - exercised only without [viz] + raise typer.BadParameter( + "The dashboard requires the 'viz' extra: install with " + "`pip install sparkparse[viz]`." + ) from exc + + from sparkparse import experiments as experiments_mod + + try: + dash_app = init_experiments_dashboard(experiment_dir) + except experiments_mod.ExperimentError as exc: + typer.echo(str(exc), err=True) + raise typer.Exit(1) from exc + run_app(app=dash_app, force_port=force_port) + + +app.add_typer(databricks_app, name="databricks", no_args_is_help=True) +app.add_typer(experiments_app, name="experiments", no_args_is_help=True) + + if __name__ == "__main__": app() diff --git a/sparkparse/databricks.py b/sparkparse/databricks.py new file mode 100644 index 0000000..82cf31c --- /dev/null +++ b/sparkparse/databricks.py @@ -0,0 +1,560 @@ +"""Databricks collection transport. + +This module owns the single boundary between sparkparse and the Databricks CLI. +Every remote call goes through :class:`DatabricksCLI`, which uses argument +arrays and structured JSON only (never shell interpolation or table parsing). +Normalization into report models lives in :mod:`sparkparse.runreport` so it can +be tested without a CLI process. + +The generated CLI does not expose filter flags for Query History listing in +tested versions, so those calls go through ``databricks api get`` with URL +encoded query parameters. +""" + +from __future__ import annotations + +import datetime +import json +import logging +import subprocess +import time +import urllib.parse +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field +from typing import Any + +from sparkparse.models import ( + CollectionStatus, + ReportDiagnostic, + WorkflowCollectionLimits, +) + +logger = logging.getLogger(__name__) + +DEFAULT_LIMITS = WorkflowCollectionLimits( + deadline_seconds=120.0, + max_pages=20, + max_requests=60, + max_retries=3, + retry_backoff_seconds=1.0, + outputs="failed", +) + +_QUERY_HISTORY_PATH = "/api/2.0/sql/history/queries" +_JOBS_RUNS_GET_PATH = "/api/2.1/jobs/runs/get" +_JOBS_RUNS_LIST_PATH = "/api/2.1/jobs/runs/list" + +_RATE_LIMIT_MARKERS = ( + "429", + "rate limit", + "too many requests", + "resource_exhausted", + "retry after", +) +_AUTH_MARKERS = ( + "invalid access token", + "invalid refresh token", + "authentication", + "unauthorized", + "401", + "403", + "no valid credential", +) + +Runner = Callable[[Sequence[str], float], "subprocess.CompletedProcess[str]"] + + +class DatabricksError(Exception): + """A collection failure with a stable code for the report and CLI.""" + + def __init__(self, message: str, *, code: str = "collection_error") -> None: + super().__init__(message) + self.code = code + + +class CollectionDeadlineExceeded(DatabricksError): + def __init__(self, message: str = "collection deadline exceeded") -> None: + super().__init__(message, code="deadline_exceeded") + + +def _default_runner( + args: Sequence[str], timeout: float +) -> subprocess.CompletedProcess[str]: + return subprocess.run( + list(args), capture_output=True, text=True, timeout=timeout, check=False + ) + + +def _now_ms() -> int: + return int(datetime.datetime.now(datetime.UTC).timestamp() * 1000) + + +class _Budget: + """Tracks a single collection's deadline and request/page counts.""" + + def __init__(self, limits: WorkflowCollectionLimits) -> None: + self.limits = limits + self.requests = 0 + self.pages = 0 + self.retries = 0 + self.deadline = time.monotonic() + limits.deadline_seconds + self.exhausted = False + self.reason: str | None = None + + def remaining(self) -> float: + return self.deadline - time.monotonic() + + def take_request(self) -> float: + if self.remaining() <= 0: + self.exhausted = True + self.reason = "deadline" + raise CollectionDeadlineExceeded() + if self.requests >= self.limits.max_requests: + self.exhausted = True + self.reason = "max_requests" + raise DatabricksError( + f"request budget exhausted ({self.limits.max_requests})", + code="budget_exhausted", + ) + self.requests += 1 + return max(self.remaining(), 0.1) + + def take_page(self) -> None: + if self.pages >= self.limits.max_pages: + self.exhausted = True + self.reason = "max_pages" + raise DatabricksError( + f"page budget exhausted ({self.limits.max_pages})", + code="budget_exhausted", + ) + self.pages += 1 + + +class DatabricksCLI: + """Small, injectable client over the Databricks CLI.""" + + def __init__( + self, + profile: str | None = None, + *, + cli_path: str = "databricks", + runner: Runner | None = None, + limits: WorkflowCollectionLimits | None = None, + ) -> None: + self.profile = profile + self.cli_path = cli_path + self.runner = runner or _default_runner + self.limits = limits or DEFAULT_LIMITS + self.budget = _Budget(self.limits) + + def _base_args(self, args: Sequence[str]) -> list[str]: + command = [self.cli_path, *args] + if self.profile is not None: + command += ["--profile", self.profile] + command += ["--output", "json"] + return command + + def _run_json(self, args: Sequence[str]) -> Any: + while True: + timeout = self.budget.take_request() + command = self._base_args(args) + logger.debug("databricks call: %s", command) + try: + proc = self.runner(command, timeout) + except subprocess.TimeoutExpired as exc: + raise CollectionDeadlineExceeded() from exc + if proc.returncode == 0: + return self._parse_stdout(proc.stdout, args) + stderr = (proc.stderr or "").strip() + if ( + self._is_rate_limit(stderr) + and self.budget.retries < self.limits.max_retries + ): + self.budget.retries += 1 + delay = self.limits.retry_backoff_seconds * ( + 2 ** (self.budget.retries - 1) + ) + time.sleep(min(delay, max(self.budget.remaining(), 0.0))) + continue + raise self._error_from_stderr(stderr) + + @staticmethod + def _parse_stdout(stdout: str, args: Sequence[str]) -> Any: + text = stdout.strip() + if not text: + return {} + try: + return json.loads(text) + except json.JSONDecodeError as exc: + raise DatabricksError( + f"Databricks CLI returned non-JSON output for {' '.join(args)}", + code="invalid_json", + ) from exc + + @staticmethod + def _is_rate_limit(stderr: str) -> bool: + lowered = stderr.lower() + return any(marker in lowered for marker in _RATE_LIMIT_MARKERS) + + @classmethod + def _error_from_stderr(cls, stderr: str) -> DatabricksError: + lowered = stderr.lower() + if any(marker in lowered for marker in _AUTH_MARKERS): + return DatabricksError( + f"Databricks authentication failed: {stderr}", + code="auth_error", + ) + return DatabricksError(stderr or "Databricks CLI call failed") + + def jobs_get_run( + self, run_id: str, *, page_token: str | None = None + ) -> dict[str, Any]: + query = [ + f"run_id={urllib.parse.quote(str(run_id))}", + "include_resolved_values=true", + ] + if page_token: + query.append(f"page_token={urllib.parse.quote(page_token)}") + raw = self._run_json(["api", "get", f"{_JOBS_RUNS_GET_PATH}?{'&'.join(query)}"]) + assert isinstance(raw, dict) + return raw + + def jobs_list_runs( + self, + *, + job_id: str | None = None, + page_token: str | None = None, + limit: int = 25, + active_only: bool | None = None, + ) -> dict[str, Any]: + query = [f"limit={limit}"] + if job_id is not None: + query.append(f"job_id={urllib.parse.quote(str(job_id))}") + if active_only is not None: + query.append(f"active_only={'true' if active_only else 'false'}") + if page_token: + query.append(f"page_token={urllib.parse.quote(page_token)}") + raw = self._run_json( + ["api", "get", f"{_JOBS_RUNS_LIST_PATH}?{'&'.join(query)}"] + ) + assert isinstance(raw, dict) + return raw + + def run_output(self, run_id: str) -> dict[str, Any]: + raw = self._run_json(["jobs", "get-run-output", str(run_id)]) + assert isinstance(raw, dict) + return raw + + def query_history_page( + self, + *, + start_time_ms: int, + end_time_ms: int, + max_results: int = 100, + page_token: str | None = None, + ) -> dict[str, Any]: + query = [ + "include_metrics=true", + f"max_results={max_results}", + f"filter_by.query_start_time_range.start_time_ms={start_time_ms}", + f"filter_by.query_start_time_range.end_time_ms={end_time_ms}", + ] + if page_token: + query.append(f"page_token={urllib.parse.quote(page_token)}") + raw = self._run_json(["api", "get", f"{_QUERY_HISTORY_PATH}?{'&'.join(query)}"]) + assert isinstance(raw, dict) + return raw + + def cluster_get(self, cluster_id: str) -> dict[str, Any]: + raw = self._run_json( + ["api", "get", f"/api/2.1/clusters/get?cluster_id={cluster_id}"] + ) + assert isinstance(raw, dict) + return raw + + +@dataclass +class RawRunCollection: + """Raw API responses plus collection bookkeeping, before normalization.""" + + run: dict[str, Any] + tasks: list[dict[str, Any]] + queries: list[dict[str, Any]] + cluster_lookups: dict[str, dict[str, Any]] = field(default_factory=dict) + outputs: dict[str, dict[str, Any]] = field(default_factory=dict) + diagnostics: list[ReportDiagnostic] = field(default_factory=list) + collection_status: CollectionStatus = CollectionStatus.complete + run_pagination_complete: bool = True + query_discovery_complete: bool = True + query_denied: bool = False + discovered_query_count: int = 0 + request_count: int = 0 + page_count: int = 0 + collection_time_ms: int | None = None + limits: WorkflowCollectionLimits = field(default_factory=lambda: DEFAULT_LIMITS) + + +def _merge_array(target: dict[str, Any], page: dict[str, Any]) -> None: + """Merge a paginated get-run page into the first page's array properties.""" + for key, value in page.items(): + if isinstance(value, list): + existing = target.get(key) + if isinstance(existing, list): + existing.extend(value) + else: + target[key] = value + elif key not in target: + target[key] = value + + +def fetch_full_run(cli: DatabricksCLI, run_id: str) -> tuple[dict[str, Any], bool]: + """Fetch a run, following page tokens for paginated array properties. + + Returns the merged run plus whether every page was retrieved. Budget or + transport exhaustion stops pagination but keeps the pages already merged, + so a partial run is never discarded. + """ + run = cli.jobs_get_run(run_id) + token = run.get("next_page_token") + complete = True + while token: + try: + cli.budget.take_page() + page = cli.jobs_get_run(run_id, page_token=token) + except DatabricksError: + complete = False + break + _merge_array(run, page) + token = page.get("next_page_token") + return run, complete + + +def resolve_run_id(cli: DatabricksCLI, job_id: str) -> tuple[str, bool]: + """Return the latest run for a job and whether it is still active. + + Runs are returned newest first, so the first terminal run encountered is the + latest terminal run. Pagination continues past a page of active runs until a + terminal run is found or the page budget is exhausted. + """ + token: str | None = None + fallback: dict[str, Any] | None = None + while True: + try: + cli.budget.take_page() + response = cli.jobs_list_runs(job_id=job_id, page_token=token, limit=25) + except DatabricksError: + break + runs = response.get("runs") or [] + for run in runs: + if fallback is None: + fallback = run + state = (run.get("state") or {}).get("life_cycle_state") + if state not in (None, "RUNNING", "PENDING", "QUEUED", "BLOCKED"): + return str(run["run_id"]), False + token = response.get("next_page_token") + if not token: + break + if fallback is None: + raise DatabricksError(f"job {job_id} has no retained runs", code="no_runs") + return str(fallback["run_id"]), True + + +def _task_failure(task: dict[str, Any]) -> bool: + result_state = (task.get("state") or {}).get("result_state") + return result_state in ("FAILED", "TIMEDOUT", "CANCELED", "UPSTREAM_FAILED") + + +def collect_run( + cli: DatabricksCLI, + *, + run_id: str | None = None, + job_id: str | None = None, + outputs: str = "failed", + max_outputs: int = 5, +) -> RawRunCollection: + """Collect a run's metadata, tasks, compute, outputs and query metrics. + + Network failures raise :class:`DatabricksError`; bounded exhaustion is + reported as a partial collection rather than raising, so a useful report is + still produced. + """ + started = _now_ms() + active = False + if run_id is None: + assert job_id is not None + run_id, active = resolve_run_id(cli, job_id) + + diagnostics: list[ReportDiagnostic] = [] + status = CollectionStatus.complete + + run, run_complete = fetch_full_run(cli, run_id) + if not run_complete: + diagnostics.append( + ReportDiagnostic( + code="run_pagination_incomplete", + message="run metadata pagination stopped early; task list may be partial", + severity="warning", + ) + ) + parent_run_id = run.get("job_run_id") + if parent_run_id is not None and str(parent_run_id) != str(run.get("run_id")): + raise DatabricksError( + f"run {run.get('run_id')} is a task run; its parent job run is " + f"{parent_run_id}. Pass the parent job-run ID with --run-id.", + code="task_run_id", + ) + life_cycle = (run.get("state") or {}).get("life_cycle_state") + if life_cycle in ("RUNNING", "PENDING", "QUEUED", "BLOCKED"): + active = True + tasks = list(run.get("tasks") or []) + + # Compute snapshots embedded in the run; only missing classic clusters + # trigger a lookup, which is labelled as a current observation by the + # normalizer. Distinct lookups are cached. + cluster_lookups: dict[str, dict[str, Any]] = {} + for task in tasks: + cluster_id = task.get("existing_cluster_id") + if not cluster_id or cluster_id in cluster_lookups: + continue + try: + cluster_lookups[str(cluster_id)] = cli.cluster_get(str(cluster_id)) + except DatabricksError as exc: + diagnostics.append( + ReportDiagnostic( + code="cluster_lookup_failed", + message=str(exc), + scope=str(cluster_id), + ) + ) + + # Task output is enrichment, not the timing source. Retrieve bounded output + # for failed tasks by default, or every task with ``--outputs all``. + outputs_by_run: dict[str, dict[str, Any]] = {} + if outputs != "none": + wanted = ( + [t for t in tasks if _task_failure(t)] if outputs == "failed" else tasks + ) + for task in wanted[:max_outputs]: + task_run_id = task.get("run_id") + if task_run_id is None: + continue + try: + outputs_by_run[str(task_run_id)] = cli.run_output(str(task_run_id)) + except DatabricksError as exc: + diagnostics.append( + ReportDiagnostic( + code="output_unavailable", + message=str(exc), + scope=task.get("task_key"), + ) + ) + if len(wanted) > max_outputs: + diagnostics.append( + ReportDiagnostic( + code="outputs_truncated", + message=( + f"{len(wanted) - max_outputs} task output(s) not retrieved " + f"(bounded at {max_outputs})" + ), + ) + ) + + queries: list[dict[str, Any]] = [] + query_discovery_complete = True + query_denied = False + discovered_query_count = 0 + start_ms = run.get("start_time") + end_ms = run.get("end_time") or _now_ms() + if start_ms is not None: + ( + queries, + query_discovery_complete, + discovered_query_count, + query_error, + ) = _collect_queries(cli, start_time_ms=int(start_ms), end_time_ms=int(end_ms)) + if query_error is not None: + query_denied = query_error.code == "auth_error" + diagnostics.append( + ReportDiagnostic( + code="query_history_unavailable", + message=str(query_error), + severity="warning", + ) + ) + query_discovery_complete = False + if active: + status = CollectionStatus.partial + diagnostics.append( + ReportDiagnostic( + code="active_run", + message="run is still active; report is a labeled partial snapshot", + severity="info", + ) + ) + if not run_complete or not query_discovery_complete: + status = CollectionStatus.partial + if cli.budget.exhausted: + status = CollectionStatus.partial + diagnostics.append( + ReportDiagnostic( + code="collection_budget_exhausted", + message=f"collection stopped early ({cli.budget.reason})", + severity="warning", + ) + ) + + return RawRunCollection( + run=run, + tasks=tasks, + queries=queries, + cluster_lookups=cluster_lookups, + outputs=outputs_by_run, + diagnostics=diagnostics, + collection_status=status, + run_pagination_complete=run_complete, + query_discovery_complete=query_discovery_complete, + query_denied=query_denied, + discovered_query_count=discovered_query_count, + request_count=cli.budget.requests, + page_count=cli.budget.pages, + collection_time_ms=_now_ms() - started, + limits=cli.limits, + ) + + +def _collect_queries( + cli: DatabricksCLI, *, start_time_ms: int, end_time_ms: int +) -> tuple[list[dict[str, Any]], bool, int, DatabricksError | None]: + """Page through Query History, returning everything collected so far. + + A later-page failure is reported through the returned error while the pages + already fetched are preserved, so a useful partial snapshot survives budget + exhaustion or a transport error. + """ + queries: list[dict[str, Any]] = [] + token: str | None = None + complete = True + discovered = 0 + error: DatabricksError | None = None + while True: + try: + cli.budget.take_page() + page = cli.query_history_page( + start_time_ms=start_time_ms, + end_time_ms=end_time_ms, + page_token=token, + ) + except DatabricksError as exc: + complete = False + error = exc + break + results = page.get("res") or [] + discovered += len(results) + queries.extend(results) + token = page.get("next_page_token") + if not page.get("has_next_page"): + break + if not token: + complete = False + break + return queries, complete, discovered, error diff --git a/sparkparse/experiments.py b/sparkparse/experiments.py new file mode 100644 index 0000000..85acf74 --- /dev/null +++ b/sparkparse/experiments.py @@ -0,0 +1,1037 @@ +"""Explicit performance experiments: portable snapshots and comparisons. + +An experiment is a local directory holding an editable ``experiment.json`` +manifest plus immutable per-run snapshots. Collection is append-only; a +recollection of the same run and variant appends a snapshot and updates the +trial reference rather than counting as a new trial. Comparisons consume saved +snapshots only and never require the network. + +Storage is local by design. Cloud URIs are rejected clearly, and manifest +updates use atomic local file replacement (one writer per experiment). +""" + +from __future__ import annotations + +import datetime +import os +import statistics +import uuid +from pathlib import Path +from typing import Any + +from sparkparse.models import ( + ComparabilityNote, + ComparisonGroup, + ExperimentComparison, + ExperimentManifest, + MetricCompleteness, + MetricDelta, + MetricUnit, + OutputAttachment, + RunReport, + TaskDelta, + TrialMetadata, + TrialRef, + TrialSnapshot, +) +from sparkparse.storage import is_cloud_path, list_files + +MANIFEST_NAME = "experiment.json" +RUNS_DIR = "runs" +SNAPSHOT_SUFFIX = ".json" + +_RUN_METRICS = ( + "workflow_elapsed_ms", + "task_execution_ms", +) + + +class ExperimentError(Exception): + """A user-facing experiment store or comparison problem.""" + + +def _utcnow() -> datetime.datetime: + return datetime.datetime.now(datetime.UTC) + + +def new_snapshot_id(now: datetime.datetime | None = None) -> str: + stamp = (now or _utcnow()).strftime("%Y%m%dT%H%M%S") + return f"{stamp}-{uuid.uuid4().hex[:8]}" + + +def _require_local(path: str, label: str) -> str: + if is_cloud_path(path): + raise ExperimentError( + f"{label} must be a local path; cloud URIs are not supported for experiments" + ) + return path + + +def _atomic_write(path: Path, text: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(path.suffix + ".tmp") + tmp.write_text(text) + os.replace(tmp, path) + + +def snapshot_path(base_dir: str | Path, run_id: str, snapshot_id: str) -> Path: + return Path(base_dir) / RUNS_DIR / str(run_id) / f"{snapshot_id}{SNAPSHOT_SUFFIX}" + + +def save_snapshot(base_dir: str | Path, snapshot: TrialSnapshot) -> Path: + """Write one immutable snapshot, plus any bounded output attachments.""" + base = Path(_require_local(str(base_dir), "snapshot directory")) + # Attachments are stored relative to the snapshot directory so copying the + # experiment preserves offline access. + for attachment in snapshot.report.outputs: + if attachment.excerpt is None: + continue + safe_key = "".join( + c if c.isalnum() or c in "-_." else "_" for c in attachment.task_key + ) + rel = ( + f"{RUNS_DIR}/{snapshot.run_id}/" + f"{snapshot.snapshot_id}.outputs/{safe_key}.txt" + ) + _atomic_write(base / rel, attachment.excerpt) + attachment.path = rel + for task in snapshot.report.tasks: + attachment = _attachment_for_task(snapshot.report.outputs, task.run_id) + if attachment is not None: + task.output_path = attachment.path + path = snapshot_path(base, snapshot.run_id, snapshot.snapshot_id) + _atomic_write(path, snapshot.model_dump_json(indent=2) + "\n") + return path + + +def _attachment_for_task( + attachments: list[OutputAttachment], run_id: str | None +) -> OutputAttachment | None: + if run_id is None: + return None + for attachment in attachments: + if attachment.run_id == run_id: + return attachment + return None + + +def load_snapshot(path: str | Path) -> TrialSnapshot: + path_obj = Path(path) + try: + return TrialSnapshot.model_validate_json(path_obj.read_text()) + except (OSError, ValueError) as exc: + raise ExperimentError(f"invalid snapshot {path}: {exc}") from exc + + +def list_snapshot_paths(base_dir: str | Path) -> list[str]: + base = Path(base_dir) + if not base.exists(): + return [] + return sorted(list_files(base, f"{RUNS_DIR}/*/*{SNAPSHOT_SUFFIX}")) + + +def trial_snapshot_path(base_dir: str | Path, run_id: str, snapshot_id: str) -> Path: + return snapshot_path(base_dir, run_id, snapshot_id) + + +def manifest_path(exp_dir: str | Path) -> Path: + return Path(exp_dir) / MANIFEST_NAME + + +def load_manifest(exp_dir: str | Path) -> ExperimentManifest: + path = manifest_path(exp_dir) + if not path.exists(): + raise ExperimentError(f"no experiment manifest at {path}") + try: + manifest = ExperimentManifest.model_validate_json(path.read_text()) + except ValueError as exc: + raise ExperimentError(f"invalid experiment manifest {path}: {exc}") from exc + if ( + manifest.schema_version + != ExperimentManifest.model_fields["schema_version"].default + ): + raise ExperimentError( + f"unsupported experiment schema {manifest.schema_version!r}" + ) + return manifest + + +def save_manifest(exp_dir: str | Path, manifest: ExperimentManifest) -> Path: + manifest.updated_at = _utcnow() + path = manifest_path(exp_dir) + _atomic_write(path, manifest.model_dump_json(indent=2) + "\n") + return path + + +def build_snapshot( + report: RunReport, + *, + trial: TrialMetadata, + snapshot_id: str | None = None, +) -> TrialSnapshot: + return TrialSnapshot( + snapshot_id=snapshot_id or new_snapshot_id(), + collected_at=report.collected_at, + run_id=report.identity.run_id, + workspace_host=report.identity.workspace_host, + job_id=report.identity.job_id, + trial=trial, + report=report, + ) + + +def record_trial( + exp_dir: str | Path, + report: RunReport, + *, + variant: str, + revision: str | None = None, + artifact_digest: str | None = None, + config_fingerprint: str | None = None, + input_snapshot: str | None = None, + warmup: bool | None = None, + correctness: str | None = None, + hypothesis: str | None = None, + notes: str | None = None, + snapshot_id: str | None = None, +) -> tuple[TrialRef, Path]: + """Save a snapshot and add or update its trial in the manifest. + + A trial is identified by its run. Recollecting a run updates its existing + trial instead of adding another, so relabeling a run with a new variant + moves the trial rather than duplicating the execution. Metadata supplied on + a recollection overrides the stored value; omitted metadata is preserved. + """ + exp = Path(_require_local(str(exp_dir), "experiment directory")) + if manifest_path(exp).exists(): + manifest = load_manifest(exp) + _check_compatible(manifest, report) + else: + now = _utcnow() + manifest = ExperimentManifest( + experiment_id=uuid.uuid4().hex, + workspace_host=report.identity.workspace_host, + job_id=report.identity.job_id, + created_at=now, + updated_at=now, + ) + existing = _find_trial(manifest, report.identity.run_id) + + derived_config = _compute_fingerprint(report) + if config_fingerprint is not None: + config_value = config_fingerprint + config_derived = False + elif ( + existing is not None + and not existing.config_fingerprint_derived + and existing.config_fingerprint + ): + # An explicitly asserted fingerprint is authoritative and preserved. + config_value = existing.config_fingerprint + config_derived = False + else: + # Derived identity is recomputed so late enrichment cannot leave a + # stale fingerprint attached to an updated report. + config_value = derived_config + config_derived = True + trial = TrialMetadata( + variant=variant, + revision=revision if revision is not None else _existing(existing, "revision"), + artifact_digest=artifact_digest, + config_fingerprint=config_value, + config_fingerprint_derived=config_derived, + input_snapshot=( + input_snapshot + if input_snapshot is not None + else _existing(existing, "input_snapshot") + ), + warmup=warmup if warmup is not None else _existing(existing, "warmup", False), + correctness=( + correctness + if correctness is not None + else _existing(existing, "correctness") + ), + hypothesis=hypothesis, + notes=notes, + workspace_host=report.identity.workspace_host, + job_id=report.identity.job_id, + ) + snapshot = build_snapshot(report, trial=trial, snapshot_id=snapshot_id) + path = save_snapshot(exp, snapshot) + + if existing is not None: + existing.variant = variant + existing.snapshot_id = snapshot.snapshot_id + existing.revision = trial.revision + existing.config_fingerprint = trial.config_fingerprint + existing.config_fingerprint_derived = trial.config_fingerprint_derived + existing.input_snapshot = trial.input_snapshot + existing.warmup = trial.warmup + existing.correctness = trial.correctness + ref = existing + else: + ref = TrialRef( + trial_id=uuid.uuid4().hex[:12], + variant=variant, + run_id=report.identity.run_id, + snapshot_id=snapshot.snapshot_id, + revision=trial.revision, + config_fingerprint=trial.config_fingerprint, + config_fingerprint_derived=trial.config_fingerprint_derived, + input_snapshot=trial.input_snapshot, + workspace_host=report.identity.workspace_host, + job_id=report.identity.job_id, + created_at=snapshot.collected_at, + warmup=trial.warmup, + correctness=trial.correctness, + ) + manifest.trials.append(ref) + if manifest.baseline_trial_id is None: + manifest.baseline_trial_id = ref.trial_id + save_manifest(exp, manifest) + return ref, path + + +def _existing(trial: TrialRef | None, field: str, default: Any = None) -> Any: + return getattr(trial, field) if trial is not None else default + + +def _check_compatible(manifest: ExperimentManifest, report: RunReport) -> None: + workspace = report.identity.workspace_host + job_id = report.identity.job_id + if manifest.workspace_host and workspace and manifest.workspace_host != workspace: + raise ExperimentError( + f"run is from workspace {workspace!r}, experiment is pinned to " + f"{manifest.workspace_host!r}" + ) + if manifest.job_id and job_id and manifest.job_id != job_id: + raise ExperimentError( + f"run is for job {job_id!r}, experiment is pinned to {manifest.job_id!r}" + ) + + +def _find_trial(manifest: ExperimentManifest, run_id: str) -> TrialRef | None: + for trial in manifest.trials: + if trial.run_id == run_id: + return trial + return None + + +def load_experiment( + exp_dir: str | Path, +) -> tuple[ExperimentManifest, dict[str, TrialSnapshot]]: + """Load a manifest and every snapshot it references, keyed by trial ID.""" + manifest = load_manifest(exp_dir) + snapshots: dict[str, TrialSnapshot] = {} + for trial in manifest.trials: + path = trial_snapshot_path(exp_dir, trial.run_id, trial.snapshot_id) + if not path.exists(): + raise ExperimentError( + f"trial {trial.trial_id} references missing snapshot {path}" + ) + snapshots[trial.trial_id] = load_snapshot(path) + return manifest, snapshots + + +def _select_group( + manifest: ExperimentManifest, + snapshots: dict[str, TrialSnapshot], + *, + run_id: str | None, + variant: str | None, + label: str, +) -> ComparisonGroup: + if bool(run_id) == bool(variant): + raise ExperimentError( + f"{label}: pass exactly one of a run ID or a variant label" + ) + if run_id is not None: + trials = [t for t in manifest.trials if t.run_id == str(run_id)] + if not trials: + raise ExperimentError( + f"{label}: run {run_id} is not recorded in this experiment" + ) + else: + trials = [t for t in manifest.trials if t.variant == variant] + if not trials: + raise ExperimentError(f"{label}: no trials with variant {variant!r}") + fingerprints = { + t.config_fingerprint for t in trials if t.config_fingerprint is not None + } + if len(fingerprints) > 1: + raise ExperimentError( + f"{label}: variant {variant!r} mixes {len(fingerprints)} " + "configurations; record runs of one configuration under one " + "variant or compare explicit run IDs" + ) + + group = ComparisonGroup(label=label, variant=variant) + for trial in trials: + snapshot = snapshots.get(trial.trial_id) + if snapshot is None: + continue + status = _trial_status(snapshot) + if status == "failed": + group.excluded_failed += 1 + continue + if status == "active": + group.excluded_active += 1 + continue + if trial.warmup: + group.excluded_warmup += 1 + continue + if _is_incomplete_objective(trial): + group.excluded_incomplete += 1 + continue + group.trial_ids.append(trial.trial_id) + group.run_ids.append(trial.run_id) + group.snapshot_ids.append(snapshot.snapshot_id) + group.sample_count += 1 + return group + + +def _trial_status(snapshot: TrialSnapshot) -> str: + result = snapshot.report.status.result_state or "" + if snapshot.report.collection_status.value == "partial" and result in ( + "", + "RUNNING", + "PENDING", + ): + return "active" + if result in ("FAILED", "TIMEDOUT", "CANCELED", "UPSTREAM_FAILED"): + return "failed" + return "ok" + + +def _is_incomplete_objective(trial: TrialRef) -> bool: + return trial.correctness in ("pending", "incomplete") + + +def compare_experiment( + exp_dir: str | Path, + *, + baseline_run_id: str | None = None, + candidate_run_id: str | None = None, + baseline_variant: str | None = None, + candidate_variant: str | None = None, +) -> ExperimentComparison: + """Compare a pinned baseline group with a candidate group from saved data.""" + manifest, snapshots = load_experiment(exp_dir) + baseline = _select_group( + manifest, + snapshots, + run_id=baseline_run_id, + variant=baseline_variant, + label="baseline", + ) + candidate = _select_group( + manifest, + snapshots, + run_id=candidate_run_id, + variant=candidate_variant, + label="candidate", + ) + if not baseline.trial_ids and not candidate.trial_ids: + raise ExperimentError( + "both groups have no eligible trials; see excluded counts and notes" + ) + + metrics, partial_metrics = _compare_metrics(baseline, candidate, snapshots) + tasks = _compare_tasks(baseline, candidate, snapshots) + comparability, revision_confidence = _compare_context( + manifest, baseline, candidate, snapshots + ) + if partial_metrics: + comparability.append( + ComparabilityNote( + aspect="metric_completeness", + status="partial", + detail=( + "partial observations retained separately from eligible " + "statistics: " + ", ".join(sorted(partial_metrics)) + ), + ) + ) + if any( + not _task_coverage_complete(snapshots[trial_id].report) + for trial_id in baseline.trial_ids + candidate.trial_ids + ): + comparability.append( + ComparabilityNote( + aspect="task_coverage", + status="partial", + detail=( + "task list was page-limited; summed task time is a partial " + "subtotal and added/removed classification is unreliable" + ), + ) + ) + notes = _group_notes(baseline, candidate) + return ExperimentComparison( + experiment_id=manifest.experiment_id, + baseline=baseline, + candidate=candidate, + metrics=metrics, + tasks=tasks, + comparability=comparability, + revision_confidence=revision_confidence, + notes=notes, + ) + + +def _task_coverage_complete(report: RunReport) -> bool: + """Whether the collected task list is complete for the run. + + A page-limited run reports only the tasks that were fetched; its summed task + duration is a partial subtotal, not a comparable total. + """ + return report.coverage.get("task_list", "complete") != "partial" + + +def _measure_values( + group: ComparisonGroup, snapshots: dict[str, TrialSnapshot] +) -> tuple[dict[str, list[float]], dict[str, list[float]]]: + """Collect per-trial values split into complete and partial observations. + + Eligible sample statistics use the complete observations only: an observed + subtotal over queries that merely reported the metric is not a total. The + partial observations are still returned so a raw delta can be reported with + a coverage caveat rather than the metric disappearing. + """ + complete: dict[str, list[float]] = {} + partial: dict[str, list[float]] = {} + for trial_id in group.trial_ids: + snapshot = snapshots[trial_id] + report = snapshot.report + raw: dict[str, float | None] = { + "workflow_elapsed_ms": report.timing.workflow_elapsed_ms, + } + raw_partial: dict[str, float] = {} + if report.timing.summed_task_execution_ms is not None: + if _task_coverage_complete(report): + raw["task_execution_ms"] = report.timing.summed_task_execution_ms + else: + raw_partial["task_execution_ms"] = ( + report.timing.summed_task_execution_ms + ) + if report.query_metrics is not None: + for aggregate in report.query_metrics.aggregates: + if aggregate.value is None: + continue + if aggregate.completeness is MetricCompleteness.complete: + raw[aggregate.metric] = aggregate.value + else: + raw_partial[aggregate.metric] = aggregate.value + for metric, value in raw.items(): + if value is not None: + complete.setdefault(metric, []).append(float(value)) + for metric, value in raw_partial.items(): + partial.setdefault(metric, []).append(float(value)) + return complete, partial + + +def _compare_metrics( + baseline: ComparisonGroup, + candidate: ComparisonGroup, + snapshots: dict[str, TrialSnapshot], +) -> tuple[list[MetricDelta], set[str]]: + b_complete, b_partial = _measure_values(baseline, snapshots) + c_complete, c_partial = _measure_values(candidate, snapshots) + partial = set(b_partial) | set(c_partial) + deltas: list[MetricDelta] = [] + for metric in sorted(set(b_complete) | set(c_complete) | partial): + bc = b_complete.get(metric, []) + cc = c_complete.get(metric, []) + bp = b_partial.get(metric, []) + cp = c_partial.get(metric, []) + # Prefer complete observations for eligible statistics; fall back to a + # partial side so a partially observed metric is still reported. + b = bc or bp + c = cc or cp + if not b or not c: + continue + b_median = statistics.median(b) + c_median = statistics.median(c) + delta = c_median - b_median + is_partial = bool(bp or cp) + if is_partial and not (bc and cc): + caveat = "raw delta uses observed subtotals, not complete totals" + elif is_partial: + caveat = ( + "statistics use complete observations; partial observations " + "are retained separately" + ) + else: + caveat = _metric_caveat(metric, snapshots) + deltas.append( + MetricDelta( + metric=metric, + unit=_metric_unit(metric, snapshots), + aggregation=_metric_aggregation(metric, snapshots), + baseline_values=b, + candidate_values=c, + baseline_partial_values=bp, + candidate_partial_values=cp, + baseline_median=b_median, + candidate_median=c_median, + baseline_min=min(b), + baseline_max=max(b), + candidate_min=min(c), + candidate_max=max(c), + delta=delta, + pct_change=None if b_median == 0 else delta / b_median, + single_observation=len(b) == 1 and len(c) == 1, + partial=is_partial, + caveat=caveat, + ) + ) + return deltas, partial + + +def _metric_unit(metric: str, snapshots: dict[str, TrialSnapshot]) -> MetricUnit: + if metric in ("workflow_elapsed_ms", "task_execution_ms"): + return MetricUnit.milliseconds + for snapshot in snapshots.values(): + if snapshot.report.query_metrics is None: + continue + for aggregate in snapshot.report.query_metrics.aggregates: + if aggregate.metric == metric: + return aggregate.unit + return MetricUnit.none + + +def _metric_aggregation(metric: str, snapshots: dict[str, TrialSnapshot]) -> str: + if metric in ("workflow_elapsed_ms", "task_execution_ms"): + return "value" + for snapshot in snapshots.values(): + if snapshot.report.query_metrics is None: + continue + for aggregate in snapshot.report.query_metrics.aggregates: + if aggregate.metric == metric: + return aggregate.aggregation + return "unknown" + + +def _metric_caveat(metric: str, snapshots: dict[str, TrialSnapshot]) -> str | None: + if metric in _RUN_METRICS: + return None + for snapshot in snapshots.values(): + if snapshot.report.query_metrics is None: + continue + for aggregate in snapshot.report.query_metrics.aggregates: + if ( + aggregate.metric == metric + and aggregate.completeness.value != "complete" + ): + return ( + "observed subtotal over queries that reported the metric; " + "not a complete total" + ) + return None + + +def _task_values(report: RunReport) -> dict[str, dict[str, float | None]]: + values: dict[str, dict[str, float | None]] = {} + for task in report.tasks: + values[task.task_key] = { + "execution_ms": float(task.timing.execution_ms) + if task.timing.execution_ms is not None + else None, + "setup_ms": float(task.timing.setup_ms) + if task.timing.setup_ms is not None + else None, + "cleanup_ms": float(task.timing.cleanup_ms) + if task.timing.cleanup_ms is not None + else None, + } + return values + + +def _task_samples( + group: ComparisonGroup, snapshots: dict[str, TrialSnapshot] +) -> dict[str, dict[str, list[float]]]: + """Per-task, per-measure samples across every eligible trial in a group.""" + samples: dict[str, dict[str, list[float]]] = {} + for trial_id in group.trial_ids: + for key, measures in _task_values(snapshots[trial_id].report).items(): + for measure, value in measures.items(): + if value is None: + continue + samples.setdefault(key, {}).setdefault(measure, []).append(value) + return samples + + +def _median_map(samples: dict[str, list[float]]) -> dict[str, float]: + return {measure: statistics.median(values) for measure, values in samples.items()} + + +def _compare_tasks( + baseline: ComparisonGroup, + candidate: ComparisonGroup, + snapshots: dict[str, TrialSnapshot], +) -> list[TaskDelta]: + base = _task_samples(baseline, snapshots) + cand = _task_samples(candidate, snapshots) + coverage_partial = any( + not _task_coverage_complete(snapshots[trial_id].report) + for trial_id in baseline.trial_ids + candidate.trial_ids + ) + deltas: list[TaskDelta] = [] + for key in sorted(set(base) | set(cand)): + b_samples = base.get(key, {}) + c_samples = cand.get(key, {}) + b = _median_map(b_samples) + c = _median_map(c_samples) + if key not in base: + change = "unknown" if coverage_partial else "added" + elif key not in cand: + change = "unknown" if coverage_partial else "removed" + elif _measures_equal(b, c): + change = "same" + else: + change = "changed" + delta_map: dict[str, float | None] = {} + for measure in set(b) | set(c): + base_value = b.get(measure) + cand_value = c.get(measure) + if base_value is None or cand_value is None: + delta_map[measure] = None + else: + delta_map[measure] = cand_value - base_value + deltas.append( + TaskDelta( + task_key=key, + change=change, + baseline=b, + candidate=c, + deltas=delta_map, + baseline_values=b_samples, + candidate_values=c_samples, + baseline_min={ + measure: min(values) for measure, values in b_samples.items() + }, + baseline_max={ + measure: max(values) for measure, values in b_samples.items() + }, + candidate_min={ + measure: min(values) for measure, values in c_samples.items() + }, + candidate_max={ + measure: max(values) for measure, values in c_samples.items() + }, + ) + ) + return deltas + + +def _measures_equal(a: dict[str, float], b: dict[str, float]) -> bool: + return all(a.get(k) == b.get(k) for k in set(a) | set(b)) + + +def _compute_fingerprint(report: RunReport) -> str: + """Derive configuration identity from measured run facts, not user labels. + + Compute snapshot, run environments, job parameters, task parameters, + cluster Spark configuration and the effective performance target all + participate so that two different configurations cannot be mistaken for one. + """ + compute_parts = [ + "|".join( + str(value) + for value in ( + reference.cluster_id, + reference.environment_key, + reference.runtime_engine, + reference.spark_version, + reference.node_type_id, + reference.num_workers, + reference.performance_target, + ) + ) + for reference in sorted(report.compute, key=lambda r: r.task_key) + ] + environment_parts = [ + "|".join( + [ + environment.environment_key, + environment.client or "", + ",".join(sorted(environment.dependencies)), + ] + ) + for environment in sorted(report.environments, key=lambda e: e.environment_key) + ] + parameter_parts = [ + f"{name}={value}" for name, value in sorted(report.job_parameters.items()) + ] + task_parts = [ + "|".join( + [ + task.task_key, + ",".join(f"{k}={v}" for k, v in sorted(task.parameters.items())), + ",".join(f"{k}={v}" for k, v in sorted(task.spark_conf.items())), + ] + ) + for task in sorted(report.tasks, key=lambda t: t.task_key) + ] + return "||".join( + [ + ";".join(compute_parts), + "&".join(environment_parts), + "&".join(parameter_parts), + ";".join(task_parts), + report.effective_performance_target or "", + ] + ) + + +def _compare_context( + manifest: ExperimentManifest, + baseline: ComparisonGroup, + candidate: ComparisonGroup, + snapshots: dict[str, TrialSnapshot], +) -> tuple[list[ComparabilityNote], str]: + notes: list[ComparabilityNote] = [] + + def side(group: ComparisonGroup) -> list[TrialSnapshot]: + return [snapshots[t] for t in group.trial_ids] + + b_snaps = side(baseline) + c_snaps = side(candidate) + + b_rev = { + s.report.revision.executed_commit + for s in b_snaps + if s.report.revision.executed_commit + } + c_rev = { + s.report.revision.executed_commit + for s in c_snaps + if s.report.revision.executed_commit + } + if b_rev and c_rev: + status = "different" if b_rev != c_rev else "match" + detail = f"baseline={sorted(b_rev)} candidate={sorted(c_rev)}" + else: + status = "unknown" + detail = "one or both sides lack an executed Git snapshot" + notes.append(ComparabilityNote(aspect="revision", status=status, detail=detail)) + + b_inputs = {s.trial.input_snapshot for s in b_snaps} + c_inputs = {s.trial.input_snapshot for s in c_snaps} + if None in b_inputs or None in c_inputs: + input_status = "unknown" + elif b_inputs != c_inputs: + input_status = "different" + else: + input_status = "match" + notes.append( + ComparabilityNote( + aspect="input_snapshot", + status=input_status, + detail=f"baseline={sorted(x for x in b_inputs if x)} candidate={sorted(x for x in c_inputs if x)}", + ) + ) + + b_config = {s.trial.config_fingerprint for s in b_snaps} + c_config = {s.trial.config_fingerprint for s in c_snaps} + if None in b_config or None in c_config: + config_status = "unknown" + elif b_config != c_config: + config_status = "different" + else: + config_status = "match" + notes.append( + ComparabilityNote( + aspect="configuration", + status=config_status, + detail="configuration differences may be the intended experiment variable", + ) + ) + + b_compute = {_compute_fingerprint(s.report) for s in b_snaps} + c_compute = {_compute_fingerprint(s.report) for s in c_snaps} + compute_status = ( + "match" if b_compute == c_compute and b_compute != {""} else "different" + ) + notes.append( + ComparabilityNote( + aspect="compute", + status=compute_status, + detail="changed compute can be the intended experiment variable", + ) + ) + + b_correct = {s.trial.correctness for s in b_snaps} + c_correct = {s.trial.correctness for s in c_snaps} + correctness_status = ( + "unknown" + if None in b_correct or None in c_correct + else "different" + if b_correct != c_correct + else "match" + ) + notes.append( + ComparabilityNote( + aspect="correctness", + status=correctness_status, + detail="job success alone is not proof of correctness", + ) + ) + + if len(b_rev) == 1 and len(c_rev) == 1 and b_rev == c_rev: + confidence = "executed" + elif b_rev and c_rev: + confidence = "mixed" + elif any(s.report.revision.confidence == "asserted" for s in b_snaps + c_snaps): + confidence = "asserted" + else: + confidence = "unknown" + return notes, confidence + + +def _group_notes(baseline: ComparisonGroup, candidate: ComparisonGroup) -> list[str]: + notes: list[str] = [] + for group in (baseline, candidate): + excluded = { + "failed": group.excluded_failed, + "active": group.excluded_active, + "warmup": group.excluded_warmup, + "incomplete": group.excluded_incomplete, + } + for reason, count in excluded.items(): + if count: + notes.append( + f"{group.label}: {count} {reason} trial(s) excluded from aggregates" + ) + if baseline.sample_count == 1 and candidate.sample_count == 1: + notes.append( + "single observation per group; treat deltas as indicative, not significant" + ) + return notes + + +def _compute_summary(report: RunReport) -> str: + parts: list[str] = [] + for reference in sorted(report.compute, key=lambda r: r.task_key): + bits = [ + reference.runtime_engine or "?", + reference.spark_version or "?", + reference.node_type_id or "?", + f"{reference.num_workers}w" if reference.num_workers is not None else "?w", + ] + suffix = " (current lookup)" if reference.source == "current_lookup" else "" + parts.append(f"{reference.task_key}: {'/'.join(bits)}{suffix}") + if report.effective_performance_target: + parts.append(f"perf_target={report.effective_performance_target}") + return "; ".join(parts) + + +def _top_queries(report: RunReport, limit: int = 5) -> list[dict[str, Any]]: + if report.query_metrics is None: + return [] + queries: list[dict[str, Any]] = [] + for observation in report.query_metrics.observations: + total = observation.metrics.get("total_time_ms") + if not isinstance(total, int | float) or isinstance(total, bool): + total = observation.metrics.get("execution_time_ms") + if not isinstance(total, int | float) or isinstance(total, bool): + continue + queries.append( + { + "query_id": observation.query_id, + "status": observation.status, + "total_ms": int(total), + "read_bytes": observation.metrics.get("read_bytes"), + } + ) + queries.sort(key=lambda q: q["total_ms"], reverse=True) + return queries[:limit] + + +def _failure_excerpts(report: RunReport) -> list[dict[str, Any]]: + return [ + {"task_key": task.task_key, "excerpt": task.error_excerpt} + for task in report.tasks + if task.result_state in ("FAILED", "TIMEDOUT", "CANCELED") + and task.error_excerpt + ] + + +def _eligibility(status: str, warmup: bool, correctness: str | None) -> str: + if status == "failed": + return "failed" + if status == "active": + return "active" + if warmup: + return "warmup" + if correctness in ("pending", "incomplete"): + return "incomplete" + return "ok" + + +def trial_rows( + manifest: ExperimentManifest, snapshots: dict[str, TrialSnapshot] +) -> list[dict[str, Any]]: + """Per-trial metric rows for tabular and plotted views, in execution order.""" + rows: list[dict[str, Any]] = [] + for trial in manifest.trials: + snapshot = snapshots.get(trial.trial_id) + if snapshot is None: + continue + report = snapshot.report + values: dict[str, float] = {} + if report.timing.workflow_elapsed_ms is not None: + values["workflow_elapsed_ms"] = float(report.timing.workflow_elapsed_ms) + partial: list[str] = [] + if report.timing.summed_task_execution_ms is not None: + values["task_execution_ms"] = float(report.timing.summed_task_execution_ms) + if not _task_coverage_complete(report): + partial.append("task_execution_ms") + if report.query_metrics is not None: + for aggregate in report.query_metrics.aggregates: + if aggregate.value is None: + continue + values[aggregate.metric] = float(aggregate.value) + if aggregate.completeness.value != "complete": + partial.append(aggregate.metric) + status = _trial_status(snapshot) + task_values = { + task.task_key: ( + float(task.timing.execution_ms) + if task.timing.execution_ms is not None + else None + ) + for task in report.tasks + } + rows.append( + { + "trial_id": trial.trial_id, + "run_id": trial.run_id, + "variant": trial.variant, + "revision": trial.revision or report.revision.executed_commit, + "short_revision": ( + report.revision.executed_commit or trial.revision or "unknown" + )[:8], + "collected_at": trial.created_at.isoformat(), + "snapshot_collected_at": snapshot.collected_at.isoformat(), + "status": status, + "eligibility": _eligibility(status, trial.warmup, trial.correctness), + "warmup": trial.warmup, + "correctness": trial.correctness, + "config_fingerprint": trial.config_fingerprint, + "input_snapshot": trial.input_snapshot, + "run_page_url": report.identity.run_page_url, + "values": values, + "task_values": task_values, + "partial_metrics": partial, + "coverage": dict(report.coverage), + "compute_summary": _compute_summary(report), + "top_queries": _top_queries(report), + "failure_excerpts": _failure_excerpts(report), + "failed_tasks": [ + t.task_key for t in report.tasks if t.result_state == "FAILED" + ], + } + ) + rows.sort(key=lambda row: row["collected_at"]) + return rows diff --git a/sparkparse/experiments_dashboard.py b/sparkparse/experiments_dashboard.py new file mode 100644 index 0000000..acadaff --- /dev/null +++ b/sparkparse/experiments_dashboard.py @@ -0,0 +1,840 @@ +"""Dash/Plotly view for a saved experiment directory. + +This is a standalone, read-only app: it loads the manifest and immutable +snapshots once, binds to loopback on launch, and reloads from disk on request so +a refresh reflects newly recorded trials. The same comparison functions used by +the CLI and JSON output back every number shown here. +""" + +from __future__ import annotations + +import statistics +from typing import Any + +import dash +import dash_bootstrap_components as dbc +import plotly.graph_objects as go +from dash import Input, Output, State, dcc, html +from plotly.subplots import make_subplots + +from sparkparse import experiments as exp +from sparkparse.models import ExperimentComparison, MetricUnit + +_TIME_METRICS = { + "workflow_elapsed_ms": "Workflow elapsed", + "task_execution_ms": "Aggregate task time", + "aggregate_task_time_ms": "Query task time", + "query_execution_ms": "Query execution", + "query_total_ms": "Query total", +} +_BYTE_METRICS = { + "read_bytes": "Read bytes", + "read_remote_bytes": "Read remote bytes", + "read_cache_bytes": "Read cache bytes", + "written_bytes_remote": "Written bytes (remote)", + "disk_spill_bytes": "Disk spill", + "shuffle_network_bytes": "Shuffle network bytes", +} +_OTHER_METRICS = { + "rows_read": "Rows read", + "rows_produced": "Rows produced", + "queries_result_from_cache": "Queries served from cache", + "read_files": "Files read", + "read_partitions": "Partitions read", +} +_ALL_METRICS = {**_TIME_METRICS, **_BYTE_METRICS, **_OTHER_METRICS} + +_ELIGIBILITY = ("ok", "warmup", "failed", "active", "incomplete") + +_SCALE = { + MetricUnit.milliseconds: ("seconds", 1000.0), + MetricUnit.bytes: ("MB", 1024.0 * 1024.0), + MetricUnit.rows: ("rows", 1.0), + MetricUnit.items: ("count", 1.0), + MetricUnit.none: ("value", 1.0), +} + + +def _unit_for(metric: str) -> MetricUnit: + if metric in _TIME_METRICS: + return MetricUnit.milliseconds + if metric in _BYTE_METRICS: + return MetricUnit.bytes + return MetricUnit.rows + + +def _run_options(rows: list[dict[str, Any]]) -> list[dict[str, str]]: + return [ + { + "label": f"{row['variant']} @ {row['short_revision']} ({row['run_id']})", + "value": row["run_id"], + } + for row in rows + ] + + +def _multi_options(values: list[str]) -> list[dict[str, str]]: + return [{"label": value, "value": value} for value in values] + + +def _filter_rows( + rows: list[dict[str, Any]], + variants: list[str] | None, + revisions: list[str] | None, + eligibility: list[str] | None, +) -> list[dict[str, Any]]: + return [ + row + for row in rows + if (not variants or row["variant"] in variants) + and (not revisions or row["short_revision"] in revisions) + and (not eligibility or row["eligibility"] in eligibility) + ] + + +def _trend_figure( + rows: list[dict[str, Any]], + metric: str, + baseline_run: str | None, + candidate_run: str | None, +) -> go.Figure: + unit = _unit_for(metric) + scale_label, scale = _SCALE[unit] + fig = make_subplots( + rows=2, + cols=1, + shared_xaxes=True, + subplot_titles=(_ALL_METRICS.get(metric, metric), f"Context ({scale_label})"), + ) + complete = [row for row in rows if metric not in row["partial_metrics"]] + partial = [row for row in rows if metric in row["partial_metrics"]] + xs = [row["run_id"] for row in complete] + primary = [ + row["values"].get(metric) / scale if metric in row["values"] else None + for row in complete + ] + fig.add_trace( + go.Scatter( + x=xs, + y=primary, + mode="markers+lines", + name=_ALL_METRICS.get(metric, metric), + text=[f"{row['variant']} {row['short_revision']}" for row in complete], + hovertemplate="%{text}
%{y:.3f} " + scale_label + "", + connectgaps=False, + ), + row=1, + col=1, + ) + if partial: + fig.add_trace( + go.Scatter( + x=[row["run_id"] for row in partial], + y=[ + row["values"].get(metric) / scale + if metric in row["values"] + else None + for row in partial + ], + mode="markers", + marker={"symbol": "circle-open", "size": 11}, + name=f"{_ALL_METRICS.get(metric, metric)} (partial)", + text=[f"{row['variant']} {row['short_revision']}" for row in partial], + hovertemplate=( + "%{text}
%{y:.3f} " + scale_label + " partial" + ), + ), + row=1, + col=1, + ) + context_metric = "workflow_elapsed_ms" if metric in _BYTE_METRICS else "read_bytes" + if context_metric in _ALL_METRICS: + _, context_scale = _SCALE[ + MetricUnit.milliseconds + if context_metric in _TIME_METRICS + else MetricUnit.bytes + ] + context_values = [ + row["values"].get(context_metric) / context_scale + if context_metric in row["values"] + else None + for row in rows + ] + context_label = ( + "seconds" + if context_metric in _TIME_METRICS + else "MB" + if context_metric in _BYTE_METRICS + else "value" + ) + fig.add_trace( + go.Bar( + x=[row["run_id"] for row in rows], + y=context_values, + name=_ALL_METRICS.get(context_metric, context_metric), + hovertemplate="%{y:.3f} " + context_label + "", + ), + row=2, + col=1, + ) + for group, run_id, color in ( + ("baseline", baseline_run, "#1f77b4"), + ("candidate", candidate_run, "#ff7f0e"), + ): + if not run_id: + continue + fig.add_vline(x=run_id, line_dash="dot", line_color=color, row=1, col=1) + _add_median_line(fig, rows, metric, run_id, color, scale) + fig.update_layout( + margin={"l": 40, "r": 20, "t": 40, "b": 30}, + height=460, + legend={"orientation": "h"}, + paper_bgcolor="rgba(0,0,0,0)", + plot_bgcolor="rgba(0,0,0,0)", + ) + return fig + + +def _add_median_line( + fig: go.Figure, + rows: list[dict[str, Any]], + metric: str, + run_id: str | None, + color: str, + scale: float, +) -> None: + if run_id is None: + return + values = [ + row["values"][metric] / scale + for row in rows + if row["run_id"] == run_id + and metric in row["values"] + and metric not in row["partial_metrics"] + ] + if not values: + return + median = statistics.median(values) + fig.add_hline(y=median, line_dash="dash", line_color=color, row=1, col=1) + + +def _task_samples_for_variant( + rows: list[dict[str, Any]], + task_key: str, + variant: str, + config_fingerprint: str | None, +) -> list[float]: + """Eligible, same-configuration samples behind a variant's displayed spread. + + Excluded trials (failed, active, warmup, incomplete) and trials of another + configuration must not move the median or spread shown next to a comparison. + """ + return [ + row["task_values"][task_key] + for row in rows + if row["variant"] == variant + and row["config_fingerprint"] == config_fingerprint + and row["eligibility"] == "ok" + and row["task_values"].get(task_key) is not None + ] + + +def _task_trend_figure( + rows: list[dict[str, Any]], + task_key: str | None, + baseline_run: str | None, + candidate_run: str | None, +) -> go.Figure: + fig = go.Figure() + if task_key is None: + return fig + xs = [row["run_id"] for row in rows] + ys = [ + row["task_values"].get(task_key) + if row["task_values"].get(task_key) is not None + else None + for row in rows + ] + fig.add_trace( + go.Scatter( + x=xs, + y=ys, + mode="markers+lines", + name=f"{task_key} execution", + text=[f"{row['variant']} {row['short_revision']}" for row in rows], + hovertemplate="%{text}
%{y:.0f} ms", + connectgaps=False, + ) + ) + meta_by_run = { + row["run_id"]: (row["variant"], row["config_fingerprint"]) for row in rows + } + for run_id, color, label in ( + (baseline_run, "#1f77b4", "baseline"), + (candidate_run, "#ff7f0e", "candidate"), + ): + if not run_id: + continue + meta = meta_by_run.get(run_id) + if meta is None: + continue + variant, config = meta + samples = _task_samples_for_variant(rows, task_key, variant, config) + if not samples: + continue + median = statistics.median(samples) + fig.add_trace( + go.Scatter( + x=[run_id], + y=[median], + mode="markers", + marker={"size": 14, "color": color, "symbol": "diamond"}, + error_y={ + "type": "data", + "symmetric": False, + "array": [max(samples) - median], + "arrayminus": [median - min(samples)], + }, + name=f"{label} variant spread (n={len(samples)})", + hovertemplate=( + f"{label} variant median %{{y:.0f}} ms" + f"
min {min(samples):.0f} / max {max(samples):.0f}" + "" + ), + ) + ) + fig.update_layout( + margin={"l": 40, "r": 20, "t": 40, "b": 30}, + height=360, + yaxis_title="ms", + legend={"orientation": "h"}, + paper_bgcolor="rgba(0,0,0,0)", + plot_bgcolor="rgba(0,0,0,0)", + ) + return fig + + +def _metric_table(comparison: ExperimentComparison) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for metric in comparison.metrics: + rows.append( + { + "metric": metric.metric, + "baseline": metric.baseline_median, + "candidate": metric.candidate_median, + "delta": metric.delta, + "pct_change": metric.pct_change, + "baseline_n": len(metric.baseline_values), + "candidate_n": len(metric.candidate_values), + "partial": metric.partial, + "caveat": metric.caveat, + } + ) + return rows + + +def _fmt_samples(values: list[float] | None) -> str: + if not values: + return "" + return ", ".join(f"{value:.0f}" for value in values) + + +def _task_table(comparison: ExperimentComparison) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for task in comparison.tasks: + rows.append( + { + "task_key": task.task_key, + "change": task.change, + "baseline_execution_ms": task.baseline.get("execution_ms"), + "candidate_execution_ms": task.candidate.get("execution_ms"), + "delta_execution_ms": task.deltas.get("execution_ms"), + "baseline_samples": _fmt_samples( + task.baseline_values.get("execution_ms") + ), + "candidate_samples": _fmt_samples( + task.candidate_values.get("execution_ms") + ), + "baseline_range": _fmt_range( + task.baseline_min.get("execution_ms"), + task.baseline_max.get("execution_ms"), + ), + "candidate_range": _fmt_range( + task.candidate_min.get("execution_ms"), + task.candidate_max.get("execution_ms"), + ), + } + ) + return rows + + +def _fmt_range(low: float | None, high: float | None) -> str: + if low is None or high is None: + return "" + return f"{low:.0f}–{high:.0f}" + + +def _detail_table(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: + return [ + { + "run_id": row["run_id"], + "variant": row["variant"], + "revision": row["short_revision"], + "collected_at": row["collected_at"], + "status": row["status"], + "eligibility": row["eligibility"], + "config": row["config_fingerprint"], + "input": row["input_snapshot"], + "compute": row["compute_summary"], + "coverage": ", ".join( + f"{key}={value}" for key, value in sorted(row["coverage"].items()) + ), + "top_queries": " | ".join( + f"{q['query_id'][:8]} {q['total_ms']}ms" for q in row["top_queries"] + ), + "failure_excerpts": " || ".join( + f"{f['task_key']}: {f['excerpt']}" for f in row["failure_excerpts"] + ), + "run_url": row["run_page_url"], + } + for row in rows + ] + + +_CONTEXT_KEYS = { + "revision": "short_revision", + "input_snapshot": "input_snapshot", + "configuration": "config_fingerprint", + "compute": "compute_summary", +} + + +def _context_table( + comparison: ExperimentComparison, + baseline_row: dict[str, Any] | None, + candidate_row: dict[str, Any] | None, +) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for note in comparison.comparability: + key = _CONTEXT_KEYS.get(note.aspect) + rows.append( + { + "aspect": note.aspect, + "status": note.status, + "baseline": (baseline_row or {}).get(key) if key else None, + "candidate": (candidate_row or {}).get(key) if key else None, + "detail": note.detail, + } + ) + return rows + + +def flatten_trend_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Flatten list-valued trend fields so CSV export has no nested columns.""" + return [ + { + key: ", ".join(str(item) for item in value) + if isinstance(value, list) + else value + for key, value in row.items() + } + for row in rows + ] + + +def _table(columns: list[str], data: list[dict[str, Any]], table_id: str) -> html.Div: + if not data: + return html.Div("No data.", style={"padding": "8px"}) + return html.Div( + dash.dash_table.DataTable( + id=table_id, + columns=[{"name": c.replace("_", " "), "id": c} for c in columns], + data=data, + sort_action="native", + page_action="native", + page_size=12, + style_table={"overflowX": "auto"}, + ) + ) + + +def build_layout( + rows: list[dict[str, Any]], + default_baseline: str | None = None, + default_candidate: str | None = None, +) -> html.Div: + options = _run_options(rows) + variants = sorted({row["variant"] for row in rows}) + revisions = sorted({row["short_revision"] for row in rows}) + task_keys = sorted({key for row in rows for key in row["task_values"]}) + metric_options = [ + {"label": _ALL_METRICS[m], "value": m} + for m in _ALL_METRICS + if any(m in row["values"] for row in rows) + ] or [{"label": "Workflow elapsed", "value": "workflow_elapsed_ms"}] + return dbc.Container( + [ + html.H3("Experiment"), + html.Div(id="experiment-summary"), + dbc.Row( + [ + dbc.Col( + [ + html.Label("Variants"), + dcc.Dropdown( + id="variant-filter", + options=_multi_options(variants), + value=variants, + multi=True, + ), + ], + width=4, + ), + dbc.Col( + [ + html.Label("Commits"), + dcc.Dropdown( + id="revision-filter", + options=_multi_options(revisions), + value=revisions, + multi=True, + ), + ], + width=4, + ), + dbc.Col( + [ + html.Label("Eligibility"), + dcc.Dropdown( + id="eligibility-filter", + options=_multi_options(list(_ELIGIBILITY)), + value=list(_ELIGIBILITY), + multi=True, + ), + ], + width=4, + ), + ] + ), + html.Br(), + dbc.Row( + [ + dbc.Col( + [ + html.Label("Baseline run"), + dcc.Dropdown( + id="baseline-select", + options=options, + value=default_baseline, + clearable=False, + ), + ], + width=3, + ), + dbc.Col( + [ + html.Label("Candidate run"), + dcc.Dropdown( + id="candidate-select", + options=options, + value=default_candidate, + clearable=False, + ), + ], + width=3, + ), + dbc.Col( + [ + html.Label("Metric"), + dcc.Dropdown( + id="metric-select", + options=metric_options, + value=metric_options[0]["value"], + clearable=False, + ), + ], + width=3, + ), + dbc.Col( + [ + html.Label("Task"), + dcc.Dropdown( + id="task-select", + options=_multi_options(task_keys), + value=task_keys[0] if task_keys else None, + clearable=False, + ), + ], + width=3, + ), + ] + ), + html.Br(), + dcc.Graph(id="trend-figure"), + html.H5("Task trend and group spread"), + dcc.Graph(id="task-trend-figure"), + html.Div(id="comparison-notes", style={"padding": "8px"}), + html.H5("Metrics"), + html.Div(id="metrics-table-container"), + html.H5("Tasks"), + html.Div(id="tasks-table-container"), + html.H5("Run context"), + html.Div(id="context-table-container"), + html.H5("Run detail"), + html.Div(id="detail-table-container"), + html.Br(), + dbc.ButtonGroup( + [ + dbc.Button("Download CSV", id="download-csv", n_clicks=0), + dbc.Button("Download JSON", id="download-json", n_clicks=0), + ] + ), + dcc.Store(id="export-payload"), + dcc.Download(id="download"), + ], + fluid=True, + ) + + +def default_runs( + manifest: exp.ExperimentManifest, rows: list[dict[str, Any]] +) -> tuple[str | None, str | None]: + """Default selectors: the manifest's pinned baseline, and the newest trial.""" + pinned = next( + (t.run_id for t in manifest.trials if t.trial_id == manifest.baseline_trial_id), + None, + ) + return pinned or (rows[0]["run_id"] if rows else None), ( + rows[-1]["run_id"] if rows else None + ) + + +def init_experiments_dashboard(exp_dir: str) -> dash.Dash: + manifest, snapshots = exp.load_experiment(exp_dir) + rows = exp.trial_rows(manifest, snapshots) + default_baseline, default_candidate = default_runs(manifest, rows) + + app = dash.Dash( + __name__, + external_stylesheets=[dbc.themes.BOOTSTRAP, dbc.icons.FONT_AWESOME], + suppress_callback_exceptions=True, + ) + app.layout = build_layout(rows, default_baseline, default_candidate) + app.server.config["EXPERIMENT_DIR"] = exp_dir + + @app.callback( # type: ignore[missing-attribute] + [ + Output("trend-figure", "figure"), + Output("task-trend-figure", "figure"), + Output("metrics-table-container", "children"), + Output("tasks-table-container", "children"), + Output("context-table-container", "children"), + Output("detail-table-container", "children"), + Output("comparison-notes", "children"), + Output("experiment-summary", "children"), + Output("export-payload", "data"), + ], + [ + Input("baseline-select", "value"), + Input("candidate-select", "value"), + Input("metric-select", "value"), + Input("task-select", "value"), + Input("variant-filter", "value"), + Input("revision-filter", "value"), + Input("eligibility-filter", "value"), + ], + ) + def refresh( + baseline_run: str | None, + candidate_run: str | None, + metric: str, + task_key: str | None, + variants: list[str] | None, + revisions: list[str] | None, + eligibility: list[str] | None, + ): + manifest, snapshots = exp.load_experiment(exp_dir) + rows = exp.trial_rows(manifest, snapshots) + filtered = _filter_rows(rows, variants, revisions, eligibility) + + comparison: ExperimentComparison | None = None + error: str | None = None + if baseline_run and candidate_run: + try: + comparison = exp.compare_experiment( + exp_dir, + baseline_run_id=baseline_run, + candidate_run_id=candidate_run, + ) + except exp.ExperimentError as exc: + error = str(exc) + + baseline_row = next( + (row for row in rows if row["run_id"] == baseline_run), None + ) + candidate_row = next( + (row for row in rows if row["run_id"] == candidate_run), None + ) + figure = _trend_figure(filtered, metric, baseline_run, candidate_run) + task_figure = _task_trend_figure( + filtered, task_key, baseline_run, candidate_run + ) + + if comparison is None: + notes = html.Div(error or "Select two runs to compare.") + return ( + figure, + task_figure, + html.Div("No comparison.", style={"padding": "8px"}), + html.Div("No comparison.", style={"padding": "8px"}), + html.Div("No comparison.", style={"padding": "8px"}), + _table( + [ + "run_id", + "variant", + "revision", + "collected_at", + "status", + "eligibility", + "config", + "input", + "compute", + "coverage", + "top_queries", + "failure_excerpts", + "run_url", + ], + _detail_table(filtered), + "detail-table", + ), + notes, + html.Div(error or ""), + { + "trend": [ + { + "run_id": row["run_id"], + "variant": row["variant"], + "partial_metrics": row["partial_metrics"], + **row["values"], + } + for row in filtered + ], + "error": error, + }, + ) + + notes_children = ( + html.Ul([html.Li(note) for note in comparison.notes]) + if comparison.notes + else "" + ) + summary = html.Div( + [ + html.Span(f"experiment_id: {comparison.experiment_id} "), + html.Span(f"revision confidence: {comparison.revision_confidence} "), + html.Span( + f"baseline n={comparison.baseline.sample_count} " + f"candidate n={comparison.candidate.sample_count} " + ), + html.Span(f"showing {len(filtered)}/{len(rows)} trial(s)"), + ] + ) + payload = { + "trend": [ + { + "run_id": row["run_id"], + "variant": row["variant"], + "partial_metrics": row["partial_metrics"], + **row["values"], + } + for row in filtered + ], + "metrics": [m.model_dump(mode="json") for m in comparison.metrics], + "tasks": [t.model_dump(mode="json") for t in comparison.tasks], + "comparability": [ + note.model_dump(mode="json") for note in comparison.comparability + ], + } + return ( + figure, + task_figure, + _table( + [ + "metric", + "baseline", + "candidate", + "delta", + "pct_change", + "baseline_n", + "candidate_n", + "partial", + "caveat", + ], + _metric_table(comparison), + "metrics-table", + ), + _table( + [ + "task_key", + "change", + "baseline_execution_ms", + "candidate_execution_ms", + "delta_execution_ms", + "baseline_samples", + "candidate_samples", + "baseline_range", + "candidate_range", + ], + _task_table(comparison), + "tasks-table", + ), + _table( + ["aspect", "status", "baseline", "candidate", "detail"], + _context_table(comparison, baseline_row, candidate_row), + "context-table", + ), + _table( + [ + "run_id", + "variant", + "revision", + "collected_at", + "status", + "eligibility", + "config", + "input", + "compute", + "coverage", + "top_queries", + "failure_excerpts", + "run_url", + ], + _detail_table(filtered), + "detail-table", + ), + notes_children, + summary, + payload, + ) + + @app.callback( # type: ignore[missing-attribute] + Output("download", "data"), + [ + Input("download-csv", "n_clicks"), + Input("download-json", "n_clicks"), + ], + State("export-payload", "data"), + prevent_initial_call=True, + ) + def download(csv_clicks: int, json_clicks: int, payload: dict[str, Any]): + triggered = dash.callback_context.triggered_id + if triggered == "download-json": + return dcc.send_string( + __import__("json").dumps(payload, indent=2), "experiment.json" + ) + import polars as pl + + frame = pl.DataFrame(flatten_trend_rows(payload.get("trend", []))) + return dcc.send_string(frame.write_csv(), "experiment-trend.csv") + + return app diff --git a/sparkparse/models.py b/sparkparse/models.py index 9c2b1af..ec08318 100644 --- a/sparkparse/models.py +++ b/sparkparse/models.py @@ -1970,3 +1970,374 @@ class AnalysisReport(BaseModel): log_name: str findings: list[Finding] = Field(default_factory=list) assessments: list[RuleAssessment] = Field(default_factory=list) + + +DATABRICKS_REPORT_VERSION = "1" +EXPERIMENT_SCHEMA_VERSION = "1" + + +class CollectionStatus(StrEnum): + """Whether collection finished within its configured bounds.""" + + complete = "complete" + partial = "partial" + failed = "failed" + + +class MetricCompleteness(StrEnum): + """Coverage of an aggregate over the discovered query population.""" + + complete = "complete" + observed_subtotal = "observed_subtotal" + unavailable = "unavailable" + + +class WorkflowCollectionLimits(BaseModel): + """Bounds persisted with every report so partial collection is explainable.""" + + deadline_seconds: float + max_pages: int + max_requests: int + max_retries: int + retry_backoff_seconds: float + outputs: str = "failed" + + +class WorkflowRunIdentity(BaseModel): + run_id: str + job_id: str | None = None + run_name: str | None = None + run_type: str | None = None + number_in_job: int | None = None + original_attempt_run_id: str | None = None + run_page_url: str | None = None + workspace_host: str | None = None + creator_user_name: str | None = None + + +class WorkflowRunStatus(BaseModel): + life_cycle_state: str | None = None + result_state: str | None = None + state_message: str | None = None + termination_code: str | None = None + termination_type: str | None = None + user_cancelled_or_timedout: bool | None = None + + +class RunEnvironment(BaseModel): + """A run-level environment specification, part of configuration identity.""" + + environment_key: str + client: str | None = None + dependencies: list[str] = Field(default_factory=list) + + +class WorkflowTiming(BaseModel): + """Run-level timing. + + ``workflow_elapsed_ms`` is wall-clock and is never the sum of task durations: + workflow tasks may overlap. ``summed_task_execution_ms`` is an aggregate, not + an elapsed measure. + """ + + workflow_elapsed_ms: int | None = None + setup_ms: int | None = None + execution_ms: int | None = None + cleanup_ms: int | None = None + summed_task_execution_ms: int | None = None + tasks_may_overlap: bool = True + + +class ComputeReference(BaseModel): + """A task's compute snapshot. + + ``source`` is ``run_snapshot`` for fields embedded in the run response and + ``current_lookup`` for a later cluster query, which describes the cluster + now rather than at run time. + """ + + task_key: str + environment_key: str | None = None + cluster_id: str | None = None + source: str = "run_snapshot" + runtime_engine: str | None = None + spark_version: str | None = None + node_type_id: str | None = None + num_workers: int | None = None + data_security_mode: str | None = None + performance_target: str | None = None + + +class RevisionEvidence(BaseModel): + """Executed-code provenance, kept separate from asserted revisions. + + ``confidence`` is ``executed`` when a Git snapshot reports the used commit, + ``asserted`` when only ``--revision`` was supplied, ``mixed`` when tasks + disagree, and ``unknown`` otherwise. Branch names and local HEAD are never + treated as proof of executed code. + """ + + executed_commit: str | None = None + repository_url: str | None = None + repository_provider: str | None = None + branch: str | None = None + commit_timestamp_ms: int | None = None + scope: str | None = None + asserted_revision: str | None = None + artifact_digest: str | None = None + confidence: str = "unknown" + + +class TaskTiming(BaseModel): + setup_ms: int | None = None + execution_ms: int | None = None + cleanup_ms: int | None = None + start_time_ms: int | None = None + end_time_ms: int | None = None + + +class WorkflowTask(BaseModel): + """One task run. Identity is retained even when its task type is unknown.""" + + task_key: str + run_id: str | None = None + attempt_number: int | None = None + original_attempt_run_id: str | None = None + run_if: str | None = None + state: str | None = None + result_state: str | None = None + timing: TaskTiming = Field(default_factory=TaskTiming) + compute: ComputeReference | None = None + task_kind: str | None = None + is_nested_job: bool = False + git_commit: str | None = None + parameters: dict[str, str] = Field(default_factory=dict) + spark_conf: dict[str, str] = Field(default_factory=dict) + output_path: str | None = None + output_status: str | None = None + error_excerpt: str | None = None + + +class OutputAttachment(BaseModel): + task_key: str + run_id: str | None = None + status: str + path: str | None = None + error: str | None = None + excerpt: str | None = None + + +class QueryObservation(BaseModel): + """One query's facts, kept once per query ID across recollections.""" + + query_id: str + status: str | None = None + is_final: bool | None = None + execution_end_time_ms: int | None = None + collected_at: datetime.datetime + job_id: str | None = None + job_run_id: str | None = None + job_task_run_id: str | None = None + attribution: str = "unattributed" + task_key: str | None = None + warehouse_id: str | None = None + metrics: dict[str, Any] = Field(default_factory=dict) + + +class MetricAggregate(BaseModel): + """One canonical measure aggregated over a matched query population. + + A subtotal over only the queries that reported a metric is an *observed + subtotal*, never a complete total. ``observed_query_count`` is the matched + query population; ``counted_query_count`` is how many reported the metric. + """ + + metric: str + unit: MetricUnit + source: str + source_version: str + aggregation: str + observed_query_count: int + counted_query_count: int + value: float | None + value_exact: int | None = None + completeness: MetricCompleteness + detail: str | None = None + + +class QueryMetrics(BaseModel): + discovered_query_count: int + matched_query_count: int + attribution_counts: dict[str, int] = Field(default_factory=dict) + aggregates: list[MetricAggregate] = Field(default_factory=list) + discovery_complete: bool = False + denied: bool = False + observations: list[QueryObservation] = Field(default_factory=list) + unavailable_measures: list[str] = Field(default_factory=list) + + +class ReportDiagnostic(BaseModel): + code: str + message: str + severity: str = "warning" + scope: str | None = None + details: dict[str, Any] = Field(default_factory=dict) + + +class ReportObservation(BaseModel): + """A measured fact plus an optional proposed cause, kept distinct.""" + + kind: str + summary: str + evidence: dict[str, Any] = Field(default_factory=dict) + hypothesis: str | None = None + + +class RunReport(BaseModel): + """Versioned, portable workflow-run report envelope.""" + + schema_version: str = DATABRICKS_REPORT_VERSION + collection_status: CollectionStatus = CollectionStatus.complete + collected_at: datetime.datetime + collection_time_ms: int | None = None + identity: WorkflowRunIdentity + status: WorkflowRunStatus + timing: WorkflowTiming + tasks: list[WorkflowTask] = Field(default_factory=list) + query_metrics: QueryMetrics | None = None + compute: list[ComputeReference] = Field(default_factory=list) + effective_performance_target: str | None = None + environments: list[RunEnvironment] = Field(default_factory=list) + job_parameters: dict[str, str] = Field(default_factory=dict) + revision: RevisionEvidence + coverage: dict[str, str] = Field(default_factory=dict) + limits: WorkflowCollectionLimits + diagnostics: list[ReportDiagnostic] = Field(default_factory=list) + observations: list[ReportObservation] = Field(default_factory=list) + next_checks: list[str] = Field(default_factory=list) + outputs: list[OutputAttachment] = Field(default_factory=list) + + +class TrialMetadata(BaseModel): + """The experiment-specific metadata attached to one collected snapshot.""" + + variant: str + revision: str | None = None + artifact_digest: str | None = None + config_fingerprint: str | None = None + config_fingerprint_derived: bool = False + input_snapshot: str | None = None + warmup: bool = False + correctness: str | None = None + hypothesis: str | None = None + notes: str | None = None + workspace_host: str | None = None + job_id: str | None = None + + +class TrialSnapshot(BaseModel): + """A report plus its experiment context, stored immutably on disk.""" + + schema_version: str = EXPERIMENT_SCHEMA_VERSION + snapshot_id: str + collected_at: datetime.datetime + run_id: str + workspace_host: str | None = None + job_id: str | None = None + trial: TrialMetadata + report: RunReport + + +class TrialRef(BaseModel): + trial_id: str + variant: str + run_id: str + snapshot_id: str + revision: str | None = None + config_fingerprint: str | None = None + config_fingerprint_derived: bool = False + input_snapshot: str | None = None + workspace_host: str | None = None + job_id: str | None = None + created_at: datetime.datetime + warmup: bool = False + correctness: str | None = None + + +class ExperimentManifest(BaseModel): + schema_version: str = EXPERIMENT_SCHEMA_VERSION + experiment_id: str + name: str | None = None + objective: str | None = None + workspace_host: str | None = None + job_id: str | None = None + baseline_trial_id: str | None = None + created_at: datetime.datetime + updated_at: datetime.datetime + trials: list[TrialRef] = Field(default_factory=list) + + +class MetricDelta(BaseModel): + metric: str + unit: MetricUnit + aggregation: str + baseline_values: list[float] + candidate_values: list[float] + baseline_partial_values: list[float] = Field(default_factory=list) + candidate_partial_values: list[float] = Field(default_factory=list) + baseline_median: float | None = None + candidate_median: float | None = None + baseline_min: float | None = None + baseline_max: float | None = None + candidate_min: float | None = None + candidate_max: float | None = None + delta: float | None = None + pct_change: float | None = None + single_observation: bool = False + partial: bool = False + caveat: str | None = None + + +class TaskDelta(BaseModel): + task_key: str + change: str + baseline: dict[str, float | None] = Field(default_factory=dict) + candidate: dict[str, float | None] = Field(default_factory=dict) + deltas: dict[str, float | None] = Field(default_factory=dict) + baseline_values: dict[str, list[float]] = Field(default_factory=dict) + candidate_values: dict[str, list[float]] = Field(default_factory=dict) + baseline_min: dict[str, float | None] = Field(default_factory=dict) + baseline_max: dict[str, float | None] = Field(default_factory=dict) + candidate_min: dict[str, float | None] = Field(default_factory=dict) + candidate_max: dict[str, float | None] = Field(default_factory=dict) + + +class ComparisonGroup(BaseModel): + label: str + variant: str | None = None + trial_ids: list[str] = Field(default_factory=list) + run_ids: list[str] = Field(default_factory=list) + snapshot_ids: list[str] = Field(default_factory=list) + sample_count: int = 0 + excluded_failed: int = 0 + excluded_active: int = 0 + excluded_warmup: int = 0 + excluded_incomplete: int = 0 + + +class ComparabilityNote(BaseModel): + aspect: str + status: str + detail: str | None = None + + +class ExperimentComparison(BaseModel): + schema_version: str = EXPERIMENT_SCHEMA_VERSION + experiment_id: str + baseline: ComparisonGroup + candidate: ComparisonGroup + metrics: list[MetricDelta] = Field(default_factory=list) + tasks: list[TaskDelta] = Field(default_factory=list) + comparability: list[ComparabilityNote] = Field(default_factory=list) + revision_confidence: str = "unknown" + notes: list[str] = Field(default_factory=list) diff --git a/sparkparse/runreport.py b/sparkparse/runreport.py new file mode 100644 index 0000000..dcd8b3a --- /dev/null +++ b/sparkparse/runreport.py @@ -0,0 +1,963 @@ +"""Pure normalization of raw Databricks responses into a :class:`RunReport`. + +Nothing here spawns a process or touches the network; it only transforms the +raw dictionaries produced by :mod:`sparkparse.databricks`. That keeps the +report contract testable offline and independent of the CLI. + +Measure semantics are preserved per source. Workflow elapsed time is never +derived from summed task durations, and a Query History subtotal over only the +queries that reported a metric is labelled an *observed subtotal*, never a +complete total. +""" + +from __future__ import annotations + +import datetime +import re +import statistics +from dataclasses import dataclass +from typing import Any + +from sparkparse.databricks import RawRunCollection +from sparkparse.models import ( + ComputeReference, + MetricAggregate, + MetricCompleteness, + MetricUnit, + OutputAttachment, + QueryMetrics, + QueryObservation, + ReportDiagnostic, + ReportObservation, + RevisionEvidence, + RunEnvironment, + RunReport, + TaskTiming, + WorkflowRunIdentity, + WorkflowRunStatus, + WorkflowTask, + WorkflowTiming, +) + +QUERY_METRIC_SOURCE_VERSION = "query_history.v1" +MAX_OUTPUT_CHARS = 4000 + +# Configuration metadata is captured only when its key is not credential-like. +# Keys are normalized so separators (fs.s3a.secret.key, api-token) still match. +_SENSITIVE_KEY_MARKERS = ( + "password", + "passwd", + "secret", + "token", + "credential", + "access_key", + "private_key", + "api_key", + "apikey", + "connection_string", + "client_secret", + "authorization", + "sas", +) + + +def is_sensitive_key(name: str) -> bool: + normalized = re.sub(r"[^a-z0-9]+", "_", name.lower()).strip("_") + return any(marker in normalized for marker in _SENSITIVE_KEY_MARKERS) + + +SHORT_EXCERPT_CHARS = 280 + + +@dataclass(frozen=True) +class _MetricSpec: + canonical: str + raw: str + unit: MetricUnit + aggregation: str + detail: str | None = None + + +QUERY_METRIC_SPECS: tuple[_MetricSpec, ...] = ( + _MetricSpec("read_bytes", "read_bytes", MetricUnit.bytes, "sum"), + _MetricSpec("read_remote_bytes", "read_remote_bytes", MetricUnit.bytes, "sum"), + _MetricSpec("read_cache_bytes", "read_cache_bytes", MetricUnit.bytes, "sum"), + _MetricSpec( + "written_bytes_remote", + "write_remote_bytes", + MetricUnit.bytes, + "sum", + "remote persistent bytes written", + ), + _MetricSpec("disk_spill_bytes", "spill_to_disk_bytes", MetricUnit.bytes, "sum"), + _MetricSpec( + "aggregate_task_time_ms", + "task_total_time_ms", + MetricUnit.milliseconds, + "sum", + "sum of query task execution time, not elapsed", + ), + _MetricSpec( + "query_execution_ms", "execution_time_ms", MetricUnit.milliseconds, "sum" + ), + _MetricSpec("query_total_ms", "total_time_ms", MetricUnit.milliseconds, "sum"), + _MetricSpec("rows_read", "rows_read_count", MetricUnit.rows, "sum"), + _MetricSpec("rows_produced", "rows_produced_count", MetricUnit.rows, "sum"), + _MetricSpec( + "shuffle_network_bytes", + "network_sent_bytes", + MetricUnit.bytes, + "sum", + ), + _MetricSpec("read_files", "read_files_count", MetricUnit.items, "sum"), + _MetricSpec("read_partitions", "read_partitions_count", MetricUnit.items, "sum"), +) + +_TASK_KIND_KEYS = ( + "notebook_task", + "spark_python_task", + "python_wheel_task", + "sql_task", + "dbt_task", + "pipeline_task", + "run_job_task", + "for_each_task", + "condition_task", + "spark_jar_task", + "spark_submit_task", +) + + +def _first_present(mapping: dict[str, Any], *keys: str) -> Any: + for key in keys: + value = mapping.get(key) + if value is not None: + return value + return None + + +def _string_map(mapping: Any) -> dict[str, str]: + if not isinstance(mapping, dict): + return {} + return { + str(key): "" if value is None else str(value) + for key, value in mapping.items() + if not is_sensitive_key(str(key)) + } + + +def normalize_job_parameters(run: dict[str, Any]) -> dict[str, str]: + """Flatten job parameters into a name/value map for configuration identity. + + Parameter values select runtime behavior (for example shuffle partitions), so + two runs that differ only here are two configurations, not one. + """ + parameters: dict[str, str] = {} + for entry in run.get("job_parameters") or []: + if not isinstance(entry, dict) or entry.get("name") is None: + continue + name = str(entry["name"]) + if is_sensitive_key(name): + continue + value = entry.get("value") + parameters[name] = "" if value is None else str(value) + return parameters + + +def _task_parameters(task: dict[str, Any]) -> dict[str, str]: + parameters: dict[str, str] = {} + for kind in _TASK_KIND_KEYS: + spec = task.get(kind) + if not isinstance(spec, dict): + continue + base = spec.get("base_parameters") + if isinstance(base, dict): + parameters.update(_string_map(base)) + args = spec.get("parameters") + if isinstance(args, dict): + parameters.update(_string_map(args)) + elif isinstance(args, list): + parameters[f"{kind}.parameters"] = _sanitize_arguments(args) + return parameters + + +def _sanitize_arguments(args: list[Any]) -> str: + """Join an argument list, redacting values of recognizable secret flags. + + Only flags whose name is credential-like are treated as secrets; an opaque + positional value under an arbitrary name is left alone. + """ + parts: list[str] = [] + redact_next = False + for arg in args: + text = str(arg) + if redact_next: + parts.append("") + redact_next = False + continue + flag, separator, value = text.partition("=") + if flag.startswith("-") and is_sensitive_key(flag.lstrip("-")): + if separator: + parts.append(f"{flag}=") + else: + parts.append(text) + redact_next = True + continue + parts.append(text) + return ",".join(parts) + + +def normalize_identity( + run: dict[str, Any], workspace_host: str | None +) -> WorkflowRunIdentity: + run_page_url = run.get("run_page_url") + return WorkflowRunIdentity( + run_id=str(run.get("run_id")), + job_id=_as_str(run.get("job_id")), + run_name=run.get("run_name"), + run_type=run.get("run_type"), + number_in_job=run.get("number_in_job"), + original_attempt_run_id=_as_str(run.get("original_attempt_run_id")), + run_page_url=run_page_url, + workspace_host=workspace_host or _host_from_url(run_page_url), + creator_user_name=run.get("creator_user_name"), + ) + + +def normalize_status(run: dict[str, Any]) -> WorkflowRunStatus: + state = run.get("state") or {} + termination = (run.get("status") or {}).get("termination_details") or {} + return WorkflowRunStatus( + life_cycle_state=state.get("life_cycle_state"), + result_state=state.get("result_state"), + state_message=state.get("state_message"), + termination_code=termination.get("code"), + termination_type=termination.get("type"), + user_cancelled_or_timedout=state.get("user_cancelled_or_timedout"), + ) + + +def normalize_timing( + run: dict[str, Any], tasks: list[dict[str, Any]] +) -> WorkflowTiming: + summed = [t.get("execution_duration") for t in tasks] + measured = [int(v) for v in summed if isinstance(v, int | float)] + return WorkflowTiming( + workflow_elapsed_ms=_as_int(run.get("run_duration")), + setup_ms=_as_int(run.get("setup_duration")), + execution_ms=_as_int(run.get("execution_duration")), + cleanup_ms=_as_int(run.get("cleanup_duration")), + summed_task_execution_ms=sum(measured) if measured else None, + ) + + +def normalize_tasks( + raw: RawRunCollection, +) -> tuple[list[WorkflowTask], list[ComputeReference]]: + tasks: list[WorkflowTask] = [] + compute: list[ComputeReference] = [] + for task in raw.tasks: + state = task.get("state") or {} + task_key = str(task.get("task_key", "unknown")) + kind_key = next((k for k in _TASK_KIND_KEYS if task.get(k) is not None), None) + reference = _compute_reference(task_key, task, raw) + if reference is not None: + compute.append(reference) + attachment = _attachment_for(task, raw) + git_source = task.get("git_source") + git_commit = None + if isinstance(git_source, dict): + git_commit = _as_str( + (git_source.get("git_snapshot") or {}).get("used_commit") + ) + tasks.append( + WorkflowTask( + task_key=task_key, + run_id=_as_str(task.get("run_id")), + attempt_number=task.get("attempt_number"), + original_attempt_run_id=_as_str(task.get("original_attempt_run_id")), + run_if=task.get("run_if"), + state=state.get("life_cycle_state"), + result_state=state.get("result_state"), + timing=TaskTiming( + setup_ms=_as_int(task.get("setup_duration")), + execution_ms=_as_int(task.get("execution_duration")), + cleanup_ms=_as_int(task.get("cleanup_duration")), + start_time_ms=_as_int(task.get("start_time")), + end_time_ms=_as_int(task.get("end_time")), + ), + compute=reference, + task_kind=kind_key.removesuffix("_task") if kind_key else None, + is_nested_job=kind_key == "run_job_task", + git_commit=git_commit, + parameters=_task_parameters(task), + spark_conf=_string_map( + (task.get("new_cluster") or {}).get("spark_conf") + ), + output_path=attachment.path, + output_status=attachment.status, + error_excerpt=attachment.excerpt, + ) + ) + return tasks, compute + + +def _compute_reference( + task_key: str, task: dict[str, Any], raw: RawRunCollection +) -> ComputeReference | None: + new_cluster = task.get("new_cluster") + if isinstance(new_cluster, dict): + return ComputeReference( + task_key=task_key, + environment_key=task.get("environment_key"), + cluster_id=_as_str(new_cluster.get("cluster_id")), + source="run_snapshot", + runtime_engine=new_cluster.get("runtime_engine"), + spark_version=new_cluster.get("spark_version"), + node_type_id=new_cluster.get("node_type_id"), + num_workers=new_cluster.get("num_workers"), + data_security_mode=new_cluster.get("data_security_mode"), + performance_target=task.get("effective_performance_target"), + ) + existing_cluster_id = task.get("existing_cluster_id") + if existing_cluster_id is not None: + lookup = raw.cluster_lookups.get(str(existing_cluster_id), {}) + return ComputeReference( + task_key=task_key, + environment_key=task.get("environment_key"), + cluster_id=str(existing_cluster_id), + source="current_lookup", + runtime_engine=lookup.get("runtime_engine"), + spark_version=lookup.get("spark_version"), + node_type_id=( + lookup.get("node_type_id") or lookup.get("driver_node_type_id") + ), + num_workers=lookup.get("num_workers"), + data_security_mode=lookup.get("data_security_mode"), + performance_target=task.get("effective_performance_target"), + ) + if task.get("environment_key") is not None: + return ComputeReference( + task_key=task_key, + environment_key=task.get("environment_key"), + source="run_snapshot", + performance_target=task.get("effective_performance_target"), + ) + return None + + +def _attachment_for(task: dict[str, Any], raw: RawRunCollection) -> OutputAttachment: + task_key = str(task.get("task_key", "unknown")) + task_run_id = task.get("run_id") + if task_run_id is None: + return OutputAttachment(task_key=task_key, status="missing") + response = raw.outputs.get(str(task_run_id)) + if response is None: + return OutputAttachment( + task_key=task_key, run_id=str(task_run_id), status="not_requested" + ) + error = response.get("error") + if error: + return OutputAttachment( + task_key=task_key, + run_id=str(task_run_id), + status="error", + error=str(error), + excerpt=_excerpt(str(error)), + ) + notebook_output = response.get("notebook_output") or {} + result = notebook_output.get("result") + if result is not None: + text, truncated = _bounded(str(result)) + return OutputAttachment( + task_key=task_key, + run_id=str(task_run_id), + status="truncated" + if notebook_output.get("truncated") or truncated + else "saved", + excerpt=text, + ) + logs = response.get("logs") + if logs: + text, truncated = _bounded(str(logs)) + return OutputAttachment( + task_key=task_key, + run_id=str(task_run_id), + status="truncated" if truncated else "saved", + excerpt=text, + ) + return OutputAttachment( + task_key=task_key, run_id=str(task_run_id), status="unsupported" + ) + + +def normalize_revision( + run: dict[str, Any], asserted_revision: str | None, artifact_digest: str | None +) -> RevisionEvidence: + """Build run revision evidence with the scope the evidence actually covers. + + A Git snapshot on the run describes the whole run; one carried only by an + individual task describes that task. The two are never conflated, so a + single task's commit cannot be presented as the run's executed revision. + """ + run_source = run.get("git_source") + run_has_git = isinstance(run_source, dict) + task_sources = [ + task["git_source"] + for task in run.get("tasks") or [] + if isinstance(task.get("git_source"), dict) + ] + sources = ([run_source] if run_has_git else []) + task_sources + + commits = { + str((source.get("git_snapshot") or {}).get("used_commit")) + for source in sources + if (source.get("git_snapshot") or {}).get("used_commit") + } + first = sources[0] if sources else {} + snapshot = first.get("git_snapshot") or {} + confidence = "unknown" + executed_commit: str | None = None + scope: str | None = None + if len(commits) == 1: + executed_commit = next(iter(commits)) + confidence = "executed" + if run_has_git and task_sources: + scope = "mixed" + elif run_has_git: + scope = "run" + else: + scope = "task" + elif len(commits) > 1: + confidence = "mixed" + scope = "mixed" + elif asserted_revision is not None: + confidence = "asserted" + + return RevisionEvidence( + executed_commit=executed_commit, + repository_url=first.get("git_url"), + repository_provider=first.get("git_provider"), + branch=first.get("git_branch"), + commit_timestamp_ms=_as_int(snapshot.get("commit_timestamp")), + scope=scope, + asserted_revision=asserted_revision, + artifact_digest=artifact_digest, + confidence=confidence, + ) + + +def normalize_outputs(raw: RawRunCollection) -> list[OutputAttachment]: + return [_attachment_for(task, raw) for task in raw.tasks] + + +def normalize_environments(run: dict[str, Any]) -> list[RunEnvironment]: + environments: list[RunEnvironment] = [] + for environment in run.get("environments") or []: + if not isinstance(environment, dict): + continue + spec = environment.get("spec") or {} + environments.append( + RunEnvironment( + environment_key=str(environment.get("environment_key", "default")), + client=_as_str(spec.get("client")), + dependencies=[ + str(dependency) + for dependency in spec.get("dependencies") or [] + if dependency is not None + ], + ) + ) + return environments + + +def aggregate_queries(raw: RawRunCollection, tasks: list[WorkflowTask]) -> QueryMetrics: + task_run_to_key = { + task.run_id: task.task_key for task in tasks if task.run_id is not None + } + run_id = str(raw.run.get("run_id")) + attribution_counts: dict[str, int] = {} + observations: dict[str, QueryObservation] = {} + collected_at = datetime.datetime.now(datetime.UTC) + + for query in raw.queries: + source = query.get("query_source") or {} + job_info = source.get("job_info") or {} + task_run_id = job_info.get("job_task_run_id") + parent_run_id = job_info.get("job_run_id") + if task_run_id is not None and str(task_run_id) in task_run_to_key: + attribution = "task_run" + elif parent_run_id is not None and str(parent_run_id) == run_id: + attribution = "run_parent" + elif task_run_id is not None or parent_run_id is not None: + attribution = "unrelated" + else: + attribution = "unattributed" + attribution_counts[attribution] = attribution_counts.get(attribution, 0) + 1 + if attribution not in ("task_run", "run_parent"): + continue + observation = QueryObservation( + query_id=str(query.get("query_id")), + status=query.get("status"), + is_final=query.get("is_final"), + execution_end_time_ms=_as_int(query.get("execution_end_time_ms")), + collected_at=collected_at, + job_id=_as_str(job_info.get("job_id")), + job_run_id=_as_str(parent_run_id), + job_task_run_id=_as_str(task_run_id), + attribution=attribution, + task_key=task_run_to_key.get(str(task_run_id)), + warehouse_id=_as_str(query.get("warehouse_id")), + metrics=_numeric_metrics(query.get("metrics")), + ) + existing = observations.get(observation.query_id) + if existing is None or _observation_is_fresher(observation, existing): + observations[observation.query_id] = observation + + kept = list(observations.values()) + final = bool(kept) and all(observation.is_final is True for observation in kept) + attribution_certain = attribution_counts.get("unattributed", 0) == 0 + aggregates = _aggregate_metrics( + kept, + raw.query_discovery_complete, + observations_final=final, + attribution_certain=attribution_certain, + ) + unavailable = [ + "memory_bytes_spilled", + "per_executor_distributions", + "query_profiles", + ] + return QueryMetrics( + discovered_query_count=raw.discovered_query_count, + matched_query_count=len(kept), + attribution_counts=attribution_counts, + aggregates=aggregates, + discovery_complete=raw.query_discovery_complete, + denied=raw.query_denied, + observations=kept, + unavailable_measures=unavailable, + ) + + +def _aggregate_metrics( + observations: list[QueryObservation], + discovery_complete: bool, + *, + observations_final: bool, + attribution_certain: bool, +) -> list[MetricAggregate]: + total = len(observations) + aggregates: list[MetricAggregate] = [] + for spec in QUERY_METRIC_SPECS: + values: list[int] = [] + for observation in observations: + value = observation.metrics.get(spec.raw) + if isinstance(value, bool): + continue + if isinstance(value, int | float): + values.append(int(value)) + if not values: + completeness = MetricCompleteness.unavailable + value_exact: int | None = None + value: float | None = None + else: + value_exact = sum(values) + value = float(value_exact) + completeness = ( + MetricCompleteness.complete + if len(values) == total + and discovery_complete + and total > 0 + and observations_final + and attribution_certain + else MetricCompleteness.observed_subtotal + ) + aggregates.append( + MetricAggregate( + metric=spec.canonical, + unit=spec.unit, + source="query_history", + source_version=QUERY_METRIC_SOURCE_VERSION, + aggregation=spec.aggregation, + observed_query_count=total, + counted_query_count=len(values), + value=value, + value_exact=value_exact, + completeness=completeness, + detail=spec.detail, + ) + ) + cached = sum( + 1 + for observation in observations + if observation.metrics.get("result_from_cache") + ) + cached_counted = sum( + 1 for observation in observations if "result_from_cache" in observation.metrics + ) + if total == 0 or cached_counted == 0: + cache_value: float | None = None + cache_exact: int | None = None + cache_completeness = MetricCompleteness.unavailable + else: + cache_value = float(cached) + cache_exact = cached + cache_completeness = ( + MetricCompleteness.complete + if cached_counted == total + and discovery_complete + and observations_final + and attribution_certain + else MetricCompleteness.observed_subtotal + ) + aggregates.append( + MetricAggregate( + metric="queries_result_from_cache", + unit=MetricUnit.items, + source="query_history", + source_version=QUERY_METRIC_SOURCE_VERSION, + aggregation="count_true", + observed_query_count=total, + counted_query_count=cached_counted, + value=cache_value, + value_exact=cache_exact, + completeness=cache_completeness, + detail="count of queries whose result was served from cache", + ) + ) + return aggregates + + +def build_report( + raw: RawRunCollection, + *, + asserted_revision: str | None = None, + artifact_digest: str | None = None, + workspace_host: str | None = None, +) -> RunReport: + """Normalize a raw collection into a versioned report envelope.""" + identity = normalize_identity(raw.run, workspace_host) + tasks, compute = normalize_tasks(raw) + query_metrics = aggregate_queries(raw, tasks) + revision = normalize_revision(raw.run, asserted_revision, artifact_digest) + diagnostics = list(raw.diagnostics) + if revision.confidence == "mixed": + diagnostics.append( + ReportDiagnostic( + code="mixed_revision", + message="tasks report different Git commits; no single SHA represents the run", + severity="info", + ) + ) + observations, next_checks = _derive_observations( + tasks, raw, query_metrics, diagnostics + ) + coverage = _coverage(raw, tasks, query_metrics) + return RunReport( + collection_status=raw.collection_status, + collected_at=datetime.datetime.now(datetime.UTC), + collection_time_ms=raw.collection_time_ms, + identity=identity, + status=normalize_status(raw.run), + timing=normalize_timing(raw.run, raw.tasks), + tasks=tasks, + query_metrics=query_metrics, + compute=compute, + effective_performance_target=_as_str( + raw.run.get("effective_performance_target") + ), + environments=normalize_environments(raw.run), + job_parameters=normalize_job_parameters(raw.run), + revision=revision, + coverage=coverage, + limits=raw.limits, + diagnostics=diagnostics, + observations=observations, + next_checks=next_checks, + outputs=normalize_outputs(raw), + ) + + +def _coverage( + raw: RawRunCollection, tasks: list[WorkflowTask], query_metrics: QueryMetrics +) -> dict[str, str]: + def status(condition: bool) -> str: + return "available" if condition else "unavailable" + + return { + "jobs": "available" if raw.run else "unavailable", + "tasks": status(bool(tasks)), + "task_list": ("complete" if raw.run_pagination_complete else "partial"), + "task_timing": status(any(t.timing.execution_ms is not None for t in tasks)), + "query_history": ( + "denied" + if query_metrics.denied + else ( + "available" + if query_metrics.matched_query_count + else ("unavailable" if not raw.queries else "unattributed") + ) + ), + "query_metrics": ( + "complete" + if query_metrics.discovery_complete and query_metrics.matched_query_count + else "partial" + if query_metrics.matched_query_count + else "unavailable" + ), + "git_revision": ( + "available" + if raw.run.get("git_source") or any(t.git_commit for t in tasks) + else "unknown" + ), + "compute": status(bool(raw.run.get("tasks"))), + "child_tasks": ( + "incomplete" + if any(t.is_nested_job or t.task_kind == "for_each" for t in tasks) + else "not_applicable" + ), + "memory_spill": "unavailable", + } + + +def _derive_observations( + tasks: list[WorkflowTask], + raw: RawRunCollection, + query_metrics: QueryMetrics, + diagnostics: list[ReportDiagnostic], +) -> tuple[list[ReportObservation], list[str]]: + observations: list[ReportObservation] = [] + next_checks: list[str] = [] + + failed = [t for t in tasks if t.result_state in ("FAILED", "TIMEDOUT", "CANCELED")] + if failed: + observations.append( + ReportObservation( + kind="failed_tasks", + summary=f"{len(failed)} task(s) did not succeed", + evidence={"task_keys": [t.task_key for t in failed]}, + ) + ) + next_checks.append( + "Inspect the failed task output excerpts for the root error." + ) + + nested = [t for t in tasks if t.is_nested_job or t.task_kind == "for_each"] + if nested: + diagnostics.append( + ReportDiagnostic( + code="incomplete_child_coverage", + message=( + f"{len(nested)} for-each/nested-job task(s) were not unrolled in v1; " + "their child runs are counted as one task and excluded from " + "aggregates that require child detail" + ), + severity="info", + ) + ) + + executions = [t.timing.execution_ms for t in tasks if t.timing.execution_ms] + if len(executions) >= 3: + median = statistics.median(executions) + slow = sorted( + (t for t in tasks if (t.timing.execution_ms or 0) > 2 * median), + key=lambda t: t.timing.execution_ms or 0, + reverse=True, + ) + if slow: + observations.append( + ReportObservation( + kind="slow_tasks", + summary=f"{len(slow)} task(s) exceeded twice the median execution time", + evidence={ + "median_execution_ms": median, + "tasks": {t.task_key: t.timing.execution_ms for t in slow[:5]}, + }, + hypothesis="The slow tasks may reflect data skew or an under-provisioned stage.", + ) + ) + next_checks.append( + "Compare the slow tasks' input sizes and partition counts." + ) + + spill = _aggregate_value(query_metrics, "disk_spill_bytes") + if spill: + observations.append( + ReportObservation( + kind="disk_spill", + summary="Query History reports bytes temporarily spilled to disk", + evidence={"disk_spill_bytes": spill}, + hypothesis="Spill may indicate insufficient memory or shuffle skew.", + ) + ) + next_checks.append( + "Open the query profile for the spilling query to find the operator." + ) + + read_bytes = _aggregate_value(query_metrics, "read_bytes") + if read_bytes: + observations.append( + ReportObservation( + kind="read_bytes", + summary="Query History reports bytes read", + evidence={"read_bytes": read_bytes}, + ) + ) + + if query_metrics.attribution_counts.get("unattributed"): + diagnostics.append( + ReportDiagnostic( + code="unattributed_queries", + message=( + f"{query_metrics.attribution_counts['unattributed']} query/queries in the " + "run window had no source identity and were not attributed" + ), + severity="info", + ) + ) + if not query_metrics.discovery_complete: + diagnostics.append( + ReportDiagnostic( + code="query_discovery_incomplete", + message="Query History pagination did not complete; metrics may be partial", + severity="warning", + ) + ) + return observations, next_checks + + +def _aggregate_value(metrics: QueryMetrics, canonical: str) -> int | None: + for aggregate in metrics.aggregates: + if aggregate.metric == canonical and aggregate.value_exact is not None: + return aggregate.value_exact + return None + + +def _numeric_metrics(metrics: Any) -> dict[str, Any]: + if not isinstance(metrics, dict): + return {} + return { + key: value + for key, value in metrics.items() + if isinstance(value, int | float | bool) and not isinstance(value, dict) + } + + +def _observation_is_fresher( + candidate: QueryObservation, existing: QueryObservation +) -> bool: + if candidate.is_final and not existing.is_final: + return True + candidate_end = candidate.execution_end_time_ms or 0 + existing_end = existing.execution_end_time_ms or 0 + return candidate_end >= existing_end + + +def _bounded(text: str) -> tuple[str, bool]: + if len(text) <= MAX_OUTPUT_CHARS: + return text, False + return text[:MAX_OUTPUT_CHARS], True + + +def _excerpt(text: str) -> str: + text = " ".join(text.split()) + if len(text) <= SHORT_EXCERPT_CHARS: + return text + return text[:SHORT_EXCERPT_CHARS] + "…" + + +def _as_int(value: Any) -> int | None: + if isinstance(value, bool): + return None + if isinstance(value, int | float): + return int(value) + return None + + +def _as_str(value: Any) -> str | None: + if value is None: + return None + return str(value) + + +def _host_from_url(url: str | None) -> str | None: + if not url: + return None + return url.split("//", 1)[-1].split("/", 1)[0] + + +def format_report_text(report: RunReport) -> str: + """Human-readable rendering of the same facts the JSON envelope carries.""" + lines: list[str] = [] + identity = report.identity + lines.append( + f"Run {identity.run_id} job={identity.job_id} {identity.run_name or ''}" + ) + status = report.status + lines.append( + f"Status: {status.life_cycle_state or '?'} / {status.result_state or '?'}" + f" collection={report.collection_status.value}" + ) + timing = report.timing + lines.append( + "Timing: " + f"elapsed={_fmt_ms(timing.workflow_elapsed_ms)} " + f"setup={_fmt_ms(timing.setup_ms)} " + f"execution={_fmt_ms(timing.execution_ms)} " + f"cleanup={_fmt_ms(timing.cleanup_ms)} " + f"summed_task_execution={_fmt_ms(timing.summed_task_execution_ms)}" + ) + lines.append( + f"Revision: {report.revision.confidence} {report.revision.executed_commit or ''}" + ) + lines.append(f"Tasks ({len(report.tasks)}):") + for task in report.tasks: + lines.append( + f" {task.task_key}\t{task.result_state or task.state or '?'}" + f"\texecution={_fmt_ms(task.timing.execution_ms)}" + f"\tsetup={_fmt_ms(task.timing.setup_ms)}" + + (f"\t{task.error_excerpt}" if task.error_excerpt else "") + ) + if report.query_metrics is not None: + metrics = report.query_metrics + lines.append( + f"Queries: discovered={metrics.discovered_query_count} " + f"matched={metrics.matched_query_count} " + f"discovery_complete={metrics.discovery_complete} " + f"denied={metrics.denied}" + ) + for aggregate in metrics.aggregates: + if aggregate.value is None: + continue + lines.append( + f" {aggregate.metric}={aggregate.value:g} {aggregate.unit.value}" + f" [{aggregate.completeness.value}" + f" {aggregate.counted_query_count}/{aggregate.observed_query_count}]" + ) + if report.observations: + lines.append("Observations:") + for observation in report.observations: + lines.append(f" [{observation.kind}] {observation.summary}") + if observation.hypothesis: + lines.append(f" hypothesis: {observation.hypothesis}") + if report.next_checks: + lines.append("Next checks:") + for check in report.next_checks: + lines.append(f" - {check}") + if report.diagnostics: + lines.append("Diagnostics:") + for diagnostic in report.diagnostics: + lines.append( + f" [{diagnostic.severity}] {diagnostic.code}: {diagnostic.message}" + ) + lines.append( + f"Coverage: {', '.join(f'{k}={v}' for k, v in sorted(report.coverage.items()))}" + ) + return "\n".join(lines) + + +def _fmt_ms(value: int | None) -> str: + if value is None: + return "n/a" + return f"{value}ms" diff --git a/tests/data/databricks/job_run_multi.json b/tests/data/databricks/job_run_multi.json new file mode 100644 index 0000000..54b498a --- /dev/null +++ b/tests/data/databricks/job_run_multi.json @@ -0,0 +1,198 @@ +{ + "job_id": 456, + "run_id": 123, + "job_run_id": 123, + "number_in_job": 7, + "original_attempt_run_id": 122, + "run_name": "join-tuning", + "run_type": "JOB_RUN", + "run_page_url": "https://example.cloud.databricks.com/?o=0000000000000000#job/456/run/123", + "creator_user_name": "user@example.com", + "trigger": "ONE_TIME", + "start_time": 1700000000000, + "end_time": 1700000060000, + "setup_duration": 4000, + "execution_duration": 42000, + "cleanup_duration": 1000, + "run_duration": 60000, + "effective_performance_target": "PERFORMANCE_OPTIMIZED", + "git_source": { + "git_url": "https://github.com/example/analytics", + "git_provider": "gitHub", + "git_branch": "main", + "git_snapshot": { + "used_commit": "abc123def4567890abc123def4567890abc1234", + "commit_timestamp": 1699990000000 + } + }, + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "SUCCESS_WITH_FAILURES", + "state_message": "", + "user_cancelled_or_timedout": false + }, + "status": { + "state": "TERMINATED", + "termination_details": { + "code": "SUCCESS_WITH_FAILURES", + "type": "SUCCESS" + } + }, + "environments": [ + { + "environment_key": "default", + "spec": { + "client": "2", + "dependencies": [ + "polars>=1.24.0" + ] + } + } + ], + "tasks": [ + { + "task_key": "ingest", + "run_id": 1001, + "attempt_number": 0, + "run_if": "ALL_SUCCESS", + "environment_key": "default", + "start_time": 1700000000000, + "end_time": 1700000010000, + "setup_duration": 1000, + "execution_duration": 8000, + "cleanup_duration": 0, + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "SUCCESS", + "state_message": "" + }, + "status": { + "state": "TERMINATED", + "termination_details": { + "code": "SUCCESS", + "type": "SUCCESS" + } + }, + "new_cluster": { + "cluster_id": "0101-010101-abcd", + "spark_version": "15.4.x-scala2.12", + "node_type_id": "i3.xlarge", + "num_workers": 4, + "runtime_engine": "PHOTON", + "data_security_mode": "SINGLE_USER" + }, + "notebook_task": { + "notebook_path": "/Jobs/ingest", + "source": "WORKSPACE" + } + }, + { + "task_key": "join", + "run_id": 1002, + "attempt_number": 1, + "original_attempt_run_id": 1000, + "run_if": "ALL_SUCCESS", + "environment_key": "default", + "start_time": 1700000010000, + "end_time": 1700000050000, + "setup_duration": 1000, + "execution_duration": 38000, + "cleanup_duration": 0, + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "SUCCESS", + "state_message": "" + }, + "status": { + "state": "TERMINATED", + "termination_details": { + "code": "SUCCESS", + "type": "SUCCESS" + } + }, + "new_cluster": { + "cluster_id": "0101-010101-abcd", + "spark_version": "15.4.x-scala2.12", + "node_type_id": "i3.xlarge", + "runtime_engine": "PHOTON" + }, + "notebook_task": { + "notebook_path": "/Jobs/join", + "source": "WORKSPACE" + } + }, + { + "task_key": "optimize", + "run_id": 1003, + "attempt_number": 0, + "run_if": "ALL_DONE", + "environment_key": "default", + "start_time": 1700000050000, + "end_time": 1700000059000, + "setup_duration": 0, + "execution_duration": 9000, + "cleanup_duration": 0, + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "SKIPPED", + "state_message": "" + }, + "status": { + "state": "TERMINATED", + "termination_details": { + "code": "SKIPPED", + "type": "SUCCESS" + } + }, + "notebook_task": { + "notebook_path": "/Jobs/optimize", + "source": "WORKSPACE" + } + }, + { + "task_key": "foreach_copy", + "run_id": 1004, + "attempt_number": 0, + "run_if": "ALL_DONE", + "environment_key": "default", + "start_time": 1700000059000, + "end_time": 1700000060000, + "execution_duration": 1000, + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "SUCCESS", + "state_message": "" + }, + "status": { + "state": "TERMINATED", + "termination_details": { + "code": "SUCCESS", + "type": "SUCCESS" + } + }, + "for_each_task": { + "inputs": "[{\"table\":\"a\"},{\"table\":\"b\"}]", + "concurrency": 2 + } + }, + { + "task_key": "unknown_type", + "run_id": 1005, + "attempt_number": 0, + "run_if": "ALL_DONE", + "execution_duration": 500, + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "SUCCESS", + "state_message": "" + }, + "status": { + "state": "TERMINATED", + "termination_details": { + "code": "SUCCESS", + "type": "SUCCESS" + } + } + } + ] +} diff --git a/tests/data/databricks/job_run_single.json b/tests/data/databricks/job_run_single.json new file mode 100644 index 0000000..470e1e1 --- /dev/null +++ b/tests/data/databricks/job_run_single.json @@ -0,0 +1,81 @@ +{ + "cleanup_duration": 0, + "creator_user_name": "user@example.com", + "effective_performance_target": "PERFORMANCE_OPTIMIZED", + "end_time": 1789249110616, + "environments": [ + { + "environment_key": "default", + "spec": { + "client": "2", + "dependencies": [ + "/Workspace/Users/user@example.com/.bundle/sparkparse/dev/files/dist/sparkparse-0.1.0-py3-none-any.whl", + "polars>=1.24.0", + "pydantic>=2.10.6" + ] + } + } + ], + "execution_duration": 36000, + "format": "MULTI_TASK", + "job_id": 240570464625850, + "job_run_id": 915049514672868, + "number_in_job": 915049514672868, + "original_attempt_run_id": 915049514672868, + "run_duration": 41745, + "run_id": 915049514672868, + "run_name": "[dev user] [dev user@example.com] sparkparse-validate-connect", + "run_page_url": "https://example.cloud.databricks.com/?o=0000000000000000#job/240570464625850/run/915049514672868", + "run_type": "JOB_RUN", + "setup_duration": 5000, + "start_time": 1789249068871, + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "SUCCESS", + "state_message": "", + "user_cancelled_or_timedout": false + }, + "status": { + "state": "TERMINATED", + "termination_details": { + "code": "SUCCESS", + "message": "", + "type": "SUCCESS" + } + }, + "tasks": [ + { + "attempt_number": 0, + "cleanup_duration": 0, + "effective_performance_target": "PERFORMANCE_OPTIMIZED", + "end_time": 1789249110191, + "environment_key": "default", + "execution_duration": 36000, + "notebook_task": { + "notebook_path": "/Workspace/Users/user@example.com/.bundle/sparkparse/dev/files/notebooks/validate_connect_capture", + "source": "WORKSPACE" + }, + "run_id": 60603881665866, + "run_if": "ALL_SUCCESS", + "run_page_url": "https://example.cloud.databricks.com/?o=0000000000000000#job/240570464625850/run/60603881665866", + "setup_duration": 5000, + "start_time": 1789249068906, + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "SUCCESS", + "state_message": "", + "user_cancelled_or_timedout": false + }, + "status": { + "state": "TERMINATED", + "termination_details": { + "code": "SUCCESS", + "message": "", + "type": "SUCCESS" + } + }, + "task_key": "validate_connect_capture" + } + ], + "trigger": "ONE_TIME" +} diff --git a/tests/data/databricks/jobs_runs_list.json b/tests/data/databricks/jobs_runs_list.json new file mode 100644 index 0000000..14027e7 --- /dev/null +++ b/tests/data/databricks/jobs_runs_list.json @@ -0,0 +1,109 @@ +{ + "has_more": false, + "next_page_token": "CAEQmsmJvIk0IK-77cG3h08=", + "prev_page_token": "CAAQx764vIk0IOTFmIq5h9AB", + "runs": [ + { + "cleanup_duration": 0, + "creator_user_name": "user@example.com", + "effective_performance_target": "PERFORMANCE_OPTIMIZED", + "end_time": 1789249110616, + "execution_duration": 0, + "format": "MULTI_TASK", + "job_id": 240570464625850, + "job_run_id": 915049514672868, + "number_in_job": 915049514672868, + "original_attempt_run_id": 915049514672868, + "run_duration": 41745, + "run_id": 915049514672868, + "run_name": "[dev user] [dev user@example.com] sparkparse-validate-connect", + "run_page_url": "https://example.cloud.databricks.com/?o=0000000000000000#job/240570464625850/run/915049514672868", + "run_type": "JOB_RUN", + "setup_duration": 0, + "start_time": 1789249068871, + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "SUCCESS", + "state_message": "", + "user_cancelled_or_timedout": false + }, + "status": { + "state": "TERMINATED", + "termination_details": { + "code": "SUCCESS", + "message": "", + "type": "SUCCESS" + } + }, + "trigger": "ONE_TIME" + }, + { + "cleanup_duration": 0, + "creator_user_name": "user@example.com", + "effective_performance_target": "PERFORMANCE_OPTIMIZED", + "end_time": 1789248802671, + "execution_duration": 0, + "format": "MULTI_TASK", + "job_id": 240570464625850, + "job_run_id": 786861458138844, + "number_in_job": 786861458138844, + "original_attempt_run_id": 786861458138844, + "run_duration": 84980, + "run_id": 786861458138844, + "run_name": "[dev user] [dev user@example.com] sparkparse-validate-connect", + "run_page_url": "https://example.cloud.databricks.com/?o=0000000000000000#job/240570464625850/run/786861458138844", + "run_type": "JOB_RUN", + "setup_duration": 0, + "start_time": 1789248717691, + "state": { + "life_cycle_state": "INTERNAL_ERROR", + "result_state": "FAILED", + "state_message": "Task validate_connect_capture failed with message: Workload failed, see run output for details.", + "user_cancelled_or_timedout": false + }, + "status": { + "state": "TERMINATED", + "termination_details": { + "code": "RUN_EXECUTION_ERROR", + "message": "Task validate_connect_capture failed with message: Workload failed, see run output for details.", + "type": "CLIENT_ERROR" + } + }, + "trigger": "ONE_TIME" + }, + { + "cleanup_duration": 0, + "creator_user_name": "user@example.com", + "effective_performance_target": "PERFORMANCE_OPTIMIZED", + "end_time": 1789248423328, + "execution_duration": 0, + "format": "MULTI_TASK", + "job_id": 240570464625850, + "job_run_id": 347701094604207, + "number_in_job": 347701094604207, + "original_attempt_run_id": 347701094604207, + "run_duration": 123142, + "run_id": 347701094604207, + "run_name": "[dev user] [dev user@example.com] sparkparse-validate-connect", + "run_page_url": "https://example.cloud.databricks.com/?o=0000000000000000#job/240570464625850/run/347701094604207", + "run_type": "JOB_RUN", + "setup_duration": 0, + "start_time": 1789248300186, + "state": { + "life_cycle_state": "INTERNAL_ERROR", + "result_state": "FAILED", + "state_message": "Task validate_connect_capture failed with message: Workload failed, see run output for details.", + "user_cancelled_or_timedout": false + }, + "status": { + "state": "TERMINATED", + "termination_details": { + "code": "RUN_EXECUTION_ERROR", + "message": "Task validate_connect_capture failed with message: Workload failed, see run output for details.", + "type": "CLIENT_ERROR" + } + }, + "trigger": "ONE_TIME" + } + ] +} diff --git a/tests/data/databricks/query_history_mixed.json b/tests/data/databricks/query_history_mixed.json new file mode 100644 index 0000000..e393cdb --- /dev/null +++ b/tests/data/databricks/query_history_mixed.json @@ -0,0 +1,75 @@ +{ + "has_next_page": false, + "res": [ + { + "query_id": "aaaa1111-0000-0000-0000-000000000001", + "status": "FINISHED", + "execution_end_time_ms": 1700000040000, + "query_source": { + "job_info": { + "job_id": "456", + "job_task_run_id": "1002" + }, + "notebook_id": "1234567890" + }, + "metrics": { + "read_bytes": 1000, + "read_cache_bytes": 400, + "write_remote_bytes": 200, + "spill_to_disk_bytes": 0, + "task_total_time_ms": 5000, + "execution_time_ms": 6000, + "total_time_ms": 7000, + "rows_read_count": 100, + "rows_produced_count": 50, + "network_sent_bytes": 300, + "result_from_cache": true + } + }, + { + "query_id": "aaaa1111-0000-0000-0000-000000000002", + "status": "FINISHED", + "execution_end_time_ms": 1700000045000, + "query_source": { + "job_info": { + "job_task_run_id": "1002" + } + }, + "metrics": { + "read_bytes": 2000, + "spill_to_disk_bytes": 900, + "task_total_time_ms": 8000, + "execution_time_ms": 9000, + "total_time_ms": 10000, + "rows_read_count": 200, + "rows_produced_count": 80, + "network_sent_bytes": 500 + } + }, + { + "query_id": "aaaa1111-0000-0000-0000-000000000003", + "status": "FAILED", + "execution_end_time_ms": 1700000048000, + "query_source": { + "notebook_id": "1234567890" + }, + "error_message": "AnalysisException: Table not found: missing_table", + "metrics": {} + }, + { + "query_id": "unrelated000-0000-0000-0000-000000000004", + "status": "FINISHED", + "execution_end_time_ms": 1700000049000, + "query_source": { + "job_info": { + "job_id": "999", + "job_task_run_id": "9999" + } + }, + "metrics": { + "read_bytes": 999999, + "task_total_time_ms": 999999 + } + } + ] +} diff --git a/tests/data/databricks/query_history_window.json b/tests/data/databricks/query_history_window.json new file mode 100644 index 0000000..bf78c16 --- /dev/null +++ b/tests/data/databricks/query_history_window.json @@ -0,0 +1,661 @@ +{ + "has_next_page": false, + "res": [ + { + "canSubscribeToLiveQuery": true, + "channel_used": { + "dbsql_version": "", + "name": "CHANNEL_NAME_UNSPECIFIED" + }, + "client_application": "Databricks Jobs", + "duration": 240, + "endpoint_id": "0000000000000000", + "executed_as_user_display_name": "", + "executed_as_user_id": 0, + "executed_as_user_name": "user@example.com", + "execution_end_time_ms": 1789249106069, + "is_cancelable": false, + "is_final": true, + "last_updated_at_ms": 1789249106119, + "lookup_key": "CiRmNDhjMTNjNy1kMTRjLTRlMTEtYjcxMS1lYWFlNGI0YmZlMzYQpd+6vIk0GhAwMDAwMDAwMDAwMDAwMDAwIJjxne2Tgd4B", + "metrics": { + "bytes_read_from_cache_percentage": 0, + "compilation_time_ms": 52, + "execution_time_ms": 188, + "metadata_time_ms": 0, + "network_sent_bytes": 0, + "photon_total_time_ms": 3, + "pruned_files_count": 0, + "query_compilation_start_timestamp": 1789249105829, + "query_execution_time_ms": 240, + "read_bytes": 0, + "read_cache_bytes": 0, + "read_files_count": 0, + "read_partitions_count": 0, + "read_remote_bytes": 0, + "result_fetch_time_ms": 49, + "result_from_cache": false, + "rows_produced_count": 1, + "rows_read_count": 0, + "spill_to_disk_bytes": 0, + "task_progress": { + "active_task_count": 0, + "completed_partition_count": 9, + "completed_task_count": 9, + "failed_task_count": 0, + "pending_task_count": 0, + "planned_task_count": 9 + }, + "task_total_time_ms": 43, + "total_time_ms": 240, + "write_remote_bytes": 0, + "write_remote_files": 0, + "write_remote_rows": 0 + }, + "plans_state": "EXISTS", + "query_end_time_ms": 1789249106069, + "query_id": "f48c13c7-d14c-4e11-b711-eaae4b4bfe36", + "query_source": { + "job_info": { + "job_id": "240570464625850", + "job_task_run_id": "60603881665866" + }, + "notebook_id": "1234567890" + }, + "query_start_time_ms": 1789249105829, + "query_text": "", + "rows_produced": 1, + "session_id": "dc9ff815-f6ed-4952-87fe-4c96a07cd988", + "statement_type": "SELECT", + "status": "FINISHED", + "user_display_name": "", + "user_id": 0, + "user_name": "user@example.com", + "warehouse_id": "0000000000000000", + "workspace_id": 0 + }, + { + "canSubscribeToLiveQuery": true, + "channel_used": { + "dbsql_version": "", + "name": "CHANNEL_NAME_UNSPECIFIED" + }, + "client_application": "Databricks Jobs", + "duration": 250, + "endpoint_id": "0000000000000000", + "executed_as_user_display_name": "", + "executed_as_user_id": 0, + "executed_as_user_name": "user@example.com", + "execution_end_time_ms": 1789249105745, + "is_cancelable": false, + "is_final": true, + "last_updated_at_ms": 1789249105804, + "lookup_key": "CiRmYzk4MDdlZi02MTQ4LTQ2M2YtODE2NC1mOGM3YjhiMGQxNGYQ19y6vIk0GhAwMDAwMDAwMDAwMDAwMDAwIJjxne2Tgd4B", + "metrics": { + "bytes_read_from_cache_percentage": 0, + "compilation_time_ms": 55, + "execution_time_ms": 195, + "metadata_time_ms": 0, + "network_sent_bytes": 0, + "photon_total_time_ms": 2, + "pruned_files_count": 0, + "query_compilation_start_timestamp": 1789249105495, + "query_execution_time_ms": 250, + "read_bytes": 0, + "read_cache_bytes": 0, + "read_files_count": 0, + "read_partitions_count": 0, + "read_remote_bytes": 0, + "result_fetch_time_ms": 58, + "result_from_cache": false, + "rows_produced_count": 1, + "rows_read_count": 0, + "spill_to_disk_bytes": 0, + "task_progress": { + "active_task_count": 0, + "completed_partition_count": 9, + "completed_task_count": 9, + "failed_task_count": 0, + "pending_task_count": 0, + "planned_task_count": 9 + }, + "task_total_time_ms": 36, + "total_time_ms": 250, + "write_remote_bytes": 0, + "write_remote_files": 0, + "write_remote_rows": 0 + }, + "plans_state": "EXISTS", + "query_end_time_ms": 1789249105745, + "query_id": "fc9807ef-6148-463f-8164-f8c7b8b0d14f", + "query_source": { + "job_info": { + "job_id": "240570464625850", + "job_task_run_id": "60603881665866" + }, + "notebook_id": "1234567890" + }, + "query_start_time_ms": 1789249105495, + "query_text": "", + "rows_produced": 1, + "session_id": "dc9ff815-f6ed-4952-87fe-4c96a07cd988", + "statement_type": "SELECT", + "status": "FINISHED", + "user_display_name": "", + "user_id": 0, + "user_name": "user@example.com", + "warehouse_id": "0000000000000000", + "workspace_id": 0 + }, + { + "canSubscribeToLiveQuery": true, + "channel_used": { + "dbsql_version": "", + "name": "CHANNEL_NAME_UNSPECIFIED" + }, + "client_application": "Databricks Jobs", + "duration": 447, + "endpoint_id": "0000000000000000", + "executed_as_user_display_name": "", + "executed_as_user_id": 0, + "executed_as_user_name": "user@example.com", + "execution_end_time_ms": 1789249104702, + "is_cancelable": false, + "is_final": true, + "last_updated_at_ms": 1789249104783, + "lookup_key": "CiQ1MWQ4ZDAxOC1kYWYyLTRhZGUtOWE2Zi01N2VkYjQzMjJjNzkQ/9K6vIk0GhAwMDAwMDAwMDAwMDAwMDAwIJjxne2Tgd4B", + "metrics": { + "bytes_read_from_cache_percentage": 0, + "compilation_time_ms": 125, + "execution_time_ms": 322, + "metadata_time_ms": 0, + "network_sent_bytes": 0, + "photon_total_time_ms": 47, + "pruned_files_count": 0, + "query_compilation_start_timestamp": 1789249104255, + "query_execution_time_ms": 447, + "read_bytes": 0, + "read_cache_bytes": 0, + "read_files_count": 0, + "read_partitions_count": 0, + "read_remote_bytes": 0, + "result_fetch_time_ms": 78, + "result_from_cache": false, + "rows_produced_count": 1, + "rows_read_count": 0, + "spill_to_disk_bytes": 0, + "task_progress": { + "active_task_count": 0, + "completed_partition_count": 17, + "completed_task_count": 17, + "failed_task_count": 0, + "pending_task_count": 0, + "planned_task_count": 17 + }, + "task_total_time_ms": 119, + "total_time_ms": 447, + "write_remote_bytes": 0, + "write_remote_files": 0, + "write_remote_rows": 0 + }, + "plans_state": "EXISTS", + "query_end_time_ms": 1789249104702, + "query_id": "51d8d018-daf2-4ade-9a6f-57edb4322c79", + "query_source": { + "job_info": { + "job_id": "240570464625850", + "job_task_run_id": "60603881665866" + }, + "notebook_id": "1234567890" + }, + "query_start_time_ms": 1789249104255, + "query_text": "", + "rows_produced": 1, + "session_id": "dc9ff815-f6ed-4952-87fe-4c96a07cd988", + "statement_type": "SELECT", + "status": "FINISHED", + "user_display_name": "", + "user_id": 0, + "user_name": "user@example.com", + "warehouse_id": "0000000000000000", + "workspace_id": 0 + }, + { + "canSubscribeToLiveQuery": true, + "channel_used": { + "dbsql_version": "", + "name": "CHANNEL_NAME_UNSPECIFIED" + }, + "client_application": "Databricks Jobs", + "duration": 438, + "endpoint_id": "0000000000000000", + "executed_as_user_display_name": "", + "executed_as_user_id": 0, + "executed_as_user_name": "user@example.com", + "execution_end_time_ms": 1789249104129, + "is_cancelable": false, + "is_final": true, + "last_updated_at_ms": 1789249104254, + "lookup_key": "CiQ0MGQ4MWM3MC1mZmI1LTQ0MGQtOTdmZi1iYWZmN2I5ZTYyYTcQy866vIk0GhAwMDAwMDAwMDAwMDAwMDAwIJjxne2Tgd4B", + "metrics": { + "bytes_read_from_cache_percentage": 0, + "compilation_time_ms": 131, + "execution_time_ms": 307, + "metadata_time_ms": 0, + "network_sent_bytes": 0, + "photon_total_time_ms": 83, + "pruned_files_count": 0, + "query_compilation_start_timestamp": 1789249103691, + "query_execution_time_ms": 438, + "read_bytes": 0, + "read_cache_bytes": 0, + "read_files_count": 0, + "read_partitions_count": 0, + "read_remote_bytes": 0, + "result_fetch_time_ms": 123, + "result_from_cache": false, + "rows_produced_count": 1, + "rows_read_count": 0, + "spill_to_disk_bytes": 0, + "task_progress": { + "active_task_count": 0, + "completed_partition_count": 17, + "completed_task_count": 17, + "failed_task_count": 0, + "pending_task_count": 0, + "planned_task_count": 17 + }, + "task_total_time_ms": 144, + "total_time_ms": 438, + "write_remote_bytes": 0, + "write_remote_files": 0, + "write_remote_rows": 0 + }, + "plans_state": "EXISTS", + "query_end_time_ms": 1789249104129, + "query_id": "40d81c70-ffb5-440d-97ff-baff7b9e62a7", + "query_source": { + "job_info": { + "job_id": "240570464625850", + "job_task_run_id": "60603881665866" + }, + "notebook_id": "1234567890" + }, + "query_start_time_ms": 1789249103691, + "query_text": "", + "rows_produced": 1, + "session_id": "dc9ff815-f6ed-4952-87fe-4c96a07cd988", + "statement_type": "SELECT", + "status": "FINISHED", + "user_display_name": "", + "user_id": 0, + "user_name": "user@example.com", + "warehouse_id": "0000000000000000", + "workspace_id": 0 + }, + { + "canSubscribeToLiveQuery": true, + "channel_used": { + "dbsql_version": "", + "name": "CHANNEL_NAME_UNSPECIFIED" + }, + "client_application": "Databricks Jobs", + "duration": 296, + "endpoint_id": "0000000000000000", + "executed_as_user_display_name": "", + "executed_as_user_id": 0, + "executed_as_user_name": "user@example.com", + "execution_end_time_ms": 1789249102649, + "is_cancelable": false, + "is_final": true, + "last_updated_at_ms": 1789249102702, + "lookup_key": "CiRkZmQ5YzFmZS0zYzkxLTQwZDEtYWNjMy1lMmY0MjY4N2FjNTUQkcS6vIk0GhAwMDAwMDAwMDAwMDAwMDAwIJjxne2Tgd4B", + "metrics": { + "bytes_read_from_cache_percentage": 0, + "compilation_time_ms": 96, + "execution_time_ms": 200, + "metadata_time_ms": 0, + "network_sent_bytes": 0, + "photon_total_time_ms": 0, + "pruned_files_count": 0, + "query_compilation_start_timestamp": 1789249102353, + "query_execution_time_ms": 296, + "read_bytes": 0, + "read_cache_bytes": 0, + "read_files_count": 0, + "read_partitions_count": 0, + "read_remote_bytes": 0, + "result_fetch_time_ms": 52, + "result_from_cache": false, + "rows_produced_count": 0, + "rows_read_count": 0, + "spill_to_disk_bytes": 0, + "task_progress": { + "active_task_count": 0, + "completed_partition_count": 8, + "completed_task_count": 8, + "failed_task_count": 0, + "pending_task_count": 0, + "planned_task_count": 8 + }, + "task_total_time_ms": 552, + "total_time_ms": 296, + "write_remote_bytes": 0, + "write_remote_files": 0, + "write_remote_rows": 0 + }, + "plans_state": "EXISTS", + "query_end_time_ms": 1789249102649, + "query_id": "dfd9c1fe-3c91-40d1-acc3-e2f42687ac55", + "query_source": { + "job_info": { + "job_id": "240570464625850", + "job_task_run_id": "60603881665866" + }, + "notebook_id": "1234567890" + }, + "query_start_time_ms": 1789249102353, + "query_text": "", + "session_id": "dc9ff815-f6ed-4952-87fe-4c96a07cd988", + "statement_type": "INSERT", + "status": "FINISHED", + "user_display_name": "", + "user_id": 0, + "user_name": "user@example.com", + "warehouse_id": "0000000000000000", + "workspace_id": 0 + }, + { + "canSubscribeToLiveQuery": true, + "channel_used": { + "dbsql_version": "", + "name": "CHANNEL_NAME_UNSPECIFIED" + }, + "client_application": "Databricks Jobs", + "duration": 211, + "endpoint_id": "0000000000000000", + "executed_as_user_display_name": "", + "executed_as_user_id": 0, + "executed_as_user_name": "user@example.com", + "execution_end_time_ms": 1789249102287, + "is_cancelable": false, + "is_final": true, + "last_updated_at_ms": 1789249102336, + "lookup_key": "CiQ4YzNlNTg2Mi1hMTVjLTQ5YjEtODE1YS1mYzFiNzBmMzMxNzgQ/MG6vIk0GhAwMDAwMDAwMDAwMDAwMDAwIJjxne2Tgd4B", + "metrics": { + "bytes_read_from_cache_percentage": 0, + "compilation_time_ms": 79, + "execution_time_ms": 132, + "metadata_time_ms": 0, + "network_sent_bytes": 0, + "photon_total_time_ms": 0, + "pruned_files_count": 0, + "query_compilation_start_timestamp": 1789249102076, + "query_execution_time_ms": 211, + "read_bytes": 0, + "read_cache_bytes": 0, + "read_files_count": 0, + "read_partitions_count": 0, + "read_remote_bytes": 0, + "result_fetch_time_ms": 48, + "result_from_cache": false, + "rows_produced_count": 1, + "rows_read_count": 1, + "spill_to_disk_bytes": 0, + "task_progress": { + "active_task_count": 0, + "completed_partition_count": 1, + "completed_task_count": 1, + "failed_task_count": 0, + "pending_task_count": 0, + "planned_task_count": 1 + }, + "task_total_time_ms": 60, + "total_time_ms": 211, + "write_remote_bytes": 0, + "write_remote_files": 0, + "write_remote_rows": 0 + }, + "plans_state": "EXISTS", + "query_end_time_ms": 1789249102287, + "query_id": "8c3e5862-a15c-49b1-815a-fc1b70f33178", + "query_source": { + "job_info": { + "job_id": "240570464625850", + "job_task_run_id": "60603881665866" + }, + "notebook_id": "1234567890" + }, + "query_start_time_ms": 1789249102076, + "query_text": "", + "rows_produced": 1, + "session_id": "dc9ff815-f6ed-4952-87fe-4c96a07cd988", + "statement_type": "SELECT", + "status": "FINISHED", + "user_display_name": "", + "user_id": 0, + "user_name": "user@example.com", + "warehouse_id": "0000000000000000", + "workspace_id": 0 + }, + { + "canSubscribeToLiveQuery": true, + "channel_used": { + "dbsql_version": "", + "name": "CHANNEL_NAME_UNSPECIFIED" + }, + "client_application": "Databricks Jobs", + "duration": 351, + "endpoint_id": "0000000000000000", + "executed_as_user_display_name": "", + "executed_as_user_id": 0, + "executed_as_user_name": "user@example.com", + "execution_end_time_ms": 1789249101851, + "is_cancelable": false, + "is_final": true, + "last_updated_at_ms": 1789249101949, + "lookup_key": "CiQ3ZDk1ZTU0NC05YTc1LTQ5YTYtOWYyMi1lMWFhNDRmYjJjMWUQvL26vIk0GhAwMDAwMDAwMDAwMDAwMDAwIJjxne2Tgd4B", + "metrics": { + "bytes_read_from_cache_percentage": 0, + "compilation_time_ms": 80, + "execution_time_ms": 271, + "metadata_time_ms": 0, + "network_sent_bytes": 0, + "photon_total_time_ms": 0, + "pruned_files_count": 0, + "query_compilation_start_timestamp": 1789249101500, + "query_execution_time_ms": 351, + "read_bytes": 0, + "read_cache_bytes": 0, + "read_files_count": 0, + "read_partitions_count": 0, + "read_remote_bytes": 0, + "result_fetch_time_ms": 84, + "result_from_cache": false, + "rows_produced_count": 50, + "rows_read_count": 0, + "spill_to_disk_bytes": 0, + "task_progress": { + "active_task_count": 0, + "completed_partition_count": 9, + "completed_task_count": 9, + "failed_task_count": 0, + "pending_task_count": 0, + "planned_task_count": 9 + }, + "task_total_time_ms": 52, + "total_time_ms": 351, + "write_remote_bytes": 0, + "write_remote_files": 0, + "write_remote_rows": 0 + }, + "plans_state": "EXISTS", + "query_end_time_ms": 1789249101851, + "query_id": "7d95e544-9a75-49a6-9f22-e1aa44fb2c1e", + "query_source": { + "job_info": { + "job_id": "240570464625850", + "job_task_run_id": "60603881665866" + }, + "notebook_id": "1234567890" + }, + "query_start_time_ms": 1789249101500, + "query_text": "", + "rows_produced": 50, + "session_id": "dc9ff815-f6ed-4952-87fe-4c96a07cd988", + "statement_type": "SELECT", + "status": "FINISHED", + "user_display_name": "", + "user_id": 0, + "user_name": "user@example.com", + "warehouse_id": "0000000000000000", + "workspace_id": 0 + }, + { + "canSubscribeToLiveQuery": true, + "channel_used": { + "dbsql_version": "", + "name": "CHANNEL_NAME_UNSPECIFIED" + }, + "client_application": "Databricks Jobs", + "duration": 363, + "endpoint_id": "0000000000000000", + "executed_as_user_display_name": "", + "executed_as_user_id": 0, + "executed_as_user_name": "user@example.com", + "execution_end_time_ms": 1789249101402, + "is_cancelable": false, + "is_final": true, + "last_updated_at_ms": 1789249101450, + "lookup_key": "CiQzOTMyYWMxMy05MWFhLTQ5ZTItODIyNy00ZmY1YjQwZjc1NTMQ77m6vIk0GhAwMDAwMDAwMDAwMDAwMDAwIJjxne2Tgd4B", + "metrics": { + "bytes_read_from_cache_percentage": 0, + "compilation_time_ms": 115, + "execution_time_ms": 248, + "metadata_time_ms": 0, + "network_sent_bytes": 0, + "photon_total_time_ms": 74, + "pruned_files_count": 0, + "query_compilation_start_timestamp": 1789249101039, + "query_execution_time_ms": 363, + "read_bytes": 0, + "read_cache_bytes": 0, + "read_files_count": 0, + "read_partitions_count": 0, + "read_remote_bytes": 0, + "result_fetch_time_ms": 44, + "result_from_cache": false, + "rows_produced_count": 10, + "rows_read_count": 0, + "spill_to_disk_bytes": 0, + "task_progress": { + "active_task_count": 0, + "completed_partition_count": 17, + "completed_task_count": 17, + "failed_task_count": 0, + "pending_task_count": 0, + "planned_task_count": 17 + }, + "task_total_time_ms": 137, + "total_time_ms": 363, + "write_remote_bytes": 0, + "write_remote_files": 0, + "write_remote_rows": 0 + }, + "plans_state": "EXISTS", + "query_end_time_ms": 1789249101402, + "query_id": "3932ac13-91aa-49e2-8227-4ff5b40f7553", + "query_source": { + "job_info": { + "job_id": "240570464625850", + "job_task_run_id": "60603881665866" + }, + "notebook_id": "1234567890" + }, + "query_start_time_ms": 1789249101039, + "query_text": "", + "rows_produced": 10, + "session_id": "dc9ff815-f6ed-4952-87fe-4c96a07cd988", + "statement_type": "SELECT", + "status": "FINISHED", + "user_display_name": "", + "user_id": 0, + "user_name": "user@example.com", + "warehouse_id": "0000000000000000", + "workspace_id": 0 + }, + { + "canSubscribeToLiveQuery": true, + "channel_used": { + "dbsql_version": "", + "name": "CHANNEL_NAME_UNSPECIFIED" + }, + "client_application": "Databricks Jobs", + "duration": 708, + "endpoint_id": "0000000000000000", + "executed_as_user_display_name": "", + "executed_as_user_id": 0, + "executed_as_user_name": "user@example.com", + "execution_end_time_ms": 1789249100885, + "is_cancelable": false, + "is_final": true, + "last_updated_at_ms": 1789249100960, + "lookup_key": "CiRhYmFjNDZkOC1mNWU0LTQ5NWQtOWQ0Yi0yN2U4ZTkyM2ExZTIQkbO6vIk0GhAwMDAwMDAwMDAwMDAwMDAwIJjxne2Tgd4B", + "metrics": { + "bytes_read_from_cache_percentage": 0, + "compilation_time_ms": 99, + "execution_time_ms": 609, + "metadata_time_ms": 0, + "network_sent_bytes": 0, + "photon_total_time_ms": 95, + "pruned_files_count": 0, + "query_compilation_start_timestamp": 1789249100177, + "query_execution_time_ms": 708, + "read_bytes": 0, + "read_cache_bytes": 0, + "read_files_count": 0, + "read_partitions_count": 0, + "read_remote_bytes": 0, + "result_fetch_time_ms": 74, + "result_from_cache": false, + "rows_produced_count": 1, + "rows_read_count": 0, + "spill_to_disk_bytes": 0, + "task_progress": { + "active_task_count": 0, + "completed_partition_count": 17, + "completed_task_count": 17, + "failed_task_count": 0, + "pending_task_count": 0, + "planned_task_count": 17 + }, + "task_total_time_ms": 195, + "total_time_ms": 708, + "write_remote_bytes": 0, + "write_remote_files": 0, + "write_remote_rows": 0 + }, + "plans_state": "EMPTY", + "query_end_time_ms": 1789249100885, + "query_id": "abac46d8-f5e4-495d-9d4b-27e8e923a1e2", + "query_source": { + "job_info": { + "job_id": "240570464625850", + "job_task_run_id": "60603881665866" + }, + "notebook_id": "1234567890" + }, + "query_start_time_ms": 1789249100177, + "query_text": "", + "rows_produced": 1, + "session_id": "dc9ff815-f6ed-4952-87fe-4c96a07cd988", + "statement_type": "SELECT", + "status": "FINISHED", + "user_display_name": "", + "user_id": 0, + "user_name": "user@example.com", + "warehouse_id": "0000000000000000", + "workspace_id": 0 + } + ] +} diff --git a/tests/test_databricks.py b/tests/test_databricks.py new file mode 100644 index 0000000..5e4a7c9 --- /dev/null +++ b/tests/test_databricks.py @@ -0,0 +1,748 @@ +import json +import subprocess +from pathlib import Path +from typing import Any + +import pytest + +from sparkparse.databricks import ( + DatabricksCLI, + DatabricksError, + RawRunCollection, + collect_run, + fetch_full_run, +) +from sparkparse.models import WorkflowCollectionLimits +from sparkparse.runreport import build_report + +FIXTURES = Path(__file__).parent / "data" / "databricks" + + +def load_fixture(name: str) -> Any: + return json.loads((FIXTURES / name).read_text()) + + +def _proc(args: Any, returncode: int = 0, stdout: str = "{}", stderr: str = "") -> Any: + return subprocess.CompletedProcess(list(args), returncode, stdout, stderr) + + +class ScriptedRunner: + """Returns canned responses based on the command path.""" + + def __init__( + self, + run: dict[str, Any] | None = None, + run_pages: list[dict[str, Any]] | None = None, + query_pages: list[dict[str, Any]] | None = None, + outputs: dict[str, Any] | None = None, + clusters: dict[str, Any] | None = None, + rate_limit_once: bool = False, + auth_fail: bool = False, + runs_list: dict[str, Any] | None = None, + ) -> None: + self.run = run + self.run_pages = list(run_pages or []) + self.query_pages = list(query_pages or []) + self.outputs = outputs or {} + self.clusters = clusters or {} + self.rate_limit_once = rate_limit_once + self.auth_fail = auth_fail + self.runs_list = runs_list + self.calls: list[list[str]] = [] + + def __call__(self, args: Any, timeout: float) -> Any: + args = list(args) + self.calls.append(args) + joined = " ".join(args) + if self.auth_fail: + return _proc(args, 1, "", "Invalid access token") + if self.rate_limit_once: + self.rate_limit_once = False + return _proc(args, 1, "", "429 Too Many Requests") + if "runs/list" in joined: + if isinstance(self.runs_list, list): + page = self.runs_list.pop(0) if self.runs_list else {"runs": []} + return _proc(args, 0, json.dumps(page)) + return _proc(args, 0, json.dumps(self.runs_list or {"runs": []})) + if "runs/get" in joined: + if self.run_pages: + page = self.run_pages.pop(0) + return _proc(args, 0, json.dumps(page)) + return _proc(args, 0, json.dumps(self.run or {})) + if "history/queries" in joined: + if self.query_pages: + return _proc(args, 0, json.dumps(self.query_pages.pop(0))) + return _proc(args, 0, json.dumps({"has_next_page": False, "res": []})) + if "get-run-output" in joined: + index = args.index("get-run-output") + run_id = args[index + 1] + return _proc( + args, 0, json.dumps(self.outputs.get(str(run_id), {"logs": "line"})) + ) + if "clusters/get" in joined: + cluster_id = joined.split("cluster_id=")[-1] + return _proc(args, 0, json.dumps(self.clusters.get(cluster_id, {}))) + return _proc(args, 0, "{}") + + +def make_runner(**kwargs: Any) -> ScriptedRunner: + return ScriptedRunner(**kwargs) + + +def client(runner: ScriptedRunner, **limit_overrides: Any) -> DatabricksCLI: + limits = WorkflowCollectionLimits( + deadline_seconds=limit_overrides.get("deadline_seconds", 30.0), + max_pages=limit_overrides.get("max_pages", 10), + max_requests=limit_overrides.get("max_requests", 20), + max_retries=limit_overrides.get("max_retries", 2), + retry_backoff_seconds=0.0, + ) + return DatabricksCLI(profile="free", runner=runner, limits=limits) + + +def test_multi_task_normalization_keeps_timing_measures_separate(): + runner = make_runner(run=load_fixture("job_run_multi.json")) + raw = collect_run(client(runner), run_id="123", outputs="none") + report = build_report(raw, asserted_revision=None) + + assert [t.task_key for t in report.tasks] == [ + "ingest", + "join", + "optimize", + "foreach_copy", + "unknown_type", + ] + assert report.timing.workflow_elapsed_ms == 60000 + assert report.timing.summed_task_execution_ms == 56500 + # Workflow elapsed is never the summed task time. + assert report.timing.workflow_elapsed_ms != report.timing.summed_task_execution_ms + # Repair/retry identity and skip state survive. + join = next(t for t in report.tasks if t.task_key == "join") + assert join.original_attempt_run_id == "1000" + assert join.attempt_number == 1 + optimize = next(t for t in report.tasks if t.task_key == "optimize") + assert optimize.result_state == "SKIPPED" + foreach = next(t for t in report.tasks if t.task_key == "foreach_copy") + assert foreach.task_kind == "for_each" + unknown = next(t for t in report.tasks if t.task_key == "unknown_type") + assert unknown.task_kind is None + assert unknown.compute is None + + +def test_paginated_get_run_merges_array_properties(): + first = {"run_id": 123, "tasks": [{"task_key": "a"}], "next_page_token": "t1"} + second = {"run_id": 123, "tasks": [{"task_key": "b"}]} + runner = make_runner(run_pages=[first, second]) + merged, complete = fetch_full_run(client(runner), "123") + assert complete is True + assert [t["task_key"] for t in merged["tasks"]] == ["a", "b"] + + +def test_run_pagination_exhaustion_keeps_partial_tasks(): + first = { + "run_id": 123, + "tasks": [{"task_key": "a"}], + "next_page_token": "t1", + "state": {"life_cycle_state": "TERMINATED"}, + } + runner = make_runner( + run_pages=[first, {"run_id": 123, "tasks": [{"task_key": "b"}]}] + ) + # max_pages=0 leaves no room for the follow-up page. + merged, complete = fetch_full_run(client(runner, max_pages=0), "123") + assert complete is False + assert [t["task_key"] for t in merged["tasks"]] == ["a"] + + +def test_executed_git_snapshot_beats_asserted_revision(): + raw = collect_run( + client(make_runner(run=load_fixture("job_run_multi.json"))), + run_id="123", + outputs="none", + ) + report = build_report(raw, asserted_revision="zzz") + assert report.revision.confidence == "executed" + assert (report.revision.executed_commit or "").startswith("abc123") + + +def test_asserted_and_unknown_revision(): + run = load_fixture("job_run_single.json") + run.pop("git_source", None) + raw = collect_run(client(make_runner(run=run)), run_id="1", outputs="none") + assert build_report(raw).revision.confidence == "unknown" + assert ( + build_report(raw, asserted_revision="deadbeef").revision.confidence + == "asserted" + ) + + +def test_mixed_task_revisions_are_not_attributed_to_one_sha(): + run = load_fixture("job_run_multi.json") + run["tasks"][0]["git_source"] = { + "git_url": "https://example/repo", + "git_snapshot": {"used_commit": "othercommit"}, + } + raw = collect_run(client(make_runner(run=run)), run_id="123", outputs="none") + report = build_report(raw) + assert report.revision.confidence == "mixed" + assert report.revision.executed_commit is None + assert any(d.code == "mixed_revision" for d in report.diagnostics) + + +def test_task_run_id_reports_parent(): + run = {"run_id": 999, "job_run_id": 123, "tasks": []} + with pytest.raises(DatabricksError) as excinfo: + collect_run(client(make_runner(run=run)), run_id="999", outputs="none") + assert excinfo.value.code == "task_run_id" + assert "123" in str(excinfo.value) + + +def test_resolve_latest_terminal_run_including_failures(): + runs = { + "runs": [ + {"run_id": 2, "state": {"life_cycle_state": "RUNNING"}}, + { + "run_id": 3, + "state": {"life_cycle_state": "TERMINATED", "result_state": "FAILED"}, + }, + ] + } + from sparkparse.databricks import resolve_run_id + + selected, active = resolve_run_id(client(make_runner(runs_list=runs)), "42") + assert selected == "3" + assert active is False + + +def test_active_run_yields_labeled_partial_snapshot(): + run = load_fixture("job_run_multi.json") + run["state"] = {"life_cycle_state": "RUNNING"} + run["end_time"] = None + raw = collect_run(client(make_runner(run=run)), run_id="123", outputs="none") + assert raw.collection_status.value == "partial" + report = build_report(raw) + assert any(d.code == "active_run" for d in report.diagnostics) + + +def test_rate_limit_retry_then_success(): + runner = make_runner(run=load_fixture("job_run_multi.json"), rate_limit_once=True) + raw = collect_run(client(runner), run_id="123", outputs="none") + assert raw.run["run_id"] == 123 + + +def test_auth_failure_is_clear(): + with pytest.raises(DatabricksError) as excinfo: + collect_run( + client(make_runner(run=load_fixture("job_run_multi.json"), auth_fail=True)), + run_id="123", + outputs="none", + ) + assert excinfo.value.code == "auth_error" + + +def test_resolve_run_id_finds_terminal_run_on_a_later_page(): + runs_pages = [ + { + "runs": [{"run_id": 2, "state": {"life_cycle_state": "RUNNING"}}], + "next_page_token": "p2", + }, + { + "runs": [ + { + "run_id": 3, + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "SUCCESS", + }, + } + ] + }, + ] + from sparkparse.databricks import resolve_run_id + + selected, active = resolve_run_id(client(make_runner(runs_list=runs_pages)), "42") + assert selected == "3" + assert active is False + + +def test_query_pagination_exhaustion_preserves_collected_queries(): + first = { + "has_next_page": True, + "next_page_token": "t", + "res": [{"query_id": "q1"}], + } + second = {"has_next_page": False, "res": [{"query_id": "q2"}]} + runner = make_runner( + run=load_fixture("job_run_multi.json"), query_pages=[first, second] + ) + raw = collect_run(client(runner, max_requests=2), run_id="123", outputs="none") + # The first page was fetched and must survive the follow-up request failure. + assert [q["query_id"] for q in raw.queries] == ["q1"] + assert raw.discovered_query_count == 1 + assert raw.query_discovery_complete is False + assert raw.collection_status.value == "partial" + assert any(d.code == "collection_budget_exhausted" for d in raw.diagnostics) + + +def test_query_budget_exhaustion_is_partial_not_fatal(): + pages = [ + {"has_next_page": True, "next_page_token": "t", "res": [{"query_id": "q1"}]}, + {"has_next_page": False, "res": [{"query_id": "q2"}]}, + ] + runner = make_runner(run=load_fixture("job_run_multi.json"), query_pages=pages) + raw = collect_run(client(runner, max_requests=2), run_id="123", outputs="none") + assert raw.collection_status.value == "partial" + assert any(d.code == "collection_budget_exhausted" for d in raw.diagnostics) + + +def test_bounded_outputs_are_truncated_and_recorded(): + run = load_fixture("job_run_multi.json") + raw = collect_run(client(make_runner(run=run)), run_id="123", outputs="all") + report = build_report(raw) + assert len(report.outputs) == len(report.tasks) + assert all(o.status in ("saved", "truncated") for o in report.outputs) + + +def test_output_truncation_flag_is_preserved(): + run = load_fixture("job_run_multi.json") + outputs = { + "1001": {"notebook_output": {"result": "x" * 10, "truncated": True}}, + } + raw = collect_run( + client(make_runner(run=run, outputs=outputs)), run_id="123", outputs="all" + ) + report = build_report(raw) + first = next(o for o in report.outputs if o.run_id == "1001") + assert first.status == "truncated" + + +def test_cluster_lookup_is_cached_and_labeled_current(): + run = load_fixture("job_run_multi.json") + run["tasks"][0].pop("new_cluster") + run["tasks"][0]["existing_cluster_id"] = "cluster-1" + run["tasks"][1].pop("new_cluster") + run["tasks"][1]["existing_cluster_id"] = "cluster-1" + clusters = {"cluster-1": {"spark_version": "15.4", "runtime_engine": "PHOTON"}} + runner = make_runner(run=run, clusters=clusters) + raw = collect_run(client(runner), run_id="123", outputs="none") + lookup_calls = [c for c in runner.calls if "clusters/get" in " ".join(c)] + assert len(lookup_calls) == 1 + report = build_report(raw) + references = [c for c in report.compute if c.cluster_id == "cluster-1"] + assert references and all(c.source == "current_lookup" for c in references) + + +def _report_with_queries(queries: list[dict[str, Any]]): + run = load_fixture("job_run_multi.json") + raw = RawRunCollection( + run=run, + tasks=run["tasks"], + queries=queries, + discovered_query_count=len(queries), + ) + return build_report(raw) + + +def test_query_metrics_observed_subtotal_not_complete_total(): + report = _report_with_queries(load_fixture("query_history_mixed.json")["res"]) + metrics = report.query_metrics + assert metrics is not None + assert metrics.matched_query_count == 2 + assert metrics.attribution_counts == { + "task_run": 2, + "unattributed": 1, + "unrelated": 1, + } + read = next(a for a in metrics.aggregates if a.metric == "read_bytes") + assert read.value_exact == 3000 + # Queries lack is_final and one query in the window is unattributed, so the + # read total cannot be presented as a complete total. + assert read.completeness.value == "observed_subtotal" + cache = next(a for a in metrics.aggregates if a.metric == "read_cache_bytes") + assert cache.value_exact == 400 + assert cache.completeness.value == "observed_subtotal" + assert cache.counted_query_count == 1 + assert cache.observed_query_count == 2 + + +def test_non_final_query_prevents_complete_total(): + queries = [ + { + "query_id": "q-1", + "is_final": False, + "status": "RUNNING", + "query_source": {"job_info": {"job_task_run_id": "1001"}}, + "metrics": {"read_bytes": 100}, + } + ] + report = _report_with_queries(queries) + metrics = report.query_metrics + assert metrics is not None + read = next(a for a in metrics.aggregates if a.metric == "read_bytes") + assert read.value_exact == 100 + assert read.completeness.value == "observed_subtotal" + + +def test_unattributed_query_prevents_complete_total(): + queries = [ + { + "query_id": "q-1", + "is_final": True, + "status": "FINISHED", + "query_source": {"job_info": {"job_task_run_id": "1001"}}, + "metrics": {"read_bytes": 100}, + }, + { + "query_id": "q-2", + "is_final": True, + "status": "FINISHED", + "query_source": {"notebook_id": "1"}, + "metrics": {"read_bytes": 50}, + }, + ] + report = _report_with_queries(queries) + metrics = report.query_metrics + assert metrics is not None + read = next(a for a in metrics.aggregates if a.metric == "read_bytes") + assert read.completeness.value == "observed_subtotal" + + +def test_task_scoped_git_evidence_is_not_the_whole_run(): + run = load_fixture("job_run_multi.json") + run.pop("git_source", None) + run["tasks"][1]["git_source"] = { + "git_url": "https://example/repo", + "git_snapshot": {"used_commit": "task-only-sha"}, + } + raw = collect_run(client(make_runner(run=run)), run_id="123", outputs="none") + report = build_report(raw) + assert report.revision.confidence == "executed" + assert report.revision.scope == "task" + join = next(t for t in report.tasks if t.task_key == "join") + assert join.git_commit == "task-only-sha" + assert any(t.git_commit is None for t in report.tasks) + + +def test_run_scoped_git_evidence_covers_the_run(): + raw = collect_run( + client(make_runner(run=load_fixture("job_run_multi.json"))), + run_id="123", + outputs="none", + ) + report = build_report(raw) + assert report.revision.scope == "run" + + +def test_job_parameters_are_normalized_for_config_identity(): + run = load_fixture("job_run_multi.json") + run["job_parameters"] = [ + {"name": "shuffle_partitions", "value": "200"}, + {"name": "unused", "value": None}, + ] + raw = collect_run(client(make_runner(run=run)), run_id="123", outputs="none") + report = build_report(raw) + assert report.job_parameters == {"shuffle_partitions": "200", "unused": ""} + + +def test_credentials_are_excluded_from_parameter_capture(): + run = load_fixture("job_run_multi.json") + run["job_parameters"] = [ + {"name": "shuffle_partitions", "value": "200"}, + {"name": "api_token", "value": "dummy-token-value"}, + ] + run["tasks"][0]["new_cluster"]["spark_conf"] = { + "spark.sql.shuffle.partitions": "800", + "fs.s3a.secret.key": "dummy-s3-secret", + "fs.s3a.access.key": "dummy-s3-access", + } + run["tasks"][0]["notebook_task"]["base_parameters"] = { + "widget": "visible", + "db_password": "dummy-password", + } + raw = collect_run(client(make_runner(run=run)), run_id="123", outputs="none") + report = build_report(raw) + assert report.job_parameters == {"shuffle_partitions": "200"} + task = next(t for t in report.tasks if t.task_key == "ingest") + assert task.spark_conf == {"spark.sql.shuffle.partitions": "800"} + assert task.parameters == {"widget": "visible"} + serialized = report.model_dump_json() + for secret in ( + "dummy-token-value", + "dummy-s3-secret", + "dummy-s3-access", + "dummy-password", + ): + assert secret not in serialized + + +def test_secret_arguments_are_redacted_from_parameter_capture(): + run = load_fixture("job_run_multi.json") + run["tasks"][0]["spark_python_task"] = { + "python_file": "task.py", + "parameters": [ + "--verbose", + "--api-token", + "dummy-token", + "--shuffle", + "200", + "--password=dummy-password", + ], + } + raw = collect_run(client(make_runner(run=run)), run_id="123", outputs="none") + report = build_report(raw) + task = next(t for t in report.tasks if t.task_key == "ingest") + captured = task.parameters["spark_python_task.parameters"] + assert "dummy-token" not in captured + assert "dummy-password" not in captured + assert "" in captured + assert "--shuffle,200" in captured + + +def test_run_environments_and_performance_mode_are_normalized(): + raw = collect_run( + client(make_runner(run=load_fixture("job_run_single.json"))), + run_id="915049514672868", + outputs="none", + ) + report = build_report(raw) + assert report.effective_performance_target == "PERFORMANCE_OPTIMIZED" + assert report.environments + default = report.environments[0] + assert default.environment_key == "default" + assert any("polars" in dep for dep in default.dependencies) + + +def test_missing_metric_is_none_not_zero(): + report = _report_with_queries(load_fixture("query_history_mixed.json")["res"]) + metrics = report.query_metrics + assert metrics is not None + remote = next(a for a in metrics.aggregates if a.metric == "read_remote_bytes") + assert remote.value is None + assert remote.value_exact is None + assert remote.completeness.value == "unavailable" + # A measured zero survives as zero. + spill = next(a for a in metrics.aggregates if a.metric == "disk_spill_bytes") + assert spill.value_exact == 900 + + +def test_missing_parent_id_with_valid_task_id_is_attributed(): + queries = [ + { + "query_id": "q-valid", + "status": "FINISHED", + "execution_end_time_ms": 2000, + "query_source": {"job_info": {"job_task_run_id": "1001"}}, + "metrics": {"read_bytes": 10}, + } + ] + report = _report_with_queries(queries) + metrics = report.query_metrics + assert metrics is not None + assert metrics.matched_query_count == 1 + observation = metrics.observations[0] + assert observation.attribution == "task_run" + assert observation.task_key == "ingest" + + +def test_unrelated_queries_are_not_persisted(): + queries = [ + { + "query_id": "q-other", + "query_source": {"job_info": {"job_task_run_id": "999"}}, + "metrics": {"read_bytes": 1}, + } + ] + report = _report_with_queries(queries) + metrics = report.query_metrics + assert metrics is not None + assert metrics.matched_query_count == 0 + assert metrics.observations == [] + + +def test_duplicate_query_pages_keep_the_fresher_observation(): + queries = [ + { + "query_id": "q-1", + "is_final": False, + "execution_end_time_ms": 1000, + "query_source": {"job_info": {"job_task_run_id": "1001"}}, + "metrics": {"read_bytes": 5}, + }, + { + "query_id": "q-1", + "is_final": True, + "execution_end_time_ms": 2000, + "query_source": {"job_info": {"job_task_run_id": "1001"}}, + "metrics": {"read_bytes": 7}, + }, + ] + report = _report_with_queries(queries) + metrics = report.query_metrics + assert metrics is not None + assert len(metrics.observations) == 1 + assert metrics.observations[0].metrics["read_bytes"] == 7 + + +def test_incomplete_run_pagination_marks_task_list_partial(): + run = load_fixture("job_run_multi.json") + raw = RawRunCollection( + run=run, + tasks=run["tasks"], + queries=[], + run_pagination_complete=False, + ) + report = build_report(raw) + assert report.coverage["task_list"] == "partial" + + +def test_incomplete_discovery_prevents_complete_coverage(): + run = load_fixture("job_run_multi.json") + raw = RawRunCollection( + run=run, + tasks=run["tasks"], + queries=[ + { + "query_id": "q-1", + "query_source": {"job_info": {"job_task_run_id": "1001"}}, + "metrics": {"read_bytes": 5}, + } + ], + discovered_query_count=1, + query_discovery_complete=False, + ) + report = build_report(raw) + metrics = report.query_metrics + assert metrics is not None + read = next(a for a in metrics.aggregates if a.metric == "read_bytes") + assert read.completeness.value == "observed_subtotal" + assert any(d.code == "query_discovery_incomplete" for d in report.diagnostics) + + +# --- CLI contracts ------------------------------------------------------- + +from typer.testing import CliRunner # noqa: E402 + +from sparkparse.app import app # noqa: E402 + +runner = CliRunner() + + +def _stub_collect(monkeypatch: pytest.MonkeyPatch) -> None: + run = load_fixture("job_run_multi.json") + + def fake_collect(cli: Any, **kwargs: Any) -> RawRunCollection: + return RawRunCollection( + run=run, + tasks=run["tasks"], + queries=load_fixture("query_history_mixed.json")["res"], + discovered_query_count=4, + ) + + monkeypatch.setattr("sparkparse.app.collect_run", fake_collect) + monkeypatch.setattr("sparkparse.app.DatabricksCLI", lambda **kwargs: object()) + + +def test_analyze_requires_exactly_one_id(): + assert runner.invoke(app, ["databricks", "analyze"]).exit_code == 2 + both = runner.invoke( + app, ["databricks", "analyze", "--run-id", "1", "--job-id", "2"] + ) + assert both.exit_code == 2 + + +def test_analyze_rejects_out_and_experiment_together(): + result = runner.invoke( + app, + ["databricks", "analyze", "--run-id", "1", "--out", "a", "--experiment", "b"], + ) + assert result.exit_code == 2 + + +def test_analyze_stdout_is_pure_json(monkeypatch, tmp_path): + _stub_collect(monkeypatch) + out_dir = tmp_path / "runs" + separate = CliRunner(mix_stderr=False) + result = separate.invoke( + app, ["databricks", "analyze", "--run-id", "123", "--out", str(out_dir)] + ) + assert result.exit_code == 0, result.output + payload = json.loads(result.stdout) + assert payload["identity"]["run_id"] == "123" + assert payload["query_metrics"]["matched_query_count"] == 2 + assert result.stderr # progress goes to stderr + assert list(out_dir.glob("runs/*/*.json")) + + +def test_analyze_records_experiment_trial(monkeypatch, tmp_path): + _stub_collect(monkeypatch) + exp_dir = tmp_path / "exp" + result = runner.invoke( + app, + [ + "databricks", + "analyze", + "--run-id", + "123", + "--experiment", + str(exp_dir), + "--variant", + "broadcast", + "--revision", + "deadbeef", + ], + ) + assert result.exit_code == 0, result.output + manifest = json.loads((exp_dir / "experiment.json").read_text()) + assert manifest["trials"][0]["variant"] == "broadcast" + + +def test_analyze_collection_failure_exits_1(monkeypatch): + def boom(cli: Any, **kwargs: Any) -> RawRunCollection: + raise DatabricksError("no runs", code="no_runs") + + monkeypatch.setattr("sparkparse.app.collect_run", boom) + monkeypatch.setattr("sparkparse.app.DatabricksCLI", lambda **kwargs: object()) + result = runner.invoke(app, ["databricks", "analyze", "--job-id", "9"]) + assert result.exit_code == 1 + + +def test_analyze_text_format(monkeypatch): + _stub_collect(monkeypatch) + result = runner.invoke( + app, ["databricks", "analyze", "--run-id", "123", "--format", "text"] + ) + assert result.exit_code == 0, result.output + assert "Run 123" in result.stdout + + +def test_nested_and_foreach_tasks_mark_incomplete_child_coverage(): + raw = collect_run( + client(make_runner(run=load_fixture("job_run_multi.json"))), + run_id="123", + outputs="none", + ) + report = build_report(raw) + assert report.coverage["child_tasks"] == "incomplete" + assert any(d.code == "incomplete_child_coverage" for d in report.diagnostics) + + +def test_retained_serverless_query_history_fixture_is_complete(): + run = load_fixture("job_run_single.json") + queries = load_fixture("query_history_window.json")["res"] + raw = RawRunCollection( + run=run, + tasks=run["tasks"], + queries=queries, + discovered_query_count=len(queries), + ) + report = build_report(raw) + metrics = report.query_metrics + assert metrics is not None + assert metrics.discovered_query_count == 9 + assert metrics.matched_query_count == 9 + assert report.tasks[0].run_id == "60603881665866" + task_time = next( + a for a in metrics.aggregates if a.metric == "aggregate_task_time_ms" + ) + assert task_time.completeness.value == "complete" + assert task_time.counted_query_count == 9 diff --git a/tests/test_experiments.py b/tests/test_experiments.py new file mode 100644 index 0000000..79ec4a9 --- /dev/null +++ b/tests/test_experiments.py @@ -0,0 +1,532 @@ +import json +from pathlib import Path +from typing import Any + +import pytest + +from sparkparse import experiments as exp +from sparkparse.databricks import RawRunCollection +from sparkparse.models import RunReport +from sparkparse.runreport import build_report + +FIXTURES = Path(__file__).parent / "data" / "databricks" + + +def load_fixture(name: str) -> Any: + return json.loads((FIXTURES / name).read_text()) + + +def make_report( + run_id: str = "1", + *, + execution_ms: int = 1000, + read_bytes: int = 100, + result_state: str = "SUCCESS", + commit: str = "aaaabbbbcccc", + variant_git: bool = True, + config: str | None = None, + query_final: bool = True, + task_list_partial: bool = False, + task_count: int | None = None, +) -> RunReport: + run = load_fixture("job_run_multi.json") + run["run_id"] = int(run_id) + run["job_run_id"] = int(run_id) + run["state"] = {"life_cycle_state": "TERMINATED", "result_state": result_state} + if task_count is not None: + run["tasks"] = run["tasks"][:task_count] + if variant_git: + run["git_source"] = { + "git_url": "https://example/repo", + "git_snapshot": {"used_commit": commit}, + } + else: + run.pop("git_source", None) + tasks = run["tasks"] + for task in tasks: + task["execution_duration"] = execution_ms + queries = [ + { + "query_id": f"q-{run_id}", + "status": "FINISHED", + "is_final": query_final, + "execution_end_time_ms": 2000, + "query_source": {"job_info": {"job_task_run_id": "1001"}}, + "metrics": {"read_bytes": read_bytes, "task_total_time_ms": execution_ms}, + } + ] + raw = RawRunCollection( + run=run, + tasks=tasks, + queries=queries, + discovered_query_count=1, + run_pagination_complete=not task_list_partial, + ) + report = build_report(raw) + if config is not None: + report.coverage["config"] = config + return report + + +def record( + exp_dir: Path, report: RunReport, variant: str, **kwargs: Any +) -> tuple[Any, Path]: + return exp.record_trial(exp_dir, report, variant=variant, **kwargs) + + +def test_recollection_updates_reference_not_trial_count(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1"), "baseline") + record(exp_dir, make_report("1"), "baseline") + ref3, _ = record(exp_dir, make_report("2"), "baseline") + manifest, _ = exp.load_experiment(exp_dir) + assert [t.run_id for t in manifest.trials] == ["1", "2"] + assert len(manifest.trials) == 2 + # Three snapshots remain on disk; the recollection did not erase the first. + assert len(exp.list_snapshot_paths(exp_dir)) == 3 + assert ref3.snapshot_id + + +def test_incompatible_job_is_rejected(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1"), "baseline") + other = make_report("9") + other.identity.job_id = "999" + with pytest.raises(exp.ExperimentError): + record(exp_dir, other, "baseline") + + +def test_cloud_experiment_path_is_rejected(): + with pytest.raises(exp.ExperimentError): + exp.save_snapshot( + "s3://bucket/exp", exp.build_snapshot(make_report(), trial=_trial()) + ) + + +def _trial(): + from sparkparse.models import TrialMetadata + + return TrialMetadata(variant="baseline") + + +def test_corrupt_manifest_is_an_error(tmp_path): + exp_dir = tmp_path / "exp" + exp_dir.mkdir() + (exp_dir / exp.MANIFEST_NAME).write_text("{not json") + with pytest.raises(exp.ExperimentError): + exp.load_manifest(exp_dir) + + +def test_zero_baseline_has_undefined_pct_change(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1", read_bytes=0), "baseline") + record(exp_dir, make_report("2", read_bytes=500), "candidate") + comparison = exp.compare_experiment( + exp_dir, baseline_run_id="1", candidate_run_id="2" + ) + read = next(m for m in comparison.metrics if m.metric == "read_bytes") + assert read.baseline_median == 0 + assert read.pct_change is None + assert read.delta == 500 + + +def test_failed_and_warmup_trials_are_excluded_but_counted(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1"), "baseline") + record(exp_dir, make_report("2", result_state="FAILED"), "baseline") + record(exp_dir, make_report("3", execution_ms=50), "baseline", warmup=True) + comparison = exp.compare_experiment( + exp_dir, baseline_variant="baseline", candidate_variant="baseline" + ) + assert comparison.baseline.excluded_failed == 1 + assert comparison.baseline.excluded_warmup == 1 + assert comparison.baseline.sample_count == 1 + assert any("failed" in note for note in comparison.notes) + + +def test_changed_compute_is_displayed_not_forbidden(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1"), "baseline") + changed = make_report("2") + for reference in changed.compute: + reference.num_workers = 16 + record(exp_dir, changed, "candidate") + comparison = exp.compare_experiment( + exp_dir, baseline_run_id="1", candidate_run_id="2" + ) + compute_note = next(n for n in comparison.comparability if n.aspect == "compute") + assert compute_note.status == "different" + + +def test_added_and_removed_tasks_are_shown(tmp_path): + exp_dir = tmp_path / "exp" + baseline = make_report("1") + candidate = make_report("2") + candidate.tasks = candidate.tasks[:-1] + record(exp_dir, baseline, "baseline") + record(exp_dir, candidate, "candidate") + comparison = exp.compare_experiment( + exp_dir, baseline_run_id="1", candidate_run_id="2" + ) + changes = {t.task_key: t.change for t in comparison.tasks} + assert changes["unknown_type"] == "removed" + + +def test_reduced_read_bytes_is_not_treated_as_incomparable(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1", read_bytes=1000), "baseline") + record(exp_dir, make_report("2", read_bytes=100), "candidate") + comparison = exp.compare_experiment( + exp_dir, baseline_run_id="1", candidate_run_id="2" + ) + read = next(m for m in comparison.metrics if m.metric == "read_bytes") + assert read.delta == -900 + assert read.pct_change == pytest.approx(-0.9) + + +def test_same_sha_different_configurations_are_separate_variants(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1"), "baseline", config_fingerprint="c1") + record(exp_dir, make_report("2"), "broadcast", config_fingerprint="c2") + manifest, _ = exp.load_experiment(exp_dir) + assert {t.variant for t in manifest.trials} == {"baseline", "broadcast"} + comparison = exp.compare_experiment( + exp_dir, baseline_variant="baseline", candidate_variant="broadcast" + ) + config_note = next( + n for n in comparison.comparability if n.aspect == "configuration" + ) + assert config_note.status == "different" + + +def test_relabeling_a_run_does_not_duplicate_the_sample(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1"), "baseline") + record(exp_dir, make_report("1"), "broadcast") + manifest, _ = exp.load_experiment(exp_dir) + assert len(manifest.trials) == 1 + assert manifest.trials[0].variant == "broadcast" + # An explicit run-ID comparison selects the one execution, not two. + comparison = exp.compare_experiment( + exp_dir, baseline_run_id="1", candidate_run_id="1" + ) + assert comparison.baseline.sample_count == 1 + assert comparison.candidate.sample_count == 1 + + +def test_recollection_preserves_existing_metadata(tmp_path): + exp_dir = tmp_path / "exp" + record( + exp_dir, + make_report("1"), + "baseline", + warmup=True, + correctness="verified", + input_snapshot="snapshot-a", + config_fingerprint="cfg-explicit", + ) + record(exp_dir, make_report("1"), "baseline") + manifest, _ = exp.load_experiment(exp_dir) + trial = manifest.trials[0] + assert trial.warmup is True + assert trial.correctness == "verified" + assert trial.input_snapshot == "snapshot-a" + assert trial.config_fingerprint == "cfg-explicit" + + +def test_variant_with_mixed_configurations_is_rejected(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1"), "baseline", config_fingerprint="c1") + record(exp_dir, make_report("2"), "baseline", config_fingerprint="c2") + with pytest.raises(exp.ExperimentError): + exp.compare_experiment( + exp_dir, baseline_variant="baseline", candidate_variant="baseline" + ) + + +def test_config_fingerprint_is_derived_when_not_supplied(tmp_path): + exp_dir = tmp_path / "exp" + ref, _ = record(exp_dir, make_report("1"), "baseline") + assert ref.config_fingerprint + + +def test_partial_metric_is_reported_with_caveat_not_dropped(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1", read_bytes=1000), "baseline") + record(exp_dir, make_report("2", read_bytes=500, query_final=False), "candidate") + comparison = exp.compare_experiment( + exp_dir, baseline_run_id="1", candidate_run_id="2" + ) + # The raw delta survives so the metric does not disappear entirely. + read = next(m for m in comparison.metrics if m.metric == "read_bytes") + assert read.partial is True + assert read.delta == -500 + assert read.candidate_partial_values == [500.0] + assert read.caveat and "subtotal" in read.caveat + note = next( + n for n in comparison.comparability if n.aspect == "metric_completeness" + ) + assert note.status == "partial" + + +def test_execution_parameters_change_config_fingerprint(tmp_path): + exp_dir = tmp_path / "exp" + baseline = make_report("1") + baseline.job_parameters = {"shuffle_partitions": "200"} + candidate = make_report("2") + candidate.job_parameters = {"shuffle_partitions": "800"} + ref_baseline, _ = record(exp_dir, baseline, "baseline") + ref_candidate, _ = record(exp_dir, candidate, "candidate") + assert ref_baseline.config_fingerprint != ref_candidate.config_fingerprint + + +def test_partial_task_list_marks_timing_partial_and_suppresses_added(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1", execution_ms=1000), "baseline") + record( + exp_dir, + make_report("2", execution_ms=1000, task_list_partial=True, task_count=1), + "candidate", + ) + comparison = exp.compare_experiment( + exp_dir, baseline_run_id="1", candidate_run_id="2" + ) + timing = next(m for m in comparison.metrics if m.metric == "task_execution_ms") + assert timing.partial is True + assert timing.caveat + # Missing pages must not be presented as removed tasks. + assert [t for t in comparison.tasks if t.change == "removed"] == [] + assert any(t.change == "unknown" for t in comparison.tasks) + assert any( + note.aspect == "task_coverage" and note.status == "partial" + for note in comparison.comparability + ) + + +def test_partial_task_list_marks_trial_row_metric_partial(tmp_path): + exp_dir = tmp_path / "exp" + record( + exp_dir, + make_report("1", task_list_partial=True, task_count=1), + "baseline", + ) + manifest, snapshots = exp.load_experiment(exp_dir) + rows = exp.trial_rows(manifest, snapshots) + assert "task_execution_ms" in rows[0]["partial_metrics"] + assert "task_execution_ms" in rows[0]["values"] + + +def test_derived_fingerprint_refreshes_on_recollection(tmp_path): + exp_dir = tmp_path / "exp" + first, _ = record(exp_dir, make_report("1"), "baseline") + assert first.config_fingerprint_derived is True + enriched = make_report("1") + enriched.job_parameters = {"shuffle_partitions": "800"} + refreshed, _ = record(exp_dir, enriched, "baseline") + assert refreshed.config_fingerprint != first.config_fingerprint + assert refreshed.config_fingerprint_derived is True + + +def test_asserted_fingerprint_is_preserved_on_recollection(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1"), "baseline", config_fingerprint="asserted-cfg") + enriched = make_report("1") + enriched.job_parameters = {"shuffle_partitions": "800"} + refreshed, _ = record(exp_dir, enriched, "baseline") + assert refreshed.config_fingerprint == "asserted-cfg" + assert refreshed.config_fingerprint_derived is False + + +def test_recollection_keeps_first_trial_ordering(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1"), "baseline") + record(exp_dir, make_report("2"), "candidate") + record(exp_dir, make_report("1"), "baseline") + manifest, snapshots = exp.load_experiment(exp_dir) + rows = exp.trial_rows(manifest, snapshots) + assert [row["run_id"] for row in rows] == ["1", "2"] + assert rows[0]["snapshot_collected_at"] is not None + + +def test_per_task_comparison_uses_medians_and_keeps_spread(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1", execution_ms=1000), "baseline") + record(exp_dir, make_report("2", execution_ms=3000), "baseline") + record(exp_dir, make_report("3", execution_ms=500), "candidate") + comparison = exp.compare_experiment( + exp_dir, baseline_variant="baseline", candidate_variant="candidate" + ) + task = next(t for t in comparison.tasks if t.task_key == "ingest") + assert task.baseline["execution_ms"] == 2000 + assert task.baseline_values["execution_ms"] == [1000.0, 3000.0] + assert task.baseline_min["execution_ms"] == 1000 + assert task.baseline_max["execution_ms"] == 3000 + assert task.candidate["execution_ms"] == 500 + assert task.deltas["execution_ms"] == -1500 + + +def test_unequal_sample_counts_use_medians(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1", execution_ms=1000), "baseline") + record(exp_dir, make_report("2", execution_ms=3000), "baseline") + record(exp_dir, make_report("3", execution_ms=500), "candidate") + comparison = exp.compare_experiment( + exp_dir, baseline_variant="baseline", candidate_variant="candidate" + ) + task = next(m for m in comparison.metrics if m.metric == "aggregate_task_time_ms") + assert task.baseline_median == 2000 + assert task.candidate_median == 500 + assert len(task.baseline_values) == 2 + assert len(task.candidate_values) == 1 + + +def test_missing_input_snapshot_is_a_caveat_not_a_refusal(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1"), "baseline") + record(exp_dir, make_report("2"), "candidate", input_snapshot="user-declared") + comparison = exp.compare_experiment( + exp_dir, baseline_run_id="1", candidate_run_id="2" + ) + input_note = next( + n for n in comparison.comparability if n.aspect == "input_snapshot" + ) + assert input_note.status == "unknown" + assert comparison.metrics # raw deltas are still shown + + +def test_trial_rows_are_in_execution_order(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("2"), "baseline") + record(exp_dir, make_report("1"), "candidate") + manifest, snapshots = exp.load_experiment(exp_dir) + rows = exp.trial_rows(manifest, snapshots) + assert len(rows) == 2 + timestamps = [row["collected_at"] for row in rows] + assert timestamps == sorted(timestamps) + assert all("workflow_elapsed_ms" in row["values"] for row in rows) + + +def test_out_of_order_runs_compare_by_selector(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("2", read_bytes=200), "baseline") + record(exp_dir, make_report("1", read_bytes=100), "baseline") + comparison = exp.compare_experiment( + exp_dir, baseline_run_id="2", candidate_run_id="1" + ) + read = next(m for m in comparison.metrics if m.metric == "read_bytes") + assert read.baseline_median == 200 + assert read.candidate_median == 100 + + +def test_manifest_rejects_unknown_schema(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1"), "baseline") + manifest_data = json.loads((exp_dir / exp.MANIFEST_NAME).read_text()) + manifest_data["schema_version"] = "999" + (exp_dir / exp.MANIFEST_NAME).write_text(json.dumps(manifest_data)) + with pytest.raises(exp.ExperimentError): + exp.load_manifest(exp_dir) + + +def test_error_text_is_preserved_as_data(tmp_path): + from sparkparse.models import OutputAttachment + + exp_dir = tmp_path / "exp" + report = make_report("1") + report.outputs = [ + OutputAttachment( + task_key="join", + run_id="1001", + status="error", + error="", + excerpt="", + ) + ] + path = exp.save_snapshot( + exp_dir, exp.build_snapshot(report, trial=exp.TrialMetadata(variant="v")) + ) + saved = path.read_text() + assert "" + assert attachment.path is not None + assert (exp_dir / attachment.path).read_text() == "" + + +def test_snapshot_records_identity_and_schema(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1"), "baseline") + paths = exp.list_snapshot_paths(exp_dir) + snapshot = exp.load_snapshot(paths[0]) + assert snapshot.run_id == "1" + assert snapshot.schema_version == "1" + assert snapshot.workspace_host is not None + + +# --- CLI contracts ------------------------------------------------------- + +from typer.testing import CliRunner # noqa: E402 + +from sparkparse.app import app # noqa: E402 + +runner = CliRunner() + + +def test_compare_cli_requires_one_selector_per_side(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1"), "baseline") + result = runner.invoke(app, ["experiments", "compare", str(exp_dir)]) + assert result.exit_code == 2 + + +def test_compare_cli_json(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1"), "baseline") + record(exp_dir, make_report("2"), "candidate") + result = runner.invoke( + app, + [ + "experiments", + "compare", + str(exp_dir), + "--baseline", + "1", + "--candidate", + "2", + ], + ) + assert result.exit_code == 0, result.output + payload = json.loads(result.stdout) + assert payload["baseline"]["run_ids"] == ["1"] + assert payload["candidate"]["run_ids"] == ["2"] + + +def test_compare_cli_missing_experiment_exits_1(tmp_path): + result = runner.invoke( + app, + [ + "experiments", + "compare", + str(tmp_path / "nope"), + "--baseline", + "1", + "--candidate", + "2", + ], + ) + assert result.exit_code == 1 + + +def test_comparison_records_pinned_snapshot_ids(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1"), "baseline") + record(exp_dir, make_report("2"), "candidate") + comparison = exp.compare_experiment( + exp_dir, baseline_run_id="1", candidate_run_id="2" + ) + assert len(comparison.baseline.snapshot_ids) == 1 + assert len(comparison.candidate.snapshot_ids) == 1 + # Snapshot IDs are stable references, not the mutable latest pointer. + manifest, _ = exp.load_experiment(exp_dir) + by_run = {t.run_id: t.snapshot_id for t in manifest.trials} + assert comparison.baseline.snapshot_ids[0] == by_run["1"] diff --git a/tests/test_experiments_dashboard.py b/tests/test_experiments_dashboard.py new file mode 100644 index 0000000..18a955e --- /dev/null +++ b/tests/test_experiments_dashboard.py @@ -0,0 +1,195 @@ +from pathlib import Path +from typing import Any + +import pytest + +from sparkparse import experiments as exp +from sparkparse.experiments_dashboard import ( + _context_table, + _filter_rows, + _task_trend_figure, + _trend_figure, + build_layout, + default_runs, + init_experiments_dashboard, +) +from tests.test_experiments import make_report, record + +pytest.importorskip("dash") + + +def _experiment(tmp_path: Path) -> Path: + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1", execution_ms=1000), "baseline") + record(exp_dir, make_report("2", execution_ms=3000), "baseline") + record(exp_dir, make_report("3", execution_ms=500), "candidate") + return exp_dir + + +def test_default_runs_honor_pinned_baseline(tmp_path): + exp_dir = _experiment(tmp_path) + manifest, snapshots = exp.load_experiment(exp_dir) + rows = exp.trial_rows(manifest, snapshots) + baseline, candidate = default_runs(manifest, rows) + assert manifest.baseline_trial_id is not None + assert baseline == "1" + assert candidate == "3" + + +def test_filter_rows_by_variant_commit_and_eligibility(tmp_path): + exp_dir = _experiment(tmp_path) + manifest, snapshots = exp.load_experiment(exp_dir) + rows = exp.trial_rows(manifest, snapshots) + assert len(_filter_rows(rows, ["baseline"], None, None)) == 2 + assert len(_filter_rows(rows, ["candidate"], None, None)) == 1 + assert _filter_rows(rows, None, None, ["failed"]) == [] + + +def test_trend_figure_keeps_missing_context_as_a_gap(): + rows: list[dict[str, Any]] = [ + { + "run_id": "1", + "variant": "v", + "short_revision": "abc", + "values": {"workflow_elapsed_ms": 1000.0}, + "partial_metrics": [], + } + ] + figure = _trend_figure(rows, "workflow_elapsed_ms", None, None) + context_trace = figure.data[1] + assert list(context_trace.y) == [None] # type: ignore[missing-attribute] + + +def test_trend_figure_marks_partial_values_distinctly(): + rows: list[dict[str, Any]] = [ + { + "run_id": "1", + "variant": "v", + "short_revision": "abc", + "values": {"read_bytes": 100.0}, + "partial_metrics": ["read_bytes"], + } + ] + figure = _trend_figure(rows, "read_bytes", None, None) + partial = [ + trace + for trace in figure.data + if "partial" in (trace.name or "") # type: ignore[missing-attribute] + ] + assert partial + assert partial[0].marker.symbol == "circle-open" # type: ignore[missing-attribute,union-attr] + + +def test_context_table_surfaces_compute_and_revision(tmp_path): + exp_dir = _experiment(tmp_path) + manifest, snapshots = exp.load_experiment(exp_dir) + rows = exp.trial_rows(manifest, snapshots) + comparison = exp.compare_experiment( + exp_dir, baseline_run_id="1", candidate_run_id="3" + ) + baseline_row = next(r for r in rows if r["run_id"] == "1") + candidate_row = next(r for r in rows if r["run_id"] == "3") + table = _context_table(comparison, baseline_row, candidate_row) + aspects = {row["aspect"] for row in table} + assert {"revision", "compute", "configuration"} <= aspects + compute = next(row for row in table if row["aspect"] == "compute") + assert compute["baseline"] == baseline_row["compute_summary"] + + +def test_trial_rows_expose_detail_fields(tmp_path): + exp_dir = _experiment(tmp_path) + manifest, snapshots = exp.load_experiment(exp_dir) + rows = exp.trial_rows(manifest, snapshots) + row = rows[0] + assert row["eligibility"] == "ok" + assert "compute_summary" in row + assert "top_queries" in row + assert "failure_excerpts" in row + assert "ingest" in row["task_values"] + + +def test_trend_figure_context_aligns_with_all_runs(): + rows: list[dict[str, Any]] = [ + { + "run_id": "1", + "variant": "v", + "short_revision": "abc", + "values": {"read_bytes": 100.0, "workflow_elapsed_ms": 1000.0}, + "partial_metrics": ["read_bytes"], + }, + { + "run_id": "2", + "variant": "v", + "short_revision": "abc", + "values": {"read_bytes": 200.0, "workflow_elapsed_ms": 2000.0}, + "partial_metrics": [], + }, + ] + figure = _trend_figure(rows, "read_bytes", None, None) + context = next( + t + for t in figure.data + if t.type == "bar" # type: ignore[missing-attribute] + ) + assert list(context.x) == ["1", "2"] # type: ignore[missing-attribute] + assert len(context.y) == 2 # type: ignore[missing-attribute] + + +def test_task_spread_excludes_ineligible_trials(tmp_path): + exp_dir = tmp_path / "exp" + record(exp_dir, make_report("1", execution_ms=1000), "baseline") + record(exp_dir, make_report("2", execution_ms=3000), "baseline") + record( + exp_dir, + make_report("4", execution_ms=50000, result_state="FAILED"), + "baseline", + ) + record(exp_dir, make_report("3", execution_ms=500), "candidate") + manifest, snapshots = exp.load_experiment(exp_dir) + rows = exp.trial_rows(manifest, snapshots) + figure = _task_trend_figure(rows, "ingest", "1", "3") + spread = [ + trace + for trace in figure.data + if trace.name and "spread" in trace.name # type: ignore[missing-attribute,operator] + ] + baseline_spread = spread[0] + # Median of 1000 and 3000 is 2000; max - median is 1000, not 48000. + assert baseline_spread.error_y.array[0] == 1000 # type: ignore[missing-attribute] + + +def test_flatten_trend_rows_supports_csv(): + pl = pytest.importorskip("polars") + from sparkparse.experiments_dashboard import flatten_trend_rows + + rows = [ + {"run_id": "1", "workflow_elapsed_ms": 1.0, "partial_metrics": ["read_bytes"]}, + {"run_id": "2", "workflow_elapsed_ms": 2.0, "partial_metrics": []}, + ] + frame = pl.DataFrame(flatten_trend_rows(rows)) + assert frame.height == 2 + assert "read_bytes" in frame["partial_metrics"].to_list()[0] + + +def test_task_trend_figure_has_group_spread_error_bars(tmp_path): + exp_dir = _experiment(tmp_path) + manifest, snapshots = exp.load_experiment(exp_dir) + rows = exp.trial_rows(manifest, snapshots) + figure = _task_trend_figure(rows, "ingest", "1", "3") + spread = [ + trace + for trace in figure.data + if trace.name and "spread" in trace.name # type: ignore[missing-attribute,operator] + ] + assert spread + # Run 1 is in the baseline variant with samples 1000ms and 3000ms. + assert spread[0].error_y.array[0] > 0 # type: ignore[missing-attribute] + + +def test_dashboard_builds_without_caching_rows(tmp_path): + exp_dir = _experiment(tmp_path) + app = init_experiments_dashboard(str(exp_dir)) + assert app.layout is not None + assert "EXPERIMENT_ROWS" not in app.server.config + layout = build_layout(exp.trial_rows(*exp.load_experiment(exp_dir)), "1", "3") + assert layout is not None