Large scale benchmarking works.

Author: atharvas

scratch/merged_context_registry_2025-09-04T23:54:53.035665.json

@@ -10,14 +10,16 @@
 import pickle
 import shutil
 from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
 
 import asv
 import pandas as pd
 
 from datasmith.benchmark.collection import BenchmarkCollection
-from datasmith.docker.context import ContextRegistry, DockerContext, Task
+from datasmith.docker.context import ContextRegistry, DockerContext, Task, build_base_image
 from datasmith.docker.orchestrator import (
+    build_repo_sha_image,
     get_docker_client,
     orchestrate,
 )
@@ -131,12 +133,16 @@ def process_inputs(args: argparse.Namespace) -> dict[tuple[str, str], set[tuple[
     return all_states
 
 
-def main(args: argparse.Namespace) -> None:
-    client = get_docker_client()
+def main(args: argparse.Namespace) -> None:  # noqa: C901
+    client = get_docker_client(args.max_concurrency)
     all_states = process_inputs(args)
     context_registry = ContextRegistry.load_from_file(path=args.context_registry)
     interim_path = Path(os.environ["CACHE_LOCATION"]).parent / "interim"  # Look here for cached docker contexts
 
+    logger.info("Building base image...")
+    base_tag = build_base_image(client, DockerContext())
+    os.environ["DOCKER_CACHE_FROM"] = base_tag
+
     # Prepare tasks
     tasks: list[tuple[Task, DockerContext]] = []
     repo_commit_pairs = defaultdict(list)
@@ -199,6 +205,18 @@ def main(args: argparse.Namespace) -> None:
     already_benchmarked = list(filter(lambda x: (interim_path / f"{x[0].get_container_name()}.json").exists(), tasks))
     logger.info("Skipping %d tasks that have already been benchmarked", len(already_benchmarked))
 
+    # build the containers.
+    builds = []
+    with ThreadPoolExecutor(max_workers=args.max_concurrency) as pool:
+        futures = [
+            pool.submit(build_repo_sha_image, client, ctx, task, args.force_rebuild, run_id="CANARY-BUILD")
+            for (task, ctx) in tasks
+        ]
+        for fut in as_completed(futures):
+            builds.append(fut.result())
+
+    to_benchmark = [t for (t, b) in zip(tasks, builds) if b.rc == 0]
+
     machine_args: dict[str, str] = asv.machine.Machine.get_defaults()  # pyright: ignore[reportAttributeAccessIssue]
     machine_args["num_cpu"] = str(args.num_cores)
     files_by_image: dict[Task, dict[str, str]] = asyncio.run(
@@ -227,6 +245,12 @@ def main(args: argparse.Namespace) -> None:
 
     logger.info("Benchmark results saved to %s", output_file)
 
+    # remove all images with CANARY-BUILD tag.
+    try:
+        client.images.prune(filters={"label": "datasmith.run=CANARY-BUILD"})
+    except Exception:
+        logger.exception("Failed to prune images with CANARY-BUILD tag")
+
 
 if __name__ == "__main__":
     args = parse_args()

scratch/scripts/benchmark_commits.py

@@ -3,6 +3,7 @@
 import argparse
 import datetime
 import json
+import os
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
 
@@ -12,7 +13,7 @@
 from datasmith.agents.config import configure_agent_backends
 from datasmith.agents.context_synthesis import agent_build_and_validate
 from datasmith.benchmark.collection import BenchmarkCollection
-from datasmith.docker.context import ContextRegistry
+from datasmith.docker.context import ContextRegistry, DockerContext, build_base_image
 from datasmith.docker.orchestrator import get_docker_client
 from datasmith.docker.validation import Task, _err_lock
 from datasmith.logging_config import configure_logging
@@ -131,6 +132,10 @@ def main(args: argparse.Namespace) -> None:
         else ContextRegistry()
     )
 
+    logger.info("Building base image...")
+    base_tag = build_base_image(client, DockerContext())
+    os.environ["DOCKER_CACHE_FROM"] = base_tag
+
     # Prepare tasks
     tasks = prepare_tasks(all_states, args.limit_per_repo, context_registry)
 

scratch/scripts/synthesize_contexts.py

@@ -75,7 +75,7 @@ def fast_cleanup_run_artifacts(  # noqa: C901
         for iid in img_ids:
             try:
                 logger.debug("Removing image id=%s", iid[:20])
-                client.images.remove(iid, force=True, noprune=False)
+                client.images.remove(iid, force=True, noprune=True)
             except (ImageNotFound, NotFound):
                 pass
             except APIError as e:
@@ -400,7 +400,7 @@ def build_once_with_context(
     run_labels: dict[str, str],
     probe: bool = False,
     pull: bool = False,
-    force: bool = True,
+    force: bool = False,
 ) -> BuildResult:
     logger.info("build_once_with_context: registering context key=%s", task.get_image_name())
     logger.debug(
@@ -487,7 +487,7 @@ def agent_build_and_validate(  # noqa: C901
             tail_chars=args.tail_chars,
             probe=True,
             pull=True,
-            force=True,  # If the env is already present, don't rebuild (saves time)
+            force=False,  # If the env is already present, don't rebuild (saves time)
             run_labels=run_labels,
         )
         if not env_res.ok:
@@ -582,7 +582,7 @@ def agent_build_and_validate(  # noqa: C901
                 sha=task.sha,
                 timeout_s=args.build_timeout,
                 tail_chars=args.tail_chars * 2,
-                force=True,  # Always rebuild package image to pick up new script
+                force=False,  # Always rebuild package image to pick up new script
                 run_labels=run_labels,
             )
             attempts.append(AttemptRecord(attempt_idx=i, building_data=script, build_result=build_res))

src/datasmith/agents/context_synthesis.py

@@ -1,6 +1,7 @@
 # syntax=docker/dockerfile:1.7
 
-FROM buildpack-deps:jammy AS base
+ARG BASE_IMAGE=buildpack-deps:jammy
+FROM ${BASE_IMAGE} AS base
 
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \

src/datasmith/docker/Dockerfile

@@ -24,6 +24,25 @@
 logger = get_logger("docker.context")
 
 
+def build_base_image(client: docker.DockerClient, ctx: DockerContext) -> str:
+    base_key = hash(ctx)
+    base_tag = f"asv-base-rev-{base_key}:base"
+
+    res = ctx.build_container_streaming(
+        client=client,
+        image_name=base_tag,
+        build_args={},
+        probe=True,
+        force=False,
+        pull=False,
+        timeout_s=1800,
+    )
+    if not res.ok:
+        logger.exception("Failed to build base image %s error=%s", base_tag, res.stderr_tail)
+        raise RuntimeError("Failed to build base image")
+    return base_tag
+
+
 @dataclass
 class BuildResult:
     ok: bool
@@ -141,43 +160,32 @@ def __init__(
         self.base_building_data = base_building_data
         self.building_data = building_data
 
+    @staticmethod
+    def add_bytes(tar: tarfile.TarFile, name: str, data: bytes, mode: int = 0o644) -> None:
+        info = tarfile.TarInfo(name=name)
+        info.size = len(data)
+        info.mode = mode
+        info.mtime = 0
+        info.uid = info.gid = 0
+        info.uname = info.gname = ""
+        tar.addfile(info, io.BytesIO(data))
+
     def build_tarball_stream(self, probe: bool = False) -> io.BytesIO:
         tar_stream = io.BytesIO()
         with tarfile.open(fileobj=tar_stream, mode="w") as tar:
             # Add Dockerfile
-            dockerfile_bytes = self.dockerfile_data.encode("utf-8")
-            dockerfile_info = tarfile.TarInfo(name="Dockerfile")
-            dockerfile_info.size = len(dockerfile_bytes)
-            tar.addfile(dockerfile_info, io.BytesIO(dockerfile_bytes))
-
+            DockerContext.add_bytes(tar, "Dockerfile", self.dockerfile_data.encode("utf-8"))
             # Add entrypoint.sh
-            entrypoint_data = self.entrypoint_data.encode("utf-8")
-            entrypoint_info = tarfile.TarInfo(name="entrypoint.sh")
-            entrypoint_info.size = len(entrypoint_data)
-            entrypoint_info.mode = 0o755  # Make it executable
-            tar.addfile(entrypoint_info, io.BytesIO(entrypoint_data))
-
+            DockerContext.add_bytes(tar, "entrypoint.sh", self.entrypoint_data.encode("utf-8"), mode=0o755)
             # Add docker_build_env.sh
-            env_building_data = self.env_building_data.encode("utf-8")
-            env_building_info = tarfile.TarInfo(name="docker_build_env.sh")
-            env_building_info.size = len(env_building_data)
-            env_building_info.mode = 0o755  # Make it executable
-            tar.addfile(env_building_info, io.BytesIO(env_building_data))
+            DockerContext.add_bytes(tar, "docker_build_env.sh", self.env_building_data.encode("utf-8"), mode=0o755)
 
             # Add docker_build_base.sh
-            base_building_data = self.base_building_data.encode("utf-8")
-            base_building_info = tarfile.TarInfo(name="docker_build_base.sh")
-            base_building_info.size = len(base_building_data)
-            base_building_info.mode = 0o755  # Make it executable
-            tar.addfile(base_building_info, io.BytesIO(base_building_data))
+            DockerContext.add_bytes(tar, "docker_build_base.sh", self.base_building_data.encode("utf-8"), mode=0o755)
 
             if not probe:
                 # Add docker_build_pkg.sh
-                building_data = self.building_data.encode("utf-8")
-                building_info = tarfile.TarInfo(name="docker_build_pkg.sh")
-                building_info.size = len(building_data)
-                building_info.mode = 0o755  # Make it executable
-                tar.addfile(building_info, io.BytesIO(building_data))
+                DockerContext.add_bytes(tar, "docker_build_pkg.sh", self.building_data.encode("utf-8"), mode=0o755)
 
         # Reset the stream position to the beginning
         tar_stream.seek(0)
@@ -216,6 +224,11 @@ def build_container(
             pass  # Image doesn't exist or was removed, proceed to build
 
         if not image_exists:
+            cache_from = None
+            if base_image := os.environ.get("DOCKER_CACHE_FROM", None):
+                build_args = {**build_args, "BASE_IMAGE": base_image}
+                cache_from = [base_image]
+
             if len(build_args):
                 build_args_str = " --build-arg ".join(f"{k}={v}" for k, v in build_args.items())
                 logger.info("$ docker build -t %s src/datasmith/docker/ --build-arg %s", image_name, build_args_str)
@@ -229,6 +242,7 @@ def build_container(
                         rm=True,
                         labels=run_labels,
                         network_mode=os.environ.get("DOCKER_NETWORK_MODE", None),
+                        cache_from=cache_from,
                     )
                 except DockerException:
                     logger.exception("Failed to build Docker image '%s'", image_name)
@@ -289,6 +303,11 @@ def build_container_streaming(  # noqa: C901
             stdout_buf: deque[str] = deque(maxlen=2000)  # chunk-tail buffers
             stderr_buf: deque[str] = deque(maxlen=2000)
 
+            cache_from = None
+            if base_image := os.environ.get("DOCKER_CACHE_FROM", None):
+                build_args = {**build_args, "BASE_IMAGE": base_image}
+                cache_from = [base_image]
+
             # Pretty log line for transparency
             if build_args:
                 build_args_str = " --build-arg ".join(f"{k}={v}" for k, v in build_args.items())
@@ -308,6 +327,7 @@ def build_container_streaming(  # noqa: C901
                     target=target,
                     labels=run_labels,
                     network_mode=os.environ.get("DOCKER_NETWORK_MODE", None),
+                    cache_from=cache_from,
                 )
             except DockerException:
                 logger.exception("Failed to initiate build for '%s'", image_name)
@@ -424,6 +444,15 @@ def from_dict(cls, data: Mapping[str, Any]) -> DockerContext:
             base_building_data=data.get("base_building_data", None),
         )
 
+    def __hash__(self) -> int:
+        return hash((
+            self.dockerfile_data,
+            self.entrypoint_data,
+            self.building_data,
+            self.env_building_data,
+            self.base_building_data,
+        ))
+
 
 class ContextRegistry:
     """Registry for Docker contexts keyed by owner/repo[/sha], independent of tag.

src/datasmith/docker/context.py

@@ -16,7 +16,7 @@
 from docker.errors import APIError, DockerException, ImageNotFound, NotFound
 from docker.models.containers import Container
 
-from datasmith.docker.context import BuildResult, ContextRegistry, DockerContext, Task
+from datasmith.docker.context import BuildResult, DockerContext, Task
 from datasmith.logging_config import get_logger
 
 logger = get_logger("docker.orchestrator")
@@ -118,10 +118,10 @@ async def _guard_loop(
         await asyncio.wait_for(stop_event.wait(), timeout=interval_s)
 
 
-def get_docker_client() -> docker.DockerClient:
+def get_docker_client(max_concurrency: int = 10) -> docker.DockerClient:
     """Return an authenticated Docker client or exit with an error."""
     try:
-        return docker.from_env(timeout=60)
+        return docker.from_env(timeout=60, max_pool_size=max_concurrency)
     except DockerException as exc:
         sys.exit(f"Could not connect to Docker daemon: {exc}")
 
@@ -151,21 +151,20 @@ def build_repo_image(client: docker.DockerClient, image_name: str, repo_url: str
 
 
 def build_repo_sha_image(
-    client: docker.DockerClient, context_registry: ContextRegistry, task: Task, force: bool = False
+    client: docker.DockerClient, docker_ctx: DockerContext, task: Task, force: bool = False, run_id: str | None = None
 ) -> BuildResult:
     assert task.sha is not None, "Task.sha must be set"  # noqa: S101
-    image_name = f"asv/{task.owner}/{task.repo}/{task.sha}".lower()
-    docker_ctx = context_registry[image_name]
+    repo_url = f"https://www.github.com/{task.owner}/{task.repo}"
     build_res: BuildResult = docker_ctx.build_container_streaming(
         client=client,
-        image_name=image_name,
-        build_args={
-            "REPO_URL": f"https://www.github.com/{task.owner}/{task.repo}",
-            "COMMIT_SHA": task.sha,
-        },
-        force=force,
+        image_name=task.get_image_name(),
+        build_args={"REPO_URL": repo_url, "COMMIT_SHA": task.sha},
+        probe=False,
+        force=False,
+        timeout_s=1800,  # 30 minutes
         tail_chars=10_000,
         pull=False,
+        run_labels=gen_run_labels(task, runid="unknown" if run_id is None else run_id),
     )
     return build_res
 
@@ -214,7 +213,7 @@ def _launch() -> tuple[int, dict[str, str]]:  # noqa: C901
             image_name=task.get_image_name(),
             build_args={"REPO_URL": repo_url, "COMMIT_SHA": task.sha},
             probe=False,
-            force=True,
+            force=False,
             timeout_s=1800,  # 30 minutes
             tail_chars=10_000,
             pull=False,
@@ -299,7 +298,7 @@ def _launch() -> tuple[int, dict[str, str]]:  # noqa: C901
             logger.exception("Failed to remove container %s", task.get_container_name())
             pass
         try:
-            client.images.remove(image=task.get_image_name(), force=True, noprune=False)
+            client.images.remove(image=task.get_image_name(), force=True, noprune=True)
         except Exception:
             logger.exception("Failed to remove image %s", task.get_image_name())
             pass

src/datasmith/docker/orchestrator.py

@@ -193,7 +193,7 @@ def validate_one(  # noqa: C901
             "REPO_URL": f"https://www.github.com/{task.owner}/{task.repo}",
             "COMMIT_SHA": task.sha,
         },
-        force=True,  # preserve your original behavior
+        force=False,
         timeout_s=args.build_timeout,
         tail_chars=args.tail_chars,
         pull=False,