def test_get_nodes_respects_max_limit():
types = {
"m4.large": {
"resources": {
"CPU": 2
},
"max_workers": 10,
},
"gpu": {
"resources": {
"GPU": 1
},
"max_workers": 99999,
},
}
assert get_nodes_for(types, {}, 2, [{"CPU": 1}] * 10) == \
[("m4.large", 2)]
assert get_nodes_for(types, {"m4.large": 9999}, 9999, [{
"CPU": 1
}] * 10) == []
assert get_nodes_for(types, {"m4.large": 0}, 9999, [{
"CPU": 1
}] * 10) == [("m4.large", 5)]
assert get_nodes_for(types, {"m4.large": 7}, 4, [{
"CPU": 1
}] * 10) == [("m4.large", 3)]
assert get_nodes_for(types, {"m4.large": 7}, 2, [{
"CPU": 1
}] * 10) == [("m4.large", 2)]
def test_get_nodes_packing_heuristic():
assert get_nodes_for(TYPES_A, {}, 9999, [{"GPU": 8}]) == \
[("p2.8xlarge", 1)]
assert get_nodes_for(TYPES_A, {}, 9999, [{"GPU": 1}] * 6) == \
[("p2.8xlarge", 1)]
assert get_nodes_for(TYPES_A, {}, 9999, [{"GPU": 1}] * 4) == \
[("p2.xlarge", 4)]
assert get_nodes_for(TYPES_A, {}, 9999, [{"CPU": 32, "GPU": 1}] * 3) \
== [("p2.8xlarge", 3)]
assert get_nodes_for(TYPES_A, {}, 9999, [{"CPU": 64, "GPU": 1}] * 3) \
== []
assert get_nodes_for(TYPES_A, {}, 9999, [{"CPU": 64}] * 3) == \
[("m4.16xlarge", 3)]
assert get_nodes_for(TYPES_A, {}, 9999, [{"CPU": 64}, {"CPU": 1}]) \
== [("m4.16xlarge", 1), ("m4.large", 1)]
assert get_nodes_for(
TYPES_A, {}, 9999, [{"CPU": 64}, {"CPU": 9}, {"CPU": 9}]) == \
[("m4.16xlarge", 1), ("m4.4xlarge", 2)]
assert get_nodes_for(TYPES_A, {}, 9999, [{"CPU": 16}] * 5) == \
[("m4.16xlarge", 1), ("m4.4xlarge", 1)]
assert get_nodes_for(TYPES_A, {}, 9999, [{"CPU": 8}] * 10) == \
[("m4.16xlarge", 1), ("m4.4xlarge", 1)]
assert get_nodes_for(TYPES_A, {}, 9999, [{"CPU": 1}] * 100) == \
[("m4.16xlarge", 1), ("m4.4xlarge", 2), ("m4.large", 2)]
assert get_nodes_for(
TYPES_A, {}, 9999, [{"GPU": 1}] + ([{"CPU": 1}] * 64)) == \
[("m4.16xlarge", 1), ("p2.xlarge", 1)]
assert get_nodes_for(
TYPES_A, {}, 9999, ([{"GPU": 1}] * 8) + ([{"CPU": 1}] * 64)) == \
[("m4.16xlarge", 1), ("p2.8xlarge", 1)]
