slabterminator 0.2.0__py3-none-any.whl

slabterminator/__init__.py

@@ -0,0 +1,7 @@
+"""SlabTerminator: enumerate the symmetrically unique slab terminations of a bulk crystal."""
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("slabterminator")
+except PackageNotFoundError:  # not installed (e.g. running from a source tree)
+    __version__ = "0.0.0"

slabterminator/batch.py

@@ -0,0 +1,224 @@
+# coding: utf-8
+"""Parallel, multi-material orchestration for the slab pipeline.
+
+run_batch fans build_slabs_for_material out across worker processes, supervises
+them, and streams each finished MaterialResult to a caller-supplied sink. It is
+scheduler- and output-agnostic: the caller provides the (id, structure) stream,
+the SlabGenConfig, and an on_result callback (e.g. a CSV or database writer). See
+scripts/build_slab_dataset_slurm.py for a SLURM entry point that wires in a CSV sink.
+
+multiprocessing (not asyncio or a thread pool) because slab generation is
+CPU-bound and, crucially, a single pathological material can spin far longer than
+its share of the budget: only a separate OS process can be forcibly killed to
+reclaim the slot, which futures / pool executors do not allow (future.cancel()
+cannot interrupt a task already running). run_batch therefore manages raw
+Process objects so it can terminate() a material that overruns max_material_seconds
+and launch the next one in its place.
+"""
+import multiprocessing as mp
+import os
+import threading
+import time
+from collections.abc import Callable, Iterable
+from multiprocessing import Manager
+from queue import Empty as QueueEmpty
+
+from pymatgen.core import Structure
+
+from slabterminator.pipeline import MaterialResult, SlabGenConfig, build_slabs_for_material
+
+# Hard per-material wall-clock budget (seconds): a worker running longer than this
+# is killed and its material recorded as a failure. Guards against a material stuck
+# in a pathological / non-terminating code path holding a worker slot forever.
+_DEFAULT_MAX_MATERIAL_SECONDS = 3 * 60 * 60  # 3 hours
+
+# Grace period (seconds) after a worker process is first seen not-alive before it is
+# reaped as a crash. A worker that finishes normally exits right after enqueueing its
+# result; this grace lets that in-flight result drain first, so a clean finish is
+# never misreported as a crash.
+_DEAD_WORKER_GRACE_SECONDS = 5.0
+
+
+def _worker(material_id: str, structure: Structure | dict, config: SlabGenConfig,
+            log_queue, results_queue) -> None:
+    """Worker-process entry point: build one material's slabs and enqueue the result.
+
+    Runs build_slabs_for_material for a single material and puts (material_id,
+    MaterialResult) on results_queue. Progress messages go to the parent-owned
+    log_queue rather than being printed, so concurrent workers never interleave
+    output and all file writes stay in the parent. build_slabs_for_material catches
+    its own per-material errors, so this normally enqueues a MaterialResult(ok=False)
+    rather than raising; a hard crash or hang is handled by the parent (see
+    run_batch).
+
+    Args:
+        material_id (str): id for this material (stamped onto the result).
+        structure (Structure or dict): the bulk structure, or its as_dict() form
+            (rebuilt here, so callers can stream cheap dicts to workers).
+        config (SlabGenConfig): generation parameters.
+        log_queue: parent-owned queue for progress messages.
+        results_queue: parent-owned queue the (material_id, MaterialResult) is put on.
+
+    """
+    if isinstance(structure, dict):
+        structure = Structure.from_dict(structure)
+    result = build_slabs_for_material(structure, config, material_id=material_id, log=log_queue.put)
+    results_queue.put((material_id, result))
+
+
+def _drain_log_queue(log_queue, log: Callable[[str], None]) -> None:
+    """Forwards queued log messages to the log callback until a None sentinel.
+
+    Args:
+        log_queue: queue of message strings (None ends the loop).
+        log (Callable[[str], None]): sink for each message (e.g. print).
+
+    """
+    while True:
+        msg = log_queue.get()
+        if msg is None:
+            break
+        log(msg)
+
+
+def run_batch(materials: Iterable[tuple[str, Structure | dict]],
+              config: SlabGenConfig | None = None, *, n_workers: int | None = None,
+              max_material_seconds: float = _DEFAULT_MAX_MATERIAL_SECONDS,
+              on_result: Callable[[MaterialResult], None],
+              log: Callable[[str], None] = print) -> int:
+    """Builds slabs for many materials in parallel, streaming each result to a sink.
+
+    Keeps up to n_workers worker processes busy; a slot frees (and the next material
+    launches) only once the finished material's result has been handed to on_result,
+    so on_result is called exactly once per material, one at a time in the parent
+    process -- concurrent completions can never interleave writes to a shared sink.
+    A material whose worker overruns max_material_seconds (or crashes) is recorded as
+    a MaterialResult(ok=False) and its slot reclaimed.
+
+    Args:
+        materials (Iterable[tuple[str, Structure | dict]]): (id, structure) pairs.
+            id is used as the result's material_id and as the supervisor key, so it
+            must be unique across the batch (e.g. a Materials Project id); structure
+            may be a Structure or its as_dict() form.
+        config (SlabGenConfig or None): generation parameters; defaults to
+            SlabGenConfig().
+        n_workers (int or None): number of worker processes; defaults to os.cpu_count().
+        max_material_seconds (float): hard per-material wall-clock budget; a worker
+            running longer is killed and the material recorded as failed.
+        on_result (Callable[[MaterialResult], None]): called once per material with
+            its result (ok True or False). The sink decides what to persist.
+        log (Callable[[str], None]): progress-message sink; defaults to print. All
+            messages (workers' and the supervisor's) are serialized through one
+            drain thread, so the callback is only ever called from that thread.
+
+    Returns:
+        int: the number of materials processed (every material is counted exactly
+        once, whether it succeeded, failed, timed out, or crashed).
+
+    """
+    config = config or SlabGenConfig()
+    n_workers = n_workers or os.cpu_count() or 1
+    try:
+        total = len(materials)  # a sized collection; a bare generator falls through
+    except TypeError:
+        total = None
+
+    manager = Manager()
+    log_queue = manager.Queue()
+    results_queue = manager.Queue()
+    drain_thread = threading.Thread(target=_drain_log_queue, args=(log_queue, log), daemon=True)
+    drain_thread.start()
+
+    mat_iter = iter(materials)
+    running: dict[str, tuple[mp.Process, float]] = {}  # id -> (Process, start_time)
+    dead_since: dict[str, float] = {}  # id -> first time its process was seen not-alive
+    completed = 0
+
+    def progress() -> str:
+        return (f"  ({completed} of {total} materials complete)" if total is not None
+                else f"  ({completed} materials complete)")
+
+    def launch_next() -> None:
+        try:
+            material_id, structure = next(mat_iter)
+        except StopIteration:
+            return
+        p = mp.Process(target=_worker,
+                       args=(material_id, structure, config, log_queue, results_queue),
+                       daemon=True)
+        p.start()
+        running[material_id] = (p, time.time())
+
+    for _ in range(n_workers):
+        launch_next()
+
+    while running:
+        # Drain every result that has arrived, without blocking. A slot only frees
+        # (and the next material only launches) once its result has been delivered to
+        # on_result, so results are handled one at a time here.
+        drained_any = False
+        while True:
+            try:
+                material_id, result = results_queue.get_nowait()
+            except QueueEmpty:
+                break
+            drained_any = True
+            entry = running.pop(material_id, None)
+            dead_since.pop(material_id, None)
+            if entry is None:
+                # This slot was already reclaimed by the timeout/crash check (the
+                # worker finished right as it was being reaped); its result raced in
+                # afterward. Skip it -- it was already counted there.
+                log_queue.put(f"Note: late result for {material_id!r} arrived after its "
+                              f"slot was reclaimed; ignoring (already counted).")
+                continue
+            p, _ = entry
+            p.join()
+            completed += 1
+            on_result(result)
+            log_queue.put(progress())
+            launch_next()
+
+        # Reclaim any worker that overran its budget or died without a result.
+        now = time.time()
+        for material_id, (p, start_time) in list(running.items()):
+            if now - start_time > max_material_seconds:
+                log_queue.put(f"TIMEOUT: material {material_id!r} exceeded "
+                              f"{max_material_seconds:.0f}s; killing its worker and skipping it")
+                p.terminate()
+                p.join(10)
+                if p.is_alive():
+                    p.kill()
+                    p.join()
+                del running[material_id]
+                dead_since.pop(material_id, None)
+                completed += 1
+                on_result(MaterialResult(material_id=material_id, ok=False,
+                                         error=f"timeout after {max_material_seconds:.0f}s"))
+                log_queue.put(progress())
+                launch_next()
+            elif not p.is_alive():
+                # Exited without a result on the queue. Give a brief grace for an
+                # in-flight result to drain (a clean finish exits right after
+                # enqueueing); if none appears, treat it as a crash.
+                first_seen = dead_since.setdefault(material_id, now)
+                if now - first_seen > _DEAD_WORKER_GRACE_SECONDS:
+                    log_queue.put(f"CRASH: worker for material {material_id!r} exited "
+                                  f"without producing a result; recording as failed")
+                    p.join()
+                    del running[material_id]
+                    dead_since.pop(material_id, None)
+                    completed += 1
+                    on_result(MaterialResult(material_id=material_id, ok=False,
+                                             error="worker exited without producing a result"))
+                    log_queue.put(progress())
+                    launch_next()
+            else:
+                dead_since.pop(material_id, None)
+
+        if not drained_any:
+            time.sleep(1)
+
+    log_queue.put(None)
+    drain_thread.join()
+    return completed