[CI] Slim JIT kernel unit tests (#36887)

python/sglang/test/ci/ci_utils.py

@@ -23,6 +23,66 @@ class TestFile:
     estimated_time: float = 60
 
 
+class _ForkTestWorker:
+    """Preloaded interpreter that forks an isolated child for each test file."""
+
+    def __init__(self):
+        result_read_fd, result_write_fd = os.pipe()
+        worker_path = os.path.join(os.path.dirname(__file__), "fork_test_worker.py")
+        self.process = subprocess.Popen(
+            ["python3", worker_path, "--result-fd", str(result_write_fd)],
+            stdin=subprocess.PIPE,
+            stdout=None,
+            stderr=None,
+            text=True,
+            pass_fds=(result_write_fd,),
+        )
+        os.close(result_write_fd)
+        self.result_stream = os.fdopen(result_read_fd)
+        self.files_run = 0
+
+    def run(self, filename: str) -> tuple[int, float]:
+        tic = time.perf_counter()
+        if self.process.poll() is not None or self.process.stdin is None:
+            return 1, 0.0
+        try:
+            self.process.stdin.write(json.dumps({"filename": filename}) + "\n")
+            self.process.stdin.flush()
+            result_line = self.result_stream.readline()
+        except (BrokenPipeError, OSError):
+            return 1, time.perf_counter() - tic
+        if not result_line:
+            return 1, time.perf_counter() - tic
+        try:
+            result = json.loads(result_line)
+        except json.JSONDecodeError:
+            return 1, time.perf_counter() - tic
+        self.files_run += 1
+        return int(result["returncode"]), float(result["elapsed"])
+
+    def close(self, terminate: bool = False):
+        if self.process.poll() is None:
+            if terminate:
+                kill_process_tree(self.process.pid)
+            elif self.process.stdin is not None:
+                try:
+                    self.process.stdin.write(json.dumps({"command": "stop"}) + "\n")
+                    self.process.stdin.flush()
+                    self.process.wait(timeout=10)
+                except (BrokenPipeError, subprocess.TimeoutExpired):
+                    kill_process_tree(self.process.pid)
+        if self.process.poll() is None:
+            self.process.kill()
+        try:
+            self.process.wait(timeout=10)
+        except subprocess.TimeoutExpired:
+            self.process.kill()
+            self.process.wait()
+        if self.process.stdin is not None:
+            self.process.stdin.close()
+        self.result_stream.close()
+
+
 # Patterns that indicate retriable accuracy/performance failures
 RETRIABLE_PATTERNS = [
     r"AssertionError:.*not greater than",
@@ -158,6 +218,7 @@ def run_unittest_files(
     enable_retry: bool = False,
     max_attempts: int = 2,
     retry_wait_seconds: int = 60,
+    fork_worker_batch_size: int = 1,
 ):
     """
     Run a list of test files.
@@ -172,6 +233,9 @@ def run_unittest_files(
                      assertion failures (not code errors).
         max_attempts: Maximum number of attempts per file including initial run (default: 2).
         retry_wait_seconds: Seconds to wait between retries (default: 60).
+        fork_worker_batch_size: Number of files served by one preloaded fork
+                                worker. Each file still runs in a fresh child
+                                process. One keeps the existing exec behavior.
     """
     coredump_enabled = cuda_coredump.is_enabled()
     if coredump_enabled:
@@ -185,6 +249,8 @@ def run_unittest_files(
     # Per-file elapsed seconds, latest attempt wins. Consumed by the
     # TIMINGS block emitted at the end of this function.
     file_elapsed: Dict[str, float] = {}
+    fork_worker = None
+    use_fork_worker = fork_worker_batch_size > 1 and not enable_retry
 
     for i, file in enumerate(files):
         if isinstance(file, CIRegistry):
@@ -203,7 +269,7 @@ def run_unittest_files(
         output_lines = []
 
         def run_one_file(filename, capture_output=False):
-            nonlocal process, output_lines
+            nonlocal process, output_lines, fork_worker
 
             full_path = os.path.join(os.getcwd(), filename)
             logger.info(
@@ -211,10 +277,24 @@ def run_unittest_files(
             )
             file_tic = time.perf_counter()
 
-            cmd = ["python3", full_path, "-f"]
-
-            if capture_output:
+            if use_fork_worker:
+                if (
+                    fork_worker is None
+                    or fork_worker.files_run >= fork_worker_batch_size
+                ):
+                    if fork_worker is not None:
+                        fork_worker.close()
+                    fork_worker = _ForkTestWorker()
+                process = fork_worker.process
+                ret_code, _ = fork_worker.run(full_path)
+                if ret_code != 0 or fork_worker.files_run >= fork_worker_batch_size:
+                    fork_worker.close()
+                    fork_worker = None
+                    process = None
+                elapsed = time.perf_counter() - file_tic
+            elif capture_output:
                 # Capture output for retry decision
+                cmd = ["python3", full_path, "-f"]
                 process = subprocess.Popen(
                     cmd,
                     stdout=subprocess.PIPE,
@@ -236,16 +316,19 @@ def run_unittest_files(
                 # Bounded wait for the reader to finish.
                 reader_thread.join(timeout=60)
             else:
+                cmd = ["python3", full_path, "-f"]
                 process = subprocess.Popen(cmd, stdout=None, stderr=None)
                 process.wait()
 
-            elapsed = time.perf_counter() - file_tic
+            if not use_fork_worker:
+                elapsed = time.perf_counter() - file_tic
+                ret_code = process.returncode
             file_elapsed[filename] = elapsed
 
             logger.info(
                 f".\n.\nEnd ({i}/{len(files) - 1}):\n{filename=}, {elapsed=:.0f}, {estimated_time=}\n.\n.\n"
             )
-            return process.returncode
+            return ret_code
 
         # Retry loop for each file
         attempt = 1
@@ -305,7 +388,11 @@ def run_unittest_files(
                     break
 
             except TimeoutError:
-                kill_process_tree(process.pid)
+                if fork_worker is not None:
+                    fork_worker.close(terminate=True)
+                    fork_worker = None
+                elif process is not None:
+                    kill_process_tree(process.pid)
                 time.sleep(5)
                 # TimeoutError aborts run_one_file before its elapsed write;
                 # record the timeout cap as an upper bound so the file still
@@ -333,6 +420,9 @@ def run_unittest_files(
             if not continue_on_error:
                 break
 
+    if fork_worker is not None:
+        fork_worker.close()
+
     elapsed_total = time.perf_counter() - tic
 
     if coredump_enabled and not success:

python/sglang/test/ci/fork_test_worker.py

@@ -0,0 +1,105 @@
+"""Execute test files in isolated fork children of a preloaded interpreter."""
+
+import argparse
+import json
+import os
+import runpy
+import signal
+import sys
+import time
+import traceback
+
+
+def _preload_common_modules() -> None:
+    """Load the expensive common stack without creating a CUDA context."""
+    import numpy  # noqa: F401
+    import pytest  # noqa: F401
+    import scipy  # noqa: F401
+    import torch
+    import triton  # noqa: F401
+
+    if torch.cuda.is_initialized():
+        raise RuntimeError("fork test worker initialized CUDA before forking")
+
+
+def _normalize_exit_code(code) -> int:
+    if code is None:
+        return 0
+    if isinstance(code, int):
+        return code
+    print(code, file=sys.stderr, flush=True)
+    return 1
+
+
+def _run_file(filename: str) -> int:
+    sys.argv = [filename, "-f"]
+    try:
+        runpy.run_path(filename, run_name="__main__")
+        return 0
+    except SystemExit as exc:
+        return _normalize_exit_code(exc.code)
+    except BaseException:
+        traceback.print_exc()
+        return 1
+
+
+def _wait_status_to_returncode(status: int) -> int:
+    if os.WIFEXITED(status):
+        return os.WEXITSTATUS(status)
+    if os.WIFSIGNALED(status):
+        return 128 + os.WTERMSIG(status)
+    return 1
+
+
+def run_file_in_fork(filename: str) -> tuple[int, float]:
+    sys.stdout.flush()
+    sys.stderr.flush()
+    tic = time.perf_counter()
+    child_pid = os.fork()
+    if child_pid == 0:
+        signal.signal(signal.SIGINT, signal.SIG_DFL)
+        returncode = _run_file(filename)
+        sys.stdout.flush()
+        sys.stderr.flush()
+        os._exit(max(0, min(returncode, 255)))
+
+    _, status = os.waitpid(child_pid, 0)
+    return _wait_status_to_returncode(status), time.perf_counter() - tic
+
+
+def main() -> int:
+    if not hasattr(os, "fork"):
+        raise RuntimeError("fork test worker requires a POSIX platform")
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--result-fd", type=int, required=True)
+    args = parser.parse_args()
+    _preload_common_modules()
+
+    with os.fdopen(args.result_fd, "w", buffering=1) as result_stream:
+        for line in sys.stdin:
+            command = json.loads(line)
+            if command.get("command") == "stop":
+                break
+
+            filename = command["filename"]
+            returncode, elapsed = run_file_in_fork(filename)
+            result_stream.write(
+                json.dumps(
+                    {
+                        "filename": filename,
+                        "returncode": returncode,
+                        "elapsed": elapsed,
+                    }
+                )
+                + "\n"
+            )
+
+    return 0
+
+
+if __name__ == "__main__":
+    # Ignore Ctrl-C in the preloader. The active child keeps the default
+    # handler, while suite timeouts terminate the complete process tree.
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
+    raise SystemExit(main())