Commit 69151813 authored by Alejandro Homs Puron's avatar Alejandro Homs Puron Committed by Samuel Debionne
Browse files

[CLIENT] Alloc multiple receivers to use the available GPUs

parent 026d5cb7
Pipeline #71834 passed with stages
in 14 minutes and 57 seconds
......@@ -6,6 +6,7 @@ import string
import copy
import itertools
import argparse
import glob
import numpy as np
from enum import Enum, auto
from functools import partial, reduce
......@@ -43,6 +44,11 @@ def cpp_2_python_name(n):
return string.capwords(n.replace('_', ' ')).replace(' ', '')
def get_gpu_devs():
# TODO: should be configurable
return glob.glob('/dev/nvidia[0-9]')
TypeData = namedtuple('Detector_TypeData', ['name', 'schema', 'klass'])
class ParamBase:
......@@ -499,6 +505,7 @@ class Detector:
if 'acq_params' in attrs:
acq_params = copy.deepcopy(attrs['acq_params'])
# Configure Round-Robin xfer parameters
xfer = acq_params.xfer
nb_recvs = len(self.__recvs)
has_round_robin = True
......@@ -517,6 +524,15 @@ class Detector:
xfer.time_slice.stride = 1
attrs['acq_params'] = acq_params
if 'proc_params' in attrs:
proc_params = copy.deepcopy(attrs['proc_params'])
# Configure receivers GPUs
nb_recvs = len(self.__recvs)
nb_gpus_per_system = len(get_gpu_devs())
if nb_recvs > 1 and nb_gpus_per_system > 1 and 'gpu' in proc_params:
proc_params.gpu.device_idx = dev.recv_rank % nb_gpus_per_system
attrs['proc_params'] = proc_params
return attrs
def _getDevs(self):
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment