[Feature] Support heuristic-aware adaptive polling to avoid getting hit by HTTP 429

[mdrakiburrahman]

Is this your first time submitting a feature request?

• I have searched the existing issues, and I could not find an existing issue for this feature
• I am requesting a straightforward extension of existing dbt functionality, rather than a Big Idea better suited to a discussion

Describe the feature

The Livy endpoint throttles because it goes through Fabric's front end which has a heavy throttle finger:

For smaller models, this is fine, but for larger, heavy models, this results in wasted compute where the Spark Cluster is not bin packed.

Basically this:

https://github.com/dbt-labs/dbt-adapters/blob/860da89225e2ecf1bf47038f5ac40d4eaa4019a2/dbt-spark/src/dbt/adapters/spark/session.py#L21

Like this:

dbt-labs/dbt-spark#279
dbt-labs/dbt-spark#272

The 2 ways Customers could schedule this is:

1. Through a Spark Job Definition as a whl file that actually ends up invoking DbtRunner
2. A notebook, which does the same thing

Describe alternatives you've considered

Livy, but it throttles and is COGS prohibitive

Who will this benefit?

Super bin pack conscious customers

Are you interested in contributing this feature?

Yes

Anything else?

Probably just need to lift and shift the Spark adapter code and make it opt-in.

The Spark dep could be assumed to be provided (since Fabric Spark clusters already have it).

I'd personally end up scheduling an SJD with a whl or perhaps just have a Spark Notebook with minimal code (better logs).

[mdrakiburrahman]

An alternative is this: https://github.com/mdrakiburrahman/privy - but the user needs a Relay

[mdrakiburrahman]

An alternative could be telemetry driven backoff.

E.g., if we can somehow figure out from the Livy response how many tasks are remaining on an associated query, that can be used to intelligently poll for query completion.

[mdrakiburrahman]

I'm going to let this simmer a little bit and see if we can get Livy Retry-After to be more intelligent. Putting dbt on a Spark driver node kind of defeats the purpose of client-server decoupling

[mdrakiburrahman]

So this is Spark code that allows us to put a query into a jobGroup and get the status back.

Great blog from Miles Cole (@mwc360): https://milescole.dev/data-engineering/2024/09/13/Mastering-Spark-Enhancing-Job-Visibility.html

from concurrent.futures import ThreadPoolExecutor
import time
import uuid
from pyspark.sql import functions as F

# Increase DATA_SIZE if the queries finish too quickly.
DATA_SIZE = 5_000_000
SHUFFLE_PARTITIONS = 200
POLL_SECONDS = 2

sc = spark.sparkContext
tracker = sc.statusTracker()

spark.conf.set("spark.sql.shuffle.partitions", SHUFFLE_PARTITIONS)

# Lazy fake dataset: no data is materialized yet.
sales = (
    spark.range(DATA_SIZE)
    .select(
        F.col("id").alias("transaction_id"),
        F.pmod(F.xxhash64("id"), F.lit(500_000)).alias("customer_id"),
        F.pmod(F.xxhash64("id", F.lit(1)), F.lit(10_000)).alias("product_id"),
        F.pmod(F.xxhash64("id", F.lit(2)), F.lit(50)).alias("region_id"),
        (
            F.pmod(F.xxhash64("id", F.lit(3)), F.lit(10_000)) / 100.0
        ).alias("amount")
    )
    .repartition(SHUFFLE_PARTITIONS, "customer_id")
)

sales.createOrReplaceTempView("fake_sales")

queries = {
    "product-sales": """
        SELECT
            region_id,
            product_id,
            COUNT(*) AS transaction_count,
            SUM(amount) AS total_sales,
            AVG(amount) AS average_sale
        FROM fake_sales
        GROUP BY region_id, product_id
    """,
    "customer-summary": """
        SELECT
            customer_id,
            COUNT(*) AS transaction_count,
            COUNT(DISTINCT product_id) AS distinct_products,
            SUM(amount) AS total_spend
        FROM fake_sales
        GROUP BY customer_id
    """
}

groups = {
    name: f"{name}-{uuid.uuid4()}"
    for name in queries
}


