Add wheel URL generation and caching for pre-built wheels in setup.py #156
Conversation
There was a problem hiding this comment.
Pull Request Overview
This PR introduces functionality to efficiently handle pre-built wheels by adding URL generation and caching capabilities to the setup process. Instead of always building from source, the system can now attempt to download compatible pre-built wheels based on environment parameters before falling back to compilation.
- Adds URL generation for pre-built wheels based on CUDA version, Python version, and platform
- Implements a custom
CachedWheelsCommandthat attempts to download compatible wheels before building - Updates the setup configuration to use the new caching command
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
| try: | ||
| urllib.request.urlretrieve(wheel_url, wheel_filename) |
There was a problem hiding this comment.
Using urllib.request.urlretrieve without SSL verification or timeout could pose security risks. Consider adding timeout parameter and validating the downloaded file integrity with checksums.
| wheel_path = os.path.join(self.dist_dir, archive_basename + ".whl") | ||
| print("Raw wheel path", wheel_path) | ||
| os.rename(wheel_filename, wheel_path) | ||
| except (urllib.error.HTTPError, urllib.error.URLError): |
There was a problem hiding this comment.
The error message is too generic and doesn't provide useful debugging information. Consider logging the specific HTTP status code or error details to help users understand why the wheel download failed.
| except (urllib.error.HTTPError, urllib.error.URLError): | |
| except (urllib.error.HTTPError, urllib.error.URLError) as e: | |
| print(f"Failed to download precompiled wheel from {wheel_url}.") | |
| print(f"Error type: {type(e).__name__}") | |
| if hasattr(e, 'code'): | |
| print(f"HTTP status code: {getattr(e, 'code', 'N/A')}") | |
| if hasattr(e, 'reason'): | |
| print(f"Reason: {getattr(e, 'reason', 'N/A')}") |
| return super().run() | ||
|
|
||
| wheel_url, wheel_filename = get_wheel_url() | ||
| print("Guessing wheel URL: ", wheel_url) |
There was a problem hiding this comment.
Using print statements for logging in production code is not ideal. Consider using the logging module for better control over log levels and output formatting.
| archive_basename = f"{self.wheel_dist_name}-{impl_tag}-{abi_tag}-{plat_tag}" | ||
|
|
||
| wheel_path = os.path.join(self.dist_dir, archive_basename + ".whl") | ||
| print("Raw wheel path", wheel_path) |
There was a problem hiding this comment.
Using print statements for logging in production code is not ideal. Consider using the logging module for better control over log levels and output formatting.
Introduce functionality to generate URLs for pre-built wheels and implement caching to avoid unnecessary builds. This enhancement allows for efficient retrieval of compatible wheels based on environment parameters.