diff --git a/python_modules/libraries/dagster-aws/dagster_aws/pipes/clients/emr_serverless.py b/python_modules/libraries/dagster-aws/dagster_aws/pipes/clients/emr_serverless.py index a90c944336fc8..f4339cbd4be0d 100644 --- a/python_modules/libraries/dagster-aws/dagster_aws/pipes/clients/emr_serverless.py +++ b/python_modules/libraries/dagster-aws/dagster_aws/pipes/clients/emr_serverless.py @@ -58,6 +58,7 @@ def __init__( message_reader: PipesMessageReader | None = None, forward_termination: bool = True, poll_interval: float = 5.0, + dashboard_refresh_interval: Optional[float] = 3300.0, ): self._client: EMRServerlessClient = cast( "EMRServerlessClient", client or boto3.client("emr-serverless") @@ -66,6 +67,7 @@ def __init__( self._message_reader = message_reader or PipesCloudWatchMessageReader() self.forward_termination = check.bool_param(forward_termination, "forward_termination") self.poll_interval = poll_interval + self.dashboard_refresh_interval = check.opt_numeric_param(dashboard_refresh_interval, "dashboard_refresh_interval") @property def client(self) -> "EMRServerlessClient": @@ -182,6 +184,7 @@ def _wait_for_completion( running_dashboard_url = None completed_dashboard_url = None + last_dashboard_fetch_time = None while response := self.client.get_job_run( applicationId=start_response["applicationId"], @@ -189,11 +192,19 @@ def _wait_for_completion( ): state: JobRunStateType = response["jobRun"]["state"] - # get dashboard url when it's ready (but only once) - if state == "RUNNING" and running_dashboard_url is None: + # get dashboard url when it's ready, and refresh it according to dashboard_refresh_interval (default 55 minutes) + if state == "RUNNING" and ( + running_dashboard_url is None + or ( + self.dashboard_refresh_interval is not None + and last_dashboard_fetch_time is not None + and time.time() - last_dashboard_fetch_time >= self.dashboard_refresh_interval + ) + ): running_dashboard_url = self.client.get_dashboard_for_job_run( applicationId=application_id, jobRunId=job_run_id ) + last_dashboard_fetch_time = time.time() context.log.info( f"[pipes] {self.AWS_SERVICE_NAME} job is running. Dashboard URL: {running_dashboard_url}" ) diff --git a/python_modules/libraries/dagster-aws/dagster_aws_tests/pipes_tests/test_emr_serverless.py b/python_modules/libraries/dagster-aws/dagster_aws_tests/pipes_tests/test_emr_serverless.py index 56d1ef8595607..44041daa0e955 100644 --- a/python_modules/libraries/dagster-aws/dagster_aws_tests/pipes_tests/test_emr_serverless.py +++ b/python_modules/libraries/dagster-aws/dagster_aws_tests/pipes_tests/test_emr_serverless.py @@ -81,3 +81,100 @@ def my_asset(context: AssetExecutionContext, emr_serverless_client: PipesEMRServ }, instance=instance, ) + + +def test_emr_serverless_dashboard_refresh(): + from unittest.mock import MagicMock, patch + from dagster_aws.pipes import PipesEMRServerlessClient + + client = MagicMock() + client.start_job_run.return_value = { + "jobRunId": "test-job-run-id", + "applicationId": "test-app-id", + } + + current_time = [0.0] + + def mock_time(): + return current_time[0] + + def mock_sleep(seconds): + current_time[0] += seconds + + mock_time_module = MagicMock() + mock_time_module.time.side_effect = mock_time + mock_time_module.sleep.side_effect = mock_sleep + + client.get_job_run.side_effect = [ + {"jobRun": {"state": "RUNNING", "applicationId": "test-app-id", "jobRunId": "test-job-run-id"}}, + {"jobRun": {"state": "RUNNING", "applicationId": "test-app-id", "jobRunId": "test-job-run-id"}}, + {"jobRun": {"state": "RUNNING", "applicationId": "test-app-id", "jobRunId": "test-job-run-id"}}, + {"jobRun": {"state": "SUCCESS", "applicationId": "test-app-id", "jobRunId": "test-job-run-id"}}, + ] + + client.get_dashboard_for_job_run.return_value = "http://mock-spark-ui-url" + + @asset + def my_asset(context: AssetExecutionContext, emr_serverless_client: PipesEMRServerlessClient): + return emr_serverless_client.run( + context=context, + start_job_run_params={ + "applicationId": "test-app-id", + "executionRoleArn": "arn:aws:iam::123456789012:role/EMRServerlessRole", + "jobDriver": { + "sparkSubmit": { + "entryPoint": "s3://my-bucket/my-script.py", + } + }, + } + ).get_results() + + with patch.object(PipesEMRServerlessClient, "_read_messages"), \ + patch.object(PipesEMRServerlessClient, "_extract_dagster_metadata", return_value={}), \ + patch("dagster_aws.pipes.clients.emr_serverless.time", mock_time_module), \ + instance_for_test() as instance: + + materialize( + [my_asset], + resources={ + "emr_serverless_client": PipesEMRServerlessClient( + client=client, + poll_interval=600.0, + ) + }, + instance=instance, + ) + + assert client.get_dashboard_for_job_run.call_count == 1 + + # Reset mock and time for second run with custom interval + client.get_dashboard_for_job_run.reset_mock() + current_time[0] = 0.0 + client.get_job_run.side_effect = [ + {"jobRun": {"state": "RUNNING", "applicationId": "test-app-id", "jobRunId": "test-job-run-id"}}, + {"jobRun": {"state": "RUNNING", "applicationId": "test-app-id", "jobRunId": "test-job-run-id"}}, + {"jobRun": {"state": "RUNNING", "applicationId": "test-app-id", "jobRunId": "test-job-run-id"}}, + {"jobRun": {"state": "SUCCESS", "applicationId": "test-app-id", "jobRunId": "test-job-run-id"}}, + ] + + with patch.object(PipesEMRServerlessClient, "_read_messages"), \ + patch.object(PipesEMRServerlessClient, "_extract_dagster_metadata", return_value={}), \ + patch("dagster_aws.pipes.clients.emr_serverless.time", mock_time_module), \ + instance_for_test() as instance: + + materialize( + [my_asset], + resources={ + "emr_serverless_client": PipesEMRServerlessClient( + client=client, + poll_interval=600.0, + dashboard_refresh_interval=900.0, + ) + }, + instance=instance, + ) + + assert client.get_dashboard_for_job_run.call_count == 2 + + +