def patch_gputil_get_gpus(self): return [ Object({ "memoryFree": gsize, "load": 0 }) for gsize in self.virtual_gpus ]
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__ = {}
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})