fix: don't run uv cache clean in parallel with uv pip install

Refactor the code so `uv cache` is cleaned up after each batch of build
wheels. Before there was a possibility that `uv cache clean` and `uv pip
install` were running at the same time.

Fixes: #756
Signed-off-by: Christian Heimes <[email protected]>

Author: tiran

src/fromager/bootstrapper.py

@@ -422,6 +422,8 @@ def _build_wheel(
             version=resolved_version,
             build_env=build_env,
         )
+        # invalidate uv cache
+        self.ctx.uv_clean_cache(req)
         server.update_wheel_mirror(self.ctx)
         # When we update the mirror, the built file moves to the
         # downloads directory.

src/fromager/build_environment.py

@@ -240,17 +240,6 @@ def get_distributions(self) -> typing.Mapping[str, Version]:
         mapping = json.loads(result.strip())
         return {name: Version(version) for name, version in sorted(mapping.items())}
 
-    def clean_cache(self, req: Requirement) -> None:
-        """Invalidate and clean uv cache for requirement
-
-        uv caches package metadata and unpacked wheels for faster dependency
-        resolution and installation. ``uv pip install`` hardlinks files from
-        cache location. This function removes a package from all caches, so
-        subsequent installations use a new built.
-        """
-        logger.debug("invalidate uv cache for %s", req.name)
-        self.run(["uv", "cache", "clean", req.name])
-
 
 @metrics.timeit(description="prepare build environment")
 def prepare_build_environment(

src/fromager/commands/build.py

@@ -125,6 +125,8 @@ def build(
             force=True,
             cache_wheel_server_url=None,
         )
+        # invalidate uv cache
+        wkctx.uv_clean_cache(req)
     print(wheel_filename)
 
 
@@ -435,6 +437,8 @@ def _build(
             version=resolved_version,
             build_env=build_env,
         )
+        # uv cache is cleaned in build / build_parallel commands
+        # wkctx.uv_clean_cache(req)
 
         hooks.run_post_build_hooks(
             ctx=wkctx,
@@ -687,11 +691,13 @@ def update_progressbar_cb(future: concurrent.futures.Future) -> None:
                 max_workers=max_workers
             ) as executor:
                 futures: list[concurrent.futures.Future[tuple[pathlib.Path, bool]]] = []
+                reqs: list[Requirement] = []
                 logger.info(
                     "starting to build: %s", sorted(n.key for n in buildable_nodes)
                 )
                 for node in buildable_nodes:
                     req = Requirement(f"{node.canonicalized_name}=={node.version}")
+                    reqs.append(req)
                     future = executor.submit(
                         _build_parallel,
                         wkctx=wkctx,
@@ -724,6 +730,9 @@ def update_progressbar_cb(future: concurrent.futures.Future) -> None:
                         logger.error(f"Failed to build {node.key}: {e}")
                         raise
 
+                # invalidate uv cache
+                wkctx.uv_clean_cache(*reqs)
+
     metrics.summarize(wkctx, "Building in parallel")
     _summary(wkctx, entries)
 

src/fromager/commands/step.py

@@ -224,5 +224,6 @@ def build_wheel(
         version=dist_version,
         build_env=build_env,
     )
+    wkctx.uv_clean_cache(req)
     requirement_ctxvar.reset(token)
     print(wheel_filename)

src/fromager/context.py

@@ -13,7 +13,13 @@
 from packaging.utils import NormalizedName, canonicalize_name
 from packaging.version import Version
 
-from . import constraints, dependency_graph, packagesettings, request_session
+from . import (
+    constraints,
+    dependency_graph,
+    external_commands,
+    packagesettings,
+    request_session,
+)
 
 if typing.TYPE_CHECKING:
     from . import build_environment
@@ -126,6 +132,26 @@ def pip_constraint_args(self) -> list[str]:
         path_to_constraints_file = path_to_constraints_file.absolute()
         return ["--constraint", os.fspath(path_to_constraints_file)]
 
+    def uv_clean_cache(self, *reqs: Requirement) -> None:
+        """Invalidate and clean uv cache for requirements
+
+        uv caches package metadata and unpacked wheels for faster dependency
+        resolution and installation. ``uv pip install`` hardlinks files from
+        cache location. This function removes a package from all caches, so
+        subsequent installations use a new built.
+
+        WARNING: 'uv clean cache' is not concurrency safe with 'uv pip install'.
+        """
+        if not reqs:
+            raise ValueError("no requirements")
+
+        extra_environ: dict[str, str] = {"UV_CACHE_DIR": str(self.uv_cache)}
+        cmd = ["uv", "clean", "cache"]
+        req_list: list[str] = sorted(set(req.name for req in reqs))
+        logger.debug("invalidate uv cache for %s", req_list)
+        cmd.extend(req_list)
+        external_commands.run(cmd, extra_environ=extra_environ)
+
     def write_to_graph_to_file(self):
         with self.graph_file.open("w", encoding="utf-8") as f:
             self.dependency_graph.serialize(f)

src/fromager/wheels.py

@@ -336,9 +336,6 @@ def build_wheel(
         build_dir=pbi.build_dir(sdist_root_dir),
     )
 
-    # invalidate uv's cache
-    build_env.clean_cache(req)
-
     wheels = list(ctx.wheels_build.glob("*.whl"))
     if len(wheels) != 1:
         raise FileNotFoundError(