diff --git a/src/taskgraph/generator.py b/src/taskgraph/generator.py index 1368ef01c..45cdfc32b 100644 --- a/src/taskgraph/generator.py +++ b/src/taskgraph/generator.py @@ -8,6 +8,7 @@ import multiprocessing import os import platform +import time from concurrent.futures import ( FIRST_COMPLETED, ProcessPoolExecutor, @@ -34,6 +35,19 @@ logger = logging.getLogger(__name__) +def phase_finished(name, start, end): + """Called after each generation phase, with `time.monotonic()` values + taken around it. `name` is None if the phase raised an exception. + + Projects can replace this from their `register` function; a replacement + made then also applies to the phase that is running. + """ + if name is None: + logger.debug(f"Generation phase failed after {end - start:.2f}s") + else: + logger.debug(f"Generated {name} in {end - start:.2f}s") + + class KindNotFound(Exception): """ Raised when trying to load kind from a directory without a kind.yml. @@ -603,10 +617,15 @@ def load_tasks(): def _run_until(self, name): while name not in self._run_results: + start = time.monotonic() try: k, v = next(self._run) # type: ignore except StopIteration: raise AttributeError(f"No such run result {name}") + except BaseException: + phase_finished(None, start, time.monotonic()) + raise + phase_finished(k, start, time.monotonic()) self._run_results[k] = v return self._run_results[name] diff --git a/test/test_generator.py b/test/test_generator.py index 8a31c186a..fa65f2302 100644 --- a/test/test_generator.py +++ b/test/test_generator.py @@ -457,3 +457,103 @@ def test_kind_graph_with_target_kinds(maketgg): # _fake3 and _other should not be included assert "_fake3" not in kind_graph.nodes assert "_other" not in kind_graph.nodes + + +PHASES = [ + "graph_config", + "parameters", + "kind_graph", + "full_task_set", + "full_task_graph", + "target_task_set", + "target_task_graph", + "optimized_task_graph", + "label_to_taskid", + "morphed_task_graph", +] + + +def test_phase_finished(monkeypatch, maketgg): + calls = [] + monkeypatch.setattr(generator, "phase_finished", lambda *args: calls.append(args)) + tgg = maketgg(["_fake-t-2"]) + tgg.target_task_set + assert [name for name, _, _ in calls] == PHASES[ + : PHASES.index("target_task_set") + 1 + ] + tgg.target_task_set + tgg.morphed_task_graph + assert [name for name, _, _ in calls] == PHASES + + prev_end = float("-inf") + for _, start, end in calls: + assert prev_end <= start <= end + prev_end = end + + +def test_phase_finished_replaced_during_generation(monkeypatch, maketgg): + first_phase = [] + register_phase = [] + tgg = maketgg(["_fake-t-2"]) + + def load_graph_config(root_dir): + graph_config = fake_load_graph_config(root_dir) + monkeypatch.setattr( + generator, "phase_finished", lambda *args: first_phase.append(args[0]) + ) + monkeypatch.setattr( + type(graph_config), + "register", + lambda self: monkeypatch.setattr( + generator, + "phase_finished", + lambda *args: register_phase.append(args[0]), + ), + ) + return graph_config + + monkeypatch.setattr(generator, "load_graph_config", load_graph_config) + tgg.morphed_task_graph + assert first_phase == ["graph_config"] + assert register_phase == PHASES[1:] + + +def test_phase_finished_on_error(monkeypatch, mocker, maketgg): + calls = [] + monkeypatch.setattr(generator, "phase_finished", lambda *args: calls.append(args)) + mocker.patch.object(Kind, "load_tasks", side_effect=RuntimeError("bug")) + tgg = maketgg() + with pytest.raises(RuntimeError): + tgg.full_task_set + assert [name for name, _, _ in calls] == [ + "graph_config", + "parameters", + "kind_graph", + None, + ] + assert calls[-2][2] <= calls[-1][1] <= calls[-1][2] + + with pytest.raises(AttributeError): + tgg.full_task_set + assert len(calls) == 4 + + +def test_phase_finished_default(caplog, maketgg): + tgg = maketgg(["_fake-t-2"]) + with caplog.at_level("DEBUG", logger="taskgraph.generator"): + tgg.morphed_task_graph + messages = [r.getMessage() for r in caplog.records] + for phase in PHASES: + assert any(m.startswith(f"Generated {phase} in ") for m in messages) + + +def test_phase_finished_default_on_error(caplog, mocker, maketgg): + mocker.patch.object(Kind, "load_tasks", side_effect=RuntimeError("bug")) + tgg = maketgg() + with caplog.at_level("DEBUG", logger="taskgraph.generator"): + with pytest.raises(RuntimeError): + tgg.full_task_set + assert any( + r.getMessage().startswith("Generation phase failed after ") + for r in caplog.records + )