diff --git a/.github/workflows/yarn-resource-cost-python.yml b/.github/workflows/yarn-resource-cost-python.yml index 627fd13..7beb29f 100644 --- a/.github/workflows/yarn-resource-cost-python.yml +++ b/.github/workflows/yarn-resource-cost-python.yml @@ -29,5 +29,29 @@ jobs: 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: Installed console error handling + run: | + python -m pip install ./yarn-resource-cost + python - <<'PY' + import shutil + import subprocess + + command = shutil.which("yarn-resource-cost") + assert command is not None + completed = subprocess.run( + [ + command, + "--adapter", + "on-prem", + "--event-log-root", + "/definitely/not/a/yarn/event/log", + ], + capture_output=True, + text=True, + ) + assert completed.returncode == 2, completed + assert completed.stderr.startswith("error: "), completed.stderr + assert "Traceback" not in completed.stderr, completed.stderr + 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/yarn-resource-cost/README.md b/yarn-resource-cost/README.md index bed3114..c07edb2 100644 --- a/yarn-resource-cost/README.md +++ b/yarn-resource-cost/README.md @@ -38,6 +38,50 @@ DominantResourceCalculator: NodeEquivalentSeconds = ContainerSeconds * NodeShare ``` +## Python API + +Install the subproject and its optional AWS dependency: + +```bash +python3 -m pip install './yarn-resource-cost[aws]' +``` + +The application-scoped API accepts injected boto3 clients and returns resource +usage without imposing a pricing policy: + +```python +import boto3 + +from yarn_resource_cost import ( + EmrApplicationUsageRequest, + calculate_emr_application_usage, +) + +session = boto3.Session(region_name="us-west-2") +usage = calculate_emr_application_usage( + EmrApplicationUsageRequest( + cluster_id="j-EXAMPLE", + application_id="application_123_0001", + event_log_uri="s3://example-bucket/spark-events/eventlog_v2_application_123_0001/", + region="us-west-2", + ), + emr_client=session.client("emr"), + s3_client=session.client("s3"), +) +print(usage.instance_seconds_by_type) +``` + +`event_log_uri` accepts an S3 URI, a plain local path, or a local `file://` URI. +Passing a pre-materialized local file or rolling-event-log directory avoids an +S3 download; the selected event log is still streamed to extract accounting +metadata. + +Incomplete archived logs return `complete=False` and indicate whether a later +retry can help. Missing summaries, allocations, terminal transitions, or node +registration metadata are retryable; contradictory or unsupported accounting +policies are not. Authentication and transport errors propagate from boto3. +The caller decides whether and how to translate instance-seconds into currency. + 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: diff --git a/yarn-resource-cost/calculate_yarn_job_cost.py b/yarn-resource-cost/calculate_yarn_job_cost.py index 200bb6c..e2d8349 100644 --- a/yarn-resource-cost/calculate_yarn_job_cost.py +++ b/yarn-resource-cost/calculate_yarn_job_cost.py @@ -754,6 +754,17 @@ def calculate_applications( ) if unknown_instance_type: warnings.append("One or more allocated containers have an unknown instance type") + transient_incomplete_evidence = ( + incomplete > 0 + or not coverage_complete + or nm_start_fallbacks > 0 + or nm_finish_fallbacks > 0 + or unknown_instance_type + or missing_gpu_capacity + or bool(resource_capacity_errors) + ) + permanent_incomplete_evidence = evidence.accounting_policy_ambiguous + complete = not transient_incomplete_evidence and not permanent_incomplete_evidence starts = [container.start_ms for container in containers] finishes = [ container.finish_ms for container in complete_containers if container.finish_ms is not None @@ -789,16 +800,11 @@ def calculate_applications( "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 - ), + "complete": complete, + # A fresh archive snapshot can resolve missing summaries, allocations, + # terminal transitions, and incomplete node registration metadata. It + # cannot resolve conflicting accounting-policy evidence already present. + "retryable": not complete and not permanent_incomplete_evidence, "warnings": warnings, } results.append(result) diff --git a/yarn-resource-cost/package_yarn_job_cost.py b/yarn-resource-cost/package_yarn_job_cost.py index ad8b526..a86bccd 100644 --- a/yarn-resource-cost/package_yarn_job_cost.py +++ b/yarn-resource-cost/package_yarn_job_cost.py @@ -22,7 +22,9 @@ PROJECT_DIR = Path(__file__).resolve().parent DEFAULT_LICENSE_FILE = PROJECT_DIR.parent / "LICENSE" PACKAGE_SOURCES = { + PROJECT_DIR / "pyproject.toml": "pyproject.toml", PROJECT_DIR / "yarn_resource_cost.py": "yarn_resource_cost.py", + PROJECT_DIR / "yarn_job_cost_api.py": "yarn_job_cost_api.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", @@ -36,6 +38,7 @@ 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 / "test_yarn_job_cost_api.py": "test_yarn_job_cost_api.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", diff --git a/yarn-resource-cost/pyproject.toml b/yarn-resource-cost/pyproject.toml new file mode 100644 index 0000000..9683099 --- /dev/null +++ b/yarn-resource-cost/pyproject.toml @@ -0,0 +1,36 @@ +[build-system] +requires = ["setuptools>=77", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "yarn-resource-cost" +version = "0.1.0" +description = "Portable YARN resource accounting for Spark applications" +readme = "README.md" +requires-python = ">=3.10,<3.13" +license = "Apache-2.0" +classifiers = [ + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", +] + +[project.optional-dependencies] +aws = ["boto3>=1.34"] + +[project.scripts] +yarn-resource-cost = "yarn_resource_cost:cli_main" + +[tool.setuptools] +py-modules = [ + "calculate_yarn_job_cost", + "yarn_job_cost_adapters", + "yarn_job_cost_api", + "yarn_job_cost_core", + "yarn_job_cost_dataproc", + "yarn_job_cost_defaults", + "yarn_job_cost_discovery", + "yarn_job_cost_eventlog", + "yarn_resource_cost", +] diff --git a/yarn-resource-cost/test_calculate_yarn_job_cost.py b/yarn-resource-cost/test_calculate_yarn_job_cost.py index aabd7c2..45c149e 100644 --- a/yarn-resource-cost/test_calculate_yarn_job_cost.py +++ b/yarn-resource-cost/test_calculate_yarn_job_cost.py @@ -328,8 +328,38 @@ def test_rm_terminal_precedes_nm_done_and_nm_fallback_is_not_final(self): self.assertEqual(10.0, result["container_seconds"]) self.assertEqual(1, result["nodemanager_finish_fallback_container_count"]) self.assertFalse(result["complete"]) + self.assertTrue(result["retryable"]) self.assertIn("NodeManager DONE fallback", " | ".join(result["warnings"])) + def test_ambiguous_accounting_policy_is_not_retryable(self): + container = MODULE.Container( + container_id="container_123_0002_01_000002", + application_id=APP_ID, + node_id="worker", + start_ms=1000, + finish_ms=2000, + memory_mb=40, + node_memory_mb=100, + vcores=1, + node_vcores=4, + source="resourcemanager", + finish_source="resourcemanager", + ) + summary = MODULE.ApplicationSummary(APP_ID, "test", "SUCCEEDED", 1) + + evidence = MODULE.YarnEvidence( + nodes={"worker": MODULE.Node("worker", "cpu.test", 100, 4, 0)}, + containers={container.container_id: container}, + calculator_class="DefaultResourceCalculator", + accounting_policy_ambiguous=True, + application_summaries={APP_ID: summary}, + ) + + result = MODULE.calculate_applications(evidence, "default", {}, False)[0] + + self.assertFalse(result["complete"]) + self.assertFalse(result["retryable"]) + def test_emr_log_cache_can_be_refreshed(self): with tempfile.TemporaryDirectory() as directory: cache = Path(directory) diff --git a/yarn-resource-cost/test_portable_yarn_resource_cost.py b/yarn-resource-cost/test_portable_yarn_resource_cost.py index d84841b..5854aba 100644 --- a/yarn-resource-cost/test_portable_yarn_resource_cost.py +++ b/yarn-resource-cost/test_portable_yarn_resource_cost.py @@ -106,6 +106,26 @@ def test_missing_catalog_rate_suppresses_final_cost(self): class PortableCliTest(unittest.TestCase): + def test_entry_point_formats_expected_errors(self): + completed = subprocess.run( + [ + sys.executable, + "-c", + "import yarn_resource_cost; yarn_resource_cost.cli_main()", + "--adapter", + "on-prem", + "--event-log-root", + "/definitely/not/a/yarn/event/log", + ], + cwd=ROOT, + capture_output=True, + text=True, + ) + + self.assertEqual(2, completed.returncode) + self.assertTrue(completed.stderr.startswith("error: "), completed.stderr) + self.assertNotIn("Traceback", completed.stderr) + def test_on_prem_fixture_end_to_end_with_catalog(self): with tempfile.TemporaryDirectory() as directory: output = Path(directory) / "result.json" diff --git a/yarn-resource-cost/test_yarn_job_cost_api.py b/yarn-resource-cost/test_yarn_job_cost_api.py new file mode 100644 index 0000000..c06cdba --- /dev/null +++ b/yarn-resource-cost/test_yarn_job_cost_api.py @@ -0,0 +1,337 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the application-scoped YARN resource accounting API.""" + +from __future__ import annotations + +import io +import unittest +from pathlib import Path + +from yarn_job_cost_api import ( + EmrApplicationUsageRequest, + calculate_emr_application_usage, +) + + +FIXTURE = Path(__file__).parent / "tests" / "fixtures" / "on_prem" + + +class FakePaginator: + def __init__(self, objects: dict[str, bytes]): + self.objects = objects + + def paginate(self, **kwargs): + prefix = kwargs["Prefix"] + yield { + "Contents": [ + {"Key": key, "Size": len(value)} + for key, value in self.objects.items() + if key.startswith(prefix) + ] + } + + +class FakeS3Client: + def __init__(self, objects: dict[str, bytes]): + self.objects = objects + + def get_paginator(self, operation: str): + if operation != "list_objects_v2": + raise AssertionError(operation) + return FakePaginator(self.objects) + + def get_object(self, **kwargs): + return {"Body": io.BytesIO(self.objects[kwargs["Key"]])} + + +class FakeEmrClient: + def describe_cluster(self, **kwargs): + return { + "Cluster": { + "Id": kwargs["ClusterId"], + "LogUri": "s3://test-bucket/emr-logs/", + } + } + + +class YarnJobCostApiTest(unittest.TestCase): + def request(self, event_log_uri=None): + return EmrApplicationUsageRequest( + cluster_id="j-TEST", + application_id="application_1_0001", + event_log_uri=event_log_uri + or str(FIXTURE / "eventlog_v2_application_1_0001"), + region="us-west-2", + ) + + @staticmethod + def yarn_log_with_instance_type(): + log = (FIXTURE / "yarn" / "hadoop-yarn-resourcemanager-rm.log").read_text( + encoding="utf-8" + ) + return log.replace( + "registered with capability: ", + "registered with capability: " + "instanceType(STRING)=m5.xlarge", + ) + + def test_calculates_one_application_from_boto_clients(self): + s3 = FakeS3Client( + { + "emr-logs/j-TEST/node/i-1/applications/" + "hadoop-yarn-resourcemanager-rm.log": self.yarn_log_with_instance_type().encode() + } + ) + + result = calculate_emr_application_usage( + self.request(), emr_client=FakeEmrClient(), s3_client=s3 + ) + + self.assertTrue(result.complete) + self.assertFalse(result.retryable) + self.assertEqual("default", result.resource_calculator) + self.assertEqual(16.0, result.vcore_seconds) + self.assertEqual(16384.0, result.memory_mb_seconds) + self.assertEqual({"m5.xlarge": 2.0}, result.instance_seconds_by_type) + self.assertEqual(1, result.container_count) + self.assertEqual(2, result.expected_container_count) + + def test_missing_archived_logs_is_retryable(self): + result = calculate_emr_application_usage( + self.request(), + emr_client=FakeEmrClient(), + s3_client=FakeS3Client({}), + ) + + self.assertFalse(result.complete) + self.assertTrue(result.retryable) + self.assertIsNone(result.vcore_seconds) + self.assertIn("No archived", result.warnings[0]) + + def test_missing_calculator_evidence_is_retryable(self): + log = "\n".join( + line + for line in self.yarn_log_with_instance_type().splitlines() + if "Initialized CapacityScheduler" not in line + ) + s3 = FakeS3Client( + { + "emr-logs/j-TEST/node/i-1/applications/" + "hadoop-yarn-resourcemanager-rm.log": log.encode() + } + ) + + result = calculate_emr_application_usage( + self.request(), emr_client=FakeEmrClient(), s3_client=s3 + ) + + self.assertFalse(result.complete) + self.assertTrue(result.retryable) + self.assertIn("Could not detect", result.warnings[0]) + + def test_conflicting_calculator_evidence_is_structured_and_not_retryable(self): + log = self.yarn_log_with_instance_type() + log += ( + "\n2026-01-01 00:00:00,050 INFO CapacityScheduler: " + "resource-calculator=DominantResourceCalculator\n" + ) + s3 = FakeS3Client( + { + "emr-logs/j-TEST/node/i-1/applications/" + "hadoop-yarn-resourcemanager-rm.log": log.encode() + } + ) + + result = calculate_emr_application_usage( + self.request(), emr_client=FakeEmrClient(), s3_client=s3 + ) + + self.assertFalse(result.complete) + self.assertFalse(result.retryable) + self.assertIn("Conflicting ResourceCalculators", result.warnings[0]) + + def test_missing_node_instance_type_is_retryable(self): + log = (FIXTURE / "yarn" / "hadoop-yarn-resourcemanager-rm.log").read_bytes() + s3 = FakeS3Client( + { + "emr-logs/j-TEST/node/i-1/applications/" + "hadoop-yarn-resourcemanager-rm.log": log + } + ) + + result = calculate_emr_application_usage( + self.request(), emr_client=FakeEmrClient(), s3_client=s3 + ) + + self.assertFalse(result.complete) + self.assertTrue(result.retryable) + self.assertIn("unknown instance type", " | ".join(result.warnings)) + + def test_s3_rolling_event_log_segments_stay_grouped(self): + event_prefix = "spark-events/eventlog_v2_application_1_0001" + segment_1 = "\n".join( + ( + '{"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}', + ) + ) + segment_2 = "\n".join( + ( + '{"Event":"SparkListenerExecutorAdded","Executor ID":"driver",' + '"Executor Info":{"Total Cores":1,"Resource Profile Id":0,' + '"Attributes":{"CONTAINER_ID":"container_1_0001_01_000001"}}}', + '{"Event":"SparkListenerApplicationEnd","Timestamp":11000}', + ) + ) + s3 = FakeS3Client( + { + f"{event_prefix}/events_1_application_1_0001": segment_1.encode(), + f"{event_prefix}/events_2_application_1_0001": segment_2.encode(), + "emr-logs/j-TEST/node/i-1/applications/" + "hadoop-yarn-resourcemanager-rm.log": self.yarn_log_with_instance_type().encode(), + } + ) + + result = calculate_emr_application_usage( + self.request(f"s3://test-bucket/{event_prefix}"), + emr_client=FakeEmrClient(), + s3_client=s3, + ) + + self.assertTrue(result.complete) + self.assertEqual(2, result.container_count) + self.assertEqual(24.0, result.vcore_seconds) + self.assertEqual(24576.0, result.memory_mb_seconds) + self.assertEqual({"m5.xlarge": 3.0}, result.instance_seconds_by_type) + + def test_s3_single_file_event_log_is_materialized_as_a_file(self): + event_key = "spark-events/application_1_0001" + event_log = ( + FIXTURE + / "eventlog_v2_application_1_0001" + / "events_1_application_1_0001" + ).read_bytes() + s3 = FakeS3Client( + { + event_key: event_log, + "emr-logs/j-TEST/node/i-1/applications/" + "hadoop-yarn-resourcemanager-rm.log": self.yarn_log_with_instance_type().encode(), + } + ) + + result = calculate_emr_application_usage( + self.request(f"s3://test-bucket/{event_key}"), + emr_client=FakeEmrClient(), + s3_client=s3, + ) + + self.assertTrue(result.complete) + self.assertFalse(result.retryable) + self.assertEqual(1, result.container_count) + self.assertEqual({"m5.xlarge": 2.0}, result.instance_seconds_by_type) + + def test_file_uri_event_log_file_and_directory(self): + s3 = FakeS3Client( + { + "emr-logs/j-TEST/node/i-1/applications/" + "hadoop-yarn-resourcemanager-rm.log": self.yarn_log_with_instance_type().encode() + } + ) + event_log_dir = FIXTURE / "eventlog_v2_application_1_0001" + event_log_file = event_log_dir / "events_1_application_1_0001" + + for event_log in (event_log_dir, event_log_file): + with self.subTest(event_log=event_log): + result = calculate_emr_application_usage( + self.request(event_log.resolve().as_uri()), + emr_client=FakeEmrClient(), + s3_client=s3, + ) + + self.assertTrue(result.complete) + self.assertFalse(result.retryable) + self.assertEqual({"m5.xlarge": 2.0}, result.instance_seconds_by_type) + + def test_remote_file_uri_is_rejected(self): + with self.assertRaisesRegex(ValueError, "remote authority"): + calculate_emr_application_usage( + self.request("file://remote-host/tmp/events"), + emr_client=FakeEmrClient(), + s3_client=FakeS3Client({}), + ) + + def test_later_node_registration_completes_missing_or_partial_metadata(self): + full_log = self.yarn_log_with_instance_type() + registration = next( + line for line in full_log.splitlines() if "registered with capability" in line + ) + application_log = "\n".join( + line for line in full_log.splitlines() if line != registration + ) + partial_registration = registration.replace( + " instanceType(STRING)=m5.xlarge", "" + ) + yarn_key = ( + "emr-logs/j-TEST/node/i-1/applications/" + "hadoop-yarn-resourcemanager-a-application.log" + ) + registration_key = ( + "emr-logs/j-TEST/node/i-1/applications/" + "hadoop-yarn-resourcemanager-z-registration.log" + ) + + for name, initial_log in ( + ("embedded max capacity", application_log), + ( + "missing max capacity", + application_log.replace(", max memory:8192", "").replace( + ", max vCores:8", "" + ), + ), + ( + "partial registration missing instance type", + full_log.replace(registration, partial_registration), + ), + ): + with self.subTest(name=name): + first = calculate_emr_application_usage( + self.request(), + emr_client=FakeEmrClient(), + s3_client=FakeS3Client({yarn_key: initial_log.encode()}), + ) + second = calculate_emr_application_usage( + self.request(), + emr_client=FakeEmrClient(), + s3_client=FakeS3Client( + { + yarn_key: initial_log.encode(), + registration_key: registration.encode(), + } + ), + ) + + self.assertFalse(first.complete) + self.assertTrue(first.retryable) + self.assertTrue(second.complete) + self.assertFalse(second.retryable) + self.assertEqual({"m5.xlarge": 2.0}, second.instance_seconds_by_type) + + def test_request_rejects_missing_identity(self): + with self.assertRaisesRegex(ValueError, "application_id is required"): + EmrApplicationUsageRequest( + cluster_id="j-TEST", + application_id="", + event_log_uri="s3://bucket/events", + region="us-west-2", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/yarn-resource-cost/yarn_job_cost_api.py b/yarn-resource-cost/yarn_job_cost_api.py new file mode 100644 index 0000000..6e32caa --- /dev/null +++ b/yarn-resource-cost/yarn_job_cost_api.py @@ -0,0 +1,303 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed library API for application-scoped YARN resource accounting.""" + +from __future__ import annotations + +import shutil +import tempfile +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any, Iterable +from urllib.parse import urlparse +from urllib.request import url2pathname + +import calculate_yarn_job_cost as reporting +from yarn_job_cost_core import calculator_mode, parse_yarn_logs +from yarn_job_cost_discovery import read_event_log_metadata + + +@dataclass(frozen=True) +class EmrApplicationUsageRequest: + """Inputs needed to attribute one EMR YARN application.""" + + cluster_id: str + application_id: str + event_log_uri: str + region: str + include_application_master: bool = False + + def __post_init__(self) -> None: + for name in ("cluster_id", "application_id", "event_log_uri", "region"): + if not getattr(self, name).strip(): + raise ValueError(f"{name} is required") + + +@dataclass(frozen=True) +class YarnApplicationUsageResult: + """Resource usage and evidence for one YARN application.""" + + application_id: str + complete: bool + retryable: bool + resource_calculator: str = "" + detected_resource_calculator_class: str = "" + vcore_seconds: float | None = None + memory_mb_seconds: float | None = None + instance_seconds_by_type: dict[str, float] = field(default_factory=dict) + container_count: int = 0 + expected_container_count: int | None = None + incomplete_container_count: int = 0 + warnings: tuple[str, ...] = () + + def as_dict(self) -> dict[str, Any]: + """Return a JSON-safe representation.""" + + return asdict(self) + + +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, separator, key = bucket_and_key.partition("/") + if not bucket: + raise ValueError(f"S3 URI has no bucket: {uri}") + return bucket, key if separator else "" + + +def _list_s3_objects(s3_client: Any, bucket: str, prefix: str) -> Iterable[dict[str, Any]]: + paginator = s3_client.get_paginator("list_objects_v2") + for page in paginator.paginate(Bucket=bucket, Prefix=prefix): + yield from page.get("Contents") or () + + +def _safe_relative_key(key: str, prefix: str) -> Path: + relative = key[len(prefix) :].lstrip("/") if key.startswith(prefix) else key + parts = [part for part in Path(relative).parts if part not in ("", ".", "..")] + if not parts: + parts = [Path(key).name or "download"] + return Path(*parts) + + +def _download_objects( + s3_client: Any, + bucket: str, + prefix: str, + objects: Iterable[dict[str, Any]], + destination: Path, +) -> list[Path]: + downloaded = [] + for item in objects: + key = str(item.get("Key") or "") + if not key or key.endswith("/"): + continue + target = destination / _safe_relative_key(key, prefix) + target.parent.mkdir(parents=True, exist_ok=True) + body = s3_client.get_object(Bucket=bucket, Key=key)["Body"] + try: + with target.open("wb") as output: + shutil.copyfileobj(body, output) + finally: + close = getattr(body, "close", None) + if close: + close() + downloaded.append(target) + return downloaded + + +def _materialize_event_log(s3_client: Any, uri: str, destination: Path) -> Path: + if not uri.startswith(("s3://", "s3a://", "s3n://")): + if uri.lower().startswith("file:"): + parsed = urlparse(uri) + if parsed.netloc and parsed.netloc.lower() != "localhost": + raise ValueError(f"File URI has a remote authority: {uri}") + if not parsed.path or parsed.query or parsed.fragment: + raise ValueError(f"Invalid file URI: {uri}") + path = Path(url2pathname(parsed.path)).expanduser() + else: + path = Path(uri).expanduser() + if not path.exists(): + raise ValueError(f"Event log does not exist: {uri}") + return path + + bucket, key = _split_s3_uri(uri) + prefix = key.rstrip("/") + objects = list(_list_s3_objects(s3_client, bucket, prefix)) + exact_object = next( + ( + item + for item in objects + if str(item.get("Key") or "") == key and not key.endswith("/") + ), + None, + ) + if exact_object is not None: + downloaded = _download_objects( + s3_client, bucket, prefix, (exact_object,), destination + ) + return downloaded[0] if downloaded else destination + selected = [ + item + for item in objects + if Path(str(item.get("Key") or "")).name.startswith("events_") + or str(item.get("Key") or "") == key + ] + if not selected: + return destination + prefix_name = Path(prefix).name + event_destination = ( + destination / prefix_name if prefix_name.startswith("eventlog_") else destination + ) + _download_objects(s3_client, bucket, prefix, selected, event_destination) + return destination + + +def _cluster_log_uri(emr_client: Any, cluster_id: str) -> str: + cluster = emr_client.describe_cluster(ClusterId=cluster_id).get("Cluster") or {} + log_uri = str(cluster.get("LogUri") or "").strip() + if not log_uri: + raise ValueError(f"EMR cluster {cluster_id} has no LogUri") + _scheme, separator, bucket_and_key = log_uri.partition("://") + if not separator or not bucket_and_key: + raise ValueError(f"EMR cluster {cluster_id} has an invalid LogUri") + normalized = "s3://" + bucket_and_key + if normalized.rstrip("/").endswith("/" + cluster_id): + return normalized.rstrip("/") + "/" + return normalized.rstrip("/") + f"/{cluster_id}/" + + +def _materialize_yarn_logs( + emr_client: Any, + s3_client: Any, + cluster_id: str, + destination: Path, +) -> list[Path]: + bucket, prefix = _split_s3_uri(_cluster_log_uri(emr_client, cluster_id)) + objects = [ + item + for item in _list_s3_objects(s3_client, bucket, prefix) + if any( + marker in Path(str(item.get("Key") or "")).name + for marker in ("hadoop-yarn-resourcemanager", "hadoop-yarn-nodemanager") + ) + ] + return _download_objects(s3_client, bucket, prefix, objects, destination) + + +def _empty_result( + request: EmrApplicationUsageRequest, + warning: str, + *, + retryable: bool, + detected_calculator: str = "", +) -> YarnApplicationUsageResult: + return YarnApplicationUsageResult( + application_id=request.application_id, + complete=False, + retryable=retryable, + detected_resource_calculator_class=detected_calculator, + warnings=(warning,), + ) + + +def calculate_emr_application_usage( + request: EmrApplicationUsageRequest, + *, + emr_client: Any, + s3_client: Any, +) -> YarnApplicationUsageResult: + """Calculate YARN resource usage for one EMR application. + + Missing or not-yet-complete archived logs are returned as retryable incomplete + results. Provider authentication and transport errors are allowed to propagate. + """ + + with tempfile.TemporaryDirectory(prefix="yarn-resource-cost-") as directory: + root = Path(directory) + event_root = _materialize_event_log(s3_client, request.event_log_uri, root / "events") + event_metadata = read_event_log_metadata(event_root) + metadata = event_metadata.get(request.application_id) + if metadata is None: + return _empty_result( + request, + f"Application {request.application_id} is missing from the Spark event log", + retryable=True, + ) + + yarn_root = root / "yarn" + downloaded = _materialize_yarn_logs(emr_client, s3_client, request.cluster_id, yarn_root) + if not downloaded: + return _empty_result( + request, + f"No archived ResourceManager or NodeManager logs found for cluster {request.cluster_id}", + retryable=True, + ) + + try: + evidence = parse_yarn_logs(yarn_root) + except ValueError as error: + return _empty_result(request, str(error), retryable=False) + try: + mode = calculator_mode(evidence.calculator_class) + except ValueError as error: + return _empty_result( + request, + str(error), + retryable=not bool(evidence.calculator_class), + detected_calculator=evidence.calculator_class, + ) + + executor_containers = { + executor.container_id + for executor in metadata.executors.values() + if executor.container_id + } + applications = reporting.calculate_applications( + evidence, + mode, + {request.application_id: metadata.as_metadata()}, + request.include_application_master, + executor_containers, + ) + application = next( + (item for item in applications if item["application_id"] == request.application_id), + None, + ) + if application is None: + return _empty_result( + request, + f"Application {request.application_id} is missing from archived YARN logs", + retryable=True, + detected_calculator=evidence.calculator_class, + ) + + expected = application.get("expected_total_allocated_containers") + return YarnApplicationUsageResult( + application_id=request.application_id, + complete=bool(application["complete"]), + retryable=bool(application["retryable"]), + resource_calculator=mode, + detected_resource_calculator_class=evidence.calculator_class, + vcore_seconds=float(application["vcore_seconds"]), + memory_mb_seconds=float(application["memory_mb_seconds"]), + instance_seconds_by_type={ + str(instance_type): float(seconds) + for instance_type, seconds in application[ + "node_equivalent_seconds_by_instance_type" + ].items() + }, + container_count=int(application["container_count"]), + expected_container_count=int(expected) if expected != "" else None, + incomplete_container_count=int(application["incomplete_container_count"]), + warnings=tuple(str(warning) for warning in application["warnings"]), + ) + + +__all__ = [ + "EmrApplicationUsageRequest", + "YarnApplicationUsageResult", + "calculate_emr_application_usage", +] diff --git a/yarn-resource-cost/yarn_job_cost_core.py b/yarn-resource-cost/yarn_job_cost_core.py index f7c24f7..a17d65e 100644 --- a/yarn-resource-cost/yarn_job_cost_core.py +++ b/yarn-resource-cost/yarn_job_cost_core.py @@ -158,7 +158,10 @@ def capacity(value: str | None, fallback: int | None, name: str) -> int: return int(value) if fallback is not None: return fallback - raise ValueError(f"Could not determine node {name} capacity") + # An allocation can be archived before the corresponding node-registration + # record. Preserve the container with an unknown capacity so the reporting + # layer can return a structured, retryable incomplete result. + return 0 def open_log(path: Path) -> TextIO: @@ -357,6 +360,19 @@ def parse_yarn_logs(path: Path) -> YarnEvidence: ) for container_id, container in evidence.containers.items(): + # Registration evidence may be stored in a later file than the + # allocation. Reconcile containers after every archived log was parsed. + node = evidence.nodes.get(container.node_id) + if node: + if container.node_memory_mb <= 0 and node.memory_mb: + container.node_memory_mb = node.memory_mb + if container.node_vcores <= 0 and node.vcores: + container.node_vcores = node.vcores + if container.node_gpus <= 0 and node.gpus: + container.node_gpus = node.gpus + for name, node_capacity in node.resources.items(): + if container.node_resources.get(name, 0) <= 0: + container.node_resources[name] = node_capacity if container_id in rm_finishes: container.finish_ms = rm_finishes[container_id] container.finish_source = "resourcemanager" @@ -380,6 +396,11 @@ def calculator_mode(detected_class: str) -> str: def container_node_share(container: Container, mode: str) -> float: """Return the YARN-scheduled node share without using Spark core counts.""" if mode == "default": + if container.node_memory_mb <= 0: + raise ValueError( + f"Container {container.container_id} allocates memory-mb but its " + "node capacity is missing or zero" + ) return container.memory_mb / container.node_memory_mb if mode != "dominant": raise ValueError(f"Unsupported detected calculator mode {mode}") diff --git a/yarn-resource-cost/yarn_resource_cost.py b/yarn-resource-cost/yarn_resource_cost.py index 847e431..0a007c8 100644 --- a/yarn-resource-cost/yarn_resource_cost.py +++ b/yarn-resource-cost/yarn_resource_cost.py @@ -26,6 +26,11 @@ materialize_hdfs, materialize_yarn_logs, ) +from yarn_job_cost_api import ( # noqa: F401 + EmrApplicationUsageRequest, + YarnApplicationUsageResult, + calculate_emr_application_usage, +) 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 ( @@ -433,9 +438,14 @@ def main() -> int: return 0 -if __name__ == "__main__": +def cli_main() -> None: + """Run the CLI with concise diagnostics for expected input failures.""" try: raise SystemExit(main()) except (AdapterCommandError, AwsCliError, ValueError) as error: print(f"error: {error}", file=sys.stderr) raise SystemExit(2) from None + + +if __name__ == "__main__": + cli_main()