def test_dedup_input(self):
os.environ['CGO'] = 'true'
lp, models = self._build_logical_with_mnist(3)
opt = DedupInputOptimizer()
opt.convert(lp)
with open('dedup_logical_graph.json', 'r') as fp:
correct_dump = fp.readlines()
lp_dump = lp.logical_graph._dump()
self.assertTrue(correct_dump[0] == json.dumps(lp_dump))
advisor = RetiariiAdvisor()
cgo = CGOExecutionEngine()
phy_models = cgo._assemble(lp)
self.assertTrue(len(phy_models) == 1)
# logging.info(phy_models[0][0]._dump())
# script=model_to_pytorch_script(phy_models[0][0], placement = phy_models[0][1])
# logging.info(script)
# with open('generated/debug_dedup_input.py', 'w') as fp:
# fp.write(script)
# sys.path.insert(0, 'generated')
# multi_model = import_('debug_dedup_input.logical_0')
# trainer = PyTorchMultiModelTrainer(
# multi_model(), phy_models[0][0].training_config.kwargs
# )
# trainer.fit()
advisor.stopping = True
advisor.default_worker.join()
advisor.assessor_worker.join()
def test_dedup_input_two_devices(self):
_reset()
lp, models = self._build_logical_with_mnist(3)
opt = DedupInputOptimizer()
opt.convert(lp)
advisor = RetiariiAdvisor()
available_devices = [GPUDevice("test", 0), GPUDevice("test", 1)]
cgo = CGOExecutionEngine(devices=available_devices,
batch_waiting_time=0)
phy_models = cgo._assemble(lp)
self.assertTrue(len(phy_models) == 2)
advisor.stopping = True
advisor.default_worker.join()
advisor.assessor_worker.join()
cgo.join()
def test_dedup_input_two_devices(self):
_reset()
lp, models = self._build_logical_with_mnist(3)
opt = DedupInputOptimizer()
opt.convert(lp)
advisor = RetiariiAdvisor('ws://_unittest_placeholder_')
advisor._channel = protocol.LegacyCommandChannel()
advisor.default_worker.start()
advisor.assessor_worker.start()
remote = RemoteConfig(machine_list=[])
remote.machine_list.append(
RemoteMachineConfig(host='test', gpu_indices=[0, 1]))
cgo = CGOExecutionEngine(training_service=remote, batch_waiting_time=0)
phy_models = cgo._assemble(lp)
self.assertTrue(len(phy_models) == 2)
advisor.stopping = True
advisor.default_worker.join()
advisor.assessor_worker.join()
cgo.join()