dew.training.runtime

Process setup every recipe runs before it builds anything.

rlimits, the XLA flags, the compilation cache, the JAX distributed pool and the env vars wandb and the tokenizers read are the same in every recipe. That makes them library wiring, and the recipes call this once at the top of main().

Name	Summary
EXECUTION_TIMEOUT	How long one device execution of a process pool may run before XLA ends its process.
prepare_process	Raise the fd/core limits, set the env vars, join the JAX process pool.
cuda_plugin	Whether JAX’s CUDA plugin is installed, the one reader of XLA’s GPU flags; asked before the backend opens, which no other question can be.
Preempted	Trainer.fit stopped at a preemption notice, at step, and wrote that step’s checkpoint and data position.
PreemptionNotice	Whether a preemption notice reached the run, asked once a step, from its creation until close.
run_timestamp	Return process 0’s wall clock as %Y-%m-%d_%H:%M:%S, on every process.

EXECUTION_TIMEOUT

attribute source

EXECUTION_TIMEOUT = '30m'

How long one device execution of a process pool may run before XLA ends its process.

A rank that stalls without failing, blocked on a read or on a compile that waits for peers, leaves the other ranks inside a collective or a communicator’s setup. No GPU backend times that out and no process reports it, so the pool would hang for ever with every process alive. XLA’s execution watchdog ends a process whose execution runs past this, and dew launch, srun or the scheduler then stops the rest. An execution is a whole step or sampling loop, which can run for minutes, so the bound is generous; --xla_gpu_execution_terminate_timeout in XLA_FLAGS or xla_flags sets another.

prepare_process

function source

def prepare_process(
    wandb: Wandb | None = None,
    multi_host: bool | None = None,
    xla_flags: str | None = None,
    compilation_cache_dir: str | None = None,
    *,
    layout: Layout | None = None,
) -> None

Raise the fd/core limits, set the env vars, join the JAX process pool.

wandb is the run’s dew.config.Wandb, or None for a run without a tracker. Only its offline switch is read, and it has to be read before wandb opens a run.

jax.distributed.initialize() finds the coordinator from the environment on TPU pods and Slurm/GKE/Open MPI clusters. dew launch leaves the process count and rank in DEW_PROCESS_COUNT and DEW_PROCESS_ID, which jax has no variable for, and those are passed to it with JAX’s cluster detection off: the launcher placed its processes, and a Slurm step around it would otherwise pin each to the GPU at SLURM_LOCALID. On a machine with no cluster environment it raises a ValueError naming the missing coordinator address, the single-host signature. Every other failure propagates, since a pod run would otherwise continue on one host. multi_host=True requires the pool, multi_host=False never asks for it. A Slurm step of one task forms no pool unless the run asks for one with multi_host=True or mpirun started its ranks there, which JAX’s detection reads before Slurm’s: JAX would still start a pool of that one task, at a coordinator named after the node, which a container on the node need not resolve. A Slurm step of several tasks with fewer on this node than the GPUs its task sees is refused: JAX gives each task the GPU at its SLURM_LOCALID, and the others would sit idle.

xla_flags reaches XLA through the environment, which XLA reads when it opens a backend. So this call has to come before the first JAX call in the process, which makes it a recipe’s first line. A library user, who never runs a recipe, sets XLA_FLAGS in the environment.

The same Layout passed to Trainer selects CPU transaction ownership when host includes params. JAX_PLATFORMS must then permit CPU beside the accelerator. JAX_NUM_CPU_DEVICES, or the existing XLA flags, must establish one CPU device per local accelerator before this call. Validation never changes backend configuration after initialization.

A GPU pool keeps the persistent compilation cache when its jax keys a computation that spans processes alike on every one of them, as the jax Dew pins does (_pool_keys_alike). With another jax it compiles without the cache: some ranks would load a step that the others compile, and that compile waits for every rank for ever.

cuda_plugin

function source

def cuda_plugin() -> bool

Whether JAX’s CUDA plugin is installed, the one reader of XLA’s GPU flags; asked before the backend opens, which no other question can be.

Preempted

class source

class Preempted(step: int)

Trainer.fit stopped at a preemption notice, at step, and wrote that step’s checkpoint and data position. Uncaught, it ends the program with PREEMPTED_EXIT and no traceback, the way SIGTERM itself would have; the same program run again resumes from the checkpoint.

PreemptionNotice

class source

class PreemptionNotice()

Whether a preemption notice reached the run, asked once a step, from its creation until close.

A scheduler stops a job with SIGTERM and SIGKILLs it a grace period later: Slurm’s KillWait, Kubernetes’ termination grace period, a spot VM’s notice. In a pool, JAX’s preemption service takes the SIGTERM (XLA’s notifier replaces the handler, so the process runs on), shares the notice through the coordination service, and reached_preemption_sync_point agrees one step on every process, where the checkpoint is whole. A process outside any pool has no such service, and the notice is SIGTERM itself, caught until close. A pool whose preemption service is off (jax_enable_preemption_service) gets no notice, and SIGTERM ends it as it always did.

PreemptionNotice.close

def close() -> None

Put back the SIGTERM handler this notice replaced.

PreemptionNotice.reached

def reached(step: int) -> bool

Whether to stop at step: in a pool, whether every process agreed on it; alone, whether SIGTERM arrived.

run_timestamp

function source

def run_timestamp() -> str

Return process 0’s wall clock as %Y-%m-%d_%H:%M:%S, on every process.

A default run name carries it, and the name is the checkpoint directory every process writes into, so a process that read its own clock a second later would write into a different directory.