Skip to content

_forkserver_worker

genlm.eval.domains.ds1000._forkserver_worker

Warm fork-server for DS-1000 harness scripts: pre-imports the science libs (not tensorflow, fork-unsafe) and forks a child per request, or from a warm per-task session child. JSON lines in/out.

run_child(script, timeout, out_path, td)

Executed in the forked child. Never returns.

Source code in genlm/eval/domains/ds1000/_forkserver_worker.py
def run_child(script: str, timeout: float, out_path: str, td: str) -> None:
    """Executed in the forked child. Never returns."""
    rc = 0
    try:
        fd = os.open(out_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
        os.dup2(fd, 1)
        os.dup2(fd, 2)
        os.close(fd)
        _sandbox_into(td)
        signal.signal(signal.SIGALRM, signal.SIG_DFL)
        signal.alarm(max(1, int(timeout)))
        g = {"__name__": "__main__", "__builtins__": __builtins__}
        try:
            exec(compile(script, "<harness>", "exec"), g, g)
        except SystemExit:
            pass
        except BaseException:  # noqa: BLE001
            traceback.print_exc()
            rc = 1
    except BaseException:  # noqa: BLE001
        rc = 2
    finally:
        try:
            sys.stdout.flush()
            sys.stderr.flush()
        except Exception:  # noqa: BLE001
            pass
        os._exit(rc)

run_session(setup, cmd_r, res_w, td)

Session child: run setup once, then fork one grandchild per body.

Source code in genlm/eval/domains/ds1000/_forkserver_worker.py
def run_session(setup: str, cmd_r: int, res_w: int, td: str) -> None:
    """Session child: run setup once, then fork one grandchild per body."""

    def send(obj):
        os.write(res_w, (json.dumps(obj) + "\n").encode())

    try:
        os.setsid()
        log_fd = os.open(os.path.join(td, "__session__.log"),
                         os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
        os.dup2(log_fd, 1)
        os.dup2(log_fd, 2)
        os.close(log_fd)
        _sandbox_into(td)
        signal.signal(signal.SIGALRM, signal.SIG_DFL)
        signal.alarm(SETUP_TIMEOUT)
        ns = {"__name__": "__main__", "__builtins__": __builtins__}
        try:
            exec(compile(setup, "<session-setup>", "exec"), ns, ns)
        except BaseException:  # noqa: BLE001
            traceback.print_exc()
            send({"setup_failed": 1})
            os._exit(0)
        signal.alarm(0)
        send({"ready": 1})

        buf = b""
        while True:
            chunk = os.read(cmd_r, 1 << 20)
            if not chunk:
                os._exit(0)
            buf += chunk
            while b"\n" in buf:
                line, buf = buf.split(b"\n", 1)
                if not line.strip():
                    continue
                req = json.loads(line)
                sys.stdout.flush()
                sys.stderr.flush()
                try:
                    pid = os.fork()
                except OSError:
                    send({"id": req["id"], "forkfail": 1})
                    continue
                if pid == 0:
                    os.close(cmd_r)
                    os.close(res_w)
                    run_child_in_ns(req["body"], req["timeout"], req["out"], ns)
                _, status = os.waitpid(pid, 0)
                send({"id": req["id"], "status": status})
    except BaseException:  # noqa: BLE001
        try:
            send({"setup_failed": 1})
        except Exception:  # noqa: BLE001
            pass
        os._exit(2)

run_child_in_ns(body, timeout, out_path, ns)

Grandchild: execute a body script in the session namespace.

Source code in genlm/eval/domains/ds1000/_forkserver_worker.py
def run_child_in_ns(body: str, timeout: float, out_path: str, ns) -> None:
    """Grandchild: execute a body script in the session namespace."""
    rc = 0
    try:
        fd = os.open(out_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
        os.dup2(fd, 1)
        os.dup2(fd, 2)
        os.close(fd)
        signal.signal(signal.SIGALRM, signal.SIG_DFL)
        signal.alarm(max(1, int(timeout)))
        try:
            exec(compile(body, "<harness-body>", "exec"), ns, ns)
        except SystemExit:
            pass
        except BaseException:  # noqa: BLE001
            traceback.print_exc()
            rc = 1
    except BaseException:  # noqa: BLE001
        rc = 2
    finally:
        try:
            sys.stdout.flush()
            sys.stderr.flush()
        except Exception:  # noqa: BLE001
            pass
        os._exit(rc)