Skip to content

feat(operator-wandb): use deployment-wide Azure storage identity - #606

Merged
zacharyblasczyk merged 22 commits into
mainfrom
zb/azure-instance-byob-identity
Aug 5, 2026
Merged

feat(operator-wandb): use deployment-wide Azure storage identity#606
zacharyblasczyk merged 22 commits into
mainfrom
zb/azure-instance-byob-identity

Conversation

@zacharyblasczyk

@zacharyblasczyk zacharyblasczyk commented Jun 4, 2026

Copy link
Copy Markdown
Contributor

Summary

This chart adds a deployment-wide Azure storage identity while preserving existing key-based and legacy BYOB behavior.

  • global.azureStorageIdentity holds the deployment tenant, client, principal, and shared ServiceAccount contract
  • the chart creates one wandb-bucket-access ServiceAccount by default, or uses an externally created account with create: false
  • wandb-base, Weave, and Bufstream storage consumers select the shared account and keep their RoleBindings aligned
  • literal client IDs are applied to the ServiceAccount annotation automatically; reference-valued client IDs require an explicit literal annotation or an external account
  • global.bucket.azureAuthMethod lets an Azure customer bucket select accessKey or workloadIdentity
  • existing bucket-scoped tenant/client IDs and explicit storage-account keys remain supported
  • both Azure credential fields are omitted from the Secret when workload identity is selected

Weave file storage remains off by default. global.weave-trace.fileStorage.useDefaultBucket maps only the managed global.defaultBucket to WF_FILE_STORAGE_URI; it rejects bucket paths containing an object prefix. Mixed access-key BYOB plus Weave default-bucket identity is covered explicitly.

Grant the deployment identity Reader and Storage Blob Data Contributor on the storage account. Existing Azure BYOB continues to use storage-account keys unless it explicitly selects workloadIdentity.

The operator chart remains at feature release 0.44.0.

Verification

  • focused Azure storage Helm suite: 25 passed
  • template formatter and 16 Python tests
  • template maintainability check
  • ct lint
  • full operator-wandb snapshot suite with Helm 3.20.1

Azure QA validation

Validated on 2026-07-30 with chart prerelease 0.44.0-PR606-13c71a2b and Core prerelease 0.84.0-azure-sas-blob-test.1.

  • W&B SDK run files and artifacts used Azure user-delegation SAS for GET and PUT
  • sibling blob access returned HTTP 403
  • allowlisted Weave storage wrote to the managed Azure default bucket
  • non-allowlisted Weave storage remained inline
  • disabling the feature restored the original account-key SAS behavior; every wandb verify check passed on the restored runtime

The final follow-up commit adds regression coverage only; it does not change rendered chart output.

Rollout

Release this chart before enabling wandb/core#49026. Keep the Core migration disabled for existing installations until they are explicitly selected for rollout.

Jira: WB-38011

Companion PRs: wandb/core#49026, wandb/docs#2967

Summary by CodeRabbit

  • New Features

    • Added Azure Storage workload identity support, including shared service accounts, credential references, and automatic workload identity labeling.
    • Added configuration options for Azure tenant/client IDs and access-key or workload-identity authentication.
    • Added Azure BYOB and workload identity deployment configurations.
    • Improved Weave Trace storage URI handling for Azure-backed storage.
  • Bug Fixes

    • Prevented unnecessary Azure access keys and secrets from being rendered when workload identity is enabled.
  • Tests

    • Added comprehensive coverage for Azure authentication, service accounts, RBAC, and storage configurations.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
charts/operator-wandb/tests/azure_storage_auth_test.yaml (1)

36-58: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win

Missing scenario: customer bucket with azureAuthMethod: workloadIdentity and no bucket-scoped IDs.

