feat(spark_submit): extra_args kwarg for non-conf spark-submit flags
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.
This commit is contained in:
@@ -50,6 +50,82 @@ def test_build_command_uses_provided_master_and_deploy_mode():
|
||||
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",
|
||||
|
||||
Reference in New Issue
Block a user