def test_one_bucket_no_match_htc(): bindings = MockClusterBinding() bindings.add_nodearray("twocpus", {"ncpus": 2}) bindings.add_bucket("twocpus", "Standard_F2", 100, 100) dc = _new_dc(bindings) result = dc._add_job(_htc_job(resources={"ncpus": 3})) _assert_insufficient_resource(result)
def test_choice_ordering() -> None: bindings = MockClusterBinding() bindings.add_nodearray("array-a", {"nodetype": "A"}) bindings.add_bucket("array-a", "Standard_F4", 10, 10) bindings.add_nodearray("array-b", {"nodetype": "B"}) bindings.add_bucket("array-b", "Standard_F4s", 10, 10) register_result_handler(DefaultContextHandler("[test_or_ordering]")) for ordering in [["A", "B"], ["B", "A"]]: node_mgr = _node_mgr(bindings) hi, lo = node_mgr.get_buckets() if hi.resources["nodetype"] != ordering[0]: hi, lo = lo, hi assert hi.available_count == 10 assert lo.available_count == 10 result = node_mgr.allocate( { "nodetype": ordering, "exclusive": True, }, node_count=15, # noqa: E231 ) assert hi.available_count == 0 assert lo.available_count == 5 assert result by_array = partition(result.nodes, lambda n: n.resources["nodetype"]) assert len(by_array[ordering[0]]) == 10 assert len(by_array[ordering[1]]) == 5
def nopgbindings(): bindings = MockClusterBinding() bindings.add_nodearray("twocpus", {"ncpus": 2}) bindings.add_nodearray("fourcpus", {"ncpus": 4}) bindings.add_bucket("twocpus", "Standard_F2", 100, 100) bindings.add_bucket("fourcpus", "Standard_F2", 100, 100) return bindings
def test_one_bucket_match_mpi(): bindings = MockClusterBinding() bindings.add_nodearray("twocpus", {"ncpus": 2}) bindings.add_bucket("twocpus", "Standard_F2", 100, 100, placement_groups=["pg0"]) dc = _new_dc(bindings) result = dc._add_job(_mpi_job(nodes=2, resources={"ncpus": 2})) _assert_success(result, ["twocpus"] * 2)
def test_one_bucket_no_match_mpi() -> None: bindings = MockClusterBinding() bindings.add_nodearray("twocpus", {"ncpus": 2}) bindings.add_bucket("twocpus", "Standard_F2", 100, 100) dc = _new_dc(bindings) result = dc._add_job(_mpi_job(resources={"ncpus": 3})) assert not result
def pgbindings(): bindings = MockClusterBinding() bindings.add_nodearray("hpc2", {"ncpus": 2}) bindings.add_nodearray("hpc4", {"ncpus": 4}) bindings.add_bucket("hpc2", "Standard_F2", 100, 100, placement_groups=["pg0"]) bindings.add_bucket("hpc4", "Standard_F4", 100, 100, placement_groups=["pg0"]) return bindings
def test_vm_family_limit(bindings: MockClusterBinding) -> None: bindings = MockClusterBinding("clusty") bindings.add_nodearray("htc", {"nodetype": "A"}) bindings.add_bucket( "htc", "Standard_F4", available_count=20, max_count=20, family_quota_count=30, family_quota_core_count=120, family_consumed_core_count=0, ) bindings.add_bucket( "htc", "Standard_F2", available_count=20, max_count=20, family_quota_count=30, family_quota_core_count=120, family_consumed_core_count=0, ) nm = _node_mgr(bindings) result = nm.allocate({}, node_count=100, all_or_nothing=False) assert len(result.nodes) == 40 pass
def test_custom_node_attrs_and_node_config() -> None: b = MockClusterBinding() b.add_nodearray("htc", {}, software_configuration={"myscheduler": { "A": 1 }}) b.add_bucket("htc", "Standard_F2", 10, 10) b.add_node("htc-1", "htc") node_mgr = new_node_manager({"_mock_bindings": b}) (existing_node, ) = node_mgr.get_nodes() try: existing_node.node_attribute_overrides["willfail"] = 123 assert False except TypeError: pass result = node_mgr.allocate({"exclusive": True}, node_count=2) assert result (node, ) = [n for n in result.nodes if not n.exists] assert node.software_configuration.get("test_thing") is None node.node_attribute_overrides["Configuration"] = {"test_thing": "is set"} assert node.software_configuration.get("test_thing") == "is set" try: node.software_configuration["willfail"] = 123 assert not node.software_configuration.get("willfail") except TypeError: pass # we won't handle dict merges here. assert node.software_configuration.get("myscheduler") == {"A": 1} node.node_attribute_overrides["Configuration"] = {"myscheduler": {"B": 2}} assert node.software_configuration.get("myscheduler") == {"B": 2} # if you want to add to the existing software_configuration, use # the node.software_configuration node.node_attribute_overrides["Configuration"][ "myscsheduler"] = node.software_configuration.get("myscheduler", {}) node.node_attribute_overrides["Configuration"]["myscheduler"]["B"] = 2 node.node_attribute_overrides["Configuration"] = { "myscheduler": { "A": 1, "B": 2 } } node.software_configuration["willsucceed"] = 123 node.exists = True try: node.software_configuration["willfail"] = 123 assert False except TypeError: pass
def test_one_bucket_match_htc_reuse(): bindings = MockClusterBinding() bindings.add_nodearray("twocpus", {"ncpus": 2}) bindings.add_bucket("twocpus", "Standard_F2", 100, 100) dc = _new_dc(bindings) result = dc._add_job(_htc_job(job_name="1", resources={"ncpus": 1})) _assert_success(result, ["twocpus"]) result = dc._add_job(_htc_job(job_name="2", resources={"ncpus": 1})) _assert_success(result, ["twocpus"])
def _bindings() -> MockClusterBinding: if previous_dcalc: return previous_dcalc.node_mgr.cluster_bindings mock_bindings = MockClusterBinding() mock_bindings.add_nodearray("hpc", {}, max_placement_group_size=5) mock_bindings.add_bucket("hpc", "Standard_F4", 100, 100) mock_bindings.add_nodearray("htc", {}, max_count=10) mock_bindings.add_bucket("htc", "Standard_F4", 10, 10) return mock_bindings
def next_dc(existing_nodes: List[Node]) -> DemandCalculator: bindings = MockClusterBinding() for_region = list(vm_sizes.VM_SIZES["southcentralus"].keys()) for n in range(num_arrays): nodearray = "nodearray{}".format(n) bindings.add_nodearray(nodearray, {}, location="southcentralus") for b in range(num_buckets): vm_size = for_region[vm_indices[n * num_buckets + b]] bindings.add_bucket( nodearray, vm_size, max_count=10, available_count=10, ) return _new_dc(bindings)
def next_node_mgr(existing_nodes: List[Node]) -> NodeManager: bindings = MockClusterBinding() for n in range(num_arrays): nodearray = "nodearray{}".format(n) bindings.add_nodearray(nodearray, {}, location="southcentralus") for b in range(num_buckets): vm_size = vm_size_choices[n * num_buckets + b] bindings.add_bucket( nodearray, vm_size, max_count=10, available_count=10, ) return _node_mgr(bindings)
def test_mock_bindings2() -> None: bindings = MockClusterBinding() bindings.add_nodearray("w", {}, location="westus2", max_count=8) bindings.add_bucket( "w", "Standard_E2_v3", max_count=80, available_count=8, family_consumed_core_count=72 * 2, family_quota_core_count=160, family_quota_count=80, ) bindings.add_bucket( "w", "Standard_E4_v3", max_count=40, available_count=4, family_consumed_core_count=72 * 2, family_quota_core_count=160, family_quota_count=80, ) bindings.add_bucket("w", "Standard_D8s_v3", max_count=80, available_count=8) bindings.add_nodearray("e", {}, location="eastus") bindings.add_bucket("e", "Standard_E2_v3", max_count=20, available_count=4) node_mgr = _node_mgr(bindings) by_size = partition_single(node_mgr.get_buckets(), lambda b: (b.location, b.vm_size)) assert by_size[("westus2", "Standard_E2_v3")].available_count == 8 assert by_size[("westus2", "Standard_E2_v3")].limits.nodearray_available_count == 8 assert by_size[("westus2", "Standard_E2_v3")].limits.family_max_count == 80 assert by_size[("westus2", "Standard_E4_v3")].available_count == 4 assert by_size[("westus2", "Standard_D8s_v3")].available_count == 8 assert by_size[("eastus", "Standard_E2_v3")].available_count == 4 result = node_mgr.allocate( { "node.vm_size": "Standard_E4_v3", "exclusive": True, "node.location": "westus2", }, node_count=1, ) assert result, "\n".join(result.reasons) assert by_size[("westus2", "Standard_E2_v3")].limits.nodearray_available_count == 7 assert by_size[("westus2", "Standard_E2_v3")].available_count == 6 assert by_size[("westus2", "Standard_E4_v3")].available_count == 3 assert by_size[("westus2", "Standard_D8s_v3")].available_count == 7 assert by_size[("eastus", "Standard_E2_v3")].available_count == 4
def _bindings(): mock_bindings = MockClusterBinding() mock_bindings.add_nodearray( "htc2", { "slots": 2, "customer_htc_flag": True }, ) mock_bindings.add_nodearray( "htc4", { "slots": 4, "customer_htc_flag": True }, ) mock_bindings.add_bucket("htc2", "Standard_F2", 10, 8) mock_bindings.add_bucket("htc4", "Standard_F4", 5, 4) return mock_bindings
def test_basic() -> None: binding = MockClusterBinding() binding.add_nodearray("hpc", {"ncpus": "node.vcpu_count"}) binding.add_bucket("hpc", "Standard_F4", max_count=100, available_count=100) node_mgr = new_node_manager({"_mock_bindings": binding}) bucket = node_mgr.get_buckets()[0] assert 100 == bucket.available_count bucket.decrement(5) assert 95 == bucket.available_count bucket.rollback() assert 100 == bucket.available_count bucket.decrement(5) assert 95 == bucket.available_count bucket.commit() assert 95 == bucket.available_count bucket.decrement(5) assert 90 == bucket.available_count bucket.rollback() assert 95 == bucket.available_count
def _bindings() -> MockClusterBinding: bindings = MockClusterBinding("clusty") bindings.add_nodearray( "htc", { "nodetype": "A", "pcpus": 2 }, software_configuration={ "custom_int": 100, "custom_str": "abc" }, ) bindings.add_bucket( "htc", "Standard_F4s", max_count=20, available_count=10, family_consumed_core_count=0, family_quota_core_count=80, family_quota_count=80 // 4, regional_consumed_core_count=0, regional_quota_core_count=80, regional_quota_count=80 // 4, ) bindings.add_nodearray("hpc", {"nodetype": "B", "pcpus": 4}) bindings.add_bucket( "hpc", "Standard_F8s", max_count=20, available_count=10, family_consumed_core_count=0, family_quota_core_count=80, family_quota_count=80 // 8, regional_consumed_core_count=0, regional_quota_core_count=80, regional_quota_count=80 // 8, ) return bindings
def test_mock_bindings3() -> None: bindings = MockClusterBinding() bindings.add_nodearray("w", {}, location="westus2", max_count=8) bindings.add_bucket("w", "Standard_E2_v3", max_count=80, available_count=8) _node_mgr(bindings)
def test_overscaling_error() -> None: bindings = MockClusterBinding() bindings.add_nodearray("htc", {}) bindings.add_bucket("htc", "Standard_D16s_v3", 10, 10) bindings.add_bucket("htc", "Standard_E16_v3", 10, 1) bindings.add_bucket("htc", "Standard_E2s_v3", max_count=80, available_count=4) node_mgr = _node_mgr(bindings) node_mgr.set_system_default_resources() result = node_mgr.allocate( { "ncpus": 1, "node.vm_size": ["Standard_E16_v3", "Standard_E2s_v3"] }, slot_count=20, assignment_id="slots", ) assert result assert len(result.nodes) == 3 result = node_mgr.allocate( { "exclusive": True, "node.vm_size": "Standard_D16s_v3" }, node_count=10, assignment_id="nodes", ) assert result by_size = partition(node_mgr.new_nodes, lambda b: b.vm_size) assert "Standard_E2s_v3" in by_size assert len(by_size["Standard_E2s_v3"]) == 2 assert len(by_size["Standard_E16_v3"]) == 1 assert len(by_size["Standard_D16s_v3"]) == 10 for node in node_mgr.get_nodes(): print(node.name, node.vm_size, node.assignments) node_mgr.bootup() # recreate it - the bindings 'remembers' that we already created nodes node_mgr = _node_mgr(bindings) assert len(node_mgr.get_nodes()) == 13 result = node_mgr.allocate({"ncpus": 1}, slot_count=100, assignment_id="slots") assert result result = node_mgr.allocate( { "exclusive": True, "node.vm_size": "Standard_D16s_v3" }, node_count=10, assignment_id="nodes", ) assert result print() print() for node in node_mgr.get_nodes(): print(node.name, node.vm_size, node.assignments) assert len(node_mgr.new_nodes) == 0 assert len(node_mgr.get_nodes()) == 13 by_size = partition(node_mgr.get_nodes(), lambda b: b.vm_size) assert "Standard_E2s_v3" in by_size assert len(by_size["Standard_E2s_v3"]) == 2 assert len(by_size["Standard_E16_v3"]) == 1 assert len(by_size["Standard_D16s_v3"]) == 10
def test_family_and_spots() -> None: bindings = MockClusterBinding("clusty") bindings.add_nodearray("htc", {}, spot=False, max_count=10, max_core_count=400) bindings.add_nodearray("hpc", {}, spot=False, max_placement_group_size=7) bindings.add_bucket( "htc", "Standard_F4s", max_count=20, available_count=10, family_consumed_core_count=40, family_quota_core_count=80, family_quota_count=20, regional_consumed_core_count=45, regional_quota_core_count=100, regional_quota_count=25, ) bindings.add_bucket( "htc", "Standard_D4s_v3", max_count=20, available_count=10, family_consumed_core_count=40, family_quota_core_count=80, family_quota_count=20, regional_consumed_core_count=45, regional_quota_core_count=100, regional_quota_count=25, ) bindings.add_bucket( "hpc", "Standard_D4s_v3", max_count=20, available_count=10, family_consumed_core_count=40, family_quota_core_count=80, family_quota_count=20, regional_consumed_core_count=45, regional_quota_core_count=100, regional_quota_count=25, ) bindings.add_bucket( "hpc", "Standard_D4s_v3", max_count=20, available_count=10, family_consumed_core_count=40, family_quota_core_count=80, family_quota_count=20, regional_consumed_core_count=45, regional_quota_core_count=100, regional_quota_count=25, placement_groups=["123"], ) bindings.add_nodearray("htcspot", {}, spot=True) bindings.add_bucket( "htcspot", "Standard_F4s", max_count=20, available_count=10, family_consumed_core_count=0, family_quota_core_count=0, family_quota_count=0, regional_consumed_core_count=45, regional_quota_core_count=100, regional_quota_count=25, ) node_mgr = new_node_manager({"_mock_bindings": bindings}) by_key: Dict[str, NodeBucket] = partition(node_mgr.get_buckets(), lambda b: (b.nodearray, b.vm_size)) htc = by_key[("htc", "Standard_F4s")][0] htc2 = by_key[("htc", "Standard_D4s_v3")][0] htcspot = by_key[("htcspot", "Standard_F4s")][0] hpcs = by_key[("hpc", "Standard_D4s_v3")] hpc_pg = [x for x in hpcs if x.placement_group][0] # ondemand instances use actual family quota assert htc.limits.family_max_count == 20 assert htc2.limits.family_max_count == 20 assert htc.limits.family_available_count == 10 assert htc2.limits.family_available_count == 10 # spot instances replace family with regional assert htcspot.limits.family_max_count == 25 assert htcspot.limits.family_available_count == 13 assert node_mgr.allocate( { "node.nodearray": "htc", "node.vm_size": "Standard_F4s" }, node_count=1) # ondemand instances use actual family quota assert htc.limits.family_max_count == 20 assert htc2.limits.family_max_count == 20 assert htc.limits.family_available_count == 9 assert htc2.limits.family_available_count == 10 assert htc.limits.nodearray_available_count == 9 assert htc2.limits.nodearray_available_count == 9 assert htc.available_count == 9 # nodearray limit affects htc2 since max_count=10 assert htc2.available_count == 9 # now the regional is affected by our allocation assert htcspot.limits.family_max_count == 25 assert htcspot.limits.family_available_count == 13 - 1 assert hpc_pg.available_count == 7
def test_mock_bindings(bindings: MockClusterBinding) -> None: ctx = register_result_handler(DefaultContextHandler("[test]")) hpc, htc = _node_mgr(bindings).get_buckets() if hpc.nodearray != "hpc": hpc, htc = htc, hpc assert hpc.nodearray == "hpc" assert htc.nodearray == "htc" assert hpc.family_available_count == 10 assert hpc.available_count == 10 assert hpc.family_available_count == 10 assert htc.family_available_count == 20 hpc.decrement(1) assert hpc.family_available_count == 9 assert htc.family_available_count == 20 hpc.commit() assert hpc.family_available_count == 9 assert htc.family_available_count == 18 hpc.increment(1) hpc.commit() assert hpc.family_available_count == 10 assert htc.family_available_count == 20 ctx.set_context("[failure]") nm = _node_mgr(bindings) b = MockClusterBinding() b.add_nodearray("haspgs", {}, max_placement_group_size=20) b.add_bucket( "haspgs", "Standard_F4", 100, 100, placement_groups=["pg0", "pg1"], ) # make sure we take the max_placement_group_size (20) into account # and that we have the non-pg and 2 pg buckets. nm = _node_mgr(b) no_pg, pg0, pg1 = sorted(nm.get_buckets(), key=lambda b: b.placement_group or "") assert no_pg.available_count == 100 assert pg0.available_count == 20 assert pg1.available_count == 20 # let's add a node to pg0 (100 - 1, 20 - 1, 20) b.add_node("haspgs-pg0-1", "haspgs", "Standard_F4", placement_group="pg0") nm = _node_mgr(b) no_pg, pg0, pg1 = sorted(nm.get_buckets(), key=lambda b: b.placement_group or "") assert no_pg.available_count == 99 assert pg0.available_count == 19 assert pg1.available_count == 20 # let's add a node to pg1 (100 - 2, 20 - 1, 20 - 1) b.add_node("haspgs-pg1-1", "haspgs", "Standard_F4", placement_group="pg1") nm = _node_mgr(b) no_pg, pg0, pg1 = sorted(nm.get_buckets(), key=lambda b: b.placement_group or "") assert no_pg.available_count == 98 assert pg0.available_count == 19 assert pg1.available_count == 19 # let's add 90 htc nodes so that our pg available counts are floored # by the overall available_count for i in range(90): b.add_node("haspgs-{}".format(i + 1), "haspgs", "Standard_F4") nm = _node_mgr(b) no_pg, pg0, pg1 = sorted(nm.get_buckets(), key=lambda b: b.placement_group or "") assert no_pg.available_count == 8 assert pg0.available_count == 8 assert pg1.available_count == 8 # lastly, add a nother node to a pg and see that all of avail go down b.add_node("haspgs-pg1-2", "haspgs", "Standard_F4", placement_group="pg1") nm = _node_mgr(b) no_pg, pg0, pg1 = sorted(nm.get_buckets(), key=lambda b: b.placement_group or "") assert no_pg.available_count == 7 assert pg0.available_count == 7 assert pg1.available_count == 7