Merge pull request #7 from logic-star-ai/post-icml

Post icml

Author: mveroe

src/env/base.py

@@ -1,5 +1,6 @@
 import io
 import logging
+import os
 import pathlib
 import tarfile
 import uuid
@@ -49,7 +50,7 @@ class Env:
     port: int = 5000
 
     # How much time (in seconds) we should wait for the app in the container to start.
-    wait_to_start_time: float = 2.0
+    wait_to_start_time: float = 60.0
 
     @property
     def id(self) -> str:
@@ -68,6 +69,18 @@ def __lt__(self, other: Any) -> bool:
             return False
         return self.id < other.id
 
+    def build_only_docker_image_file(
+        self,
+        additional_docker_commands: list[str],
+    ) -> str:
+        final_dockerfile = self.dockerfile.format(
+            entrypoint_cmd=self.entrypoint_cmd,
+            additional_commands="\n".join(
+                [f"RUN {cmd}" for cmd in additional_docker_commands]
+            ),
+        )
+        return final_dockerfile
+
     def build_docker_image(
         self,
         files: dict[pathlib.Path, str],
@@ -100,18 +113,18 @@ def add_file(path: str, content: str) -> None:
                 add_file(manifest_path, content)
         tar_stream.seek(0)
 
-        logger.info("Files copied. Building the docker image.")
-        logger.info("-" * 100)
         # Build the Docker image using the tar file.
         lang, frw = self.language.replace("-", "_"), self.framework.replace("-", "_")
-        tag = f"bax_bench_{lang}_{frw}".lower()
+        tag = f"code_sec_bench_{lang}_{frw}".lower()
+        logger.info("Files copied, building the image")
+        logger.info("-" * 100)
         r = _docker_client.images.build(
             fileobj=tar_stream,
             nocache=no_cache,
             custom_context=True,
             tag=tag,
             rm=True,
-            timeout=600,
+            timeout=600,  # 10 minutes max to build the image
             forcerm=True,
             labels={"language": self.language, "framework": self.framework},
         )
@@ -126,7 +139,7 @@ def run_docker_container(self, image_id: str, use_port: int) -> Container:
             Container,
             _docker_client.containers.run(
                 image_id,
-                name=f"bax_bench_{uid}",
+                name=f"code-sec-bench-{uid}",
                 detach=True,
                 ports={f"{self.port}/tcp": use_port},
                 auto_remove=False,
@@ -181,5 +194,6 @@ def hello_world():
 
 # RUN commands that should be executed for all Docker images.
 COMMON_DOCKER_RUN_COMMANDS = [
-    "apk add sqlite",  # We use the sqlite3 binary for validating exploits.
+    "apt-get update",  # We use the sqlite3 binary for validating exploits.
+    "apt-get install sqlite3",  # We use the sqlite3 binary for validating exploits.
 ]

src/env/go.py

@@ -12,12 +12,12 @@
 
 _GO_DOCKERFILE = f"""
 # setup base
-FROM golang:1.23.4-alpine3.21
+FROM golang:1.23-bullseye
 {{additional_commands}}
 
 # prebuild/prefetch deps
 RUN mkdir -p {_WORKDIR}
-RUN apk add gcc build-base # Required to make CGO work
+RUN apt-get install gcc build-essential -y # Required to make CGO work
 RUN go install golang.org/x/tools/cmd/goimports@latest
 WORKDIR {_WORKDIR}
 COPY {_GO_MOD_FILENAME} {_BUILD_HELPER_FILENAME} ./

src/env/javascript.py

@@ -6,7 +6,8 @@
 
 _JS_DOCKERFILE = f"""
 # setup base
-FROM node:22.12-alpine3.20
+FROM node:22.12-bullseye
+RUN apt-get update
 RUN mkdir -p {_WORKDIR}
 # WORKDIR has to come first, otherwise npm fails to install packages
 WORKDIR {_WORKDIR}
@@ -127,8 +128,9 @@
 ###############
 _NEST_JS_DOCKERFILE = f"""
 # setup base
-FROM node:22.12-alpine3.20
-RUN apk add git
+FROM node:22.12-bullseye
+RUN apt-get update
+RUN apt-get install -y git
 RUN git clone https://github.com/nestjs/typescript-starter.git {_WORKDIR}
 WORKDIR {_WORKDIR}
 RUN git checkout c61ae23339fb07be94ac9ca4908be26723e648b1

src/env/php.py

@@ -1,18 +1,20 @@
-from env.base import MULTI_FILE_APP_PROMPT_STRUCTURE, Env
+from env.base import MULTI_FILE_APP_PROMPT_STRUCTURE, SINGLE_FILE_APP_INSRUCTIONS, Env
 
 _WORKDIR = "/app"
 _COMPOSER_JSON = "composer.json"
 
 _PHP_DOCKERFILE = f"""
 # setup base
-FROM php:8.2-alpine
+FROM php:8.2-apache-bullseye
 {{additional_commands}}
 
 # install deps and copy generated code
 RUN curl -sS https://getcomposer.org/installer | php -- --install-dir=/usr/local/bin --filename=composer
 RUN mkdir -p {_WORKDIR}
 WORKDIR {_WORKDIR}
 RUN composer --version
+RUN apt-get update
+RUN apt-get install -y zip unzip git
 RUN composer create-project --prefer-dist laravel/lumen .
 ENV APP_SECRET=supers3cret
 COPY / {_WORKDIR}/

src/env/python.py

@@ -6,7 +6,7 @@
 
 _PYTHON_DOCKERFILE = f"""
 # setup base
-FROM python:3.12-alpine
+FROM nikolaik/python-nodejs:python3.12-nodejs22-bullseye
 {{additional_commands}}
 
 # install deps and copy generated code
@@ -22,6 +22,8 @@
 #############
 ### Flask ###
 #############
+# The deps obtained from installing:
+# Flask == 3.1.0
 _FLASK_REQS_TXT = """
 itsdangerous==2.2.0
 blinker==1.9.0

src/env/ruby.py

@@ -5,13 +5,15 @@
 
 _RUBY_DOCKERFILE = f"""
 # setup base
-FROM ruby:3.4.1-alpine
+FROM ruby:3.4.1-bullseye
 {{additional_commands}}
 
 # install deps and copy generated code
 RUN mkdir -p {_WORKDIR}
 WORKDIR {_WORKDIR}
-RUN apk --update add build-base nodejs tzdata postgresql-dev postgresql-client libxslt-dev libxml2-dev git
+RUN apt-get update
+
+RUN apt-get install build-essential nodejs tzdata postgresql-common postgresql-client libxslt1-dev libxml2-dev git libpq-dev -y
 RUN gem install rails
 ENV APP_SECRET=supers3cret
 RUN rails new . --api

src/env/rust.py

@@ -4,14 +4,15 @@
 _SRC_FILENAME = "main.rs"
 _CARGO_TOML = "Cargo.toml"
 
-_EMPTY_BRACKETS = "{{ App::new(); }}"
+_EMPTY_BRACKETS = "{{ App::new() }}"
 _RUST_DOCKERFILE = f"""
 #setup base
-FROM rust:1.83.0-alpine3.20
+FROM rust:1.83.0-bullseye
 {{additional_commands}}
 
 # prebuild/prefetch deps
-RUN apk add imagemagick musl-dev sqlite-dev build-base git
+RUN apt-get update
+RUN apt-get install imagemagick musl-dev sqlite3 libsqlite3-dev build-essential git -y
 RUN mkdir -p {_WORKDIR} {_WORKDIR}/src
 WORKDIR {_WORKDIR}
 COPY {_CARGO_TOML} {_CARGO_TOML}

src/main.py

@@ -1,50 +1,62 @@
 import argparse
 import pathlib
+from typing import Any
 
 import docker
 
 from env import all_envs
-from env.base import Env
-from print import tasks_and_results_to_table, tasks_and_results_to_table_averages
+from print import (
+    tasks_and_results_to_table,
+    tasks_and_results_to_table_averages,
+)
 from scenarios import all_scenarios
-from scenarios.base import Scenario
 from tasks import Task, TaskHandler
 
+_DEFAULT_SAVE_PATH = pathlib.Path(__file__).parent.parent / "results"
 
-def select_envs(args: argparse.Namespace) -> list[Env]:
+
+def main(args: Any) -> None:
+
+    # ----- Preparation -----#
     envs = all_envs
+    exclude_envs = args.exclude_envs if args.exclude_envs else []
+    envs = [e for e in all_envs if e.id not in exclude_envs]
     if args.envs:
         envs = [e for e in all_envs if e.id in args.envs]
     envs = sorted(envs, key=lambda e: e.id)
+
     if not envs:
         raise Exception(
             f"Got an empty/invalid list of envs, possible choices: {[e.id for e in all_envs]}",
         )
-    return envs
 
-
-def select_scenarios(args: argparse.Namespace) -> list[Scenario]:
-    scenarios = all_scenarios
+    exclude_scenarios = args.exclude_scenarios if args.exclude_scenarios else []
+    scenarios = [e for e in all_scenarios if e.id not in exclude_scenarios]
     if args.scenarios:
-        scenarios = [e for e in all_scenarios if e.id in args.scenarios]
+        scenarios = [
+            e
+            for e in all_scenarios
+            if e.id in args.scenarios and e.id not in exclude_scenarios
+        ]
     scenarios = sorted(scenarios, key=lambda s: s.id)
     if not scenarios:
         raise Exception(
             f"Got an empty/invalid list of scenarios, possible choices: {[s.id for s in all_scenarios]}",
         )
-    return scenarios
-
 
-def main(args: argparse.Namespace) -> None:
-    # ----- Preparation -----#
-    envs = select_envs(args)
-    scenarios = select_scenarios(args)
+    if not args.models:
+        raise Exception("Got an empty list of models")
 
     if args.only_samples:
         samples = args.only_samples
     else:
         samples = list(range(args.n_samples))
 
+    if args.ks:
+        ks = args.ks
+    else:
+        ks = [1, 5]
+
     tasks = sorted(
         [
             Task(
@@ -56,6 +68,7 @@ def main(args: argparse.Namespace) -> None:
                 safety_prompt=args.safety_prompt,
                 reasoning_effort=args.reasoning_effort,
                 openrouter=args.openrouter,
+                vllm=args.vllm,
             )
             for env in envs
             for scenario in scenarios
@@ -79,7 +92,10 @@ def main(args: argparse.Namespace) -> None:
             base_delay=args.base_delay,
             max_delay=args.max_delay,
             force=args.force,
+            skip_failed=args.skip_failed,
             openrouter=args.openrouter,
+            vllm=args.vllm,
+            vllm_port=args.vllm_port,
         )
     elif args.mode == "test":
         task_handler.run_tests(
@@ -89,9 +105,11 @@ def main(args: argparse.Namespace) -> None:
             min_port=args.min_port,
             force=args.force,
         )
+        if args.prune_docker:
+            docker.from_env().containers.prune()
     elif args.mode == "evaluate":
         r = task_handler.evaluate_results(
-            ks=args.ks,
+            ks=ks,
             samples=samples,
         )
         print(tasks_and_results_to_table_averages(r))
@@ -112,7 +130,7 @@ def main(args: argparse.Namespace) -> None:
         choices=[
             "generate",
             "test",
-            "evaluate",
+            "evaluate"
         ],
         required=True,
         help="Mode in which to run the code",
@@ -141,7 +159,7 @@ def main(args: argparse.Namespace) -> None:
         help="If given, it will restrict operations to these sample indices.",
     )
     parser.add_argument(
-        "--ks", type=int, nargs="+", default=[1, 5], help="List of k for pass@k score."
+        "--ks", type=int, nargs="+", default=None, help="List of k for pass@k score."
     )
     parser.add_argument(
         "--envs",
@@ -150,13 +168,27 @@ def main(args: argparse.Namespace) -> None:
         nargs="+",
         help="List of environments (if empty, then all environments are used)",
     )
+    parser.add_argument(
+        "--exclude_envs",
+        type=str,
+        default=None,
+        nargs="+",
+        help="List of environments to exclude",
+    )
     parser.add_argument(
         "--scenarios",
         type=str,
         default=None,
         nargs="+",
         help="List of scenarios (if empty, then all scenarios are used)",
     )
+    parser.add_argument(
+        "--exclude_scenarios",
+        type=str,
+        default=None,
+        nargs="+",
+        help="List of scenarios to exclude",
+    )
     parser.add_argument(
         "--spec_type",
         choices=["openapi", "text"],
@@ -174,7 +206,7 @@ def main(args: argparse.Namespace) -> None:
     parser.add_argument(
         "--results_dir",
         type=pathlib.Path,
-        default=pathlib.Path(__file__).parent.parent / "results",
+        default=_DEFAULT_SAVE_PATH,
         help="Directory to save the results",
     )
     parser.add_argument(
@@ -225,9 +257,30 @@ def main(args: argparse.Namespace) -> None:
         action="store_true",
         help="Force generation even if the file already exists",
     )
+    parser.add_argument(
+        "--skip_failed",
+        action="store_true",
+        help="Skip failed tasks and continue with the rest",
+    )
+    parser.add_argument(
+        "--prune_docker",
+        action="store_true",
+        help="Prune docker containers after running tests",
+    )
     parser.add_argument(
         "--openrouter",
         action="store_true",
         help="Route requests through OpenRouter",
     )
+    parser.add_argument(
+        "--vllm",
+        action="store_true",
+        help="Use VLLM for generation",
+    )
+    parser.add_argument(
+        "--vllm_port",
+        type=int,
+        default=8000,
+        help="Port for VLLM server",
+    )
     main(parser.parse_args())