fix(job_store): persist Jobs to disk and accept either ID in job tools

Two user-reported bugs, same root cause: the in-memory JobStore + the
'job_id must be the 12-char hex' tool contract.

Bug 1: 'Unknown job_id' reported frequently
  JobStore was a process-local dict (spark_executor/core/job_store.py).
  Under gunicorn workers > 1, a job created by confirm_submit_job
  landing on worker A was invisible to worker B, so a follow-up
  get_job_status / get_job_result / get_job_logs / kill_job landing on
  a different worker returned 'Unknown job_id'. Same multi-worker
  problem that bit the MCP session layer; only the affected data was
  different.

Bug 2: 'get_job_logs frequently confuses job_id and application_id'
  confirm_submit_job returns BOTH identifiers in SubmitResult, but
  get_job_logs (and friends) only accepted the local 12-char job_id
  and never said so in their description. When the agent passed the
  YARN application_id, the error message itself was misleading:
  'Unknown job_id: application_17400000001_0001' — the agent had
  passed an id, just the wrong kind.

This change fixes both at the root:

  * JobStore is now JSON-backed at data/jobs.json (atomic tempfile +
    os.replace), with cross-process safety via fcntl.flock on a sibling
    .lock file. Stage 3's SQLite migration is still planned; the file
    format is intentionally simple so it is a straight
    'for j in read_all(): db.insert(j)'.

  * New JobStore.get_either(uid) looks up by job_id first, then
    application_id. All four job-lifecycle tools (get_job_status,
    get_job_result, get_job_logs, kill_job) call get_either instead
    of get(job_id), so the agent can pass either identifier and get
    the same answer.

  * The 'neither matched' KeyError now spells out both id forms and
    what they look like, so the agent isn't left guessing.

  * server.py tool descriptions for the four job tools explicitly
    state 'job_id accepts BOTH identifiers' so this is visible to the
    LLM at tool-selection time, not only at error time.

Tests:
  * test_job_store.py: tmp_path isolation, persistence across
    instances, human-readable JSON, corrupt-file resilience,
    get_either (by job_id, by application_id, collision preference,
    unknown), put idempotency.
  * test_{logs,status,kill,result}_tool.py: per-test tmp_path fixture,
    'accepts application_id' regression for each tool, and an
    explicit assertion that the unknown-id error message mentions
    BOTH id forms. test_result_raises_keyerror_for_unknown_job's
    match pattern updated for the new message.

242 tests pass (was 226; +16 new). Zero regressions.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Claude
2026-06-29 18:57:31 +08:00
co-authored by Claude Fable 5
parent 565f6a9c9d
commit 0fef77c0b8
12 changed files with 541 additions and 42 deletions
+25 -3
View File
@@ -11,11 +11,33 @@ from spark_executor.tools.connections import store as conn_store
store = JobStore()
def _unknown_job_error(uid: str) -> KeyError:
"""Standard "we tried both IDs and found nothing" message.
The agent gets this from confirm_submit_job's response:
{"job_id": "a1b2c3d4e5f6", "application_id": "application_...", ...}
and routinely confuses which to pass here. Spelling out that BOTH
IDs were tried (and what they look like) saves a round trip.
"""
return KeyError(
f"No Job found for id={uid!r} (neither as job_id nor as "
f"application_id). Pass the job_id from confirm_submit_job's "
f"response — it is a 12-char hex like 'a1b2c3d4e5f6'. The "
f"application_id is the YARN ID, e.g. 'application_17400000001_0001'."
)
def get_job_logs(job_id: str, tail_chars: int = 5000) -> str:
"""Fetch aggregated container logs for a Spark job.
`job_id` accepts either the local job_id (returned by
confirm_submit_job) or the YARN application_id — both are looked up
against the same Job record.
"""
logger.debug(f"get_job_logs enter job_id={job_id} tail_chars={tail_chars}")
job = store.get(job_id)
job = store.get_either(job_id)
if job is None:
raise KeyError(f"Unknown job_id: {job_id}")
raise _unknown_job_error(job_id)
conn = conn_store.get(job.connection)
if conn is None:
raise KeyError(f"Connection not found: {job.connection}")
@@ -23,7 +45,7 @@ def get_job_logs(job_id: str, tail_chars: int = 5000) -> str:
full = get_application_logs(job.application_id, config)
tailed = full[-tail_chars:] if len(full) > tail_chars else full
logger.info(
f"get_job_logs ok job_id={job_id} application_id={job.application_id} "
f"get_job_logs ok job_id={job.job_id} application_id={job.application_id} "
f"full_chars={len(full)} returned_chars={len(tailed)}"
)
return tailed