Feature/auto enable precompiled wheels GitHub #24706

setup.py

@@ -32,6 +32,219 @@ def load_module_from_path(module_name, path):
 ROOT_DIR = Path(__file__).parent
 logger = logging.getLogger(__name__)
 
+
+class GitUtils:
+    """Utility class for Git operations."""
+
+    @staticmethod
+    def run_git_command(args: list, cwd: str = None) -> tuple[bool, str]:
+        """Run a git command and return success status and output."""
+        try:
+            result = subprocess.run(['git'] + args,
+                                    capture_output=True,
+                                    text=True,
+                                    cwd=cwd or ROOT_DIR)
+            return result.returncode == 0, result.stdout.strip()
+        except Exception:
+            return False, ""
+
+    @staticmethod
+    def get_origin_url() -> str:
+        """Get the origin URL of the current git repository."""
+        success, output = GitUtils.run_git_command(
+            ['remote', 'get-url', 'origin'])
+        return output if success else ""
+
+    @staticmethod
+    def get_config_url() -> str:
+        """Get the origin URL from git config."""
+        success, output = GitUtils.run_git_command(
+            ['config', '--get', 'remote.origin.url'])
+        return output if success else ""
+
+    @staticmethod
+    def get_log_refs() -> str:
+        """Get git log references."""
+        success, output = GitUtils.run_git_command(
+            ['log', '--format=%D', '-1'])
+        return output if success else ""
+
+
+class DebugLogger:
+    """Base class for components that need debug logging."""
+
+    DEFAULT_LOG_PATH = os.path.join(
+        __import__("tempfile").gettempdir(), "vllm_debug.log")
+
+    def __init__(self, debug_log_path: str = None):
+        self.debug_log_path = debug_log_path or self.DEFAULT_LOG_PATH
+
+    def _log_debug(self, message: str) -> None:
+        """Write debug message to log file."""
+        with open(self.debug_log_path, "a") as f:
+            f.write(f"DEBUG: {message}\n")
+
+    @classmethod
+    def write_initial_log(cls) -> None:
+        """Write initial debug information to log file."""
+        with open(cls.DEFAULT_LOG_PATH, "w") as f:
+            f.write("DEBUG: SETUP.PY IS RUNNING - "
+                    "Starting GitHub detection...\n")
+
+    @classmethod
+    def log_github_detection_result(cls, is_github: bool) -> None:
+        """Log the result of GitHub detection."""
+        with open(cls.DEFAULT_LOG_PATH, "a") as f:
+            if is_github:
+                f.write(f"DEBUG: VLLM_USE_PRECOMPILED set to: "
+                        f"{os.environ.get('VLLM_USE_PRECOMPILED')}\n")
+            else:
+                f.write("DEBUG: No GitHub installation detected\n")
+
+
+class GitHubDetector(DebugLogger):
+    """Detects if vLLM is being installed from a GitHub repository."""
+
+    # Original vLLM repository patterns
+    ORIGINAL_PATTERNS = [
+        'github.com/vllm-project/vllm', '[email protected]:vllm-project/vllm.git',
+        'git+https://github.com/vllm-project/vllm.git',
+        'git+ssh://[email protected]/vllm-project/vllm.git'
+    ]
+
+    # Temporary directory patterns
+    TEMP_DIR_PATTERNS = ['/tmp/', '\\temp\\', 'cache', 'builds', '.cache']
+
+    # UV cache directory pattern
+    UV_CACHE_PATTERN = os.path.join('.cache', 'uv', 'git-v0', 'checkouts')
+
+    def _is_vllm_github_repo(self, origin_url: str) -> bool:
+        """Check given URL is a vLLM GitHub repository (original or fork)."""
+        if not origin_url:
+            return False
+
+        # Check for original repository
+        if any(pattern in origin_url for pattern in self.ORIGINAL_PATTERNS):
+            return True
+
+        # Check for any fork: github.com/*/vllm.git
+        return ('github.com/' in origin_url and '/vllm' in origin_url
+                and origin_url.endswith('.git'))
+
+    def _get_git_origin_url(self) -> str:
+        """Get the origin URL of the current git repository."""
+        origin_url = GitUtils.get_origin_url()
+
+        # Handle UV cache directory case
+        uv_cache_db_path = os.path.join(os.path.expanduser('~'), '.cache',
+                                        'uv', 'git-v0', 'db') + os.sep
+        if origin_url.startswith(uv_cache_db_path):
+            return self._resolve_uv_cache_url()
+
+        return origin_url
+
+    def _resolve_uv_cache_url(self) -> str:
+        """Resolve original URL from UV cache directory."""
+        # Try to get the original URL from git config
+        config_url = GitUtils.get_config_url()
+        uv_cache_db_path = os.path.join(os.path.expanduser('~'), '.cache',
+                                        'uv', 'git-v0', 'db') + os.sep
+        if config_url and not config_url.startswith(uv_cache_db_path):
+            return config_url
+
+        # Try to get URL from git log
+        log_output = GitUtils.get_log_refs()
+        if 'github.com' in log_output and '/vllm' in log_output:
+            import re
+            url_match = re.search(r'https://github\.com/[^/]+/vllm',
+                                  log_output)
+            if url_match:
+                return url_match.group(0) + '.git'
+
+        return ""
+
+    def _check_command_line_args(self) -> bool:
+        """Check command line arguments for git+ URLs."""
+        if len(sys.argv) > 1:
+            for arg in sys.argv[1:]:
+                if arg.startswith('git+') and self._is_vllm_github_repo(arg):
+                    return True
+        return False
+
+    def _check_uv_cache_directory(self) -> bool:
+        """Check if we're in a UV cache directory with vLLM repo."""
+        current_dir = os.getcwd()
+        expected_cache_path = os.path.join(os.path.expanduser("~"),
+                                           self.UV_CACHE_PATTERN)
+        if not os.path.abspath(current_dir).startswith(expected_cache_path):
+            return False
+
+        has_setup_py = os.path.exists(os.path.join(ROOT_DIR, 'setup.py'))
+        has_vllm_dir = os.path.exists(os.path.join(ROOT_DIR, 'vllm'))
+        is_vllm_project = has_setup_py and has_vllm_dir
+
+        if is_vllm_project:
+            return True
+
+        # Fallback to git remote check
+        origin_url = self._get_git_origin_url()
+        return self._is_vllm_github_repo(origin_url)
+
+        return False
+
+    def _check_temp_directory(self) -> bool:
+        """Check if we're in a temporary directory with vLLM repo."""
+        current_dir = os.getcwd()
+        is_temp_dir = any(pattern in current_dir.lower()
+                          for pattern in self.TEMP_DIR_PATTERNS)
+
+        if not is_temp_dir:
+            return False
+
+        origin_url = self._get_git_origin_url()
+        return self._is_vllm_github_repo(origin_url)
+
+    def _check_local_repository(self) -> bool:
+        """Check if we're in the actual vLLM repository (for local dev)."""
+        origin_url = self._get_git_origin_url()
+        return self._is_vllm_github_repo(origin_url)
+
+    def is_installing_from_github(self) -> bool:
+        """Check if vLLM is being installed from GitHub repository."""
+        # Method 1: Check command line arguments
+        if self._check_command_line_args():
+            return True
+
+        # Method 2: Check UV cache directory
+        if self._check_uv_cache_directory():
+            return True
+
+        # Method 3: Check temporary directory
+        if self._check_temp_directory():
+            return True
+
+        # Method 4: Check local repository
+        return self._check_local_repository()
+
+
+def is_installing_from_github() -> bool:
+    """Check if vLLM is being installed from GitHub repository."""
+    detector = GitHubDetector()
+    return detector.is_installing_from_github()
+
+
+# Auto-enable precompiled wheels when installing from GitHub
+# Write debug info to a file to ensure we can see it
+DebugLogger.write_initial_log()
+
+if is_installing_from_github():
+    print("Detected installation from GitHub repository. "
+          "Automatically enabling precompiled wheels.")
+    os.environ["VLLM_USE_PRECOMPILED"] = "1"
+    DebugLogger.log_github_detection_result(True)
+else:
+    DebugLogger.log_github_detection_result(False)
+
 # cannot import envs directly because it depends on vllm,
 #  which is not installed yet
 envs = load_module_from_path('envs', os.path.join(ROOT_DIR, 'vllm', 'envs.py'))