_bucket.tpl assigns global.azureStorageIdentity.tenantId/clientId for customer buckets when azureAuthMethod is workloadIdentity, but azure_storage_auth_test.yaml only covers this with bucket-scoped legacy IDs or access-key BYOB. Add a test analogous to the JSON helper serialization case, using global.bucket.provider: az, no azureTenantId/azureClientId, and literal or valueFrom credentials under global.azureStorageIdentity to assert AZURE_STORAGE_TENANT_ID and AZURE_STORAGE_CLIENT_ID.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@charts/operator-wandb/tests/azure_storage_auth_test.yaml` around lines 36 -
58, Add a workload-identity customer-bucket test in azure_storage_auth_test.yaml
using global.bucket.provider az, azureAuthMethod workloadIdentity, no
bucket-scoped tenant/client IDs, and credentials from
global.azureStorageIdentity (literal or valueFrom). Render templates/bucket.yaml
and assert AZURE_STORAGE_TENANT_ID and AZURE_STORAGE_CLIENT_ID are populated,
matching the JSON helper serialization scenario.
🧹 Nitpick comments (1)
charts/wandb-base/templates/_helpers.tpl (1)

102-125: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy lift

Same $usesDeploymentIdentity condition duplicated across wandb-base and bufstream. Both helpers independently recompute "customer bucket + azureAuthMethod: workloadIdentity, or default bucket + provider: az" to decide whether a workload should use the deployment-wide Azure storage identity, with no shared source of truth.

  • charts/wandb-base/templates/_helpers.tpl#L102-L125: extract the $usesDeploymentIdentity computation (lines 108-111) into a single named template that both this helper and the bufstream one can call.
  • charts/operator-wandb/charts/bufstream/templates/_helpers.tpl#L28-L39: call the shared template instead of recomputing $usesDeploymentIdentity (lines 34-37) inline.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@charts/wandb-base/templates/_helpers.tpl` around lines 102 - 125, Extract the
duplicated deployment-identity condition from wandb-base’s
azureStorageServiceAccountEnabled helper into a shared named template,
preserving the existing customer-bucket and default-bucket checks. Update
charts/wandb-base/templates/_helpers.tpl lines 102-125 to call that template,
and update charts/operator-wandb/charts/bufstream/templates/_helpers.tpl lines
28-39 to replace its inline $usesDeploymentIdentity calculation with the shared
helper call.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@charts/operator-wandb/tests/azure_storage_auth_test.yaml`:
- Around line 36-58: Add a workload-identity customer-bucket test in
azure_storage_auth_test.yaml using global.bucket.provider az, azureAuthMethod
workloadIdentity, no bucket-scoped tenant/client IDs, and credentials from
global.azureStorageIdentity (literal or valueFrom). Render templates/bucket.yaml
and assert AZURE_STORAGE_TENANT_ID and AZURE_STORAGE_CLIENT_ID are populated,
matching the JSON helper serialization scenario.

---

Nitpick comments:
In `@charts/wandb-base/templates/_helpers.tpl`:
- Around line 102-125: Extract the duplicated deployment-identity condition from
wandb-base’s azureStorageServiceAccountEnabled helper into a shared named
template, preserving the existing customer-bucket and default-bucket checks.
Update charts/wandb-base/templates/_helpers.tpl lines 102-125 to call that
template, and update
charts/operator-wandb/charts/bufstream/templates/_helpers.tpl lines 28-39 to
replace its inline $usesDeploymentIdentity calculation with the shared helper
call.

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: b294bb59-cb46-4b23-817c-94cd0a598f1f

📥 Commits

Reviewing files that changed from the base of the PR and between ee6ed29 and 201e70a.

⛔ Files ignored due to path filters (2)
  • test-configs/operator-wandb/__snapshots__/azure-workload-identity-refs.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/azure-workload-identity.snap is excluded by !**/*.snap
📒 Files selected for processing (11)
  • charts/operator-wandb/charts/bufstream/templates/_helpers.tpl
  • charts/operator-wandb/charts/bufstream/templates/cluster-role-bind.yaml
  • charts/operator-wandb/charts/bufstream/templates/deployment.yaml
  • charts/operator-wandb/charts/bufstream/templates/service-account.yaml
  • charts/operator-wandb/templates/_bucket.tpl
  • charts/operator-wandb/templates/azure-storage-serviceaccount.yaml
  • charts/operator-wandb/tests/azure_storage_auth_test.yaml
  • charts/operator-wandb/values.yaml
  • charts/wandb-base/templates/_helpers.tpl
  • charts/wandb-base/templates/serviceaccount.yaml
  • test-configs/operator-wandb/azure-workload-identity-refs.yaml

@zacharyblasczyk

Copy link
Copy Markdown
Contributor Author

@coderabbitai full review.

@coderabbitai

coderabbitai Bot commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

@zacharyblasczyk I’ll perform a full review of PR #606, including the Azure workload identity and Weave storage changes.

✅ Action performed

Full review finished.


Your included review limit is currently reached under our Fair Usage Limits Policy. This review may still proceed through usage-based billing if eligible. Your next included review will be available in 55 minutes.

zacharyblasczyk and others added 22 commits August 5, 2026 08:42
Surface azureTenantId / azureClientId from the bucket config and, when both are
set, inject AZURE_STORAGE_TENANT_ID + AZURE_STORAGE_CLIENT_ID into the app
deployments alongside the (now optional) AZURE_STORAGE_KEY. This lets gorilla
authenticate to Azure Blob storage with an Entra ID workload identity
(user-delegation SAS) instead of a storage account access key, which dedicated
Azure customers (e.g. Novo Nordisk) must adopt before losing Account SAS.

- _bucket.tpl: read azureTenantId/azureClientId from bucket / defaultBucket and
  add them to the wandb.bucket helper output.
- _env.tpl (wandb.bucketEnvs): emit the two env vars only when present, so
  existing key-only deployments are unchanged.
- values.yaml: document the fields and the workload-identity prerequisites.

Pairs with core (reads these env vars) and console (writes the fields into the
user spec).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@coderabbitai

coderabbitai Bot commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (2)
charts/operator-wandb/tests/azure_storage_auth_test.yaml (2)

320-361: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Assert shared-account and pod-label wiring for workload-identity BYOB.

This case verifies only the Azure environment variables. Also assert spec.template.spec.serviceAccountName and spec.template.metadata.labels["azure.workload.identity/use"]. The test can otherwise pass if the app keeps its component account or disables Azure Workload Identity injection.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@charts/operator-wandb/tests/azure_storage_auth_test.yaml` around lines 320 -
361, Extend the workload-identity BYOB test identified by “uses deployment
identity references for workload-identity BYOB” to assert that
spec.template.spec.serviceAccountName uses the configured shared Azure identity
service account and that
spec.template.metadata.labels["azure.workload.identity/use"] is enabled. Keep
the existing Azure environment-variable assertions unchanged.

