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