Coverage for lunar_toolkit/cli/_jobs.py: 62%
329 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-19 02:58 +0000
1"""Background jobs for ``lunar serve`` / ``lunar app``: index builds and downloads with progress."""
3from __future__ import annotations
5import itertools
6import logging
7import re
8import threading
9import time
10from collections import deque
11from collections.abc import Callable
12from concurrent.futures import ThreadPoolExecutor
13from dataclasses import dataclass
14from pathlib import Path
15from typing import Any
17log = logging.getLogger(__name__)
19TERMINAL = ("done", "error", "cancelled")
20MAX_IDS = 1000
21_RELEASE_RE = re.compile(r"^LROLRC_\d{4}[A-Z]?$")
24class Cancelled(Exception):
25 """Raised inside a job's thread once the job was cancelled."""
28@dataclass(frozen=True)
29class JobContext:
30 """Where jobs read and write: the Parquet data dir, image downloads and reference data."""
32 data_dir: Path
33 download_dir: Path
34 reference_dir: Path
37class Job:
38 """One background task: status, progress (0-1), current file, bytes, rate, ETA and a log tail."""
40 def __init__(self, manager: JobManager, jid: str, kind: str, title: str, key: str, params: dict) -> None:
41 self._manager = manager
42 self.id, self.kind, self.title, self.key, self.params = jid, kind, title, key, params
43 self.status = "queued"
44 self.progress = 0.0
45 self.busy = False
46 self.cancellable = True
47 self.stage = ""
48 self.file = ""
49 self.file_index = self.file_count = 0
50 self.bytes_done = 0
51 self.bytes_total: int | None = None
52 self.rate: float | None = None
53 self.eta: float | None = None
54 self.error: str | None = None
55 self.result: Any = None
56 self.created = time.time()
57 self.started: float | None = None
58 self.finished: float | None = None
59 self.log: deque[str] = deque(maxlen=40)
60 self._span = (0.0, 1.0)
61 self._cancel = threading.Event()
62 self._last: tuple[float, int] | None = None
64 @property
65 def cancelled(self) -> bool:
66 return self._cancel.is_set()
68 def check(self) -> None:
69 """Raise `Cancelled` if the job was cancelled."""
70 if self._cancel.is_set():
71 raise Cancelled()
73 def update(self, **fields: Any) -> None:
74 with self._manager._cond:
75 for k, v in fields.items():
76 setattr(self, k, v)
77 self._manager._changed()
79 def info(self, msg: str) -> None:
80 log.info("job %s: %s", self.id, msg)
81 with self._manager._cond:
82 self.log.append(f"{time.strftime('%H:%M:%S')} {msg}")
83 self._manager._changed()
85 def set_stage(self, stage: str, span: tuple[float, float] | None = None, busy: bool = False) -> None:
86 """Start a stage; with *span*, byte progress maps onto that part of the overall bar."""
87 self.check()
88 if span is not None:
89 self._span = span
90 self.info(stage)
91 self.update(stage=stage, busy=busy, progress=max(self.progress, self._span[0]))
93 def set_file(self, k: int, n: int, name: str) -> None:
94 """File *k* of *n* (1-based) starts; it covers ``[(k-1)/n, k/n]`` of the bar."""
95 self.check()
96 self._span = ((k - 1) / n, k / n)
97 self._last = None
98 self.update(
99 file=name, file_index=k, file_count=n, bytes_done=0, bytes_total=None, rate=None, eta=None,
100 progress=max(self.progress, self._span[0]), stage=f"File {k} of {n}",
101 ) # fmt: skip
103 def on_bytes(self, name: str, done: int, total: int | None) -> None:
104 """Download progress callback
105 ([`lunar_toolkit.downloader.Progress`][lunar_toolkit.downloader.Progress])."""
106 self.check()
107 if done < 0: 107 ↛ 108line 107 didn't jump to line 108 because the condition on line 107 was never true
108 self.update(stage="Waiting: the host asked us to slow down", rate=None, eta=None)
109 return
110 now = time.monotonic()
111 rate = self.rate
112 if self._last is None or done < self._last[1]:
113 self._last = (now, done)
114 elif now - self._last[0] >= 0.5: 114 ↛ 115line 114 didn't jump to line 115 because the condition on line 114 was never true
115 inst = (done - self._last[1]) / (now - self._last[0])
116 rate = inst if rate is None else 0.6 * rate + 0.4 * inst
117 self._last = (now, done)
118 lo, hi = self._span
119 frac = min(1.0, done / total) if total else 0.0
120 eta = (total - done) / rate if total and rate else None
121 self.update(
122 file=name, bytes_done=done, bytes_total=total, rate=rate, eta=eta,
123 progress=max(self.progress, lo + (hi - lo) * frac),
124 ) # fmt: skip
126 def to_dict(self, admin: bool = True) -> dict[str, Any]:
127 d: dict[str, Any] = {
128 k: getattr(self, k)
129 for k in (
130 "id", "kind", "title", "status", "progress", "busy", "cancellable", "stage", "file",
131 "file_index", "file_count", "bytes_done", "bytes_total", "rate", "eta", "created",
132 "started", "finished",
133 )
134 } # fmt: skip
135 if admin:
136 d.update(params=self.params, error=self.error, result=self.result, log=list(self.log))
137 else:
138 d.update(error="failed" if self.error else None)
139 return d
142class JobManager:
143 """Runs jobs in a small thread pool and lets readers wait for any change (``version``)."""
145 def __init__(
146 self, workers: int = 2, keep: int = 50, on_finish: Callable[[Job], None] | None = None
147 ) -> None:
148 self._cond = threading.Condition()
149 self._on_finish = on_finish
150 self._jobs: dict[str, Job] = {}
151 self._pool = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="lunar-job")
152 self._ids = itertools.count(1)
153 self._keep = keep
154 self.version = 0
155 self.closed = False
157 def _changed(self) -> None:
158 self.version += 1
159 self._cond.notify_all()
161 def submit(self, kind: str, title: str, key: str, params: dict, func: Callable[[Job], Any]) -> Job:
162 """Queue *func(job)*; ValueError when a job with the same *key* is still active."""
163 with self._cond:
164 if self.closed: 164 ↛ 165line 164 didn't jump to line 165 because the condition on line 164 was never true
165 raise RuntimeError("shutting down")
166 for other in self._jobs.values():
167 if other.key == key and other.status not in TERMINAL:
168 raise ValueError(f"{other.title} is already {other.status} (job {other.id})")
169 job = Job(self, str(next(self._ids)), kind, title, key, params)
170 self._jobs[job.id] = job
171 done = [j for j in self._jobs.values() if j.status in TERMINAL]
172 for old in done[: max(0, len(self._jobs) - self._keep)]: 172 ↛ 173line 172 didn't jump to line 173 because the loop on line 172 never started
173 del self._jobs[old.id]
174 self._changed()
175 self._pool.submit(self._run, job, func)
176 return job
178 def _run(self, job: Job, func: Callable[[Job], Any]) -> None:
179 if job.cancelled: 179 ↛ 180line 179 didn't jump to line 180 because the condition on line 179 was never true
180 return
181 job.update(status="running", started=time.time())
182 try:
183 result = func(job)
184 except Cancelled:
185 job.info("Cancelled")
186 job.update(status="cancelled", finished=time.time(), busy=False)
187 except Exception as exc:
188 log.exception("job %s (%s) failed", job.id, job.kind)
189 job.info(f"Error: {exc}")
190 job.update(status="error", error=str(exc) or type(exc).__name__, finished=time.time(), busy=False)
191 else:
192 job.update(status="done", result=result, progress=1.0, finished=time.time(), busy=False, eta=None)
193 if self._on_finish:
194 self._on_finish(job)
196 def get(self, jid: str) -> Job | None:
197 with self._cond:
198 return self._jobs.get(jid)
200 def jobs(self) -> list[Job]:
201 with self._cond:
202 return list(self._jobs.values())
204 def cancel(self, jid: str) -> Job | None:
205 """Ask job *jid* to stop (a queued job is cancelled at once); None if unknown."""
206 with self._cond:
207 job = self._jobs.get(jid)
208 if job is None or job.status in TERMINAL:
209 return job
210 job._cancel.set()
211 if job.status == "queued": 211 ↛ 212line 211 didn't jump to line 212 because the condition on line 211 was never true
212 job.status, job.finished = "cancelled", time.time()
213 self._changed()
214 return job
216 def wait(self, version: int, timeout: float) -> int:
217 """Block until `version` differs from *version* (or *timeout* / close); the new version."""
218 with self._cond:
219 self._cond.wait_for(lambda: self.version != version or self.closed, timeout)
220 return self.version
222 def close(self) -> None:
223 """Cancel everything and stop accepting jobs (running ones stop at their next check)."""
224 with self._cond:
225 self.closed = True
226 for job in self._jobs.values():
227 if job.status not in TERMINAL:
228 job._cancel.set()
229 self._changed()
230 self._pool.shutdown(wait=False, cancel_futures=True)
233def _flag(p: dict, name: str) -> bool:
234 v = p.get(name, False)
235 if not isinstance(v, bool): 235 ↛ 236line 235 didn't jump to line 236 because the condition on line 235 was never true
236 raise ValueError(f"{name} must be true or false")
237 return v
240def _str(p: dict, name: str, default: str | None = None) -> str | None:
241 v = p.get(name, default)
242 if v is not None and not isinstance(v, str): 242 ↛ 243line 242 didn't jump to line 243 because the condition on line 242 was never true
243 raise ValueError(f"{name} must be a string")
244 return v.strip() if v else default
247def _strings(p: dict, name: str, limit: int = 100) -> list[str]:
248 v = p.get(name)
249 if isinstance(v, str): 249 ↛ 250line 249 didn't jump to line 250 because the condition on line 249 was never true
250 v = v.split()
251 if not isinstance(v, list) or not all(isinstance(x, str) and x.strip() for x in v) or not v: 251 ↛ 252line 251 didn't jump to line 252 because the condition on line 251 was never true
252 raise ValueError(f"{name} must be a non-empty list of strings")
253 if len(v) > limit: 253 ↛ 254line 253 didn't jump to line 254 because the condition on line 253 was never true
254 raise ValueError(f"at most {limit} {name}")
255 return [x.strip() for x in v]
258def _region(p: dict):
259 from .. import reference as ref
261 key = p.get("region")
262 if not isinstance(key, (str, int)) or isinstance(key, bool):
263 raise ValueError("region must be a region id or name")
264 try:
265 return ref.get_region(key)
266 except KeyError as exc:
267 raise ValueError(exc.args[0] if exc.args else str(exc)) from None
270def _build_db(p: dict, ctx: JobContext) -> tuple[str, str, Callable[[Job], Any]]:
271 from ..constants import DATASETS
273 dataset = p.get("dataset")
274 if dataset not in DATASETS:
275 raise ValueError(f"dataset must be one of {', '.join(DATASETS)}")
276 update, force, volume_only = _flag(p, "update"), _flag(p, "force"), _flag(p, "volume_only")
277 release = _str(p, "release")
278 if release and not _RELEASE_RE.match(release.upper()):
279 raise ValueError("release must look like LROLRC_0066C")
281 def run(job: Job) -> dict[str, Any]:
282 from ..downloader import report_progress
283 from ..lroc.cumindex import CumIndex
285 ci = CumIndex(dataset, data_dir=ctx.data_dir)
287 def on_stage(stage: str) -> None:
288 if stage == "download":
289 job.set_stage("Downloading the index label and table", span=(0.0, 0.7))
290 else:
291 job.set_stage("Building the Parquet database (DuckDB)", busy=True)
292 job.update(cancellable=False)
294 with report_progress(job.on_bytes):
295 if update: 295 ↛ 296line 295 didn't jump to line 296 because the condition on line 295 was never true
296 added = ci.update_db(on_stage=on_stage)
297 job.info(f"Appended {', '.join(added)}" if added else "Already up to date")
298 else:
299 ci.build_db(force=force, release=release, volume_only=volume_only, on_stage=on_stage)
300 info = ci.build_info() or {}
301 job.info(f"{info.get('rows', 0):,} rows, release {info.get('release')}")
302 return {"dataset": dataset, "rows": info.get("rows"), "release": info.get("release")}
304 verb = "Update" if update else "Build"
305 return f"{verb} {dataset.upper()} index", f"build-db:{dataset}", run
308def _pgda(p: dict, ctx: JobContext) -> tuple[str, str, Callable[[Job], Any]]:
309 from .. import reference as ref
311 reg = _region(p)
312 if reg.pgda_site is None:
313 raise ValueError(f"Region {reg.id} {reg.name} has no PGDA LOLA 5 m site")
314 kinds = _strings(p, "kinds", len(ref.PGDA_KINDS))
315 bad = [k for k in kinds if k not in ref.PGDA_KINDS]
316 if bad:
317 raise ValueError(f"unknown PGDA kind(s) {bad}; choose from {', '.join(ref.PGDA_KINDS)}")
318 clones = p.get("clones")
319 if clones is not None and (not isinstance(clones, int) or isinstance(clones, bool) or clones < 1):
320 raise ValueError("clones must be a positive integer")
321 site = reg.pgda_site
323 def run(job: Job) -> dict[str, Any]:
324 from ..downloader import report_progress, url_filename
326 job.set_stage("Reading the PGDA file index")
327 files = [(k, u) for k, urls in ref.pgda_urls(site, kinds, clones).items() for u in urls]
328 out = ctx.reference_dir / "pgda" / site
329 with report_progress(job.on_bytes):
330 for i, (kind, url) in enumerate(files, 1):
331 name = url_filename(url)
332 job.set_file(i, len(files), name)
333 ref.fetch(url, out / ("Clones" if kind == "clones" else "") / name, show_progress=False)
334 job.info(f"{len(files)} file(s) for {site}")
335 return {"site": site, "files": len(files)}
337 return f"PGDA {site} ({', '.join(kinds)})", f"pgda:{site}", run
340def _sdem(p: dict, ctx: JobContext) -> tuple[str, str, Callable[[Job], Any]]:
341 reg = _region(p)
342 members = _strings(p, "members")
344 def run(job: Job) -> dict[str, Any]:
345 from .. import reference as ref
346 from ..tools import reuse_existing
348 job.set_stage("Reading the zip directory (HTTP Range)")
349 out_dir = ctx.reference_dir / "sdem" / reg.zenodo_file.removesuffix(".zip")
350 done = []
351 with ref._open_zip(reg) as zf:
352 names = ref.select_members(zf.namelist(), members)
353 if not names:
354 raise ValueError(f"No zip member matches {members}")
355 for i, name in enumerate(names, 1):
356 info = zf.getinfo(name)
357 dest = out_dir / (name.split("/", 1)[1] if "/" in name else name)
358 job.set_file(i, len(names), dest.name)
359 if not reuse_existing(dest, False):
360 dest.parent.mkdir(parents=True, exist_ok=True)
361 tmp = dest.with_suffix(dest.suffix + ".part")
362 n = 0
363 with zf.open(name) as src, open(tmp, "wb") as dst:
364 while buf := src.read(1 << 20):
365 dst.write(buf)
366 n += len(buf)
367 job.on_bytes(dest.name, n, info.file_size)
368 tmp.replace(dest)
369 done.append(dest.name)
370 job.info(f"{len(done)} member(s) of {reg.zenodo_file}")
371 return {"region": reg.id, "files": len(done)}
373 return f"SDEM {reg.name} ({', '.join(members)})", f"sdem:{reg.id}", run
376def _download(p: dict, ctx: JobContext) -> tuple[str, str, Callable[[Job], Any]]:
377 from ..constants import NAC_ID_RE, WAC_ID_RE
379 ids = [i.upper() for i in _strings(p, "ids", MAX_IDS)]
380 bad = [i for i in ids if not (NAC_ID_RE.match(i) or WAC_ID_RE.match(i))]
381 if bad: 381 ↛ 383line 381 didn't jump to line 383 because the condition on line 381 was always true
382 raise ValueError(f"not LROC product IDs: {' '.join(bad[:5])}")
383 product = p.get("product", "edr")
384 if product not in ("edr", "cdr"):
385 raise ValueError("product must be edr or cdr")
386 dest_s = _str(p, "dest")
387 dest = Path(dest_s).expanduser() if dest_s else ctx.download_dir
388 if not dest.is_absolute():
389 raise ValueError("dest must be an absolute directory")
391 def run(job: Job) -> dict[str, Any]:
392 from ..downloader import report_progress
393 from ..lroc.nac import NacImageCollection
395 job.set_stage("Looking up the product IDs")
396 images = list(NacImageCollection.from_sources(ids, data_dir=ctx.data_dir))
397 missing = sorted(set(ids) - {img.product_id for img in images})
398 if missing:
399 job.info(f"Not in the index: {' '.join(missing[:20])}{' ...' if len(missing) > 20 else ''}")
400 ok, failed = [], []
401 with report_progress(job.on_bytes):
402 for i, img in enumerate(images, 1):
403 job.set_file(i, len(images), img.product_id)
404 try:
405 path = img.download_cdr(dest) if product == "cdr" else img.download_edr(dest)
406 except Cancelled:
407 raise
408 except Exception as exc:
409 job.info(f"{img.product_id}: {exc}")
410 failed.append(img.product_id)
411 else:
412 ok.append(path.name)
413 job.info(f"{len(ok)} downloaded, {len(failed)} failed, {len(missing)} not found")
414 if not ok and (failed or missing):
415 raise RuntimeError("nothing was downloaded")
416 return {"downloaded": len(ok), "failed": failed, "missing": missing}
418 return f"Download {len(ids)} {product.upper()} image(s)", f"download:{product}:{' '.join(ids)}", run
421KINDS: dict[str, Callable[[dict, JobContext], tuple[str, str, Callable[[Job], Any]]]] = {
422 "build-db": _build_db,
423 "pgda": _pgda,
424 "sdem": _sdem,
425 "download": _download,
426}
429def submit(manager: JobManager, payload: Any, ctx: JobContext) -> Job:
430 """Validate a ``POST /api/jobs`` body and queue the job (ValueError for a bad request)."""
431 if not isinstance(payload, dict): 431 ↛ 432line 431 didn't jump to line 432 because the condition on line 431 was never true
432 raise ValueError("expected a JSON object")
433 kind = payload.get("kind")
434 if kind not in KINDS:
435 raise ValueError(f"kind must be one of {', '.join(KINDS)}")
436 title, key, func = KINDS[kind](payload, ctx)
437 params = {k: v for k, v in payload.items() if k != "kind"}
438 return manager.submit(kind, title, key, params, func)