File size: 2,491 Bytes
079c32c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
import pytest
import os
from ding.entry.cli_parsers.k8s_parser import k8s_parser
@pytest.fixture
def set_k8s_env():
os.environ["DI_NODES"] = 'SH-0,SH-1,SH-2,SH-3,SH-4,SH-5' # All the nodes
os.environ["DI_RANK"] = '3' # Proc order, start from 0, can not be modified by config
yield
del os.environ["DI_NODES"]
del os.environ["DI_RANK"]
@pytest.mark.unittest
@pytest.mark.usefixtures('set_k8s_env')
def test_k8s_parser():
# With platform_spec
platform_spec = {
"tasks": [
{
"labels": "league,collect",
"node_ids": 10
}, {
"labels": "league,collect",
"node_ids": 11
}, {
"labels": "evaluate",
"node_ids": 20,
"attach_to": "$node.10,$node.11"
}, {
"labels": "learn",
"node_ids": 31,
"ports": 50000,
"attach_to": "$node.10,$node.11,$node.20"
}, {
"labels": "learn",
"node_ids": 32,
"attach_to": "$node.10,$node.11,$node.20"
}, {
"labels": "learn",
"node_ids": 33,
"attach_to": "$node.10,$node.11,$node.20"
}
]
}
all_args = k8s_parser(platform_spec, mq_type="nng")
assert all_args["labels"] == "learn"
assert all_args["address"] == "SH-3"
assert all_args["ports"] == 50000
assert all_args["node_ids"] == 31
assert all_args["parallel_workers"] == 1
assert all_args[
"attach_to"
] == "tcp://SH-0:50515," +\
"tcp://SH-1:50515," +\
"tcp://SH-2:50515"
# Without platform_spec, parse by global config
all_args = k8s_parser(None, topology="mesh", mq_type="nng")
assert all_args["address"] == "SH-3"
assert all_args["node_ids"] == 3
assert all_args["parallel_workers"] == 1
assert all_args[
"attach_to"
] == "tcp://SH-0:50515," +\
"tcp://SH-1:50515," +\
"tcp://SH-2:50515"
# With multiple parallel workers
all_args = k8s_parser(None, topology="mesh", parallel_workers=2)
assert all_args["address"] == "SH-3"
assert all_args["node_ids"] == 6
assert all_args["parallel_workers"] == 2
assert all_args[
"attach_to"
] == "tcp://SH-0:50515,tcp://SH-0:50516," +\
"tcp://SH-1:50515,tcp://SH-1:50516," +\
"tcp://SH-2:50515,tcp://SH-2:50516"
|