diff --git a/.github/workflows/yarn-resource-cost-python.yml b/.github/workflows/yarn-resource-cost-python.yml new file mode 100644 index 0000000..627fd13 --- /dev/null +++ b/.github/workflows/yarn-resource-cost-python.yml @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +name: YARN resource cost Python checks + +on: + pull_request: + paths: + - "yarn-resource-cost/**" + - ".github/workflows/yarn-resource-cost-python.yml" + push: + branches: [dev] + paths: + - "yarn-resource-cost/**" + +jobs: + test: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.11", "3.12"] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Compile + run: python -m py_compile yarn-resource-cost/*.py + - name: Unit tests + run: python -m unittest discover -s yarn-resource-cost -p 'test*.py' + - name: Standalone bundle smoke test + run: python yarn-resource-cost/package_yarn_job_cost.py --output-dir /tmp/yarn-resource-cost-dist diff --git a/README.md b/README.md index 9b3da8f..97e8866 100644 --- a/README.md +++ b/README.md @@ -9,3 +9,9 @@ A repo for Spark related benchmark sets and utilities using the Please see README in each benchmark set for more details including building instructions and usage descriptions. + +## Utilities + +- [Portable YARN resource cost](./yarn-resource-cost/) attributes Spark + application worker consumption from Spark and YARN logs across EMR, + Dataproc, and on-premises deployments. diff --git a/yarn-resource-cost/.gitignore b/yarn-resource-cost/.gitignore new file mode 100644 index 0000000..84d346b --- /dev/null +++ b/yarn-resource-cost/.gitignore @@ -0,0 +1,4 @@ +__pycache__/ +*.py[cod] +dist/ +.cache/ diff --git a/yarn-resource-cost/CONTRIBUTING.md b/yarn-resource-cost/CONTRIBUTING.md new file mode 100644 index 0000000..ab17037 --- /dev/null +++ b/yarn-resource-cost/CONTRIBUTING.md @@ -0,0 +1,16 @@ +# Contributing + +Contributions follow the repository-level `CONTRIBUTING.md`, including the +Developer's Certificate of Origin 1.1 and signed-off commits. + +Do not contribute customer event logs, cluster logs, bucket names, account +identifiers, internal URLs, or pricing agreements. Tests must use synthetic or +explicitly sanitized fixtures. + +Before submitting a change, run: + +```bash +python3 -m py_compile yarn-resource-cost/*.py +python3 -m unittest discover -s yarn-resource-cost -p 'test*.py' +python3 yarn-resource-cost/package_yarn_job_cost.py +``` diff --git a/yarn-resource-cost/README.md b/yarn-resource-cost/README.md new file mode 100644 index 0000000..bed3114 --- /dev/null +++ b/yarn-resource-cost/README.md @@ -0,0 +1,180 @@ +# Portable YARN resource cost + +Spark elapsed time, task duration, YARN vcore-seconds, and Spark +executor-core-seconds answer different questions. None of them consistently +describes the fraction of worker nodes that YARN allocated to an application. +For example, `DefaultResourceCalculator` schedules by memory even when Spark +advertises several executor cores. Comparing raw vcore-seconds across clusters +can therefore undercount memory-heavy containers or compare unrelated YARN +accounting units. + +This tool reconstructs container lifetimes and allocations from ResourceManager +logs, converts them into node-equivalent seconds, and optionally applies an +auditable hourly rate. Spark event logs select applications and provide names, +wall-clock duration, executor/container joins, and task-packing diagnostics; +they are not treated as the allocation ledger. + +## Requirements + +- Python 3.10 or newer; the accounting code has no third-party Python packages. +- Archived Spark rolling event logs. +- ResourceManager logs containing node registrations, allocations, terminal + transitions, application summaries, and ResourceCalculator evidence. +- NodeManager logs are accepted only as incomplete fallback evidence. +- The provider CLI is needed only for remote discovery: AWS CLI for EMR, + `gcloud` for Dataproc, or `hdfs` for on-premises HDFS paths. + +## Accounting model + +For a container lasting `ContainerSeconds` on a node: + +```text +DefaultResourceCalculator: + NodeShare = ContainerMemoryMB / NodeMemoryMB + +DominantResourceCalculator: + NodeShare = max(ContainerResource[r] / NodeResource[r]) for every r + +NodeEquivalentSeconds = ContainerSeconds * NodeShare +``` + +Memory, vcores, `yarn.io/gpu`, and arbitrary numeric custom resources are +parsed generically. Heterogeneous node classes remain separate in structured +output and expressions such as: + +```text +123.4 aws:ec2:r7a.24xlarge-seconds + 50.0 aws:ec2:g6.4xlarge-seconds +``` + +The calculator comes from archived scheduler evidence. There is intentionally +no calculator override. Built-in FairScheduler policies map to their actual +memory or dominant-resource calculators; ambiguous or conflicting policy +evidence is not guessed. + +Final worker cost is emitted only for complete ledgers. Missing RM allocations, +terminal transitions, node capacities, application summaries, event-log +segments, node classifications, or price-catalog entries suppress final cost. + +## Usage + +Resource accounting is offline by default: + +```bash +python3 yarn-resource-cost/yarn_resource_cost.py \ + --adapter emr \ + --event-log-root s3://example-bucket/run/spark-events/ \ + --pricing none +``` + +EMR live on-demand pricing includes EC2 and EMR worker components and records +the lookup result and timestamp: + +```bash +python3 yarn-resource-cost/yarn_resource_cost.py \ + --adapter emr \ + --event-log-root s3://example-bucket/run/spark-events/ \ + --pricing live \ + --aws-profile example-profile \ + --aws-region us-west-2 \ + --output-json emr-cost.json +``` + +Dataproc accepts `gs://` Spark logs and exported RM/NM daemon logs. A node-class +map separates machine and accelerator shapes: + +```bash +python3 yarn-resource-cost/yarn_resource_cost.py \ + --adapter dataproc \ + --event-log-root gs://example-bucket/spark-events/ \ + --yarn-log-root gs://example-bucket/cluster-daemon-logs/ \ + --node-class-map dataproc-node-classes.json \ + --pricing catalog \ + --price-catalog dataproc-prices.json +``` + +On-premises inputs can be directories, ZIP/tar archives, or HDFS paths: + +```bash +python3 yarn-resource-cost/yarn_resource_cost.py \ + --adapter on-prem \ + --event-log-root /archive/spark-events \ + --yarn-log-root /archive/yarn-daemon-logs.tar.gz \ + --node-class-map node-classes.json \ + --pricing catalog \ + --price-catalog internal-rates.json \ + --output-csv application-cost.csv +``` + +Compare a baseline with one or more test reruns. Later test roots replace earlier +ones with the same comparison key: + +```bash +python3 yarn-resource-cost/yarn_resource_cost.py \ + --adapter emr \ + --event-log-root s3://example-bucket/baseline/events/ \ + --test-event-log-root s3://example-bucket/test/events/ \ + --comparison-key regex \ + --comparison-key-regex 'job-(?P[0-9]+)' \ + --pricing live \ + --sort-by cost-factor +``` + +For non-EMR comparisons, pair every `--test-event-log-root` with a +`--test-yarn-log-root` in the same order. + +## Input schemas + +Node classes are adapter-stable worker shapes, not hostnames: + +```json +{ + "schema_version": 1, + "nodes": { + "worker-01.example.net": "onprem:gpu-a10-16c-128g" + }, + "default_node_class": "onprem:cpu-32c-256g" +} +``` + +Pricing is deliberately separate from accounting: + +```json +{ + "schema_version": 1, + "currency": "USD", + "effective_at": "2026-08-01T00:00:00Z", + "source": "approved internal rate card", + "rates": [ + {"node_class": "onprem:cpu-32c-256g", "hourly_rate": 4.25} + ] +} +``` + +JSON output uses `schema_version: 1` and retains input roots, discovery source, +calculator evidence, node capacities, resource expressions, completeness, +warnings, pricing provenance, applications, and run summaries. + +## Reproducible experiment capture + +Archive these together for every benchmark run: + +1. All Spark event-log segments through `SparkListenerApplicationEnd`. +2. All rolled ResourceManager and NodeManager daemon logs for the cluster. +3. `yarn-site.xml`, `resource-types.xml`, and the active + `capacity-scheduler.xml` or `fair-scheduler.xml`. +4. Provider cluster descriptions and host-to-node-class mappings. +5. A frozen price catalog when results must be reproducible later. +6. A small manifest linking each event-log root to its YARN log root and node + map. Never rely on a live cluster remaining available. + +## Compatibility command + +`calculate_yarn_job_cost.py` preserves the original EMR-oriented command and +output columns for existing integrations. New integrations should use +`yarn_resource_cost.py` and its provider-neutral versioned JSON output. + +## License and contribution + +The tool is licensed under Apache License 2.0 as part of this repository. See +the repository `LICENSE` and `CONTRIBUTING.md`. Generated standalone bundles +include both documents and a checksum manifest. diff --git a/yarn-resource-cost/RESOURCE_COST_MODEL.md b/yarn-resource-cost/RESOURCE_COST_MODEL.md new file mode 100644 index 0000000..397f2aa --- /dev/null +++ b/yarn-resource-cost/RESOURCE_COST_MODEL.md @@ -0,0 +1,43 @@ +# YARN resource-proportional worker cost model + +For every completed executor container `c` on node class `h`: + +```text +DurationSeconds(c) = (authoritative RM finish - RM allocation) / 1000 +``` + +With `DefaultResourceCalculator`: + +```text +Share(c, h) = AllocatedMemoryMB(c) / AdvertisedMemoryMB(h) +``` + +With `DominantResourceCalculator`: + +```text +Share(c, h) = max over allocated resources r of + Allocated(c, r) / Advertised(h, r) +``` + +The unpriced result retains every heterogeneous node class: + +```text +NodeClassSeconds(h) = sum(DurationSeconds(c) * Share(c, h)) +``` + +Given an auditable hourly worker rate: + +```text +WorkerCost = sum(NodeClassSeconds(h) * HourlyRate(h) / 3600) +``` + +ApplicationMaster containers are excluded by default. Wall-clock duration, +successful-task duration, perfect-packing estimates, Spark executor cores, and +raw YARN vcore-seconds are diagnostics only. They never replace the allocation +ledger or determine the node share. + +An application is complete only when the selected Spark log is complete and +the RM evidence accounts for every allocated executor container, its terminal +time, its node, all allocated resource capacities, and its stable node class. +An incomplete application retains resource evidence and warnings but has no +final price or price factor. diff --git a/yarn-resource-cost/calculate_yarn_job_cost.py b/yarn-resource-cost/calculate_yarn_job_cost.py new file mode 100644 index 0000000..200bb6c --- /dev/null +++ b/yarn-resource-cost/calculate_yarn_job_cost.py @@ -0,0 +1,2107 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Calculate per-application cost from a Spark event-log run and EMR YARN logs. + +The preferred input is one Spark event-log root, for example: + + s3://bucket/experiments/RUN/spark-events/benchmark/ + +The event logs identify the selected applications and EMR cluster. The EMR API +then resolves the cluster archived log URI. ResourceManager logs remain the +authoritative source for container accounting. An exact EMR log URI is also +accepted for offline or diagnostic use. +""" + +from __future__ import annotations + +import argparse +import csv +import gzip +import hashlib +import json +import os +import re +import subprocess +import sys +import tempfile +from collections import Counter, defaultdict +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import TextIO + +SCRIPT_DIR = Path(__file__).resolve().parent +sys.path.insert(0, str(SCRIPT_DIR)) +from yarn_job_cost_discovery import ( # noqa: E402 + AwsCliError, + EventLogApplication, + load_csv_metadata, + materialize_event_metadata_files, + read_event_log_metadata, + resolve_emr_log_uri, + run_aws, +) +from yarn_job_cost_defaults import ( # noqa: E402 + DEFAULT_AWS_PROFILE, + DEFAULT_AWS_REGION, +) + +from yarn_job_cost_adapters import ( # noqa: E402 + ADAPTERS, + AdapterCommandError, + apply_catalog_costs, + apply_node_class_map, + load_price_catalog, + materialize_gcs, + materialize_hdfs, + materialize_yarn_logs, +) + +TIMESTAMP = r"(?P\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2},\d{3})" +CONTAINER = r"(?Pcontainer_\d+_\d+_\d+_\d+)" +START_RE = re.compile( + rf"^{TIMESTAMP} .*Start request for {CONTAINER} .* resource " + r"\d+)(?:, max memory:(?P\d+))?, " + r"vCores:(?P\d+)(?:, max vCores:(?P\d+))?" + r"(?P[^>]*)>" +) +DONE_RE = re.compile(rf"^{TIMESTAMP} .*Container {CONTAINER} transitioned from .* to DONE\b") +RM_ASSIGN_RE = re.compile( + rf"^{TIMESTAMP} .*Assigned container {CONTAINER} of capacity " + r"\d+)(?:, max memory:(?P\d+))?, " + r"vCores:(?P\d+)(?:, max vCores:(?P\d+))?" + r"(?P[^>]*)> " + r"on host (?P[^:,\s]+):\d+" +) +RM_TERMINAL_RE = re.compile( + rf"^{TIMESTAMP} .*{CONTAINER} Container Transitioned from .* to (?:COMPLETED|RELEASED|KILLED|EXPIRED)\b" +) +APPLICATION_SUMMARY_RE = re.compile( + r"appId=(?Papplication_\d+_\d+),name=(?P.*?),user=.*?" + r"finalStatus=(?P[^,]+).*?totalAllocatedContainers=(?P\d+)" +) +NODE_RE = re.compile( + r"Registered with ResourceManager .* total resource of " + r"\d+), vCores:(?P\d+)(?P[^>]*)>.*" + r"instanceType\(STRING\)=(?P[^}\]]+)" +) +RM_NODE_RE = re.compile( + r"NodeManager from node (?P[^(]+)\(cmPort: \d+ httpPort: \d+\) " + r"registered with capability: \d+), " + r"vCores:(?P\d+)(?P[^>]*)>.*" + r"instanceType\(STRING\)=(?P[^}\]\s]+)" +) +GPU_RESOURCE_RE = re.compile(r"(?:^|,\s*)yarn\.io/gpu:\s*(?P\d+)") +CALCULATOR_RE = re.compile( + r"Initialized CapacityScheduler with calculator=class " + r"org\.apache\.hadoop\.yarn\.util\.resource\.(?P\w+)" +) +NODE_ID_RE = re.compile(r"(?:^|/)node/(?Pi-[^/]+)/") +CONTAINER_PARTS_RE = re.compile( + r"container_(?P\d+)_(?P\d+)_(?P\d+)_(?P\d+)" +) +OUTPUT_FIELDS = ( + "job id", + "job name", + "application_id", + "application_name", + "spark_duration_seconds", + "emr_cluster_id", + "emr_release_label", + "spark_version", + "configured_spark_executor_cores", + "spark_task_cpus", + "final_status", + "resource_calculator", + "container_count", + "expected_total_allocated_containers", + "observed_total_allocated_containers", + "incomplete_container_count", + "nodemanager_start_fallback_container_count", + "nodemanager_finish_fallback_container_count", + "container_seconds", + "memory_mb_seconds", + "vcore_seconds", + "gpu_seconds", + "node_equivalent_seconds", + "instance_vcore_seconds", + "instance_vcore_seconds_expression", + "task_metrics_complete", + "perfect_packing_complete", + "successful_task_attempt_count", + "task_duration_sum_seconds", + "perfect_packing_node_seconds", + "perfect_packing_cost_expression", + "perfect_packing_ec2_ondemand_usd", + "perfect_packing_emr_usd", + "perfect_packing_ec2_plus_emr_usd", + "packing_efficiency", + "actual_to_perfect_cost_factor", + "actual_to_perfect_cost_overhead_percent", + "cost_expression", + "ec2_ondemand_usd", + "emr_usd", + "ec2_plus_emr_usd", + "first_container_start_utc", + "last_container_finish_utc", + "complete", + "task_metric_warnings", + "warnings", +) + + +CONSOLE_FIELDS = ( + ("Job ID", "job id"), + ("Application", "application_id"), + ("Status", "final_status"), + ("Spark sec", "spark_duration_seconds"), + ("Containers", "container_count"), + ("GPU sec", "gpu_seconds"), + ("Node-equivalent sec", "node_equivalent_seconds"), + ("Instance-vcore", "instance_vcore_seconds_expression"), + ("Task sec", "task_duration_sum_seconds"), + ("Perfect USD", "perfect_packing_ec2_plus_emr_usd"), + ("Packing", "packing_efficiency"), + ("Actual/perfect", "actual_to_perfect_cost_factor"), + ("Cost", "cost_expression"), + ("EC2+EMR USD", "ec2_plus_emr_usd"), + ("Complete", "complete"), + ("Warnings", "warnings"), +) + + +COMPARISON_FIELDS = ( + "job_id", + "baseline_application_id", + "test_application_id", + "baseline_instance_type", + "test_instance_type", + "baseline_event_log_root", + "test_event_log_root", + "baseline_final_status", + "test_final_status", + "baseline_complete", + "test_complete", + "baseline_warnings", + "test_warnings", + "baseline_wall_clock_seconds", + "test_wall_clock_seconds", + "wall_clock_factor", + "wall_clock_delta_seconds", + "baseline_task_metrics_complete", + "test_task_metrics_complete", + "baseline_perfect_packing_complete", + "test_perfect_packing_complete", + "baseline_task_duration_sum_seconds", + "test_task_duration_sum_seconds", + "task_duration_factor", + "baseline_perfect_packing_node_seconds", + "test_perfect_packing_node_seconds", + "baseline_perfect_packing_ec2_plus_emr_usd", + "test_perfect_packing_ec2_plus_emr_usd", + "perfect_packing_cost_factor", + "perfect_packing_cost_delta_usd", + "baseline_packing_efficiency", + "test_packing_efficiency", + "packing_efficiency_delta", + "baseline_actual_to_perfect_cost_factor", + "test_actual_to_perfect_cost_factor", + "baseline_actual_to_perfect_cost_overhead_percent", + "test_actual_to_perfect_cost_overhead_percent", + "baseline_node_equivalent_seconds", + "test_node_equivalent_seconds", + "baseline_instance_vcore_seconds", + "test_instance_vcore_seconds", + "instance_vcore_seconds_factor", + "instance_vcore_seconds_delta", + "node_equivalent_delta_seconds", + "baseline_ec2_ondemand_usd", + "test_ec2_ondemand_usd", + "ec2_ondemand_cost_factor", + "ec2_ondemand_delta_usd", + "baseline_emr_usd", + "test_emr_usd", + "baseline_ec2_plus_emr_usd", + "test_ec2_plus_emr_usd", + "ec2_plus_emr_cost_factor", + "ec2_plus_emr_delta_usd", +) +COMPARISON_CONSOLE_FIELDS = ( + ("Job ID", "job_id"), + ("Base instance", "baseline_instance_type"), + ("Test instance", "test_instance_type"), + ("Base complete", "baseline_complete"), + ("Test complete", "test_complete"), + ("Base sec", "baseline_wall_clock_seconds"), + ("Test sec", "test_wall_clock_seconds"), + ("Wall factor", "wall_clock_factor"), + ("Wall delta", "wall_clock_delta_seconds"), + ("Task factor", "task_duration_factor"), + ("Perfect factor", "perfect_packing_cost_factor"), + ("Base packing", "baseline_packing_efficiency"), + ("Test packing", "test_packing_efficiency"), + ("Base actual/perfect", "baseline_actual_to_perfect_cost_factor"), + ("Test actual/perfect", "test_actual_to_perfect_cost_factor"), + ("Base node-sec", "baseline_node_equivalent_seconds"), + ("Test node-sec", "test_node_equivalent_seconds"), + ("Node-sec delta", "node_equivalent_delta_seconds"), + ("Base instance-vcore sec", "baseline_instance_vcore_seconds"), + ("Test instance-vcore sec", "test_instance_vcore_seconds"), + ("Instance-vcore factor", "instance_vcore_seconds_factor"), + ("Base EC2+EMR USD", "baseline_ec2_plus_emr_usd"), + ("Test EC2+EMR USD", "test_ec2_plus_emr_usd"), + ("Cost factor", "ec2_plus_emr_cost_factor"), + ("USD delta", "ec2_plus_emr_delta_usd"), +) + + +@dataclass +class Node: + node_id: str + instance_type: str = "" + memory_mb: int | None = None + vcores: int | None = None + gpus: int = 0 + + +@dataclass +class Container: + container_id: str + application_id: str + node_id: str + start_ms: int + memory_mb: int + node_memory_mb: int + vcores: int + node_vcores: int + gpus: int = 0 + node_gpus: int = 0 + finish_ms: int | None = None + finish_source: str = "" + source: str = "nodemanager" + + @property + def sequence(self) -> int: + match = CONTAINER_PARTS_RE.fullmatch(self.container_id) + if not match: + raise ValueError(f"Unexpected container ID {self.container_id}") + return int(match.group("sequence")) + + +@dataclass +class ApplicationSummary: + application_id: str + name: str + final_status: str + total_allocated_containers: int + + + + +@dataclass +class YarnEvidence: + nodes: dict[str, Node] = field(default_factory=dict) + containers: dict[str, Container] = field(default_factory=dict) + calculator_class: str = "" + application_summaries: dict[str, ApplicationSummary] = field(default_factory=dict) + warnings: list[str] = field(default_factory=list) + + +def gpu_amount(resources: str | None) -> int: + match = GPU_RESOURCE_RE.search(resources or "") + return int(match.group("gpus")) if match else 0 + + +def capacity(value: str | None, fallback: int | None, name: str) -> int: + if value is not None: + return int(value) + if fallback is not None: + return fallback + raise ValueError(f"Could not determine node {name} capacity") + + +def parse_timestamp(value: str) -> int: + parsed = datetime.strptime(value, "%Y-%m-%d %H:%M:%S,%f").replace(tzinfo=timezone.utc) + return int(parsed.timestamp() * 1000) + + +def iso_utc(timestamp_ms: int | None) -> str: + if timestamp_ms is None: + return "" + return datetime.fromtimestamp(timestamp_ms / 1000, tz=timezone.utc).isoformat().replace("+00:00", "Z") + + +def application_id(container_id: str) -> str: + match = CONTAINER_PARTS_RE.fullmatch(container_id) + if not match: + raise ValueError(f"Unexpected container ID {container_id}") + return f"application_{match.group('cluster')}_{match.group('application')}" + + +def node_id_from_path(path: Path) -> str: + match = NODE_ID_RE.search(path.as_posix()) + return match.group("node_id") if match else path.parent.name + + +def open_log(path: Path) -> TextIO: + if path.suffix == ".gz": + return gzip.open(path, mode="rt", encoding="utf-8", errors="replace") + return path.open(encoding="utf-8", errors="replace") + + +def relevant_log_files(path: Path) -> list[Path]: + if path.is_file(): + return [path] + return sorted( + file + for file in path.rglob("*") + if file.is_file() + and ( + "hadoop-yarn-nodemanager" in file.name + or "hadoop-yarn-resourcemanager" in file.name + ) + and ".log" in file.name + ) + + +def parse_yarn_logs(path: Path) -> YarnEvidence: + evidence = YarnEvidence() + rm_finishes: dict[str, int] = {} + nm_finishes: dict[str, int] = {} + files = relevant_log_files(path) + if not files: + raise ValueError(f"No NodeManager or ResourceManager log files found under {path}") + + for file in files: + node_id = node_id_from_path(file) + node = evidence.nodes.setdefault(node_id, Node(node_id=node_id)) + with open_log(file) as handle: + for line in handle: + if "ApplicationSummary" in line and "totalAllocatedContainers=" in line: + match = APPLICATION_SUMMARY_RE.search(line) + if match: + app_id = match.group("application") + evidence.application_summaries[app_id] = ApplicationSummary( + application_id=app_id, + name=match.group("name"), + final_status=match.group("final_status"), + total_allocated_containers=int(match.group("containers")), + ) + continue + if "Initialized CapacityScheduler with calculator=" in line: + match = CALCULATOR_RE.search(line) + if match: + calculator = match.group("calculator") + if evidence.calculator_class and evidence.calculator_class != calculator: + raise ValueError( + f"Conflicting ResourceCalculators: " + f"{evidence.calculator_class}, {calculator}" + ) + evidence.calculator_class = calculator + continue + if "Registered with ResourceManager" in line: + match = NODE_RE.search(line) + if match: + node.instance_type = match.group("instance_type").strip() + node.memory_mb = int(match.group("memory")) + node.vcores = int(match.group("vcores")) + node.gpus = gpu_amount(match.group("resources")) + continue + if "registered with capability:" in line and "NodeManager from node " in line: + match = RM_NODE_RE.search(line) + if match: + host = match.group("host").strip() + evidence.nodes[host] = Node( + node_id=host, + instance_type=match.group("instance_type").strip(), + memory_mb=int(match.group("memory")), + vcores=int(match.group("vcores")), + gpus=gpu_amount(match.group("resources")), + ) + continue + if "Assigned container container_" in line: + match = RM_ASSIGN_RE.search(line) + if not match: + continue + container_id = match.group("container") + host = match.group("host") + registered_node = evidence.nodes.get(host) + evidence.containers[container_id] = Container( + container_id=container_id, + application_id=application_id(container_id), + node_id=host, + start_ms=parse_timestamp(match.group("timestamp")), + memory_mb=int(match.group("memory")), + node_memory_mb=capacity( + match.group("max_memory"), + registered_node.memory_mb if registered_node else None, + "memory", + ), + vcores=int(match.group("vcores")), + node_vcores=capacity( + match.group("max_vcores"), + registered_node.vcores if registered_node else None, + "vcore", + ), + gpus=gpu_amount(match.group("resources")), + node_gpus=registered_node.gpus if registered_node else 0, + source="resourcemanager", + ) + continue + if " Container Transitioned from " in line: + match = RM_TERMINAL_RE.search(line) + if match: + container_id = match.group("container") + finish_ms = parse_timestamp(match.group("timestamp")) + existing = rm_finishes.get(container_id) + rm_finishes[container_id] = ( + finish_ms if existing is None else min(existing, finish_ms) + ) + continue + if "Start request for container_" in line: + match = START_RE.search(line) + if not match: + continue + container_id = match.group("container") + candidate = Container( + container_id=container_id, + application_id=application_id(container_id), + node_id=node_id, + start_ms=parse_timestamp(match.group("timestamp")), + memory_mb=int(match.group("memory")), + node_memory_mb=capacity( + match.group("max_memory"), node.memory_mb, "memory" + ), + vcores=int(match.group("vcores")), + node_vcores=capacity( + match.group("max_vcores"), node.vcores, "vcore" + ), + gpus=gpu_amount(match.group("resources")), + node_gpus=node.gpus, + ) + existing = evidence.containers.get(container_id) + if existing is None: + evidence.containers[container_id] = candidate + continue + if " to DONE" in line and "Container container_" in line: + match = DONE_RE.search(line) + if match: + container_id = match.group("container") + finish_ms = parse_timestamp(match.group("timestamp")) + existing = nm_finishes.get(container_id) + nm_finishes[container_id] = ( + finish_ms if existing is None else min(existing, finish_ms) + ) + + for container_id, container in evidence.containers.items(): + if container_id in rm_finishes: + container.finish_ms = rm_finishes[container_id] + container.finish_source = "resourcemanager" + elif container_id in nm_finishes: + container.finish_ms = nm_finishes[container_id] + container.finish_source = "nodemanager" + for node_id, node in evidence.nodes.items(): + if any(container.node_id == node_id for container in evidence.containers.values()): + if not node.instance_type: + evidence.warnings.append(f"{node_id}: instance type was not found in registration logs") + return evidence + + +def materialize_emr_logs( + uri: str, + cache_dir: Path, + aws_profile: str | None, + refresh: bool = False, +) -> Path: + if not uri.startswith(("s3://", "s3a://", "s3n://")): + path = Path(uri).expanduser() + if not path.exists(): + raise FileNotFoundError(path) + return path + + s3_uri = "s3://" + uri.split("://", 1)[1] + digest = hashlib.sha256(s3_uri.rstrip("/").encode()).hexdigest()[:16] + cluster_name = s3_uri.rstrip("/").rsplit("/", 1)[-1] + target = cache_dir / f"{cluster_name}-{digest}" + marker = target / ".download-complete" + if marker.is_file() and not refresh: + return target + if refresh: + marker.unlink(missing_ok=True) + target.mkdir(parents=True, exist_ok=True) + command = ["aws"] + if aws_profile: + command += ["--profile", aws_profile] + command += [ + "s3", + "cp", + "--recursive", + s3_uri.rstrip("/") + "/", + str(target), + "--exclude", + "*", + "--include", + "*hadoop-yarn-nodemanager*.log*", + "--include", + "*hadoop-yarn-resourcemanager*.log*", + ] + subprocess.run(command, check=True) + if not relevant_log_files(target): + raise ValueError(f"No YARN logs downloaded from {s3_uri}") + marker.write_text(s3_uri + "\n") + return target + + +def aws_command(aws_profile: str | None, aws_region: str | None = None) -> list[str]: + command = ["aws"] + if aws_profile: + command += ["--profile", aws_profile] + if aws_region: + command += ["--region", aws_region] + return command + + +def resolve_aws_region( + explicit_region: str | None, aws_profile: str | None +) -> str: + if explicit_region: + return explicit_region + for variable in ("AWS_REGION", "AWS_DEFAULT_REGION"): + value = os.environ.get(variable, "").strip() + if value: + return value + command = aws_command(aws_profile) + ["configure", "get", "region"] + completed = subprocess.run(command, capture_output=True, text=True) + configured = completed.stdout.strip() if completed.returncode == 0 else "" + if configured: + return configured + profile_hint = f" for profile {aws_profile!r}" if aws_profile else "" + raise ValueError( + "AWS region is not configured" + profile_hint + "; pass --aws-region " + "or set AWS_REGION/AWS_DEFAULT_REGION" + ) + +def calculator_mode(detected_class: str) -> str: + if detected_class == "DefaultResourceCalculator": + return "default" + if detected_class == "DominantResourceCalculator": + return "dominant" + raise ValueError( + "Could not auto-detect DefaultResourceCalculator or " + "DominantResourceCalculator from the ResourceManager log" + ) + + +def container_node_share(container: Container, mode: str) -> float: + memory_share = container.memory_mb / container.node_memory_mb + vcore_share = container.vcores / container.node_vcores + gpu_share = ( + container.gpus / container.node_gpus if container.node_gpus else 0.0 + ) + if mode == "default": + return memory_share + if mode == "dominant": + return max(memory_share, vcore_share, gpu_share) + raise ValueError(f"Unsupported detected calculator mode {mode}") + + +def container_instance_type(evidence: YarnEvidence, container: Container) -> str: + node = evidence.nodes.get(container.node_id) + if node and node.instance_type: + return node.instance_type + return f"unknown:{container.node_id}" + + +# Use the provider-neutral ledger while retaining the established output layer. +from yarn_job_cost_core import ( # noqa: E402 + Container as CoreContainer, + Node as CoreNode, + YarnEvidence as CoreYarnEvidence, + calculator_mode as core_calculator_mode, + container_instance_type as core_container_instance_type, + container_node_share as core_container_node_share, + parse_yarn_logs as core_parse_yarn_logs, + relevant_log_files as core_relevant_log_files, +) + +Container = CoreContainer +Node = CoreNode +YarnEvidence = CoreYarnEvidence +calculator_mode = core_calculator_mode +container_instance_type = core_container_instance_type +container_node_share = core_container_node_share +parse_yarn_logs = core_parse_yarn_logs +relevant_log_files = core_relevant_log_files + +def calculate_applications( + evidence: YarnEvidence, + mode: str, + csv_metadata: dict[str, dict[str, str]], + include_application_master: bool, + known_executor_container_ids: set[str] | None = None, +) -> list[dict]: + executor_container_ids = known_executor_container_ids or set() + by_application: dict[str, list[Container]] = defaultdict(list) + total_containers_by_application = Counter( + container.application_id for container in evidence.containers.values() + ) + for container in evidence.containers.values(): + if ( + not include_application_master + and container.sequence == 1 + and container.container_id not in executor_container_ids + ): + continue + by_application[container.application_id].append(container) + + results = [] + for app_id, containers in sorted(by_application.items()): + complete_containers = [ + container for container in containers if container.finish_ms is not None + ] + warnings = list(evidence.warnings) + sequence_one_executors = [ + container.container_id + for container in containers + if container.sequence == 1 + and container.container_id in executor_container_ids + ] + if sequence_one_executors: + warnings.append( + "Included sequence-1 container because Spark identifies it " + "as an executor: " + ", ".join(sequence_one_executors) + ) + summary = evidence.application_summaries.get(app_id) + observed_total = total_containers_by_application[app_id] + coverage_complete = ( + summary is not None and summary.total_allocated_containers == observed_total + ) + if summary is None: + warnings.append("ResourceManager ApplicationSummary is missing") + elif summary.total_allocated_containers != observed_total: + warnings.append( + f"ResourceManager summary reports {summary.total_allocated_containers} " + f"allocated containers but logs contain {observed_total}" + ) + incomplete = len(containers) - len(complete_containers) + if incomplete: + warnings.append( + f"{incomplete} container(s) have no terminal timestamp and were omitted" + ) + nm_start_fallbacks = sum( + container.source != "resourcemanager" for container in containers + ) + if nm_start_fallbacks: + warnings.append( + f"{nm_start_fallbacks} container allocation(s) use a " + "NodeManager start fallback" + ) + nm_finish_fallbacks = sum( + container.finish_source == "nodemanager" + for container in complete_containers + ) + if nm_finish_fallbacks: + warnings.append( + f"{nm_finish_fallbacks} container terminal timestamp(s) use a " + "NodeManager DONE fallback" + ) + + by_instance_type: dict[str, float] = defaultdict(float) + by_instance_vcore_seconds: dict[str, float] = defaultdict(float) + container_seconds = 0.0 + memory_mb_seconds = 0.0 + vcore_seconds = 0.0 + gpu_seconds = 0.0 + missing_gpu_capacity = any( + container.gpus > 0 and container.node_gpus <= 0 + for container in complete_containers + ) + if missing_gpu_capacity: + warnings.append("A GPU allocation has no registered node GPU capacity") + resource_capacity_errors: list[str] = [] + for container in complete_containers: + duration = max(0.0, (container.finish_ms - container.start_ms) / 1000.0) + try: + share = container_node_share(container, mode) + except ValueError as error: + resource_capacity_errors.append(str(error)) + share = 0.0 + instance_type = container_instance_type(evidence, container) + by_instance_type[instance_type] += duration * share + by_instance_vcore_seconds[instance_type] += ( + duration * share * container.node_vcores + ) + container_seconds += duration + memory_mb_seconds += duration * container.memory_mb + vcore_seconds += duration * container.vcores + gpu_seconds += duration * container.gpus + + warnings.extend(resource_capacity_errors) + expression = " + ".join( + f"{seconds:.6f} {instance_type}-seconds" + for instance_type, seconds in sorted(by_instance_type.items()) + ) + instance_vcore_expression = " + ".join( + f"{seconds:.6f} {instance_type}-vcore-seconds" + for instance_type, seconds in sorted( + by_instance_vcore_seconds.items() + ) + ) + unknown_instance_type = any( + instance_type.startswith("unknown:") for instance_type in by_instance_type + ) + if unknown_instance_type: + warnings.append("One or more allocated containers have an unknown instance type") + starts = [container.start_ms for container in containers] + finishes = [ + container.finish_ms for container in complete_containers if container.finish_ms is not None + ] + result = { + **csv_metadata.get(app_id, {"job id": "", "job name": ""}), + "application_id": app_id, + "application_name": summary.name if summary else "", + "final_status": summary.final_status if summary else "", + "resource_calculator": mode, + "detected_resource_calculator_class": evidence.calculator_class, + "container_count": len(containers), + "expected_total_allocated_containers": ( + summary.total_allocated_containers if summary else "" + ), + "observed_total_allocated_containers": observed_total, + "incomplete_container_count": incomplete, + "nodemanager_start_fallback_container_count": nm_start_fallbacks, + "nodemanager_finish_fallback_container_count": nm_finish_fallbacks, + "container_seconds": round(container_seconds, 6), + "memory_mb_seconds": round(memory_mb_seconds, 6), + "vcore_seconds": round(vcore_seconds, 6), + "gpu_seconds": round(gpu_seconds, 6), + "node_equivalent_seconds": round(sum(by_instance_type.values()), 6), + "node_equivalent_seconds_by_instance_type": dict(sorted(by_instance_type.items())), + "instance_vcore_seconds": round( + sum(by_instance_vcore_seconds.values()), 6 + ), + "instance_vcore_seconds_by_instance_type": dict( + sorted(by_instance_vcore_seconds.items()) + ), + "instance_vcore_seconds_expression": instance_vcore_expression, + "cost_expression": expression, + "first_container_start_utc": iso_utc(min(starts) if starts else None), + "last_container_finish_utc": iso_utc(max(finishes) if finishes else None), + "complete": ( + incomplete == 0 + and coverage_complete + and nm_start_fallbacks == 0 + and nm_finish_fallbacks == 0 + and not unknown_instance_type + and not missing_gpu_capacity + and not resource_capacity_errors + and not evidence.accounting_policy_ambiguous + ), + "warnings": warnings, + } + results.append(result) + return results + + +def add_task_packing_metrics( + applications: list[dict], + event_metadata: dict[str, EventLogApplication], + evidence: YarnEvidence, + mode: str, +) -> None: + for application in applications: + application.update( + { + "task_metrics_complete": False, + "perfect_packing_complete": False, + "successful_task_attempt_count": "", + "task_duration_sum_seconds": "", + "perfect_packing_node_seconds": "", + "perfect_packing_node_seconds_by_instance_type": {}, + "perfect_packing_cost_expression": "", + "perfect_packing_ec2_ondemand_usd": "", + "perfect_packing_emr_usd": "", + "perfect_packing_ec2_plus_emr_usd": "", + "packing_efficiency": "", + "actual_to_perfect_cost_factor": "", + "actual_to_perfect_cost_overhead_percent": "", + "task_metric_warnings": [], + } + ) + app_id = application["application_id"] + metadata = event_metadata.get(app_id) + warnings: list[str] = [] + if metadata is None: + warnings.append("Spark event log is required for task metrics") + else: + application["successful_task_attempt_count"] = ( + metadata.successful_task_attempt_count + ) + warnings.extend(metadata.task_metric_warnings) + if metadata.event_task_metrics_complete(): + application["task_metrics_complete"] = True + application["task_duration_sum_seconds"] = round( + metadata.task_duration_sum_ms / 1000.0, 6 + ) + if not application.get("complete"): + warnings.append( + "Complete YARN accounting is required for " + "perfect-packing metrics" + ) + else: + by_instance_type: dict[str, float] = defaultdict(float) + mapping_warnings = [] + for executor_id, duration_ms in ( + metadata.task_duration_ms_by_executor.items() + ): + executor = metadata.executors.get(executor_id) + if executor is None: + mapping_warnings.append( + f"Executor {executor_id} has tasks but no " + "ExecutorAdded event" + ) + continue + if not executor.container_id: + mapping_warnings.append( + f"Executor {executor_id} has no YARN container ID" + ) + continue + if executor.total_cores is None or executor.total_cores <= 0: + mapping_warnings.append( + f"Executor {executor_id} has invalid total cores" + ) + continue + container = evidence.containers.get(executor.container_id) + if container is None: + mapping_warnings.append( + f"Executor {executor_id} container " + f"{executor.container_id} is missing from YARN logs" + ) + continue + if container.application_id != app_id: + mapping_warnings.append( + f"Executor {executor_id} maps to a non-executor " + "YARN container" + ) + continue + instance_type = container_instance_type( + evidence, container + ) + if instance_type.startswith("unknown:"): + mapping_warnings.append( + f"Executor {executor_id} has an unknown instance type" + ) + continue + perfect_executor_seconds = ( + duration_ms + / 1000.0 + * metadata.task_cpus + / executor.total_cores + ) + by_instance_type[instance_type] += ( + perfect_executor_seconds + * container_node_share(container, mode) + ) + warnings.extend(mapping_warnings) + if not mapping_warnings: + perfect_node_seconds = sum(by_instance_type.values()) + actual_node_seconds = float( + application["node_equivalent_seconds"] + ) + application["perfect_packing_complete"] = True + application[ + "perfect_packing_node_seconds_by_instance_type" + ] = { + instance_type: round(seconds, 6) + for instance_type, seconds in sorted( + by_instance_type.items() + ) + } + application["perfect_packing_node_seconds"] = round( + perfect_node_seconds, 6 + ) + application["perfect_packing_cost_expression"] = ( + " + ".join( + f"{seconds:.6f} {instance_type}-seconds" + for instance_type, seconds in sorted( + by_instance_type.items() + ) + ) + ) + if actual_node_seconds > 0: + efficiency = ( + perfect_node_seconds / actual_node_seconds + ) + application["packing_efficiency"] = round( + efficiency, 8 + ) + if efficiency > 1.05: + warnings.append( + "Packing efficiency exceeds 1.05; Spark " + "and YARN clock boundaries may be inconsistent" + ) + warnings = list(dict.fromkeys(warnings)) + application["task_metric_warnings"] = warnings + application["warnings"].extend( + f"Task metrics: {warning}" for warning in warnings + ) + + +def current_ondemand_hourly_price( + instance_type: str, product_region: str, aws_profile: str | None +) -> dict: + filters = [ + f"Type=TERM_MATCH,Field=instanceType,Value={instance_type}", + f"Type=TERM_MATCH,Field=regionCode,Value={product_region}", + "Type=TERM_MATCH,Field=operatingSystem,Value=Linux", + "Type=TERM_MATCH,Field=tenancy,Value=Shared", + "Type=TERM_MATCH,Field=preInstalledSw,Value=NA", + "Type=TERM_MATCH,Field=capacitystatus,Value=Used", + ] + command = aws_command(aws_profile, "us-east-1") + [ + "pricing", + "get-products", + "--service-code", + "AmazonEC2", + "--filters", + *filters, + "--max-results", + "100", + "--output", + "json", + ] + completed = run_aws(command) + payload = json.loads(completed.stdout) + matches = [] + for encoded_product in payload.get("PriceList") or []: + product = json.loads(encoded_product) + if product.get("product", {}).get("productFamily") != "Compute Instance": + continue + for term in product.get("terms", {}).get("OnDemand", {}).values(): + for dimension in term.get("priceDimensions", {}).values(): + if dimension.get("unit") != "Hrs" or dimension.get("beginRange") != "0": + continue + matches.append( + { + "usd_per_hour": float(dimension["pricePerUnit"]["USD"]), + "effective_date": term.get("effectiveDate", ""), + "description": dimension.get("description", ""), + } + ) + rates = {match["usd_per_hour"] for match in matches} + if len(rates) != 1: + raise ValueError( + f"Expected one current Linux on-demand rate for {instance_type} in " + f"{product_region}; found {sorted(rates)}" + ) + return matches[0] + + +def current_emr_hourly_price( + instance_type: str, product_region: str, aws_profile: str | None +) -> dict: + filters = [ + f"Type=TERM_MATCH,Field=instanceType,Value={instance_type}", + f"Type=TERM_MATCH,Field=regionCode,Value={product_region}", + ] + command = aws_command(aws_profile, "us-east-1") + [ + "pricing", + "get-products", + "--service-code", + "ElasticMapReduce", + "--filters", + *filters, + "--max-results", + "100", + "--output", + "json", + ] + completed = run_aws(command) + payload = json.loads(completed.stdout) + matches = [] + for encoded_product in payload.get("PriceList") or []: + product = json.loads(encoded_product) + for term in product.get("terms", {}).get("OnDemand", {}).values(): + for dimension in term.get("priceDimensions", {}).values(): + if ( + dimension.get("unit") != "Hrs" + or dimension.get("beginRange") != "0" + ): + continue + matches.append( + { + "usd_per_hour": float( + dimension["pricePerUnit"]["USD"] + ), + "effective_date": term.get("effectiveDate", ""), + "description": dimension.get("description", ""), + } + ) + rates = {match["usd_per_hour"] for match in matches} + if len(rates) != 1: + raise ValueError( + f"Expected one current EMR rate for {instance_type} in " + f"{product_region}; found {sorted(rates)}" + ) + return matches[0] + + +def instance_price_catalog( + instance_types: set[str], product_region: str, aws_profile: str | None +) -> dict[str, dict[str, dict]]: + return { + "ec2": { + instance_type: current_ondemand_hourly_price( + instance_type, product_region, aws_profile + ) + for instance_type in sorted(instance_types) + }, + "emr": { + instance_type: current_emr_hourly_price( + instance_type, product_region, aws_profile + ) + for instance_type in sorted(instance_types) + }, + } + + +def application_resource_cost( + application: dict, prices: dict[str, dict], seconds_field: str +) -> float: + return sum( + float(seconds) * prices[instance_type]["usd_per_hour"] / 3600.0 + for instance_type, seconds in application[seconds_field].items() + ) + + +def application_ondemand_cost(application: dict, prices: dict[str, dict]) -> float: + return application_resource_cost( + application, prices, "node_equivalent_seconds_by_instance_type" + ) + + +def perfect_packing_ondemand_cost( + application: dict, prices: dict[str, dict] +) -> float: + return application_resource_cost( + application, + prices, + "perfect_packing_node_seconds_by_instance_type", + ) + + +def add_ondemand_costs( + applications: list[dict], product_region: str, aws_profile: str +) -> dict[str, dict[str, dict]]: + instance_types = { + instance_type + for application in applications + if application.get("complete") is True + for instance_type in application[ + "node_equivalent_seconds_by_instance_type" + ] + } + prices = instance_price_catalog( + instance_types, product_region, aws_profile + ) + for application in applications: + if application.get("complete") is True: + ec2_cost = application_ondemand_cost(application, prices["ec2"]) + emr_cost = application_ondemand_cost(application, prices["emr"]) + application["ec2_ondemand_usd"] = round(ec2_cost, 8) + application["emr_usd"] = round(emr_cost, 8) + application["ec2_plus_emr_usd"] = round( + ec2_cost + emr_cost, 8 + ) + if application.get("perfect_packing_complete") is True: + perfect_ec2 = perfect_packing_ondemand_cost( + application, prices["ec2"] + ) + perfect_emr = perfect_packing_ondemand_cost( + application, prices["emr"] + ) + perfect_total = perfect_ec2 + perfect_emr + actual_total = ec2_cost + emr_cost + application["perfect_packing_ec2_ondemand_usd"] = round( + perfect_ec2, 8 + ) + application["perfect_packing_emr_usd"] = round( + perfect_emr, 8 + ) + application["perfect_packing_ec2_plus_emr_usd"] = round( + perfect_total, 8 + ) + if perfect_total > 0: + factor = actual_total / perfect_total + application["actual_to_perfect_cost_factor"] = round( + factor, 8 + ) + application[ + "actual_to_perfect_cost_overhead_percent" + ] = round((factor - 1.0) * 100.0, 6) + else: + application["ec2_ondemand_usd"] = "" + application["emr_usd"] = "" + application["ec2_plus_emr_usd"] = "" + return prices + + +def application_warnings(application: dict | None) -> str: + if not application: + return "" + warnings = application.get("warnings", []) + return " | ".join(warnings) if isinstance(warnings, list) else str(warnings) + + +def application_instance_types(application: dict | None) -> str: + if not application: + return "" + by_instance_type = application.get( + "node_equivalent_seconds_by_instance_type", {} + ) + return ", ".join(sorted(by_instance_type)) + + +def index_applications_by_job_id(applications: list[dict], label: str) -> dict[str, dict]: + indexed = {} + for application in applications: + job_id = str(application.get("job id") or "") + if not job_id: + raise ValueError( + f"{label} application {application.get('application_id')} has no derived Job ID" + ) + if job_id in indexed: + raise ValueError(f"{label} has duplicate Job ID {job_id}") + indexed[job_id] = application + return indexed + + +def rounded_delta(test: object, baseline: object, digits: int = 6) -> float | str: + if test == "" or baseline == "" or test is None or baseline is None: + return "" + return round(float(test) - float(baseline), digits) + + +def rounded_factor( + test: object, baseline: object, digits: int = 6 +) -> float | str: + if ( + test == "" + or baseline == "" + or test is None + or baseline is None + ): + return "" + baseline_value = float(baseline) + if baseline_value == 0: + return "" + return round(float(test) / baseline_value, digits) + + +def sort_comparison_rows(rows: list[dict], sort_by: str) -> list[dict]: + if sort_by == "job-id": + return sorted(rows, key=lambda row: int(row["job_id"])) + field = { + "wall-clock-factor": "wall_clock_factor", + "cost-factor": "ec2_plus_emr_cost_factor", + "task-duration-factor": "task_duration_factor", + "perfect-packing-cost-factor": "perfect_packing_cost_factor", + }[sort_by] + return sorted( + rows, + key=lambda row: ( + row[field] == "" + or row["baseline_final_status"] != "SUCCEEDED" + or row["test_final_status"] != "SUCCEEDED", + float(row[field]) if row[field] != "" else float("inf"), + int(row["job_id"]), + ), + ) + + +def build_comparison_rows( + baseline_applications: list[dict], + test_applications: list[dict], + prices: dict[str, dict], +) -> list[dict]: + baseline = index_applications_by_job_id(baseline_applications, "baseline") + test = index_applications_by_job_id(test_applications, "test") + job_ids = sorted(set(baseline) | set(test), key=lambda value: int(value)) + rows = [] + for job_id in job_ids: + base = baseline.get(job_id) + other = test.get(job_id) + base_wall = base.get("spark_duration_seconds", "") if base else "" + other_wall = other.get("spark_duration_seconds", "") if other else "" + base_node = base.get("node_equivalent_seconds", "") if base else "" + other_node = other.get("node_equivalent_seconds", "") if other else "" + base_instance_vcore = ( + base.get("instance_vcore_seconds", "") if base else "" + ) + test_instance_vcore = ( + other.get("instance_vcore_seconds", "") if other else "" + ) + base_complete = base.get("complete") is True if base else False + test_complete = other.get("complete") is True if other else False + base_task_complete = ( + base.get("task_metrics_complete") is True if base else False + ) + test_task_complete = ( + other.get("task_metrics_complete") is True if other else False + ) + base_perfect_complete = ( + base.get("perfect_packing_complete") is True if base else False + ) + test_perfect_complete = ( + other.get("perfect_packing_complete") is True if other else False + ) + base_task_duration = ( + base.get("task_duration_sum_seconds", "") if base else "" + ) + test_task_duration = ( + other.get("task_duration_sum_seconds", "") if other else "" + ) + base_perfect_node = ( + base.get("perfect_packing_node_seconds", "") if base else "" + ) + test_perfect_node = ( + other.get("perfect_packing_node_seconds", "") if other else "" + ) + base_ec2_cost = ( + application_ondemand_cost(base, prices["ec2"]) + if base_complete + else "" + ) + other_ec2_cost = ( + application_ondemand_cost(other, prices["ec2"]) + if test_complete + else "" + ) + base_emr_cost = ( + application_ondemand_cost(base, prices["emr"]) + if base_complete + else "" + ) + other_emr_cost = ( + application_ondemand_cost(other, prices["emr"]) + if test_complete + else "" + ) + base_total_cost = ( + base_ec2_cost + base_emr_cost if base_complete else "" + ) + other_total_cost = ( + other_ec2_cost + other_emr_cost if test_complete else "" + ) + base_perfect_cost = ( + perfect_packing_ondemand_cost(base, prices["ec2"]) + + perfect_packing_ondemand_cost(base, prices["emr"]) + if base_complete and base_perfect_complete + else "" + ) + test_perfect_cost = ( + perfect_packing_ondemand_cost(other, prices["ec2"]) + + perfect_packing_ondemand_cost(other, prices["emr"]) + if test_complete and test_perfect_complete + else "" + ) + base_actual_to_perfect = ( + base_total_cost / base_perfect_cost + if base_perfect_cost not in ("", 0) + else "" + ) + test_actual_to_perfect = ( + other_total_cost / test_perfect_cost + if test_perfect_cost not in ("", 0) + else "" + ) + rows.append( + { + "job_id": job_id, + "baseline_application_id": base.get("application_id", "") if base else "", + "test_application_id": other.get("application_id", "") if other else "", + "baseline_instance_type": application_instance_types(base), + "test_instance_type": application_instance_types(other), + "baseline_event_log_root": ( + base.get("source_event_log_root", "") if base else "" + ), + "test_event_log_root": ( + other.get("source_event_log_root", "") if other else "" + ), + "baseline_final_status": base.get("final_status", "") if base else "MISSING", + "test_final_status": other.get("final_status", "") if other else "MISSING", + "baseline_complete": base_complete if base else "", + "test_complete": test_complete if other else "", + "baseline_warnings": application_warnings(base), + "test_warnings": application_warnings(other), + "baseline_wall_clock_seconds": base_wall, + "test_wall_clock_seconds": other_wall, + "wall_clock_factor": rounded_factor(other_wall, base_wall), + "wall_clock_delta_seconds": rounded_delta(other_wall, base_wall), + "baseline_task_metrics_complete": ( + base_task_complete if base else "" + ), + "test_task_metrics_complete": ( + test_task_complete if other else "" + ), + "baseline_perfect_packing_complete": ( + base_perfect_complete if base else "" + ), + "test_perfect_packing_complete": ( + test_perfect_complete if other else "" + ), + "baseline_task_duration_sum_seconds": base_task_duration, + "test_task_duration_sum_seconds": test_task_duration, + "task_duration_factor": rounded_factor( + test_task_duration, base_task_duration + ), + "baseline_perfect_packing_node_seconds": base_perfect_node, + "test_perfect_packing_node_seconds": test_perfect_node, + "baseline_perfect_packing_ec2_plus_emr_usd": ( + round(base_perfect_cost, 8) + if base_perfect_cost != "" + else "" + ), + "test_perfect_packing_ec2_plus_emr_usd": ( + round(test_perfect_cost, 8) + if test_perfect_cost != "" + else "" + ), + "perfect_packing_cost_factor": rounded_factor( + test_perfect_cost, base_perfect_cost + ), + "perfect_packing_cost_delta_usd": rounded_delta( + test_perfect_cost, base_perfect_cost, 8 + ), + "baseline_packing_efficiency": ( + base.get("packing_efficiency", "") if base else "" + ), + "test_packing_efficiency": ( + other.get("packing_efficiency", "") if other else "" + ), + "packing_efficiency_delta": rounded_delta( + other.get("packing_efficiency", "") if other else "", + base.get("packing_efficiency", "") if base else "", + 8, + ), + "baseline_actual_to_perfect_cost_factor": ( + round(base_actual_to_perfect, 8) + if base_actual_to_perfect != "" + else "" + ), + "test_actual_to_perfect_cost_factor": ( + round(test_actual_to_perfect, 8) + if test_actual_to_perfect != "" + else "" + ), + "baseline_actual_to_perfect_cost_overhead_percent": ( + round((base_actual_to_perfect - 1.0) * 100.0, 6) + if base_actual_to_perfect != "" + else "" + ), + "test_actual_to_perfect_cost_overhead_percent": ( + round((test_actual_to_perfect - 1.0) * 100.0, 6) + if test_actual_to_perfect != "" + else "" + ), + "baseline_node_equivalent_seconds": base_node, + "test_node_equivalent_seconds": other_node, + "node_equivalent_delta_seconds": rounded_delta(other_node, base_node), + "baseline_instance_vcore_seconds": base_instance_vcore, + "test_instance_vcore_seconds": test_instance_vcore, + "instance_vcore_seconds_factor": rounded_factor( + test_instance_vcore if test_complete else "", + base_instance_vcore if base_complete else "", + ), + "instance_vcore_seconds_delta": rounded_delta( + test_instance_vcore if test_complete else "", + base_instance_vcore if base_complete else "", + ), + "baseline_ec2_ondemand_usd": ( + round(base_ec2_cost, 8) if base_ec2_cost != "" else "" + ), + "test_ec2_ondemand_usd": ( + round(other_ec2_cost, 8) if other_ec2_cost != "" else "" + ), + "ec2_ondemand_cost_factor": rounded_factor( + other_ec2_cost, base_ec2_cost + ), + "ec2_ondemand_delta_usd": rounded_delta( + other_ec2_cost, base_ec2_cost, 8 + ), + "baseline_emr_usd": ( + round(base_emr_cost, 8) if base_emr_cost != "" else "" + ), + "test_emr_usd": ( + round(other_emr_cost, 8) if other_emr_cost != "" else "" + ), + "baseline_ec2_plus_emr_usd": ( + round(base_total_cost, 8) + if base_total_cost != "" + else "" + ), + "test_ec2_plus_emr_usd": ( + round(other_total_cost, 8) + if other_total_cost != "" + else "" + ), + "ec2_plus_emr_cost_factor": rounded_factor( + other_total_cost, base_total_cost + ), + "ec2_plus_emr_delta_usd": rounded_delta( + other_total_cost, base_total_cost, 8 + ), + } + ) + return rows + + +def summarize_applications(applications: list[dict]) -> dict: + successful_costed = [ + application + for application in applications + if application.get("final_status") == "SUCCEEDED" + and application.get("complete") is True + ] + successful_incomplete_count = sum( + application.get("final_status") == "SUCCEEDED" + and application.get("complete") is not True + for application in applications + ) + instance_vcore_seconds_by_type: dict[str, float] = defaultdict(float) + for application in successful_costed: + for instance_type, seconds in application.get( + "instance_vcore_seconds_by_instance_type", {} + ).items(): + instance_vcore_seconds_by_type[instance_type] += float(seconds) + total_instance_vcore_seconds = sum( + instance_vcore_seconds_by_type.values() + ) + total_instance_vcore_expression = " + ".join( + f"{seconds:.6f} {instance_type}-vcore-seconds" + for instance_type, seconds in sorted( + instance_vcore_seconds_by_type.items() + ) + ) + + eligible: list[dict] = [] + excluded = [] + for application in applications: + reasons = [] + if application.get("final_status") != "SUCCEEDED": + reasons.append("application did not succeed") + if application.get("complete") is not True: + reasons.append("YARN accounting is incomplete") + if application.get("task_metrics_complete") is not True: + reasons.append("task metrics are incomplete") + if application.get("perfect_packing_complete") is not True: + reasons.append("perfect-packing metrics are incomplete") + if reasons: + excluded.append( + { + "job_id": str(application.get("job id") or ""), + "application_id": application.get("application_id", ""), + "reasons": reasons, + } + ) + else: + eligible.append(application) + + actual_node_seconds = sum( + float(application["node_equivalent_seconds"]) + for application in eligible + ) + perfect_node_seconds = sum( + float(application["perfect_packing_node_seconds"]) + for application in eligible + ) + actual_cost = sum( + float(application["ec2_plus_emr_usd"]) for application in eligible + ) + perfect_cost = sum( + float(application["perfect_packing_ec2_plus_emr_usd"]) + for application in eligible + ) + actual_to_perfect = ( + actual_cost / perfect_cost if perfect_cost > 0 else None + ) + return { + "eligible_job_count": len(eligible), + "excluded_job_count": len(excluded), + "excluded_jobs": excluded, + "successful_complete_job_count": len(successful_costed), + "successful_incomplete_job_count": successful_incomplete_count, + "total_instance_vcore_seconds": round( + total_instance_vcore_seconds, 6 + ), + "total_instance_vcore_seconds_by_instance_type": { + instance_type: round(seconds, 6) + for instance_type, seconds in sorted( + instance_vcore_seconds_by_type.items() + ) + }, + "total_instance_vcore_seconds_expression": ( + total_instance_vcore_expression or "0 instance-vcore-seconds" + ), + "successful_task_attempt_count": sum( + int(application["successful_task_attempt_count"]) + for application in eligible + ), + "task_duration_sum_seconds": round( + sum( + float(application["task_duration_sum_seconds"]) + for application in eligible + ), + 6, + ), + "perfect_packing_node_seconds": round(perfect_node_seconds, 6), + "actual_node_equivalent_seconds": round(actual_node_seconds, 6), + "perfect_packing_ec2_plus_emr_usd": round(perfect_cost, 8), + "actual_ec2_plus_emr_usd": round(actual_cost, 8), + "packing_efficiency": ( + round(perfect_node_seconds / actual_node_seconds, 8) + if actual_node_seconds > 0 + else None + ), + "actual_to_perfect_cost_factor": ( + round(actual_to_perfect, 8) + if actual_to_perfect is not None + else None + ), + "actual_to_perfect_cost_overhead_percent": ( + round((actual_to_perfect - 1.0) * 100.0, 6) + if actual_to_perfect is not None + else None + ), + } + + +def build_comparison_summary(rows: list[dict]) -> dict: + eligible = [] + excluded = [] + for row in rows: + reasons = [] + if row.get("baseline_final_status") != "SUCCEEDED": + reasons.append("baseline application did not succeed") + if row.get("test_final_status") != "SUCCEEDED": + reasons.append("test application did not succeed") + if row.get("baseline_complete") is not True: + reasons.append("baseline YARN accounting is incomplete") + if row.get("test_complete") is not True: + reasons.append("test YARN accounting is incomplete") + if row.get("baseline_task_metrics_complete") is not True: + reasons.append("baseline task metrics are incomplete") + if row.get("test_task_metrics_complete") is not True: + reasons.append("test task metrics are incomplete") + if row.get("baseline_perfect_packing_complete") is not True: + reasons.append("baseline perfect-packing metrics are incomplete") + if row.get("test_perfect_packing_complete") is not True: + reasons.append("test perfect-packing metrics are incomplete") + if reasons: + excluded.append({"job_id": row["job_id"], "reasons": reasons}) + else: + eligible.append(row) + + def side(prefix: str) -> dict: + task_seconds = sum( + float(row[f"{prefix}_task_duration_sum_seconds"]) + for row in eligible + ) + perfect_node_seconds = sum( + float(row[f"{prefix}_perfect_packing_node_seconds"]) + for row in eligible + ) + actual_node_seconds = sum( + float(row[f"{prefix}_node_equivalent_seconds"]) + for row in eligible + ) + perfect_cost = sum( + float(row[f"{prefix}_perfect_packing_ec2_plus_emr_usd"]) + for row in eligible + ) + actual_cost = sum( + float(row[f"{prefix}_ec2_plus_emr_usd"]) + for row in eligible + ) + factor = actual_cost / perfect_cost if perfect_cost > 0 else None + return { + "task_duration_sum_seconds": round(task_seconds, 6), + "perfect_packing_node_seconds": round(perfect_node_seconds, 6), + "actual_node_equivalent_seconds": round(actual_node_seconds, 6), + "perfect_packing_ec2_plus_emr_usd": round(perfect_cost, 8), + "actual_ec2_plus_emr_usd": round(actual_cost, 8), + "packing_efficiency": ( + round(perfect_node_seconds / actual_node_seconds, 8) + if actual_node_seconds > 0 + else None + ), + "actual_to_perfect_cost_factor": ( + round(factor, 8) if factor is not None else None + ), + "actual_to_perfect_cost_overhead_percent": ( + round((factor - 1.0) * 100.0, 6) + if factor is not None + else None + ), + } + + baseline = side("baseline") + test = side("test") + return { + "eligible_job_count": len(eligible), + "excluded_job_count": len(excluded), + "excluded_jobs": excluded, + "baseline": baseline, + "test": test, + "task_duration_factor": rounded_factor( + test["task_duration_sum_seconds"], + baseline["task_duration_sum_seconds"], + ), + "perfect_packing_cost_factor": rounded_factor( + test["perfect_packing_ec2_plus_emr_usd"], + baseline["perfect_packing_ec2_plus_emr_usd"], + ), + "actual_cost_factor": rounded_factor( + test["actual_ec2_plus_emr_usd"], + baseline["actual_ec2_plus_emr_usd"], + ), + } + + +def print_summary(summary: dict, heading: str = "Summary") -> None: + print( + "{}: {} packing-eligible job(s), {} excluded".format( + heading, + summary["eligible_job_count"], + summary["excluded_job_count"], + ) + ) + for name in ( + "successful_complete_job_count", + "successful_incomplete_job_count", + "total_instance_vcore_seconds", + "total_instance_vcore_seconds_expression", + "task_duration_sum_seconds", + "perfect_packing_node_seconds", + "actual_node_equivalent_seconds", + "perfect_packing_ec2_plus_emr_usd", + "actual_ec2_plus_emr_usd", + "packing_efficiency", + "actual_to_perfect_cost_factor", + "actual_to_perfect_cost_overhead_percent", + ): + print(" {}: {}".format(name, summary[name])) + + +def print_comparison_summary(summary: dict) -> None: + print( + "Comparison summary: {} eligible job(s), {} excluded".format( + summary["eligible_job_count"], + summary["excluded_job_count"], + ) + ) + for name in ( + "task_duration_factor", + "perfect_packing_cost_factor", + "actual_cost_factor", + ): + print(f" {name}: {summary[name]}") + for label, side in (("Baseline", "baseline"), ("Test", "test")): + print(f" {label} totals:") + for name, value in summary[side].items(): + print(f" {name}: {value}") + + +def analyze_event_root_for_comparison(root: str, args: argparse.Namespace) -> dict: + with tempfile.TemporaryDirectory() as directory: + output = Path(directory) / "analysis.json" + command = [ + sys.executable, + str(Path(__file__).resolve()), + "--event-log-root", + root, + "--output-json", + str(output), + "--aws-profile", + args.aws_profile or "", + "--aws-region", + args.aws_region, + "--cache-dir", + str(args.cache_dir), + "--event-cache-dir", + str(args.event_cache_dir), + ] + if args.include_application_master: + command.append("--include-application-master") + if args.refresh_emr_log_cache: + command.append("--refresh-emr-log-cache") + completed = subprocess.run(command, capture_output=True, text=True) + if completed.returncode: + raise RuntimeError( + f"Analysis failed for {root}:\n{completed.stderr.strip()}" + ) + return json.loads(output.read_text()) + + +def print_comparison_table(rows: list[dict]) -> None: + rendered = [ + [str(row.get(field, "")) for _, field in COMPARISON_CONSOLE_FIELDS] + for row in rows + ] + widths = [ + max(len(heading), *(len(row[index]) for row in rendered)) + for index, (heading, _) in enumerate(COMPARISON_CONSOLE_FIELDS) + ] + print( + " ".join( + heading.ljust(widths[index]) + for index, (heading, _) in enumerate(COMPARISON_CONSOLE_FIELDS) + ) + ) + print(" ".join("-" * width for width in widths)) + for row in rendered: + print( + " ".join( + value.ljust(widths[index]) for index, value in enumerate(row) + ) + ) + + +def merge_test_analyses( + roots: list[str], analyses: list[dict] +) -> tuple[list[dict], list[dict]]: + merged: dict[str, dict] = {} + overrides = [] + for root, analysis in zip(roots, analyses, strict=True): + for application in analysis["applications"]: + job_id = str(application.get("job id") or "") + if not job_id: + raise ValueError( + f"Test application {application.get('application_id')} " + f"from {root} has no derived Job ID" + ) + enriched = dict(application) + enriched["source_event_log_root"] = root + previous = merged.get(job_id) + if previous is not None: + overrides.append( + { + "job_id": job_id, + "replaced_application_id": previous.get("application_id", ""), + "replaced_event_log_root": previous["source_event_log_root"], + "winning_application_id": enriched.get("application_id", ""), + "winning_event_log_root": root, + } + ) + merged[job_id] = enriched + return ( + [merged[job_id] for job_id in sorted(merged, key=lambda value: int(value))], + overrides, + ) + + +def run_comparison(args: argparse.Namespace) -> int: + if not args.event_log_root: + raise ValueError("--test-event-log-root requires --event-log-root") + baseline = analyze_event_root_for_comparison(args.event_log_root, args) + test_roots = args.test_event_log_root + test_runs = [ + analyze_event_root_for_comparison(root, args) for root in test_roots + ] + baseline_applications = [] + for application in baseline["applications"]: + enriched = dict(application) + enriched["source_event_log_root"] = args.event_log_root + baseline_applications.append(enriched) + test_applications, test_overrides = merge_test_analyses( + test_roots, test_runs + ) + instance_types = { + instance_type + for application in baseline_applications + test_applications + if application.get("complete") is True + for instance_type in application[ + "node_equivalent_seconds_by_instance_type" + ] + } + prices = instance_price_catalog( + instance_types, args.aws_region, args.aws_profile + ) + rows = build_comparison_rows( + baseline_applications, test_applications, prices + ) + rows = sort_comparison_rows(rows, args.sort_by) + summary = build_comparison_summary(rows) + test_merged_summary = summarize_applications(test_applications) + if args.output_csv: + args.output_csv.parent.mkdir(parents=True, exist_ok=True) + with args.output_csv.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=COMPARISON_FIELDS) + writer.writeheader() + writer.writerows(rows) + if args.output_json: + args.output_json.parent.mkdir(parents=True, exist_ok=True) + payload = { + "baseline_event_log_root": args.event_log_root, + "test_event_log_roots": test_roots, + "pricing_product_region": args.aws_region, + "pricing_queried_at_utc": datetime.now(timezone.utc).isoformat(), + "ec2_ondemand_prices": prices["ec2"], + "emr_prices": prices["emr"], + "comparison": rows, + "comparison_summary": summary, + "baseline_run_summary": baseline["summary"], + "baseline": baseline, + "test_runs": [ + {"event_log_root": root, "analysis": analysis} + for root, analysis in zip( + test_roots, test_runs, strict=True + ) + ], + "test_merged_applications": test_applications, + "test_merged_summary": test_merged_summary, + "test_overrides": test_overrides, + } + args.output_json.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + if not args.output_csv and not args.output_json: + print_comparison_table(rows) + elif args.output_csv: + print(f"Wrote {len(rows)} job comparison row(s) to {args.output_csv}") + else: + print(f"Wrote {len(rows)} job comparison row(s) to {args.output_json}") + print_comparison_summary(summary) + print_summary(baseline["summary"], "Baseline run summary") + print_summary(test_merged_summary, "Test overlay summary") + return 0 + + +def print_console_table(results: list[dict]) -> None: + rows = [] + for result in results: + row = [] + for _, field in CONSOLE_FIELDS: + value = result.get(field, "") + if field == "warnings": + value = " | ".join(value) + row.append(str(value)) + rows.append(row) + + widths = [ + max(len(heading), *(len(row[index]) for row in rows)) + for index, (heading, _) in enumerate(CONSOLE_FIELDS) + ] + print( + " ".join( + heading.ljust(widths[index]) + for index, (heading, _) in enumerate(CONSOLE_FIELDS) + ) + ) + print(" ".join("-" * width for width in widths)) + for row in rows: + print( + " ".join( + value.ljust(widths[index]) for index, value in enumerate(row) + ) + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument("--event-log-root", help="Spark event-log run directory (preferred)") + source.add_argument( + "--yarn-log-root", "--emr-log-uri", dest="yarn_log_root", + help="Exact archived YARN daemon-log URI, directory, or archive", + ) + parser.add_argument( + "--adapter", choices=ADAPTERS, default="emr", + help="Platform discovery and node-classification adapter (default: emr)", + ) + parser.add_argument( + "--pricing", choices=("none", "catalog", "live"), default="live", + help="Pricing source; live is currently supported by the EMR adapter", + ) + parser.add_argument("--price-catalog", type=Path) + parser.add_argument("--node-class-map", type=Path) + parser.add_argument( + "--test-event-log-root", + action="append", + help=( + "Test event-log root; repeat in overlay order, with the last " + "root containing a Job ID taking precedence" + ), + ) + parser.add_argument("--output-csv", type=Path) + parser.add_argument("--output-json", type=Path) + parser.add_argument( + "--sort-by", + choices=( + "job-id", + "wall-clock-factor", + "cost-factor", + "task-duration-factor", + "perfect-packing-cost-factor", + ), + default="job-id", + help="Comparison row order (default: job-id)", + ) + parser.add_argument("--input-csv", type=Path, help="Optional job metadata/event-log CSV") + parser.add_argument("--event-log-column") + parser.add_argument("--include-application-master", action="store_true") + parser.add_argument( + "--refresh-emr-log-cache", + action="store_true", + help="Refresh cached ResourceManager and NodeManager logs from S3", + ) + parser.add_argument( + "--aws-profile", + default=DEFAULT_AWS_PROFILE, + help="AWS CLI profile; omit to use the standard credential chain", + ) + parser.add_argument( + "--aws-region", + default=DEFAULT_AWS_REGION, + help="AWS product region; defaults to the effective AWS CLI region", + ) + parser.add_argument("--cache-dir", type=Path, default=Path(".cache/yarn-resource-cost/yarn-logs")) + parser.add_argument( + "--event-cache-dir", + type=Path, + default=Path(".cache/yarn-job-cost/event-metadata"), + ) + args = parser.parse_args() + args.emr_log_uri = args.yarn_log_root + return args + + +def main() -> int: + args = parse_args() + args.aws_region = resolve_aws_region(args.aws_region, args.aws_profile) + if args.test_event_log_root: + return run_comparison(args) + event_metadata: dict[str, EventLogApplication] = {} + selected_application_ids: set[str] | None = None + cluster_id = "" + emr_log_uri = args.emr_log_uri + local_event_path: Path | None = None + + if args.event_log_root: + local_event_path, selected_application_ids = materialize_event_metadata_files( + args.event_log_root, args.event_cache_dir, args.aws_profile, args.aws_region + ) + event_metadata = read_event_log_metadata(local_event_path) + if not selected_application_ids: + selected_application_ids = set(event_metadata) + missing_metadata = selected_application_ids - set(event_metadata) + if missing_metadata: + raise ValueError( + "Could not read application metadata for: " + ", ".join(sorted(missing_metadata)) + ) + cluster_ids = {app.cluster_id for app in event_metadata.values() if app.cluster_id} + if len(cluster_ids) != 1: + raise ValueError( + "Event-log root must identify exactly one EMR cluster; found: " + + (", ".join(sorted(cluster_ids)) or "none") + ) + cluster_id = next(iter(cluster_ids)) + emr_log_uri = resolve_emr_log_uri(cluster_id, args.aws_profile, args.aws_region) + + local_logs = materialize_emr_logs( + emr_log_uri, + args.cache_dir, + args.aws_profile, + refresh=args.refresh_emr_log_cache, + ) + evidence = parse_yarn_logs(local_logs) + mode = calculator_mode(evidence.calculator_class) + metadata = {app_id: app.as_metadata() for app_id, app in event_metadata.items()} + csv_metadata = load_csv_metadata(args.input_csv, args.event_log_column) + for app_id, csv_values in csv_metadata.items(): + target = metadata.setdefault(app_id, {}) + target.update({key: value for key, value in csv_values.items() if value}) + known_executor_container_ids = { + executor.container_id + for application in event_metadata.values() + for executor in application.executors.values() + if executor.container_id + } + results = calculate_applications( + evidence, + mode, + metadata, + args.include_application_master, + known_executor_container_ids, + ) + if selected_application_ids is not None: + results = [row for row in results if row["application_id"] in selected_application_ids] + missing_costs = selected_application_ids - {row["application_id"] for row in results} + if missing_costs: + raise ValueError( + "Selected applications are missing from YARN logs: " + + ", ".join(sorted(missing_costs)) + ) + add_task_packing_metrics(results, event_metadata, evidence, mode) + prices = add_ondemand_costs(results, args.aws_region, args.aws_profile) + summary = summarize_applications(results) + + if args.output_csv: + args.output_csv.parent.mkdir(parents=True, exist_ok=True) + with args.output_csv.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter( + handle, fieldnames=OUTPUT_FIELDS, extrasaction="ignore" + ) + writer.writeheader() + for result in results: + row = dict(result) + row["warnings"] = " | ".join(result["warnings"]) + row["task_metric_warnings"] = " | ".join( + result["task_metric_warnings"] + ) + writer.writerow(row) + if args.output_json: + args.output_json.parent.mkdir(parents=True, exist_ok=True) + payload = { + "event_log_root": args.event_log_root or "", + "local_event_metadata_path": str(local_event_path) if local_event_path else "", + "emr_cluster_id": cluster_id, + "emr_log_uri": emr_log_uri, + "local_log_path": str(local_logs), + "resource_calculator": mode, + "detected_resource_calculator_class": evidence.calculator_class, + "pricing_product_region": args.aws_region, + "pricing_queried_at_utc": datetime.now(timezone.utc).isoformat(), + "ec2_ondemand_prices": prices["ec2"], + "emr_prices": prices["emr"], + "nodes": { + node_id: { + "instance_type": node.instance_type, + "memory_mb": node.memory_mb, + "vcores": node.vcores, + "gpus": node.gpus, + } + for node_id, node in sorted(evidence.nodes.items()) + if node.instance_type + and any(container.node_id == node_id for container in evidence.containers.values()) + }, + "applications": results, + "summary": summary, + } + args.output_json.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + if not args.output_csv and not args.output_json: + print_console_table(results) + elif args.output_csv: + print(f"Wrote {len(results)} YARN application cost row(s) to {args.output_csv}") + else: + print(f"Wrote {len(results)} YARN application cost row(s) to {args.output_json}") + print_summary(summary) + return 0 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except AwsCliError as error: + print(f"error: {error}", file=sys.stderr) + raise SystemExit(2) from None diff --git a/yarn-resource-cost/package_yarn_job_cost.py b/yarn-resource-cost/package_yarn_job_cost.py new file mode 100644 index 0000000..ad8b526 --- /dev/null +++ b/yarn-resource-cost/package_yarn_job_cost.py @@ -0,0 +1,247 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Build and verify the standalone YARN job cost source bundle.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import shutil +import subprocess +import sys +import tarfile +import tempfile +import zipfile +from datetime import datetime, timezone +from pathlib import Path + + +PROJECT_DIR = Path(__file__).resolve().parent +DEFAULT_LICENSE_FILE = PROJECT_DIR.parent / "LICENSE" +PACKAGE_SOURCES = { + PROJECT_DIR / "yarn_resource_cost.py": "yarn_resource_cost.py", + PROJECT_DIR / "yarn_job_cost_core.py": "yarn_job_cost_core.py", + PROJECT_DIR / "yarn_job_cost_adapters.py": "yarn_job_cost_adapters.py", + PROJECT_DIR / "yarn_job_cost_dataproc.py": "yarn_job_cost_dataproc.py", + PROJECT_DIR / "calculate_yarn_job_cost.py": "calculate_yarn_job_cost.py", + PROJECT_DIR / "yarn_job_cost_discovery.py": "yarn_job_cost_discovery.py", + PROJECT_DIR / "yarn_job_cost_eventlog.py": "yarn_job_cost_eventlog.py", + PROJECT_DIR / "yarn_job_cost_defaults.py": "yarn_job_cost_defaults.py", + PROJECT_DIR / "test_calculate_yarn_job_cost.py": "test_calculate_yarn_job_cost.py", + PROJECT_DIR / "test_portable_yarn_resource_cost.py": "test_portable_yarn_resource_cost.py", + PROJECT_DIR / "test_fair_scheduler_policy.py": "test_fair_scheduler_policy.py", + PROJECT_DIR / "test_dataproc_log_normalization.py": "test_dataproc_log_normalization.py", + PROJECT_DIR / "test_dataproc_adapter.py": "test_dataproc_adapter.py", + PROJECT_DIR / "test_portable_comparison.py": "test_portable_comparison.py", + PROJECT_DIR / "RESOURCE_COST_MODEL.md": "RESOURCE_COST_MODEL.md", + PROJECT_DIR / "tests/fixtures/on_prem/eventlog_v2_application_1_0001/events_1_application_1_0001": "tests/fixtures/on_prem/eventlog_v2_application_1_0001/events_1_application_1_0001", + PROJECT_DIR / "tests/fixtures/on_prem/yarn/hadoop-yarn-resourcemanager-rm.log": "tests/fixtures/on_prem/yarn/hadoop-yarn-resourcemanager-rm.log", + PROJECT_DIR / "tests/fixtures/on_prem/node-classes.json": "tests/fixtures/on_prem/node-classes.json", + PROJECT_DIR / "tests/fixtures/on_prem/prices.json": "tests/fixtures/on_prem/prices.json", + PROJECT_DIR / "README.md": "README.md", + PROJECT_DIR / "CONTRIBUTING.md": "CONTRIBUTING.md", +} +BANNED_TEXT: tuple[str, ...] = () + + +def run(command: list[str], cwd: Path | None = None) -> str: + completed = subprocess.run( + command, cwd=cwd, capture_output=True, text=True + ) + if completed.returncode: + diagnostic = completed.stderr.strip() or completed.stdout.strip() + raise RuntimeError( + f"Command failed with exit code {completed.returncode}: " + f"{' '.join(command)}\n{diagnostic}" + ) + return completed.stdout.strip() + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def aggregate_sha256(paths: tuple[Path, ...]) -> str: + digest = hashlib.sha256() + for path in sorted(paths): + digest.update(path.name.encode("utf-8")) + digest.update(b"\0") + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def scan_bundle(root: Path) -> None: + failures = [] + for path in sorted(candidate for candidate in root.rglob("*") if candidate.is_file()): + if path.suffix not in {".py", ".md", ".json"}: + continue + text = path.read_text(encoding="utf-8") + for forbidden in BANNED_TEXT: + if forbidden.lower() in text.lower(): + failures.append(f"{path.name}: contains {forbidden!r}") + if failures: + raise ValueError("Bundle content scan failed:\n" + "\n".join(failures)) + + +def validate_bundle(root: Path) -> None: + license_path = root / "LICENSE" + if not license_path.is_file() or license_path.stat().st_size == 0: + raise ValueError("Bundle LICENSE is missing or empty") + contributing_path = root / "CONTRIBUTING.md" + if ( + not contributing_path.is_file() + or contributing_path.stat().st_size == 0 + ): + raise ValueError("Bundle CONTRIBUTING.md is missing or empty") + if ( + "Developer's Certificate of Origin 1.1" + not in contributing_path.read_text() + ): + raise ValueError("Bundle CONTRIBUTING.md does not contain DCO 1.1") + run([sys.executable, "-m", "unittest", "discover", "-s", ".", "-p", "test*.py"], cwd=root) + run([sys.executable, "yarn_resource_cost.py", "--help"], cwd=root) + + +def remove_bytecode(root: Path) -> None: + for cache in root.rglob("__pycache__"): + shutil.rmtree(cache) + for bytecode in root.rglob("*.py[co]"): + bytecode.unlink() + + +def create_archive( + staging: Path, archive: Path, package_name: str, archive_format: str +) -> None: + if archive_format == "zip": + with zipfile.ZipFile( + archive, "w", compression=zipfile.ZIP_DEFLATED + ) as bundle: + for path in sorted(staging.rglob("*")): + if path.is_file(): + arcname = Path(package_name) / path.relative_to(staging) + bundle.write(path, arcname=arcname) + return + with tarfile.open(archive, "w:gz") as bundle: + bundle.add(staging, arcname=package_name) + + +def extract_archive( + archive: Path, destination: Path, archive_format: str +) -> None: + if archive_format == "zip": + with zipfile.ZipFile(archive, "r") as bundle: + bundle.extractall(destination) + return + with tarfile.open(archive, "r:gz") as bundle: + bundle.extractall(destination) + + +def build( + output_dir: Path, license_file: Path, archive_format: str +) -> tuple[Path, Path]: + if archive_format not in {"zip", "tar.gz"}: + raise ValueError(f"Unsupported archive format: {archive_format}") + if not license_file.is_file() or license_file.stat().st_size == 0: + raise ValueError(f"License file is missing or empty: {license_file}") + license_digest = sha256(license_file) + source_digest = aggregate_sha256(tuple(PACKAGE_SOURCES)) + package_name = ( + f"yarn-job-cost-{source_digest[:12]}-{license_digest[:12]}" + ) + output_dir.mkdir(parents=True, exist_ok=True) + extension = ".zip" if archive_format == "zip" else ".tar.gz" + archive = output_dir / f"{package_name}{extension}" + + with tempfile.TemporaryDirectory() as temporary: + staging = Path(temporary) / package_name + staging.mkdir() + for source, destination in PACKAGE_SOURCES.items(): + target = staging / destination + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source, target) + shutil.copy2(license_file, staging / "LICENSE") + scan_bundle(staging) + validate_bundle(staging) + remove_bytecode(staging) + + files = { + path.relative_to(staging).as_posix(): sha256(path) + for path in sorted(staging.rglob("*")) + if path.is_file() + } + manifest = { + "package": "yarn-resource-cost", + "archive_format": archive_format, + "source_sha256": source_digest, + "built_at_utc": datetime.now(timezone.utc).isoformat(), + "python_requires": ">=3.10", + "license_file": "LICENSE", + "license_sha256": license_digest, + "files": files, + } + (staging / "PACKAGE-MANIFEST.json").write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + create_archive(staging, archive, package_name, archive_format) + + checksum = archive.with_suffix(archive.suffix + ".sha256") + checksum.write_text(f"{sha256(archive)} {archive.name}\n", encoding="utf-8") + + with tempfile.TemporaryDirectory() as temporary: + extracted = Path(temporary) + extract_archive(archive, extracted, archive_format) + validate_bundle(extracted / package_name) + return archive, checksum + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--license-file", + type=Path, + default=DEFAULT_LICENSE_FILE, + help=( + "Approved license text to embed verbatim as LICENSE " + f"(default: {DEFAULT_LICENSE_FILE})" + ), + ) + parser.add_argument( + "--archive-format", + choices=("zip", "tar.gz"), + default="zip", + help="Archive format (default: zip for Slack compatibility)", + ) + parser.add_argument( + "--output-dir", + type=Path, + default=PROJECT_DIR / "dist", + ) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + try: + archive, checksum = build( + args.output_dir, args.license_file, args.archive_format + ) + except (RuntimeError, ValueError) as error: + print(f"error: {error}", file=sys.stderr) + return 1 + print(archive) + print(checksum) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/yarn-resource-cost/test_calculate_yarn_job_cost.py b/yarn-resource-cost/test_calculate_yarn_job_cost.py new file mode 100644 index 0000000..aabd7c2 --- /dev/null +++ b/yarn-resource-cost/test_calculate_yarn_job_cost.py @@ -0,0 +1,1179 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +import ast +import contextlib +import csv +import gzip +import importlib.util +import io +import json +import os +import tarfile +import subprocess +import sys +import tempfile +import unittest +from unittest import mock +from pathlib import Path + + +SCRIPT = Path(__file__).with_name("calculate_yarn_job_cost.py") +DISCOVERY = Path(__file__).with_name("yarn_job_cost_discovery.py") +EVENTLOG_SCRIPT = Path(__file__).with_name("yarn_job_cost_eventlog.py") +DEFAULTS_SCRIPT = Path(__file__).with_name("yarn_job_cost_defaults.py") +APP_ID = "application_123_0002" +SPEC = importlib.util.spec_from_file_location("calculate_yarn_job_cost", SCRIPT) +MODULE = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = MODULE +SPEC.loader.exec_module(MODULE) +import yarn_job_cost_eventlog as EVENTLOG +import yarn_job_cost_discovery as DISCOVERY_MODULE + + +class CalculateYarnJobCostTest(unittest.TestCase): + def test_scripts_parse_as_python_310_and_311(self): + for script in (SCRIPT, DISCOVERY, EVENTLOG_SCRIPT, DEFAULTS_SCRIPT): + source = script.read_text() + ast.parse(source, filename=str(script), feature_version=(3, 10)) + ast.parse(source, filename=str(script), feature_version=(3, 11)) + + def test_aws_cli_error_surfaces_captured_stderr(self): + failure = subprocess.CompletedProcess( + ["aws", "--profile", "example-profile", "s3api", "list-objects-v2"], + 255, + stdout="", + stderr="Error loading SSO Token: Token has expired", + ) + with mock.patch.object(MODULE.subprocess, "run", return_value=failure): + with self.assertRaisesRegex( + MODULE.AwsCliError, "Error loading SSO Token: Token has expired" + ) as raised: + MODULE.run_aws(failure.args) + self.assertIn("exit code 255", str(raised.exception)) + self.assertIn("--profile example-profile", str(raised.exception)) + + def test_emr_logs_are_authoritative_and_csv_only_enriches(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + yarn_dir = root / "node" / "i-core" / "applications" / "hadoop-yarn" + yarn_dir.mkdir(parents=True) + node_log = yarn_dir / "hadoop-yarn-nodemanager-worker.log.gz" + lines = [ + "2026-01-01 00:00:00,000 INFO X: Registered with ResourceManager " + "as worker:8041 with total resource of " + "and with following Node attribute(s) : " + "{[nm.yarn.io/instanceType(STRING)=test.1xlarge]}", + "2026-01-01 00:00:00,500 INFO ApplicationSummary: " + "appId=application_123_0002,name=test,user=hadoop,queue=root.default," + "state=FINISHED,finalStatus=SUCCEEDED,totalAllocatedContainers=3", + "2026-01-01 00:00:01,000 INFO X: Start request for " + "container_123_0002_01_000001 by user appattempt_123_0002_000001 " + "with resource ", + "2026-01-01 00:00:03,000 INFO X: Container " + "container_123_0002_01_000001 transitioned from RUNNING to DONE", + "2026-01-01 00:00:10,000 INFO X: Start request for " + "container_123_0002_01_000002 by user appattempt_123_0002_000001 " + "with resource ", + "2026-01-01 00:00:20,000 INFO X: Container " + "container_123_0002_01_000002 transitioned from RUNNING to DONE", + "2026-01-01 00:00:30,000 INFO X: Start request for " + "container_123_0002_01_000003 by user appattempt_123_0002_000001 " + "with resource ", + "2026-01-01 00:00:50,000 INFO X: Container " + "container_123_0002_01_000003 transitioned from RUNNING to DONE", + ] + with gzip.open(node_log, "wt") as handle: + handle.write("\n".join(lines) + "\n") + rm_log = yarn_dir / "hadoop-yarn-resourcemanager-master.log.gz" + with gzip.open(rm_log, "wt") as handle: + rm_lines = [ + "2026-01-01 00:00:00,000 INFO CapacityScheduler: " + "Initialized CapacityScheduler with calculator=class " + "org.apache.hadoop.yarn.util.resource.DefaultResourceCalculator, " + "minimumAllocation=", + "2026-01-01 00:00:00,100 INFO X: NodeManager from node " + "worker(cmPort: 8041 httpPort: 8042) registered with capability: " + " and attributes " + "[nm.yarn.io/instanceType(STRING)=test.1xlarge]", + ] + allocations = [ + ("000001", "00:00:01,000", "00:00:03,000", 10, 1), + ("000002", "00:00:10,000", "00:00:20,000", 40, 2), + ("000003", "00:00:30,000", "00:00:50,000", 40, 2), + ] + for sequence, start, finish, memory, vcores in allocations: + container_id = f"container_123_0002_01_{sequence}" + rm_lines.extend( + [ + f"2026-01-01 {start} INFO X: Assigned container " + f"{container_id} of capacity " + "on host worker:8041", + f"2026-01-01 {finish} INFO X: {container_id} Container " + "Transitioned from RUNNING to COMPLETED", + ] + ) + handle.write("\n".join(rm_lines) + "\n") + input_csv = root / "input.csv" + with input_csv.open("w", newline="") as handle: + writer = csv.DictWriter( + handle, fieldnames=["job id", "job name", "eventlog\n(benchmark)"] + ) + writer.writeheader() + writer.writerow( + { + "job id": "6", + "job name": "test job", + "eventlog\n(benchmark)": f"s3://bucket/eventlog_v2_{APP_ID}", + } + ) + + output_csv = root / "cost.csv" + output_json = root / "cost.json" + price = { + "usd_per_hour": 3.0, + "effective_date": "2026-01-01", + "description": "EC2 fixture", + } + emr_price = { + "usd_per_hour": 0.75, + "effective_date": "2026-01-01", + "description": "EMR fixture", + } + with mock.patch.object( + sys, + "argv", + [ + str(SCRIPT), + "--emr-log-uri", + str(root), + "--aws-region", + "us-west-2", + "--input-csv", + str(input_csv), + "--output-csv", + str(output_csv), + "--output-json", + str(output_json), + ], + ), mock.patch.object( + MODULE, "current_ondemand_hourly_price", return_value=price + ), mock.patch.object( + MODULE, "current_emr_hourly_price", return_value=emr_price + ): + self.assertEqual(0, MODULE.main()) + with output_csv.open(newline="") as handle: + row = next(csv.DictReader(handle)) + self.assertEqual(row["application_id"], APP_ID) + self.assertEqual(row["job id"], "6") + self.assertEqual(row["job name"], "test job") + self.assertEqual(int(row["container_count"]), 2) + self.assertEqual(float(row["container_seconds"]), 30.0) + self.assertEqual(float(row["node_equivalent_seconds"]), 12.0) + self.assertEqual(float(row["instance_vcore_seconds"]), 48.0) + self.assertEqual( + row["instance_vcore_seconds_expression"], + "48.000000 test.1xlarge-vcore-seconds", + ) + self.assertEqual(float(row["ec2_ondemand_usd"]), 0.01) + self.assertEqual(float(row["emr_usd"]), 0.0025) + self.assertEqual(float(row["ec2_plus_emr_usd"]), 0.0125) + self.assertEqual(row["resource_calculator"], "default") + self.assertEqual(row["complete"], "True") + payload = json.loads(output_json.read_text()) + self.assertEqual("us-west-2", payload["pricing_product_region"]) + self.assertEqual( + 3.0, + payload["ec2_ondemand_prices"]["test.1xlarge"]["usd_per_hour"], + ) + self.assertEqual( + 0.75, + payload["emr_prices"]["test.1xlarge"]["usd_per_hour"], + ) + self.assertEqual( + 0.01, payload["applications"][0]["ec2_ondemand_usd"] + ) + self.assertEqual( + 0.0025, payload["applications"][0]["emr_usd"] + ) + self.assertEqual( + 0.0125, payload["applications"][0]["ec2_plus_emr_usd"] + ) + self.assertEqual( + 48.0, payload["summary"]["total_instance_vcore_seconds"] + ) + self.assertEqual( + "48.000000 test.1xlarge-vcore-seconds", + payload["summary"][ + "total_instance_vcore_seconds_expression" + ], + ) + self.assertEqual( + "dominant", + MODULE.calculator_mode("DominantResourceCalculator"), + ) + with self.assertRaises(ValueError): + MODULE.calculator_mode("") + + stdout = io.StringIO() + with mock.patch.object( + sys, + "argv", + [ + str(SCRIPT), + "--emr-log-uri", + str(root), + "--aws-region", + "us-west-2", + ], + ), mock.patch.object( + MODULE, "current_ondemand_hourly_price", return_value=price + ), mock.patch.object( + MODULE, "current_emr_hourly_price", return_value=emr_price + ), contextlib.redirect_stdout(stdout): + self.assertEqual(0, MODULE.main()) + console = stdout.getvalue() + self.assertIn("Job ID", console) + self.assertIn("Node-equivalent sec", console) + self.assertIn("EC2+EMR USD", console) + self.assertIn("12.000000 test.1xlarge-seconds", console) + self.assertIn( + "48.000000 test.1xlarge-vcore-seconds", console + ) + self.assertIn("0.01", console) + self.assertIn(APP_ID, console) + + def test_single_run_does_not_price_incomplete_ledgers(self): + applications = [ + { + "complete": False, + "node_equivalent_seconds_by_instance_type": { + "test.1xlarge": 12.0 + }, + } + ] + with mock.patch.object( + MODULE, "current_ondemand_hourly_price" + ) as ec2_lookup, mock.patch.object( + MODULE, "current_emr_hourly_price" + ) as emr_lookup: + prices = MODULE.add_ondemand_costs( + applications, "us-west-2", "example-profile" + ) + self.assertEqual({"ec2": {}, "emr": {}}, prices) + self.assertEqual("", applications[0]["ec2_ondemand_usd"]) + self.assertEqual("", applications[0]["emr_usd"]) + self.assertEqual("", applications[0]["ec2_plus_emr_usd"]) + ec2_lookup.assert_not_called() + emr_lookup.assert_not_called() + + def test_rm_terminal_precedes_nm_done_and_nm_fallback_is_not_final(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + yarn_dir = root / "node" / "i-core" / "applications" / "hadoop-yarn" + yarn_dir.mkdir(parents=True) + rm_log = yarn_dir / "hadoop-yarn-resourcemanager-master.log.gz" + rm_lines = [ + "2026-01-01 00:00:00,000 INFO X: NodeManager from node " + "worker(cmPort: 8041 httpPort: 8042) registered with capability: " + " and attributes " + "[nm.yarn.io/instanceType(STRING)=test.1xlarge]", + "2026-01-01 00:00:00,100 INFO ApplicationSummary: " + f"appId={APP_ID},name=test,user=hadoop,queue=root.default," + "state=FINISHED,finalStatus=SUCCEEDED,totalAllocatedContainers=2", + "2026-01-01 00:00:00,200 INFO CapacityScheduler: " + "Initialized CapacityScheduler with calculator=class " + "org.apache.hadoop.yarn.util.resource.DefaultResourceCalculator, " + "minimumAllocation=", + "2026-01-01 00:00:01,000 INFO X: Assigned container " + "container_123_0002_01_000002 of capacity on host worker:8041", + "2026-01-01 00:00:05,000 INFO X: " + "container_123_0002_01_000002 Container Transitioned from " + "RUNNING to RELEASED", + "2026-01-01 00:00:02,000 INFO X: Assigned container " + "container_123_0002_01_000003 of capacity on host worker:8041", + ] + with gzip.open(rm_log, "wt") as handle: + handle.write("\n".join(rm_lines) + "\n") + nm_log = yarn_dir / "hadoop-yarn-nodemanager-worker.log.gz" + nm_lines = [ + "2026-01-01 00:00:10,000 INFO X: Container " + "container_123_0002_01_000002 transitioned from RUNNING to DONE", + "2026-01-01 00:00:08,000 INFO X: Container " + "container_123_0002_01_000003 transitioned from RUNNING to DONE", + ] + with gzip.open(nm_log, "wt") as handle: + handle.write("\n".join(nm_lines) + "\n") + + evidence = MODULE.parse_yarn_logs(root) + rm_container = evidence.containers["container_123_0002_01_000002"] + nm_container = evidence.containers["container_123_0002_01_000003"] + self.assertEqual( + MODULE.parse_timestamp("2026-01-01 00:00:05,000"), + rm_container.finish_ms, + ) + self.assertEqual("resourcemanager", rm_container.finish_source) + self.assertEqual( + MODULE.parse_timestamp("2026-01-01 00:00:08,000"), + nm_container.finish_ms, + ) + self.assertEqual("nodemanager", nm_container.finish_source) + + result = MODULE.calculate_applications(evidence, "default", {}, False)[0] + self.assertEqual(10.0, result["container_seconds"]) + self.assertEqual(1, result["nodemanager_finish_fallback_container_count"]) + self.assertFalse(result["complete"]) + self.assertIn("NodeManager DONE fallback", " | ".join(result["warnings"])) + + def test_emr_log_cache_can_be_refreshed(self): + with tempfile.TemporaryDirectory() as directory: + cache = Path(directory) + uri = "s3://bucket/logs/j-TEST" + digest = MODULE.hashlib.sha256(uri.encode()).hexdigest()[:16] + target = cache / f"j-TEST-{digest}" + target.mkdir(parents=True) + marker = target / ".download-complete" + marker.write_text(uri + "\n") + (target / "hadoop-yarn-resourcemanager.log").write_text("fixture\n") + + with mock.patch.object(MODULE.subprocess, "run") as run: + self.assertEqual( + target, + MODULE.materialize_emr_logs(uri, cache, None), + ) + run.assert_not_called() + + with mock.patch.object(MODULE.subprocess, "run") as run: + self.assertEqual( + target, + MODULE.materialize_emr_logs(uri, cache, None, refresh=True), + ) + run.assert_called_once() + self.assertTrue(marker.is_file()) + + with mock.patch.object( + MODULE.subprocess, + "run", + side_effect=subprocess.CalledProcessError(1, ["aws"]), + ): + with self.assertRaises(subprocess.CalledProcessError): + MODULE.materialize_emr_logs(uri, cache, None, refresh=True) + self.assertFalse(marker.exists()) + + def test_dominant_resource_cost_includes_gpu_share(self): + assignment = MODULE.RM_ASSIGN_RE.search( + "2026-07-24 23:39:05,028 INFO X: Assigned container " + "container_123_0002_01_000002 of capacity " + " on host worker:8041" + ) + self.assertIsNotNone(assignment) + self.assertEqual(1, MODULE.gpu_amount(assignment.group("resources"))) + + container = MODULE.Container( + container_id="container_123_0002_01_000002", + application_id=APP_ID, + node_id="worker", + start_ms=1000, + finish_ms=11000, + memory_mb=49152, + node_memory_mb=53248, + vcores=15, + node_vcores=16, + gpus=1, + node_gpus=1, + source="resourcemanager", + finish_source="resourcemanager", + ) + evidence = MODULE.YarnEvidence( + nodes={"worker": MODULE.Node("worker", "g6.4xlarge", 53248, 16, 1)}, + containers={container.container_id: container}, + calculator_class="DominantResourceCalculator", + application_summaries={ + APP_ID: MODULE.ApplicationSummary(APP_ID, "test", "SUCCEEDED", 1) + }, + ) + result = MODULE.calculate_applications(evidence, "dominant", {}, False)[0] + self.assertEqual(10.0, result["gpu_seconds"]) + self.assertEqual(10.0, result["node_equivalent_seconds"]) + self.assertEqual(160.0, result["instance_vcore_seconds"]) + self.assertEqual( + "160.000000 g6.4xlarge-vcore-seconds", + result["instance_vcore_seconds_expression"], + ) + self.assertEqual("10.000000 g6.4xlarge-seconds", result["cost_expression"]) + + def test_comparison_joins_job_ids_and_calculates_deltas(self): + baseline = [ + { + "job id": "6", + "application_id": "application_base_6", + "final_status": "SUCCEEDED", + "spark_duration_seconds": 100.0, + "node_equivalent_seconds": 200.0, + "node_equivalent_seconds_by_instance_type": {"cpu.test": 200.0}, + "instance_vcore_seconds": 19200.0, + "complete": True, + "warnings": [], + }, + { + "job id": "8", + "application_id": "application_base_8", + "final_status": "SUCCEEDED", + "spark_duration_seconds": 50.0, + "node_equivalent_seconds": 50.0, + "node_equivalent_seconds_by_instance_type": {"cpu.test": 50.0}, + "complete": True, + "warnings": [], + }, + ] + test = [ + { + "job id": "6", + "application_id": "application_test_6", + "final_status": "SUCCEEDED", + "spark_duration_seconds": 80.0, + "node_equivalent_seconds": 300.0, + "node_equivalent_seconds_by_instance_type": {"gpu.test": 300.0}, + "instance_vcore_seconds": 4800.0, + "complete": True, + "warnings": [], + }, + { + "job id": "9", + "application_id": "application_test_9", + "final_status": "KILLED", + "spark_duration_seconds": "", + "node_equivalent_seconds": 10.0, + "node_equivalent_seconds_by_instance_type": {"gpu.test": 10.0}, + "complete": True, + "warnings": [], + }, + ] + prices = { + "ec2": { + "cpu.test": {"usd_per_hour": 7.2}, + "gpu.test": {"usd_per_hour": 1.2}, + }, + "emr": { + "cpu.test": {"usd_per_hour": 1.8}, + "gpu.test": {"usd_per_hour": 0.3}, + }, + } + rows = MODULE.build_comparison_rows(baseline, test, prices) + self.assertEqual(["6", "8", "9"], [row["job_id"] for row in rows]) + matched = rows[0] + self.assertEqual("cpu.test", matched["baseline_instance_type"]) + self.assertEqual("gpu.test", matched["test_instance_type"]) + self.assertEqual(-20.0, matched["wall_clock_delta_seconds"]) + self.assertEqual(0.8, matched["wall_clock_factor"]) + self.assertEqual(100.0, matched["node_equivalent_delta_seconds"]) + self.assertEqual(19200.0, matched["baseline_instance_vcore_seconds"]) + self.assertEqual(4800.0, matched["test_instance_vcore_seconds"]) + self.assertEqual(0.25, matched["instance_vcore_seconds_factor"]) + self.assertEqual(-14400.0, matched["instance_vcore_seconds_delta"]) + self.assertEqual(0.4, matched["baseline_ec2_ondemand_usd"]) + self.assertEqual(0.1, matched["test_ec2_ondemand_usd"]) + self.assertEqual(0.25, matched["ec2_ondemand_cost_factor"]) + self.assertEqual(-0.3, matched["ec2_ondemand_delta_usd"]) + self.assertEqual(0.1, matched["baseline_emr_usd"]) + self.assertEqual(0.025, matched["test_emr_usd"]) + self.assertEqual(0.5, matched["baseline_ec2_plus_emr_usd"]) + self.assertEqual(0.125, matched["test_ec2_plus_emr_usd"]) + self.assertEqual(0.25, matched["ec2_plus_emr_cost_factor"]) + self.assertEqual(-0.375, matched["ec2_plus_emr_delta_usd"]) + self.assertEqual("MISSING", rows[1]["test_final_status"]) + self.assertEqual("MISSING", rows[2]["baseline_final_status"]) + self.assertEqual( + ["6", "8", "9"], + [ + row["job_id"] + for row in MODULE.sort_comparison_rows(rows, "wall-clock-factor") + ], + ) + self.assertEqual( + ["6", "8", "9"], + [ + row["job_id"] + for row in MODULE.sort_comparison_rows(rows, "cost-factor") + ], + ) + sortable = [ + { + "job_id": "1", + "baseline_final_status": "SUCCEEDED", + "test_final_status": "SUCCEEDED", + "wall_clock_factor": 2.0, + "ec2_plus_emr_cost_factor": 0.5, + }, + { + "job_id": "2", + "baseline_final_status": "SUCCEEDED", + "test_final_status": "SUCCEEDED", + "wall_clock_factor": 0.5, + "ec2_plus_emr_cost_factor": 2.0, + }, + { + "job_id": "3", + "baseline_final_status": "SUCCEEDED", + "test_final_status": "KILLED", + "wall_clock_factor": 0.1, + "ec2_plus_emr_cost_factor": 0.1, + }, + ] + self.assertEqual( + ["2", "1", "3"], + [ + row["job_id"] + for row in MODULE.sort_comparison_rows( + sortable, "wall-clock-factor" + ) + ], + ) + self.assertEqual( + ["1", "2", "3"], + [ + row["job_id"] + for row in MODULE.sort_comparison_rows(sortable, "cost-factor") + ], + ) + self.assertEqual( + "a.test, z.test", + MODULE.application_instance_types( + { + "node_equivalent_seconds_by_instance_type": { + "z.test": 1.0, + "a.test": 2.0, + } + } + ), + ) + + def test_comparison_does_not_price_incomplete_ledgers(self): + def application(job_id, instance_type, complete, warning): + return { + "job id": str(job_id), + "application_id": f"application_{instance_type}_{job_id}", + "final_status": "SUCCEEDED", + "spark_duration_seconds": 10.0, + "node_equivalent_seconds": 20.0, + "node_equivalent_seconds_by_instance_type": {instance_type: 20.0}, + "complete": complete, + "warnings": [warning] if warning else [], + } + + baseline = [ + application(1, "cpu.test", False, "baseline partial"), + application(2, "cpu.test", True, ""), + ] + test = [ + application(1, "gpu.test", True, ""), + application(2, "gpu.test", False, "test partial"), + ] + prices = { + "ec2": { + "cpu.test": {"usd_per_hour": 7.2}, + "gpu.test": {"usd_per_hour": 1.8}, + }, + "emr": { + "cpu.test": {"usd_per_hour": 1.8}, + "gpu.test": {"usd_per_hour": 0.45}, + }, + } + + rows = MODULE.build_comparison_rows(baseline, test, prices) + + self.assertFalse(rows[0]["baseline_complete"]) + self.assertEqual("baseline partial", rows[0]["baseline_warnings"]) + self.assertEqual("", rows[0]["baseline_ec2_ondemand_usd"]) + self.assertNotEqual("", rows[0]["test_ec2_ondemand_usd"]) + self.assertEqual("", rows[0]["ec2_ondemand_cost_factor"]) + self.assertEqual("", rows[0]["baseline_emr_usd"]) + self.assertNotEqual("", rows[0]["test_emr_usd"]) + self.assertEqual("", rows[0]["baseline_ec2_plus_emr_usd"]) + self.assertNotEqual("", rows[0]["test_ec2_plus_emr_usd"]) + self.assertEqual("", rows[0]["ec2_plus_emr_cost_factor"]) + self.assertFalse(rows[1]["test_complete"]) + self.assertEqual("test partial", rows[1]["test_warnings"]) + self.assertNotEqual("", rows[1]["baseline_ec2_ondemand_usd"]) + self.assertEqual("", rows[1]["test_ec2_ondemand_usd"]) + self.assertEqual("", rows[1]["ec2_ondemand_delta_usd"]) + self.assertNotEqual("", rows[1]["baseline_emr_usd"]) + self.assertEqual("", rows[1]["test_emr_usd"]) + self.assertNotEqual("", rows[1]["baseline_ec2_plus_emr_usd"]) + self.assertEqual("", rows[1]["test_ec2_plus_emr_usd"]) + self.assertEqual("", rows[1]["ec2_plus_emr_delta_usd"]) + + def test_roots_overlay_in_order_and_last_root_wins(self): + roots = ["s3://bucket/base", "s3://bucket/patch"] + analyses = [ + { + "applications": [ + {"job id": "6", "application_id": "application_old_6"}, + {"job id": "8", "application_id": "application_8"}, + ] + }, + { + "applications": [ + {"job id": "6", "application_id": "application_new_6"}, + {"job id": "9", "application_id": "application_9"}, + ] + }, + ] + + merged, overrides = MODULE.merge_test_analyses(roots, analyses) + + self.assertEqual(["6", "8", "9"], [row["job id"] for row in merged]) + self.assertEqual("application_new_6", merged[0]["application_id"]) + self.assertEqual("s3://bucket/patch", merged[0]["source_event_log_root"]) + self.assertEqual( + [ + { + "job_id": "6", + "replaced_application_id": "application_old_6", + "replaced_event_log_root": "s3://bucket/base", + "winning_application_id": "application_new_6", + "winning_event_log_root": "s3://bucket/patch", + } + ], + overrides, + ) + + def test_event_logs_derive_identity_cluster_versions_and_duration(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + fixtures = [ + ( + "application_123_0002", + "/run/005_hitch.source_triplify/job.py", + "plain plan without a benchmark table", + ), + ( + "application_123_0003", + "/run/006_hitch.source_triplify__rewrite/job.py", + "scan benchmark_100pct.j0380__input", + ), + ] + for offset, (app_id, app_name, detail) in enumerate(fixtures): + app_dir = root / f"eventlog_v2_{app_id}" + app_dir.mkdir() + events = [ + {"Event": "SparkListenerLogStart", "Spark Version": "4.0.2-amzn-0"}, + { + "Event": "SparkListenerEnvironmentUpdate", + "Spark Properties": { + "spark.emr.clusterId": "j-TEST", + "spark.emr.releaseLabel": "emr-spark-8.0.0", + }, + }, + { + "Event": "SparkListenerApplicationStart", + "App ID": app_id, + "App Name": app_name, + "Timestamp": 1000 + offset, + }, + {"Event": "SyntheticPlan", "detail": detail}, + {"Event": "SparkListenerApplicationEnd", "Timestamp": 6000 + offset}, + ] + (app_dir / f"events_1_{app_id}").write_text( + "\n".join(json.dumps(event) for event in events) + "\n" + ) + + metadata = MODULE.read_event_log_metadata(root) + self.assertEqual("380", metadata["application_123_0002"].job_id) + self.assertEqual("900380", metadata["application_123_0003"].job_id) + self.assertEqual("j-TEST", metadata["application_123_0002"].cluster_id) + self.assertEqual( + "emr-spark-8.0.0", metadata["application_123_0002"].emr_release_label + ) + self.assertEqual("4.0.2-amzn-0", metadata["application_123_0002"].spark_version) + self.assertEqual( + 5.0, + metadata["application_123_0002"].as_metadata()["spark_duration_seconds"], + ) + + + def test_event_logs_sum_successful_task_duration_and_require_all_segments(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + app_dir = root / f"eventlog_v2_{APP_ID}" + app_dir.mkdir() + first = [ + { + "Event": "SparkListenerEnvironmentUpdate", + "Spark Properties": { + "spark.executor.cores": "4", + "spark.task.cpus": "1", + }, + }, + { + "Event": "SparkListenerApplicationStart", + "App ID": APP_ID, + "App Name": "/run/j0144__job.py", + "Timestamp": 1000, + }, + { + "Event": "SparkListenerExecutorAdded", + "Executor ID": "7", + "Executor Info": { + "Total Cores": 4, + "Resource Profile Id": 0, + "Attributes": { + "CONTAINER_ID": "container_123_0002_01_000002" + }, + }, + }, + { + "Event": "SparkListenerTaskEnd", + "Task End Reason": {"Reason": "Success"}, + "Task Info": { + "Executor ID": "7", + "Launch Time": 2000, + "Finish Time": 10000, + }, + }, + { + "Event": "SparkListenerTaskEnd", + "Task End Reason": {"Reason": "ExceptionFailure"}, + "Task Info": { + "Executor ID": "7", + "Launch Time": 2000, + "Finish Time": 12000, + "Failed": True, + }, + }, + ] + second = [ + { + "Event": "SparkListenerApplicationEnd", + "Timestamp": 11000, + } + ] + for number, events in ((1, first), (2, second)): + (app_dir / f"events_{number}_{APP_ID}").write_text( + "\n".join(json.dumps(event) for event in events) + "\n" + ) + + metadata = MODULE.read_event_log_metadata(root)[APP_ID] + self.assertTrue(metadata.event_task_metrics_complete()) + self.assertEqual(1, metadata.successful_task_attempt_count) + self.assertEqual(8000, metadata.task_duration_sum_ms) + self.assertEqual({"7": 8000}, metadata.task_duration_ms_by_executor) + self.assertEqual(4, metadata.configured_executor_cores) + self.assertEqual(1, metadata.task_cpus) + self.assertEqual( + "container_123_0002_01_000002", + metadata.executors["7"].container_id, + ) + + (app_dir / f"events_2_{APP_ID}").rename( + app_dir / f"events_3_{APP_ID}" + ) + incomplete = MODULE.read_event_log_metadata(root)[APP_ID] + self.assertFalse(incomplete.event_task_metrics_complete()) + self.assertIn( + "Missing event-log segments: 2", + incomplete.task_metric_warnings, + ) + + def test_task_packing_metrics_use_executor_cores_and_yarn_node_share(self): + container = MODULE.Container( + container_id="container_123_0002_01_000002", + application_id=APP_ID, + node_id="worker", + start_ms=1000, + finish_ms=11000, + memory_mb=40, + node_memory_mb=100, + vcores=1, + node_vcores=4, + source="resourcemanager", + finish_source="resourcemanager", + ) + evidence = MODULE.YarnEvidence( + nodes={"worker": MODULE.Node("worker", "cpu.test", 100, 4, 0)}, + containers={container.container_id: container}, + calculator_class="DefaultResourceCalculator", + application_summaries={ + APP_ID: MODULE.ApplicationSummary(APP_ID, "test", "SUCCEEDED", 1) + }, + ) + application = MODULE.calculate_applications( + evidence, "default", {}, False + )[0] + event = MODULE.EventLogApplication( + application_id=APP_ID, + application_ended=True, + task_cpus=1, + successful_task_attempt_count=1, + task_duration_sum_ms=8000, + task_duration_ms_by_executor={"7": 8000}, + executors={ + "7": DISCOVERY_MODULE.SparkExecutor( + "7", container.container_id, 4, 0 + ) + }, + event_segments={1}, + ) + + partial = dict(application) + partial["complete"] = False + partial["warnings"] = list(application["warnings"]) + MODULE.add_task_packing_metrics( + [partial], {APP_ID: event}, evidence, "default" + ) + self.assertTrue(partial["task_metrics_complete"]) + self.assertFalse(partial["perfect_packing_complete"]) + self.assertEqual(8.0, partial["task_duration_sum_seconds"]) + self.assertEqual("", partial["perfect_packing_node_seconds"]) + + MODULE.add_task_packing_metrics( + [application], {APP_ID: event}, evidence, "default" + ) + self.assertTrue(application["task_metrics_complete"]) + self.assertTrue(application["perfect_packing_complete"]) + self.assertEqual(8.0, application["task_duration_sum_seconds"]) + self.assertEqual(0.8, application["perfect_packing_node_seconds"]) + self.assertEqual(0.2, application["packing_efficiency"]) + + ec2 = { + "usd_per_hour": 3.6, + "effective_date": "2026-01-01", + "description": "fixture", + } + emr = { + "usd_per_hour": 0.9, + "effective_date": "2026-01-01", + "description": "fixture", + } + with mock.patch.object( + MODULE, "current_ondemand_hourly_price", return_value=ec2 + ), mock.patch.object( + MODULE, "current_emr_hourly_price", return_value=emr + ): + MODULE.add_ondemand_costs( + [application], "us-west-2", "example-profile" + ) + self.assertEqual(0.005, application["ec2_plus_emr_usd"]) + self.assertEqual(0.001, application["perfect_packing_ec2_plus_emr_usd"]) + self.assertEqual(5.0, application["actual_to_perfect_cost_factor"]) + self.assertEqual(400.0, application["actual_to_perfect_cost_overhead_percent"]) + + def test_dominant_gpu_task_packing_charges_full_node_share(self): + container = MODULE.Container( + container_id="container_123_0002_01_000001", + application_id=APP_ID, + node_id="worker", + start_ms=1000, + finish_ms=11000, + memory_mb=49152, + node_memory_mb=53248, + vcores=15, + node_vcores=16, + gpus=1, + node_gpus=1, + source="resourcemanager", + finish_source="resourcemanager", + ) + evidence = MODULE.YarnEvidence( + nodes={"worker": MODULE.Node("worker", "g6.4xlarge", 53248, 16, 1)}, + containers={container.container_id: container}, + calculator_class="DominantResourceCalculator", + application_summaries={ + APP_ID: MODULE.ApplicationSummary(APP_ID, "test", "SUCCEEDED", 1) + }, + ) + application = MODULE.calculate_applications( + evidence, "dominant", {}, False, {container.container_id} + )[0] + event = MODULE.EventLogApplication( + application_id=APP_ID, + application_ended=True, + task_cpus=1, + successful_task_attempt_count=15, + task_duration_sum_ms=15000, + task_duration_ms_by_executor={"7": 15000}, + executors={ + "7": DISCOVERY_MODULE.SparkExecutor( + "7", container.container_id, 15, 0 + ) + }, + event_segments={1}, + ) + + MODULE.add_task_packing_metrics( + [application], {APP_ID: event}, evidence, "dominant" + ) + self.assertTrue(application["perfect_packing_complete"]) + self.assertEqual(1.0, application["perfect_packing_node_seconds"]) + self.assertEqual(10.0, application["node_equivalent_seconds"]) + self.assertIn( + "Included sequence-1 container", " | ".join(application["warnings"]) + ) + self.assertEqual(0.1, application["packing_efficiency"]) + + def test_comparison_summary_uses_only_complete_matched_jobs(self): + complete = { + "job_id": "1", + "baseline_final_status": "SUCCEEDED", + "test_final_status": "SUCCEEDED", + "baseline_complete": True, + "test_complete": True, + "baseline_task_metrics_complete": True, + "test_task_metrics_complete": True, + "baseline_perfect_packing_complete": True, + "test_perfect_packing_complete": True, + "baseline_task_duration_sum_seconds": 100.0, + "test_task_duration_sum_seconds": 50.0, + "baseline_perfect_packing_node_seconds": 20.0, + "test_perfect_packing_node_seconds": 10.0, + "baseline_node_equivalent_seconds": 40.0, + "test_node_equivalent_seconds": 25.0, + "baseline_perfect_packing_ec2_plus_emr_usd": 4.0, + "test_perfect_packing_ec2_plus_emr_usd": 3.0, + "baseline_ec2_plus_emr_usd": 8.0, + "test_ec2_plus_emr_usd": 6.0, + } + incomplete = dict(complete, job_id="2", test_task_metrics_complete=False) + + summary = MODULE.build_comparison_summary([complete, incomplete]) + + self.assertEqual(1, summary["eligible_job_count"]) + self.assertEqual(1, summary["excluded_job_count"]) + self.assertEqual(0.5, summary["task_duration_factor"]) + self.assertEqual(0.75, summary["perfect_packing_cost_factor"]) + self.assertEqual(0.75, summary["actual_cost_factor"]) + self.assertEqual(0.5, summary["baseline"]["packing_efficiency"]) + self.assertEqual(0.4, summary["test"]["packing_efficiency"]) + self.assertEqual(2.0, summary["baseline"]["actual_to_perfect_cost_factor"]) + + def test_instance_vcore_total_uses_all_successful_complete_jobs(self): + applications = [ + { + "job id": "1", + "application_id": "application_success_1", + "final_status": "SUCCEEDED", + "complete": True, + "task_metrics_complete": False, + "perfect_packing_complete": False, + "instance_vcore_seconds_by_instance_type": { + "r7a.24xlarge": 100.0 + }, + }, + { + "job id": "2", + "application_id": "application_success_2", + "final_status": "SUCCEEDED", + "complete": True, + "task_metrics_complete": False, + "perfect_packing_complete": False, + "instance_vcore_seconds_by_instance_type": { + "r7a.24xlarge": 20.0, + "g6.4xlarge": 30.0, + }, + }, + { + "job id": "3", + "application_id": "application_failed", + "final_status": "FAILED", + "complete": True, + "task_metrics_complete": False, + "perfect_packing_complete": False, + "instance_vcore_seconds_by_instance_type": { + "r7a.24xlarge": 1000.0 + }, + }, + { + "job id": "4", + "application_id": "application_partial", + "final_status": "SUCCEEDED", + "complete": False, + "task_metrics_complete": False, + "perfect_packing_complete": False, + "instance_vcore_seconds_by_instance_type": { + "r7a.24xlarge": 2000.0 + }, + }, + ] + + summary = MODULE.summarize_applications(applications) + + self.assertEqual(2, summary["successful_complete_job_count"]) + self.assertEqual(1, summary["successful_incomplete_job_count"]) + self.assertEqual(150.0, summary["total_instance_vcore_seconds"]) + self.assertEqual( + {"g6.4xlarge": 30.0, "r7a.24xlarge": 120.0}, + summary["total_instance_vcore_seconds_by_instance_type"], + ) + self.assertEqual( + "30.000000 g6.4xlarge-vcore-seconds + " + "120.000000 r7a.24xlarge-vcore-seconds", + summary["total_instance_vcore_seconds_expression"], + ) + + def test_new_comparison_sort_modes(self): + rows = [ + { + "job_id": "1", + "baseline_final_status": "SUCCEEDED", + "test_final_status": "SUCCEEDED", + "task_duration_factor": 2.0, + "perfect_packing_cost_factor": 0.5, + }, + { + "job_id": "2", + "baseline_final_status": "SUCCEEDED", + "test_final_status": "SUCCEEDED", + "task_duration_factor": 0.5, + "perfect_packing_cost_factor": 2.0, + }, + ] + self.assertEqual( + ["2", "1"], + [ + row["job_id"] + for row in MODULE.sort_comparison_rows( + rows, "task-duration-factor" + ) + ], + ) + self.assertEqual( + ["1", "2"], + [ + row["job_id"] + for row in MODULE.sort_comparison_rows( + rows, "perfect-packing-cost-factor" + ) + ], + ) + + def test_event_materializer_downloads_every_segment(self): + objects = { + APP_ID: [ + f"s3://bucket/root/eventlog_v2_{APP_ID}/events_{number}_{APP_ID}" + for number in (1, 2, 3) + ] + } + with tempfile.TemporaryDirectory() as directory, mock.patch( + "yarn_job_cost_discovery.list_event_log_objects", + return_value=objects, + ), mock.patch( + "yarn_job_cost_discovery.run_aws", + return_value=subprocess.CompletedProcess([], 0, "", ""), + ) as run_aws: + _, application_ids = MODULE.materialize_event_metadata_files( + "s3://bucket/root", Path(directory), None, "us-west-2" + ) + self.assertEqual({APP_ID}, application_ids) + self.assertEqual(3, run_aws.call_count) + self.assertEqual( + [ + objects[APP_ID][0], + objects[APP_ID][1], + objects[APP_ID][2], + ], + [call.args[0][-3] for call in run_aws.call_args_list], + ) + + def test_cli_defaults_come_from_defaults_module(self): + with mock.patch.object( + sys, "argv", ["calculate_yarn_job_cost.py", "--emr-log-uri", "/tmp/logs"] + ): + args = MODULE.parse_args() + self.assertEqual(MODULE.DEFAULT_AWS_PROFILE, args.aws_profile) + self.assertEqual(MODULE.DEFAULT_AWS_REGION, args.aws_region) + + def test_resolve_aws_region_precedence(self): + with mock.patch.object(MODULE.subprocess, "run") as run: + self.assertEqual( + "eu-west-1", MODULE.resolve_aws_region("eu-west-1", None) + ) + run.assert_not_called() + + with mock.patch.dict(os.environ, {"AWS_REGION": "ap-southeast-2"}, clear=False): + self.assertEqual( + "ap-southeast-2", MODULE.resolve_aws_region(None, None) + ) + + environment = dict(os.environ) + environment.pop("AWS_REGION", None) + environment.pop("AWS_DEFAULT_REGION", None) + configured = subprocess.CompletedProcess( + ["aws", "configure", "get", "region"], 0, stdout="us-east-2\n", stderr="" + ) + with mock.patch.dict(os.environ, environment, clear=True), mock.patch.object( + MODULE.subprocess, "run", return_value=configured + ) as run: + self.assertEqual( + "us-east-2", MODULE.resolve_aws_region(None, "example-profile") + ) + self.assertEqual( + ["aws", "--profile", "example-profile", "configure", "get", "region"], + run.call_args.args[0], + ) + + def test_resolve_aws_region_rejects_missing_region(self): + environment = dict(os.environ) + environment.pop("AWS_REGION", None) + environment.pop("AWS_DEFAULT_REGION", None) + missing = subprocess.CompletedProcess( + ["aws", "configure", "get", "region"], 1, stdout="", stderr="" + ) + with mock.patch.dict(os.environ, environment, clear=True), mock.patch.object( + MODULE.subprocess, "run", return_value=missing + ): + with self.assertRaisesRegex(ValueError, "pass --aws-region"): + MODULE.resolve_aws_region(None, None) + + def test_portable_eventlog_reader_handles_raw_lz4block(self): + payload = b"first\nsecond\n" + header = ( + EVENTLOG.LZ4_BLOCK_MAGIC + + bytes([EVENTLOG.RAW_BLOCK]) + + len(payload).to_bytes(4, "little") + + len(payload).to_bytes(4, "little") + + bytes(4) + ) + self.assertEqual( + ["first", "second"], + list(EVENTLOG.iter_text_lines(io.BytesIO(header + payload), True)), + ) + + def test_portable_eventlog_reader_handles_literal_lz4_block(self): + payload = b"spark" + compressed = bytes([len(payload) << 4]) + payload + self.assertEqual( + payload, EVENTLOG.lz4_decompress_block(compressed, len(payload)) + ) + + def test_portable_eventlog_reader_handles_tar_bundle(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source = root / "events_1_application_123_0002" + source.write_text("one\ntwo\n") + archive = root / "events.tar.gz" + with tarfile.open(archive, "w:gz") as bundle: + bundle.add( + source, + arcname=( + "eventlog_v2_application_123_0002/" + "events_1_application_123_0002" + ), + ) + streams = [] + for app_dir, _, stream, compressed in EVENTLOG.iter_eventlog_streams( + archive + ): + streams.append( + (app_dir, list(EVENTLOG.iter_text_lines(stream, compressed))) + ) + self.assertEqual( + [("eventlog_v2_application_123_0002", ["one", "two"])], streams + ) + +if __name__ == "__main__": + unittest.main() diff --git a/yarn-resource-cost/test_dataproc_adapter.py b/yarn-resource-cost/test_dataproc_adapter.py new file mode 100644 index 0000000..b7e6847 --- /dev/null +++ b/yarn-resource-cost/test_dataproc_adapter.py @@ -0,0 +1,56 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import subprocess +import unittest +from unittest import mock + +import yarn_job_cost_dataproc as dataproc + + +class DataprocAdapterTest(unittest.TestCase): + @mock.patch.object(dataproc, "run_command") + def test_primary_and_secondary_worker_shapes_remain_distinct(self, run): + run.return_value = subprocess.CompletedProcess( + ["gcloud"], + 0, + stdout=json.dumps( + { + "clusterUuid": "fixture-uuid", + "config": { + "workerConfig": { + "machineTypeUri": "zones/us-west1-a/machineTypes/n2-standard-16", + "instanceNames": ["sample-w-0"], + }, + "secondaryWorkerConfig": { + "machineTypeUri": "zones/us-west1-a/machineTypes/g2-standard-16", + "instanceNames": ["sample-sw-0"], + "accelerators": [ + { + "acceleratorTypeUri": ( + "zones/us-west1-a/acceleratorTypes/nvidia-l4" + ), + "acceleratorCount": 1, + } + ], + }, + }, + } + ), + stderr="", + ) + mappings, provenance = dataproc.describe_node_classes( + "sample", "us-west1", "example-project" + ) + self.assertEqual("gcp:dataproc:n2-standard-16", mappings["sample-w-0"]) + self.assertEqual( + "gcp:dataproc:g2-standard-16+1xnvidia-l4", + mappings["sample-sw-0"], + ) + self.assertEqual("fixture-uuid", provenance["cluster_uuid"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/yarn-resource-cost/test_dataproc_log_normalization.py b/yarn-resource-cost/test_dataproc_log_normalization.py new file mode 100644 index 0000000..3e34988 --- /dev/null +++ b/yarn-resource-cost/test_dataproc_log_normalization.py @@ -0,0 +1,26 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import unittest + +from yarn_job_cost_core import normalize_log_line + + +class DataprocLogNormalizationTest(unittest.TestCase): + def test_cloud_logging_text_payload_is_unwrapped(self): + message = ( + "2026-01-01 00:00:00,000 INFO CapacityScheduler: " + "resource-calculator=DefaultResourceCalculator" + ) + line = json.dumps({"textPayload": message}) + self.assertEqual(message, normalize_log_line(line).rstrip()) + + def test_plain_daemon_log_line_is_unchanged(self): + line = "2026-01-01 00:00:00,000 INFO ResourceManager: started\n" + self.assertEqual(line, normalize_log_line(line)) + + +if __name__ == "__main__": + unittest.main() diff --git a/yarn-resource-cost/test_fair_scheduler_policy.py b/yarn-resource-cost/test_fair_scheduler_policy.py new file mode 100644 index 0000000..e8eb04a --- /dev/null +++ b/yarn-resource-cost/test_fair_scheduler_policy.py @@ -0,0 +1,29 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import tempfile +import unittest +from pathlib import Path + +import yarn_job_cost_core as core + + +class FairSchedulerPolicyTest(unittest.TestCase): + def test_conflicting_built_in_policies_are_ambiguous(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "hadoop-yarn-resourcemanager.log" + path.write_text( + "2026-01-01 00:00:00,000 INFO FairScheduler: " + "queue=root.a policy=FairSharePolicy\n" + "2026-01-01 00:00:01,000 INFO FairScheduler: " + "queue=root.b policy=DominantResourceFairnessPolicy\n", + encoding="utf-8", + ) + evidence = core.parse_yarn_logs(path) + self.assertTrue(evidence.accounting_policy_ambiguous) + self.assertIn("Multiple FairScheduler policies", " ".join(evidence.warnings)) + + +if __name__ == "__main__": + unittest.main() diff --git a/yarn-resource-cost/test_portable_comparison.py b/yarn-resource-cost/test_portable_comparison.py new file mode 100644 index 0000000..7493009 --- /dev/null +++ b/yarn-resource-cost/test_portable_comparison.py @@ -0,0 +1,35 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import unittest + +import yarn_resource_cost as cli + + +def application(app_id: str, complete: bool, node_seconds: float, cost: object) -> dict: + return { + "comparison_key": "job-1", + "application_id": app_id, + "node_equivalent_seconds": node_seconds, + "node_equivalent_seconds_by_instance_type": {"node-class": node_seconds}, + "spark_duration_seconds": 10.0, + "worker_cost": cost, + "complete": complete, + } + + +class PortableComparisonTest(unittest.TestCase): + def test_incomplete_ledger_has_no_resource_or_cost_factor(self): + args = argparse.Namespace(sort_by="comparison-key") + baseline = {"applications": [application("application_1_1", True, 10.0, 1.0)]} + test = {"applications": [application("application_2_1", False, 5.0, "")]} + row = cli.compare_runs(baseline, [test], args)[0] + self.assertEqual("", row["node_equivalent_factor"]) + self.assertEqual("", row["worker_cost_factor"]) + self.assertEqual(1.0, row["wall_clock_factor"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/yarn-resource-cost/test_portable_yarn_resource_cost.py b/yarn-resource-cost/test_portable_yarn_resource_cost.py new file mode 100644 index 0000000..d84841b --- /dev/null +++ b/yarn-resource-cost/test_portable_yarn_resource_cost.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +import yarn_job_cost_adapters as adapters +import yarn_job_cost_core as core + + +ROOT = Path(__file__).resolve().parent +FIXTURE = ROOT / "tests" / "fixtures" / "on_prem" +SCRIPT = ROOT / "yarn_resource_cost.py" + + +class ProviderNeutralAccountingTest(unittest.TestCase): + def container(self, **updates) -> core.Container: + values = { + "container_id": "container_1_0001_01_000002", + "application_id": "application_1_0001", + "node_id": "worker1", + "start_ms": 0, + "finish_ms": 1000, + "memory_mb": 2048, + "node_memory_mb": 8192, + "vcores": 2, + "node_vcores": 8, + "resources": {"memory-mb": 2048, "vcores": 2}, + "node_resources": {"memory-mb": 8192, "vcores": 8}, + } + values.update(updates) + return core.Container(**values) + + def test_default_calculator_uses_memory_only(self): + container = self.container( + vcores=8, + resources={"memory-mb": 2048, "vcores": 8, "yarn.io/gpu": 1}, + node_resources={"memory-mb": 8192, "vcores": 8, "yarn.io/gpu": 1}, + ) + self.assertEqual(0.25, core.container_node_share(container, "default")) + + def test_dominant_calculator_uses_arbitrary_custom_resource(self): + container = self.container( + resources={"memory-mb": 2048, "vcores": 2, "vendor/device": 3}, + node_resources={"memory-mb": 8192, "vcores": 8, "vendor/device": 4}, + ) + self.assertEqual(0.75, core.container_node_share(container, "dominant")) + + def test_dominant_rejects_missing_allocated_resource_capacity(self): + container = self.container( + resources={"memory-mb": 2048, "vcores": 2, "yarn.io/gpu": 1} + ) + with self.assertRaisesRegex(ValueError, "yarn.io/gpu"): + core.container_node_share(container, "dominant") + + def test_fair_scheduler_drf_policy_is_detected_from_evidence(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "hadoop-yarn-resourcemanager.log" + path.write_text( + "2026-01-01 00:00:00,000 INFO FairScheduler: " + "policy=DominantResourceFairnessPolicy\n", + encoding="utf-8", + ) + evidence = core.parse_yarn_logs(path) + self.assertEqual("DominantResourceCalculator", evidence.calculator_class) + self.assertEqual("DominantResourceFairnessPolicy", evidence.scheduler_policy) + + +class AdapterTest(unittest.TestCase): + def test_catalog_cost_and_node_mapping(self): + evidence = core.parse_yarn_logs(FIXTURE / "yarn") + adapters.apply_node_class_map(evidence, FIXTURE / "node-classes.json") + self.assertEqual("onprem:worker-8", evidence.nodes["worker1"].node_class) + catalog = adapters.load_price_catalog(FIXTURE / "prices.json") + applications = [ + { + "complete": True, + "node_equivalent_seconds_by_instance_type": { + "onprem:worker-8": 10.0 + }, + "warnings": [], + } + ] + adapters.apply_catalog_costs(applications, catalog) + self.assertEqual(0.01, applications[0]["worker_cost"]) + + def test_missing_catalog_rate_suppresses_final_cost(self): + catalog = adapters.load_price_catalog(FIXTURE / "prices.json") + applications = [ + { + "complete": True, + "node_equivalent_seconds_by_instance_type": {"unknown": 1.0}, + "warnings": [], + } + ] + adapters.apply_catalog_costs(applications, catalog) + self.assertFalse(applications[0]["complete"]) + self.assertEqual("", applications[0]["worker_cost"]) + + +class PortableCliTest(unittest.TestCase): + def test_on_prem_fixture_end_to_end_with_catalog(self): + with tempfile.TemporaryDirectory() as directory: + output = Path(directory) / "result.json" + completed = subprocess.run( + [ + sys.executable, + str(SCRIPT), + "--adapter", + "on-prem", + "--event-log-root", + str(FIXTURE), + "--yarn-log-root", + str(FIXTURE / "yarn"), + "--node-class-map", + str(FIXTURE / "node-classes.json"), + "--pricing", + "catalog", + "--price-catalog", + str(FIXTURE / "prices.json"), + "--output-json", + str(output), + ], + capture_output=True, + text=True, + ) + self.assertEqual(0, completed.returncode, completed.stderr) + payload = json.loads(output.read_text(encoding="utf-8")) + self.assertEqual(1, payload["schema_version"]) + self.assertEqual("on-prem", payload["adapter"]) + application = payload["applications"][0] + self.assertTrue(application["complete"]) + self.assertEqual(2.0, application["node_equivalent_seconds"]) + self.assertEqual(0.002, application["worker_cost"]) + self.assertEqual("USD", application["worker_cost_currency"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/yarn-resource-cost/tests/fixtures/on_prem/eventlog_v2_application_1_0001/events_1_application_1_0001 b/yarn-resource-cost/tests/fixtures/on_prem/eventlog_v2_application_1_0001/events_1_application_1_0001 new file mode 100644 index 0000000..2754b0b --- /dev/null +++ b/yarn-resource-cost/tests/fixtures/on_prem/eventlog_v2_application_1_0001/events_1_application_1_0001 @@ -0,0 +1,5 @@ +{"Event":"SparkListenerLogStart","Spark Version":"3.5.1"} +{"Event":"SparkListenerEnvironmentUpdate","Spark Properties":{"spark.executor.cores":"2","spark.task.cpus":"1"}} +{"Event":"SparkListenerApplicationStart","App Name":"portable-cost-sample","App ID":"application_1_0001","Timestamp":1000} +{"Event":"SparkListenerExecutorAdded","Executor ID":"1","Executor Info":{"Total Cores":2,"Resource Profile Id":0,"Attributes":{"CONTAINER_ID":"container_1_0001_01_000002"}}} +{"Event":"SparkListenerApplicationEnd","Timestamp":11000} diff --git a/yarn-resource-cost/tests/fixtures/on_prem/node-classes.json b/yarn-resource-cost/tests/fixtures/on_prem/node-classes.json new file mode 100644 index 0000000..cc546fc --- /dev/null +++ b/yarn-resource-cost/tests/fixtures/on_prem/node-classes.json @@ -0,0 +1,6 @@ +{ + "schema_version": 1, + "nodes": { + "worker1": "onprem:worker-8" + } +} diff --git a/yarn-resource-cost/tests/fixtures/on_prem/prices.json b/yarn-resource-cost/tests/fixtures/on_prem/prices.json new file mode 100644 index 0000000..024dc29 --- /dev/null +++ b/yarn-resource-cost/tests/fixtures/on_prem/prices.json @@ -0,0 +1,12 @@ +{ + "schema_version": 1, + "currency": "USD", + "effective_at": "2026-01-01T00:00:00Z", + "source": "synthetic unit-test rate", + "rates": [ + { + "node_class": "onprem:worker-8", + "hourly_rate": 3.6 + } + ] +} diff --git a/yarn-resource-cost/tests/fixtures/on_prem/yarn/hadoop-yarn-resourcemanager-rm.log b/yarn-resource-cost/tests/fixtures/on_prem/yarn/hadoop-yarn-resourcemanager-rm.log new file mode 100644 index 0000000..117fca8 --- /dev/null +++ b/yarn-resource-cost/tests/fixtures/on_prem/yarn/hadoop-yarn-resourcemanager-rm.log @@ -0,0 +1,7 @@ +2026-01-01 00:00:00,000 INFO CapacityScheduler: Initialized CapacityScheduler with calculator=class org.apache.hadoop.yarn.util.resource.DefaultResourceCalculator +2026-01-01 00:00:00,100 INFO RMNodeImpl: NodeManager from node worker1(cmPort: 8041 httpPort: 8042) registered with capability: +2026-01-01 00:00:01,000 INFO CapacityScheduler: Assigned container container_1_0001_01_000001 of capacity on host worker1:8041 +2026-01-01 00:00:02,000 INFO CapacityScheduler: Assigned container container_1_0001_01_000002 of capacity on host worker1:8041 +2026-01-01 00:00:09,000 INFO RMContainerImpl: container_1_0001_01_000001 Container Transitioned from RUNNING to COMPLETED +2026-01-01 00:00:10,000 INFO RMContainerImpl: container_1_0001_01_000002 Container Transitioned from RUNNING to COMPLETED +2026-01-01 00:00:10,100 INFO RMAppManager$ApplicationSummary: appId=application_1_0001,name=portable-cost-sample,user=test,queue=default,state=FINISHED,finalStatus=SUCCEEDED,totalAllocatedContainers=2 diff --git a/yarn-resource-cost/yarn_job_cost_adapters.py b/yarn-resource-cost/yarn_job_cost_adapters.py new file mode 100644 index 0000000..c175990 --- /dev/null +++ b/yarn-resource-cost/yarn_job_cost_adapters.py @@ -0,0 +1,222 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Platform adapters for discovery, node classification, and pricing inputs.""" + +from __future__ import annotations + +import hashlib +import json +import shlex +import subprocess +import tarfile +import zipfile +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Callable + +from yarn_job_cost_core import YarnEvidence, relevant_log_files + + +class AdapterCommandError(RuntimeError): + """External adapter command failure with its diagnostic text preserved.""" + + +def run_command(command: list[str]) -> subprocess.CompletedProcess[str]: + completed = subprocess.run(command, capture_output=True, text=True) + if completed.returncode: + diagnostic = completed.stderr.strip() or completed.stdout.strip() + raise AdapterCommandError( + f"Command failed with exit code {completed.returncode}:\n" + f"{diagnostic or '(no diagnostic output)'}\n" + f"Command: {shlex.join(command)}" + ) + return completed + + +def _safe_archive_member(destination: Path, member_name: str) -> Path: + member = (destination / member_name).resolve() + if destination.resolve() not in member.parents and member != destination.resolve(): + raise ValueError(f"Archive member escapes destination: {member_name}") + return member + + +def extract_archive(path: Path, destination: Path) -> Path: + destination.mkdir(parents=True, exist_ok=True) + if zipfile.is_zipfile(path): + with zipfile.ZipFile(path) as archive: + for info in archive.infolist(): + _safe_archive_member(destination, info.filename) + archive.extractall(destination) + elif tarfile.is_tarfile(path): + with tarfile.open(path) as archive: + members = archive.getmembers() + for member in members: + _safe_archive_member(destination, member.name) + if member.issym() or member.islnk(): + raise ValueError(f"Archive links are not accepted: {member.name}") + archive.extractall(destination, members=members) + else: + raise ValueError(f"Unsupported log archive: {path}") + return destination + + +def materialize_local_or_archive(uri: str, cache_dir: Path) -> Path: + path = Path(uri).expanduser() + if not path.exists(): + raise FileNotFoundError(path) + if path.is_dir(): + return path + digest = hashlib.sha256(str(path.resolve()).encode()).hexdigest()[:16] + target = cache_dir / f"archive-{digest}" + if not target.exists(): + extract_archive(path, target) + return target + + +def materialize_hdfs(uri: str, cache_dir: Path, refresh: bool) -> Path: + digest = hashlib.sha256(uri.encode()).hexdigest()[:16] + target = cache_dir / f"hdfs-{digest}" + marker = target / ".download-complete" + if marker.exists() and not refresh: + return target + target.mkdir(parents=True, exist_ok=True) + run_command(["hdfs", "dfs", "-copyToLocal", "-f", uri, str(target)]) + marker.write_text(uri + "\n", encoding="utf-8") + return target + + +def materialize_gcs(uri: str, cache_dir: Path, refresh: bool) -> Path: + digest = hashlib.sha256(uri.rstrip("/").encode()).hexdigest()[:16] + target = cache_dir / f"gcs-{digest}" + marker = target / ".download-complete" + if marker.exists() and not refresh: + return target + target.mkdir(parents=True, exist_ok=True) + run_command( + ["gcloud", "storage", "cp", "--recursive", uri.rstrip("/") + "/*", str(target)] + ) + marker.write_text(uri + "\n", encoding="utf-8") + return target + + +def materialize_s3( + uri: str, cache_dir: Path, refresh: bool, aws_profile: str | None +) -> Path: + normalized = "s3://" + uri.split("://", 1)[1] + digest = hashlib.sha256(normalized.rstrip("/").encode()).hexdigest()[:16] + target = cache_dir / f"s3-{digest}" + marker = target / ".download-complete" + if marker.exists() and not refresh: + return target + target.mkdir(parents=True, exist_ok=True) + command = ["aws"] + if aws_profile: + command += ["--profile", aws_profile] + command += [ + "s3", "cp", "--recursive", normalized.rstrip("/") + "/", str(target), + "--exclude", "*", "--include", "*hadoop-yarn-nodemanager*.log*", + "--include", "*hadoop-yarn-resourcemanager*.log*", + ] + run_command(command) + marker.write_text(uri + "\n", encoding="utf-8") + return target + + +def materialize_yarn_logs( + adapter: str, + uri: str, + cache_dir: Path, + refresh: bool = False, + aws_profile: str | None = None, +) -> Path: + if uri.startswith(("s3://", "s3a://", "s3n://")): + if adapter != "emr": + raise ValueError("S3 YARN log discovery is only provided by the EMR adapter") + result = materialize_s3(uri, cache_dir, refresh, aws_profile) + elif uri.startswith("gs://"): + if adapter != "dataproc": + raise ValueError("gs:// YARN logs require --adapter dataproc") + result = materialize_gcs(uri, cache_dir, refresh) + elif uri.startswith("hdfs://"): + if adapter != "on-prem": + raise ValueError("HDFS YARN logs require --adapter on-prem") + result = materialize_hdfs(uri, cache_dir, refresh) + else: + result = materialize_local_or_archive(uri, cache_dir) + if not relevant_log_files(result): + raise ValueError(f"No ResourceManager or NodeManager logs found in {uri}") + return result + + +@dataclass(frozen=True) +class PriceCatalog: + currency: str + rates_per_hour: dict[str, float] + provenance: dict[str, object] + + +def load_price_catalog(path: Path) -> PriceCatalog: + payload = json.loads(path.read_text(encoding="utf-8")) + if payload.get("schema_version") != 1: + raise ValueError("Price catalog schema_version must be 1") + rates: dict[str, float] = {} + for entry in payload.get("rates") or []: + node_class = str(entry.get("node_class") or "") + hourly = float(entry.get("hourly_rate")) + if not node_class or hourly < 0: + raise ValueError("Each price entry needs node_class and nonnegative hourly_rate") + rates[node_class] = hourly + return PriceCatalog( + currency=str(payload.get("currency") or "USD"), + rates_per_hour=rates, + provenance={ + "mode": "catalog", + "source": payload.get("source") or str(path), + "effective_at": payload.get("effective_at") or "", + "loaded_at_utc": datetime.now(timezone.utc).isoformat(), + }, + ) + + +def apply_node_class_map(evidence: YarnEvidence, path: Path | None) -> None: + if path is None: + return + payload = json.loads(path.read_text(encoding="utf-8")) + mappings = payload.get("nodes") or {} + default = str(payload.get("default_node_class") or "") + for node_id, node in evidence.nodes.items(): + mapped = str(mappings.get(node_id) or "") + if mapped: + node.instance_type = mapped + elif not node.instance_type and default: + node.instance_type = default + + +def apply_catalog_costs(applications: list[dict], catalog: PriceCatalog) -> None: + for application in applications: + application["worker_cost_currency"] = catalog.currency + if application.get("complete") is not True: + application["worker_cost"] = "" + continue + seconds_by_class = application["node_equivalent_seconds_by_instance_type"] + missing = sorted(set(seconds_by_class) - set(catalog.rates_per_hour)) + if missing: + application["complete"] = False + application["worker_cost"] = "" + application["warnings"].append( + "Price catalog has no rate for node class(es): " + ", ".join(missing) + ) + continue + application["worker_cost"] = round( + sum( + float(seconds) * catalog.rates_per_hour[node_class] / 3600.0 + for node_class, seconds in seconds_by_class.items() + ), + 8, + ) + + +ADAPTERS = ("emr", "dataproc", "on-prem") diff --git a/yarn-resource-cost/yarn_job_cost_core.py b/yarn-resource-cost/yarn_job_cost_core.py new file mode 100644 index 0000000..f7c24f7 --- /dev/null +++ b/yarn-resource-cost/yarn_job_cost_core.py @@ -0,0 +1,416 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Provider-neutral YARN allocation ledger and resource-share accounting.""" + +from __future__ import annotations + +import gzip +import json +import re +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import TextIO + + +TIMESTAMP = r"(?P\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2},\d{3})" +CONTAINER = r"(?Pcontainer_\d+_\d+_\d+_\d+)" +RESOURCE = ( + r"\d+)(?:, max memory:(?P\d+))?, " + r"vCores:(?P\d+)(?:, max vCores:(?P\d+))?" + r"(?P[^>]*)>" +) +START_RE = re.compile( + rf"^{TIMESTAMP} .*Start request for {CONTAINER} .* resource {RESOURCE}" +) +DONE_RE = re.compile( + rf"^{TIMESTAMP} .*Container {CONTAINER} transitioned from .* to DONE\b" +) +RM_ASSIGN_RE = re.compile( + rf"^{TIMESTAMP} .*Assigned container {CONTAINER} of capacity {RESOURCE} " + r"on host (?P[^:,\s]+):\d+" +) +RM_TERMINAL_RE = re.compile( + rf"^{TIMESTAMP} .*{CONTAINER} Container Transitioned from .* to " + r"(?:COMPLETED|RELEASED|KILLED|EXPIRED)\b" +) +APPLICATION_SUMMARY_RE = re.compile( + r"appId=(?Papplication_\d+_\d+),name=(?P.*?),user=.*?" + r"finalStatus=(?P[^,]+).*?" + r"totalAllocatedContainers=(?P\d+)" +) +NODE_RE = re.compile( + r"Registered with ResourceManager .* total resource of " + r"\d+), vCores:(?P\d+)(?P[^>]*)>" +) +RM_NODE_RE = re.compile( + r"NodeManager from node (?P[^ (]+)(?:\([^)]*\))? registered with " + r"capability: \d+), vCores:(?P\d+)" + r"(?P[^>]*)>" +) +INSTANCE_TYPE_RE = re.compile(r"instanceType\(STRING\)=(?P[^}\]\s]+)") +RESOURCE_ENTRY_RE = re.compile( + r"(?:^|,\s*)(?P[A-Za-z0-9_.\-/]+):\s*(?P\d+)" +) +CALCULATOR_RE = re.compile( + r"(?:calculator=class |resource-calculator(?:=|: )\s*(?:class )?)" + r"(?:org\.apache\.hadoop\.yarn\.util\.resource\.)?" + r"(?PDefaultResourceCalculator|DominantResourceCalculator)" +) +FAIR_POLICY_RE = re.compile( + r"(?PDominantResourceFairnessPolicy|FairSharePolicy|FifoPolicy)" +) +CONTAINER_PARTS_RE = re.compile( + r"container_(?P\d+)_(?P\d+)_" + r"(?P\d+)_(?P\d+)" +) + + +@dataclass +class Node: + node_id: str + instance_type: str = "" + memory_mb: int | None = None + vcores: int | None = None + gpus: int = 0 + resources: dict[str, int] = field(default_factory=dict) + + @property + def node_class(self) -> str: + return self.instance_type or self.node_id + + +@dataclass +class Container: + container_id: str + application_id: str + node_id: str + start_ms: int + memory_mb: int + node_memory_mb: int + vcores: int + node_vcores: int + gpus: int = 0 + node_gpus: int = 0 + resources: dict[str, int] = field(default_factory=dict) + node_resources: dict[str, int] = field(default_factory=dict) + finish_ms: int | None = None + finish_source: str = "" + source: str = "nodemanager" + + @property + def sequence(self) -> int: + match = CONTAINER_PARTS_RE.fullmatch(self.container_id) + if not match: + raise ValueError(f"Unexpected container ID {self.container_id}") + return int(match.group("sequence")) + + +@dataclass +class ApplicationSummary: + application_id: str + name: str + final_status: str + total_allocated_containers: int + + +@dataclass +class YarnEvidence: + nodes: dict[str, Node] = field(default_factory=dict) + containers: dict[str, Container] = field(default_factory=dict) + calculator_class: str = "" + calculator_source: str = "" + scheduler_policy: str = "" + accounting_policy_ambiguous: bool = False + application_summaries: dict[str, ApplicationSummary] = field(default_factory=dict) + warnings: list[str] = field(default_factory=list) + + +def parse_timestamp(value: str) -> int: + parsed = datetime.strptime(value, "%Y-%m-%d %H:%M:%S,%f").replace( + tzinfo=timezone.utc + ) + return int(parsed.timestamp() * 1000) + + +def application_id(container_id: str) -> str: + match = CONTAINER_PARTS_RE.fullmatch(container_id) + if not match: + raise ValueError(f"Unexpected container ID {container_id}") + return f"application_{match.group('cluster')}_{match.group('application')}" + + +def resource_map(memory: int, vcores: int, suffix: str | None) -> dict[str, int]: + values = {"memory-mb": int(memory), "vcores": int(vcores)} + for match in RESOURCE_ENTRY_RE.finditer(suffix or ""): + values[match.group("name")] = int(match.group("value")) + return values + + +def gpu_amount(resources: str | None) -> int: + return resource_map(0, 0, resources).get("yarn.io/gpu", 0) + + +def capacity(value: str | None, fallback: int | None, name: str) -> int: + if value is not None: + return int(value) + if fallback is not None: + return fallback + raise ValueError(f"Could not determine node {name} capacity") + + +def open_log(path: Path) -> TextIO: + if path.suffix == ".gz": + return gzip.open(path, mode="rt", encoding="utf-8", errors="replace") + return path.open(encoding="utf-8", errors="replace") + + +def relevant_log_files(path: Path) -> list[Path]: + if path.is_file(): + return [path] + return sorted( + file + for file in path.rglob("*") + if file.is_file() + and ( + "hadoop-yarn-nodemanager" in file.name + or "hadoop-yarn-resourcemanager" in file.name + or file.suffix in {".jsonl", ".log"} + ) + ) + + +def normalize_log_line(line: str) -> str: + stripped = line.strip() + if not stripped.startswith("{"): + return line + try: + record = json.loads(stripped) + except json.JSONDecodeError: + return line + payload = record.get("textPayload") + if payload is None and isinstance(record.get("jsonPayload"), dict): + payload = record["jsonPayload"].get("message") + return str(payload) + "\n" if payload is not None else line + + +def _record_calculator(evidence: YarnEvidence, calculator: str, source: str) -> None: + if evidence.calculator_class and evidence.calculator_class != calculator: + raise ValueError( + "Conflicting ResourceCalculators: " + f"{evidence.calculator_class}, {calculator}" + ) + evidence.calculator_class = calculator + evidence.calculator_source = source + + +def _node_instance_type(line: str) -> str: + match = INSTANCE_TYPE_RE.search(line) + return match.group("value").strip() if match else "" + + +def _container_from_match( + match: re.Match[str], node_id: str, node: Node | None, source: str +) -> Container: + allocated = resource_map( + int(match.group("memory")), + int(match.group("vcores")), + match.group("resources"), + ) + node_memory = capacity( + match.group("max_memory"), node.memory_mb if node else None, "memory" + ) + node_vcores = capacity( + match.group("max_vcores"), node.vcores if node else None, "vcore" + ) + node_resources = dict(node.resources) if node else { + "memory-mb": node_memory, + "vcores": node_vcores, + } + node_resources.setdefault("memory-mb", node_memory) + node_resources.setdefault("vcores", node_vcores) + return Container( + container_id=match.group("container"), + application_id=application_id(match.group("container")), + node_id=node_id, + start_ms=parse_timestamp(match.group("timestamp")), + memory_mb=allocated["memory-mb"], + node_memory_mb=node_memory, + vcores=allocated["vcores"], + node_vcores=node_vcores, + gpus=allocated.get("yarn.io/gpu", 0), + node_gpus=node_resources.get("yarn.io/gpu", 0), + resources=allocated, + node_resources=node_resources, + source=source, + ) + + +def parse_yarn_logs(path: Path) -> YarnEvidence: + """Parse RM/NM daemon logs without relying on a cloud-provider layout.""" + evidence = YarnEvidence() + rm_finishes: dict[str, int] = {} + nm_finishes: dict[str, int] = {} + files = relevant_log_files(path) + if not files: + raise ValueError(f"No ResourceManager or NodeManager log files found under {path}") + + for file in files: + path_node_id = file.parent.name + path_node = evidence.nodes.setdefault(path_node_id, Node(path_node_id)) + with open_log(file) as handle: + for raw_line in handle: + line = normalize_log_line(raw_line) + summary = APPLICATION_SUMMARY_RE.search(line) + if summary: + app_id = summary.group("application") + evidence.application_summaries[app_id] = ApplicationSummary( + app_id, + summary.group("name"), + summary.group("final_status"), + int(summary.group("containers")), + ) + calculator = CALCULATOR_RE.search(line) + if calculator: + _record_calculator( + evidence, calculator.group("calculator"), file.name + ) + policy = FAIR_POLICY_RE.search(line) + if policy: + name = policy.group("policy") + if evidence.scheduler_policy and evidence.scheduler_policy != name: + evidence.accounting_policy_ambiguous = True + evidence.warnings.append( + "Multiple FairScheduler policies were observed; archive " + "queue-specific policy evidence for exact accounting" + ) + else: + evidence.scheduler_policy = name + mapped = ( + "DominantResourceCalculator" + if name == "DominantResourceFairnessPolicy" + else "DefaultResourceCalculator" + ) + _record_calculator(evidence, mapped, file.name) + + rm_node = RM_NODE_RE.search(line) + if rm_node: + host = rm_node.group("host").strip() + resources = resource_map( + int(rm_node.group("memory")), + int(rm_node.group("vcores")), + rm_node.group("resources"), + ) + evidence.nodes[host] = Node( + node_id=host, + instance_type=_node_instance_type(line), + memory_mb=resources["memory-mb"], + vcores=resources["vcores"], + gpus=resources.get("yarn.io/gpu", 0), + resources=resources, + ) + continue + node_match = NODE_RE.search(line) + if node_match: + resources = resource_map( + int(node_match.group("memory")), + int(node_match.group("vcores")), + node_match.group("resources"), + ) + path_node.instance_type = _node_instance_type(line) + path_node.memory_mb = resources["memory-mb"] + path_node.vcores = resources["vcores"] + path_node.gpus = resources.get("yarn.io/gpu", 0) + path_node.resources = resources + continue + assignment = RM_ASSIGN_RE.search(line) + if assignment: + host = assignment.group("host") + candidate = _container_from_match( + assignment, host, evidence.nodes.get(host), "resourcemanager" + ) + evidence.containers[candidate.container_id] = candidate + continue + terminal = RM_TERMINAL_RE.search(line) + if terminal: + container_id = terminal.group("container") + finish = parse_timestamp(terminal.group("timestamp")) + rm_finishes[container_id] = min( + finish, rm_finishes.get(container_id, finish) + ) + continue + start = START_RE.search(line) + if start: + candidate = _container_from_match( + start, path_node_id, path_node, "nodemanager" + ) + evidence.containers.setdefault(candidate.container_id, candidate) + continue + done = DONE_RE.search(line) + if done: + container_id = done.group("container") + finish = parse_timestamp(done.group("timestamp")) + nm_finishes[container_id] = min( + finish, nm_finishes.get(container_id, finish) + ) + + for container_id, container in evidence.containers.items(): + if container_id in rm_finishes: + container.finish_ms = rm_finishes[container_id] + container.finish_source = "resourcemanager" + elif container_id in nm_finishes: + container.finish_ms = nm_finishes[container_id] + container.finish_source = "nodemanager" + return evidence + + +def calculator_mode(detected_class: str) -> str: + if detected_class == "DefaultResourceCalculator": + return "default" + if detected_class == "DominantResourceCalculator": + return "dominant" + raise ValueError( + "Could not detect DefaultResourceCalculator or " + "DominantResourceCalculator from archived scheduler evidence" + ) + + +def container_node_share(container: Container, mode: str) -> float: + """Return the YARN-scheduled node share without using Spark core counts.""" + if mode == "default": + return container.memory_mb / container.node_memory_mb + if mode != "dominant": + raise ValueError(f"Unsupported detected calculator mode {mode}") + allocated_resources = container.resources or { + "memory-mb": container.memory_mb, + "vcores": container.vcores, + "yarn.io/gpu": container.gpus, + } + node_resources = container.node_resources or { + "memory-mb": container.node_memory_mb, + "vcores": container.node_vcores, + "yarn.io/gpu": container.node_gpus, + } + shares = [] + for name, allocated in allocated_resources.items(): + if allocated <= 0: + continue + node_capacity = node_resources.get(name, 0) + if node_capacity <= 0: + raise ValueError( + f"Container {container.container_id} allocates {name} but its " + "node capacity is missing or zero" + ) + shares.append(allocated / node_capacity) + if not shares: + raise ValueError(f"Container {container.container_id} has no resources") + return max(shares) + + +def container_instance_type(evidence: YarnEvidence, container: Container) -> str: + node = evidence.nodes.get(container.node_id) + if node and node.instance_type: + return node.instance_type + return f"unknown:{container.node_id}" diff --git a/yarn-resource-cost/yarn_job_cost_dataproc.py b/yarn-resource-cost/yarn_job_cost_dataproc.py new file mode 100644 index 0000000..1c3b566 --- /dev/null +++ b/yarn-resource-cost/yarn_job_cost_dataproc.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Dataproc cluster metadata adapter.""" + +from __future__ import annotations + +import json + +from yarn_job_cost_adapters import run_command +from yarn_job_cost_core import YarnEvidence + + +def _last_uri_component(value: object) -> str: + return str(value or "").rstrip("/").rsplit("/", 1)[-1] + + +def _node_class(config: dict) -> str: + machine = _last_uri_component(config.get("machineTypeUri")) + accelerators = [] + for accelerator in config.get("accelerators") or []: + name = _last_uri_component(accelerator.get("acceleratorTypeUri")) + count = int(accelerator.get("acceleratorCount") or 0) + if name and count: + accelerators.append(f"{count}x{name}") + suffix = "+" + "+".join(sorted(accelerators)) if accelerators else "" + return f"gcp:dataproc:{machine}{suffix}" if machine else "" + + +def describe_node_classes( + cluster: str, region: str, project: str | None +) -> tuple[dict[str, str], dict[str, object]]: + command = [ + "gcloud", "dataproc", "clusters", "describe", cluster, + "--region", region, "--format", "json", + ] + if project: + command += ["--project", project] + payload = json.loads(run_command(command).stdout) + config = payload.get("config") or {} + mappings: dict[str, str] = {} + for group_name in ("workerConfig", "secondaryWorkerConfig"): + group = config.get(group_name) or {} + node_class = _node_class(group) + for name in group.get("instanceNames") or []: + if node_class: + mappings[str(name)] = node_class + provenance = { + "cluster": cluster, + "region": region, + "project": project or payload.get("projectId") or "", + "cluster_uuid": payload.get("clusterUuid") or "", + "source": "gcloud dataproc clusters describe", + } + return mappings, provenance + + +def classify_nodes( + evidence: YarnEvidence, cluster: str, region: str, project: str | None +) -> dict[str, object]: + mappings, provenance = describe_node_classes(cluster, region, project) + for node_id, node in evidence.nodes.items(): + short = node_id.split(".", 1)[0] + node.instance_type = mappings.get(node_id) or mappings.get(short) or node.instance_type + return provenance diff --git a/yarn-resource-cost/yarn_job_cost_defaults.py b/yarn-resource-cost/yarn_job_cost_defaults.py new file mode 100644 index 0000000..1956e58 --- /dev/null +++ b/yarn-resource-cost/yarn_job_cost_defaults.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Standalone defaults for the YARN job cost command.""" + +DEFAULT_AWS_PROFILE: str | None = None +DEFAULT_AWS_REGION: str | None = None diff --git a/yarn-resource-cost/yarn_job_cost_discovery.py b/yarn-resource-cost/yarn_job_cost_discovery.py new file mode 100644 index 0000000..8e68398 --- /dev/null +++ b/yarn-resource-cost/yarn_job_cost_discovery.py @@ -0,0 +1,490 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Discover benchmark applications and EMR logs for YARN cost accounting.""" + +from __future__ import annotations + +import csv +import hashlib +import json +import re +import shlex +import subprocess +from collections import defaultdict +from dataclasses import dataclass, field +from pathlib import Path +from typing import Iterable + + +from yarn_job_cost_eventlog import iter_eventlog_streams, iter_text_lines + + +APPLICATION_ID_RE = re.compile(r"application_\d+_\d+") +EVENTLOG_DIRECTORY_RE = re.compile(r"(eventlog_v2_(application_\d+_\d+))/") +EVENT_SEGMENT_RE = re.compile(r"(?:^|/)events_(?P\d+)(?:_|\.)") +CONTAINER_ID_RE = re.compile(r"container_\d+_\d+_\d+_\d+") +JOB_ID_RE = re.compile(r"(?:^|__)j0*(?P\d+)(?:__|\Z)", re.IGNORECASE) +BENCHMARK_OBJECT_JOB_ID_RE = re.compile(r"\bj0*(?P\d+)__", re.IGNORECASE) +EVENT_LOG_COLUMN_CANDIDATES = ( + "eventlog benchmark", + "event log benchmark", + "eventlog", + "event log", + "eventlog uri", + "event log uri", +) + + +@dataclass +class SparkExecutor: + executor_id: str + container_id: str = "" + total_cores: int | None = None + resource_profile_id: int = 0 + + +@dataclass +class EventLogApplication: + application_id: str + application_name: str = "" + job_id: str = "" + cluster_id: str = "" + emr_release_label: str = "" + spark_version: str = "" + start_ms: int | None = None + end_ms: int | None = None + application_ended: bool = False + configured_executor_cores: int | None = None + task_cpus: int = 1 + successful_task_attempt_count: int = 0 + task_duration_sum_ms: int = 0 + task_duration_ms_by_executor: dict[str, int] = field(default_factory=dict) + executors: dict[str, SparkExecutor] = field(default_factory=dict) + event_segments: set[int] = field(default_factory=set) + task_metric_warnings: list[str] = field(default_factory=list) + unsupported_resource_profile_ids: set[int] = field(default_factory=set) + + def event_task_metrics_complete(self) -> bool: + if not self.application_ended or self.task_metric_warnings: + return False + if not self.event_segments: + return False + expected = set(range(1, max(self.event_segments) + 1)) + return self.event_segments == expected + + def as_metadata(self) -> dict[str, str | float]: + duration: str | float = "" + if self.start_ms is not None and self.end_ms is not None: + duration = round((self.end_ms - self.start_ms) / 1000.0, 6) + return { + "job id": self.job_id, + "job name": self.application_name, + "spark_duration_seconds": duration, + "emr_cluster_id": self.cluster_id, + "emr_release_label": self.emr_release_label, + "spark_version": self.spark_version, + "configured_spark_executor_cores": ( + self.configured_executor_cores + if self.configured_executor_cores is not None + else "" + ), + "spark_task_cpus": self.task_cpus, + } + + +def aws_command( + aws_profile: str | None, aws_region: str | None = None +) -> list[str]: + command = ["aws"] + if aws_profile: + command += ["--profile", aws_profile] + if aws_region: + command += ["--region", aws_region] + return command + + +class AwsCliError(RuntimeError): + """An AWS CLI failure with the captured diagnostic output preserved.""" + + +def run_aws(command: list[str]) -> subprocess.CompletedProcess[str]: + completed = subprocess.run(command, capture_output=True, text=True) + if completed.returncode: + diagnostic = completed.stderr.strip() or completed.stdout.strip() + if not diagnostic: + diagnostic = "(AWS CLI produced no diagnostic output)" + raise AwsCliError( + f"AWS CLI failed with exit code {completed.returncode}:\n" + f"{diagnostic}\n" + f"Command: {shlex.join(command)}" + ) + return completed + + +def split_s3_uri(uri: str) -> tuple[str, str]: + if not uri.startswith(("s3://", "s3a://", "s3n://")): + raise ValueError(f"Expected an S3 URI, got {uri}") + bucket_and_key = uri.split("://", 1)[1] + bucket, _, key = bucket_and_key.partition("/") + if not bucket: + raise ValueError(f"S3 URI has no bucket: {uri}") + return bucket, key.rstrip("/") + "/" + + +def event_segment_number(uri: str) -> int: + match = EVENT_SEGMENT_RE.search(uri) + if not match: + raise ValueError(f"Cannot determine event segment number from {uri}") + return int(match.group("segment")) + + +def list_event_log_objects( + event_log_root: str, aws_profile: str | None, aws_region: str | None +) -> dict[str, list[str]]: + bucket, prefix = split_s3_uri(event_log_root) + command = aws_command(aws_profile, aws_region) + [ + "s3api", + "list-objects-v2", + "--bucket", + bucket, + "--prefix", + prefix, + "--output", + "json", + ] + completed = run_aws(command) + payload = json.loads(completed.stdout) + objects: dict[str, list[str]] = defaultdict(list) + for item in payload.get("Contents") or []: + key = str(item.get("Key") or "") + relative = key[len(prefix) :] if key.startswith(prefix) else key + match = EVENTLOG_DIRECTORY_RE.match(relative) + if match and Path(key).name.startswith("events_"): + objects[match.group(2)].append(f"s3://{bucket}/{key}") + if not objects: + raise ValueError( + f"No eventlog_v2_application_* directories found under {event_log_root}" + ) + return { + app_id: sorted(uris, key=event_segment_number) + for app_id, uris in objects.items() + } + + +def materialize_event_metadata_files( + event_log_root: str, + cache_dir: Path, + aws_profile: str | None, + aws_region: str | None, +) -> tuple[Path, set[str]]: + if not event_log_root.startswith(("s3://", "s3a://", "s3n://")): + path = Path(event_log_root).expanduser() + if not path.exists(): + raise FileNotFoundError(path) + application_ids: set[str] = set() + for match in APPLICATION_ID_RE.finditer(path.as_posix()): + application_ids.add(match.group()) + for child in path.glob("eventlog_v2_application_*"): + match = APPLICATION_ID_RE.search(child.name) + if match: + application_ids.add(match.group()) + return path, application_ids + + objects = list_event_log_objects(event_log_root, aws_profile, aws_region) + digest = hashlib.sha256(event_log_root.rstrip("/").encode()).hexdigest()[:16] + target = cache_dir / f"event-root-{digest}" + target.mkdir(parents=True, exist_ok=True) + for app_id, uris in objects.items(): + for uri in uris: + app_dir = target / f"eventlog_v2_{app_id}" + app_dir.mkdir(parents=True, exist_ok=True) + destination = app_dir / Path(uri).name + if destination.is_file(): + continue + command = aws_command(aws_profile, aws_region) + [ + "s3", + "cp", + uri, + str(destination), + "--only-show-errors", + ] + run_aws(command) + return target, set(objects) + + +def derive_job_id(application_name: str) -> str: + match = JOB_ID_RE.search(application_name) + if not match: + return "" + job_id = int(match.group("job_id")) + if "rewrite" in application_name.lower(): + job_id += 900000 + return str(job_id) + + +def benchmark_application_signature(application_name: str) -> str: + parent = Path(application_name).parent.name.lower() + parent = re.sub(r"^\d+_", "", parent) + return parent.replace("__rewrite", "") + + +def task_end_succeeded(event: dict) -> bool: + reason = event.get("Task End Reason") + if isinstance(reason, dict): + return reason.get("Reason") == "Success" + return reason == "Success" + + +def integer_property(properties: dict, name: str, default: int | None) -> int | None: + value = properties.get(name) + if value in (None, ""): + return default + try: + return int(value) + except (TypeError, ValueError): + return None + + +def read_event_log_metadata(path: Path) -> dict[str, EventLogApplication]: + applications: dict[str, EventLogApplication] = {} + by_directory: dict[str, EventLogApplication] = {} + for app_dir, member_name, stream, compressed in iter_eventlog_streams(path): + current = by_directory.setdefault( + app_dir, EventLogApplication(application_id="") + ) + segment_match = EVENT_SEGMENT_RE.search(member_name) + if segment_match: + current.event_segments.add(int(segment_match.group("segment"))) + for line in iter_text_lines(stream, compressed): + if not line: + continue + try: + event = json.loads(line) + except json.JSONDecodeError: + continue + event_name = event.get("Event") + if event_name == "SparkListenerLogStart": + current.spark_version = str(event.get("Spark Version") or "") + elif event_name == "SparkListenerEnvironmentUpdate": + properties = event.get("Spark Properties") or {} + current.cluster_id = str(properties.get("spark.emr.clusterId") or "") + current.emr_release_label = str( + properties.get("spark.emr.releaseLabel") or "" + ) + current.configured_executor_cores = integer_property( + properties, "spark.executor.cores", None + ) + task_cpus = integer_property(properties, "spark.task.cpus", 1) + if task_cpus is None or task_cpus <= 0: + current.task_metric_warnings.append( + "spark.task.cpus is missing or invalid" + ) + else: + current.task_cpus = task_cpus + elif event_name == "SparkListenerApplicationStart": + current.application_id = str(event.get("App ID") or "") + current.application_name = str(event.get("App Name") or "") + current.job_id = derive_job_id(current.application_name) + try: + current.start_ms = int(event.get("Timestamp")) + except (TypeError, ValueError): + pass + elif event_name == "SparkListenerApplicationEnd": + current.application_ended = True + try: + current.end_ms = int(event.get("Timestamp")) + except (TypeError, ValueError): + pass + elif event_name == "SparkListenerExecutorAdded": + executor_id = str(event.get("Executor ID") or "") + info = event.get("Executor Info") or {} + attributes = info.get("Attributes") or {} + container_id = str(attributes.get("CONTAINER_ID") or "") + if not container_id: + match = CONTAINER_ID_RE.search(line) + container_id = match.group() if match else "" + try: + total_cores = int(info.get("Total Cores")) + except (TypeError, ValueError): + total_cores = None + try: + resource_profile_id = int(info.get("Resource Profile Id") or 0) + except (TypeError, ValueError): + resource_profile_id = -1 + if resource_profile_id != 0: + current.unsupported_resource_profile_ids.add(resource_profile_id) + if executor_id: + current.executors[executor_id] = SparkExecutor( + executor_id=executor_id, + container_id=container_id, + total_cores=total_cores, + resource_profile_id=resource_profile_id, + ) + elif event_name == "SparkListenerResourceProfileAdded": + try: + profile_id = int(event.get("Resource Profile Id") or 0) + except (TypeError, ValueError): + profile_id = -1 + if profile_id != 0: + current.unsupported_resource_profile_ids.add(profile_id) + elif event_name == "SparkListenerStageSubmitted": + stage_info = event.get("Stage Info") or {} + try: + profile_id = int(stage_info.get("Resource Profile Id") or 0) + except (TypeError, ValueError): + profile_id = -1 + if profile_id != 0: + current.unsupported_resource_profile_ids.add(profile_id) + elif event_name == "SparkListenerTaskEnd" and task_end_succeeded(event): + task_info = event.get("Task Info") or {} + if task_info.get("Failed") or task_info.get("Killed"): + current.task_metric_warnings.append( + "A successful task event is marked failed or killed" + ) + continue + executor_id = str(task_info.get("Executor ID") or "") + try: + launch_ms = int(task_info.get("Launch Time")) + finish_ms = int(task_info.get("Finish Time")) + except (TypeError, ValueError): + current.task_metric_warnings.append( + "A successful task event has invalid timestamps" + ) + continue + if not executor_id or finish_ms < launch_ms: + current.task_metric_warnings.append( + "A successful task event has invalid executor or duration" + ) + continue + duration_ms = finish_ms - launch_ms + current.successful_task_attempt_count += 1 + current.task_duration_sum_ms += duration_ms + current.task_duration_ms_by_executor[executor_id] = ( + current.task_duration_ms_by_executor.get(executor_id, 0) + + duration_ms + ) + if not current.job_id: + object_match = BENCHMARK_OBJECT_JOB_ID_RE.search(line) + if object_match: + job_id = int(object_match.group("job_id")) + if "rewrite" in current.application_name.lower(): + job_id += 900000 + current.job_id = str(job_id) + if current.application_id: + applications[current.application_id] = current + + for current in applications.values(): + if not current.application_ended: + current.task_metric_warnings.append( + "SparkListenerApplicationEnd is missing" + ) + if not current.event_segments: + current.task_metric_warnings.append( + "Event-log segment numbers are unavailable" + ) + else: + expected = set(range(1, max(current.event_segments) + 1)) + missing = sorted(expected - current.event_segments) + if missing: + current.task_metric_warnings.append( + "Missing event-log segments: " + ", ".join(map(str, missing)) + ) + if current.unsupported_resource_profile_ids: + current.task_metric_warnings.append( + "Unsupported non-default Spark resource profiles: " + + ", ".join( + map(str, sorted(current.unsupported_resource_profile_ids)) + ) + ) + current.task_metric_warnings = list( + dict.fromkeys(current.task_metric_warnings) + ) + + known_by_signature: dict[str, set[int]] = defaultdict(set) + for application in applications.values(): + if application.job_id: + signature = benchmark_application_signature( + application.application_name + ) + known_by_signature[signature].add(int(application.job_id) % 900000) + for application in applications.values(): + candidates = known_by_signature.get( + benchmark_application_signature(application.application_name), set() + ) + if not application.job_id and len(candidates) == 1: + application.job_id = str(next(iter(candidates))) + return applications + + +def resolve_emr_log_uri( + cluster_id: str, aws_profile: str | None, aws_region: str | None +) -> str: + command = aws_command(aws_profile, aws_region) + [ + "emr", + "describe-cluster", + "--cluster-id", + cluster_id, + "--query", + "Cluster.LogUri", + "--output", + "text", + ] + completed = run_aws(command) + log_root = completed.stdout.strip() + if not log_root or log_root == "None": + raise ValueError(f"EMR cluster {cluster_id} has no LogUri") + normalized = "s3://" + log_root.split("://", 1)[1] + if normalized.rstrip("/").endswith("/" + cluster_id): + return normalized.rstrip("/") + "/" + return normalized.rstrip("/") + f"/{cluster_id}/" + + +def normalize_header(value: str) -> str: + return " ".join(re.sub(r"[^a-z0-9]+", " ", value.lower()).split()) + + +def find_event_log_column(headers: Iterable[str], explicit: str | None) -> str: + headers = list(headers) + if explicit: + if explicit not in headers: + raise ValueError(f"CSV has no event-log column {explicit!r}") + return explicit + normalized = {normalize_header(header): header for header in headers} + for candidate in EVENT_LOG_COLUMN_CANDIDATES: + if candidate in normalized: + return normalized[candidate] + raise ValueError( + "Could not find the event-log column. Pass --event-log-column; " + f"available columns: {', '.join(headers)}" + ) + + +def row_value(row: dict[str, str], normalized_name: str) -> str: + for key, value in row.items(): + if normalize_header(key) == normalized_name and str(value).strip(): + return str(value).strip() + return "" + + +def load_csv_metadata( + path: Path | None, event_log_column: str | None +) -> dict[str, dict[str, str]]: + if path is None: + return {} + with path.open(newline="", encoding="utf-8-sig") as handle: + reader = csv.DictReader(handle) + if not reader.fieldnames: + raise ValueError("CSV has no header") + column = find_event_log_column(reader.fieldnames, event_log_column) + metadata = {} + for row in reader: + match = APPLICATION_ID_RE.search(str(row.get(column, ""))) + if not match: + continue + metadata[match.group()] = { + "job id": row_value(row, "job id"), + "job name": row_value(row, "job name"), + } + return metadata diff --git a/yarn-resource-cost/yarn_job_cost_eventlog.py b/yarn-resource-cost/yarn_job_cost_eventlog.py new file mode 100644 index 0000000..efc38a6 --- /dev/null +++ b/yarn-resource-cost/yarn_job_cost_eventlog.py @@ -0,0 +1,134 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Read Spark rolling event logs without requiring PySpark or python-lz4.""" + +from __future__ import annotations + +import tarfile +from pathlib import Path +from typing import BinaryIO, Iterable + + +LZ4_BLOCK_MAGIC = b"LZ4Block" +RAW_BLOCK = 0x10 +LZ4_COMPRESSED_BLOCK = 0x20 + + +def lz4_decompress_block(src: bytes, expected_len: int) -> bytes: + out = bytearray() + index = 0 + while index < len(src): + token = src[index] + index += 1 + literal_len = token >> 4 + if literal_len == 15: + while True: + extra = src[index] + index += 1 + literal_len += extra + if extra != 255: + break + out.extend(src[index : index + literal_len]) + index += literal_len + if index >= len(src): + break + + offset = src[index] | (src[index + 1] << 8) + index += 2 + match_len = token & 0x0F + if match_len == 15: + while True: + extra = src[index] + index += 1 + match_len += extra + if extra != 255: + break + match_len += 4 + if offset <= 0 or offset > len(out): + raise ValueError(f"Invalid LZ4 offset {offset} at compressed offset {index}") + start = len(out) - offset + for copy_index in range(match_len): + out.append(out[start + copy_index]) + if len(out) != expected_len: + raise ValueError(f"LZ4 block length mismatch: got {len(out)}, expected {expected_len}") + return bytes(out) + + +def iter_lz4block_chunks(stream: BinaryIO) -> Iterable[bytes]: + while True: + header = stream.read(21) + if not header: + return + if len(header) != 21: + raise ValueError("Truncated LZ4Block header") + if header[: len(LZ4_BLOCK_MAGIC)] != LZ4_BLOCK_MAGIC: + raise ValueError(f"Bad LZ4Block magic: {header[:8]!r}") + token = header[8] + compressed_len = int.from_bytes(header[9:13], "little") + decompressed_len = int.from_bytes(header[13:17], "little") + if compressed_len == 0 and decompressed_len == 0: + return + block = stream.read(compressed_len) + if len(block) != compressed_len: + raise ValueError("Truncated LZ4Block payload") + method = token & 0xF0 + if method == RAW_BLOCK: + yield block + elif method == LZ4_COMPRESSED_BLOCK: + yield lz4_decompress_block(block, decompressed_len) + else: + raise ValueError(f"Unsupported LZ4Block token {token:#x}") + + +def iter_text_lines(stream: BinaryIO, compressed: bool) -> Iterable[str]: + pending = "" + chunks = iter_lz4block_chunks(stream) if compressed else iter(lambda: stream.read(1024 * 1024), b"") + for chunk in chunks: + text = pending + chunk.decode("utf-8", errors="replace") + lines = text.splitlines(keepends=True) + pending = "" + for line in lines: + if line.endswith("\n") or line.endswith("\r"): + yield line.strip() + else: + pending = line + if pending.strip(): + yield pending.strip() + + +def normalized_member_name(name: str) -> str: + while name.startswith("./"): + name = name[2:] + return name + + +def iter_eventlog_streams(path: Path) -> Iterable[tuple[str, str, BinaryIO, bool]]: + if path.is_file() and tarfile.is_tarfile(path): + with tarfile.open(path, mode="r:*") as tar: + members = sorted( + (member for member in tar.getmembers() if member.isfile()), + key=lambda member: normalized_member_name(member.name), + ) + for member in members: + name = normalized_member_name(member.name) + if not member.isfile() or "/events_" not in name: + continue + extracted = tar.extractfile(member) + if extracted is None: + continue + app_dir = name.split("/", 1)[0] + with extracted: + yield app_dir, name, extracted, name.endswith(".lz4") + return + + if path.is_dir(): + files = sorted(p for p in path.rglob("events_*") if p.is_file()) + else: + files = [path] + for file_path in files: + rel = file_path.as_posix() + app_dir = file_path.parent.name if file_path.parent.name.startswith("eventlog_") else file_path.stem + with file_path.open("rb") as handle: + yield app_dir, rel, handle, file_path.suffix == ".lz4" diff --git a/yarn-resource-cost/yarn_resource_cost.py b/yarn-resource-cost/yarn_resource_cost.py new file mode 100644 index 0000000..847e431 --- /dev/null +++ b/yarn-resource-cost/yarn_resource_cost.py @@ -0,0 +1,441 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Attribute Spark application worker consumption from Spark and YARN logs.""" + +from __future__ import annotations + +import argparse +import csv +import json +import re +import sys +from datetime import datetime, timezone +from pathlib import Path + +import calculate_yarn_job_cost as reporting +from yarn_job_cost_adapters import ( + ADAPTERS, + AdapterCommandError, + PriceCatalog, + apply_catalog_costs, + apply_node_class_map, + load_price_catalog, + materialize_gcs, + materialize_hdfs, + materialize_yarn_logs, +) +from yarn_job_cost_core import calculator_mode, parse_yarn_logs +from yarn_job_cost_dataproc import classify_nodes as classify_dataproc_nodes +from yarn_job_cost_discovery import ( + AwsCliError, + EventLogApplication, + materialize_event_metadata_files, + read_event_log_metadata, + resolve_emr_log_uri, +) + + +PORTABLE_FIELDS = ( + "comparison_key", + "application_id", + "application_name", + "final_status", + "spark_duration_seconds", + "resource_calculator", + "container_count", + "container_seconds", + "memory_mb_seconds", + "vcore_seconds", + "gpu_seconds", + "node_equivalent_seconds", + "resource_expression", + "worker_cost", + "worker_cost_currency", + "complete", + "warnings", +) + +COMPARISON_FIELDS = ( + "comparison_key", + "baseline_application_id", + "test_application_id", + "baseline_instance_types", + "test_instance_types", + "baseline_wall_clock_seconds", + "test_wall_clock_seconds", + "wall_clock_factor", + "baseline_node_equivalent_seconds", + "test_node_equivalent_seconds", + "node_equivalent_factor", + "baseline_worker_cost", + "test_worker_cost", + "worker_cost_factor", + "baseline_complete", + "test_complete", + "warnings", +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--adapter", choices=ADAPTERS, required=True) + parser.add_argument("--event-log-root", required=True) + parser.add_argument("--yarn-log-root") + parser.add_argument("--test-event-log-root", action="append", default=[]) + parser.add_argument("--test-yarn-log-root", action="append", default=[]) + parser.add_argument("--node-class-map", type=Path) + parser.add_argument("--pricing", choices=("none", "catalog", "live"), default="none") + parser.add_argument("--price-catalog", type=Path) + parser.add_argument("--aws-profile") + parser.add_argument("--aws-region") + parser.add_argument("--gcp-project") + parser.add_argument("--gcp-region") + parser.add_argument("--dataproc-cluster") + parser.add_argument("--refresh-cache", action="store_true") + parser.add_argument( + "--cache-dir", type=Path, default=Path(".cache/yarn-resource-cost") + ) + parser.add_argument("--include-application-master", action="store_true") + parser.add_argument( + "--comparison-key", choices=("application-name", "application-id", "regex"), + default="application-name", + ) + parser.add_argument("--comparison-key-regex") + parser.add_argument( + "--sort-by", choices=("comparison-key", "wall-clock-factor", "cost-factor"), + default="comparison-key", + ) + parser.add_argument("--output-csv", type=Path) + parser.add_argument("--output-json", type=Path) + return parser.parse_args() + + +def validate_args(args: argparse.Namespace) -> None: + if args.pricing == "catalog" and not args.price_catalog: + raise ValueError("--pricing catalog requires --price-catalog") + if args.pricing != "catalog" and args.price_catalog: + raise ValueError("--price-catalog requires --pricing catalog") + if args.pricing == "live" and args.adapter != "emr": + raise ValueError("Live pricing is currently supported only by --adapter emr") + if args.dataproc_cluster and args.adapter != "dataproc": + raise ValueError("--dataproc-cluster requires --adapter dataproc") + if args.dataproc_cluster and not args.gcp_region: + raise ValueError("--dataproc-cluster requires --gcp-region") + if args.comparison_key == "regex" and not args.comparison_key_regex: + raise ValueError("--comparison-key regex requires --comparison-key-regex") + if args.test_yarn_log_root and ( + len(args.test_yarn_log_root) != len(args.test_event_log_root) + ): + raise ValueError( + "--test-yarn-log-root must be supplied once per --test-event-log-root" + ) + if args.adapter != "emr" and not args.yarn_log_root: + raise ValueError(f"--adapter {args.adapter} requires --yarn-log-root") + + +def materialize_event_root(root: str, args: argparse.Namespace) -> tuple[Path, set[str]]: + event_cache = args.cache_dir / "spark-events" + if root.startswith("gs://"): + if args.adapter != "dataproc": + raise ValueError("gs:// Spark event logs require --adapter dataproc") + local = materialize_gcs(root, event_cache, args.refresh_cache) + return materialize_event_metadata_files(local.as_posix(), event_cache, None, None) + if root.startswith("hdfs://"): + if args.adapter != "on-prem": + raise ValueError("HDFS Spark event logs require --adapter on-prem") + local = materialize_hdfs(root, event_cache, args.refresh_cache) + return materialize_event_metadata_files(local.as_posix(), event_cache, None, None) + return materialize_event_metadata_files( + root, event_cache, args.aws_profile, args.aws_region + ) + + +def resolve_yarn_root( + explicit: str | None, + applications: dict[str, EventLogApplication], + args: argparse.Namespace, +) -> tuple[str, str]: + if explicit: + return explicit, "explicit" + if args.adapter != "emr": + raise ValueError(f"--adapter {args.adapter} requires an explicit YARN log root") + cluster_ids = {app.cluster_id for app in applications.values() if app.cluster_id} + if len(cluster_ids) != 1: + raise ValueError( + "Spark event logs must identify exactly one EMR cluster; found: " + + (", ".join(sorted(cluster_ids)) or "none") + ) + cluster_id = next(iter(cluster_ids)) + return ( + resolve_emr_log_uri(cluster_id, args.aws_profile, args.aws_region), + f"spark.emr.clusterId={cluster_id}", + ) + + +def application_key(application: dict, args: argparse.Namespace) -> str: + if args.comparison_key == "application-id": + return str(application.get("application_id") or "") + name = str(application.get("application_name") or application.get("job name") or "") + if args.comparison_key == "application-name": + return name + match = re.search(args.comparison_key_regex, name) + if not match: + return "" + if "key" in match.groupdict(): + return str(match.group("key")) + return str(match.group(1) if match.groups() else match.group()) + + +def summarize_applications(applications: list[dict]) -> dict: + complete = [app for app in applications if app.get("complete") is True] + return { + "application_count": len(applications), + "complete_application_count": len(complete), + "incomplete_application_count": len(applications) - len(complete), + "node_equivalent_seconds": round( + sum(float(app.get("node_equivalent_seconds") or 0) for app in complete), 6 + ), + "worker_cost": round( + sum(float(app.get("worker_cost") or 0) for app in complete), 8 + ), + } + + +def analyze_run( + event_root: str, + yarn_root: str | None, + args: argparse.Namespace, + catalog: PriceCatalog | None, +) -> dict: + local_events, selected_ids = materialize_event_root(event_root, args) + event_metadata = read_event_log_metadata(local_events) + missing = selected_ids - set(event_metadata) + if missing: + raise ValueError("Missing Spark event-log metadata for: " + ", ".join(sorted(missing))) + resolved_yarn_root, discovery_source = resolve_yarn_root( + yarn_root, event_metadata, args + ) + local_yarn = materialize_yarn_logs( + args.adapter, + resolved_yarn_root, + args.cache_dir / "yarn-logs", + args.refresh_cache, + args.aws_profile, + ) + evidence = parse_yarn_logs(local_yarn) + node_classification: dict[str, object] = { + "source": "node-class-map" if args.node_class_map else "daemon-log attributes" + } + if args.adapter == "dataproc" and args.dataproc_cluster: + node_classification = classify_dataproc_nodes( + evidence, args.dataproc_cluster, args.gcp_region, args.gcp_project + ) + apply_node_class_map(evidence, args.node_class_map) + mode = calculator_mode(evidence.calculator_class) + metadata = {app_id: app.as_metadata() for app_id, app in event_metadata.items()} + executor_containers = { + executor.container_id + for app in event_metadata.values() + for executor in app.executors.values() + if executor.container_id + } + applications = reporting.calculate_applications( + evidence, + mode, + metadata, + args.include_application_master, + executor_containers, + ) + applications = [app for app in applications if app["application_id"] in selected_ids] + missing_yarn = selected_ids - {app["application_id"] for app in applications} + if missing_yarn: + raise ValueError("Applications missing from YARN logs: " + ", ".join(sorted(missing_yarn))) + reporting.add_task_packing_metrics(applications, event_metadata, evidence, mode) + pricing_provenance: dict[str, object] = {"mode": args.pricing} + if args.pricing == "catalog": + apply_catalog_costs(applications, catalog) + pricing_provenance.update(catalog.provenance) + elif args.pricing == "live": + region = reporting.resolve_aws_region(args.aws_region, args.aws_profile) + prices = reporting.add_ondemand_costs(applications, region, args.aws_profile) + for app in applications: + app["worker_cost"] = app.get("ec2_plus_emr_usd", "") + app["worker_cost_currency"] = "USD" + pricing_provenance = { + "mode": "live", + "provider": "aws", + "region": region, + "queried_at_utc": datetime.now(timezone.utc).isoformat(), + "components": prices, + } + else: + for app in applications: + app["worker_cost"] = "" + app["worker_cost_currency"] = "" + for app in applications: + app["comparison_key"] = application_key(app, args) + app["resource_expression"] = app.get("cost_expression", "") + return { + "schema_version": 1, + "adapter": args.adapter, + "event_log_root": event_root, + "yarn_log_root": resolved_yarn_root, + "yarn_log_discovery_source": discovery_source, + "resource_calculator": mode, + "calculator_evidence": { + "class": evidence.calculator_class, + "source": evidence.calculator_source, + "fair_scheduler_policy": evidence.scheduler_policy, + }, + "pricing": pricing_provenance, + "node_classification": node_classification, + "nodes": { + node_id: { + "node_class": node.node_class, + "resources": node.resources or { + "memory-mb": node.memory_mb, + "vcores": node.vcores, + "yarn.io/gpu": node.gpus, + }, + } + for node_id, node in sorted(evidence.nodes.items()) + }, + "applications": applications, + "summary": summarize_applications(applications), + } + + +def factor(test: object, baseline: object) -> float | str: + if baseline in ("", None, 0) or test in ("", None): + return "" + return round(float(test) / float(baseline), 8) + + +def instance_types(application: dict) -> str: + return " + ".join(sorted(application["node_equivalent_seconds_by_instance_type"])) + + +def compare_runs(baseline: dict, tests: list[dict], args: argparse.Namespace) -> list[dict]: + baseline_index = { + app["comparison_key"]: app for app in baseline["applications"] if app["comparison_key"] + } + test_index: dict[str, dict] = {} + for run in tests: + for app in run["applications"]: + if app["comparison_key"]: + test_index[app["comparison_key"]] = app + rows = [] + for key in sorted(set(baseline_index) | set(test_index)): + base = baseline_index.get(key) + test = test_index.get(key) + complete = bool(base and test and base.get("complete") and test.get("complete")) + base_cost = base.get("worker_cost", "") if complete else "" + test_cost = test.get("worker_cost", "") if complete else "" + row = { + "comparison_key": key, + "baseline_application_id": base.get("application_id", "") if base else "", + "test_application_id": test.get("application_id", "") if test else "", + "baseline_instance_types": instance_types(base) if base else "", + "test_instance_types": instance_types(test) if test else "", + "baseline_wall_clock_seconds": base.get("spark_duration_seconds", "") if base else "", + "test_wall_clock_seconds": test.get("spark_duration_seconds", "") if test else "", + "baseline_node_equivalent_seconds": ( + base.get("node_equivalent_seconds", "") if base else "" + ), + "test_node_equivalent_seconds": test.get("node_equivalent_seconds", "") if test else "", + "baseline_worker_cost": base_cost, + "test_worker_cost": test_cost, + "baseline_complete": base.get("complete", False) if base else False, + "test_complete": test.get("complete", False) if test else False, + "warnings": " | ".join( + (["missing baseline application"] if not base else []) + + (["missing test application"] if not test else []) + ), + } + row["wall_clock_factor"] = factor( + row["test_wall_clock_seconds"], row["baseline_wall_clock_seconds"] + ) + row["node_equivalent_factor"] = ( + factor( + row["test_node_equivalent_seconds"], + row["baseline_node_equivalent_seconds"], + ) + if complete + else "" + ) + row["worker_cost_factor"] = factor(test_cost, base_cost) + rows.append(row) + sort_field = { + "comparison-key": "comparison_key", + "wall-clock-factor": "wall_clock_factor", + "cost-factor": "worker_cost_factor", + }[args.sort_by] + return sorted(rows, key=lambda row: (row[sort_field] == "", row[sort_field])) + + +def write_csv(path: Path, rows: list[dict], fields: tuple[str, ...]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fields, extrasaction="ignore") + writer.writeheader() + for source in rows: + row = dict(source) + if isinstance(row.get("warnings"), list): + row["warnings"] = " | ".join(row["warnings"]) + writer.writerow(row) + + +def print_table(rows: list[dict], fields: tuple[str, ...]) -> None: + rendered = [[str(row.get(field, "")) for field in fields] for row in rows] + widths = [ + max(len(field), *(len(row[index]) for row in rendered)) + for index, field in enumerate(fields) + ] + print(" ".join(field.ljust(widths[i]) for i, field in enumerate(fields))) + print(" ".join("-" * width for width in widths)) + for row in rendered: + print(" ".join(value.ljust(widths[i]) for i, value in enumerate(row))) + + +def main() -> int: + args = parse_args() + validate_args(args) + catalog = load_price_catalog(args.price_catalog) if args.price_catalog else None + baseline = analyze_run(args.event_log_root, args.yarn_log_root, args, catalog) + test_runs = [] + for index, root in enumerate(args.test_event_log_root): + explicit = args.test_yarn_log_root[index] if args.test_yarn_log_root else None + test_runs.append(analyze_run(root, explicit, args, catalog)) + if test_runs: + rows = compare_runs(baseline, test_runs, args) + fields = COMPARISON_FIELDS + payload: dict = { + "schema_version": 1, + "baseline": baseline, + "test_runs": test_runs, + "comparison": rows, + } + else: + rows = baseline["applications"] + fields = PORTABLE_FIELDS + payload = baseline + if args.output_csv: + write_csv(args.output_csv, rows, fields) + if args.output_json: + args.output_json.parent.mkdir(parents=True, exist_ok=True) + args.output_json.write_text( + json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + if not args.output_csv and not args.output_json: + print_table(rows, fields) + return 0 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except (AdapterCommandError, AwsCliError, ValueError) as error: + print(f"error: {error}", file=sys.stderr) + raise SystemExit(2) from None