2323import gzip
2424import re
2525from io import BytesIO
26+ from typing import NamedTuple , Optional
2627from urllib .parse import urlparse
2728
2829import boto3
3334from dagster_aws .emr .types import EMR_CLUSTER_TERMINATED_STATES , EmrClusterState , EmrStepState
3435from dagster_aws .utils .mrjob .utils import _boto3_now , _wrap_aws_client , strip_microseconds
3536
37+ # Default configuration for failure log retrieval
38+ DEFAULT_STDERR_LINES = 50
39+ DEFAULT_STDOUT_LINES = 20
40+ DEFAULT_LOG_WAIT_TIMEOUT = 600 # 10 minutes
41+
42+
43+ class FailureLogConfig (NamedTuple ):
44+ """Configuration for log retrieval on EMR step failure.
45+
46+ Args:
47+ retrieve_logs: Whether to retrieve logs from S3 on failure. Defaults to True.
48+ stderr_lines: Number of stderr lines to include in error message. Defaults to 50.
49+ stdout_lines: Number of stdout lines to include in error message. Defaults to 20.
50+ wait_timeout: Seconds to wait for logs to appear in S3. Defaults to 600.
51+ """
52+
53+ retrieve_logs : bool = True
54+ stderr_lines : int = DEFAULT_STDERR_LINES
55+ stdout_lines : int = DEFAULT_STDOUT_LINES
56+ wait_timeout : int = DEFAULT_LOG_WAIT_TIMEOUT
57+
58+
59+ class StepFailureDiagnostics (NamedTuple ):
60+ """Diagnostics information for a failed EMR step."""
61+
62+ step_id : str
63+ cluster_id : str
64+ state_change_reason : str
65+ stderr_excerpt : Optional [str ] = None
66+ stdout_excerpt : Optional [str ] = None
67+ logs_s3_uri : Optional [str ] = None
68+ error_retrieving_logs : Optional [str ] = None
69+
70+
3671# if we can't create or find our own service role, use the one
3772# created by the AWS console and CLI
3873_FALLBACK_SERVICE_ROLE = "EMR_DefaultRole"
4378
4479
4580class EmrError (Exception ):
46- pass
81+ """Exception raised when an EMR operation fails.
82+
83+ Attributes:
84+ diagnostics: Optional diagnostics information for step failures.
85+ """
86+
87+ def __init__ (
88+ self ,
89+ message : str = "EMR operation failed" ,
90+ diagnostics : Optional [StepFailureDiagnostics ] = None ,
91+ ) -> None :
92+ self .diagnostics = diagnostics
93+ if diagnostics :
94+ full_message = self ._build_diagnostic_message (message , diagnostics )
95+ else :
96+ full_message = message
97+ super ().__init__ (full_message )
98+
99+ @staticmethod
100+ def _build_diagnostic_message (base_message : str , diagnostics : StepFailureDiagnostics ) -> str :
101+ parts = [base_message ]
102+ parts .append (f"\n \n Cluster ID: { diagnostics .cluster_id } " )
103+ parts .append (f"Step ID: { diagnostics .step_id } " )
104+
105+ if diagnostics .state_change_reason :
106+ parts .append (f"Reason: { diagnostics .state_change_reason } " )
107+
108+ if diagnostics .logs_s3_uri :
109+ parts .append (f"Full logs: { diagnostics .logs_s3_uri } " )
110+
111+ if diagnostics .error_retrieving_logs :
112+ parts .append (f"Note: { diagnostics .error_retrieving_logs } " )
113+
114+ if diagnostics .stderr_excerpt :
115+ parts .append (f"\n --- stderr (last lines) ---\n { diagnostics .stderr_excerpt } " )
116+
117+ if diagnostics .stdout_excerpt :
118+ parts .append (f"\n --- stdout (last lines) ---\n { diagnostics .stdout_excerpt } " )
119+
120+ return "\n " .join (parts )
47121
48122
49123class EmrJobRunner :
@@ -53,6 +127,7 @@ def __init__(
53127 check_cluster_every = 30 ,
54128 aws_access_key_id = None ,
55129 aws_secret_access_key = None ,
130+ failure_log_config : Optional [FailureLogConfig ] = None ,
56131 ):
57132 """This object encapsulates various utilities for interacting with EMR clusters and invoking
58133 steps (jobs) on them.
@@ -68,6 +143,8 @@ def __init__(
68143 use the default boto3 credentials chain.
69144 aws_secret_access_key ([type], optional): AWS secret access key. Defaults to None, which
70145 will use the default boto3 credentials chain.
146+ failure_log_config (FailureLogConfig, optional): Configuration for log retrieval on
147+ step failure. Defaults to FailureLogConfig() which enables log retrieval.
71148 """
72149 self .region = check .str_param (region , "region" )
73150
@@ -77,6 +154,7 @@ def __init__(
77154 self .aws_secret_access_key = check .opt_str_param (
78155 aws_secret_access_key , "aws_secret_access_key"
79156 )
157+ self .failure_log_config = failure_log_config or FailureLogConfig ()
80158
81159 def make_emr_client (self ):
82160 """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):
302380 # was it caused by IAM roles?
303381 self ._check_for_missing_default_iam_roles (log , cluster )
304382
305- # TODO: extract logs here to surface failure reason
306- # See: https://github.com/dagster-io/dagster/issues/1954
383+ # Retrieve failure diagnostics including logs
384+ diagnostics = self ._retrieve_step_failure_diagnostics (
385+ log , cluster_id , emr_step_id , _get_reason (step )
386+ )
307387
308388 if step_state == EmrStepState .Failed :
309- log .error (f"EMR step { emr_step_id } failed" )
389+ log .error (
390+ f"EMR step { emr_step_id } failed" ,
391+ extra = {
392+ "emr_step_id" : diagnostics .step_id ,
393+ "emr_cluster_id" : diagnostics .cluster_id ,
394+ "state_change_reason" : diagnostics .state_change_reason ,
395+ "stderr_excerpt" : diagnostics .stderr_excerpt ,
396+ "stdout_excerpt" : diagnostics .stdout_excerpt ,
397+ "logs_s3_uri" : diagnostics .logs_s3_uri ,
398+ },
399+ )
400+
401+ raise EmrError (f"EMR step { emr_step_id } failed" , diagnostics = diagnostics )
310402
311- raise EmrError (f"EMR step { emr_step_id } failed" )
403+ def _retrieve_step_failure_diagnostics (
404+ self , log , cluster_id : str , step_id : str , state_change_reason : str
405+ ) -> StepFailureDiagnostics :
406+ """Retrieve diagnostics information for a failed EMR step.
407+
408+ Args:
409+ log: DagsterLogManager for logging
410+ cluster_id: EMR cluster ID
411+ step_id: EMR step ID
412+ state_change_reason: The reason for the step state change
413+
414+ Returns:
415+ StepFailureDiagnostics with available information
416+ """
417+ stderr_excerpt : Optional [str ] = None
418+ stdout_excerpt : Optional [str ] = None
419+ logs_s3_uri : Optional [str ] = None
420+ error_retrieving_logs : Optional [str ] = None
421+
422+ if not self .failure_log_config .retrieve_logs :
423+ return StepFailureDiagnostics (
424+ step_id = step_id ,
425+ cluster_id = cluster_id ,
426+ state_change_reason = state_change_reason ,
427+ error_retrieving_logs = "Log retrieval disabled in configuration" ,
428+ )
429+
430+ try :
431+ log_bucket , log_key_prefix = self .log_location_for_cluster (cluster_id )
432+ logs_s3_uri = f"s3://{ log_bucket } /{ log_key_prefix } { cluster_id } /steps/{ step_id } /"
433+
434+ # Calculate waiter settings from timeout
435+ waiter_delay = 30
436+ waiter_max_attempts = max (1 , self .failure_log_config .wait_timeout // waiter_delay )
437+
438+ stdout_log , stderr_log = self .retrieve_logs_for_step_id (
439+ log ,
440+ cluster_id ,
441+ step_id ,
442+ waiter_delay = waiter_delay ,
443+ waiter_max_attempts = waiter_max_attempts ,
444+ )
445+
446+ # Extract last N lines as excerpts
447+ if stderr_log :
448+ stderr_lines = stderr_log .strip ().split ("\n " )
449+ stderr_excerpt = "\n " .join (stderr_lines [- self .failure_log_config .stderr_lines :])
450+ if stdout_log :
451+ stdout_lines = stdout_log .strip ().split ("\n " )
452+ stdout_excerpt = "\n " .join (stdout_lines [- self .failure_log_config .stdout_lines :])
453+
454+ log .info (
455+ f"Retrieved failure logs for EMR step { step_id } " ,
456+ extra = {"logs_s3_uri" : logs_s3_uri },
457+ )
458+
459+ except EmrError as e :
460+ if "Log URI not specified" in str (e ):
461+ error_retrieving_logs = "S3 logging not enabled for this cluster"
462+ else :
463+ error_retrieving_logs = f"Could not retrieve logs: { e } "
464+ log .warning (error_retrieving_logs )
465+ except Exception as e :
466+ error_retrieving_logs = f"Error retrieving logs: { e } "
467+ log .warning (
468+ f"Unable to retrieve EMR step logs for { step_id } on cluster { cluster_id } : { e } "
469+ )
470+
471+ return StepFailureDiagnostics (
472+ step_id = step_id ,
473+ cluster_id = cluster_id ,
474+ state_change_reason = state_change_reason ,
475+ stderr_excerpt = stderr_excerpt ,
476+ stdout_excerpt = stdout_excerpt ,
477+ logs_s3_uri = logs_s3_uri ,
478+ error_retrieving_logs = error_retrieving_logs ,
479+ )
312480
313481 def _check_for_missing_default_iam_roles (self , log , cluster ):
314482 """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):
353521 log_key_prefix = log_uri_parsed .path .lstrip ("/" )
354522 return log_bucket , log_key_prefix
355523
356- def retrieve_logs_for_step_id (self , log , cluster_id , step_id ):
524+ def retrieve_logs_for_step_id (
525+ self , log , cluster_id , step_id , waiter_delay = 30 , waiter_max_attempts = 20
526+ ):
357527 """Retrieves stdout and stderr logs for the given step ID.
358528
359529 Args:
360530 log (DagsterLogManager): Log manager, for logging
361531 cluster_id (str): EMR cluster ID
362532 step_id (str): EMR step ID for the job that was submitted.
533+ waiter_delay (int): How long to wait between attempts to check S3 for the log file
534+ waiter_max_attempts (int): Number of attempts before giving up on waiting
363535
364536 Returns:
365537 (str, str): Tuple of stdout log string contents, and stderr log string contents
366538 """
367539 check .str_param (cluster_id , "cluster_id" )
368540 check .str_param (step_id , "step_id" )
541+ check .int_param (waiter_delay , "waiter_delay" )
542+ check .int_param (waiter_max_attempts , "waiter_max_attempts" )
369543
370544 log_bucket , log_key_prefix = self .log_location_for_cluster (cluster_id )
371545
372546 prefix = f"{ log_key_prefix } { cluster_id } /steps/{ step_id } "
373- stdout_log = self .wait_for_log (log , log_bucket , f"{ prefix } /stdout.gz" )
374- stderr_log = self .wait_for_log (log , log_bucket , f"{ prefix } /stderr.gz" )
547+ stdout_log = self .wait_for_log (
548+ log ,
549+ log_bucket ,
550+ f"{ prefix } /stdout.gz" ,
551+ waiter_delay = waiter_delay ,
552+ waiter_max_attempts = waiter_max_attempts ,
553+ )
554+ stderr_log = self .wait_for_log (
555+ log ,
556+ log_bucket ,
557+ f"{ prefix } /stderr.gz" ,
558+ waiter_delay = waiter_delay ,
559+ waiter_max_attempts = waiter_max_attempts ,
560+ )
375561 return stdout_log , stderr_log
376562
377563 def wait_for_log (self , log , log_bucket , log_key , waiter_delay = 30 , waiter_max_attempts = 20 ):
0 commit comments