hpo_vnncomp.py

import os
import sys
import time
import yaml
import random
import json
import traceback
import argparse
import pickle
import subprocess
from functools import partial
import numpy as np
from smac import Scenario, AlgorithmConfigurationFacade
import torch

from util import count_nonlinear_activations, count_onnx_inputs, count_onnx_outputs, count_onnx_parameters, smac_to_abcrown_config, NumpyEncoder, nested_set
from configspace_abcrown import build_abcrown_config_space

MAX_LOSS = 2**30

def get_hash_from_dict(dict_to_hash):
    return hash(json.dumps(dict_to_hash, sort_keys=True))

def limited_abcrown_vnncomp_smac(config, seed, instance, benchmark_dir='/tmp', work_dir='/tmp', log_path='/tmp/logs', no_cores=28, par_factor=10, fixed_config_keys=None):
    timestamp = time.time()
    _, _, timeout_str = instance.split('|')
    timeout = float(timeout_str)
    outer_timeout = timeout * 1.5 + 60 # Give ABCROWN some buffer to timeout internally first
    
    config = json.loads(json.dumps(dict(config), cls=NumpyEncoder))
    if fixed_config_keys is not None:
        for key, value in fixed_config_keys.items():
            config[key] = value

    abcrown_args = {
        "config": config,
        "seed": seed,
        "instance": instance,
        "benchmark_dir": benchmark_dir,
        "no_cores": no_cores,
        "par_factor": par_factor
    }

    config_hash = get_hash_from_dict(json.loads(json.dumps(dict(config), cls=NumpyEncoder)))
    
    args_pkl_path = f'{work_dir}/args_vnncomp_{timestamp}.pkl'
    result_path = f"{work_dir}/result_vnncomp_{timestamp}.pkl"
    
    with open(args_pkl_path, "wb") as f:
        pickle.dump(abcrown_args, f)
        
    verification_ok = False
    
    runner_args = [args_pkl_path, result_path]
    
    try:
        print(f"Running ['python3', 'abcrown_vnncomp_smac_runner.py'] + runner_args")
        process = subprocess.Popen(
            ["python3", "abcrown_vnncomp_smac_runner.py"] + runner_args,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )

        try:
            stdout, stderr = process.communicate(timeout=outer_timeout)
            print("Function finished successfully.")
            verification_ok = True

        except subprocess.TimeoutExpired:
            print(f"Function exceeded timeout of {outer_timeout} seconds. Terminating...")
            process.terminate()

            try:
                # Give the process time to exit cleanly
                stdout, stderr = process.communicate(timeout=10)
            except subprocess.TimeoutExpired:
                print("Function did not terminate after SIGTERM. Killing...")
                process.kill()
                stdout, stderr = process.communicate()

                print("Partial Output:", stdout.decode())
                print("Partial Error Output:", stderr.decode())
                verification_ok = False

    except Exception as e:
        print(f"Error running the process: {e}")
        stdout, stderr = b"", b""
    
    instance_name = instance.replace('/', '_').replace('|', '_').replace('.vnnlib', '')

    os.makedirs(f'{log_path}/{config_hash}/{instance_name}', exist_ok=True)
    with open(f'{log_path}/{config_hash}/{instance_name}/stdout.log', 'wb') as f:
        f.write(stdout)
    with open(f'{log_path}/{config_hash}/{instance_name}/stderr.log', 'wb') as f:
        f.write(stderr)
    
    if verification_ok:
        try:
            with open(result_path, 'rb') as f:
                running_time, result = pickle.load(f)
                if isinstance(result, dict):
                    result['config_hash'] = config_hash
                elif isinstance(result, str):
                    result = {'result': result, 'config_hash': config_hash}
                else:
                    result = {'result': 'unknown (bad format)', 'config_hash': config_hash}

            return running_time, result
        except:
            return MAX_LOSS, {'result': 'unknown (crashed)', 'config_hash': config_hash}
    
    return MAX_LOSS, {'result': 'unknown (crashed)', 'config_hash': config_hash}

def read_instances_csv(csv_path):
    instances = []
    with open(csv_path, 'r', encoding='u8') as f:
        for line in f:
            line = line.strip()
            if not line: continue
            parts = line.split(',')
            if len(parts) >= 3:
                # Format: onnx_path, vnnlib_path, timeout
                onnx_path, vnnlib_path, timeout = parts[0], parts[1], parts[2]
                instances.append(f"{onnx_path}|{vnnlib_path}|{timeout}")
    return instances

