def test_set_attr(self):
# set data keys
obj = Object({"a": 1, "const": 0}, const_attrs={"const"})
obj.a = 1
assert obj.a == 1
obj.b = 1
assert obj.b == 1
# set const keys
with pytest.raises(RuntimeError, match="is const"):
obj.const = 1
# set .call attribute
obj.call = lambda _: "pong"
assert obj("ping") == "pong"
# set .data attribute
obj.data = {}
assert obj.a is None
with pytest.raises(ValueError, match="must be a dictionary"):
obj.data = None
# set other attributes
with pytest.raises(RuntimeError, match="should not set"):
obj.__dict__ = {}
def patch_gputil_get_gpus(self):
return [
Object({
"memoryFree": gsize,
"load": 0
}) for gsize in self.virtual_gpus
]
import os
import json
from unittest import mock
def get_config():
c = Config()
c.conf1 = 1
c.conf2 = 2
return c
@mock.patch(
"machin.utils.conf.argparse.ArgumentParser.parse_args",
return_value=Object(data={"conf": ["conf1=2", "conf3=3"]}),
)
def test_load_config_cmd(*_mock_classes):
conf = load_config_cmd()
assert conf["conf1"] == 2
assert conf["conf2"] is None
assert conf["conf3"] == 3
conf = load_config_cmd(get_config())
# configs from commandline precedes configs from the config file
assert conf["conf1"] == 2
assert conf["conf2"] == 2
assert conf["conf3"] == 3
def test_load_config_file(tmpdir):
def patch_model_size_estimator(mocked_model, multiplier):
return Object({"estimate_size": lambda: mocked_model.size * multiplier})
def patch_psutil_virtual_memory(self):
return Object({"available": self.virtual_cpu * 1024 ** 2})
class TestModelAssigner:
# unit of size is MB
virtual_gpus = []
virtual_cpu = 0
def patch_gputil_get_available(self, order):
return list(range(len(self.virtual_gpus)))
def patch_gputil_get_gpus(self):
return [Object({"memoryFree": gsize, "load": 0}) for gsize in self.virtual_gpus]
def patch_psutil_virtual_memory(self):
return Object({"available": self.virtual_cpu * 1024 ** 2})
@staticmethod
def patch_model_size_estimator(mocked_model, multiplier):
return Object({"estimate_size": lambda: mocked_model.size * multiplier})
@pytest.mark.parametrize(
"params,gpus,cpu,assignment,exception,match",
[
(
{
"models": [
Object({"size": 10, "to": lambda *_: None}),
Object({"size": 10, "to": lambda *_: None}),
],
"model_connection": {(0, 1): 1},
"devices": ["cuda:0", "cuda:1", "cpu"],
"model_size_multiplier": 1,
"max_mem_ratio": 0.7,
"cpu_weight": 0.1,
"connection_weight": 1,
"size_match_weight": 1e-2,
"complexity_match_weight": 10,
"entropy_weight": 1,
"iterations": 500,
"update_rate": 0.01,
"gpu_gpu_distance": 1,
"cpu_gpu_distance": 10,
"move_models": True,
},
[1000, 1000],
1000,
[["cuda:0", "cuda:1"], ["cuda:1", "cuda:0"]],
None,
None,
),
(
{
"models": [
Object({"size": 10, "to": lambda *_: None}),
Object({"size": 10, "to": lambda *_: None}),
],
"model_connection": {(0, 1): 1},
"devices": ["cuda:0", "cuda:1", "cpu"],
"model_size_multiplier": 1,
"max_mem_ratio": 0.7,
"cpu_weight": 0.1,
"connection_weight": 10,
"size_match_weight": 1e-2,
"complexity_match_weight": 1,
"entropy_weight": 1,
"iterations": 500,
"update_rate": 0.01,
"gpu_gpu_distance": 1,
"cpu_gpu_distance": 10,
"move_models": True,
},
[1000, 1000],
1000,
[["cuda:0", "cuda:0"], ["cuda:1", "cuda:1"]],
None,
None,
),
(
{
"models": [
Object({"size": 600, "to": lambda *_: None}),
Object({"size": 600, "to": lambda *_: None}),
],
"model_connection": {(0, 1): 1},
"devices": ["cuda:0", "cuda:1", "cpu"],
"model_size_multiplier": 1,
"max_mem_ratio": 0.7,
"cpu_weight": 0.1,
"connection_weight": 1,
"size_match_weight": 1e-2,
"complexity_match_weight": 10,
"entropy_weight": 1,
"iterations": 500,
"update_rate": 0.01,
"gpu_gpu_distance": 1,
"cpu_gpu_distance": 10,
"move_models": True,
},
[1000, 1000],
1000,
[["cuda:1", "cuda:0"], ["cuda:0", "cuda:1"]],
None,
None,
),
(
{
"models": [
Object({"size": 10, "to": lambda *_: None}),
Object({"size": 10, "to": lambda *_: None}),
],
"model_connection": {(0, 1): 1},
"devices": ["cuda:0", "cuda:1", "cpu"],
"model_size_multiplier": 1,
"max_mem_ratio": 0.7,
"cpu_weight": 0.1,
"connection_weight": 1,
"size_match_weight": 1e-2,
"complexity_match_weight": 10,
"entropy_weight": 1,
"iterations": 500,
"update_rate": 0.01,
"gpu_gpu_distance": 1,
"cpu_gpu_distance": 10,
"move_models": True,
},
[1000],
1000,
[["cuda:0", "cuda:0"]],
None,
None,
),
(
{
"models": [
Object({"size": 10, "to": lambda *_: None}),
Object({"size": 10, "to": lambda *_: None}),
],
"model_connection": {(0, 1): 1},
"devices": ["cuda:0", "cuda:1", "cpu"],
"model_size_multiplier": 1,
"max_mem_ratio": 0.7,
"cpu_weight": 0.1,
"connection_weight": 1,
"size_match_weight": 1e-2,
"complexity_match_weight": 10,
"entropy_weight": 1,
"iterations": 500,
"update_rate": 0.01,
"gpu_gpu_distance": 1,
"cpu_gpu_distance": 10,
"move_models": True,
},
[],
1000,
[["cpu", "cpu"]],
None,
None,
),
(
{
"models": [
Object({"size": 10, "to": lambda *_: None}),
Object({"size": 10, "to": lambda *_: None}),
],
"model_connection": {(0, 1): 1},
"devices": None,
"model_size_multiplier": 1,
"max_mem_ratio": 0.7,
"cpu_weight": 0.1,
"connection_weight": 1,
"size_match_weight": 1e-2,
"complexity_match_weight": 10,
"entropy_weight": 1,
"iterations": 500,
"update_rate": 0.01,
"gpu_gpu_distance": 1,
"cpu_gpu_distance": 10,
"move_models": True,
},
[],
1000,
[["cpu", "cpu"]],
None,
None,
),
(
{
"models": [
Object({"size": 100, "to": lambda *_: None}),
Object({"size": 100, "to": lambda *_: None}),
],
"model_connection": {(0, 1): 1},
"devices": ["cuda:0", "cuda:1", "cpu"],
"model_size_multiplier": 1,
"max_mem_ratio": 0.7,
"cpu_weight": 0.1,
"connection_weight": 1,
"size_match_weight": 1e-2,
"complexity_match_weight": 10,
"entropy_weight": 1,
"iterations": 500,
"update_rate": 0.01,
"gpu_gpu_distance": 1,
"cpu_gpu_distance": 10,
"move_models": True,
},
[10],
10,
None,
RuntimeError,
"Estimated model will use",
),
],
)
def test_assigner(self, params, gpus, cpu, assignment, exception, match):
t.manual_seed(0)
self.virtual_gpus = gpus
self.virtual_cpu = cpu
with mock.patch(
"machin.parallel.assigner.GPUtil.getAvailable",
self.patch_gputil_get_available,
) as _p1, mock.patch(
"machin.parallel.assigner.GPUtil.getGPUs", self.patch_gputil_get_gpus
) as _p2, mock.patch(
"machin.parallel.assigner.psutil.virtual_memory",
self.patch_psutil_virtual_memory,
) as _p3, mock.patch(
"machin.parallel.assigner.ModelSizeEstimator",
self.patch_model_size_estimator,
) as _p4:
if exception is not None:
with pytest.raises(exception, match=match):
ModelAssigner(**params)
else:
assigner = ModelAssigner(**params)
real_assignment = [str(dev) for dev in assigner.assignment]
assert real_assignment in assignment
def test_set_item(self):
obj = Object({"a": 1})
obj["a"] = 2
assert obj.a == 2
def test_get_item(self):
obj = Object({"a": 1})
assert obj["a"] == 1
def test_get_attr(self):
obj = Object({"a": 1})
with pytest.raises(AttributeError, match="Failed to find"):
_ = obj.__some_invalid_special_attr__
assert obj.a == 1
def test_call(self):
obj = Object()
obj("original_call")
obj.call = lambda _: "pong"
assert obj("ping") == "pong"
def test_init(self):
obj = Object()
assert obj.data == {}
obj = Object({"a": 1})
assert obj.data == {"a": 1}
def __init__(self):
super().__init__()
self.nn_model = NNModule()
self.frame = Object({"optimizers": None, "lr_schedulers": None})