def do_allocate(
self,
node_mgr: nodemanager.NodeManager,
allow_existing: bool,
all_or_nothing: bool,
) -> AllocationResult:
if self.__node_count > 0:
return node_mgr.allocate(
self._constraints,
node_count=self.__node_count,
allow_existing=allow_existing,
all_or_nothing=self.__colocated,
assignment_id=self.name,
)
assert self.iterations_remaining > 0
result = node_mgr.allocate(
self._constraints,
slot_count=self.iterations_remaining,
allow_existing=allow_existing,
# colocated is always all or nothing
all_or_nothing=all_or_nothing or self.colocated,
assignment_id=self.name,
)
if result:
self.iterations_remaining -= result.total_slots
return result
def test_over_allocate(node_mgr: NodeManager) -> None:
assert node_mgr.allocate({"node.nodearray": "htc"}, node_count=1)
# can't allocate 10, because there are only 9 left
assert not node_mgr.allocate(
{"node.nodearray": "htc"}, node_count=10, all_or_nothing=True)
result = node_mgr.allocate({"node.nodearray": "htc"},
node_count=10,
all_or_nothing=False)
assert result and len(result.nodes) == 9
assert result.nodes[0].nodearray == "htc"
def make_requests(node_mgr: NodeManager) -> None:
for n, mag in enumerate(magnitudes):
node_count = None if slots_or_nodes[n] else mag
slot_count = None if node_count else mag
node_mgr.allocate(
{
"ncpus": ncpus_per_job[n],
"exclusive": exclusivity[n]
},
node_count=node_count,
slot_count=slot_count,
)
def test_top_level_limits(node_mgr: NodeManager) -> None:
assert node_mgr.cluster_max_core_count == 10_000
assert node_mgr.cluster_consumed_core_count == 0
assert ["westus2"] == node_mgr.get_locations()
assert node_mgr.get_regional_consumed_core_count("westus2") == 0
assert node_mgr.get_regional_max_core_count("westus2") == 80
assert node_mgr.allocate({"node.vcpu_count": 4}, node_count=1)
assert node_mgr.cluster_consumed_core_count == 4
assert node_mgr.get_regional_consumed_core_count("westus2") == 4
def test_packing(node_mgr: NodeManager) -> None:
# htc node can fit 4 ncpus, so only allocate one node
result = node_mgr.allocate({
"node.nodearray": "htc",
"ncpus": 1
},
slot_count=2)
assert result, str(result)
assert len(result.nodes) == 1, result.nodes
assert result.nodes[0].name == "htc-1"
assert result.nodes[0].resources["ncpus"] == 4
assert result.nodes[0].available["ncpus"] == 2, result.nodes[0].available[
"ncpus"]
assert len(node_mgr.new_nodes) == 1, len(node_mgr.new_nodes)
# htc node can fit 4 ncpus, but 2 are used up on the first node, so allocate a second
result = node_mgr.allocate({
"node.nodearray": "htc",
"ncpus": 1
},
slot_count=4)
assert result
assert len(result.nodes) == 2, result.nodes
assert result.nodes[0].name == "htc-1"
assert result.nodes[1].name == "htc-2"
assert len(node_mgr.new_nodes) == 2, [n.name for n in node_mgr.new_nodes]
assert len(set([n.name for n in node_mgr.new_nodes])) == 2
result = node_mgr.allocate({
"node.nodearray": "htc",
"ncpus": 1
},
slot_count=2)
assert len(result.nodes) == 1
assert result.nodes[0].name == "htc-2"
assert len(node_mgr.new_nodes) == 2
def test_single_alloc(node_mgr: NodeManager) -> None:
result = node_mgr.allocate({"node.nodearray": "htc"}, node_count=1)
assert result and len(result.nodes) == 1
assert result.nodes[0].nodearray == "htc"