"""Live progress logging for vibe-qc SCF and post-SCF stages.
Designed to defeat output buffering (every write is flushed) so that
the typical remote-job pattern just works::
nohup python my_calc.py > job.log 2>&1 &
tail -f job.log
… and you see each SCF iteration land as it happens, instead of a
long silence followed by a wall of text at the end.
Every SCF entry point in :mod:`vibeqc` accepts a ``progress=`` keyword
argument that is one of:
* ``True`` — a fresh :class:`ProgressLogger` writing to ``sys.stdout``.
* ``False`` / ``None`` — silent (the historical behavior).
* a :class:`ProgressLogger` instance — fully under caller control,
including a tee to a persistent ``.out`` file.
Verbosity (v0.5.3)
------------------
The amount of detail is controlled by an integer ``verbose`` level
following the PySCF convention (0..9). Each level is a strict
superset of the one below, so increasing ``verbose`` only adds
output:
===== ==========================================================
level what is emitted
===== ==========================================================
0 silent — every method is a no-op
1 banner + warnings + final SCF status only
2 add per-stage start lines + ``info()`` milestones
3 add per-stage timing on stage exit (current pre-v0.5.3
default behavior)
4 add per-iteration SCF rows (DEFAULT)
5 add inline memory snapshots
6+ debug — phase-level wall-clock breakdown (overlaps the
post-mortem ``.perf`` log)
===== ==========================================================
Logging integration (v0.5.3)
----------------------------
Pass ``use_logging=True`` to route every emit through
:mod:`logging` instead of the bare stream. Banners, milestones,
and ``converged()`` summaries land at ``INFO`` on the
``vibeqc.run_job`` logger; per-iteration SCF rows at ``DEBUG``;
warnings at ``WARNING``. This composes naturally with stdlib
handlers — ``RotatingFileHandler``, syslog, ``dictConfig`` — so
the canonical example::
import logging
logging.basicConfig(level=logging.INFO)
vq.run_job(mol, basis="6-31g*", method="rhf", output="x",
use_logging=True)
… mirrors progress through the user's logging stack rather than
``sys.stdout``. The verbose-level gate still applies before the
logging call, so ``verbose=2`` + ``use_logging=True`` will not
emit per-iteration ``DEBUG`` records.
Public surface
--------------
.. autoclass:: ProgressLogger
:members:
.. autofunction:: resolve_progress
"""
from __future__ import annotations
import logging
import os
import sys
import time
from contextlib import contextmanager
from pathlib import Path
from typing import IO, Any, Iterator, Optional, Union
__all__ = ["ProgressLogger", "resolve_progress"]
_INDENT = " "
#: Default verbose level when neither caller nor environment
#: override is set. Matches PySCF's convention (4 = "verbose"):
#: banner + stage milestones + per-iteration SCF rows.
DEFAULT_VERBOSE: int = 4
#: Logger name used when ``use_logging=True``. Users can attach
#: their own handlers via ``logging.getLogger("vibeqc.run_job")``.
_LOGGER_NAME: str = "vibeqc.run_job"
def _coerce_verbose(verbose: Union[int, bool, None]) -> int:
"""Normalize the ``verbose`` argument into a non-negative int.
* ``True`` → :data:`DEFAULT_VERBOSE` (back-compat with the
pre-v0.5.3 boolean-only API).
* ``False`` → 0 (silent — same back-compat).
* ``None`` → :data:`DEFAULT_VERBOSE`.
* ``int`` → clamped at 0.
"""
if verbose is None:
return DEFAULT_VERBOSE
if verbose is True:
return DEFAULT_VERBOSE
if verbose is False:
return 0
return max(0, int(verbose))
[docs]
class ProgressLogger:
"""Per-stage progress emitter for long-running calculations.
Writes plain-ASCII, line-flushed output to ``stream`` (default
``sys.stdout``) and, optionally, to ``log_path`` (truncated on
construction, appended-and-flushed thereafter). Every write goes
through ``flush()`` immediately so ``tail -f`` against a redirected
stdout shows the run unfolding in real time.
Parameters
----------
stream
Where to write live progress. Default ``sys.stdout``.
Ignored when ``use_logging=True``.
log_path
Optional path. Truncated on construction; appended-to with a
flush after each write. Useful when callers want both an
interactive stdout view and a persistent on-disk record.
Honored regardless of ``use_logging``.
verbose
Integer verbosity level (0..9, default 4). Higher values
add more detail; level 0 makes every method a no-op. The
level table is in the module docstring. ``True`` / ``False``
are accepted for back-compat (``True`` → 4, ``False`` → 0).
use_logging
If ``True``, route emits through ``logging.getLogger(
"vibeqc.run_job")`` instead of the stream. Composes with
stdlib handlers (rotating files, syslog, dictConfig) — no
special integration required. The ``log_path`` tee still
applies if set.
Notes
-----
The current implementation deliberately stays plain ASCII — the
goal is the ``nohup`` + ``tail -f`` workflow, where color codes
would only inject control bytes that no log viewer asked for. TTY
detection is exposed via :attr:`is_tty` for downstream tooling
that wants to layer color or progress bars on top.
"""
[docs]
def __init__(
self,
stream: Optional[IO[str]] = None,
log_path: Optional[Union[str, os.PathLike]] = None,
verbose: Union[int, bool, None] = DEFAULT_VERBOSE,
use_logging: bool = False,
) -> None:
self._stream = stream if stream is not None else sys.stdout
self._log_path: Optional[Path] = (
Path(log_path) if log_path is not None else None
)
self._level = _coerce_verbose(verbose)
self._use_logging = bool(use_logging)
self._logger = (
logging.getLogger(_LOGGER_NAME) if self._use_logging else None
)
self._t_start = time.perf_counter()
if self._log_path is not None and self._level > 0:
# Truncate so a fresh run doesn't accumulate alongside last
# week's tail. Subsequent writes are open(..., 'a') so an
# external 'tail -f' sees lines appear as they're written.
self._log_path.parent.mkdir(parents=True, exist_ok=True)
with open(self._log_path, "w", encoding="utf-8") as fh:
fh.write("")
# ----- low-level write -------------------------------------------------
def _emit(self, line: str, log_level: int = logging.INFO) -> None:
"""Push ``line`` to whichever sink is active.
Stream mode (default): write + flush to ``self._stream`` and
(if configured) ``self._log_path``.
Logging mode (``use_logging=True``): hand the line to
``logging.getLogger("vibeqc.run_job").log(log_level, ...)``,
which lets users plug in stdlib handlers (rotating files,
syslog, dictConfig). The ``log_path`` tee still applies so
callers can have both a stdlib-managed log stream AND a
verbatim per-job ``.out`` companion.
"""
if self._level <= 0:
return
if self._use_logging:
try:
# logging.<level> already handles its own newlines.
self._logger.log(log_level, line.lstrip())
except Exception:
# A logging failure must never tank the calculation.
pass
else:
try:
self._stream.write(line + "\n")
self._stream.flush()
except Exception:
pass
if self._log_path is not None:
try:
with open(self._log_path, "a", encoding="utf-8") as fh:
fh.write(line + "\n")
fh.flush()
except Exception:
pass
# ----- public API ------------------------------------------------------
[docs]
def info(self, message: str) -> None:
"""Emit a one-line informational message, indented.
Visible at verbose level 2 and above.
"""
if self._level < 2:
return
self._emit(f"{_INDENT}{message}", logging.INFO)
[docs]
def warn(self, message: str) -> None:
"""Emit a one-line warning, prefixed with ``WARN:``.
Visible at verbose level 1 and above (warnings should not
be hidden by quiet runs short of full silence).
"""
if self._level < 1:
return
self._emit(f"{_INDENT}WARN: {message}", logging.WARNING)
[docs]
def banner(self, title: str) -> None:
"""Section banner: blank line + indented title + dashed rule.
Visible at verbose level 1 and above. Under ``use_logging``
the three lines emit as three INFO records.
"""
if self._level < 1:
return
self._emit("", logging.INFO)
self._emit(f"{_INDENT}{title}", logging.INFO)
self._emit(f"{_INDENT}{'-' * max(len(title), 32)}", logging.INFO)
[docs]
def write_raw(self, text: str) -> None:
"""Write ``text`` exactly as given (newlines preserved), with a
single flush. Used by callers that want to splice in a
pre-formatted block (e.g. the geometry table from
:mod:`vibeqc.scf_log`).
Visible at verbose level 2 and above.
"""
if self._level < 2:
return
if self._use_logging:
# Preserve internal newlines as one logging record per
# line, so handlers that prefix each record stay sane.
for line in text.splitlines() or [""]:
try:
self._logger.log(logging.INFO, line)
except Exception:
pass
else:
try:
self._stream.write(text)
if not text.endswith("\n"):
self._stream.write("\n")
self._stream.flush()
except Exception:
pass
if self._log_path is not None:
try:
with open(self._log_path, "a", encoding="utf-8") as fh:
fh.write(text)
if not text.endswith("\n"):
fh.write("\n")
fh.flush()
except Exception:
pass
[docs]
@contextmanager
def stage(
self, name: str, *, detail: Optional[str] = None,
) -> Iterator[None]:
"""Context manager: emit ``[name]`` on enter, ``[name] done
(X.XXs)`` on exit — regardless of which branch ran inside.
The start line emits at verbose level 2; the timing
``done`` line emits at level 3 (so ``verbose=2`` shows
which stages ran but suppresses per-stage wall-time noise,
and ``verbose>=3`` adds the timing detail).
Parameters
----------
name
Short stage identifier (printed in brackets).
detail
Optional one-line addendum appended to the start line
(e.g. ``"4x4x4 -> 8 IBZ k-points"``).
"""
if self._level < 2:
yield
return
suffix = f" - {detail}" if detail else ""
self._emit(f"{_INDENT}[{name}]{suffix}", logging.INFO)
t0 = time.perf_counter()
try:
yield
finally:
if self._level >= 3:
dt = time.perf_counter() - t0
self._emit(
f"{_INDENT}[{name}] done ({_fmt_seconds(dt)})",
logging.INFO,
)
[docs]
def iteration(self, n: int, **fields: Any) -> None:
"""One SCF iteration line.
Recognized fields: ``energy`` (Ha), ``dE`` (Ha),
``grad`` (``||[F,DS]||``), ``diis`` (subspace dim). Any other
key/value pair is appended at the end. Wall time since logger
construction is appended automatically.
At ``n == 1``, ``dE`` is rendered as a placeholder ``--`` so
the column layout matches :func:`vibeqc.scf_log.format_scf_trace`.
Visible at verbose level 4 and above. Under ``use_logging``
emits at ``DEBUG`` (per-iter detail belongs below INFO so
``logging.basicConfig(level=INFO)`` keeps a quiet stream
while ``level=DEBUG`` exposes the trace).
"""
if self._level < 4:
return
cells = [f"iter {n:4d}"]
e = fields.pop("energy", None)
if e is not None:
cells.append(f"E = {float(e):18.10f} Ha")
dE = fields.pop("dE", None)
if dE is not None:
cells.append(
"dE = -- " if n <= 1 else f"dE = {float(dE):+.3e}"
)
grad = fields.pop("grad", None)
if grad is not None:
cells.append(f"||[F,DS]|| = {float(grad):.3e}")
diis = fields.pop("diis", None)
if diis is not None:
cells.append(
f"DIIS={int(diis):2d}" if int(diis) > 0 else "DIIS= -"
)
for key, value in fields.items():
cells.append(f"{key}={value}")
wall = time.perf_counter() - self._t_start
cells.append(f"[{_fmt_seconds(wall)}]")
self._emit(_INDENT + " ".join(cells), logging.DEBUG)
[docs]
def converged(
self, *, n_iter: int, energy: float, converged: bool,
) -> None:
"""Final SCF status line; called once after the iteration loop.
Visible at verbose level 1 and above (the final summary is
almost always wanted — only a fully silent ``verbose=0``
suppresses it).
"""
if self._level < 1:
return
status = "converged" if converged else "NOT converged"
self._emit(
f"{_INDENT}SCF {status} in {n_iter} iterations; "
f"E = {energy:.10f} Ha",
logging.INFO,
)
[docs]
def memory(self, label: str, rss_mib: float) -> None:
"""Inline RSS-memory snapshot.
Visible at verbose level 5 and above. Pairs with the
post-mortem ``.perf`` log's ``Memory snapshots`` section —
same data, but live.
"""
if self._level < 5:
return
self._emit(
f"{_INDENT}[memory] {label}: {float(rss_mib):.1f} MiB",
logging.INFO,
)
[docs]
def debug(self, message: str) -> None:
"""Phase-level wall-clock breakdown / debug detail.
Visible at verbose level 6 and above. Overlaps the
post-mortem ``.perf`` log on purpose — the perf log shows
the same numbers when the run is over; ``verbose>=6``
streams them live for users debugging an in-flight job.
Under ``use_logging`` emits at ``DEBUG`` so a stdlib
``RotatingFileHandler`` set to ``DEBUG`` captures the
trace without lifting the bar for INFO-only handlers.
"""
if self._level < 6:
return
self._emit(f"{_INDENT}[debug] {message}", logging.DEBUG)
@property
def enabled(self) -> bool:
"""``True`` iff :attr:`level` is non-zero (i.e. anything
emits). Kept as the legacy alias for ``verbose=True/False``
callers that predate the integer level."""
return self._level > 0
@property
def level(self) -> int:
"""The integer verbosity level this logger was constructed
with (0..9)."""
return self._level
@property
def use_logging(self) -> bool:
"""Whether emits route through :mod:`logging` (``True``) or
the raw stream (``False``)."""
return self._use_logging
@property
def stream(self) -> IO[str]:
return self._stream
@property
def is_tty(self) -> bool:
"""Whether the underlying stream is a TTY. Useful for callers
that want to layer color / spinners over the plain output."""
try:
return bool(self._stream.isatty())
except Exception:
return False
_NULL_LOGGER: Optional["ProgressLogger"] = None
def _null_logger() -> "ProgressLogger":
"""Singleton no-op logger reused when callers pass ``progress=False``."""
global _NULL_LOGGER
if _NULL_LOGGER is None:
_NULL_LOGGER = ProgressLogger(stream=sys.stdout, verbose=0)
return _NULL_LOGGER
[docs]
def resolve_progress(
progress: Union[bool, "ProgressLogger", None],
*,
verbose: Union[int, bool, None] = None,
use_logging: bool = False,
) -> "ProgressLogger":
"""Normalize the ``progress=`` argument every SCF entry point takes.
* ``True`` → a fresh :class:`ProgressLogger` writing to stdout
at the requested ``verbose`` level (default 4).
* ``False`` or ``None`` → a shared, silent no-op logger.
* :class:`ProgressLogger` → returned unchanged (``verbose`` /
``use_logging`` kwargs are ignored — the caller-supplied
logger is fully under the caller's control).
This indirection lets every SCF driver accept any of the four
forms via a one-line normalisation at the top of the function,
and threads the ``verbose=N`` knob through without each entry
point having to construct its own logger by hand.
"""
if isinstance(progress, ProgressLogger):
return progress
if progress is True:
level = (
DEFAULT_VERBOSE if verbose is None else _coerce_verbose(verbose)
)
return ProgressLogger(
stream=sys.stdout, verbose=level, use_logging=use_logging,
)
return _null_logger()
def _fmt_seconds(s: float) -> str:
"""Compact wall-clock formatter: ms / s / m+s, never negative width."""
if s < 1.0:
return f"{s * 1e3:6.1f}ms"
if s < 60.0:
return f"{s:6.2f}s"
m, sec = divmod(s, 60.0)
return f"{int(m)}m{sec:05.2f}s"