ClaudeandClaude Fable 5
523e6a9c76
The MCP tool was named `generate_job_file` from Stage 2 but it does
NOT generate PySpark code — the calling LLM writes the code in its own
context, and this tool only persists it to a file under
SPARK_EXECUTOR_JOBS_DIR so `spark-submit` can see it. The misleading
`generate_` prefix sent agents (and humans) looking for a code
generator that doesn't exist.
This commit folds three related polish changes into one (split later
with rebase -i if you want them as separate history):
1. The rename itself:
- `tools/generate.py` → `tools/write_job.py`
- `generate_job_file` → `write_job_file`
- `GenerateJobFileRequest` → `WriteJobFileRequest`
- `/generate_job_file` route → `/write_job_file`
- `operation_id="generate_job_file"` → `operation_id="write_job_file"`
The internal helper `core.job_writer.write_job_file` (which just
writes bytes to disk with no SQL guard) is imported with an
`_write_to_disk` alias to avoid the name collision with the
MCP-exposed function in the same module.
The description for the tool now explicitly states 'this tool
does NOT generate PySpark code. The calling LLM is expected to
have already written the code; this tool only persists it.'
2. Skill for LLM agents operating the service
(`docs/superpowers/skills/spark-executor-mcp-operate/SKILL.md`,
449 lines). Covers the 16 tools, the two-step prepare/confirm
flow, the dual-ID contract (job_id vs application_id), the
PendingSubmission state machine, the Connection profile, the
job-file workflow, the error reference, common pitfalls, and a
full end-to-end word-count example.
3. Default `executor_memory` lowered 4G → 2G
(`_DEFAULTS_TO_CONFIRM` in `server.py`). Mirrors the matching
change in `test_mcp_routes.py` and the 5 unit tests that
reference the default. Aligns with the lighter workloads the
service is sized for in its current container profile.
Also tracked in git for the first time:
- `docs/superpowers/plans/2026-06-24-spark-executor-mcp.md`
(the original Stage 1/2/3 design plan, updated to use the new
tool name throughout).
Test rename:
- `tests/unit/test_generate_tool.py` → `test_write_job_tool.py`
- the new test file picks up an extra assertion that the SQL guard
rejects a `DROP TABLE` statement at write time.
243 tests pass (was 242; +1 new SQL-guard assertion). Zero regressions.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
214 lines
6.6 KiB
Python
214 lines
6.6 KiB
Python
# coding=utf-8
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from common import config
|
|
from spark_executor.core.spark_submit import (
|
|
SparkSubmitError,
|
|
build_spark_submit_command,
|
|
run_spark_submit,
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _restore_settings():
|
|
snapshot = config.Settings(
|
|
data_dir=config.settings.data_dir,
|
|
jobs_dir=config.settings.jobs_dir,
|
|
log_dir=config.settings.log_dir,
|
|
yarn_resource_manager_url=config.settings.yarn_resource_manager_url,
|
|
log_level=config.settings.log_level,
|
|
spark_submit_bin=config.settings.spark_submit_bin,
|
|
)
|
|
yield
|
|
config.settings.data_dir = snapshot.data_dir
|
|
config.settings.jobs_dir = snapshot.jobs_dir
|
|
config.settings.log_dir = snapshot.log_dir
|
|
config.settings.yarn_resource_manager_url = snapshot.yarn_resource_manager_url
|
|
config.settings.log_level = snapshot.log_level
|
|
config.settings.spark_submit_bin = snapshot.spark_submit_bin
|
|
|
|
|
|
def test_build_command_uses_provided_master_and_deploy_mode():
|
|
cmd = build_spark_submit_command(
|
|
master="yarn",
|
|
deploy_mode="cluster",
|
|
script_path="/tmp/jobs/job_001.py",
|
|
queue="default",
|
|
executor_memory="2G",
|
|
executor_cores=2,
|
|
num_executors=2,
|
|
)
|
|
assert cmd[:2] == ["spark-submit", "--master"] # default binary
|
|
assert "yarn" in cmd
|
|
assert "cluster" in cmd
|
|
assert "--queue" in cmd and "default" in cmd
|
|
assert "--executor-memory" in cmd and "2G" in cmd
|
|
assert "--executor-cores" in cmd and "2" in cmd
|
|
assert "--num-executors" in cmd and "2" in cmd
|
|
assert cmd[-1] == "/tmp/jobs/job_001.py"
|
|
|
|
|
|
def test_extra_args_none_default():
|
|
cmd = build_spark_submit_command(
|
|
master="yarn",
|
|
deploy_mode="cluster",
|
|
script_path="/tmp/jobs/job_001.py",
|
|
queue="default",
|
|
executor_memory="2G",
|
|
executor_cores=2,
|
|
num_executors=2,
|
|
)
|
|
# No extra flags between --conf block and script_path.
|
|
assert cmd[-1] == "/tmp/jobs/job_001.py"
|
|
|
|
|
|
def test_extra_args_single_pair():
|
|
cmd = build_spark_submit_command(
|
|
master="yarn",
|
|
deploy_mode="cluster",
|
|
script_path="/tmp/jobs/job_001.py",
|
|
queue="default",
|
|
executor_memory="2G",
|
|
executor_cores=2,
|
|
num_executors=2,
|
|
extra_args={"jars": "hdfs:///libs/foo.jar"},
|
|
)
|
|
script_idx = cmd.index("/tmp/jobs/job_001.py")
|
|
assert cmd[script_idx - 2 : script_idx] == [
|
|
"--jars",
|
|
"hdfs:///libs/foo.jar",
|
|
]
|
|
|
|
|
|
def test_extra_args_multiple_pairs_preserve_order():
|
|
cmd = build_spark_submit_command(
|
|
master="yarn",
|
|
deploy_mode="cluster",
|
|
script_path="/tmp/jobs/job_001.py",
|
|
queue="default",
|
|
executor_memory="2G",
|
|
executor_cores=2,
|
|
num_executors=2,
|
|
extra_args={
|
|
"driver-memory": "2G",
|
|
"name": "my-job",
|
|
"py-files": "hdfs:///libs/udf.py",
|
|
},
|
|
)
|
|
script_idx = cmd.index("/tmp/jobs/job_001.py")
|
|
extras = cmd[script_idx - 6 : script_idx]
|
|
assert extras == [
|
|
"--driver-memory", "2G",
|
|
"--name", "my-job",
|
|
"--py-files", "hdfs:///libs/udf.py",
|
|
]
|
|
|
|
|
|
def test_extra_args_does_not_collide_with_spark_conf():
|
|
cmd = build_spark_submit_command(
|
|
master="yarn",
|
|
deploy_mode="cluster",
|
|
script_path="/tmp/jobs/job_001.py",
|
|
queue="default",
|
|
executor_memory="2G",
|
|
executor_cores=2,
|
|
num_executors=2,
|
|
spark_conf={"spark.executor.memory": "8G"},
|
|
extra_args={"driver-memory": "2G"},
|
|
)
|
|
script_idx = cmd.index("/tmp/jobs/job_001.py")
|
|
# --conf entries come before --extra_args entries.
|
|
assert cmd[script_idx - 4 : script_idx] == [
|
|
"--conf", "spark.executor.memory=8G",
|
|
"--driver-memory", "2G",
|
|
]
|
|
|
|
|
|
def test_build_command_supports_standalone_master():
|
|
cmd = build_spark_submit_command(
|
|
master="spark://master:7077",
|
|
deploy_mode="client",
|
|
script_path="/tmp/j.py",
|
|
queue="default",
|
|
executor_memory="1G",
|
|
executor_cores=1,
|
|
num_executors=1,
|
|
)
|
|
assert "spark://master:7077" in cmd
|
|
assert "client" in cmd
|
|
|
|
|
|
def test_build_command_appends_spark_conf_entries():
|
|
cmd = build_spark_submit_command(
|
|
master="yarn",
|
|
deploy_mode="cluster",
|
|
script_path="/tmp/j.py",
|
|
queue="default",
|
|
executor_memory="2G",
|
|
executor_cores=2,
|
|
num_executors=2,
|
|
spark_conf={"spark.sql.shuffle.partitions": "200", "spark.executor.memoryOverhead": "1G"},
|
|
)
|
|
# every k=v is emitted as --conf k=v
|
|
assert "--conf" in cmd
|
|
assert "spark.sql.shuffle.partitions=200" in cmd
|
|
assert "spark.executor.memoryOverhead=1G" in cmd
|
|
# script_path is still last
|
|
assert cmd[-1] == "/tmp/j.py"
|
|
|
|
|
|
def test_build_command_uses_settings_spark_submit_bin():
|
|
"""The Spark CLI binary is env-configurable via settings.spark_submit_bin
|
|
(env var SPARK_EXECUTOR_SPARK_SUBMIT_BIN). Default is 'spark-submit';
|
|
override to 'spark2-submit', 'pyspark', etc. for mixed-version hosts."""
|
|
config.settings.spark_submit_bin = "spark2-submit"
|
|
cmd = build_spark_submit_command(
|
|
master="yarn",
|
|
deploy_mode="cluster",
|
|
script_path="/tmp/j.py",
|
|
queue="default",
|
|
executor_memory="2G",
|
|
executor_cores=2,
|
|
num_executors=2,
|
|
)
|
|
assert cmd[0] == "spark2-submit"
|
|
|
|
|
|
def test_build_command_supports_pyspark_binary():
|
|
config.settings.spark_submit_bin = "pyspark"
|
|
cmd = build_spark_submit_command(
|
|
master="yarn",
|
|
deploy_mode="cluster",
|
|
script_path="/tmp/j.py",
|
|
queue="default",
|
|
executor_memory="2G",
|
|
executor_cores=2,
|
|
num_executors=2,
|
|
)
|
|
assert cmd[0] == "pyspark"
|
|
|
|
|
|
def test_run_spark_submit_returns_completed_process(monkeypatch):
|
|
fake = MagicMock()
|
|
fake.returncode = 0
|
|
fake.stderr = "Submitted application application_1\n"
|
|
with patch("spark_executor.core.spark_submit.subprocess.run", return_value=fake) as m:
|
|
result = run_spark_submit(["spark-submit", "/tmp/x.py"])
|
|
assert result is fake
|
|
m.assert_called_once()
|
|
# capture_output=True and text=True must be set
|
|
kwargs = m.call_args.kwargs
|
|
assert kwargs["capture_output"] is True
|
|
assert kwargs["text"] is True
|
|
|
|
|
|
def test_run_spark_submit_raises_on_nonzero_return():
|
|
fake = MagicMock()
|
|
fake.returncode = 1
|
|
fake.stderr = "boom"
|
|
with patch("spark_executor.core.spark_submit.subprocess.run", return_value=fake):
|
|
with pytest.raises(SparkSubmitError, match="spark-submit failed"):
|
|
run_spark_submit(["spark-submit", "/tmp/x.py"])
|