Currently, tf-yarn is only able to reserve GPUs with node-level granularity, i.e. it assumes that a GPU node has a capacity of a single container, and then uses all of the GPUs on that node. It is possible to restrict tf-yarn container to a subset of GPUs:
- augment
TaskSpec with num_gpus field,
- prior to running
_dispatch_task discover which GPUs are not in use, and list them explicitly in CUDA_VISIBLE_DEVICES.
Currently, tf-yarn is only able to reserve GPUs with node-level granularity, i.e. it assumes that a GPU node has a capacity of a single container, and then uses all of the GPUs on that node. It is possible to restrict tf-yarn container to a subset of GPUs:
TaskSpecwithnum_gpusfield,_dispatch_taskdiscover which GPUs are not in use, and list them explicitly in CUDA_VISIBLE_DEVICES.