363-491: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Cover Weave URI mapping and object-prefix rejection.

These cases verify ServiceAccount selection only. Add an assertion for WF_FILE_STORAGE_URI that uses the managed default bucket, not global.bucket. Add a failedTemplate case for useDefaultBucket: true with a path such as managed-container/object-prefix.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@charts/operator-wandb/tests/azure_storage_auth_test.yaml` around lines 363 -
491, Extend the Weave deployment tests around the useDefaultBucket cases to
assert WF_FILE_STORAGE_URI is built from global.defaultBucket’s name and path,
not global.bucket. Add a failedTemplate test for useDefaultBucket: true with a
defaultBucket path containing an object prefix such as
managed-container/object-prefix, and verify rendering is rejected.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@charts/operator-wandb/tests/azure_storage_auth_test.yaml`:
- Around line 320-361: Extend the workload-identity BYOB test identified by
“uses deployment identity references for workload-identity BYOB” to assert that
spec.template.spec.serviceAccountName uses the configured shared Azure identity
service account and that
spec.template.metadata.labels["azure.workload.identity/use"] is enabled. Keep
the existing Azure environment-variable assertions unchanged.
- Around line 363-491: Extend the Weave deployment tests around the
useDefaultBucket cases to assert WF_FILE_STORAGE_URI is built from
global.defaultBucket’s name and path, not global.bucket. Add a failedTemplate
test for useDefaultBucket: true with a defaultBucket path containing an object
prefix such as managed-container/object-prefix, and verify rendering is
rejected.

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 2501072d-f61c-4f9c-9fa6-9d176ac0ccdc

📥 Commits

Reviewing files that changed from the base of the PR and between e49742b and 72f42aa.

