AI-Hypercomputer / AI-Hypercomputer/google-cloud-mldiagnostics
ControlPlaneClient uses google.auth.default() without scopes → 400 INVALID_ARGUMENT under Workload Identity Federation (external_account)
- Linguagem predominante
- Python
- Estrelas
- 14
- Forks
- 4
- Métricas de merge de PRs
- Nenhum PR com merge em 30d
Descrição
### Summary
`ControlPlaneClient.__init__` calls `google.auth.default()` **without scopes**, which produces an unscoped credential for `external_account` (Workload Identity Federation) principals. The subsequent STS token exchange / service-account impersonation is then malformed and the Hypercompute Cluster control-plane API rejects every call with **HTTP 400 `INVALID_ARGUMENT`**. This breaks MLRun create/get for any GKE workload that authenticates via Workload Identity Federation (i.e. an `external_account` ADC).
A one-line fix (pass `scopes=[".../cloud-platform"]`) resolves it for both user ADC and `external_account` credentials. Proposed patch attached below.
### Affected code
`src/google_cloud_mldiagnostics/clients/control_plane_client.py` (current `main`, commit `08d554d`), line ~92:
```python
# Initialize Google Cloud credentials
self.credentials, _ = google.auth.default()
```
`self.credentials` is then used in `_get_access_token()` to mint the bearer token for all control-plane REST calls (`create_ml_run`, `get_ml_run`, `update_ml_run`, …).
### Root cause
`google.auth.default()` returns a credential **without an effective scope** unless one is requested. For most default credential types this is benign (the backend applies a default scope), but for **`external_account` credentials** (Workload Identity Federation) the scope is required to form a well-formed STS token exchange and service-account impersonation request. With no scope:
- the impersonation/STS call is malformed, and
- the control-plane API returns `400 INVALID_ARGUMENT` on the very first request.
This is deterministic (not transient/propagation-related) and reproduces on every call.
### Reproduction
Environment where we hit this:
- GKE workload (TPU v7x, plain-SPMD JobSet) authenticating as a GSA via **Workload Identity Federation** — the pod's ADC is an `external_account` config (subject token = a WIF-audience KSA token, with `service_account_impersonation_url`).
- SDK `google-cloud-mldiagnostics` 1.0.3, JAX 0.10.0.
Steps:
1. Run a workload whose ADC is an `external_account` credential (WIF), i.e. `GOOGLE_APPLICATION_CREDENTIALS` points at an external_account JSON.
2. Let the SDK create an MLRun (`managed_mldiagnostics=True`, or call `ControlPlaneClient` directly).
3. Observe every control-plane call fail with `400 INVALID_ARGUMENT`.
Note: this is **masked in most manual testing** because interactive `gcloud`/user ADC carries an implicit scope, and any manual `google.auth.default(scopes=[...])` call also masks it. It only surfaces with an unscoped `external_account` credential — exactly the GKE + Workload Identity Federation path the SDK targets.
### Proposed fix
```diff
--- a/src/google_cloud_mldiagnostics/clients/control_plane_client.py
+++ b/src/google_cloud_mldiagnostics/clients/control_plane_client.py
@@ -88,8 +88,19 @@ class ControlPlaneClient:
self.base_url = base_url
self.ml_runs_path = f"{base_url}/projects/{project_id}/locations/{location}/machineLearningRuns"
- # Initialize Google Cloud credentials
- self.credentials, _ = google.auth.default()
+ # Initialize Google Cloud credentials.
+ #
+ # Request the cloud-platform scope explicitly. On credential types that do
+ # not carry an implicit scope -- notably external_account (Workload
+ # Identity Federation) credentials, e.g. a GKE workload authenticating as a
+ # service account via WIF -- google.auth.default() with no scopes yields an
+ # unscoped credential. The subsequent STS token exchange / service account
+ # impersonation is then malformed and the control-plane API rejects it with
+ # HTTP 400 INVALID_ARGUMENT. Passing scopes makes token minting well-formed
+ # for both user ADC and external_account credentials.
+ self.credentials, _ = google.auth.default(
+ scopes=["https://www.googleapis.com/auth/cloud-platform"]
+ )
def _get_access_token(self) -> str:
"""Get Google Cloud access token for authentication."""
```
- One call site, no new imports.
- Verified: patch applies cleanly on `main` (`08d554d`), `ast.parse` OK.
- With the scope passed, MLRun create/get succeed under an `external_account` credential; behavior is unchanged for user ADC.
### Suggested follow-ups (optional, separate)
Other clients that build a Google client under the same WIF path may want the same explicit scoping if they ever exhibit unscoped `external_account` behavior (e.g. the logging/storage clients) — the control-plane client is the one that deterministically breaks today.
Happy to open a PR with this change if the project prefers (noting CONTRIBUTING.md currently routes contributions through issues).
Guia de contribuição
Avaliação
Esta issue ainda não foi avaliada.