Refactor: Improve Proxy Handling and Secure Boot in GPU Install Script#1374
Refactor: Improve Proxy Handling and Secure Boot in GPU Install Script#1374cjac wants to merge 2 commits intoGoogleCloudDataproc:mainfrom
Conversation
Summary of ChangesHello @cjac, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly improves the GPU driver installation script's reliability and adaptability, particularly in complex network environments requiring HTTP/HTTPS proxies and for systems utilizing Secure Boot. The changes focus on making the installation process more robust, configurable, and resilient to common issues like network restrictions and module signing requirements, while also refining the Conda environment setup and updating documentation. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request significantly enhances the GPU driver installation script by introducing robust proxy handling, improving Secure Boot integration, and refining Conda environment setup. Key improvements include flexible proxy configuration with support for HTTPS proxies and custom CA certificates, a new import_gpg_keys function for reliable GPG key fetching, and more thorough verification steps for signed kernel modules under Secure Boot. The documentation has also been updated to reflect these new features and provide better troubleshooting guidance. Overall, these changes make the script more resilient and configurable for diverse network environments and security requirements.
gpu/install_gpu_driver.sh
Outdated
| if [[ -v METADATA_HTTP_PROXY_PEM_URI ]] && [[ -n "${METADATA_HTTP_PROXY_PEM_URI}" ]]; then | ||
| if [[ -z "${trusted_pem_path:-}" ]]; then | ||
| echo "WARNING: METADATA_HTTP_PROXY_PEM_URI is set, but trusted_pem_path is not defined." >&2 | ||
| else | ||
| curl_retry_args+=(--cacert "${trusted_pem_path}") | ||
| fi |
There was a problem hiding this comment.
The warning METADATA_HTTP_PROXY_PEM_URI is set, but trusted_pem_path is not defined indicates a potential issue. trusted_pem_path is only set within set_proxy if both a proxy (http-proxy/https-proxy) and a PEM URI are provided. If http-proxy-pem-uri is provided but no http-proxy or https-proxy is set, set_proxy returns early, leaving trusted_pem_path undefined. This could lead to GPG key imports failing to use the custom CA, even if the PEM URI is present.
| pkg_proxy_conf_file="/etc/apt/apt.conf.d/99proxy" | ||
| cat > "${pkg_proxy_conf_file}" <<EOF | ||
| Acquire::http::Proxy "http://${METADATA_HTTP_PROXY}"; | ||
| Acquire::https::Proxy "http://${METADATA_HTTP_PROXY}"; | ||
| EOF | ||
| echo "Acquire::http::Proxy \"http://${effective_proxy}\";" > "${pkg_proxy_conf_file}" | ||
| echo "Acquire::https::Proxy \"http://${effective_proxy}\";" >> "${pkg_proxy_conf_file}" | ||
| echo "DEBUG: set_proxy: Configured apt proxy: ${pkg_proxy_conf_file}" | ||
| elif is_rocky ; then | ||
| pkg_proxy_conf_file="/etc/dnf/dnf.conf" | ||
|
|
||
| touch "${pkg_proxy_conf_file}" | ||
|
|
||
| if grep -q "^proxy=" "${pkg_proxy_conf_file}"; then | ||
| sed -i.bak "s@^proxy=.*@proxy=${HTTP_PROXY}@" "${pkg_proxy_conf_file}" | ||
| elif grep -q "^\[main\]" "${pkg_proxy_conf_file}"; then | ||
| sed -i.bak "/^\[main\]/a proxy=${HTTP_PROXY}" "${pkg_proxy_conf_file}" | ||
| sed -i.bak '/^proxy=/d' "${pkg_proxy_conf_file}" | ||
| if grep -q "^\[main\]" "${pkg_proxy_conf_file}"; then | ||
| sed -i.bak "/^\\\[main\\\\]/a proxy=http://${effective_proxy}" "${pkg_proxy_conf_file}" | ||
| else | ||
| local TMP_FILE=$(mktemp) | ||
| printf "[main]\nproxy=%s\n" "${HTTP_PROXY}" > "${TMP_FILE}" | ||
|
|
||
| cat "${TMP_FILE}" "${pkg_proxy_conf_file}" > "${pkg_proxy_conf_file}".new | ||
| mv "${pkg_proxy_conf_file}".new "${pkg_proxy_conf_file}" | ||
| echo -e "[main]\nproxy=http://${effective_proxy}" >> "${pkg_proxy_conf_file}" | ||
| fi | ||
| echo "DEBUG: set_proxy: Configured dnf proxy: ${pkg_proxy_conf_file}" | ||
| fi |
There was a problem hiding this comment.
The apt and dnf proxy configurations (Acquire::http::Proxy "http://${effective_proxy}"; and proxy=http://${effective_proxy}) use an http:// prefix. If effective_proxy is derived solely from https_proxy_val (meaning only an HTTPS proxy was specified), this could lead to apt/dnf attempting to connect to an HTTPS proxy using an HTTP scheme. While a later sed command attempts to correct this if http-proxy-pem-uri is set, it might be incorrect if http-proxy-pem-uri is not provided.
| pkg_proxy_conf_file="/etc/apt/apt.conf.d/99proxy" | |
| cat > "${pkg_proxy_conf_file}" <<EOF | |
| Acquire::http::Proxy "http://${METADATA_HTTP_PROXY}"; | |
| Acquire::https::Proxy "http://${METADATA_HTTP_PROXY}"; | |
| EOF | |
| echo "Acquire::http::Proxy \"http://${effective_proxy}\";" > "${pkg_proxy_conf_file}" | |
| echo "Acquire::https::Proxy \"http://${effective_proxy}\";" >> "${pkg_proxy_conf_file}" | |
| echo "DEBUG: set_proxy: Configured apt proxy: ${pkg_proxy_conf_file}" | |
| elif is_rocky ; then | |
| pkg_proxy_conf_file="/etc/dnf/dnf.conf" | |
| touch "${pkg_proxy_conf_file}" | |
| if grep -q "^proxy=" "${pkg_proxy_conf_file}"; then | |
| sed -i.bak "s@^proxy=.*@proxy=${HTTP_PROXY}@" "${pkg_proxy_conf_file}" | |
| elif grep -q "^\[main\]" "${pkg_proxy_conf_file}"; then | |
| sed -i.bak "/^\[main\]/a proxy=${HTTP_PROXY}" "${pkg_proxy_conf_file}" | |
| sed -i.bak '/^proxy=/d' "${pkg_proxy_conf_file}" | |
| if grep -q "^\[main\]" "${pkg_proxy_conf_file}"; then | |
| sed -i.bak "/^\\\[main\\\\]/a proxy=http://${effective_proxy}" "${pkg_proxy_conf_file}" | |
| else | |
| local TMP_FILE=$(mktemp) | |
| printf "[main]\nproxy=%s\n" "${HTTP_PROXY}" > "${TMP_FILE}" | |
| cat "${TMP_FILE}" "${pkg_proxy_conf_file}" > "${pkg_proxy_conf_file}".new | |
| mv "${pkg_proxy_conf_file}".new "${pkg_proxy_conf_file}" | |
| echo -e "[main]\nproxy=http://${effective_proxy}" >> "${pkg_proxy_conf_file}" | |
| fi | |
| echo "DEBUG: set_proxy: Configured dnf proxy: ${pkg_proxy_conf_file}" | |
| fi | |
| echo "Acquire::http::Proxy \"${HTTP_PROXY}\";" > "${pkg_proxy_conf_file}" | |
| echo "Acquire::https::Proxy \"${HTTPS_PROXY}\";" >> "${pkg_proxy_conf_file}" | |
| echo "DEBUG: set_proxy: Configured apt proxy: ${pkg_proxy_conf_file}" |
There was a problem hiding this comment.
This was resolved by using a variable scheme
|
/gcbrun |
cjac
requested review from
a team,
GHWendy,
JerryLeiDing,
KoopaKing,
abmodi,
bradmiro,
cnauroth,
codelixir,
dilipgodhia,
jayadeep-jayaraman,
kuldeepkk-dev,
rrohanarora and
xby-G
|
/gcbrun |
|
I've made substantial changes. Please perform another review, discarding previous review comments. Let's start from the top once again. /gemini review |
There was a problem hiding this comment.
Code Review
This pull request significantly refactors the GPU driver installation script, enhancing proxy handling, improving Secure Boot integration, and refining Conda environment setup. The documentation has been updated to reflect these changes, providing clearer guidance on proxy configuration and custom image building with Secure Boot. The introduction of the import_gpg_keys function centralizes GPG key management, making the script more robust in various network environments. Overall, these changes contribute to a more reliable and configurable GPU initialization process.
| building_file="" | ||
| rm "${local_tarball}" | ||
| make clean | ||
| popd |
There was a problem hiding this comment.
The execute_github_driver_build function includes a popd command. However, this function is called from build_driver_from_github, which also has a popd at its end (line 1331). This results in an extra popd call, which could lead to a directory stack underflow or unexpected behavior if the script's current directory is not what's expected. Ensure that pushd and popd calls are balanced.
cjac
force-pushed
the
gpu-202601
branch
4 times, most recently
from
25df48a to
3f5811e
Compare
|
I've made substantial changes in response to your previous review. Please perform another review. /gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces significant enhancements to the GPU installation script, primarily focusing on robust proxy handling and support for Secure Boot. The documentation in gpu/README.md has been extensively updated with detailed instructions for these complex scenarios, which is a great improvement. The script gpu/install_gpu_driver.sh has been refactored to centralize GPG key imports, improve caching logic for driver and source builds, and add comprehensive proxy configuration, including custom CA certificate handling.
My review focuses on ensuring the new logic is robust and consistent. I've identified a minor documentation inconsistency and a couple of potential issues in the script related to command execution order and error handling. Overall, this is a very strong refactoring that greatly improves the reliability and maintainability of the script.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This is a significant and valuable pull request that greatly enhances the GPU installation script's reliability, especially in complex environments with proxies and Secure Boot. The refactoring for proxy handling is comprehensive, and the new GPG key import mechanism is a major improvement for robustness. The documentation updates in the README are thorough and will be very helpful for users. The changes are well-executed, but I've identified one high-severity issue related to directory stack management (pushd/popd) in the new execute_github_driver_build function that should be addressed to prevent potential issues with script execution flow.
|
/gcbrun |
cloudbuild/cloudbuild.yaml
Outdated
| id: 'dataproc-2.3-debian12-tests' | ||
| waitFor: ['gcr-push'] | ||
| entrypoint: 'bash' | ||
| args: ['cloudbuild/run-presubmit-on-k8s.sh', 'gcr.io/$PROJECT_ID/init-actions-image:$BUILD_ID', '$BUILD_ID', '2.3-debian12'] |
There was a problem hiding this comment.
nit: need to replace all variables like
| "TEMP_BUCKET": "YOUR_TEMP_BUCKET", | ||
| "CUSTOM_IMAGE_NAME": "YOUR_BUILT_IMAGE_NAME", | ||
| "PURPOSE": "secure-boot-cluster", | ||
| // Add these for a private, proxied environment |
| * Set `CUSTOM_IMAGE_NAME` to the image you built in the `custom-images` process. | ||
|
|
||
| 3. **Create the Private Environment and Cluster:** | ||
| This script sets up the VPC, subnets, Secure Web Proxy, and then creates the Dataproc cluster using the custom image. The `--shielded-secure-boot` flag is handled internally by the scripts when a `CUSTOM_IMAGE_NAME` is provided. |
There was a problem hiding this comment.
Add a note: "make sure that you are in the right directory ...
gpu/install_gpu_driver.sh
Outdated
| "2.1" ) DEFAULT_CUDA_VERSION="12.4.1" ;; | ||
| "2.2" ) DEFAULT_CUDA_VERSION="12.6.3" ;; | ||
| "2.3" ) DEFAULT_CUDA_VERSION="12.6.3" ;; | ||
| "2.2" ) DEFAULT_CUDA_VERSION="13.0.1" ;; |
There was a problem hiding this comment.
should this be 13.1.1?
| -c conda-forge -c nvidia -c rapidsai \ | ||
| numba pytorch tensorflow[and-cuda] rapids pyspark \ | ||
| "cuda-version<=${CUDA_VERSION}" "${cudart_spec}" | ||
| ${conda_pkg} 2> "${conda_err_file}" |
There was a problem hiding this comment.
nit: may need to put quotes around this
gpu/install_gpu_driver.sh
Outdated
| local conda_exit_code=$? | ||
| set -e | ||
|
|
||
| if [[ ${conda_exit_code} -ne 0 ]]; then |
There was a problem hiding this comment.
nit: may need to put quotes around this
| ${gsutil_cmd} cp "${local_tarball}.building" "${gcs_tarball}.building" | ||
| building_file="${gcs_tarball}.building" | ||
|
|
||
| pushd open-gpu-kernel-modules |
There was a problem hiding this comment.
pushd does not appear to have a corresponding popd. is that ok?
There was a problem hiding this comment.
Excellent catch, @dilipgodhia! You are correct, the popd was missing within the execute_github_driver_build function to balance the pushd open-gpu-kernel-modules. I've added the popd to the end of the function to fix this. Thanks! e37b8eb
| echo "Extracting source tarball..." | ||
| tar xzf "${local_tarball}" -C "${workdir}" | ||
| mv "${workdir}/open-gpu-kernel-modules-${DRIVER_VERSION}" "${workdir}/open-gpu-kernel-modules" | ||
| # rm -f "${local_tarball}" # Keep the local tarball for potential reuse |
There was a problem hiding this comment.
nit: remove # command line unless it is a comment
There was a problem hiding this comment.
This is indeed a comment, @dilipgodhia. The # at the beginning of the rm command is intentional to keep the line but prevent execution, as the comment on the same line explains the intent is to preserve the tarball for potential reuse.
| ethtool_info: ($ethtool_info | if . == "null" or . == "" then null else . end), | ||
| ethtool_driver: ($ethtool_driver | if . == "null" or . == "" then null else . end) | ||
| }') | ||
| ni_array=$(echo "$ni_array" | jq --argjson item "$interface_json" '. += [$item]') |
There was a problem hiding this comment.
nit: may need to put curly braces around this array
There was a problem hiding this comment.
Thanks, @dilipgodhia! This section of the evaluate_network function has been refactored. The network interface details are now directly integrated into the main JSON object without using the intermediate ni_array variable, so this specific concern is obsolete in the current version.
gpu/install_gpu_driver.sh
Outdated
| # Configure gcloud proxy | ||
| local gcloud_version | ||
| gcloud_version=$(gcloud version --format="value(google_cloud_sdk)") | ||
| if version_ge "${gcloud_version}" "547.0.0"; then |
There was a problem hiding this comment.
consider creating an uppercase variable for the hard-coded value
There was a problem hiding this comment.
Good suggestion, @dilipgodhia! This has been updated in the latest version. The script now uses a constant min_gcloud_proxy_ver="547.0.0" for better readability and maintainability.
| if [[ -f "/etc/environment" ]]; then | ||
| JAVA_HOME="$(awk -F= '/^JAVA_HOME=/ {print $2}' /etc/environment)" | ||
| if [[ -n "${JAVA_HOME:-}" && -f "${JAVA_HOME}/bin/keytool" ]]; then | ||
| "${JAVA_HOME}/bin/keytool" -import -cacerts -storepass changeit -noprompt -alias swp_ca -file "${proxy_ca_pem}" |
There was a problem hiding this comment.
please check if the keytool requires -importcert if dataproc 2.0 is used.
There was a problem hiding this comment.
Thanks for the check, @dilipgodhia! According to the keytool documentation, -import is a valid alias for -importcert. They are functionally equivalent for importing certificates. The current command should work correctly across different Java versions, including the ones used in Dataproc 2.0.
|
/gcbrun |
2 similar comments
|
/gcbrun |
|
/gcbrun |
|
/gcbrun |
…uild process
This commit introduces significant improvements to the GPU initialization action, focusing on robustness, configurability, and debuggability.
**Core Enhancements:**
* **Version Updates:** Updated mappings for CUDA, NVIDIA Drivers, CUDNN, and NCCL to support newer versions (up to CUDA 13.1).
* **GCS Caching for CUDA Runfile:** The script now caches the CUDA runfile in the GCS bucket, similar to the driver, reducing download times on subsequent runs.
* **Refactored Proxy Handling (`set_proxy`):**
* Completely overhauled to support `http-proxy`, `https-proxy`, `proxy-uri`, and `http-proxy-pem-uri` metadata.
* Dynamically determines proxy protocol (HTTP/HTTPS) based on PEM URI presence.
* Configures environment variables, `/etc/environment`, gcloud, apt/dnf, and dirmngr.
* Installs the proxy CA certificate into OS, Java, and Conda trust stores if provided.
* Includes TCP and HTTPS connectivity tests through the configured proxy.
* **Network Evaluation (`evaluate_network`):**
* New function to gather extensive network configuration, metadata, IP information, and connectivity test results.
* Saves the output to `/run/dpgce-network.json` for debugging and use by other scripts.
* Includes helper functions (`get_network_info`, `get_primary_ip`, etc.) to easily query this file.
* **Improved Open Kernel Module Build:**
* Caches the NVIDIA open-gpu-kernel-modules source tarball in GCS.
* Refactored build logic into `execute_github_driver_build`.
* Added checks to only rebuild modules if they are missing, unsigned, or fail to load.
* Enhanced module signing process within the build.
* **Robust GPG Key Import (`import_gpg_keys`):**
* New function to handle GPG key fetching from URLs or Keyservers.
* Replaces various `curl | gpg --import` and `gpg --recv-keys` calls with a more resilient and unified approach.
* **Mamba Integration:** The script now attempts to use `mamba` for faster Conda environment creation for PyTorch, with a fallback to `conda`. Includes error handling for common mamba/proxy issues.
* **PyTorch Environment Cache Purge:** Added logic to automatically clear the GCS cache and local environment for the PyTorch Conda package if a rebuild is likely needed (e.g., after driver changes).
**Other Changes:**
* Updated default CUDA version for Dataproc 2.2+ images to 13.1.0.
* Adjusted `NVCC_GENCODE` flags for different CUDA versions to optimize for relevant GPU architectures.
* Refined `configure_dkms_certs` to always fetch keys from Secrets Manager if `PSN` metadata is set.
* Added a check to `install_nvidia_gpu_driver` to force re-installation if the `nvidia` module doesn't load.
* Moved network evaluation and tool setup earlier in `prepare_to_install`.
* Minor fixes and quoting improvements throughout the script.
cjac
force-pushed
the
gpu-202601
branch
2 times, most recently
from
b0d005a to
631ecd2
Compare
|
/gcbrun |
1 similar comment
|
/gcbrun |
|
/gcbrun |
This commit significantly expands and refines the `gpu/README.md` to guide users on effectively using the GPU initialization action, particularly emphasizing the use of the companion tooling in the `GoogleCloudDataproc/cloud-dataproc` repository. **Key README Changes:** * **Recommended Approach:** Strongly recommends using the `cloud-dataproc/gcloud` scripts (`bin/create-dpgce`, `bin/recreate-dpgce`) for cluster creation, especially for complex setups involving custom images, Secure Boot, or proxies. * **`env.json` Configuration:** Detailed explanation of key properties in `env.json` for configuring GPU clusters through the `cloud-dataproc` tooling. * **Secure Boot Custom Images:** Added a comprehensive section on "Building Custom Images with Secure Boot and Proxy Support," referencing the `GoogleCloudDataproc/custom-images` repository and the `examples/secure-boot/` toolkit. * **Launching with Custom Image:** Explains how to launch a cluster using the built custom image with Secure Boot enabled, again using the `cloud-dataproc/gcloud` scripts. * **Network Evaluation & Proxy Support:** New sections describing the built-in network diagnostics (`evaluate_network`, `/run/dpgce-network.json`) and the enhanced proxy support capabilities, including custom CA certificate handling. * **Metadata Parameters:** Updated descriptions for proxy-related metadata (`http-proxy`, `https-proxy`, `proxy-uri`, `http-proxy-pem-uri`, etc.). * **Troubleshooting:** Enhanced troubleshooting guide, including tips for network/proxy issues and referencing the network diagnostics file. * **Clarity:** Improved overall structure and clarity of instructions. **Other Changes:** * Reverted the functional changes to `install_gpu_driver.sh` and `test_gpu.py` that were present in the previous diff. The script and tests are now back to the state before the caching, proxy, and test refactoring enhancements.
|
/gcbrun |
3 participants
GPU Initialization Action Enhancements for Secure Boot, Proxy, and Reliability
This large update significantly improves the
install_gpu_driver.shscript and its accompanying documentation, focusing on robust support for complex environments involving Secure Boot and HTTP/S proxies, and increasing overall reliability and maintainability.1.
gpu/README.md:GoogleCloudDataproc/custom-imagesrepository to create Dataproc images with NVIDIA drivers signed for Secure Boot. It covers environment setup, key management in GCP Secret Manager, Docker builder image creation, and running the image generation process.--shielded-secure-boot. It includes instructions for private network setups using Google Cloud Secure Web Proxy, leveraging scripts from theGoogleCloudDataproc/cloud-dataprocrepository for VPC, subnet, and proxy configuration.http-proxy,https-proxy,proxy-uri,no-proxy, andhttp-proxy-pem-uri.2.
gpu/install_gpu_driver.sh:set_proxy):http-proxy,https-proxy, andproxy-urimetadata, determining the correct proxy values for HTTP and HTTPS.HTTP_PROXY,HTTPS_PROXY, andNO_PROXYenvironment variables./etc/environmentwith the current proxy settings.gcloudproxy settings only if the gcloud SDK version is 547.0.0 or greater.aptanddnfto use the proxy.dirmngrorgnupg2-smimeis installed and configuresdirmngr.confto use the HTTP proxy.http-proxy-pem-uriinto system, Java, and Conda trust stores. Switches to HTTPS for proxy communications when a CA cert is provided.import_gpg_keys):import_gpg_keysto handle GPG key fetching and importing in a proxy-aware manner usingcurlover HTTPS, replacing directgpg --recv-keyscalls to keyservers.install_pytorch):numba,pytorch,tensorflow[and-cuda],rapids,pyspark, andcuda-version<=${CUDA_VERSION}. Explicit CUDA runtime (e.g.,cudart_spec) is no longer added, allowing the solver more flexibility.install_gpu_driver-mainandpytorchsentinels to allow forced refreshes.set_driver_version: Usescurl -Ifor a more lightweight HEAD request to check URL validity.build_driver_from_github: Caches the open kernel module source tarball from GitHub to GCS. Checks for existing signed and loadable modules to avoid unnecessary rebuilds.execute_github_driver_build: Refactored to accept tarball paths.popdremoved to balancepushdin caller. Removed a debug echo of thesign-fileexit code.make -j$(nproc)tomodules_installfor parallelization.modinfoforsigner:to confirm modules are signed.prepare_to_install: Movedcurl_retry_argsdefinition earlier.install_nvidia_gpu_driver: Checks ifnvidiamodule loads at the start and marks incomplete if not.main: Addedmark_complete install_gpu_driver-mainat the end.configure_dkms_certs: Always fetches keys from secret manager ifPSNis set to ensuremodulus_md5sumis available.install_gpu_agent: Checks ifMETADATA_HTTP_PROXY_PEM_URIis non-empty before using it.