MLPlatformAgentBench is a benchmark for evaluating large language model coding agents on machine learning platform tasks. It measures how the interface given to the agent, a command-line interface (CLI) or a Python SDK, affects how well the agent completes those tasks.
This repository accompanies the KTH master's thesis MLPlatformAgentBench: Can coding agents optimize machine learning platform interfaces? The thesis is located in
docs/thesis/thesis.pdf.
Each task is run across a grid of conditions so the conditions can be compared directly.
- Platform. Hopsworks, Databricks, AWS SageMaker, Azure ML, GCP Vertex, and a
localbaseline (none). - Interface. The platform CLI or the Python SDK (or
none). - Skills. With or without a bundle of platform skills (Claude Code slash-commands and docs).
- Agent engine and model.
claude-*(Claude Code),gpt-*(Codex),mistral-*(Mistral Vibe).
The agent's work is graded by reading the result back through the platform's own data paths, not by inspecting the agent transcript. Tasks are generated as fresh, seeded instances of the Feature, Training, and Inference (FTI) lifecycle, plus Ops and Capstone tasks. Each run uses a new seed, so runs are reproducible but not identical. The agent runs in a sandbox confined to its run directory and is restricted to the one interface under test. Per-platform setup and teardown scripts provision and remove cloud resources around each run.
MLPAB is the testbed for a KTH master's thesis on how interface design and
supporting skills affect agent performance on ML platform operations. The
framework, configs, evals, and results live at the repository root, and the KTH
LaTeX thesis is in docs/thesis/ (docs/thesis/thesis.tex).
- A treatment config (
configs/treatments/<n>_<name>.yaml, a numbered flat file) expands into a grid of runs over models, interfaces, skills, and tasks. - A session-start check confirms that every platform credential, skill bundle, and model is available and responds to a live probe before any work begins.
- Each interface is built once and installed into a per-interface prepared virtual environment. Every run then clones that environment read-only, so runs do not reinstall anything or mutate shared state, and parallel sessions do not interfere.
- For each run: a fresh seeded task instance is generated,
setup.pyprovisions the platform and averifystep confirms it is ready, the sandboxed agent attempts the task using only its interface, thenteardown.pyremoves what the agent created and averifystep confirms nothing was left behind. - A platform adapter (
evals/adapters/<platform>.py) reads the deliverable back through the platform and runs the task's assertion suite. - One row per run is appended to
results/results.csv.
Every task has a documentation page under
docs/tasks/, auto-generated from the
task packages with make task-docs: the generator's design notes, the literal
seed-1 prompt, the staged files, the assert suite, and the diagnosed naive
variants all come straight from the code, so the pages cannot drift from the
tasks.
| Family | Tasks |
|---|---|
| feature | ingest, backfill, mit, validate, incremental_load, full_reload, pit (point-in-time correct), leakage |
| training | train, mdt, register, llm_finetuning |
| inference | batch, online, odt, skew, llm_serving, recsys, vector_search |
| ops | drift, prediction_monitoring, scheduled_jobs, alerting, lineage |
| capstone | ccfraud, airquality |
The committed treatment configs are the experiments behind the thesis results.
Each is a numbered flat yaml; its results live under results/<n>_<name>/.
| Treatments | Experiment |
|---|---|
| 1–4 | Hopsworks, full grid (CLI + SDK, with and without skills) across models: Opus, Mistral Large, Sonnet, Mistral Medium |
| 5–10 | Hopsworks CLI optimization arms opt1–opt6 (batch, session-reuse, compact-json, idempotent, quiet, stable-output), Opus, no skills |
| 11 | Hopsworks CLI with an AI-optimized skills bundle, Opus |
| 12–17 | Hopsworks CLI optimization arms opt1–opt6 combined with the optimized skills bundle, Opus |
| 18–21 | Databricks, full grid across models: Opus, Sonnet, Mistral Large, Mistral Medium |
| 22 | GCP Vertex, full grid, Opus |
| 23–24 | Hopsworks and Databricks, full grid, Fable |
make install # create .venv, install mlpab + evals + dev tools, link `mlpab` onto PATH
make setup # interactive: authenticate agent engine(s) and set up platform(s)
# Verify a config is runnable: platform reachable, and each model answers a live probe
make check CONFIG=configs/treatments/1_hw-full-cli-sdk-skills-opus.yaml
# Run a treatment session in tmux, detached from any terminal
mlpab start configs/treatments/18_db-full-cli-sdk-skills-opus.yaml
mlpab status # list running sessions
mlpab attach <config.yaml> # watch live (detach with Ctrl-b d)
# Or run inline
mlpab run configs/treatments/22_gcp-full-cli-sdk-skills-opus.yaml
# Resume: re-running a config skips combos already completed — a valid row whose
# run folder still exists on disk. --retry also re-runs FAILED combos (no valid
# row, or a valid row whose folder was deleted), purging their rows + folders
# and running clean; --no-skip re-runs every combo, accumulating attempts.
mlpab run --retry configs/treatments/18_db-full-cli-sdk-skills-opus.yamlCredentials are read from .env. macOS is the primary target. It uses
APFS copy-on-write clones for the per-run virtual environments and the macOS
Keychain for the Claude OAuth token. Other systems fall back to plain copies.
make test # unit tests
make lint # ruff check plus isort and format checks (no changes)
make fmt # auto-fix: isort imports, then ruff format, then ruff --fixresults/results.csv one row per run (the single results table)
results/<config>/... per-run artifacts: agent.log, submission/, grading, commands
results/*.ipynb hand-curated analysis notebooks
Each row records the identity of the cell (model, platform, interface,
version, skills, category, task, n) and its outcome (valid,
success, asserts_passed, asserts_total), along with cost, latency, and the
agent's tool-call counts. This is enough to reproduce and compare any value in
the table.
A platform lives in configs/platforms/<platform>/. It has one flat manifest per
interface (cli.yaml, sdk.yaml) describing how to build, install,
authenticate, and test that interface. It also has setup.py and teardown.py
(each supporting a verify mode) and a skills/ bundle. A matching grader
adapter goes in evals/adapters/<platform>.py and implements the read-back
contract (get_feature_table, read_rows, read_training_dataset, and the
state reads).
