run_process.py

#!/usr/bin/env python3
# batch_run_process.py
"""
Batch runner for Process(project, seed). 

Usage:
    1. Put this file alongside the module that defines `Process`.
    2. Edit PROCESS_MODULE_PATH to point to that module (no .py).
    3. Edit the `cwe_list` below to define your CWEs.
    4. Run: python3 batch_run_process.py
"""

import os
import sys
import time
import traceback
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed

# -------------- CONFIG --------------
# 模块名（不带 .py），该模块中必须定义 Process 类（和你的 ROOT 常量等）
PROCESS_MODULE_PATH = "build_dataset"  # <- 修改为你的模块名，例如："build_dataset" 或 "your_script_name"

# 使用这个内置列表作为要处理的CWEs
cwe_list = [
    'CWE-119',
    'CWE-189',
    'CWE-190',
    'CWE-20',
    'CWE-200',
    'CWE-254',
    'CWE-264',
    'CWE-310',
    'CWE-362',
    'CWE-399',
    'CWE-400',
    'CWE-416',
    'CWE-426',
    'CWE-476',
    'CWE-59',
    'CWE-667',
    'CWE-787',
    'CWE-79',
    'CWE-94'
]

# 输出检查：如果该文件已存在，则默认跳过（避免重复计算）
SKIP_IF_OUTPUT_EXISTS = True

# 输出文件后缀/路径模式（根据你的 Process 类中 save_g 的构造）
# 你的 Process.save_g = f'{ROOT}/processed_data/{project}_all_dataset.pt'
OUTPUT_DIR_IN_MODULE = "processed_data"  # 如果你的模块中ROOT不同，请留意
OUTPUT_FILENAME_TEMPLATE = "{project}_all_dataset.pt"

# 并发设置（1=串行，>1并发）
MAX_WORKERS = 1

# 随机种子起始值（会按索引+seed_offset传给Process）
SEED_OFFSET = 0

# 日志设置
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s: %(message)s",
    datefmt="%H:%M:%S"
)
logger = logging.getLogger(__name__)
# ------------------------------------

# 动态导入 Process
try:
    module = __import__(PROCESS_MODULE_PATH, fromlist=["Process", "ROOT"])
except Exception as e:
    logger.error(f"Failed to import module '{PROCESS_MODULE_PATH}': {e}")
    logger.error("Please set PROCESS_MODULE_PATH to the module file (without .py) that defines `Process`.")
    sys.exit(1)

if not hasattr(module, "Process"):
    logger.error(f"Module '{PROCESS_MODULE_PATH}' does not define `Process` class.")
    sys.exit(1)

Process = getattr(module, "Process")


# helper: compute expected output path for a project (to check exist)
def expected_output_path(project):
    # try read ROOT from module if exists
    ROOT = getattr(module, "ROOT", None)
    if ROOT is None:
        # fallback: current working dir
        ROOT = os.getcwd()
    out_dir = os.path.join(ROOT, OUTPUT_DIR_IN_MODULE)
    os.makedirs(out_dir, exist_ok=True)
    fname = OUTPUT_FILENAME_TEMPLATE.format(project=project)
    return os.path.join(out_dir, fname)

# task runner
def run_project(project, seed=0, force=False):
    logger.info(f"[{project}] start (seed={seed})")
    out_path = expected_output_path(project)
    if SKIP_IF_OUTPUT_EXISTS and (not force) and os.path.exists(out_path):
        logger.info(f"[{project}] output {out_path} already exists -> skip")
        return {"project": project, "status": "skipped", "path": out_path}

    try:
        # instantiate and run
        p = Process(project, seed)
        p.save_all_set()
        logger.info(f"[{project}] finished, saved to {out_path} (if used module's save path)")
        return {"project": project, "status": "ok", "path": out_path}
    except Exception as e:
        tb = traceback.format_exc()
        logger.error(f"[{project}] ERROR: {e}\n{tb}")
        return {"project": project, "status": "error", "error": str(e)}

def main():
    results = []
    
    # --- MODIFIED: Directly use the predefined cwe_list for projects ---
    projects = list(cwe_list) # Initialize projects directly from the global cwe_list

    if not projects:
        logger.error("No CWE projects to process. Exiting.")
        sys.exit(1)
    # --- END MODIFIED SECTION ---

    if MAX_WORKERS <= 1:
        # serial
        for i, project in enumerate(projects):
            seed = SEED_OFFSET + i
            res = run_project(project, seed=seed)
            results.append(res)
    else:
        # parallel with ThreadPoolExecutor (safer than ProcessPool for heavy imports)
        with ThreadPoolExecutor(max_workers=MAX_WORKERS) as exe:
            futures = {}
            for i, project in enumerate(projects):
                seed = SEED_OFFSET + i
                futures[exe.submit(run_project, project, seed)] = project
            for fut in as_completed(futures):
                res = fut.result()
                results.append(res)

    # summary
    ok = [r for r in results if r["status"] == "ok"]
    skipped = [r for r in results if r["status"] == "skipped"]
    err = [r for r in results if r["status"] == "error"]
    logger.info(f"Done. OK: {len(ok)}, Skipped: {len(skipped)}, Error: {len(err)}")
    if err:
        logger.info("Errors:")
        for r in err:
            logger.info(f"  - {r['project']}: {r.get('error')}")
    # write summary
    with open("batch_run_summary.txt", "w", encoding="utf-8") as f:
        for r in results:
            f.write(str(r) + "\n")

if __name__ == "__main__":
    main()