class JobQueue:
"""File-state job queue over a local directory or s3:// prefix."""
def __init__(self, base: str | Path, *, profile: str = ""):
self.base = str(base).rstrip("/")
self.profile = profile or os.environ.get(JOBS_PROFILE_ENV, "")
self.is_s3 = self.base.startswith("s3://")
self._client = None
# -- backend primitives ---------------------------------------------
def _s3(self):
if self._client is None:
try:
import boto3
except ImportError as exc:
raise RuntimeError(
"boto3 is required for the S3 job queue (pip install boto3)."
) from exc
session = (
boto3.Session(profile_name=self.profile)
if self.profile
else boto3.Session()
)
self._client = session.client("s3")
return self._client
def _bucket_key(self, rel: str) -> tuple[str, str]:
without = self.base[len("s3://"):]
bucket, _, prefix = without.partition("/")
key = f"{prefix}/{rel}".strip("/") if prefix else rel
return bucket, key
def _local_path(self, rel: str) -> Path:
"""Resolve a relative queue path, asserting it stays under ``base``.
Defense-in-depth on top of :func:`_safe_job_id`: even if an unsafe id
slipped through, ``safe_join`` refuses to escape the queue directory.
"""
return safe_join(self.base, *rel.split("/"))
def _read(self, rel: str) -> str | None:
if self.is_s3:
bucket, key = self._bucket_key(rel)
try:
response = self._s3().get_object(Bucket=bucket, Key=key)
except self._s3().exceptions.NoSuchKey:
return None
except Exception as exc: # boto's ClientError for 404s
if getattr(exc, "response", {}).get("Error", {}).get("Code") in ("404", "NoSuchKey"):
return None
raise
return response["Body"].read().decode("utf-8")
path = self._local_path(rel)
if not path.exists():
return None
return path.read_text(encoding="utf-8")
def _write(self, rel: str, content: str) -> None:
if self.is_s3:
bucket, key = self._bucket_key(rel)
self._s3().put_object(Bucket=bucket, Key=key, Body=content.encode("utf-8"))
return
path = self._local_path(rel)
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(content, encoding="utf-8")
def _delete(self, rel: str) -> None:
if self.is_s3:
bucket, key = self._bucket_key(rel)
self._s3().delete_object(Bucket=bucket, Key=key)
return
path = self._local_path(rel)
path.unlink(missing_ok=True)
def _list(self, rel_prefix: str) -> list[str]:
"""Return relative file names (without directory) under a prefix."""
if self.is_s3:
bucket, key_prefix = self._bucket_key(rel_prefix)
paginator = self._s3().get_paginator("list_objects_v2")
names = []
for page in paginator.paginate(Bucket=bucket, Prefix=f"{key_prefix}/"):
for obj in page.get("Contents", []):
names.append(obj["Key"].rsplit("/", 1)[-1])
return sorted(names)
directory = Path(self.base) / rel_prefix
if not directory.exists():
return []
return sorted(p.name for p in directory.iterdir() if p.is_file())
# -- queue operations -------------------------------------------------
def enqueue(self, job: Mapping[str, Any]) -> str:
job_id = str(job.get("job_id") or "").strip()
if not job_id:
raise ValueError("Job spec must include a job_id.")
job_id = _safe_job_id(job_id)
self._write(f"pending/{job_id}.json", json.dumps(dict(job), indent=2, default=str))
return job_id
def list_jobs(self, state: str) -> list[str]:
if state not in JOB_STATES:
raise ValueError(f"Unknown job state: {state}")
return [name[:-5] for name in self._list(state) if name.endswith(".json")]
def get_job(self, job_id: str) -> tuple[str, dict[str, Any]] | None:
"""Return (state, job) for a job id, searching all states."""
job_id = _safe_job_id(job_id)
for state in JOB_STATES:
content = self._read(f"{state}/{job_id}.json")
if content is not None:
try:
return state, json.loads(content)
except json.JSONDecodeError:
return state, {}
return None
def find_jobs(
self,
*,
experiment: str = "",
states: Sequence[str] | None = None,
max_reads_per_state: int = 20,
) -> list[dict[str, Any]]:
"""Return jobs annotated with their state, newest first.
Job ids embed a UTC timestamp, so listing names is enough to pick the
newest files; only those are read (bounded S3 GETs per call). When
``experiment`` is given, only jobs whose family or experiment id
matches are returned.
"""
matches: list[dict[str, Any]] = []
for state in states or JOB_STATES:
job_ids = self.list_jobs(state)
for job_id in reversed(job_ids[-max_reads_per_state:]):
content = self._read(f"{state}/{job_id}.json")
if content is None:
continue
try:
job = json.loads(content)
except json.JSONDecodeError:
continue
if experiment and experiment not in (
str(job.get("experiment") or ""),
str(job.get("experiment_id") or ""),
):
continue
job["state"] = state
job.setdefault("job_id", job_id)
matches.append(job)
matches.sort(key=lambda job: str(job.get("requested_at") or ""), reverse=True)
return matches
def claim_next(self) -> dict[str, Any] | None:
"""Claim the oldest pending job by moving it to running/."""
for job_id in self.list_jobs("pending"):
content = self._read(f"pending/{job_id}.json")
if content is None:
continue
try:
job = json.loads(content)
except json.JSONDecodeError:
continue
job["claimed_at"] = _utc_now()
job["heartbeat_at"] = job["claimed_at"]
self._write(f"running/{job_id}.json", json.dumps(job, indent=2, default=str))
self._delete(f"pending/{job_id}.json")
return job
return None
def heartbeat(self, job_id: str) -> None:
job_id = _safe_job_id(job_id)
content = self._read(f"running/{job_id}.json")
if content is None:
return
try:
job = json.loads(content)
except json.JSONDecodeError:
return
job["heartbeat_at"] = _utc_now()
self._write(f"running/{job_id}.json", json.dumps(job, indent=2, default=str))
def complete(self, job_id: str, result: Mapping[str, Any] | None = None) -> None:
self._finish(job_id, "done", result or {})
def fail(self, job_id: str, error: str, result: Mapping[str, Any] | None = None) -> None:
merged = dict(result or {})
merged["error"] = error
self._finish(job_id, "failed", merged)
def _finish(self, job_id: str, state: str, result: Mapping[str, Any]) -> None:
job_id = _safe_job_id(job_id)
content = self._read(f"running/{job_id}.json")
job = {}
if content is not None:
try:
job = json.loads(content)
except json.JSONDecodeError:
job = {}
job["job_id"] = job.get("job_id") or job_id
job["finished_at"] = _utc_now()
job["result"] = dict(result)
self._write(f"{state}/{job_id}.json", json.dumps(job, indent=2, default=str))
self._delete(f"running/{job_id}.json")
# -- logs --------------------------------------------------------------
def append_log(self, job_id: str, text: str) -> None:
if not text:
return
job_id = _safe_job_id(job_id)
existing = self._read(f"logs/{job_id}.log") or ""
self._write(f"logs/{job_id}.log", existing + text)
def read_log(self, job_id: str, *, tail_chars: int = 0) -> str:
job_id = _safe_job_id(job_id)
content = self._read(f"logs/{job_id}.log") or ""
if tail_chars and len(content) > tail_chars:
return content[-tail_chars:]
return content