mllm_make_task.py

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

def parse_task_prompt(raw_task_prompt, base_path):
    task_prompt = []
    files = []
    difficulty = None
    for line in raw_task_prompt.splitlines():
        if line.strip().startswith('files:'):
            file_patterns = [f.strip() for f in line.split(':', 1)[1].split(',')]
            for pattern in file_patterns:
                matched_files = glob.glob(str(base_path / pattern), recursive=True)
                if not matched_files:
                    raise Exception('pattern %r doesn\'t match any file' % pattern)
                for f in matched_files:
                    rel_path = os.path.relpath(f, base_path)
                    files.append(rel_path)
        else:
            stripped_line = line.strip()
            if stripped_line.startswith("EASY"):
                difficulty = "EASY"
                line_without_annotation = stripped_line[4:].strip()
                if line_without_annotation:
                    task_prompt.append(line_without_annotation)
            elif stripped_line.startswith("HARD"):
                difficulty = "HARD"
                line_without_annotation = stripped_line[4:].strip()
                if line_without_annotation:
                    task_prompt.append(line_without_annotation)
            else:
                task_prompt.append(line)

    if not files:
        raise Exception('no files provided')
    return files, '\n'.join(task_prompt), difficulty

def run(env_config_path: Path, task_file_path: Path, task_id):
    
    mount_path = env_config_path.parent
    task_items = mllm_task_file.parse_text(Path(task_file_path).read_text())

    raw_task_prompt = None
    for item in task_items:
        if item.task_id == task_id:
            raw_task_prompt = '\n'.join(item.content)
    if raw_task_prompt is None:
        raise Exception('task with id=%r missing' % task_id)

    with open(env_config_path, 'r') as f: 
        env_config_data = yaml.safe_load(f)
    env_config = EnvConfig.parse_obj(env_config_data)

    db = reactivedb.Db(env_config.state_root / "db.sqlite3")
    generation_env_table = db.table(GenerationEnv)
    
    generation_set_id = shortuuid.uuid()
    base_revision = shortuuid.uuid()

    base_folder = env_config.state_root / "base"

    subprocess.check_call(['rsync', '--delete', '-a', str(mount_path) + '/', str(base_folder) + '/'])

    subprocess.check_call(['btrfs', 'subvolume', 'snapshot',
                           str(base_folder), env_config.state_root / base_revision])
    
    build_result = subprocess.run(mllm_sandbox.command_in_env(
        env_config, mount_path, base_revision,
        ["bash", "-c", env_config.build_script]), stdin=subprocess.DEVNULL, capture_output=True, timeout=120) 
    print('Build return code: ', build_result)
    
    files, task_prompt, difficulty = parse_task_prompt(raw_task_prompt, env_config.state_root / base_revision)
    
    env = GenerationEnv(
        config=env_config,
        task_id=task_id,
        generation_set_id=generation_set_id,
        base_revision=base_revision,
        mount_path=mount_path,
        task_prompt=task_prompt,
        files=files,
        build_successful = build_result.returncode == 0,
        build_output = (build_result.stdout + build_result.stderr).decode('utf8', 'replace'),
        created=time.time(),
        difficulty=difficulty
    )
    generation_env_table.insert(env)
    db.commit()
    return generation_set_id

def start_prompt(generation_set_id, env_config_path, task_id):    
    with open(env_config_path, 'r') as f: 
        env_config_data = yaml.safe_load(f)
    env_config = EnvConfig.parse_obj(env_config_data)

    db = reactivedb.Db(env_config.state_root / "db.sqlite3")
    generation_env_table = db.table(GenerationEnv)
    env = generation_env_table.get(generation_set_id=generation_set_id)
    max_phases = '2'
    if env.difficulty == 'EASY':
        models = [('o1-mini', 'claude-sonnet', 'gpt-4o-mini'),
                  ('claude-sonnet', 'o1-mini', 'gpt-4o-mini'),]
    elif env.difficulty == 'HARD':
        models = [('o1:high', 'o1,claude-sonnet', 'gpt-4o-mini'),
                  ('claude-sonnet', 'o1,claude-sonnet', 'gpt-4o-mini'),
                  #('gemini-2.0-flash-thinking-exp-1219', 'gemini-2.0-flash-thinking-exp-1219', 'gpt-4o-mini'),
                  ('gemini-exp-1206', 'gemini-exp-1206', 'gpt-4o-mini'),
                  ]
        max_phases = '3'
    else:
        models = [('o1', 'o1,claude-sonnet', 'gpt-4o-mini'),
                  ('claude-sonnet', 'o1,claude-sonnet', 'gpt-4o-mini')]
        
    for model, second_stage_model, diff_model in models:
        cmd = [sys.executable, 'mllm_prompt.py',
               '--state-root', env_config.state_root,
               '--id', generation_set_id,
               '--model', model, '--second-stage-model', second_stage_model,
               '--diff-model', diff_model, '--max-phases', max_phases]
        
        cwd = Path(__file__).parent
        subprocess.check_call(['systemd-run', '--user',
                               '--unit=mllm-task-%s-%s-%s-prompt-%s'
                               % (task_id,
                                  generation_set_id,
                                  model, second_stage_model),
                               '--slice=mllm.slice',
                               '--working-directory=%s' % cwd,
                               ] + cmd)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Process MLLM tasks')
    parser.add_argument('--env-config', type=Path, required=True,
                      help='Path to environment configuration file')
    parser.add_argument('--task-file', type=Path, required=True,
                      help='Path to task file')
    parser.add_argument('--task-id', type=str, required=True,
                      help='ID of the task to run')
    parser.add_argument('--and-start-prompt', action='store_true')
    
    args = parser.parse_args()
    
    generation_set_id = run(
        env_config_path=args.env_config,
        task_file_path=args.task_file,
        task_id=args.task_id
    )
    print(generation_set_id)
    if args.and_start_prompt:
        start_prompt(generation_set_id, env_config_path=args.env_config, task_id=args.task_id)