AI_DIAGNOSE: Automating Spark Job Debugging With LLM Agents

Summary

AI_DIAGNOSE is a new Spark SQL procedure in Databricks Runtime that uses a ReAct-style LLM agent to automatically diagnose Spark job failures, reducing the manual work of sifting through hundreds of lines of distributed stack traces.
The agent autonomously gathers distributed logs, retrieves task metrics, and produces grounded evidence-based root cause analysis—then generates actionable fixes such as AQE configuration, OPTIMIZE commands, and statistics computation.
This video demonstrates AI_DIAGNOSE diagnosing two real failure types: a data skew issue responsible for a 40% runtime impact and a streaming Auto Loader memory failure caused by misconfigured executor memory.

AI_DIAGNOSE: Automating Spark Job Debugging With LLM Agents

Debugging distributed Spark jobs is notoriously difficult, often requiring hours of sifting through massive logs to find root causes. We introduce AI_DIAGNOSE, a new Spark SQL procedure in Databricks Runtime that automates this process using large language models (LLMs). this video explores the architecture of AI_DIAGNOSE, detailing how we implemented a ReAct-style agent within the Spark engine. We will demonstrate how the agent autonomously gathers context, retrieves distributed logs, and analyzes job failures. We'll cover technical challenges like managing LLM context windows with large logs, ensuring structured actionable outputs, and leveraging Databricks Model Serving for secure inference. Join us to learn how AI_DIAGNOSE works under the hood to drastically reduce Mean Time To Resolution (MTTR) for data pipelines.
Talk By: Allison Wang, Staff Software Engineer, Databricks ; Shujing Yang, Software Engineer, Databricks ;

Chapters

FAQs

What is AI_DIAGNOSE and how does it work?

AI_DIAGNOSE is a Spark SQL procedure built into Databricks Runtime that uses a ReAct-style LLM agent to automatically debug failing Spark jobs. The agent retrieves distributed logs, queries task metrics, and reasons through job failures step-by-step to produce a grounded root cause analysis with evidence from actual execution data rather than guesses.

What kinds of Spark failures can AI_DIAGNOSE detect?

AI_DIAGNOSE can detect a range of Spark failure patterns including data skew, out-of-memory executor failures, streaming pipeline crashes, and jobs that stall without producing an error. This video demonstrates it diagnosing a data skew issue with a 40% runtime impact and a streaming Auto Loader failure caused by misconfigured executor memory.

How does AI_DIAGNOSE handle large log files and context window limitations?

Managing LLM context windows with large distributed logs is a core engineering challenge addressed in this video. AI_DIAGNOSE addresses this by having the agent selectively retrieve and prioritize the most relevant log segments and task metrics rather than passing entire log files to the model.

Why is debugging distributed Spark jobs so difficult without AI assistance?

Debugging Spark jobs typically requires deep expertise: opening the Spark UI, reading query plans, cross-referencing distributed logs and stage metrics, and assembling root cause analysis manually. Even experienced engineers find this slow and tedious, and the problem is compounded in 2026 by the fact that most Spark code is now written by coding agents that may not fully understand the execution context.

Full transcript