@@ -604,35 +817,90 @@ def _read_requirements(filename: str) -> list[str]:
     ]
 }
 
-# If using precompiled, extract and patch package_data (in advance of setup)
-if envs.VLLM_USE_PRECOMPILED:
-    assert _is_cuda(), "VLLM_USE_PRECOMPILED is only supported for CUDA builds"
-    wheel_location = os.getenv("VLLM_PRECOMPILED_WHEEL_LOCATION", None)
-    if wheel_location is not None:
-        wheel_url = wheel_location
-    else:
+
+class PrecompiledWheelManager(DebugLogger):
+    """Manages precompiled wheel extraction and installation."""
+
+    WHEEL_BASE_URL = "https://wheels.vllm.ai"
+    WHEEL_VERSION = "1.0.0.dev-cp38-abi3"
+
+    # Architecture to wheel tag mapping
+    ARCH_WHEEL_TAGS = {
+        "x86_64": "manylinux1_x86_64",
+        "aarch64": "manylinux2014_aarch64"
+    }
+
+    def _get_wheel_tag(self) -> str:
+        """Get the appropriate wheel tag for the current architecture."""
         import platform
         arch = platform.machine()
-        if arch == "x86_64":
-            wheel_tag = "manylinux1_x86_64"
-        elif arch == "aarch64":
-            wheel_tag = "manylinux2014_aarch64"
-        else:
+
+        if arch not in self.ARCH_WHEEL_TAGS:
             raise ValueError(f"Unsupported architecture: {arch}")
