warp_drive.training.utils.device_child_process package¶

Submodules¶

warp_drive.training.utils.device_child_process.child_process_numba module¶

class warp_drive.training.utils.device_child_process.child_process_numba.NumbaDeviceContextProcessWrapper(*args, **kwargs)¶

Bases: warp_drive.training.utils.child_process_base.DeviceContextProcessWrapper

A worker process wrapper that will (1) open up a GPU context for both Numba and PyTorch (2) under the current GPU context, call the kernel to run on this GPU,

and get the data via DataManager under the current GPU context

warp_drive.training.utils.device_child_process.child_process_pycuda module¶

class warp_drive.training.utils.device_child_process.child_process_pycuda.PyCUDADeviceContextProcessWrapper(*args, **kwargs)¶