Runtime

substrax.runtime declares the settings a JAX process starts with and applies them at the point where they still take effect.

import os
import subprocess
import sys

from substrax.runtime import JaxRuntime, apply_runtime, runtime_environment

runtime = JaxRuntime(platforms=("cpu",), cpu_devices=8, enable_x64=True)

# A child process on eight emulated CPU devices.
env = {**os.environ, **runtime_environment(runtime, os.environ)}
subprocess.run([sys.executable, "train.py"], env=env, check=True)

# This process, before it imports jax.
apply_runtime(runtime)

import jax

assert jax.device_count() == 8

runtime_environment renders each field as the variable jax or XLA reads when it starts:

Field	Variable
platforms	JAX_PLATFORMS
cpu_devices	JAX_NUM_CPU_DEVICES, which jax prefers over --xla_force_host_platform_device_count
enable_x64	JAX_ENABLE_X64
matmul_precision	JAX_DEFAULT_MATMUL_PRECISION
compilation_cache_dir	JAX_COMPILATION_CACHE_DIR
xla_flags	XLA_FLAGS, merged by flag name
preallocate	XLA_PYTHON_CLIENT_PREALLOCATE
memory_fraction	XLA_CLIENT_MEM_FRACTION

A requested XLA flag that is already set to a different value raises XlaFlagConflictError, so a device count or GPU flag exported by the caller is never silently replaced. A memory fraction raises while the deprecated XLA_PYTHON_CLIENT_MEM_FRACTION is still set, because jax refuses both at once. JaxRuntime applies no per-backend flag presets: xla_flags is empty unless a run names the flags it needs.

Once jax is imported, apply_runtime applies the CPU device count, 64-bit types, matmul precision and compilation cache directory through jax.config; jax refuses a new device count after its first operation, and that refusal is raised as RuntimeConfigurationError. Platforms, XLA flags and accelerator memory settings are read only when jax starts its backends, and jax gives no public signal of whether that has happened, so after the import they raise instead of being written where they may change nothing.

A test suite chooses its backend and emulated CPU devices before anything imports jax:

# conftest.py
import importlib.util
import os

from substrax.runtime import resolve_test_runtime, runtime_environment

runtime = resolve_test_runtime(
    os.environ,
    prefix="DATARAX_TEST_",
    cuda_plugin_available=importlib.util.find_spec("jax_cuda12_plugin") is not None,
)
os.environ.update(runtime_environment(runtime, os.environ))

Tests run on eight emulated CPU devices unless DATARAX_TEST_JAX_PLATFORMS asks for an accelerator or DATARAX_TEST_DEVICE_COUNT names another count; an inherited JAX_PLATFORMS does not move them onto a GPU.

Entry points configure logging inside main(), never at import:

from substrax.runtime import configure_entry_point_logging


def main() -> None:
    configure_entry_point_logging()

A second call replaces only the handler the first one installed, so it needs no force=True and leaves pytest's log capture in place. Importing substrax.runtime imports no jax and changes no process state.

The managed environment file

A repository's setup.sh writes the file its activate.sh sources once the project environment is synced:

uv run --no-sync python -m substrax.runtime.managed_env write --prefix DATARAX \
    --backend auto --output .datarax.env --set TF_CPP_MIN_LOG_LEVEL=1
python -m substrax.runtime.managed_env show --prefix DATARAX --env-file .datarax.env \
    --user-env .env --user-env .env.local

auto resolves to cuda12 on Linux when nvidia-smi lists a GPU, metal on Apple Silicon, and cpu otherwise. The file names every variable it manages in <PREFIX>_MANAGED_ENV_VARS, so activate.sh unsets them before sourcing a newer file, then exports or unsets each:

Variable	cpu	cuda12, metal
<PREFIX>_BACKEND, <PREFIX>_ENV_ROOT	the resolved backend and the project root	the same
JAX_PLATFORMS	cpu	unset, so jax picks the accelerator
JAX_ENABLE_X64	false	false
XLA_PYTHON_CLIENT_PREALLOCATE	unset	false
XLA_CLIENT_MEM_FRACTION	unset	0.75
XLA_PYTHON_CLIENT_MEM_FRACTION	unset	unset, because jax refuses it beside the current name

jaxlib reads preallocation and the memory fraction only when it builds the GPU plugin's options, so they are unset on the CPU. Each --set NAME=VALUE adds a repository variable to the file; a --set naming a variable the file manages itself raises.

substrax.runtime

JAX process configuration, declared once and applied before or after jax is imported.

JaxRuntime describes the settings. runtime_environment renders them as the environment of a process that has not imported jax yet (a child process, or a conftest before its first import jax), child_environment builds a child's whole environment from them, apply_runtime applies them to the current process, and resolve_test_runtime chooses the backend and emulated CPU devices of a test run. configure_entry_point_logging sets up the root logger from an entry point. Importing this package imports no jax and changes no process state.

JaxRuntime(*, platforms=None, cpu_devices=None, enable_x64=None, matmul_precision=None, compilation_cache_dir=None, xla_flags=(), preallocate=None, memory_fraction=None) dataclass

Declarative JAX process settings; None leaves the process's current value alone.

:func:~substrax.runtime.runtime_environment renders the settings as the environment variables jax and XLA read when a process starts, and :func:~substrax.runtime.apply_runtime applies them to the current process.

Attributes:

Name	Type	Description
platforms	tuple[str, ...] | None	The backends jax initialises, the first being the default (JAX_PLATFORMS).
cpu_devices	int | None	The number of CPU devices (JAX_NUM_CPU_DEVICES). jax prefers it over the --xla_force_host_platform_device_count flag.
enable_x64	bool | None	Whether 64-bit types are enabled (JAX_ENABLE_X64).
matmul_precision	MatmulPrecision | None	The default matrix-multiplication precision (JAX_DEFAULT_MATMUL_PRECISION).
compilation_cache_dir	Path | None	The absolute directory of the persistent compilation cache (JAX_COMPILATION_CACHE_DIR).
xla_flags	tuple[str, ...]	Flags merged into XLA_FLAGS by name, each --name or --name=value.
preallocate	bool | None	Whether the accelerator client preallocates its memory (XLA_PYTHON_CLIENT_PREALLOCATE).
memory_fraction	float | None	The fraction of accelerator memory the client may use, in (0, 1] (XLA_CLIENT_MEM_FRACTION).

__post_init__()

Validate every field.

Raises:

Type	Description
TypeError	If a field holds a value of the wrong type.
ValueError	If a field holds a value outside its domain.
XlaFlagConflictError	If one requested flag is given two different values.

RuntimeConfigurationError

Bases: RuntimeError

A JAX process setting cannot be applied, or conflicts with one already in place.

XlaFlagConflictError(flag, existing, requested)

Bases: RuntimeConfigurationError

One XLA flag was given two different values.

Attributes:

Name	Type	Description
flag	str	The flag name, with its leading dashes.
existing	str | None	The value already in place; None for a flag given without a value.
requested	str | None	The conflicting value; None for a flag given without a value.

Record the flag and both of its values.

Parameters:

Name	Type	Description	Default
flag	str	The flag name, with its leading dashes.	required
existing	str | None	The value already in place.	required
requested	str | None	The conflicting value.	required

apply_runtime(runtime, *, environ=None)

Apply runtime to the current process.

Before jax is imported, every field is written to environ as the variables jax and XLA read at start-up (see :func:~substrax.runtime.runtime_environment). Once jax is imported, the CPU device count, 64-bit types, matmul precision and compilation cache directory go through jax.config.update; jax refuses a new device count once its backends have started. Platforms, XLA flags and client memory settings are read only when the backends start, and jax offers no public signal of whether that has happened, so once jax is imported they raise instead of being written where they might change nothing.

Parameters:

Name	Type	Description	Default
runtime	JaxRuntime	The settings to apply.	required
environ	MutableMapping[str, str] | None	The environment written before jax is imported; os.environ when None.	None

Raises:

Type	Description
RuntimeConfigurationError	If jax is imported and runtime sets a platform, an XLA flag or a client memory setting, or jax refuses the CPU device count.
XlaFlagConflictError	If a requested flag has a different value in the current XLA_FLAGS.

configure_entry_point_logging(level=logging.INFO, *, stream=None, fmt='%(message)s')

Send log records to stream without replacing any other handler.

Call it inside main(), never at import. A second call replaces only the handler the first one installed and leaves every other handler, pytest's log capture included, in place, so it needs no force=True.

Parameters:

Name	Type	Description	Default
level	int	The level set on the root logger.	INFO
stream	TextIO | None	Where records are written; standard error when None.	None
fmt	str	The logging.Formatter format string.	'%(message)s'

child_environment(runtime, env=None, *, parent=None)

Return the complete environment of a child process whose JAX settings the caller chooses.

The child gets parent without any JAX_* or XLA_* variable, so settings exported in the shell that started this process cannot change what the child runs, then env, then runtime as :func:runtime_environment renders it on top of the two.

Parameters:

Name	Type	Description	Default
runtime	JaxRuntime	The child's JAX settings.	required
env	Mapping[str, str] | None	Variables set in the child before runtime is applied.	None
parent	Mapping[str, str] | None	The environment the child inherits; this process's when None.	None

Returns:

Type	Description
dict[str, str]	Variable names mapped to their values.

Raises:

Type	Description
XlaFlagConflictError	If a flag in runtime conflicts with XLA_FLAGS in env.
RuntimeConfigurationError	If runtime sets a memory fraction and env holds the deprecated XLA_PYTHON_CLIENT_MEM_FRACTION.

resolve_test_runtime(env, *, prefix, cuda_plugin_available, default_cpu_devices=8)

Choose the backend and the emulated CPU devices of a test run.

