Change to allow passing s3 uri to pick the logs from a separate s3 directory and limit to just resource manager logs

yarn-resource-cost/README.md

@@ -72,6 +72,13 @@ usage = calculate_emr_application_usage(
 print(usage.instance_seconds_by_type)
 ```
 
+ResourceManager and NodeManager logs are read from the cluster's EMR `LogUri`
+archive by default, which the EMR log pusher updates about every five minutes.
+Set `yarn_log_uri` to an S3 prefix holding this cluster's ResourceManager logs
+when a log shipper uploads them sooner. Only file names containing
+`hadoop-yarn-resourcemanager` are read from that prefix; when none exist yet,
+the `LogUri` archive is used instead.
+
 `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

yarn-resource-cost/calculate_yarn_job_cost.py

@@ -30,7 +30,7 @@
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import TextIO
+from typing import Callable, TextIO
 
 SCRIPT_DIR = Path(__file__).resolve().parent
 sys.path.insert(0, str(SCRIPT_DIR))
@@ -514,6 +514,7 @@ def materialize_emr_logs(
     cache_dir: Path,
     aws_profile: str | None,
     refresh: bool = False,
+    markers: tuple[str, ...] = ("hadoop-yarn-nodemanager", "hadoop-yarn-resourcemanager"),
 ) -> Path:
     if not uri.startswith(("s3://", "s3a://", "s3n://")):
         path = Path(uri).expanduser()
@@ -534,26 +535,47 @@ def materialize_emr_logs(
     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*",
-    ]
+    command += ["s3", "cp", "--recursive", s3_uri.rstrip("/") + "/", str(target), "--exclude", "*"]
+    for marker_name in markers:
+        command += ["--include", f"*{marker_name}*.log*"]
     subprocess.run(command, check=True)
-    if not relevant_log_files(target):
+    if not any(file.stat().st_size for file in relevant_log_files(target)):
         raise ValueError(f"No YARN logs downloaded from {s3_uri}")
     marker.write_text(s3_uri + "\n")
     return target
 
 
+def materialize_yarn_logs(
+    emr_log_uri: str | Callable[[], str],
+    cache_dir: Path,
+    aws_profile: str | None,
+    refresh: bool = False,
+    yarn_log_shipper_uri: str | None = None,
+) -> Path:
+    """Materialize ResourceManager/NodeManager logs, preferring a log shipper.
+
+    Use yarn_log_shipper_uri when a log shipper uploads ResourceManager logs
+    faster than the EMR log archive. When it contains no ResourceManager
+    logs, the EMR log URI is used instead; pass it as a callable to defer
+    resolving it until that fallback is needed. Shipped logs keep growing while the
+    cluster runs, so they are always re-synced rather than served from cache.
+    """
+    if yarn_log_shipper_uri:
+        try:
+            return materialize_emr_logs(
+                yarn_log_shipper_uri,
+                cache_dir,
+                aws_profile,
+                refresh=True,
+                markers=("hadoop-yarn-resourcemanager",),
+            )
+        except (ValueError, FileNotFoundError):
+            pass
+    if callable(emr_log_uri):
+        emr_log_uri = emr_log_uri()
+    return materialize_emr_logs(emr_log_uri, cache_dir, aws_profile, refresh=refresh)
+
+
 def aws_command(aws_profile: str | None, aws_region: str | None = None) -> list[str]:
     command = ["aws"]
     if aws_profile:
@@ -1811,6 +1833,11 @@ def merge_test_analyses(
 def run_comparison(args: argparse.Namespace) -> int:
     if not args.event_log_root:
         raise ValueError("--test-event-log-root requires --event-log-root")
+    if args.yarn_log_shipper_uri:
+        raise ValueError(
+            "--yarn-log-shipper-uri names one cluster's logs and cannot be applied to "
+            "a comparison; run each event-log root separately"
+        )
     baseline = analyze_event_root_for_comparison(args.event_log_root, args)
     test_roots = args.test_event_log_root
     test_runs = [
@@ -1961,6 +1988,16 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Refresh cached ResourceManager and NodeManager logs from S3",
     )
+    parser.add_argument(
+        "--yarn-log-shipper-uri",
+        help=(
+            "Optional S3 prefix holding this cluster's ResourceManager logs from a "
+            "log shipper. Use this when the shipper uploads faster than the EMR log "
+            "archive; falls back to the resolved EMR log URI when it contains no "
+            "ResourceManager logs. Always re-downloaded, never served from cache. "
+            "Not supported with --test-event-log-root."
+        ),
+    )
     parser.add_argument(
         "--aws-profile",
         default=DEFAULT_AWS_PROFILE,
@@ -1991,6 +2028,7 @@ def main() -> int:
     selected_application_ids: set[str] | None = None
     cluster_id = ""
     emr_log_uri = args.emr_log_uri
+    resolved_emr_log_uri: list[str] = []
     local_event_path: Path | None = None
 
     if args.event_log_root:
@@ -2012,14 +2050,27 @@ def main() -> int:
                 + (", ".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,
+        def emr_log_uri_resolver() -> str:
+            resolved = resolve_emr_log_uri(cluster_id, args.aws_profile, args.aws_region)
+            resolved_emr_log_uri.append(resolved)
+            return resolved
+
+        log_source: str | Callable[[], str] = emr_log_uri_resolver
+    else:
+        log_source = emr_log_uri
+
+    local_logs = materialize_yarn_logs(
+        log_source,
         args.cache_dir,
         args.aws_profile,
         refresh=args.refresh_emr_log_cache,
+        yarn_log_shipper_uri=args.yarn_log_shipper_uri,
     )
+    if resolved_emr_log_uri:
+        emr_log_uri = resolved_emr_log_uri[0]
+    elif args.event_log_root:
+        emr_log_uri = ""  # LogUri lookup skipped: the log shipper supplied the logs
     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()}
@@ -2073,6 +2124,7 @@ def main() -> int:
             "local_event_metadata_path": str(local_event_path) if local_event_path else "",
             "emr_cluster_id": cluster_id,
             "emr_log_uri": emr_log_uri,
+            "yarn_log_shipper_uri": args.yarn_log_shipper_uri or "",
             "local_log_path": str(local_logs),
             "resource_calculator": mode,
             "detected_resource_calculator_class": evidence.calculator_class,

yarn-resource-cost/test_calculate_yarn_job_cost.py

@@ -380,7 +380,10 @@ def test_emr_log_cache_can_be_refreshed(self):
                 )
                 run.assert_not_called()
 
-            with mock.patch.object(MODULE.subprocess, "run") as run:
+            def fake_download(command, check):
+                (target / "hadoop-yarn-resourcemanager.log").write_text("fixture\n")
+
+            with mock.patch.object(MODULE.subprocess, "run", side_effect=fake_download) as run:
                 self.assertEqual(
                     target,
                     MODULE.materialize_emr_logs(uri, cache, None, refresh=True),

yarn-resource-cost/test_yarn_job_cost_api.py

@@ -8,6 +8,7 @@
 
 import io
 import unittest
+from dataclasses import replace
 from pathlib import Path
 from types import SimpleNamespace
 from unittest import mock
@@ -115,6 +116,82 @@ def test_missing_archived_logs_is_retryable(self):
         self.assertIsNone(result.vcore_seconds)
         self.assertIn("No archived", result.warnings[0])
 
+    def test_shipped_resourcemanager_logs_are_preferred(self):
+        s3 = FakeS3Client(
+            {
+                "yarn-rm-logs/j-TEST/ip-10-0-0-1/"
+                "hadoop-yarn-resourcemanager-20260928-120000-a.log": self.yarn_log_with_instance_type().encode(),
+                "emr-logs/j-TEST/node/i-1/applications/"
+                "hadoop-yarn-resourcemanager-rm.log": b"not a yarn log",
+            }
+        )
+        request = replace(self.request(), yarn_log_uri="s3://test-bucket/yarn-rm-logs/j-TEST/")
+
+        with mock.patch.object(FakeEmrClient, "describe_cluster") as describe_cluster:
+            result = calculate_emr_application_usage(
+                request, emr_client=FakeEmrClient(), s3_client=s3
+            )
+
+        describe_cluster.assert_not_called()
+        self.assertTrue(result.complete)
+        self.assertEqual({"m5.xlarge": 2.0}, result.instance_seconds_by_type)
+
+    def test_empty_shipped_log_prefix_falls_back_to_cluster_log_uri(self):
+        s3 = FakeS3Client(
+            {
+                "emr-logs/j-TEST/node/i-1/applications/"
+                "hadoop-yarn-resourcemanager-rm.log": self.yarn_log_with_instance_type().encode()
+            }
+        )
+        request = replace(self.request(), yarn_log_uri="s3://test-bucket/yarn-rm-logs/j-TEST/")
+
+        result = calculate_emr_application_usage(
+            request, emr_client=FakeEmrClient(), s3_client=s3
+        )
+
+        self.assertTrue(result.complete)
+        self.assertEqual({"m5.xlarge": 2.0}, result.instance_seconds_by_type)
+
+    def test_shipped_log_prefix_without_slash_excludes_sibling_clusters(self):
+        s3 = FakeS3Client(
+            {
+                "yarn-rm-logs/j-TEST2/ip-10-0-0-1/"
+                "hadoop-yarn-resourcemanager-20260928-120000-a.log": b"other cluster",
+                "emr-logs/j-TEST/node/i-1/applications/"
+                "hadoop-yarn-resourcemanager-rm.log": self.yarn_log_with_instance_type().encode(),
+            }
+        )
+        request = replace(self.request(), yarn_log_uri="s3://test-bucket/yarn-rm-logs/j-TEST")
+
+        result = calculate_emr_application_usage(
+            request, emr_client=FakeEmrClient(), s3_client=s3
+        )
+
+        self.assertTrue(result.complete)
+        self.assertEqual({"m5.xlarge": 2.0}, result.instance_seconds_by_type)
+
+    def test_empty_shipped_log_object_falls_back_to_cluster_log_uri(self):
+        s3 = FakeS3Client(
+            {
+                "yarn-rm-logs/j-TEST/ip-10-0-0-1/"
+                "hadoop-yarn-resourcemanager-20260928-120000-a.log": b"",
+                "emr-logs/j-TEST/node/i-1/applications/"
+                "hadoop-yarn-resourcemanager-rm.log": self.yarn_log_with_instance_type().encode(),
+            }
+        )
+        request = replace(self.request(), yarn_log_uri="s3://test-bucket/yarn-rm-logs/j-TEST/")
+
+        result = calculate_emr_application_usage(
+            request, emr_client=FakeEmrClient(), s3_client=s3
+        )
+
+        self.assertTrue(result.complete)
+        self.assertEqual({"m5.xlarge": 2.0}, result.instance_seconds_by_type)
+
+    def test_request_rejects_non_s3_yarn_log_uri(self):
+        with self.assertRaisesRegex(ValueError, "Expected an S3 URI"):
+            replace(self.request(), yarn_log_uri="/var/log/hadoop-yarn")
+
     def test_missing_calculator_evidence_is_retryable(self):
         log = "\n".join(
             line

yarn-resource-cost/yarn_job_cost_api.py

@@ -28,11 +28,20 @@ class EmrApplicationUsageRequest:
     event_log_uri: str
     region: str
     include_application_master: bool = False
+    yarn_log_uri: str | None = None
+    """Optional S3 prefix holding this cluster's ResourceManager logs.
+
+    Use this when a log shipper uploads ResourceManager logs faster than the EMR
+    log pusher. When it contains no ResourceManager logs, the cluster LogUri
+    archive is used instead.
+    """
 
     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")
+        if self.yarn_log_uri is not None:
+            _split_s3_uri(self.yarn_log_uri)
 
 
 @dataclass(frozen=True)
@@ -169,24 +178,46 @@ def _cluster_log_uri(emr_client: Any, cluster_id: str) -> str:
     return normalized.rstrip("/") + f"/{cluster_id}/"
 
 
-def _materialize_yarn_logs(
-    emr_client: Any,
+def _download_yarn_logs(
     s3_client: Any,
-    cluster_id: str,
+    log_uri: str,
+    markers: tuple[str, ...],
     destination: Path,
 ) -> list[Path]:
-    bucket, prefix = _split_s3_uri(_cluster_log_uri(emr_client, cluster_id))
+    bucket, prefix = _split_s3_uri(log_uri)
+    if prefix and not prefix.endswith("/"):
+        # Directory boundary: "j-TEST" must not also match "j-TEST2/".
+        prefix += "/"
     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")
-        )
+        if item.get("Size") != 0  # a zero-byte placeholder is not usable evidence
+        and any(marker in Path(str(item.get("Key") or "")).name for marker in markers)
     ]
     return _download_objects(s3_client, bucket, prefix, objects, destination)
 
 
+def _materialize_yarn_logs(
+    emr_client: Any,
+    s3_client: Any,
+    cluster_id: str,
+    destination: Path,
+    yarn_log_uri: str | None = None,
+) -> list[Path]:
+    if yarn_log_uri:
+        downloaded = _download_yarn_logs(
+            s3_client, yarn_log_uri, ("hadoop-yarn-resourcemanager",), destination
+        )
+        if downloaded:
+            return downloaded
+    return _download_yarn_logs(
+        s3_client,
+        _cluster_log_uri(emr_client, cluster_id),
+        ("hadoop-yarn-resourcemanager", "hadoop-yarn-nodemanager"),
+        destination,
+    )
+
+
 def _empty_result(
     request: EmrApplicationUsageRequest,
     warning: str,
@@ -234,7 +265,13 @@ def calculate_emr_application_usage(
             )
 
         yarn_root = root / "yarn"
-        downloaded = _materialize_yarn_logs(emr_client, s3_client, request.cluster_id, yarn_root)
+        downloaded = _materialize_yarn_logs(
+            emr_client,
+            s3_client,
+            request.cluster_id,
+            yarn_root,
+            request.yarn_log_uri,
+        )
         if not downloaded:
             return _empty_result(
                 request,