feat(athena): track Spark Connect DPU usage against an account-wide budget

Author: dtaniwaki

dbt-athena/.changes/unreleased/Features-20260530-160000.yaml

@@ -0,0 +1,6 @@
+kind: Features
+body: Track per-session DPU usage against an account-wide budget so the Spark Connect pool throttles before AWS rejects with `Maximum allowed sessions`, and back off on transient `required capacity not being available` errors.
+time: 2026-05-30T16:00:00.000000+09:00
+custom:
+    Author: dtaniwaki
+    Issue: "1854"

dbt-athena/src/dbt/adapters/athena/connections.py

@@ -80,6 +80,7 @@ class AthenaCredentials(Credentials):
     spark_work_group: Optional[str] = None
     spark_connect_max_sessions: Optional[int] = None
     spark_connect_session_concurrency: Optional[int] = None
+    spark_connect_dpu_budget: Optional[int] = None
     s3_tmp_table_dir: Optional[str] = None
     # Unfortunately we can not just use dict, must be Dict because we'll get the following error:
     # Credentials in profile "athena", target "athena" invalid: Unable to create schema for 'dict'
@@ -91,7 +92,11 @@ def __post_init__(self) -> None:
         # time rather than waiting until a python model runs — the misconfig
         # would otherwise only manifest once a Spark 3.5 model is submitted,
         # potentially minutes into a long dbt run.
-        for field_name in ("spark_connect_max_sessions", "spark_connect_session_concurrency"):
+        for field_name in (
+            "spark_connect_max_sessions",
+            "spark_connect_session_concurrency",
+            "spark_connect_dpu_budget",
+        ):
             raw = getattr(self, field_name)
             if raw is None:
                 continue
@@ -135,6 +140,7 @@ def _connection_keys(self) -> Tuple[str, ...]:
             "spark_work_group",
             "spark_connect_max_sessions",
             "spark_connect_session_concurrency",
+            "spark_connect_dpu_budget",
         )
 
 

dbt-athena/src/dbt/adapters/athena/constants.py

@@ -3,6 +3,9 @@
 DEFAULT_THREAD_COUNT = 4
 DEFAULT_SPARK_CONNECT_MAX_SESSIONS = 4
 DEFAULT_SPARK_CONNECT_SESSION_CONCURRENCY = 1
+# AWS account-level quota for concurrent Athena Spark DPUs
+# (service-quotas code L-5A8D5237, default 160, not adjustable without AWS Support).
+DEFAULT_SPARK_CONNECT_DPU_BUDGET = 160
 DEFAULT_RETRY_ATTEMPTS = 3
 DEFAULT_POLLING_INTERVAL = 5
 DEFAULT_SPARK_COORDINATOR_DPU_SIZE = 1

dbt-athena/src/dbt/adapters/athena/spark_connect/job.py

@@ -31,6 +31,7 @@
 from dbt.adapters.athena.config import AthenaSparkSessionConfig
 from dbt.adapters.athena.connections import AthenaCredentials
 from dbt.adapters.athena.constants import (
+    DEFAULT_SPARK_CONNECT_DPU_BUDGET,
     DEFAULT_SPARK_CONNECT_MAX_SESSIONS,
     DEFAULT_SPARK_CONNECT_SESSION_CONCURRENCY,
     LOGGER,
@@ -59,6 +60,26 @@ class _EndpointNotReady(Exception):
     """Internal sentinel: GetSessionEndpoint should be polled again."""
 
 
+def _spark_max_executors(engine_config: EngineConfigurationTypeDef) -> Optional[int]:
+    """Return ``spark.dynamicAllocation.maxExecutors`` from Spark Connect engine_config.
+
+    Spark Connect uses Classifications (not SparkProperties) to carry
+    spark-defaults; the value is a string and must be parsed.
+    """
+    classifications = engine_config.get("Classifications") or []
+    for entry in classifications:
+        if entry.get("Name") != "spark-defaults":
+            continue
+        raw = (entry.get("Properties") or {}).get("spark.dynamicAllocation.maxExecutors")
+        if raw is None:
+            return None
+        try:
+            return int(raw)
+        except (TypeError, ValueError):
+            return None
+    return None
+
+
 class _AttemptResult(NamedTuple):
     result: Optional[SparkConnectResult]
     error: Optional[BaseException]
@@ -126,6 +147,24 @@ def _session_concurrency(self) -> int:
             or DEFAULT_SPARK_CONNECT_SESSION_CONCURRENCY
         )
 
