fix/emr-serverless-configurable-spark-ui-dashboard-refresh-interval

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,18 +184,27 @@ 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"],
             jobRunId=job_run_id,
         ):
             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}"
                 )

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
+
+
+