[00:08] Okay. Hi everyone. Uh thank you all for coming today. Uh my name is Allison and this is my co-orker Shujene. So uh we are both engineers on the Spark experience team at data bricks and uh so today we're going to talk about how to
[00:22] debug spark uh using LMS. Okay. So, quick show of hands. How many of you have ever tried Spark?
[00:35] Okay, great. Most most of you. Okay. So, how many of you actually run into issues when using Spark? Okay, 80% of them. Yeah. So, I love to talk to people who didn't raise your
[00:47] hands. Uh, love to know your secrets. Um, but yeah, so Spark can go wrong sometimes. And here is but here's what changed in 2026. Most of us are not
[01:00] actually writing code by hand anymore. So we're actually uh pointing a coding agent and ask it to do the work. So let's start here.
[01:14] Here is um a very common workflow in 2026. So imagine that you are a data engineer and you ask your coding agent to write this pispark code for you. So it reads a couple of tables. It does some filtering, joins, aggregations, and
[01:28] it writes to another table. So honestly, if you look at the code, it looks completely correct. But in reality, depending on your input data, this can go wrong. So sometimes it can fail with
[01:42] a hundreds lines of stack trace and you're left digging into that one line that actually matters. And even if you find it, you still can't understand what exactly it means. and how to fix it. And
[01:56] sometimes it's running out of memory. You're getting this womb exception and some executive died in the middle of the stage. Sometimes uh it finishes but it took a bit longer than you expected. And
[02:09] what worse, the job just stuck forever and there's no error at all. So how how do you debug? And here are um two ways to debug Spark today. Um the first one, let's just call
[02:23] it the expert path. You open the Spark UI, you go through all the stages, you read the query plans, you cross reference the logs, the metrics, and the configurations until you actually figure
[02:36] out what's going on and assemble your uh root cost analysis yourself. And on the right, there's another path that everybody knows. Uh if you think about the first path, right? So it's um it
[02:48] works uh but it can be slow. It can be tedious and it basically assumes that you know everything already. You you know everything you're looking for. You know what are the errors, how to debug an issue. But when we don't actually
[03:02] know what's going on, we move to the right path. So this is the path that almost everybody would take nowadays. You point in your coding agent at your code. uh you copy paste the error messages into the coding agent and you
[03:16] just ask it how to fix my job it's failing right so this path is certainly faster when it works the problem is that oftent time it just doesn't give you the right answer
[03:29] so here is what it looks like when it doesn't work so you copy paste this fetch fail exception which you have no idea what's going on and agent come back and it very confidently tells tells you
[03:42] that bump your executive memory uh add a cache before the join and then you should be good to go. If you look at this answer, it's plausible but the only problem is that it's wrong. So in this
[03:56] case, the shuffle was not failing because of uh the memory. It was failing because of a data skew problem. So adding memory and adding a cache layer doesn't really fix anything. It just
[04:08] costs you more money and hides the real problem for a little bit longer. So here's the trap, right? So the models sound exactly like an expert, but it sends you down to a completely wrong
[04:20] path. It can give you such suggestions, but they're so generic and they just don't tell you what exactly is going on. What's the root cause of the problem?
[04:32] Okay, so I want to be a bit uh more fair here because you will be easy just to blame the model uh not doing a good job but that's not right. Okay, so it's not model's fault. So the truth is that
[04:46] spark debugging is very hard and we just handed the model almost nothing to work with except for the source code and also the stack trace. Right? So there are a couple of issues, right? So first of
[04:58] all, there just too little context like a code snippet or a stack trace just don't tell you what hap actually happened across the entire cluster, right? So the evidence simply isn't in
[05:11] the code. And second of all, there just too many root causes for a same spar problem. For example, if you see an out of memory issue or a hanging issue, it could be because of skew. It could be
[05:23] because of spills. It could be because of slow UDFs uh or even bad partitioning. Right? So there are dozens of different causes that can are look identical from the outside. So you have
[05:35] to dig into what's exactly causing this issue. And the third one is that even knowing the cause doesn't mean you can fix it. So there are already years of dice talks and many blog posts online
[05:49] mentioning how to use the Spark UI or how to do performance tuning for Spark, but the knowledge just lives in everybody's heads, right? They're not written down. Uh they might be included in the pre-training phase of the model,
[06:02] but they never make it to the context, the model's existing context. So the model is basically guessing the right answer. So, so I want to reframe this problem
[06:15] for us. So, the issue here is not how we get a smarter model. Um, we don't control that, right? It's out of outside of our control. Uh, and it's not the bottleneck anyways because model nowadays are very smart. So, given this,
[06:30] how can we build a better debugging agents for Spark that's reliable, that's accurate, and is grounded in how Spark works under the hood. So now let's talk about how we approached it.
[06:44] So the first question where do we even start? It seems to be a very vague and very big problem because Spark can just fail in so many different ways. So here's how we draw the line. Um so the
[06:57] things that we go after are the execution errors. So the problem that actually show up while the job is actually running, right? So for example, you're running out of memory. uh you have a data skew and you have Python
[07:09] UDFs that just stuck forever. You have lost executors, Python storage streaming. So all these execution errors and the reason why that we focus on these errors are is because there are
[07:22] evidence in the log right there in the metrics and the agents once they can read the evidence they can reason about the root cause and they can give you the right fix. And we're not targeting these
[07:36] non-execution failures. For example, logic like the job succeeded but the data is actually wrong or it's just not a spar problem at all like a third party issue like upstream issue. So these are still important but they're just not our
[07:50] immediate focus for now. So this boundary is very important because we want agent to uh agent reasoning to be rooted in the evidence that it can find because there's one
[08:02] symptom that can have different causes. The agent need to lean on everything available to it like the spark UI metrics logs query plans uh to find the actual root cause instead of guessing
[08:14] the most common issues which we will demo later in um by Shujim. Okay. So here is the core insight of the entire talk. We can make the sp the model smarter but we can give it the
[08:28] right context is been missing. Oh sorry. Okay. Yeah. So there's the core insight. Um we can make this model smarter but we can give the right context. We know exactly which metrics
[08:41] and signal reveal what actually went wrong. So because we have debugged like everybody here work with spark. We have debugged some spark issues before, right? And we have been solving this problem by hands and now it's a time to
[08:55] transfer this knowledge to agents. So instead of hoping the model guesses well, uh we feed the evidence. Um that's the whole idea behind AI diagnose. Um and everything else is just execution.
[09:10] So now let me um show you how the system actually works and uh show you two case studies to help you understand better.
[09:25] So this is uh the diagram that showing how AI diagnose how the debugging agent works. Um so the center is the agent. Um so it's running essentially running a react loop. It gets the initial spark failure or symptoms uh reasons about it.
[09:38] uh fetch the evidence uh look at what comes back and just repeat this process. So the agent uh if you look at the uh the top uh the top one so the agent uh talks to an uh LLM endpoint. Uh so it
[09:52] can be any LM model but here for the most accurate reasoning uh we typically use either GPT model or the cloud model. Um so that is the reasoning engine. So the agents hand the model a right
[10:05] context that it gathered and then the model decide what to do next. And where's this context coming from? So it's actually coming from two different places. So on the left hand side uh we
[10:18] have the tools. So these are the tools that can be used by the agent to fetch sparks information. So for example, what if you want to fetch the metrics uh from spark logs? What if you want to fetch
[10:31] the uh the query plans or you want to fetch the uh standard out or standard error from your driver and executive nodes and um so we give this tools we provide them as the as different tools
[10:43] for the agent so that um it's like uh it's able to get all the informations um directly from Spark UI right so the agent can use the tool uh selectively uh based on the playbook uh on the right
[10:57] hand side so the playbook are the debugging skills right so these are the markdown files that shows how to debug various issues so for example how to debug skew shuffle fetch failure out of
[11:11] memory python udfs and so much uh so much more so this is the knowledge base uh where the agent can use um to uh to debug a certain issue and the good thing is the agent can only load a skill uh
[11:25] once it's actually relevant to the problem so you wouldn't load everything together which will just blow up the context window here. So this is the entire loop. Um agent loads the skill uh gather the evidence
[11:37] using the tools and reason with the LM model uh and then it repeats until it's confident that it can give you a real answer uh with the real root cause the fix and also provide the evidence behind
[11:50] it. or if there's no much signal where it cannot debug the issue, uh it will just tell you that uh I cannot debug it instead of guessing or hallucinating a wrong answer.
[12:09] Okay, so here is our first case study. Uh let's imagine that you have a Python UDF. So it's showing this code on the right hand side. Um so this Python UDF uh you observe that it just never finishes, right? This is kind of thing
[12:21] that's worse to debug because there's nothing to grab onto. There's no logs. You don't see anything. You're just stuck there. So like you have no idea what's going on and why just hangs for like two hours. And so basically if we
[12:36] take a closer look at the code, it's basically sending a request uh to an external service. So if that service is unresponsive, uh then the Python worker will just block forever and waiting for
[12:49] the call. So if you're debugging using an generic LLM, uh what it typically does is that you will assume that the job is under resourced and you will try to bump the memory. Uh you will try to
[13:02] add a cache or you just guess a couple reasons for you to try it out and you can validate what's going on. And here is what AI diagnose does. Okay. So, Python UDF hanging is a very common
[13:15] issue in Spark and since it's a very common issue uh so Spark developers um so we have introduced a tool for debugging issues like this. So we added
[13:27] this in Spark 4.0 uh and we add this uh we add this concept or this uh this instruction into AI diagonal skill so it knows how to search for this log. So essentially we
[13:40] added this warning log for example this idle timeout reach for Python worker into uh into spark. So whenever uh spark encounters an errors or or udf hanging this warning log will show up in exeutor
[13:55] and this is by default enabled is added in spark 4 but almost nobody knows it right there are thousands of words hundreds of different spark configs and it's just impossible to remember every
[14:07] single one of them and when you encounter issues like this like how do you remember to search for this idle timeout log right but nowadays since we agents uh and is agents debugging the
[14:19] issue. You don't really need to remember like our human don't even need to remember this flag. We just encode this knowledge into agent skill and the agent will do the right thing for us. And if
[14:31] you can see that in the fake section uh it suggested two spark configurations which is also nobody knows or understand what's going on. But essentially these are the two things are are critical
[14:43] because the first thing is to tell spark to kill the udf uh when once the timeout uh limit reached so it won't stuck there forever for example if you set to be five minutes then the spark job will
[14:56] fail after five minutes so the udf won't stuck forever and the second one is actually to turn on the trace back dump for python python udfs python workers so you actually know it will tell you
[15:08] exactly from the stack trace which line in your UDF is stuck then you understand what's going on in your UDF and how to debug it. So these two configurations again they are just hard to remember and
[15:21] nowadays we just don't ask human to remember we just ask agent to suggest them as the next fix right then another coding agent your coding agent sees the fix you will set these two configs you will run the job again and then the
[15:34] agent can relook at the logs to see okay what's going on and then pinpoint the issue where your UDF is stuck on. So another use case uh for here is okay
[15:47] so this is another very common problem is data skew um so your code is just doing a very simple repartitioning into 2,000 partitions um and we can say that um so let's assume there's one key which
[16:00] is 40 times bigger than the rest and this operation is running very slowly uh a generic LM again it doesn't have too much context so it just can guess and ask you to add more memory uh which is
[16:13] not really the right fix. So AI diagnose here it doesn't guess right we want the truth we want the metrics and it just pull the right metrics and run a skew detection algorithm to detect the skew
[16:27] statistically. So for example in this case uh it pulled the metrics from spark event logs and it discovered that there are three strangular tasks uh the longest one is about 38 times than the medium and executor are sitting idle for
[16:41] 62% of the time and the most importantly it just quantifies the impact. So the skew is responsible for roughly 40% of the queries runtime. So this is
[16:54] basically links the skew stage um back to the operator is able to actually find uh which operator caused the skill and then identify the key the hot key behind it. So it just doesn't tell you, okay,
[17:08] you have a skill. Um, you'll hear from AI diagnose which key is the hot key, how bad it is, and how much is costing you all backed by the real numbers.
[17:26] But um, context engineering is certainly helpful. Uh, it can drastically increase the accuracy for your agents. Uh, but it's not easy. Like if you take like one engineering lesson from this talk uh is this. So getting the right context to
[17:40] agent reliably is smart uh much harder uh than it sounds. So there are pre pretty much uh three challenges. So first of all it's the size right
[17:52] everybody knows here like uh spark law can be very big. So agents has very limited uh limited context window. It's usually around 200 or 250K.
[18:04] And since Spark Locks is big, right, you can't really just dump everything into the agent's context window. You have to be very selective on what actually we want to send to LMS.
[18:16] And second one is it's nondeterministic. The agent can take drastically different path on the same problem uh from one round to another. So how can you make it more deterministic? um and the more debugging surface that
[18:29] it covers um the more is allowed to perform uh in a very very differently. So it's very hard to actually keep the behavior stable. And the third one is um a very real
[18:42] challenge um is how to ensure agent quality at scale. It's easy pretty easy to build agent that can just tackle one or two ad hoc problems. But how can you
[18:54] actually reliably build agents that can cover the entire Spark debugging and how can we make sure that we don't introduce any regressions when we tweak the prompt or the skill files.
[19:11] Okay, so here is how we keep agents reliable. Um there are essentially two pieces. So first of all, you need a eval framework. This is literally like a unit test or integration test for your agents. Um so we have over hundreds of
[19:26] eval test cases end to end eval test cases for our debugging agents and every single one is reproduced uh from a real customer issue not something we just made up. Uh and every prompt change and
[19:39] every context change uh gets scored against all the eval cases before we actually ship them. So we use internally we use LM as a judge uh to compare the results against the ground truth. So
[19:52] reliability is no longer a vibe uh it's a number that we can track. And the second piece is the agent observability. So we not only debug spark uh we also
[20:05] need to understand why agents fail to debug spark. So to do that we need to capture the full trace of the agent's reasoning and every single tool call that it used um so that when it does not
[20:19] get something right uh we can check all the agent traces all the reasoning process and make sure and understand what's going on and debug the agent. So we're basically debugging the debugger.
[20:31] Uh so there's no black box here and um and I also want to mention that um the real long-term goal here is to make spark agent friendly. Uh so here's
[20:44] the key pieces. So as I said raw spark logs are simply too large and too inconvenient for agent to debug. So we are uh thinking and adding a new kind of spark logs that can be more compact that
[20:58] can be more structured and that can have more signal of what the query actually does so that the agent can easily use these new type of logs to get the right information right away.
[21:10] And log is not the only piece. Um we are also thinking about uh giving agent more uh userfriendly error messages uh so the agent can fix them right away. um and
[21:22] also um basically give them the programmatic access to the plan and the metrics so the agent doesn't have to scrape a UI uh that's built for human. Yeah. So if you have more ideas on how
[21:36] to make Spark agent friendly, feel free to talk to us offline. Okay. Now Shujin will give a demo on AI diagnose.
[21:57] Hi everyone, I'm excited to present how AI diagnose works today. So imagine you own this injection pipeline and it used to works but one day when you wake up you felt that okay the recent drop fails. You
[22:11] click into the drop run and you are facing with this long um arrow messages
[22:23] after. So for this injection job what it does is very simple um and very common in our production job. What it does is that first of all it start a spark streaming job. It's an autoloader ingesting JSON file.
[22:45] And then after the injection you will do some custom transformation for here um it is you are doing conting summary and maxing and also this collecting list and then write it to a bronze table or parket files after start the spark
[22:59] streaming job if else the arrow message here is is a for each batch arrow which is too dramatic and you scroll all
[23:11] the way down to find the underlying cause for it. Okay, it's so hard to find but um yeah if you scroll all the way down the
[23:24] issue is it is a metadata fetch failure. If you are dealing with terabytes of data is one of the most common issue we have seen in shuffle fat failure. What it says is that it's missing an
[23:37] output location for shuffle zero partition 200. And nowadays what people basically uh would do is that we'll pay uh copy all the error messages sending to your favorite coding agent and also getting
[23:51] your codes also sending to your um favorite coding agent and also let's try to do it now. That's a very long stack trace.
[24:05] Yeah, we finally get your copy of the error message and then copy your code if your agent doesn't have context to that
[24:17] and then it will send to one of the most frontier LM. Paste your error message here and ask it to just debug this spark drop failure.
[24:49] It's a pretty long input and LM is trying its best to generate answer.
[25:01] How many lines are there in the stack trace? >> Definitely over 100. >> Yeah. Okay. >> So, yeah, let's see what LMS will usually outputs. So, first of all, it said that it's an executive die during
[25:13] the shuffle. It's pretty reasonable because if you are familiar with the shuffle fetch failure when it says it's missing an output, it means that spark lost the exeutor while it's trying to fetch the metadata. So what LM asks you
[25:27] to do is that you need to look into the spark UI and you should figure out why the exeutor w there could be many reason that executive w um it could be the disk view or spot instance preeemption or
[25:39] just network issues or hot bit timeouts. And besides this executive die um dying issue, I'm is also guessing that it's probably the collect list. This operator
[25:51] can be super dangerous because while trying to do the collect list, you're basically pulling all the data into one reducer which can cause
[26:05] executive womb or massive shuffle spill. Besides the single operator, it also suggests that okay the for each batch itself can be also problematic. You should inspect the payload size
[26:20] and also you should rerun it to check what exactly the payload size is.
[26:33] Besides this um for each batch operation, it could also be the microbatch size itself because you are running a spark streaming job and um each microbatch size. If the each microbatch size is huge, it can also run into issues
[26:47] and also it could be a skill issue. If you have a skill key, if you have a hot key that can also cause this um shuffle fetch failure. So all the reasonings and all the suggested next steps sounds
[26:59] pretty reasonable and this is just what um Allison mentioned in the previous talk in the previous um slides. So all the M's output sounds plausible well formatted and generic but it is not
[27:12] grounded in what spark actually does. So let's see how AI diagnose works today. You can also go into your um datab bricks workspace and go into the
[27:25] top right click into the genie code bottom and ask genic code to debug this spark job failures. Gen code will spend some time to pull your drop run configurations, pull your drop run outputs and gather all the context and
[27:39] gen code will also trigger AI diagnose. Genico was said that okay this is um this is not um this is not only a metadata fetch failure it is actually
[27:52] because of executive womb doing the aggregation in this for each batch function and here's what actually happened so first of all it has a massive micro batch in the autoloader you have specify autoloader was
[28:05] available now and unfortunately today there's just almost thousands of files today in the single batch which is far too much data for the configurable uh for the config memory
[28:17] and it also pull the spark configs to get um your executive memory configuration which only has three gigabytes configured and this is a cascading failure. This womb cute
[28:30] exeutor and it lost shuffle outputs and then finally it causes the metadata fetch failure. So based on this um statistic and reasoning what it proposed is that first of all you need to limit
[28:42] your microbatch size and you can possibly also increase your executive memory to 12 GB and possibly increase the shuffle partition. The good thing is uh for that is that it
[28:55] not only just suggest this but you can simply apply the fix by clicking the button in genic code and then genode will apply the fix redraw for you.
[29:12] Um, I also want to show you another demo. What if AI uh what if the spark job doesn't fail but it just runs slower than I expected. I wanted to debug the performance issue. For instance, here we have a very simple SQL query which
[29:25] basically create a table by doing some transformation based on a bronze table. Um given that it only outputs around a 100 rows of table, I
[29:37] didn't expect it to run over like 12 minutes. So I also go into the genie code ask genie code to debug why my spark job runs slower than my expect. So first of all, Genie finds that okay,
[29:50] it spends eight minutes in spinning up the cluster and it spends four minutes in the actual spark execution and in the spark execution it identifies there's a severe data skew and also you
[30:02] have this enable uh disable AQE for that. What he found that is a particularly single hotkey customer partition that is processed almost 90% 98% of the rows in the sort mer join and
[30:16] the single task running 82 longer than the average task. Besides this, the joint strategy is also not the optimal one because of lacking
[30:28] the um because of disabling the AQE and another issues is that the table itself contain too minus more files uh which can also be slow. So based on this
[30:41] first of all um you know there's a data skill issue and you should um enable AQE for that. Besides this, you can also run this optimize query to compact the small files
[30:54] and also compute the statics for costbased um optimization. So what it does is that it not only pull over all the drop run arrows, it does
[31:06] what basically a human spark engineer will will usually search for. So first of all it will pull the task duration distribution and then you'll find the strangular task find the stage that
[31:18] contains the strangular task and then find the operator and consequently it find the hotkey. So AI diagnut is not just guessing it is grounded in the spark evidence.
[31:31] Yes that's basically everything I wanted to demo today.

Learn more about the Databricks Data and AI platform.

The information provided herein is for general informational purposes only and may not reflect the most current product capabilities or configurations.