TensorFlowAgentParallelRunner

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

TensorFlow Agent Parallel runner class. This is experimental.

This runner is pams.runners.MultiProcessAgentParallelRunner whose worker processes are prepared for agents that use TensorFlow in pams.agents.Agent.submit_orders(), e.g., Keras models. The simulation results are identical to those of pams.runners.SequentialRunner with the same settings and the same seed as long as the agents are deterministic and TensorFlow uses the same number of threads (see below).

TensorFlow is an optional dependency of PAMS: import pams does not import it, and this runner imports it when it is created. If TensorFlow is not installed, an ImportError is raised. Install it by pip install tensorflow (or pip install tensorflow-cpu for the CPU-only build).

The worker processes are started by spawn by default, i.e., default_start_method is "spawn", because TensorFlow is not fork-safe. Once TensorFlow runs an operation, the process has thread pools (and CUDA contexts for GPUs), but a process started by fork has only a copy of the thread calling fork. In the forked process, operations that use the thread pools of TensorFlow, e.g., a large matrix multiplication, hang, GPUs cannot be used, and TensorFlow cannot be configured anymore. Such a fork happens easily, e.g., when agents build their models on the main process or a pams.runners.SequentialRunner is run before on the same process. A spawned process starts a new interpreter and initializes TensorFlow by itself. simulation.startMethod in the config still takes precedence: forkserver is also safe if TensorFlow is not used at the top level of the main module, and fork works only if TensorFlow is not used on the main process before the simulation.

Each worker process configures TensorFlow once before it runs any task, by the following keys of simulation in the config:

• tensorflowIntraOpThreads (int ≥ 1): number of threads that TensorFlow uses to run one operation, e.g., a matrix multiplication, in each worker process. The default is the number of CPUs divided by the number of worker processes running tasks at the same time, i.e., the smaller of numParallel and the largest maxNormalOrders of the sessions (at least 1), so that the worker processes do not use more threads than the CPUs in total. TensorFlow uses all the CPUs in each process by default. The results of TensorFlow can depend on this number (see below).

• tensorflowInterOpThreads (int ≥ 1): number of threads that TensorFlow uses to run independent operations at the same time in each worker process. The default is 1.

• tensorflowGpuMemoryGrowth (bool): whether to enable memory growth for all the visible GPUs in each worker process. The default is true. Without memory growth, TensorFlow allocates almost all the memory of a GPU in the first process using it, so the other worker processes fail with out of memory errors when they share the GPU. Set it to false to keep the default behavior of TensorFlow, e.g., when each worker process uses its own GPU.

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. These keys do not configure TensorFlow on the main process. If tensorflowGpuMemoryGrowth is true, TensorFlow lists the GPUs before the initializer is called, so the initializer cannot select the GPUs of a worker process by the environment variable CUDA_VISIBLE_DEVICES; use tf.config.set_visible_devices instead.

Note

As pams.runners.MultiProcessAgentParallelRunner, this runner pickles the agents and the markets, i.e., the whole simulation, for each task. tf.Tensor and tf.Variable held by agents are pickled as copies of their values, and Keras 3 models are pickled by saving them in the .keras format, which takes tens of milliseconds even for a small model, in every task and even in the tasks of other agents. Objects made by tf.function cannot be pickled even if they are defined at the top level of a module. Therefore, agents should not hold models as their attributes. Instead, an agent should hold a key of its model, such as the path of a saved model, and get the model from a cache at the top level of a module, e.g., a function decorated by functools.lru_cache() that loads the model. Then, each worker process loads the model once, at its first task or in the worker initializer, and reuses it in the later tasks, and pams.runners.SequentialRunner can run the same agents.

Note

Only the orders and the state of the agent’s pseudo random number generator (prng) are returned from the worker processes. The random states of TensorFlow, e.g., the global seed and tf.random.Generator, are not returned, and which worker process runs an agent changes from run to run. For the same results as pams.runners.SequentialRunner, agents should draw random numbers from prng, e.g., the seeds of stateless random operations such as tf.random.stateless_normal.

Models must also give the same outputs on every process. On CPUs, the outputs of the operations that TensorFlow splits among threads, e.g., reductions and matrix multiplications over long axes, depend on the number of intra-op threads, even with tf.config.experimental.enable_op_determinism. The default number of this runner depends on the number of CPUs, numParallel and maxNormalOrders, and it can differ from that of the main process, where TensorFlow uses all the CPUs. Therefore, for results identical to pams.runners.SequentialRunner, or across different numParallel, set tensorflowIntraOpThreads explicitly, and set the same number on the process running pams.runners.SequentialRunner by tf.config.threading.set_intra_op_parallelism_threads before TensorFlow is initialized there. On GPUs, some operations are not deterministic unless tf.config.experimental.enable_op_determinism is called on the main process and in the worker processes.

Warning

This runner makes the simulation faster only if the TensorFlow computations in pams.agents.Agent.submit_orders() are much heavier than the cost of pickling the simulation in every task. Moreover, each worker process imports and initializes TensorFlow at its start, which takes seconds and hundreds of megabytes of memory per process.

default_start_method: Optional[str] = 'spawn'

start method of the worker processes used when simulation.startMethod is not set in the config. TensorFlow is not fork-safe, so it is "spawn".

Type:

Optional[str]