Switch to UV (#319)

Author: alexreinking

master/master.cfg

@@ -92,7 +92,6 @@ class Version:
 class VersionedBranch:
     ref: str
     version: Version
-    pip_version: str = ""  # halide-llvm version on pypi.halide-lang.org
 
     @property
     def shortversion(self):
@@ -107,26 +106,10 @@ LLVM_RELEASE_21 = "release_21"
 LLVM_RELEASE_20 = "release_20"
 
 LLVM_BRANCHES = {
-    LLVM_MAIN: VersionedBranch(
-        ref="main",
-        version=Version(23, 0, 0),
-        pip_version="23.0.0.dev0+g80f627e6",
-    ),
-    LLVM_RELEASE_22: VersionedBranch(
-        ref="llvmorg-22.1.0-rc3",
-        version=Version(22, 1, 0),
-        pip_version="22.1.0rc3",
-    ),
-    LLVM_RELEASE_21: VersionedBranch(
-        ref="llvmorg-21.1.8",
-        version=Version(21, 1, 8),
-        pip_version="21.1.8",
-    ),
-    LLVM_RELEASE_20: VersionedBranch(
-        ref="llvmorg-20.1.8",
-        version=Version(20, 1, 8),
-        pip_version="20.1.8",
-    ),
+    LLVM_MAIN: VersionedBranch(ref="main", version=Version(23, 0, 0)),
+    LLVM_RELEASE_22: VersionedBranch(ref="llvmorg-22.1.0-rc3", version=Version(22, 1, 0)),
+    LLVM_RELEASE_21: VersionedBranch(ref="llvmorg-21.1.8", version=Version(21, 1, 8)),
+    LLVM_RELEASE_20: VersionedBranch(ref="llvmorg-20.1.8", version=Version(20, 1, 8)),
 }
 
 # At any given time, Halide has a main branch, which supports (at least)
@@ -465,6 +448,18 @@ def save_environment(prop, allowlist):
     return extract_fn
 
 
+def extract_llvm_pip_version(llvm_ver):
+    def extract_fn(_rc, stdout, _stderr):
+        for line in stdout.splitlines():
+            if f"ci-llvm-{llvm_ver}" in line:
+                m = re.search(r"halide-llvm==([\w.+]+)", line)
+                if m:
+                    return {"HALIDE_LLVM_PIP_VERSION": m.group(1)}
+        return {}
+
+    return extract_fn
+
+
 def get_msvc_config_steps(factory, builder_type):
     # ensure that we use the x64 host compiler, not the x86 host compiler
     arch_for_bits = {32: "x64_x86", 64: "x64"}
@@ -586,28 +581,24 @@ def get_build_command(build_dir, targets=None):
 
 
 def add_create_venv_step(factory, builder_type):
+    llvm_ver = LLVM_BRANCHES[builder_type.llvm_branch].shortversion
+    platform_tag = _PIP_PLATFORM_TAGS[(builder_type.arch, builder_type.bits, builder_type.os)]
+
+    group = f"ci-llvm-{llvm_ver}" if platform_tag is None else "ci-base"
+
     if builder_type.os == "windows":
-        activate = r".venv\Scripts\activate"
-        env = "set"
-        remove = "rmdir /s /q"
+        activate = r'"%UV_PROJECT_ENVIRONMENT%\Scripts\activate"'
+        env_cmd = "set"
     else:
-        activate = ". .venv/bin/activate"
-        env = "env"
-        remove = "rm -rf"
+        activate = '. "$UV_PROJECT_ENVIRONMENT/bin/activate"'
+        env_cmd = "env"
 
     factory.addStep(
         SetPropertyFromCommand(
             name="Create virtual environment",
             locks=[performance_lock.access("counting")],
             workdir=get_source_path(),
-            command=(
-                f"{remove} .venv && "
-                "python -m venv .venv && "
-                f"{activate} && "
-                'python -m pip install -U pip "setuptools[core]" wheel && '
-                "python -m pip install -r requirements.txt && "
-                f"{env}"
-            ),
+            command=f"uv sync --group {group} --no-install-project && {activate} && {env_cmd}",
             env=Property("env", {}),
             extract_fn=save_environment("venv", ["VIRTUAL_ENV", "PATH"]),
         )
@@ -638,70 +629,72 @@ _PIP_PLATFORM_TAGS = {
 
 
 def add_install_llvm_step(factory, builder_type):
-    pip_version = LLVM_BRANCHES[builder_type.llvm_branch].pip_version
-    llvm_version = LLVM_BRANCHES[builder_type.llvm_branch].shortversion
-    pip_spec = f"halide-llvm=={pip_version}"
-    extra_index = "https://pypi.halide-lang.org/simple/"
-
+    llvm_ver = LLVM_BRANCHES[builder_type.llvm_branch].shortversion
     platform_tag = _PIP_PLATFORM_TAGS[(builder_type.arch, builder_type.bits, builder_type.os)]
 
     if platform_tag is None:
-        # Native install: pip install into the venv, then query the prefix
-        if builder_type.os == "windows":
-            activate = r".venv\Scripts\activate"
-        else:
-            activate = ". .venv/bin/activate"
+        # Native install: halide-llvm was already installed by uv sync
+        factory.addStep(
+            SetPropertyFromCommand(
+                name=f"Get LLVM {llvm_ver} prefix",
+                locks=[performance_lock.access("counting")],
+                workdir=get_source_path(),
+                haltOnFailure=True,
+                env=Property("env"),
+                command="halide-llvm --prefix",
+                property="HALIDE_LLVM_ROOT",
+            )
+        )
+    else:
+        # Cross install: use pip to install halide-llvm for a foreign platform
+        extra_index = "https://pypi.halide-lang.org/simple/"
+        target_dir = get_builddir_subpath("halide-llvm-cross")
 
         factory.addStep(
             ShellCommand(
-                name=f"Install halide-llvm {llvm_version}",
+                name="Install pip",
                 locks=[performance_lock.access("counting")],
                 workdir=get_source_path(),
                 haltOnFailure=True,
                 env=Property("env"),
-                command=f'{activate} && python -m pip install --pre "{pip_spec}" --extra-index-url {extra_index}',
+                command="uv pip install pip",
             )
         )
+
+        cat = "type" if builder_type.os == "windows" else "cat"
         factory.addStep(
             SetPropertyFromCommand(
-                name=f"Get LLVM {llvm_version} prefix",
+                name=f"Get halide-llvm {llvm_ver} version",
                 locks=[performance_lock.access("counting")],
                 workdir=get_source_path(),
                 haltOnFailure=True,
-                env=Property("env"),
-                command=f"{activate} && halide-llvm --prefix",
-                property="HALIDE_LLVM_ROOT",
+                command=f"{cat} pyproject.toml",
+                extract_fn=extract_llvm_pip_version(llvm_ver),
             )
         )
-    else:
-        # Cross install: use --target and --platform to fetch the right wheel
-        target_dir = get_builddir_subpath("halide-llvm-cross")
 
         factory.addStep(
             RemoveDirectory(
-                name=f"Remove old halide-llvm {llvm_version} cross dir",
+                name=f"Remove old halide-llvm {llvm_ver} cross dir",
                 locks=[performance_lock.access("counting")],
                 dir=target_dir,
                 haltOnFailure=False,
             )
         )
 
-        if builder_type.os == "windows":
-            activate = r".venv\Scripts\activate"
-        else:
-            activate = ". .venv/bin/activate"
-
         factory.addStep(
             ShellCommand(
-                name=f"Install halide-llvm {llvm_version} ({platform_tag})",
+                name=f"Install halide-llvm {llvm_ver} ({platform_tag})",
                 locks=[performance_lock.access("counting")],
                 workdir=get_source_path(),
                 haltOnFailure=True,
                 env=Property("env"),
-                command=Interpolate(f"{activate} && python -m pip install --pre --target ")
+                command=Interpolate("python -m pip install --pre --target ")
                 + target_dir
                 + Interpolate(
-                    f' --platform {platform_tag} --only-binary=:all: "{pip_spec}" --extra-index-url {extra_index}'
+                    f" --platform {platform_tag} --only-binary=:all: "
+                    '"halide-llvm==%(prop:HALIDE_LLVM_PIP_VERSION)s" '
+                    f"--extra-index-url {extra_index}"
                 ),
             )
         )
@@ -710,7 +703,7 @@ def add_install_llvm_step(factory, builder_type):
         llvm_root = get_builddir_subpath("halide-llvm-cross/halide_llvm/data")
         factory.addStep(
             SetProperties(
-                name=f"Set LLVM {llvm_version} prefix (cross)",
+                name=f"Set LLVM {llvm_ver} prefix (cross)",
                 properties={"HALIDE_LLVM_ROOT": llvm_root},
             )
         )