diff --git a/spark_executor/core/spark_submit.py b/spark_executor/core/spark_submit.py index 4f374dc..bffab06 100644 --- a/spark_executor/core/spark_submit.py +++ b/spark_executor/core/spark_submit.py @@ -23,6 +23,7 @@ def build_spark_submit_command( executor_cores: int, num_executors: int, spark_conf: dict[str, str] | None = None, + extra_args: dict[str, str] | None = None, ) -> list[str]: # cmd[0] is the Spark CLI binary name, configurable via # SPARK_EXECUTOR_SPARK_SUBMIT_BIN. Defaults to 'spark-submit' but @@ -39,6 +40,8 @@ def build_spark_submit_command( ] for key, value in (spark_conf or {}).items(): cmd.extend(["--conf", f"{key}={value}"]) + for key, value in (extra_args or {}).items(): + cmd.extend([f"--{key}", str(value)]) cmd.append(script_path) logger.debug(f"build_spark_submit_command -> {cmd}") return cmd diff --git a/tests/unit/test_spark_submit.py b/tests/unit/test_spark_submit.py index 9890ba7..276ebec 100644 --- a/tests/unit/test_spark_submit.py +++ b/tests/unit/test_spark_submit.py @@ -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",