Testing Argon on
Dask Shuffle Recovery

I built an automated setup to test Argon on Dask shuffle recovery. Each attempt ran in a separate Docker container, an isolated environment for running programs. A set of tests hidden from Argon checked its code afterward. A fresh attempt finished in 57 minutes 3 seconds and passed all 15 tests of the task requirements, 15 selected checks for existing Dask behavior, and two checks of the new setting that controls how many times Dask can restart a failing shuffle.

The project had two questions: could Argon build the recovery behavior we requested, and could our tests reliably catch mistakes? To answer them, we defined the task and built a reference solution to show how it could be solved. We also built a grader, the program that runs the tests and scores submitted code. We checked it using bugs we deliberately introduced, then used it to evaluate Argon’s changes.

A realistic engineering task

Polaris’s description of carefully constructed software engineering evaluations prompted me to build one myself. I wanted a task based on a specific failure, involving parts of a system that need to work together, with tests that could check the resulting code.

I chose Dask Distributed, using a fixed version of its code and a recovery scenario motivated by issue 8595. Dask already retries failed communications and can recover when workers disappear. Our task proposed an additional recovery feature that stays off unless someone enables it.

Dask splits large data jobs across workers, programs that handle parts of the work and can run on different machines. Sometimes those workers need to exchange data so the right rows end up together. That exchange is called a shuffle. A coordination step, called a barrier, checks that the transfers have finished before the results are used.

In our scenario, one worker is still running and can do other work, but keeps failing a communication step during the shuffle. Dask retries that step until it reaches its retry limit. We asked Argon to add a way to restart the shuffle on other workers, recreating the input data as needed. The faulty worker should stay out of that shuffle, even if some of the shuffle’s results later need to be recreated. It should remain available for unrelated work.

We asked Argon to add a setting that limits how many times the shuffle can restart, with the new recovery feature off by default. The code also had to keep data in the correct output groups, ignore outdated messages from earlier attempts, stop when the job was cancelled, and remove temporary data and recovery records when no longer needed.

Across shuffle restarts, the code had to keep track of how many restarts had been used and which workers should stay excluded. A worker having trouble with this shuffle could still do other work. The fix had to leave it out of that shuffle while allowing it to handle unrelated tasks.

Testing beyond the first recovery

One important test introduced another problem after the shuffle had recovered from the first one. It checked this sequence:

1. Worker A repeatedly fails a communication step during the shuffle. It is still running.
2. The shuffle recovers without A. The application receives the correct result and still needs Dask to keep it available.
3. Worker B, holding part of the result, disappears. Dask must rerun the necessary work to recreate the missing data.
4. A must stay excluded while Dask recreates the missing data. The test checks both the returned data and whether Dask contacts A again.

Auditing the grader

The initial trial passed all 12 tests we used to check the solution. Before finalizing the tests for the next attempts, I reviewed what they checked and what they might miss.

We added three checks: whether each row ended up in the correct output group, whether faulty workers stayed excluded when missing results had to be recreated, and whether disk errors caused the shuffle to fail instead of triggering this recovery. All three behaviors were already required by the task.

We checked the grader using our corrected reference solution, Argon’s code from the initial trial, and code we deliberately modified to introduce bugs. The reference solution and Argon’s initial solution both passed all 15 tests. One version we intentionally broke sent rows to the wrong output groups. The grader caught that mistake and gave it 90/100, which failed the requirement to pass every mandatory check.

Argon’s code from the initial trial passed the three additional checks without any changes. We finalized this 15-test version of the grader before running further attempts and kept those grading rules unchanged.

Running the evaluation with Docker

The evaluation had three parts: the Dask task, a script that starts and records each attempt, and a separate grader. We used a fixed version of Dask and the software packages it needed, wrote the task instructions and reference solution, and checked that the grader could detect mistakes.

For each independent attempt, our script created a new Docker container with the prepared Dask code and no saved agent history. It checked that the code was unchanged and the access restrictions were working, then started Gemini CLI with the saved task instructions. Each container was limited to three CPUs and 3 GiB of memory, with a 60-minute time limit.

Gemini CLI is Google’s command-line coding agent. Argon chose which files to inspect and which tools to use. The CLI carried out permitted reads, edits, and tests inside Docker, then sent the results back to the model so it could choose its next action. Argon itself ran remotely through the Gemini API. Our scripts controlled the environment, access restrictions, time limits, recording, and separate grading.

The container could not access files on our computer, the hidden grader, the reference solution, previous submissions, or previous conversations. Its external network access was limited to the Gemini API. Argon could read Dask’s existing tests and write its own, but it received no results from our hidden tests while working.

Our script saved the conversation, tool activity, and how the attempt ended. It then collected the code changes into a patch, a file recording what was added, removed, or modified. We applied those changes to a clean copy of the original Dask code and ran our tests in a separate Docker container with no network access or model credentials. Those tests determined the score.

We selected gemini-4-argon as the main model through Gemini CLI 0.64.0. We used Dask code version (Git commit) dc182bda54af488d029f2c40db0893d5f48e69d5, Python 3.12, and dask[dataframe]==2026.8.0.

Argon completed the task

The fresh attempt finished normally in 57 minutes 3 seconds, within the 60-minute limit. Argon made 159 tool calls through Gemini CLI, such as reading files, editing code, and running commands. It had no access to our hidden tests, reference solution, or previous attempts, and received no results from the hidden tests while working.

After the attempt ended, we tested its code separately. It scored 100/100 and passed every required check:

The completed solution kept Worker A excluded when Worker B disappeared and Dask recreated missing results. It also returned the correct data. The API connection remained working throughout the attempt, and no session restart was needed.

What the results show

Argon, running through Gemini CLI, completed a solution within the one-hour limit and passed every required check in our finalized evaluation. We also verified that the grader caught bugs we deliberately introduced into the code.

These results cover one task in a controlled test environment, with all workers running inside a single Python program. We tested Dask’s peer-to-peer (P2P) shuffle for table-like data on its own, with workers exchanging rows through real TCP network connections and Dask able to recreate the input data when needed.

The useful engineering lesson is to test what happens after the first successful recovery. Results may need to be recreated, messages may arrive late, and jobs may be cancelled. A recovery feature needs to handle those events as well as return the correct data.

This is the work I want to do more of: turn realistic failure scenarios into clear requirements, build tests that catch mistakes, and provide enough evidence for others to inspect the results. I would welcome feedback on cases the evaluation should cover or weaknesses in how it grades solutions.