ae5a1a192e
- make_test_spec() for all 500 samples was the real bottleneck that made swe_bench_verified appear stuck at 'Processing records: 0%' even when all instance images were already loaded locally. - Now image names are computed directly from instance IDs, matching the swebench naming convention. Local images are listed once and missing ones are reported immediately. make_test_spec is only called for images that actually need to be built. - With 500 pre-loaded images, build_images() now completes in ~45s instead of hanging indefinitely. Also add bash/diagnose_swe_images.py for quickly checking which images are missing locally.
211 lines
9.2 KiB
Python
211 lines
9.2 KiB
Python
import traceback
|
|
from typing import TYPE_CHECKING, Dict, List, Literal
|
|
|
|
from evalscope.api.dataset import Dataset
|
|
from evalscope.utils.function_utils import run_in_threads_with_progress
|
|
from evalscope.utils.logger import get_logger
|
|
|
|
if TYPE_CHECKING:
|
|
from swebench.harness.test_spec.test_spec import TestSpec
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
def build_images(
|
|
samples: Dataset,
|
|
max_workers: int = 4,
|
|
force_rebuild: bool = False,
|
|
use_remote_images: bool = True,
|
|
force_arch: Literal['', 'arm64', 'x86_64'] = '',
|
|
dockerhub_username: str = 'swebench',
|
|
) -> Dict[str, str]:
|
|
"""This function uses the swe_bench library to build the docker images for the SWE-bench dataset.
|
|
|
|
It can also try to pull images from a registry before building them locally.
|
|
|
|
Args:
|
|
samples (Dataset): The dataset to build the images for
|
|
max_workers (int): The maximum number of workers to use for building images. Defaults to 4.
|
|
force_rebuild (bool, optional): Whether to force a rebuild of the images. Defaults to False.
|
|
use_remote_images (bool, optional): Whether to try pulling images from Docker Hub before building locally. Defaults to True. See https://hub.docker.com/u/swebench
|
|
force_arch (str, optional): Optionally force the docker images to be pulled/built for a specific architecture. Defaults to "".
|
|
dockerhub_username (str, optional): DockerHub user/org namespace for remote images. Defaults to "swebench".
|
|
""" # noqa: E501
|
|
from docker.client import DockerClient # type: ignore
|
|
from swebench.harness.constants import LATEST, SWEbenchInstance # type: ignore
|
|
from swebench.harness.docker_build import build_instance_images # type: ignore
|
|
from swebench.harness.test_spec.test_spec import make_test_spec # type: ignore
|
|
|
|
from .utils import resolve_swebench_arch
|
|
|
|
extra_build_instance_images_kwargs = {'tag': LATEST, 'env_image_tag': LATEST}
|
|
|
|
# Code copied from the swe_bench repository
|
|
docker_client = DockerClient.from_env()
|
|
|
|
def _get_available_docker_images() -> List[str]:
|
|
return [tag for image in docker_client.images.list() for tag in image.tags]
|
|
|
|
# The swebench library requires a huggingface version of the code to be loaded in order to build the images.
|
|
# We load the dataset and then use the library to build the images.
|
|
samples_hf: List[SWEbenchInstance] = [s.metadata for s in samples]
|
|
|
|
# We also keep a mapping from instance_ids to the name of the docker image
|
|
id_to_docker_image: Dict[str, str] = {}
|
|
id_to_test_spec: Dict[str, TestSpec] = {}
|
|
|
|
# Note that remote images are named eg "sphinx-doc_1776_sphinx-11502"
|
|
namespace = (dockerhub_username or 'swebench') if use_remote_images else None
|
|
|
|
def _get_instance_image_name(instance_id: str) -> str:
|
|
"""Fast image name computation that matches make_test_spec output."""
|
|
arch = resolve_swebench_arch(instance_id, force_arch)
|
|
if namespace:
|
|
# Remote / pre-built image format: swebench/sweb.eval.x86_64.repo_1776_name-123:latest
|
|
updated_id = instance_id.replace('__', '_1776_')
|
|
return f'{namespace}/sweb.eval.{arch}.{updated_id}:latest'
|
|
# Local build format: sweb.eval.x86_64.repo__name-123:latest
|
|
return f'sweb.eval.{arch}.{instance_id}:latest'
|
|
|
|
# Fast path: compute image names directly without calling make_test_spec.
|
|
# make_test_spec is very slow for 500 samples and is only needed when we
|
|
# actually have to build missing images locally.
|
|
logger.info('Computing SWE-Bench image names from instance IDs...')
|
|
for swebench_instance in samples_hf:
|
|
instance_id = swebench_instance['instance_id']
|
|
id_to_docker_image[instance_id] = _get_instance_image_name(instance_id)
|
|
|
|
# Get list of locally available Docker images
|
|
logger.info('Listing local Docker images...')
|
|
available_docker_images = _get_available_docker_images()
|
|
samples_to_build_images_for = [
|
|
s for s in samples_hf if id_to_docker_image[s['instance_id']] not in available_docker_images
|
|
]
|
|
logger.info(
|
|
f'SWE-Bench images: {len(samples_hf)} needed, '
|
|
f'{len(samples_hf) - len(samples_to_build_images_for)} already local, '
|
|
f'{len(samples_to_build_images_for)} missing'
|
|
)
|
|
|
|
# Only build expensive test_specs for images that are actually missing.
|
|
if samples_to_build_images_for:
|
|
logger.info('Building test specs for missing images (this may take a while)...')
|
|
for swebench_instance in samples_to_build_images_for:
|
|
arch = resolve_swebench_arch(swebench_instance['instance_id'], force_arch)
|
|
test_spec = make_test_spec(swebench_instance, namespace=namespace, arch=arch)
|
|
id_to_test_spec[swebench_instance['instance_id']] = test_spec
|
|
|
|
# Try to pull images from Docker Hub first if requested
|
|
if use_remote_images and len(samples_to_build_images_for) > 0:
|
|
logger.info(f'Attempting to pull {len(samples_to_build_images_for)} SWE-BENCH images from Docker Hub')
|
|
successfully_pulled: List[str] = []
|
|
|
|
def _pull_image(sample) -> str:
|
|
"""Pull a single SWE-Bench image; returns instance_id on success."""
|
|
instance_id: str = sample['instance_id']
|
|
image_name: str = id_to_docker_image[instance_id]
|
|
image_base_name: str = image_name.split(':')[0]
|
|
logger.info(f'Pulling {image_name}...')
|
|
for line in docker_client.api.pull(image_name, stream=True, decode=True):
|
|
status = line.get('status')
|
|
progress = line.get('progress')
|
|
if progress:
|
|
logger.info(f'{image_name} {status} {progress}')
|
|
elif status:
|
|
logger.info(f'{image_name} {status}')
|
|
docker_client.api.tag(image_name, image_base_name, 'latest')
|
|
logger.info(f'Successfully pulled {image_name}')
|
|
return instance_id
|
|
|
|
def _on_error(sample, exc: Exception) -> None:
|
|
image_name = id_to_docker_image[sample['instance_id']]
|
|
logger.warning(f'Failed to pull {image_name}: {exc}')
|
|
|
|
def _on_result(sample, instance_id: str) -> None:
|
|
successfully_pulled.append(instance_id)
|
|
|
|
run_in_threads_with_progress(
|
|
items=samples_to_build_images_for,
|
|
worker=_pull_image,
|
|
desc='Pulling SWE-Bench images',
|
|
max_workers=max_workers,
|
|
log_interval=30,
|
|
on_result=_on_result,
|
|
on_error=_on_error,
|
|
skip_failed=True,
|
|
)
|
|
logger.info(f'Pulled {len(successfully_pulled)} images from Docker Hub')
|
|
|
|
# Remove successfully pulled images from the build list
|
|
samples_to_build_images_for = [
|
|
s for s in samples_to_build_images_for if s['instance_id'] not in successfully_pulled
|
|
]
|
|
|
|
# Update available images list
|
|
available_docker_images = _get_available_docker_images()
|
|
|
|
# Build any remaining images locally
|
|
if len(samples_to_build_images_for) > 0:
|
|
logger.warning('BUILDING SWE-BENCH IMAGES. NOTE: This can take a long time.')
|
|
test_specs_to_build = [id_to_test_spec[s['instance_id']] for s in samples_to_build_images_for]
|
|
build_instance_images(
|
|
client=docker_client,
|
|
dataset=test_specs_to_build,
|
|
force_rebuild=force_rebuild,
|
|
max_workers=max_workers,
|
|
**extra_build_instance_images_kwargs,
|
|
)
|
|
|
|
# Check that all the images were built
|
|
available_docker_images = _get_available_docker_images()
|
|
missing_images = [
|
|
id_to_docker_image[s['instance_id']]
|
|
for s in samples_hf
|
|
if id_to_docker_image[s['instance_id']] not in available_docker_images
|
|
]
|
|
assert len(missing_images) == 0, (f'Not all images were built: {missing_images}, {id_to_docker_image}')
|
|
|
|
return id_to_docker_image
|
|
|
|
|
|
def build_container(
|
|
test_spec: 'TestSpec',
|
|
client,
|
|
):
|
|
"""
|
|
Builds the instance image for the given test spec and creates a container from the image.
|
|
|
|
Args:
|
|
test_spec (TestSpec): Test spec to build the instance image and container for
|
|
client (docker.DockerClient): Docker client for building image + creating the container
|
|
"""
|
|
from swebench.harness.constants import DOCKER_USER
|
|
from swebench.harness.docker_utils import cleanup_container
|
|
|
|
# Build corresponding instance image
|
|
container = None
|
|
try:
|
|
# Create the container
|
|
logger.info(f'Creating container for {test_spec.instance_id}...')
|
|
|
|
# Define arguments for running the container
|
|
run_args = test_spec.docker_specs.get('run_args', {})
|
|
cap_add = run_args.get('cap_add', [])
|
|
|
|
container = client.containers.create(
|
|
image=test_spec.instance_image_key,
|
|
user=DOCKER_USER,
|
|
detach=True,
|
|
command='tail -f /dev/null',
|
|
platform=test_spec.platform,
|
|
cap_add=cap_add,
|
|
)
|
|
logger.info(f'Container for {test_spec.instance_id} created: {container.id}')
|
|
return container
|
|
except Exception as e:
|
|
# If an error occurs, clean up the container and raise an exception
|
|
logger.error(f'Error creating container for {test_spec.instance_id}: {e}')
|
|
logger.info(traceback.format_exc())
|
|
cleanup_container(client, container, logger)
|
|
raise e
|