Skip to content

Logging

Roxxel's Logger implements a high-performance, queue-based logging architecture. To prevent system I/O (like writes to stdout, files, and CSVs) from slowing down high-throughput TPU/GPU training loops, all log writing is offloaded to background worker threads.

Why the Context Manager is Critical

Because logging happens asynchronously on a separate background thread, standard print statements or un-managed logs are highly vulnerable. If your TPU/GPU throws an Out of Memory (OOM) error or JAX crashes: 1. The main thread terminates instantly. 2. The logging queue gets cut off. 3. The most critical debug messages/tracebacks at the end of the run are lost.

Using the with context manager solves this completely:

import time
from roxxel import Logger

# Initialize the async logger context
with Logger(log_dir="./run_directory") as logger:
    logger.log_message("Initializing deep pre-training cluster...")

    # Under the hood, any exceptions raised here are caught by the context manager.
    # It logs the traceback, drains/flushes the async queue to disk, and then propagates the error.
    time.sleep(1)
    raise RuntimeError("TPU Device Out of Memory!")

# The background thread is safely joined and shut down here.

Automatic Crash Traceback Capture

When a crash occurs inside the Logger context: 1. The traceback is immediately intercepted. 2. It writes the traceback cleanly to both stdout and {log_dir}/{prefix}_system.log. 3. It forces the queue to block and drain entirely, guaranteeing that every single log line is written to disk before the program terminates.

Multi-Host TPU Rank-Zero Filter

When scaling JAX code across TPU Pods or multi-node GPU clusters, standard print statements are executed by every worker node simultaneously, resulting in corrupted, duplicate log files.

Roxxel's Logger detects JAX rank automatically: * Only Rank 0 writes messages to stdout, log files, or CSVs. * Other ranks (1..N) execute logging statements as safe noop operations, preventing file conflicts and terminal pollution.

Asynchronous Metrics CSV Logging

You can record training metrics (like loss, learning rate, and perplexity) directly to a CSV file without blocking JAX JIT execution:

from roxxel import Logger

with Logger(log_dir="./logs", filename_prefix="run_alpha") as logger:
    for step in range(100):
        # Your JAX training loop here
        loss = 2.5 - (step * 0.01)
        lr = 3e-4

        # This pushes metrics to a background queue instantly (0ms overhead)
        logger.log_metrics_summary(
            step=step,
            metrics={"loss": loss, "lr": lr}
        )
This produces logs/run_alpha_metrics.csv automatically with properly aligned column headers on step resumption.

API Reference

roxxel.logging.Logger

A lightweight, asynchronous-friendly Rank-0 Logger for Roxxel.

Provides interactive terminal progress bars via tqdm and optional cloud-based run tracking via wandb (Weights & Biases). Coordinates stdout printing exclusively on Rank 0 to prevent terminal clutter in distributed JAX/Flax training runs.

Source code in roxxel/logging.py 
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
class Logger:
    """
    A lightweight, asynchronous-friendly Rank-0 Logger for Roxxel.

    Provides interactive terminal progress bars via `tqdm` and optional 
    cloud-based run tracking via `wandb` (Weights & Biases). Coordinates 
    stdout printing exclusively on Rank 0 to prevent terminal clutter in 
    distributed JAX/Flax training runs.
    """
    def __init__(self, log_dir: str = None, project: str = None, name: str = None, config: dict = None):
        """
        Args:
            log_dir (str, optional): Root directory to save logs (no-op here, kept for backward compatibility).
            project (str, optional): Weights & Biases project name. If provided, initializes wandb.
            name (str, optional): Display name for the Weights & Biases run.
            config (dict, optional): Hyperparameter dictionary to save to the wandb run configuration.
        """
        import jax
        self.is_rank_zero = (jax.process_index() == 0)
        self.project = project
        self.pbar = None
        self._log_dir = log_dir

        if self.is_rank_zero:
            if log_dir:
                os.makedirs(log_dir, exist_ok=True)
            if project:
                import wandb
                wandb.init(project=project, name=name, config=config)

    def init_pbar(self, total_steps: int, initial_step: int = 0):
        """
        Lazily initializes the tqdm progress bar on Rank 0 once the 
        total optimization horizon is determined.
        """
        if self.is_rank_zero and self.pbar is None:
            from tqdm.auto import tqdm
            self.pbar = tqdm(total=total_steps, initial=initial_step, desc="Training")

    def update_pbar(self, step: int):
        """Updates the progress bar step count on Rank 0."""
        if self.is_rank_zero and self.pbar is not None:
            self.pbar.update(step - self.pbar.n)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.is_rank_zero:
            if exc_type is not None:
                # Intercept crash tracebacks to log them cleanly to stderr
                import traceback
                print("\n❌ CRITICAL: Uncaught exception occurred during execution!", file=sys.stderr)
                traceback.print_exception(exc_type, exc_val, exc_tb, file=sys.stderr)

                if self.project:
                    import wandb
                    wandb.finish(exit_code=1)
            else:
                self.close()
        return False

    def log_message(self, message: str, level: int = None):
        """Prints a message to stdout on Rank 0, pushing it safely above the active progress bar."""
        if self.is_rank_zero:
            if self.pbar is not None:
                from tqdm.auto import tqdm
                tqdm.write(message)
            else:
                print(message)

    def log_metrics_summary(self, step: int, metrics: dict):
        """Updates the progress bar postfix metrics and pushes summaries asynchronously to WandB."""
        if self.is_rank_zero:
            if self.pbar is not None:
                self.pbar.update(step - self.pbar.n)
                # Format floating points to avoid float representation clutter in terminal
                formatted_metrics = {
                    k: f"{v:.4f}" if isinstance(v, float) else str(v)
                    for k, v in metrics.items()
                }
                self.pbar.set_postfix(**formatted_metrics)
            if self.project:
                import wandb
                wandb.log(metrics, step=step)

    def close(self):
        """Cleans up and finalizes progress bars and wandb runs."""
        if self.is_rank_zero:
            if self.pbar is not None:
                self.pbar.close()
                self.pbar = None
            if self.project:
                import wandb
                wandb.finish()