def run_hpo(config_file):
    with open(config_file, 'r', encoding='u8') as f:
        config = json.load(f)
        
    current_slurm_id = os.environ.get('SLURM_JOB_ID', str(random.choice(range(1000, 9999))))

    experiment_name = config.get('experiment_name', 'vnncomp_cifar100')
    benchmark_dir = config['benchmark_dir']
    
    tune_on_test = config.get('tune_on_test', False)
    print(f'Tuning on {"test" if tune_on_test else "train"} set.')
    
    instances_file_name = config.get('test_instances', 'test_instances.csv') if tune_on_test else config.get('train_instances', 'train_instances.csv')
    instances_csv_path = os.path.join(benchmark_dir, instances_file_name)
    print(f'Loading instances from {instances_csv_path}')
    
    if not os.path.isfile(instances_csv_path):
        print(f"Error: Could not find instances csv at {instances_csv_path}")
        sys.exit(1)
    
    # eps = config['eps']
    # no_classes = config['no_classes']
    eps = 0.01 # dummy value (only needed in optimising attack parameters, which we dont do either way)
    no_classes = 100 # dummy value
    walltime_limit = config['walltime_limit']
    trial_limit = config['trial_limit']
    no_cores = config['no_cores']
    results_path = config['results_path']
    seed = config.get('seed', 42)
    
    random.seed(seed)
    torch.manual_seed(seed)
    
    results_path = f'{results_path}/{experiment_name}_{current_slurm_id}'
    os.makedirs(f'{results_path}', exist_ok=True)
        
    work_dir = f'/tmp/{current_slurm_id}_vnncomp'
    os.makedirs(work_dir, exist_ok=True)
    
    instances_list = read_instances_csv(instances_csv_path)

    onnx_paths = set()
    for inst in instances_list:
        onnx_path, _, _ = inst.split('|')
        onnx_paths.add(benchmark_dir + '/' + onnx_path)
    
    max_inputs = 0
    max_outputs = 0
    for onnx_path in onnx_paths:
        param_count = count_onnx_parameters(onnx_path)
        # activation_count, activation_elements = count_nonlinear_activations(onnx_path)
        no_inputs = count_onnx_inputs(onnx_path)
        no_outputs = count_onnx_outputs(onnx_path)
        max_inputs = max(max_inputs, no_inputs)
        max_outputs = max(max_outputs, no_outputs)
        print(f"Model {onnx_path} has {param_count} parameters, {no_inputs} inputs, and {no_outputs} outputs.")
    
    print(f'Running optimization on {len(instances_list)} instances.')
    
    # Define dummy features for instances
    smac_instance_features = {
        inst: [0.0, 0.0, 0.0, 0.0] for inst in instances_list
    }

    
        
    tune_flag = config.get('tune_attack_extensive', False)
    cs = build_abcrown_config_space(
        eps=eps,
        no_classes=no_classes, 
        include_bab=config.get('include_bab', True), 
        include_mip=config.get('include_mip', True),
        include_bab_refine=config.get('include_bab_refine', True),
        include_input_split=config.get('include_input_split', True),
        include_attack=config.get('include_attack', True),
        tune_attack_extensive=tune_flag,
    )
    print(f"[hpo_vnncomp] build config space with tune_attack_extensive={tune_flag}")

    fixed_config_keys = config.get('fixed_config_keys', None)

    scenario = Scenario(
        configspace=cs,
        deterministic=True,
        walltime_limit=walltime_limit,
        instances=instances_list,
        instance_features=smac_instance_features,
        n_trials=trial_limit,
        output_directory=f'{results_path}',
    )
    
    smac = AlgorithmConfigurationFacade(
        scenario=scenario,
        target_function=partial(limited_abcrown_vnncomp_smac, benchmark_dir=benchmark_dir, work_dir=work_dir, no_cores=no_cores, log_path=f'{results_path}/logs', fixed_config_keys=fixed_config_keys),
        overwrite=True
    )
    
    inc = smac.optimize()
    
    np.save(f'{results_path}/results.npy', [dict(inc), smac_instance_features])
    
    with open(f'{results_path}/vnncomp_hpo_conf.json', 'w', encoding='u8') as f:
        json.dump(dict(inc), f, indent=2)
    
    hpo_abcrown_conf = smac_to_abcrown_config(dict(inc))
    
    with open(f'{results_path}/vnncomp_hpo_conf.yaml', 'w', encoding='u8') as f:
        yaml.dump(hpo_abcrown_conf, f, Dumper=yaml.SafeDumper)
    
    print("Optimization Completed!")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, required=True, help='Path to hpo_vnncomp json config')
    args = parser.parse_args()
    print(args)
    run_hpo(config_file=args.config)