AI-Hypercomputer / AI-Hypercomputer/google-cloud-mldiagnostics

ControlPlaneClient uses google.auth.default() without scopes → 400 INVALID_ARGUMENT under Workload Identity Federation (external_account)

Abierto
#4 0 comentarios 0 reacciones 0 asignados Ver en GitHub
Lenguaje dominante
Python
Estrellas
14
Forks
4
Métricas de merge de PR
Sin PR fusionados en 30 d

Descripción

### 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).

Guía de contribución

Abrir la guía de contribución

Evaluación

Este issue todavía no se ha evaluado.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.