def test_fork_non_existing_trial(self):
lineages = Lineages()
trial = TrialStub(id="stub")
new_trial = TrialStub(id="fork")
with pytest.raises(KeyError):
new_lineage = lineages.fork(trial, new_trial)
def test_fork_existing_trial(self, tmp_path):
lineages = Lineages()
trial = TrialStub(id="stub", working_dir=os.path.join(tmp_path, "stub"))
os.makedirs(trial.working_dir)
lineage = lineages.add(trial)
assert len(lineages) == 1
new_trial = TrialStub(id="fork", working_dir=os.path.join(tmp_path, "fork"))
new_lineage = lineages.fork(trial, new_trial)
assert len(lineages) == 1
assert lineages._lineage_roots[0].children[0] is new_lineage
assert lineages._trial_to_lineages["fork"] is new_lineage
def test_set_jump_existing_trial(self):
lineages = Lineages()
root_1 = TrialStub(id="root-1")
lineage_1 = lineages.add(root_1)
root_2 = TrialStub(id="root-2")
lineage_2 = lineages.add(root_2)
child_trial = TrialStub(id="child")
child_lineage = lineages.fork(root_1, child_trial)
lineages.set_jump(root_2, child_trial)
assert child_lineage.base is lineage_2
assert lineage_2.jumps == [child_lineage]
assert child_lineage.jumps == []
assert lineage_2.base is None
assert lineage_1.jumps == []
assert lineage_1.base is None
def test_get_lineage_existing_node_trial(self):
lineages = Lineages()
for root_index in range(2):
trial = TrialStub(id=f"lineage-{root_index}-0")
lineage = lineages.add(trial)
for depth in range(1, 10):
new_trial = TrialStub(id=f"lineage-{root_index}-{depth}")
lineage = lineages.fork(trial, new_trial)
trial = new_trial
lineage = lineages.get_lineage(TrialStub(id="lineage-0-2"))
assert lineage.root is lineages._lineage_roots[0]
assert lineage.node_depth == 2
lineage = lineages.get_lineage(TrialStub(id="lineage-1-5"))
assert lineage.root is lineages._lineage_roots[1]
assert lineage.node_depth == 5