+
+        return self.ARCH_WHEEL_TAGS[arch]
+
+    def _get_wheel_urls(self) -> tuple[str, str]:
+        """Get base and nightly wheel URLs."""
+        wheel_tag = self._get_wheel_tag()
         base_commit = precompiled_wheel_utils.get_base_commit_in_main_branch()
-        wheel_url = f"https://wheels.vllm.ai/{base_commit}/vllm-1.0.0.dev-cp38-abi3-{wheel_tag}.whl"
-        nightly_wheel_url = f"https://wheels.vllm.ai/nightly/vllm-1.0.0.dev-cp38-abi3-{wheel_tag}.whl"
+
+        base_url = (f"{self.WHEEL_BASE_URL}/{base_commit}/"
+                    f"vllm-{self.WHEEL_VERSION}-{wheel_tag}.whl")
+        nightly_url = (f"{self.WHEEL_BASE_URL}/nightly/"
+                       f"vllm-{self.WHEEL_VERSION}-{wheel_tag}.whl")
+
+        # Log key information for debugging
+        self._log_debug(f"Base commit: {base_commit}")
+
+        return base_url, nightly_url
+
+    def _select_wheel_url(self) -> str:
+        """Select the appropriate wheel URL (custom, base, or nightly)."""
+        # Check for custom wheel location
+        wheel_location = os.getenv("VLLM_PRECOMPILED_WHEEL_LOCATION")
+        if wheel_location:
+            return wheel_location
+
+        # Get base and nightly URLs
+        base_url, nightly_url = self._get_wheel_urls()
+
+        # Try base URL first
         from urllib.request import urlopen
         try:
-            with urlopen(wheel_url) as resp:
-                if resp.status != 200:
-                    wheel_url = nightly_wheel_url
+            with urlopen(base_url) as resp:
+                if resp.status == 200:
+                    return base_url
         except Exception as e:
-            print(f"[warn] Falling back to nightly wheel: {e}")
-            wheel_url = nightly_wheel_url
+            self._log_debug(
+                f"Failed to access base wheel URL '{base_url}': {e}. "
+                "Falling back to nightly.")
+
+        # Fall back to nightly
+        return nightly_url
+
+    def extract_and_patch(self) -> dict:
+        """Extract precompiled wheel and return package data patch."""
+        wheel_url = self._select_wheel_url()
+
+        try:
+            patch = precompiled_wheel_utils.\
+               extract_precompiled_and_patch_package(wheel_url)
+            self._log_debug("Wheel extraction successful")
+            return patch
+        except Exception as e:
+            print(f"[error] Failed to extract precompiled wheel: {e}")
+            return {}
+
+
+# If using precompiled, extract and patch package_data (in advance of setup)
+if envs.VLLM_USE_PRECOMPILED:
+    assert _is_cuda(), "VLLM_USE_PRECOMPILED is only supported for CUDA builds"
+
+    wheel_manager = PrecompiledWheelManager()
+    patch = wheel_manager.extract_and_patch()
 
-    patch = precompiled_wheel_utils.extract_precompiled_and_patch_package(
-        wheel_url)
+    # Apply patch to package_data
     for pkg, files in patch.items():
         package_data.setdefault(pkg, []).extend(files)