def check_clean():
assert alloc_map.allocations[SlotName('cuda.device:1g.5gb-mig')][
DeviceId('a0')] == Decimal('0')
assert alloc_map.allocations[SlotName('cuda.device:1g.5gb-mig')][
DeviceId('a1')] == Decimal('0')
assert alloc_map.allocations[SlotName('cuda.shares')][DeviceId(
'a2')] == Decimal('0')
assert alloc_map.allocations[SlotName('cuda.shares')][DeviceId(
'a3')] == Decimal('0')
assert alloc_map.allocations[SlotName('cuda.device:3g.20gb-mig')][
DeviceId('a4')] == Decimal('0')
def test_fraction_alloc_map_random_generated_allocations():
alloc_map = FractionAllocMap(
device_slots={
DeviceId('a0'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(1.0)),
DeviceId('a1'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(1.0)),
},
allocation_strategy=FractionAllocationStrategy.FILL,
)
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == Decimal('0')
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal('0')
quantum = Decimal('.01')
for _ in range(5):
allocations = []
for _ in range(10):
result = alloc_map.allocate({
SlotName('x'):
Decimal(random.uniform(0, 0.1)).quantize(quantum, ROUND_DOWN),
})
allocations.append(result)
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] >= Decimal(
'0')
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] >= Decimal(
'0')
for a in allocations:
alloc_map.free(a)
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == Decimal(
'0')
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal(
'0')
async def list_devices(self) -> Collection[CUDADevice]:
if not self.enabled:
return []
all_devices = []
num_devices = libcudart.get_device_count()
for dev_id in map(lambda idx: DeviceId(str(idx)), range(num_devices)):
if dev_id in self.device_mask:
continue
raw_info = libcudart.get_device_props(int(dev_id))
sysfs_node_path = "/sys/bus/pci/devices/" \
f"{raw_info['pciBusID_str'].lower()}/numa_node"
node: Optional[int]
try:
node = int(Path(sysfs_node_path).read_text().strip())
except OSError:
node = None
dev_uuid, raw_dev_uuid = None, raw_info.get('uuid', None)
if raw_dev_uuid is not None:
dev_uuid = str(uuid.UUID(bytes=raw_dev_uuid))
else:
dev_uuid = '00000000-0000-0000-0000-000000000000'
dev_info = CUDADevice(
device_id=dev_id,
hw_location=raw_info['pciBusID_str'],
numa_node=node,
memory_size=raw_info['totalGlobalMem'],
processing_units=raw_info['multiProcessorCount'],
model_name=raw_info['name'],
uuid=dev_uuid,
)
all_devices.append(dev_info)
return all_devices
async def gather_node_measures(
self, ctx: StatContext) -> Sequence[NodeMeasurement]:
_cstat = psutil.cpu_times(True)
q = Decimal('0.000')
total_cpu_used = cast(
Decimal,
sum((Decimal(c.user + c.system) * 1000).quantize(q)
for c in _cstat))
now, raw_interval = ctx.update_timestamp('cpu-node')
interval = Decimal(raw_interval * 1000).quantize(q)
return [
NodeMeasurement(
MetricKey('cpu_util'),
MetricTypes.UTILIZATION,
unit_hint='msec',
current_hook=lambda metric: metric.stats.diff,
per_node=Measurement(total_cpu_used, interval),
per_device={
DeviceId(str(idx)): Measurement(
(Decimal(c.user + c.system) * 1000).quantize(q),
interval,
)
for idx, c in enumerate(_cstat)
},
),
]
def test_exclusive_resource_slots():
alloc_map = DiscretePropertyAllocMap(
device_slots={
DeviceId('a0'):
DeviceSlotInfo(SlotTypes.UNIQUE,
SlotName('cuda.device:1g.5gb-mig'),
Decimal(1)), # noqa
DeviceId('a1'):
DeviceSlotInfo(SlotTypes.UNIQUE,
SlotName('cuda.device:1g.5gb-mig'),
Decimal(1)), # noqa
DeviceId('a2'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('cuda.device'),
Decimal(1)),
DeviceId('a3'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('cuda.device'),
Decimal(1)),
DeviceId('a4'):
DeviceSlotInfo(SlotTypes.UNIQUE,
SlotName('cuda.device:3g.20gb-mig'),
Decimal(1)), # noqa
},
exclusive_slot_types={
'cuda.device:*-mig', 'cuda.device', 'cuda.shares'
},
)
def check_clean():
assert alloc_map.allocations[SlotName('cuda.device:1g.5gb-mig')][
DeviceId('a0')] == Decimal('0')
assert alloc_map.allocations[SlotName('cuda.device:1g.5gb-mig')][
DeviceId('a1')] == Decimal('0')
assert alloc_map.allocations[SlotName('cuda.device')][DeviceId(
'a2')] == Decimal('0')
assert alloc_map.allocations[SlotName('cuda.device')][DeviceId(
'a3')] == Decimal('0')
assert alloc_map.allocations[SlotName('cuda.device:3g.20gb-mig')][
DeviceId('a4')] == Decimal('0')
with pytest.raises(InvalidResourceCombination):
alloc_map.allocate({
SlotName('cuda.device'): Decimal('2'),
SlotName('cuda.device:1g.5gb-mig'): Decimal('1'),
})
check_clean()
async def restore_from_container(
self,
container: Container,
alloc_map: AbstractAllocMap,
) -> None:
assert isinstance(alloc_map, DiscretePropertyAllocMap)
memory_limit = container.backend_obj['HostConfig']['Memory']
alloc_map.apply_allocation({
SlotName('mem'): {
DeviceId('root'): memory_limit
},
})
async def list_devices(self) -> Collection[MemoryDevice]:
# TODO: support NUMA?
memory_size = psutil.virtual_memory().total
return [
MemoryDevice(
device_id=DeviceId('root'),
hw_location='root',
numa_node=0,
memory_size=memory_size,
processing_units=0,
)
]
async def init(self, context: Any = None) -> None:
rx_triple_version = re.compile(r'(\d+\.\d+\.\d+)')
# Check nvidia-docker and docker versions
try:
proc = await asyncio.create_subprocess_exec(
'nvidia-docker', 'version', '-f', '{{json .}}',
stdout=asyncio.subprocess.PIPE,
)
stdout, _ = await proc.communicate()
lines = stdout.decode().splitlines()
except FileNotFoundError:
log.warning('nvidia-docker is not installed.')
log.info('CUDA acceleration is disabled.')
self.enabled = False
return
m = rx_triple_version.search(lines[0])
if m:
self.nvdocker_version = tuple(map(int, m.group(1).split('.')))
else:
log.error('could not detect nvidia-docker version!')
log.info('CUDA acceleration is disabled.')
self.enabled = False
return
docker_version_data = json.loads(lines[1])
m = rx_triple_version.search(docker_version_data['Server']['Version'])
if m:
self.docker_version = tuple(map(int, m.group(1).split('.')))
else:
log.error('could not detect docker version!')
log.info('CUDA acceleration is disabled.')
self.enabled = False
return
raw_device_mask = self.plugin_config.get('device_mask')
if raw_device_mask is not None:
self.device_mask = [
*map(lambda dev_id: DeviceId(dev_id), raw_device_mask.split(','))
]
try:
detected_devices = await self.list_devices()
log.info('detected devices:\n' + pformat(detected_devices))
log.info('nvidia-docker version: {}', self.nvdocker_version)
log.info('CUDA acceleration is enabled.')
except ImportError:
log.warning('could not load the CUDA runtime library.')
log.info('CUDA acceleration is disabled.')
self.enabled = False
except RuntimeError as e:
log.warning('CUDA init error: {}', e)
log.info('CUDA acceleration is disabled.')
self.enabled = False
async def list_devices(self) -> Collection[CPUDevice]:
cores = await libnuma.get_available_cores()
overcommit_factor = int(
os.environ.get('BACKEND_CPU_OVERCOMMIT_FACTOR', '1'))
assert 1 <= overcommit_factor <= 4
return [
CPUDevice(
device_id=DeviceId(str(core_idx)),
hw_location='root',
numa_node=libnuma.node_of_cpu(core_idx),
memory_size=0,
processing_units=1 * overcommit_factor,
) for core_idx in sorted(cores)
]
async def gather_node_measures(
self,
ctx: StatContext,
) -> Sequence[NodeMeasurement]:
dev_count = 0
mem_avail_total = 0
mem_used_total = 0
mem_stats = {}
util_total = 0
util_stats = {}
if self.enabled:
try:
dev_count = libnvml.get_device_count()
for dev_id in map(lambda idx: DeviceId(str(idx)), range(dev_count)):
if dev_id in self.device_mask:
continue
dev_stat = libnvml.get_device_stats(int(dev_id))
mem_avail_total += dev_stat.mem_total
mem_used_total += dev_stat.mem_used
mem_stats[dev_id] = Measurement(Decimal(dev_stat.mem_used),
Decimal(dev_stat.mem_total))
util_total += dev_stat.gpu_util
util_stats[dev_id] = Measurement(Decimal(dev_stat.gpu_util), Decimal(100))
except ImportError:
log.warning('gather_node_measure(): NVML library is not found')
except LibraryError as e:
log.warning('gather_node_measure(): {!r}', e)
return [
NodeMeasurement(
MetricKey('cuda_mem'),
MetricTypes.USAGE,
unit_hint='bytes',
stats_filter=frozenset({'max'}),
per_node=Measurement(Decimal(mem_used_total), Decimal(mem_avail_total)),
per_device=mem_stats,
),
NodeMeasurement(
MetricKey('cuda_util'),
MetricTypes.USAGE,
unit_hint='percent',
stats_filter=frozenset({'avg', 'max'}),
per_node=Measurement(Decimal(util_total), Decimal(dev_count * 100)),
per_device=util_stats,
),
]
def read_from_string(cls, text: str) -> 'KernelResourceSpec':
kvpairs = {}
for line in text.split('\n'):
if '=' not in line:
continue
key, val = line.strip().split('=', maxsplit=1)
kvpairs[key] = val
allocations = cast(
MutableMapping[DeviceName, MutableMapping[SlotName,
Mapping[DeviceId,
Decimal]]],
defaultdict(lambda: defaultdict(Decimal)),
)
for key, val in kvpairs.items():
if key.endswith('_SHARES'):
slot_name = SlotName(key[:-7].lower())
device_name = DeviceName(slot_name.split('.')[0])
per_device_alloc: MutableMapping[DeviceId, Decimal] = {}
for entry in val.split(','):
raw_dev_id, _, raw_alloc = entry.partition(':')
if not raw_dev_id or not raw_alloc:
continue
dev_id = DeviceId(raw_dev_id)
try:
if known_slot_types.get(slot_name, 'count') == 'bytes':
alloc = Decimal(BinarySize.from_str(raw_alloc))
else:
alloc = Decimal(raw_alloc)
except KeyError as e:
log.warning(
'A previously launched container has '
'unknown slot type: {}. Ignoring it.', e.args[0])
continue
per_device_alloc[dev_id] = alloc
allocations[device_name][slot_name] = per_device_alloc
mounts = [Mount.from_str(m) for m in kvpairs['MOUNTS'].split(',') if m]
return cls(
container_id=kvpairs.get('CID', 'unknown'),
scratch_disk_size=BinarySize.finite_from_str(
kvpairs['SCRATCH_SIZE']),
allocations=dict(allocations),
slots=ResourceSlot(json.loads(kvpairs['SLOTS'])),
mounts=mounts,
)
async def gather_node_measures(
self, ctx: StatContext) -> Sequence[NodeMeasurement]:
_mstat = psutil.virtual_memory()
total_mem_used_bytes = Decimal(_mstat.total - _mstat.available)
total_mem_capacity_bytes = Decimal(_mstat.total)
_nstat = psutil.net_io_counters()
net_rx_bytes = _nstat.bytes_recv
net_tx_bytes = _nstat.bytes_sent
def get_disk_stat():
pruned_disk_types = frozenset(['squashfs', 'vfat', 'tmpfs'])
total_disk_usage = Decimal(0)
total_disk_capacity = Decimal(0)
per_disk_stat = {}
for disk_info in psutil.disk_partitions():
if disk_info.fstype not in pruned_disk_types:
dstat = os.statvfs(disk_info.mountpoint)
disk_usage = Decimal(dstat.f_frsize *
(dstat.f_blocks - dstat.f_bavail))
disk_capacity = Decimal(dstat.f_frsize * dstat.f_blocks)
per_disk_stat[disk_info.device] = Measurement(
disk_usage, disk_capacity)
total_disk_usage += disk_usage
total_disk_capacity += disk_capacity
return total_disk_usage, total_disk_capacity, per_disk_stat
loop = current_loop()
total_disk_usage, total_disk_capacity, per_disk_stat = \
await loop.run_in_executor(None, get_disk_stat)
return [
NodeMeasurement(
MetricKey('mem'),
MetricTypes.USAGE,
unit_hint='bytes',
stats_filter=frozenset({'max'}),
per_node=Measurement(total_mem_used_bytes,
total_mem_capacity_bytes),
per_device={
DeviceId('root'):
Measurement(total_mem_used_bytes, total_mem_capacity_bytes)
},
),
NodeMeasurement(
MetricKey('disk'),
MetricTypes.USAGE,
unit_hint='bytes',
per_node=Measurement(total_disk_usage, total_disk_capacity),
per_device=per_disk_stat,
),
NodeMeasurement(
MetricKey('net_rx'),
MetricTypes.RATE,
unit_hint='bps',
current_hook=lambda metric: metric.stats.rate,
per_node=Measurement(Decimal(net_rx_bytes)),
per_device={
DeviceId('node'): Measurement(Decimal(net_rx_bytes))
},
),
NodeMeasurement(
MetricKey('net_tx'),
MetricTypes.RATE,
unit_hint='bps',
current_hook=lambda metric: metric.stats.rate,
per_node=Measurement(Decimal(net_tx_bytes)),
per_device={
DeviceId('node'): Measurement(Decimal(net_tx_bytes))
},
),
]
def test_fraction_alloc_map():
alloc_map = FractionAllocMap(
device_slots={
DeviceId('a0'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(1.0)),
DeviceId('a1'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(1.0)),
},
allocation_strategy=FractionAllocationStrategy.FILL,
)
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == Decimal('0')
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal('0')
result = alloc_map.allocate({
SlotName('x'): Decimal('1.5'),
})
assert result[SlotName('x')][DeviceId('a0')] == Decimal('1.0')
assert result[SlotName('x')][DeviceId('a1')] == Decimal('0.5')
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == Decimal(
'1.0')
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal(
'0.5')
with pytest.raises(InsufficientResource):
alloc_map.allocate({
SlotName('x'): Decimal('1.5'),
})
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == Decimal(
'1.0')
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal(
'0.5')
alloc_map.free(result)
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == Decimal('0')
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal('0')
def test_fraction_alloc_map_many_device():
alloc_map = FractionAllocMap(
device_slots={
DeviceId('a0'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(1.0)),
DeviceId('a1'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(1.0)),
DeviceId('a2'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(1.0)),
DeviceId('a3'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(1.0)),
DeviceId('a4'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(1.0)),
DeviceId('a5'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(1.0)),
DeviceId('a6'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(1.0)),
DeviceId('a7'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(1.0)),
},
allocation_strategy=FractionAllocationStrategy.FILL,
)
for idx in range(8):
assert alloc_map.allocations[SlotName('x')][DeviceId(
f'a{idx}')] == Decimal('0')
result = alloc_map.allocate({
SlotName('x'): Decimal('7.95'),
})
for idx in range(7):
assert result[SlotName('x')][DeviceId(f'a{idx}')] == Decimal('1.0')
assert result[SlotName('x')][DeviceId('a7')] == Decimal('0.95')
for idx in range(7):
assert alloc_map.allocations[SlotName('x')][DeviceId(
f'a{idx}')] == Decimal('1.0')
assert alloc_map.allocations[SlotName('x')][DeviceId('a7')] == Decimal(
'0.95')
with pytest.raises(InsufficientResource):
alloc_map.allocate({
SlotName('x'): Decimal('1.0'),
})
for idx in range(7):
assert alloc_map.allocations[SlotName('x')][DeviceId(
f'a{idx}')] == Decimal('1.0')
assert alloc_map.allocations[SlotName('x')][DeviceId('a7')] == Decimal(
'0.95')
alloc_map.free(result)
for idx in range(8):
assert alloc_map.allocations[SlotName('x')][DeviceId(
f'a{idx}')] == Decimal('0')
def test_heterogeneous_resource_slots_with_fractional_alloc_map():
alloc_map = FractionAllocMap(
device_slots={
DeviceId('a0'):
DeviceSlotInfo(SlotTypes.UNIQUE,
SlotName('cuda.device:1g.5gb-mig'),
Decimal(1)), # noqa
DeviceId('a1'):
DeviceSlotInfo(SlotTypes.UNIQUE,
SlotName('cuda.device:1g.5gb-mig'),
Decimal(1)), # noqa
DeviceId('a2'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('cuda.shares'),
Decimal('1.0')),
DeviceId('a3'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('cuda.shares'),
Decimal('1.0')),
DeviceId('a4'):
DeviceSlotInfo(SlotTypes.UNIQUE,
SlotName('cuda.device:3g.20gb-mig'),
Decimal(1)), # noqa
},
exclusive_slot_types={
'cuda.device:*-mig', 'cuda.device', 'cuda.shares'
},
allocation_strategy=FractionAllocationStrategy.FILL,
)
def check_clean():
assert alloc_map.allocations[SlotName('cuda.device:1g.5gb-mig')][
DeviceId('a0')] == Decimal('0')
assert alloc_map.allocations[SlotName('cuda.device:1g.5gb-mig')][
DeviceId('a1')] == Decimal('0')
assert alloc_map.allocations[SlotName('cuda.shares')][DeviceId(
'a2')] == Decimal('0')
assert alloc_map.allocations[SlotName('cuda.shares')][DeviceId(
'a3')] == Decimal('0')
assert alloc_map.allocations[SlotName('cuda.device:3g.20gb-mig')][
DeviceId('a4')] == Decimal('0')
check_clean()
# check allocation of non-unique slots
result = alloc_map.allocate({SlotName('cuda.shares'): Decimal('2.0')})
assert alloc_map.allocations[SlotName('cuda.device:1g.5gb-mig')][DeviceId(
'a0')] == Decimal('0')
assert alloc_map.allocations[SlotName('cuda.device:1g.5gb-mig')][DeviceId(
'a1')] == Decimal('0')
assert alloc_map.allocations[SlotName('cuda.shares')][DeviceId(
'a2')] == Decimal('1.0')
assert alloc_map.allocations[SlotName('cuda.shares')][DeviceId(
'a3')] == Decimal('1.0')
assert alloc_map.allocations[SlotName('cuda.device:3g.20gb-mig')][DeviceId(
'a4')] == Decimal('0')
alloc_map.free(result)
check_clean()
with pytest.raises(InsufficientResource):
alloc_map.allocate({SlotName('cuda.shares'): Decimal('2.5')})
check_clean()
# allocating zero means no-op.
alloc_map.allocate({SlotName('cuda.device:1g.5gb-mig'): Decimal('0')})
check_clean()
# any allocation request for unique slots should specify the amount 1.
with pytest.raises(InvalidResourceArgument):
alloc_map.allocate(
{SlotName('cuda.device:1g.5gb-mig'): Decimal('0.3')})
with pytest.raises(InvalidResourceArgument):
alloc_map.allocate(
{SlotName('cuda.device:1g.5gb-mig'): Decimal('1.5')})
check_clean()
# test alloaction of unique slots
result1 = alloc_map.allocate(
{SlotName('cuda.device:1g.5gb-mig'): Decimal('1')})
assert alloc_map.allocations[SlotName('cuda.device:1g.5gb-mig')][DeviceId(
'a0')] == Decimal('1')
assert alloc_map.allocations[SlotName('cuda.device:1g.5gb-mig')][DeviceId(
'a1')] == Decimal('0')
result2 = alloc_map.allocate(
{SlotName('cuda.device:1g.5gb-mig'): Decimal('1')})
assert alloc_map.allocations[SlotName('cuda.device:1g.5gb-mig')][DeviceId(
'a0')] == Decimal('1')
assert alloc_map.allocations[SlotName('cuda.device:1g.5gb-mig')][DeviceId(
'a1')] == Decimal('1')
with pytest.raises(InsufficientResource):
alloc_map.allocate({SlotName('cuda.device:1g.5gb-mig'): Decimal('1')})
alloc_map.free(result1)
alloc_map.free(result2)
check_clean()
def test_quantum_size(alloc_strategy):
alloc_map = FractionAllocMap(
device_slots={
DeviceId('a0'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(1)), # noqa
DeviceId('a1'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(1)), # noqa
},
quantum_size=Decimal("0.25"),
allocation_strategy=alloc_strategy,
)
result = alloc_map.allocate({
SlotName('x'): Decimal("0.5"),
})
assert sum(alloc_map.allocations[SlotName('x')].values()) == Decimal("0.5")
alloc_map.free(result)
result = alloc_map.allocate({
SlotName('x'): Decimal("1.5"),
})
assert sum(alloc_map.allocations[SlotName('x')].values()) == Decimal("1.5")
if alloc_strategy == FractionAllocationStrategy.EVENLY:
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == Decimal(
"0.75")
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal(
"0.75")
else:
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == Decimal(
"1.00")
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal(
"0.50")
alloc_map.free(result)
# inputs are not multiple of 0.25
with pytest.raises(NotMultipleOfQuantum):
alloc_map.allocate({
SlotName('x'): Decimal("0.52"),
})
with pytest.raises(NotMultipleOfQuantum):
alloc_map.allocate({
SlotName('x'): Decimal("0.42"),
})
with pytest.raises(NotMultipleOfQuantum):
alloc_map.allocate({
SlotName('x'): Decimal("3.99"),
})
if alloc_strategy == FractionAllocationStrategy.EVENLY:
# input IS multiple of 0.25 but the CALCULATED allocations are not multiple of 0.25
with pytest.raises(InsufficientResource, match="multiple-of-quantum"):
alloc_map.allocate({
SlotName('x'): Decimal("1.75"), # divided to 0.88 and 0.87
})
else:
# In this case, it satisfies the quantum condition, because the capacity of devices are
# multiples of the quantum.
alloc_map.allocate({
SlotName('x'): Decimal("1.75"),
})
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == Decimal(
"1.00")
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal(
"0.75")
# So let's change the situation.
alloc_map = FractionAllocMap(
device_slots={
DeviceId('a0'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'),
Decimal(1)), # noqa
DeviceId('a1'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'),
Decimal(1)), # noqa
},
quantum_size=Decimal("0.3"),
allocation_strategy=alloc_strategy,
)
with pytest.raises(NotMultipleOfQuantum):
alloc_map.allocate({
SlotName('x'): Decimal("0.5"),
})
with pytest.raises(InsufficientResource, match="multiple-of-quantum"):
alloc_map.allocate({
SlotName('x'): Decimal("1.2"),
})
def test_discrete_alloc_map_large_number():
alloc_map = DiscretePropertyAllocMap(device_slots={
DeviceId('a0'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(100)),
DeviceId('a1'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(100)),
}, )
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == 0
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == 0
result = alloc_map.allocate({
SlotName('x'): Decimal('130'),
})
assert result[SlotName('x')][DeviceId('a0')] == 100
assert result[SlotName('x')][DeviceId('a1')] == 30
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == 100
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == 30
with pytest.raises(InsufficientResource):
alloc_map.allocate({
SlotName('x'): Decimal('71'),
})
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == 100
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == 30
alloc_map.free(result)
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == 0
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == 0
def test_fraction_alloc_map_even_allocation_many_devices_2():
alloc_map = FractionAllocMap(
device_slots={
DeviceId('a0'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal('1.0')),
DeviceId('a1'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal('1.0')),
DeviceId('a2'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal('1.0')),
DeviceId('a3'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal('1.0')),
DeviceId('a4'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal('1.0')),
DeviceId('a5'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal('1.0')),
DeviceId('a6'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal('1.0')),
DeviceId('a7'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal('1.0')),
},
allocation_strategy=FractionAllocationStrategy.EVENLY,
)
result = alloc_map.allocate({SlotName('x'): Decimal('6')})
count_0 = 0
count_1 = 0
# NOTE: the even allocator favors the tail of device list when it fills up.
# So we rely on the counting of desire per-device allocations instead of matching
# the device index and the allocations.
for idx in range(8):
if alloc_map.allocations[SlotName('x')][DeviceId(
f'a{idx}')] == Decimal('1.0'):
count_1 += 1
if alloc_map.allocations[SlotName('x')][DeviceId(
f'a{idx}')] == Decimal('0'):
count_0 += 1
assert count_0 == 2
assert count_1 == 6
alloc_map.free(result)
for idx in range(8):
assert alloc_map.allocations[SlotName('x')][DeviceId(
f'a{idx}')] == Decimal('0')
def test_fraction_alloc_map_even_allocation_many_devices():
alloc_map = FractionAllocMap(
device_slots={
DeviceId('a0'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(2)),
DeviceId('a1'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(3)),
DeviceId('a2'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(3)),
DeviceId('a3'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(5)),
},
allocation_strategy=FractionAllocationStrategy.EVENLY,
)
result = alloc_map.allocate({SlotName('x'): Decimal('6')})
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal('3')
assert alloc_map.allocations[SlotName('x')][DeviceId('a2')] == Decimal('3')
alloc_map.free(result)
for idx in range(4):
assert alloc_map.allocations[SlotName('x')][DeviceId(
f'a{idx}')] == Decimal('0')
alloc_map = FractionAllocMap(
device_slots={
DeviceId('a0'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(1)),
DeviceId('a1'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(1.5)),
DeviceId('a2'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(2)),
DeviceId('a3'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(3)),
DeviceId('a4'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(3)),
DeviceId('a5'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(4)),
DeviceId('a6'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(4.5)),
DeviceId('a7'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(5)),
DeviceId('a8'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(5)),
},
allocation_strategy=FractionAllocationStrategy.EVENLY,
)
result = alloc_map.allocate({SlotName('x'): Decimal('6')},
min_memory=Decimal('2.5'))
assert alloc_map.allocations[SlotName('x')][DeviceId('a3')] == Decimal('3')
assert alloc_map.allocations[SlotName('x')][DeviceId('a4')] == Decimal('3')
alloc_map.free(result)
for idx in range(9):
assert alloc_map.allocations[SlotName('x')][DeviceId(
f'a{idx}')] == Decimal('0')
result = alloc_map.allocate({SlotName('x'): Decimal('11')},
min_memory=Decimal('0.84'))
assert alloc_map.allocations[SlotName('x')][DeviceId('a3')] == Decimal(
'2.75')
assert alloc_map.allocations[SlotName('x')][DeviceId('a4')] == Decimal(
'2.75')
assert alloc_map.allocations[SlotName('x')][DeviceId('a5')] == Decimal(
'2.75')
assert alloc_map.allocations[SlotName('x')][DeviceId('a5')] == Decimal(
'2.75')
alloc_map.free(result)
for idx in range(9):
assert alloc_map.allocations[SlotName('x')][DeviceId(
f'a{idx}')] == Decimal('0')
def test_fraction_alloc_map_even_allocation_fractions():
alloc_map = FractionAllocMap(
device_slots={
DeviceId('a0'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal('0.8')),
DeviceId('a1'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal('0.75')),
DeviceId('a2'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal('0.7')),
DeviceId('a3'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal('0.3')),
DeviceId('a4'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal('0.0')),
},
allocation_strategy=FractionAllocationStrategy.EVENLY,
)
result = alloc_map.allocate({SlotName('x'): Decimal('2.31')})
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == Decimal(
'0.67')
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal(
'0.67')
assert alloc_map.allocations[SlotName('x')][DeviceId('a2')] == Decimal(
'0.67')
assert alloc_map.allocations[SlotName('x')][DeviceId('a3')] == Decimal(
'0.3')
alloc_map.free(result)
for idx in range(4):
assert alloc_map.allocations[SlotName('x')][DeviceId(
f'a{idx}')] == Decimal('0')
result = alloc_map.allocate({SlotName('x'): Decimal('2')})
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == Decimal(
'0.67')
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal(
'0.67')
assert alloc_map.allocations[SlotName('x')][DeviceId('a2')] == Decimal(
'0.66')
alloc_map.free(result)
for idx in range(3):
assert alloc_map.allocations[SlotName('x')][DeviceId(
f'a{idx}')] == Decimal('0')
def test_fraction_alloc_map_even_allocation():
alloc_map = FractionAllocMap(
device_slots={
DeviceId('a0'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(0.05)),
DeviceId('a1'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(0.1)),
DeviceId('a2'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(0.2)),
DeviceId('a3'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(0.3)),
DeviceId('a4'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(0.0)),
},
allocation_strategy=FractionAllocationStrategy.EVENLY,
)
for idx in range(5):
assert alloc_map.allocations[SlotName('x')][DeviceId(
f'a{idx}')] == Decimal('0')
with pytest.raises(InsufficientResource):
alloc_map.allocate({
SlotName('x'): Decimal('0.66'),
})
with pytest.raises(InsufficientResource):
alloc_map.allocate({
SlotName('x'): Decimal('0.06'),
},
min_memory=Decimal(0.6))
for _ in range(20):
alloc_map.allocate({
SlotName('x'): Decimal('0.01'),
})
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == Decimal(
'0.05')
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal(
'0.1')
assert alloc_map.allocations[SlotName('x')][DeviceId('a2')] == Decimal(
'0.05')
alloc_map.free({
SlotName('x'): {
DeviceId('a0'): Decimal('0.05'),
DeviceId('a1'): Decimal('0.1'),
DeviceId('a2'): Decimal('0.05')
}
})
for idx in range(0):
assert alloc_map.allocations[SlotName('x')][DeviceId(
f'a{idx}')] == Decimal('0')
result = alloc_map.allocate({SlotName('x'): Decimal('0.2')})
assert alloc_map.allocations[SlotName('x')][DeviceId('a2')] == Decimal(
'0.2')
alloc_map.free(result)
assert alloc_map.allocations[SlotName('x')][DeviceId('a2')] == Decimal('0')
result = alloc_map.allocate({SlotName('x'): Decimal('0.2')},
min_memory=Decimal('0.25'))
assert alloc_map.allocations[SlotName('x')][DeviceId('a3')] == Decimal(
'0.2')
alloc_map.free(result)
for idx in range(5):
assert alloc_map.allocations[SlotName('x')][DeviceId(
f'a{idx}')] == Decimal('0')
result = alloc_map.allocate({SlotName('x'): Decimal('0.5')})
assert alloc_map.allocations[SlotName('x')][DeviceId('a2')] == Decimal(
'0.2')
assert alloc_map.allocations[SlotName('x')][DeviceId('a3')] == Decimal(
'0.3')
alloc_map.free(result)
for idx in range(5):
assert alloc_map.allocations[SlotName('x')][DeviceId(
f'a{idx}')] == Decimal('0')
result = alloc_map.allocate({SlotName('x'): Decimal('0.65')})
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == Decimal(
'0.05')
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal(
'0.1')
assert alloc_map.allocations[SlotName('x')][DeviceId('a2')] == Decimal(
'0.2')
assert alloc_map.allocations[SlotName('x')][DeviceId('a3')] == Decimal(
'0.3')
alloc_map.free(result)
for idx in range(5):
assert alloc_map.allocations[SlotName('x')][DeviceId(
f'a{idx}')] == Decimal('0')
result = alloc_map.allocate({SlotName('x'): Decimal('0.6')},
min_memory=Decimal('0.1'))
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal(
'0.1')
assert alloc_map.allocations[SlotName('x')][DeviceId('a2')] == Decimal(
'0.2')
assert alloc_map.allocations[SlotName('x')][DeviceId('a3')] == Decimal(
'0.3')
alloc_map.free(result)
for idx in range(5):
assert alloc_map.allocations[SlotName('x')][DeviceId(
f'a{idx}')] == Decimal('0')
alloc_map = FractionAllocMap(device_slots={
DeviceId('a0'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal('0.3')),
DeviceId('a1'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal('0.3')),
DeviceId('a2'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal('0.9')),
}, )
result = alloc_map.allocate({SlotName('x'): Decimal('1')})
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == Decimal(
'0.3')
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal(
'0.3')
assert alloc_map.allocations[SlotName('x')][DeviceId('a2')] == Decimal(
'0.4')
def test_fraction_alloc_map_iteration():
alloc_map = FractionAllocMap(
device_slots={
DeviceId('a0'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(1.0)),
DeviceId('a1'):
DeviceSlotInfo(SlotTypes.COUNT, SlotName('x'), Decimal(1.0)),
},
allocation_strategy=FractionAllocationStrategy.FILL,
quantum_size=Decimal("0.00001"))
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == Decimal('0')
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal('0')
for _ in range(1000):
alloc_map.allocate({
SlotName('x'): Decimal('0.00001'),
})
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == Decimal(
'0.005')
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal(
'0.005')
alloc_map.free({SlotName('x'): {DeviceId('a0'): Decimal('0.00001')}})
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == Decimal(
'0.00499')
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal(
'0.005')
for _ in range(499):
alloc_map.free({SlotName('x'): {DeviceId('a0'): Decimal('0.00001')}})
assert alloc_map.allocations[SlotName('x')][DeviceId('a0')] == Decimal('0')
assert alloc_map.allocations[SlotName('x')][DeviceId('a1')] == Decimal(
'0.005')
