JaxAgentParallelRunner

class pams.runners.JaxAgentParallelRunner(settings, prng=None, logger=None, simulator_class=<class 'pams.simulator.Simulator'>)

JAX Agent Parallel runner class. This is experimental.

This runner is pams.runners.MultiProcessAgentParallelRunner for agents whose pams.agents.Agent.submit_orders() uses JAX, e.g., neural network models of Flax. The simulation results are identical to those of pams.runners.SequentialRunner with the same settings and the same seed if the agents are deterministic given their pseudo random number generators.

JAX is an optional dependency of pams. It is imported when this runner is created, and an ImportError is raised if it is not installed. import pams does not import it.

The worker processes are started by spawn by default, i.e., default_start_method is "spawn", because fork is unsafe for JAX. Once JAX is initialized on the main process, e.g., when the agents create their models, JAX runs its own threads. A forked worker process has only a copy of the thread calling fork, so a lock held by another thread is never released and the worker process can deadlock. JAX warns about it when fork is called. simulation.startMethod in the config still takes precedence.

Each worker process configures JAX once before it runs any task, by the following keys in simulation of the config. Each key overrides the corresponding environment variable, which the worker processes otherwise inherit from the main process:

• jaxPlatforms (str, Optional): platforms that JAX uses on the worker processes, e.g., "cpu" or "cuda", i.e., jax_platforms of JAX. If it is not set, JAX_PLATFORMS in the environment is used if it is set; otherwise, JAX chooses the platform.

• jaxPreallocate (bool, Optional): whether JAX preallocates the memory of GPUs on each worker process, i.e., XLA_PYTHON_CLIENT_PREALLOCATE. If it is not set, XLA_PYTHON_CLIENT_PREALLOCATE in the environment is used if it is set; otherwise, it is false, so that several worker processes can share a GPU, because JAX preallocates 75% of the memory of the GPU by default.

• jaxMemoryFraction (float, Optional): fraction of the memory of the GPU that JAX on each worker process can use, in (0, 1], i.e., XLA_PYTHON_CLIENT_MEM_FRACTION. If it is not set, XLA_PYTHON_CLIENT_MEM_FRACTION or XLA_CLIENT_MEM_FRACTION in the environment is used if it is set; otherwise, the default of JAX (0.75) is used.

Then, the initializer given by _get_worker_initializer is called as in pams.runners.MultiProcessAgentParallelRunner, e.g., to load a read-only model once per worker process. The environment variables set on the worker processes are given by _get_worker_environment, which subclasses can override.

Note

The worker processes persist during the simulation. Thus, a function compiled by jax.jit() is compiled once per worker process (for each shape and type of the arguments) at its first call, and the compiled function is reused in the following steps. The functions compiled by jax.jit() cannot be pickled, so define them at the top level of a module instead of keeping them as attributes of the agents.

Note

The agents are pickled and copied to a worker process for each task, including their jax arrays and Flax modules and parameters, which are copied to the default device of the worker process. Large models should be loaded once per worker process by the worker initializer instead (see pams.runners.MultiProcessAgentParallelRunner). Only the state of agent.prng is sent back to the main process, so the agents should derive their random numbers, including the keys of jax.random, from agent.prng.

Note

Do not run JAX, e.g., create jax arrays, at the top level of the modules imported by the worker processes, such as the main module, and do not pass jax arrays to the worker initializer; otherwise, JAX is initialized before it is configured, and the settings above are ignored. The main process is not configured by the settings above, so set XLA_PYTHON_CLIENT_PREALLOCATE=false in the environment too if the agents use JAX on a GPU when they are created on the main process.

default_start_method: Optional[str] = 'spawn'

start method of the worker processes used when simulation.startMethod is not set in the config. None means the default start method of the platform. Subclasses can override it, e.g., with "spawn" for libraries that are not fork-safe.

Type:

Optional[str]