mllm_sandbox.py

import mllm_client, pathlib, concurrent.futures, os, subprocess, json
from pydantic import BaseModel
from pathlib import Path
import argparse
import yaml, logging, shortuuid, sqlite3, reactivedb, time
from mllm_types import EnvConfig, GenerationOutput, GenerationEnv

base_readonly_locations = [
    '/usr', '/lib', '/lib64', '/lib32', '/bin', '/etc'
]

def command_in_env(env_config, mount_path, tip_revision, command):
    mounts = [('--bind', env_config.state_root / tip_revision, mount_path),
              ] + [ ('--ro-bind', loc, loc) for loc in base_readonly_locations + env_config.additional_readonly_locations ]

    mount_args = []
    for cmd, src, target in mounts:
        src = Path(src)
        target = Path(target)
        parent = target.parent
        while parent != Path('/'):
            mount_args += ['--dir', str(parent)]
            parent = parent.parent
        mount_args += [cmd, str(src), str(target)]

    return ['bwrap',
            '--unshare-pid', '--unshare-ipc', '--unshare-net', '--unshare-uts',
            '--dir', '/proc', '--proc', '/proc',
            '--dir', '/dev', '--dev', '/dev',
            '--dir', '/tmp', '--tmpfs', '/tmp',
            '--setenv', "MLLM_ISOLATED_ENV", "1",
            ] + mount_args + [
                '--chdir', mount_path, '--'
            ] + command

class TimeoutResponse():
    def __init__(self):
        self.returncode = 1
        self.stdout = b''
        self.stderr = b'Timed out.'

def run_command_in_env(env: GenerationEnv,
                       base_revision: str,
                       command: list[str], timeout: float) -> tuple[subprocess.CompletedProcess[bytes] | TimeoutResponse, str]:    
    new_revision = shortuuid.uuid()
    new_root_path = env.config.state_root / new_revision
    base_root_path = env.config.state_root / base_revision
    subprocess.check_call(['btrfs', 'subvolume', 'snapshot', base_root_path, new_root_path])

    cmd = command_in_env(env.config, env.mount_path, new_revision, command)
    try:
        resp: subprocess.CompletedProcess[bytes] | TimeoutResponse = \
            subprocess.run(cmd, stdin=subprocess.DEVNULL, capture_output=True, timeout=timeout) 
    except subprocess.TimeoutExpired:
        resp = TimeoutResponse()
    subprocess.check_call(['btrfs', 'property', 'set', new_root_path, 'ro', 'true'])

    return resp, new_revision

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='MLLM Sandbox Utilities')
    subparsers = parser.add_subparsers(dest='command', required=True)
    
    # Define the 'build' subcommand
    build_parser = subparsers.add_parser('build', help='Run the build script in the sandbox')
    build_parser.add_argument('--env-config', type=Path, required=True,
                              help='Path to environment configuration file')
    
    args = parser.parse_args()
    
    if args.command == 'build':
        # Load the environment configuration
        with open(args.env_config, 'r') as f:
            env_config_data = yaml.safe_load(f)
        env_config = EnvConfig.parse_obj(env_config_data)
        
        # Create a base snapshot
        base_revision = shortuuid.uuid()
        base_folder = env_config.state_root / "base"
        mount_path = args.env_config.parent.resolve()
        
        subprocess.check_call([
            'rsync', '--delete', '-a',
            str(mount_path) + '/', str(base_folder) + '/'
        ])
        subprocess.check_call([
            'btrfs', 'subvolume', 'snapshot',
            str(base_folder), str(env_config.state_root / base_revision)
        ])
        
        # Create a GenerationEnv object with minimal required fields
        env = GenerationEnv(
            config=env_config,
            task_id='build_test',  # Dummy value for testing
            generation_set_id='build_test_set',  # Dummy value for testing
            base_revision=base_revision,
            mount_path=mount_path,
            task_prompt='',  # Empty since we're testing the build
            files=[],  # Empty list since no specific files are needed
            created=time.time(),
        )
        
        # Run the build script in the sandbox
        build_command = ["bash", "-c", env_config.build_script]
        build_result, new_revision = run_command_in_env(
            env=env,
            base_revision=base_revision,
            command=build_command,
            timeout=60
        )
        
        # Print the build output and status
        output = (build_result.stdout + build_result.stderr).decode('utf8', 'replace')
        print('Build Output:')
        print(output)
        print(f'Build Return Code: {build_result.returncode}')
        if build_result.returncode == 0:
            print('Build succeeded.')
        else:
            print('Build failed.')