Claude
5566d55a7c
The function had hardcoded --master / --deploy-mode / --queue /
--executor-memory / --executor-cores / --num-executors plus a
`spark_conf` dict for --conf, with no way to pass other flags like
--jars, --py-files, --files, --driver-memory, --name, etc.
Add `extra_args: dict[str, str] | None = None` that emits `--{key}
{value}` pairs, placed after the --conf loop and before the script
path. Default None preserves the existing cmd shape exactly.
Structured (dict) instead of raw list[str] so LLM agents writing
tool-call payloads can't accidentally pack flag+value into a single
string. Key order in the dict is preserved.
Tests: 4 cases — None default, single pair, multiple pairs in order,
no collision with spark_conf. uv run pytest -> 175 passed.
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="4G",
|
|
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 "4G" 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="4G",
|
|
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="4G",
|
|
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="4G",
|
|
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="4G",
|
|
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="4G",
|
|
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="4G",
|
|
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="4G",
|
|
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"])
|