tao.chen/mcp-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
spark_executor/core/job_store.py
+171 -7
View File
@@ -2,27 +2,191 @@
"""
@Time :2026/6/24
@Author :tao.chen
JSON-backed Job registry. Each Job carries both a local `job_id` (12-char
hex generated in confirm_submit_job) and the YARN `application_id`
(`application_<ts>_<n>`).
Why file-backed instead of in-memory:
- The MCP service runs under gunicorn with workers > 1 by default; an
in-memory dict lives in worker A's process, so a follow-up
`get_job_logs(job_id=...)` landing on worker B would return
"Unknown job_id" for jobs created elsewhere. The same problem that
bit us on MCP sessions.
- Restarting the server should not lose track of running YARN apps.
Why fcntl.flock on writes:
- In-process `threading.Lock` is per-process. Two gunicorn workers can
call `put()` concurrently and lose updates (last-writer-wins over
the other's data). `fcntl.flock(LOCK_EX)` on the data file gives
cross-process mutual exclusion.
Why `get_either` (accept job_id OR application_id):
- The MCP client gets both IDs back from `confirm_submit_job` and
routinely confuses which to pass to `get_job_logs` / `get_job_status`
/ `kill_job` / `get_job_result`. The tool contract used to be
"must be job_id" but the description never said so, so agents
passed application_id and got a misleading "Unknown job_id" error.
Accepting either ID lets us blame the lookup, not the user.
This is a Step-1 persistence fix; the Stage-3 plan still calls for
SQLite (queryable, transactional). The file format is intentionally
simple so a future migration is just `for j in read_all(): db.insert(j)`.
"""
import fcntl
import json
import os
import tempfile
from contextlib import contextmanager
from pathlib import Path
from threading import Lock
from common.config import settings
from common.logging import logger
from spark_executor.models import Job
DEFAULT_DATA_DIR = settings.data_dir
DEFAULT_FILE_NAME = "jobs.json"
class JobStore:
"""In-memory job registry. Stage-3 will swap this for SQLite."""
"""JSON-backed CRUD for Job records, indexed by job_id.
def __init__(self) -> None:
File path: <data_dir>/<file_name>. Format: {job_id: <Job.model_dump()>}.
Writes are atomic (tempfile + os.replace) and cross-process safe
(fcntl.flock around the load-modify-dump critical section).
"""
def __init__(self, data_dir: str | None = None, file_name: str = DEFAULT_FILE_NAME) -> None:
self._data_dir = data_dir or DEFAULT_DATA_DIR
self._file_name = file_name
# In-process lock: serializes threads within one worker.
# Cross-process serialization is handled by fcntl.flock in _locked_dump.
self._lock = Lock()
self._jobs: dict[str, Job] = {}
@property
def path(self) -> Path:
return Path(self._data_dir) / self._file_name
# --- Low-level I/O ---
def _load(self) -> dict[str, Job]:
"""Read all jobs from disk. Returns {} on missing/corrupt file."""
if not self.path.exists():
logger.debug(f"JobStore._load: file {self.path} absent, returning empty")
return {}
try:
raw = json.loads(self.path.read_text(encoding="utf-8"))
except json.JSONDecodeError:
logger.exception(
f"JobStore._load: file {self.path} is corrupt; treating as empty. "
f"Inspect and either repair or delete to recover."
)
return {}
return {jid: Job.model_validate(j) for jid, j in raw.items()}
def _dump(self, records: dict[str, Job]) -> None:
"""Atomic write. Caller must hold `self._lock`."""
os.makedirs(self._data_dir, exist_ok=True)
payload = {jid: j.model_dump() for jid, j in records.items()}
fd, tmp_path = tempfile.mkstemp(prefix=self._file_name + ".", dir=self._data_dir)
try:
with os.fdopen(fd, "w", encoding="utf-8") as f:
# default=str handles datetime — same trick pending_store uses.
json.dump(payload, f, indent=2, ensure_ascii=False, default=str)
os.replace(tmp_path, self.path)
logger.debug(f"JobStore._dump: wrote {len(records)} records to {self.path}")
except Exception:
if os.path.exists(tmp_path):
os.unlink(tmp_path)
raise
@contextmanager
def _locked_dump(self, mutate):
"""load -> mutate (in-memory) -> dump, under both in-process and
cross-process exclusive locks.
This is the only sanctioned way to write. It guarantees:
1. Two threads in the same worker see consistent state.
2. Two gunicorn workers can't both load X, both compute
X+their-change, both dump — one update would be lost.
"""
with self._lock:
# Open the data file for flocking; if it doesn't exist yet,
# open the parent dir so we can still take an exclusive lock
# before the first write creates the file.
lock_path = self.path
lock_path.parent.mkdir(parents=True, exist_ok=True)
# Use a stable lock file (separate from data file) so an
# os.replace on the data file doesn't break our flock.
flock_path = lock_path.with_suffix(lock_path.suffix + ".lock")
f = open(flock_path, "w")
try:
fcntl.flock(f.fileno(), fcntl.LOCK_EX)
records = self._load()
mutate(records)
self._dump(records)
finally:
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
f.close()
# --- Public CRUD ---
def put(self, job: Job) -> None:
with self._lock:
self._jobs[job.job_id] = job
"""Insert or replace a Job by job_id. Persists to disk."""
def _mutate(records: dict[str, Job]) -> None:
records[job.job_id] = job
self._locked_dump(_mutate)
logger.info(
f"job_store put job_id={job.job_id} application_id={job.application_id}"
)
def get(self, job_id: str) -> Job | None:
"""Look up a Job by its local job_id only.
Prefer `get_either` in tool code — it also accepts application_id,
which is what the LLM agent sometimes has on hand.
"""
with self._lock:
return self._jobs.get(job_id)
return self._load().get(job_id)
def get_by_application_id(self, application_id: str) -> Job | None:
"""Look up a Job by its YARN application_id. O(n) over the file."""
with self._lock:
for job in self._load().values():
if job.application_id == application_id:
return job
return None
def get_either(self, job_id_or_application_id: str) -> Job | None:
"""Look up a Job by either identifier. Try job_id first (O(1)),
then fall back to application_id scan (O(n)).
This is what every job tool (get_job_status / get_job_logs /
kill_job / get_job_result) should call — it eliminates the
"agent passed the wrong ID and got a misleading 'Unknown job_id'"
failure mode.
"""
with self._lock:
records = self._load()
direct = records.get(job_id_or_application_id)
if direct is not None:
return direct
for job in records.values():
if job.application_id == job_id_or_application_id:
return job
return None
def list(self) -> list[Job]:
with self._lock:
return list(self._jobs.values())
return list(self._load().values())
def clear(self) -> None:
"""Test helper: drop everything. Not part of the public API."""
def _mutate(records: dict[str, Job]) -> None:
records.clear()
self._locked_dump(_mutate)
# Module-level singleton; replaced in tests.
store = JobStore()