subprocess

class SubprocessExecutor(Executor):
    """A plain child process: killed after the time limit plus grace, pinned to its CPUs (on Linux). A CPU
    time limit is the kernel's (`RLIMIT_CPU`), of the main process only.

    No memory limit, and memory is the peak resident size of the main process only, so measurements are
    not reliable. For laptops and CI without cgroups.
    """

    name = "subprocess"
    pin = hasattr(os, "sched_setaffinity")

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._running: set[int] = set()

    def execute(self, job: WorkerJob) -> Measurement:
        cpus, cputime = job.cpus, job.limits.cputime_s
        hard_cpu = math.ceil(cputime + self.grace_s) if cputime is not None else None

        def limit():  # in the child, before it starts the worker
            if cpus:
                os.sched_setaffinity(0, cpus)
            if hard_cpu is not None:
                import resource

                resource.setrlimit(resource.RLIMIT_CPU, (hard_cpu, hard_cpu))  # then SIGKILL

        with open(job.log, "ab") as log:
            start = time.monotonic()
            proc = subprocess.Popen(
                job.cmd,
                stdout=log,
                stderr=subprocess.STDOUT,
                env={**os.environ, **job.env},
                start_new_session=True,  # its own process group, so a kill gets its children too
                preexec_fn=limit if cpus or hard_cpu is not None else None,
            )
            self._running.add(proc.pid)
            killed = threading.Event()

            def kill():
                killed.set()
                _killpg(proc.pid)

            timer = threading.Timer(job.limits.time_s + self.grace_s, kill)
            timer.start()
            try:
                _, status, usage = os.wait4(proc.pid, 0)
            finally:
                timer.cancel()
                self._running.discard(proc.pid)
            proc.returncode = os.waitstatus_to_exitcode(status)
        cputime_s = usage.ru_utime + usage.ru_stime
        termination = None
        if killed.is_set():
            termination = "walltime"
        elif (
            hard_cpu is not None and proc.returncode in (-signal.SIGKILL, -signal.SIGXCPU) and cputime_s >= hard_cpu - 1
        ):
            termination = "cputime"
        return Measurement(
            walltime_s=time.monotonic() - start,
            cputime_s=cputime_s,
            memory_mib=usage.ru_maxrss / 1024,  # KiB on Linux
            termination=termination,
            exitcode=proc.returncode,
            executor=self.name,
            reliable=self.reliable,
        )

    def close(self) -> None:
        for pid in list(self._running):  # only left if interrupted
            _killpg(pid)
        super().close()