add utils

f4f5ef6 verified 21 days ago

2.49 kB

"""Run CREDIT rollouts from the notebook in a subprocess, with live progress output."""

import codecs
import fcntl
import os
import pty
import select
import struct
import subprocess
import sys
import termios
from pathlib import Path

LAUNCHER = Path(__file__).with_name("credit_launch.py")


def run_configs(configs, fast_math=False):
    """Run rollout configs one after another in ONE subprocess (the model is loaded once; see credit_launch.py).
    The subprocess uses this kernel's interpreter and gets the kernel's sys.path as PYTHONPATH, so it finds packages
    the kernel finds (e.g. a `pip install --target` directory added to sys.path). Output goes through a pseudo-terminal
    so CREDIT's tqdm bars refresh in place instead of printing a line per update.
    """
    configs = [str(c) for c in configs]
    cmd = [sys.executable, str(LAUNCHER), *(["--fast"] if fast_math else []), *configs]
    print(f"$ credit rollout x {len(configs)} run(s){' (FAST_MATH)' if fast_math else ''}", flush=True)
    env = {**os.environ, "PYTHONPATH": os.pathsep.join(p for p in sys.path if p)}

    master, slave = pty.openpty()
    # A fresh pty reports 0 columns, and tqdm then draws no bar; give it a normal terminal size.
    fcntl.ioctl(slave, termios.TIOCSWINSZ, struct.pack("HHHH", 24, 120, 0, 0))
    proc = subprocess.Popen(cmd, stdout=slave, stderr=slave, close_fds=True, env=env)
    os.close(slave)  # the parent only reads the master end
    decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")  # reads can split multi-byte characters

    def forward(timeout):
        # Copy whatever the child wrote to the notebook. False when nothing was read.
        if not select.select([master], [], [], timeout)[0]:
            return False
        try:
            data = os.read(master, 1024)
        except OSError:  # Linux raises EIO once the child has exited and the pty is empty
            return False
        if not data:
            return False
        sys.stdout.write(decoder.decode(data))
        sys.stdout.flush()
        return True

    try:
        while True:
            if forward(0.1):
                continue
            if proc.poll() is not None:
                while forward(0):  # drain output still buffered after the child exited
                    pass
                break
    finally:
        os.close(master)
        proc.wait()
    if proc.returncode != 0:
        raise RuntimeError(f"credit rollout exited with code {proc.returncode}")

	"""Run CREDIT rollouts from the notebook in a subprocess, with live progress output."""

	import codecs
	import fcntl
	import os
	import pty
	import select
	import struct
	import subprocess
	import sys
	import termios
	from pathlib import Path

	LAUNCHER = Path(__file__).with_name("credit_launch.py")


	def run_configs(configs, fast_math=False):
	"""Run rollout configs one after another in ONE subprocess (the model is loaded once; see credit_launch.py).

	The subprocess uses this kernel's interpreter and gets the kernel's sys.path as PYTHONPATH, so it finds packages
	the kernel finds (e.g. a `pip install --target` directory added to sys.path). Output goes through a pseudo-terminal
	so CREDIT's tqdm bars refresh in place instead of printing a line per update.
	"""
	configs = [str(c) for c in configs]
	cmd = [sys.executable, str(LAUNCHER), (["--fast"] if fast_math else []), configs]
	print(f"$ credit rollout x {len(configs)} run(s){' (FAST_MATH)' if fast_math else ''}", flush=True)
	env = {**os.environ, "PYTHONPATH": os.pathsep.join(p for p in sys.path if p)}

	master, slave = pty.openpty()
	# A fresh pty reports 0 columns, and tqdm then draws no bar; give it a normal terminal size.
	fcntl.ioctl(slave, termios.TIOCSWINSZ, struct.pack("HHHH", 24, 120, 0, 0))
	proc = subprocess.Popen(cmd, stdout=slave, stderr=slave, close_fds=True, env=env)
	os.close(slave) # the parent only reads the master end
	decoder = codecs.getincrementaldecoder("utf-8")(errors="replace") # reads can split multi-byte characters

	def forward(timeout):
	# Copy whatever the child wrote to the notebook. False when nothing was read.
	if not select.select([master], [], [], timeout)[0]:
	return False
	try:
	data = os.read(master, 1024)
	except OSError: # Linux raises EIO once the child has exited and the pty is empty
	return False
	if not data:
	return False
	sys.stdout.write(decoder.decode(data))
	sys.stdout.flush()
	return True

	try:
	while True:
	if forward(0.1):
	continue
	if proc.poll() is not None:
	while forward(0): # drain output still buffered after the child exited
	pass
	break
	finally:
	os.close(master)
	proc.wait()
	if proc.returncode != 0:
	raise RuntimeError(f"credit rollout exited with code {proc.returncode}")