def run_query(name, sql_text):
    group_id = groups[name]

    # Spark local properties are thread-local, so set the group here.
    sc.setJobGroup(
        group_id,
        f"Demo query: {name}",
        interruptOnCancel=True
    )

    try:
        # The noop sink executes the complete query without storing output.
        (
            spark.sql(sql_text)
            .write
            .format("noop")
            .mode("overwrite")
            .save()
        )
        return f"{name}: completed"
    finally:
        # Compatible cleanup for Spark versions without clearJobGroup().
        sc.setLocalProperty("spark.jobGroup.id", None)
        sc.setLocalProperty("spark.job.description", None)
        sc.setLocalProperty("spark.job.interruptOnCancel", None)


def print_progress():
    print(f"\n--- {time.strftime('%H:%M:%S')} ---")

    for query_name, group_id in groups.items():
        job_ids = list(tracker.getJobIdsForGroup(group_id))

        if not job_ids:
            print(f"{query_name}: waiting for job submission")
            continue

        for job_id in job_ids:
            job = tracker.getJobInfo(job_id)
            if job is None:
                continue

            print(f"{query_name} | job={job_id} | status={job.status}")

            for stage_id in job.stageIds:
                stage = tracker.getStageInfo(stage_id)
                if stage is None:
                    continue

                total = stage.numTasks
                completed = stage.numCompletedTasks
                percent = 100 * completed / total if total else 0

                print(
                    f"  stage={stage_id:<4} "
                    f"progress={completed}/{total} ({percent:5.1f}%) "
                    f"active={stage.numActiveTasks:<3} "
                    f"failed={stage.numFailedTasks:<3} "
                    f"name={stage.name}"
                )


with ThreadPoolExecutor(max_workers=len(queries)) as executor:
    futures = {
        name: executor.submit(run_query, name, sql_text)
        for name, sql_text in queries.items()
    }

    while not all(future.done() for future in futures.values()):
        print_progress()
        time.sleep(POLL_SECONDS)

    # Capture final statuses while scheduler metadata is still retained.
    print_progress()

    for name, future in futures.items():
        print(future.result())

spark.catalog.dropTempView("fake_sales")

--- 11:15:38 ---
product-sales: waiting for job submission
customer-summary: waiting for job submission

--- 11:15:40 ---
product-sales | job=10 | status=RUNNING
  stage=13   progress=2/8 ( 25.0%) active=6   failed=0   name=save at NativeMethodAccessorImpl.java:0
customer-summary | job=9 | status=RUNNING
  stage=14   progress=0/8 (  0.0%) active=2   failed=0   name=save at NativeMethodAccessorImpl.java:0

