From f9d927af95816fc9e3c952952891722809d598a1 Mon Sep 17 00:00:00 2001 From: Vamsi-klu Date: Sat, 14 Feb 2026 19:19:07 -0800 Subject: [PATCH] [dagster-aws] Add EMR step failure diagnostics with automatic log retrieval Co-Authored-By: Claude Opus 4.6 --- .../dagster-aws/dagster_aws/emr/__init__.py | 2 + .../dagster-aws/dagster_aws/emr/emr.py | 202 +++++++++++++++++- .../dagster_aws/pipes/clients/emr.py | 50 ++++- .../dagster_aws_tests/emr_tests/test_emr.py | 149 ++++++++++++- 4 files changed, 391 insertions(+), 12 deletions(-) diff --git a/python_modules/libraries/dagster-aws/dagster_aws/emr/__init__.py b/python_modules/libraries/dagster-aws/dagster_aws/emr/__init__.py index 626ffb80d497d..1ba1114809323 100644 --- a/python_modules/libraries/dagster-aws/dagster_aws/emr/__init__.py +++ b/python_modules/libraries/dagster-aws/dagster_aws/emr/__init__.py @@ -1,6 +1,8 @@ from dagster_aws.emr.emr import ( EmrError as EmrError, EmrJobRunner as EmrJobRunner, + FailureLogConfig as FailureLogConfig, + StepFailureDiagnostics as StepFailureDiagnostics, ) from dagster_aws.emr.pyspark_step_launcher import ( emr_pyspark_step_launcher as emr_pyspark_step_launcher, diff --git a/python_modules/libraries/dagster-aws/dagster_aws/emr/emr.py b/python_modules/libraries/dagster-aws/dagster_aws/emr/emr.py index 19c01952a71e1..7b2140244c711 100644 --- a/python_modules/libraries/dagster-aws/dagster_aws/emr/emr.py +++ b/python_modules/libraries/dagster-aws/dagster_aws/emr/emr.py @@ -23,6 +23,7 @@ import gzip import re from io import BytesIO +from typing import NamedTuple, Optional from urllib.parse import urlparse import boto3 @@ -33,6 +34,40 @@ from dagster_aws.emr.types import EMR_CLUSTER_TERMINATED_STATES, EmrClusterState, EmrStepState from dagster_aws.utils.mrjob.utils import _boto3_now, _wrap_aws_client, strip_microseconds +# Default configuration for failure log retrieval +DEFAULT_STDERR_LINES = 50 +DEFAULT_STDOUT_LINES = 20 +DEFAULT_LOG_WAIT_TIMEOUT = 600 # 10 minutes + + +class FailureLogConfig(NamedTuple): + """Configuration for log retrieval on EMR step failure. + + Args: + retrieve_logs: Whether to retrieve logs from S3 on failure. Defaults to True. + stderr_lines: Number of stderr lines to include in error message. Defaults to 50. + stdout_lines: Number of stdout lines to include in error message. Defaults to 20. + wait_timeout: Seconds to wait for logs to appear in S3. Defaults to 600. + """ + + retrieve_logs: bool = True + stderr_lines: int = DEFAULT_STDERR_LINES + stdout_lines: int = DEFAULT_STDOUT_LINES + wait_timeout: int = DEFAULT_LOG_WAIT_TIMEOUT + + +class StepFailureDiagnostics(NamedTuple): + """Diagnostics information for a failed EMR step.""" + + step_id: str + cluster_id: str + state_change_reason: str + stderr_excerpt: Optional[str] = None + stdout_excerpt: Optional[str] = None + logs_s3_uri: Optional[str] = None + error_retrieving_logs: Optional[str] = None + + # if we can't create or find our own service role, use the one # created by the AWS console and CLI _FALLBACK_SERVICE_ROLE = "EMR_DefaultRole" @@ -43,7 +78,46 @@ class EmrError(Exception): - pass + """Exception raised when an EMR operation fails. + + Attributes: + diagnostics: Optional diagnostics information for step failures. + """ + + def __init__( + self, + message: str = "EMR operation failed", + diagnostics: Optional[StepFailureDiagnostics] = None, + ) -> None: + self.diagnostics = diagnostics + if diagnostics: + full_message = self._build_diagnostic_message(message, diagnostics) + else: + full_message = message + super().__init__(full_message) + + @staticmethod + def _build_diagnostic_message(base_message: str, diagnostics: StepFailureDiagnostics) -> str: + parts = [base_message] + parts.append(f"\n\nCluster ID: {diagnostics.cluster_id}") + parts.append(f"Step ID: {diagnostics.step_id}") + + if diagnostics.state_change_reason: + parts.append(f"Reason: {diagnostics.state_change_reason}") + + if diagnostics.logs_s3_uri: + parts.append(f"Full logs: {diagnostics.logs_s3_uri}") + + if diagnostics.error_retrieving_logs: + parts.append(f"Note: {diagnostics.error_retrieving_logs}") + + if diagnostics.stderr_excerpt: + parts.append(f"\n--- stderr (last lines) ---\n{diagnostics.stderr_excerpt}") + + if diagnostics.stdout_excerpt: + parts.append(f"\n--- stdout (last lines) ---\n{diagnostics.stdout_excerpt}") + + return "\n".join(parts) class EmrJobRunner: @@ -53,6 +127,7 @@ def __init__( check_cluster_every=30, aws_access_key_id=None, aws_secret_access_key=None, + failure_log_config: Optional[FailureLogConfig] = None, ): """This object encapsulates various utilities for interacting with EMR clusters and invoking steps (jobs) on them. @@ -68,6 +143,8 @@ def __init__( use the default boto3 credentials chain. aws_secret_access_key ([type], optional): AWS secret access key. Defaults to None, which will use the default boto3 credentials chain. + failure_log_config (FailureLogConfig, optional): Configuration for log retrieval on + step failure. Defaults to FailureLogConfig() which enables log retrieval. """ self.region = check.str_param(region, "region") @@ -77,6 +154,7 @@ def __init__( self.aws_secret_access_key = check.opt_str_param( aws_secret_access_key, "aws_secret_access_key" ) + self.failure_log_config = failure_log_config or FailureLogConfig() def make_emr_client(self): """Creates a boto3 EMR client. Construction is wrapped in retries in case client connection @@ -302,13 +380,103 @@ def is_emr_step_complete(self, log, cluster_id, emr_step_id): # was it caused by IAM roles? self._check_for_missing_default_iam_roles(log, cluster) - # TODO: extract logs here to surface failure reason - # See: https://github.com/dagster-io/dagster/issues/1954 + # Retrieve failure diagnostics including logs + diagnostics = self._retrieve_step_failure_diagnostics( + log, cluster_id, emr_step_id, _get_reason(step) + ) if step_state == EmrStepState.Failed: - log.error(f"EMR step {emr_step_id} failed") + log.error( + f"EMR step {emr_step_id} failed", + extra={ + "emr_step_id": diagnostics.step_id, + "emr_cluster_id": diagnostics.cluster_id, + "state_change_reason": diagnostics.state_change_reason, + "stderr_excerpt": diagnostics.stderr_excerpt, + "stdout_excerpt": diagnostics.stdout_excerpt, + "logs_s3_uri": diagnostics.logs_s3_uri, + }, + ) + + raise EmrError(f"EMR step {emr_step_id} failed", diagnostics=diagnostics) - raise EmrError(f"EMR step {emr_step_id} failed") + def _retrieve_step_failure_diagnostics( + self, log, cluster_id: str, step_id: str, state_change_reason: str + ) -> StepFailureDiagnostics: + """Retrieve diagnostics information for a failed EMR step. + + Args: + log: DagsterLogManager for logging + cluster_id: EMR cluster ID + step_id: EMR step ID + state_change_reason: The reason for the step state change + + Returns: + StepFailureDiagnostics with available information + """ + stderr_excerpt: Optional[str] = None + stdout_excerpt: Optional[str] = None + logs_s3_uri: Optional[str] = None + error_retrieving_logs: Optional[str] = None + + if not self.failure_log_config.retrieve_logs: + return StepFailureDiagnostics( + step_id=step_id, + cluster_id=cluster_id, + state_change_reason=state_change_reason, + error_retrieving_logs="Log retrieval disabled in configuration", + ) + + try: + log_bucket, log_key_prefix = self.log_location_for_cluster(cluster_id) + logs_s3_uri = f"s3://{log_bucket}/{log_key_prefix}{cluster_id}/steps/{step_id}/" + + # Calculate waiter settings from timeout + waiter_delay = 30 + waiter_max_attempts = max(1, self.failure_log_config.wait_timeout // waiter_delay) + + stdout_log, stderr_log = self.retrieve_logs_for_step_id( + log, + cluster_id, + step_id, + waiter_delay=waiter_delay, + waiter_max_attempts=waiter_max_attempts, + ) + + # Extract last N lines as excerpts + if stderr_log: + stderr_lines = stderr_log.strip().split("\n") + stderr_excerpt = "\n".join(stderr_lines[-self.failure_log_config.stderr_lines :]) + if stdout_log: + stdout_lines = stdout_log.strip().split("\n") + stdout_excerpt = "\n".join(stdout_lines[-self.failure_log_config.stdout_lines :]) + + log.info( + f"Retrieved failure logs for EMR step {step_id}", + extra={"logs_s3_uri": logs_s3_uri}, + ) + + except EmrError as e: + if "Log URI not specified" in str(e): + error_retrieving_logs = "S3 logging not enabled for this cluster" + else: + error_retrieving_logs = f"Could not retrieve logs: {e}" + log.warning(error_retrieving_logs) + except Exception as e: + error_retrieving_logs = f"Error retrieving logs: {e}" + log.warning( + f"Unable to retrieve EMR step logs for {step_id} on cluster {cluster_id}: {e}" + ) + + return StepFailureDiagnostics( + step_id=step_id, + cluster_id=cluster_id, + state_change_reason=state_change_reason, + stderr_excerpt=stderr_excerpt, + stdout_excerpt=stdout_excerpt, + logs_s3_uri=logs_s3_uri, + error_retrieving_logs=error_retrieving_logs, + ) def _check_for_missing_default_iam_roles(self, log, cluster): """If cluster couldn't start due to missing IAM roles, tell user what to do.""" @@ -353,25 +521,43 @@ def log_location_for_cluster(self, cluster_id): log_key_prefix = log_uri_parsed.path.lstrip("/") return log_bucket, log_key_prefix - def retrieve_logs_for_step_id(self, log, cluster_id, step_id): + def retrieve_logs_for_step_id( + self, log, cluster_id, step_id, waiter_delay=30, waiter_max_attempts=20 + ): """Retrieves stdout and stderr logs for the given step ID. Args: log (DagsterLogManager): Log manager, for logging cluster_id (str): EMR cluster ID step_id (str): EMR step ID for the job that was submitted. + waiter_delay (int): How long to wait between attempts to check S3 for the log file + waiter_max_attempts (int): Number of attempts before giving up on waiting Returns: (str, str): Tuple of stdout log string contents, and stderr log string contents """ check.str_param(cluster_id, "cluster_id") check.str_param(step_id, "step_id") + check.int_param(waiter_delay, "waiter_delay") + check.int_param(waiter_max_attempts, "waiter_max_attempts") log_bucket, log_key_prefix = self.log_location_for_cluster(cluster_id) prefix = f"{log_key_prefix}{cluster_id}/steps/{step_id}" - stdout_log = self.wait_for_log(log, log_bucket, f"{prefix}/stdout.gz") - stderr_log = self.wait_for_log(log, log_bucket, f"{prefix}/stderr.gz") + stdout_log = self.wait_for_log( + log, + log_bucket, + f"{prefix}/stdout.gz", + waiter_delay=waiter_delay, + waiter_max_attempts=waiter_max_attempts, + ) + stderr_log = self.wait_for_log( + log, + log_bucket, + f"{prefix}/stderr.gz", + waiter_delay=waiter_delay, + waiter_max_attempts=waiter_max_attempts, + ) return stdout_log, stderr_log def wait_for_log(self, log, log_bucket, log_key, waiter_delay=30, waiter_max_attempts=20): diff --git a/python_modules/libraries/dagster-aws/dagster_aws/pipes/clients/emr.py b/python_modules/libraries/dagster-aws/dagster_aws/pipes/clients/emr.py index f852c1b35b7b3..82e3b007cdcb1 100644 --- a/python_modules/libraries/dagster-aws/dagster_aws/pipes/clients/emr.py +++ b/python_modules/libraries/dagster-aws/dagster_aws/pipes/clients/emr.py @@ -182,8 +182,13 @@ def _wait_for_completion( context.log.info(f"[pipes] EMR cluster {cluster_id} completed with state: {state}") if state in EMR_CLUSTER_TERMINATED_STATES: - context.log.error(f"[pipes] EMR job {cluster_id} failed") - raise Exception(f"[pipes] EMR job {cluster_id} failed:\n{cluster}") + # Get step-level failure information + failure_details = self._get_step_failure_details(context, cluster_id) + context.log.error( + f"[pipes] EMR job {cluster_id} failed", + extra={"failure_details": failure_details}, + ) + raise Exception(f"[pipes] EMR job {cluster_id} failed:\n{cluster}\n\n{failure_details}") return cluster @@ -334,3 +339,44 @@ def _terminate( cluster_id = start_response["JobFlowId"] context.log.info(f"[pipes] Terminating EMR job {cluster_id}") self._client.terminate_job_flows(JobFlowIds=[cluster_id]) + + def _get_step_failure_details( + self, + context: Union[OpExecutionContext, AssetExecutionContext], + cluster_id: str, + ) -> str: + """Retrieve failure details for all failed steps in the cluster. + + Args: + context: The execution context for logging + cluster_id: The EMR cluster ID + + Returns: + A formatted string with step failure details + """ + try: + steps_response = self._client.list_steps(ClusterId=cluster_id) + steps = steps_response.get("Steps", []) + + failed_steps = [] + for step in steps: + step_state = step.get("Status", {}).get("State", "") + if step_state in ("FAILED", "CANCELLED", "INTERRUPTED"): + step_id = step.get("Id", "unknown") + step_name = step.get("Name", "unknown") + reason = ( + step.get("Status", {}) + .get("StateChangeReason", {}) + .get("Message", "No reason provided") + ) + failed_steps.append( + f" - Step '{step_name}' ({step_id}): {step_state}\n Reason: {reason}" + ) + + if failed_steps: + return "Failed steps:\n" + "\n".join(failed_steps) + return "No step-level failure information available" + + except Exception as e: + context.log.warning(f"[pipes] Could not retrieve step failure details: {e}") + return f"Could not retrieve step details: {e}" diff --git a/python_modules/libraries/dagster-aws/dagster_aws_tests/emr_tests/test_emr.py b/python_modules/libraries/dagster-aws/dagster_aws_tests/emr_tests/test_emr.py index 650ddb8fe92b3..146a7274a594a 100644 --- a/python_modules/libraries/dagster-aws/dagster_aws_tests/emr_tests/test_emr.py +++ b/python_modules/libraries/dagster-aws/dagster_aws_tests/emr_tests/test_emr.py @@ -195,11 +195,156 @@ def get_step_dict(step_id, step_state): get_step_dict(emr_step_id, "COMPLETED"), get_step_dict(emr_step_id, "FAILED"), ] - with mock.patch.object(EmrJobRunner, "describe_step", side_effect=describe_step_returns): + with ( + mock.patch.object(EmrJobRunner, "describe_step", side_effect=describe_step_returns), + mock.patch.object( + EmrJobRunner, "retrieve_logs_for_step_id", return_value=("stdout", "stderr") + ) as retrieve_logs, + ): assert not emr.is_emr_step_complete(context.log, cluster_id, emr_step_id) assert not emr.is_emr_step_complete(context.log, cluster_id, emr_step_id) assert emr.is_emr_step_complete(context.log, cluster_id, emr_step_id) with pytest.raises(EmrError) as exc_info: emr.is_emr_step_complete(context.log, cluster_id, emr_step_id) - assert "step failed" in str(exc_info.value) + assert "failed" in str(exc_info.value) + # Check diagnostics are included in error + assert exc_info.value.diagnostics is not None + assert exc_info.value.diagnostics.step_id == emr_step_id + assert exc_info.value.diagnostics.cluster_id == cluster_id + assert exc_info.value.diagnostics.stderr_excerpt is not None + assert exc_info.value.diagnostics.stdout_excerpt is not None + retrieve_logs.assert_called_once() + + +@mock_emr +def test_emr_step_failure_diagnostics_with_config(emr_cluster_config): + """Test that failure diagnostics use FailureLogConfig settings.""" + from dagster_aws.emr import FailureLogConfig + + context = create_test_pipeline_execution_context() + config = FailureLogConfig( + retrieve_logs=True, + stderr_lines=10, + stdout_lines=5, + wait_timeout=60, + ) + emr = EmrJobRunner(region=REGION, check_cluster_every=1, failure_log_config=config) + + cluster_id = emr.run_job_flow(context.log, emr_cluster_config) + step_ids = emr.add_job_flow_steps( + context.log, + cluster_id, + [emr.construct_step_dict_for_command("test", ["ls"])], + ) + emr_step_id = step_ids[0] + + # Create multi-line logs to test excerpt truncation + long_stderr = "\n".join([f"stderr line {i}" for i in range(20)]) + long_stdout = "\n".join([f"stdout line {i}" for i in range(20)]) + + failed_step = { + "Step": { + "Id": emr_step_id, + "Name": "test", + "Config": {"Jar": "command-runner.jar", "Properties": {}, "Args": ["ls"]}, + "ActionOnFailure": "CONTINUE", + "Status": { + "State": "FAILED", + "StateChangeReason": {"Message": "Step failed due to error"}, + "Timeline": {"StartDateTime": _boto3_now()}, + }, + }, + } + + with ( + mock.patch.object(EmrJobRunner, "describe_step", return_value=failed_step), + mock.patch.object( + EmrJobRunner, "retrieve_logs_for_step_id", return_value=(long_stdout, long_stderr) + ), + ): + with pytest.raises(EmrError) as exc_info: + emr.is_emr_step_complete(context.log, cluster_id, emr_step_id) + + diagnostics = exc_info.value.diagnostics + assert diagnostics is not None + # Check stderr was truncated to 10 lines + assert diagnostics.stderr_excerpt.count("\n") == 9 # 10 lines = 9 newlines + assert "stderr line 19" in diagnostics.stderr_excerpt + assert "stderr line 10" in diagnostics.stderr_excerpt + # Check stdout was truncated to 5 lines + assert diagnostics.stdout_excerpt.count("\n") == 4 # 5 lines = 4 newlines + assert "stdout line 19" in diagnostics.stdout_excerpt + assert "stdout line 15" in diagnostics.stdout_excerpt + + +@mock_emr +def test_emr_step_failure_logs_disabled(emr_cluster_config): + """Test that log retrieval can be disabled via config.""" + from dagster_aws.emr import FailureLogConfig + + context = create_test_pipeline_execution_context() + config = FailureLogConfig(retrieve_logs=False) + emr = EmrJobRunner(region=REGION, check_cluster_every=1, failure_log_config=config) + + cluster_id = emr.run_job_flow(context.log, emr_cluster_config) + step_ids = emr.add_job_flow_steps( + context.log, + cluster_id, + [emr.construct_step_dict_for_command("test", ["ls"])], + ) + emr_step_id = step_ids[0] + + failed_step = { + "Step": { + "Id": emr_step_id, + "Name": "test", + "Config": {"Jar": "command-runner.jar", "Properties": {}, "Args": ["ls"]}, + "ActionOnFailure": "CONTINUE", + "Status": { + "State": "FAILED", + "StateChangeReason": {"Message": "Test failure"}, + "Timeline": {"StartDateTime": _boto3_now()}, + }, + }, + } + + with ( + mock.patch.object(EmrJobRunner, "describe_step", return_value=failed_step), + mock.patch.object(EmrJobRunner, "retrieve_logs_for_step_id") as retrieve_logs_mock, + ): + with pytest.raises(EmrError) as exc_info: + emr.is_emr_step_complete(context.log, cluster_id, emr_step_id) + + # Logs should not be retrieved when disabled + retrieve_logs_mock.assert_not_called() + diagnostics = exc_info.value.diagnostics + assert diagnostics.stderr_excerpt is None + assert diagnostics.stdout_excerpt is None + assert "Log retrieval disabled" in diagnostics.error_retrieving_logs + + +def test_emr_error_message_formatting(): + """Test that EmrError formats diagnostics correctly.""" + from dagster_aws.emr import StepFailureDiagnostics + + diagnostics = StepFailureDiagnostics( + step_id="s-123", + cluster_id="j-456", + state_change_reason="Out of memory", + stderr_excerpt="Error: OOM killed", + stdout_excerpt="Processing...", + logs_s3_uri="s3://bucket/logs/j-456/steps/s-123/", + ) + + error = EmrError("EMR step s-123 failed", diagnostics=diagnostics) + + error_str = str(error) + assert "Cluster ID: j-456" in error_str + assert "Step ID: s-123" in error_str + assert "Reason: Out of memory" in error_str + assert "Full logs: s3://bucket/logs/j-456/steps/s-123/" in error_str + assert "--- stderr (last lines) ---" in error_str + assert "Error: OOM killed" in error_str + assert "--- stdout (last lines) ---" in error_str + assert "Processing..." in error_str