⛔ Files ignored due to path filters (57)
  • charts/lumen/Chart.lock is excluded by !**/*.lock
  • charts/operator-wandb/Chart.lock is excluded by !**/*.lock
  • charts/orchestrator/Chart.lock is excluded by !**/*.lock
  • test-configs/operator-wandb/__snapshots__/activity-store-disabled.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/activity-store-serve-backfill-off.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/azure-byob-access-key.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/azure-workload-identity-refs.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/azure-workload-identity.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/default.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/fmb.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/glue-leader-election.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/history-reader.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/historystore-parquet-grpc.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/historystore-parquet-only.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/keda-api-prometheus.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/keda-executor.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/keda-frfu.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/keda-parquet.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/local-bypass-no-app.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/local-bypass-with-app.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/mcp-server-empty-trace-url.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/mcp-server-external-trace.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/mcp-server-no-weave.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/mcp-server.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/no-local-bypass.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/oidc-secret-from-k8s.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/olap-features-enabled.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/olap-multi-feature.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/run-store-accelerator-run-updater.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/runs-v2-bufstream.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-api-rate-limits.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-ch-migration-job-no-olap.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-ch-migration-job.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-console-env-defaults.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-console-env.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-console-extraEnv.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-enable-backfill.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-image-digest-override.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-image-tag-override.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-lumen-aws.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-lumen-azure.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-lumen-gcp.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-lumen-onprem-custom-bucket.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-lumen-onprem-default-bucket.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-mysql-cacert-inline.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-no-oidc-settings-extra-cors.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-oidc-settings-default.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-oidc-settings-extra-cors.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-priority-classes.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-smtp-mail-from-secret.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-smtp-mail-from.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/snap-tolerations-and-selectors.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/url-encoded-password.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/user-defined-clickhouse-secret.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/user-defined-secrets.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/weave-trace-with-worker.snap is excluded by !**/*.snap
  • test-configs/operator-wandb/__snapshots__/weave-trace.snap is excluded by !**/*.snap
📒 Files selected for processing (21)
  • .github/workflows/test-operator-wandb.yaml
  • charts/lumen/Chart.yaml
  • charts/operator-wandb/Chart.yaml
  • charts/operator-wandb/charts/bufstream/templates/_helpers.tpl
  • charts/operator-wandb/charts/bufstream/templates/cluster-role-bind.yaml
  • charts/operator-wandb/charts/bufstream/templates/deployment.yaml
  • charts/operator-wandb/charts/bufstream/templates/service-account.yaml
  • charts/operator-wandb/templates/_bucket.tpl
  • charts/operator-wandb/templates/_env.tpl
  • charts/operator-wandb/templates/azure-storage-serviceaccount.yaml
  • charts/operator-wandb/templates/bucket.yaml
  • charts/operator-wandb/templates/weave-trace.yaml
  • charts/operator-wandb/tests/azure_storage_auth_test.yaml
  • charts/operator-wandb/values.yaml
  • charts/orchestrator/Chart.yaml
  • charts/wandb-base/Chart.yaml
  • charts/wandb-base/templates/_helpers.tpl
  • charts/wandb-base/templates/serviceaccount.yaml
  • test-configs/operator-wandb/azure-byob-access-key.yaml
  • test-configs/operator-wandb/azure-workload-identity-refs.yaml
  • test-configs/operator-wandb/azure-workload-identity.yaml
🚧 Files skipped from review as they are similar to previous changes (15)
  • charts/operator-wandb/charts/bufstream/templates/deployment.yaml
  • .github/workflows/test-operator-wandb.yaml
  • charts/wandb-base/templates/serviceaccount.yaml
  • charts/operator-wandb/charts/bufstream/templates/_helpers.tpl
  • test-configs/operator-wandb/azure-workload-identity.yaml
  • charts/operator-wandb/charts/bufstream/templates/cluster-role-bind.yaml
  • charts/operator-wandb/templates/weave-trace.yaml
  • charts/operator-wandb/templates/bucket.yaml
  • test-configs/operator-wandb/azure-byob-access-key.yaml
  • test-configs/operator-wandb/azure-workload-identity-refs.yaml
  • charts/operator-wandb/templates/_bucket.tpl
  • charts/operator-wandb/templates/_env.tpl
  • charts/operator-wandb/charts/bufstream/templates/service-account.yaml
  • charts/operator-wandb/templates/azure-storage-serviceaccount.yaml
  • charts/operator-wandb/values.yaml

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants