limit stdout and stderr capture to 16KiB

scripts/ci_test_cuda.py

@@ -109,3 +109,33 @@ def test_cuda_correct():
     assert "warming up..." in run.stdout
     assert run.exit_code == ExitCode.SUCCESS
     assert run.result["check"] == "pass"
+
+
+def test_huge_output():
+    sub = """
+#include "reference.cuh"
+#include <iostream>
+output_t custom_kernel(input_t data)
+{
+    for(int i = 0; i < 10000; ++i) {
+        std::cout << "blah blah\\n";
+    }
+    return data;
+}
+    """
+
+    comp, run = run_cuda_script(
+        {"eval.cu": cu_eval}, {"reference.cuh": ref.read_text(), "submission.cuh": sub}, arch=None
+    )
+    assert run.success
+    assert len(run.stdout) < 16384
+    assert "[...]" in run.stdout
+
+    sub = sub.replace("std::cout", "std::cerr")
+
+    comp, run = run_cuda_script(
+        {"eval.cu": cu_eval}, {"reference.cuh": ref.read_text(), "submission.cuh": sub}, arch=None
+    )
+    assert run.success
+    assert len(run.stderr) < 16384
+    assert "[...]" in run.stderr

scripts/ci_test_python.py

@@ -61,3 +61,29 @@ def test_correct():
     assert "warming up..." in run.stdout
     assert run.exit_code == ExitCode.SUCCESS
     assert run.result["check"] == "pass"
+
+
+def test_huge_output():
+    sub = """
+import sys
+def custom_kernel(input):
+    print("blah blah\\n" * 10000, file=sys.stdout)
+    return input
+"""
+    run = run_pytorch_script(
+        {"eval.py": py_eval, "reference.py": ref.read_text(), "submission.py": sub},
+        "eval.py",
+    )
+    assert run.success
+    assert len(run.stdout) < 16384
+    assert "[...]" in run.stdout
+
+    sub = sub.replace("sys.stdout", "sys.stderr")
+
+    run = run_pytorch_script(
+        {"eval.py": py_eval, "reference.py": ref.read_text(), "submission.py": sub},
+        "eval.py",
+    )
+    assert run.success
+    assert len(run.stderr) < 16384
+    assert "[...]" in run.stderr

src/discord-cluster-manager/run_eval.py

@@ -49,6 +49,17 @@ def _make_cmd(args: list[str]):
     return " ".join(map(shlex.quote, args))
 
 
+def _limit_length(text: str, max_len: int = 16384):
+    lines = text.split("\n")
+    size = 0
+    for i, line in enumerate(lines):
+        size += len(line) + 1
+        if size + 100 > max_len:
+            lines = lines[:i] + [f"[...] {len(lines)-i} lines omitted"]
+            return "\n".join(lines)
+    return text
+
+
 def compile_cuda_script(  # # noqa: C901
     files: list[str],
     arch: int = None,
@@ -88,8 +99,8 @@ def compile_cuda_script(  # # noqa: C901
             success=False,
             nvcc_version="",
             command=_make_cmd(e.cmd),
-            stdout=e.stdout,
-            stderr=e.stderr,
+            stdout=_limit_length(e.stdout),
+            stderr=_limit_length(e.stderr),
             exit_code=e.returncode,
         )
 
@@ -109,8 +120,8 @@ def compile_cuda_script(  # # noqa: C901
             success=False,
             nvcc_version=nvcc_version,
             command=_make_cmd(e.cmd),
-            stdout=e.stdout,
-            stderr=e.stderr,
+            stdout=_limit_length(e.stdout),
+            stderr=_limit_length(e.stderr),
             exit_code=e.returncode,
         )
 
@@ -119,8 +130,8 @@ def compile_cuda_script(  # # noqa: C901
         success=True,
         nvcc_version=nvcc_version,
         command=_make_cmd(compile_process.args),
-        stdout=compile_process.stdout,
-        stderr=compile_process.stderr,
+        stdout=_limit_length(compile_process.stdout),
+        stderr=_limit_length(compile_process.stderr),
         exit_code=compile_process.returncode,
     )
 
@@ -161,8 +172,8 @@ def run_program(args: list[str], seed: int) -> RunResult:
         ),
         passed=result_dict.get("check", None) == "pass",
         command=_make_cmd(run_process.args),
-        stdout=run_process.stdout,
-        stderr=run_process.stderr,
+        stdout=_limit_length(run_process.stdout),
+        stderr=_limit_length(run_process.stderr),
         exit_code=run_process.returncode,
         duration=execution_end_time - execution_start_time,
         result=result_dict,