__init__(log_dir=None, project=None, name=None, config=None)

Parameters:

Name Type Description Default
log_dir str

Root directory to save logs (no-op here, kept for backward compatibility).

 None
project str

Weights & Biases project name. If provided, initializes wandb.

 None
name str

Display name for the Weights & Biases run.

 None
config dict

Hyperparameter dictionary to save to the wandb run configuration.

 None
Source code in roxxel/logging.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
def __init__(self, log_dir: str = None, project: str = None, name: str = None, config: dict = None):
    """
    Args:
        log_dir (str, optional): Root directory to save logs (no-op here, kept for backward compatibility).
        project (str, optional): Weights & Biases project name. If provided, initializes wandb.
        name (str, optional): Display name for the Weights & Biases run.
        config (dict, optional): Hyperparameter dictionary to save to the wandb run configuration.
    """
    import jax
    self.is_rank_zero = (jax.process_index() == 0)
    self.project = project
    self.pbar = None
    self._log_dir = log_dir

    if self.is_rank_zero:
        if log_dir:
            os.makedirs(log_dir, exist_ok=True)
        if project:
            import wandb
            wandb.init(project=project, name=name, config=config)

close()

Cleans up and finalizes progress bars and wandb runs.

Source code in roxxel/logging.py 
91
92
93
94
95
96
97
98
99
def close(self):
    """Cleans up and finalizes progress bars and wandb runs."""
    if self.is_rank_zero:
        if self.pbar is not None:
            self.pbar.close()
            self.pbar = None
        if self.project:
            import wandb
            wandb.finish()

init_pbar(total_steps, initial_step=0)

Lazily initializes the tqdm progress bar on Rank 0 once the total optimization horizon is determined.

Source code in roxxel/logging.py 
35
36
37
38
39
40
41
42
def init_pbar(self, total_steps: int, initial_step: int = 0):
    """
    Lazily initializes the tqdm progress bar on Rank 0 once the 
    total optimization horizon is determined.
    """
    if self.is_rank_zero and self.pbar is None:
        from tqdm.auto import tqdm
        self.pbar = tqdm(total=total_steps, initial=initial_step, desc="Training")

log_message(message, level=None)

Prints a message to stdout on Rank 0, pushing it safely above the active progress bar.

Source code in roxxel/logging.py 
67
68
69
70
71
72
73
74
def log_message(self, message: str, level: int = None):
    """Prints a message to stdout on Rank 0, pushing it safely above the active progress bar."""
    if self.is_rank_zero:
        if self.pbar is not None:
            from tqdm.auto import tqdm
            tqdm.write(message)
        else:
            print(message)

log_metrics_summary(step, metrics)

Updates the progress bar postfix metrics and pushes summaries asynchronously to WandB.

Source code in roxxel/logging.py 
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def log_metrics_summary(self, step: int, metrics: dict):
    """Updates the progress bar postfix metrics and pushes summaries asynchronously to WandB."""
    if self.is_rank_zero:
        if self.pbar is not None:
            self.pbar.update(step - self.pbar.n)
            # Format floating points to avoid float representation clutter in terminal
            formatted_metrics = {
                k: f"{v:.4f}" if isinstance(v, float) else str(v)
                for k, v in metrics.items()
            }
            self.pbar.set_postfix(**formatted_metrics)
        if self.project:
            import wandb
            wandb.log(metrics, step=step)

update_pbar(step)

Updates the progress bar step count on Rank 0.

Source code in roxxel/logging.py 
44
45
46
47
def update_pbar(self, step: int):
    """Updates the progress bar step count on Rank 0."""
    if self.is_rank_zero and self.pbar is not None:
        self.pbar.update(step - self.pbar.n)