diff --git a/spark_executor/models.py b/spark_executor/models.py index b0e8427..8291609 100644 --- a/spark_executor/models.py +++ b/spark_executor/models.py @@ -86,7 +86,7 @@ class Connection(BaseModel): class PendingSubmission(BaseModel): pending_id: str - app_name: str | None = None + app_name: str connection: str master: str deploy_mode: str @@ -97,6 +97,7 @@ class PendingSubmission(BaseModel): executor_cores: int num_executors: int spark_conf: dict[str, str] = Field(default_factory=dict) + extra_args: dict[str, str] = Field(default_factory=dict) created_at: datetime status: str = "PENDING" # PENDING | SUBMITTED | CANCELLED | FAILED error: str | None = None diff --git a/spark_executor/tools/requests.py b/spark_executor/tools/requests.py index e316981..8f342af 100644 --- a/spark_executor/tools/requests.py +++ b/spark_executor/tools/requests.py @@ -33,9 +33,9 @@ class SaveConnectionRequest(BaseModel): class PrepareSubmitJobRequest(BaseModel): connection: str - app_name: str | None = Field( - default=None, - description="Optional human-readable application name for tracking the pending submission.", + app_name: str = Field( + ..., + description="Human-readable application name for tracking the pending submission.", ) script_path: str = Field( ..., @@ -49,10 +49,26 @@ class PrepareSubmitJobRequest(BaseModel): "if the path is missing or not a file." ), ) - queue: str = "default" - executor_memory: str = "4G" - executor_cores: int = 2 - num_executors: int = 2 + queue: str = Field( + ..., + description="YARN queue to submit to. Must be explicitly confirmed by the caller.", + ) + executor_memory: str = Field( + ..., + description="Executor memory, e.g. '4G'. Must be explicitly confirmed by the caller.", + ) + executor_cores: int = Field( + ..., + description="Number of cores per executor. Must be explicitly confirmed by the caller.", + ) + num_executors: int = Field( + ..., + description="Total number of executors. Must be explicitly confirmed by the caller.", + ) + extra_args: dict[str, str] | None = Field( + default=None, + description="Additional spark-submit flags (e.g. jars, py-files) confirmed at prepare time.", + ) class PendingIdRequest(BaseModel): diff --git a/spark_executor/tools/submit.py b/spark_executor/tools/submit.py index f1c8b2a..5fb33b1 100644 --- a/spark_executor/tools/submit.py +++ b/spark_executor/tools/submit.py @@ -64,17 +64,18 @@ def prepare_submit_job( *, connection: str, script_path: str, - queue: str = "default", - executor_memory: str = "4G", - executor_cores: int = 2, - num_executors: int = 2, - app_name: str | None = None, + queue: str, + executor_memory: str, + executor_cores: int, + num_executors: int, + app_name: str, + extra_args: dict[str, str] | None = None, ) -> dict[str, object]: """Snapshot connection params and persist a PendingSubmission. Does NOT submit.""" logger.debug( f"prepare_submit_job enter connection={connection} script_path={script_path} " f"queue={queue} executor_memory={executor_memory} executor_cores={executor_cores} " - f"num_executors={num_executors}" + f"num_executors={num_executors} app_name={app_name}" ) # Order of checks matters for the error the agent sees: # 1. Unknown connection -> 404 (KeyError -> 404 in server.py) @@ -121,6 +122,7 @@ def prepare_submit_job( executor_cores=executor_cores, num_executors=num_executors, spark_conf=dict(conn.spark_conf), + extra_args=dict(extra_args or {}), created_at=datetime.utcnow(), status="PENDING", ) @@ -164,6 +166,7 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult: executor_cores=pending.executor_cores, num_executors=pending.num_executors, spark_conf=pending.spark_conf, + extra_args=pending.extra_args, ) logger.info( f"confirm_submit_job start pending_id={pending_id} " diff --git a/tests/integration/test_mcp_routes.py b/tests/integration/test_mcp_routes.py index 973b4e2..c49a77e 100644 --- a/tests/integration/test_mcp_routes.py +++ b/tests/integration/test_mcp_routes.py @@ -86,7 +86,15 @@ def test_prepare_submit_job_works_via_body(tmp_path): script.write_text("print('hi')\n") r = c.post( "/prepare_submit_job", - json={"connection": "prod", "script_path": str(script), "queue": "research"}, + json={ + "connection": "prod", + "script_path": str(script), + "queue": "research", + "executor_memory": "4G", + "executor_cores": 2, + "num_executors": 2, + "app_name": "test-app", + }, ) assert r.status_code == 200, r.text body = r.json() @@ -102,7 +110,15 @@ def test_prepare_rejects_nonexistent_script_path_with_400(): c.post("/save_connection", json={"name": "prod", "master": "yarn"}) r = c.post( "/prepare_submit_job", - json={"connection": "prod", "script_path": "/nope/does_not_exist.py"}, + json={ + "connection": "prod", + "script_path": "/nope/does_not_exist.py", + "queue": "default", + "executor_memory": "4G", + "executor_cores": 2, + "num_executors": 2, + "app_name": "test-app", + }, ) assert r.status_code == 400 detail = r.json()["detail"] @@ -121,7 +137,15 @@ def test_prepare_rejects_sql_policy_violation_with_400(tmp_path): bad_script.write_text('spark.sql("DROP TABLE users")\n') r = c.post( "/prepare_submit_job", - json={"connection": "prod", "script_path": str(bad_script)}, + json={ + "connection": "prod", + "script_path": str(bad_script), + "queue": "default", + "executor_memory": "4G", + "executor_cores": 2, + "num_executors": 2, + "app_name": "test-app", + }, ) assert r.status_code == 400 detail = r.json()["detail"] @@ -199,7 +223,15 @@ def test_list_and_get_pending_job_roundtrip_via_body(tmp_path): script.write_text("print('hi')\n") prep = c.post( "/prepare_submit_job", - json={"connection": "prod", "script_path": str(script)}, + json={ + "connection": "prod", + "script_path": str(script), + "queue": "default", + "executor_memory": "4G", + "executor_cores": 2, + "num_executors": 2, + "app_name": "test-app", + }, ).json() pid = prep["pending_id"] @@ -261,7 +293,15 @@ def test_unknown_connection_in_prepare_returns_404(): c = TestClient(app) r = c.post( "/prepare_submit_job", - json={"connection": "nope", "script_path": "/tmp/x.py"}, + json={ + "connection": "nope", + "script_path": "/tmp/x.py", + "queue": "default", + "executor_memory": "4G", + "executor_cores": 2, + "num_executors": 2, + "app_name": "test-app", + }, ) assert r.status_code == 404 assert "nope" in r.json()["detail"] @@ -275,7 +315,15 @@ def test_confirm_non_pending_returns_400(tmp_path): script.write_text("print('hi')\n") prep = c.post( "/prepare_submit_job", - json={"connection": "prod", "script_path": str(script)}, + json={ + "connection": "prod", + "script_path": str(script), + "queue": "default", + "executor_memory": "4G", + "executor_cores": 2, + "num_executors": 2, + "app_name": "test-app", + }, ).json() pid = prep["pending_id"] c.post("/cancel_pending_job", json={"pending_id": pid}) @@ -292,7 +340,15 @@ def test_cancel_already_submitted_returns_400(tmp_path): script.write_text("print('hi')\n") prep = c.post( "/prepare_submit_job", - json={"connection": "prod", "script_path": str(script)}, + json={ + "connection": "prod", + "script_path": str(script), + "queue": "default", + "executor_memory": "4G", + "executor_cores": 2, + "num_executors": 2, + "app_name": "test-app", + }, ).json() pid = prep["pending_id"] # Simulate a SUBMITTED state by mutating the pending directly diff --git a/tests/unit/test_logging.py b/tests/unit/test_logging.py index 9913abe..2d73a23 100644 --- a/tests/unit/test_logging.py +++ b/tests/unit/test_logging.py @@ -47,7 +47,15 @@ def test_prepare_submit_job_emits_debug_and_info(log_capture, tmp_path): script = tmp_path / "demo.py" script.write_text("print('hi')\n") log_capture.truncate(0); log_capture.seek(0) - submit.prepare_submit_job(connection="prod", script_path=str(script), queue="research") + submit.prepare_submit_job( + connection="prod", + script_path=str(script), + queue="research", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) text = log_capture.getvalue() assert "DEBUG|prepare_submit_job enter" in text assert "INFO|prepare_submit_job ok" in text diff --git a/tests/unit/test_models.py b/tests/unit/test_models.py index 52ed2b0..cab35a4 100644 --- a/tests/unit/test_models.py +++ b/tests/unit/test_models.py @@ -89,6 +89,7 @@ def test_connection_rejects_bad_master(): def test_pending_submission_defaults_to_pending_status(): p = PendingSubmission( pending_id="p_1", + app_name="test-app", connection="prod", master="yarn", deploy_mode="cluster", @@ -105,13 +106,14 @@ def test_pending_submission_defaults_to_pending_status(): assert p.job_id is None assert p.application_id is None assert p.yarn_rm_url is None # default for connections without one set - assert p.app_name is None + assert p.app_name == "test-app" def test_pending_submission_carries_yarn_rm_url_snapshot(): """snapshotting yarn_rm_url lets prepare/confirm survive connection edits.""" p = PendingSubmission( pending_id="p_x", + app_name="test-app", connection="prod", master="yarn", deploy_mode="cluster", @@ -122,6 +124,7 @@ def test_pending_submission_carries_yarn_rm_url_snapshot(): executor_cores=2, num_executors=2, spark_conf={}, + extra_args={}, created_at=datetime(2026, 6, 24), ) assert p.yarn_rm_url == "http://rm-prod:8088" @@ -130,6 +133,7 @@ def test_pending_submission_carries_yarn_rm_url_snapshot(): def test_pending_submission_can_record_outcome(): p = PendingSubmission( pending_id="p_2", + app_name="test-app", connection="prod", master="yarn", deploy_mode="cluster", @@ -139,6 +143,7 @@ def test_pending_submission_can_record_outcome(): executor_cores=2, num_executors=2, spark_conf={}, + extra_args={}, created_at=datetime(2026, 6, 24), status="SUBMITTED", job_id="j_abc", diff --git a/tests/unit/test_pending_store.py b/tests/unit/test_pending_store.py index 3aedaf7..73c3b0f 100644 --- a/tests/unit/test_pending_store.py +++ b/tests/unit/test_pending_store.py @@ -11,6 +11,7 @@ from spark_executor.models import PendingSubmission def _pending(pid: str = "p_1", created_at: datetime | None = None) -> PendingSubmission: return PendingSubmission( pending_id=pid, + app_name="test", connection="prod", master="yarn", deploy_mode="cluster", @@ -20,6 +21,7 @@ def _pending(pid: str = "p_1", created_at: datetime | None = None) -> PendingSub executor_cores=2, num_executors=2, spark_conf={}, + extra_args={}, created_at=created_at or datetime(2026, 6, 24), ) @@ -80,10 +82,10 @@ def test_delete_removes_from_date_file_and_deletes_empty_file(): def _legacy_record_text() -> str: return ( - '{"p_legacy": {"pending_id": "p_legacy", "connection": "prod", ' - '"master": "yarn", "deploy_mode": "cluster", "script_path": "/tmp/j.py", ' - '"queue": "default", "executor_memory": "4G", "executor_cores": 2, ' - '"num_executors": 2, "spark_conf": {}, ' + '{"p_legacy": {"pending_id": "p_legacy", "app_name": "legacy", ' + '"connection": "prod", "master": "yarn", "deploy_mode": "cluster", ' + '"script_path": "/tmp/j.py", "queue": "default", "executor_memory": "4G", ' + '"executor_cores": 2, "num_executors": 2, "spark_conf": {}, ' '"created_at": "2026-06-23T00:00:00"}}' ) diff --git a/tests/unit/test_requests.py b/tests/unit/test_requests.py new file mode 100644 index 0000000..1e1f51e --- /dev/null +++ b/tests/unit/test_requests.py @@ -0,0 +1,23 @@ +# coding=utf-8 +import pytest + +from spark_executor.tools.requests import PrepareSubmitJobRequest + + +def test_prepare_submit_job_request_requires_all_confirmed_parameters(): + with pytest.raises(ValueError): + PrepareSubmitJobRequest(connection="prod", script_path="/tmp/x.py") + + +def test_prepare_submit_job_request_accepts_extra_args(): + req = PrepareSubmitJobRequest( + connection="prod", + script_path="/tmp/x.py", + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + extra_args={"jars": "hdfs:///lib/foo.jar"}, + ) + assert req.extra_args == {"jars": "hdfs:///lib/foo.jar"} diff --git a/tests/unit/test_submit_tool.py b/tests/unit/test_submit_tool.py index f5349dd..1a87e2e 100644 --- a/tests/unit/test_submit_tool.py +++ b/tests/unit/test_submit_tool.py @@ -45,7 +45,15 @@ def _last_pending_id() -> str: def test_prepare_does_not_invoke_spark_submit(monkeypatch, real_script): called = [] monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: called.append(cmd)) - out = submit.prepare_submit_job(connection="prod", script_path=str(real_script)) + out = submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) assert called == [] # never called assert out["status"] == "PENDING" assert out["pending_id"].startswith("p_") @@ -62,7 +70,15 @@ def test_prepare_persists_pending_with_snapshot(monkeypatch, real_script): spark_conf={"spark.sql.shuffle.partitions": "200"}, ) ) - submit.prepare_submit_job(connection="prod", script_path=str(real_script), queue="research") + submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="research", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) p = submit.pending_store.get(_last_pending_id()) assert p is not None assert p.connection == "prod" @@ -77,24 +93,65 @@ def test_prepare_persists_pending_with_snapshot(monkeypatch, real_script): def test_prepare_persists_app_name(monkeypatch, real_script): monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) submit.prepare_submit_job( - connection="prod", script_path=str(real_script), app_name="my-etl-job" + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="my-etl-job", ) p = submit.pending_store.get(_last_pending_id()) assert p is not None assert p.app_name == "my-etl-job" -def test_prepare_app_name_defaults_to_none(monkeypatch, real_script): +def test_prepare_requires_explicit_app_name(monkeypatch, real_script): monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) - submit.prepare_submit_job(connection="prod", script_path=str(real_script)) - p = submit.pending_store.get(_last_pending_id()) + with pytest.raises(TypeError, match="app_name"): + submit.prepare_submit_job(connection="prod", script_path=str(real_script)) + + +def test_prepare_snapshots_extra_args(monkeypatch, real_script): + monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) + submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + extra_args={"jars": "hdfs:///lib/foo.jar"}, + ) + pid = _last_pending_id() + p = submit.pending_store.get(pid) assert p is not None - assert p.app_name is None + assert p.extra_args == {"jars": "hdfs:///lib/foo.jar"} + + fake_proc = type("P", (), { + "returncode": 0, + "stderr": "tracking URL: http://rm:8088/proxy/application_17400000002/\n", + })() + with patch("spark_executor.tools.submit.run_spark_submit", return_value=fake_proc) as m: + submit.confirm_submit_job(pending_id=pid) + cmd = m.call_args.args[0] + assert ["--jars", "hdfs:///lib/foo.jar"] in [ + cmd[i : i + 2] for i in range(len(cmd) - 1) + ] def test_prepare_raises_for_unknown_connection(real_script): with pytest.raises(KeyError, match="missing"): - submit.prepare_submit_job(connection="missing", script_path=str(real_script)) + submit.prepare_submit_job( + connection="missing", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) # --- script_path validation --- @@ -102,25 +159,55 @@ def test_prepare_raises_for_unknown_connection(real_script): def test_prepare_rejects_nonexistent_script_path(tmp_path): missing = str(tmp_path / "does_not_exist.py") with pytest.raises(ValueError, match="does not exist or is not a file"): - submit.prepare_submit_job(connection="prod", script_path=missing) + submit.prepare_submit_job( + connection="prod", + script_path=missing, + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) def test_prepare_rejects_directory_as_script_path(tmp_path): """A directory is not a file, even if it exists.""" with pytest.raises(ValueError, match="does not exist or is not a file"): - submit.prepare_submit_job(connection="prod", script_path=str(tmp_path)) + submit.prepare_submit_job( + connection="prod", + script_path=str(tmp_path), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) def test_prepare_rejects_empty_script_path(): with pytest.raises(ValueError, match="does not exist or is not a file"): - submit.prepare_submit_job(connection="prod", script_path="") + submit.prepare_submit_job( + connection="prod", + script_path="", + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) def test_prepare_error_message_mentions_generate_job_file(tmp_path): """The agent must be told to call generate_job_file first.""" with pytest.raises(ValueError, match="generate_job_file"): submit.prepare_submit_job( - connection="prod", script_path=str(tmp_path / "x.py") + connection="prod", + script_path=str(tmp_path / "x.py"), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", ) @@ -128,7 +215,15 @@ def test_confirm_rejects_if_script_was_deleted_after_prepare(tmp_path, monkeypat """Defense in depth: file present at prepare, gone by confirm -> 400.""" script = tmp_path / "demo.py" script.write_text("print('hi')\n") - submit.prepare_submit_job(connection="prod", script_path=str(script)) + submit.prepare_submit_job( + connection="prod", + script_path=str(script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) pid = _last_pending_id() # Simulate the file being deleted (e.g. cleanup job) between prepare and confirm script.unlink() @@ -139,7 +234,15 @@ def test_confirm_rejects_if_script_was_deleted_after_prepare(tmp_path, monkeypat def test_prepare_snapshots_connection_at_prepare_time(monkeypatch, real_script): """Editing the connection between prepare and confirm must NOT silently retarget.""" monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) - submit.prepare_submit_job(connection="prod", script_path=str(real_script)) + submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) p = submit.pending_store.get(_last_pending_id()) # Now mutate the connection connection_store.store.save( @@ -154,7 +257,15 @@ def test_prepare_snapshots_connection_at_prepare_time(monkeypatch, real_script): # --- confirm_submit_job --- def test_confirm_invokes_spark_submit_and_marks_submitted(monkeypatch, real_script): - submit.prepare_submit_job(connection="prod", script_path=str(real_script)) + submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) pid = _last_pending_id() fake_proc = type("P", (), { "returncode": 0, @@ -182,7 +293,15 @@ def test_confirm_raises_for_unknown_pending_id(): def test_confirm_refuses_non_pending_status(monkeypatch, real_script): - submit.prepare_submit_job(connection="prod", script_path=str(real_script)) + submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) pid = _last_pending_id() # Mark it CANCELLED first p = submit.pending_store.get(pid) @@ -194,7 +313,15 @@ def test_confirm_refuses_non_pending_status(monkeypatch, real_script): def test_confirm_marks_failed_on_spark_submit_error(monkeypatch, real_script): from spark_executor.core.spark_submit import SparkSubmitError - submit.prepare_submit_job(connection="prod", script_path=str(real_script)) + submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) pid = _last_pending_id() def _raise(_cmd): raise SparkSubmitError("boom") @@ -217,8 +344,24 @@ def test_list_pending_jobs_returns_all(monkeypatch, tmp_path): b = tmp_path / "b.py" a.write_text("a\n"); b.write_text("b\n") monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) - submit.prepare_submit_job(connection="prod", script_path=str(a)) - submit.prepare_submit_job(connection="prod", script_path=str(b)) + submit.prepare_submit_job( + connection="prod", + script_path=str(a), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) + submit.prepare_submit_job( + connection="prod", + script_path=str(b), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) out = submit.list_pending_jobs() assert {p["script_path"] for p in out} == {str(a), str(b)} assert all(p["status"] == "PENDING" for p in out) @@ -228,7 +371,15 @@ def test_list_pending_jobs_returns_all(monkeypatch, tmp_path): def test_get_pending_job_returns_dump(monkeypatch, real_script): monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) - submit.prepare_submit_job(connection="prod", script_path=str(real_script)) + submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) pid = _last_pending_id() out = submit.get_pending_job(pid) assert out["pending_id"] == pid @@ -245,7 +396,15 @@ def test_get_pending_job_unknown_raises(): def test_cancel_pending_job_marks_cancelled(monkeypatch, real_script): monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) - submit.prepare_submit_job(connection="prod", script_path=str(real_script)) + submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) pid = _last_pending_id() out = submit.cancel_pending_job(pid) assert out == {"pending_id": pid, "status": "CANCELLED"} @@ -258,7 +417,15 @@ def test_cancel_pending_job_unknown_raises(): def test_cancel_pending_job_refuses_submitted(monkeypatch, real_script): - submit.prepare_submit_job(connection="prod", script_path=str(real_script)) + submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) pid = _last_pending_id() p = submit.pending_store.get(pid) p.status = "SUBMITTED"