+    @cached_property
+    def _dpu_budget(self) -> int:
+        return self.credentials.spark_connect_dpu_budget or DEFAULT_SPARK_CONNECT_DPU_BUDGET
+
+    @cached_property
+    def _dpu_request(self) -> int:
+        """DPUs reserved against the budget when starting a session.
+
+        ``MaxConcurrentDpus`` is the AWS-side hard cap; with dynamic
+        allocation, Spark scales up to ``maxExecutors + 1`` (executors +
+        driver). The true peak is the smaller of the two.
+        """
+        max_concurrent = int(self.engine_config["MaxConcurrentDpus"])
+        max_executors = _spark_max_executors(self.engine_config)
+        if max_executors is None:
+            return max_concurrent
+        return min(max_concurrent, max_executors + 1)
+
     @cached_property
     def _session_description(self) -> str:
         return f"dbt: {get_invocation_id()} - {self._session_fingerprint}"
@@ -199,6 +238,8 @@ def _acquire_session(self) -> str:
             timeout=self.timeout,
             polling_interval=self.polling_interval,
             session_concurrency=self._session_concurrency,
+            dpu_request=self._dpu_request,
+            dpu_budget=self._dpu_budget,
         )
 
     def _wait_for_endpoint(

dbt-athena/src/dbt/adapters/athena/spark_connect/session.py

@@ -19,12 +19,21 @@ class _SessionInfo(TypedDict):
     key: SessionKey
     client: AthenaClient
     load: int
+    dpu: int
 
 
 class _GlobalSessionLimitReached(Exception):
     """Raised when Athena returns ``Maximum allowed sessions reached``."""
 
 
+class _AccountCapacityUnavailable(Exception):
+    """Raised when Athena returns ``required capacity not being available``.
+
+    Distinct from the account session-limit signal: this is a transient
+    region-level capacity shortage that the DPU budget cannot predict.
+    """
+
+
 class SparkConnectSessionPool:
     """Singleton pool of Athena Spark Connect sessions.
 
@@ -65,36 +74,64 @@ def acquire(
         timeout: float,
         polling_interval: float,
         session_concurrency: int,
+        dpu_request: int,
+        dpu_budget: int,
     ) -> str:
         """Acquire a session for ``key``, reusing or starting one.
 
         Reuses when load < ``session_concurrency``; starts new when
-        per-key count < ``max_sessions``; waits up to ``timeout``.
+        per-key count < ``max_sessions`` AND
+        ``used_dpu + dpu_request <= dpu_budget``; waits up to ``timeout``.
+
+        Raises immediately when ``dpu_request > dpu_budget`` — no future
+        release could ever satisfy the request, so waiting would deadlock.
         """
+        if dpu_request > dpu_budget:
+            raise DbtRuntimeError(
+                f"Spark Connect session for key {key} requests {dpu_request} DPUs but "
+                f"spark_connect_dpu_budget is {dpu_budget}; the session can never start. "
+                f"Raise spark_connect_dpu_budget, or lower MaxConcurrentDpus / "
+                f"spark.dynamicAllocation.maxExecutors for the model."
+            )
+        if dpu_request == dpu_budget:
+            LOGGER.warning(
+                f"Spark Connect session for key {key} consumes the full DPU budget "
+                f"({dpu_request}/{dpu_budget}); other sessions will block until it releases."
+            )
+
         invocation_id = key[0]
         deadline = time.monotonic() + timeout
         time_since_eviction = self._EVICTION_INTERVAL  # evict on first pass
 
         while True:
             new_session_id: Optional[str] = None
             global_limit_hit = False
+            capacity_unavailable = False
             start_error: Optional[BaseException] = None
+            budget_used = 0
+            budget_ok = False
             with self._lock:
                 stale_entries = self._collect_stale_invocations(invocation_id)
                 reuse_candidate = self._attach(key, session_concurrency)
-                if reuse_candidate is None and self._has_room(key, max_sessions):
-                    try:
-                        new_session_id = self._start(
-                            key,
-                            athena_client,
-                            spark_work_group,
-                            engine_config,
-                            session_description,
-                        )
-                    except _GlobalSessionLimitReached:
-                        global_limit_hit = True
-                    except Exception as e:  # noqa: BLE001 - re-raised after cleanup
-                        start_error = e
+                if reuse_candidate is None:
+                    budget_used = self._used_dpu()
+                    budget_ok = budget_used + dpu_request <= dpu_budget
+                    if budget_ok and self._has_room(key, max_sessions):
+                        try:
+                            new_session_id = self._start(
+                                key,
+                                athena_client,
+                                spark_work_group,
+                                engine_config,
+                                session_description,
+                                dpu_request,
+                            )
+                        except _GlobalSessionLimitReached:
+                            global_limit_hit = True
+                        except _AccountCapacityUnavailable:
+                            capacity_unavailable = True
+                        except Exception as e:  # noqa: BLE001 - re-raised after cleanup
+                            start_error = e
 
             # Slow Athena calls run outside the lock. Stale cleanup runs
             # even on start_session failure to avoid leaking prior sessions.
@@ -103,6 +140,21 @@ def acquire(
             if start_error is not None:
                 raise start_error
 
+            if global_limit_hit:
+                # AWS rejected even though our client-side accounting said
+                # the budget had room — usually means another dbt process
+                # shares the account-wide DPU quota.
+                LOGGER.warning(
+                    f"Athena rejected StartSession (account session limit) for key {key}; "
+                    f"client-side accounting saw used={budget_used} + request={dpu_request} "
+                    f"<= budget={dpu_budget}. Another process may share the account quota."
+                )
+            if capacity_unavailable:
+                LOGGER.warning(
+                    f"Athena rejected StartSession for key {key}: AWS region capacity "
+                    f"unavailable. Backing off; this is transient and budget cannot predict it."
+                )
+
             if reuse_candidate is not None:
                 # Athena may have killed the session while it sat in the pool.
                 if self._is_session_alive(athena_client, reuse_candidate):
@@ -127,12 +179,16 @@ def acquire(
             if time.monotonic() >= deadline:
                 raise DbtRuntimeError(
                     f"No Spark Connect session available for key {key} within {timeout}s "
-                    f"(max_sessions={max_sessions})"
+                    f"(max_sessions={max_sessions}, dpu_request={dpu_request}, "
+                    f"dpu_budget={dpu_budget}, last used_dpu={budget_used})"
                 )
 
-            # Longer wait on account-level limit so we don't hammer StartSession.
+            # Longer wait when AWS pushed back (limit or capacity) so we don't
+            # hammer StartSession during a region-wide event.
             sleep_for = (
-                self._GLOBAL_LIMIT_BACKOFF_SECONDS if global_limit_hit else polling_interval
+                self._GLOBAL_LIMIT_BACKOFF_SECONDS
+                if (global_limit_hit or capacity_unavailable)
+                else polling_interval
             )
             time.sleep(sleep_for)
             time_since_eviction += sleep_for
@@ -169,18 +225,24 @@ def _has_room(self, key: SessionKey, max_sessions: int) -> bool:
         count = sum(1 for info in self._sessions.values() if info["key"] == key)
         return count < max_sessions
 
+    def _used_dpu(self) -> int:
+        """Sum of DPUs reserved by registered sessions. Caller must hold ``self._lock``."""
+        return sum(info["dpu"] for info in self._sessions.values())
+
     def _start(
         self,
         key: SessionKey,
         athena_client: AthenaClient,
         spark_work_group: str,
         engine_config: EngineConfigurationTypeDef,
         session_description: str,
+        dpu: int,
     ) -> str:
         """Start a session and register it. Caller must hold ``self._lock``.
 
-        Translates the account-level session limit into
-        ``_GlobalSessionLimitReached``; other errors propagate.
+        Translates two transient AWS rejections into typed exceptions for the
+        caller's backoff loop: account session limit and region capacity
+        unavailable. Other errors propagate.
         """
         try:
             response = athena_client.start_session(
@@ -189,14 +251,21 @@ def _start(
                 EngineConfiguration=engine_config,
                 SessionIdleTimeoutInMinutes=SESSION_IDLE_TIMEOUT_MIN,
             )
-        except Exception as e:  # noqa: BLE001 - global-limit handled below
-            if "Maximum allowed sessions" in str(e):
-                LOGGER.warning(f"Athena session limit reached, will retry: {e}")
+        except Exception as e:  # noqa: BLE001 - transient errors handled below
+            message = str(e)
+            if "Maximum allowed sessions" in message:
                 raise _GlobalSessionLimitReached() from e
+            if "required capacity not being available" in message:
+                raise _AccountCapacityUnavailable() from e
             raise
 
         session_id = str(response["SessionId"])
-        self._sessions[session_id] = {"key": key, "client": athena_client, "load": 1}
+        self._sessions[session_id] = {
+            "key": key,
+            "client": athena_client,
+            "load": 1,
+            "dpu": dpu,
+        }
         return session_id
 
     def _get_session_state(self, athena_client: AthenaClient, session_id: str) -> str: