Refactor: Improve Proxy Handling and Secure Boot in GPU Install Script

WORKSPACE

@@ -2,8 +2,8 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
 http_archive(
     name = "rules_python",
-    sha256 = "aa96a691d3a8177f3215b14b0edc9641787abaaa30363a080165d06ab65e1161",
-    url = "https://github.com/bazelbuild/rules_python/releases/download/0.0.1/rules_python-0.0.1.tar.gz",
+    sha256 = "6251b118b690835697ebc0b080552b827f8376916c526d135c36952cb9109038",
+    url = "https://github.com/bazelbuild/rules_python/releases/download/0.35.0/rules_python-0.35.0.tar.gz",
 )
 
 http_archive(

cloudbuild/Dockerfile

@@ -22,7 +22,7 @@ RUN /usr/bin/curl -s https://bazel.build/bazel-release.pub.gpg | \
     dd of="${bazel_repo_file}" status=none && \
     apt-get update -qq
 RUN apt-get autoremove -y -qq > /dev/null 2>&1 && \
-    apt-get install -y -qq default-jdk python3-setuptools python3-pip bazel-${bazel_version} > /dev/null 2>&1 && \
+    apt-get install -y -qq jq default-jdk python3-setuptools python3-pip bazel-${bazel_version} > /dev/null 2>&1 && \
     apt-get clean
 
 # Set bazel-${bazel_version} as the default bazel alternative in this container

cloudbuild/cloudbuild.yaml

@@ -108,6 +108,39 @@ steps:
       - 'CLOUDSDK_COMPUTE_REGION=us-central1'
       - 'CLOUDSDK_CONTAINER_CLUSTER=init-actions-presubmit'
 
+  # Run presubmit tests in parallel for 2.3 Debian image
+  - name: 'gcr.io/cloud-builders/kubectl'
+    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']
+    env:
+      - 'COMMIT_SHA=$COMMIT_SHA'
+      - 'CLOUDSDK_COMPUTE_REGION=us-central1'
+      - 'CLOUDSDK_CONTAINER_CLUSTER=init-actions-presubmit'
+
+  # Run presubmit tests in parallel for 2.3 Rocky Linux image
+  - name: 'gcr.io/cloud-builders/kubectl'
+    id: 'dataproc-2.3-rocky9-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-rocky9']
+    env:
+      - 'COMMIT_SHA=$COMMIT_SHA'
+      - 'CLOUDSDK_COMPUTE_REGION=us-central1'
+      - 'CLOUDSDK_CONTAINER_CLUSTER=init-actions-presubmit'
+
+  # Run presubmit tests in parallel for 2.3 Ubuntu image
+  - name: 'gcr.io/cloud-builders/kubectl'
+    id: 'dataproc-2.3-ubuntu22-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-ubuntu22']
+    env:
+      - 'COMMIT_SHA=$COMMIT_SHA'
+      - 'CLOUDSDK_COMPUTE_REGION=us-central1'
+      - 'CLOUDSDK_CONTAINER_CLUSTER=init-actions-presubmit'
+
   # Delete Docker image from GCR
   - name: 'gcr.io/cloud-builders/gcloud'
     args: ['container', 'images', 'delete', 'gcr.io/$PROJECT_ID/init-actions-image:$BUILD_ID']

gpu/BUILD

@@ -2,6 +2,14 @@ package(default_visibility = ["//visibility:public"])
 
 exports_files(["install_gpu_driver.sh", "mig.sh"])
 
+filegroup(
+    name = "gpu_test_data",
+    srcs = [
+        "install_gpu_driver.sh",
+        "mig.sh",
+    ],
+)
+
 py_library(
     name = "gpu_test_case_base",
     srcs = ["gpu_test_case_base.py"],
@@ -17,12 +25,25 @@ py_test(
     name = "test_gpu",
     size = "enormous",
     srcs = ["test_gpu.py"],
-    data = ["install_gpu_driver.sh", "mig.sh"],
+    data = [":gpu_test_data"],
     local = True,
     shard_count = 15,
     deps = [
         ":gpu_test_case_base",  # Add this dependency
         "//integration_tests:dataproc_test_case",
+        "@io_abseil_py//absl/testing:absltest",
         "@io_abseil_py//absl/testing:parameterized",
     ],
-)
+)
+
+py_test(
+    name = "test_gpu_single",
+    size = "enormous",
+    srcs = ["test_gpu_single.py"],
+    data = [":gpu_test_data"],
+    local = True,
+    deps = [
+        "//integration_tests:dataproc_test_case",
+        "@io_abseil_py//absl/testing:absltest",
+    ], 
+)

