Progress & Logging Utilities

console

console() -> Console

Return a shared Rich Console instance with basic theming.

Source code in src/spark_fuse/utils/progress.py

def console() -> Console:
    """Return a shared Rich Console instance with basic theming."""
    global _console
    if _console is None:
        _console = Console(theme=Theme({"info": "cyan", "warn": "yellow", "error": "bold red"}))
    return _console

create_progress_tracker

create_progress_tracker(total_steps: int, *, event_sinks: Optional[LogEventSinks] = None) -> Dict[str, object]

Return a simple progress tracker structure.

Source code in src/spark_fuse/utils/progress.py

def create_progress_tracker(
    total_steps: int,
    *,
    event_sinks: Optional[LogEventSinks] = None,
) -> Dict[str, object]:
    """Return a simple progress tracker structure."""
    tracker = {
        "current": 0,
        "total": float(total_steps),
        "start": time.perf_counter(),
        "last": None,
        "last_progress": None,
        "bar": None,
    }
    if event_sinks is not None:
        tracker["event_sinks"] = _normalize_event_sinks(event_sinks)
    return tracker

enable_spark_logging

enable_spark_logging(spark: SparkSession, *, level: str = 'INFO', categories: Optional[Iterable[str]] = None) -> None

Promote Spark log verbosity so shuffle spilling and scheduler details surface.

Spark's default log level is WARN, which hides shuffle spill diagnostics, broadcast cache messages, and other executor hints. This helper raises the log level both through the public SparkContext.setLogLevel API and directly on the underlying Log4j loggers that emit the spill messages.

Parameters:

Name	Type	Description	Default
spark	SparkSession	Active SparkSession instance.	required
level	str	Target log level (case insensitive), defaults to "INFO".	'INFO'
categories	Optional[Iterable[str]]	Optional iterable of Log4j logger names to tune. When omitted, a curated set covering storage, scheduler, and shuffle components is used.	None

Source code in src/spark_fuse/utils/progress.py

def enable_spark_logging(
    spark: SparkSession,
    *,
    level: str = "INFO",
    categories: Optional[Iterable[str]] = None,
) -> None:
    """Promote Spark log verbosity so shuffle spilling and scheduler details surface.

    Spark's default log level is ``WARN``, which hides shuffle spill diagnostics,
    broadcast cache messages, and other executor hints. This helper raises the log
    level both through the public ``SparkContext.setLogLevel`` API and directly on
    the underlying Log4j loggers that emit the spill messages.

    Args:
        spark: Active ``SparkSession`` instance.
        level: Target log level (case insensitive), defaults to ``"INFO"``.
        categories: Optional iterable of Log4j logger names to tune. When omitted,
            a curated set covering storage, scheduler, and shuffle components is used.
    """

    sc = spark.sparkContext
    sc.setLogLevel(level.upper())

    jvm = getattr(spark, "_jvm", None)
    if jvm is None:
        return

    log_manager = jvm.org.apache.log4j.LogManager
    log4j_level = jvm.org.apache.log4j.Level.toLevel(level.upper())

    for name in categories or _DEFAULT_SPARK_LOGGERS:
        logger = log_manager.getLogger(name)
        if logger is not None:
            logger.setLevel(log4j_level)

log_event

log_event(tracker: Dict[str, object], logger: Console, label: str, *, event: Optional[str] = None, advance: int = 1, sinks: Optional[LogEventSinks] = None, show_html: bool = False) -> None

Advance a progress tracker, emit a validated event record, and log with tqdm.

Source code in src/spark_fuse/utils/progress.py

def log_event(
    tracker: Dict[str, object],
    logger: Console,
    label: str,
    *,
    event: Optional[str] = None,
    advance: int = 1,
    sinks: Optional[LogEventSinks] = None,
    show_html: bool = False,
) -> None:
    """Advance a progress tracker, emit a validated event record, and log with tqdm."""

    now = time.perf_counter()
    last_event = tracker.get("last") or tracker["start"]
    last_progress = tracker.get("last_progress") or tracker["start"]

    advance_by = int(advance)
    current = int(tracker.get("current", 0))
    if advance_by:
        current += advance_by

    total = int(tracker.get("total") or 1)

    elapsed_event = now - float(last_event)
    elapsed_progress = now - float(last_progress)
    total_elapsed = now - float(tracker["start"])

    record = LogEventRecord(
        label=label,
        event=event,
        step=current,
        total=total,
        elapsed_seconds=elapsed_event,
        total_elapsed_seconds=total_elapsed,
        timestamp=time.time(),
    )

    event_sinks = _normalize_event_sinks(sinks if sinks is not None else tracker.get("event_sinks"))
    if event_sinks:
        _emit_log_event(record, event_sinks)
    if show_html:
        event_type = (record.event or "info").lower()
        style = _EVENT_STYLES.get(event_type, _EVENT_STYLES["info"])

        html = f"""
        <div style="
            font-family: ui-sans-serif, system-ui, -apple-system, Segoe UI, Roboto, Helvetica, Arial;
            padding: 10px 12px;
            border-radius: 8px;
            margin: 6px 0;
            background-color: {style["bg"]};
            border: 1px solid {style["border"]};
            color: {style["color"]};
        ">
            <div style="font-weight: 600; margin-bottom: 6px; display: flex; align-items: center; gap: 6px;">
                <span>{style["icon"]}</span>
                <span>{record.label}</span>
                {f"<span style='font-weight:400;'>— {record.event}</span>" if record.event else ""}
            </div>

            <div style="font-size: 12px; line-height: 1.4;">
                <div><strong>Step:</strong> {record.step} / {record.total}</div>
                <div><strong>Δ since progress:</strong> {elapsed_progress:.2f}s</div>
                <div><strong>Total:</strong> {total_elapsed:.2f}s</div>
            </div>
        </div>
        """.strip()

        _display_html_databricks(html)

    tracker["current"] = current
    tracker["last"] = now
    if advance_by:
        tracker["last_progress"] = now

    bar = tracker.get("bar")
    if bar is None:
        bar = tqdm(
            total=total,
            file=logger.file,
            dynamic_ncols=True,
        )
        tracker["bar"] = bar
    elif bar.total != total:
        bar.total = total

    bar.set_description_str(_format_event_label(label, record.event))
    bar.set_postfix_str(f"+{elapsed_progress:.2f}s, total {total_elapsed:.2f}s")
    if advance_by:
        bar.update(advance_by)
    else:
        bar.refresh()

    if current >= total:
        bar.close()
        tracker["bar"] = None

log_exception

log_exception(tracker: Dict[str, object], logger: Console, label: str, exc: BaseException, *, advance: int = 1, sinks: Optional[LogEventSinks] = None, include_traceback: bool = True, show_html: bool = False) -> None

Log an exception as an error event and optionally emit its traceback.

Source code in src/spark_fuse/utils/progress.py

def log_exception(
    tracker: Dict[str, object],
    logger: Console,
    label: str,
    exc: BaseException,
    *,
    advance: int = 1,
    sinks: Optional[LogEventSinks] = None,
    include_traceback: bool = True,
    show_html: bool = False,
) -> None:
    """Log an exception as an error event and optionally emit its traceback."""
    message = f"{label}: {exc}" if str(exc) else label
    log_error(tracker, logger, message, advance=advance, sinks=sinks, show_html=show_html)

    if include_traceback:
        trace = "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)).rstrip()
        if trace:
            logger.print(trace, style="error", markup=False)