Extend bt fuzz / update signal handling (#455)

* new signal handling

* add variation of old handler

* less restrictice

* move asserts outside of actual preexec function

* move asserts outside of actual preexec function

* move asserts outside of actual preexec function

* directly reference the python executeable

* clean up ans generation/validation

* use proper max_len

* lock output

Author: mzuenni

bin/fuzz.py

@@ -1,16 +1,19 @@
 import config
 import problem
-import run
 import random
 import generate
+import shutil
 import signal
+import sys
 import time
 import threading
+from colorama import Style
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, Optional, TextIO
 
 import parallel
 from util import *
+from run import Run, Submission
 from testcase import Testcase
 from validate import OutputValidator, Mode
 from verdicts import Verdict
@@ -27,6 +30,7 @@
 class GeneratorTask:
     def __init__(self, fuzz: "Fuzz", t: generate.TestcaseRule, i: int, tmp_id: int):
         self.fuzz = fuzz
+        self.rule = t
         generator = t.generator
         assert generator is not None
         self.generator = generator
@@ -52,6 +56,7 @@ def _run(self, bar: ProgressBar) -> bool:
         # GENERATE THE TEST DATA
         dir = Path("fuzz") / f"tmp_id_{str(self.tmp_id)}"
         cwd = self.fuzz.problem.tmpdir / "tool_runs" / dir
+        shutil.rmtree(cwd, ignore_errors=True)
         cwd.mkdir(parents=True, exist_ok=True)
         name = "testcase"
         infile = cwd / (name + ".in")
@@ -63,51 +68,64 @@ def _run(self, bar: ProgressBar) -> bool:
 
         localbar = bar.start(f"{self.i}: generate")
         result = self.generator.run(localbar, cwd, name, self.seed)
+        self.fuzz.queue.ensure_alive()
         if not result.status:
             return False  # No need to call bar.done() in this case, because the Generator calls bar.error()
+        if ".ans" in self.rule.hardcoded:
+            ansfile.write_text(self.rule.hardcoded[".ans"])
         localbar.done()
 
         testcase = Testcase(self.fuzz.problem, infile, short_path=dir / (name + ".in"))
 
         # Validate the generated .in.
         localbar = bar.start(f"{self.i}: validate input")
         if not testcase.validate_format(Mode.INPUT, bar=localbar, constraints=None):
+            self.fuzz.queue.ensure_alive()
             localbar.done(False)
             return False
+        self.fuzz.queue.ensure_alive()
         localbar.done()
 
         # Generate .ans.
-        if not self.fuzz.problem.interactive and not self.fuzz.problem.multi_pass:
-            if self.solution and not testcase.ans_path.is_file():
-                if testcase.ans_path.is_file():
-                    testcase.ans_path.unlink()
-                # Run the solution and validate the generated .ans.
-                localbar = bar.start(f"{self.i}: generate ans")
-                if not self.solution.run(bar, cwd).status:
+        if not ansfile.is_file():
+            if self.fuzz.problem.settings.ans_is_output:
+                if self.solution:
+                    # Run the solution and validate the generated .ans.
+                    localbar = bar.start(f"{self.i}: generate ans")
+                    if not self.solution.run(bar, cwd).status:
+                        self.fuzz.queue.ensure_alive()
+                        localbar.done()
+                        return False
+                    self.fuzz.queue.ensure_alive()
                     localbar.done()
-                    return False
-                localbar.done()
-
-            if ansfile.is_file():
-                localbar = bar.start(f"{self.i}: validate output")
-                if not testcase.validate_format(Mode.ANSWER, bar=localbar):
-                    localbar.done(False)
-                    return False
-                localbar.done()
-            else:
-                bar.error(f"{self.i}: {ansfile.name} was not generated.")
+            elif self.fuzz.problem.interactive or self.fuzz.problem.multi_pass:
+                ansfile.write_text("")
+
+        if ansfile.is_file():
+            localbar = bar.start(f"{self.i}: validate output")
+            if not testcase.validate_format(Mode.ANSWER, bar=localbar):
+                self.fuzz.queue.ensure_alive()
+                localbar.done(False)
                 return False
+            self.fuzz.queue.ensure_alive()
+            localbar.done()
         else:
-            if not testcase.ans_path.is_file():
-                testcase.ans_path.write_text("")
+            bar.error(f"{self.i}: {ansfile.name} was not generated.")
+            return False
 
         # Run all submissions against the testcase.
         with self.fuzz.queue:
             for submission in self.fuzz.submissions:
                 self.fuzz.queue.put(SubmissionTask(self, submission, testcase, self.tmp_id))
         return True
 
-    def save_test(self, bar: ProgressBar) -> None:
+    def get_command(self) -> dict[str, str] | str:
+        if not self.fuzz.problem.settings.ans_is_output and ".ans" in self.rule.hardcoded:
+            return {"generate": self.command, "ans": self.rule.hardcoded[".ans"]}
+        else:
+            return self.command
+
+    def save_test(self, bar: ProgressBar, submission: Submission, verdict: Verdict) -> None:
         if self.saved:
             return
         save = False
@@ -116,19 +134,21 @@ def save_test(self, bar: ProgressBar) -> None:
             if not self.saved:
                 self.saved = True
                 save = True
+        self.fuzz.queue.ensure_alive()
         # only save rule if we set self.saved to True
-        if save and not self.fuzz.queue.aborted:
+        if save:
             localbar = bar.start(f"{self.i}: {self.command}")
             localbar.log("Saving testcase in generators.yaml.")
+            self.fuzz.save_test(self.get_command(), submission, verdict)
+            self.fuzz.queue.ensure_alive()
             localbar.done()
-            self.fuzz.save_test(self.command)
 
 
 class SubmissionTask:
     def __init__(
         self,
         generator_task: GeneratorTask,
-        submission: run.Submission,
+        submission: Submission,
         testcase: Testcase,
         tmp_id: int,
     ):
@@ -142,38 +162,67 @@ def run(self, bar: ProgressBar) -> None:
         self.generator_task.fuzz.finish_task(self.tmp_id)
 
     def _run(self, bar: ProgressBar) -> None:
-        r = run.Run(self.generator_task.fuzz.problem, self.submission, self.testcase)
+        r = Run(self.generator_task.fuzz.problem, self.submission, self.testcase)
         localbar = bar.start(f"{self.generator_task.i}: {self.submission.name}")
         result = r.run(localbar)
+        self.generator_task.fuzz.queue.ensure_alive()
         if result.verdict != Verdict.ACCEPTED:
-            self.generator_task.save_test(bar)
+            self.generator_task.save_test(bar, self.submission, result.verdict)
             localbar.done(False, f"{result.verdict}!")
         else:
             localbar.done()
 
 
+class FuzzProgressBar(ProgressBar):
+    def __init__(self, queue: parallel.AbstractQueue, prefix: str, max_len: int):
+        super().__init__(prefix, max_len)
+        self.queue = queue
+
+    def _print(
+        self,
+        *objects,
+        sep: str = "",
+        end: str = "\n",
+        file: TextIO = sys.stderr,
+        flush: bool = True,
+    ):
+        self.queue.ensure_alive()
+        super()._print(*objects, sep=sep, end=end, file=file, flush=flush)
+
+
 class Fuzz:
     def __init__(self, problem: problem.Problem):
         self.generators_yaml_mutex = threading.Lock()
         self.problem = problem
+        self.summary: dict[Submission, set[Verdict]] = {}
+        self.added = 0
 
         # GENERATOR INVOCATIONS
         generator_config = generate.GeneratorConfig(self.problem, config.args.testcases)
         self.testcase_rules: list[generate.TestcaseRule] = []
 
         # Filter to only keep valid rules depending on seed without duplicates from count
-        added_testcase_rules = set()
+        added_testcase_rule_data = set()
 
         def add_testcase(t: generate.TestcaseRule) -> None:
             if (
-                t.in_is_generated
-                and t.parse_error is None
-                and t.generator is not None
-                and t.generator.uses_seed
-                and t.generator.command_string.strip() not in added_testcase_rules
+                not t.in_is_generated
+                or t.root in config.INVALID_CASE_DIRECTORIES
+                or t.parse_error is not None
+                or t.generator is None
+                or not t.generator.uses_seed
             ):
-                self.testcase_rules.append(t)
-                added_testcase_rules.add(t.generator.command_string.strip())
+                return
+
+            testcase_rule_data = t.generator.command_string.strip()
+            if not problem.settings.ans_is_output and ".ans" in t.hardcoded:
+                testcase_rule_data += t.hardcoded[".ans"]
+
+            if testcase_rule_data in added_testcase_rule_data:
+                return
+
+            self.testcase_rules.append(t)
+            added_testcase_rule_data.add(testcase_rule_data)
 
         generator_config.root_dir.walk(add_testcase, dir_f=None)
         if len(self.testcase_rules) == 0:
@@ -205,33 +254,42 @@ def run(self) -> bool:
         def runner(task: GeneratorTask | SubmissionTask) -> None:
             task.run(bar)
 
-        # config.args.no_bar = True
-        # max(len(s.name) for s in self.submissions)
-        bar = ProgressBar("Fuzz", max_len=60)
         self.start_time = time.monotonic()
         self.iteration = 0
         self.tasks = 0
         self.queue = parallel.new_queue(runner, pin=True)
 
+        # pool of ids used for generators
+        self.tmp_ids = 2 * max(1, self.queue.num_threads) + 1
+        self.free_tmp_id = {*range(self.tmp_ids)}
+        self.tmp_id_count = [0] * self.tmp_ids
+
+        max_len = max(
+            25,
+            *[len(s.name) for s in self.submissions],
+            *[
+                len(t.generator.cache_command(seed=2**32))
+                for t in self.testcase_rules
+                if t.generator is not None
+            ],
+        )
+        max_len += len(f"{self.tmp_ids}: ")
+        bar = FuzzProgressBar(self.queue, "Fuzz", max_len=max_len)
+
         def soft_exit(sig: Any, frame: Any) -> None:
             if self.queue.aborted:
                 fatal("Running interrupted", force=True)
             else:
                 self.queue.abort()
                 with bar:
-                    bar.clearline()
+                    print(bar.carriage_return, file=sys.stderr)
                     message(
                         "Running interrupted (waiting on remaining tasks)\n",
                         "\nFuzz",
                         color_type=MessageType.ERROR,
                     )
 
-        signal.signal(signal.SIGINT, soft_exit)
-
-        # pool of ids used for generators
-        self.tmp_ids = 2 * max(1, self.queue.num_threads) + 1
-        self.free_tmp_id = {*range(self.tmp_ids)}
-        self.tmp_id_count = [0] * self.tmp_ids
+        old_handler = signal.signal(signal.SIGINT, soft_exit)
 
         # add first generator task
         self.finish_task()
@@ -241,11 +299,19 @@ def soft_exit(sig: Any, frame: Any) -> None:
         # At this point, no new tasks may be started anymore.
         self.queue.done()
 
+        signal.signal(signal.SIGINT, old_handler)
+
+        for submission, verdicts in self.summary.items():
+            msg = ", ".join(f"{v.color()}{v.short()}{Style.RESET_ALL}" for v in sorted(verdicts))
+            message(msg, "Fuzz", submission.name)
+        message(f"Found {self.added} testcases in total.", "Fuzz")
+
         if self.queue.aborted:
-            fatal("Running interrupted", force=True)
+            fatal("Running interrupted")
 
         bar.done()
         bar.finalize()
+
         return True
 
     # finish task from generator with tmp_id
@@ -280,7 +346,9 @@ def finish_task(self, tmp_id: Optional[int] = None, count: int = 1) -> None:
 
     # Write new rule to yaml
     # lock between read and write to ensure that no rule gets lost
-    def save_test(self, command: str) -> None:
+    def save_test(
+        self, command: dict[str, str] | str, submission: Submission, verdict: Verdict
+    ) -> None:
         with self.generators_yaml_mutex:
             generators_yaml = self.problem.path / "generators/generators.yaml"
             data = None
@@ -298,3 +366,6 @@ def save_test(self, command: str) -> None:
 
             # Overwrite generators.yaml.
             write_yaml(data, generators_yaml)
+
+            self.summary.setdefault(submission, set()).add(verdict)
+            self.added += 1

bin/interactive.py

@@ -242,7 +242,7 @@ def get_validator_command():
 
         if interaction:
             team_tee = subprocess.Popen(
-                ["python3", "-c", TEE_CODE, ">"],
+                [sys.executable, "-c", TEE_CODE, ">"],
                 stdin=subprocess.PIPE,
                 stdout=validator.stdin,
                 stderr=interaction_file,
@@ -251,7 +251,7 @@ def get_validator_command():
             )
             team_tee_pid = team_tee.pid
             val_tee = subprocess.Popen(
-                ["python3", "-c", TEE_CODE, "<"],
+                [sys.executable, "-c", TEE_CODE, "<"],
                 stdin=validator.stdout,
                 stdout=subprocess.PIPE,
                 stderr=interaction_file,