protect service accounts for ongoing Cloud Builds and queued Kueue jobs - #6253
protect service accounts for ongoing Cloud Builds and queued Kueue jobs#6253aslam-quad wants to merge 4 commits into
Conversation
Summary of ChangesHello, 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 introduces a safety mechanism to the cleanup process by identifying and protecting service accounts currently in use by active Cloud Build builds and GKE Kueue jobs. By querying active builds and inspecting GKE clusters for running or queued jobs, the script ensures that critical service accounts are not prematurely removed during automated cleanup tasks. Highlights
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 the 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 counterproductive. 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. 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 introduces a new function populate_active_build_exclusions in tools/cleanup.sh to protect service accounts associated with ongoing Cloud Build builds and active GKE Kueue batch jobs from being cleaned up. Feedback on this change highlights issues with inconsistent indentation (spaces instead of tabs) and a potential resource leak where the temporary KUBECONFIG file is not cleaned up if the script is interrupted, suggesting the use of a trap on EXIT for robust cleanup.
27ceb2e to
22e3d60
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a new function populate_active_build_exclusions in tools/cleanup.sh to protect service accounts associated with active Cloud Build builds and GKE Kueue batch jobs from being cleaned up. The review feedback highlights two important issues: first, the gcloud builds list command uses an unsupported --ongoing flag which will cause it to fail, and second, a failure in mktemp is not handled, which could inadvertently lead to overwriting the user's default kubeconfig file.
22e3d60 to
39b1109
Compare
39b1109 to
caa0088
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a new function populate_active_build_exclusions in tools/cleanup.sh to protect service accounts associated with active Cloud Build builds and GKE Kueue batch jobs from being deleted during cleanup. The review feedback highlights a potential race condition in Kueue job filtering where provisioning jobs could be missed, suggesting a robust filtering approach that avoids masking kubectl errors. Additionally, the feedback recommends preserving and restoring the user's original KUBECONFIG environment variable instead of unconditionally unsetting it to prevent unintended side effects.
c1269be to
5dfb560
Compare
| if [[ "$succeeded" =~ ^(<none>|0)$ && "$failed" =~ ^(<none>|0)$ ]]; then | ||
| if [[ -z "${EXCLUSION_MAP[${job_name}]:-}" ]]; then | ||
| log "INFO" "Protecting active/queued Kueue job: ${job_name} (cluster: ${cluster_name})" | ||
| EXCLUSION_MAP["${job_name}"]=1 |
There was a problem hiding this comment.
There is a critical mismatch here that will cause service accounts for active/queued Kueue jobs to still be deleted.
The Problem:
process_service_accounts() checks if any key in EXCLUSION_MAP is a substring of the service account email ("$email" == *"$protected_name"*).
Adding the raw Kubernetes job name EXCLUSION_MAP["${job_name}"]=1 does not match the actual service account email because many integration tests use a DEPLOYMENT_NAME (and thus service account name) that diverges from the Kubernetes Job name:
-
vm-storage.yaml: Job name isvm-storage-a1b2c3, but deployment name isvmstorage-a1b2c3(no hyphen)$\rightarrow$ SA email isvmstorage-a1b2c3-compute@.... -
ml-a3-ultragpu-onspot-jbvms.yaml: Job name isml-a3-ultragpu-onspot-jbvms-a1b2c3, but deployment name isa3u-spot-jbvms-a1b2c3$\rightarrow$ SA email isa3u-spot-jbvms-a1b2c3-sa@.... -
slurm-gcp-v6-ubuntu.yaml: Job name isslurm-gcp-v6-ubuntu-a1b2c3, but deployment name isubun-v6-a1b2c3$\rightarrow$ SA email isubun-v6-a1b2c3-sa@.... -
gke-managed-hyperdisk.yaml: Job name isgke-managed-hyperdisk-a1b2c3, but deployment name isgke-hyperdisk-a1b2c3. -
packer.yaml: Job name ispacker-a1b2c3, but deployment name ispkrv6a1b2c3.
Because "vmstorage-a1b2c3-compute@..." does not contain "vm-storage-a1b2c3", the substring match fails and the service account is deleted during cleanup.
Solution:
In addition to the job name, extract the 6-character short build ID from the job (from the job name suffix ${job_name##*-} or the build-id label) and add that prefix to EXCLUSION_MAP, matching how Part A works:
| EXCLUSION_MAP["${job_name}"]=1 | |
| # If a job hasn't succeeded and hasn't failed, it MUST be Running or Queued! | |
| if [[ "$succeeded" =~ ^(<none>|0)$ && "$failed" =~ ^(<none>|0)$ ]]; then | |
| EXCLUSION_MAP["${job_name}"]=1 | |
| # Extract short build ID prefix to protect SAs derived from DEPLOYMENT_NAME | |
| local short_id="${job_name##*-}" | |
| if [[ -n "$short_id" && -z "${EXCLUSION_MAP[${short_id}]:-}" ]]; then | |
| log "INFO" "Protecting active/queued Kueue job prefix: ${short_id} (job: ${job_name}, cluster: ${cluster_name})" | |
| EXCLUSION_MAP["${short_id}"]=1 | |
| fi | |
| fi |
| fi | ||
|
|
||
| # --- Part B: Protect service accounts tied to active GKE Kueue batch jobs --- | ||
| if ! command -v kubectl &>/dev/null; then |
There was a problem hiding this comment.
In tools/cloud-build/project-cleanup.yaml, the scheduled cleanup step executes inside gcr.io/cloud-builders/gcloud:
- name: gcr.io/cloud-builders/gcloud
entrypoint: /bin/bash
...
/workspace/tools/cleanup.shgcr.io/cloud-builders/gcloud does not have kubectl pre-installed. In production CI, command -v kubectl will evaluate to false, log a warning, and silently skip Part B on every scheduled run.
To ensure Part B actually protects Kueue jobs in CI, we should update project-cleanup.yaml in this PR to install kubectl and gke-gcloud-auth-plugin before calling cleanup.sh, e.g.:
- |
set -euo pipefail
gcloud components install kubectl gke-gcloud-auth-plugin --quiet || true
...
/workspace/tools/cleanup.shThere was a problem hiding this comment.
Updated the cleanup step to use the pre-built Cloud SDK image (gcr.io/google.com/cloudsdktool/cloud-sdk) and installed kubectl and gke-gcloud-auth-plugin before running cleanup.sh. This ensures the required tools are available in the CI environment.
Neelabh94
left a comment
Neelabh94
left a comment
There was a problem hiding this comment.
Requesting changes to address empty key poisoning, TOCTOU race condition, and kubeconfig trap safety.
| fi | ||
| local old_kubeconfig="${KUBECONFIG:-}" | ||
| export KUBECONFIG="$temp_kubeconfig" | ||
| trap 'rm -f "$temp_kubeconfig"' EXIT |
There was a problem hiding this comment.
Modifying global KUBECONFIG and setting trap 'rm -f "$temp_kubeconfig"' EXIT inside a helper function clobbers any script-level EXIT handlers, and trap - EXIT at line 180 permanently unsets all exit handlers for the remainder of the script.
Both gcloud container clusters get-credentials and kubectl accept --kubeconfig="$temp_kubeconfig" directly as a flag. Using the flag eliminates export KUBECONFIG, trap EXIT, and variable restoring entirely:
if ! gcloud container clusters get-credentials "$cluster_name" \
--location="$cluster_location" --project="$PROJECT_ID" \
--kubeconfig="$temp_kubeconfig" &>/dev/null; then
log "WARNING" "Failed to authenticate kubectl against cluster ${cluster_name}; skipping this cluster for Kueue job protection."
continue
fi
local raw_jobs
if ! raw_jobs=$(kubectl get jobs -A -l "kueue.x-k8s.io/queue-name" \
--kubeconfig="$temp_kubeconfig" \
--request-timeout=10s \
-o custom-columns="NAME:.metadata.name,SUCCEEDED:.status.succeeded,FAILED:.status.failed" \
--no-headers 2>/dev/null); then|
|
||
| # In kubectl, empty numbers show up as "<none>" | ||
| # If a job hasn't succeeded and hasn't failed, it MUST be Running or Queued! | ||
| if [[ "$succeeded" =~ ^(<none>|0)$ && "$failed" =~ ^(<none>|0)$ ]]; then |
There was a problem hiding this comment.
Two critical bugs in this block:
-
Empty Key Project-Wide Poisoning: If
linecontains leading whitespace or empty fields,job_namecan be empty, settingEXCLUSION_MAP[""]=1. Downstream inprocess_service_accounts(),[[ "$email" == *"$protected_name"* ]]matches every service account in the GCP project whenprotected_name="", causing the script to silently abort deleting all orphaned service accounts project-wide. Please guard with[[ -z "$job_name" ]] && continue. -
Kueue Job vs. Service Account Name Mismatch: Matching on full
job_namefails because integration test jobs derive service account emails fromDEPLOYMENT_NAME(e.g.vmstorage-a1b2c3-compute@...), which diverges from the k8s Job name (vm-storage-a1b2c3). In addition tojob_name, extract the short build ID suffix (${job_name##*-}) and add that prefix toEXCLUSION_MAP.
[[ -z "$job_name" ]] && continue
if [[ "$succeeded" =~ ^(<none>|0)$ && "$failed" =~ ^(<none>|0)$ ]]; then
if [[ -z "${EXCLUSION_MAP[${job_name}]:-}" ]]; then
log "INFO" "Protecting active/queued Kueue job: ${job_name} (cluster: ${cluster_name})"
EXCLUSION_MAP["${job_name}"]=1
fi
local short_id="${job_name##*-}"
if [[ -n "$short_id" && -z "${EXCLUSION_MAP[${short_id}]:-}" ]]; then
log "INFO" "Protecting service accounts matching Kueue build prefix: ${short_id} (job: ${job_name})"
EXCLUSION_MAP["${short_id}"]=1
fi
fi|
|
||
| check_dependencies | ||
| load_exclusions | ||
| populate_active_build_exclusions |
There was a problem hiding this comment.
There is a 15–20 minute TOCTOU (Time-of-Check to Time-of-Use) race window here. populate_active_build_exclusions is called at script start in Phase 1, but process_service_accounts() runs in Phase 5 after synchronous deletion of clusters, Filestores, disks, and networks. Any integration test triggered while earlier phases are executing will have its newly created service accounts deleted in Phase 5 because EXCLUSION_MAP was populated 20 minutes earlier.
Please move populate_active_build_exclusions to execute immediately before process_service_accounts in Phase 5.
0d8621d to
6139189
Compare
6139189 to
b19db43
Compare
3 participants
Description
This PR fixes a bug where the cleanup script accidentally deleted Service Accounts for tests that were still running or waiting in the queue.
Changes:
Cloud Build Protection: Finds ongoing Cloud Builds and extracts their 6-character ID, adding it to the script's EXCLUSION_MAP to protect their Service Accounts.
Kueue Job Protection: Safely checks clusters using custom-columns to find jobs that haven't succeeded or failed yet (catching both Running and Queued tests), and adds them to the EXCLUSION_MAP.
Submission Checklist
NOTE: Community submissions can take up to 2 weeks to be reviewed.
Please take the following actions before submitting this pull request.