gpu/Dockerfile

@@ -1,47 +1,45 @@
 # This Dockerfile builds the container from which manual tests are run
-# This process needs to be executed manually from a git clone
-#
-# See manual-test-runner.sh for instructions
-
 FROM gcr.io/cloud-builders/gcloud
 
 RUN useradd -m -d /home/ia-tests -s /bin/bash ia-tests
 
-# Installed here are packages on which the tests depend
-RUN apt-get -qq update \
-  && apt-get -y -qq install \
-     apt-transport-https apt-utils \
-     ca-certificates libmime-base64-perl gnupg \
-     curl jq less screen > /dev/null 2>&1  && apt-get clean
-
-# Install bazel signing key, repo and package
+# Install Bazel:
 ENV bazel_kr_path=/usr/share/keyrings/bazel-keyring.gpg \
     bazel_version=7.4.0 \
     bazel_repo_data="http://storage.googleapis.com/bazel-apt stable jdk1.8" \
+    bazel_repo_file="/etc/apt/sources.list.d/bazel.list" \
     DEBIAN_FRONTEND=noninteractive
-
-RUN /usr/bin/curl -s https://bazel.build/bazel-release.pub.gpg \
-      | gpg --dearmor -o "${bazel_kr_path}" \
-    && echo "deb [arch=amd64 signed-by=${bazel_kr_path}] ${bazel_repo_data}" \
-      | dd of=/etc/apt/sources.list.d/bazel.list status=none \
-    && apt-get update -qq
-
+RUN apt-get update -qq && apt-get install -y -qq curl >/dev/null 2>&1 && \
+    apt-get clean
+RUN /usr/bin/curl -s https://bazel.build/bazel-release.pub.gpg | \
+    gpg --import --no-default-keyring --keyring "${bazel_kr_path}" && \
+    echo "deb [arch=amd64 signed-by=${bazel_kr_path}] ${bazel_repo_data}" | \
+    dd of="${bazel_repo_file}" status=none && \
+    apt-get update -qq
 RUN apt-get autoremove -y -qq > /dev/null 2>&1 && \
-    apt-get install -y -qq default-jdk python3-setuptools bazel-${bazel_version} > /dev/null 2>&1 && \
+    apt-get install -y -qq default-jdk python3-setuptools python3-pip bazel-${bazel_version} > /dev/null 2>&1 && \
     apt-get clean
-
-# Set bazel-${bazel_version} as the default bazel alternative in this container
 RUN update-alternatives --install /usr/bin/bazel bazel /usr/bin/bazel-${bazel_version} 1 && \
     update-alternatives --set bazel /usr/bin/bazel-${bazel_version}
 
-# Install here any utilities you find useful when troubleshooting
-RUN apt-get -y -qq install emacs-nox vim uuid-runtime > /dev/null 2>&1 && apt-get clean
+# Install other needed tools
+RUN apt-get install -y -qq jq less screen vim uuid-runtime > /dev/null 2>&1 && \
+    apt-get clean
 
 WORKDIR /init-actions
 
+# Copy necessary files
+COPY --chown=ia-tests:ia-tests WORKSPACE BUILD requirements.txt env.json key.json /init-actions/
+COPY --chown=ia-tests:ia-tests gpu /init-actions/gpu
+COPY --chown=ia-tests:ia-tests integration_tests /init-actions/integration_tests
+
 USER ia-tests
-COPY --chown=ia-tests:ia-tests "cloudbuild/key.json" /key.json
-COPY --chown=ia-tests:ia-tests . ${WORKDIR}
 
-ENTRYPOINT ["/bin/bash"]
-#CMD ["/bin/bash"]
+# Install Python dependencies
+RUN python3 -m pip install --upgrade pip && \
+    python3 -m pip install -r /init-actions/requirements.txt
+# Removed --user and --break-system-packages
+
+#ENTRYPOINT ["/bin/bash"]
+
+CMD ["/bin/bash"]