Tests run on the CPU unless {prefix}JAX_PLATFORMS names an accelerator, which is then used without CPU emulation. An inherited JAX_PLATFORMS, such as the cuda,cpu a developer shell exports, does not move tests onto an accelerator, so a local run matches CI. The CPU device count is {prefix}DEVICE_COUNT when set (0 for no emulation), otherwise the count JAX_NUM_CPU_DEVICES or the XLA device-count flag already names, otherwise default_cpu_devices.

Parameters:

Name	Type	Description	Default
env	Mapping[str, str]	The process environment.	required
prefix	str	The prefix of the repository's test variables, such as "DATARAX_TEST_".	required
cuda_plugin_available	bool	Whether a JAX CUDA plugin is installed.	required
default_cpu_devices	int	The emulated device count when nothing else names one.	8

Returns:

Type	Description
JaxRuntime	The runtime of the test process.

Raises:

Type	Description
RuntimeConfigurationError	If CUDA is requested and no JAX CUDA plugin is installed.
ValueError	If {prefix}DEVICE_COUNT is not a non-negative integer.

runtime_environment(runtime, base)

Return the environment variables that apply runtime to a process not yet running jax.

Only the variables runtime sets are returned, and base is read, never changed. A field replaces the variable base holds, except xla_flags, which merges into base's XLA_FLAGS by flag name. Pass the result to a child process on top of its environment, or write it to os.environ before importing jax.

Parameters:

Name	Type	Description	Default
runtime	JaxRuntime	The settings to render.	required
base	Mapping[str, str]	The environment the variables will be added to.	required

Returns:

Type	Description
dict[str, str]	Variable names mapped to their values.

Raises:

Type	Description
RuntimeConfigurationError	If runtime sets a memory fraction and base still holds the deprecated XLA_PYTHON_CLIENT_MEM_FRACTION, which jax refuses beside the current name.
XlaFlagConflictError	If a requested flag has a different value in base's XLA_FLAGS.

merge_xla_flags(existing, requested)

Merge requested flags into an XLA_FLAGS value by flag name.

Flags already present keep their position, and new flags follow in the order requested. A requested flag already set to the same value changes nothing, and one set to a different value raises, so a value already in the environment is never silently replaced.

Parameters:

Name	Type	Description	Default
existing	str	The current XLA_FLAGS value.	required
requested	Iterable[str]	Flags to add, each --name or --name=value.	required

Returns:

Type	Description
str	The merged XLA_FLAGS value, with flags separated by single spaces.

Raises:

Type	Description
ValueError	If a token is not a flag.
XlaFlagConflictError	If a requested flag has a different value from an existing one.

parse_xla_flags(value)

Parse an XLA_FLAGS value into flag names mapped to their values, in order.

Flags are separated by whitespace. --name=value maps --name to value, and a bare --name maps it to None. A flag given twice with the same value counts once.

Parameters:

Name	Type	Description	Default
value	str	The XLA_FLAGS string.	required

Returns:

Type	Description
dict[str, str | None]	Each flag name, with its leading dashes, mapped to its value.

Raises:

Type	Description
ValueError	If a token is not a flag.
XlaFlagConflictError	If one flag is given two different values.

substrax.runtime.managed_env

The managed environment file a repository's setup.sh writes and its activate.sh sources.

Run as python -m substrax.runtime.managed_env write --prefix DATARAX --backend auto --output .datarax.env after the project environment is synced. The file declares every variable it manages in <PREFIX>_MANAGED_ENV_VARS, so activate.sh can unset them before sourcing a new one, then exports or unsets each. The JAX settings come from :class:JaxRuntime through :func:runtime_environment, so the file uses the names jax reads today: the memory fraction is XLA_CLIENT_MEM_FRACTION, and the deprecated XLA_PYTHON_CLIENT_MEM_FRACTION, which jaxlib refuses beside it, is always unset. Preallocation and the memory fraction apply only to an accelerator backend: jaxlib reads them when it builds the GPU plugin's options, so on the CPU they are unset. A repository adds its own variables (for example TensorFlow's log level) with --set.

Importing this module imports no jax.

nvidia_gpu_visible()

Whether nvidia-smi lists at least one GPU.

resolve_backend(requested)

The concrete backend for requested, resolving auto from the host.

managed_environment(*, prefix, backend, project_root, extra)

Every managed variable mapped to its value, or None for a variable the file unsets.

Parameters:

Name	Type	Description	Default
prefix	str	The repository's variable prefix, such as "DATARAX".	required
backend	str	One of :data:BACKENDS; auto is resolved first.	required
project_root	Path	The repository root, written as <PREFIX>_ENV_ROOT.	required
extra	Mapping[str, str]	The repository's own variables, written as given.	required

Returns:

Type	Description
dict[str, str | None]	The managed variables in the order they are written.

Raises:

Type	Description
ValueError	If the prefix is not an upper-case shell name, the backend is unknown, or extra names a variable the file manages itself.

render(*, prefix, values)

The shell text of a managed environment file: the manifest, then one line per variable.

main(argv=None)

Run the write or show command.

Parameters:

Name	Type	Description	Default
argv	Sequence[str] | None	The arguments after the module name; the process arguments when None.	None

Returns:

Type	Description
int	The exit status.