--- 11:15:42 ---
product-sales | job=11 | status=RUNNING
  stage=15   progress=0/8 (  0.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
  stage=16   progress=0/200 (  0.0%) active=6   failed=0   name=save at NativeMethodAccessorImpl.java:0
product-sales | job=10 | status=SUCCEEDED
  stage=13   progress=8/8 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
customer-summary | job=9 | status=SUCCEEDED
  stage=14   progress=8/8 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0

--- 11:15:44 ---
product-sales | job=11 | status=RUNNING
  stage=15   progress=0/8 (  0.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
  stage=16   progress=11/200 (  5.5%) active=5   failed=0   name=save at NativeMethodAccessorImpl.java:0
product-sales | job=10 | status=SUCCEEDED
  stage=13   progress=8/8 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
customer-summary | job=12 | status=RUNNING
  stage=17   progress=0/8 (  0.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
  stage=18   progress=7/200 (  3.5%) active=5   failed=0   name=save at NativeMethodAccessorImpl.java:0
customer-summary | job=9 | status=SUCCEEDED
  stage=14   progress=8/8 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0

--- 11:15:46 ---
product-sales | job=11 | status=RUNNING
  stage=15   progress=0/8 (  0.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
  stage=16   progress=27/200 ( 13.5%) active=5   failed=0   name=save at NativeMethodAccessorImpl.java:0
product-sales | job=10 | status=SUCCEEDED
  stage=13   progress=8/8 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
customer-summary | job=12 | status=RUNNING
  stage=17   progress=0/8 (  0.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
  stage=18   progress=56/200 ( 28.0%) active=5   failed=0   name=save at NativeMethodAccessorImpl.java:0
customer-summary | job=9 | status=SUCCEEDED
  stage=14   progress=8/8 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0

--- 11:15:48 ---
product-sales | job=11 | status=RUNNING
  stage=15   progress=0/8 (  0.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
  stage=16   progress=46/200 ( 23.0%) active=5   failed=0   name=save at NativeMethodAccessorImpl.java:0
product-sales | job=10 | status=SUCCEEDED
  stage=13   progress=8/8 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
customer-summary | job=12 | status=RUNNING
  stage=17   progress=0/8 (  0.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
  stage=18   progress=119/200 ( 59.5%) active=5   failed=0   name=save at NativeMethodAccessorImpl.java:0
customer-summary | job=9 | status=SUCCEEDED
  stage=14   progress=8/8 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0

--- 11:15:50 ---
product-sales | job=11 | status=RUNNING
  stage=15   progress=0/8 (  0.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
  stage=16   progress=71/200 ( 35.5%) active=9   failed=0   name=save at NativeMethodAccessorImpl.java:0
product-sales | job=10 | status=SUCCEEDED
  stage=13   progress=8/8 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
customer-summary | job=12 | status=SUCCEEDED
  stage=17   progress=0/8 (  0.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
  stage=18   progress=200/200 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
customer-summary | job=9 | status=SUCCEEDED
  stage=14   progress=8/8 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0

--- 11:15:52 ---
product-sales | job=11 | status=RUNNING
  stage=15   progress=0/8 (  0.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
  stage=16   progress=144/200 ( 72.0%) active=9   failed=0   name=save at NativeMethodAccessorImpl.java:0
product-sales | job=10 | status=SUCCEEDED
  stage=13   progress=8/8 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
customer-summary | job=12 | status=SUCCEEDED
  stage=17   progress=0/8 (  0.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
  stage=18   progress=200/200 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
customer-summary | job=9 | status=SUCCEEDED
  stage=14   progress=8/8 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0

--- 11:15:54 ---
product-sales | job=13 | status=SUCCEEDED
  stage=19   progress=0/8 (  0.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
  stage=20   progress=0/200 (  0.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
  stage=21   progress=9/9 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
product-sales | job=11 | status=SUCCEEDED
  stage=15   progress=0/8 (  0.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
  stage=16   progress=200/200 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
product-sales | job=10 | status=SUCCEEDED
  stage=13   progress=8/8 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
customer-summary | job=12 | status=SUCCEEDED
  stage=17   progress=0/8 (  0.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
  stage=18   progress=200/200 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
customer-summary | job=9 | status=SUCCEEDED
  stage=14   progress=8/8 (100.0%) active=0   failed=0   name=save at NativeMethodAccessorImpl.java:0
product-sales: completed
customer-summary: completed
True

Technically speaking, if we could do this with Fabric Livy somehow, some sort of trivial linear regressor lets us retry more intelligently

[mdrakiburrahman]

Proposed implementation: HC Livy + Job Groups + per-query adaptive polling

This can stay client/server and does not require dbt to run on the Spark driver. The adapter can inject a very small amount of PySpark through the existing High Concurrency Livy REPLs, while all scheduling, rate limiting, heuristics, and state remain in the adapter process.

Why High Concurrency makes this possible

The current HC backend already gives each dbt thread an isolated REPL while packing all REPLs onto the same underlying Spark session:

dbt worker A -> repl A --\
dbt worker B -> repl B ----> shared sessionId / SparkContext
dbt worker C -> repl C --/
monitor      -> repl M --/

That allows model statements to execute concurrently while one additional HC REPL reads SparkStatusTracker for every active model.

Minimal Spark-side surface

Bootstrap this helper once in each worker REPL:

def __dbt_run(group_id, description, sql):
    sc.setJobGroup(group_id, description, True)
    try:
        return spark.sql(sql).collect()
    finally:
        sc.clearJobGroup()

Subsequent model execution sends only:

__dbt_run(<json-encoded-group-id>, <json-encoded-description>, <json-encoded-sql>)

Use json.dumps() in the adapter to generate Python string literals safely. Setting the group and executing the query must occur in the same PySpark statement because Spark local properties are thread-local.

Keep ordinary kind: "sql" for adapter metadata queries that need the existing Livy application/json schema/row envelope. Initially enable wrapped PySpark only for dbt model/materialization statements where completion, rather than a returned rowset, is required. Multi-statement materializations should receive child IDs:

<invocation>.<node unique_id>.01.create_staging
<invocation>.<node unique_id>.02.merge
<invocation>.<node unique_id>.03.cleanup

The adapter aggregates those under the model's unique_id.

Bootstrap one monitor helper in a dedicated HC REPL:

def __dbt_progress(group_ids):
    tracker = sc.statusTracker()
    result = {}
    for group_id in group_ids:
        jobs = []
        for job_id in tracker.getJobIdsForGroup(group_id):
            job = tracker.getJobInfo(job_id)
            if job is None:
                continue
            stages = []
            for stage_id in job.stageIds:
                stage = tracker.getStageInfo(stage_id)
                if stage is not None:
                    stages.append((
                        stage_id,
                        stage.numTasks,
                        stage.numActiveTasks,
                        stage.numCompletedTasks,
                        stage.numFailedTasks,
                    ))
            jobs.append((job_id, str(job.status), stages))
        result[group_id] = jobs
    return result

Each monitoring statement can request only the queries currently due:

__dbt_progress(["group-a", "group-c"])

Adapter-side execution state

Add a process-wide coordinator and immutable snapshots:

@dataclass
class QueryProgress:
    execution_id: str
    node_unique_id: str
    statement_url: str
    submitted_at: float
    last_progress_poll_at: float | None
    last_statement_poll_at: float | None
    next_progress_poll_at: float
    next_statement_poll_at: float
    total_tasks: int
    completed_tasks: int
    active_tasks: int
    failed_tasks: int
    completion_rate: float | None
    estimated_remaining_seconds: float | None
    stage_ids: frozenset[int]
    terminal_event: threading.Event

The worker thread submits its model and waits on terminal_event. It must no longer run its own 1.5-second _poll() loop. A single coordinator owns:

• The active-query registry.
• A priority queue ordered by next_progress_poll_at.
• Sparse authoritative Livy statement polling.
• A workspace/lakehouse-wide rate-limit deadline.
• Delivery of terminal results back to waiting cursors.

Per-query adaptive cadence

Derive the next sample from observed task throughput, not table size:

elapsed = now - previous.last_progress_poll_at
delta = current.completed_tasks - previous.completed_tasks
instantaneous_rate = delta / elapsed if elapsed > 0 and delta > 0 else None

# Exponentially weighted moving average.
rate = (
    instantaneous_rate
    if previous.completion_rate is None
    else 0.3 * instantaneous_rate + 0.7 * previous.completion_rate
) if instantaneous_rate is not None else previous.completion_rate

remaining = max(current.total_tasks - current.completed_tasks, 0)
eta = remaining / rate if rate and rate > 0 else None

When throughput is known, schedule the next sample around the earlier of:

• The time required for approximately another 1% of tasks to finish.
• 10% of the estimated remaining duration.

Clamp it:

interval = clamp(min(tasks_for_one_percent / rate, eta * 0.10), 0.25, 60.0)

Suggested event overrides:

Observation	Next progress sample
No job registered yet	Back off toward 5 seconds
New job or stage appears	0.5 seconds to learn its size/rate
Tiny stage (<=100 tasks)	0.25 seconds
Progress observed	Throughput/ETA formula
Active tasks but no completion	Multiply prior interval by 1.5, cap 30 seconds
No active tasks and no transition	Multiply by 2, cap 60 seconds
All known jobs terminal	Poll original Livy statement immediately

The 5 TB query therefore naturally moves toward tens of seconds between samples, while a tiny query stays subsecond and receives its terminal result quickly.

Do not assume all stages are known up front: AQE and later Spark jobs can introduce new stages. Progress is a scheduling signal, not a fixed completion percentage.

Batch only queries that are due

Use a small coalescing window:

due = [
    query.execution_id
    for query in active_queries
    if query.next_progress_poll_at <= now + 0.25
]

Example:

t=0.0   monitor [tiny, medium, five_tb]
t=0.3   monitor [tiny]
t=1.0   monitor [tiny, medium]
t=2.1   tiny Spark jobs finish -> GET tiny Livy statement now
t=8.0   monitor [medium]
t=30.0  monitor [five_tb]
t=60.0  monitor [five_tb]

Do not include the 5 TB query merely because the tiny query is due.

Separate progress polling from terminal polling

statusTracker is advisory and recent-only; Livy remains authoritative for statement output and terminal state.

GET the original model statement when:

• Its Job Group reports all known jobs terminal.
• A job reports failure.
• No Spark job appears within a bounded registration window.
• A maximum liveness interval expires (for example, 60 seconds).
• Cancellation is requested.
• The statement timeout is approaching.

This replaces a GET every 1.5 seconds per model with targeted GETs near likely completion.

Retry-After must be global

The existing _poll() loops honor Retry-After only in the thread that receives the 429. With many HC REPLs, the remaining threads continue calling Fabric.

Store a shared deadline keyed by (workspace_id, lakehouse_id):

rate_limit_until[key] = max(
    rate_limit_until.get(key, 0),
    monotonic() + retry_after + random_jitter,
)

Every submit, monitor, statement GET, cancel, and session request should wait behind this gate. The server-provided deadline always overrides heuristic due times. Add jitter so all deferred calls do not resume simultaneously.

Retry and correctness rules

• Once a statement ID exists, polling failures and 429s must always retry that same ID; never resubmit the SQL.
• waiting and running prove acceptance and must never trigger resubmission.
• Spark progress stopping is not proof of failure; it only lengthens the next sample.
• An engine error is not a transport retry.
• A lost REPL or ambiguous POST may have an unknown outcome. Do not automatically replay DDL/DML.
• Incremental models contain independently side-effecting staging, DELETE/INSERT/MERGE, and cleanup statements. Automatic replay requires operation-specific idempotency, not merely missing telemetry.
• Cancellation should call both the HC statement cancel endpoint and, when needed, sc.cancelJobGroup(execution_id) through the monitor REPL.

Suggested integration points

• connections.py:add_query: capture invocation/node identity and classify whether the statement needs returned rows.
• concurrent_livy.py:HighConcurrencyCursor._submit: choose SQL versus wrapped PySpark and register the returned statement ID.
• Replace HighConcurrencyCursor._poll with coordinator registration plus waiting for its terminal result.
• HighConcurrencySessionManager: lazily acquire one monitor HC session with the same sessionTag.
• New module, e.g. adaptive_polling.py: registry, priority queue, global rate-limit gate, throughput calculation, and coordinator lifecycle.

Keep this behind an opt-in profile flag for the first release, for example:

high_concurrency: true
adaptive_polling: true

Required tests

1. Two worker REPLs and one monitor REPL share sessionId but have distinct replIds.
2. Wrapped statements create distinct Job Groups for concurrent models.
3. Tiny jobs receive short intervals; long-running high-task jobs back off.
4. A newly discovered stage resets the interval briefly.
5. Only due group IDs are sent to the monitor.
6. Successful/failed Spark jobs force immediate authoritative Livy polling.
7. A 120-second Retry-After blocks traffic from every worker and the monitor.
8. Jitter prevents a thundering herd after cooldown.
9. Poll errors never resubmit an accepted statement.
10. Terminal events wake only the correct waiting cursor.
11. Multi-statement incremental models retain separate execution IDs and are never blindly replayed.
12. Monitor failure degrades safely to slow Livy polling rather than failing active models.

This design keeps dbt outside the driver, injects only two small helpers, uses HC packing for concurrent visibility, and makes each query's own measured Spark progress determine its next poll time.