diff --git a/.github/actions/setup-kind-cluster/action.yaml b/.github/actions/setup-kind-cluster/action.yaml index 57d591c75f..90d3115d3e 100644 --- a/.github/actions/setup-kind-cluster/action.yaml +++ b/.github/actions/setup-kind-cluster/action.yaml @@ -1,38 +1,42 @@ name: Setup Kind cluster with NeMo Platform description: > Creates a Kind cluster, installs tooling (kind, kubectl, Helm, uv), - pre-pulls images, deploys NeMo Platform via Helm, and waits for the - API to become healthy. + optionally deploys NeMo Platform via Helm, and waits for the API to become + healthy. inputs: - kind_cluster_name: + kind-cluster-name: description: Kind cluster name required: true - kube_namespace: + install-nemo-platform: + description: Deploy NeMo Platform into the Kind cluster + required: false + default: "true" + kube-namespace: description: Kubernetes namespace for NeMo Platform required: false default: nemo-platform - kube_gateway_name: + kube-gateway-name: description: Gateway resource name required: false default: nmp-e2e-gateway - image_registry: + image-registry: description: Container image registry required: true - image_tag: + image-tag: description: Container image tag required: true - helm_values: + helm-values: description: Helm values file path (relative to repo root) required: false default: e2e/k8s/values/kind.yaml - kind_image_pull_token: + kind-image-pull-token: description: Token for pulling images into Kind required: true - kind_image_pull_user: + kind-image-pull-user: description: Username for pulling images into Kind required: true - ngc_api_key: + ngc-api-key: description: NGC API key (can be a placeholder for CPU-only) required: false default: not-used @@ -108,22 +112,23 @@ runs: - name: Start kind cluster shell: bash env: - KIND_CLUSTER_NAME: ${{ inputs.kind_cluster_name }} - KUBE_NAMESPACE: ${{ inputs.kube_namespace }} - NGC_API_KEY: ${{ inputs.ngc_api_key }} + KIND_CLUSTER_NAME: ${{ inputs['kind-cluster-name'] }} + KUBE_NAMESPACE: ${{ inputs['kube-namespace'] }} + NGC_API_KEY: ${{ inputs['ngc-api-key'] }} + GITHUB_TOKEN: ${{ inputs['kind-image-pull-token'] }} run: bash e2e/k8s/scripts/setup_local_kind_cpu.sh - name: Set default kubectl namespace shell: bash env: - NAMESPACE: ${{ inputs.kube_namespace }} + NAMESPACE: ${{ inputs['kube-namespace'] }} run: kubectl config set-context --current --namespace="${NAMESPACE}" - name: Verify Gateway API setup shell: bash env: - NAMESPACE: ${{ inputs.kube_namespace }} - KUBE_GATEWAY_NAME: ${{ inputs.kube_gateway_name }} + NAMESPACE: ${{ inputs['kube-namespace'] }} + KUBE_GATEWAY_NAME: ${{ inputs['kube-gateway-name'] }} run: | set -euo pipefail kubectl wait --for=condition=Established crd/gateways.gateway.networking.k8s.io --timeout=2m @@ -132,27 +137,30 @@ runs: kubectl -n "${NAMESPACE}" get gateway "${KUBE_GATEWAY_NAME}" - name: Pre-pull GHCR images into kind + if: ${{ inputs['install-nemo-platform'] == 'true' }} shell: bash env: - KIND_IMAGE_PULL_TOKEN: ${{ inputs.kind_image_pull_token }} - KIND_IMAGE_PULL_USER: ${{ inputs.kind_image_pull_user }} - NMP_E2E_REGISTRY: ${{ inputs.image_registry }} - NMP_E2E_TAG: ${{ inputs.image_tag }} + KIND_IMAGE_PULL_TOKEN: ${{ inputs['kind-image-pull-token'] }} + KIND_IMAGE_PULL_USER: ${{ inputs['kind-image-pull-user'] }} + NMP_E2E_REGISTRY: ${{ inputs['image-registry'] }} + NMP_E2E_TAG: ${{ inputs['image-tag'] }} run: | e2e/k8s/scripts/prepull_kind_images.sh \ "${NMP_E2E_REGISTRY}/nmp-api:${NMP_E2E_TAG}" \ "${NMP_E2E_REGISTRY}/nmp-cpu-tasks:${NMP_E2E_TAG}" - name: Install NeMo Platform + if: ${{ inputs['install-nemo-platform'] == 'true' }} shell: bash env: - NAMESPACE: ${{ inputs.kube_namespace }} - NMP_E2E_REGISTRY: ${{ inputs.image_registry }} - NMP_E2E_TAG: ${{ inputs.image_tag }} - HELM_VALUES: ${{ inputs.helm_values }} + NAMESPACE: ${{ inputs['kube-namespace'] }} + NMP_E2E_REGISTRY: ${{ inputs['image-registry'] }} + NMP_E2E_TAG: ${{ inputs['image-tag'] }} + HELM_VALUES: ${{ inputs['helm-values'] }} REQUIRE_NMP_E2E_IMAGES: "true" POSTGRES_IMAGE: docker.io/library/postgres BUSYBOX_IMAGE: docker.io/library/busybox + GITHUB_TOKEN: ${{ inputs['kind-image-pull-token'] }} run: | if ! e2e/k8s/scripts/install_helm_e2e.sh; then echo "--- helm list -A ---" @@ -165,6 +173,7 @@ runs: fi - name: Wait for API + if: ${{ inputs['install-nemo-platform'] == 'true' }} shell: bash env: NMP_E2E_CLUSTER_URL: ${{ env.NMP_E2E_CLUSTER_URL }} diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 03fd538bcc..3bb03b3ecc 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -430,11 +430,11 @@ jobs: - name: Setup Kind cluster with NeMo Platform uses: ./.github/actions/setup-kind-cluster with: - kind_cluster_name: ${{ env.KIND_CLUSTER_NAME }} - image_registry: ${{ needs.build-cpu-smoke-images.outputs.image_registry }} - image_tag: ${{ needs.build-cpu-smoke-images.outputs.image_tag }} - kind_image_pull_token: ${{ github.token }} - kind_image_pull_user: ${{ github.actor }} + kind-cluster-name: ${{ env.KIND_CLUSTER_NAME }} + image-registry: ${{ needs.build-cpu-smoke-images.outputs.image_registry }} + image-tag: ${{ needs.build-cpu-smoke-images.outputs.image_tag }} + kind-image-pull-token: ${{ github.token }} + kind-image-pull-user: ${{ github.actor }} - name: Run CPU job e2e smoke test shell: bash @@ -513,11 +513,11 @@ jobs: - name: Setup Kind cluster with NeMo Platform uses: ./.github/actions/setup-kind-cluster with: - kind_cluster_name: ${{ env.KIND_CLUSTER_NAME }} - image_registry: ${{ needs.build-cpu-smoke-images.outputs.image_registry }} - image_tag: ${{ needs.build-cpu-smoke-images.outputs.image_tag }} - kind_image_pull_token: ${{ github.token }} - kind_image_pull_user: ${{ github.actor }} + kind-cluster-name: ${{ env.KIND_CLUSTER_NAME }} + image-registry: ${{ needs.build-cpu-smoke-images.outputs.image_registry }} + image-tag: ${{ needs.build-cpu-smoke-images.outputs.image_tag }} + kind-image-pull-token: ${{ github.token }} + kind-image-pull-user: ${{ github.actor }} - name: Run jobs and data-designer e2e tests shell: bash @@ -949,9 +949,51 @@ jobs: coverage.xml coverage.json - python-auth-idp-test: - name: Python auth-idp tests - needs: [changes, policy-wasm, build-cpu-smoke-images] + python-auth-idp-static-test: + name: Python auth-idp static tests + needs: [changes] + if: > + !cancelled() && ( + github.event_name == 'workflow_dispatch' || + needs.changes.outputs.cpu-smoke == 'true' || + needs.changes.outputs.auth-idp == 'true' + ) + runs-on: ubuntu-latest + permissions: + contents: read + steps: + - name: Checkout code + uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 + with: + persist-credentials: false + - name: Install uv + uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + with: + python-version: "3.11" + enable-cache: true + cache-dependency-glob: uv.lock + - name: Run auth-idp static tests + run: | + helm dependency build k8s/helm + helm dependency build contrib/auth/authentik/helm + uv run --frozen pytest tests/auth_idp/static -v + env: + _TYPER_FORCE_DISABLE_TERMINAL: "1" + - name: Upload test artifacts + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: python-auth-idp-static-test-results + retention-days: 30 + if-no-files-found: ignore + path: | + report.xml + coverage.xml + coverage.json + + python-auth-idp-e2e-test: + name: Python auth-idp ${{ matrix.runtime }} tests + needs: [changes, policy-wasm, build-cpu-smoke-images, python-auth-idp-static-test] if: > !cancelled() && ( github.event_name == 'workflow_dispatch' || @@ -959,11 +1001,33 @@ jobs: needs.changes.outputs.auth-idp == 'true' ) && needs.policy-wasm.result == 'success' && - needs.build-cpu-smoke-images.result == 'success' + needs.build-cpu-smoke-images.result == 'success' && + needs.build-cpu-smoke-images.outputs.publish_images == 'true' && + needs.python-auth-idp-static-test.result == 'success' runs-on: ubuntu-latest + timeout-minutes: 90 permissions: contents: read packages: read + strategy: + fail-fast: false + matrix: + include: + - provider: authentik + runtime: authentik-compose + backend: compose + command: compose + - provider: authentik + runtime: authentik-kubernetes + backend: kubernetes + command: k8s + env: + NMP_AUTHENTIK_K8S_RUNTIME: kind + NMP_AUTHENTIK_K8S_NAMESPACE: nemo-authentik + NMP_AUTHENTIK_K8S_IMAGE_PULL_SECRET: ghcr-pull + NMP_AUTHENTIK_K8S_NGC_EXISTING_SECRET: ngc-api + NMP_AUTHENTIK_K8S_CLUSTER_NAME: gha-${{ github.run_id }}-${{ github.run_attempt }}-${{ matrix.runtime }} + NMP_AUTHENTIK_K8S_REUSE_CLUSTER: "1" steps: - name: Checkout code uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 @@ -971,39 +1035,85 @@ jobs: persist-credentials: false - name: Free disk space uses: ./.github/actions/free-disk-space - - name: Download policy WASM - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 - with: - name: policy-wasm - path: services/core/auth/src/nmp/core/auth/assets - name: Install uv uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 with: python-version: "3.11" enable-cache: true cache-dependency-glob: uv.lock + - name: Setup Kind cluster + if: matrix.backend == 'kubernetes' + uses: ./.github/actions/setup-kind-cluster + with: + kind-cluster-name: ${{ env.NMP_AUTHENTIK_K8S_CLUSTER_NAME }} + install-nemo-platform: "false" + kube-namespace: ${{ env.NMP_AUTHENTIK_K8S_NAMESPACE }} + image-registry: ${{ needs.build-cpu-smoke-images.outputs.image_registry }} + image-tag: ${{ needs.build-cpu-smoke-images.outputs.image_tag }} + kind-image-pull-token: ${{ github.token }} + kind-image-pull-user: ${{ github.actor }} - name: Log in to GHCR - if: needs.build-cpu-smoke-images.outputs.publish_images == 'true' + if: matrix.backend == 'compose' && needs.build-cpu-smoke-images.outputs.publish_images == 'true' uses: docker/login-action@650006c6eb7dba73a995cc03b0b2d7f5ca915bee # v4.2.0 with: registry: ghcr.io username: ${{ github.actor }} password: ${{ github.token }} - - name: Run auth-idp tests - run: make test-auth-idp + - name: Install Helm + if: matrix.backend == 'kubernetes' + uses: azure/setup-helm@dda3372f752e03dde6b3237bc9431cdc2f7a02a2 # v5.0.0 + - name: Lint Authentik umbrella Helm chart + if: matrix.provider == 'authentik' && matrix.backend == 'kubernetes' + shell: bash + run: | + helm repo add authentik https://charts.goauthentik.io --force-update + helm dependency build k8s/helm + helm dependency build contrib/auth/authentik/helm + helm lint --strict contrib/auth/authentik/helm + - name: Verify official Authentik Helm chart is reachable + if: matrix.provider == 'authentik' && matrix.backend == 'kubernetes' + shell: bash + run: | + helm show chart authentik --repo https://charts.goauthentik.io --version 2026.5.4 + - name: Download policy WASM + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: policy-wasm + path: services/core/auth/src/nmp/core/auth/assets + - name: Run auth-idp ${{ matrix.runtime }} tests + run: > + contrib/auth/${{ matrix.provider }}/run.sh + ${{ matrix.command }} + --image "${IMAGE_REGISTRY}/nmp-api:${BAKE_TAG}" env: _TYPER_FORCE_DISABLE_TERMINAL: "1" E2E_SERVICES_LOG_DIR: ${{ runner.temp }}/e2e-services-logs IMAGE_REGISTRY: ${{ needs.build-cpu-smoke-images.outputs.image_registry }} BAKE_TAG: ${{ needs.build-cpu-smoke-images.outputs.image_tag }} + NMP_AUTHENTIK_K8S_JUNIT_XML: report-auth-idp-${{ matrix.runtime }}.xml + NMP_AUTHENTIK_K8S_IMAGE_PULL_SECRET: ${{ env.NMP_AUTHENTIK_K8S_IMAGE_PULL_SECRET }} + NMP_AUTHENTIK_K8S_SKIP_IMAGE_LOAD: ${{ needs.build-cpu-smoke-images.outputs.publish_images == 'true' && '1' || '0' }} + - name: Collect Kubernetes logs + if: always() && matrix.backend == 'kubernetes' + shell: bash + run: e2e/k8s/scripts/collect_k8s_logs.sh - name: Upload test artifacts if: always() uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: - name: python-auth-idp-test-results + name: python-auth-idp-${{ matrix.runtime }}-test-results retention-days: 30 + if-no-files-found: ignore path: | + k8s-logs/ + report-auth-idp-${{ matrix.runtime }}.xml ${{ runner.temp }}/e2e-services-logs/ + - name: Delete kind cluster + if: always() && matrix.backend == 'kubernetes' + shell: bash + run: | + docker rm -f "cloud-provider-kind-${NMP_AUTHENTIK_K8S_CLUSTER_NAME}" || true + kind delete cluster --name "${NMP_AUTHENTIK_K8S_CLUSTER_NAME}" || true # Build wheels for all packages × python versions. Downstream jobs # (wheel-test, python-e2e-test) download these artifacts instead diff --git a/Makefile b/Makefile index 8dccf7c7b6..90c735ae52 100644 --- a/Makefile +++ b/Makefile @@ -122,10 +122,6 @@ docs-watch: ## Start Fern docs dev plus a repo-level watcher for docs/** changes docs-check: ## Validate the Fern docs (fern check + validate-mdx + gated-link check) cd docs/fern && npm run check -.PHONY: test-auth-idp -test-auth-idp: ## Run the auth-idp test suite - bash contrib/auth/authentik/run.sh test $(ARGS) - .PHONY: docs-check-python-snippets docs-check-python-snippets: ## Syntax-check and type-check Python snippets in one doc (DOCS_PATH=...) @if [ -z "$(strip $(DOCS_PATH))" ]; then echo "Usage: make docs-check-python-snippets DOCS_PATH=docs/customizer/tutorials/import-hf-model.mdx" >&2; exit 2; fi diff --git a/conftest.py b/conftest.py index 195ad9bac6..94b596944d 100644 --- a/conftest.py +++ b/conftest.py @@ -213,6 +213,8 @@ def pytest_collection_modifyitems(config, items): "unit", "e2e", "auth_idp", + "auth_idp_docker", + "auth_idp_k8s", "smoke_gpu_tasks", "smoke_nmp_customizer_tasks", "smoke_nmp_automodel_training", @@ -308,7 +310,9 @@ def pytest_runtest_setup(item): if not item.config.getoption("--run-slow"): skip_test("Skipping slow test (use --run-slow to run)") if "e2e" in [marker.name for marker in item.iter_markers()]: - if not item.config.getoption("--run-e2e"): + auth_idp_runtime = item.config.getoption("--auth-idp-runtime", default=None) + auth_idp_runtime_selected = auth_idp_runtime and "auth_idp" in [marker.name for marker in item.iter_markers()] + if not item.config.getoption("--run-e2e") and not auth_idp_runtime_selected: skip_test("Skipping e2e test (use --run-e2e to run)") if "subprocess_only" in [marker.name for marker in item.iter_markers()]: if os.environ.get("NMP_BASE_URL"): diff --git a/contrib/auth/authentik/.gitignore b/contrib/auth/authentik/.gitignore new file mode 100644 index 0000000000..e5d8ac351b --- /dev/null +++ b/contrib/auth/authentik/.gitignore @@ -0,0 +1,3 @@ +.generated/ +helm/Chart.lock +helm/charts/ diff --git a/contrib/auth/authentik/README.md b/contrib/auth/authentik/README.md index 1a66ea4f2f..2aa86ec0c4 100644 --- a/contrib/auth/authentik/README.md +++ b/contrib/auth/authentik/README.md @@ -1,263 +1,92 @@ # Authentik Reference Example This directory contains a local Authentik-backed NeMo Platform example. Use it -to verify three user-visible flows: +to validate three user-visible flows: - log in to NeMo with Authentik - call NeMo APIs through the Authentik gateway -- run a NeMo job whose workload uses a real Authentik workload token +- run a NeMo job whose workload exchanges a real Authentik workload subject token All credentials in this example are for local development only. -## Prerequisites +## Tutorial -- Docker with `docker compose` -- a bootstrapped NeMo Platform checkout -- a shell from the repo root +Use the shared tutorial when you want to understand or manually debug the +reference deployment: -## Demo Identities - -The stack seeds these local-only identities: - -- Human user: `nemo-user` -- Human password: `nemo-user-password-dev` -- Human email: `nemo-user@example.com` -- CLI OIDC client: `nemo-platform-cli` -- Workload identity: `svc-nemo` -- Workload group: `nemo-editors` - -## Start The Stack - -From the repo root: - -```bash -contrib/auth/authentik/run.sh stack -``` - -This starts NeMo, Authentik, and the local gateway with the existing default -NeMo API image, `my-registry/nmp-api:local`. The `stack` action does not build -images. -Leave this process running. Stop it with `Ctrl-C` when you are done; the script -removes the Compose stack and volumes on exit. - -To use a different prebuilt image for the example, pass it explicitly: - -```bash -export IMAGE_REGISTRY=registry.example.com/nemo -export BAKE_TAG= - -contrib/auth/authentik/run.sh stack --image "$IMAGE_REGISTRY/nmp-api:$BAKE_TAG" -``` - -Use the same `IMAGE_REGISTRY` and `BAKE_TAG` values in the shell where you -submit the workload job so the job container image matches the running NeMo API -image. - -The auth-idp test suite uses the same helper script: - -```bash -contrib/auth/authentik/run.sh test -``` - -For iteration or prebuilt images, pass options to the script directly. See -`contrib/auth/authentik/run.sh --help` for the full option list. - -```bash -contrib/auth/authentik/run.sh test --lifecycle reuse -contrib/auth/authentik/run.sh test --image registry.example.com/nemo/nmp-api: -``` - -Wait until the platform is ready through the gateway: +- [Authentik Reference Tutorial](tutorial.md) -```bash -until curl -sf http://127.0.0.1:18080/health/ready >/dev/null; do - sleep 2 -done -echo "NeMo Platform Ready" -``` +At the start, choose Docker Compose or Kubernetes. The rest of the tutorial uses +the same CLI and workload commands for both runtimes. Generated local material +lives under `.generated/`, including the workload-token signing key and gateway +TLS material when those files are created locally. -The local gateway URL is: +Runtime details live separately: -```text -http://127.0.0.1:18080 -``` +- [Compose Implementation Details](compose/implementation-details.md) +- [Kubernetes Implementation Details](kubernetes/implementation-details.md) -## Log In With Authentik +## Test Harness -Point the CLI at the Authentik gateway: +`run.sh` is the automation entrypoint for CI-style validation and repeatable +local test runs. It can run the local Compose stack, run the Compose auth-idp +contract tests, run the Kubernetes auth-idp contract tests, and clean up local +resources. ```bash -nemo config set --context authentik-human --base-url http://127.0.0.1:18080 --activate +contrib/auth/authentik/run.sh --help ``` -Start browser login: +Common commands: ```bash -nemo auth login --context authentik-human --base-url http://127.0.0.1:18080 +contrib/auth/authentik/run.sh compose +contrib/auth/authentik/run.sh k8s +contrib/auth/authentik/run.sh prepare-local +contrib/auth/authentik/run.sh run-local +contrib/auth/authentik/run.sh down ``` -Log in with: +The harness keeps generated local inputs in `contrib/auth/authentik/.generated` +so Compose and Kubernetes test runs can reuse the same workload-token signing +key. Diagnostics are written under `docker/logs/authentik-*` by default, or +under `E2E_SERVICES_LOG_DIR` when that environment variable is set. -- username: `nemo-user` -- password: `nemo-user-password-dev` - -Verify the saved session: - -```bash -nemo --context authentik-human auth status -nemo --context authentik-human workspaces list -``` - -Expected result: `auth status` shows `Auth Type: oauth`, the email -`nemo-user@example.com`, and a refresh token. `workspaces list` should return -without an auth error. - -## Create A Demo Workspace - -```bash -export WORKSPACE=authentik-demo - -nemo --context authentik-human workspaces create "$WORKSPACE" \ - --description "Authentik reference example" \ - --wait-role-propagation -``` - -Grant the demo workload group access to the workspace: - -```bash -nemo --context authentik-human workspaces members create \ - --workspace "$WORKSPACE" \ - --principal nemo-editors \ - --roles Viewer \ - --roles JobRunner \ - --wait-role-propagation -``` +For manual startup and walkthroughs, prefer the shared tutorial above. -Expected result: the human user can manage the workspace, and the workload -group can read the workspace from a job. - -## Run A Workload Job - -```bash -export JOB_NAME=authentik-workload-demo -``` - -Fetch a local demo token for the seeded workload identity: - -```bash -export WORKLOAD_ACCESS_TOKEN="$( - curl -fsS http://127.0.0.1:18080/application/o/token/ \ - -d grant_type=password \ - -d client_id=nemo-platform \ - -d client_secret=nemo-platform-secret-dev \ - -d username=svc-nemo \ - -d password=svc-nemo-token-secret-dev \ - -d scope="openid email groups" \ - | python -c 'import json, sys; print(json.load(sys.stdin)["access_token"])' -)" -``` - -Keep this token in a non-reserved shell variable. Do not export it as -`NEMO_WORKLOAD_TOKEN` in your shell. The NeMo CLI uses that variable as a -runtime credential override, which would make later CLI commands run as the -workload identity instead of `authentik-human`. - -Submit a job that runs the built-in hello-world workload auth task: - -```bash -export NMP_API_IMAGE="${NMP_API_IMAGE:-${IMAGE_REGISTRY:-my-registry}/nmp-api:${BAKE_TAG:-local}}" - -cat <` to the Helm +upgrade command. + +The chart creates or reuses these additional local-demo Secrets during Helm +rendering: + +- `shared-postgresql` for the shared PostgreSQL superuser, Authentik, and NeMo + database passwords. +- `shared-postgresql-nemo` for the NeMo Platform external database password. +- `nemo-platform-envoy-tls` for the demo Envoy TLS certificate and CA. +- `nemo-workload-token-signing-key` for the NeMo-issued workload access token + signing key. + +The chart generates `Secret/nemo-platform-envoy-tls` during Helm rendering and +reuses an existing Secret on upgrade. The Envoy TLS private key is never checked +into the repository and is not supplied through `values.yaml`. + +All credentials and generated keys in this example are for local development +only. + +## Authentik Blueprint + +The umbrella chart packages the shared blueprint from +`helm/files/blueprints/nemo.yaml` into `ConfigMap/authentik-nemo-blueprint` and +configures the official Authentik chart to mount it. A Helm +`post-install,post-upgrade` hook Job runs `ak apply_blueprint` against that +mounted file. + +Use `--wait --wait-for-jobs` when installing the chart so Helm only returns +after the blueprint has been applied. + +## NeMo Kubernetes Override + +The umbrella chart passes the Kubernetes-specific NeMo Platform configuration +through `nemo-platform.platformConfig` values. It also configures +`nemo-platform.envoyProxy.configOverride` so the NeMo Platform chart's Envoy +deployment keeps the Authentik path split and validates both Authentik-issued +tokens and NeMo workload-exchange tokens. + +Kubernetes projected service account token expiration defaults to `600` seconds +in the jobs backend. Override it through the NeMo Platform chart values if you +need a longer projected token lifetime. + +## Workload Token Exchange + +The Kubernetes jobs backend projects a Kubernetes service account token into +workload pods and injects: + +```text +NMP_WORKLOAD_IDENTITY_TOKEN_FILE=/var/run/secrets/nemo-platform/workload/token +``` + +The SDK reads that file and sends an RFC 8693 token exchange request to the NeMo +auth service over HTTPS. The chart mounts `ca.crt` from +`Secret/nemo-platform-envoy-tls` into Kubernetes workload pods and sets +`SSL_CERT_FILE` and `REQUESTS_CA_BUNDLE` so in-pod Python HTTP clients verify +the demo Envoy certificate. Host-side `nemo` commands should use +`NMP_CLIENT_SSL_CERT_FILE` instead so unrelated tools keep their normal trust +store. + +The NeMo auth service validates projected service account tokens with the +TokenReview API and returns a NeMo-signed JWT trusted by the NeMo Platform +Envoy. The useful end-to-end validation is the workload job in the tutorial: +the job pod uses the exchanged token to call the NeMo Platform API and read the +workspace. + +The workload job request should not include workload auth environment +variables. The Kubernetes jobs backend owns `NMP_WORKLOAD_IDENTITY_TOKEN_FILE`; +users must not set `NEMO_WORKLOAD_TOKEN` or `NEMO_WORKLOAD_TOKEN_FILE`. + +To inspect the projected token mount for a submitted job: + +```bash +kubectl --context "${KUBE_CONTEXT}" -n "${NAMESPACE}" get pods \ + -l "nmp.nvidia.com/job_id=${JOB_NAME}" + +kubectl --context "${KUBE_CONTEXT}" -n "${NAMESPACE}" describe pod \ + -l "nmp.nvidia.com/job_id=${JOB_NAME}" +``` + +## Workload Token Signing Key + +Workload identity token exchange requires the NeMo auth service to sign the +access token it mints from a Kubernetes projected service account subject +token. The chart creates `Secret/nemo-workload-token-signing-key` by default, +mounts `private-key.pem` into the NeMo Platform API pod at +`/etc/nmp/workload-token/private-key.pem`, and sets +`auth.oidc.workload_token_private_key_file` to that path. The matching public +key is served from `/apis/auth/jwks`; the NeMo Platform chart's Envoy deployment +uses that JWKS endpoint to validate exchanged workload tokens. + +The manual walkthrough can rely on the Helm chart to create and preserve this +Secret. + +If you need the same deterministic key for manual debugging, generate one and +add the `--set-file` line to the Helm upgrade command: + +```bash +mkdir -p contrib/auth/authentik/.generated +openssl genrsa -out contrib/auth/authentik/.generated/workload-token-private-key.pem 2048 +chmod 600 contrib/auth/authentik/.generated/workload-token-private-key.pem +--set-file workloadTokenSigningKey.privateKeyPem=contrib/auth/authentik/.generated/workload-token-private-key.pem +``` + +For production-style deployments, provide an externally managed RSA private-key +Secret instead of relying on the demo-generated key. Set +`workloadTokenSigningKey.create=false`, keep +`workloadTokenSigningKey.secretName` and +`nemo-platform.api.extraVolumes[].secret.secretName` aligned, and keep +`auth.oidc.workload_token_private_key_file` pointed at the mounted file path. diff --git a/contrib/auth/authentik/manifest.yaml b/contrib/auth/authentik/manifest.yaml index 69d7e713f2..2ac6588a85 100644 --- a/contrib/auth/authentik/manifest.yaml +++ b/contrib/auth/authentik/manifest.yaml @@ -1,53 +1,85 @@ provider: authentik mode: compose-ci -compose_file: docker-compose.yml -gateway_base_url: http://127.0.0.1:18080 +compose_file: compose/docker-compose.yml +gateway_base_url: https://127.0.0.1:18080 issuer_url: http://authentik-server:9000/application/o/nemo/ -discovery_url: http://127.0.0.1:18080/application/o/nemo/.well-known/openid-configuration +discovery_url: https://127.0.0.1:18080/application/o/nemo/.well-known/openid-configuration nemo_config: config/platform-compose-authentik.yaml principal_contract: subject_claim: sub groups_claim: groups external_machine_principals_use_service_prefix: false internal_service_prefix_reserved: "service:" -human_identity: +interactive_user_identity: username: nemo-user + password: nemo-user-password-dev expected_email: nemo-user@example.com workload_identity: principal_id: svc-nemo expected_groups: - - nemo-editors + - nemo-workloads workload_contract: audience: nemo-platform principal_claim: sub groups_claim: groups groups_format: comma_string token_env_vars: - - NEMO_WORKLOAD_TOKEN - - NEMO_WORKLOAD_TOKEN_FILE + - NMP_WORKLOAD_IDENTITY_TOKEN_FILE forwarded_headers: principal_id: X-NMP-Principal-Id principal_groups: X-NMP-Principal-Groups token_acquisition: - token_endpoint: http://127.0.0.1:18080/application/o/token/ - human_grant: + token_endpoint: https://127.0.0.1:18080/application/o/token/ + # E2E TEST HARNESS ONLY. This password grant gives auth-idp contract tests a + # setup principal for creating temporary workspaces and role grants. It is not + # a user login flow, not a workload identity flow, and must not be copied into + # production deployments. + e2e_setup_password_grant: grant_type: password client_id: nemo-platform client_secret: nemo-platform-secret-dev - username: nemo-user - password: nemo-user-token-secret-dev - scope: "openid profile email offline_access groups" - machine_grant: + username: nemo-setup + password: nemo-setup-token-secret-dev + scope: "openid email groups" + workload_provider_password_grant: grant_type: password - client_id: nemo-platform - client_secret: nemo-platform-secret-dev + client_id: nemo-platform-workload username: svc-nemo - password: svc-nemo-token-secret-dev + password_env_var: AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD scope: "openid email groups" healthchecks: - kind: http - url: http://127.0.0.1:18080/application/o/nemo/.well-known/openid-configuration + url: https://127.0.0.1:18080/application/o/nemo/.well-known/openid-configuration startup_timeouts: healthchecks_seconds: 600 gateway_seconds: 30 token_endpoint_seconds: 180 +test_runtimes: + - id: authentik-compose + backend: compose + command: compose + capabilities: + - gateway_discovery + - gateway_authn + - spoofed_header_rejection + - workload_provider_token + - workload_subject_token + - workload_token_exchange + - workspace_rbac + - workload_job + - device_flow + - docker_subject_token_refresh + - id: authentik-kubernetes + backend: kubernetes + command: k8s + capabilities: + - gateway_discovery + - gateway_authn + - spoofed_header_rejection + - workload_provider_token + - workload_subject_token + - workload_token_exchange + - workspace_rbac + - workload_job + - device_flow + - kubernetes_token_review diff --git a/contrib/auth/authentik/run.sh b/contrib/auth/authentik/run.sh index e87449c76c..8d0a181220 100755 --- a/contrib/auth/authentik/run.sh +++ b/contrib/auth/authentik/run.sh @@ -5,56 +5,151 @@ set -euo pipefail SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +AUTHENTIK_ROOT="${SCRIPT_DIR}" REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../../.." && pwd)" ACTION="" -COMPOSE_DIR="${SCRIPT_DIR}" +COMPOSE_DIR="${AUTHENTIK_ROOT}/compose" DRY_RUN="false" -IMAGE_SELECTED="$([[ -n "${IMAGE_REGISTRY:-}" || -n "${BAKE_TAG:-}" ]] && printf "true" || printf "false")" +IMAGE_SELECTED="false" IMAGE_REGISTRY="${IMAGE_REGISTRY:-my-registry}" BAKE_TAG="${BAKE_TAG:-local}" TEST_LIFECYCLE="fresh" -TEST_LIFECYCLE_SET="false" +REUSE_SET="false" TEST_PLATFORM="" TEST_PLATFORM_SET="false" TEST_DOCKER_TARGET="nmp-api-docker" COMPOSE_DIR_SET="false" +REUSE_COMPOSE_PROJECT_NAME="${NMP_AUTHENTIK_COMPOSE_PROJECT_NAME:-authentik-e2e-reuse}" +REUSE_COMPOSE_GATEWAY_PORT="${NMP_AUTHENTIK_COMPOSE_GATEWAY_PORT:-18083}" +REUSE_COMPOSE_GATEWAY_TLS_VOLUME="${NMP_AUTHENTIK_COMPOSE_GATEWAY_TLS_VOLUME:-authentik-e2e-${REUSE_COMPOSE_GATEWAY_PORT}-gateway-tls}" +REUSE_COMPOSE_WORKLOAD_NETWORK_NAME="${NMP_AUTHENTIK_COMPOSE_WORKLOAD_NETWORK_NAME:-authentik-e2e-${REUSE_COMPOSE_GATEWAY_PORT}-workload}" +REUSE_K8S_CLUSTER_NAME="${NMP_AUTHENTIK_K8S_REUSE_CLUSTER_NAME:-nmp-authentik-reuse}" +K8S_GATEWAY_PORT="${NMP_AUTHENTIK_K8S_GATEWAY_PORT:-18082}" +K8S_JUNIT_XML="${NMP_AUTHENTIK_K8S_JUNIT_XML:-report-auth-idp-kubernetes.xml}" +HELM_NAMESPACE="${HELM_NAMESPACE:-${NMP_AUTHENTIK_K8S_NAMESPACE:-nemo-authentik}}" +HELM_RELEASE="${HELM_RELEASE:-${NMP_AUTHENTIK_K8S_HELM_RELEASE:-authentik-demo}}" +K8S_CLUSTER_NAME="${NMP_AUTHENTIK_K8S_CLUSTER_NAME:-}" +K8S_RUNTIME="${NMP_AUTHENTIK_K8S_RUNTIME:-kind}" +K8S_RUNTIME_SET="false" +K8S_KEEP_CLUSTER="${NMP_AUTHENTIK_K8S_KEEP_CLUSTER:-0}" +K8S_REUSE_CLUSTER="${NMP_AUTHENTIK_K8S_REUSE_CLUSTER:-0}" +K8S_SKIP_IMAGE_LOAD="${NMP_AUTHENTIK_K8S_SKIP_IMAGE_LOAD:-0}" +K8S_SKIP_IMAGE_LOAD_SET="false" +K8S_NGC_EXISTING_SECRET="${NMP_AUTHENTIK_K8S_NGC_EXISTING_SECRET:-}" +K8S_IMAGE_PULL_SECRET="${NMP_AUTHENTIK_K8S_IMAGE_PULL_SECRET:-}" + +diagnostics_dir() { + local mode="$1" + local configured="${E2E_SERVICES_LOG_DIR:-}" + local timestamp + + if [[ -n "${configured}" ]]; then + if [[ "${configured}" = /* ]]; then + printf "%s" "${configured}" + else + printf "%s/%s" "${REPO_ROOT}" "${configured#./}" + fi + return + fi + + timestamp="$(date -u +%Y%m%dT%H%M%SZ)" + printf "%s/docker/logs/authentik-%s/%s" "${REPO_ROOT}" "${mode}" "${timestamp}" +} + +prepare_diagnostics_dir() { + local mode="$1" + local output + + output="$(diagnostics_dir "${mode}")" + if [[ "${DRY_RUN}" != "true" ]]; then + mkdir -p "${output}" + fi + printf "%s" "${output}" +} + +write_diagnostics_metadata() { + local mode="$1" + local output="$2" + + if [[ "${DRY_RUN}" == "true" ]]; then + return + fi + + { + printf "mode=%s\n" "${mode}" + printf "timestamp_utc=%s\n" "$(date -u +%Y-%m-%dT%H:%M:%SZ)" + printf "image=%s\n" "$(image_ref)" + printf "test_lifecycle=%s\n" "${TEST_LIFECYCLE}" + printf "k8s_runtime=%s\n" "${K8S_RUNTIME}" + printf "k8s_cluster_name=%s\n" "${K8S_CLUSTER_NAME}" + printf "k8s_gateway_port=%s\n" "${K8S_GATEWAY_PORT}" + printf "k8s_namespace=%s\n" "${HELM_NAMESPACE}" + printf "k8s_release=%s\n" "${HELM_RELEASE}" + } >"${output}/run-metadata.txt" +} usage() { cat <<'EOF' Usage: - contrib/auth/authentik/run.sh stack [options] + contrib/auth/authentik/run.sh run-local [options] contrib/auth/authentik/run.sh down [options] - contrib/auth/authentik/run.sh test [options] + contrib/auth/authentik/run.sh prepare-local [options] + contrib/auth/authentik/run.sh compose [options] + contrib/auth/authentik/run.sh render-blueprint [options] + contrib/auth/authentik/run.sh k8s [options] Runs the local Authentik reference example or its auth-idp test suite. Actions: - stack Start NeMo, Authentik, and the gateway in the foreground. - down Remove the example Compose stack and volumes. - test Build the local test image if needed, then run auth-idp tests. + run-local Start NeMo, Authentik, and the gateway in the foreground. + down Remove the example Compose stack, reused test Compose stack, + and reused Kubernetes test cluster. + prepare-local Create generated local inputs without starting Compose. + compose Build the local test image if needed, then run Compose auth-idp tests. + render-blueprint Copy the checked-in Authentik blueprint into generated inputs. + k8s Run the Helm-based kind/k3d Kubernetes E2E test. Image options: --image IMAGE Use an existing nmp-api image. Expected format: /nmp-api: Test options: - --lifecycle MODE Docker Compose lifecycle for tests: fresh or reuse. - Default: fresh. + --reuse Reuse a deterministic test environment. + For compose, use Compose project authentik-e2e-reuse + on gateway port 18083. + For k8s, use cluster nmp-authentik-reuse with + the selected runtime, creating it if needed and + keeping it after the run. + The k8s runner uses gateway port 18082 by default + to avoid the tutorial's 18081 port. Override with + NMP_AUTHENTIK_K8S_GATEWAY_PORT. --platform PLATFORM Platform for the default local test image build. + Applies to compose and k8s. Default: current machine architecture. + --runtime RUNTIME Kubernetes runtime for k8s/down: kind or k3d. + Default: kind. + --skip-image-load Do not load the nmp-api image into the reused cluster. + For a fresh cluster, use only with an explicit pullable + --image. Other options: - --compose-dir DIR Compose directory for stack/down. Default: this script's directory. + --compose-dir DIR Compose directory for run-local/down. Default: contrib/auth/authentik/compose. --dry-run Print commands without running them. -h, --help Show this help. Examples: - contrib/auth/authentik/run.sh stack - contrib/auth/authentik/run.sh stack --image my-registry/nmp-api:local - contrib/auth/authentik/run.sh test - contrib/auth/authentik/run.sh test --lifecycle reuse - contrib/auth/authentik/run.sh test --image my-registry/nmp-api:local + contrib/auth/authentik/run.sh run-local + contrib/auth/authentik/run.sh run-local --image my-registry/nmp-api:local + contrib/auth/authentik/run.sh prepare-local + contrib/auth/authentik/run.sh compose + contrib/auth/authentik/run.sh compose --reuse + contrib/auth/authentik/run.sh compose --image my-registry/nmp-api:local + contrib/auth/authentik/run.sh render-blueprint + contrib/auth/authentik/run.sh k8s + contrib/auth/authentik/run.sh k8s --runtime k3d + contrib/auth/authentik/run.sh k8s --reuse + contrib/auth/authentik/run.sh k8s --reuse --skip-image-load contrib/auth/authentik/run.sh down EOF } @@ -105,7 +200,17 @@ validate_test_lifecycle() { fresh | reuse) ;; *) - die "--lifecycle must be fresh or reuse" + die "test lifecycle must be fresh or reuse" + ;; + esac +} + +validate_k8s_runtime() { + case "${K8S_RUNTIME}" in + kind | k3d) + ;; + *) + die "--runtime must be kind or k3d" ;; esac } @@ -127,18 +232,201 @@ print_command_in_dir() { printf "\n" } -run_with_image_env_in_dir() { +print_command() { + printf "+ " + quote_args "$@" + printf "\n" +} + +blueprint_output_dir() { + local configured="${AUTHENTIK_BLUEPRINT_DIR:-./.generated/blueprints}" + if [[ "${configured}" = /* ]]; then + printf "%s" "${configured}" + else + printf "%s/%s" "${AUTHENTIK_ROOT}" "${configured#./}" + fi +} + +workload_token_private_key_file() { + printf "%s/.generated/workload-token-private-key.pem" "${AUTHENTIK_ROOT}" +} + +gateway_tls_dir() { + local configured="${AUTHENTIK_GATEWAY_TLS_DIR:-./.generated/gateway-tls}" + if [[ "${configured}" = /* ]]; then + printf "%s" "${configured}" + else + printf "%s/%s" "${AUTHENTIK_ROOT}" "${configured#./}" + fi +} + +gateway_tls_cert_file() { + printf "%s/tls.crt" "$(gateway_tls_dir)" +} + +gateway_tls_key_file() { + printf "%s/tls.key" "$(gateway_tls_dir)" +} + +ensure_gateway_tls_certificate() { + local output_dir + local cert + local key + local openssl_config + + output_dir="$(gateway_tls_dir)" + cert="$(gateway_tls_cert_file)" + key="$(gateway_tls_key_file)" + openssl_config="${output_dir}/openssl.cnf" + + if [[ -f "${cert}" && -f "${key}" ]]; then + echo "Using gateway TLS certificate: ${cert}" + return + fi + + if [[ "${DRY_RUN}" == "true" ]]; then + printf "+ mkdir -p %q\n" "${output_dir}" + printf "+ write %q\n" "${openssl_config}" + printf "+ openssl req -x509 -newkey rsa:2048 -nodes -keyout %q -out %q -days 365 -sha256 -config %q -extensions v3_req\n" "${key}" "${cert}" "${openssl_config}" + printf "+ chmod 600 %q\n" "${key}" + printf "+ chmod 644 %q\n" "${cert}" + return + fi + + if ! command -v openssl >/dev/null 2>&1; then + die "openssl is required to generate the gateway TLS certificate" + fi + + mkdir -p "${output_dir}" + cat >"${openssl_config}" <<'EOF' +[req] +prompt = no +distinguished_name = dn +x509_extensions = v3_req + +[dn] +CN = nemo-gateway + +[v3_req] +basicConstraints = critical, CA:TRUE +keyUsage = critical, digitalSignature, keyEncipherment, keyCertSign +extendedKeyUsage = serverAuth +subjectAltName = @alt_names + +[alt_names] +DNS.1 = localhost +DNS.2 = nemo-gateway +IP.1 = 127.0.0.1 +EOF + openssl req -x509 -newkey rsa:2048 -nodes -keyout "${key}" -out "${cert}" -days 365 -sha256 -config "${openssl_config}" -extensions v3_req >/dev/null 2>&1 + chmod 600 "${key}" + chmod 644 "${cert}" + echo "Generated gateway TLS certificate: ${cert}" +} + +ensure_workload_token_private_key() { + local output + local output_dir + + output="$(workload_token_private_key_file)" + output_dir="$(dirname -- "${output}")" + + if [[ -f "${output}" ]]; then + echo "Using workload token signing key: ${output}" + return + fi + + if [[ "${DRY_RUN}" == "true" ]]; then + printf "+ mkdir -p %q\n" "${output_dir}" + printf "+ openssl genrsa -out %q 2048\n" "${output}" + printf "+ chmod 600 %q\n" "${output}" + return + fi + + if ! command -v openssl >/dev/null 2>&1; then + die "openssl is required to generate the workload token signing key" + fi + + mkdir -p "${output_dir}" + openssl genrsa -out "${output}" 2048 >/dev/null 2>&1 + chmod 600 "${output}" + echo "Generated workload token signing key: ${output}" +} + +render_blueprint() { + local source="${AUTHENTIK_ROOT}/helm/files/blueprints/nemo.yaml" + local output_dir + local output + + output_dir="$(blueprint_output_dir)" + output="${output_dir}/nemo.yaml" + + if [[ ! -f "${source}" ]]; then + die "missing Authentik blueprint: ${source}" + fi + + if [[ "${DRY_RUN}" == "true" ]]; then + printf "+ mkdir -p %q\n" "${output_dir}" + printf "+ cp %q %q\n" "${source}" "${output}" + return + fi + + mkdir -p "${output_dir}" + cp "${source}" "${output}" + echo "Copied Authentik blueprint: ${output}" +} + +authentik_workload_identity_password() { + printf "%s" "${AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD:-svc-nemo-token-secret-e2e}" +} + +run_with_compose_env_in_dir() { local dir="$1" shift if [[ "${DRY_RUN}" == "true" ]]; then - printf "+ cd %q && IMAGE_REGISTRY=%q BAKE_TAG=%q " "${dir}" "${IMAGE_REGISTRY}" "${BAKE_TAG}" + printf "+ cd %q && IMAGE_REGISTRY=%q BAKE_TAG=%q AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD= " "${dir}" "${IMAGE_REGISTRY}" "${BAKE_TAG}" quote_args "$@" printf "\n" return fi - (cd "${dir}" && IMAGE_REGISTRY="${IMAGE_REGISTRY}" BAKE_TAG="${BAKE_TAG}" "$@") + ( + cd "${dir}" && + IMAGE_REGISTRY="${IMAGE_REGISTRY}" \ + BAKE_TAG="${BAKE_TAG}" \ + AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD="$(authentik_workload_identity_password)" \ + "$@" + ) +} + +run_with_reuse_compose_env_in_dir() { + local dir="$1" + shift + + if [[ "${DRY_RUN}" == "true" ]]; then + printf "+ cd %q && IMAGE_REGISTRY=%q BAKE_TAG=%q AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD= " "${dir}" "${IMAGE_REGISTRY}" "${BAKE_TAG}" + printf "COMPOSE_PROJECT_NAME=%q AUTHENTIK_GATEWAY_PORT=%q AUTHENTIK_GATEWAY_TLS_VOLUME=%q AUTHENTIK_WORKLOAD_NETWORK_NAME=%q " \ + "${REUSE_COMPOSE_PROJECT_NAME}" \ + "${REUSE_COMPOSE_GATEWAY_PORT}" \ + "${REUSE_COMPOSE_GATEWAY_TLS_VOLUME}" \ + "${REUSE_COMPOSE_WORKLOAD_NETWORK_NAME}" + quote_args "$@" + printf "\n" + return + fi + + ( + cd "${dir}" && + IMAGE_REGISTRY="${IMAGE_REGISTRY}" \ + BAKE_TAG="${BAKE_TAG}" \ + AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD="$(authentik_workload_identity_password)" \ + COMPOSE_PROJECT_NAME="${REUSE_COMPOSE_PROJECT_NAME}" \ + AUTHENTIK_GATEWAY_PORT="${REUSE_COMPOSE_GATEWAY_PORT}" \ + AUTHENTIK_GATEWAY_TLS_VOLUME="${REUSE_COMPOSE_GATEWAY_TLS_VOLUME}" \ + AUTHENTIK_WORKLOAD_NETWORK_NAME="${REUSE_COMPOSE_WORKLOAD_NETWORK_NAME}" \ + "$@" + ) } run_in_repo() { @@ -150,20 +438,74 @@ run_in_repo() { (cd "${REPO_ROOT}" && "$@") } -stack_up() { +prepare_local() { + render_blueprint + ensure_workload_token_private_key + ensure_gateway_tls_certificate +} + +run_local() { echo "Using existing NeMo API image: $(image_ref)" + prepare_local if [[ "${DRY_RUN}" == "true" ]]; then - run_with_image_env_in_dir "${COMPOSE_DIR}" docker compose up + run_with_compose_env_in_dir "${COMPOSE_DIR}" docker compose up return fi - trap 'run_with_image_env_in_dir "${COMPOSE_DIR}" docker compose down -v' EXIT INT TERM - run_with_image_env_in_dir "${COMPOSE_DIR}" docker compose up + trap 'run_with_compose_env_in_dir "${COMPOSE_DIR}" docker compose down -v' EXIT INT TERM + run_with_compose_env_in_dir "${COMPOSE_DIR}" docker compose up } compose_down() { - run_with_image_env_in_dir "${COMPOSE_DIR}" docker compose down -v + local status=0 + + run_with_compose_env_in_dir "${COMPOSE_DIR}" docker compose down -v --remove-orphans || status="$?" + run_with_reuse_compose_env_in_dir "${COMPOSE_DIR}" docker compose down -v --remove-orphans || status="$?" + return "${status}" +} + +delete_reuse_k8s_cluster() { + local cluster_name="${K8S_CLUSTER_NAME:-${REUSE_K8S_CLUSTER_NAME}}" + local -a delete_command + local status + + validate_k8s_runtime + + case "${K8S_RUNTIME}" in + kind) + delete_command=(kind delete cluster --name "${cluster_name}") + ;; + k3d) + delete_command=(k3d cluster delete "${cluster_name}") + ;; + esac + + if [[ "${DRY_RUN}" == "true" ]]; then + print_command "${delete_command[@]}" + return + fi + + if ! command -v "${delete_command[0]}" >/dev/null 2>&1; then + echo "Skipping reused Kubernetes test cluster cleanup: ${delete_command[0]} not found" + return + fi + + set +e + "${delete_command[@]}" + status="$?" + set -e + if [[ "${status}" -ne 0 ]]; then + echo "Warning: failed to delete reused Kubernetes test cluster ${cluster_name} with ${K8S_RUNTIME}" >&2 + fi +} + +down() { + local status=0 + + compose_down || status="$?" + delete_reuse_k8s_cluster || status="$?" + return "${status}" } build_default_test_image() { @@ -172,8 +514,46 @@ build_default_test_image() { run_in_repo make docker-load "DOCKER_TARGET=${TEST_DOCKER_TARGET}" "DOCKER_PLATFORMS=${platform}" } +run_pytest_with_diagnostics() { + local diagnostics="$1" + local status + shift + + if [[ "${DRY_RUN}" == "true" ]]; then + print_command_in_dir "${REPO_ROOT}" "$@" + return + fi + + set +e + ( + cd "${REPO_ROOT}" + set +e + "$@" 2>&1 | tee "${diagnostics}/pytest.log" + status="${PIPESTATUS[0]}" + set -e + exit "${status}" + ) + status="$?" + set -e + return "${status}" +} + run_tests() { + local workload_identity_password="${AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD:-svc-nemo-token-secret-e2e}" + local diagnostics + local compose_project_name="" + local compose_gateway_port="" + local status + validate_test_lifecycle + ensure_workload_token_private_key + if [[ "${TEST_LIFECYCLE}" == "reuse" ]]; then + compose_project_name="${REUSE_COMPOSE_PROJECT_NAME}" + compose_gateway_port="${REUSE_COMPOSE_GATEWAY_PORT}" + fi + diagnostics="$(prepare_diagnostics_dir compose)" + write_diagnostics_metadata compose "${diagnostics}" + echo "Auth-idp Compose diagnostics: ${diagnostics}" if [[ "${IMAGE_SELECTED}" == "true" ]]; then echo "Using prebuilt auth-idp test image: $(image_ref)" @@ -181,9 +561,101 @@ run_tests() { build_default_test_image fi - run_in_repo \ - env "IMAGE_REGISTRY=${IMAGE_REGISTRY}" "BAKE_TAG=${BAKE_TAG}" "NMP_E2E_COMPOSE_LIFECYCLE=${TEST_LIFECYCLE}" \ - uv run --frozen pytest tests/auth_idp -v --run-e2e + if [[ "${DRY_RUN}" == "true" ]]; then + run_pytest_with_diagnostics "${diagnostics}" \ + env "IMAGE_REGISTRY=${IMAGE_REGISTRY}" "BAKE_TAG=${BAKE_TAG}" \ + "AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD=" \ + "E2E_SERVICES_LOG_DIR=${diagnostics}" \ + "NMP_E2E_COMPOSE_LIFECYCLE=${TEST_LIFECYCLE}" \ + "NMP_AUTHENTIK_COMPOSE_PROJECT_NAME=${compose_project_name}" \ + "NMP_AUTHENTIK_COMPOSE_GATEWAY_PORT=${compose_gateway_port}" \ + "NMP_CLIENT_SSL_CERT_FILE=$(gateway_tls_cert_file)" \ + uv run --frozen pytest tests/auth_idp/contracts -v --auth-idp-runtime authentik-compose -m auth_idp_runtime + return + fi + + echo "Writing Authentik Compose diagnostics to: ${diagnostics}" + run_pytest_with_diagnostics "${diagnostics}" \ + env "IMAGE_REGISTRY=${IMAGE_REGISTRY}" \ + "BAKE_TAG=${BAKE_TAG}" \ + "AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD=${workload_identity_password}" \ + "E2E_SERVICES_LOG_DIR=${diagnostics}" \ + "NMP_E2E_COMPOSE_LIFECYCLE=${TEST_LIFECYCLE}" \ + "NMP_AUTHENTIK_COMPOSE_PROJECT_NAME=${compose_project_name}" \ + "NMP_AUTHENTIK_COMPOSE_GATEWAY_PORT=${compose_gateway_port}" \ + "NMP_CLIENT_SSL_CERT_FILE=$(gateway_tls_cert_file)" \ + uv run --frozen pytest tests/auth_idp/contracts -v --auth-idp-runtime authentik-compose -m auth_idp_runtime + status="$?" + echo "Auth-idp Compose diagnostics: ${diagnostics}" + return "${status}" +} + +run_k8s_tests() { + local diagnostics + local k8s_diagnostics + local workload_token_private_key + local status + + validate_k8s_runtime + ensure_workload_token_private_key + workload_token_private_key="$(workload_token_private_key_file)" + diagnostics="$(prepare_diagnostics_dir kubernetes)" + write_diagnostics_metadata kubernetes "${diagnostics}" + echo "Auth-idp Kubernetes diagnostics: ${diagnostics}" + + if [[ "${IMAGE_SELECTED}" == "true" ]]; then + echo "Using prebuilt auth-idp Kubernetes test image: $(image_ref)" + elif [[ "${K8S_REUSE_CLUSTER}" == "1" && "${K8S_SKIP_IMAGE_LOAD}" == "1" ]]; then + echo "Reusing Kubernetes cluster without rebuilding or loading image: $(image_ref)" + else + build_default_test_image + fi + + k8s_diagnostics="${diagnostics}/kubernetes" + if [[ "${DRY_RUN}" != "true" ]]; then + mkdir -p "${k8s_diagnostics}" + fi + + if [[ "${DRY_RUN}" == "true" ]]; then + run_pytest_with_diagnostics "${diagnostics}" \ + env "IMAGE_REGISTRY=${IMAGE_REGISTRY}" "BAKE_TAG=${BAKE_TAG}" \ + "E2E_SERVICES_LOG_DIR=${diagnostics}" \ + "NMP_AUTHENTIK_K8S_LOG_DIR=${k8s_diagnostics}" \ + "NMP_AUTHENTIK_K8S_HELM_RELEASE=${HELM_RELEASE}" \ + "NMP_AUTHENTIK_K8S_NAMESPACE=${HELM_NAMESPACE}" \ + "NMP_AUTHENTIK_K8S_RUNTIME=${K8S_RUNTIME}" \ + "NMP_AUTHENTIK_K8S_CLUSTER_NAME=${K8S_CLUSTER_NAME}" \ + "NMP_AUTHENTIK_K8S_GATEWAY_PORT=${K8S_GATEWAY_PORT}" \ + "NMP_AUTHENTIK_K8S_KEEP_CLUSTER=${K8S_KEEP_CLUSTER}" \ + "NMP_AUTHENTIK_K8S_REUSE_CLUSTER=${K8S_REUSE_CLUSTER}" \ + "NMP_AUTHENTIK_K8S_SKIP_IMAGE_LOAD=${K8S_SKIP_IMAGE_LOAD}" \ + "NMP_AUTHENTIK_K8S_NGC_EXISTING_SECRET=${K8S_NGC_EXISTING_SECRET}" \ + "NMP_AUTHENTIK_K8S_IMAGE_PULL_SECRET=${K8S_IMAGE_PULL_SECRET}" \ + "NMP_AUTHENTIK_K8S_WORKLOAD_TOKEN_PRIVATE_KEY_FILE=${workload_token_private_key}" \ + uv run --frozen pytest tests/auth_idp/contracts -v --auth-idp-runtime authentik-kubernetes -m auth_idp_runtime --junitxml="${K8S_JUNIT_XML}" + return + fi + + echo "Writing Authentik Kubernetes diagnostics to: ${diagnostics}" + run_pytest_with_diagnostics "${diagnostics}" \ + env "IMAGE_REGISTRY=${IMAGE_REGISTRY}" "BAKE_TAG=${BAKE_TAG}" \ + "E2E_SERVICES_LOG_DIR=${diagnostics}" \ + "NMP_AUTHENTIK_K8S_LOG_DIR=${k8s_diagnostics}" \ + "NMP_AUTHENTIK_K8S_HELM_RELEASE=${HELM_RELEASE}" \ + "NMP_AUTHENTIK_K8S_NAMESPACE=${HELM_NAMESPACE}" \ + "NMP_AUTHENTIK_K8S_RUNTIME=${K8S_RUNTIME}" \ + "NMP_AUTHENTIK_K8S_CLUSTER_NAME=${K8S_CLUSTER_NAME}" \ + "NMP_AUTHENTIK_K8S_GATEWAY_PORT=${K8S_GATEWAY_PORT}" \ + "NMP_AUTHENTIK_K8S_KEEP_CLUSTER=${K8S_KEEP_CLUSTER}" \ + "NMP_AUTHENTIK_K8S_REUSE_CLUSTER=${K8S_REUSE_CLUSTER}" \ + "NMP_AUTHENTIK_K8S_SKIP_IMAGE_LOAD=${K8S_SKIP_IMAGE_LOAD}" \ + "NMP_AUTHENTIK_K8S_NGC_EXISTING_SECRET=${K8S_NGC_EXISTING_SECRET}" \ + "NMP_AUTHENTIK_K8S_IMAGE_PULL_SECRET=${K8S_IMAGE_PULL_SECRET}" \ + "NMP_AUTHENTIK_K8S_WORKLOAD_TOKEN_PRIVATE_KEY_FILE=${workload_token_private_key}" \ + uv run --frozen pytest tests/auth_idp/contracts -v --auth-idp-runtime authentik-kubernetes -m auth_idp_runtime --junitxml="${K8S_JUNIT_XML}" + status="$?" + echo "Auth-idp Kubernetes diagnostics: ${diagnostics}" + return "${status}" } if [[ $# -eq 0 ]]; then @@ -193,7 +665,7 @@ fi while [[ $# -gt 0 ]]; do case "$1" in - stack | down | test) + run-local | down | prepare-local | compose | render-blueprint | k8s) if [[ -n "${ACTION}" ]]; then die "only one action can be specified" fi @@ -205,18 +677,32 @@ while [[ $# -gt 0 ]]; do parse_image "$2" shift 2 ;; - --lifecycle) - [[ $# -ge 2 ]] || die "--lifecycle requires a value" - TEST_LIFECYCLE="$2" - TEST_LIFECYCLE_SET="true" - shift 2 - ;; --platform) [[ $# -ge 2 ]] || die "--platform requires a value" TEST_PLATFORM="$2" TEST_PLATFORM_SET="true" shift 2 ;; + --runtime) + [[ $# -ge 2 ]] || die "--runtime requires a value" + K8S_RUNTIME="$2" + K8S_RUNTIME_SET="true" + shift 2 + ;; + --runtime=*) + K8S_RUNTIME="${1#*=}" + K8S_RUNTIME_SET="true" + shift + ;; + --reuse) + REUSE_SET="true" + shift + ;; + --skip-image-load) + K8S_SKIP_IMAGE_LOAD="1" + K8S_SKIP_IMAGE_LOAD_SET="true" + shift + ;; --compose-dir) [[ $# -ge 2 ]] || die "--compose-dir requires a value" COMPOSE_DIR="$2" @@ -241,31 +727,79 @@ if [[ -z "${ACTION}" ]]; then die "missing action" fi +if [[ "${REUSE_SET}" == "true" ]]; then + case "${ACTION}" in + compose) + TEST_LIFECYCLE="reuse" + ;; + k8s) + K8S_REUSE_CLUSTER="1" + K8S_KEEP_CLUSTER="1" + if [[ -z "${K8S_CLUSTER_NAME}" ]]; then + K8S_CLUSTER_NAME="${REUSE_K8S_CLUSTER_NAME}" + fi + ;; + *) + die "--reuse is only valid with compose or k8s" + ;; + esac +fi + +if [[ "${ACTION}" == "k8s" && "${K8S_REUSE_CLUSTER}" == "1" && -z "${K8S_CLUSTER_NAME}" ]]; then + K8S_CLUSTER_NAME="${REUSE_K8S_CLUSTER_NAME}" +fi + if [[ -z "${IMAGE_REGISTRY}" || -z "${BAKE_TAG}" ]]; then die "image registry and tag must be non-empty" fi -if [[ "${ACTION}" != "test" ]]; then - if [[ "${TEST_LIFECYCLE_SET}" == "true" ]]; then - die "--lifecycle is only valid with the test action" - fi +if [[ "${ACTION}" != "compose" && "${ACTION}" != "k8s" ]]; then if [[ "${TEST_PLATFORM_SET}" == "true" ]]; then - die "--platform is only valid with the test action" + die "--platform is only valid with the compose or k8s action" fi fi -if [[ "${ACTION}" == "test" && "${COMPOSE_DIR_SET}" == "true" ]]; then - die "--compose-dir is only valid with stack or down" +if [[ "${ACTION}" != "k8s" && "${ACTION}" != "down" ]]; then + if [[ "${K8S_RUNTIME_SET}" == "true" ]]; then + die "--runtime is only valid with k8s or down" + fi +fi + +if [[ "${ACTION}" != "k8s" ]]; then + if [[ "${K8S_SKIP_IMAGE_LOAD_SET}" == "true" ]]; then + die "--skip-image-load is only valid with k8s" + fi +fi + +if [[ "${ACTION}" == "k8s" && + "${K8S_SKIP_IMAGE_LOAD}" == "1" && + "${K8S_REUSE_CLUSTER}" != "1" && + "${IMAGE_SELECTED}" != "true" ]]; then + die "--skip-image-load with a fresh k8s cluster requires a pullable image selected with" \ + "--image, or a reused cluster via --reuse" +fi + +if [[ "${ACTION}" != "run-local" && "${ACTION}" != "down" && "${COMPOSE_DIR_SET}" == "true" ]]; then + die "--compose-dir is only valid with run-local or down" fi case "${ACTION}" in - stack) - stack_up + run-local) + run_local ;; down) - compose_down + down + ;; + prepare-local) + prepare_local ;; - test) + compose) run_tests ;; + render-blueprint) + render_blueprint + ;; + k8s) + run_k8s_tests + ;; esac diff --git a/contrib/auth/authentik/tutorial.md b/contrib/auth/authentik/tutorial.md new file mode 100644 index 0000000000..ad28fd67fc --- /dev/null +++ b/contrib/auth/authentik/tutorial.md @@ -0,0 +1,395 @@ +# Authentik Reference Tutorial + +This tutorial validates the Authentik reference deployment against one runtime: +Docker Compose or Kubernetes. Choose a runtime in the first section, then run +the remaining sections exactly the same way for both. + +The tutorial covers: + +- NeMo CLI login through Authentik. +- NeMo API calls through the Authentik gateway. +- Workload identity token exchange through a workload job. + +For shared identities, token lifetimes, and automated test harness commands, +see [the top-level Authentik README](README.md). For runtime internals, see the +[Compose details](compose/implementation-details.md) or +[Kubernetes details](kubernetes/implementation-details.md). + +All credentials in this example are for local development only. + +## Choose A Runtime + +Use either Docker Compose or Kubernetes. After the chosen runtime is running, +continue with [Wait For The Gateway](#wait-for-the-gateway). + +From the repo root, prepare shared generated inputs once: + +```bash +contrib/auth/authentik/run.sh prepare-local +``` + +This creates the shared workload-token signing key, gateway TLS material, and +rendered Authentik blueprint under `contrib/auth/authentik/.generated`. + +### Docker Compose + +Prerequisites: + +- Docker with `docker compose` +- `openssl` +- `curl` +- a bootstrapped NeMo Platform checkout +- a shell from the repo root + +Start Compose in one terminal: + +```bash +docker compose -f contrib/auth/authentik/compose/docker-compose.yml up +``` + +This starts NeMo, Authentik, and the local gateway with the default NeMo API +image, `my-registry/nmp-api:local`. The Compose stack does not build images. + +In another terminal from the repo root, export the runtime variables used by the +rest of the tutorial: + +```bash +export AUTHENTIK_RUNTIME=compose +export AUTHENTIK_CONTEXT=authentik-compose +export AUTHENTIK_BASE_URL=https://127.0.0.1:18080 +export AUTHENTIK_GATEWAY_CA=contrib/auth/authentik/.generated/gateway-tls/tls.crt +export NMP_CLIENT_SSL_CERT_FILE="$AUTHENTIK_GATEWAY_CA" +export AUTHENTIK_WORKLOAD_GROUP=nemo-workloads +export WORKSPACE=authentik-demo +export JOB_NAME=authentik-workload-demo +export IMAGE_REGISTRY="${IMAGE_REGISTRY:-my-registry}" +export BAKE_TAG="${BAKE_TAG:-local}" +export NMP_API_IMAGE="${NMP_API_IMAGE:-${IMAGE_REGISTRY}/nmp-api:${BAKE_TAG}}" +``` + +`AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD` has a local-development default. If you +override it, keep the same value until you remove the Compose volumes with +`docker compose down -v`. + +### Kubernetes + +Prerequisites: + +- `helm` +- `kubectl` +- `kind` +- Docker +- `curl` +- outbound image pull access for the third-party Authentik, Envoy, and + PostgreSQL images +- a NeMo Platform service image loaded into the kind cluster, or pushed to a + registry the cluster can pull + +From the repo root: + +```bash +export AUTHENTIK_RUNTIME=kubernetes +export KIND_CLUSTER=nmp-authentik-dev +export KUBE_CONTEXT="kind-${KIND_CLUSTER}" +export NAMESPACE=nemo-authentik +export HELM_RELEASE=authentik-demo +export IMAGE_REGISTRY="${IMAGE_REGISTRY:-my-registry}" +export BAKE_TAG="${BAKE_TAG:-local}" +export NMP_API_IMAGE="${IMAGE_REGISTRY}/nmp-api:${BAKE_TAG}" +export NEMO_AUTHENTIK_TMP_DIR="${TMPDIR:-/tmp}/nemo-authentik" +export KUBECONFIG="${NEMO_AUTHENTIK_TMP_DIR}/kubeconfig.yaml" + +mkdir -p "${NEMO_AUTHENTIK_TMP_DIR}" + +if kind get clusters | grep -qx "${KIND_CLUSTER}"; then + kind export kubeconfig --name "${KIND_CLUSTER}" --kubeconfig "${KUBECONFIG}" +else + kind create cluster --name "${KIND_CLUSTER}" --kubeconfig "${KUBECONFIG}" +fi + +kubectl --context "${KUBE_CONTEXT}" create namespace "${NAMESPACE}" \ + --dry-run=client -o yaml | \ + kubectl --context "${KUBE_CONTEXT}" apply -f - +``` + +Build the local NeMo Platform image and load it into kind: + +```bash +make docker-load DOCKER_TARGET=nmp-api-docker +docker image inspect "${NMP_API_IMAGE}" >/dev/null +kind load docker-image "${NMP_API_IMAGE}" --name "${KIND_CLUSTER}" +``` + +Install the chart: + +```bash +helm repo add nvidia https://helm.ngc.nvidia.com/nvidia --force-update +helm repo add authentik https://charts.goauthentik.io --force-update +helm repo update + +helm dependency build k8s/helm +helm dependency build contrib/auth/authentik/helm +helm --kube-context "${KUBE_CONTEXT}" upgrade --install "${HELM_RELEASE}" contrib/auth/authentik/helm \ + --namespace "${NAMESPACE}" \ + --create-namespace \ + --wait \ + --wait-for-jobs \ + --timeout 10m \ + --set-string nemo-platform.api.image.repository="${IMAGE_REGISTRY}/nmp-api" \ + --set-string nemo-platform.api.image.tag="${BAKE_TAG}" \ + --set-string nemo-platform.core.image.repository="${IMAGE_REGISTRY}/nmp-api" \ + --set-string nemo-platform.core.image.tag="${BAKE_TAG}" \ + --set-string nemo-platform.platformConfig.platform.image_registry="${IMAGE_REGISTRY}" \ + --set-string nemo-platform.platformConfig.platform.image_tag="${BAKE_TAG}" \ + --set-file workloadTokenSigningKey.privateKeyPem=contrib/auth/authentik/.generated/workload-token-private-key.pem +``` + +Wait for the main workloads: + +```bash +kubectl --context "${KUBE_CONTEXT}" -n "${NAMESPACE}" rollout status statefulset/shared-postgresql +kubectl --context "${KUBE_CONTEXT}" -n "${NAMESPACE}" rollout status deploy/authentik-server +kubectl --context "${KUBE_CONTEXT}" -n "${NAMESPACE}" rollout status deploy/authentik-worker +kubectl --context "${KUBE_CONTEXT}" -n "${NAMESPACE}" rollout status deploy/nemo-platform-api +kubectl --context "${KUBE_CONTEXT}" -n "${NAMESPACE}" rollout status deploy/nemo-platform-envoy +``` + +Port-forward the NeMo Platform Envoy service in a separate terminal with the +same `KUBECONFIG`, `KUBE_CONTEXT`, and `NAMESPACE` exports: + +```bash +kubectl --context "${KUBE_CONTEXT}" -n "${NAMESPACE}" port-forward svc/nemo-platform-envoy 18081:8080 +``` + +In the original terminal, export the demo CA and runtime variables used by the +rest of the tutorial: + +```bash +kubectl --context "${KUBE_CONTEXT}" -n "${NAMESPACE}" get secret nemo-platform-envoy-tls \ + -o jsonpath='{.data.ca\.crt}' | base64 -d \ + > "${NEMO_AUTHENTIK_TMP_DIR}/ca.crt" + +touch "${NEMO_AUTHENTIK_TMP_DIR}/config.yaml" + +export AUTHENTIK_CONTEXT=authentik-k8s +export AUTHENTIK_BASE_URL=https://127.0.0.1:18081 +export AUTHENTIK_GATEWAY_CA="${NEMO_AUTHENTIK_TMP_DIR}/ca.crt" +export NMP_CLIENT_SSL_CERT_FILE="$AUTHENTIK_GATEWAY_CA" +export NMP_CONFIG_FILE="${NEMO_AUTHENTIK_TMP_DIR}/config.yaml" +export AUTHENTIK_WORKLOAD_GROUP="system:serviceaccounts:${NAMESPACE}" +export WORKSPACE=authentik-demo +export JOB_NAME=authentik-workload-demo +``` + +## Wait For The Gateway + +From this point on, the commands are the same for Compose and Kubernetes. + +```bash +until curl --cacert "$AUTHENTIK_GATEWAY_CA" -sf "${AUTHENTIK_BASE_URL}/health/gateway/ready" >/dev/null; do + sleep 2 +done +echo "NeMo Platform and Authentik Ready" +``` + +## Log In With Authentik + +Start the Authentik device-code login: + +```bash +uv run nemo auth login \ + --context "$AUTHENTIK_CONTEXT" \ + --base-url "$AUTHENTIK_BASE_URL" +``` + +If a browser opens, log in with: + +- username: `nemo-user` +- password: `nemo-user-password-dev` + +If the browser does not open automatically, the CLI prints a URL and code. Open +the URL promptly, enter the code, and log in with the same demo credentials. + +Verify the saved session: + +```bash +uv run nemo --context "$AUTHENTIK_CONTEXT" auth status +``` + +Expected result: `auth status` shows `Auth Type: oauth`, the email +`nemo-user@example.com`, a refresh token, and an access token. + +Wait for the short-lived access token to get close to expiry, then run a normal +authenticated command: + +```bash +sleep 70 +uv run nemo --context "$AUTHENTIK_CONTEXT" workspaces list +uv run nemo --context "$AUTHENTIK_CONTEXT" auth status +``` + +Expected result: `workspaces list` returns without an auth error. The CLI uses +the saved refresh token before the request and might print +`[Auto-refreshed expired token]` before the workspace output. + +## Create A Demo Workspace + +Create the workspace: + +```bash +uv run nemo --context "$AUTHENTIK_CONTEXT" workspaces create "$WORKSPACE" \ + --description "Authentik reference example (${AUTHENTIK_RUNTIME})" \ + --wait-role-propagation +``` + +Grant the demo human Authentik group access: + +```bash +uv run nemo --context "$AUTHENTIK_CONTEXT" workspaces members create \ + --workspace "$WORKSPACE" \ + --principal nemo-editors \ + --roles Viewer \ + --wait-role-propagation +``` + +Grant the workload identity group read access and permission to upload workload +logs. In Compose this is the dedicated `nemo-workloads` Authentik group; in +Kubernetes this is the projected service-account group: + +```bash +uv run nemo --context "$AUTHENTIK_CONTEXT" workspaces members create \ + --workspace "$WORKSPACE" \ + --principal "$AUTHENTIK_WORKLOAD_GROUP" \ + --roles Viewer \ + --roles JobRunner \ + --wait-role-propagation +``` + +Expected result: the human user can manage the workspace, and the workload +identity can read the workspace from a job and upload the job logs. + +## Run A Workload Job + +Submit a workload job that reads the workspace through the public SDK: + +```bash +cat <- + E2E test harness only. Do not use this grant as a production user, + CLI, or workload authentication pattern. + required: + - grant_type + - client_id + - username + - scope + oneOf: + - required: + - password + - required: + - password_env_var + properties: + grant_type: + type: string + client_id: + type: string + client_secret: + type: string + username: + type: string + password: + type: string + password_env_var: + type: string + scope: + type: string + interactive_user_password_grant: + type: object + required: + - grant_type + - client_id + - username + - scope + oneOf: + - required: + - password + - required: + - password_env_var + properties: + grant_type: + type: string + client_id: + type: string + client_secret: + type: string + username: + type: string + password: + type: string + password_env_var: + type: string + scope: + type: string + workload_provider_password_grant: + type: object + required: + - grant_type + - client_id + - username + - scope + oneOf: + - required: + - password + - required: + - password_env_var + properties: + grant_type: + type: string + client_id: + type: string + client_secret: + type: string + username: + type: string + password: + type: string + password_env_var: + type: string + scope: + type: string healthchecks: type: array startup_timeouts: @@ -64,3 +167,28 @@ properties: - healthchecks_seconds - gateway_seconds - token_endpoint_seconds + test_runtimes: + type: array + minItems: 1 + items: + type: object + required: + - id + - backend + - capabilities + properties: + id: + type: string + backend: + type: string + enum: + - compose + - kubernetes + - external + command: + type: string + capabilities: + type: array + minItems: 1 + items: + type: string diff --git a/docs/auth/authentication/idp-integration.mdx b/docs/auth/authentication/idp-integration.mdx index 8e017a165e..7242ab47b9 100644 --- a/docs/auth/authentication/idp-integration.mdx +++ b/docs/auth/authentication/idp-integration.mdx @@ -21,7 +21,10 @@ the first-pass workload mapping: - `sub` -> `X-NMP-Principal-Id` - `groups` -> `X-NMP-Principal-Groups` -The demo workload token contract is: +The demo workload identity contract is: -- `NEMO_WORKLOAD_TOKEN` -- `NEMO_WORKLOAD_TOKEN_FILE` +- `NMP_WORKLOAD_IDENTITY_TOKEN_FILE` + +Managed job backends inject this variable and mount the token file. Job specs +must not set `NEMO_WORKLOAD_TOKEN`, `NEMO_WORKLOAD_TOKEN_FILE`, or +`NMP_WORKLOAD_IDENTITY_TOKEN_FILE` themselves. diff --git a/docs/auth/deployment/configuration.mdx b/docs/auth/deployment/configuration.mdx index 19373f94c6..d8e7dd5a73 100644 --- a/docs/auth/deployment/configuration.mdx +++ b/docs/auth/deployment/configuration.mdx @@ -82,6 +82,69 @@ NMP_AUTH_EMBEDDED_PDP_AUTO_BUILD_WASM=true Nested keys (e.g., OIDC) use double underscore: `NMP_AUTH_OIDC__ISSUER`, `NMP_AUTH_OIDC__CLIENT_ID`. +## OIDC Workload Identity Exchange + +`auth.oidc` can also advertise SDK workload identity token exchange metadata: + +```yaml +auth: + oidc: + enabled: true + token_endpoint: "https://idp.example.com/oauth/token" + workload_token_exchange_enabled: true + workload_client_id: "nemo-platform-workload" + workload_token_endpoint: "https://idp.example.com/oauth/token" + workload_audience: "nemo-platform" + workload_scope: "openid email groups" +``` + +When `NMP_WORKLOAD_IDENTITY_TOKEN_FILE` is present, the SDK reads the subject +token from that file and sends an RFC 8693 exchange request using +`subject_token`, fixed JWT subject/access-token token types, optional `audience`, +and optional `scope`. `workload_token_endpoint` is optional; when it is unset, +the SDK uses `token_endpoint`. This is useful when host CLI login and workload +containers need different network-reachable IdP URLs. + +For Docker-backed job runtimes, the executor may also need a controller-side +subject-token issuer. Configure only Docker-specific fields under the Docker +executor profile: + +```yaml +jobs: + executors: + - provider: cpu + profile: workload + backend: docker + config: + workload_identity: + token_endpoint: "https://idp.example.com/oauth/token" + client_id: "nemo-platform-workload" + username: "svc-nemo" + password_env_var: "WORKLOAD_IDENTITY_PASSWORD" + scope: "openid email groups" +``` + +The password value is read from the controller process environment using +`password_env_var`. The config does not support an inline `password` field. + +For Kubernetes-backed job runtimes, the jobs backend can project a Kubernetes +service account token into each workload pod for token exchange. The projected +token expiration defaults to `600` seconds and must be at least `600` seconds. +In platform configuration, set +`jobs.executors[].config.workload_identity_token_expiration_seconds` on the +Kubernetes job executor profile. For example, this changes the projected token +lifetime from the 600-second default to one hour: + +```yaml +jobs: + executors: + - provider: cpu + profile: workload + backend: kubernetes_job + config: + workload_identity_token_expiration_seconds: 3600 +``` + ## Example Configurations ### Quickstart / development (auth disabled) diff --git a/docs/auth/deployment/credential-propagation.mdx b/docs/auth/deployment/credential-propagation.mdx index 1b76d1c847..afd353f360 100644 --- a/docs/auth/deployment/credential-propagation.mdx +++ b/docs/auth/deployment/credential-propagation.mdx @@ -4,18 +4,65 @@ description: "" --- How credentials flow through the system when NeMo Platform runs jobs or serves inference. -## Job Credential Propagation +## Managed Job Workload Identity -When a user submits a job (customization, evaluation, data generation), the job runs in a Kubernetes pod that needs to call NeMo Platform APIs — to download datasets, upload results, and read secrets. The platform propagates the submitting user's identity into the job container so it operates with the user's permissions, not elevated service credentials. +When a user submits a job (customization, evaluation, data generation), the job +runs in a workload container that needs to call NeMo Platform APIs to download +datasets, upload results, and read secrets. Managed job credentials are +propagated through workload identity token exchange, not by using a serialized +principal as an API credential. The flow: -1. The user submits a job via the API. The platform records the user's principal identity. -2. The platform creates a Kubernetes job with the `NMP_PRINCIPAL` environment variable set to the submitting user's identity. -3. Secrets needed by the job are fetched on behalf of the user (the platform checks that the user has access before injecting them). -4. The job container uses the propagated principal to authenticate API calls back to NeMo Platform. +1. The user submits a job via the API. The control plane authorizes job creation + and records the job metadata. +2. When workload identity token exchange is enabled, the managed backend creates + a workload with `NMP_WORKLOAD_IDENTITY_TOKEN_FILE` set to a subject-token file + path. +3. The backend owns that subject-token file. Kubernetes uses a projected service + account token volume, while Docker uses a controller-managed volume that is + refreshed by the jobs controller. +4. The SDK in the workload detects `NMP_WORKLOAD_IDENTITY_TOKEN_FILE`, reads the + subject token, discovers the workload token exchange endpoint, and exchanges + the subject token for a NeMo Platform access token using OAuth 2.0 Token + Exchange (RFC 8693). +5. API calls from the workload use the exchanged access token. When that access + token nears expiry, the SDK rereads the subject-token file and performs + another exchange. -Job containers need to run inside the trust boundary so that their `X-NMP-Principal-*` headers are accepted by downstream services. Network policies and gateway configuration enforce this boundary. For the full architecture, see [Security Model](/documentation/access-control/security-model#job-credential-propagation). +Job containers still need to run inside the platform trust boundary. Network +policies and gateway configuration should prevent workloads from reaching +internal endpoints they do not need and should prevent external callers from +forging trusted identity headers. For the full architecture, see [Security +Model](/documentation/access-control/security-model#job-credential-propagation). + +## Workload Identity Token Files + +For external IdP-backed workload identity, managed backends inject +`NMP_WORKLOAD_IDENTITY_TOKEN_FILE`. The file contains a workload identity +subject token, not a final NeMo API access token. + +The SDK reads that file immediately before OAuth 2.0 Token Exchange (RFC 8693), +exchanges the subject token at the discovered IdP token endpoint, caches the +returned access token until near expiry, and then rereads the file for the next +exchange. + +Backend ownership: + +- Kubernetes uses a projected service account token volume. Kubelet rotates the + file. +- Docker uses a dedicated controller-managed workload identity volume. The + Docker backend writes and refreshes the file. + +Users must not provide `NEMO_WORKLOAD_TOKEN`, `NEMO_WORKLOAD_TOKEN_FILE`, or +`NMP_WORKLOAD_IDENTITY_TOKEN_FILE` in managed job requests. Direct SDK users may +set `NMP_WORKLOAD_IDENTITY_TOKEN_FILE` only when they own the refreshed subject +token file. + +This follows the same shape as common cloud SDKs: AWS uses +`AWS_WEB_IDENTITY_TOKEN_FILE`, Azure uses `AZURE_FEDERATED_TOKEN_FILE`, and +Google Workload Identity Federation uses `GOOGLE_APPLICATION_CREDENTIALS` to +point at a credential config that references a projected token file. ## Inference Auth Context @@ -23,7 +70,10 @@ When a model is deployed as an inference endpoint, incoming requests are authent ## Trust Implications -The `NMP_PRINCIPAL` environment variable is trusted by NeMo Platform services. If a user can exec into a job pod and read this variable, they have the submitting user's identity for the duration of the job. +The workload identity subject-token file is not the final NeMo Platform access +token, but it can be exchanged for one. Protect it as credential material, mount +it read-only into the workload container where possible, and keep its lifetime +short. diff --git a/docs/set-up/config-reference.mdx b/docs/set-up/config-reference.mdx index ada9944a93..718caac044 100644 --- a/docs/set-up/config-reference.mdx +++ b/docs/set-up/config-reference.mdx @@ -105,6 +105,34 @@ auth: subject_claim: sub # Space-separated OAuth scopes to request during authentication. For Azure AD with custom API, use: 'api://{app-id}/.default openid profile email' | default: 'openid profile email offline_access' default_scopes: openid profile email offline_access + # Enable SDK workload identity token exchange using NMP_WORKLOAD_IDENTITY_TOKEN_FILE subject tokens. | default: False + workload_token_exchange_enabled: false + # OAuth client ID to use for workload identity token exchange. Defaults to client_id when unset. + workload_client_id: + # OAuth token endpoint to use for workload identity token exchange. Defaults to token_endpoint. + workload_token_endpoint: + # RFC 8693 audience requested for workload identity token exchange. + workload_audience: + # Space-separated OAuth scopes requested for workload identity token exchange. + workload_scope: + # Issuer to stamp on workload identity access tokens minted by the NeMo auth service. Defaults to the platform auth endpoint origin serving the token exchange request. + workload_token_issuer: + # Lifetime in seconds for workload identity access tokens minted by the NeMo auth service. | default: 300 + workload_token_ttl_seconds: 300 + # JWT key id advertised by the NeMo auth service workload identity JWKS endpoint. | default: 'nemo-workload-exchange' + workload_token_key_id: nemo-workload-exchange + # Path to a PEM-encoded RSA private key used by the NeMo auth service to sign workload identity access tokens. Intended for mounted shared secrets. + workload_token_private_key_file: + # Additional RFC 8693 audience values accepted by the NeMo auth service workload token exchange endpoint. The configured workload_audience is always accepted. + workload_allowed_audiences: [] + # JWKS URI used by the NeMo auth service to validate JWT subject tokens for workload token exchange. Leave unset when only Kubernetes TokenReview subject validation is enabled. + workload_subject_jwks_uri: + # Allowed JWT subject token issuers for workload token exchange. Required when workload_subject_jwks_uri is set. + workload_subject_issuers: [] + # TTL in seconds for caching workload subject JWKS responses. Set to 0 to disable caching. | default: 3600 + workload_subject_jwks_cache_ttl_seconds: 3600 + # Allow the NeMo auth service workload token exchange endpoint to validate Kubernetes projected service account subject tokens using the TokenReview API. | default: False + workload_kubernetes_token_review_enabled: false # Prefix to strip from token scopes before authorization. For example, if IdP returns 'api://my-app/models:read', set prefix to 'api://my-app/' to normalize to 'models:read'. If not set, scopes are used as-is. scope_prefix: # TTL in seconds for caching IdP discovery document responses. Used by the discovery endpoint to avoid per-request IdP calls. Set to 0 to disable caching. | default: 300 @@ -250,6 +278,26 @@ jobs: networking: # Docker network for the job container | default: 'host' job_container_network: host + # Docker workload identity subject-token issuer configuration. + workload_identity: + # Enable Docker workload identity token-file injection. Defaults to auth.oidc.workload_token_exchange_enabled. + enabled: + # OAuth token endpoint used by the Docker demo issuer. Defaults to auth.oidc.token_endpoint. + token_endpoint: + # OAuth client ID used by the Docker demo issuer. Defaults to auth.oidc.workload_client_id or auth.oidc.client_id. + client_id: + # OAuth client secret for the Docker demo issuer. + client_secret: + # Username for the Docker demo issuer password grant. + username: + # Controller environment variable that contains the Docker demo issuer password grant shared secret. | default: 'AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD' + password_env_var: AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD + # OAuth scope for the Docker demo issuer. + scope: + # Fallback subject-token lifetime when the Docker demo issuer response omits expires_in. + subject_token_ttl_seconds: 600 + # Seconds before subject-token expiry when the Docker refresher issues a replacement token. | default: 60 + refresh_margin_seconds: 60 # Default Kubernetes execution profile configuration kubernetes_job: # default: 1800 @@ -324,6 +372,10 @@ jobs: scheduler_name: '' # Container image that contains the jobs-launcher binary. | default: 'nvcr.io/nvidia/nemo-microservices/jobs-launcher:latest' launcher_image: nvcr.io/nvidia/nemo-microservices/jobs-launcher:latest + # Requested expirationSeconds for the projected service account token used as the workload identity subject token. | default: 600 + workload_identity_token_expiration_seconds: 600 + # Audience for the projected service account token. Defaults to auth.oidc.workload_client_id, auth.oidc.client_id, then 'nemo-platform'. + workload_identity_token_audience: # Default Volcano execution profile configuration volcano_job: # default: 1800 @@ -398,6 +450,10 @@ jobs: scheduler_name: volcano # Container image that contains the jobs-launcher binary. | default: 'nvcr.io/nvidia/nemo-microservices/jobs-launcher:latest' launcher_image: nvcr.io/nvidia/nemo-microservices/jobs-launcher:latest + # Requested expirationSeconds for the projected service account token used as the workload identity subject token. | default: 600 + workload_identity_token_expiration_seconds: 600 + # Audience for the projected service account token. Defaults to auth.oidc.workload_client_id, auth.oidc.client_id, then 'nemo-platform'. + workload_identity_token_audience: # The Volcano queue to submit the job to. | default: 'default' queue: default # maxRetry indicates the maximum number of retries allowed by the job | default: 0 diff --git a/e2e/authz_oidc/conftest.py b/e2e/authz_oidc/conftest.py index 2262dbc191..aa26f40f50 100644 --- a/e2e/authz_oidc/conftest.py +++ b/e2e/authz_oidc/conftest.py @@ -129,16 +129,21 @@ def _uninstall_fixture_plugins(names: list[str]) -> None: logger.info("Uninstalled fixture plugins: %s", ", ".join(names)) -def _platform_env(issuer_url: str, data_dir: Path, extra: dict[str, str]) -> dict[str, str]: +def _platform_env(issuer_url: str, base_url: str, data_dir: Path, extra: dict[str, str]) -> dict[str, str]: env = {k: v for k, v in os.environ.items() if not k.startswith(("NMP_", "DATABASE_"))} env.update( { "NMP_CONFIG_FILE_PATH": str(_PLATFORM_CONFIG), "NMP_CONFIG_WARNINGS_DISABLED": "1", + "NMP_BASE_URL": base_url, "NMP_DATA_DIR": str(data_dir), "NMP_SEED_ON_STARTUP": "true", + # Authz rows need auth seeding, but the default model-provider seed + # requires a real NGC_API_KEY and is unrelated to this matrix. + "NMP_PLATFORM_SEED_MODEL_PROVIDER_ENABLED": "false", "NMP_AUTH_ENABLED": "true", "NMP_AUTH_ALLOW_UNSIGNED_JWT": "false", # defaults are true; signed JWTs only + "NMP_AUTH_POLICY_DECISION_POINT_BASE_URL": base_url, "NMP_AUTH_OIDC_ENABLED": "true", "NMP_AUTH_OIDC_ISSUER": issuer_url, "NMP_AUTH_OIDC_AUDIENCE": DEFAULT_AUDIENCE, @@ -186,7 +191,7 @@ def _spawn_platform( "--port", str(port), ] - env = _platform_env(issuer.issuer_url, data_dir, extra_env) + env = _platform_env(issuer.issuer_url, base_url, data_dir, extra_env) logger.info("Spawning platform [%s] on %s (log: %s)", label, base_url, log_path) with open(log_path, "w") as log_file: diff --git a/e2e/backends/docker_compose.py b/e2e/backends/docker_compose.py index 325382d732..a6f3345d66 100644 --- a/e2e/backends/docker_compose.py +++ b/e2e/backends/docker_compose.py @@ -8,11 +8,13 @@ import subprocess import time from pathlib import Path -from typing import Literal +from typing import Literal, TextIO import httpx +from nemo_platform_ext.client.tls import NMP_CLIENT_SSL_CERT_FILE_ENVVAR ComposeLifecycle = Literal["fresh", "reuse"] +_DIAGNOSTIC_COMMAND_TIMEOUT_SECONDS = 60 def _compose_env(env: dict[str, str] | None) -> dict[str, str]: @@ -41,6 +43,18 @@ def _parse_compose_ps_json(output: str) -> list[dict[str, object]]: raise ValueError("docker compose ps did not return JSON objects") +def _compose_exit_code(entry: dict[str, object]) -> int | None: + exit_code = entry.get("ExitCode") + if isinstance(exit_code, int): + return exit_code + if isinstance(exit_code, str) and exit_code: + try: + return int(exit_code) + except ValueError: + return None + return None + + def _compose_stack_readiness(entries: list[dict[str, object]], expected_services: set[str]) -> tuple[bool, list[str]]: entries_by_service = { str(entry.get("Service") or entry.get("Name")): entry @@ -58,6 +72,8 @@ def _compose_stack_readiness(entries: list[dict[str, object]], expected_services if health: if health != "healthy": not_ready.append(f"{service} (state={state or 'unknown'}, health={health})") + elif service.endswith("-init") and state == "exited" and _compose_exit_code(entry) == 0: + continue elif state != "running": not_ready.append(f"{service} (state={state or 'unknown'})") return not not_ready, not_ready @@ -72,6 +88,7 @@ def __init__( project_name: str, service_url: str, wait_url: str | None = None, + wait_urls: list[str] | None = None, env: dict[str, str] | None = None, lifecycle: ComposeLifecycle = "fresh", wait_timeout_seconds: int = 180, @@ -80,7 +97,8 @@ def __init__( self.config_path = config_path self.project_name = project_name self.service_url = service_url - self.wait_url = wait_url or service_url + self.wait_urls = wait_urls or [wait_url or service_url] + self.wait_url = self.wait_urls[0] self.env = { "NEMO_COMPOSE_CONFIG_PATH": str(config_path.resolve()), **(env or {}), @@ -99,6 +117,9 @@ def _run(self, *extra_args: str, capture_output: bool = False) -> subprocess.Com env=_compose_env(self.env), ) + def exec(self, service: str, *command: str, capture_output: bool = False) -> subprocess.CompletedProcess[str]: + return self._run("exec", "-T", service, *command, capture_output=capture_output) + def _services(self, *extra_args: str) -> set[str]: result = self._run(*extra_args, capture_output=True) return {line for line in result.stdout.splitlines() if line} @@ -143,16 +164,29 @@ def start(self) -> None: self._wait_ready() def _wait_ready(self) -> None: + verify = ( + self.env.get(NMP_CLIENT_SSL_CERT_FILE_ENVVAR) + or self.env.get("REQUESTS_CA_BUNDLE") + or self.env.get("SSL_CERT_FILE") + or True + ) deadline = time.monotonic() + self.wait_timeout_seconds - while time.monotonic() < deadline: - try: - response = httpx.get(self.wait_url, timeout=5) - if response.status_code == 200: - return - except httpx.HTTPError: - pass + pending = list(dict.fromkeys(self.wait_urls)) + last_results: dict[str, str] = {} + while time.monotonic() < deadline and pending: + for wait_url in list(pending): + try: + response = httpx.get(wait_url, timeout=5, verify=verify) + if response.status_code == 200: + pending.remove(wait_url) + else: + last_results[wait_url] = f"HTTP {response.status_code}" + except httpx.HTTPError as exc: + last_results[wait_url] = str(exc) + if not pending: + return time.sleep(2) - raise TimeoutError(f"compose backend did not become ready: {self.wait_url}") + raise TimeoutError(f"compose backend did not become ready: {pending}; last_results={last_results}") def stop(self) -> None: if self.lifecycle == "reuse": @@ -161,9 +195,25 @@ def stop(self) -> None: def write_logs(self, log_path: Path) -> None: log_path.parent.mkdir(parents=True, exist_ok=True) - args = _compose_base_args(self.compose_file, self.project_name) - args.extend(["logs", "--no-color", "--timestamps"]) with log_path.open("w", encoding="utf-8") as log_file: + self._write_diagnostic_command(log_file, "docker compose config --services", ["config", "--services"]) + self._write_diagnostic_command( + log_file, + "docker compose ps --all --format json", + ["ps", "--all", "--format", "json"], + ) + self._write_diagnostic_command(log_file, "docker compose ps --all", ["ps", "--all"]) + self._write_diagnostic_command( + log_file, + "docker compose logs --no-color --timestamps", + ["logs", "--no-color", "--timestamps"], + ) + + def _write_diagnostic_command(self, log_file: TextIO, title: str, extra_args: list[str]) -> None: + args = _compose_base_args(self.compose_file, self.project_name) + args.extend(extra_args) + log_file.write(f"\n===== {title} =====\n") + try: result = subprocess.run( args, check=False, @@ -171,6 +221,9 @@ def write_logs(self, log_path: Path) -> None: stdout=log_file, stderr=subprocess.STDOUT, env=_compose_env(self.env), + timeout=_DIAGNOSTIC_COMMAND_TIMEOUT_SECONDS, ) - if result.returncode != 0: - log_file.write(f"\n[docker compose logs exited with status {result.returncode}]\n") + except subprocess.TimeoutExpired: + log_file.write(f"\n[command timed out after {_DIAGNOSTIC_COMMAND_TIMEOUT_SECONDS}s]\n") + return + log_file.write(f"\n[command exited with status {result.returncode}]\n") diff --git a/e2e/conftest.py b/e2e/conftest.py index 7a54810afa..14938026bb 100644 --- a/e2e/conftest.py +++ b/e2e/conftest.py @@ -113,8 +113,6 @@ def pytest_collection_modifyitems(session: pytest.Session, config: pytest.Config _TAIL_LINES_ON_FAILURE = 100 _services_log_key = pytest.StashKey[Path]() -_active_services_log_key = pytest.StashKey[Path]() -_active_services_metadata_key = pytest.StashKey[dict[str, str]]() NGC_API_KEY_ENV = "NGC_API_KEY" @@ -161,14 +159,25 @@ def pytest_runtest_makereport(item: pytest.Item, call: pytest.CallInfo): # noqa if not report.failed: return - log_path = item.stash.get(_active_services_log_key, None) or item.session.stash.get(_services_log_key, None) + metadata: dict[str, str] | None = None + module = item.getparent(pytest.Module) + if module is not None: + manager = item.config.stash[_services_pool_manager_key] + active_metadata = manager.describe_active_module_binding(module.nodeid) + if active_metadata: + metadata = {key: str(value) for key, value in active_metadata.items() if value is not None} + + log_path: Path | None = None + if metadata and metadata.get("service_log_path"): + log_path = Path(metadata["service_log_path"]) + if log_path is None: + log_path = item.session.stash.get(_services_log_key, None) if log_path and log_path.exists(): lines = log_path.read_text().splitlines(keepends=True) tail = lines[-_TAIL_LINES_ON_FAILURE:] if tail: header = f"--- services log (last {len(tail)} lines) [{log_path}] ---" report.sections.append(("Services Log", f"{header}\n{''.join(tail)}")) - metadata = item.stash.get(_active_services_metadata_key, None) if metadata: report.sections.append( ( @@ -217,30 +226,6 @@ def _services_instance( _services_pool_manager.release_for_module(module) -@pytest.fixture(autouse=True) -def _bind_services_log_to_test(request: pytest.FixtureRequest, _services_instance: RunningServices) -> None: - if _services_instance.log_path is not None: - request.node.stash[_active_services_log_key] = _services_instance.log_path - module = request.node.getparent(pytest.Module) - if module is None: - return - manager = request.config.stash[_services_pool_manager_key] - metadata = { - key: str(value) - for key, value in manager.describe_module_binding(module.nodeid, _services_instance).items() - if value is not None - } - request.node.stash[_active_services_metadata_key] = metadata - if _E2E_HARNESS_DEBUG: - logger.info( - "E2E test binding", - extra={ - **metadata, - "test": request.node.nodeid, - }, - ) - - @pytest.fixture(scope="module") def _services(_services_instance: RunningServices) -> Iterator[str]: yield _services_instance.url diff --git a/e2e/services_pool.py b/e2e/services_pool.py index 2d3cb08192..444f975660 100644 --- a/e2e/services_pool.py +++ b/e2e/services_pool.py @@ -12,11 +12,13 @@ import sys import time import uuid +from collections.abc import Mapping, Sequence from copy import deepcopy from dataclasses import dataclass from importlib.metadata import entry_points from pathlib import Path -from typing import Any, Callable, Literal, NotRequired, TypedDict +from string import Template +from typing import Any, Callable, Literal, NotRequired, TypedDict, cast import httpx import pytest @@ -51,6 +53,12 @@ class ServicesPoolKey: config_hash: str +class DynamicPortConfig(TypedDict, total=False): + host: str + port: str | int + scheme: str + + class E2EHarnessConfig(TypedDict, total=False): backend: Literal["subprocess", "docker", "docker_compose"] compose_file: str @@ -58,9 +66,11 @@ class E2EHarnessConfig(TypedDict, total=False): service_url: str auth_ready_url: str wait_url: str + wait_urls: list[str] lifecycle: Literal["fresh", "reuse"] compose_project_prefix: str env: dict[str, str] + dynamic_ports: dict[str, DynamicPortConfig] # Subprocess backend only; container-reachable host for platform.base_url (see # _start_services_subprocess for how it's applied). container_base_url_host: str @@ -78,6 +88,7 @@ class RunningServices: docker_network_name: str | None = None docker_container_alias: str | None = None docker_container_port: int | None = None + compose_project_name: str | None = None @dataclass(frozen=True) @@ -124,6 +135,23 @@ def register_collected_items(self, items: list[pytest.Item]) -> None: def acquire_for_module(self, module: pytest.Module) -> RunningServices: self._ensure_module_registered(module) state = self._module_states[module.nodeid] + return self._acquire_for_registered_owner(module.nodeid, state) + + def acquire_for_config( + self, + owner_id: str, + config_layers: Sequence[str | dict[str, Any]], + harness_config: Mapping[str, Any], + ) -> RunningServices: + self._ensure_config_registered( + owner_id, + config_layers=config_layers, + harness_config=_normalize_e2e_harness_config(harness_config), + ) + state = self._module_states[owner_id] + return self._acquire_for_registered_owner(owner_id, state) + + def _acquire_for_registered_owner(self, owner_id: str, state: ModuleConfigState) -> RunningServices: external_url = os.environ.get("NMP_BASE_URL") if external_url: return RunningServices( @@ -135,7 +163,7 @@ def acquire_for_module(self, module: pytest.Module) -> RunningServices: ) if state.config_path is None: state = self._materialize_config_path(state) - self._module_states[module.nodeid] = state + self._module_states[owner_id] = state assert state.config_path is not None services = self._running_by_key.get(state.key) if services is None: @@ -148,43 +176,50 @@ def acquire_for_module(self, module: pytest.Module) -> RunningServices: log_path, ) self._running_by_key[state.key] = services - previous_key = self._active_service_key_by_module.get(module.nodeid) + previous_key = self._active_service_key_by_module.get(owner_id) if previous_key is not None and previous_key != state.key: logger.error( "E2E module rebound to a different services pool key", extra={ - "e2e_module": module.nodeid, + "e2e_module": owner_id, "previous_config_hash": previous_key.config_hash, "new_config_hash": state.key.config_hash, "new_url": services.url, "new_pid": services.proc.pid if services.proc is not None else None, }, ) - self._active_service_key_by_module[module.nodeid] = state.key - self._log_debug("E2E services acquire", **self.describe_module_binding(module.nodeid, services)) + self._remaining_modules_by_key.setdefault(state.key, set()).add(owner_id) + self._active_service_key_by_module[owner_id] = state.key + self._log_debug("E2E services acquire", **self.describe_module_binding(owner_id, services)) return services def release_for_module(self, module: pytest.Module) -> None: + self._release_for_registered_owner(module.nodeid) + + def release_for_config(self, owner_id: str) -> None: + self._release_for_registered_owner(owner_id) + + def _release_for_registered_owner(self, owner_id: str) -> None: + self._active_service_key_by_module.pop(owner_id, None) if os.environ.get("NMP_BASE_URL"): return - state = self._module_states.get(module.nodeid) + state = self._module_states.get(owner_id) if state is None: return remaining = self._remaining_modules_by_key.get(state.key) - if remaining is None or module.nodeid not in remaining: + if remaining is None or owner_id not in remaining: return - remaining.remove(module.nodeid) + remaining.remove(owner_id) self._log_debug( "E2E services release", **{ - **self.describe_module_binding(module.nodeid), + **self.describe_module_binding(owner_id), "remaining_modules_for_hash": sorted(remaining), }, ) if remaining: return self._remaining_modules_by_key.pop(state.key, None) - self._active_service_key_by_module.pop(module.nodeid, None) services = self._running_by_key.pop(state.key, None) if services is not None: self._terminate_services(services) @@ -200,10 +235,42 @@ def _ensure_module_registered(self, module: pytest.Module) -> None: return resolved_paths, config_data = _load_effective_e2e_config_from_node(module) harness_config = _resolve_e2e_harness_config_from_node(module) + self._register_config_state( + module.nodeid, + resolved_paths=resolved_paths, + config_data=config_data, + harness_config=harness_config, + ) + + def _ensure_config_registered( + self, + owner_id: str, + *, + config_layers: Sequence[str | dict[str, Any]], + harness_config: E2EHarnessConfig, + ) -> None: + if owner_id in self._module_states: + return + resolved_paths, config_data = _load_effective_e2e_config_from_layers(config_layers) + self._register_config_state( + owner_id, + resolved_paths=resolved_paths, + config_data=config_data, + harness_config=harness_config, + ) + + def _register_config_state( + self, + owner_id: str, + *, + resolved_paths: Sequence[Path], + config_data: dict[str, Any], + harness_config: E2EHarnessConfig, + ) -> None: key = _services_pool_key(_canonical_services_hash(config_data, harness_config)) auth_enabled = _e2e_auth_enabled(config_data) - self._module_states[module.nodeid] = ModuleConfigState( - module_id=module.nodeid, + self._module_states[owner_id] = ModuleConfigState( + module_id=owner_id, key=key, config_path=None, config_data=config_data, @@ -211,13 +278,13 @@ def _ensure_module_registered(self, module: pytest.Module) -> None: config_layers=tuple(str(path) for path in resolved_paths), auth_enabled=auth_enabled, ) - self._remaining_modules_by_key.setdefault(key, set()).add(module.nodeid) + self._remaining_modules_by_key.setdefault(key, set()).add(owner_id) self._log_debug( "Registered E2E module config", - e2e_module=module.nodeid, + e2e_module=owner_id, config_hash=key.config_hash, harness_backend=harness_config["backend"], - config_layers=list(self._module_states[module.nodeid].config_layers), + config_layers=list(self._module_states[owner_id].config_layers), auth_enabled=auth_enabled, ) @@ -304,10 +371,23 @@ def describe_module_binding( "docker_network_name": services.docker_network_name, "docker_container_alias": services.docker_container_alias, "docker_container_port": services.docker_container_port, + "compose_project_name": services.compose_project_name, } ) return details + def describe_active_module_binding(self, module_id: str) -> dict[str, Any] | None: + key = self._active_service_key_by_module.get(module_id) + if key is None: + return None + active_modules = self._remaining_modules_by_key.get(key) + if active_modules is None or module_id not in active_modules: + return None + services = self._running_by_key.get(key) + if services is None: + return None + return self.describe_module_binding(module_id, services) + def _services_log_dir(tmp_path_factory: pytest.TempPathFactory) -> Path: log_dir = os.environ.get("E2E_SERVICES_LOG_DIR") @@ -341,9 +421,13 @@ def _resolve_e2e_harness_config_from_node(node: Node) -> E2EHarnessConfig: harness = marker.kwargs.get("harness") if harness is None: return {"backend": "subprocess"} - if not isinstance(harness, dict): + return _normalize_e2e_harness_config(harness) + + +def _normalize_e2e_harness_config(harness: Mapping[str, Any]) -> E2EHarnessConfig: + if not isinstance(harness, Mapping): raise pytest.UsageError("pytest.mark.e2e_config harness must be a mapping") - normalized = _normalize_config(harness) + normalized = _normalize_config(dict(harness)) backend = normalized.get("backend", "subprocess") if backend not in {"subprocess", "docker", "docker_compose"}: raise pytest.UsageError(f"unsupported e2e harness backend: {backend}") @@ -390,6 +474,30 @@ def _normalize_config(value: Any, path: tuple[str, ...] = ()) -> Any: return value +def _render_template_value(value: Any, context: Mapping[str, str]) -> Any: + if isinstance(value, str): + return Template(value).safe_substitute(context) + if isinstance(value, Mapping): + return {key: _render_template_value(item, context) for key, item in value.items()} + if isinstance(value, Sequence) and not isinstance(value, str): + return [_render_template_value(item, context) for item in value] + return value + + +def _dynamic_port_template_context(harness_config: E2EHarnessConfig) -> dict[str, str]: + context: dict[str, str] = {} + for name, port_config in harness_config.get("dynamic_ports", {}).items(): + configured_port = port_config.get("port") + port = _find_free_port() if configured_port is None else int(configured_port) + host = port_config.get("host", "127.0.0.1") + scheme = port_config.get("scheme", "http") + context[f"{name}_host"] = host + context[f"{name}_port"] = str(port) + context[f"{name}_scheme"] = scheme + context[f"{name}_url"] = f"{scheme}://{host}:{port}" + return context + + def _canonical_config_hash(config_data: dict[str, Any]) -> str: normalized = _normalize_config(config_data) payload = json.dumps(normalized, sort_keys=True, separators=(",", ":"), ensure_ascii=True) @@ -410,10 +518,16 @@ def _canonical_services_hash(config_data: dict[str, Any], harness_config: E2EHar def _load_effective_e2e_config_from_node(node: Node) -> tuple[list[Path], dict[str, Any]]: + return _load_effective_e2e_config_from_layers(_resolve_e2e_config_layers_from_node(node)) + + +def _load_effective_e2e_config_from_layers( + layers: Sequence[str | dict[str, Any]], +) -> tuple[list[Path], dict[str, Any]]: effective_config: dict[str, Any] = {} resolved_paths: list[Path] = [] - for layer in _resolve_e2e_config_layers_from_node(node): + for layer in layers: if isinstance(layer, str): config_path = _resolve_config_path(layer) if not config_path.is_file(): @@ -632,6 +746,60 @@ def _wait_for_auth_ready(url: str, proc: subprocess.Popen[Any] | None, timeout: return False +def _request_verify_from_env(env: Mapping[str, str] | None = None) -> str | bool: + source = dict(os.environ) + if env: + source.update(env) + return ( + source.get("NMP_CLIENT_SSL_CERT_FILE") + or source.get("REQUESTS_CA_BUNDLE") + or source.get("SSL_CERT_FILE") + or True + ) + + +def _wait_for_auth_ready_url( + url: str, + proc: subprocess.Popen[Any] | None, + *, + env: Mapping[str, str] | None = None, + timeout: float = _AUTH_READY_TIMEOUT, +) -> bool: + deadline = time.monotonic() + timeout + verify = _request_verify_from_env(env) + while time.monotonic() < deadline: + if proc is not None and _process_exited(proc): + return False + try: + response = httpx.get(url, timeout=5.0, verify=verify) + if response.status_code == 200: + return True + except httpx.RequestError as exc: + logger.debug("Auth readiness URL probe failed; will retry: %s", exc) + if proc is not None and _process_exited(proc): + return False + time.sleep(_HEALTH_POLL_INTERVAL) + return False + + +def _auth_ready_url(harness_config: E2EHarnessConfig, service_url: str) -> str | None: + url = harness_config.get("auth_ready_url") + if not url: + return None + return Template(url).safe_substitute({"service_url": service_url}) + + +def _wait_for_configured_auth_ready( + service_url: str, + proc: subprocess.Popen[Any] | None, + harness_config: E2EHarnessConfig, +) -> bool: + ready_url = _auth_ready_url(harness_config, service_url) + if ready_url is not None: + return _wait_for_auth_ready_url(ready_url, proc, env=harness_config.get("env")) + return _wait_for_auth_ready(service_url, proc) + + def _start_services( config_path: Path, config_data: dict[str, Any], @@ -641,7 +809,7 @@ def _start_services( ) -> RunningServices: backend = _e2e_backend(harness_config) if backend == "docker": - return _start_services_docker(config_path, config_data, config_hash) + return _start_services_docker(config_path, config_data, harness_config, config_hash) if backend == "docker_compose": return _start_services_docker_compose(config_path, config_data, harness_config, config_hash, log_path) return _start_services_subprocess(config_path, config_data, harness_config, config_hash, log_path) @@ -708,7 +876,7 @@ def _start_services_subprocess( f"nemo services run did not become healthy within {_HEALTH_TIMEOUT}s.\nlog:\n{log_path.read_text()}" ) auth_enabled = _e2e_auth_enabled(config_data) - if auth_enabled and not _wait_for_auth_ready(url, proc): + if auth_enabled and not _wait_for_configured_auth_ready(url, proc, harness_config): try: proc.terminate() proc.wait(timeout=10) @@ -730,7 +898,12 @@ def _start_services_subprocess( ) -def _start_services_docker(config_path: Path, config_data: dict[str, Any], config_hash: str) -> RunningServices: +def _start_services_docker( + config_path: Path, + config_data: dict[str, Any], + harness_config: E2EHarnessConfig, + config_hash: str, +) -> RunningServices: backend = DockerE2EBackend(config_path=config_path, **_docker_backend_overrides()) try: backend.start() @@ -739,7 +912,7 @@ def _start_services_docker(config_path: Path, config_data: dict[str, Any], confi raise auth_enabled = _e2e_auth_enabled(config_data) - if auth_enabled and not _wait_for_auth_ready(backend.base_url, None): + if auth_enabled and not _wait_for_configured_auth_ready(backend.base_url, None, harness_config): backend.stop() pytest.fail(f"Platform auth seed did not become ready within {_AUTH_READY_TIMEOUT}s.") @@ -771,14 +944,35 @@ def _start_services_docker_compose( if project_name is None: project_prefix = harness_config.get("compose_project_prefix", "e2e-compose") project_name = f"{project_prefix}-{config_hash}" + if harness_config["lifecycle"] == "fresh": + project_name = f"{project_name}-{uuid.uuid4().hex[:8]}" + if harness_config["lifecycle"] != "fresh": + missing_fixed_ports = sorted( + name for name, port_config in harness_config.get("dynamic_ports", {}).items() if "port" not in port_config + ) + if missing_fixed_ports: + raise pytest.UsageError( + f"docker_compose dynamic_ports require explicit ports with lifecycle='reuse': {missing_fixed_ports}" + ) + template_context = _dynamic_port_template_context(harness_config) + if template_context: + config_data = cast(dict[str, Any], _render_template_value(config_data, template_context)) + rendered_harness_config = cast(E2EHarnessConfig, _render_template_value(dict(harness_config), template_context)) + rendered_harness_config.pop("dynamic_ports", None) + harness_config = rendered_harness_config + config_path.write_text(yaml.safe_dump(config_data, default_flow_style=False, sort_keys=True)) + auth_enabled = _e2e_auth_enabled(config_data) service_url = harness_config["service_url"] - wait_url = harness_config.get("wait_url") + auth_ready_url = _auth_ready_url(harness_config, service_url) if auth_enabled else None + wait_url = auth_ready_url or harness_config.get("wait_url") + wait_urls = [auth_ready_url] if auth_ready_url is not None else harness_config.get("wait_urls") backend = DockerComposeE2EBackend( compose_file=compose_file, config_path=config_path, project_name=project_name, service_url=service_url, wait_url=wait_url, + wait_urls=wait_urls, env=harness_config.get("env"), lifecycle=harness_config["lifecycle"], ) @@ -789,15 +983,6 @@ def _start_services_docker_compose( backend.stop() raise - auth_enabled = _e2e_auth_enabled(config_data) - auth_ready_url = harness_config.get("auth_ready_url", backend.service_url) - if auth_enabled and not _wait_for_auth_ready(auth_ready_url, None): - _write_docker_compose_logs(backend, log_path) - backend.stop() - pytest.fail( - f"Platform auth seed did not become ready within {_AUTH_READY_TIMEOUT}s.\nlog:\n{_read_log_text(log_path)}" - ) - def close() -> None: try: _write_docker_compose_logs(backend, log_path) @@ -812,6 +997,7 @@ def close() -> None: close=close, auth_enabled=auth_enabled, key=_services_pool_key(config_hash), + compose_project_name=project_name, ) diff --git a/k8s/helm/README.md b/k8s/helm/README.md index 88cd83cb5d..e4cd7e650f 100644 --- a/k8s/helm/README.md +++ b/k8s/helm/README.md @@ -44,6 +44,8 @@ secrets will not decrypt with a new key. | api.autoscaling.targetCPUUtilizationPercentage | int | `80` | The target CPU utilization percentage. | | api.enabled | bool | `true` | Specifies whether to enable the api deployment. | | api.extraArgs | list | `[]` | Additional arguments to pass to the Platform API service | +| api.extraVolumeMounts | list | `[]` | Additional volume mounts to add to the Platform API container. | +| api.extraVolumes | list | `[]` | Additional volumes to add to the Platform API pod. | | api.image | object | This object has the following default values for the image configuration. | Container image configuration for the api deployment. | | api.image.pullPolicy | string | `"IfNotPresent"` | The image pull policy determining when to pull new images. | | api.image.repository | string | `"nvcr.io/nvidia/nemo-platform/nmp-api"` | The registry where the NeMo Platform image is located. | @@ -166,8 +168,12 @@ secrets will not decrypt with a new key. | envoyProxy.autoscaling.maxReplicas | int | `10` | The maximum number of replicas for the deployment. | | envoyProxy.autoscaling.minReplicas | int | `1` | The minimum number of replicas for the deployment. | | envoyProxy.autoscaling.targetCPUUtilizationPercentage | int | `80` | The target CPU utilization percentage. | +| envoyProxy.configOverride | string | `""` | Full Envoy config override. When set, this replaces the chart's default passthrough Envoy config. | | envoyProxy.enabled | bool | `true` | Specifies whether to enable the Envoy proxy deployment. Rendered only when platform config has auth.enabled: true. | | envoyProxy.extraArgs | list | `[]` | Extra arguments to append to the envoy container command. Useful for passing server flags such as concurrency. Example: ["--concurrency", "4"] | +| envoyProxy.extraVolumeMounts | list | `[]` | Additional volume mounts to add to the Envoy container. | +| envoyProxy.extraVolumes | list | `[]` | Additional volumes to add to the Envoy pod. | +| envoyProxy.image.digest | string | `""` | Optional image digest. When set, the Envoy image renders as repository@digest. | | envoyProxy.livenessProbe | object | `{"failureThreshold":3,"httpGet":{"path":"/ready","port":"admin"},"periodSeconds":10,"timeoutSeconds":5}` | Liveness probe for the Envoy container (admin interface /ready). | | envoyProxy.nodeSelector | object | `{}` | Node selector configuration for the Envoy pods. | | envoyProxy.podAnnotations | object | `{}` | Annotations to add to the Envoy service pod. | diff --git a/k8s/helm/ci/21-api-extra-volumes.yaml b/k8s/helm/ci/21-api-extra-volumes.yaml new file mode 100644 index 0000000000..da2f7737f4 --- /dev/null +++ b/k8s/helm/ci/21-api-extra-volumes.yaml @@ -0,0 +1,18 @@ +# CI validates by running: helm template nemo-platform . -f ci/21-api-extra-volumes.yaml +# Exercises API extra volume/mount rendering with workload token exchange config. + +api: + extraVolumes: + - name: workload-token-signing-key + secret: + secretName: nemo-workload-token-signing-key + extraVolumeMounts: + - name: workload-token-signing-key + mountPath: /etc/nmp/workload-token + readOnly: true +platformConfig: + auth: + oidc: + workload_token_exchange_enabled: true + workload_token_key_id: nemo-workload-exchange + workload_token_private_key_file: /etc/nmp/workload-token/private-key.pem diff --git a/k8s/helm/ci/22-envoy-config-override.yaml b/k8s/helm/ci/22-envoy-config-override.yaml new file mode 100644 index 0000000000..c931443390 --- /dev/null +++ b/k8s/helm/ci/22-envoy-config-override.yaml @@ -0,0 +1,51 @@ +# CI validates by running: helm template nemo-platform . -f ci/22-envoy-config-override.yaml +# Exercises the full-replacement envoyProxy.configOverride path with a deliberately +# small custom Envoy config, so this fixture does not duplicate the chart default. + +platformConfig: + auth: + enabled: true +envoyProxy: + image: + digest: sha256:f91c972d5c99bc133233a079b5663b903e7d56b3b0b0216398924f7b80d09e47 + securityContext: + readOnlyRootFilesystem: true + extraVolumeMounts: + - name: tmp + mountPath: /tmp + extraVolumes: + - name: tmp + emptyDir: {} + configOverride: |- + admin: + address: + socket_address: + address: 127.0.0.1 + port_value: 9901 + static_resources: + listeners: + - name: ci_override_listener + address: + socket_address: + address: 0.0.0.0 + port_value: 8080 + filter_chains: + - filters: + - name: envoy.filters.network.http_connection_manager + typed_config: + "@type": type.googleapis.com/envoy.extensions.filters.network.http_connection_manager.v3.HttpConnectionManager + stat_prefix: ci_override + route_config: + name: ci_override_route + virtual_hosts: + - name: ci + domains: ["*"] + routes: + - match: + prefix: "/" + direct_response: + status: 204 + http_filters: + - name: envoy.filters.http.router + typed_config: + "@type": type.googleapis.com/envoy.extensions.filters.http.router.v3.Router diff --git a/k8s/helm/templates/api/api-deployment.yaml b/k8s/helm/templates/api/api-deployment.yaml index c8515a19d5..b3bff3c8b3 100644 --- a/k8s/helm/templates/api/api-deployment.yaml +++ b/k8s/helm/templates/api/api-deployment.yaml @@ -138,6 +138,9 @@ spec: - name: files-storage mountPath: {{ include "nmp-core.localStoragePath" . }} {{- end }} + {{- with .Values.api.extraVolumeMounts }} + {{- toYaml . | nindent 12 }} + {{- end }} volumes: - name: data-designer-tmp emptyDir: {} @@ -149,6 +152,9 @@ spec: persistentVolumeClaim: claimName: {{ include "nmp-core.persistentVolumeClaim" . }} {{- end }} + {{- with .Values.api.extraVolumes }} + {{- toYaml . | nindent 8 }} + {{- end }} {{- with .Values.api.nodeSelector }} nodeSelector: {{- toYaml . | nindent 8 }} diff --git a/k8s/helm/templates/core/controller-deployment.yaml b/k8s/helm/templates/core/controller-deployment.yaml index bba63e47e2..0b9204afd6 100644 --- a/k8s/helm/templates/core/controller-deployment.yaml +++ b/k8s/helm/templates/core/controller-deployment.yaml @@ -59,6 +59,10 @@ spec: value: /etc/nmp/config.yaml - name: NMP_BASE_URL value: {{ include "nemo-platform.internalBaseUrl" . | quote }} + {{- if include "nemo-platform.embeddedPdpEnabled" . }} + - name: NMP_AUTH_POLICY_DECISION_POINT_BASE_URL + value: {{ include "nemo-platform.internalBaseUrl" . | quote }} + {{- end }} - name: OTEL_SERVICE_NAME value: {{ include "nmp-core.controller-servicename" . }} {{- include "nemo-common.otel-env" (dict "root" $ "local" .Values) | indent 12 }} diff --git a/k8s/helm/templates/platform-seed-job.yaml b/k8s/helm/templates/platform-seed-job.yaml index d831eda3f0..2990992248 100644 --- a/k8s/helm/templates/platform-seed-job.yaml +++ b/k8s/helm/templates/platform-seed-job.yaml @@ -62,6 +62,8 @@ spec: {{- end }} - name: NMP_CONFIG_FILE_PATH value: /etc/nmp/config.yaml + - name: NMP_BASE_URL + value: {{ include "nemo-platform.internalBaseUrl" . | quote }} - name: OTEL_SERVICE_NAME value: {{ include "nemo-platform.fullname" . }}-platform-seed {{ include "nemo-platform.env" . | nindent 12 | trim }} diff --git a/k8s/helm/templates/proxy/_helpers.tpl b/k8s/helm/templates/proxy/_helpers.tpl index ec092087f0..10adf5ad1f 100644 --- a/k8s/helm/templates/proxy/_helpers.tpl +++ b/k8s/helm/templates/proxy/_helpers.tpl @@ -5,6 +5,17 @@ Create a named Envoy service name which can be included from parent chart {{- printf "%s-envoy" ( include "nemo-platform.fullname" . | trunc 57 ) }} {{- end }} +{{/* +Create the Envoy image reference. +*/}} +{{- define "nmp-envoy.image" -}} +{{- if .Values.envoyProxy.image.digest -}} +{{ printf "%s@%s" .Values.envoyProxy.image.repository .Values.envoyProxy.image.digest }} +{{- else -}} +{{ printf "%s:%s" .Values.envoyProxy.image.repository .Values.envoyProxy.image.tag }} +{{- end -}} +{{- end }} + {{/* Labels for Envoy proxy resources (component + platform labels). */}} diff --git a/k8s/helm/templates/proxy/envoy-configmap.yaml b/k8s/helm/templates/proxy/envoy-configmap.yaml index 575bee1016..b446d29672 100644 --- a/k8s/helm/templates/proxy/envoy-configmap.yaml +++ b/k8s/helm/templates/proxy/envoy-configmap.yaml @@ -7,6 +7,9 @@ metadata: {{- include "nmp-envoy.labels" . | nindent 4 }} data: envoy.yaml: | +{{- if .Values.envoyProxy.configOverride }} +{{ tpl .Values.envoyProxy.configOverride . | nindent 4 }} +{{- else }} admin: address: socket_address: @@ -66,3 +69,4 @@ data: address: {{ include "nmp-api.api-servicename" . }} port_value: {{ .Values.api.service.port }} {{- end }} +{{- end }} diff --git a/k8s/helm/templates/proxy/envoy-deployment.yaml b/k8s/helm/templates/proxy/envoy-deployment.yaml index 25fbdb9af7..94b68d0738 100644 --- a/k8s/helm/templates/proxy/envoy-deployment.yaml +++ b/k8s/helm/templates/proxy/envoy-deployment.yaml @@ -42,7 +42,7 @@ spec: - name: envoy securityContext: {{- include "nemo-common.securityContext" (dict "global" .Values.securityContext "local" .Values.envoyProxy.securityContext) | nindent 12 }} - image: "{{ .Values.envoyProxy.image.repository }}:{{ .Values.envoyProxy.image.tag }}" + image: {{ include "nmp-envoy.image" . | quote }} imagePullPolicy: {{ .Values.envoyProxy.image.pullPolicy }} args: - -c @@ -63,6 +63,9 @@ spec: - name: config mountPath: /etc/envoy readOnly: true + {{- with .Values.envoyProxy.extraVolumeMounts }} + {{- toYaml . | nindent 12 }} + {{- end }} startupProbe: {{- toYaml .Values.envoyProxy.startupProbe | nindent 12 }} livenessProbe: @@ -75,6 +78,9 @@ spec: - name: config configMap: name: {{ include "nmp-envoy.servicename" . }} + {{- with .Values.envoyProxy.extraVolumes }} + {{- toYaml . | nindent 8 }} + {{- end }} {{- with .Values.envoyProxy.nodeSelector }} nodeSelector: {{- toYaml . | nindent 8 }} diff --git a/k8s/helm/values.yaml b/k8s/helm/values.yaml index 031cda827c..2e7deceb0c 100644 --- a/k8s/helm/values.yaml +++ b/k8s/helm/values.yaml @@ -529,6 +529,10 @@ api: replicaCount: 1 # -- Additional arguments to pass to the Platform API service extraArgs: [] + # -- Additional volume mounts to add to the Platform API container. + extraVolumeMounts: [] + # -- Additional volumes to add to the Platform API pod. + extraVolumes: [] # -- Service account configuration for the API service. # @default -- This object has the following default values for the service account configuration. serviceAccount: @@ -872,6 +876,8 @@ envoyProxy: image: repository: envoyproxy/envoy tag: v1.37.0 + # -- Optional image digest. When set, the Envoy image renders as repository@digest. + digest: "" pullPolicy: IfNotPresent # -- Service account configuration for the Envoy service. @@ -998,6 +1004,15 @@ envoyProxy: # Example: ["--concurrency", "4"] extraArgs: [] + # -- Full Envoy config override. When set, this replaces the chart's default passthrough Envoy config. + configOverride: "" + + # -- Additional volume mounts to add to the Envoy container. + extraVolumeMounts: [] + + # -- Additional volumes to add to the Envoy pod. + extraVolumes: [] + # ServiceMonitor configuration for Prometheus Operator serviceMonitor: # -- Enable ServiceMonitor resources for Prometheus Operator diff --git a/openapi/ga/individual/platform.openapi.yaml b/openapi/ga/individual/platform.openapi.yaml index 3db8849574..b0c624f38e 100644 --- a/openapi/ga/individual/platform.openapi.yaml +++ b/openapi/ga/individual/platform.openapi.yaml @@ -20,8 +20,13 @@ paths: \ CLI)\n - `userinfo_endpoint`: UserInfo endpoint\n - `client_id`: OAuth\ \ client ID to use\n - `default_scopes`: OAuth scopes to request during authentication\n\ \ - `scope_prefix`: Prefix to prepend to custom scopes (those with ':' or\ - \ '.default')" - operationId: get_auth_discovery_apis_auth_discovery_get + \ '.default')\n - `workload_token_exchange_enabled`: Whether SDK workload\ + \ identity token exchange is enabled\n - `workload_client_id`: OAuth client\ + \ ID to use for workload identity token exchange\n - `workload_token_endpoint`:\ + \ Token endpoint to use only for workload identity token exchange\n - `workload_audience`:\ + \ RFC 8693 audience for exchanged workload tokens\n - `workload_scope`: OAuth\ + \ scopes for exchanged workload tokens" + operationId: get_auth_discovery_endpoint_apis_auth_discovery_get responses: '200': description: Successful Response @@ -29,6 +34,89 @@ paths: application/json: schema: $ref: '#/components/schemas/AuthDiscoveryResponse' + /apis/auth/jwks: + get: + tags: + - Workload Identity + summary: Workload identity token exchange JWKS + description: Return the public signing key for workload identity access tokens + minted by the NeMo auth service. + operationId: jwks_apis_auth_jwks_get + responses: + '200': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/JsonWebKeySetResponse' + /apis/auth/token: + post: + tags: + - Workload Identity + summary: Exchange a workload identity subject token + description: Exchange a configured workload identity subject token for a NeMo + Platform access token. + operationId: token_exchange_apis_auth_token_post + requestBody: + content: + application/x-www-form-urlencoded: + schema: + properties: + grant_type: + type: string + enum: + - urn:ietf:params:oauth:grant-type:token-exchange + description: OAuth 2.0 token exchange grant type. + client_id: + type: string + description: Workload token exchange OAuth client ID. + subject_token: + type: string + description: JWT subject token to exchange. + subject_token_type: + type: string + enum: + - urn:ietf:params:oauth:token-type:jwt + description: Token type identifier for the subject token. + requested_token_type: + type: string + enum: + - urn:ietf:params:oauth:token-type:access_token + description: Requested token type identifier for the issued token. + default: urn:ietf:params:oauth:token-type:access_token + audience: + type: string + description: Requested audience for the issued access token. + scope: + type: string + description: Space-separated scopes requested for the issued access + token. + type: object + required: + - grant_type + - client_id + - subject_token + - subject_token_type + required: true + responses: + '200': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/WorkloadTokenExchangeResponse' + '400': + description: RFC 8693 token exchange error + content: + application/json: + schema: + $ref: '#/components/schemas/WorkloadTokenExchangeErrorResponse' + '401': + description: OAuth 2.0 invalid_client error + content: + application/json: + schema: + $ref: '#/components/schemas/WorkloadTokenExchangeErrorResponse' /apis/auth/v2/iam/role-bindings: get: tags: @@ -9979,6 +10067,10 @@ components: allOf: - $ref: '#/components/schemas/DockerJobNetworkConfig' description: Docker networking configuration + workload_identity: + allOf: + - $ref: '#/components/schemas/DockerWorkloadIdentityConfig' + description: Docker workload identity subject-token issuer configuration. type: object title: DockerJobExecutionProfileConfig description: Configuration for Docker Job execution profile. @@ -10048,6 +10140,60 @@ components: - volume_name - mount_path title: DockerVolumeMount + DockerWorkloadIdentityConfig: + properties: + enabled: + title: Enabled + description: Enable Docker workload identity token-file injection. Defaults + to auth.oidc.workload_token_exchange_enabled. + type: boolean + token_endpoint: + title: Token Endpoint + description: OAuth token endpoint used by the Docker demo issuer. Defaults + to auth.oidc.token_endpoint. + type: string + client_id: + title: Client Id + description: OAuth client ID used by the Docker demo issuer. Defaults to + auth.oidc.workload_client_id or auth.oidc.client_id. + type: string + client_secret: + format: password + title: Client Secret + description: OAuth client secret for the Docker demo issuer. + writeOnly: true + type: string + username: + title: Username + description: Username for the Docker demo issuer password grant. + type: string + password_env_var: + type: string + title: Password Env Var + description: Controller environment variable that contains the Docker demo + issuer password grant shared secret. + default: AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD + scope: + title: Scope + description: OAuth scope for the Docker demo issuer. + type: string + subject_token_ttl_seconds: + type: integer + minimum: 1.0 + title: Subject Token Ttl Seconds + description: Fallback subject-token lifetime when the Docker demo issuer + response omits expires_in. + refresh_margin_seconds: + type: integer + minimum: 0.0 + title: Refresh Margin Seconds + description: Seconds before subject-token expiry when the Docker refresher + issues a replacement token. + default: 60 + additionalProperties: false + type: object + title: DockerWorkloadIdentityConfig + description: Docker-only subject token issuer configuration for workload identity. E2EJobExecutionProfile: properties: provider: @@ -12591,6 +12737,25 @@ components: these variables. type: object title: JobExecutionProfileConfig + JsonWebKey: + properties: {} + additionalProperties: true + type: object + title: JsonWebKey + description: JSON Web Key object. + JsonWebKeySetResponse: + properties: + keys: + items: + $ref: '#/components/schemas/JsonWebKey' + type: array + title: Keys + description: Public signing keys in the JWKS document. + type: object + required: + - keys + title: JsonWebKeySetResponse + description: JSON Web Key Set document. K8sNIMOperatorConfig: properties: resources: @@ -12630,6 +12795,31 @@ components: These fields provide typed access to commonly-used NIMService Spec fields and are applied before override_config in the compilation precedence.' + KubernetesConfigMapVolume: + properties: + name: + type: string + title: Name + description: ConfigMap name to mount + default_mode: + title: Default Mode + description: Optional default file mode + type: integer + optional: + title: Optional + description: Whether the ConfigMap is optional + type: boolean + items: + items: + $ref: '#/components/schemas/KubernetesKeyToPath' + type: array + title: Items + description: Optional ConfigMap keys to project + type: object + required: + - name + title: KubernetesConfigMapVolume + description: Kubernetes ConfigMap volume definition. KubernetesEmptyDirVolume: properties: medium: @@ -12787,6 +12977,18 @@ components: title: Launcher Image description: Container image that contains the jobs-launcher binary. default: nvcr.io/nvidia/nemo-microservices/jobs-launcher:latest + workload_identity_token_expiration_seconds: + type: integer + minimum: 600.0 + title: Workload Identity Token Expiration Seconds + description: Requested expirationSeconds for the projected service account + token used as the workload identity subject token. + default: 600 + workload_identity_token_audience: + title: Workload Identity Token Audience + description: Audience for the projected service account token. Defaults + to auth.oidc.workload_client_id, auth.oidc.client_id, then 'nemo-platform'. + type: string type: object title: KubernetesJobExecutionProfileConfig description: Configuration for Kubernetes execution environment. @@ -12817,6 +13019,26 @@ components: type: object title: KubernetesJobStorageConfig description: Configuration for persistent storage in Kubernetes jobs. + KubernetesKeyToPath: + properties: + key: + type: string + title: Key + description: Source key to project from the volume source + path: + type: string + title: Path + description: Relative file path to write the key to + mode: + title: Mode + description: Optional file mode for this key + type: integer + type: object + required: + - key + - path + title: KubernetesKeyToPath + description: Kubernetes volume key-to-path mapping. KubernetesObjectMetadata: properties: labels: @@ -12847,7 +13069,57 @@ components: - claim_name title: KubernetesPersistentVolumeClaim description: Kubernetes Persistent Volume Claim definition. + KubernetesSecretVolume: + properties: + secret_name: + type: string + title: Secret Name + description: Secret name to mount + default_mode: + title: Default Mode + description: Optional default file mode + type: integer + optional: + title: Optional + description: Whether the Secret is optional + type: boolean + items: + items: + $ref: '#/components/schemas/KubernetesKeyToPath' + type: array + title: Items + description: Optional Secret keys to project + type: object + required: + - secret_name + title: KubernetesSecretVolume + description: Kubernetes Secret volume definition. KubernetesVolume: + oneOf: + - properties: + persistent_volume_claim: + not: + type: 'null' + required: + - persistent_volume_claim + - properties: + empty_dir: + not: + type: 'null' + required: + - empty_dir + - properties: + secret: + not: + type: 'null' + required: + - secret + - properties: + config_map: + not: + type: 'null' + required: + - config_map properties: name: type: string @@ -12861,11 +13133,19 @@ components: allOf: - $ref: '#/components/schemas/KubernetesEmptyDirVolume' description: EmptyDir Volume configuration + secret: + allOf: + - $ref: '#/components/schemas/KubernetesSecretVolume' + description: Secret Volume configuration + config_map: + allOf: + - $ref: '#/components/schemas/KubernetesConfigMapVolume' + description: ConfigMap Volume configuration type: object required: - name title: KubernetesVolume - description: Kubernetes Volume definition. + description: Kubernetes Volume definition with secret and config_map support. KubernetesVolumeMount: properties: name: @@ -14759,6 +15039,22 @@ components: scope_prefix: title: Scope Prefix type: string + workload_token_exchange_enabled: + type: boolean + title: Workload Token Exchange Enabled + default: false + workload_client_id: + title: Workload Client Id + type: string + workload_token_endpoint: + title: Workload Token Endpoint + type: string + workload_audience: + title: Workload Audience + type: string + workload_scope: + title: Workload Scope + type: string type: object required: - issuer @@ -19078,6 +19374,18 @@ components: title: Launcher Image description: Container image that contains the jobs-launcher binary. default: nvcr.io/nvidia/nemo-microservices/jobs-launcher:latest + workload_identity_token_expiration_seconds: + type: integer + minimum: 600.0 + title: Workload Identity Token Expiration Seconds + description: Requested expirationSeconds for the projected service account + token used as the workload identity subject token. + default: 600 + workload_identity_token_audience: + title: Workload Identity Token Audience + description: Audience for the projected service account token. Defaults + to auth.oidc.workload_client_id, auth.oidc.client_id, then 'nemo-platform'. + type: string queue: type: string title: Queue @@ -19104,6 +19412,58 @@ components: type: object title: VolcanoJobExecutionProfileConfig description: Configuration for Volcano Job Execution Profile + WorkloadTokenExchangeErrorResponse: + properties: + error: + type: string + title: Error + description: OAuth 2.0 or RFC 8693 token exchange error code, such as invalid_client, + invalid_request, invalid_grant, invalid_scope, or invalid_target. + error_description: + title: Error Description + description: Human-readable ASCII text providing additional information + about the error. + type: string + error_uri: + title: Error Uri + description: URI identifying a human-readable web page with information + about the error. + type: string + type: object + required: + - error + title: WorkloadTokenExchangeErrorResponse + description: RFC 8693 token exchange error response. + WorkloadTokenExchangeResponse: + properties: + access_token: + type: string + title: Access Token + description: JWT access token minted for the workload identity. + issued_token_type: + type: string + title: Issued Token Type + description: Token type identifier for the issued token. + token_type: + type: string + title: Token Type + description: OAuth token type used in Authorization headers. + expires_in: + type: integer + title: Expires In + description: Lifetime of the access token in seconds. + scope: + title: Scope + description: Space-separated scopes granted to the access token. + type: string + type: object + required: + - access_token + - issued_token_type + - token_type + - expires_in + title: WorkloadTokenExchangeResponse + description: RFC 8693 token exchange response for workload identity access tokens. Workspace: properties: id: diff --git a/openapi/ga/openapi.yaml b/openapi/ga/openapi.yaml index 3db8849574..b0c624f38e 100644 --- a/openapi/ga/openapi.yaml +++ b/openapi/ga/openapi.yaml @@ -20,8 +20,13 @@ paths: \ CLI)\n - `userinfo_endpoint`: UserInfo endpoint\n - `client_id`: OAuth\ \ client ID to use\n - `default_scopes`: OAuth scopes to request during authentication\n\ \ - `scope_prefix`: Prefix to prepend to custom scopes (those with ':' or\ - \ '.default')" - operationId: get_auth_discovery_apis_auth_discovery_get + \ '.default')\n - `workload_token_exchange_enabled`: Whether SDK workload\ + \ identity token exchange is enabled\n - `workload_client_id`: OAuth client\ + \ ID to use for workload identity token exchange\n - `workload_token_endpoint`:\ + \ Token endpoint to use only for workload identity token exchange\n - `workload_audience`:\ + \ RFC 8693 audience for exchanged workload tokens\n - `workload_scope`: OAuth\ + \ scopes for exchanged workload tokens" + operationId: get_auth_discovery_endpoint_apis_auth_discovery_get responses: '200': description: Successful Response @@ -29,6 +34,89 @@ paths: application/json: schema: $ref: '#/components/schemas/AuthDiscoveryResponse' + /apis/auth/jwks: + get: + tags: + - Workload Identity + summary: Workload identity token exchange JWKS + description: Return the public signing key for workload identity access tokens + minted by the NeMo auth service. + operationId: jwks_apis_auth_jwks_get + responses: + '200': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/JsonWebKeySetResponse' + /apis/auth/token: + post: + tags: + - Workload Identity + summary: Exchange a workload identity subject token + description: Exchange a configured workload identity subject token for a NeMo + Platform access token. + operationId: token_exchange_apis_auth_token_post + requestBody: + content: + application/x-www-form-urlencoded: + schema: + properties: + grant_type: + type: string + enum: + - urn:ietf:params:oauth:grant-type:token-exchange + description: OAuth 2.0 token exchange grant type. + client_id: + type: string + description: Workload token exchange OAuth client ID. + subject_token: + type: string + description: JWT subject token to exchange. + subject_token_type: + type: string + enum: + - urn:ietf:params:oauth:token-type:jwt + description: Token type identifier for the subject token. + requested_token_type: + type: string + enum: + - urn:ietf:params:oauth:token-type:access_token + description: Requested token type identifier for the issued token. + default: urn:ietf:params:oauth:token-type:access_token + audience: + type: string + description: Requested audience for the issued access token. + scope: + type: string + description: Space-separated scopes requested for the issued access + token. + type: object + required: + - grant_type + - client_id + - subject_token + - subject_token_type + required: true + responses: + '200': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/WorkloadTokenExchangeResponse' + '400': + description: RFC 8693 token exchange error + content: + application/json: + schema: + $ref: '#/components/schemas/WorkloadTokenExchangeErrorResponse' + '401': + description: OAuth 2.0 invalid_client error + content: + application/json: + schema: + $ref: '#/components/schemas/WorkloadTokenExchangeErrorResponse' /apis/auth/v2/iam/role-bindings: get: tags: @@ -9979,6 +10067,10 @@ components: allOf: - $ref: '#/components/schemas/DockerJobNetworkConfig' description: Docker networking configuration + workload_identity: + allOf: + - $ref: '#/components/schemas/DockerWorkloadIdentityConfig' + description: Docker workload identity subject-token issuer configuration. type: object title: DockerJobExecutionProfileConfig description: Configuration for Docker Job execution profile. @@ -10048,6 +10140,60 @@ components: - volume_name - mount_path title: DockerVolumeMount + DockerWorkloadIdentityConfig: + properties: + enabled: + title: Enabled + description: Enable Docker workload identity token-file injection. Defaults + to auth.oidc.workload_token_exchange_enabled. + type: boolean + token_endpoint: + title: Token Endpoint + description: OAuth token endpoint used by the Docker demo issuer. Defaults + to auth.oidc.token_endpoint. + type: string + client_id: + title: Client Id + description: OAuth client ID used by the Docker demo issuer. Defaults to + auth.oidc.workload_client_id or auth.oidc.client_id. + type: string + client_secret: + format: password + title: Client Secret + description: OAuth client secret for the Docker demo issuer. + writeOnly: true + type: string + username: + title: Username + description: Username for the Docker demo issuer password grant. + type: string + password_env_var: + type: string + title: Password Env Var + description: Controller environment variable that contains the Docker demo + issuer password grant shared secret. + default: AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD + scope: + title: Scope + description: OAuth scope for the Docker demo issuer. + type: string + subject_token_ttl_seconds: + type: integer + minimum: 1.0 + title: Subject Token Ttl Seconds + description: Fallback subject-token lifetime when the Docker demo issuer + response omits expires_in. + refresh_margin_seconds: + type: integer + minimum: 0.0 + title: Refresh Margin Seconds + description: Seconds before subject-token expiry when the Docker refresher + issues a replacement token. + default: 60 + additionalProperties: false + type: object + title: DockerWorkloadIdentityConfig + description: Docker-only subject token issuer configuration for workload identity. E2EJobExecutionProfile: properties: provider: @@ -12591,6 +12737,25 @@ components: these variables. type: object title: JobExecutionProfileConfig + JsonWebKey: + properties: {} + additionalProperties: true + type: object + title: JsonWebKey + description: JSON Web Key object. + JsonWebKeySetResponse: + properties: + keys: + items: + $ref: '#/components/schemas/JsonWebKey' + type: array + title: Keys + description: Public signing keys in the JWKS document. + type: object + required: + - keys + title: JsonWebKeySetResponse + description: JSON Web Key Set document. K8sNIMOperatorConfig: properties: resources: @@ -12630,6 +12795,31 @@ components: These fields provide typed access to commonly-used NIMService Spec fields and are applied before override_config in the compilation precedence.' + KubernetesConfigMapVolume: + properties: + name: + type: string + title: Name + description: ConfigMap name to mount + default_mode: + title: Default Mode + description: Optional default file mode + type: integer + optional: + title: Optional + description: Whether the ConfigMap is optional + type: boolean + items: + items: + $ref: '#/components/schemas/KubernetesKeyToPath' + type: array + title: Items + description: Optional ConfigMap keys to project + type: object + required: + - name + title: KubernetesConfigMapVolume + description: Kubernetes ConfigMap volume definition. KubernetesEmptyDirVolume: properties: medium: @@ -12787,6 +12977,18 @@ components: title: Launcher Image description: Container image that contains the jobs-launcher binary. default: nvcr.io/nvidia/nemo-microservices/jobs-launcher:latest + workload_identity_token_expiration_seconds: + type: integer + minimum: 600.0 + title: Workload Identity Token Expiration Seconds + description: Requested expirationSeconds for the projected service account + token used as the workload identity subject token. + default: 600 + workload_identity_token_audience: + title: Workload Identity Token Audience + description: Audience for the projected service account token. Defaults + to auth.oidc.workload_client_id, auth.oidc.client_id, then 'nemo-platform'. + type: string type: object title: KubernetesJobExecutionProfileConfig description: Configuration for Kubernetes execution environment. @@ -12817,6 +13019,26 @@ components: type: object title: KubernetesJobStorageConfig description: Configuration for persistent storage in Kubernetes jobs. + KubernetesKeyToPath: + properties: + key: + type: string + title: Key + description: Source key to project from the volume source + path: + type: string + title: Path + description: Relative file path to write the key to + mode: + title: Mode + description: Optional file mode for this key + type: integer + type: object + required: + - key + - path + title: KubernetesKeyToPath + description: Kubernetes volume key-to-path mapping. KubernetesObjectMetadata: properties: labels: @@ -12847,7 +13069,57 @@ components: - claim_name title: KubernetesPersistentVolumeClaim description: Kubernetes Persistent Volume Claim definition. + KubernetesSecretVolume: + properties: + secret_name: + type: string + title: Secret Name + description: Secret name to mount + default_mode: + title: Default Mode + description: Optional default file mode + type: integer + optional: + title: Optional + description: Whether the Secret is optional + type: boolean + items: + items: + $ref: '#/components/schemas/KubernetesKeyToPath' + type: array + title: Items + description: Optional Secret keys to project + type: object + required: + - secret_name + title: KubernetesSecretVolume + description: Kubernetes Secret volume definition. KubernetesVolume: + oneOf: + - properties: + persistent_volume_claim: + not: + type: 'null' + required: + - persistent_volume_claim + - properties: + empty_dir: + not: + type: 'null' + required: + - empty_dir + - properties: + secret: + not: + type: 'null' + required: + - secret + - properties: + config_map: + not: + type: 'null' + required: + - config_map properties: name: type: string @@ -12861,11 +13133,19 @@ components: allOf: - $ref: '#/components/schemas/KubernetesEmptyDirVolume' description: EmptyDir Volume configuration + secret: + allOf: + - $ref: '#/components/schemas/KubernetesSecretVolume' + description: Secret Volume configuration + config_map: + allOf: + - $ref: '#/components/schemas/KubernetesConfigMapVolume' + description: ConfigMap Volume configuration type: object required: - name title: KubernetesVolume - description: Kubernetes Volume definition. + description: Kubernetes Volume definition with secret and config_map support. KubernetesVolumeMount: properties: name: @@ -14759,6 +15039,22 @@ components: scope_prefix: title: Scope Prefix type: string + workload_token_exchange_enabled: + type: boolean + title: Workload Token Exchange Enabled + default: false + workload_client_id: + title: Workload Client Id + type: string + workload_token_endpoint: + title: Workload Token Endpoint + type: string + workload_audience: + title: Workload Audience + type: string + workload_scope: + title: Workload Scope + type: string type: object required: - issuer @@ -19078,6 +19374,18 @@ components: title: Launcher Image description: Container image that contains the jobs-launcher binary. default: nvcr.io/nvidia/nemo-microservices/jobs-launcher:latest + workload_identity_token_expiration_seconds: + type: integer + minimum: 600.0 + title: Workload Identity Token Expiration Seconds + description: Requested expirationSeconds for the projected service account + token used as the workload identity subject token. + default: 600 + workload_identity_token_audience: + title: Workload Identity Token Audience + description: Audience for the projected service account token. Defaults + to auth.oidc.workload_client_id, auth.oidc.client_id, then 'nemo-platform'. + type: string queue: type: string title: Queue @@ -19104,6 +19412,58 @@ components: type: object title: VolcanoJobExecutionProfileConfig description: Configuration for Volcano Job Execution Profile + WorkloadTokenExchangeErrorResponse: + properties: + error: + type: string + title: Error + description: OAuth 2.0 or RFC 8693 token exchange error code, such as invalid_client, + invalid_request, invalid_grant, invalid_scope, or invalid_target. + error_description: + title: Error Description + description: Human-readable ASCII text providing additional information + about the error. + type: string + error_uri: + title: Error Uri + description: URI identifying a human-readable web page with information + about the error. + type: string + type: object + required: + - error + title: WorkloadTokenExchangeErrorResponse + description: RFC 8693 token exchange error response. + WorkloadTokenExchangeResponse: + properties: + access_token: + type: string + title: Access Token + description: JWT access token minted for the workload identity. + issued_token_type: + type: string + title: Issued Token Type + description: Token type identifier for the issued token. + token_type: + type: string + title: Token Type + description: OAuth token type used in Authorization headers. + expires_in: + type: integer + title: Expires In + description: Lifetime of the access token in seconds. + scope: + title: Scope + description: Space-separated scopes granted to the access token. + type: string + type: object + required: + - access_token + - issued_token_type + - token_type + - expires_in + title: WorkloadTokenExchangeResponse + description: RFC 8693 token exchange response for workload identity access tokens. Workspace: properties: id: diff --git a/openapi/openapi.yaml b/openapi/openapi.yaml index 3db8849574..b0c624f38e 100644 --- a/openapi/openapi.yaml +++ b/openapi/openapi.yaml @@ -20,8 +20,13 @@ paths: \ CLI)\n - `userinfo_endpoint`: UserInfo endpoint\n - `client_id`: OAuth\ \ client ID to use\n - `default_scopes`: OAuth scopes to request during authentication\n\ \ - `scope_prefix`: Prefix to prepend to custom scopes (those with ':' or\ - \ '.default')" - operationId: get_auth_discovery_apis_auth_discovery_get + \ '.default')\n - `workload_token_exchange_enabled`: Whether SDK workload\ + \ identity token exchange is enabled\n - `workload_client_id`: OAuth client\ + \ ID to use for workload identity token exchange\n - `workload_token_endpoint`:\ + \ Token endpoint to use only for workload identity token exchange\n - `workload_audience`:\ + \ RFC 8693 audience for exchanged workload tokens\n - `workload_scope`: OAuth\ + \ scopes for exchanged workload tokens" + operationId: get_auth_discovery_endpoint_apis_auth_discovery_get responses: '200': description: Successful Response @@ -29,6 +34,89 @@ paths: application/json: schema: $ref: '#/components/schemas/AuthDiscoveryResponse' + /apis/auth/jwks: + get: + tags: + - Workload Identity + summary: Workload identity token exchange JWKS + description: Return the public signing key for workload identity access tokens + minted by the NeMo auth service. + operationId: jwks_apis_auth_jwks_get + responses: + '200': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/JsonWebKeySetResponse' + /apis/auth/token: + post: + tags: + - Workload Identity + summary: Exchange a workload identity subject token + description: Exchange a configured workload identity subject token for a NeMo + Platform access token. + operationId: token_exchange_apis_auth_token_post + requestBody: + content: + application/x-www-form-urlencoded: + schema: + properties: + grant_type: + type: string + enum: + - urn:ietf:params:oauth:grant-type:token-exchange + description: OAuth 2.0 token exchange grant type. + client_id: + type: string + description: Workload token exchange OAuth client ID. + subject_token: + type: string + description: JWT subject token to exchange. + subject_token_type: + type: string + enum: + - urn:ietf:params:oauth:token-type:jwt + description: Token type identifier for the subject token. + requested_token_type: + type: string + enum: + - urn:ietf:params:oauth:token-type:access_token + description: Requested token type identifier for the issued token. + default: urn:ietf:params:oauth:token-type:access_token + audience: + type: string + description: Requested audience for the issued access token. + scope: + type: string + description: Space-separated scopes requested for the issued access + token. + type: object + required: + - grant_type + - client_id + - subject_token + - subject_token_type + required: true + responses: + '200': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/WorkloadTokenExchangeResponse' + '400': + description: RFC 8693 token exchange error + content: + application/json: + schema: + $ref: '#/components/schemas/WorkloadTokenExchangeErrorResponse' + '401': + description: OAuth 2.0 invalid_client error + content: + application/json: + schema: + $ref: '#/components/schemas/WorkloadTokenExchangeErrorResponse' /apis/auth/v2/iam/role-bindings: get: tags: @@ -9979,6 +10067,10 @@ components: allOf: - $ref: '#/components/schemas/DockerJobNetworkConfig' description: Docker networking configuration + workload_identity: + allOf: + - $ref: '#/components/schemas/DockerWorkloadIdentityConfig' + description: Docker workload identity subject-token issuer configuration. type: object title: DockerJobExecutionProfileConfig description: Configuration for Docker Job execution profile. @@ -10048,6 +10140,60 @@ components: - volume_name - mount_path title: DockerVolumeMount + DockerWorkloadIdentityConfig: + properties: + enabled: + title: Enabled + description: Enable Docker workload identity token-file injection. Defaults + to auth.oidc.workload_token_exchange_enabled. + type: boolean + token_endpoint: + title: Token Endpoint + description: OAuth token endpoint used by the Docker demo issuer. Defaults + to auth.oidc.token_endpoint. + type: string + client_id: + title: Client Id + description: OAuth client ID used by the Docker demo issuer. Defaults to + auth.oidc.workload_client_id or auth.oidc.client_id. + type: string + client_secret: + format: password + title: Client Secret + description: OAuth client secret for the Docker demo issuer. + writeOnly: true + type: string + username: + title: Username + description: Username for the Docker demo issuer password grant. + type: string + password_env_var: + type: string + title: Password Env Var + description: Controller environment variable that contains the Docker demo + issuer password grant shared secret. + default: AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD + scope: + title: Scope + description: OAuth scope for the Docker demo issuer. + type: string + subject_token_ttl_seconds: + type: integer + minimum: 1.0 + title: Subject Token Ttl Seconds + description: Fallback subject-token lifetime when the Docker demo issuer + response omits expires_in. + refresh_margin_seconds: + type: integer + minimum: 0.0 + title: Refresh Margin Seconds + description: Seconds before subject-token expiry when the Docker refresher + issues a replacement token. + default: 60 + additionalProperties: false + type: object + title: DockerWorkloadIdentityConfig + description: Docker-only subject token issuer configuration for workload identity. E2EJobExecutionProfile: properties: provider: @@ -12591,6 +12737,25 @@ components: these variables. type: object title: JobExecutionProfileConfig + JsonWebKey: + properties: {} + additionalProperties: true + type: object + title: JsonWebKey + description: JSON Web Key object. + JsonWebKeySetResponse: + properties: + keys: + items: + $ref: '#/components/schemas/JsonWebKey' + type: array + title: Keys + description: Public signing keys in the JWKS document. + type: object + required: + - keys + title: JsonWebKeySetResponse + description: JSON Web Key Set document. K8sNIMOperatorConfig: properties: resources: @@ -12630,6 +12795,31 @@ components: These fields provide typed access to commonly-used NIMService Spec fields and are applied before override_config in the compilation precedence.' + KubernetesConfigMapVolume: + properties: + name: + type: string + title: Name + description: ConfigMap name to mount + default_mode: + title: Default Mode + description: Optional default file mode + type: integer + optional: + title: Optional + description: Whether the ConfigMap is optional + type: boolean + items: + items: + $ref: '#/components/schemas/KubernetesKeyToPath' + type: array + title: Items + description: Optional ConfigMap keys to project + type: object + required: + - name + title: KubernetesConfigMapVolume + description: Kubernetes ConfigMap volume definition. KubernetesEmptyDirVolume: properties: medium: @@ -12787,6 +12977,18 @@ components: title: Launcher Image description: Container image that contains the jobs-launcher binary. default: nvcr.io/nvidia/nemo-microservices/jobs-launcher:latest + workload_identity_token_expiration_seconds: + type: integer + minimum: 600.0 + title: Workload Identity Token Expiration Seconds + description: Requested expirationSeconds for the projected service account + token used as the workload identity subject token. + default: 600 + workload_identity_token_audience: + title: Workload Identity Token Audience + description: Audience for the projected service account token. Defaults + to auth.oidc.workload_client_id, auth.oidc.client_id, then 'nemo-platform'. + type: string type: object title: KubernetesJobExecutionProfileConfig description: Configuration for Kubernetes execution environment. @@ -12817,6 +13019,26 @@ components: type: object title: KubernetesJobStorageConfig description: Configuration for persistent storage in Kubernetes jobs. + KubernetesKeyToPath: + properties: + key: + type: string + title: Key + description: Source key to project from the volume source + path: + type: string + title: Path + description: Relative file path to write the key to + mode: + title: Mode + description: Optional file mode for this key + type: integer + type: object + required: + - key + - path + title: KubernetesKeyToPath + description: Kubernetes volume key-to-path mapping. KubernetesObjectMetadata: properties: labels: @@ -12847,7 +13069,57 @@ components: - claim_name title: KubernetesPersistentVolumeClaim description: Kubernetes Persistent Volume Claim definition. + KubernetesSecretVolume: + properties: + secret_name: + type: string + title: Secret Name + description: Secret name to mount + default_mode: + title: Default Mode + description: Optional default file mode + type: integer + optional: + title: Optional + description: Whether the Secret is optional + type: boolean + items: + items: + $ref: '#/components/schemas/KubernetesKeyToPath' + type: array + title: Items + description: Optional Secret keys to project + type: object + required: + - secret_name + title: KubernetesSecretVolume + description: Kubernetes Secret volume definition. KubernetesVolume: + oneOf: + - properties: + persistent_volume_claim: + not: + type: 'null' + required: + - persistent_volume_claim + - properties: + empty_dir: + not: + type: 'null' + required: + - empty_dir + - properties: + secret: + not: + type: 'null' + required: + - secret + - properties: + config_map: + not: + type: 'null' + required: + - config_map properties: name: type: string @@ -12861,11 +13133,19 @@ components: allOf: - $ref: '#/components/schemas/KubernetesEmptyDirVolume' description: EmptyDir Volume configuration + secret: + allOf: + - $ref: '#/components/schemas/KubernetesSecretVolume' + description: Secret Volume configuration + config_map: + allOf: + - $ref: '#/components/schemas/KubernetesConfigMapVolume' + description: ConfigMap Volume configuration type: object required: - name title: KubernetesVolume - description: Kubernetes Volume definition. + description: Kubernetes Volume definition with secret and config_map support. KubernetesVolumeMount: properties: name: @@ -14759,6 +15039,22 @@ components: scope_prefix: title: Scope Prefix type: string + workload_token_exchange_enabled: + type: boolean + title: Workload Token Exchange Enabled + default: false + workload_client_id: + title: Workload Client Id + type: string + workload_token_endpoint: + title: Workload Token Endpoint + type: string + workload_audience: + title: Workload Audience + type: string + workload_scope: + title: Workload Scope + type: string type: object required: - issuer @@ -19078,6 +19374,18 @@ components: title: Launcher Image description: Container image that contains the jobs-launcher binary. default: nvcr.io/nvidia/nemo-microservices/jobs-launcher:latest + workload_identity_token_expiration_seconds: + type: integer + minimum: 600.0 + title: Workload Identity Token Expiration Seconds + description: Requested expirationSeconds for the projected service account + token used as the workload identity subject token. + default: 600 + workload_identity_token_audience: + title: Workload Identity Token Audience + description: Audience for the projected service account token. Defaults + to auth.oidc.workload_client_id, auth.oidc.client_id, then 'nemo-platform'. + type: string queue: type: string title: Queue @@ -19104,6 +19412,58 @@ components: type: object title: VolcanoJobExecutionProfileConfig description: Configuration for Volcano Job Execution Profile + WorkloadTokenExchangeErrorResponse: + properties: + error: + type: string + title: Error + description: OAuth 2.0 or RFC 8693 token exchange error code, such as invalid_client, + invalid_request, invalid_grant, invalid_scope, or invalid_target. + error_description: + title: Error Description + description: Human-readable ASCII text providing additional information + about the error. + type: string + error_uri: + title: Error Uri + description: URI identifying a human-readable web page with information + about the error. + type: string + type: object + required: + - error + title: WorkloadTokenExchangeErrorResponse + description: RFC 8693 token exchange error response. + WorkloadTokenExchangeResponse: + properties: + access_token: + type: string + title: Access Token + description: JWT access token minted for the workload identity. + issued_token_type: + type: string + title: Issued Token Type + description: Token type identifier for the issued token. + token_type: + type: string + title: Token Type + description: OAuth token type used in Authorization headers. + expires_in: + type: integer + title: Expires In + description: Lifetime of the access token in seconds. + scope: + title: Scope + description: Space-separated scopes granted to the access token. + type: string + type: object + required: + - access_token + - issued_token_type + - token_type + - expires_in + title: WorkloadTokenExchangeResponse + description: RFC 8693 token exchange response for workload identity access tokens. Workspace: properties: id: diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/auth/device_flow.py b/packages/nemo_platform_ext/src/nemo_platform_ext/auth/device_flow.py index 2293b60ff5..d45ed21ea5 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/auth/device_flow.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/auth/device_flow.py @@ -13,6 +13,7 @@ from rich.panel import Panel from nemo_platform_ext.auth.token_provider import refresh_token_grant +from nemo_platform_ext.client.tls import client_verify_from_env console = Console() @@ -73,7 +74,7 @@ def __init__( async def start_device_authorization(self) -> DeviceCodeResponse: """Start the device authorization flow.""" - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(verify=client_verify_from_env()) as client: response = await client.post( self.device_authorization_endpoint, data={ @@ -103,7 +104,7 @@ async def poll_for_token( """Poll the token endpoint until authorization is complete.""" start_time = time.time() - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(verify=client_verify_from_env()) as client: while time.time() - start_time < expires_in: await _async_pause(interval) @@ -283,7 +284,7 @@ def authenticate_with_password_grant( "password": password, "scope": scope, } - with httpx.Client() as client: + with httpx.Client(verify=client_verify_from_env()) as client: response = client.post(token_endpoint, data=data, timeout=30.0) if response.status_code != 200: diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/auth/helpers.py b/packages/nemo_platform_ext/src/nemo_platform_ext/auth/helpers.py index 8155ab6aa6..9c3281afe0 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/auth/helpers.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/auth/helpers.py @@ -25,6 +25,8 @@ import httpx +from nemo_platform_ext.client.tls import client_verify_from_env + DEFAULT_OAUTH_SCOPES = "openid profile email offline_access" @@ -143,6 +145,11 @@ class NMPOIDCConfig: device_authorization_endpoint: str | None = None default_scopes: str = DEFAULT_OAUTH_SCOPES scope_prefix: str | None = None + workload_token_exchange_enabled: bool = False + workload_client_id: str | None = None + workload_token_endpoint: str | None = None + workload_audience: str | None = None + workload_scope: str | None = None def discover_nmp_config(base_url: str, timeout: float = 10.0) -> NMPOIDCConfig: @@ -150,6 +157,7 @@ def discover_nmp_config(base_url: str, timeout: float = 10.0) -> NMPOIDCConfig: response = httpx.get( f"{base_url.rstrip('/')}/apis/auth/discovery", timeout=timeout, + verify=client_verify_from_env(), ) response.raise_for_status() data = response.json() @@ -163,6 +171,11 @@ def discover_nmp_config(base_url: str, timeout: float = 10.0) -> NMPOIDCConfig: device_authorization_endpoint=oidc.get("device_authorization_endpoint"), default_scopes=oidc.get("default_scopes", DEFAULT_OAUTH_SCOPES), scope_prefix=oidc.get("scope_prefix"), + workload_token_exchange_enabled=oidc.get("workload_token_exchange_enabled", False), + workload_client_id=oidc.get("workload_client_id"), + workload_token_endpoint=oidc.get("workload_token_endpoint"), + workload_audience=oidc.get("workload_audience"), + workload_scope=oidc.get("workload_scope"), ) diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/auth/token_provider.py b/packages/nemo_platform_ext/src/nemo_platform_ext/auth/token_provider.py index 17ef2fab28..bcc7b7c5c8 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/auth/token_provider.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/auth/token_provider.py @@ -16,6 +16,7 @@ from typing_extensions import Self from nemo_platform_ext.auth.helpers import decode_jwt_claims +from nemo_platform_ext.client.tls import client_verify_from_env logger = logging.getLogger(__name__) @@ -32,6 +33,12 @@ def __init__(self, *, error: str, error_description: str) -> None: super().__init__(f"Token refresh failed: {error} - {error_description}") +def _validate_expires_in(expires_in: object) -> int | float | None: + if isinstance(expires_in, bool): + return None + return expires_in if isinstance(expires_in, int | float) else None + + def refresh_token_grant( token_endpoint: str, client_id: str, @@ -49,7 +56,7 @@ def refresh_token_grant( if scope: data["scope"] = scope - response = httpx.post(token_endpoint, data=data, timeout=timeout) + response = httpx.post(token_endpoint, data=data, timeout=timeout, verify=client_verify_from_env()) if response.status_code != 200: error_data: dict[str, str] = {} @@ -77,12 +84,16 @@ class TokenSet: def from_access_token( access_token: str, refresh_token: str | None = None, + expires_in: object = None, ) -> Self: """Create a TokenSet, extracting expiry from the JWT's `exp` claim.""" expires_at = None claims = decode_jwt_claims(access_token) if claims: expires_at = claims.get("exp") + validated_expires_in = _validate_expires_in(expires_in) + if expires_at is None and validated_expires_in is not None: + expires_at = time.time() + float(validated_expires_in) return TokenSet( access_token=access_token, refresh_token=refresh_token, @@ -223,7 +234,11 @@ def _refresh(self, *, force: bool = False) -> None: # The IdP may rotate the refresh token. new_refresh_token = token_data.get("refresh_token", self.tokens.refresh_token) - self.tokens = TokenSet.from_access_token(new_access_token, new_refresh_token) + self.tokens = TokenSet.from_access_token( + new_access_token, + new_refresh_token, + expires_in=token_data.get("expires_in"), + ) logger.debug("Access token refreshed successfully (expires_at=%s)", self.tokens.expires_at) if self.on_tokens_refreshed: diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/auth/workload_exchange.py b/packages/nemo_platform_ext/src/nemo_platform_ext/auth/workload_exchange.py new file mode 100644 index 0000000000..2e9fd8a332 --- /dev/null +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/auth/workload_exchange.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Workload identity token exchange for SDK authentication.""" + +from __future__ import annotations + +import asyncio +import json +import logging +import math +import threading +from dataclasses import dataclass, field +from ipaddress import ip_address +from pathlib import Path +from urllib.parse import urlparse + +import httpx +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR + +from nemo_platform_ext.auth.token_provider import DEFAULT_REFRESH_MARGIN_SECONDS, TokenSet +from nemo_platform_ext.client.tls import client_verify_from_env + +logger = logging.getLogger(__name__) + +TOKEN_EXCHANGE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange" +JWT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt" +ACCESS_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:access_token" + + +class WorkloadTokenExchangeError(RuntimeError): + """Structured error raised for RFC 8693 workload token exchange failures.""" + + def __init__(self, *, error: str, error_description: str) -> None: + self.error = error + self.error_description = error_description + super().__init__(f"Workload token exchange failed: {error} - {error_description}") + + +def read_subject_token_file(path: Path) -> str: + """Read a subject token from a workload identity token file.""" + try: + token = path.read_text(encoding="utf-8").strip() + except OSError as exc: + raise ValueError(f"Unable to read {WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR} at {path}: {exc}") from exc + if not token: + raise ValueError(f"{WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR} at {path} is empty") + return token + + +def _is_loopback_host(hostname: str | None) -> bool: + if hostname == "localhost": + return True + if hostname is None: + return False + try: + return ip_address(hostname).is_loopback + except ValueError: + return False + + +def _validate_token_endpoint(token_endpoint: str) -> None: + """Reject non-HTTPS token endpoints (except loopback for local dev).""" + parsed = urlparse(token_endpoint) + if parsed.scheme == "https": + return + if parsed.scheme == "http" and _is_loopback_host(parsed.hostname): + return + raise ValueError( + f"OIDC token endpoint must use HTTPS (got {token_endpoint!r}). " + "HTTP is only allowed for loopback addresses (localhost, 127.0.0.1, ::1)." + ) + + +def token_exchange_grant( + *, + token_endpoint: str, + client_id: str, + subject_token: str, + audience: str | None = None, + scope: str | None = None, + timeout: float = 30.0, +) -> dict[str, object]: + """Execute RFC 8693 token exchange and return token response JSON.""" + _validate_token_endpoint(token_endpoint) + data: dict[str, str] = { + "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, + "client_id": client_id, + "subject_token": subject_token, + "subject_token_type": JWT_TOKEN_TYPE, + "requested_token_type": ACCESS_TOKEN_TYPE, + } + if audience: + data["audience"] = audience + if scope: + data["scope"] = scope + + response = httpx.post(token_endpoint, data=data, timeout=timeout, verify=client_verify_from_env()) + + if response.status_code != 200: + error_data: dict[str, object] = {} + if response.headers.get("content-type", "").startswith("application/json"): + error_data = _response_json_object( + response, + error_description="Token endpoint error response was not a JSON object", + ) + error = _response_string(error_data, "error", "unknown_error") + error_description = _response_string(error_data, "error_description", response.text) + raise WorkloadTokenExchangeError(error=error, error_description=error_description) + + token_data = _response_json_object( + response, + error_description="Token endpoint response was not a JSON object", + ) + _access_token_from_response(token_data) + return token_data + + +def _response_json_object(response: httpx.Response, *, error_description: str) -> dict[str, object]: + try: + payload = response.json() + except (json.JSONDecodeError, ValueError) as exc: + raise WorkloadTokenExchangeError(error="invalid_response", error_description=error_description) from exc + if not isinstance(payload, dict): + raise WorkloadTokenExchangeError(error="invalid_response", error_description=error_description) + return payload + + +def _response_string(payload: dict[str, object], key: str, default: str) -> str: + value = payload.get(key) + return value if isinstance(value, str) and value else default + + +def _access_token_from_response(token_data: dict[str, object]) -> str: + access_token = token_data.get("access_token") + if not isinstance(access_token, str) or not access_token.strip(): + raise WorkloadTokenExchangeError( + error="invalid_response", + error_description="Token endpoint response did not include a non-empty access_token", + ) + return access_token + + +def _expires_in_from_response(token_data: dict[str, object]) -> int | float | None: + expires_in = token_data.get("expires_in") + if isinstance(expires_in, bool): + return None + return expires_in if isinstance(expires_in, int | float) else None + + +@dataclass +class WorkloadTokenExchangeProvider: + """Provides access tokens by exchanging a workload identity subject token file.""" + + token_endpoint: str + client_id: str + subject_token_file: Path + audience: str | None = None + scope: str | None = None + refresh_margin_seconds: float = DEFAULT_REFRESH_MARGIN_SECONDS + tokens: TokenSet = field(default_factory=lambda: TokenSet(access_token="")) + _lock: threading.Lock = field(default_factory=threading.Lock, repr=False) + + def get_access_token(self) -> str: + """Return a valid access token, exchanging the current subject token if needed.""" + with self._lock: + if not self.tokens.access_token or self.tokens.is_expired(self.refresh_margin_seconds): + self._exchange() + return self.tokens.access_token + + async def get_access_token_async(self) -> str: + """Return a valid access token in async contexts.""" + return await asyncio.to_thread(self.get_access_token) + + def _exchange(self) -> None: + subject_token = read_subject_token_file(self.subject_token_file) + logger.debug("Exchanging workload identity token via %s", self.token_endpoint) + token_data = token_exchange_grant( + token_endpoint=self.token_endpoint, + client_id=self.client_id, + subject_token=subject_token, + audience=self.audience, + scope=self.scope, + ) + access_token = _access_token_from_response(token_data) + try: + tokens = TokenSet.from_access_token( + access_token, + refresh_token=None, + expires_in=_expires_in_from_response(token_data), + ) + except (OverflowError, TypeError, ValueError) as exc: + raise WorkloadTokenExchangeError( + error="invalid_response", + error_description="Token endpoint response did not include a usable access_token lifetime", + ) from exc + if tokens.expires_at is None or not math.isfinite(tokens.expires_at): + raise WorkloadTokenExchangeError( + error="invalid_response", + error_description="Token endpoint response did not include a usable access_token lifetime", + ) + if tokens.is_expired(0): + raise WorkloadTokenExchangeError( + error="invalid_response", + error_description="Token endpoint response returned an expired access_token", + ) + self.tokens = tokens diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/auth.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/auth.py index 888dd8ebc0..235f2e520c 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/auth.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/auth.py @@ -10,6 +10,7 @@ from __future__ import annotations import asyncio +import logging import time from typing import Annotated, cast @@ -37,17 +38,19 @@ help="Manage authentication for NeMo Platform.", ) +logger = logging.getLogger(__name__) -def is_auth_disabled(base_url: str, timeout: float = 3.0) -> bool | None: + +def is_auth_disabled(base_url: str, timeout: float = 3.0) -> bool: """Check whether authentication is disabled on the cluster. Returns: - True if auth is definitely disabled, False if enabled, None if unreachable. + True if auth is disabled, False if enabled. """ try: return not discover_nmp_config(base_url, timeout=timeout).auth_enabled - except httpx.HTTPError: - return None + except httpx.HTTPError as exc: + raise AuthError(f"Failed to discover auth configuration: {exc}") from exc def _runtime_token_source_label() -> str | None: @@ -57,7 +60,8 @@ def _runtime_token_source_label() -> str | None: try: return Config.runtime_access_token_source_label() except ValueError: - return "NEMO_WORKLOAD_TOKEN_FILE environment override could not be read" + logger.debug("Failed to resolve runtime token override source label", exc_info=True) + return None def ensure_valid_token(context: Context, refresh_buffer_seconds: int = 300) -> bool: @@ -131,13 +135,13 @@ def ensure_valid_token(context: Context, refresh_buffer_seconds: int = 300) -> b ) provider.force_refresh() - config_params = { + config_params: ConfigParams = { "access_token": provider.tokens.access_token, } if provider.tokens.refresh_token: config_params["refresh_token"] = provider.tokens.refresh_token - Config.write(config_params, context_name=context.context_name) # type: ignore[arg-type] + Config.write(config_params, context_name=context.context_name) typer.echo("[Auto-refreshed expired token]", err=True) return True @@ -282,7 +286,7 @@ def login( from nemo_platform_ext.config.config import Config cli_context: CLIContext = ctx.obj - selected_context = cast(str | None, cli_context.overrides.get("current_context")) + selected_context = cli_context.overrides.get("current_context") if context_name is not None: selected_context = context_name @@ -537,9 +541,13 @@ def logout(ctx: typer.Context) -> None: console = Console() base_url = str(context.cluster.base_url).rstrip("/") - if is_auth_disabled(base_url) is True: - console.print("[yellow]Authentication is disabled on this cluster — nothing to log out from.[/]") - return + try: + if is_auth_disabled(base_url) is True: + console.print("[yellow]Authentication is disabled on this cluster — nothing to log out from.[/]") + return + except AuthError as exc: + logger.debug("Failed to discover auth configuration during logout", exc_info=True) + console.print(f"[yellow]Warning:[/] {exc}; continuing to clear local credentials.") logout_params: ConfigParams = {"access_token": None, "refresh_token": None} updated_config = Config.write(logout_params, context_name=context.context_name) @@ -699,10 +707,10 @@ def refresh(ctx: typer.Context) -> None: raise AuthError(f"Token refresh failed: {e}") from e # Save new tokens (refresh token may be rotated) - config_params = {"access_token": provider.tokens.access_token} + config_params: ConfigParams = {"access_token": provider.tokens.access_token} if provider.tokens.refresh_token: config_params["refresh_token"] = provider.tokens.refresh_token - Config.write(config_params, context_name=context.context_name) # type: ignore[arg-type] + Config.write(config_params, context_name=context.context_name) # Show new token info claims = decode_jwt_claims(provider.tokens.access_token) @@ -763,14 +771,21 @@ def status(ctx: typer.Context) -> None: # Check whether the cluster has auth enabled before showing token details. base_url = str(context.cluster.base_url).rstrip("/") - if is_auth_disabled(base_url) is True: - console.print() - console.print(f"[cyan]Cluster:[/] {base_url}") - console.print(f"[cyan]Context:[/] {context.context_name}") - console.print() - console.print("[green]Authentication is disabled on this cluster.[/]") - console.print("All API requests are accepted without credentials.") - return + auth_discovery_error: AuthError | None = None + try: + auth_disabled = is_auth_disabled(base_url) + except AuthError as exc: + logger.debug("Failed to discover auth configuration during status", exc_info=True) + auth_discovery_error = exc + else: + if auth_disabled is True: + console.print() + console.print(f"[cyan]Cluster:[/] {base_url}") + console.print(f"[cyan]Context:[/] {context.context_name}") + console.print() + console.print("[green]Authentication is disabled on this cluster.[/]") + console.print("All API requests are accepted without credentials.") + return table = Table(title="Authentication Status", show_header=False) table.add_column("Property", style="cyan") @@ -779,6 +794,8 @@ def status(ctx: typer.Context) -> None: table.add_row("Cluster", str(context.cluster.base_url)) table.add_row("Context", context.context_name) table.add_row("Config File", str(Config.get_default_config_path())) + if auth_discovery_error is not None: + table.add_row("Auth Discovery", f"[yellow]unavailable[/] ({auth_discovery_error})") runtime_token_source = _runtime_token_source_label() if context.user: diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/context.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/context.py index a1e44a016b..d636645a9d 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/context.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/core/context.py @@ -6,6 +6,7 @@ from __future__ import annotations import logging +import os import typing from dataclasses import dataclass, field @@ -73,6 +74,33 @@ def get_sdk_context(self) -> Context: def reset_sdk_context(self) -> None: self._sdk_context = None + def _context_exists_in_config_file(self, context_name: str) -> bool: + from nemo_platform_ext.config.config import Config + + try: + config = Config.load(overrides=self.overrides) + except FileNotFoundError: + return False + + return any(ctx.name == context_name for ctx in config.get_config_file().contexts) + + def _client_auth_config(self, ctx: Context) -> dict[str, object]: + from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR + + from nemo_platform_ext.config.config import Config + from nemo_platform_ext.config.models import OAuthUser + + if self.overrides.get("access_token") is not None or Config.runtime_access_token_source_label(): + return ctx.user.get_client_config() if ctx.user else {} + + if os.environ.get(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR): + return {} + + if isinstance(ctx.user, OAuthUser) and self._context_exists_in_config_file(ctx.context_name): + return {"context_name": ctx.context_name} + + return ctx.user.get_client_config() if ctx.user else {} + def get_client(self, timeout: float = 60.0) -> NeMoPlatform: """ Get or create the NeMo Platform client. @@ -92,12 +120,12 @@ def get_client(self, timeout: float = 60.0) -> NeMoPlatform: f"Creating NeMoPlatform client with base_url={base_url}, workspace={ctx.workspace}, timeout={timeout}" ) - client_config = ctx.user.get_client_config() + auth_config = self._client_auth_config(ctx) self._client = NeMoPlatform( base_url=base_url, timeout=timeout, workspace=ctx.workspace, - **client_config, + **auth_config, ) return self._client @@ -120,12 +148,12 @@ def get_async_client(self, timeout: float = 60.0) -> AsyncNeMoPlatform: f"Creating AsyncNeMoPlatform client with base_url={base_url}, workspace={ctx.workspace}, timeout={timeout}" ) - client_config = ctx.user.get_client_config() + auth_config = self._client_auth_config(ctx) self._async_client = AsyncNeMoPlatform( base_url=base_url, timeout=timeout, workspace=ctx.workspace, - **client_config, + **auth_config, ) return self._async_client diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/client/enhanced.py b/packages/nemo_platform_ext/src/nemo_platform_ext/client/enhanced.py index c050c00e49..af6451b3ed 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/client/enhanced.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/client/enhanced.py @@ -5,13 +5,48 @@ from __future__ import annotations +import os from pathlib import Path from typing import Any, Mapping import httpx from httpx import Timeout -from nemo_platform import DEFAULT_MAX_RETRIES, AsyncStream, NotGiven, __version__, not_given +from nemo_platform import ( + DEFAULT_MAX_RETRIES, + AsyncStream, + DefaultAsyncHttpxClient, + DefaultHttpxClient, + NotGiven, + __version__, + not_given, +) from nemo_platform._base_client import AsyncAPIClient, SyncAPIClient +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR + +from nemo_platform_ext.client.tls import client_verify_from_env + + +def _should_bootstrap_config( + *, + http_client: object | None, + base_url: str | httpx.URL | None, + config_path: Path | None, + context_name: str | None, + access_token: str | None, +) -> bool: + """Return whether constructor arguments require config/auth bootstrap.""" + if http_client is not None: + return False + + # Backward compatibility: an explicit base_url means direct mode (no config + # bootstrap), unless config-specific overrides or workload identity are set. + return ( + base_url is None + or config_path is not None + or context_name is not None + or access_token is not None + or bool(os.environ.get(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR)) + ) class NeMoPlatform(SyncAPIClient): @@ -98,10 +133,12 @@ def __init__( http_client: Custom ``httpx.Client`` instance. When provided, the auth bootstrap is skipped entirely regardless of other parameters. """ - # Backward compatibility: an explicit base_url means direct mode (no config bootstrap), - # unless config-specific overrides are provided. - should_bootstrap = http_client is None and ( - base_url is None or config_path is not None or context_name is not None or access_token is not None + should_bootstrap = _should_bootstrap_config( + http_client=http_client, + base_url=base_url, + config_path=config_path, + context_name=context_name, + access_token=access_token, ) if should_bootstrap: try: @@ -122,6 +159,10 @@ def __init__( except Exception as e: raise RuntimeError(f"NeMoPlatform client initialization failed: {e}") + client_verify = client_verify_from_env() + if http_client is None and client_verify is not True: + http_client = DefaultHttpxClient(verify=client_verify) + self.workspace = workspace super().__init__( @@ -254,10 +295,12 @@ async def main() -> None: http_client: Custom ``httpx.AsyncClient`` instance. When provided, the auth bootstrap is skipped entirely regardless of other parameters. """ - # Backward compatibility: an explicit base_url means direct mode (no config bootstrap), - # unless config-specific overrides are provided. - should_bootstrap = http_client is None and ( - base_url is None or config_path is not None or context_name is not None or access_token is not None + should_bootstrap = _should_bootstrap_config( + http_client=http_client, + base_url=base_url, + config_path=config_path, + context_name=context_name, + access_token=access_token, ) if should_bootstrap: try: @@ -278,6 +321,10 @@ async def main() -> None: except Exception as e: raise RuntimeError(f"NeMoPlatform client initialization failed: {e}") + client_verify = client_verify_from_env() + if http_client is None and client_verify is not True: + http_client = DefaultAsyncHttpxClient(verify=client_verify) + self.workspace = workspace super().__init__( diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/client/factory.py b/packages/nemo_platform_ext/src/nemo_platform_ext/client/factory.py index 5d0505727a..99594a41cc 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/client/factory.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/client/factory.py @@ -49,13 +49,14 @@ See also: ``architecture/docs/auth/sdk-cli-oauth.md`` for a full design doc. """ +import asyncio import logging import os import threading from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path -from typing import Any, Callable, Mapping +from typing import Any, Callable, Mapping, Protocol import httpx from nemo_platform import ( @@ -65,12 +66,15 @@ NotGiven, not_given, ) +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR from nemo_platform_ext.auth.helpers import NMPOIDCConfig, build_effective_scope, discover_nmp_config from nemo_platform_ext.auth.token_provider import ( OIDCTokenProvider, TokenSet, ) +from nemo_platform_ext.auth.workload_exchange import WorkloadTokenExchangeProvider +from nemo_platform_ext.client.tls import client_verify_from_env logger = logging.getLogger(__name__) @@ -82,12 +86,17 @@ # Guards _TOKEN_PROVIDER_CACHE; acquired only during dict lookup/insert (fast). _TOKEN_PROVIDER_CACHE_LOCK = threading.Lock() - # --------------------------------------------------------------------------- # Data classes # --------------------------------------------------------------------------- +class _AccessTokenProvider(Protocol): + def get_access_token(self) -> str: ... + + async def get_access_token_async(self) -> str: ... + + @dataclass(frozen=True) class ClientInitConfig: """Everything the SDK client constructor needs after config resolution. @@ -110,7 +119,7 @@ class _ResolvedBootstrap: base_url: str workspace: str | None default_headers: dict[str, str] - token_provider: OIDCTokenProvider | None # None for non-OAuth users + token_provider: _AccessTokenProvider | None # None for non-OAuth users @dataclass(frozen=True) @@ -159,12 +168,86 @@ def _discover_oidc_client_settings(base_url: str) -> NMPOIDCConfig: ) +def _workload_identity_token_file_from_env() -> Path | None: + """Return the configured workload identity token file, if workload bootstrap is active.""" + token_file = os.environ.get(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR) + return Path(token_file) if token_file else None + + +def _create_workload_exchange_provider(base_url: str, subject_token_file: Path) -> WorkloadTokenExchangeProvider: + """Create a workload identity token exchange provider from NeMo auth discovery metadata.""" + oidc_config = _discover_oidc_client_settings(base_url) + if not oidc_config.workload_token_exchange_enabled: + raise RuntimeError( + f"{WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR} is set but workload token exchange is not enabled by auth discovery" + ) + + token_endpoint = oidc_config.workload_token_endpoint or oidc_config.token_endpoint or "" + client_id = oidc_config.workload_client_id or oidc_config.client_id or "" + if not token_endpoint: + raise RuntimeError( + "Workload token exchange is enabled but auth discovery did not return workload_token_endpoint or token_endpoint" + ) + if not client_id: + raise RuntimeError( + "Workload token exchange is enabled but auth discovery did not return workload_client_id or client_id" + ) + + return WorkloadTokenExchangeProvider( + token_endpoint=token_endpoint, + client_id=client_id, + subject_token_file=subject_token_file, + audience=oidc_config.workload_audience, + scope=oidc_config.workload_scope, + refresh_margin_seconds=_TOKEN_REFRESH_MARGIN_SECONDS, + ) + + +class _LazyWorkloadTokenExchangeProvider: + """Create the workload exchange provider on the first token request.""" + + def __init__(self, *, base_url: str, subject_token_file: Path) -> None: + self._base_url = base_url + self._subject_token_file = subject_token_file + self._provider: WorkloadTokenExchangeProvider | None = None + self._lock = threading.Lock() + + def _get_provider(self) -> WorkloadTokenExchangeProvider: + provider = self._provider + if provider is not None: + return provider + with self._lock: + provider = self._provider + if provider is None: + provider = _create_workload_exchange_provider(self._base_url, self._subject_token_file) + self._provider = provider + return provider + + def get_cached_access_token(self) -> str | None: + provider = self._provider + if provider is None: + return None + tokens = provider.tokens + if not tokens.access_token or tokens.is_expired(provider.refresh_margin_seconds): + return None + return tokens.access_token + + def get_access_token(self) -> str: + return self._get_provider().get_access_token() + + async def get_access_token_async(self) -> str: + provider = self._provider + if provider is None: + provider = await asyncio.to_thread(self._get_provider) + return await provider.get_access_token_async() + + # --------------------------------------------------------------------------- # httpx event hooks — the core of transparent token injection # --------------------------------------------------------------------------- -def _make_auth_event_hook(provider: OIDCTokenProvider): +def _make_auth_event_hook(provider: _AccessTokenProvider): """Create a **sync** httpx request event hook that injects the Bearer token. Called before every SDK HTTP request. ``provider.get_access_token()`` @@ -179,7 +262,7 @@ def inject_auth(request: httpx.Request) -> None: return inject_auth -def _make_async_auth_event_hook(provider: OIDCTokenProvider): +def _make_async_auth_event_hook(provider: _AccessTokenProvider): """Create an **async** httpx request event hook for AsyncNeMoPlatform. The actual refresh still runs in a worker thread (via @@ -193,6 +276,17 @@ async def inject_auth(request: httpx.Request) -> None: return inject_auth +def _headers_with_seeded_auth(headers: Mapping[str, str], provider: _AccessTokenProvider) -> dict[str, str]: + seeded_headers = dict(headers) + if isinstance(provider, _LazyWorkloadTokenExchangeProvider): + token = provider.get_cached_access_token() + else: + token = provider.get_access_token() + if token: + seeded_headers["Authorization"] = f"Bearer {token}" + return seeded_headers + + # --------------------------------------------------------------------------- # Callbacks wired into OIDCTokenProvider for config-file integration # --------------------------------------------------------------------------- @@ -411,6 +505,14 @@ def _resolve_bootstrap( base_url = str(resolved.cluster.base_url) headers: dict[str, str] = dict(extra_headers) if extra_headers else {} + workload_identity_token_file = _workload_identity_token_file_from_env() + if workload_identity_token_file is not None and access_token is None and not os.environ.get("NMP_ACCESS_TOKEN"): + provider = _LazyWorkloadTokenExchangeProvider( + base_url=base_url, + subject_token_file=workload_identity_token_file, + ) + return _ResolvedBootstrap(base_url, resolved.workspace, headers, provider) + # --- Non-OAuth path (no auth) --- if not isinstance(resolved.user, OAuthUser): user_config = resolved.user.get_client_config() if resolved.user else {} @@ -508,12 +610,17 @@ def build_client_init_kwargs( default_headers=bootstrap.default_headers or None, ) - # Seed the default headers with the current token so that SDK internals - # that inspect headers (e.g. auth_headers property) see a value. + # Seed the default headers with a current token so that SDK internals + # that inspect headers (e.g. auth_headers property) see a value. Workload + # identity only seeds when the request-time provider already has a token. # The event hook will overwrite it with a fresh token on each request. - headers = {**bootstrap.default_headers, "Authorization": f"Bearer {bootstrap.token_provider.get_access_token()}"} + headers = _headers_with_seeded_auth(bootstrap.default_headers, bootstrap.token_provider) hook = _make_auth_event_hook(bootstrap.token_provider) - http_client = DefaultHttpxClient(event_hooks={"request": [hook], "response": []}, follow_redirects=True) + http_client = DefaultHttpxClient( + event_hooks={"request": [hook], "response": []}, + follow_redirects=True, + verify=client_verify_from_env(), + ) return ClientInitConfig( base_url=bootstrap.base_url, workspace=bootstrap.workspace, @@ -550,9 +657,13 @@ def build_async_client_init_kwargs( default_headers=bootstrap.default_headers or None, ) - headers = {**bootstrap.default_headers, "Authorization": f"Bearer {bootstrap.token_provider.get_access_token()}"} + headers = _headers_with_seeded_auth(bootstrap.default_headers, bootstrap.token_provider) hook = _make_async_auth_event_hook(bootstrap.token_provider) - http_client = DefaultAsyncHttpxClient(event_hooks={"request": [hook], "response": []}, follow_redirects=True) + http_client = DefaultAsyncHttpxClient( + event_hooks={"request": [hook], "response": []}, + follow_redirects=True, + verify=client_verify_from_env(), + ) return ClientInitConfig( base_url=bootstrap.base_url, workspace=bootstrap.workspace, @@ -583,6 +694,9 @@ def create_client( access_token=access_token, extra_headers=extra_headers, ) + http_client = client_init_kwargs.http_client + if http_client is not None and not isinstance(http_client, httpx.Client): + raise TypeError("build_client_init_kwargs returned a non-sync HTTP client") return NeMoPlatform( config_path=config_path, @@ -591,7 +705,7 @@ def create_client( base_url=client_init_kwargs.base_url, workspace=client_init_kwargs.workspace, default_headers=client_init_kwargs.default_headers, - http_client=client_init_kwargs.http_client, + http_client=http_client, max_retries=max_retries, timeout=timeout, ) diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/client/tls.py b/packages/nemo_platform_ext/src/nemo_platform_ext/client/tls.py new file mode 100644 index 0000000000..b2bc998d2e --- /dev/null +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/client/tls.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""TLS configuration shared by NeMo Platform SDK and CLI clients.""" + +from __future__ import annotations + +import os + +NMP_CLIENT_SSL_CERT_FILE_ENVVAR = "NMP_CLIENT_SSL_CERT_FILE" + + +def client_verify_from_env() -> str | bool: + """Return the httpx verify setting for NeMo Platform client requests.""" + cert_file = os.environ.get(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "").strip() + return cert_file or True diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py b/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py index 4365b889c2..7adf87c9a6 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py @@ -29,9 +29,6 @@ logger = logging.getLogger(__name__) -_WORKLOAD_TOKEN_ENVVAR = "NEMO_WORKLOAD_TOKEN" -_WORKLOAD_TOKEN_FILE_ENVVAR = "NEMO_WORKLOAD_TOKEN_FILE" - @dataclass(frozen=True) class _RuntimeAccessTokenSource: @@ -141,18 +138,6 @@ def _migrate_legacy_api_key_users(cls, config_data: dict) -> None: def _runtime_access_token_source_from_env(cls) -> _RuntimeAccessTokenSource | None: if token := os.environ.get("NMP_ACCESS_TOKEN"): return _RuntimeAccessTokenSource(token, "NMP_ACCESS_TOKEN environment override") - if token := os.environ.get(_WORKLOAD_TOKEN_ENVVAR): - return _RuntimeAccessTokenSource(token, f"{_WORKLOAD_TOKEN_ENVVAR} environment override") - if token_path := os.environ.get(_WORKLOAD_TOKEN_FILE_ENVVAR): - try: - token = Path(token_path).read_text(encoding="utf-8").strip() - except OSError as exc: - raise ValueError(f"Unable to read {_WORKLOAD_TOKEN_FILE_ENVVAR} at {token_path}: {exc}") from exc - if token: - return _RuntimeAccessTokenSource( - token, - f"{_WORKLOAD_TOKEN_FILE_ENVVAR} environment override ({token_path})", - ) return None @classmethod diff --git a/packages/nemo_platform_ext/tests/auth/test_device_flow.py b/packages/nemo_platform_ext/tests/auth/test_device_flow.py index 002fcc4cce..4672af4d7d 100644 --- a/packages/nemo_platform_ext/tests/auth/test_device_flow.py +++ b/packages/nemo_platform_ext/tests/auth/test_device_flow.py @@ -26,6 +26,7 @@ authenticate_with_password_grant, refresh_access_token, ) +from nemo_platform_ext.client.tls import NMP_CLIENT_SSL_CERT_FILE_ENVVAR class TestDeviceCodeResponse: @@ -137,6 +138,30 @@ async def test_start_device_authorization_success(self, device_flow): timeout=30.0, ) + @pytest.mark.asyncio + async def test_start_device_authorization_uses_nemo_scoped_ca_bundle(self, device_flow, monkeypatch): + mock_response_data = { + "device_code": "device_code_123", + "user_code": "ABC-123", + "verification_uri": "https://sso.example.com/device", + "expires_in": 1800, + } + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + with patch("httpx.AsyncClient") as mock_client_class: + mock_client = AsyncMock() + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_response.json.return_value = mock_response_data + mock_client.post.return_value = mock_response + mock_client.__aenter__.return_value = mock_client + mock_client.__aexit__.return_value = None + mock_client_class.return_value = mock_client + + await device_flow.start_device_authorization() + + mock_client_class.assert_called_once_with(verify="/tmp/nemo-ca.pem") + @pytest.mark.asyncio async def test_start_device_authorization_default_interval(self, device_flow): """Test device authorization with default interval.""" @@ -403,6 +428,32 @@ def test_authenticate_with_password_grant_success(self): timeout=30.0, ) + def test_authenticate_with_password_grant_uses_nemo_scoped_ca_bundle(self, monkeypatch): + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + with patch("httpx.Client") as mock_client_class: + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "access_token": "access_123", + "token_type": "Bearer", + "expires_in": 3600, + } + mock_client.post.return_value = mock_response + mock_client.__enter__.return_value = mock_client + mock_client.__exit__.return_value = None + mock_client_class.return_value = mock_client + + authenticate_with_password_grant( + token_endpoint="https://idp/token", + client_id="client-id", + username="user", + password="secret", + ) + + mock_client_class.assert_called_once_with(verify="/tmp/nemo-ca.pem") + def test_authenticate_with_password_grant_failure(self): with patch("httpx.Client") as mock_client_class: mock_client = MagicMock() diff --git a/packages/nemo_platform_ext/tests/auth/test_token_provider.py b/packages/nemo_platform_ext/tests/auth/test_token_provider.py index bc3a4d1be2..785d14cdf3 100644 --- a/packages/nemo_platform_ext/tests/auth/test_token_provider.py +++ b/packages/nemo_platform_ext/tests/auth/test_token_provider.py @@ -13,7 +13,9 @@ from nemo_platform_ext.auth.token_provider import ( OIDCTokenProvider, TokenSet, + refresh_token_grant, ) +from nemo_platform_ext.client.tls import NMP_CLIENT_SSL_CERT_FILE_ENVVAR def _make_jwt(claims: dict, header: dict | None = None) -> str: @@ -41,6 +43,21 @@ def test_from_access_token_no_exp_claim(self): assert ts.expires_at is None + @pytest.mark.parametrize("expires_in", [120, 120.5]) + def test_from_access_token_uses_numeric_expires_in_when_no_exp_claim(self, expires_in): + token = _make_jwt({"sub": "user1"}) + before = time.time() + ts = TokenSet.from_access_token(token, expires_in=expires_in) + + assert ts.expires_at is not None + assert before + expires_in <= ts.expires_at <= time.time() + expires_in + + def test_from_access_token_rejects_bool_expires_in(self): + token = _make_jwt({"sub": "user1"}) + ts = TokenSet.from_access_token(token, expires_in=True) + + assert ts.expires_at is None + def test_from_access_token_non_jwt(self): ts = TokenSet.from_access_token("not-a-jwt", refresh_token="r") @@ -90,6 +107,23 @@ async def test_get_access_token_async_returns_current_when_not_expired(self): assert await provider.get_access_token_async() == token + @patch("nemo_platform_ext.auth.token_provider.httpx.post") + def test_refresh_token_grant_uses_nemo_scoped_ca_bundle(self, mock_post, monkeypatch): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"access_token": "new_access"} + mock_post.return_value = mock_response + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + result = refresh_token_grant( + token_endpoint="https://idp/token", + client_id="client", + refresh_token="refresh_abc", + ) + + assert result == {"access_token": "new_access"} + assert mock_post.call_args.kwargs["verify"] == "/tmp/nemo-ca.pem" + @patch("nemo_platform_ext.auth.token_provider.httpx.post") def test_get_access_token_refreshes_when_expired(self, mock_post): old_token = _make_jwt({"exp": int(time.time()) - 100}) @@ -121,6 +155,33 @@ def test_get_access_token_refreshes_when_expired(self, mock_post): assert call_kwargs[1]["data"]["client_id"] == "client" assert call_kwargs[1]["data"]["refresh_token"] == "old_refresh" + @patch("nemo_platform_ext.auth.token_provider.httpx.post") + def test_get_access_token_refreshes_opaque_token_with_expires_in(self, mock_post): + old_token = _make_jwt({"exp": int(time.time()) - 100}) + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "access_token": "opaque_access", + "expires_in": 120, + } + mock_post.return_value = mock_response + + tokens = TokenSet.from_access_token(old_token, refresh_token="old_refresh") + provider = OIDCTokenProvider( + token_endpoint="https://idp/token", + client_id="client", + tokens=tokens, + refresh_margin_seconds=0, + ) + + before = time.time() + result = provider.get_access_token() + + assert result == "opaque_access" + assert provider.tokens.expires_at is not None + assert before + 120 <= provider.tokens.expires_at <= time.time() + 120 + @patch("nemo_platform_ext.auth.token_provider.httpx.post") def test_refresh_reloads_tokens_before_request(self, mock_post): stale_token = _make_jwt({"exp": int(time.time()) - 200}) diff --git a/packages/nemo_platform_ext/tests/auth/test_utils.py b/packages/nemo_platform_ext/tests/auth/test_utils.py index 07979a85e8..6c66577652 100644 --- a/packages/nemo_platform_ext/tests/auth/test_utils.py +++ b/packages/nemo_platform_ext/tests/auth/test_utils.py @@ -3,6 +3,7 @@ import base64 import json +from unittest.mock import MagicMock, patch import httpx import pytest @@ -18,6 +19,7 @@ normalize_scope_prefix, validate_requested_scopes_granted, ) +from nemo_platform_ext.client.tls import NMP_CLIENT_SSL_CERT_FILE_ENVVAR from pytest_httpserver import HTTPServer @@ -160,6 +162,22 @@ def test_strips_trailing_slash(self, httpserver: HTTPServer): result = discover_nmp_config(httpserver.url_for("") + "/") assert result.auth_enabled is False + @patch("nemo_platform_ext.auth.helpers.httpx.get") + def test_uses_nemo_scoped_ca_bundle(self, mock_get, monkeypatch): + response = MagicMock() + response.json.return_value = {"auth_enabled": False} + mock_get.return_value = response + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + result = discover_nmp_config("https://nemo.example.com") + + assert result.auth_enabled is False + mock_get.assert_called_once_with( + "https://nemo.example.com/apis/auth/discovery", + timeout=10.0, + verify="/tmp/nemo-ca.pem", + ) + class TestBuildEffectiveScope: def test_no_prefix_returns_unchanged(self): diff --git a/packages/nemo_platform_ext/tests/auth/test_workload_exchange.py b/packages/nemo_platform_ext/tests/auth/test_workload_exchange.py new file mode 100644 index 0000000000..f64f226776 --- /dev/null +++ b/packages/nemo_platform_ext/tests/auth/test_workload_exchange.py @@ -0,0 +1,258 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for RFC 8693 workload identity token exchange.""" + +import json +import time +from base64 import urlsafe_b64encode +from unittest.mock import MagicMock, patch + +import pytest +from nemo_platform_ext.auth.workload_exchange import ( + ACCESS_TOKEN_TYPE, + JWT_TOKEN_TYPE, + TOKEN_EXCHANGE_GRANT_TYPE, + WorkloadTokenExchangeError, + WorkloadTokenExchangeProvider, + read_subject_token_file, + token_exchange_grant, +) +from nemo_platform_ext.client.tls import NMP_CLIENT_SSL_CERT_FILE_ENVVAR +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR + + +def _make_jwt(claims: dict) -> str: + header = {"alg": "RS256", "typ": "JWT"} + h = urlsafe_b64encode(json.dumps(header).encode()).rstrip(b"=").decode() + p = urlsafe_b64encode(json.dumps(claims).encode()).rstrip(b"=").decode() + s = urlsafe_b64encode(b"fake-signature").rstrip(b"=").decode() + return f"{h}.{p}.{s}" + + +def test_read_subject_token_file_strips_whitespace(tmp_path): + token_file = tmp_path / "token" + token_file.write_text("subject-token\n", encoding="utf-8") + + assert read_subject_token_file(token_file) == "subject-token" + + +def test_read_subject_token_file_rejects_empty_file(tmp_path): + token_file = tmp_path / "token" + token_file.write_text("\n", encoding="utf-8") + + with pytest.raises(ValueError, match=WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR): + read_subject_token_file(token_file) + + +@patch("nemo_platform_ext.auth.workload_exchange.httpx.post") +def test_token_exchange_grant_sends_rfc8693_request(mock_post): + access_token = _make_jwt({"exp": int(time.time()) + 3600}) + response = MagicMock() + response.status_code = 200 + response.json.return_value = {"access_token": access_token, "expires_in": 300} + mock_post.return_value = response + + result = token_exchange_grant( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token="subject-token", + audience="nemo-platform", + scope="openid email groups", + ) + + assert result["access_token"] == access_token + mock_post.assert_called_once() + assert mock_post.call_args.kwargs["data"] == { + "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, + "client_id": "nemo-platform-workload", + "subject_token": "subject-token", + "subject_token_type": JWT_TOKEN_TYPE, + "requested_token_type": ACCESS_TOKEN_TYPE, + "audience": "nemo-platform", + "scope": "openid email groups", + } + + +@patch("nemo_platform_ext.auth.workload_exchange.httpx.post") +def test_token_exchange_grant_rejects_http_non_loopback_endpoint_before_sending_subject_token(mock_post): + with pytest.raises(ValueError, match="must use HTTPS"): + token_exchange_grant( + token_endpoint="http://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token="subject-token", + ) + + mock_post.assert_not_called() + + +@patch("nemo_platform_ext.auth.workload_exchange.httpx.post") +@pytest.mark.parametrize( + "token_endpoint", + [ + "http://localhost:18080/token", + "http://127.0.0.1:18080/token", + "http://[::1]:18080/token", + ], +) +def test_token_exchange_grant_allows_http_loopback_endpoints(mock_post, token_endpoint): + access_token = _make_jwt({"exp": int(time.time()) + 3600}) + response = MagicMock() + response.status_code = 200 + response.json.return_value = {"access_token": access_token} + mock_post.return_value = response + + result = token_exchange_grant( + token_endpoint=token_endpoint, + client_id="nemo-platform-workload", + subject_token="subject-token", + ) + + assert result["access_token"] == access_token + mock_post.assert_called_once() + + +@patch("nemo_platform_ext.auth.workload_exchange.httpx.post") +def test_token_exchange_grant_uses_nemo_scoped_ca_bundle(mock_post, monkeypatch): + access_token = _make_jwt({"exp": int(time.time()) + 3600}) + response = MagicMock() + response.status_code = 200 + response.json.return_value = {"access_token": access_token} + mock_post.return_value = response + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + result = token_exchange_grant( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token="subject-token", + ) + + assert result["access_token"] == access_token + assert mock_post.call_args.kwargs["verify"] == "/tmp/nemo-ca.pem" + + +@patch("nemo_platform_ext.auth.workload_exchange.httpx.post") +def test_token_exchange_grant_surfaces_idp_error(mock_post): + response = MagicMock() + response.status_code = 400 + response.text = "invalid subject token" + response.headers = {"content-type": "application/json"} + response.json.return_value = { + "error": "invalid_request", + "error_description": "invalid subject token", + } + mock_post.return_value = response + + with pytest.raises(WorkloadTokenExchangeError, match="invalid_request - invalid subject token"): + token_exchange_grant( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token="bad-subject-token", + ) + + +@patch("nemo_platform_ext.auth.workload_exchange.httpx.post") +def test_token_exchange_grant_rejects_non_object_error_payload(mock_post): + response = MagicMock() + response.status_code = 400 + response.text = "[]" + response.headers = {"content-type": "application/json"} + response.json.return_value = [] + mock_post.return_value = response + + with pytest.raises(WorkloadTokenExchangeError, match="invalid_response - Token endpoint error response"): + token_exchange_grant( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token="bad-subject-token", + ) + + +@patch("nemo_platform_ext.auth.workload_exchange.httpx.post") +@pytest.mark.parametrize("payload", [[], {}, {"access_token": ""}, {"access_token": None}]) +def test_token_exchange_grant_rejects_success_response_without_non_empty_access_token(mock_post, payload): + response = MagicMock() + response.status_code = 200 + response.json.return_value = payload + mock_post.return_value = response + + with pytest.raises(WorkloadTokenExchangeError): + token_exchange_grant( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token="subject-token", + ) + + +@patch("nemo_platform_ext.auth.workload_exchange.token_exchange_grant") +def test_provider_rejects_exchange_response_without_access_token(mock_exchange, tmp_path): + subject_token_file = tmp_path / "token" + subject_token_file.write_text("subject-token", encoding="utf-8") + mock_exchange.return_value = {} + + provider = WorkloadTokenExchangeProvider( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token_file=subject_token_file, + ) + + with pytest.raises(WorkloadTokenExchangeError, match="non-empty access_token"): + provider.get_access_token() + + +@patch("nemo_platform_ext.auth.workload_exchange.token_exchange_grant") +@pytest.mark.parametrize( + "expires_in", + [None, "300", True, float("nan"), float("inf"), 10**400], +) +def test_provider_rejects_exchange_response_without_usable_lifetime(mock_exchange, tmp_path, expires_in): + subject_token_file = tmp_path / "token" + subject_token_file.write_text("subject-token", encoding="utf-8") + token_data = {"access_token": "opaque-access-token"} + if expires_in is not None: + token_data["expires_in"] = expires_in + mock_exchange.return_value = token_data + + provider = WorkloadTokenExchangeProvider( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token_file=subject_token_file, + ) + + with pytest.raises(WorkloadTokenExchangeError, match="usable access_token lifetime"): + provider.get_access_token() + + assert provider.tokens.access_token == "" + + +@patch("nemo_platform_ext.auth.workload_exchange.token_exchange_grant") +def test_provider_rejects_expired_exchange_response_and_retries_with_current_subject_token(mock_exchange, tmp_path): + subject_token_file = tmp_path / "token" + subject_token_file.write_text("subject-token-one", encoding="utf-8") + expired_access_token = _make_jwt({"exp": int(time.time()) - 10}) + fresh_access_token = _make_jwt({"exp": int(time.time()) + 3600}) + mock_exchange.side_effect = [ + {"access_token": expired_access_token}, + {"access_token": fresh_access_token}, + ] + + provider = WorkloadTokenExchangeProvider( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token_file=subject_token_file, + audience="nemo-platform", + scope="openid email groups", + refresh_margin_seconds=0, + ) + + with pytest.raises(WorkloadTokenExchangeError, match="expired access_token"): + provider.get_access_token() + + subject_token_file.write_text("subject-token-two", encoding="utf-8") + + assert provider.get_access_token() == fresh_access_token + assert mock_exchange.call_count == 2 + assert mock_exchange.call_args_list[0].kwargs["subject_token"] == "subject-token-one" + assert mock_exchange.call_args_list[1].kwargs["subject_token"] == "subject-token-two" + assert mock_exchange.call_args_list[1].kwargs["audience"] == "nemo-platform" + assert mock_exchange.call_args_list[1].kwargs["scope"] == "openid email groups" diff --git a/packages/nemo_platform_ext/tests/cli/commands/test_auth.py b/packages/nemo_platform_ext/tests/cli/commands/test_auth.py index f4ef89cb56..d126590641 100644 --- a/packages/nemo_platform_ext/tests/cli/commands/test_auth.py +++ b/packages/nemo_platform_ext/tests/cli/commands/test_auth.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import logging from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace @@ -10,6 +11,7 @@ import yaml from nemo_platform_ext.auth.helpers import decode_jwt_claims, generate_unsigned_jwt from nemo_platform_ext.cli.app import app +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR from typer.testing import CliRunner from ..utils import assert_exit_code @@ -67,7 +69,12 @@ def _decode_jwt_noop(token: str) -> dict: @pytest.fixture def oauth_config_file(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: - for env_key in ("NMP_ACCESS_TOKEN", "NEMO_WORKLOAD_TOKEN", "NEMO_WORKLOAD_TOKEN_FILE"): + for env_key in ( + "NMP_ACCESS_TOKEN", + "NEMO_WORKLOAD_TOKEN", + "NEMO_WORKLOAD_TOKEN_FILE", + WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, + ): monkeypatch.delenv(env_key, raising=False) config_data = { @@ -158,7 +165,7 @@ def test_auth_logout_warns_when_runtime_token_override_remains( ) -> None: monkeypatch.setattr("nemo_platform_ext.cli.commands.auth.discover_nmp_config", _discover_auth_enabled) monkeypatch.setenv( - "NEMO_WORKLOAD_TOKEN", + "NMP_ACCESS_TOKEN", generate_unsigned_jwt( principal_id="svc-nemo-ci", email="svc-nemo-ci@example.com", @@ -170,7 +177,7 @@ def test_auth_logout_warns_when_runtime_token_override_remains( assert_exit_code(result, 0) assert "Logged out successfully" in result.output - assert "NEMO_WORKLOAD_TOKEN environment override is still active" in result.output + assert "NMP_ACCESS_TOKEN environment override is still active" in result.output def test_auth_logout_fails_if_credentials_remain(oauth_config_file: Path, monkeypatch: pytest.MonkeyPatch) -> None: @@ -308,7 +315,7 @@ def test_auth_status_shows_warning_for_unsigned_token(oauth_config_file: Path, m assert "local/testing" in result.output -def test_runtime_token_source_label_handles_unreadable_token_file( +def test_runtime_token_source_label_ignores_workload_identity_token_file( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: from nemo_platform_ext.cli.commands.auth import _runtime_token_source_label @@ -317,8 +324,26 @@ def test_runtime_token_source_label_handles_unreadable_token_file( monkeypatch.delenv("NMP_ACCESS_TOKEN", raising=False) monkeypatch.delenv("NEMO_WORKLOAD_TOKEN", raising=False) monkeypatch.setenv("NEMO_WORKLOAD_TOKEN_FILE", str(token_file)) + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(token_file)) - assert _runtime_token_source_label() == "NEMO_WORKLOAD_TOKEN_FILE environment override could not be read" + assert _runtime_token_source_label() is None + + +def test_runtime_token_source_label_returns_none_when_config_label_fails( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + from nemo_platform_ext.cli.commands.auth import _runtime_token_source_label + + monkeypatch.setattr( + "nemo_platform_ext.config.config.Config.runtime_access_token_source_label", + lambda: (_ for _ in ()).throw(ValueError("invalid runtime token")), + ) + + with caplog.at_level(logging.DEBUG, logger="nemo_platform_ext.cli.commands.auth"): + assert _runtime_token_source_label() is None + + assert "Failed to resolve runtime token override source label" in caplog.text + assert "invalid runtime token" in caplog.text def test_auth_status_shows_config_file_credential_source( @@ -651,8 +676,8 @@ def test_auth_login_unsigned_token_uses_principal_id_when_provided( @dataclass class IsAuthDisabledCase: id: str - auth_enabled: bool | None # None means the cluster is unreachable (raises) - expected: bool | None + auth_enabled: bool + expected: bool @pytest.mark.parametrize( @@ -660,16 +685,11 @@ class IsAuthDisabledCase: [ IsAuthDisabledCase(id="disabled", auth_enabled=False, expected=True), IsAuthDisabledCase(id="enabled", auth_enabled=True, expected=False), - IsAuthDisabledCase(id="unreachable", auth_enabled=None, expected=None), ], ids=lambda c: c.id, ) def test_is_auth_disabled(monkeypatch: pytest.MonkeyPatch, case: IsAuthDisabledCase) -> None: - import httpx - def mock_discover(url: str, timeout: float = 10.0) -> SimpleNamespace: - if case.auth_enabled is None: - raise httpx.ConnectError("Connection refused") return SimpleNamespace(auth_enabled=case.auth_enabled) monkeypatch.setattr("nemo_platform_ext.cli.commands.auth.discover_nmp_config", mock_discover) @@ -678,6 +698,20 @@ def mock_discover(url: str, timeout: float = 10.0) -> SimpleNamespace: assert is_auth_disabled("http://localhost:8080") is case.expected +def test_is_auth_disabled_raises_when_discovery_fails(monkeypatch: pytest.MonkeyPatch) -> None: + import httpx + from nemo_platform_ext.auth.helpers import AuthError + from nemo_platform_ext.cli.commands.auth import is_auth_disabled + + def raise_connect_error(url: str, timeout: float = 10.0) -> SimpleNamespace: + raise httpx.ConnectError("Connection refused") + + monkeypatch.setattr("nemo_platform_ext.cli.commands.auth.discover_nmp_config", raise_connect_error) + + with pytest.raises(AuthError, match="Failed to discover auth configuration: Connection refused"): + is_auth_disabled("http://localhost:8080") + + # --------------------------------------------------------------------------- # auth status when auth disabled # --------------------------------------------------------------------------- @@ -702,6 +736,31 @@ def test_auth_status_when_auth_disabled_does_not_show_token_details( assert "Refresh Token" not in result.output +def test_auth_status_when_cluster_unreachable_shows_local_state( + oauth_config_file: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + import httpx + + def raise_connect_error(url: str, timeout: float = 10.0) -> SimpleNamespace: + raise httpx.ConnectError("Connection refused") + + monkeypatch.setattr("nemo_platform_ext.cli.commands.auth.discover_nmp_config", raise_connect_error) + result = runner.invoke(app, ["--context", "foo", "auth", "status"]) + output = " ".join(result.output.split()) + + assert_exit_code(result, 0) + assert "Failed to discover auth configuration:" in output + assert "Connection refused" in output + assert "Auth Discovery" in result.output + assert "unavailable" in result.output + assert "Auth Type" in result.output + assert "oauth" in result.output + assert "Credential Source" in result.output + assert "config file" in result.output + assert "Refresh Token" in result.output + assert "foo-token" not in result.output + + # --------------------------------------------------------------------------- # auth logout when auth disabled # --------------------------------------------------------------------------- @@ -719,7 +778,7 @@ def test_auth_logout_when_auth_disabled_shows_message_and_skips_credential_clear mock_write.assert_not_called() -def test_auth_logout_when_cluster_unreachable_still_clears_credentials( +def test_auth_logout_when_cluster_unreachable_clears_local_credentials( oauth_config_file: Path, monkeypatch: pytest.MonkeyPatch ) -> None: import httpx @@ -728,9 +787,23 @@ def raise_connect_error(url: str, timeout: float = 10.0) -> SimpleNamespace: raise httpx.ConnectError("Connection refused") monkeypatch.setattr("nemo_platform_ext.cli.commands.auth.discover_nmp_config", raise_connect_error) - with patch("nemo_platform_ext.config.config.Config.write") as mock_write: - result = runner.invoke(app, ["--context", "foo", "auth", "logout"]) + result = runner.invoke(app, ["--context", "foo", "auth", "logout"]) + output = " ".join(result.output.split()) assert_exit_code(result, 0) - mock_write.assert_called_once() - assert mock_write.call_args.kwargs["context_name"] == "foo" + assert "Failed to discover auth configuration: Connection refused" in output + assert "continuing to clear local credentials" in output + assert "Logged out successfully" in result.output + + with open(oauth_config_file) as f: + data = yaml.safe_load(f) + + default_user = next(user for user in data["users"] if user["name"] == "default") + foo_user = next(user for user in data["users"] if user["name"] == "foo") + + assert default_user["type"] == "oauth" + assert default_user["token"] == "default-token" + assert default_user["refresh_token"] == "default-refresh" + assert foo_user["type"] == "no-auth" + assert "token" not in foo_user + assert "refresh_token" not in foo_user diff --git a/packages/nemo_platform_ext/tests/cli/core/test_context.py b/packages/nemo_platform_ext/tests/cli/core/test_context.py index f977b086bc..27066e4c5f 100644 --- a/packages/nemo_platform_ext/tests/cli/core/test_context.py +++ b/packages/nemo_platform_ext/tests/cli/core/test_context.py @@ -1,7 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from types import SimpleNamespace +from unittest.mock import patch + from nemo_platform_ext.cli.core.context import CLIContext +from nemo_platform_ext.config.models import NoAuthUser, OAuthUser def test_context_instances_are_independent(): @@ -60,16 +64,112 @@ def test_get_no_truncate_default(): assert result is False -def test_get_client_passes_user_config_and_is_cached(): - """Test that get_client passes user's client config to the SDK client and caches it.""" - ctx = CLIContext(overrides={"base_url": "http://test.example.com", "access_token": "token-123"}) +def test_get_client_uses_config_bootstrap_for_persisted_oauth_context_and_is_cached(): + """Test that get_client lets the SDK bootstrap config-backed OAuth auth and caches it.""" + ctx = CLIContext(overrides={"base_url": "http://test.example.com"}) + resolved_context = SimpleNamespace( + cluster=SimpleNamespace(base_url="http://test.example.com"), + context_name="dev", + workspace="test-workspace", + user=OAuthUser(name="dev-user", token="token-123", refresh_token="refresh-123"), + ) + config_file = SimpleNamespace(contexts=[SimpleNamespace(name="dev")]) + + with ( + patch("nemo_platform_ext.config.config.get_context", return_value=resolved_context), + patch("nemo_platform_ext.config.config.Config.load") as mock_config_load, + patch("nemo_platform_ext.config.config.Config.runtime_access_token_source_label", return_value=None), + patch("nemo_platform.NeMoPlatform", autospec=True) as mock_client_cls, + ): + mock_config_load.return_value.get_config_file.return_value = config_file + client = ctx.get_client() + client2 = ctx.get_client() + + mock_client_cls.assert_called_once_with( + base_url="http://test.example.com", + context_name="dev", + timeout=60.0, + workspace="test-workspace", + ) + assert client is mock_client_cls.return_value - client = ctx.get_client() + # Verify the client is cached + assert client is client2 - # Verify the client has the expected headers from get_client_config() - assert "Authorization" in client.default_headers - assert client.default_headers["Authorization"] == "Bearer token-123" - # Verify the client is cached - client2 = ctx.get_client() +def test_get_client_preserves_direct_mode_for_synthetic_no_auth_context(): + """Test that synthesized default/no-auth contexts do not force SDK config bootstrap.""" + ctx = CLIContext(overrides={"base_url": "http://test.example.com"}) + resolved_context = SimpleNamespace( + cluster=SimpleNamespace(base_url="http://test.example.com"), + context_name="default", + workspace="default", + user=NoAuthUser(name="default-user"), + ) + + with ( + patch("nemo_platform_ext.config.config.get_context", return_value=resolved_context), + patch("nemo_platform.NeMoPlatform", autospec=True) as mock_client_cls, + ): + ctx.get_client() + + mock_client_cls.assert_called_once_with( + base_url="http://test.example.com", + timeout=60.0, + workspace="default", + ) + + +def test_get_client_passes_explicit_access_token_override(): + """Test that explicit access token overrides remain caller-managed static headers.""" + ctx = CLIContext(overrides={"base_url": "http://test.example.com", "access_token": "token-123"}) + resolved_context = SimpleNamespace( + cluster=SimpleNamespace(base_url="http://test.example.com"), + context_name="dev", + workspace="test-workspace", + user=OAuthUser(name="dev-user", token="token-123"), + ) + + with ( + patch("nemo_platform_ext.config.config.get_context", return_value=resolved_context), + patch("nemo_platform.NeMoPlatform", autospec=True) as mock_client_cls, + ): + ctx.get_client() + + mock_client_cls.assert_called_once_with( + base_url="http://test.example.com", + default_headers={"Authorization": "Bearer token-123"}, + timeout=60.0, + workspace="test-workspace", + ) + + +def test_get_async_client_uses_config_bootstrap_for_persisted_oauth_context_and_is_cached(): + """Test that get_async_client lets the SDK bootstrap config-backed OAuth auth and caches it.""" + ctx = CLIContext(overrides={"base_url": "http://test.example.com"}) + resolved_context = SimpleNamespace( + cluster=SimpleNamespace(base_url="http://test.example.com"), + context_name="dev", + workspace="test-workspace", + user=OAuthUser(name="dev-user", token="token-123", refresh_token="refresh-123"), + ) + config_file = SimpleNamespace(contexts=[SimpleNamespace(name="dev")]) + + with ( + patch("nemo_platform_ext.config.config.get_context", return_value=resolved_context), + patch("nemo_platform_ext.config.config.Config.load") as mock_config_load, + patch("nemo_platform_ext.config.config.Config.runtime_access_token_source_label", return_value=None), + patch("nemo_platform.AsyncNeMoPlatform", autospec=True) as mock_client_cls, + ): + mock_config_load.return_value.get_config_file.return_value = config_file + client = ctx.get_async_client() + client2 = ctx.get_async_client() + + mock_client_cls.assert_called_once_with( + base_url="http://test.example.com", + context_name="dev", + timeout=60.0, + workspace="test-workspace", + ) + assert client is mock_client_cls.return_value assert client is client2 diff --git a/packages/nemo_platform_ext/tests/client/test_client.py b/packages/nemo_platform_ext/tests/client/test_client.py index 4e1d5f478c..daf1775337 100644 --- a/packages/nemo_platform_ext/tests/client/test_client.py +++ b/packages/nemo_platform_ext/tests/client/test_client.py @@ -17,6 +17,8 @@ from nemo_platform import AsyncNeMoPlatform, DefaultHttpxClient, NeMoPlatform, not_given from nemo_platform_ext.auth.helpers import NMPOIDCConfig, decode_jwt_claims from nemo_platform_ext.client.factory import create_client +from nemo_platform_ext.client.tls import NMP_CLIENT_SSL_CERT_FILE_ENVVAR +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR def _make_jwt(claims: dict) -> str: @@ -67,6 +69,17 @@ def _write_config(tmp_path, *, user_type="oauth", token=None, refresh_token=None token_endpoint="https://idp/token", ) +_MOCK_WORKLOAD_NMP_CONFIG = NMPOIDCConfig( + auth_enabled=True, + client_id="nmp-client-id", + token_endpoint="https://idp/token", + workload_token_exchange_enabled=True, + workload_client_id="nmp-workload-client-id", + workload_token_endpoint="https://workload-idp/token", + workload_audience="nemo-platform", + workload_scope="openid email groups", +) + class TestCreateClientOAuth: @patch("nemo_platform_ext.client.factory.discover_nmp_config", return_value=_MOCK_NMP_CONFIG) @@ -116,6 +129,27 @@ def test_oauth_uses_sdk_default_httpx_client(self, _mock_discover, tmp_path): finally: client.close() + @patch("nemo_platform_ext.client.factory.DefaultHttpxClient") + @patch("nemo_platform_ext.client.factory.discover_nmp_config", return_value=_MOCK_NMP_CONFIG) + def test_oauth_uses_nemo_scoped_ca_bundle(self, _mock_discover, mock_default_httpx_client, tmp_path, monkeypatch): + token = _make_jwt({"exp": int(time.time()) + 3600, "sub": "user1"}) + config_path = _write_config( + tmp_path, + token=token, + refresh_token="refresh_abc", + ) + http_client = httpx.Client() + mock_default_httpx_client.return_value = http_client + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + client = create_client(config_path=config_path) + try: + assert client is not None + finally: + client.close() + + assert mock_default_httpx_client.call_args.kwargs["verify"] == "/tmp/nemo-ca.pem" + @patch("nemo_platform_ext.client.factory.discover_nmp_config", return_value=_MOCK_NMP_CONFIG) @patch("nemo_platform_ext.auth.token_provider.httpx.post") def test_persist_refreshed_tokens_writes_to_config(self, mock_post, _mock_discover, tmp_path): @@ -158,6 +192,90 @@ def test_env_access_token_overrides_user_auth(self, tmp_path, monkeypatch): assert request.headers["Authorization"] == "Bearer env-access-token-123" +class TestCreateClientWorkloadIdentity: + @patch("nemo_platform_ext.client.factory.discover_nmp_config", return_value=_MOCK_WORKLOAD_NMP_CONFIG) + @patch("nemo_platform_ext.auth.workload_exchange.token_exchange_grant") + def test_exchanges_workload_identity_token_file(self, mock_exchange, _mock_discover, tmp_path, monkeypatch): + subject_token_file = tmp_path / "workload-token" + subject_token_file.write_text("subject-token-one\n", encoding="utf-8") + access_token = _make_jwt({"exp": int(time.time()) + 3600, "sub": "workload-user"}) + mock_exchange.return_value = {"access_token": access_token, "expires_in": 300} + monkeypatch.setenv("NMP_BASE_URL", "https://api.example.com") + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(subject_token_file)) + + client = create_client() + + try: + assert str(client.base_url).rstrip("/") == "https://api.example.com" + assert "Authorization" not in client._custom_headers + _mock_discover.assert_not_called() + mock_exchange.assert_not_called() + + request = client._client.build_request("GET", "https://api.example.com/test") + client._client._event_hooks["request"][0](request) + assert request.headers["Authorization"] == f"Bearer {access_token}" + finally: + client.close() + + mock_exchange.assert_called_once() + assert mock_exchange.call_args.kwargs["token_endpoint"] == "https://workload-idp/token" + assert mock_exchange.call_args.kwargs["client_id"] == "nmp-workload-client-id" + assert mock_exchange.call_args.kwargs["subject_token"] == "subject-token-one" + assert mock_exchange.call_args.kwargs["audience"] == "nemo-platform" + assert mock_exchange.call_args.kwargs["scope"] == "openid email groups" + + @pytest.mark.asyncio + @patch("nemo_platform.client.factory.discover_nmp_config", return_value=_MOCK_WORKLOAD_NMP_CONFIG) + @patch("nemo_platform.auth.workload_exchange.token_exchange_grant") + async def test_async_exchanges_workload_identity_token_file_at_request_time( + self, mock_exchange, _mock_discover, tmp_path, monkeypatch + ): + subject_token_file = tmp_path / "workload-token" + subject_token_file.write_text("subject-token-one\n", encoding="utf-8") + access_token = _make_jwt({"exp": int(time.time()) + 3600, "sub": "workload-user"}) + mock_exchange.return_value = {"access_token": access_token, "expires_in": 300} + monkeypatch.setenv("NMP_BASE_URL", "https://api.example.com") + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(subject_token_file)) + + client = AsyncNeMoPlatform() + + try: + assert str(client.base_url).rstrip("/") == "https://api.example.com" + assert "Authorization" not in client._custom_headers + _mock_discover.assert_not_called() + mock_exchange.assert_not_called() + + request = client._client.build_request("GET", "https://api.example.com/test") + await client._client._event_hooks["request"][0](request) + assert request.headers["Authorization"] == f"Bearer {access_token}" + finally: + await client.close() + + mock_exchange.assert_called_once() + assert mock_exchange.call_args.kwargs["token_endpoint"] == "https://workload-idp/token" + assert mock_exchange.call_args.kwargs["client_id"] == "nmp-workload-client-id" + assert mock_exchange.call_args.kwargs["subject_token"] == "subject-token-one" + assert mock_exchange.call_args.kwargs["audience"] == "nemo-platform" + assert mock_exchange.call_args.kwargs["scope"] == "openid email groups" + + @patch("nemo_platform_ext.client.factory.discover_nmp_config", return_value=_MOCK_WORKLOAD_NMP_CONFIG) + def test_env_access_token_takes_precedence_over_workload_identity_file(self, _mock_discover, tmp_path, monkeypatch): + subject_token_file = tmp_path / "workload-token" + subject_token_file.write_text("subject-token-one\n", encoding="utf-8") + monkeypatch.setenv("NMP_BASE_URL", "https://api.example.com") + monkeypatch.setenv("NMP_ACCESS_TOKEN", "env-access-token-123") + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(subject_token_file)) + + client = create_client() + + try: + request = client._client.build_request("GET", "https://api.example.com/test") + client._client._event_hooks["request"][0](request) + assert request.headers["Authorization"] == "Bearer env-access-token-123" + finally: + client.close() + + class TestCreateClientApiKey: def test_creates_client_with_api_key(self, tmp_path): config_path = _write_config(tmp_path, user_type="api-key", api_key="nvapi-test-key-123") @@ -357,6 +475,42 @@ def test_sync_constructor_with_base_url_skips_config_bootstrap(self, mock_build_ mock_build_client_kwargs.assert_not_called() + @patch("nemo_platform._client.DefaultHttpxClient") + @patch("nemo_platform.client.factory.build_client_init_kwargs") + def test_sync_constructor_direct_mode_uses_nemo_scoped_ca_bundle( + self, mock_build_client_kwargs, mock_default_httpx_client, monkeypatch + ): + mock_build_client_kwargs.side_effect = AssertionError("bootstrap should not be called") + mock_default_httpx_client.return_value = httpx.Client() + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + client = NeMoPlatform(base_url="https://override-host:8081", workspace="test-workspace") + try: + assert str(client.base_url).rstrip("/") == "https://override-host:8081" + finally: + client.close() + + mock_build_client_kwargs.assert_not_called() + mock_default_httpx_client.assert_called_once_with(verify="/tmp/nemo-ca.pem") + + @patch("nemo_platform.client.factory.build_client_init_kwargs") + def test_sync_constructor_with_workload_file_and_base_url_bootstraps(self, mock_build_client_kwargs, monkeypatch): + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, "/var/run/secrets/nemo-platform/workload/token") + mock_build_client_kwargs.return_value = MagicMock( + base_url="http://override-host:8081", + workspace="test-workspace", + default_headers=None, + http_client=None, + ) + + client = NeMoPlatform(base_url="http://override-host:8081", workspace="test-workspace") + try: + assert str(client.base_url).rstrip("/") == "http://override-host:8081" + finally: + client.close() + + assert mock_build_client_kwargs.call_args.kwargs["base_url"] == "http://override-host:8081" + @patch("nemo_platform.client.factory.build_client_init_kwargs") def test_sync_constructor_passes_context_name_to_bootstrap(self, mock_build_client_kwargs): mock_build_client_kwargs.return_value = MagicMock( @@ -411,6 +565,46 @@ async def test_async_constructor_with_base_url_skips_config_bootstrap(self, mock mock_build_client_kwargs.assert_not_called() + @pytest.mark.asyncio + @patch("nemo_platform._client.DefaultAsyncHttpxClient") + @patch("nemo_platform.client.factory.build_async_client_init_kwargs") + async def test_async_constructor_direct_mode_uses_nemo_scoped_ca_bundle( + self, mock_build_client_kwargs, mock_default_httpx_client, monkeypatch + ): + mock_build_client_kwargs.side_effect = AssertionError("bootstrap should not be called") + mock_default_httpx_client.return_value = httpx.AsyncClient() + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + client = AsyncNeMoPlatform(base_url="https://override-host:8081", workspace="test-workspace") + try: + assert str(client.base_url).rstrip("/") == "https://override-host:8081" + finally: + await client.close() + + mock_build_client_kwargs.assert_not_called() + mock_default_httpx_client.assert_called_once_with(verify="/tmp/nemo-ca.pem") + + @pytest.mark.asyncio + @patch("nemo_platform.client.factory.build_async_client_init_kwargs") + async def test_async_constructor_with_workload_file_and_base_url_bootstraps( + self, mock_build_client_kwargs, monkeypatch + ): + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, "/var/run/secrets/nemo-platform/workload/token") + mock_build_client_kwargs.return_value = MagicMock( + base_url="http://override-host:8081", + workspace="test-workspace", + default_headers=None, + http_client=None, + ) + + client = AsyncNeMoPlatform(base_url="http://override-host:8081", workspace="test-workspace") + try: + assert str(client.base_url).rstrip("/") == "http://override-host:8081" + finally: + await client.close() + + assert mock_build_client_kwargs.call_args.kwargs["base_url"] == "http://override-host:8081" + @pytest.mark.asyncio @patch("nemo_platform.client.factory.build_async_client_init_kwargs") async def test_async_constructor_with_http_client_skips_config_bootstrap(self, mock_build_client_kwargs): diff --git a/packages/nemo_platform_ext/tests/config/test_config.py b/packages/nemo_platform_ext/tests/config/test_config.py index 87889f9955..652c879681 100644 --- a/packages/nemo_platform_ext/tests/config/test_config.py +++ b/packages/nemo_platform_ext/tests/config/test_config.py @@ -13,6 +13,7 @@ get_context, ) from nemo_platform_ext.config.models import DEFAULT_BASE_URL, ConfigFile, NoAuthUser, OAuthUser +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR @pytest.fixture @@ -337,18 +338,17 @@ def test_config_from_env_access_token_only(self, tmp_path: Path, monkeypatch: py assert config.user.refresh_token is None assert not hasattr(config.user, "token_endpoint") - def test_config_from_workload_token_env_only(self, monkeypatch: pytest.MonkeyPatch): - """NEMO_WORKLOAD_TOKEN should bootstrap OAuth auth without a config file.""" + def test_legacy_workload_token_env_is_ignored(self, monkeypatch: pytest.MonkeyPatch): + """NEMO_WORKLOAD_TOKEN is no longer a runtime auth bootstrap source.""" monkeypatch.setenv("NMP_BASE_URL", "https://api.example.com") monkeypatch.setenv("NEMO_WORKLOAD_TOKEN", "workload-token-123") config = get_context() - assert isinstance(config.user, OAuthUser) - assert config.user.token.get_secret_value() == "workload-token-123" + assert isinstance(config.user, NoAuthUser) - def test_config_from_workload_token_file_env_only(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): - """NEMO_WORKLOAD_TOKEN_FILE should bootstrap OAuth auth without a config file.""" + def test_legacy_workload_token_file_env_is_ignored(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + """NEMO_WORKLOAD_TOKEN_FILE is no longer read as a bearer-token file.""" token_path = tmp_path / "workload.token" token_path.write_text("workload-token-from-file\n", encoding="utf-8") monkeypatch.setenv("NMP_BASE_URL", "https://api.example.com") @@ -356,25 +356,37 @@ def test_config_from_workload_token_file_env_only(self, tmp_path: Path, monkeypa config = get_context() - assert isinstance(config.user, OAuthUser) - assert config.user.token.get_secret_value() == "workload-token-from-file" + assert isinstance(config.user, NoAuthUser) - def test_config_from_missing_workload_token_file_reports_configuration_error( - self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch - ): - """NEMO_WORKLOAD_TOKEN_FILE should fail clearly when the configured token file cannot be read.""" + def test_missing_legacy_workload_token_file_is_ignored(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + """Legacy workload token file env no longer triggers config-time file reads.""" token_path = tmp_path / "missing-workload.token" monkeypatch.setenv("NMP_BASE_URL", "https://api.example.com") monkeypatch.setenv("NEMO_WORKLOAD_TOKEN_FILE", str(token_path)) - with pytest.raises(ValueError, match="NEMO_WORKLOAD_TOKEN_FILE"): - get_context() + config = get_context() + + assert isinstance(config.user, NoAuthUser) + + def test_workload_identity_token_file_is_not_static_access_token( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ): + """The workload identity token file env var is handled by the client factory, not Config.""" + token_path = tmp_path / "workload.token" + token_path.write_text("subject-token\n", encoding="utf-8") + monkeypatch.setenv("NMP_BASE_URL", "https://api.example.com") + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(token_path)) + + config = get_context() + + assert isinstance(config.user, NoAuthUser) def test_nmp_access_token_precedes_workload_token_env(self, monkeypatch: pytest.MonkeyPatch): """NMP_ACCESS_TOKEN remains the highest-precedence token env var.""" monkeypatch.setenv("NMP_BASE_URL", "https://api.example.com") monkeypatch.setenv("NMP_ACCESS_TOKEN", "preferred-token") monkeypatch.setenv("NEMO_WORKLOAD_TOKEN", "workload-token-123") + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, "/var/run/secrets/nemo-platform/workload/token") config = get_context() @@ -389,6 +401,7 @@ def test_runtime_access_token_source_label_uses_config_precedence( monkeypatch.setenv("NMP_ACCESS_TOKEN", "preferred-token") monkeypatch.setenv("NEMO_WORKLOAD_TOKEN", "workload-token-123") monkeypatch.setenv("NEMO_WORKLOAD_TOKEN_FILE", str(missing_token_path)) + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(missing_token_path)) assert Config.runtime_access_token_source_label() == "NMP_ACCESS_TOKEN environment override" diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py index 23747bb25c..d67180020d 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/adapter.py @@ -17,8 +17,10 @@ def make_sync_resource(platform: NeMoPlatform) -> NemoClient: from __future__ import annotations +from collections.abc import Callable from typing import TypeVar, overload +import httpx from nemo_platform import AsyncNeMoPlatform, NeMoPlatform from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient from nemo_platform_plugin.client.types import RetryPolicy @@ -27,6 +29,14 @@ def make_sync_resource(platform: NeMoPlatform) -> NemoClient: AsyncT = TypeVar("AsyncT", bound=AsyncNemoClient) +def _url_resolver_from_platform(platform: NeMoPlatform | AsyncNeMoPlatform) -> Callable[[str], str | httpx.URL]: + router = getattr(platform, "_nmp_request_router", None) + resolver = getattr(router, "resolve", None) + if resolver is not None: + return resolver + return platform._prepare_url + + @overload def client_from_platform(platform: NeMoPlatform, client_cls: type[SyncT]) -> SyncT: ... @overload @@ -52,6 +62,7 @@ def client_from_platform( headers = {k: v for k, v in platform._client.headers.items() if k.lower() not in _skip} # type: ignore[union-attr] retry = RetryPolicy(max_retries=platform.max_retries) + url_resolver = _url_resolver_from_platform(platform) if isinstance(platform, AsyncNeMoPlatform): if not issubclass(client_cls, AsyncNemoClient): raise TypeError("AsyncNeMoPlatform requires an AsyncNemoClient class") @@ -61,6 +72,7 @@ def client_from_platform( default_headers=headers or None, retry=retry, http_client=platform._client, + url_resolver=url_resolver, ) if not issubclass(client_cls, NemoClient): raise TypeError("NeMoPlatform requires a NemoClient class") @@ -70,4 +82,5 @@ def client_from_platform( default_headers=headers or None, retry=retry, http_client=platform._client, + url_resolver=url_resolver, ) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py index daa3b90e68..feb69df85c 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py @@ -20,8 +20,9 @@ import copy import inspect import json +import os import time -from collections.abc import AsyncIterator, Iterator, Mapping +from collections.abc import AsyncIterator, Callable, Iterator, Mapping from contextlib import asynccontextmanager, contextmanager from functools import cache from pathlib import Path @@ -30,6 +31,7 @@ import httpx from nemo_platform_plugin.client.auth import ( + AsyncTokenProvider, StaticToken, TokenProvider, ) @@ -170,15 +172,17 @@ def __init__( *, base_url: str, workspace: str | None = None, - auth: TokenProvider | str | None = None, + auth: TokenProvider | AsyncTokenProvider | str | None = None, retry: RetryPolicy | None = None, default_headers: Mapping[str, str] | None = None, + url_resolver: Callable[[str], str | httpx.URL] | None = None, ) -> None: self._base_url = base_url.rstrip("/") self._workspace = workspace - self._auth: TokenProvider | None = StaticToken(auth) if isinstance(auth, str) else auth + self._auth: TokenProvider | AsyncTokenProvider | None = StaticToken(auth) if isinstance(auth, str) else auth self._retry = retry self._default_headers = dict(default_headers) if default_headers else {} + self._url_resolver = url_resolver self._timeout: float | httpx.Timeout | None = None @property @@ -215,7 +219,10 @@ def _resolve_path(self, request: PreparedRequest) -> str: path = request.path_template.format_map(encoded_params) except KeyError as exc: raise ValueError(f"Missing path parameter {exc} for {request.method} {request.path_template}") from exc - return self._base_url + path + url = self._base_url + path + if self._url_resolver is not None: + return str(self._url_resolver(url)) + return url def _request_headers(self, request: PreparedRequest) -> dict[str, str] | None: headers: dict[str, str] = {} @@ -287,6 +294,8 @@ def _resolve_query_params(self, request: PreparedRequest) -> dict[str, str | int class NemoClient(BaseNemoClient): """Sync HTTP client for NeMo Platform APIs.""" + _auth: TokenProvider | None + def __init__( self, *, @@ -297,9 +306,15 @@ def __init__( timeout: float = DEFAULT_TIMEOUT, retry: RetryPolicy | None = None, http_client: httpx.Client | None = None, + url_resolver: Callable[[str], str | httpx.URL] | None = None, ) -> None: super().__init__( - base_url=base_url, workspace=workspace, auth=auth, retry=retry, default_headers=default_headers + base_url=base_url, + workspace=workspace, + auth=auth, + retry=retry, + default_headers=default_headers, + url_resolver=url_resolver, ) self._http = http_client or httpx.Client( headers=dict(default_headers) if default_headers else None, @@ -316,6 +331,7 @@ def from_client(cls, client: NemoClient) -> Self: default_headers=client._default_headers or None, retry=client._retry, http_client=client._http, + url_resolver=client._url_resolver, ) @overload @@ -533,14 +549,20 @@ def __init__( *, base_url: str, workspace: str | None = None, - auth: TokenProvider | str | None = None, + auth: TokenProvider | AsyncTokenProvider | str | None = None, default_headers: Mapping[str, str] | None = None, timeout: float = DEFAULT_TIMEOUT, retry: RetryPolicy | None = None, http_client: httpx.AsyncClient | None = None, + url_resolver: Callable[[str], str | httpx.URL] | None = None, ) -> None: super().__init__( - base_url=base_url, workspace=workspace, auth=auth, retry=retry, default_headers=default_headers + base_url=base_url, + workspace=workspace, + auth=auth, + retry=retry, + default_headers=default_headers, + url_resolver=url_resolver, ) self._http = http_client or httpx.AsyncClient( headers=dict(default_headers) if default_headers else None, @@ -557,6 +579,7 @@ def from_client(cls, client: AsyncNemoClient) -> Self: default_headers=client._default_headers or None, retry=client._retry, http_client=client._http, + url_resolver=client._url_resolver, ) @overload @@ -771,7 +794,8 @@ def _client_from_config( """Shared implementation for NemoClient.from_config / AsyncNemoClient.from_config.""" from nemo_platform_plugin.client.config.config import Config from nemo_platform_plugin.client.config.models import ConfigParams, OAuthUser - from nemo_platform_plugin.client.oidc_factory import resolve_oidc_provider + from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR + from nemo_platform_plugin.client.oidc_factory import resolve_oidc_provider, resolve_workload_exchange_provider resolved_path = Path(config_path) if isinstance(config_path, str) else config_path overrides: ConfigParams | None = None @@ -786,8 +810,14 @@ def _client_from_config( ctx = config.resolve() auth: TokenProvider | str | None = None + workload_identity_token_file = os.environ.get(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR) - if isinstance(ctx.user, OAuthUser): + if workload_identity_token_file and not explicit_access_token: + auth = resolve_workload_exchange_provider( + base_url=str(ctx.cluster.base_url), + subject_token_file=Path(workload_identity_token_file), + ) + elif isinstance(ctx.user, OAuthUser): auth = resolve_oidc_provider( base_url=str(ctx.cluster.base_url), context_name=ctx.context_name, diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/constants.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/constants.py new file mode 100644 index 0000000000..ebd61b5bab --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/constants.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared client constants.""" + +WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR = "NMP_WORKLOAD_IDENTITY_TOKEN_FILE" diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/oidc.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/oidc.py index 8e466e3df6..ed396993a3 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/oidc.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/oidc.py @@ -18,15 +18,20 @@ import base64 import json import logging +import math import threading import time from collections.abc import Callable from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass, field +from ipaddress import ip_address +from pathlib import Path from typing import Any from urllib.parse import urlparse import httpx +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR +from nemo_platform_plugin.client.tls import client_verify_from_env logger = logging.getLogger(__name__) @@ -118,6 +123,9 @@ def generate_unsigned_jwt( # --------------------------------------------------------------------------- DEFAULT_OAUTH_SCOPES = "openid profile email offline_access" +TOKEN_EXCHANGE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange" +JWT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt" +ACCESS_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:access_token" @dataclass(frozen=True) @@ -131,6 +139,11 @@ class NMPOIDCConfig: device_authorization_endpoint: str | None = None default_scopes: str = DEFAULT_OAUTH_SCOPES scope_prefix: str | None = None + workload_token_exchange_enabled: bool = False + workload_client_id: str | None = None + workload_token_endpoint: str | None = None + workload_audience: str | None = None + workload_scope: str | None = None def discover_nmp_config(base_url: str, timeout: float = 10.0) -> NMPOIDCConfig: @@ -138,6 +151,7 @@ def discover_nmp_config(base_url: str, timeout: float = 10.0) -> NMPOIDCConfig: response = httpx.get( f"{base_url.rstrip('/')}/apis/auth/discovery", timeout=timeout, + verify=client_verify_from_env(), ) response.raise_for_status() data = response.json() @@ -151,6 +165,11 @@ def discover_nmp_config(base_url: str, timeout: float = 10.0) -> NMPOIDCConfig: device_authorization_endpoint=oidc.get("device_authorization_endpoint"), default_scopes=oidc.get("default_scopes", DEFAULT_OAUTH_SCOPES), scope_prefix=oidc.get("scope_prefix"), + workload_token_exchange_enabled=oidc.get("workload_token_exchange_enabled", False), + workload_client_id=oidc.get("workload_client_id"), + workload_token_endpoint=oidc.get("workload_token_endpoint"), + workload_audience=oidc.get("workload_audience"), + workload_scope=oidc.get("workload_scope"), ) @@ -197,12 +216,23 @@ def build_effective_scope(requested_scopes: str, scope_prefix: str | None) -> st DEFAULT_REFRESH_MARGIN_SECONDS = 60 +def _is_loopback_host(hostname: str | None) -> bool: + if hostname == "localhost": + return True + if hostname is None: + return False + try: + return ip_address(hostname).is_loopback + except ValueError: + return False + + def _validate_token_endpoint(token_endpoint: str) -> None: """Reject non-HTTPS token endpoints (except loopback for local dev).""" parsed = urlparse(token_endpoint) if parsed.scheme == "https": return - if parsed.scheme == "http" and parsed.hostname in {"localhost", "127.0.0.1", "::1"}: + if parsed.scheme == "http" and _is_loopback_host(parsed.hostname): return raise ValueError( f"OIDC token endpoint must use HTTPS (got {token_endpoint!r}). " @@ -228,7 +258,7 @@ def refresh_token_grant( if scope: data["scope"] = scope - response = httpx.post(token_endpoint, data=data, timeout=timeout) + response = httpx.post(token_endpoint, data=data, timeout=timeout, verify=client_verify_from_env()) if response.status_code != 200: error_data: dict[str, str] = {} @@ -244,6 +274,101 @@ def refresh_token_grant( return response.json() +class WorkloadTokenExchangeError(RuntimeError): + """Raised when workload identity token exchange fails.""" + + +def _workload_exchange_error(error: str, error_description: str) -> WorkloadTokenExchangeError: + return WorkloadTokenExchangeError(f"Workload token exchange failed: {error} - {error_description}") + + +def read_subject_token_file(path: Path) -> str: + try: + token = path.read_text(encoding="utf-8").strip() + except OSError as exc: + raise ValueError(f"Unable to read {WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR} at {path}: {exc}") from exc + if not token: + raise ValueError(f"{WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR} at {path} is empty") + return token + + +def token_exchange_grant( + *, + token_endpoint: str, + client_id: str, + subject_token: str, + audience: str | None = None, + scope: str | None = None, + allow_http: bool = False, + timeout: float = 30.0, +) -> dict[str, Any]: + """Execute an OAuth 2.0 Token Exchange (RFC 8693) request.""" + # Backward-compatible argument: endpoint host controls HTTP allowance. + _validate_token_endpoint(token_endpoint) + data: dict[str, str] = { + "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, + "client_id": client_id, + "subject_token": subject_token, + "subject_token_type": JWT_TOKEN_TYPE, + "requested_token_type": ACCESS_TOKEN_TYPE, + } + if audience: + data["audience"] = audience + if scope: + data["scope"] = scope + + response = httpx.post(token_endpoint, data=data, timeout=timeout, verify=client_verify_from_env()) + if response.status_code != 200: + error_data: dict[str, object] = {} + if response.headers.get("content-type", "").startswith("application/json"): + error_data = _response_json_object( + response, + error_description="Token endpoint error response was not a JSON object", + ) + error = _response_string(error_data, "error", "unknown_error") + error_description = _response_string(error_data, "error_description", response.text) + raise _workload_exchange_error(error, error_description) + + token_data = _response_json_object( + response, + error_description="Token endpoint response was not a JSON object", + ) + _access_token_from_response(token_data) + return token_data + + +def _response_json_object(response: httpx.Response, *, error_description: str) -> dict[str, object]: + try: + payload = response.json() + except (json.JSONDecodeError, ValueError) as exc: + raise _workload_exchange_error("invalid_response", error_description) from exc + if not isinstance(payload, dict): + raise _workload_exchange_error("invalid_response", error_description) + return payload + + +def _response_string(payload: dict[str, object], key: str, default: str) -> str: + value = payload.get(key) + return value if isinstance(value, str) and value else default + + +def _access_token_from_response(token_data: dict[str, object]) -> str: + access_token = token_data.get("access_token") + if not isinstance(access_token, str) or not access_token.strip(): + raise _workload_exchange_error( + "invalid_response", + "Token endpoint response did not include a non-empty access_token", + ) + return access_token + + +def _expires_in_from_response(token_data: dict[str, object]) -> int | float | None: + expires_in = token_data.get("expires_in") + if isinstance(expires_in, bool): + return None + return expires_in if isinstance(expires_in, int | float) else None + + # --------------------------------------------------------------------------- # TokenSet # --------------------------------------------------------------------------- @@ -426,3 +551,61 @@ def force_refresh(self) -> str: with self._lock: self._refresh(force=True) return self.tokens.access_token + + +@dataclass +class WorkloadTokenExchangeProvider: + """Provides access tokens by exchanging a refreshed workload subject-token file.""" + + token_endpoint: str + client_id: str + subject_token_file: Path + audience: str | None = None + scope: str | None = None + allow_http: bool = False + refresh_margin_seconds: float = DEFAULT_REFRESH_MARGIN_SECONDS + tokens: TokenSet | None = field(default=None, repr=False) + _lock: threading.Lock = field(default_factory=threading.Lock, repr=False) + + def get_access_token(self) -> str: + with self._lock: + if self.tokens is None or self.tokens.is_expired(self.refresh_margin_seconds): + self._exchange() + assert self.tokens is not None + return self.tokens.access_token + + async def get_access_token_async(self) -> str: + return await asyncio.to_thread(self.get_access_token) + + def _exchange(self) -> None: + subject_token = read_subject_token_file(self.subject_token_file) + token_data = token_exchange_grant( + token_endpoint=self.token_endpoint, + client_id=self.client_id, + subject_token=subject_token, + audience=self.audience, + scope=self.scope, + allow_http=self.allow_http, + ) + access_token = _access_token_from_response(token_data) + try: + tokens = TokenSet.from_access_token( + access_token, + expires_in=_expires_in_from_response(token_data), + ) + except (OverflowError, TypeError, ValueError) as exc: + raise _workload_exchange_error( + "invalid_response", + "Token endpoint response did not include a usable access_token lifetime", + ) from exc + if tokens.expires_at is None or not math.isfinite(tokens.expires_at): + raise _workload_exchange_error( + "invalid_response", + "Token endpoint response did not include a usable access_token lifetime", + ) + if tokens.is_expired(0): + raise _workload_exchange_error( + "invalid_response", + "Token endpoint response returned an expired access_token", + ) + self.tokens = tokens diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/oidc_factory.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/oidc_factory.py index d4244bdc07..4ea447f15c 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/oidc_factory.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/oidc_factory.py @@ -19,10 +19,12 @@ from pathlib import Path from nemo_platform_plugin.client.auth import AuthError +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR from nemo_platform_plugin.client.oidc import ( DEFAULT_REFRESH_MARGIN_SECONDS, OIDCTokenProvider, TokenSet, + WorkloadTokenExchangeProvider, _discover_oidc_client_settings, build_effective_scope, ) @@ -210,3 +212,32 @@ def resolve_oidc_provider( refresh_margin_seconds=DEFAULT_REFRESH_MARGIN_SECONDS, refresh_scope=refresh_scope, ) + + +def resolve_workload_exchange_provider(*, base_url: str, subject_token_file: Path) -> WorkloadTokenExchangeProvider: + """Build a workload identity token exchange provider from auth discovery.""" + oidc_config = _discover_oidc_client_settings(base_url) + if not oidc_config.workload_token_exchange_enabled: + raise AuthError( + f"{WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR} is set but workload token exchange is not enabled by auth discovery" + ) + + token_endpoint = oidc_config.workload_token_endpoint or oidc_config.token_endpoint or "" + client_id = oidc_config.workload_client_id or oidc_config.client_id or "" + if not token_endpoint: + raise AuthError( + "Workload token exchange is enabled but auth discovery did not return workload_token_endpoint or token_endpoint" + ) + if not client_id: + raise AuthError( + "Workload token exchange is enabled but auth discovery did not return workload_client_id or client_id" + ) + + return WorkloadTokenExchangeProvider( + token_endpoint=token_endpoint, + client_id=client_id, + subject_token_file=subject_token_file, + audience=oidc_config.workload_audience, + scope=oidc_config.workload_scope, + refresh_margin_seconds=DEFAULT_REFRESH_MARGIN_SECONDS, + ) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/tls.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/tls.py new file mode 100644 index 0000000000..86f8e5871f --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/client/tls.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""TLS configuration shared by NeMo Platform plugin clients.""" + +from __future__ import annotations + +import os + +NMP_CLIENT_SSL_CERT_FILE_ENVVAR = "NMP_CLIENT_SSL_CERT_FILE" + + +def client_verify_from_env() -> str | bool: + """Return the httpx verify setting for NeMo Platform client requests.""" + cert_file = os.environ.get(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "").strip() + return cert_file or True diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/config.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/config.py index b6f8b40496..4703772145 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/config.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/config.py @@ -169,6 +169,7 @@ def global_settings_key() -> str: ... _T_config = TypeVar("_T_config", bound=ServiceConfig) +_T_platform_config = TypeVar("_T_platform_config", bound="NemoPlatformConfig") def create_service_config_class(service_name: str) -> Type[ServiceConfig]: @@ -208,6 +209,7 @@ class Configuration: """Singleton config loader from YAML files and environment variables.""" _overrides: ClassVar[dict[Type[ServiceConfig], ServiceConfig]] = {} + _platform_config_class: ClassVar[Type["NemoPlatformConfig"] | None] = None @classmethod def set_override(cls, config: ServiceConfig) -> None: @@ -284,9 +286,37 @@ def _get_cached_config(service_config: Type[_T_config]) -> _T_config: @classmethod def get_service_config(cls, service_config: Type[_T_config]) -> _T_config: if service_config in cls._overrides: - return cls._overrides[service_config] # type: ignore[return-value] + override = cls._overrides[service_config] + if not isinstance(override, service_config): + raise TypeError(f"Config override for {service_config.__name__} has type {type(override).__name__}") + return override return cls._get_cached_config(service_config) + @classmethod + def register_platform_config_class(cls, platform_config_class: Type[_T_platform_config]) -> None: + """Register the config class returned by get_platform_config(). + + Platform packages can use this to provide an extended platform config + without monkey-patching Configuration. + """ + if not issubclass(platform_config_class, NemoPlatformConfig): + raise TypeError("platform_config_class must be a NemoPlatformConfig subclass") + if platform_config_class.global_settings_key() != "platform": + raise ValueError("platform_config_class must use global settings key 'platform'") + cls._platform_config_class = platform_config_class + cls.clear_cache() + + @classmethod + def get_platform_config_class(cls) -> Type["NemoPlatformConfig"]: + """Return the currently registered platform config class.""" + if cls._platform_config_class is None: + raise RuntimeError("No platform config class is registered.") + return cls._platform_config_class + + @classmethod + def get_platform_config(cls) -> "NemoPlatformConfig": + return cls.get_service_config(cls.get_platform_config_class()) + @staticmethod def get_service_config_from_file(filename: str, service_config: Type[_T_config]) -> _T_config: return Configuration.global_settings_to_service_config( @@ -299,6 +329,16 @@ def get_service_config(service_config: Type[_T_config]) -> _T_config: return Configuration.get_service_config(service_config) +def register_platform_config_class(platform_config_class: Type[_T_platform_config]) -> None: + """Register the config class returned by get_platform_config().""" + Configuration.register_platform_config_class(platform_config_class) + + +def get_platform_config_class() -> Type[NemoPlatformConfig]: + """Return the currently registered platform config class.""" + return Configuration.get_platform_config_class() + + # --------------------------------------------------------------------------- # Platform configuration types # --------------------------------------------------------------------------- @@ -456,6 +496,15 @@ def global_settings_key() -> str: @classmethod def get(cls) -> NemoPlatformConfig: + """Return the active platform config. + + ``NemoPlatformConfig`` itself is a public handle for the registered + platform config class. Platform packages may register a subclass with + extra behavior, so callers of the base class should receive that active + subclass. Calls on a concrete subclass still load that subclass directly. + """ + if cls is NemoPlatformConfig: + return Configuration.get_platform_config() return Configuration.get_service_config(cls) services: str = internal_field( @@ -637,15 +686,15 @@ def get_nemo_platform_config() -> NemoPlatformConfig: return Configuration.get_platform_config() -# Default implementation — subclasses (nmp-common) override this to return their extended PlatformConfig. -Configuration.get_platform_config = classmethod(lambda cls: cls.get_service_config(NemoPlatformConfig)) # type: ignore[attr-defined] +# Default implementation; platform packages can register an extended class. +Configuration.register_platform_config_class(NemoPlatformConfig) # Aliases — nmp-common and services use PlatformConfig / get_platform_config PlatformConfig = NemoPlatformConfig get_platform_config = get_nemo_platform_config -class CommonServiceConfig(create_service_config_class("service")): +class CommonServiceConfig(create_service_config_class("service")): # ty: ignore[unsupported-base] """Common configuration shared by all services. Reads from env vars with prefix ``NMP_SERVICE_`` and YAML key ``service``. @@ -687,10 +736,12 @@ def get_common_service_config() -> CommonServiceConfig: "get_nemo_config", "get_nemo_platform_config", "get_platform_config", + "get_platform_config_class", "get_service_config", "get_service_config_prefix", "internal_field", "nmp_user_data_dir", + "register_platform_config_class", "set_nemo_config_override", ] @@ -836,7 +887,7 @@ def __new__( # --------------------------------------------------------------------------- -class NemoConfig(EnvironmentFirstSettings, metaclass=_NemoConfigMeta): +class NemoConfig(ServiceConfig, metaclass=_NemoConfigMeta): """Base class for plugin configuration. Subclasses declare :attr:`plugin_name` and :attr:`plugin_description` as diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/execution_profiles.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/execution_profiles.py index 54319cc166..20c9b17025 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/execution_profiles.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/execution_profiles.py @@ -18,6 +18,7 @@ from typing import Any, Literal +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR from nemo_platform_plugin.config import NMP_CONFIG_WARNINGS_DISABLED_ENV_VAR from nemo_platform_plugin.jobs.constants import ( CONFIG_TASK_STORAGE_PATH_ENVVAR, @@ -39,10 +40,11 @@ # Default image used to set filesystem permissions on job storage volumes. DEFAULT_VOLUME_PERMISSIONS_IMAGE = "busybox" +JOB_LOGS_ENDPOINT_ENVVAR = "NMP_JOB_LOGS_ENDPOINT" # Env var names set by the platform during job creation; user-provided profile # environment must not conflict. The job-scoped names come from the shared -# ``jobs.constants`` leaf; the auth / config / telemetry names are stable env +# ``jobs.constants`` leaf; the auth / config / logging / telemetry names are stable env # var strings kept here to avoid importing server-side auth/config modules. RESERVED_JOB_ENVIRONMENT_VARIABLE_NAMES: frozenset[str] = frozenset( { @@ -61,6 +63,9 @@ TASK_CONFIG_ENVVAR, # Auth "NMP_PRINCIPAL", + WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, + # Platform launcher logs + JOB_LOGS_ENDPOINT_ENVVAR, # OTEL (telemetry) "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT", "OTEL_LOGS_EXPORTER", @@ -68,6 +73,7 @@ "OTEL_EXPORTER_OTLP_LOGS_HEADERS", # Platform shared envvars (to_shared_envvars with NMP_ prefix) NMP_CONFIG_WARNINGS_DISABLED_ENV_VAR, + "NMP_AUTH_URL", "NMP_BASE_URL", "NMP_JOBS_URL", "NMP_FILES_URL", diff --git a/packages/nemo_platform_plugin/tests/client/test_adapter.py b/packages/nemo_platform_plugin/tests/client/test_adapter.py index 48ca65656e..e51c4fd19d 100644 --- a/packages/nemo_platform_plugin/tests/client/test_adapter.py +++ b/packages/nemo_platform_plugin/tests/client/test_adapter.py @@ -6,6 +6,7 @@ import httpx from nemo_platform import NeMoPlatform from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.jobs import endpoints from nemo_platform_plugin.jobs.client import JobsClient @@ -23,3 +24,57 @@ def test_client_from_platform_preserves_retry_count_with_nemoclient_defaults() - assert client.retry is not None assert client.retry.max_retries == 4 assert client.retry.retryable_status_codes == (502, 503, 504, 429) + + +def test_client_from_platform_prefers_platform_request_router() -> None: + http_client = httpx.Client(transport=httpx.MockTransport(lambda request: httpx.Response(200, request=request))) + platform = NeMoPlatform( + base_url="http://gateway", + workspace="default", + http_client=http_client, + ) + + class RequestRouter: + def resolve(self, url: str) -> str: + return url.replace("http://gateway/apis/jobs", "http://127.0.0.1:8080/apis/jobs") + + platform._nmp_request_router = RequestRouter() # type: ignore[attr-defined] + + client = client_from_platform(platform, JobsClient) + + request = endpoints.list_steps(workspace="default", name="job-1") + assert client._resolve_path(request) == ("http://127.0.0.1:8080/apis/jobs/v2/workspaces/default/jobs/job-1/steps") + + +def test_client_from_platform_falls_back_to_sdk_prepare_url() -> None: + http_client = httpx.Client(transport=httpx.MockTransport(lambda request: httpx.Response(200, request=request))) + platform = NeMoPlatform( + base_url="http://gateway", + workspace="default", + http_client=http_client, + ) + + def prepare_url(url: str) -> str: + return url.replace("http://gateway/apis/jobs", "http://127.0.0.1:8080/apis/jobs") + + platform._prepare_url = prepare_url # type: ignore[method-assign] + + client = client_from_platform(platform, JobsClient) + + request = endpoints.list_steps(workspace="default", name="job-1") + assert client._resolve_path(request) == ("http://127.0.0.1:8080/apis/jobs/v2/workspaces/default/jobs/job-1/steps") + + +def test_from_client_preserves_url_resolver() -> None: + http_client = httpx.Client(transport=httpx.MockTransport(lambda request: httpx.Response(200, request=request))) + client = JobsClient( + base_url="http://gateway", + workspace="default", + http_client=http_client, + url_resolver=lambda url: url.replace("http://gateway/apis/jobs", "http://127.0.0.1:8080/apis/jobs"), + ) + + clone = JobsClient.from_client(client) + + request = endpoints.list_steps(workspace="default", name="job-1") + assert clone._resolve_path(request) == ("http://127.0.0.1:8080/apis/jobs/v2/workspaces/default/jobs/job-1/steps") diff --git a/packages/nemo_platform_plugin/tests/jobs/test_execution_profiles.py b/packages/nemo_platform_plugin/tests/jobs/test_execution_profiles.py new file mode 100644 index 0000000000..02de3bd299 --- /dev/null +++ b/packages/nemo_platform_plugin/tests/jobs/test_execution_profiles.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR +from nemo_platform_plugin.jobs.execution_profiles import JOB_LOGS_ENDPOINT_ENVVAR, JobExecutionProfileConfig +from pydantic import ValidationError + + +@pytest.mark.parametrize( + "envvar", + [ + JOB_LOGS_ENDPOINT_ENVVAR, + WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, + ], +) +def test_job_execution_profile_config_rejects_platform_injected_env_vars(envvar: str) -> None: + with pytest.raises(ValidationError, match=envvar): + JobExecutionProfileConfig(env={envvar: "override"}) diff --git a/packages/nemo_platform_plugin/tests/test_client_auth.py b/packages/nemo_platform_plugin/tests/test_client_auth.py index 5ad150ad39..ddb0c81c04 100644 --- a/packages/nemo_platform_plugin/tests/test_client_auth.py +++ b/packages/nemo_platform_plugin/tests/test_client_auth.py @@ -23,11 +23,22 @@ NoAuthUser, OAuthUser, ) +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR from nemo_platform_plugin.client.oidc import ( + ACCESS_TOKEN_TYPE, + JWT_TOKEN_TYPE, + TOKEN_EXCHANGE_GRANT_TYPE, + NMPOIDCConfig, OIDCTokenProvider, TokenSet, + WorkloadTokenExchangeError, + WorkloadTokenExchangeProvider, + discover_nmp_config, generate_unsigned_jwt, + refresh_token_grant, + token_exchange_grant, ) +from nemo_platform_plugin.client.tls import NMP_CLIENT_SSL_CERT_FILE_ENVVAR from nemo_platform_plugin.client_provider import ( get_async_nemo_client, get_nemo_client, @@ -179,6 +190,25 @@ def _make_jwt(exp: float | None = None, sub: str = "user") -> str: class TestOIDCTokenProvider: + def test_discover_nmp_config_uses_nemo_scoped_ca_bundle(self, monkeypatch): + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + with patch("nemo_platform_plugin.client.oidc.httpx.get") as mock_get: + mock_get.return_value = httpx.Response( + 200, + json={"auth_enabled": False}, + request=httpx.Request("GET", "https://nemo.example.com/apis/auth/discovery"), + ) + + result = discover_nmp_config("https://nemo.example.com") + + assert result.auth_enabled is False + mock_get.assert_called_once_with( + "https://nemo.example.com/apis/auth/discovery", + timeout=10.0, + verify="/tmp/nemo-ca.pem", + ) + def test_returns_token_when_not_expired(self): token = _make_jwt(exp=time.time() + 3600) provider = OIDCTokenProvider( @@ -257,6 +287,204 @@ def test_invalid_grant_recovery_with_shared_tokens(self): # Should have recovered by reloading fresh tokens from the shared store assert result == fresh_token + def test_refresh_token_grant_uses_nemo_scoped_ca_bundle(self, monkeypatch): + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + with patch("nemo_platform_plugin.client.oidc.httpx.post") as mock_post: + mock_post.return_value = httpx.Response(200, json={"access_token": "new-access"}) + + result = refresh_token_grant( + token_endpoint="https://idp.example.com/token", + client_id="client", + refresh_token="refresh-token", + ) + + assert result == {"access_token": "new-access"} + assert mock_post.call_args.kwargs["verify"] == "/tmp/nemo-ca.pem" + + +class TestWorkloadTokenExchangeProvider: + def test_token_exchange_grant_sends_rfc8693_request(self, monkeypatch): + monkeypatch.delenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, raising=False) + + with patch("nemo_platform_plugin.client.oidc.httpx.post") as mock_post: + mock_post.return_value = httpx.Response(200, json={"access_token": "exchanged-token", "expires_in": 300}) + + token_data = token_exchange_grant( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token="subject-token", + audience="nemo-platform", + scope="openid email groups", + timeout=5.0, + ) + + assert token_data["access_token"] == "exchanged-token" + mock_post.assert_called_once_with( + "https://idp.example.com/token", + data={ + "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, + "client_id": "nemo-platform-workload", + "subject_token": "subject-token", + "subject_token_type": JWT_TOKEN_TYPE, + "requested_token_type": ACCESS_TOKEN_TYPE, + "audience": "nemo-platform", + "scope": "openid email groups", + }, + timeout=5.0, + verify=True, + ) + + def test_token_exchange_grant_uses_nemo_scoped_ca_bundle(self, monkeypatch): + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + with patch("nemo_platform_plugin.client.oidc.httpx.post") as mock_post: + mock_post.return_value = httpx.Response(200, json={"access_token": "exchanged-token", "expires_in": 300}) + + token_data = token_exchange_grant( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token="subject-token", + ) + + assert token_data["access_token"] == "exchanged-token" + assert mock_post.call_args.kwargs["verify"] == "/tmp/nemo-ca.pem" + + def test_token_exchange_grant_allows_http_loopback_endpoint(self): + with patch("nemo_platform_plugin.client.oidc.httpx.post") as mock_post: + mock_post.return_value = httpx.Response(200, json={"access_token": "exchanged-token", "expires_in": 300}) + + token_data = token_exchange_grant( + token_endpoint="http://localhost:8080/apis/auth/token", + client_id="nemo-platform-workload", + subject_token="subject-token", + ) + + assert token_data["access_token"] == "exchanged-token" + + def test_token_exchange_grant_rejects_http_non_loopback_endpoint_when_allow_http_is_set(self): + with patch("nemo_platform_plugin.client.oidc.httpx.post") as mock_post: + with pytest.raises(ValueError, match="HTTPS"): + token_exchange_grant( + token_endpoint="http://nemo-gateway:8080/apis/auth/token", + client_id="nemo-platform-workload", + subject_token="subject-token", + allow_http=True, + ) + + mock_post.assert_not_called() + + def test_token_exchange_grant_rejects_non_object_error_payload(self): + with patch("nemo_platform_plugin.client.oidc.httpx.post") as mock_post: + mock_post.return_value = httpx.Response( + 400, + json=[], + headers={"content-type": "application/json"}, + ) + + with pytest.raises(WorkloadTokenExchangeError, match="invalid_response - Token endpoint error response"): + token_exchange_grant( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token="bad-subject-token", + ) + + @pytest.mark.parametrize("payload", [[], {}, {"access_token": ""}, {"access_token": None}]) + def test_token_exchange_grant_rejects_success_response_without_non_empty_access_token(self, payload): + with patch("nemo_platform_plugin.client.oidc.httpx.post") as mock_post: + mock_post.return_value = httpx.Response(200, json=payload) + + with pytest.raises(WorkloadTokenExchangeError, match="invalid_response"): + token_exchange_grant( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token="subject-token", + ) + + def test_token_exchange_grant_rejects_non_json_success_response(self): + with patch("nemo_platform_plugin.client.oidc.httpx.post") as mock_post: + mock_post.return_value = httpx.Response(200, content=b"not-json") + + with pytest.raises(WorkloadTokenExchangeError, match="Token endpoint response was not a JSON object"): + token_exchange_grant( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token="subject-token", + ) + + def test_provider_rejects_exchange_response_without_access_token(self, tmp_path): + subject_token_file = tmp_path / "token" + subject_token_file.write_text("subject-token", encoding="utf-8") + provider = WorkloadTokenExchangeProvider( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token_file=subject_token_file, + ) + + with ( + patch("nemo_platform_plugin.client.oidc.token_exchange_grant", return_value={}), + pytest.raises(WorkloadTokenExchangeError, match="non-empty access_token"), + ): + provider.get_access_token() + + @pytest.mark.parametrize( + "expires_in", + [None, "300", True, float("nan"), float("inf"), 10**400], + ) + def test_provider_rejects_exchange_response_without_usable_lifetime(self, tmp_path, expires_in): + subject_token_file = tmp_path / "token" + subject_token_file.write_text("subject-token", encoding="utf-8") + token_data = {"access_token": "opaque-access-token"} + if expires_in is not None: + token_data["expires_in"] = expires_in + provider = WorkloadTokenExchangeProvider( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token_file=subject_token_file, + ) + + with ( + patch("nemo_platform_plugin.client.oidc.token_exchange_grant", return_value=token_data), + pytest.raises(WorkloadTokenExchangeError, match="usable access_token lifetime"), + ): + provider.get_access_token() + + assert provider.tokens is None + + def test_provider_rejects_expired_exchange_response_and_retries_with_current_subject_token(self, tmp_path): + subject_token_file = tmp_path / "token" + subject_token_file.write_text("subject-token-one", encoding="utf-8") + expired_access_token = _make_jwt(exp=time.time() - 10) + fresh_access_token = _make_jwt(exp=time.time() + 3600) + provider = WorkloadTokenExchangeProvider( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token_file=subject_token_file, + audience="nemo-platform", + scope="openid email groups", + refresh_margin_seconds=0, + ) + + with patch( + "nemo_platform_plugin.client.oidc.token_exchange_grant", + side_effect=[ + {"access_token": expired_access_token}, + {"access_token": fresh_access_token}, + ], + ) as mock_exchange: + with pytest.raises(WorkloadTokenExchangeError, match="expired access_token"): + provider.get_access_token() + + assert provider.tokens is None + subject_token_file.write_text("subject-token-two", encoding="utf-8") + + assert provider.get_access_token() == fresh_access_token + assert mock_exchange.call_count == 2 + assert mock_exchange.call_args_list[0].kwargs["subject_token"] == "subject-token-one" + assert mock_exchange.call_args_list[1].kwargs["subject_token"] == "subject-token-two" + assert mock_exchange.call_args_list[1].kwargs["audience"] == "nemo-platform" + assert mock_exchange.call_args_list[1].kwargs["scope"] == "openid email groups" + # --------------------------------------------------------------------------- # Config @@ -395,6 +623,104 @@ def test_from_config_with_no_auth(self, tmp_path): assert client.base_url == "http://localhost:9090" assert client._auth is None + @respx.mock + def test_from_config_with_workload_identity_token_file(self, tmp_path, monkeypatch): + subject_token_file = tmp_path / "workload-token" + subject_token_file.write_text("subject-token\n", encoding="utf-8") + exchanged_token = _make_jwt() + config_data = { + "current_context": "test", + "clusters": [{"name": "test-cluster", "base_url": "http://localhost:9090"}], + "users": [{"name": "test-user", "type": "no-auth"}], + "contexts": [{"name": "test", "cluster": "test-cluster", "user": "test-user"}], + } + config_file = tmp_path / "config.yaml" + config_file.write_text(yaml.safe_dump(config_data)) + route = respx.get("http://localhost:9090/test").mock(return_value=httpx.Response(200, json={"ok": True})) + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(subject_token_file)) + monkeypatch.delenv("NMP_ACCESS_TOKEN", raising=False) + + with ( + patch("nemo_platform_plugin.client.oidc_factory._discover_oidc_client_settings") as mock_discover, + patch("nemo_platform_plugin.client.oidc.token_exchange_grant") as mock_exchange, + ): + mock_discover.return_value = NMPOIDCConfig( + auth_enabled=True, + client_id="nemo-platform-cli", + token_endpoint="https://idp.example.com/token", + workload_token_exchange_enabled=True, + workload_client_id="nemo-platform-workload", + workload_token_endpoint="https://workload-idp.example.com/token", + workload_audience="nemo-platform", + workload_scope="openid email groups", + ) + mock_exchange.return_value = {"access_token": exchanged_token, "expires_in": 300} + + client = NemoClient.from_config(config_path=config_file) + + from nemo_platform_plugin.client.types import PreparedRequest + + req = PreparedRequest( + method="GET", + path_template="/test", + path_params={}, + content=None, + content_type=None, + response_type=None, + ) + client.send(req) + + assert route.calls[0].request.headers["Authorization"] == f"Bearer {exchanged_token}" + assert mock_exchange.call_args.kwargs["token_endpoint"] == "https://workload-idp.example.com/token" + assert mock_exchange.call_args.kwargs["client_id"] == "nemo-platform-workload" + assert mock_exchange.call_args.kwargs["subject_token"] == "subject-token" + assert mock_exchange.call_args.kwargs["audience"] == "nemo-platform" + assert mock_exchange.call_args.kwargs["scope"] == "openid email groups" + assert mock_exchange.call_args.kwargs["allow_http"] is False + + def test_from_config_rejects_discovered_http_non_loopback_token_endpoint(self, tmp_path, monkeypatch): + subject_token_file = tmp_path / "workload-token" + subject_token_file.write_text("subject-token\n", encoding="utf-8") + config_data = { + "current_context": "test", + "clusters": [{"name": "test-cluster", "base_url": "http://localhost:9090"}], + "users": [{"name": "test-user", "type": "no-auth"}], + "contexts": [{"name": "test", "cluster": "test-cluster", "user": "test-user"}], + } + config_file = tmp_path / "config.yaml" + config_file.write_text(yaml.safe_dump(config_data)) + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(subject_token_file)) + monkeypatch.delenv("NMP_ACCESS_TOKEN", raising=False) + + with ( + patch("nemo_platform_plugin.client.oidc_factory._discover_oidc_client_settings") as mock_discover, + patch("nemo_platform_plugin.client.oidc.httpx.post") as mock_post, + ): + mock_discover.return_value = NMPOIDCConfig( + auth_enabled=True, + client_id="nemo-platform-cli", + token_endpoint="https://idp.example.com/token", + workload_token_exchange_enabled=True, + workload_client_id="nemo-platform-workload", + workload_token_endpoint="http://idp.example.com/token", + ) + client = NemoClient.from_config(config_path=config_file) + + from nemo_platform_plugin.client.types import PreparedRequest + + req = PreparedRequest( + method="GET", + path_template="/test", + path_params={}, + content=None, + content_type=None, + response_type=None, + ) + with pytest.raises(ValueError, match="HTTPS"): + client.send(req) + + mock_post.assert_not_called() + def test_from_config_selects_context(self, tmp_path): """from_config(context='staging') uses the staging context, not the default.""" config_data = { diff --git a/packages/nemo_platform_plugin/tests/test_config.py b/packages/nemo_platform_plugin/tests/test_config.py index 2f10c21884..e1a7307446 100644 --- a/packages/nemo_platform_plugin/tests/test_config.py +++ b/packages/nemo_platform_plugin/tests/test_config.py @@ -9,6 +9,7 @@ import pytest from nemo_platform_plugin.config import ( + Configuration, NemoConfig, NemoPlatformConfig, PlatformConfig, @@ -17,6 +18,8 @@ get_nemo_config, get_nemo_platform_config, get_platform_config, + get_platform_config_class, + register_platform_config_class, set_nemo_config_override, ) from pydantic import Field @@ -245,6 +248,23 @@ def test_get_platform_config_is_callable() -> None: assert isinstance(result, PlatformConfig) +def test_register_platform_config_class_controls_platform_config_accessors() -> None: + """A registered platform config class is used by all platform config accessors.""" + + class _RegisteredPlatformConfig(NemoPlatformConfig): + registered_marker: str = "registered" + + previous_platform_config_class = get_platform_config_class() + register_platform_config_class(_RegisteredPlatformConfig) + try: + assert isinstance(Configuration.get_platform_config(), _RegisteredPlatformConfig) + assert isinstance(get_platform_config(), _RegisteredPlatformConfig) + assert isinstance(NemoPlatformConfig.get(), _RegisteredPlatformConfig) + assert isinstance(_RegisteredPlatformConfig.get(), _RegisteredPlatformConfig) + finally: + register_platform_config_class(previous_platform_config_class) + + def test_get_nemo_platform_config_alias() -> None: """get_nemo_platform_config is an alias for get_platform_config.""" assert get_nemo_platform_config is get_platform_config diff --git a/packages/nmp_common/src/nmp/common/auth/middleware.py b/packages/nmp_common/src/nmp/common/auth/middleware.py index a224f47f23..ed4bc11b38 100644 --- a/packages/nmp_common/src/nmp/common/auth/middleware.py +++ b/packages/nmp_common/src/nmp/common/auth/middleware.py @@ -69,6 +69,8 @@ def _embedded_pdp_base_url_hint(config: AuthConfig) -> str: "/health/ready", "/metrics", "/apis/auth/discovery", # Discovery endpoint for CLI/SDK + "/apis/auth/jwks", # Workload identity exchange signing keys + "/apis/auth/token", # Workload identity token exchange validates the subject token itself } # GET requests to these paths bypass authentication (e.g. / -> /studio redirect). diff --git a/packages/nmp_common/src/nmp/common/config/__init__.py b/packages/nmp_common/src/nmp/common/config/__init__.py index ce20e635b7..68807956df 100644 --- a/packages/nmp_common/src/nmp/common/config/__init__.py +++ b/packages/nmp_common/src/nmp/common/config/__init__.py @@ -25,9 +25,11 @@ get_auth_config, get_common_service_config, get_platform_config, + get_platform_config_class, get_service_config, get_service_config_prefix, internal_field, + register_platform_config_class, ) from nmp.common.config.paths import ( NMP_DATA_DIR_ENV_VAR, @@ -59,8 +61,10 @@ "get_auth_config", "get_common_service_config", "get_platform_config", + "get_platform_config_class", "get_service_config", "get_service_config_prefix", "internal_field", "nmp_user_data_dir", + "register_platform_config_class", ] diff --git a/packages/nmp_common/src/nmp/common/config/base.py b/packages/nmp_common/src/nmp/common/config/base.py index fc87ec9f9d..1bb654c6da 100644 --- a/packages/nmp_common/src/nmp/common/config/base.py +++ b/packages/nmp_common/src/nmp/common/config/base.py @@ -29,9 +29,11 @@ from nemo_platform_plugin.config import ServiceConfig as ServiceConfig from nemo_platform_plugin.config import create_service_config_class as create_service_config_class from nemo_platform_plugin.config import determine_loopback_override as determine_loopback_override +from nemo_platform_plugin.config import get_platform_config_class as get_platform_config_class from nemo_platform_plugin.config import get_service_config as get_service_config from nemo_platform_plugin.config import get_service_config_prefix as get_service_config_prefix from nemo_platform_plugin.config import internal_field as internal_field +from nemo_platform_plugin.config import register_platform_config_class as register_platform_config_class from nmp.common.config.paths import nmp_user_data_dir from pydantic import Field from pydantic_settings import BaseSettings, SettingsConfigDict @@ -60,9 +62,9 @@ def get_service_url(self, api_name: str) -> str: return super().get_service_url(api_name) -# Re-register PlatformConfig on Configuration so get_platform_config() returns -# the extended version (with local-service routing). -Configuration.get_platform_config = classmethod(lambda cls: cls.get_service_config(PlatformConfig)) # type: ignore[attr-defined] +# Register PlatformConfig so get_platform_config() returns the extended version +# with local-service routing. +Configuration.register_platform_config_class(PlatformConfig) class OIDCConfig(BaseSettings): @@ -143,6 +145,96 @@ class OIDCConfig(BaseSettings): "For Azure AD with custom API, use: 'api://{app-id}/.default openid profile email'", ) + workload_token_exchange_enabled: bool = Field( + default=False, + description="Enable SDK workload identity token exchange using NMP_WORKLOAD_IDENTITY_TOKEN_FILE subject tokens.", + ) + + workload_client_id: str | None = Field( + default=None, + description="OAuth client ID to use for workload identity token exchange. Defaults to client_id when unset.", + ) + + workload_token_endpoint: str | None = Field( + default=None, + description="OAuth token endpoint to use for workload identity token exchange. Defaults to token_endpoint.", + ) + + workload_audience: str | None = Field( + default=None, + description="RFC 8693 audience requested for workload identity token exchange.", + ) + + workload_scope: str | None = Field( + default=None, + description="Space-separated OAuth scopes requested for workload identity token exchange.", + ) + + workload_token_issuer: str | None = Field( + default=None, + description=( + "Issuer to stamp on workload identity access tokens minted by the NeMo auth service. " + "Defaults to the platform auth endpoint origin serving the token exchange request." + ), + ) + + workload_token_ttl_seconds: int = Field( + default=300, + ge=1, + description="Lifetime in seconds for workload identity access tokens minted by the NeMo auth service.", + ) + + workload_token_key_id: str = Field( + default="nemo-workload-exchange", + description="JWT key id advertised by the NeMo auth service workload identity JWKS endpoint.", + ) + + workload_token_private_key_file: str | None = Field( + default=None, + description=( + "Path to a PEM-encoded RSA private key used by the NeMo auth service to sign workload identity " + "access tokens. Intended for mounted shared secrets." + ), + ) + + workload_allowed_audiences: list[str] = Field( + default_factory=list, + description=( + "Additional RFC 8693 audience values accepted by the NeMo auth service workload token exchange endpoint. " + "The configured workload_audience is always accepted." + ), + ) + + workload_subject_jwks_uri: str | None = Field( + default=None, + description=( + "JWKS URI used by the NeMo auth service to validate JWT subject tokens for workload token exchange. " + "Leave unset when only Kubernetes TokenReview subject validation is enabled." + ), + ) + + workload_subject_issuers: list[str] = Field( + default_factory=list, + description=( + "Allowed JWT subject token issuers for workload token exchange. " + "Required when workload_subject_jwks_uri is set." + ), + ) + + workload_subject_jwks_cache_ttl_seconds: int = Field( + default=3600, + ge=0, + description=("TTL in seconds for caching workload subject JWKS responses. Set to 0 to disable caching."), + ) + + workload_kubernetes_token_review_enabled: bool = Field( + default=False, + description=( + "Allow the NeMo auth service workload token exchange endpoint to validate Kubernetes projected " + "service account subject tokens using the TokenReview API." + ), + ) + scope_prefix: str | None = Field( default=None, description="Prefix to strip from token scopes before authorization. " @@ -159,7 +251,7 @@ class OIDCConfig(BaseSettings): ) -class AuthConfig(create_service_config_class("auth")): +class AuthConfig(create_service_config_class("auth")): # ty: ignore[unsupported-base] """ Shared authorization configuration read from the 'auth' key in config.yaml. diff --git a/packages/nmp_common/src/nmp/common/entities/client.py b/packages/nmp_common/src/nmp/common/entities/client.py index 037b058a15..f2214cb3d3 100644 --- a/packages/nmp_common/src/nmp/common/entities/client.py +++ b/packages/nmp_common/src/nmp/common/entities/client.py @@ -47,10 +47,11 @@ def as_service(self, service_name: str, *, internal: bool = False) -> "EntityCli """ from nemo_platform.resources.entities import AsyncEntitiesResource from nmp.common.observability import MARK_INTERNAL_REQUEST_HEADERS + from nmp.common.sdk_factory import with_options_preserving_request_router underlying_sdk = self.entities_api._client headers: dict[str, str] = {"X-NMP-Principal-Id": f"service:{service_name}"} if internal: headers.update(MARK_INTERNAL_REQUEST_HEADERS) - service_sdk = underlying_sdk.with_options(set_default_headers=headers) + service_sdk = with_options_preserving_request_router(underlying_sdk, set_default_headers=headers) return EntityClient(AsyncEntitiesResource(service_sdk)) diff --git a/packages/nmp_common/src/nmp/common/sdk_factory.py b/packages/nmp_common/src/nmp/common/sdk_factory.py index 0d63504039..0a6545397c 100644 --- a/packages/nmp_common/src/nmp/common/sdk_factory.py +++ b/packages/nmp_common/src/nmp/common/sdk_factory.py @@ -4,17 +4,19 @@ """SDK factory functions for creating NeMo Platform SDK instances.""" import logging -from typing import Callable, Optional +from dataclasses import dataclass +from typing import Any, Callable, Optional, TypeVar, cast import httpx from nemo_platform import AsyncNeMoPlatform, NeMoPlatform from nmp.common.auth import Principal, get_principal_auth_headers, principal_from_env -from nmp.common.config import Configuration +from nmp.common.config import Configuration, PlatformConfig from nmp.common.http_clients import shared_async_http_client, shared_sync_http_client from nmp.common.observability import MARK_INTERNAL_REQUEST_HEADERS from nmp.common.observability.otel import get_otel_headers logger = logging.getLogger(__name__) +PlatformSDKT = TypeVar("PlatformSDKT", NeMoPlatform, AsyncNeMoPlatform) # Test-only: HTTP clients to use for SDK requests in test context. # Set by test fixtures to route requests through the in-process test transport. @@ -29,45 +31,93 @@ def _base_url_from_config() -> str: return Configuration.get_platform_config().base_url -def _create_url_router( - original: Callable[[str], httpx.URL], -) -> Callable[[str], httpx.URL]: - """Create a URL routing function that routes requests based on API name. +def resolve_platform_request_url( + url: str, + *, + platform_config: PlatformConfig, + default_resolver: Callable[[str], httpx.URL], +) -> httpx.URL: + """Resolve the destination URL for an SDK request. - Returns: - A function that routes URLs based on path segments in service_urls. - """ + The generated SDK builds requests from relative paths like + ``/apis/entities/v2/...`` and then calls its private ``_prepare_url`` hook. + Keep the routing policy here, not in that private hook: - platform_config = Configuration.get_platform_config() + - resolve the SDK URL normally against ``platform.base_url``; + - if the path targets ``/apis/{api_name}/...``, replace only the origin + with ``platform.get_service_url(api_name)``; + - preserve the original path and query string. + """ + request_url = default_resolver(url) service_pattern = platform_config.create_service_pattern() - - def route_url(url: str) -> httpx.URL: - # Try to match the API name in the URL - if service_pattern: - match = service_pattern.search(url) - if match: - api_name = match.group(1) - svc_url = httpx.URL(platform_config.get_service_url(api_name)) - request_url = httpx.URL(url) - logger.debug( - "Routing URL to matched service URL", - extra={"service": api_name, "url": url, "host": svc_url.host, "port": svc_url.port}, - ) - # Use scheme/host/port from service URL, path/params from request URL - return request_url.copy_with( - scheme=svc_url.scheme, - host=svc_url.host, - port=svc_url.port, - ) - request_url = original(url) + if service_pattern is None: + return request_url + match = service_pattern.search(request_url.path) + if match is None: logger.debug( "Routing URL to original URL", - extra={"service": "unknown", "url": url, "host": request_url.host, "port": request_url.port}, + extra={"service": "unknown", "path": request_url.path, "host": request_url.host, "port": request_url.port}, ) - # Default: route to original URL if no service pattern is found return request_url - return route_url + api_name = match.group(1) + service_url = httpx.URL(platform_config.get_service_url(api_name)) + routed_url = request_url.copy_with( + scheme=service_url.scheme, + host=service_url.host, + port=service_url.port, + ) + logger.debug( + "Routing URL to service URL", + extra={ + "service": api_name, + "path": request_url.path, + "host": routed_url.host, + "port": routed_url.port, + }, + ) + return routed_url + + +@dataclass(frozen=True) +class PlatformRequestRouter: + """Routes SDK requests to the platform gateway or a service-specific origin.""" + + platform_config: PlatformConfig + default_resolver: Callable[[str], httpx.URL] + + def resolve(self, url: str) -> httpx.URL: + return resolve_platform_request_url( + url, + platform_config=self.platform_config, + default_resolver=self.default_resolver, + ) + + +def attach_platform_request_router(sdk: PlatformSDKT) -> PlatformSDKT: + """Attach the platform request router to a generated SDK instance. + + Stainless sends every request through ``_prepare_url``. Assigning the hook is + the SDK integration point; the routing policy itself lives in + :class:`PlatformRequestRouter`. + """ + router = PlatformRequestRouter( + platform_config=Configuration.get_platform_config(), + default_resolver=sdk._prepare_url, + ) + setattr(sdk, "_nmp_request_router", router) + sdk._prepare_url = router.resolve + return sdk + + +def with_options_preserving_request_router(base_sdk: PlatformSDKT, **kwargs: Any) -> PlatformSDKT: + """Return ``base_sdk.with_options(...)`` while preserving platform request routing.""" + scoped_sdk = cast(PlatformSDKT, base_sdk.with_options(**kwargs)) + router = getattr(base_sdk, "_nmp_request_router", None) + if isinstance(router, PlatformRequestRouter): + setattr(scoped_sdk, "_nmp_request_router", router) + scoped_sdk._prepare_url = router.resolve + return scoped_sdk def _get_default_headers( @@ -159,8 +209,7 @@ def get_platform_sdk( http_client=http_client or shared_sync_http_client(), default_headers=headers if headers else None, ) - sdk._prepare_url = _create_url_router(sdk._prepare_url) - return sdk + return attach_platform_request_router(sdk) def get_task_sdk(as_service: str, http_client: httpx.Client | None = None) -> NeMoPlatform: @@ -252,8 +301,7 @@ def get_async_platform_sdk( http_client=effective_client, default_headers=headers if headers else None, ) - sdk._prepare_url = _create_url_router(sdk._prepare_url) - return sdk + return attach_platform_request_router(sdk) def get_request_scoped_sdk( @@ -284,7 +332,7 @@ def get_request_scoped_sdk( # If we have headers to add, create a new SDK with them # This reuses the underlying HTTP client (lightweight operation) if headers: - return base_sdk.with_options(set_default_headers=headers) + return with_options_preserving_request_router(base_sdk, set_default_headers=headers) return base_sdk @@ -340,7 +388,7 @@ def get_sdk_on_behalf_of( merged_headers = {**headers, "X-NMP-Principal-On-Behalf-Of": on_behalf_of} merged_headers.pop("X-NMP-Principal-On-Behalf-Of-Groups", None) merged_headers.pop("X-NMP-Principal-On-Behalf-Of-Email", None) - return base_sdk.with_options(set_default_headers=merged_headers) + return with_options_preserving_request_router(base_sdk, set_default_headers=merged_headers) def get_entity_parts(name: str, default_workspace: str | None = None) -> tuple[str, str]: diff --git a/packages/nmp_common/src/nmp/common/service/base.py b/packages/nmp_common/src/nmp/common/service/base.py index 6560e8d18d..f1d59d573c 100644 --- a/packages/nmp_common/src/nmp/common/service/base.py +++ b/packages/nmp_common/src/nmp/common/service/base.py @@ -145,12 +145,13 @@ def _get_entity_sdk_on_behalf_of(self) -> AsyncNeMoPlatform: Uses the cached base SDK and applies per-request headers via .with_options() (lightweight — reuses the HTTP connection pool). """ + from nmp.common.sdk_factory import with_options_preserving_request_router from nmp.common.service.headers import build_downstream_service_headers base_sdk = self.get_sdk_client() headers = build_downstream_service_headers(self._service_name) - return base_sdk.with_options(set_default_headers=headers) + return with_options_preserving_request_router(base_sdk, set_default_headers=headers) def get_platform_config(self) -> PlatformConfig: """Return the PlatformConfig (lazily initialized).""" diff --git a/packages/nmp_common/tests/nmp_common/test_common_config.py b/packages/nmp_common/tests/nmp_common/test_common_config.py index c2752f84d2..a0bcac90c5 100644 --- a/packages/nmp_common/tests/nmp_common/test_common_config.py +++ b/packages/nmp_common/tests/nmp_common/test_common_config.py @@ -102,6 +102,19 @@ def test_get_service_url_local_uses_common_service_config_host_port(self): assert url == common.get_host_url() assert config.get_service_url("entities") == "http://localhost:8080" # not local, base_url + def test_get_service_url_local_overrides_gateway_base_url(self): + """Local services bypass the gateway even when base_url points at the gateway.""" + config = PlatformConfig( + base_url="https://nemo-gateway:8080", + services="auth,jobs", + service_discovery={}, + ) + common = get_common_service_config() + + assert config.get_service_url("auth") == common.get_host_url() + assert config.get_service_url("jobs") == common.get_host_url() + assert config.get_service_url("files") == "https://nemo-gateway:8080" + def test_get_service_url_local_overrides_service_discovery(self): """When a service is both local and in service_discovery, prefer local URL (CommonServiceConfig).""" config = PlatformConfig( diff --git a/packages/nmp_common/tests/sdk_factory/test_sdk.py b/packages/nmp_common/tests/sdk_factory/test_sdk.py index 65742fae58..efb24ba120 100644 --- a/packages/nmp_common/tests/sdk_factory/test_sdk.py +++ b/packages/nmp_common/tests/sdk_factory/test_sdk.py @@ -2,26 +2,30 @@ # SPDX-License-Identifier: Apache-2.0 import json +import logging from unittest.mock import patch import httpx import pytest from nmp.common.config import Configuration, PlatformConfig from nmp.common.sdk_factory import ( + PlatformRequestRouter, get_async_platform_sdk, get_entity_parts, get_platform_sdk, get_request_scoped_sdk, + get_sdk_on_behalf_of, get_task_sdk, + resolve_platform_request_url, ) @pytest.fixture(autouse=True) def _clear_sdk_factory_test_client(): - """Clear _test_http_client before each test so config-based SDK behavior is asserted. + """Clear SDK factory state before each test so config-based SDK behavior is asserted. When _test_http_client is set (e.g. by another test's create_test_client), the SDK - is created with base_url='http://testserver' and no URL router, which breaks tests + is created with base_url='http://testserver' and no request router, which breaks tests that assert on base_url or service routing. Clearing it keeps tests order-independent and ensures sdk_factory tests always exercise the config path. """ @@ -29,10 +33,12 @@ def _clear_sdk_factory_test_client(): old = sdk_factory_module._test_http_client sdk_factory_module._test_http_client = None + Configuration.clear_cache() try: yield finally: sdk_factory_module._test_http_client = old + Configuration.clear_cache() def test_get_platform_sdk(): @@ -46,6 +52,73 @@ def test_get_platform_sdk(): assert sdk.base_url == Configuration.get_platform_config().base_url +def test_get_platform_sdk_keeps_platform_base_url_for_local_services(monkeypatch: pytest.MonkeyPatch): + """The SDK base URL remains the platform entrypoint; per-service routing handles local APIs.""" + monkeypatch.setenv("NMP_BASE_URL", "https://nemo-gateway:8080") + monkeypatch.setenv("NMP_SERVICES", "auth") + monkeypatch.setenv("NMP_SERVICE_HOST", "127.0.0.1") + monkeypatch.setenv("NMP_SERVICE_PORT", "8080") + Configuration.clear_cache() + + sdk = get_platform_sdk() + + assert str(sdk.base_url).rstrip("/") == "https://nemo-gateway:8080" + + +def test_get_platform_sdk_preserves_api_base_url_for_controller_only_pods(monkeypatch: pytest.MonkeyPatch): + """Controller-only pods must call the API service, not their own health listener.""" + captured_requests: list[httpx.Request] = [] + + def capture_request(request: httpx.Request) -> httpx.Response: + captured_requests.append(request) + return httpx.Response( + 200, + json={ + "data": [], + "pagination": { + "current_page_size": 0, + "page": 1, + "page_size": 0, + "total_pages": 1, + "total_results": 0, + }, + }, + ) + + monkeypatch.setenv("NMP_BASE_URL", "http://nemo-platform-api:8080") + monkeypatch.setenv("NMP_CONTROLLERS", "jobs") + monkeypatch.delenv("NMP_SERVICES", raising=False) + monkeypatch.setenv("NMP_SERVICE_HOST", "127.0.0.1") + monkeypatch.setenv("NMP_SERVICE_PORT", "8080") + Configuration.clear_cache() + + with httpx.Client(transport=httpx.MockTransport(capture_request)) as http_client: + sdk = get_platform_sdk(http_client=http_client) + + assert str(sdk.base_url).rstrip("/") == "http://nemo-platform-api:8080" + sdk.jobs.list(workspace="default") + + assert len(captured_requests) == 1 + assert str(captured_requests[0].url) == "http://nemo-platform-api:8080/apis/jobs/v2/workspaces/default/jobs" + + +def test_get_platform_sdk_routes_local_service_path_to_process_listener(monkeypatch: pytest.MonkeyPatch): + """Requests for APIs hosted in this process bypass the platform entrypoint.""" + monkeypatch.setenv("NMP_BASE_URL", "https://nemo-gateway:8080") + monkeypatch.setenv("NMP_SERVICES", "auth") + monkeypatch.setenv("NMP_SERVICE_HOST", "127.0.0.1") + monkeypatch.setenv("NMP_SERVICE_PORT", "8080") + Configuration.clear_cache() + + sdk = get_platform_sdk() + prepared = sdk._prepare_url("https://nemo-gateway:8080/apis/auth/v2/authz/allow") + + assert prepared.scheme == "http" + assert prepared.host == "127.0.0.1" + assert prepared.port == 8080 + assert prepared.path == "/apis/auth/v2/authz/allow" + + def test_get_platform_sdk_with_service_principal(): """Test get_platform_sdk with as_service parameter.""" sdk = get_platform_sdk(as_service="my-service") @@ -194,6 +267,56 @@ def test_get_request_scoped_sdk_merges_otel_and_auth_headers(): assert scoped_sdk.default_headers["X-NMP-Principal-Groups"] == "group1,group2" +def test_get_request_scoped_sdk_preserves_request_router(monkeypatch: pytest.MonkeyPatch): + """Derived request SDKs must keep the base SDK's path-aware platform request router.""" + monkeypatch.setenv("NMP_BASE_URL", "https://nemo-gateway:8080") + monkeypatch.setenv("NMP_SERVICES", "entities") + monkeypatch.setenv("NMP_SERVICE_HOST", "127.0.0.1") + monkeypatch.setenv("NMP_SERVICE_PORT", "8080") + Configuration.clear_cache() + + try: + base_sdk = get_async_platform_sdk() + + with patch("nmp.common.sdk_factory.get_otel_headers", return_value={}): + with patch( + "nmp.common.sdk_factory.get_principal_auth_headers", + return_value={"X-NMP-Principal-Id": "service:models"}, + ): + scoped_sdk = get_request_scoped_sdk(base_sdk) + + prepared = scoped_sdk._prepare_url("https://nemo-gateway:8080/apis/entities/v2/workspaces") + + assert prepared.scheme == "http" + assert prepared.host == "127.0.0.1" + assert prepared.port == 8080 + assert prepared.path == "/apis/entities/v2/workspaces" + finally: + Configuration.clear_cache() + + +def test_get_sdk_on_behalf_of_preserves_request_router(monkeypatch: pytest.MonkeyPatch): + """SDKs derived with on-behalf-of headers must still keep platform request routing.""" + monkeypatch.setenv("NMP_BASE_URL", "https://nemo-gateway:8080") + monkeypatch.setenv("NMP_SERVICES", "entities") + monkeypatch.setenv("NMP_SERVICE_HOST", "127.0.0.1") + monkeypatch.setenv("NMP_SERVICE_PORT", "8080") + Configuration.clear_cache() + + try: + base_sdk = get_async_platform_sdk(as_service="models", internal=True) + scoped_sdk = get_sdk_on_behalf_of(base_sdk, "user@example.com") + + prepared = scoped_sdk._prepare_url("https://nemo-gateway:8080/apis/entities/v2/workspaces") + + assert prepared.scheme == "http" + assert prepared.host == "127.0.0.1" + assert prepared.port == 8080 + assert prepared.path == "/apis/entities/v2/workspaces" + finally: + Configuration.clear_cache() + + def test_get_request_scoped_sdk_returns_base_sdk_when_no_headers(): """Test that get_request_scoped_sdk returns base SDK when no headers to add.""" base_sdk = get_async_platform_sdk() @@ -373,6 +496,82 @@ def platform_config_with_service_discovery(): ) +def test_resolve_platform_request_url_routes_api_path_to_service_url(platform_config_with_service_discovery): + """The named request router policy owns per-service routing.""" + + def default_resolver(url: str) -> httpx.URL: + if url.startswith("/"): + return httpx.URL(f"http://platform:8080{url}") + return httpx.URL(url) + + prepared = resolve_platform_request_url( + "/apis/entities/v2/workspaces?limit=10", + platform_config=platform_config_with_service_discovery, + default_resolver=default_resolver, + ) + + assert prepared.scheme == "http" + assert prepared.host == "entities-service" + assert prepared.port == 8080 + assert prepared.path == "/apis/entities/v2/workspaces" + assert prepared.query == b"limit=10" + + +def test_resolve_platform_request_url_logs_path_without_raw_url( + caplog: pytest.LogCaptureFixture, + platform_config_with_service_discovery, +): + """Routing logs expose the resolved path without query parameters.""" + + def default_resolver(url: str) -> httpx.URL: + if url.startswith("/"): + return httpx.URL(f"http://platform:8080{url}") + return httpx.URL(url) + + caplog.set_level(logging.DEBUG, logger="nmp.common.sdk_factory") + + resolve_platform_request_url( + "/health/ready?token=secret", + platform_config=platform_config_with_service_discovery, + default_resolver=default_resolver, + ) + resolve_platform_request_url( + "/apis/entities/v2/workspaces?token=secret", + platform_config=platform_config_with_service_discovery, + default_resolver=default_resolver, + ) + + original_record = next(record for record in caplog.records if record.message == "Routing URL to original URL") + service_record = next(record for record in caplog.records if record.message == "Routing URL to service URL") + + assert not hasattr(original_record, "url") + assert original_record.service == "unknown" + assert original_record.path == "/health/ready" + assert original_record.host == "platform" + assert original_record.port == 8080 + + assert not hasattr(service_record, "url") + assert service_record.service == "entities" + assert service_record.path == "/apis/entities/v2/workspaces" + assert service_record.host == "entities-service" + assert service_record.port == 8080 + + for record in (original_record, service_record): + assert "token=secret" not in str(record.__dict__) + + +def test_platform_request_router_uses_default_resolver_for_non_api_paths(platform_config_with_service_discovery): + """Non-API paths follow the SDK's normal URL preparation.""" + router = PlatformRequestRouter( + platform_config=platform_config_with_service_discovery, + default_resolver=lambda url: httpx.URL(f"http://platform:8080{url}"), + ) + + prepared = router.resolve("/health/ready") + + assert str(prepared) == "http://platform:8080/health/ready" + + def test_get_platform_sdk_routes_entities_path_to_entities_service( platform_config_with_service_discovery, ): @@ -382,7 +581,6 @@ def test_get_platform_sdk_routes_entities_path_to_entities_service( return_value=platform_config_with_service_discovery, ): sdk = get_platform_sdk() - # Router calls get_platform_config when _prepare_url runs; keep patch active request_url = "http://platform:8080/apis/entities/v2/workspaces" prepared = sdk._prepare_url(request_url) diff --git a/packages/nmp_platform_runner/src/nmp/platform_runner/config.py b/packages/nmp_platform_runner/src/nmp/platform_runner/config.py index 6f7932252f..977f9bba45 100644 --- a/packages/nmp_platform_runner/src/nmp/platform_runner/config.py +++ b/packages/nmp_platform_runner/src/nmp/platform_runner/config.py @@ -6,7 +6,7 @@ from __future__ import annotations import os -from collections.abc import Callable, MutableMapping +from collections.abc import MutableMapping from dataclasses import dataclass, field from importlib.resources import files from urllib.parse import urlparse @@ -19,6 +19,7 @@ Configuration, ) from nmp.common.service import Service +from nmp.platform_runner.loader import ControllerRunFunc from nmp.platform_runner.registry import ( AVAILABLE_SIDECARS, get_available_controllers, @@ -43,7 +44,7 @@ class ResolvedRunConfiguration: port: int config_path: str available_services: dict[str, str | Service] = field(default_factory=dict) - available_controllers: dict[str, str | Callable] = field(default_factory=dict) + available_controllers: dict[str, str | ControllerRunFunc] = field(default_factory=dict) def default_config_path() -> str: @@ -189,8 +190,9 @@ def apply_run_environment( host_for_url = _bracket_ipv6(effective_host) default_base_url = f"http://{host_for_url}:{effective_port}" base_url = env.setdefault("NMP_BASE_URL", default_base_url) - # Embedded PDP is served from the same platform process; keep the auth client - # origin aligned with NMP_BASE_URL when services run on a non-default port. + # Embedded PDP is usually served from the same platform process, so its + # self-call origin must stay aligned with the resolved base URL. Deployed + # mode can still override this explicitly by pre-setting the env var. env.setdefault("NMP_AUTH_POLICY_DECISION_POINT_BASE_URL", base_url) _set_or_clear_env(env, NMP_SERVICES_ENV_VAR, config.services) _set_or_clear_env(env, NMP_CONTROLLERS_ENV_VAR, config.controllers) diff --git a/packages/nmp_platform_runner/src/nmp/platform_runner/health.py b/packages/nmp_platform_runner/src/nmp/platform_runner/health.py index 2037ba584c..6fe30ca951 100644 --- a/packages/nmp_platform_runner/src/nmp/platform_runner/health.py +++ b/packages/nmp_platform_runner/src/nmp/platform_runner/health.py @@ -6,6 +6,8 @@ from __future__ import annotations import logging +from collections.abc import Awaitable, Callable +from dataclasses import dataclass from fastapi import APIRouter, HTTPException from nmp.common.controller import ControllerManager @@ -29,6 +31,15 @@ NMP_PLATFORM_REVISION_ATTR = "nmp.platform.revision" +@dataclass(frozen=True) +class ReadinessCheck: + """Additional platform-level readiness check.""" + + name: str + is_ready: Callable[[], Awaitable[bool]] + message: Callable[[], str] | None = None + + def get_platform_resource_attributes() -> dict[str, str]: """Return OTEL resource attributes for the platform.""" return { @@ -49,8 +60,26 @@ async def _get_service_status_breakdown(services: list[Service]) -> tuple[list[s return ready, not_ready -def create_platform_health_router(services: list[Service]) -> APIRouter: +async def _get_readiness_check_status_breakdown( + readiness_checks: list[ReadinessCheck], +) -> tuple[list[str], list[dict[str, str]]]: + ready: list[str] = [] + not_ready: list[dict[str, str]] = [] + for check in readiness_checks: + if await check.is_ready(): + ready.append(check.name) + else: + message = check.message() if check.message is not None else "" + not_ready.append({"name": check.name, "message": message}) + return ready, not_ready + + +def create_platform_health_router( + services: list[Service], + readiness_checks: list[ReadinessCheck] | None = None, +) -> APIRouter: """Create the shared platform health router.""" + readiness_checks = readiness_checks or [] router = APIRouter(tags=["Health"]) @router.get("/cluster-info", operation_id="platform_cluster_info", response_model=ClusterInfo) @@ -60,9 +89,12 @@ async def cluster_info() -> ClusterInfo: @router.get("/status", operation_id="platform_status", response_model=PlatformStatusResponse) async def status() -> PlatformStatusResponse: ready, not_ready = await _get_service_status_breakdown(services) + ready_checks, not_ready_checks = await _get_readiness_check_status_breakdown(readiness_checks) + ready.extend(ready_checks) + not_ready.extend(not_ready_checks) ready_count = len(ready) not_ready_count = len(not_ready) - total = len(services) + total = len(services) + len(readiness_checks) if ready_count == total: status_value = "healthy" @@ -92,9 +124,10 @@ async def health_live() -> HealthLiveResponse: @router.get("/health/ready", operation_id="platform_health_ready", response_model=HealthReadyResponse) async def health_ready() -> HealthReadyResponse: services_ready = all([await service.is_ready() for service in services]) + readiness_checks_ready = all([await check.is_ready() for check in readiness_checks]) manager = ControllerManager.get_instance() all_controllers_healthy, _ = manager.validate_all_healthy(detailed=False) - if not services_ready or not all_controllers_healthy: + if not services_ready or not readiness_checks_ready or not all_controllers_healthy: raise HTTPException(status_code=503, detail=HealthNotReadyDetail().model_dump()) return HealthReadyResponse() diff --git a/packages/nmp_platform_runner/src/nmp/platform_runner/loader.py b/packages/nmp_platform_runner/src/nmp/platform_runner/loader.py index 0ffc3cc939..ed6ea42c37 100644 --- a/packages/nmp_platform_runner/src/nmp/platform_runner/loader.py +++ b/packages/nmp_platform_runner/src/nmp/platform_runner/loader.py @@ -7,13 +7,17 @@ import importlib import logging -from typing import Callable +import threading +from collections.abc import Callable +from typing import cast from nmp.common.service import Service from nmp.common.service.deptree import resolve_service_loading_order logger = logging.getLogger(__name__) +ControllerRunFunc = Callable[[threading.Event], object] + def load_service( service_name: str, @@ -49,7 +53,7 @@ def order_services_by_dependencies(services: list[Service]) -> list[Service]: return [services_by_name[name] for name in ordered_names if name in services_by_name] -def load_controller_run_func(controller_name: str, import_path: str) -> Callable: +def load_controller_run_func(controller_name: str, import_path: str) -> ControllerRunFunc: """Load a controller run function from an import path.""" if ":" not in import_path: raise ValueError(f"Import path must be in format 'module:function', got: {import_path}") @@ -62,4 +66,4 @@ def load_controller_run_func(controller_name: str, import_path: str) -> Callable raise TypeError(f"Controller {controller_name} must be a callable, got {type(run_func)}") logger.debug("Loaded controller %s", controller_name) - return run_func + return cast(ControllerRunFunc, run_func) diff --git a/packages/nmp_platform_runner/src/nmp/platform_runner/registry.py b/packages/nmp_platform_runner/src/nmp/platform_runner/registry.py index d9e555b8e4..bb4b0f4100 100644 --- a/packages/nmp_platform_runner/src/nmp/platform_runner/registry.py +++ b/packages/nmp_platform_runner/src/nmp/platform_runner/registry.py @@ -6,11 +6,11 @@ from __future__ import annotations import logging -from collections.abc import Callable from functools import cache from nemo_platform_plugin.discovery import discover_controllers, discover_services from nmp.common.service import Service +from nmp.platform_runner.loader import ControllerRunFunc from nmp.platform_runner.plugin_adapter import NemoServiceAdapter, make_controller_run_func logger = logging.getLogger(__name__) @@ -75,7 +75,7 @@ @cache -def get_available_controllers() -> dict[str, str | Callable]: +def get_available_controllers() -> dict[str, str | ControllerRunFunc]: """Return all available controller run functions for the current run. Merges built-in core controllers (stored as ``"module:object"`` import @@ -86,7 +86,7 @@ def get_available_controllers() -> dict[str, str | Callable]: Returns: Mapping of controller name → string import path or run callable. """ - controllers: dict[str, str | Callable] = dict(AVAILABLE_CONTROLLERS) + controllers: dict[str, str | ControllerRunFunc] = dict(AVAILABLE_CONTROLLERS) for name, controller_cls in discover_controllers().items(): try: @@ -105,7 +105,7 @@ def get_available_controllers() -> dict[str, str | Callable]: def get_controller_groups( - available_controllers: dict[str, str | Callable] | None = None, + available_controllers: dict[str, str | ControllerRunFunc] | None = None, ) -> dict[str, list[str]]: """Return dynamic controller groups for the current run.""" available_controllers = available_controllers or get_available_controllers() diff --git a/packages/nmp_platform_runner/src/nmp/platform_runner/run.py b/packages/nmp_platform_runner/src/nmp/platform_runner/run.py index 554127fa11..0dbad175ed 100644 --- a/packages/nmp_platform_runner/src/nmp/platform_runner/run.py +++ b/packages/nmp_platform_runner/src/nmp/platform_runner/run.py @@ -11,14 +11,20 @@ import threading import time from collections.abc import Callable, Mapping +from typing import cast -from nmp.common.config import get_auth_config, get_common_service_config, get_service_config +from nmp.common.config import get_auth_config, get_common_service_config, get_platform_config, get_service_config from nmp.common.observability import initialize_obs, setup_global_instrumentations from nmp.common.observability.otel import settings as otel_settings from nmp.common.service import CircularDependencyError, Service from nmp.platform_runner.config import apply_run_environment, resolve_run_configuration from nmp.platform_runner.health import get_platform_resource_attributes -from nmp.platform_runner.loader import load_controller_run_func, load_service, order_services_by_dependencies +from nmp.platform_runner.loader import ( + ControllerRunFunc, + load_controller_run_func, + load_service, + order_services_by_dependencies, +) from nmp.platform_runner.registry import AVAILABLE_SIDECARS from nmp.platform_runner.server import run_server, run_server_with_reload from nmp.platform_runner.version import get_platform_version @@ -53,7 +59,7 @@ def _database_display(db_url: str) -> str: def run_controllers_in_threads( - controller_run_funcs: dict[str, Callable], + controller_run_funcs: dict[str, ControllerRunFunc], stop_signal: threading.Event, ) -> list[threading.Thread]: """Start controller run functions in daemon threads.""" @@ -111,6 +117,7 @@ def run_platform( raise ValueError(f"Controller/sidecar name collision: {', '.join(sorted(collisions))}") service_instances = _load_service_instances(sorted(resolved.services), resolved.available_services) + get_platform_config().services = ",".join(sorted(service.name for service in service_instances)) controller_run_funcs = _load_run_functions( sorted(resolved.controllers), resolved.available_controllers, "controller" ) @@ -212,16 +219,16 @@ def _load_service_instances( def _load_run_functions( names: list[str], - registry: Mapping[str, str | Callable], + registry: Mapping[str, str | ControllerRunFunc], kind: str, -) -> dict[str, Callable]: - run_funcs: dict[str, Callable] = {} +) -> dict[str, ControllerRunFunc]: + run_funcs: dict[str, ControllerRunFunc] = {} for name in names: t0 = time.perf_counter() value = registry[name] try: if callable(value): - run_funcs[name] = value + run_funcs[name] = cast(ControllerRunFunc, value) else: run_funcs[name] = load_controller_run_func(name, value) except (ImportError, TypeError, AttributeError, ValueError) as error: diff --git a/packages/nmp_platform_runner/src/nmp/platform_runner/server.py b/packages/nmp_platform_runner/src/nmp/platform_runner/server.py index a236148b24..daac4f6182 100644 --- a/packages/nmp_platform_runner/src/nmp/platform_runner/server.py +++ b/packages/nmp_platform_runner/src/nmp/platform_runner/server.py @@ -11,6 +11,7 @@ import os import threading from contextlib import asynccontextmanager +from typing import cast import httpx import uvicorn @@ -24,8 +25,13 @@ from nmp.common.observability.context import create_app_context_dependency from nmp.common.pyleak import detect_blocking from nmp.common.service import Service -from nmp.platform_runner.health import create_platform_health_router, get_platform_resource_attributes -from nmp.platform_runner.loader import load_controller_run_func, load_service, order_services_by_dependencies +from nmp.platform_runner.health import ReadinessCheck, create_platform_health_router, get_platform_resource_attributes +from nmp.platform_runner.loader import ( + ControllerRunFunc, + load_controller_run_func, + load_service, + order_services_by_dependencies, +) from nmp.platform_runner.registry import get_available_controllers, get_available_services, get_openapi_service_names from starlette.middleware.base import BaseHTTPMiddleware from starlette.requests import Request @@ -34,6 +40,27 @@ logger = logging.getLogger(__name__) +class _StartupReadinessState: + def __init__(self, name: str) -> None: + self.name = name + self._ready = False + self._message = "pending" + + async def is_ready(self) -> bool: + return self._ready + + def message(self) -> str: + return self._message + + def mark_ready(self) -> None: + self._ready = True + self._message = "ready" + + def mark_failed(self, message: str) -> None: + self._ready = False + self._message = message + + async def platform_global_exception_handler(request: Request, exc: Exception) -> JSONResponse: """Fallback exception handler for uncaught platform errors.""" extra = { @@ -103,18 +130,32 @@ def create_platform_openapi_app() -> FastAPI: def create_app( services: list[Service] | None = None, - controller_run_funcs: dict[str, object] | None = None, + controller_run_funcs: dict[str, ControllerRunFunc] | None = None, http_client: httpx.AsyncClient | None = None, ) -> FastAPI: """Create the FastAPI app from service instances.""" services = services or [] controller_run_funcs = controller_run_funcs or {} controller_stop_signal = threading.Event() + platform_config = get_platform_config() + platform_config.services = ",".join(sorted(service.name for service in services)) + readiness_checks: list[ReadinessCheck] = [] + platform_seed_state: _StartupReadinessState | None = None + if platform_config.seed_on_startup: + platform_seed_state = _StartupReadinessState("platform-seed") + readiness_checks.append( + ReadinessCheck( + name=platform_seed_state.name, + is_ready=platform_seed_state.is_ready, + message=platform_seed_state.message, + ) + ) @asynccontextmanager async def lifespan(app: FastAPI): logger.info("Starting Nemo Platform server") controller_threads = [] + platform_seed_task: asyncio.Task[None] | None = None if controller_run_funcs: logger.info("Starting controllers in lifespan: %s", list(controller_run_funcs)) for name, run_func in controller_run_funcs.items(): @@ -127,21 +168,42 @@ async def lifespan(app: FastAPI): thread.start() controller_threads.append(thread) - platform_config = get_platform_config() - if platform_config.seed_on_startup: + if platform_seed_state is not None: try: from nmp.platform_seed import run_platform_seed_from_startup - asyncio.create_task(run_platform_seed_from_startup()) + async def run_platform_seed_and_update_readiness() -> None: + try: + ok = await run_platform_seed_from_startup() + except Exception: + logger.exception("Platform seed task failed") + platform_seed_state.mark_failed("platform seed failed") + return + + if ok: + platform_seed_state.mark_ready() + else: + platform_seed_state.mark_failed("platform seed failed") + + platform_seed_task = asyncio.create_task(run_platform_seed_and_update_readiness()) logger.info("Platform seed task scheduled") except ImportError as error: logger.warning("platform.seed_on_startup is True but platform_seed is not installed: %s", error) + platform_seed_state.mark_failed("platform seed is not installed") app.state.controller_threads = controller_threads app.state.controller_stop_signal = controller_stop_signal + app.state.platform_seed_task = platform_seed_task yield + if platform_seed_task is not None and not platform_seed_task.done(): + platform_seed_task.cancel() + try: + await platform_seed_task + except asyncio.CancelledError: + pass + controller_stop_signal.set() for thread in controller_threads: thread.join(timeout=5) @@ -174,9 +236,9 @@ async def lifespan(app: FastAPI): ) app.state.service_configs = {} - app.include_router(create_platform_health_router(services)) + app.include_router(create_platform_health_router(services, readiness_checks=readiness_checks)) - redirect_root_to_studio = get_platform_config().redirect_root_to_studio + redirect_root_to_studio = platform_config.redirect_root_to_studio @app.api_route("/", methods=["GET", "HEAD"], include_in_schema=False, response_model=None) async def root_handler() -> Response: @@ -283,7 +345,7 @@ def create_default_app() -> FastAPI: services.append(load_service(service_name, service_value)) services = order_services_by_dependencies(services) - controller_run_funcs = {} + controller_run_funcs: dict[str, ControllerRunFunc] = {} for controller_name in controller_names: controller_value = available_controllers.get(controller_name) if controller_value is None: @@ -293,7 +355,7 @@ def create_default_app() -> FastAPI: % (controller_name, controller_names_env, available) ) if callable(controller_value): - controller_run_funcs[controller_name] = controller_value + controller_run_funcs[controller_name] = cast(ControllerRunFunc, controller_value) else: controller_run_funcs[controller_name] = load_controller_run_func(controller_name, controller_value) diff --git a/packages/nmp_platform_runner/tests/test_config.py b/packages/nmp_platform_runner/tests/test_config.py index 7853a23e1b..2318957569 100644 --- a/packages/nmp_platform_runner/tests/test_config.py +++ b/packages/nmp_platform_runner/tests/test_config.py @@ -132,11 +132,41 @@ def test_sets_base_url_when_not_present(self): apply_run_environment(_make_config(host="0.0.0.0", port=8080), env=env) assert env["NMP_BASE_URL"] == "http://127.0.0.1:8080" + def test_config_file_gateway_base_url_seeds_base_url(self, tmp_path: Path): + config_path = tmp_path / "config.yaml" + config_path.write_text( + """ +platform: + base_url: "https://nemo-gateway:8080" +""", + encoding="utf-8", + ) + env: dict[str, str] = {} + apply_run_environment(_make_config(host="0.0.0.0", port=8080, config_path=str(config_path)), env=env) + assert env["NMP_BASE_URL"] == "https://nemo-gateway:8080" + assert env["NMP_SERVICE_HOST"] == "127.0.0.1" + def test_sets_embedded_pdp_base_url_from_base_url(self): env: dict[str, str] = {} apply_run_environment(_make_config(host="0.0.0.0", port=9090), env=env) assert env["NMP_AUTH_POLICY_DECISION_POINT_BASE_URL"] == "http://127.0.0.1:9090" + def test_embedded_pdp_base_url_uses_resolved_base_url_when_auth_config_is_static(self, tmp_path: Path): + config_path = tmp_path / "config.yaml" + config_path.write_text( + """ +platform: + base_url: "https://nemo-gateway:8080" +auth: + policy_decision_point_base_url: "http://127.0.0.1:8080" +""", + encoding="utf-8", + ) + env: dict[str, str] = {} + apply_run_environment(_make_config(host="0.0.0.0", port=59007, config_path=str(config_path)), env=env) + assert env["NMP_BASE_URL"] == "https://nemo-gateway:59007" + assert env["NMP_AUTH_POLICY_DECISION_POINT_BASE_URL"] == "https://nemo-gateway:59007" + def test_sets_service_host_when_not_present(self): env: dict[str, str] = {} apply_run_environment(_make_config(host="0.0.0.0", port=8080), env=env) @@ -247,6 +277,21 @@ def test_uses_actual_bind_port_not_config_port(self, tmp_path: Path): assert env["NMP_BASE_URL"] == "http://172.17.0.1:59007" assert env["NMP_AUTH_POLICY_DECISION_POINT_BASE_URL"] == "http://172.17.0.1:59007" + def test_local_config_static_pdp_url_does_not_override_actual_bind_port(self, tmp_path: Path): + config_path = self._write_config( + tmp_path, + """ +platform: + base_url: http://0.0.0.0:8080 +auth: + policy_decision_point_base_url: http://localhost:8080 +""", + ) + env: dict[str, str] = {} + apply_run_environment(_make_config(host="0.0.0.0", port=59007, config_path=config_path), env=env) + assert env["NMP_BASE_URL"] == "http://127.0.0.1:59007" + assert env["NMP_AUTH_POLICY_DECISION_POINT_BASE_URL"] == "http://127.0.0.1:59007" + def test_config_base_url_without_port_gets_bind_port(self, tmp_path: Path): config_path = self._write_config(tmp_path, "platform:\n base_url: http://172.17.0.1\n") env: dict[str, str] = {} diff --git a/packages/nmp_platform_runner/tests/test_run.py b/packages/nmp_platform_runner/tests/test_run.py index 1fa449bf07..f5cb9d4617 100644 --- a/packages/nmp_platform_runner/tests/test_run.py +++ b/packages/nmp_platform_runner/tests/test_run.py @@ -4,11 +4,20 @@ """Unit tests for platform runner startup display helpers.""" import logging +import threading import pytest +from nmp.common.config import Configuration, get_platform_config +from nmp.platform_runner import run as runner +from nmp.platform_runner.config import ResolvedRunConfiguration from nmp.platform_runner.run import _database_display +class _StubService: + def __init__(self, name: str) -> None: + self.name = name + + @pytest.mark.parametrize( ("db_url", "expected"), [ @@ -40,3 +49,48 @@ def test_database_display_logs_parse_failures(caplog: pytest.LogCaptureFixture) assert len(records) == 1 assert records[0].message == "Failed to parse database URL for startup banner" assert records[0].exc_info is not None + + +def test_run_platform_marks_loaded_services_local_before_starting_controllers(monkeypatch): + Configuration.clear_cache() + captured: dict[str, str] = {} + + resolved = ResolvedRunConfiguration( + services={"jobs", "entities"}, + controllers={"jobs"}, + sidecars=set(), + host="127.0.0.1", + port=8080, + config_path="", + ) + services = [_StubService("jobs"), _StubService("entities")] + + monkeypatch.setattr(runner, "resolve_run_configuration", lambda **_: resolved) + monkeypatch.setattr(runner, "apply_run_environment", lambda config: None) + monkeypatch.setattr(runner, "initialize_obs", lambda *, resource_attributes: None) + monkeypatch.setattr(runner, "setup_global_instrumentations", lambda: None) + monkeypatch.setattr(runner, "_load_service_instances", lambda service_names, available_services: services) + monkeypatch.setattr( + runner, + "_load_run_functions", + lambda names, registry, kind: {"jobs": lambda stop_signal: None} if kind == "controller" else {}, + ) + monkeypatch.setattr(runner, "_display_banner", lambda **_: None) + monkeypatch.setattr(runner, "run_server", lambda services, host, port: None) + monkeypatch.setattr(runner.signal, "signal", lambda *args: None) + + def capture_controller_start( + controller_run_funcs, + stop_signal: threading.Event, + ) -> list[threading.Thread]: + captured["services"] = get_platform_config().services + return [] + + monkeypatch.setattr(runner, "run_controllers_in_threads", capture_controller_start) + + try: + runner.run_platform() + finally: + Configuration.clear_cache() + + assert captured["services"] == "entities,jobs" diff --git a/packages/nmp_platform_runner/tests/test_server.py b/packages/nmp_platform_runner/tests/test_server.py index 901d54e964..f752d1c51a 100644 --- a/packages/nmp_platform_runner/tests/test_server.py +++ b/packages/nmp_platform_runner/tests/test_server.py @@ -1,15 +1,22 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import asyncio +import builtins +import sys +import threading +import time +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest from fastapi import FastAPI from fastapi.testclient import TestClient -from nmp.common.config import AuthConfig +from nmp.common.config import AuthConfig, Configuration from nmp.common.config.base import OIDCConfig from nmp.common.service import Service from nmp.platform_runner import server +from nmp.platform_runner.health import ReadinessCheck, create_platform_health_router def _make_auth_config(*, enabled: bool) -> AuthConfig: @@ -28,6 +35,67 @@ def get_routers(self): return [] +async def _ready() -> bool: + return True + + +async def _not_ready() -> bool: + return False + + +def _client_for_health_checks(checks: list[ReadinessCheck]) -> TestClient: + app = FastAPI() + app.include_router(create_platform_health_router([PluginService()], readiness_checks=checks)) + return TestClient(app) + + +def _patch_platform_app_config(monkeypatch, *, seed_on_startup: bool): + auth_cfg = _make_auth_config(enabled=False) + platform_cfg = _make_platform_config_mock() + platform_cfg.seed_on_startup = seed_on_startup + monkeypatch.setattr(server, "get_platform_config", lambda: platform_cfg) + monkeypatch.setattr(server, "get_auth_config", lambda: auth_cfg) + + import nmp.common.auth.middleware as auth_middleware + + monkeypatch.setattr(auth_middleware, "get_auth_config", lambda: auth_cfg) + return platform_cfg + + +def _wait_for_response(client: TestClient, path: str, status_code: int, timeout: float = 2.0): + deadline = time.monotonic() + timeout + last_response = None + while time.monotonic() < deadline: + last_response = client.get(path) + if last_response.status_code == status_code: + return last_response + time.sleep(0.05) + return last_response + + +def test_platform_health_ready_includes_startup_readiness_checks(): + client = _client_for_health_checks( + [ReadinessCheck(name="platform-seed", is_ready=_not_ready, message=lambda: "pending")] + ) + + response = client.get("/health/ready") + assert response.status_code == 503 + + status = client.get("/status").json() + assert "agents" in status["services"]["ready"] + assert {"name": "platform-seed", "message": "pending"} in status["services"]["not_ready"] + + +def test_platform_health_status_reports_ready_startup_checks(): + client = _client_for_health_checks([ReadinessCheck(name="platform-seed", is_ready=_ready)]) + + response = client.get("/health/ready") + assert response.status_code == 200 + + status = client.get("/status").json() + assert "platform-seed" in status["services"]["ready"] + + def test_create_platform_openapi_app_includes_explicit_service_instances(monkeypatch): plugin_service = PluginService() captured: dict[str, object] = {} @@ -72,6 +140,46 @@ def fake_create_app(services, controller_run_funcs=None, _http_client=None): assert captured["controller_run_funcs"] == {"agents-deployment": plugin_controller} +def test_create_app_marks_mounted_services_as_local(monkeypatch): + platform_cfg = _patch_platform_app_config(monkeypatch, seed_on_startup=False) + platform_cfg.services = "" + + server.create_app(services=[PluginService()]) + + assert platform_cfg.services == "agents" + + +def test_create_app_mounted_services_drive_sdk_local_routing_without_services_env(monkeypatch): + monkeypatch.delenv("NMP_SERVICES", raising=False) + monkeypatch.setenv("NMP_BASE_URL", "https://nemo-gateway:8080") + monkeypatch.setenv("NMP_SERVICE_HOST", "127.0.0.1") + monkeypatch.setenv("NMP_SERVICE_PORT", "8080") + Configuration.clear_cache() + platform_cfg = Configuration.get_platform_config() + + try: + auth_cfg = _make_auth_config(enabled=False) + monkeypatch.setattr(server, "get_platform_config", lambda: platform_cfg) + monkeypatch.setattr(server, "get_auth_config", lambda: auth_cfg) + + import nmp.common.auth.middleware as auth_middleware + from nmp.common.sdk_factory import get_platform_sdk + + monkeypatch.setattr(auth_middleware, "get_auth_config", lambda: auth_cfg) + + server.create_app(services=[PluginService()]) + + sdk = get_platform_sdk() + prepared = sdk._prepare_url("https://nemo-gateway:8080/apis/agents/v2/example") + + assert platform_cfg.services == "agents" + assert prepared.scheme == "http" + assert prepared.host == "127.0.0.1" + assert prepared.port == 8080 + finally: + Configuration.clear_cache() + + def test_embedded_auth_preflight_invokes_policy_wasm_helper(monkeypatch): calls: list[bool] = [] auth_cfg = AuthConfig( @@ -160,6 +268,81 @@ def _make_platform_config_mock(*, redirect_root_to_studio: bool = True) -> Magic return cfg +def test_create_app_without_seed_on_startup_keeps_health_ready_unchanged(monkeypatch): + _patch_platform_app_config(monkeypatch, seed_on_startup=False) + + with TestClient(server.create_app(services=[])) as client: + response = client.get("/health/ready") + status = client.get("/status").json() + + assert response.status_code == 200 + assert "platform-seed" not in status["services"]["ready"] + assert all(item["name"] != "platform-seed" for item in status["services"]["not_ready"]) + + +def test_create_app_with_seed_on_startup_blocks_readiness_until_seed_completes(monkeypatch): + _patch_platform_app_config(monkeypatch, seed_on_startup=True) + started = threading.Event() + release = threading.Event() + + async def fake_seed() -> bool: + started.set() + await asyncio.to_thread(release.wait) + return True + + monkeypatch.setitem(sys.modules, "nmp.platform_seed", SimpleNamespace(run_platform_seed_from_startup=fake_seed)) + + with TestClient(server.create_app(services=[])) as client: + assert started.wait(timeout=2) + response = client.get("/health/ready") + assert response.status_code == 503 + + status = client.get("/status").json() + assert {"name": "platform-seed", "message": "pending"} in status["services"]["not_ready"] + + release.set() + response = _wait_for_response(client, "/health/ready", 200) + + assert response is not None + assert response.status_code == 200 + + +def test_create_app_with_failed_seed_keeps_health_not_ready(monkeypatch): + _patch_platform_app_config(monkeypatch, seed_on_startup=True) + + async def fake_seed() -> bool: + return False + + monkeypatch.setitem(sys.modules, "nmp.platform_seed", SimpleNamespace(run_platform_seed_from_startup=fake_seed)) + + with TestClient(server.create_app(services=[])) as client: + response = _wait_for_response(client, "/health/ready", 503) + status = client.get("/status").json() + + assert response is not None + assert response.status_code == 503 + assert {"name": "platform-seed", "message": "platform seed failed"} in status["services"]["not_ready"] + + +def test_create_app_with_missing_seed_package_keeps_health_not_ready(monkeypatch): + _patch_platform_app_config(monkeypatch, seed_on_startup=True) + real_import = builtins.__import__ + + def fake_import(name, globals=None, locals=None, fromlist=(), level=0): + if name == "nmp.platform_seed": + raise ImportError("missing platform seed") + return real_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", fake_import) + + with TestClient(server.create_app(services=[])) as client: + response = client.get("/health/ready") + status = client.get("/status").json() + + assert response.status_code == 503 + assert {"name": "platform-seed", "message": "platform seed is not installed"} in status["services"]["not_ready"] + + @pytest.mark.parametrize("auth_enabled", [True, False]) @pytest.mark.parametrize("method", ["get", "head"]) def test_root_redirects_to_studio(auth_enabled, method): diff --git a/pytest.ini b/pytest.ini index a542b7dcba..f6ed98dc22 100644 --- a/pytest.ini +++ b/pytest.ini @@ -65,6 +65,9 @@ markers = smoke_nmp_automodel_training: Import smoke tests for the nmp-automodel-training image e2e: End-to-end tests - test complete customer workflows on deployed infrastructure (Helm/Docker Compose) auth_idp: Auth IdP e2e tests - provider-backed auth compose coverage through the gateway + auth_idp_runtime: Auth IdP tests that require a provider runtime; optional marker args list matching runtime ids + auth_idp_docker: Auth IdP tests for the Docker Compose reference deployment + auth_idp_k8s: Auth IdP tests for the Kubernetes Helm reference deployment e2e_config(*layers, harness=...): Ordered list of repo-root-relative config paths and/or inline dict overlays; harness config stays separate from platform config subprocess_only: Test only works in subprocess mode (not on Kubernetes); skipped when NMP_BASE_URL is set container_only: Test requires a container backend (Docker or Kubernetes); skipped unless NMP_BASE_URL is set diff --git a/sdk/python/nemo-platform/.nmpcontext/openapi.yaml b/sdk/python/nemo-platform/.nmpcontext/openapi.yaml index 3db8849574..b0c624f38e 100644 --- a/sdk/python/nemo-platform/.nmpcontext/openapi.yaml +++ b/sdk/python/nemo-platform/.nmpcontext/openapi.yaml @@ -20,8 +20,13 @@ paths: \ CLI)\n - `userinfo_endpoint`: UserInfo endpoint\n - `client_id`: OAuth\ \ client ID to use\n - `default_scopes`: OAuth scopes to request during authentication\n\ \ - `scope_prefix`: Prefix to prepend to custom scopes (those with ':' or\ - \ '.default')" - operationId: get_auth_discovery_apis_auth_discovery_get + \ '.default')\n - `workload_token_exchange_enabled`: Whether SDK workload\ + \ identity token exchange is enabled\n - `workload_client_id`: OAuth client\ + \ ID to use for workload identity token exchange\n - `workload_token_endpoint`:\ + \ Token endpoint to use only for workload identity token exchange\n - `workload_audience`:\ + \ RFC 8693 audience for exchanged workload tokens\n - `workload_scope`: OAuth\ + \ scopes for exchanged workload tokens" + operationId: get_auth_discovery_endpoint_apis_auth_discovery_get responses: '200': description: Successful Response @@ -29,6 +34,89 @@ paths: application/json: schema: $ref: '#/components/schemas/AuthDiscoveryResponse' + /apis/auth/jwks: + get: + tags: + - Workload Identity + summary: Workload identity token exchange JWKS + description: Return the public signing key for workload identity access tokens + minted by the NeMo auth service. + operationId: jwks_apis_auth_jwks_get + responses: + '200': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/JsonWebKeySetResponse' + /apis/auth/token: + post: + tags: + - Workload Identity + summary: Exchange a workload identity subject token + description: Exchange a configured workload identity subject token for a NeMo + Platform access token. + operationId: token_exchange_apis_auth_token_post + requestBody: + content: + application/x-www-form-urlencoded: + schema: + properties: + grant_type: + type: string + enum: + - urn:ietf:params:oauth:grant-type:token-exchange + description: OAuth 2.0 token exchange grant type. + client_id: + type: string + description: Workload token exchange OAuth client ID. + subject_token: + type: string + description: JWT subject token to exchange. + subject_token_type: + type: string + enum: + - urn:ietf:params:oauth:token-type:jwt + description: Token type identifier for the subject token. + requested_token_type: + type: string + enum: + - urn:ietf:params:oauth:token-type:access_token + description: Requested token type identifier for the issued token. + default: urn:ietf:params:oauth:token-type:access_token + audience: + type: string + description: Requested audience for the issued access token. + scope: + type: string + description: Space-separated scopes requested for the issued access + token. + type: object + required: + - grant_type + - client_id + - subject_token + - subject_token_type + required: true + responses: + '200': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/WorkloadTokenExchangeResponse' + '400': + description: RFC 8693 token exchange error + content: + application/json: + schema: + $ref: '#/components/schemas/WorkloadTokenExchangeErrorResponse' + '401': + description: OAuth 2.0 invalid_client error + content: + application/json: + schema: + $ref: '#/components/schemas/WorkloadTokenExchangeErrorResponse' /apis/auth/v2/iam/role-bindings: get: tags: @@ -9979,6 +10067,10 @@ components: allOf: - $ref: '#/components/schemas/DockerJobNetworkConfig' description: Docker networking configuration + workload_identity: + allOf: + - $ref: '#/components/schemas/DockerWorkloadIdentityConfig' + description: Docker workload identity subject-token issuer configuration. type: object title: DockerJobExecutionProfileConfig description: Configuration for Docker Job execution profile. @@ -10048,6 +10140,60 @@ components: - volume_name - mount_path title: DockerVolumeMount + DockerWorkloadIdentityConfig: + properties: + enabled: + title: Enabled + description: Enable Docker workload identity token-file injection. Defaults + to auth.oidc.workload_token_exchange_enabled. + type: boolean + token_endpoint: + title: Token Endpoint + description: OAuth token endpoint used by the Docker demo issuer. Defaults + to auth.oidc.token_endpoint. + type: string + client_id: + title: Client Id + description: OAuth client ID used by the Docker demo issuer. Defaults to + auth.oidc.workload_client_id or auth.oidc.client_id. + type: string + client_secret: + format: password + title: Client Secret + description: OAuth client secret for the Docker demo issuer. + writeOnly: true + type: string + username: + title: Username + description: Username for the Docker demo issuer password grant. + type: string + password_env_var: + type: string + title: Password Env Var + description: Controller environment variable that contains the Docker demo + issuer password grant shared secret. + default: AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD + scope: + title: Scope + description: OAuth scope for the Docker demo issuer. + type: string + subject_token_ttl_seconds: + type: integer + minimum: 1.0 + title: Subject Token Ttl Seconds + description: Fallback subject-token lifetime when the Docker demo issuer + response omits expires_in. + refresh_margin_seconds: + type: integer + minimum: 0.0 + title: Refresh Margin Seconds + description: Seconds before subject-token expiry when the Docker refresher + issues a replacement token. + default: 60 + additionalProperties: false + type: object + title: DockerWorkloadIdentityConfig + description: Docker-only subject token issuer configuration for workload identity. E2EJobExecutionProfile: properties: provider: @@ -12591,6 +12737,25 @@ components: these variables. type: object title: JobExecutionProfileConfig + JsonWebKey: + properties: {} + additionalProperties: true + type: object + title: JsonWebKey + description: JSON Web Key object. + JsonWebKeySetResponse: + properties: + keys: + items: + $ref: '#/components/schemas/JsonWebKey' + type: array + title: Keys + description: Public signing keys in the JWKS document. + type: object + required: + - keys + title: JsonWebKeySetResponse + description: JSON Web Key Set document. K8sNIMOperatorConfig: properties: resources: @@ -12630,6 +12795,31 @@ components: These fields provide typed access to commonly-used NIMService Spec fields and are applied before override_config in the compilation precedence.' + KubernetesConfigMapVolume: + properties: + name: + type: string + title: Name + description: ConfigMap name to mount + default_mode: + title: Default Mode + description: Optional default file mode + type: integer + optional: + title: Optional + description: Whether the ConfigMap is optional + type: boolean + items: + items: + $ref: '#/components/schemas/KubernetesKeyToPath' + type: array + title: Items + description: Optional ConfigMap keys to project + type: object + required: + - name + title: KubernetesConfigMapVolume + description: Kubernetes ConfigMap volume definition. KubernetesEmptyDirVolume: properties: medium: @@ -12787,6 +12977,18 @@ components: title: Launcher Image description: Container image that contains the jobs-launcher binary. default: nvcr.io/nvidia/nemo-microservices/jobs-launcher:latest + workload_identity_token_expiration_seconds: + type: integer + minimum: 600.0 + title: Workload Identity Token Expiration Seconds + description: Requested expirationSeconds for the projected service account + token used as the workload identity subject token. + default: 600 + workload_identity_token_audience: + title: Workload Identity Token Audience + description: Audience for the projected service account token. Defaults + to auth.oidc.workload_client_id, auth.oidc.client_id, then 'nemo-platform'. + type: string type: object title: KubernetesJobExecutionProfileConfig description: Configuration for Kubernetes execution environment. @@ -12817,6 +13019,26 @@ components: type: object title: KubernetesJobStorageConfig description: Configuration for persistent storage in Kubernetes jobs. + KubernetesKeyToPath: + properties: + key: + type: string + title: Key + description: Source key to project from the volume source + path: + type: string + title: Path + description: Relative file path to write the key to + mode: + title: Mode + description: Optional file mode for this key + type: integer + type: object + required: + - key + - path + title: KubernetesKeyToPath + description: Kubernetes volume key-to-path mapping. KubernetesObjectMetadata: properties: labels: @@ -12847,7 +13069,57 @@ components: - claim_name title: KubernetesPersistentVolumeClaim description: Kubernetes Persistent Volume Claim definition. + KubernetesSecretVolume: + properties: + secret_name: + type: string + title: Secret Name + description: Secret name to mount + default_mode: + title: Default Mode + description: Optional default file mode + type: integer + optional: + title: Optional + description: Whether the Secret is optional + type: boolean + items: + items: + $ref: '#/components/schemas/KubernetesKeyToPath' + type: array + title: Items + description: Optional Secret keys to project + type: object + required: + - secret_name + title: KubernetesSecretVolume + description: Kubernetes Secret volume definition. KubernetesVolume: + oneOf: + - properties: + persistent_volume_claim: + not: + type: 'null' + required: + - persistent_volume_claim + - properties: + empty_dir: + not: + type: 'null' + required: + - empty_dir + - properties: + secret: + not: + type: 'null' + required: + - secret + - properties: + config_map: + not: + type: 'null' + required: + - config_map properties: name: type: string @@ -12861,11 +13133,19 @@ components: allOf: - $ref: '#/components/schemas/KubernetesEmptyDirVolume' description: EmptyDir Volume configuration + secret: + allOf: + - $ref: '#/components/schemas/KubernetesSecretVolume' + description: Secret Volume configuration + config_map: + allOf: + - $ref: '#/components/schemas/KubernetesConfigMapVolume' + description: ConfigMap Volume configuration type: object required: - name title: KubernetesVolume - description: Kubernetes Volume definition. + description: Kubernetes Volume definition with secret and config_map support. KubernetesVolumeMount: properties: name: @@ -14759,6 +15039,22 @@ components: scope_prefix: title: Scope Prefix type: string + workload_token_exchange_enabled: + type: boolean + title: Workload Token Exchange Enabled + default: false + workload_client_id: + title: Workload Client Id + type: string + workload_token_endpoint: + title: Workload Token Endpoint + type: string + workload_audience: + title: Workload Audience + type: string + workload_scope: + title: Workload Scope + type: string type: object required: - issuer @@ -19078,6 +19374,18 @@ components: title: Launcher Image description: Container image that contains the jobs-launcher binary. default: nvcr.io/nvidia/nemo-microservices/jobs-launcher:latest + workload_identity_token_expiration_seconds: + type: integer + minimum: 600.0 + title: Workload Identity Token Expiration Seconds + description: Requested expirationSeconds for the projected service account + token used as the workload identity subject token. + default: 600 + workload_identity_token_audience: + title: Workload Identity Token Audience + description: Audience for the projected service account token. Defaults + to auth.oidc.workload_client_id, auth.oidc.client_id, then 'nemo-platform'. + type: string queue: type: string title: Queue @@ -19104,6 +19412,58 @@ components: type: object title: VolcanoJobExecutionProfileConfig description: Configuration for Volcano Job Execution Profile + WorkloadTokenExchangeErrorResponse: + properties: + error: + type: string + title: Error + description: OAuth 2.0 or RFC 8693 token exchange error code, such as invalid_client, + invalid_request, invalid_grant, invalid_scope, or invalid_target. + error_description: + title: Error Description + description: Human-readable ASCII text providing additional information + about the error. + type: string + error_uri: + title: Error Uri + description: URI identifying a human-readable web page with information + about the error. + type: string + type: object + required: + - error + title: WorkloadTokenExchangeErrorResponse + description: RFC 8693 token exchange error response. + WorkloadTokenExchangeResponse: + properties: + access_token: + type: string + title: Access Token + description: JWT access token minted for the workload identity. + issued_token_type: + type: string + title: Issued Token Type + description: Token type identifier for the issued token. + token_type: + type: string + title: Token Type + description: OAuth token type used in Authorization headers. + expires_in: + type: integer + title: Expires In + description: Lifetime of the access token in seconds. + scope: + title: Scope + description: Space-separated scopes granted to the access token. + type: string + type: object + required: + - access_token + - issued_token_type + - token_type + - expires_in + title: WorkloadTokenExchangeResponse + description: RFC 8693 token exchange response for workload identity access tokens. Workspace: properties: id: diff --git a/sdk/python/nemo-platform/.nmpcontext/stainless.yaml b/sdk/python/nemo-platform/.nmpcontext/stainless.yaml index a9a2f89d03..76037a1354 100644 --- a/sdk/python/nemo-platform/.nmpcontext/stainless.yaml +++ b/sdk/python/nemo-platform/.nmpcontext/stainless.yaml @@ -536,17 +536,21 @@ resources: docker_job_network_config: DockerJobNetworkConfig docker_job_storage_config: DockerJobStorageConfig docker_volume_mount: DockerVolumeMount + docker_workload_identity_config: DockerWorkloadIdentityConfig e2e_job_execution_profile: E2EJobExecutionProfile gpu_execution_provider: GPUExecutionProviderOutput gpu_execution_provider_param: GPUExecutionProviderInput image_pull_secret: ImagePullSecret job_execution_profile_config: JobExecutionProfileConfig + kubernetes_config_map_volume: KubernetesConfigMapVolume kubernetes_empty_dir_volume: KubernetesEmptyDirVolume kubernetes_job_execution_profile: KubernetesJobExecutionProfile kubernetes_job_execution_profile_config: KubernetesJobExecutionProfileConfig kubernetes_job_storage_config: KubernetesJobStorageConfig + kubernetes_key_to_path: KubernetesKeyToPath kubernetes_object_metadata: KubernetesObjectMetadata kubernetes_persistent_volume_claim: KubernetesPersistentVolumeClaim + kubernetes_secret_volume: KubernetesSecretVolume kubernetes_volume: KubernetesVolume kubernetes_volume_mount: KubernetesVolumeMount platform_job_environment_variable: PlatformJobEnvironmentVariable @@ -743,6 +747,10 @@ resources: tool_calling_metadata_content: ToolCallingMetadataContent backend_format: BackendFormat finetuning_type: FinetuningType + workload_token_exchange_response: WorkloadTokenExchangeResponse + workload_token_exchange_error_response: WorkloadTokenExchangeErrorResponse + json_web_key: JsonWebKey + json_web_key_set_response: JsonWebKeySetResponse iam: standalone_api: true subresources: diff --git a/sdk/python/nemo-platform/api.md b/sdk/python/nemo-platform/api.md index 3a573a2566..b235448956 100644 --- a/sdk/python/nemo-platform/api.md +++ b/sdk/python/nemo-platform/api.md @@ -16,6 +16,8 @@ from nemo_platform.types import ( GenericSortField, HTTPValidationError, InferenceParams, + JsonWebKey, + JsonWebKeySetResponse, LinearLayerSpec, MambaConfig, MoEConfig, @@ -37,6 +39,8 @@ from nemo_platform.types import ( ToolCallConfig, ToolCallingMetadataContent, ValidationError, + WorkloadTokenExchangeErrorResponse, + WorkloadTokenExchangeResponse, ) ``` diff --git a/sdk/python/nemo-platform/src/nemo_platform/_client.py b/sdk/python/nemo-platform/src/nemo_platform/_client.py index 14026ae6b7..893eae284f 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/_client.py +++ b/sdk/python/nemo-platform/src/nemo_platform/_client.py @@ -48,6 +48,9 @@ SyncAPIClient, AsyncAPIClient, ) +from nemo_platform._base_client import DefaultAsyncHttpxClient, DefaultHttpxClient +from nemo_platform.client.tls import client_verify_from_env +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR from pathlib import Path if TYPE_CHECKING: @@ -94,6 +97,29 @@ ] +def _should_bootstrap_config( + *, + http_client: object | None, + base_url: str | httpx.URL | None, + config_path: Path | None, + context_name: str | None, + access_token: str | None, +) -> bool: + """Return whether constructor arguments require config/auth bootstrap.""" + if http_client is not None: + return False + + # Backward compatibility: an explicit base_url means direct mode (no config + # bootstrap), unless config-specific overrides or workload identity are set. + return ( + base_url is None + or config_path is not None + or context_name is not None + or access_token is not None + or bool(os.environ.get(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR)) + ) + + class NeMoPlatform(SyncAPIClient): # client options workspace: str | None @@ -180,10 +206,12 @@ def __init__( http_client: Custom ``httpx.Client`` instance. When provided, the auth bootstrap is skipped entirely regardless of other parameters. """ - # Backward compatibility: an explicit base_url means direct mode (no config bootstrap), - # unless config-specific overrides are provided. - should_bootstrap = http_client is None and ( - base_url is None or config_path is not None or context_name is not None or access_token is not None + should_bootstrap = _should_bootstrap_config( + http_client=http_client, + base_url=base_url, + config_path=config_path, + context_name=context_name, + access_token=access_token, ) if should_bootstrap: try: @@ -204,6 +232,10 @@ def __init__( except Exception as e: raise RuntimeError(f"NeMoPlatform client initialization failed: {e}") + client_verify = client_verify_from_env() + if http_client is None and client_verify is not True: + http_client = DefaultHttpxClient(verify=client_verify) + self.workspace = workspace super().__init__( @@ -535,10 +567,12 @@ async def main() -> None: http_client: Custom ``httpx.AsyncClient`` instance. When provided, the auth bootstrap is skipped entirely regardless of other parameters. """ - # Backward compatibility: an explicit base_url means direct mode (no config bootstrap), - # unless config-specific overrides are provided. - should_bootstrap = http_client is None and ( - base_url is None or config_path is not None or context_name is not None or access_token is not None + should_bootstrap = _should_bootstrap_config( + http_client=http_client, + base_url=base_url, + config_path=config_path, + context_name=context_name, + access_token=access_token, ) if should_bootstrap: try: @@ -559,6 +593,10 @@ async def main() -> None: except Exception as e: raise RuntimeError(f"NeMoPlatform client initialization failed: {e}") + client_verify = client_verify_from_env() + if http_client is None and client_verify is not True: + http_client = DefaultAsyncHttpxClient(verify=client_verify) + self.workspace = workspace super().__init__( diff --git a/sdk/python/nemo-platform/src/nemo_platform/auth/device_flow.py b/sdk/python/nemo-platform/src/nemo_platform/auth/device_flow.py index 0a7b3d2fbd..06a62e1874 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/auth/device_flow.py +++ b/sdk/python/nemo-platform/src/nemo_platform/auth/device_flow.py @@ -13,6 +13,7 @@ from rich.panel import Panel from nemo_platform.auth.token_provider import refresh_token_grant +from nemo_platform.client.tls import client_verify_from_env console = Console() @@ -73,7 +74,7 @@ def __init__( async def start_device_authorization(self) -> DeviceCodeResponse: """Start the device authorization flow.""" - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(verify=client_verify_from_env()) as client: response = await client.post( self.device_authorization_endpoint, data={ @@ -103,7 +104,7 @@ async def poll_for_token( """Poll the token endpoint until authorization is complete.""" start_time = time.time() - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(verify=client_verify_from_env()) as client: while time.time() - start_time < expires_in: await _async_pause(interval) @@ -283,7 +284,7 @@ def authenticate_with_password_grant( "password": password, "scope": scope, } - with httpx.Client() as client: + with httpx.Client(verify=client_verify_from_env()) as client: response = client.post(token_endpoint, data=data, timeout=30.0) if response.status_code != 200: diff --git a/sdk/python/nemo-platform/src/nemo_platform/auth/helpers.py b/sdk/python/nemo-platform/src/nemo_platform/auth/helpers.py index 8155ab6aa6..c3ff84c17a 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/auth/helpers.py +++ b/sdk/python/nemo-platform/src/nemo_platform/auth/helpers.py @@ -25,6 +25,8 @@ import httpx +from nemo_platform.client.tls import client_verify_from_env + DEFAULT_OAUTH_SCOPES = "openid profile email offline_access" @@ -143,6 +145,11 @@ class NMPOIDCConfig: device_authorization_endpoint: str | None = None default_scopes: str = DEFAULT_OAUTH_SCOPES scope_prefix: str | None = None + workload_token_exchange_enabled: bool = False + workload_client_id: str | None = None + workload_token_endpoint: str | None = None + workload_audience: str | None = None + workload_scope: str | None = None def discover_nmp_config(base_url: str, timeout: float = 10.0) -> NMPOIDCConfig: @@ -150,6 +157,7 @@ def discover_nmp_config(base_url: str, timeout: float = 10.0) -> NMPOIDCConfig: response = httpx.get( f"{base_url.rstrip('/')}/apis/auth/discovery", timeout=timeout, + verify=client_verify_from_env(), ) response.raise_for_status() data = response.json() @@ -163,6 +171,11 @@ def discover_nmp_config(base_url: str, timeout: float = 10.0) -> NMPOIDCConfig: device_authorization_endpoint=oidc.get("device_authorization_endpoint"), default_scopes=oidc.get("default_scopes", DEFAULT_OAUTH_SCOPES), scope_prefix=oidc.get("scope_prefix"), + workload_token_exchange_enabled=oidc.get("workload_token_exchange_enabled", False), + workload_client_id=oidc.get("workload_client_id"), + workload_token_endpoint=oidc.get("workload_token_endpoint"), + workload_audience=oidc.get("workload_audience"), + workload_scope=oidc.get("workload_scope"), ) diff --git a/sdk/python/nemo-platform/src/nemo_platform/auth/token_provider.py b/sdk/python/nemo-platform/src/nemo_platform/auth/token_provider.py index ebe99e1581..947e0a87d6 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/auth/token_provider.py +++ b/sdk/python/nemo-platform/src/nemo_platform/auth/token_provider.py @@ -16,6 +16,7 @@ from typing_extensions import Self from nemo_platform.auth.helpers import decode_jwt_claims +from nemo_platform.client.tls import client_verify_from_env logger = logging.getLogger(__name__) @@ -32,6 +33,12 @@ def __init__(self, *, error: str, error_description: str) -> None: super().__init__(f"Token refresh failed: {error} - {error_description}") +def _validate_expires_in(expires_in: object) -> int | float | None: + if isinstance(expires_in, bool): + return None + return expires_in if isinstance(expires_in, int | float) else None + + def refresh_token_grant( token_endpoint: str, client_id: str, @@ -49,7 +56,7 @@ def refresh_token_grant( if scope: data["scope"] = scope - response = httpx.post(token_endpoint, data=data, timeout=timeout) + response = httpx.post(token_endpoint, data=data, timeout=timeout, verify=client_verify_from_env()) if response.status_code != 200: error_data: dict[str, str] = {} @@ -77,12 +84,16 @@ class TokenSet: def from_access_token( access_token: str, refresh_token: str | None = None, + expires_in: object = None, ) -> Self: """Create a TokenSet, extracting expiry from the JWT's `exp` claim.""" expires_at = None claims = decode_jwt_claims(access_token) if claims: expires_at = claims.get("exp") + validated_expires_in = _validate_expires_in(expires_in) + if expires_at is None and validated_expires_in is not None: + expires_at = time.time() + float(validated_expires_in) return TokenSet( access_token=access_token, refresh_token=refresh_token, @@ -223,7 +234,11 @@ def _refresh(self, *, force: bool = False) -> None: # The IdP may rotate the refresh token. new_refresh_token = token_data.get("refresh_token", self.tokens.refresh_token) - self.tokens = TokenSet.from_access_token(new_access_token, new_refresh_token) + self.tokens = TokenSet.from_access_token( + new_access_token, + new_refresh_token, + expires_in=token_data.get("expires_in"), + ) logger.debug("Access token refreshed successfully (expires_at=%s)", self.tokens.expires_at) if self.on_tokens_refreshed: diff --git a/sdk/python/nemo-platform/src/nemo_platform/auth/workload_exchange.py b/sdk/python/nemo-platform/src/nemo_platform/auth/workload_exchange.py new file mode 100644 index 0000000000..c54423fc01 --- /dev/null +++ b/sdk/python/nemo-platform/src/nemo_platform/auth/workload_exchange.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Workload identity token exchange for SDK authentication.""" + +from __future__ import annotations + +import asyncio +import json +import logging +import math +import threading +from dataclasses import dataclass, field +from ipaddress import ip_address +from pathlib import Path +from urllib.parse import urlparse + +import httpx +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR + +from nemo_platform.auth.token_provider import DEFAULT_REFRESH_MARGIN_SECONDS, TokenSet +from nemo_platform.client.tls import client_verify_from_env + +logger = logging.getLogger(__name__) + +TOKEN_EXCHANGE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange" +JWT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt" +ACCESS_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:access_token" + + +class WorkloadTokenExchangeError(RuntimeError): + """Structured error raised for RFC 8693 workload token exchange failures.""" + + def __init__(self, *, error: str, error_description: str) -> None: + self.error = error + self.error_description = error_description + super().__init__(f"Workload token exchange failed: {error} - {error_description}") + + +def read_subject_token_file(path: Path) -> str: + """Read a subject token from a workload identity token file.""" + try: + token = path.read_text(encoding="utf-8").strip() + except OSError as exc: + raise ValueError(f"Unable to read {WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR} at {path}: {exc}") from exc + if not token: + raise ValueError(f"{WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR} at {path} is empty") + return token + + +def _is_loopback_host(hostname: str | None) -> bool: + if hostname == "localhost": + return True + if hostname is None: + return False + try: + return ip_address(hostname).is_loopback + except ValueError: + return False + + +def _validate_token_endpoint(token_endpoint: str) -> None: + """Reject non-HTTPS token endpoints (except loopback for local dev).""" + parsed = urlparse(token_endpoint) + if parsed.scheme == "https": + return + if parsed.scheme == "http" and _is_loopback_host(parsed.hostname): + return + raise ValueError( + f"OIDC token endpoint must use HTTPS (got {token_endpoint!r}). " + "HTTP is only allowed for loopback addresses (localhost, 127.0.0.1, ::1)." + ) + + +def token_exchange_grant( + *, + token_endpoint: str, + client_id: str, + subject_token: str, + audience: str | None = None, + scope: str | None = None, + timeout: float = 30.0, +) -> dict[str, object]: + """Execute RFC 8693 token exchange and return token response JSON.""" + _validate_token_endpoint(token_endpoint) + data: dict[str, str] = { + "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, + "client_id": client_id, + "subject_token": subject_token, + "subject_token_type": JWT_TOKEN_TYPE, + "requested_token_type": ACCESS_TOKEN_TYPE, + } + if audience: + data["audience"] = audience + if scope: + data["scope"] = scope + + response = httpx.post(token_endpoint, data=data, timeout=timeout, verify=client_verify_from_env()) + + if response.status_code != 200: + error_data: dict[str, object] = {} + if response.headers.get("content-type", "").startswith("application/json"): + error_data = _response_json_object( + response, + error_description="Token endpoint error response was not a JSON object", + ) + error = _response_string(error_data, "error", "unknown_error") + error_description = _response_string(error_data, "error_description", response.text) + raise WorkloadTokenExchangeError(error=error, error_description=error_description) + + token_data = _response_json_object( + response, + error_description="Token endpoint response was not a JSON object", + ) + _access_token_from_response(token_data) + return token_data + + +def _response_json_object(response: httpx.Response, *, error_description: str) -> dict[str, object]: + try: + payload = response.json() + except (json.JSONDecodeError, ValueError) as exc: + raise WorkloadTokenExchangeError(error="invalid_response", error_description=error_description) from exc + if not isinstance(payload, dict): + raise WorkloadTokenExchangeError(error="invalid_response", error_description=error_description) + return payload + + +def _response_string(payload: dict[str, object], key: str, default: str) -> str: + value = payload.get(key) + return value if isinstance(value, str) and value else default + + +def _access_token_from_response(token_data: dict[str, object]) -> str: + access_token = token_data.get("access_token") + if not isinstance(access_token, str) or not access_token.strip(): + raise WorkloadTokenExchangeError( + error="invalid_response", + error_description="Token endpoint response did not include a non-empty access_token", + ) + return access_token + + +def _expires_in_from_response(token_data: dict[str, object]) -> int | float | None: + expires_in = token_data.get("expires_in") + if isinstance(expires_in, bool): + return None + return expires_in if isinstance(expires_in, int | float) else None + + +@dataclass +class WorkloadTokenExchangeProvider: + """Provides access tokens by exchanging a workload identity subject token file.""" + + token_endpoint: str + client_id: str + subject_token_file: Path + audience: str | None = None + scope: str | None = None + refresh_margin_seconds: float = DEFAULT_REFRESH_MARGIN_SECONDS + tokens: TokenSet = field(default_factory=lambda: TokenSet(access_token="")) + _lock: threading.Lock = field(default_factory=threading.Lock, repr=False) + + def get_access_token(self) -> str: + """Return a valid access token, exchanging the current subject token if needed.""" + with self._lock: + if not self.tokens.access_token or self.tokens.is_expired(self.refresh_margin_seconds): + self._exchange() + return self.tokens.access_token + + async def get_access_token_async(self) -> str: + """Return a valid access token in async contexts.""" + return await asyncio.to_thread(self.get_access_token) + + def _exchange(self) -> None: + subject_token = read_subject_token_file(self.subject_token_file) + logger.debug("Exchanging workload identity token via %s", self.token_endpoint) + token_data = token_exchange_grant( + token_endpoint=self.token_endpoint, + client_id=self.client_id, + subject_token=subject_token, + audience=self.audience, + scope=self.scope, + ) + access_token = _access_token_from_response(token_data) + try: + tokens = TokenSet.from_access_token( + access_token, + refresh_token=None, + expires_in=_expires_in_from_response(token_data), + ) + except (OverflowError, TypeError, ValueError) as exc: + raise WorkloadTokenExchangeError( + error="invalid_response", + error_description="Token endpoint response did not include a usable access_token lifetime", + ) from exc + if tokens.expires_at is None or not math.isfinite(tokens.expires_at): + raise WorkloadTokenExchangeError( + error="invalid_response", + error_description="Token endpoint response did not include a usable access_token lifetime", + ) + if tokens.is_expired(0): + raise WorkloadTokenExchangeError( + error="invalid_response", + error_description="Token endpoint response returned an expired access_token", + ) + self.tokens = tokens diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/auth.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/auth.py index a8a6b0aca5..4bf510bb51 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/auth.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/auth.py @@ -10,6 +10,7 @@ from __future__ import annotations import asyncio +import logging import time from typing import Annotated, cast @@ -37,17 +38,19 @@ help="Manage authentication for NeMo Platform.", ) +logger = logging.getLogger(__name__) -def is_auth_disabled(base_url: str, timeout: float = 3.0) -> bool | None: + +def is_auth_disabled(base_url: str, timeout: float = 3.0) -> bool: """Check whether authentication is disabled on the cluster. Returns: - True if auth is definitely disabled, False if enabled, None if unreachable. + True if auth is disabled, False if enabled. """ try: return not discover_nmp_config(base_url, timeout=timeout).auth_enabled - except httpx.HTTPError: - return None + except httpx.HTTPError as exc: + raise AuthError(f"Failed to discover auth configuration: {exc}") from exc def _runtime_token_source_label() -> str | None: @@ -57,7 +60,8 @@ def _runtime_token_source_label() -> str | None: try: return Config.runtime_access_token_source_label() except ValueError: - return "NEMO_WORKLOAD_TOKEN_FILE environment override could not be read" + logger.debug("Failed to resolve runtime token override source label", exc_info=True) + return None def ensure_valid_token(context: Context, refresh_buffer_seconds: int = 300) -> bool: @@ -131,13 +135,13 @@ def ensure_valid_token(context: Context, refresh_buffer_seconds: int = 300) -> b ) provider.force_refresh() - config_params = { + config_params: ConfigParams = { "access_token": provider.tokens.access_token, } if provider.tokens.refresh_token: config_params["refresh_token"] = provider.tokens.refresh_token - Config.write(config_params, context_name=context.context_name) # type: ignore[arg-type] + Config.write(config_params, context_name=context.context_name) typer.echo("[Auto-refreshed expired token]", err=True) return True @@ -282,7 +286,7 @@ def login( from nemo_platform.config.config import Config cli_context: CLIContext = ctx.obj - selected_context = cast(str | None, cli_context.overrides.get("current_context")) + selected_context = cli_context.overrides.get("current_context") if context_name is not None: selected_context = context_name @@ -537,9 +541,13 @@ def logout(ctx: typer.Context) -> None: console = Console() base_url = str(context.cluster.base_url).rstrip("/") - if is_auth_disabled(base_url) is True: - console.print("[yellow]Authentication is disabled on this cluster — nothing to log out from.[/]") - return + try: + if is_auth_disabled(base_url) is True: + console.print("[yellow]Authentication is disabled on this cluster — nothing to log out from.[/]") + return + except AuthError as exc: + logger.debug("Failed to discover auth configuration during logout", exc_info=True) + console.print(f"[yellow]Warning:[/] {exc}; continuing to clear local credentials.") logout_params: ConfigParams = {"access_token": None, "refresh_token": None} updated_config = Config.write(logout_params, context_name=context.context_name) @@ -699,10 +707,10 @@ def refresh(ctx: typer.Context) -> None: raise AuthError(f"Token refresh failed: {e}") from e # Save new tokens (refresh token may be rotated) - config_params = {"access_token": provider.tokens.access_token} + config_params: ConfigParams = {"access_token": provider.tokens.access_token} if provider.tokens.refresh_token: config_params["refresh_token"] = provider.tokens.refresh_token - Config.write(config_params, context_name=context.context_name) # type: ignore[arg-type] + Config.write(config_params, context_name=context.context_name) # Show new token info claims = decode_jwt_claims(provider.tokens.access_token) @@ -763,14 +771,21 @@ def status(ctx: typer.Context) -> None: # Check whether the cluster has auth enabled before showing token details. base_url = str(context.cluster.base_url).rstrip("/") - if is_auth_disabled(base_url) is True: - console.print() - console.print(f"[cyan]Cluster:[/] {base_url}") - console.print(f"[cyan]Context:[/] {context.context_name}") - console.print() - console.print("[green]Authentication is disabled on this cluster.[/]") - console.print("All API requests are accepted without credentials.") - return + auth_discovery_error: AuthError | None = None + try: + auth_disabled = is_auth_disabled(base_url) + except AuthError as exc: + logger.debug("Failed to discover auth configuration during status", exc_info=True) + auth_discovery_error = exc + else: + if auth_disabled is True: + console.print() + console.print(f"[cyan]Cluster:[/] {base_url}") + console.print(f"[cyan]Context:[/] {context.context_name}") + console.print() + console.print("[green]Authentication is disabled on this cluster.[/]") + console.print("All API requests are accepted without credentials.") + return table = Table(title="Authentication Status", show_header=False) table.add_column("Property", style="cyan") @@ -779,6 +794,8 @@ def status(ctx: typer.Context) -> None: table.add_row("Cluster", str(context.cluster.base_url)) table.add_row("Context", context.context_name) table.add_row("Config File", str(Config.get_default_config_path())) + if auth_discovery_error is not None: + table.add_row("Auth Discovery", f"[yellow]unavailable[/] ({auth_discovery_error})") runtime_token_source = _runtime_token_source_label() if context.user: diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/core/context.py b/sdk/python/nemo-platform/src/nemo_platform/cli/core/context.py index 1b267205da..e343d47a55 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/core/context.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/core/context.py @@ -6,6 +6,7 @@ from __future__ import annotations import logging +import os import typing from dataclasses import dataclass, field @@ -73,6 +74,33 @@ def get_sdk_context(self) -> Context: def reset_sdk_context(self) -> None: self._sdk_context = None + def _context_exists_in_config_file(self, context_name: str) -> bool: + from nemo_platform.config.config import Config + + try: + config = Config.load(overrides=self.overrides) + except FileNotFoundError: + return False + + return any(ctx.name == context_name for ctx in config.get_config_file().contexts) + + def _client_auth_config(self, ctx: Context) -> dict[str, object]: + from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR + + from nemo_platform.config.config import Config + from nemo_platform.config.models import OAuthUser + + if self.overrides.get("access_token") is not None or Config.runtime_access_token_source_label(): + return ctx.user.get_client_config() if ctx.user else {} + + if os.environ.get(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR): + return {} + + if isinstance(ctx.user, OAuthUser) and self._context_exists_in_config_file(ctx.context_name): + return {"context_name": ctx.context_name} + + return ctx.user.get_client_config() if ctx.user else {} + def get_client(self, timeout: float = 60.0) -> NeMoPlatform: """ Get or create the NeMo Platform client. @@ -92,12 +120,12 @@ def get_client(self, timeout: float = 60.0) -> NeMoPlatform: f"Creating NeMoPlatform client with base_url={base_url}, workspace={ctx.workspace}, timeout={timeout}" ) - client_config = ctx.user.get_client_config() + auth_config = self._client_auth_config(ctx) self._client = NeMoPlatform( base_url=base_url, timeout=timeout, workspace=ctx.workspace, - **client_config, + **auth_config, ) return self._client @@ -120,12 +148,12 @@ def get_async_client(self, timeout: float = 60.0) -> AsyncNeMoPlatform: f"Creating AsyncNeMoPlatform client with base_url={base_url}, workspace={ctx.workspace}, timeout={timeout}" ) - client_config = ctx.user.get_client_config() + auth_config = self._client_auth_config(ctx) self._async_client = AsyncNeMoPlatform( base_url=base_url, timeout=timeout, workspace=ctx.workspace, - **client_config, + **auth_config, ) return self._async_client diff --git a/sdk/python/nemo-platform/src/nemo_platform/client/factory.py b/sdk/python/nemo-platform/src/nemo_platform/client/factory.py index 5c1a696007..b99c64c75f 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/client/factory.py +++ b/sdk/python/nemo-platform/src/nemo_platform/client/factory.py @@ -49,13 +49,14 @@ See also: ``architecture/docs/auth/sdk-cli-oauth.md`` for a full design doc. """ +import asyncio import logging import os import threading from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path -from typing import Any, Callable, Mapping +from typing import Any, Callable, Mapping, Protocol import httpx from nemo_platform import ( @@ -65,12 +66,15 @@ NotGiven, not_given, ) +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR from nemo_platform.auth.helpers import NMPOIDCConfig, build_effective_scope, discover_nmp_config from nemo_platform.auth.token_provider import ( OIDCTokenProvider, TokenSet, ) +from nemo_platform.auth.workload_exchange import WorkloadTokenExchangeProvider +from nemo_platform.client.tls import client_verify_from_env logger = logging.getLogger(__name__) @@ -82,12 +86,17 @@ # Guards _TOKEN_PROVIDER_CACHE; acquired only during dict lookup/insert (fast). _TOKEN_PROVIDER_CACHE_LOCK = threading.Lock() - # --------------------------------------------------------------------------- # Data classes # --------------------------------------------------------------------------- +class _AccessTokenProvider(Protocol): + def get_access_token(self) -> str: ... + + async def get_access_token_async(self) -> str: ... + + @dataclass(frozen=True) class ClientInitConfig: """Everything the SDK client constructor needs after config resolution. @@ -110,7 +119,7 @@ class _ResolvedBootstrap: base_url: str workspace: str | None default_headers: dict[str, str] - token_provider: OIDCTokenProvider | None # None for non-OAuth users + token_provider: _AccessTokenProvider | None # None for non-OAuth users @dataclass(frozen=True) @@ -159,12 +168,86 @@ def _discover_oidc_client_settings(base_url: str) -> NMPOIDCConfig: ) +def _workload_identity_token_file_from_env() -> Path | None: + """Return the configured workload identity token file, if workload bootstrap is active.""" + token_file = os.environ.get(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR) + return Path(token_file) if token_file else None + + +def _create_workload_exchange_provider(base_url: str, subject_token_file: Path) -> WorkloadTokenExchangeProvider: + """Create a workload identity token exchange provider from NeMo auth discovery metadata.""" + oidc_config = _discover_oidc_client_settings(base_url) + if not oidc_config.workload_token_exchange_enabled: + raise RuntimeError( + f"{WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR} is set but workload token exchange is not enabled by auth discovery" + ) + + token_endpoint = oidc_config.workload_token_endpoint or oidc_config.token_endpoint or "" + client_id = oidc_config.workload_client_id or oidc_config.client_id or "" + if not token_endpoint: + raise RuntimeError( + "Workload token exchange is enabled but auth discovery did not return workload_token_endpoint or token_endpoint" + ) + if not client_id: + raise RuntimeError( + "Workload token exchange is enabled but auth discovery did not return workload_client_id or client_id" + ) + + return WorkloadTokenExchangeProvider( + token_endpoint=token_endpoint, + client_id=client_id, + subject_token_file=subject_token_file, + audience=oidc_config.workload_audience, + scope=oidc_config.workload_scope, + refresh_margin_seconds=_TOKEN_REFRESH_MARGIN_SECONDS, + ) + + +class _LazyWorkloadTokenExchangeProvider: + """Create the workload exchange provider on the first token request.""" + + def __init__(self, *, base_url: str, subject_token_file: Path) -> None: + self._base_url = base_url + self._subject_token_file = subject_token_file + self._provider: WorkloadTokenExchangeProvider | None = None + self._lock = threading.Lock() + + def _get_provider(self) -> WorkloadTokenExchangeProvider: + provider = self._provider + if provider is not None: + return provider + with self._lock: + provider = self._provider + if provider is None: + provider = _create_workload_exchange_provider(self._base_url, self._subject_token_file) + self._provider = provider + return provider + + def get_cached_access_token(self) -> str | None: + provider = self._provider + if provider is None: + return None + tokens = provider.tokens + if not tokens.access_token or tokens.is_expired(provider.refresh_margin_seconds): + return None + return tokens.access_token + + def get_access_token(self) -> str: + return self._get_provider().get_access_token() + + async def get_access_token_async(self) -> str: + provider = self._provider + if provider is None: + provider = await asyncio.to_thread(self._get_provider) + return await provider.get_access_token_async() + + # --------------------------------------------------------------------------- # httpx event hooks — the core of transparent token injection # --------------------------------------------------------------------------- -def _make_auth_event_hook(provider: OIDCTokenProvider): +def _make_auth_event_hook(provider: _AccessTokenProvider): """Create a **sync** httpx request event hook that injects the Bearer token. Called before every SDK HTTP request. ``provider.get_access_token()`` @@ -179,7 +262,7 @@ def inject_auth(request: httpx.Request) -> None: return inject_auth -def _make_async_auth_event_hook(provider: OIDCTokenProvider): +def _make_async_auth_event_hook(provider: _AccessTokenProvider): """Create an **async** httpx request event hook for AsyncNeMoPlatform. The actual refresh still runs in a worker thread (via @@ -193,6 +276,17 @@ async def inject_auth(request: httpx.Request) -> None: return inject_auth +def _headers_with_seeded_auth(headers: Mapping[str, str], provider: _AccessTokenProvider) -> dict[str, str]: + seeded_headers = dict(headers) + if isinstance(provider, _LazyWorkloadTokenExchangeProvider): + token = provider.get_cached_access_token() + else: + token = provider.get_access_token() + if token: + seeded_headers["Authorization"] = f"Bearer {token}" + return seeded_headers + + # --------------------------------------------------------------------------- # Callbacks wired into OIDCTokenProvider for config-file integration # --------------------------------------------------------------------------- @@ -411,6 +505,14 @@ def _resolve_bootstrap( base_url = str(resolved.cluster.base_url) headers: dict[str, str] = dict(extra_headers) if extra_headers else {} + workload_identity_token_file = _workload_identity_token_file_from_env() + if workload_identity_token_file is not None and access_token is None and not os.environ.get("NMP_ACCESS_TOKEN"): + provider = _LazyWorkloadTokenExchangeProvider( + base_url=base_url, + subject_token_file=workload_identity_token_file, + ) + return _ResolvedBootstrap(base_url, resolved.workspace, headers, provider) + # --- Non-OAuth path (no auth) --- if not isinstance(resolved.user, OAuthUser): user_config = resolved.user.get_client_config() if resolved.user else {} @@ -508,12 +610,17 @@ def build_client_init_kwargs( default_headers=bootstrap.default_headers or None, ) - # Seed the default headers with the current token so that SDK internals - # that inspect headers (e.g. auth_headers property) see a value. + # Seed the default headers with a current token so that SDK internals + # that inspect headers (e.g. auth_headers property) see a value. Workload + # identity only seeds when the request-time provider already has a token. # The event hook will overwrite it with a fresh token on each request. - headers = {**bootstrap.default_headers, "Authorization": f"Bearer {bootstrap.token_provider.get_access_token()}"} + headers = _headers_with_seeded_auth(bootstrap.default_headers, bootstrap.token_provider) hook = _make_auth_event_hook(bootstrap.token_provider) - http_client = DefaultHttpxClient(event_hooks={"request": [hook], "response": []}, follow_redirects=True) + http_client = DefaultHttpxClient( + event_hooks={"request": [hook], "response": []}, + follow_redirects=True, + verify=client_verify_from_env(), + ) return ClientInitConfig( base_url=bootstrap.base_url, workspace=bootstrap.workspace, @@ -550,9 +657,13 @@ def build_async_client_init_kwargs( default_headers=bootstrap.default_headers or None, ) - headers = {**bootstrap.default_headers, "Authorization": f"Bearer {bootstrap.token_provider.get_access_token()}"} + headers = _headers_with_seeded_auth(bootstrap.default_headers, bootstrap.token_provider) hook = _make_async_auth_event_hook(bootstrap.token_provider) - http_client = DefaultAsyncHttpxClient(event_hooks={"request": [hook], "response": []}, follow_redirects=True) + http_client = DefaultAsyncHttpxClient( + event_hooks={"request": [hook], "response": []}, + follow_redirects=True, + verify=client_verify_from_env(), + ) return ClientInitConfig( base_url=bootstrap.base_url, workspace=bootstrap.workspace, @@ -583,6 +694,9 @@ def create_client( access_token=access_token, extra_headers=extra_headers, ) + http_client = client_init_kwargs.http_client + if http_client is not None and not isinstance(http_client, httpx.Client): + raise TypeError("build_client_init_kwargs returned a non-sync HTTP client") return NeMoPlatform( config_path=config_path, @@ -591,7 +705,7 @@ def create_client( base_url=client_init_kwargs.base_url, workspace=client_init_kwargs.workspace, default_headers=client_init_kwargs.default_headers, - http_client=client_init_kwargs.http_client, + http_client=http_client, max_retries=max_retries, timeout=timeout, ) diff --git a/sdk/python/nemo-platform/src/nemo_platform/client/tls.py b/sdk/python/nemo-platform/src/nemo_platform/client/tls.py new file mode 100644 index 0000000000..b2bc998d2e --- /dev/null +++ b/sdk/python/nemo-platform/src/nemo_platform/client/tls.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""TLS configuration shared by NeMo Platform SDK and CLI clients.""" + +from __future__ import annotations + +import os + +NMP_CLIENT_SSL_CERT_FILE_ENVVAR = "NMP_CLIENT_SSL_CERT_FILE" + + +def client_verify_from_env() -> str | bool: + """Return the httpx verify setting for NeMo Platform client requests.""" + cert_file = os.environ.get(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "").strip() + return cert_file or True diff --git a/sdk/python/nemo-platform/src/nemo_platform/config/config.py b/sdk/python/nemo-platform/src/nemo_platform/config/config.py index ac8b2f70b9..f136ec9019 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/config/config.py +++ b/sdk/python/nemo-platform/src/nemo_platform/config/config.py @@ -29,9 +29,6 @@ logger = logging.getLogger(__name__) -_WORKLOAD_TOKEN_ENVVAR = "NEMO_WORKLOAD_TOKEN" -_WORKLOAD_TOKEN_FILE_ENVVAR = "NEMO_WORKLOAD_TOKEN_FILE" - @dataclass(frozen=True) class _RuntimeAccessTokenSource: @@ -141,18 +138,6 @@ def _migrate_legacy_api_key_users(cls, config_data: dict) -> None: def _runtime_access_token_source_from_env(cls) -> _RuntimeAccessTokenSource | None: if token := os.environ.get("NMP_ACCESS_TOKEN"): return _RuntimeAccessTokenSource(token, "NMP_ACCESS_TOKEN environment override") - if token := os.environ.get(_WORKLOAD_TOKEN_ENVVAR): - return _RuntimeAccessTokenSource(token, f"{_WORKLOAD_TOKEN_ENVVAR} environment override") - if token_path := os.environ.get(_WORKLOAD_TOKEN_FILE_ENVVAR): - try: - token = Path(token_path).read_text(encoding="utf-8").strip() - except OSError as exc: - raise ValueError(f"Unable to read {_WORKLOAD_TOKEN_FILE_ENVVAR} at {token_path}: {exc}") from exc - if token: - return _RuntimeAccessTokenSource( - token, - f"{_WORKLOAD_TOKEN_FILE_ENVVAR} environment override ({token_path})", - ) return None @classmethod diff --git a/sdk/python/nemo-platform/src/nemo_platform/resources/jobs/api.md b/sdk/python/nemo-platform/src/nemo_platform/resources/jobs/api.md index c06db14788..b269277b61 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/resources/jobs/api.md +++ b/sdk/python/nemo-platform/src/nemo_platform/resources/jobs/api.md @@ -17,17 +17,21 @@ from nemo_platform.types.jobs import ( DockerJobNetworkConfig, DockerJobStorageConfig, DockerVolumeMount, + DockerWorkloadIdentityConfig, E2EJobExecutionProfile, GPUExecutionProvider, GPUExecutionProviderParam, ImagePullSecret, JobExecutionProfileConfig, + KubernetesConfigMapVolume, KubernetesEmptyDirVolume, KubernetesJobExecutionProfile, KubernetesJobExecutionProfileConfig, KubernetesJobStorageConfig, + KubernetesKeyToPath, KubernetesObjectMetadata, KubernetesPersistentVolumeClaim, + KubernetesSecretVolume, KubernetesVolume, KubernetesVolumeMount, PlatformJobEnvironmentVariable, diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/__init__.py b/sdk/python/nemo-platform/src/nemo_platform/types/__init__.py index eb3af5c4f7..0dc61f81ec 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/types/__init__.py +++ b/sdk/python/nemo-platform/src/nemo_platform/types/__init__.py @@ -20,6 +20,7 @@ from .shared import ( ModelSpec as ModelSpec, MoEConfig as MoEConfig, + JsonWebKey as JsonWebKey, PromptData as PromptData, AuthContext as AuthContext, MambaConfig as MambaConfig, @@ -45,6 +46,7 @@ FilesetMetadataParam as FilesetMetadataParam, ModelMetadataContent as ModelMetadataContent, AuthDiscoveryResponse as AuthDiscoveryResponse, + JsonWebKeySetResponse as JsonWebKeySetResponse, OidcDiscoveryResponse as OidcDiscoveryResponse, DatasetMetadataContent as DatasetMetadataContent, PlatformJobResultResponse as PlatformJobResultResponse, @@ -53,4 +55,6 @@ PlatformJobListResultResponse as PlatformJobListResultResponse, PlatformJobStepStatusResponse as PlatformJobStepStatusResponse, PlatformJobTaskStatusResponse as PlatformJobTaskStatusResponse, + WorkloadTokenExchangeResponse as WorkloadTokenExchangeResponse, + WorkloadTokenExchangeErrorResponse as WorkloadTokenExchangeErrorResponse, ) diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/jobs/__init__.py b/sdk/python/nemo-platform/src/nemo_platform/types/jobs/__init__.py index 792e99704e..64858b7d59 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/types/jobs/__init__.py +++ b/sdk/python/nemo-platform/src/nemo_platform/types/jobs/__init__.py @@ -38,11 +38,13 @@ from .platform_job_response import PlatformJobResponse as PlatformJobResponse from .cpu_execution_provider import CPUExecutionProvider as CPUExecutionProvider from .gpu_execution_provider import GPUExecutionProvider as GPUExecutionProvider +from .kubernetes_key_to_path import KubernetesKeyToPath as KubernetesKeyToPath from .platform_job_step_spec import PlatformJobStepSpec as PlatformJobStepSpec from .compute_resources_param import ComputeResourcesParam as ComputeResourcesParam from .kubernetes_volume_mount import KubernetesVolumeMount as KubernetesVolumeMount from .platform_job_sort_field import PlatformJobSortField as PlatformJobSortField from .platform_job_spec_param import PlatformJobSpecParam as PlatformJobSpecParam +from .kubernetes_secret_volume import KubernetesSecretVolume as KubernetesSecretVolume from .docker_job_network_config import DockerJobNetworkConfig as DockerJobNetworkConfig from .docker_job_storage_config import DockerJobStorageConfig as DockerJobStorageConfig from .e2e_job_execution_profile import E2EJobExecutionProfile as E2EJobExecutionProfile @@ -55,6 +57,7 @@ from .docker_job_execution_profile import DockerJobExecutionProfile as DockerJobExecutionProfile from .gpu_execution_provider_param import GPUExecutionProviderParam as GPUExecutionProviderParam from .job_execution_profile_config import JobExecutionProfileConfig as JobExecutionProfileConfig +from .kubernetes_config_map_volume import KubernetesConfigMapVolume as KubernetesConfigMapVolume from .platform_job_list_sort_field import PlatformJobListSortField as PlatformJobListSortField from .platform_job_step_spec_param import PlatformJobStepSpecParam as PlatformJobStepSpecParam from .task_create_or_update_params import TaskCreateOrUpdateParams as TaskCreateOrUpdateParams @@ -62,6 +65,7 @@ from .subprocess_execution_provider import SubprocessExecutionProvider as SubprocessExecutionProvider from .volcano_job_execution_profile import VolcanoJobExecutionProfile as VolcanoJobExecutionProfile from .platform_job_step_with_context import PlatformJobStepWithContext as PlatformJobStepWithContext +from .docker_workload_identity_config import DockerWorkloadIdentityConfig as DockerWorkloadIdentityConfig from .platform_job_list_task_response import PlatformJobListTaskResponse as PlatformJobListTaskResponse from .platform_jobs_list_filter_param import PlatformJobsListFilterParam as PlatformJobsListFilterParam from .job_update_status_details_params import JobUpdateStatusDetailsParams as JobUpdateStatusDetailsParams diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/jobs/docker_job_execution_profile_config.py b/sdk/python/nemo-platform/src/nemo_platform/types/jobs/docker_job_execution_profile_config.py index cdce4354fa..b1b31cad62 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/types/jobs/docker_job_execution_profile_config.py +++ b/sdk/python/nemo-platform/src/nemo_platform/types/jobs/docker_job_execution_profile_config.py @@ -20,6 +20,7 @@ from ..._models import BaseModel from .docker_job_network_config import DockerJobNetworkConfig from .docker_job_storage_config import DockerJobStorageConfig +from .docker_workload_identity_config import DockerWorkloadIdentityConfig __all__ = ["DockerJobExecutionProfileConfig"] @@ -58,3 +59,6 @@ class DockerJobExecutionProfileConfig(BaseModel): ttl_seconds_after_finished: Optional[int] = None ttl_seconds_before_active: Optional[int] = None + + workload_identity: Optional[DockerWorkloadIdentityConfig] = None + """Docker-only subject token issuer configuration for workload identity.""" diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/jobs/docker_workload_identity_config.py b/sdk/python/nemo-platform/src/nemo_platform/types/jobs/docker_workload_identity_config.py new file mode 100644 index 0000000000..d4910a5138 --- /dev/null +++ b/sdk/python/nemo-platform/src/nemo_platform/types/jobs/docker_workload_identity_config.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +from typing import Optional + +from ..._models import BaseModel + +__all__ = ["DockerWorkloadIdentityConfig"] + + +class DockerWorkloadIdentityConfig(BaseModel): + """Docker-only subject token issuer configuration for workload identity.""" + + client_id: Optional[str] = None + """OAuth client ID used by the Docker demo issuer. + + Defaults to auth.oidc.workload_client_id or auth.oidc.client_id. + """ + + enabled: Optional[bool] = None + """Enable Docker workload identity token-file injection. + + Defaults to auth.oidc.workload_token_exchange_enabled. + """ + + password_env_var: Optional[str] = None + """ + Controller environment variable that contains the Docker demo issuer password + grant shared secret. + """ + + refresh_margin_seconds: Optional[int] = None + """ + Seconds before subject-token expiry when the Docker refresher issues a + replacement token. + """ + + scope: Optional[str] = None + """OAuth scope for the Docker demo issuer.""" + + subject_token_ttl_seconds: Optional[int] = None + """ + Fallback subject-token lifetime when the Docker demo issuer response omits + expires_in. + """ + + token_endpoint: Optional[str] = None + """OAuth token endpoint used by the Docker demo issuer. + + Defaults to auth.oidc.token_endpoint. + """ + + username: Optional[str] = None + """Username for the Docker demo issuer password grant.""" diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/jobs/kubernetes_config_map_volume.py b/sdk/python/nemo-platform/src/nemo_platform/types/jobs/kubernetes_config_map_volume.py new file mode 100644 index 0000000000..d2107f37f6 --- /dev/null +++ b/sdk/python/nemo-platform/src/nemo_platform/types/jobs/kubernetes_config_map_volume.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +from typing import List, Optional + +from ..._models import BaseModel +from .kubernetes_key_to_path import KubernetesKeyToPath + +__all__ = ["KubernetesConfigMapVolume"] + + +class KubernetesConfigMapVolume(BaseModel): + """Kubernetes ConfigMap volume definition.""" + + name: str + """ConfigMap name to mount""" + + default_mode: Optional[int] = None + """Optional default file mode""" + + items: Optional[List[KubernetesKeyToPath]] = None + """Optional ConfigMap keys to project""" + + optional: Optional[bool] = None + """Whether the ConfigMap is optional""" diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/jobs/kubernetes_job_execution_profile_config.py b/sdk/python/nemo-platform/src/nemo_platform/types/jobs/kubernetes_job_execution_profile_config.py index 81f5941f63..660808e16e 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/types/jobs/kubernetes_job_execution_profile_config.py +++ b/sdk/python/nemo-platform/src/nemo_platform/types/jobs/kubernetes_job_execution_profile_config.py @@ -107,3 +107,16 @@ class KubernetesJobExecutionProfileConfig(BaseModel): ttl_seconds_after_finished: Optional[int] = None ttl_seconds_before_active: Optional[int] = None + + workload_identity_token_audience: Optional[str] = None + """Audience for the projected service account token. + + Defaults to auth.oidc.workload_client_id, auth.oidc.client_id, then + 'nemo-platform'. + """ + + workload_identity_token_expiration_seconds: Optional[int] = None + """ + Requested expirationSeconds for the projected service account token used as the + workload identity subject token. + """ diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/jobs/kubernetes_key_to_path.py b/sdk/python/nemo-platform/src/nemo_platform/types/jobs/kubernetes_key_to_path.py new file mode 100644 index 0000000000..374b3d8bbc --- /dev/null +++ b/sdk/python/nemo-platform/src/nemo_platform/types/jobs/kubernetes_key_to_path.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +from typing import Optional + +from ..._models import BaseModel + +__all__ = ["KubernetesKeyToPath"] + + +class KubernetesKeyToPath(BaseModel): + """Kubernetes volume key-to-path mapping.""" + + key: str + """Source key to project from the volume source""" + + path: str + """Relative file path to write the key to""" + + mode: Optional[int] = None + """Optional file mode for this key""" diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/jobs/kubernetes_secret_volume.py b/sdk/python/nemo-platform/src/nemo_platform/types/jobs/kubernetes_secret_volume.py new file mode 100644 index 0000000000..ee66d2d441 --- /dev/null +++ b/sdk/python/nemo-platform/src/nemo_platform/types/jobs/kubernetes_secret_volume.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +from typing import List, Optional + +from ..._models import BaseModel +from .kubernetes_key_to_path import KubernetesKeyToPath + +__all__ = ["KubernetesSecretVolume"] + + +class KubernetesSecretVolume(BaseModel): + """Kubernetes Secret volume definition.""" + + secret_name: str + """Secret name to mount""" + + default_mode: Optional[int] = None + """Optional default file mode""" + + items: Optional[List[KubernetesKeyToPath]] = None + """Optional Secret keys to project""" + + optional: Optional[bool] = None + """Whether the Secret is optional""" diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/jobs/kubernetes_volume.py b/sdk/python/nemo-platform/src/nemo_platform/types/jobs/kubernetes_volume.py index 253c748502..77045edefe 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/types/jobs/kubernetes_volume.py +++ b/sdk/python/nemo-platform/src/nemo_platform/types/jobs/kubernetes_volume.py @@ -15,23 +15,80 @@ # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. -from typing import Optional +from typing import Union, Optional +from typing_extensions import TypeAlias from ..._models import BaseModel +from .kubernetes_secret_volume import KubernetesSecretVolume from .kubernetes_empty_dir_volume import KubernetesEmptyDirVolume +from .kubernetes_config_map_volume import KubernetesConfigMapVolume from .kubernetes_persistent_volume_claim import KubernetesPersistentVolumeClaim -__all__ = ["KubernetesVolume"] +__all__ = ["KubernetesVolume", "UnionMember0", "UnionMember1", "UnionMember2", "UnionMember3"] -class KubernetesVolume(BaseModel): - """Kubernetes Volume definition.""" +class UnionMember0(BaseModel): + persistent_volume_claim: object - name: str + config_map: Optional[KubernetesConfigMapVolume] = None + """Kubernetes ConfigMap volume definition.""" + + empty_dir: Optional[KubernetesEmptyDirVolume] = None + """Kubernetes EmptyDir Volume definition.""" + + name: Optional[str] = None + """Volume Name""" + + secret: Optional[KubernetesSecretVolume] = None + """Kubernetes Secret volume definition.""" + + +class UnionMember1(BaseModel): + empty_dir: object + + config_map: Optional[KubernetesConfigMapVolume] = None + """Kubernetes ConfigMap volume definition.""" + + name: Optional[str] = None + """Volume Name""" + + persistent_volume_claim: Optional[KubernetesPersistentVolumeClaim] = None + """Kubernetes Persistent Volume Claim definition.""" + + secret: Optional[KubernetesSecretVolume] = None + """Kubernetes Secret volume definition.""" + + +class UnionMember2(BaseModel): + secret: object + + config_map: Optional[KubernetesConfigMapVolume] = None + """Kubernetes ConfigMap volume definition.""" + + empty_dir: Optional[KubernetesEmptyDirVolume] = None + """Kubernetes EmptyDir Volume definition.""" + + name: Optional[str] = None """Volume Name""" + persistent_volume_claim: Optional[KubernetesPersistentVolumeClaim] = None + """Kubernetes Persistent Volume Claim definition.""" + + +class UnionMember3(BaseModel): + config_map: object + empty_dir: Optional[KubernetesEmptyDirVolume] = None """Kubernetes EmptyDir Volume definition.""" + name: Optional[str] = None + """Volume Name""" + persistent_volume_claim: Optional[KubernetesPersistentVolumeClaim] = None """Kubernetes Persistent Volume Claim definition.""" + + secret: Optional[KubernetesSecretVolume] = None + """Kubernetes Secret volume definition.""" + + +KubernetesVolume: TypeAlias = Union[UnionMember0, UnionMember1, UnionMember2, UnionMember3] diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/jobs/volcano_job_execution_profile_config.py b/sdk/python/nemo-platform/src/nemo_platform/types/jobs/volcano_job_execution_profile_config.py index c4ace352c0..699f4309f2 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/types/jobs/volcano_job_execution_profile_config.py +++ b/sdk/python/nemo-platform/src/nemo_platform/types/jobs/volcano_job_execution_profile_config.py @@ -120,3 +120,16 @@ class VolcanoJobExecutionProfileConfig(BaseModel): ttl_seconds_after_finished: Optional[int] = None ttl_seconds_before_active: Optional[int] = None + + workload_identity_token_audience: Optional[str] = None + """Audience for the projected service account token. + + Defaults to auth.oidc.workload_client_id, auth.oidc.client_id, then + 'nemo-platform'. + """ + + workload_identity_token_expiration_seconds: Optional[int] = None + """ + Requested expirationSeconds for the projected service account token used as the + workload identity subject token. + """ diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/shared/__init__.py b/sdk/python/nemo-platform/src/nemo_platform/types/shared/__init__.py index d16fead87f..e89faaf631 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/types/shared/__init__.py +++ b/sdk/python/nemo-platform/src/nemo_platform/types/shared/__init__.py @@ -19,6 +19,7 @@ from .mo_e_config import MoEConfig as MoEConfig from .prompt_data import PromptData as PromptData from .auth_context import AuthContext as AuthContext +from .json_web_key import JsonWebKey as JsonWebKey from .mamba_config import MambaConfig as MambaConfig from .string_filter import StringFilter as StringFilter from .backend_format import BackendFormat as BackendFormat @@ -44,9 +45,14 @@ from .auth_discovery_response import AuthDiscoveryResponse as AuthDiscoveryResponse from .oidc_discovery_response import OidcDiscoveryResponse as OidcDiscoveryResponse from .dataset_metadata_content import DatasetMetadataContent as DatasetMetadataContent +from .json_web_key_set_response import JsonWebKeySetResponse as JsonWebKeySetResponse from .platform_job_result_response import PlatformJobResultResponse as PlatformJobResultResponse from .platform_job_status_response import PlatformJobStatusResponse as PlatformJobStatusResponse from .tool_calling_metadata_content import ToolCallingMetadataContent as ToolCallingMetadataContent +from .workload_token_exchange_response import WorkloadTokenExchangeResponse as WorkloadTokenExchangeResponse from .platform_job_list_result_response import PlatformJobListResultResponse as PlatformJobListResultResponse from .platform_job_step_status_response import PlatformJobStepStatusResponse as PlatformJobStepStatusResponse from .platform_job_task_status_response import PlatformJobTaskStatusResponse as PlatformJobTaskStatusResponse +from .workload_token_exchange_error_response import ( + WorkloadTokenExchangeErrorResponse as WorkloadTokenExchangeErrorResponse, +) diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/shared/json_web_key.py b/sdk/python/nemo-platform/src/nemo_platform/types/shared/json_web_key.py new file mode 100644 index 0000000000..3420544b2c --- /dev/null +++ b/sdk/python/nemo-platform/src/nemo_platform/types/shared/json_web_key.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +from typing import Dict +from typing_extensions import TypeAlias + +__all__ = ["JsonWebKey"] + +JsonWebKey: TypeAlias = Dict[str, object] diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/shared/json_web_key_set_response.py b/sdk/python/nemo-platform/src/nemo_platform/types/shared/json_web_key_set_response.py new file mode 100644 index 0000000000..c5f9edc404 --- /dev/null +++ b/sdk/python/nemo-platform/src/nemo_platform/types/shared/json_web_key_set_response.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +from typing import List + +from ..._models import BaseModel +from .json_web_key import JsonWebKey + +__all__ = ["JsonWebKeySetResponse"] + + +class JsonWebKeySetResponse(BaseModel): + """JSON Web Key Set document.""" + + keys: List[JsonWebKey] + """Public signing keys in the JWKS document.""" diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/shared/oidc_discovery_response.py b/sdk/python/nemo-platform/src/nemo_platform/types/shared/oidc_discovery_response.py index 10da3bc377..a2f221e646 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/types/shared/oidc_discovery_response.py +++ b/sdk/python/nemo-platform/src/nemo_platform/types/shared/oidc_discovery_response.py @@ -40,3 +40,13 @@ class OidcDiscoveryResponse(BaseModel): token_endpoint: Optional[str] = None userinfo_endpoint: Optional[str] = None + + workload_audience: Optional[str] = None + + workload_client_id: Optional[str] = None + + workload_scope: Optional[str] = None + + workload_token_endpoint: Optional[str] = None + + workload_token_exchange_enabled: Optional[bool] = None diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/shared/workload_token_exchange_error_response.py b/sdk/python/nemo-platform/src/nemo_platform/types/shared/workload_token_exchange_error_response.py new file mode 100644 index 0000000000..e02b5d084b --- /dev/null +++ b/sdk/python/nemo-platform/src/nemo_platform/types/shared/workload_token_exchange_error_response.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +from typing import Optional + +from ..._models import BaseModel + +__all__ = ["WorkloadTokenExchangeErrorResponse"] + + +class WorkloadTokenExchangeErrorResponse(BaseModel): + """RFC 8693 token exchange error response.""" + + error: str + """ + OAuth 2.0 or RFC 8693 token exchange error code, such as invalid_client, + invalid_request, invalid_grant, invalid_scope, or invalid_target. + """ + + error_description: Optional[str] = None + """Human-readable ASCII text providing additional information about the error.""" + + error_uri: Optional[str] = None + """URI identifying a human-readable web page with information about the error.""" diff --git a/sdk/python/nemo-platform/src/nemo_platform/types/shared/workload_token_exchange_response.py b/sdk/python/nemo-platform/src/nemo_platform/types/shared/workload_token_exchange_response.py new file mode 100644 index 0000000000..5e4dd95092 --- /dev/null +++ b/sdk/python/nemo-platform/src/nemo_platform/types/shared/workload_token_exchange_response.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +from typing import Optional + +from ..._models import BaseModel + +__all__ = ["WorkloadTokenExchangeResponse"] + + +class WorkloadTokenExchangeResponse(BaseModel): + """RFC 8693 token exchange response for workload identity access tokens.""" + + access_token: str + """JWT access token minted for the workload identity.""" + + expires_in: int + """Lifetime of the access token in seconds.""" + + issued_token_type: str + """Token type identifier for the issued token.""" + + token_type: str + """OAuth token type used in Authorization headers.""" + + scope: Optional[str] = None + """Space-separated scopes granted to the access token.""" diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/auth/test_device_flow.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/auth/test_device_flow.py index 9b8db82749..05942c30fe 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/auth/test_device_flow.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/auth/test_device_flow.py @@ -26,6 +26,7 @@ authenticate_with_password_grant, refresh_access_token, ) +from nemo_platform.client.tls import NMP_CLIENT_SSL_CERT_FILE_ENVVAR class TestDeviceCodeResponse: @@ -137,6 +138,30 @@ async def test_start_device_authorization_success(self, device_flow): timeout=30.0, ) + @pytest.mark.asyncio + async def test_start_device_authorization_uses_nemo_scoped_ca_bundle(self, device_flow, monkeypatch): + mock_response_data = { + "device_code": "device_code_123", + "user_code": "ABC-123", + "verification_uri": "https://sso.example.com/device", + "expires_in": 1800, + } + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + with patch("httpx.AsyncClient") as mock_client_class: + mock_client = AsyncMock() + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_response.json.return_value = mock_response_data + mock_client.post.return_value = mock_response + mock_client.__aenter__.return_value = mock_client + mock_client.__aexit__.return_value = None + mock_client_class.return_value = mock_client + + await device_flow.start_device_authorization() + + mock_client_class.assert_called_once_with(verify="/tmp/nemo-ca.pem") + @pytest.mark.asyncio async def test_start_device_authorization_default_interval(self, device_flow): """Test device authorization with default interval.""" @@ -403,6 +428,32 @@ def test_authenticate_with_password_grant_success(self): timeout=30.0, ) + def test_authenticate_with_password_grant_uses_nemo_scoped_ca_bundle(self, monkeypatch): + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + with patch("httpx.Client") as mock_client_class: + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "access_token": "access_123", + "token_type": "Bearer", + "expires_in": 3600, + } + mock_client.post.return_value = mock_response + mock_client.__enter__.return_value = mock_client + mock_client.__exit__.return_value = None + mock_client_class.return_value = mock_client + + authenticate_with_password_grant( + token_endpoint="https://idp/token", + client_id="client-id", + username="user", + password="secret", + ) + + mock_client_class.assert_called_once_with(verify="/tmp/nemo-ca.pem") + def test_authenticate_with_password_grant_failure(self): with patch("httpx.Client") as mock_client_class: mock_client = MagicMock() diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/auth/test_token_provider.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/auth/test_token_provider.py index c8ae5913a1..cfe25bed7d 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/auth/test_token_provider.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/auth/test_token_provider.py @@ -13,7 +13,9 @@ from nemo_platform.auth.token_provider import ( OIDCTokenProvider, TokenSet, + refresh_token_grant, ) +from nemo_platform.client.tls import NMP_CLIENT_SSL_CERT_FILE_ENVVAR def _make_jwt(claims: dict, header: dict | None = None) -> str: @@ -41,6 +43,21 @@ def test_from_access_token_no_exp_claim(self): assert ts.expires_at is None + @pytest.mark.parametrize("expires_in", [120, 120.5]) + def test_from_access_token_uses_numeric_expires_in_when_no_exp_claim(self, expires_in): + token = _make_jwt({"sub": "user1"}) + before = time.time() + ts = TokenSet.from_access_token(token, expires_in=expires_in) + + assert ts.expires_at is not None + assert before + expires_in <= ts.expires_at <= time.time() + expires_in + + def test_from_access_token_rejects_bool_expires_in(self): + token = _make_jwt({"sub": "user1"}) + ts = TokenSet.from_access_token(token, expires_in=True) + + assert ts.expires_at is None + def test_from_access_token_non_jwt(self): ts = TokenSet.from_access_token("not-a-jwt", refresh_token="r") @@ -90,6 +107,23 @@ async def test_get_access_token_async_returns_current_when_not_expired(self): assert await provider.get_access_token_async() == token + @patch("nemo_platform.auth.token_provider.httpx.post") + def test_refresh_token_grant_uses_nemo_scoped_ca_bundle(self, mock_post, monkeypatch): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"access_token": "new_access"} + mock_post.return_value = mock_response + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + result = refresh_token_grant( + token_endpoint="https://idp/token", + client_id="client", + refresh_token="refresh_abc", + ) + + assert result == {"access_token": "new_access"} + assert mock_post.call_args.kwargs["verify"] == "/tmp/nemo-ca.pem" + @patch("nemo_platform.auth.token_provider.httpx.post") def test_get_access_token_refreshes_when_expired(self, mock_post): old_token = _make_jwt({"exp": int(time.time()) - 100}) @@ -121,6 +155,33 @@ def test_get_access_token_refreshes_when_expired(self, mock_post): assert call_kwargs[1]["data"]["client_id"] == "client" assert call_kwargs[1]["data"]["refresh_token"] == "old_refresh" + @patch("nemo_platform.auth.token_provider.httpx.post") + def test_get_access_token_refreshes_opaque_token_with_expires_in(self, mock_post): + old_token = _make_jwt({"exp": int(time.time()) - 100}) + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "access_token": "opaque_access", + "expires_in": 120, + } + mock_post.return_value = mock_response + + tokens = TokenSet.from_access_token(old_token, refresh_token="old_refresh") + provider = OIDCTokenProvider( + token_endpoint="https://idp/token", + client_id="client", + tokens=tokens, + refresh_margin_seconds=0, + ) + + before = time.time() + result = provider.get_access_token() + + assert result == "opaque_access" + assert provider.tokens.expires_at is not None + assert before + 120 <= provider.tokens.expires_at <= time.time() + 120 + @patch("nemo_platform.auth.token_provider.httpx.post") def test_refresh_reloads_tokens_before_request(self, mock_post): stale_token = _make_jwt({"exp": int(time.time()) - 200}) diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/auth/test_utils.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/auth/test_utils.py index 03ab0485d0..686f2a2c7e 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/auth/test_utils.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/auth/test_utils.py @@ -3,6 +3,7 @@ import base64 import json +from unittest.mock import MagicMock, patch import httpx import pytest @@ -18,6 +19,7 @@ normalize_scope_prefix, validate_requested_scopes_granted, ) +from nemo_platform.client.tls import NMP_CLIENT_SSL_CERT_FILE_ENVVAR from pytest_httpserver import HTTPServer @@ -160,6 +162,22 @@ def test_strips_trailing_slash(self, httpserver: HTTPServer): result = discover_nmp_config(httpserver.url_for("") + "/") assert result.auth_enabled is False + @patch("nemo_platform.auth.helpers.httpx.get") + def test_uses_nemo_scoped_ca_bundle(self, mock_get, monkeypatch): + response = MagicMock() + response.json.return_value = {"auth_enabled": False} + mock_get.return_value = response + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + result = discover_nmp_config("https://nemo.example.com") + + assert result.auth_enabled is False + mock_get.assert_called_once_with( + "https://nemo.example.com/apis/auth/discovery", + timeout=10.0, + verify="/tmp/nemo-ca.pem", + ) + class TestBuildEffectiveScope: def test_no_prefix_returns_unchanged(self): diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/auth/test_workload_exchange.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/auth/test_workload_exchange.py new file mode 100644 index 0000000000..311c43173b --- /dev/null +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/auth/test_workload_exchange.py @@ -0,0 +1,258 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for RFC 8693 workload identity token exchange.""" + +import json +import time +from base64 import urlsafe_b64encode +from unittest.mock import MagicMock, patch + +import pytest +from nemo_platform.auth.workload_exchange import ( + ACCESS_TOKEN_TYPE, + JWT_TOKEN_TYPE, + TOKEN_EXCHANGE_GRANT_TYPE, + WorkloadTokenExchangeError, + WorkloadTokenExchangeProvider, + read_subject_token_file, + token_exchange_grant, +) +from nemo_platform.client.tls import NMP_CLIENT_SSL_CERT_FILE_ENVVAR +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR + + +def _make_jwt(claims: dict) -> str: + header = {"alg": "RS256", "typ": "JWT"} + h = urlsafe_b64encode(json.dumps(header).encode()).rstrip(b"=").decode() + p = urlsafe_b64encode(json.dumps(claims).encode()).rstrip(b"=").decode() + s = urlsafe_b64encode(b"fake-signature").rstrip(b"=").decode() + return f"{h}.{p}.{s}" + + +def test_read_subject_token_file_strips_whitespace(tmp_path): + token_file = tmp_path / "token" + token_file.write_text("subject-token\n", encoding="utf-8") + + assert read_subject_token_file(token_file) == "subject-token" + + +def test_read_subject_token_file_rejects_empty_file(tmp_path): + token_file = tmp_path / "token" + token_file.write_text("\n", encoding="utf-8") + + with pytest.raises(ValueError, match=WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR): + read_subject_token_file(token_file) + + +@patch("nemo_platform.auth.workload_exchange.httpx.post") +def test_token_exchange_grant_sends_rfc8693_request(mock_post): + access_token = _make_jwt({"exp": int(time.time()) + 3600}) + response = MagicMock() + response.status_code = 200 + response.json.return_value = {"access_token": access_token, "expires_in": 300} + mock_post.return_value = response + + result = token_exchange_grant( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token="subject-token", + audience="nemo-platform", + scope="openid email groups", + ) + + assert result["access_token"] == access_token + mock_post.assert_called_once() + assert mock_post.call_args.kwargs["data"] == { + "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, + "client_id": "nemo-platform-workload", + "subject_token": "subject-token", + "subject_token_type": JWT_TOKEN_TYPE, + "requested_token_type": ACCESS_TOKEN_TYPE, + "audience": "nemo-platform", + "scope": "openid email groups", + } + + +@patch("nemo_platform.auth.workload_exchange.httpx.post") +def test_token_exchange_grant_rejects_http_non_loopback_endpoint_before_sending_subject_token(mock_post): + with pytest.raises(ValueError, match="must use HTTPS"): + token_exchange_grant( + token_endpoint="http://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token="subject-token", + ) + + mock_post.assert_not_called() + + +@patch("nemo_platform.auth.workload_exchange.httpx.post") +@pytest.mark.parametrize( + "token_endpoint", + [ + "http://localhost:18080/token", + "http://127.0.0.1:18080/token", + "http://[::1]:18080/token", + ], +) +def test_token_exchange_grant_allows_http_loopback_endpoints(mock_post, token_endpoint): + access_token = _make_jwt({"exp": int(time.time()) + 3600}) + response = MagicMock() + response.status_code = 200 + response.json.return_value = {"access_token": access_token} + mock_post.return_value = response + + result = token_exchange_grant( + token_endpoint=token_endpoint, + client_id="nemo-platform-workload", + subject_token="subject-token", + ) + + assert result["access_token"] == access_token + mock_post.assert_called_once() + + +@patch("nemo_platform.auth.workload_exchange.httpx.post") +def test_token_exchange_grant_uses_nemo_scoped_ca_bundle(mock_post, monkeypatch): + access_token = _make_jwt({"exp": int(time.time()) + 3600}) + response = MagicMock() + response.status_code = 200 + response.json.return_value = {"access_token": access_token} + mock_post.return_value = response + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + result = token_exchange_grant( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token="subject-token", + ) + + assert result["access_token"] == access_token + assert mock_post.call_args.kwargs["verify"] == "/tmp/nemo-ca.pem" + + +@patch("nemo_platform.auth.workload_exchange.httpx.post") +def test_token_exchange_grant_surfaces_idp_error(mock_post): + response = MagicMock() + response.status_code = 400 + response.text = "invalid subject token" + response.headers = {"content-type": "application/json"} + response.json.return_value = { + "error": "invalid_request", + "error_description": "invalid subject token", + } + mock_post.return_value = response + + with pytest.raises(WorkloadTokenExchangeError, match="invalid_request - invalid subject token"): + token_exchange_grant( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token="bad-subject-token", + ) + + +@patch("nemo_platform.auth.workload_exchange.httpx.post") +def test_token_exchange_grant_rejects_non_object_error_payload(mock_post): + response = MagicMock() + response.status_code = 400 + response.text = "[]" + response.headers = {"content-type": "application/json"} + response.json.return_value = [] + mock_post.return_value = response + + with pytest.raises(WorkloadTokenExchangeError, match="invalid_response - Token endpoint error response"): + token_exchange_grant( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token="bad-subject-token", + ) + + +@patch("nemo_platform.auth.workload_exchange.httpx.post") +@pytest.mark.parametrize("payload", [[], {}, {"access_token": ""}, {"access_token": None}]) +def test_token_exchange_grant_rejects_success_response_without_non_empty_access_token(mock_post, payload): + response = MagicMock() + response.status_code = 200 + response.json.return_value = payload + mock_post.return_value = response + + with pytest.raises(WorkloadTokenExchangeError): + token_exchange_grant( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token="subject-token", + ) + + +@patch("nemo_platform.auth.workload_exchange.token_exchange_grant") +def test_provider_rejects_exchange_response_without_access_token(mock_exchange, tmp_path): + subject_token_file = tmp_path / "token" + subject_token_file.write_text("subject-token", encoding="utf-8") + mock_exchange.return_value = {} + + provider = WorkloadTokenExchangeProvider( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token_file=subject_token_file, + ) + + with pytest.raises(WorkloadTokenExchangeError, match="non-empty access_token"): + provider.get_access_token() + + +@patch("nemo_platform.auth.workload_exchange.token_exchange_grant") +@pytest.mark.parametrize( + "expires_in", + [None, "300", True, float("nan"), float("inf"), 10**400], +) +def test_provider_rejects_exchange_response_without_usable_lifetime(mock_exchange, tmp_path, expires_in): + subject_token_file = tmp_path / "token" + subject_token_file.write_text("subject-token", encoding="utf-8") + token_data = {"access_token": "opaque-access-token"} + if expires_in is not None: + token_data["expires_in"] = expires_in + mock_exchange.return_value = token_data + + provider = WorkloadTokenExchangeProvider( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token_file=subject_token_file, + ) + + with pytest.raises(WorkloadTokenExchangeError, match="usable access_token lifetime"): + provider.get_access_token() + + assert provider.tokens.access_token == "" + + +@patch("nemo_platform.auth.workload_exchange.token_exchange_grant") +def test_provider_rejects_expired_exchange_response_and_retries_with_current_subject_token(mock_exchange, tmp_path): + subject_token_file = tmp_path / "token" + subject_token_file.write_text("subject-token-one", encoding="utf-8") + expired_access_token = _make_jwt({"exp": int(time.time()) - 10}) + fresh_access_token = _make_jwt({"exp": int(time.time()) + 3600}) + mock_exchange.side_effect = [ + {"access_token": expired_access_token}, + {"access_token": fresh_access_token}, + ] + + provider = WorkloadTokenExchangeProvider( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform-workload", + subject_token_file=subject_token_file, + audience="nemo-platform", + scope="openid email groups", + refresh_margin_seconds=0, + ) + + with pytest.raises(WorkloadTokenExchangeError, match="expired access_token"): + provider.get_access_token() + + subject_token_file.write_text("subject-token-two", encoding="utf-8") + + assert provider.get_access_token() == fresh_access_token + assert mock_exchange.call_count == 2 + assert mock_exchange.call_args_list[0].kwargs["subject_token"] == "subject-token-one" + assert mock_exchange.call_args_list[1].kwargs["subject_token"] == "subject-token-two" + assert mock_exchange.call_args_list[1].kwargs["audience"] == "nemo-platform" + assert mock_exchange.call_args_list[1].kwargs["scope"] == "openid email groups" diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_auth.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_auth.py index fc148fb5e8..eec6cc4ef1 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_auth.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_auth.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import logging from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace @@ -10,6 +11,7 @@ import yaml from nemo_platform.auth.helpers import decode_jwt_claims, generate_unsigned_jwt from nemo_platform.cli.app import app +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR from typer.testing import CliRunner from ..utils import assert_exit_code @@ -67,7 +69,12 @@ def _decode_jwt_noop(token: str) -> dict: @pytest.fixture def oauth_config_file(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: - for env_key in ("NMP_ACCESS_TOKEN", "NEMO_WORKLOAD_TOKEN", "NEMO_WORKLOAD_TOKEN_FILE"): + for env_key in ( + "NMP_ACCESS_TOKEN", + "NEMO_WORKLOAD_TOKEN", + "NEMO_WORKLOAD_TOKEN_FILE", + WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, + ): monkeypatch.delenv(env_key, raising=False) config_data = { @@ -158,7 +165,7 @@ def test_auth_logout_warns_when_runtime_token_override_remains( ) -> None: monkeypatch.setattr("nemo_platform.cli.commands.auth.discover_nmp_config", _discover_auth_enabled) monkeypatch.setenv( - "NEMO_WORKLOAD_TOKEN", + "NMP_ACCESS_TOKEN", generate_unsigned_jwt( principal_id="svc-nemo-ci", email="svc-nemo-ci@example.com", @@ -170,7 +177,7 @@ def test_auth_logout_warns_when_runtime_token_override_remains( assert_exit_code(result, 0) assert "Logged out successfully" in result.output - assert "NEMO_WORKLOAD_TOKEN environment override is still active" in result.output + assert "NMP_ACCESS_TOKEN environment override is still active" in result.output def test_auth_logout_fails_if_credentials_remain(oauth_config_file: Path, monkeypatch: pytest.MonkeyPatch) -> None: @@ -308,7 +315,7 @@ def test_auth_status_shows_warning_for_unsigned_token(oauth_config_file: Path, m assert "local/testing" in result.output -def test_runtime_token_source_label_handles_unreadable_token_file( +def test_runtime_token_source_label_ignores_workload_identity_token_file( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: from nemo_platform.cli.commands.auth import _runtime_token_source_label @@ -317,8 +324,26 @@ def test_runtime_token_source_label_handles_unreadable_token_file( monkeypatch.delenv("NMP_ACCESS_TOKEN", raising=False) monkeypatch.delenv("NEMO_WORKLOAD_TOKEN", raising=False) monkeypatch.setenv("NEMO_WORKLOAD_TOKEN_FILE", str(token_file)) + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(token_file)) - assert _runtime_token_source_label() == "NEMO_WORKLOAD_TOKEN_FILE environment override could not be read" + assert _runtime_token_source_label() is None + + +def test_runtime_token_source_label_returns_none_when_config_label_fails( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + from nemo_platform.cli.commands.auth import _runtime_token_source_label + + monkeypatch.setattr( + "nemo_platform.config.config.Config.runtime_access_token_source_label", + lambda: (_ for _ in ()).throw(ValueError("invalid runtime token")), + ) + + with caplog.at_level(logging.DEBUG, logger="nemo_platform.cli.commands.auth"): + assert _runtime_token_source_label() is None + + assert "Failed to resolve runtime token override source label" in caplog.text + assert "invalid runtime token" in caplog.text def test_auth_status_shows_config_file_credential_source( @@ -651,8 +676,8 @@ def test_auth_login_unsigned_token_uses_principal_id_when_provided( @dataclass class IsAuthDisabledCase: id: str - auth_enabled: bool | None # None means the cluster is unreachable (raises) - expected: bool | None + auth_enabled: bool + expected: bool @pytest.mark.parametrize( @@ -660,16 +685,11 @@ class IsAuthDisabledCase: [ IsAuthDisabledCase(id="disabled", auth_enabled=False, expected=True), IsAuthDisabledCase(id="enabled", auth_enabled=True, expected=False), - IsAuthDisabledCase(id="unreachable", auth_enabled=None, expected=None), ], ids=lambda c: c.id, ) def test_is_auth_disabled(monkeypatch: pytest.MonkeyPatch, case: IsAuthDisabledCase) -> None: - import httpx - def mock_discover(url: str, timeout: float = 10.0) -> SimpleNamespace: - if case.auth_enabled is None: - raise httpx.ConnectError("Connection refused") return SimpleNamespace(auth_enabled=case.auth_enabled) monkeypatch.setattr("nemo_platform.cli.commands.auth.discover_nmp_config", mock_discover) @@ -678,6 +698,20 @@ def mock_discover(url: str, timeout: float = 10.0) -> SimpleNamespace: assert is_auth_disabled("http://localhost:8080") is case.expected +def test_is_auth_disabled_raises_when_discovery_fails(monkeypatch: pytest.MonkeyPatch) -> None: + import httpx + from nemo_platform.auth.helpers import AuthError + from nemo_platform.cli.commands.auth import is_auth_disabled + + def raise_connect_error(url: str, timeout: float = 10.0) -> SimpleNamespace: + raise httpx.ConnectError("Connection refused") + + monkeypatch.setattr("nemo_platform.cli.commands.auth.discover_nmp_config", raise_connect_error) + + with pytest.raises(AuthError, match="Failed to discover auth configuration: Connection refused"): + is_auth_disabled("http://localhost:8080") + + # --------------------------------------------------------------------------- # auth status when auth disabled # --------------------------------------------------------------------------- @@ -702,6 +736,31 @@ def test_auth_status_when_auth_disabled_does_not_show_token_details( assert "Refresh Token" not in result.output +def test_auth_status_when_cluster_unreachable_shows_local_state( + oauth_config_file: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + import httpx + + def raise_connect_error(url: str, timeout: float = 10.0) -> SimpleNamespace: + raise httpx.ConnectError("Connection refused") + + monkeypatch.setattr("nemo_platform.cli.commands.auth.discover_nmp_config", raise_connect_error) + result = runner.invoke(app, ["--context", "foo", "auth", "status"]) + output = " ".join(result.output.split()) + + assert_exit_code(result, 0) + assert "Failed to discover auth configuration:" in output + assert "Connection refused" in output + assert "Auth Discovery" in result.output + assert "unavailable" in result.output + assert "Auth Type" in result.output + assert "oauth" in result.output + assert "Credential Source" in result.output + assert "config file" in result.output + assert "Refresh Token" in result.output + assert "foo-token" not in result.output + + # --------------------------------------------------------------------------- # auth logout when auth disabled # --------------------------------------------------------------------------- @@ -719,7 +778,7 @@ def test_auth_logout_when_auth_disabled_shows_message_and_skips_credential_clear mock_write.assert_not_called() -def test_auth_logout_when_cluster_unreachable_still_clears_credentials( +def test_auth_logout_when_cluster_unreachable_clears_local_credentials( oauth_config_file: Path, monkeypatch: pytest.MonkeyPatch ) -> None: import httpx @@ -728,9 +787,23 @@ def raise_connect_error(url: str, timeout: float = 10.0) -> SimpleNamespace: raise httpx.ConnectError("Connection refused") monkeypatch.setattr("nemo_platform.cli.commands.auth.discover_nmp_config", raise_connect_error) - with patch("nemo_platform.config.config.Config.write") as mock_write: - result = runner.invoke(app, ["--context", "foo", "auth", "logout"]) + result = runner.invoke(app, ["--context", "foo", "auth", "logout"]) + output = " ".join(result.output.split()) assert_exit_code(result, 0) - mock_write.assert_called_once() - assert mock_write.call_args.kwargs["context_name"] == "foo" + assert "Failed to discover auth configuration: Connection refused" in output + assert "continuing to clear local credentials" in output + assert "Logged out successfully" in result.output + + with open(oauth_config_file) as f: + data = yaml.safe_load(f) + + default_user = next(user for user in data["users"] if user["name"] == "default") + foo_user = next(user for user in data["users"] if user["name"] == "foo") + + assert default_user["type"] == "oauth" + assert default_user["token"] == "default-token" + assert default_user["refresh_token"] == "default-refresh" + assert foo_user["type"] == "no-auth" + assert "token" not in foo_user + assert "refresh_token" not in foo_user diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/core/test_context.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/core/test_context.py index f68fc38edb..1db10e7d7a 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/core/test_context.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/core/test_context.py @@ -1,7 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from types import SimpleNamespace +from unittest.mock import patch + from nemo_platform.cli.core.context import CLIContext +from nemo_platform.config.models import NoAuthUser, OAuthUser def test_context_instances_are_independent(): @@ -60,16 +64,112 @@ def test_get_no_truncate_default(): assert result is False -def test_get_client_passes_user_config_and_is_cached(): - """Test that get_client passes user's client config to the SDK client and caches it.""" - ctx = CLIContext(overrides={"base_url": "http://test.example.com", "access_token": "token-123"}) +def test_get_client_uses_config_bootstrap_for_persisted_oauth_context_and_is_cached(): + """Test that get_client lets the SDK bootstrap config-backed OAuth auth and caches it.""" + ctx = CLIContext(overrides={"base_url": "http://test.example.com"}) + resolved_context = SimpleNamespace( + cluster=SimpleNamespace(base_url="http://test.example.com"), + context_name="dev", + workspace="test-workspace", + user=OAuthUser(name="dev-user", token="token-123", refresh_token="refresh-123"), + ) + config_file = SimpleNamespace(contexts=[SimpleNamespace(name="dev")]) + + with ( + patch("nemo_platform.config.config.get_context", return_value=resolved_context), + patch("nemo_platform.config.config.Config.load") as mock_config_load, + patch("nemo_platform.config.config.Config.runtime_access_token_source_label", return_value=None), + patch("nemo_platform.NeMoPlatform", autospec=True) as mock_client_cls, + ): + mock_config_load.return_value.get_config_file.return_value = config_file + client = ctx.get_client() + client2 = ctx.get_client() + + mock_client_cls.assert_called_once_with( + base_url="http://test.example.com", + context_name="dev", + timeout=60.0, + workspace="test-workspace", + ) + assert client is mock_client_cls.return_value - client = ctx.get_client() + # Verify the client is cached + assert client is client2 - # Verify the client has the expected headers from get_client_config() - assert "Authorization" in client.default_headers - assert client.default_headers["Authorization"] == "Bearer token-123" - # Verify the client is cached - client2 = ctx.get_client() +def test_get_client_preserves_direct_mode_for_synthetic_no_auth_context(): + """Test that synthesized default/no-auth contexts do not force SDK config bootstrap.""" + ctx = CLIContext(overrides={"base_url": "http://test.example.com"}) + resolved_context = SimpleNamespace( + cluster=SimpleNamespace(base_url="http://test.example.com"), + context_name="default", + workspace="default", + user=NoAuthUser(name="default-user"), + ) + + with ( + patch("nemo_platform.config.config.get_context", return_value=resolved_context), + patch("nemo_platform.NeMoPlatform", autospec=True) as mock_client_cls, + ): + ctx.get_client() + + mock_client_cls.assert_called_once_with( + base_url="http://test.example.com", + timeout=60.0, + workspace="default", + ) + + +def test_get_client_passes_explicit_access_token_override(): + """Test that explicit access token overrides remain caller-managed static headers.""" + ctx = CLIContext(overrides={"base_url": "http://test.example.com", "access_token": "token-123"}) + resolved_context = SimpleNamespace( + cluster=SimpleNamespace(base_url="http://test.example.com"), + context_name="dev", + workspace="test-workspace", + user=OAuthUser(name="dev-user", token="token-123"), + ) + + with ( + patch("nemo_platform.config.config.get_context", return_value=resolved_context), + patch("nemo_platform.NeMoPlatform", autospec=True) as mock_client_cls, + ): + ctx.get_client() + + mock_client_cls.assert_called_once_with( + base_url="http://test.example.com", + default_headers={"Authorization": "Bearer token-123"}, + timeout=60.0, + workspace="test-workspace", + ) + + +def test_get_async_client_uses_config_bootstrap_for_persisted_oauth_context_and_is_cached(): + """Test that get_async_client lets the SDK bootstrap config-backed OAuth auth and caches it.""" + ctx = CLIContext(overrides={"base_url": "http://test.example.com"}) + resolved_context = SimpleNamespace( + cluster=SimpleNamespace(base_url="http://test.example.com"), + context_name="dev", + workspace="test-workspace", + user=OAuthUser(name="dev-user", token="token-123", refresh_token="refresh-123"), + ) + config_file = SimpleNamespace(contexts=[SimpleNamespace(name="dev")]) + + with ( + patch("nemo_platform.config.config.get_context", return_value=resolved_context), + patch("nemo_platform.config.config.Config.load") as mock_config_load, + patch("nemo_platform.config.config.Config.runtime_access_token_source_label", return_value=None), + patch("nemo_platform.AsyncNeMoPlatform", autospec=True) as mock_client_cls, + ): + mock_config_load.return_value.get_config_file.return_value = config_file + client = ctx.get_async_client() + client2 = ctx.get_async_client() + + mock_client_cls.assert_called_once_with( + base_url="http://test.example.com", + context_name="dev", + timeout=60.0, + workspace="test-workspace", + ) + assert client is mock_client_cls.return_value assert client is client2 diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/client/test_client.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/client/test_client.py index dd03f8344e..19422053fb 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/client/test_client.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/client/test_client.py @@ -17,6 +17,8 @@ from nemo_platform import AsyncNeMoPlatform, DefaultHttpxClient, NeMoPlatform, not_given from nemo_platform.auth.helpers import NMPOIDCConfig, decode_jwt_claims from nemo_platform.client.factory import create_client +from nemo_platform.client.tls import NMP_CLIENT_SSL_CERT_FILE_ENVVAR +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR def _make_jwt(claims: dict) -> str: @@ -67,6 +69,17 @@ def _write_config(tmp_path, *, user_type="oauth", token=None, refresh_token=None token_endpoint="https://idp/token", ) +_MOCK_WORKLOAD_NMP_CONFIG = NMPOIDCConfig( + auth_enabled=True, + client_id="nmp-client-id", + token_endpoint="https://idp/token", + workload_token_exchange_enabled=True, + workload_client_id="nmp-workload-client-id", + workload_token_endpoint="https://workload-idp/token", + workload_audience="nemo-platform", + workload_scope="openid email groups", +) + class TestCreateClientOAuth: @patch("nemo_platform.client.factory.discover_nmp_config", return_value=_MOCK_NMP_CONFIG) @@ -116,6 +129,27 @@ def test_oauth_uses_sdk_default_httpx_client(self, _mock_discover, tmp_path): finally: client.close() + @patch("nemo_platform.client.factory.DefaultHttpxClient") + @patch("nemo_platform.client.factory.discover_nmp_config", return_value=_MOCK_NMP_CONFIG) + def test_oauth_uses_nemo_scoped_ca_bundle(self, _mock_discover, mock_default_httpx_client, tmp_path, monkeypatch): + token = _make_jwt({"exp": int(time.time()) + 3600, "sub": "user1"}) + config_path = _write_config( + tmp_path, + token=token, + refresh_token="refresh_abc", + ) + http_client = httpx.Client() + mock_default_httpx_client.return_value = http_client + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + client = create_client(config_path=config_path) + try: + assert client is not None + finally: + client.close() + + assert mock_default_httpx_client.call_args.kwargs["verify"] == "/tmp/nemo-ca.pem" + @patch("nemo_platform.client.factory.discover_nmp_config", return_value=_MOCK_NMP_CONFIG) @patch("nemo_platform.auth.token_provider.httpx.post") def test_persist_refreshed_tokens_writes_to_config(self, mock_post, _mock_discover, tmp_path): @@ -158,6 +192,90 @@ def test_env_access_token_overrides_user_auth(self, tmp_path, monkeypatch): assert request.headers["Authorization"] == "Bearer env-access-token-123" +class TestCreateClientWorkloadIdentity: + @patch("nemo_platform.client.factory.discover_nmp_config", return_value=_MOCK_WORKLOAD_NMP_CONFIG) + @patch("nemo_platform.auth.workload_exchange.token_exchange_grant") + def test_exchanges_workload_identity_token_file(self, mock_exchange, _mock_discover, tmp_path, monkeypatch): + subject_token_file = tmp_path / "workload-token" + subject_token_file.write_text("subject-token-one\n", encoding="utf-8") + access_token = _make_jwt({"exp": int(time.time()) + 3600, "sub": "workload-user"}) + mock_exchange.return_value = {"access_token": access_token, "expires_in": 300} + monkeypatch.setenv("NMP_BASE_URL", "https://api.example.com") + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(subject_token_file)) + + client = create_client() + + try: + assert str(client.base_url).rstrip("/") == "https://api.example.com" + assert "Authorization" not in client._custom_headers + _mock_discover.assert_not_called() + mock_exchange.assert_not_called() + + request = client._client.build_request("GET", "https://api.example.com/test") + client._client._event_hooks["request"][0](request) + assert request.headers["Authorization"] == f"Bearer {access_token}" + finally: + client.close() + + mock_exchange.assert_called_once() + assert mock_exchange.call_args.kwargs["token_endpoint"] == "https://workload-idp/token" + assert mock_exchange.call_args.kwargs["client_id"] == "nmp-workload-client-id" + assert mock_exchange.call_args.kwargs["subject_token"] == "subject-token-one" + assert mock_exchange.call_args.kwargs["audience"] == "nemo-platform" + assert mock_exchange.call_args.kwargs["scope"] == "openid email groups" + + @pytest.mark.asyncio + @patch("nemo_platform.client.factory.discover_nmp_config", return_value=_MOCK_WORKLOAD_NMP_CONFIG) + @patch("nemo_platform.auth.workload_exchange.token_exchange_grant") + async def test_async_exchanges_workload_identity_token_file_at_request_time( + self, mock_exchange, _mock_discover, tmp_path, monkeypatch + ): + subject_token_file = tmp_path / "workload-token" + subject_token_file.write_text("subject-token-one\n", encoding="utf-8") + access_token = _make_jwt({"exp": int(time.time()) + 3600, "sub": "workload-user"}) + mock_exchange.return_value = {"access_token": access_token, "expires_in": 300} + monkeypatch.setenv("NMP_BASE_URL", "https://api.example.com") + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(subject_token_file)) + + client = AsyncNeMoPlatform() + + try: + assert str(client.base_url).rstrip("/") == "https://api.example.com" + assert "Authorization" not in client._custom_headers + _mock_discover.assert_not_called() + mock_exchange.assert_not_called() + + request = client._client.build_request("GET", "https://api.example.com/test") + await client._client._event_hooks["request"][0](request) + assert request.headers["Authorization"] == f"Bearer {access_token}" + finally: + await client.close() + + mock_exchange.assert_called_once() + assert mock_exchange.call_args.kwargs["token_endpoint"] == "https://workload-idp/token" + assert mock_exchange.call_args.kwargs["client_id"] == "nmp-workload-client-id" + assert mock_exchange.call_args.kwargs["subject_token"] == "subject-token-one" + assert mock_exchange.call_args.kwargs["audience"] == "nemo-platform" + assert mock_exchange.call_args.kwargs["scope"] == "openid email groups" + + @patch("nemo_platform.client.factory.discover_nmp_config", return_value=_MOCK_WORKLOAD_NMP_CONFIG) + def test_env_access_token_takes_precedence_over_workload_identity_file(self, _mock_discover, tmp_path, monkeypatch): + subject_token_file = tmp_path / "workload-token" + subject_token_file.write_text("subject-token-one\n", encoding="utf-8") + monkeypatch.setenv("NMP_BASE_URL", "https://api.example.com") + monkeypatch.setenv("NMP_ACCESS_TOKEN", "env-access-token-123") + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(subject_token_file)) + + client = create_client() + + try: + request = client._client.build_request("GET", "https://api.example.com/test") + client._client._event_hooks["request"][0](request) + assert request.headers["Authorization"] == "Bearer env-access-token-123" + finally: + client.close() + + class TestCreateClientApiKey: def test_creates_client_with_api_key(self, tmp_path): config_path = _write_config(tmp_path, user_type="api-key", api_key="nvapi-test-key-123") @@ -357,6 +475,42 @@ def test_sync_constructor_with_base_url_skips_config_bootstrap(self, mock_build_ mock_build_client_kwargs.assert_not_called() + @patch("nemo_platform._client.DefaultHttpxClient") + @patch("nemo_platform.client.factory.build_client_init_kwargs") + def test_sync_constructor_direct_mode_uses_nemo_scoped_ca_bundle( + self, mock_build_client_kwargs, mock_default_httpx_client, monkeypatch + ): + mock_build_client_kwargs.side_effect = AssertionError("bootstrap should not be called") + mock_default_httpx_client.return_value = httpx.Client() + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + client = NeMoPlatform(base_url="https://override-host:8081", workspace="test-workspace") + try: + assert str(client.base_url).rstrip("/") == "https://override-host:8081" + finally: + client.close() + + mock_build_client_kwargs.assert_not_called() + mock_default_httpx_client.assert_called_once_with(verify="/tmp/nemo-ca.pem") + + @patch("nemo_platform.client.factory.build_client_init_kwargs") + def test_sync_constructor_with_workload_file_and_base_url_bootstraps(self, mock_build_client_kwargs, monkeypatch): + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, "/var/run/secrets/nemo-platform/workload/token") + mock_build_client_kwargs.return_value = MagicMock( + base_url="http://override-host:8081", + workspace="test-workspace", + default_headers=None, + http_client=None, + ) + + client = NeMoPlatform(base_url="http://override-host:8081", workspace="test-workspace") + try: + assert str(client.base_url).rstrip("/") == "http://override-host:8081" + finally: + client.close() + + assert mock_build_client_kwargs.call_args.kwargs["base_url"] == "http://override-host:8081" + @patch("nemo_platform.client.factory.build_client_init_kwargs") def test_sync_constructor_passes_context_name_to_bootstrap(self, mock_build_client_kwargs): mock_build_client_kwargs.return_value = MagicMock( @@ -411,6 +565,46 @@ async def test_async_constructor_with_base_url_skips_config_bootstrap(self, mock mock_build_client_kwargs.assert_not_called() + @pytest.mark.asyncio + @patch("nemo_platform._client.DefaultAsyncHttpxClient") + @patch("nemo_platform.client.factory.build_async_client_init_kwargs") + async def test_async_constructor_direct_mode_uses_nemo_scoped_ca_bundle( + self, mock_build_client_kwargs, mock_default_httpx_client, monkeypatch + ): + mock_build_client_kwargs.side_effect = AssertionError("bootstrap should not be called") + mock_default_httpx_client.return_value = httpx.AsyncClient() + monkeypatch.setenv(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, "/tmp/nemo-ca.pem") + + client = AsyncNeMoPlatform(base_url="https://override-host:8081", workspace="test-workspace") + try: + assert str(client.base_url).rstrip("/") == "https://override-host:8081" + finally: + await client.close() + + mock_build_client_kwargs.assert_not_called() + mock_default_httpx_client.assert_called_once_with(verify="/tmp/nemo-ca.pem") + + @pytest.mark.asyncio + @patch("nemo_platform.client.factory.build_async_client_init_kwargs") + async def test_async_constructor_with_workload_file_and_base_url_bootstraps( + self, mock_build_client_kwargs, monkeypatch + ): + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, "/var/run/secrets/nemo-platform/workload/token") + mock_build_client_kwargs.return_value = MagicMock( + base_url="http://override-host:8081", + workspace="test-workspace", + default_headers=None, + http_client=None, + ) + + client = AsyncNeMoPlatform(base_url="http://override-host:8081", workspace="test-workspace") + try: + assert str(client.base_url).rstrip("/") == "http://override-host:8081" + finally: + await client.close() + + assert mock_build_client_kwargs.call_args.kwargs["base_url"] == "http://override-host:8081" + @pytest.mark.asyncio @patch("nemo_platform.client.factory.build_async_client_init_kwargs") async def test_async_constructor_with_http_client_skips_config_bootstrap(self, mock_build_client_kwargs): diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/config/test_config.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/config/test_config.py index dbe34f96c8..5af80018c9 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/config/test_config.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/config/test_config.py @@ -13,6 +13,7 @@ get_context, ) from nemo_platform.config.models import DEFAULT_BASE_URL, ConfigFile, NoAuthUser, OAuthUser +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR @pytest.fixture @@ -337,18 +338,17 @@ def test_config_from_env_access_token_only(self, tmp_path: Path, monkeypatch: py assert config.user.refresh_token is None assert not hasattr(config.user, "token_endpoint") - def test_config_from_workload_token_env_only(self, monkeypatch: pytest.MonkeyPatch): - """NEMO_WORKLOAD_TOKEN should bootstrap OAuth auth without a config file.""" + def test_legacy_workload_token_env_is_ignored(self, monkeypatch: pytest.MonkeyPatch): + """NEMO_WORKLOAD_TOKEN is no longer a runtime auth bootstrap source.""" monkeypatch.setenv("NMP_BASE_URL", "https://api.example.com") monkeypatch.setenv("NEMO_WORKLOAD_TOKEN", "workload-token-123") config = get_context() - assert isinstance(config.user, OAuthUser) - assert config.user.token.get_secret_value() == "workload-token-123" + assert isinstance(config.user, NoAuthUser) - def test_config_from_workload_token_file_env_only(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): - """NEMO_WORKLOAD_TOKEN_FILE should bootstrap OAuth auth without a config file.""" + def test_legacy_workload_token_file_env_is_ignored(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + """NEMO_WORKLOAD_TOKEN_FILE is no longer read as a bearer-token file.""" token_path = tmp_path / "workload.token" token_path.write_text("workload-token-from-file\n", encoding="utf-8") monkeypatch.setenv("NMP_BASE_URL", "https://api.example.com") @@ -356,25 +356,37 @@ def test_config_from_workload_token_file_env_only(self, tmp_path: Path, monkeypa config = get_context() - assert isinstance(config.user, OAuthUser) - assert config.user.token.get_secret_value() == "workload-token-from-file" + assert isinstance(config.user, NoAuthUser) - def test_config_from_missing_workload_token_file_reports_configuration_error( - self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch - ): - """NEMO_WORKLOAD_TOKEN_FILE should fail clearly when the configured token file cannot be read.""" + def test_missing_legacy_workload_token_file_is_ignored(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + """Legacy workload token file env no longer triggers config-time file reads.""" token_path = tmp_path / "missing-workload.token" monkeypatch.setenv("NMP_BASE_URL", "https://api.example.com") monkeypatch.setenv("NEMO_WORKLOAD_TOKEN_FILE", str(token_path)) - with pytest.raises(ValueError, match="NEMO_WORKLOAD_TOKEN_FILE"): - get_context() + config = get_context() + + assert isinstance(config.user, NoAuthUser) + + def test_workload_identity_token_file_is_not_static_access_token( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ): + """The workload identity token file env var is handled by the client factory, not Config.""" + token_path = tmp_path / "workload.token" + token_path.write_text("subject-token\n", encoding="utf-8") + monkeypatch.setenv("NMP_BASE_URL", "https://api.example.com") + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(token_path)) + + config = get_context() + + assert isinstance(config.user, NoAuthUser) def test_nmp_access_token_precedes_workload_token_env(self, monkeypatch: pytest.MonkeyPatch): """NMP_ACCESS_TOKEN remains the highest-precedence token env var.""" monkeypatch.setenv("NMP_BASE_URL", "https://api.example.com") monkeypatch.setenv("NMP_ACCESS_TOKEN", "preferred-token") monkeypatch.setenv("NEMO_WORKLOAD_TOKEN", "workload-token-123") + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, "/var/run/secrets/nemo-platform/workload/token") config = get_context() @@ -389,6 +401,7 @@ def test_runtime_access_token_source_label_uses_config_precedence( monkeypatch.setenv("NMP_ACCESS_TOKEN", "preferred-token") monkeypatch.setenv("NEMO_WORKLOAD_TOKEN", "workload-token-123") monkeypatch.setenv("NEMO_WORKLOAD_TOKEN_FILE", str(missing_token_path)) + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(missing_token_path)) assert Config.runtime_access_token_source_label() == "NMP_ACCESS_TOKEN environment override" diff --git a/sdk/stainless.yaml b/sdk/stainless.yaml index a9a2f89d03..76037a1354 100644 --- a/sdk/stainless.yaml +++ b/sdk/stainless.yaml @@ -536,17 +536,21 @@ resources: docker_job_network_config: DockerJobNetworkConfig docker_job_storage_config: DockerJobStorageConfig docker_volume_mount: DockerVolumeMount + docker_workload_identity_config: DockerWorkloadIdentityConfig e2e_job_execution_profile: E2EJobExecutionProfile gpu_execution_provider: GPUExecutionProviderOutput gpu_execution_provider_param: GPUExecutionProviderInput image_pull_secret: ImagePullSecret job_execution_profile_config: JobExecutionProfileConfig + kubernetes_config_map_volume: KubernetesConfigMapVolume kubernetes_empty_dir_volume: KubernetesEmptyDirVolume kubernetes_job_execution_profile: KubernetesJobExecutionProfile kubernetes_job_execution_profile_config: KubernetesJobExecutionProfileConfig kubernetes_job_storage_config: KubernetesJobStorageConfig + kubernetes_key_to_path: KubernetesKeyToPath kubernetes_object_metadata: KubernetesObjectMetadata kubernetes_persistent_volume_claim: KubernetesPersistentVolumeClaim + kubernetes_secret_volume: KubernetesSecretVolume kubernetes_volume: KubernetesVolume kubernetes_volume_mount: KubernetesVolumeMount platform_job_environment_variable: PlatformJobEnvironmentVariable @@ -743,6 +747,10 @@ resources: tool_calling_metadata_content: ToolCallingMetadataContent backend_format: BackendFormat finetuning_type: FinetuningType + workload_token_exchange_response: WorkloadTokenExchangeResponse + workload_token_exchange_error_response: WorkloadTokenExchangeErrorResponse + json_web_key: JsonWebKey + json_web_key_set_response: JsonWebKeySetResponse iam: standalone_api: true subresources: diff --git a/services/core/auth/src/nmp/core/auth/api/v2/discovery/endpoints.py b/services/core/auth/src/nmp/core/auth/api/v2/discovery/endpoints.py index 88df9bd9a6..960e1930aa 100644 --- a/services/core/auth/src/nmp/core/auth/api/v2/discovery/endpoints.py +++ b/services/core/auth/src/nmp/core/auth/api/v2/discovery/endpoints.py @@ -7,8 +7,9 @@ import time import httpx -from fastapi import APIRouter +from fastapi import APIRouter, Request from nmp.common.config import get_auth_config +from nmp.core.auth.api.v2.workload_token_exchange import workload_token_endpoint_url from pydantic import BaseModel logger = logging.getLogger(__name__) @@ -32,6 +33,11 @@ class OIDCDiscoveryResponse(BaseModel): client_id: str default_scopes: str = "openid profile email offline_access" scope_prefix: str | None = None + workload_token_exchange_enabled: bool = False + workload_client_id: str | None = None + workload_token_endpoint: str | None = None + workload_audience: str | None = None + workload_scope: str | None = None class AuthDiscoveryResponse(BaseModel): @@ -111,9 +117,19 @@ def _clear_idp_discovery_cache() -> None: - `client_id`: OAuth client ID to use - `default_scopes`: OAuth scopes to request during authentication - `scope_prefix`: Prefix to prepend to custom scopes (those with ':' or '.default') + - `workload_token_exchange_enabled`: Whether SDK workload identity token exchange is enabled + - `workload_client_id`: OAuth client ID to use for workload identity token exchange + - `workload_token_endpoint`: Token endpoint to use only for workload identity token exchange + - `workload_audience`: RFC 8693 audience for exchanged workload tokens + - `workload_scope`: OAuth scopes for exchanged workload tokens """, ) -async def get_auth_discovery() -> AuthDiscoveryResponse: +async def get_auth_discovery_endpoint(request: Request) -> AuthDiscoveryResponse: + """FastAPI route wrapper for auth configuration discovery.""" + return await get_auth_discovery(request) + + +async def get_auth_discovery(request: Request | None = None) -> AuthDiscoveryResponse: """Return auth configuration for CLI/SDK discovery. This endpoint is unauthenticated and returns the information @@ -136,6 +152,14 @@ async def get_auth_discovery() -> AuthDiscoveryResponse: client_id=config.oidc.client_id, default_scopes=config.oidc.default_scopes, scope_prefix=config.oidc.scope_prefix, + workload_token_exchange_enabled=config.oidc.workload_token_exchange_enabled, + workload_client_id=config.oidc.workload_client_id, + workload_token_endpoint=( + config.oidc.workload_token_endpoint + or (workload_token_endpoint_url(request) if config.oidc.workload_token_exchange_enabled else None) + ), + workload_audience=config.oidc.workload_audience, + workload_scope=config.oidc.workload_scope, ) return AuthDiscoveryResponse( diff --git a/services/core/auth/src/nmp/core/auth/api/v2/workload_token_exchange.py b/services/core/auth/src/nmp/core/auth/api/v2/workload_token_exchange.py new file mode 100644 index 0000000000..dd98d4e52a --- /dev/null +++ b/services/core/auth/src/nmp/core/auth/api/v2/workload_token_exchange.py @@ -0,0 +1,552 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Workload identity token exchange endpoints.""" + +from __future__ import annotations + +import json +import logging +import time +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import Any + +import httpx +import jwt +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from fastapi import APIRouter, Depends, Request +from fastapi.responses import JSONResponse +from jwt.algorithms import RSAAlgorithm +from nmp.common.config import AuthConfig, get_auth_config, get_platform_config +from pydantic import BaseModel, ConfigDict, Field + +logger = logging.getLogger(__name__) + +TOKEN_EXCHANGE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange" +JWT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt" +ACCESS_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:access_token" + +WORKLOAD_TOKEN_PATH = "/apis/auth/token" +WORKLOAD_JWKS_PATH = "/apis/auth/jwks" +DEFAULT_WORKLOAD_AUDIENCE = "nemo-platform" +DEFAULT_WORKLOAD_SCOPE = "openid email groups" + +_TOKEN_EXCHANGE_FORM_REQUEST_BODY: dict[str, Any] = { + "required": True, + "content": { + "application/x-www-form-urlencoded": { + "schema": { + "type": "object", + "required": ["grant_type", "client_id", "subject_token", "subject_token_type"], + "properties": { + "grant_type": { + "type": "string", + "description": "OAuth 2.0 token exchange grant type.", + "enum": [TOKEN_EXCHANGE_GRANT_TYPE], + }, + "client_id": { + "type": "string", + "description": "Workload token exchange OAuth client ID.", + }, + "subject_token": { + "type": "string", + "description": "JWT subject token to exchange.", + }, + "subject_token_type": { + "type": "string", + "description": "Token type identifier for the subject token.", + "enum": [JWT_TOKEN_TYPE], + }, + "requested_token_type": { + "type": "string", + "description": "Requested token type identifier for the issued token.", + "default": ACCESS_TOKEN_TYPE, + "enum": [ACCESS_TOKEN_TYPE], + }, + "audience": { + "type": "string", + "description": "Requested audience for the issued access token.", + }, + "scope": { + "type": "string", + "description": "Space-separated scopes requested for the issued access token.", + }, + }, + } + } + }, +} + +router = APIRouter(tags=["Workload Identity"]) + + +@dataclass(frozen=True) +class _WorkloadSigningKey: + private_key: rsa.RSAPrivateKey + public_key: rsa.RSAPublicKey + kid: str + + +@dataclass(frozen=True) +class _SubjectJWKSCacheEntry: + jwks: dict[str, Any] + fetched_at: float + + +@dataclass(frozen=True) +class _SubjectTokenDecoder: + name: str + decode: Callable[[], Awaitable[dict[str, Any]]] + + +class WorkloadTokenExchangeResponse(BaseModel): + """RFC 8693 token exchange response for workload identity access tokens.""" + + access_token: str = Field(description="JWT access token minted for the workload identity.") + issued_token_type: str = Field(description="Token type identifier for the issued token.") + token_type: str = Field(description="OAuth token type used in Authorization headers.") + expires_in: int = Field(description="Lifetime of the access token in seconds.") + scope: str | None = Field(default=None, description="Space-separated scopes granted to the access token.") + + +class WorkloadTokenExchangeErrorResponse(BaseModel): + """RFC 8693 token exchange error response.""" + + error: str = Field( + description=( + "OAuth 2.0 or RFC 8693 token exchange error code, such as invalid_client, " + "invalid_request, invalid_grant, invalid_scope, or invalid_target." + ), + ) + error_description: str | None = Field( + default=None, + description="Human-readable ASCII text providing additional information about the error.", + ) + error_uri: str | None = Field( + default=None, + description="URI identifying a human-readable web page with information about the error.", + ) + + +class JsonWebKey(BaseModel): + """JSON Web Key object.""" + + model_config = ConfigDict(extra="allow") + + +class JsonWebKeySetResponse(BaseModel): + """JSON Web Key Set document.""" + + keys: list[JsonWebKey] = Field(description="Public signing keys in the JWKS document.") + + +class ValidationError(BaseModel): + """FastAPI validation error item.""" + + loc: list[str | int] = Field(description="Location of the validation error.") + msg: str = Field(description="Validation error message.") + type: str = Field(description="Validation error type.") + + +class HTTPValidationError(BaseModel): + """FastAPI request validation error response.""" + + detail: list[ValidationError] | None = Field(default=None, description="Validation error details.") + + +_WORKLOAD_TOKEN_EXCHANGE_SERVICE_STATE_KEY = "workload_token_exchange_service" + + +def _platform_base_url_from_request(request: Request | None) -> str: + if request is not None: + return str(request.base_url).rstrip("/") + return get_platform_config().base_url.rstrip("/") + + +def workload_token_endpoint_url(request: Request | None = None) -> str: + """Return the externally reachable workload token exchange endpoint URL.""" + return f"{_platform_base_url_from_request(request)}{WORKLOAD_TOKEN_PATH}" + + +def workload_jwks_url(request: Request | None = None) -> str: + """Return the externally reachable workload exchange JWKS URL.""" + return f"{_platform_base_url_from_request(request)}{WORKLOAD_JWKS_PATH}" + + +def _workload_token_issuer(config: AuthConfig, request: Request | None) -> str: + return config.oidc.workload_token_issuer or f"{_platform_base_url_from_request(request)}/apis/auth" + + +def _workload_private_key_pem(config: AuthConfig) -> bytes: + private_key_file = config.oidc.workload_token_private_key_file + if private_key_file: + try: + return Path(private_key_file).read_bytes() + except OSError as exc: + raise RuntimeError(f"Could not read workload token private key file: {private_key_file}") from exc + + raise RuntimeError("workload_token_private_key_file must be configured for workload token exchange") + + +def _subject_token_key_id(subject_token: str) -> str: + key_id = jwt.get_unverified_header(subject_token).get("kid") + if not key_id: + raise jwt.InvalidTokenError("subject token did not include a signing key id") + return str(key_id) + + +def _find_subject_signing_key(key_id: str, jwks: dict[str, Any]) -> Any | None: + jwk_set = jwt.PyJWKSet.from_dict(jwks) + signing_keys = [jwk for jwk in jwk_set.keys if jwk.public_key_use in ("sig", None) and jwk.key_id] + for jwk in signing_keys: + if jwk.key_id == key_id: + return jwk.key + return None + + +def _validate_subject_jwks(jwks: dict[str, Any]) -> None: + try: + jwt.PyJWKSet.from_dict(jwks) + except (jwt.PyJWTError, TypeError, ValueError, KeyError) as exc: + raise jwt.InvalidTokenError("Subject token JWKS response was not a valid JWKS") from exc + + +class WorkloadTokenExchangeService: + """Stateful helpers for workload token exchange endpoints.""" + + def __init__(self) -> None: + self._workload_signing_key_cache: dict[tuple[str, str], _WorkloadSigningKey] = {} + self._subject_jwks_cache: dict[str, _SubjectJWKSCacheEntry] = {} + + def workload_signing_key(self, config: AuthConfig) -> _WorkloadSigningKey: + kid = config.oidc.workload_token_key_id + if not kid: + raise RuntimeError("workload_token_key_id must be configured for workload token exchange") + + private_key_pem = _workload_private_key_pem(config) + cache_key = (kid, sha256(private_key_pem).hexdigest()) + cached = self._workload_signing_key_cache.get(cache_key) + if cached is not None: + return cached + + private_key = serialization.load_pem_private_key(private_key_pem, password=None) + if not isinstance(private_key, rsa.RSAPrivateKey): + raise RuntimeError("workload token private key must be an RSA private key") + + signing_key = _WorkloadSigningKey(private_key=private_key, public_key=private_key.public_key(), kid=kid) + self._workload_signing_key_cache[cache_key] = signing_key + return signing_key + + def public_jwk(self, config: AuthConfig) -> dict[str, Any]: + signing_key = self.workload_signing_key(config) + jwk = json.loads(RSAAlgorithm.to_jwk(signing_key.public_key)) + jwk.update({"kid": signing_key.kid, "use": "sig", "alg": "RS256"}) + return jwk + + async def fetch_subject_jwks(self, config: AuthConfig, *, refresh: bool = False) -> dict[str, Any]: + jwks_uri = config.oidc.workload_subject_jwks_uri + if not jwks_uri: + raise jwt.InvalidTokenError("JWT subject token validation is disabled") + + now = time.monotonic() + cache_ttl = config.oidc.workload_subject_jwks_cache_ttl_seconds + cached = self._subject_jwks_cache.get(jwks_uri) + if cached is not None and cache_ttl > 0 and not refresh and (now - cached.fetched_at) < cache_ttl: + return cached.jwks + + try: + async with httpx.AsyncClient(timeout=30.0) as client: + response = await client.get(jwks_uri) + response.raise_for_status() + except httpx.HTTPError as exc: + raise jwt.InvalidTokenError(f"Subject token JWKS request failed: {exc}") from exc + + jwks = response.json() + if not isinstance(jwks, dict): + raise jwt.InvalidTokenError("Subject token JWKS response was not an object") + _validate_subject_jwks(jwks) + if cache_ttl > 0: + self._subject_jwks_cache[jwks_uri] = _SubjectJWKSCacheEntry(jwks=jwks, fetched_at=time.monotonic()) + return jwks + + async def get_subject_signing_key(self, config: AuthConfig, subject_token: str) -> Any: + key_id = _subject_token_key_id(subject_token) + jwks = await self.fetch_subject_jwks(config) + signing_key = _find_subject_signing_key(key_id, jwks) + if signing_key is not None: + return signing_key + + jwks = await self.fetch_subject_jwks(config, refresh=True) + signing_key = _find_subject_signing_key(key_id, jwks) + if signing_key is not None: + return signing_key + + raise jwt.InvalidTokenError(f'Unable to find a signing key that matches: "{key_id}"') + + async def decode_jwt_subject_token(self, config: AuthConfig, subject_token: str) -> dict[str, Any]: + signing_key = await self.get_subject_signing_key(config, subject_token) + claims = jwt.decode( + subject_token, + signing_key, + algorithms=["RS256"], + audience=_workload_subject_audience(config), + leeway=30, + options={"require": ["exp"]}, + ) + issuer = claims.get("iss") + if issuer not in _allowed_subject_issuers(config): + raise jwt.InvalidIssuerError(f"unexpected subject token issuer: {issuer!r}") + return claims + + async def decode_subject_token(self, config: AuthConfig, subject_token: str, audience: str) -> dict[str, Any]: + errors: list[str] = [] + for decoder in ( + _SubjectTokenDecoder("JWT subject token", lambda: self.decode_jwt_subject_token(config, subject_token)), + _SubjectTokenDecoder( + "Kubernetes TokenReview subject token", + lambda: _decode_kubernetes_subject_token(config, subject_token, audience), + ), + ): + try: + return await decoder.decode() + except jwt.InvalidTokenError as exc: + message = str(exc) or exc.__class__.__name__ + errors.append(f"{decoder.name}: {message}") + raise jwt.InvalidTokenError("; ".join(errors)) + + +def get_workload_token_exchange_service(request: Request) -> WorkloadTokenExchangeService: + service = getattr(request.app.state, _WORKLOAD_TOKEN_EXCHANGE_SERVICE_STATE_KEY, None) + if isinstance(service, WorkloadTokenExchangeService): + return service + + service = WorkloadTokenExchangeService() + setattr(request.app.state, _WORKLOAD_TOKEN_EXCHANGE_SERVICE_STATE_KEY, service) + return service + + +def _oauth_error(status_code: int, error: str, description: str) -> JSONResponse: + # Keep descriptions passed to clients fixed and non-sensitive. Detailed + # validation errors are logged at call sites instead of returned here. + return JSONResponse( + status_code=status_code, + content={ + "error": error, + "error_description": description, + }, + ) + + +def _workload_client_id(config: AuthConfig) -> str: + return config.oidc.workload_client_id or config.oidc.client_id + + +def _workload_subject_audience(config: AuthConfig) -> str: + return _workload_client_id(config) or DEFAULT_WORKLOAD_AUDIENCE + + +def _default_workload_audience(config: AuthConfig) -> str: + return config.oidc.workload_audience or config.oidc.audience or DEFAULT_WORKLOAD_AUDIENCE + + +def _allowed_audiences(config: AuthConfig) -> set[str]: + return {_default_workload_audience(config), *config.oidc.workload_allowed_audiences} + + +def _validated_audience(config: AuthConfig, requested_audience: Any) -> str: + audience = str(requested_audience or _default_workload_audience(config)) + if audience not in _allowed_audiences(config): + raise jwt.InvalidAudienceError(f"unexpected requested audience: {audience!r}") + return audience + + +def _groups_claim_for_gateway_header(groups: Any) -> str | None: + if isinstance(groups, str): + return groups + if isinstance(groups, list): + return ",".join(str(group).strip() for group in groups if str(group).strip()) + return None + + +def _allowed_subject_issuers(config: AuthConfig) -> set[str]: + return {issuer for issuer in config.oidc.workload_subject_issuers if issuer} + + +def _kubernetes_reviewer_credentials() -> tuple[str, str]: + service_account_dir = Path("/var/run/secrets/kubernetes.io/serviceaccount") + reviewer_token = (service_account_dir / "token").read_text(encoding="utf-8").strip() + ca_path = service_account_dir / "ca.crt" + return reviewer_token, str(ca_path) + + +async def _decode_kubernetes_subject_token( + config: AuthConfig, + subject_token: str, + audience: str, +) -> dict[str, Any]: + if not config.oidc.workload_kubernetes_token_review_enabled: + raise jwt.InvalidTokenError("Kubernetes TokenReview subject token validation is disabled") + + import os + + host = os.environ.get("KUBERNETES_SERVICE_HOST") + port = os.environ.get("KUBERNETES_SERVICE_PORT", "443") + if not host: + raise jwt.InvalidTokenError("Kubernetes service environment is unavailable") + + reviewer_token, ca_path = _kubernetes_reviewer_credentials() + token_review_url = f"https://{host}:{port}/apis/authentication.k8s.io/v1/tokenreviews" + try: + async with httpx.AsyncClient(timeout=10.0, verify=ca_path) as client: + response = await client.post( + token_review_url, + json={ + "apiVersion": "authentication.k8s.io/v1", + "kind": "TokenReview", + "spec": { + "token": subject_token, + "audiences": [audience], + }, + }, + headers={ + "Authorization": f"Bearer {reviewer_token}", + "Content-Type": "application/json", + }, + ) + response.raise_for_status() + except httpx.HTTPError as exc: + raise jwt.InvalidTokenError(f"Kubernetes TokenReview request failed: {exc}") from exc + + token_review = response.json() + if not isinstance(token_review, dict): + raise jwt.InvalidTokenError("Kubernetes TokenReview response was not an object") + + status = token_review.get("status", {}) + if not isinstance(status, dict): + raise jwt.InvalidTokenError("Kubernetes TokenReview response did not include a status object") + if not status.get("authenticated"): + raise jwt.InvalidTokenError(status.get("error") or "Kubernetes TokenReview rejected subject token") + + user = status.get("user", {}) + if not isinstance(user, dict): + raise jwt.InvalidTokenError("Kubernetes TokenReview response did not include a user object") + subject = user.get("username") + if not subject: + raise jwt.InvalidTokenError("Kubernetes TokenReview response did not include a username") + + return { + "sub": subject, + "groups": user.get("groups", []), + } + + +@router.get( + "/jwks", + summary="Workload identity token exchange JWKS", + description="Return the public signing key for workload identity access tokens minted by the NeMo auth service.", + response_model=JsonWebKeySetResponse, +) +async def jwks( + workload_token_exchange_service: WorkloadTokenExchangeService = Depends(get_workload_token_exchange_service), +) -> JsonWebKeySetResponse: + """Return workload identity exchange signing keys.""" + return JsonWebKeySetResponse(keys=[JsonWebKey(**workload_token_exchange_service.public_jwk(get_auth_config()))]) + + +@router.post( + "/token", + summary="Exchange a workload identity subject token", + description="Exchange a configured workload identity subject token for a NeMo Platform access token.", + response_model=WorkloadTokenExchangeResponse, + responses={ + 400: { + "description": "RFC 8693 token exchange error", + "model": WorkloadTokenExchangeErrorResponse, + }, + 401: { + "description": "OAuth 2.0 invalid_client error", + "model": WorkloadTokenExchangeErrorResponse, + }, + }, + openapi_extra={"requestBody": _TOKEN_EXCHANGE_FORM_REQUEST_BODY}, +) +async def token_exchange( + request: Request, + workload_token_exchange_service: WorkloadTokenExchangeService = Depends(get_workload_token_exchange_service), +) -> WorkloadTokenExchangeResponse | JSONResponse: + """Exchange an RFC 8693 workload identity subject token for a NeMo access token.""" + config = get_auth_config() + if not config.oidc.workload_token_exchange_enabled: + return _oauth_error(400, "invalid_request", "Workload token exchange is not enabled") + + form = await request.form() + client_id = _workload_client_id(config) + + if form.get("grant_type") != TOKEN_EXCHANGE_GRANT_TYPE: + return _oauth_error(400, "unsupported_grant_type", "Only RFC 8693 token exchange is supported") + if form.get("client_id") != client_id: + return _oauth_error(401, "invalid_client", "Unknown workload token exchange client") + if form.get("subject_token_type") != JWT_TOKEN_TYPE: + return _oauth_error(400, "invalid_request", "subject_token_type must be a JWT token type") + if form.get("requested_token_type", ACCESS_TOKEN_TYPE) != ACCESS_TOKEN_TYPE: + return _oauth_error(400, "invalid_request", "requested_token_type must be access_token") + + subject_token = form.get("subject_token") + if not subject_token: + return _oauth_error(400, "invalid_request", "subject_token is required") + + try: + audience = _validated_audience(config, form.get("audience")) + except jwt.InvalidAudienceError as exc: + logger.info("Requested token audience validation failed: %s", exc) + return _oauth_error(400, "invalid_target", "Requested audience is not allowed") + + try: + subject_claims = await workload_token_exchange_service.decode_subject_token( + config, str(subject_token), _workload_subject_audience(config) + ) + subject = subject_claims.get("sub") + if not subject: + raise jwt.InvalidTokenError("Subject token did not include a subject") + except jwt.InvalidTokenError as exc: + logger.info("Subject token validation failed: %s", exc) + # RFC 8693 clients only need a stable invalid_request response. Avoid + # returning decoder details that may include infrastructure internals. + return _oauth_error(400, "invalid_request", "Could not validate subject token") + + now = int(time.time()) + scope = str(form.get("scope") or config.oidc.workload_scope or DEFAULT_WORKLOAD_SCOPE) + exchanged_claims: dict[str, Any] = { + "iss": _workload_token_issuer(config, request), + "sub": subject, + "aud": audience, + "iat": now, + "nbf": now, + "exp": now + config.oidc.workload_token_ttl_seconds, + "scope": scope, + } + if "email" in subject_claims: + exchanged_claims["email"] = subject_claims["email"] + if "groups" in subject_claims: + groups_claim = _groups_claim_for_gateway_header(subject_claims["groups"]) + if groups_claim: + exchanged_claims["groups"] = groups_claim + + signing_key = workload_token_exchange_service.workload_signing_key(config) + access_token = jwt.encode( + exchanged_claims, + signing_key.private_key, + algorithm="RS256", + headers={"kid": signing_key.kid}, + ) + return WorkloadTokenExchangeResponse( + access_token=access_token, + issued_token_type=ACCESS_TOKEN_TYPE, + token_type="Bearer", + expires_in=config.oidc.workload_token_ttl_seconds, + scope=scope, + ) diff --git a/services/core/auth/src/nmp/core/auth/assets/static-authz.yaml b/services/core/auth/src/nmp/core/auth/assets/static-authz.yaml index f00a6adb25..7df668f042 100644 --- a/services/core/auth/src/nmp/core/auth/assets/static-authz.yaml +++ b/services/core/auth/src/nmp/core/auth/assets/static-authz.yaml @@ -427,6 +427,14 @@ authz: get: permissions: [] scopes: [] + /apis/auth/jwks: + get: + permissions: [] + scopes: [] + /apis/auth/token: + post: + permissions: [] + scopes: [] /apis/auth/v2/iam/opa-bundle.tar.gz: get: permissions: diff --git a/services/core/auth/src/nmp/core/auth/service.py b/services/core/auth/src/nmp/core/auth/service.py index 1af59512cc..2140fe38c5 100644 --- a/services/core/auth/src/nmp/core/auth/service.py +++ b/services/core/auth/src/nmp/core/auth/service.py @@ -9,6 +9,7 @@ from nmp.common.config import get_service_config from nmp.common.service import RouterConfig, Service +from nmp.core.auth.api.v2 import workload_token_exchange from nmp.core.auth.api.v2.bundle import endpoints as bundle from nmp.core.auth.api.v2.discovery import endpoints as discovery from nmp.core.auth.api.v2.iam import endpoints as iam @@ -37,6 +38,11 @@ def get_routers(self) -> List[RouterConfig]: RouterConfig(iam.router, tag="IAM", description="Identity and Access Management endpoints"), RouterConfig(bundle.router, tag="Bundle", description="OPA bundle endpoints"), RouterConfig(discovery.router, tag="Discovery", description="Platform configuration discovery endpoints"), + RouterConfig( + workload_token_exchange.router, + tag="Workload Identity", + description="Workload identity token exchange endpoints", + ), ] # Only register authz routes when auth is enabled and using embedded policy engine diff --git a/services/core/auth/tests/test_discovery.py b/services/core/auth/tests/test_discovery.py index 2cd8294a57..c5119f74d2 100644 --- a/services/core/auth/tests/test_discovery.py +++ b/services/core/auth/tests/test_discovery.py @@ -36,6 +36,11 @@ def oidc_config(): authorization_endpoint="https://sso.example.com/authorize", token_endpoint="https://sso.example.com/token", device_authorization_endpoint="https://sso.example.com/device/code", + workload_token_exchange_enabled=True, + workload_client_id="test-workload-client", + workload_token_endpoint="https://workload-idp.example.com/token", + workload_audience="nemo-platform", + workload_scope="openid email groups", ) @@ -103,6 +108,25 @@ def test_oidc_discovery_response_optional_fields(self): assert response.device_authorization_endpoint is None assert response.userinfo_endpoint is None + def test_oidc_discovery_response_includes_workload_exchange_fields(self): + """Test OIDCDiscoveryResponse includes workload identity token exchange fields.""" + response = OIDCDiscoveryResponse( + issuer="https://sso.example.com", + client_id="test-client", + token_endpoint="https://sso.example.com/token", + workload_token_exchange_enabled=True, + workload_client_id="test-workload-client", + workload_token_endpoint="https://workload-idp.example.com/token", + workload_audience="nemo-platform", + workload_scope="openid email groups", + ) + + assert response.workload_token_exchange_enabled is True + assert response.workload_client_id == "test-workload-client" + assert response.workload_token_endpoint == "https://workload-idp.example.com/token" + assert response.workload_audience == "nemo-platform" + assert response.workload_scope == "openid email groups" + class TestAuthDiscoveryResponse: """Tests for AuthDiscoveryResponse model.""" @@ -152,6 +176,11 @@ async def test_auth_enabled_oidc_enabled_with_configured_endpoints(self, auth_co assert result.oidc.authorization_endpoint == "https://sso.example.com/authorize" assert result.oidc.token_endpoint == "https://sso.example.com/token" assert result.oidc.device_authorization_endpoint == "https://sso.example.com/device/code" + assert result.oidc.workload_token_exchange_enabled is True + assert result.oidc.workload_client_id == "test-workload-client" + assert result.oidc.workload_token_endpoint == "https://workload-idp.example.com/token" + assert result.oidc.workload_audience == "nemo-platform" + assert result.oidc.workload_scope == "openid email groups" finally: Configuration.clear_overrides() @@ -168,6 +197,31 @@ async def test_auth_enabled_oidc_disabled(self, auth_config_oidc_disabled): finally: Configuration.clear_overrides() + @pytest.mark.asyncio + async def test_workload_token_endpoint_defaults_to_platform_auth_endpoint(self): + """Test workload exchange endpoint defaults to the NeMo auth service.""" + oidc_config = OIDCConfig( + enabled=True, + issuer="https://sso.example.com", + client_id="test-client", + workload_token_exchange_enabled=True, + workload_client_id="test-workload-client", + ) + auth_config = AuthConfig( + enabled=True, + policy_decision_point_base_url="http://localhost:8181", + oidc=oidc_config, + ) + Configuration.set_override(auth_config) + + try: + result = await get_auth_discovery() + + assert result.oidc is not None + assert result.oidc.workload_token_endpoint == "http://localhost:8080/apis/auth/token" + finally: + Configuration.clear_overrides() + @pytest.mark.asyncio async def test_auth_disabled(self, auth_config_disabled): """Test response when auth is disabled.""" diff --git a/services/core/auth/tests/test_workload_token_exchange.py b/services/core/auth/tests/test_workload_token_exchange.py new file mode 100644 index 0000000000..db587f9a15 --- /dev/null +++ b/services/core/auth/tests/test_workload_token_exchange.py @@ -0,0 +1,584 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import json +from typing import Any + +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from fastapi import FastAPI +from fastapi.testclient import TestClient +from nmp.common.config import AuthConfig, Configuration +from nmp.common.config.base import OIDCConfig +from nmp.core.auth.api.v2 import workload_token_exchange as exchange + + +@pytest.fixture(autouse=True) +def _clear_config_overrides(): + yield + Configuration.clear_overrides() + + +@pytest.fixture +def exchange_service() -> exchange.WorkloadTokenExchangeService: + return exchange.WorkloadTokenExchangeService() + + +@pytest.fixture +def workload_signing_key() -> rsa.RSAPrivateKey: + return rsa.generate_private_key(public_exponent=65537, key_size=2048) + + +def _private_key_pem(private_key: rsa.RSAPrivateKey) -> str: + return private_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ).decode() + + +@pytest.fixture +def exchange_config(workload_signing_key: rsa.RSAPrivateKey, tmp_path) -> AuthConfig: + private_key_file = tmp_path / "workload-token-private-key.pem" + private_key_file.write_text(_private_key_pem(workload_signing_key), encoding="utf-8") + return AuthConfig( + enabled=True, + oidc=OIDCConfig( + enabled=True, + issuer="https://idp.example.com/application/o/nemo-cli/", + additional_issuers=["https://idp.example.com/application/o/nemo/"], + client_id="nemo-platform-cli", + workload_token_exchange_enabled=True, + workload_client_id="nemo-platform-workload", + workload_audience="nemo-platform", + workload_scope="openid email groups", + workload_subject_jwks_uri="https://idp.example.com/application/o/nemo-workload/jwks/", + workload_subject_issuers=["https://idp.example.com/application/o/nemo-workload/"], + workload_token_private_key_file=str(private_key_file), + ), + ) + + +@pytest.fixture +def client(exchange_config: AuthConfig, exchange_service: exchange.WorkloadTokenExchangeService) -> TestClient: + Configuration.set_override(exchange_config) + app = FastAPI() + app.dependency_overrides[exchange.get_workload_token_exchange_service] = lambda: exchange_service + app.include_router(exchange.router) + return TestClient(app, raise_server_exceptions=False) + + +def test_jwks_publishes_workload_exchange_signing_key(client: TestClient) -> None: + response = client.get("/jwks") + + assert response.status_code == 200 + keys = response.json()["keys"] + assert keys[0]["kid"] == "nemo-workload-exchange" + assert keys[0]["use"] == "sig" + assert keys[0]["alg"] == "RS256" + + +def test_jwks_openapi_documents_jwks_response(client: TestClient) -> None: + openapi = client.app.openapi() + operation = openapi["paths"]["/jwks"]["get"] + + assert operation["responses"]["200"]["content"]["application/json"]["schema"] == { + "$ref": "#/components/schemas/JsonWebKeySetResponse" + } + jwks_schema = openapi["components"]["schemas"]["JsonWebKeySetResponse"] + assert jwks_schema["required"] == ["keys"] + assert jwks_schema["properties"]["keys"]["type"] == "array" + assert jwks_schema["properties"]["keys"]["items"] == {"$ref": "#/components/schemas/JsonWebKey"} + assert openapi["components"]["schemas"]["JsonWebKey"] == { + "additionalProperties": True, + "properties": {}, + "type": "object", + "title": "JsonWebKey", + "description": "JSON Web Key object.", + } + + +def test_token_exchange_openapi_documents_form_request_and_token_response(client: TestClient) -> None: + openapi = client.app.openapi() + operation = openapi["paths"]["/token"]["post"] + + request_schema = operation["requestBody"]["content"]["application/x-www-form-urlencoded"]["schema"] + if "$ref" in request_schema: + request_schema = openapi["components"]["schemas"][request_schema["$ref"].rsplit("/", 1)[-1]] + + assert { + "grant_type", + "client_id", + "subject_token", + "subject_token_type", + "requested_token_type", + "audience", + "scope", + } <= set(request_schema["properties"]) + assert set(operation["responses"]) == {"200", "400", "401"} + assert operation["responses"]["200"]["description"] == "Successful Response" + assert operation["responses"]["200"]["content"]["application/json"]["schema"] == { + "$ref": "#/components/schemas/WorkloadTokenExchangeResponse" + } + assert operation["responses"]["400"]["description"] == "RFC 8693 token exchange error" + assert operation["responses"]["400"]["content"]["application/json"]["schema"] == { + "$ref": "#/components/schemas/WorkloadTokenExchangeErrorResponse" + } + assert operation["responses"]["401"]["description"] == "OAuth 2.0 invalid_client error" + assert operation["responses"]["401"]["content"]["application/json"]["schema"] == { + "$ref": "#/components/schemas/WorkloadTokenExchangeErrorResponse" + } + error_schema = openapi["components"]["schemas"]["WorkloadTokenExchangeErrorResponse"] + assert error_schema["properties"]["error"]["type"] == "string" + error_description = error_schema["properties"]["error"]["description"] + for error_code in ("invalid_client", "invalid_request", "invalid_grant", "invalid_scope", "invalid_target"): + assert error_code in error_description + response_schema = openapi["components"]["schemas"]["WorkloadTokenExchangeResponse"] + assert { + "access_token", + "issued_token_type", + "token_type", + "expires_in", + "scope", + } == set(response_schema["properties"]) + assert response_schema["required"] == [ + "access_token", + "issued_token_type", + "token_type", + "expires_in", + ] + + +def test_token_exchange_rejects_subject_token_without_subject( + client: TestClient, + exchange_service: exchange.WorkloadTokenExchangeService, + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def decode_subject_token(config: AuthConfig, subject_token: str, audience: str) -> dict[str, str]: + return {"email": "svc@example.com"} + + monkeypatch.setattr(exchange_service, "decode_subject_token", decode_subject_token) + + response = client.post( + "/token", + data={ + "grant_type": exchange.TOKEN_EXCHANGE_GRANT_TYPE, + "client_id": "nemo-platform-workload", + "subject_token": "subject-token", + "subject_token_type": exchange.JWT_TOKEN_TYPE, + }, + ) + + assert response.status_code == 400 + assert response.json() == { + "error": "invalid_request", + "error_description": "Could not validate subject token", + } + + +def test_token_exchange_rejects_disallowed_audience_before_subject_validation( + client: TestClient, + exchange_service: exchange.WorkloadTokenExchangeService, + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def decode_subject_token(config: AuthConfig, subject_token: str, audience: str) -> dict[str, str]: + raise AssertionError("disallowed audience should be rejected before subject token validation") + + monkeypatch.setattr(exchange_service, "decode_subject_token", decode_subject_token) + + response = client.post( + "/token", + data={ + "grant_type": exchange.TOKEN_EXCHANGE_GRANT_TYPE, + "client_id": "nemo-platform-workload", + "subject_token": "subject-token", + "subject_token_type": exchange.JWT_TOKEN_TYPE, + "audience": "unexpected-audience", + }, + ) + + assert response.status_code == 400 + assert response.json() == { + "error": "invalid_target", + "error_description": "Requested audience is not allowed", + } + + +def test_token_exchange_mints_access_token_signed_by_configured_key( + client: TestClient, + exchange_config: AuthConfig, + exchange_service: exchange.WorkloadTokenExchangeService, + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, Any] = {} + + async def decode_subject_token(config: AuthConfig, subject_token: str, audience: str) -> dict[str, Any]: + captured["subject_audience"] = audience + return { + "sub": "workload-subject", + "email": "svc@example.com", + "groups": ["svc-group", "system:serviceaccounts"], + } + + monkeypatch.setattr(exchange_service, "decode_subject_token", decode_subject_token) + + response = client.post( + "/token", + data={ + "grant_type": exchange.TOKEN_EXCHANGE_GRANT_TYPE, + "client_id": "nemo-platform-workload", + "subject_token": "subject-token", + "subject_token_type": exchange.JWT_TOKEN_TYPE, + }, + ) + + assert response.status_code == 200 + access_token = response.json()["access_token"] + signing_key = exchange_service.workload_signing_key(exchange_config) + assert exchange.jwt.get_unverified_header(access_token)["kid"] == signing_key.kid + claims = exchange.jwt.decode(access_token, signing_key.public_key, algorithms=["RS256"], audience="nemo-platform") + assert captured["subject_audience"] == "nemo-platform-workload" + assert claims["sub"] == "workload-subject" + assert claims["email"] == "svc@example.com" + assert claims["groups"] == "svc-group,system:serviceaccounts" + + +def test_validated_audience_accepts_configured_allowlist(exchange_config: AuthConfig) -> None: + exchange_config.oidc.workload_allowed_audiences.append("extra-audience") + + assert exchange._validated_audience(exchange_config, "extra-audience") == "extra-audience" + + +class _FakeResponse: + def __init__(self, payload: dict[str, Any]) -> None: + self._payload = payload + + def raise_for_status(self) -> None: + return None + + def json(self) -> dict[str, Any]: + return self._payload + + +def _signed_subject_token( + config: AuthConfig, + exchange_service: exchange.WorkloadTokenExchangeService, + *, + audience: str, + issuer: str | None = None, + private_key: rsa.RSAPrivateKey | None = None, + key_id: str | None = None, +) -> str: + token_issuer = issuer or config.oidc.workload_subject_issuers[0] + signing_key = exchange_service.workload_signing_key(config) + return exchange.jwt.encode( + { + "iss": token_issuer, + "sub": "authentik-user", + "aud": audience, + "exp": int(exchange.time.time()) + 300, + }, + private_key or signing_key.private_key, + algorithm="RS256", + headers={"kid": key_id or signing_key.kid}, + ) + + +def _public_jwk_for_key(private_key: rsa.RSAPrivateKey, *, key_id: str) -> dict[str, Any]: + jwk = json.loads(exchange.RSAAlgorithm.to_jwk(private_key.public_key())) + jwk.update({"kid": key_id, "use": "sig", "alg": "RS256"}) + return jwk + + +def _mock_subject_jwks_client( + config: AuthConfig, + exchange_service: exchange.WorkloadTokenExchangeService, + monkeypatch: pytest.MonkeyPatch, +) -> dict[str, Any]: + captured: dict[str, Any] = {"request_count": 0} + + class FakeAsyncClient: + def __init__(self, *, timeout: float) -> None: + captured["timeout"] = timeout + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, traceback) -> None: + return None + + async def get(self, url: str) -> _FakeResponse: + captured["request_count"] += 1 + captured["url"] = url + return _FakeResponse({"keys": [exchange_service.public_jwk(config)]}) + + monkeypatch.setattr(exchange.httpx, "AsyncClient", FakeAsyncClient) + return captured + + +def test_jwt_subject_token_decoder_fetches_configured_jwks( + exchange_config: AuthConfig, + exchange_service: exchange.WorkloadTokenExchangeService, + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured = _mock_subject_jwks_client(exchange_config, exchange_service, monkeypatch) + subject_token = _signed_subject_token(exchange_config, exchange_service, audience="nemo-platform-workload") + + claims = asyncio.run(exchange_service.decode_jwt_subject_token(exchange_config, subject_token)) + + assert claims["sub"] == "authentik-user" + assert captured == { + "request_count": 1, + "timeout": 30.0, + "url": "https://idp.example.com/application/o/nemo-workload/jwks/", + } + + +def test_jwt_subject_token_decoder_caches_configured_jwks( + exchange_config: AuthConfig, + exchange_service: exchange.WorkloadTokenExchangeService, + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured = _mock_subject_jwks_client(exchange_config, exchange_service, monkeypatch) + subject_token = _signed_subject_token(exchange_config, exchange_service, audience="nemo-platform-workload") + + asyncio.run(exchange_service.decode_jwt_subject_token(exchange_config, subject_token)) + asyncio.run(exchange_service.decode_jwt_subject_token(exchange_config, subject_token)) + + assert captured["request_count"] == 1 + + +def test_jwt_subject_token_decoder_does_not_cache_invalid_jwks( + exchange_config: AuthConfig, + exchange_service: exchange.WorkloadTokenExchangeService, + monkeypatch: pytest.MonkeyPatch, +) -> None: + jwks_responses = [ + {}, + {"keys": [exchange_service.public_jwk(exchange_config)]}, + ] + captured: dict[str, Any] = {"request_count": 0} + + class FakeAsyncClient: + def __init__(self, *, timeout: float) -> None: + captured["timeout"] = timeout + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, traceback) -> None: + return None + + async def get(self, url: str) -> _FakeResponse: + captured["request_count"] += 1 + captured["url"] = url + return _FakeResponse(jwks_responses.pop(0)) + + monkeypatch.setattr(exchange.httpx, "AsyncClient", FakeAsyncClient) + subject_token = _signed_subject_token(exchange_config, exchange_service, audience="nemo-platform-workload") + + with pytest.raises(exchange.jwt.InvalidTokenError, match="valid JWKS"): + asyncio.run(exchange_service.decode_jwt_subject_token(exchange_config, subject_token)) + + claims = asyncio.run(exchange_service.decode_jwt_subject_token(exchange_config, subject_token)) + + assert claims["sub"] == "authentik-user" + assert captured["request_count"] == 2 + + +def test_jwt_subject_token_decoder_refreshes_cached_jwks_on_unknown_kid( + exchange_config: AuthConfig, + exchange_service: exchange.WorkloadTokenExchangeService, + monkeypatch: pytest.MonkeyPatch, +) -> None: + rotated_key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + subject_token = _signed_subject_token( + exchange_config, + exchange_service, + audience="nemo-platform-workload", + private_key=rotated_key, + key_id="rotated-key", + ) + jwks_responses = [ + {"keys": [exchange_service.public_jwk(exchange_config)]}, + {"keys": [_public_jwk_for_key(rotated_key, key_id="rotated-key")]}, + ] + captured: dict[str, Any] = {"request_count": 0} + + class FakeAsyncClient: + def __init__(self, *, timeout: float) -> None: + captured["timeout"] = timeout + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, traceback) -> None: + return None + + async def get(self, url: str) -> _FakeResponse: + captured["request_count"] += 1 + captured["url"] = url + return _FakeResponse(jwks_responses.pop(0)) + + monkeypatch.setattr(exchange.httpx, "AsyncClient", FakeAsyncClient) + + claims = asyncio.run(exchange_service.decode_jwt_subject_token(exchange_config, subject_token)) + + assert claims["sub"] == "authentik-user" + assert captured["request_count"] == 2 + + +def test_jwt_subject_token_decoder_rejects_wrong_audience( + exchange_config: AuthConfig, + exchange_service: exchange.WorkloadTokenExchangeService, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _mock_subject_jwks_client(exchange_config, exchange_service, monkeypatch) + subject_token = _signed_subject_token(exchange_config, exchange_service, audience="some-other-client") + + with pytest.raises(exchange.jwt.InvalidAudienceError): + asyncio.run(exchange_service.decode_jwt_subject_token(exchange_config, subject_token)) + + +def test_jwt_subject_token_decoder_rejects_unexpected_issuer( + exchange_config: AuthConfig, + exchange_service: exchange.WorkloadTokenExchangeService, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _mock_subject_jwks_client(exchange_config, exchange_service, monkeypatch) + subject_token = _signed_subject_token( + exchange_config, + exchange_service, + audience="nemo-platform-workload", + issuer="https://idp.example.com/application/o/other/", + ) + + with pytest.raises(exchange.jwt.InvalidIssuerError): + asyncio.run(exchange_service.decode_jwt_subject_token(exchange_config, subject_token)) + + +def test_jwt_subject_token_decoder_requires_explicit_workload_subject_issuers( + exchange_config: AuthConfig, + exchange_service: exchange.WorkloadTokenExchangeService, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _mock_subject_jwks_client(exchange_config, exchange_service, monkeypatch) + exchange_config.oidc.workload_subject_issuers = [] + subject_token = _signed_subject_token( + exchange_config, + exchange_service, + audience="nemo-platform-workload", + issuer=exchange_config.oidc.additional_issuers[0], + ) + + with pytest.raises(exchange.jwt.InvalidIssuerError): + asyncio.run(exchange_service.decode_jwt_subject_token(exchange_config, subject_token)) + + +def test_subject_token_decoder_reports_all_validation_failures( + exchange_config: AuthConfig, + exchange_service: exchange.WorkloadTokenExchangeService, + monkeypatch: pytest.MonkeyPatch, +) -> None: + subject_token = _signed_subject_token(exchange_config, exchange_service, audience="nemo-platform-workload") + exchange_config.oidc.workload_subject_jwks_uri = None + exchange_config.oidc.workload_kubernetes_token_review_enabled = True + monkeypatch.delenv("KUBERNETES_SERVICE_HOST", raising=False) + + with pytest.raises(exchange.jwt.InvalidTokenError) as exc_info: + asyncio.run(exchange_service.decode_subject_token(exchange_config, subject_token, "nemo-platform")) + + assert str(exc_info.value) == ( + "JWT subject token: JWT subject token validation is disabled; " + "Kubernetes TokenReview subject token: Kubernetes service environment is unavailable" + ) + + +def test_subject_token_decoder_does_not_mask_unexpected_decoder_errors( + exchange_config: AuthConfig, + exchange_service: exchange.WorkloadTokenExchangeService, + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def reject_jwt_subject_token(config: AuthConfig, subject_token: str) -> dict[str, Any]: + raise exchange.jwt.InvalidTokenError("JWT rejected") + + async def broken_kubernetes_subject_token( + config: AuthConfig, + subject_token: str, + audience: str, + ) -> dict[str, Any]: + raise RuntimeError("Kubernetes decoder broke") + + monkeypatch.setattr(exchange_service, "decode_jwt_subject_token", reject_jwt_subject_token) + monkeypatch.setattr(exchange, "_decode_kubernetes_subject_token", broken_kubernetes_subject_token) + + with pytest.raises(RuntimeError, match="Kubernetes decoder broke"): + asyncio.run(exchange_service.decode_subject_token(exchange_config, "subject-token", "nemo-platform")) + + +def test_kubernetes_subject_token_decoder_posts_token_review_with_async_client( + exchange_config: AuthConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, Any] = {} + exchange_config.oidc.workload_kubernetes_token_review_enabled = True + monkeypatch.setenv("KUBERNETES_SERVICE_HOST", "kubernetes.default.svc") + monkeypatch.setenv("KUBERNETES_SERVICE_PORT", "443") + monkeypatch.setattr(exchange, "_kubernetes_reviewer_credentials", lambda: ("reviewer-token", "/tmp/ca.crt")) + + class FakeAsyncClient: + def __init__(self, *, timeout: float, verify: str) -> None: + captured["timeout"] = timeout + captured["verify"] = verify + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, traceback) -> None: + return None + + async def post(self, url: str, **kwargs) -> _FakeResponse: + captured["url"] = url + captured["json"] = kwargs["json"] + captured["headers"] = kwargs["headers"] + return _FakeResponse( + { + "status": { + "authenticated": True, + "user": { + "username": "system:serviceaccount:default:nemo", + "groups": ["system:serviceaccounts"], + }, + } + } + ) + + monkeypatch.setattr(exchange.httpx, "AsyncClient", FakeAsyncClient) + + claims = asyncio.run( + exchange._decode_kubernetes_subject_token(exchange_config, "kubernetes-subject-token", "nemo-platform") + ) + + assert claims == { + "sub": "system:serviceaccount:default:nemo", + "groups": ["system:serviceaccounts"], + } + assert captured == { + "timeout": 10.0, + "verify": "/tmp/ca.crt", + "url": "https://kubernetes.default.svc:443/apis/authentication.k8s.io/v1/tokenreviews", + "json": { + "apiVersion": "authentication.k8s.io/v1", + "kind": "TokenReview", + "spec": { + "token": "kubernetes-subject-token", + "audiences": ["nemo-platform"], + }, + }, + "headers": { + "Authorization": "Bearer reviewer-token", + "Content-Type": "application/json", + }, + } diff --git a/services/core/jobs/jobs-launcher/cmd/otel.go b/services/core/jobs/jobs-launcher/cmd/otel.go index 72de2883f4..92b03cdabc 100644 --- a/services/core/jobs/jobs-launcher/cmd/otel.go +++ b/services/core/jobs/jobs-launcher/cmd/otel.go @@ -6,12 +6,17 @@ package cmd import ( "context" "errors" + "fmt" "log/slog" + "net/http" "os" + "sync/atomic" + "time" "go.opentelemetry.io/contrib/bridges/otelslog" "go.opentelemetry.io/contrib/exporters/autoexport" "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploghttp" "go.opentelemetry.io/otel/exporters/stdout/stdoutlog" "go.opentelemetry.io/otel/log/global" "go.opentelemetry.io/otel/sdk/log" @@ -19,12 +24,14 @@ import ( ) const ( - name = "nmp.nvidia.com/nemo-platform/jobs-launcher" - NEMO_JOB_WORKSPACE = "NEMO_JOB_WORKSPACE" - NEMO_JOB_ID_ENV = "NEMO_JOB_ID" - NEMO_JOB_ATTEMPT_ID_ENV = "NEMO_JOB_ATTEMPT_ID" - NEMO_JOB_STEP_NAME_ENV = "NEMO_JOB_STEP" - NEMO_JOB_TASK_ID_ENV = "NEMO_JOB_TASK" + name = "nmp.nvidia.com/nemo-platform/jobs-launcher" + NEMO_JOB_WORKSPACE = "NEMO_JOB_WORKSPACE" + NEMO_JOB_ID_ENV = "NEMO_JOB_ID" + NEMO_JOB_ATTEMPT_ID_ENV = "NEMO_JOB_ATTEMPT_ID" + NEMO_JOB_STEP_NAME_ENV = "NEMO_JOB_STEP" + NEMO_JOB_TASK_ID_ENV = "NEMO_JOB_TASK" + nmpJobLogsEndpointEnv = "NMP_JOB_LOGS_ENDPOINT" + otlpHTTPLogExportTimeout = 10 * time.Second ) var ( @@ -75,7 +82,7 @@ func setupOTELSDK(ctx context.Context) (func(context.Context) error, *log.Logger } // Set up logger provider. - loggerProvider, err := newLoggerProvider(res) + loggerProvider, err := newLoggerProvider(ctx, res) if err != nil { handleErr(err) return shutdown, nil, err @@ -90,9 +97,33 @@ func setupOTELSDK(ctx context.Context) (func(context.Context) error, *log.Logger } // newLoggerProvider creates a new OTEL logger provider with the given resource. -func newLoggerProvider(res *resource.Resource) (*log.LoggerProvider, error) { - logExporter, err := autoexport.NewLogExporter( - context.Background(), +func newLoggerProvider(ctx context.Context, res *resource.Resource) (*log.LoggerProvider, error) { + logExporter, err := newLogExporter(ctx) + if err != nil { + return nil, err + } + + loggerProvider := log.NewLoggerProvider( + log.WithProcessor(log.NewBatchProcessor(logExporter)), + log.WithResource(res), + ) + return loggerProvider, nil +} + +func newLogExporter(ctx context.Context) (log.Exporter, error) { + if endpoint := os.Getenv(nmpJobLogsEndpointEnv); endpoint != "" { + if os.Getenv(workloadIdentityTokenFileEnv) != "" { + tokenSource, err := newOTLPLogWorkloadAuthTokenSource(ctx) + if err != nil { + return nil, fmt.Errorf("configure workload identity auth for OTLP logs: %w", err) + } + return newRefreshableAuthLogExporter(ctx, endpoint, tokenSource, otlpHTTPLogExporter) + } + return otlploghttp.New(ctx, otlploghttp.WithEndpointURL(endpoint)) + } + + return autoexport.NewLogExporter( + ctx, // Default to a stdout log exporter if autoexport fails to configure one. autoexport.WithFallbackLogExporter( func(ctx context.Context) (log.Exporter, error) { @@ -100,13 +131,86 @@ func newLoggerProvider(res *resource.Resource) (*log.LoggerProvider, error) { }, ), ) +} + +type authHeaderSource interface { + AuthorizationHeader(context.Context) (string, error) +} + +type logExporterFactory func(context.Context, ...otlploghttp.Option) (log.Exporter, error) + +type refreshableAuthLogExporter struct { + exporter log.Exporter + stopped atomic.Bool +} + +func newRefreshableAuthLogExporter( + ctx context.Context, + endpoint string, + authSource authHeaderSource, + newExporter logExporterFactory, +) (log.Exporter, error) { + if _, err := authSource.AuthorizationHeader(ctx); err != nil { + return nil, fmt.Errorf("configure workload identity auth for OTLP logs: %w", err) + } + + exporter, err := newExporter( + ctx, + otlploghttp.WithEndpointURL(endpoint), + otlploghttp.WithHTTPClient(&http.Client{ + Transport: &authHeaderTransport{source: authSource}, + Timeout: otlpHTTPLogExportTimeout, + }), + ) if err != nil { return nil, err } + return &refreshableAuthLogExporter{exporter: exporter}, nil +} - loggerProvider := log.NewLoggerProvider( - log.WithProcessor(log.NewBatchProcessor(logExporter)), - log.WithResource(res), - ) - return loggerProvider, nil +func (e *refreshableAuthLogExporter) Export(ctx context.Context, records []log.Record) error { + if e.stopped.Load() { + return nil + } + + return e.exporter.Export(ctx, records) +} + +func (e *refreshableAuthLogExporter) Shutdown(ctx context.Context) error { + if e.stopped.Swap(true) { + return nil + } + return e.exporter.Shutdown(ctx) +} + +func (e *refreshableAuthLogExporter) ForceFlush(ctx context.Context) error { + return e.exporter.ForceFlush(ctx) +} + +type authHeaderTransport struct { + source authHeaderSource + base http.RoundTripper +} + +func (t *authHeaderTransport) RoundTrip(req *http.Request) (*http.Response, error) { + authHeader, err := t.source.AuthorizationHeader(req.Context()) + if err != nil { + return nil, fmt.Errorf("refresh workload identity auth for OTLP logs: %w", err) + } + + clonedReq := req.Clone(req.Context()) + clonedReq.Header = req.Header.Clone() + clonedReq.Header.Set("Authorization", authHeader) + return t.baseTransport().RoundTrip(clonedReq) +} + +func (t *authHeaderTransport) baseTransport() http.RoundTripper { + if t.base != nil { + return t.base + } + return http.DefaultTransport +} + +func otlpHTTPLogExporter(ctx context.Context, opts ...otlploghttp.Option) (log.Exporter, error) { + return otlploghttp.New(ctx, opts...) } diff --git a/services/core/jobs/jobs-launcher/cmd/run.go b/services/core/jobs/jobs-launcher/cmd/run.go index b16a732d69..f5ac028767 100644 --- a/services/core/jobs/jobs-launcher/cmd/run.go +++ b/services/core/jobs/jobs-launcher/cmd/run.go @@ -10,7 +10,6 @@ import ( "fmt" "io" "log/slog" - "net/url" "os" "os/exec" "os/signal" @@ -133,11 +132,10 @@ func fetchSecrets(apiBaseURL string, principal *nmpclient.Principal, secretRefs } // runExecWithStdin sets up OTEL and runs the specified command with stdin -func runExecWithStdin(args []string) (int, error) { +func runExecWithStdin(args []string) (exitCode int, err error) { ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer stop() - configureOTELHeadersFromWorkloadToken() otelShutdown, _, err := setupOTELSDK(ctx) if err != nil { return 1, err @@ -150,29 +148,6 @@ func runExecWithStdin(args []string) (int, error) { return runExec(args, os.Stdin) } -func configureOTELHeadersFromWorkloadToken() { - token := os.Getenv("NEMO_WORKLOAD_TOKEN") - if token == "" { - return - } - - const headersEnv = "OTEL_EXPORTER_OTLP_LOGS_HEADERS" - headers := os.Getenv(headersEnv) - for _, item := range strings.Split(headers, ",") { - key, _, _ := strings.Cut(strings.TrimSpace(item), "=") - if strings.EqualFold(key, "authorization") { - return - } - } - - authHeader := "Authorization=" + url.PathEscape("Bearer "+token) - if headers == "" { - os.Setenv(headersEnv, authHeader) - return - } - os.Setenv(headersEnv, headers+","+authHeader) -} - // runExec runs the specified command with arguments, injecting secrets as environment variables if specified func runExec(args []string, stdinReader io.Reader) (int, error) { // Command and arguments diff --git a/services/core/jobs/jobs-launcher/cmd/run_test.go b/services/core/jobs/jobs-launcher/cmd/run_test.go index ffc9586259..f1e696a8aa 100644 --- a/services/core/jobs/jobs-launcher/cmd/run_test.go +++ b/services/core/jobs/jobs-launcher/cmd/run_test.go @@ -306,66 +306,6 @@ func TestRunExecWithoutSecrets(t *testing.T) { } } -func TestConfigureOTELHeadersFromWorkloadToken(t *testing.T) { - testCases := []struct { - name string - token string - existingHeaders string - expectedHeaders string - }{ - { - name: "adds_authorization_header", - token: "token.with-symbols_123", - expectedHeaders: "Authorization=Bearer%20token.with-symbols_123", - }, - { - name: "preserves_existing_headers", - token: "abc.def", - existingHeaders: "X-NMP-Principal-Id=nemo-user", - expectedHeaders: "X-NMP-Principal-Id=nemo-user,Authorization=Bearer%20abc.def", - }, - { - name: "keeps_existing_authorization_header", - token: "abc.def", - existingHeaders: "authorization=Bearer+explicit", - expectedHeaders: "authorization=Bearer+explicit", - }, - { - name: "does_nothing_without_token", - existingHeaders: "X-Test=value", - expectedHeaders: "X-Test=value", - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - origEnvVars := map[string]envVarState{ - "NEMO_WORKLOAD_TOKEN": getEnvState("NEMO_WORKLOAD_TOKEN"), - "OTEL_EXPORTER_OTLP_LOGS_HEADERS": getEnvState("OTEL_EXPORTER_OTLP_LOGS_HEADERS"), - } - defer restoreEnvVars(origEnvVars) - - if tc.token != "" { - os.Setenv("NEMO_WORKLOAD_TOKEN", tc.token) - } else { - os.Unsetenv("NEMO_WORKLOAD_TOKEN") - } - if tc.existingHeaders != "" { - os.Setenv("OTEL_EXPORTER_OTLP_LOGS_HEADERS", tc.existingHeaders) - } else { - os.Unsetenv("OTEL_EXPORTER_OTLP_LOGS_HEADERS") - } - - configureOTELHeadersFromWorkloadToken() - - got := os.Getenv("OTEL_EXPORTER_OTLP_LOGS_HEADERS") - if got != tc.expectedHeaders { - t.Errorf("Expected OTEL headers %q, got %q", tc.expectedHeaders, got) - } - }) - } -} - func TestParseSecretReferences(t *testing.T) { testCases := []struct { name string diff --git a/services/core/jobs/jobs-launcher/cmd/workload_auth.go b/services/core/jobs/jobs-launcher/cmd/workload_auth.go new file mode 100644 index 0000000000..01848c9660 --- /dev/null +++ b/services/core/jobs/jobs-launcher/cmd/workload_auth.go @@ -0,0 +1,263 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package cmd + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "os" + "strings" + "sync" + "time" +) + +const ( + nmpBaseURLEnv = "NMP_BASE_URL" + workloadIdentityTokenFileEnv = "NMP_WORKLOAD_IDENTITY_TOKEN_FILE" + tokenExchangeGrantType = "urn:ietf:params:oauth:grant-type:token-exchange" + jwtTokenType = "urn:ietf:params:oauth:token-type:jwt" + accessTokenType = "urn:ietf:params:oauth:token-type:access_token" + workloadAuthRequestTimeoutSeconds = 30 + maxAuthResponseBodyBytes = 64 * 1024 + workloadAuthRefreshMarginFraction = 5 + workloadAuthMaxRefreshMargin = time.Minute +) + +type authDiscoveryResponse struct { + OIDC authDiscoveryOIDC `json:"oidc"` +} + +type authDiscoveryOIDC struct { + ClientID string `json:"client_id"` + TokenEndpoint string `json:"token_endpoint"` + WorkloadTokenExchangeEnabled bool `json:"workload_token_exchange_enabled"` + WorkloadClientID string `json:"workload_client_id"` + WorkloadTokenEndpoint string `json:"workload_token_endpoint"` + WorkloadAudience string `json:"workload_audience"` + WorkloadScope string `json:"workload_scope"` +} + +type tokenExchangeResponse struct { + AccessToken string `json:"access_token"` + ExpiresIn int64 `json:"expires_in"` +} + +type workloadAccessToken struct { + value string + refreshAt time.Time +} + +func newWorkloadAccessToken(value string, expiresIn int64, issuedAt time.Time) workloadAccessToken { + lifetime := time.Duration(expiresIn) * time.Second + refreshMargin := lifetime / workloadAuthRefreshMarginFraction + if refreshMargin > workloadAuthMaxRefreshMargin { + refreshMargin = workloadAuthMaxRefreshMargin + } + return workloadAccessToken{ + value: value, + refreshAt: issuedAt.Add(lifetime - refreshMargin), + } +} + +func (t workloadAccessToken) needsRefresh(now time.Time) bool { + return t.value == "" || !now.Before(t.refreshAt) +} + +type workloadAuthTokenSource struct { + subjectTokenFile string + tokenEndpoint string + clientID string + discovery authDiscoveryOIDC + + mu sync.Mutex + token workloadAccessToken + now func() time.Time +} + +func getOTLPLogWorkloadAuthHeaders(ctx context.Context) (map[string]string, error) { + tokenSource, err := newOTLPLogWorkloadAuthTokenSource(ctx) + if err != nil || tokenSource == nil { + return nil, err + } + return tokenSource.authHeaders(ctx) +} + +func newOTLPLogWorkloadAuthTokenSource(ctx context.Context) (*workloadAuthTokenSource, error) { + subjectTokenFile := os.Getenv(workloadIdentityTokenFileEnv) + if subjectTokenFile == "" { + return nil, nil + } + + baseURL := strings.TrimRight(os.Getenv(nmpBaseURLEnv), "/") + if baseURL == "" { + return nil, fmt.Errorf("%s is required when %s is set", nmpBaseURLEnv, workloadIdentityTokenFileEnv) + } + + ctx, cancel := context.WithTimeout(ctx, workloadAuthRequestTimeoutSeconds*time.Second) + defer cancel() + + discovery, err := discoverAuthConfig(ctx, baseURL) + if err != nil { + return nil, err + } + if !discovery.OIDC.WorkloadTokenExchangeEnabled { + return nil, fmt.Errorf("workload token exchange is not enabled by auth discovery") + } + + tokenEndpoint := discovery.OIDC.WorkloadTokenEndpoint + if tokenEndpoint == "" { + tokenEndpoint = discovery.OIDC.TokenEndpoint + } + if tokenEndpoint == "" { + return nil, fmt.Errorf("auth discovery did not return workload_token_endpoint or token_endpoint") + } + + clientID := discovery.OIDC.WorkloadClientID + if clientID == "" { + clientID = discovery.OIDC.ClientID + } + if clientID == "" { + return nil, fmt.Errorf("auth discovery did not return workload_client_id or client_id") + } + + return &workloadAuthTokenSource{ + subjectTokenFile: subjectTokenFile, + tokenEndpoint: tokenEndpoint, + clientID: clientID, + discovery: discovery.OIDC, + now: time.Now, + }, nil +} + +func (s *workloadAuthTokenSource) authHeaders(ctx context.Context) (map[string]string, error) { + authHeader, err := s.AuthorizationHeader(ctx) + if err != nil { + return nil, err + } + return map[string]string{"Authorization": authHeader}, nil +} + +func (s *workloadAuthTokenSource) AuthorizationHeader(ctx context.Context) (string, error) { + accessToken, err := s.accessToken(ctx) + if err != nil { + return "", err + } + return "Bearer " + accessToken, nil +} + +func (s *workloadAuthTokenSource) accessToken(ctx context.Context) (string, error) { + s.mu.Lock() + defer s.mu.Unlock() + + if !s.token.needsRefresh(s.now()) { + return s.token.value, nil + } + + ctx, cancel := context.WithTimeout(ctx, workloadAuthRequestTimeoutSeconds*time.Second) + defer cancel() + + subjectToken, err := readSubjectToken(s.subjectTokenFile) + if err != nil { + return "", err + } + + accessToken, err := exchangeWorkloadToken(ctx, s.tokenEndpoint, s.clientID, subjectToken, s.discovery, s.now()) + if err != nil { + return "", err + } + + s.token = accessToken + return s.token.value, nil +} + +func discoverAuthConfig(ctx context.Context, baseURL string) (*authDiscoveryResponse, error) { + request, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL+"/apis/auth/discovery", nil) + if err != nil { + return nil, err + } + + response, err := http.DefaultClient.Do(request) + if err != nil { + return nil, err + } + defer response.Body.Close() + + if response.StatusCode != http.StatusOK { + body, _ := io.ReadAll(io.LimitReader(response.Body, maxAuthResponseBodyBytes)) + return nil, fmt.Errorf("auth discovery failed: status %d: %s", response.StatusCode, strings.TrimSpace(string(body))) + } + + var discovery authDiscoveryResponse + if err := json.NewDecoder(io.LimitReader(response.Body, maxAuthResponseBodyBytes)).Decode(&discovery); err != nil { + return nil, fmt.Errorf("decode auth discovery response: %w", err) + } + return &discovery, nil +} + +func readSubjectToken(path string) (string, error) { + tokenBytes, err := os.ReadFile(path) + if err != nil { + return "", fmt.Errorf("read %s at %s: %w", workloadIdentityTokenFileEnv, path, err) + } + token := strings.TrimSpace(string(tokenBytes)) + if token == "" { + return "", fmt.Errorf("%s at %s is empty", workloadIdentityTokenFileEnv, path) + } + return token, nil +} + +func exchangeWorkloadToken( + ctx context.Context, + tokenEndpoint string, + clientID string, + subjectToken string, + discovery authDiscoveryOIDC, + issuedAt time.Time, +) (workloadAccessToken, error) { + form := url.Values{} + form.Set("grant_type", tokenExchangeGrantType) + form.Set("client_id", clientID) + form.Set("subject_token", subjectToken) + form.Set("subject_token_type", jwtTokenType) + form.Set("requested_token_type", accessTokenType) + if discovery.WorkloadAudience != "" { + form.Set("audience", discovery.WorkloadAudience) + } + if discovery.WorkloadScope != "" { + form.Set("scope", discovery.WorkloadScope) + } + + request, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenEndpoint, strings.NewReader(form.Encode())) + if err != nil { + return workloadAccessToken{}, err + } + request.Header.Set("Content-Type", "application/x-www-form-urlencoded") + + response, err := http.DefaultClient.Do(request) + if err != nil { + return workloadAccessToken{}, err + } + defer response.Body.Close() + + if response.StatusCode != http.StatusOK { + body, _ := io.ReadAll(response.Body) + return workloadAccessToken{}, fmt.Errorf("workload token exchange failed: status %d: %s", response.StatusCode, strings.TrimSpace(string(body))) + } + + var tokenResponse tokenExchangeResponse + if err := json.NewDecoder(response.Body).Decode(&tokenResponse); err != nil { + return workloadAccessToken{}, fmt.Errorf("decode workload token exchange response: %w", err) + } + if strings.TrimSpace(tokenResponse.AccessToken) == "" { + return workloadAccessToken{}, fmt.Errorf("workload token exchange response did not include a non-empty access_token") + } + if tokenResponse.ExpiresIn <= 0 { + return workloadAccessToken{}, fmt.Errorf("workload token exchange response did not include a positive expires_in") + } + return newWorkloadAccessToken(tokenResponse.AccessToken, tokenResponse.ExpiresIn, issuedAt), nil +} diff --git a/services/core/jobs/jobs-launcher/cmd/workload_auth_test.go b/services/core/jobs/jobs-launcher/cmd/workload_auth_test.go new file mode 100644 index 0000000000..cb4a72d74f --- /dev/null +++ b/services/core/jobs/jobs-launcher/cmd/workload_auth_test.go @@ -0,0 +1,407 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package cmd + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + otellog "go.opentelemetry.io/otel/log" + sdklog "go.opentelemetry.io/otel/sdk/log" +) + +const otelExporterOTLPLogsHeadersEnv = "OTEL_EXPORTER_OTLP_LOGS_HEADERS" + +func TestGetOTLPLogWorkloadAuthHeadersReturnsAuthorizationWithoutMutatingEnv(t *testing.T) { + subjectTokenPath := filepath.Join(t.TempDir(), "subject.jwt") + if err := os.WriteFile(subjectTokenPath, []byte("subject-token\n"), 0o600); err != nil { + t.Fatalf("failed to write subject token: %v", err) + } + + var serverURL string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/apis/auth/discovery": + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf( + w, + `{"auth_enabled":true,"oidc":{"workload_token_exchange_enabled":true,"workload_client_id":"nemo-platform-workload","workload_token_endpoint":%q,"workload_audience":"nemo-platform","workload_scope":"openid email groups"}}`, + serverURL+"/apis/auth/token", + ) + case "/apis/auth/token": + if r.Method != http.MethodPost { + t.Errorf("expected POST token exchange, got %s", r.Method) + } + if err := r.ParseForm(); err != nil { + t.Errorf("failed to parse token exchange form: %v", err) + } + expectedForm := url.Values{ + "grant_type": {tokenExchangeGrantType}, + "client_id": {"nemo-platform-workload"}, + "subject_token": {"subject-token"}, + "subject_token_type": {jwtTokenType}, + "requested_token_type": {accessTokenType}, + "audience": {"nemo-platform"}, + "scope": {"openid email groups"}, + } + expectedEncodedForm := expectedForm.Encode() + actualEncodedForm := r.PostForm.Encode() + if actualEncodedForm != expectedEncodedForm { + t.Errorf( + "form: expected encoded %q (%v), got %q (%v)", + expectedEncodedForm, + expectedForm, + actualEncodedForm, + r.PostForm, + ) + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{"access_token": "access.token.value", "expires_in": 120}) // nolint:errcheck + default: + http.NotFound(w, r) + } + })) + defer server.Close() + serverURL = server.URL + + t.Setenv(nmpBaseURLEnv, server.URL) + t.Setenv(workloadIdentityTokenFileEnv, subjectTokenPath) + t.Setenv(otelExporterOTLPLogsHeadersEnv, "X-NMP-Principal-Id=nemo-user") + + headers, err := getOTLPLogWorkloadAuthHeaders(context.Background()) + if err != nil { + t.Fatalf("getOTLPLogWorkloadAuthHeaders returned error: %v", err) + } + + if got := headers["Authorization"]; got != "Bearer access.token.value" { + t.Fatalf("expected returned bearer header, got %q", got) + } + if got := os.Getenv(otelExporterOTLPLogsHeadersEnv); got != "X-NMP-Principal-Id=nemo-user" { + t.Fatalf("expected OTEL headers env to remain untouched, got %s", got) + } +} + +func TestGetOTLPLogWorkloadAuthHeadersNoopsWithoutTokenFile(t *testing.T) { + t.Setenv(otelExporterOTLPLogsHeadersEnv, "X-NMP-Principal-Id=nemo-user") + + headers, err := getOTLPLogWorkloadAuthHeaders(context.Background()) + if err != nil { + t.Fatalf("getOTLPLogWorkloadAuthHeaders returned error: %v", err) + } + if headers != nil { + t.Fatalf("expected no auth headers without workload token file, got %v", headers) + } + + if headers := os.Getenv(otelExporterOTLPLogsHeadersEnv); headers != "X-NMP-Principal-Id=nemo-user" { + t.Fatalf("expected headers to be unchanged, got %s", headers) + } +} + +func TestNewLogExporterCachesWorkloadAuthAcrossExports(t *testing.T) { + subjectTokenPath := filepath.Join(t.TempDir(), "subject.jwt") + if err := os.WriteFile(subjectTokenPath, []byte("subject-token\n"), 0o600); err != nil { + t.Fatalf("failed to write subject token: %v", err) + } + + var serverURL string + var mu sync.Mutex + var tokenExchangeCount int + var exportAuthHeaders []string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/apis/auth/discovery": + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf( + w, + `{"auth_enabled":true,"oidc":{"workload_token_exchange_enabled":true,"workload_client_id":"nemo-platform-workload","workload_token_endpoint":%q}}`, + serverURL+"/apis/auth/token", + ) + case "/apis/auth/token": + mu.Lock() + tokenExchangeCount++ + token := fmt.Sprintf("access-token-%d", tokenExchangeCount) + mu.Unlock() + + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{"access_token": token, "expires_in": 120}) // nolint:errcheck + case "/v1/logs": + mu.Lock() + exportAuthHeaders = append(exportAuthHeaders, r.Header.Get("Authorization")) + mu.Unlock() + w.WriteHeader(http.StatusOK) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + serverURL = server.URL + + t.Setenv(nmpBaseURLEnv, server.URL) + t.Setenv(workloadIdentityTokenFileEnv, subjectTokenPath) + t.Setenv(nmpJobLogsEndpointEnv, server.URL+"/v1/logs") + + exporter, err := newLogExporter(context.Background()) + if err != nil { + t.Fatalf("newLogExporter returned error: %v", err) + } + + record := sdklog.Record{} + record.SetBody(otellog.StringValue("first")) + if err := exporter.Export(context.Background(), []sdklog.Record{record}); err != nil { + t.Fatalf("first export returned error: %v", err) + } + + record.SetBody(otellog.StringValue("second")) + if err := exporter.Export(context.Background(), []sdklog.Record{record}); err != nil { + t.Fatalf("second export returned error: %v", err) + } + + mu.Lock() + defer mu.Unlock() + if tokenExchangeCount != 1 { + t.Fatalf("expected cached token to avoid per-export token exchanges, got %d", tokenExchangeCount) + } + expectedAuthHeaders := []string{"Bearer access-token-1", "Bearer access-token-1"} + if strings.Join(exportAuthHeaders, ",") != strings.Join(expectedAuthHeaders, ",") { + t.Fatalf("expected export Authorization headers %v, got %v", expectedAuthHeaders, exportAuthHeaders) + } +} + +func TestWorkloadAuthTokenSourceRefreshesNearExpiry(t *testing.T) { + subjectTokenPath := filepath.Join(t.TempDir(), "subject.jwt") + if err := os.WriteFile(subjectTokenPath, []byte("subject-token\n"), 0o600); err != nil { + t.Fatalf("failed to write subject token: %v", err) + } + + var serverURL string + var tokenExchangeMu sync.Mutex + var tokenExchangeCount int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/apis/auth/discovery": + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf( + w, + `{"auth_enabled":true,"oidc":{"workload_token_exchange_enabled":true,"workload_client_id":"nemo-platform-workload","workload_token_endpoint":%q}}`, + serverURL+"/apis/auth/token", + ) + case "/apis/auth/token": + tokenExchangeMu.Lock() + tokenExchangeCount++ + token := fmt.Sprintf("access-token-%d", tokenExchangeCount) + tokenExchangeMu.Unlock() + + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ // nolint:errcheck + "access_token": token, + "expires_in": 100, + }) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + serverURL = server.URL + + t.Setenv(nmpBaseURLEnv, server.URL) + t.Setenv(workloadIdentityTokenFileEnv, subjectTokenPath) + + source, err := newOTLPLogWorkloadAuthTokenSource(context.Background()) + if err != nil { + t.Fatalf("newOTLPLogWorkloadAuthTokenSource returned error: %v", err) + } + if source == nil { + t.Fatal("expected workload auth token source") + } + + now := time.Unix(1_700_000_000, 0) + source.now = func() time.Time { return now } + + header, err := source.AuthorizationHeader(context.Background()) + if err != nil { + t.Fatalf("first AuthorizationHeader returned error: %v", err) + } + if header != "Bearer access-token-1" { + t.Fatalf("expected first token, got %q", header) + } + + now = now.Add(79 * time.Second) + header, err = source.AuthorizationHeader(context.Background()) + if err != nil { + t.Fatalf("cached AuthorizationHeader returned error: %v", err) + } + if header != "Bearer access-token-1" { + t.Fatalf("expected cached token, got %q", header) + } + + now = now.Add(2 * time.Second) + header, err = source.AuthorizationHeader(context.Background()) + if err != nil { + t.Fatalf("refreshed AuthorizationHeader returned error: %v", err) + } + if header != "Bearer access-token-2" { + t.Fatalf("expected refreshed token, got %q", header) + } + tokenExchangeMu.Lock() + actualTokenExchangeCount := tokenExchangeCount + tokenExchangeMu.Unlock() + if actualTokenExchangeCount != 2 { + t.Fatalf("expected token exchange only when missing and near expiry, got %d", actualTokenExchangeCount) + } +} + +func TestWorkloadAuthTokenSourceSynchronizesConcurrentRefresh(t *testing.T) { + subjectTokenPath := filepath.Join(t.TempDir(), "subject.jwt") + if err := os.WriteFile(subjectTokenPath, []byte("subject-token\n"), 0o600); err != nil { + t.Fatalf("failed to write subject token: %v", err) + } + + var serverURL string + var mu sync.Mutex + var tokenExchangeCount int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/apis/auth/discovery": + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf( + w, + `{"auth_enabled":true,"oidc":{"workload_token_exchange_enabled":true,"workload_client_id":"nemo-platform-workload","workload_token_endpoint":%q}}`, + serverURL+"/apis/auth/token", + ) + case "/apis/auth/token": + mu.Lock() + tokenExchangeCount++ + mu.Unlock() + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{"access_token": "access-token", "expires_in": 100}) // nolint:errcheck + default: + http.NotFound(w, r) + } + })) + defer server.Close() + serverURL = server.URL + + t.Setenv(nmpBaseURLEnv, server.URL) + t.Setenv(workloadIdentityTokenFileEnv, subjectTokenPath) + + source, err := newOTLPLogWorkloadAuthTokenSource(context.Background()) + if err != nil { + t.Fatalf("newOTLPLogWorkloadAuthTokenSource returned error: %v", err) + } + if source == nil { + t.Fatal("expected workload auth token source") + } + source.now = func() time.Time { return time.Unix(1_700_000_000, 0) } + + const callers = 8 + start := make(chan struct{}) + errs := make(chan error, callers) + var wg sync.WaitGroup + for range callers { + wg.Add(1) + go func() { + defer wg.Done() + <-start + header, err := source.AuthorizationHeader(context.Background()) + if err != nil { + errs <- err + return + } + if header != "Bearer access-token" { + errs <- fmt.Errorf("expected cached header, got %q", header) + } + }() + } + close(start) + wg.Wait() + close(errs) + + for err := range errs { + if err != nil { + t.Fatal(err) + } + } + mu.Lock() + defer mu.Unlock() + if tokenExchangeCount != 1 { + t.Fatalf("expected concurrent refreshes to share one token exchange, got %d", tokenExchangeCount) + } +} + +func TestNewLogExporterPropagatesInitialWorkloadAuthFailure(t *testing.T) { + subjectTokenPath := filepath.Join(t.TempDir(), "subject.jwt") + if err := os.WriteFile(subjectTokenPath, []byte("subject-token\n"), 0o600); err != nil { + t.Fatalf("failed to write subject token: %v", err) + } + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "discovery unavailable", http.StatusBadGateway) + })) + defer server.Close() + + t.Setenv(nmpBaseURLEnv, server.URL) + t.Setenv(workloadIdentityTokenFileEnv, subjectTokenPath) + t.Setenv(nmpJobLogsEndpointEnv, server.URL+"/v1/logs") + + _, err := newLogExporter(context.Background()) + if err == nil { + t.Fatal("expected newLogExporter to return workload auth error") + } + if !strings.Contains(err.Error(), "configure workload identity auth for OTLP logs") { + t.Fatalf("expected setup auth context in error, got %v", err) + } + if !strings.Contains(err.Error(), "auth discovery failed: status 502") { + t.Fatalf("expected discovery status in error, got %v", err) + } +} + +func TestDiscoverAuthConfigLimitsNonOKResponseBody(t *testing.T) { + largeBody := strings.Repeat("x", maxAuthResponseBodyBytes+1024) + "tail-sentinel" + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, largeBody, http.StatusBadGateway) + })) + defer server.Close() + + _, err := discoverAuthConfig(context.Background(), server.URL) + if err == nil { + t.Fatal("expected auth discovery error") + } + if !strings.Contains(err.Error(), "auth discovery failed: status 502") { + t.Fatalf("expected status error, got %v", err) + } + if strings.Contains(err.Error(), "tail-sentinel") { + t.Fatalf("expected auth discovery error body to be limited, got %v", err) + } +} + +func TestDiscoverAuthConfigLimitsSuccessfulJSONResponseBody(t *testing.T) { + padding := strings.Repeat("x", maxAuthResponseBodyBytes) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf( + w, + `{"oidc":{"workload_token_exchange_enabled":true,"workload_client_id":"client","workload_token_endpoint":"https://idp.example.com/token","padding":%q}}`, + padding, + ) + })) + defer server.Close() + + _, err := discoverAuthConfig(context.Background(), server.URL) + if err == nil { + t.Fatal("expected oversized auth discovery response to fail decoding") + } + if !strings.Contains(err.Error(), "decode auth discovery response") { + t.Fatalf("expected decode error, got %v", err) + } +} diff --git a/services/core/jobs/jobs-launcher/go.mod b/services/core/jobs/jobs-launcher/go.mod index 3a22e9e7e1..5f1dd590bb 100644 --- a/services/core/jobs/jobs-launcher/go.mod +++ b/services/core/jobs/jobs-launcher/go.mod @@ -7,6 +7,7 @@ require ( go.opentelemetry.io/contrib/bridges/otelslog v0.19.0 go.opentelemetry.io/contrib/exporters/autoexport v0.69.0 go.opentelemetry.io/otel v1.44.0 + go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploghttp v0.20.0 go.opentelemetry.io/otel/exporters/stdout/stdoutlog v0.20.0 go.opentelemetry.io/otel/log v0.20.0 go.opentelemetry.io/otel/sdk v1.44.0 @@ -32,7 +33,6 @@ require ( go.opentelemetry.io/auto/sdk v1.2.1 // indirect go.opentelemetry.io/contrib/bridges/prometheus v0.69.0 // indirect go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploggrpc v0.20.0 // indirect - go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploghttp v0.20.0 // indirect go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetricgrpc v1.44.0 // indirect go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetrichttp v1.44.0 // indirect go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.44.0 // indirect diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/base.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/base.py index 52adba9e0a..64432ba840 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/base.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/base.py @@ -6,24 +6,43 @@ import datetime import logging from abc import ABC, abstractmethod +from collections.abc import Iterable from enum import Enum from typing import Generic, Optional, TypeVar +from urllib.parse import SplitResult, urlsplit from nemo_platform import NeMoPlatform from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR from nemo_platform_plugin.client.errors import NotFoundError as ClientNotFoundError from nemo_platform_plugin.jobs import execution_profiles as _execution_profiles from nemo_platform_plugin.jobs.client import JobsClient from nemo_platform_plugin.jobs.schemas import PlatformJobStatus from nemo_platform_plugin.jobs.types import PlatformJobStepResponse, PlatformJobStepWithContext +from nmp.common.auth.models import NMP_PRINCIPAL_ENVVAR from nmp.common.config.base import ( LOOPBACK_ADDRESSES, + NMP_CONFIG_WARNINGS_DISABLED_ENV_VAR, PlatformConfig, determine_loopback_override, ) +from nmp.common.jobs.constants import ( + CONFIG_TASK_STORAGE_PATH_ENVVAR, + EPHEMERAL_TASK_STORAGE_PATH_ENVVAR, + NEMO_JOB_ATTEMPT_ID_ENVVAR, + NEMO_JOB_FILESET_ENVVAR, + NEMO_JOB_ID_ENVVAR, + NEMO_JOB_SECRETS_ENVVAR, + NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR, + NEMO_JOB_STEP_ENVVAR, + NEMO_JOB_TASK_ENVVAR, + NEMO_JOB_WORKSPACE_ENVVAR, + PERSISTENT_JOB_STORAGE_PATH_ENVVAR, + TASK_CONFIG_ENVVAR, +) from nmp.common.sdk_factory import get_entity_parts from nmp.core.jobs.app.providers import ComputeResources -from pydantic import BaseModel +from pydantic import BaseModel, model_validator logger = logging.getLogger(__name__) @@ -32,12 +51,108 @@ DEFAULT_PROFILE = "default" DEFAULT_PROVIDER = "cpu" -JobExecutionProfileConfig = _execution_profiles.JobExecutionProfileConfig -RESERVED_JOB_ENVIRONMENT_VARIABLE_NAMES = _execution_profiles.RESERVED_JOB_ENVIRONMENT_VARIABLE_NAMES -# The env-var-name reserved set and the base ``JobExecutionProfileConfig`` now -# live in the shared plugin leaf node (imported above) so that both the server -# and the typed HTTP client agree on the wire shape and validation. +# The base ``JobExecutionProfileConfig`` lives in the shared plugin leaf node +# (imported above) so that both the server and the typed HTTP client agree on +# the wire shape. + +WORKLOAD_IDENTITY_TOKEN_FILE_PATH = "/var/run/secrets/nemo-platform/workload/token" +WORKLOAD_IDENTITY_VOLUME_PATH = "/var/run/secrets/nemo-platform/workload" +WORKLOAD_IDENTITY_VOLUME_NAME = "nmp-workload-identity" +JOB_LOGS_ENDPOINT_ENVVAR = _execution_profiles.JOB_LOGS_ENDPOINT_ENVVAR + +RESERVED_MANAGED_JOB_AUTH_ENVIRONMENT_VARIABLE_NAMES: frozenset[str] = frozenset( + { + JOB_LOGS_ENDPOINT_ENVVAR, + "NMP_ACCESS_TOKEN", + "NEMO_WORKLOAD_TOKEN", + "NEMO_WORKLOAD_TOKEN_FILE", + WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, + } +) + +RESERVED_MANAGED_JOB_ENVIRONMENT_VARIABLE_PREFIXES: tuple[str, ...] = ( + "NEMO_WORKLOAD_", + "NMP_WORKLOAD_", + "NEMO_WORKFLOW_", + "NMP_WORKFLOW_", +) + +# Env var names set by the platform during job creation; user-provided profile environment must not conflict. +RESERVED_JOB_ENVIRONMENT_VARIABLE_NAMES: frozenset[str] = ( + frozenset( + { + # From nmp.common.jobs.constants + CONFIG_TASK_STORAGE_PATH_ENVVAR, + EPHEMERAL_TASK_STORAGE_PATH_ENVVAR, + NEMO_JOB_ATTEMPT_ID_ENVVAR, + NEMO_JOB_FILESET_ENVVAR, + NEMO_JOB_ID_ENVVAR, + NEMO_JOB_SECRETS_ENVVAR, + NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR, + NEMO_JOB_STEP_ENVVAR, + NEMO_JOB_TASK_ENVVAR, + NEMO_JOB_WORKSPACE_ENVVAR, + PERSISTENT_JOB_STORAGE_PATH_ENVVAR, + TASK_CONFIG_ENVVAR, + # Auth + NMP_PRINCIPAL_ENVVAR, + # Platform launcher logs + JOB_LOGS_ENDPOINT_ENVVAR, + # Platform shared envvars (to_shared_envvars with NMP_ prefix) + NMP_CONFIG_WARNINGS_DISABLED_ENV_VAR, + "NMP_AUTH_URL", + "NMP_BASE_URL", + "NMP_JOBS_URL", + "NMP_FILES_URL", + "NMP_MODELS_URL", + "NMP_SECRETS_URL", + } + ) + | RESERVED_MANAGED_JOB_AUTH_ENVIRONMENT_VARIABLE_NAMES +) + + +def find_reserved_managed_job_environment_variable_names(env_names: Iterable[str]) -> list[str]: + """Return user-provided managed-job env names reserved for platform-managed injection.""" + return sorted( + { + name + for name in env_names + if name in RESERVED_MANAGED_JOB_AUTH_ENVIRONMENT_VARIABLE_NAMES + or name.startswith(RESERVED_MANAGED_JOB_ENVIRONMENT_VARIABLE_PREFIXES) + } + ) + + +def validate_no_reserved_managed_job_environment_variable_names(env_names: Iterable[str], *, source: str) -> None: + """Reject user-provided platform-managed env names for managed jobs.""" + conflicting = find_reserved_managed_job_environment_variable_names(env_names) + if conflicting: + raise ValueError(f"{source} must not use platform-reserved managed environment keys: {conflicting}") + + +def is_workload_identity_token_exchange_enabled() -> bool: + """Return whether managed jobs should inject workload identity token-file config.""" + try: + from nmp.common.config import get_auth_config + + return bool(get_auth_config().oidc.workload_token_exchange_enabled) + except Exception: + logger.debug("Could not resolve auth config for workload identity token exchange", exc_info=True) + return False + + +def get_workload_identity_token_audience() -> str: + """Return the Kubernetes projected service-account token audience for workload identity.""" + try: + from nmp.common.config import get_auth_config + + oidc = get_auth_config().oidc + return oidc.workload_client_id or oidc.client_id or "nemo-platform" + except Exception: + logger.debug("Could not resolve auth config for workload identity audience", exc_info=True) + return "nemo-platform" class JobUpdate(BaseModel): @@ -46,7 +161,30 @@ class JobUpdate(BaseModel): error_details: dict | None = None +class JobExecutionProfileConfig(_execution_profiles.JobExecutionProfileConfig): + """Server-side extension adding workload-identity managed-auth env validation.""" + + @model_validator(mode="after") + def validate_env_no_reserved_names(self) -> JobExecutionProfileConfig: + conflicting = [k for k in self.env if k in RESERVED_JOB_ENVIRONMENT_VARIABLE_NAMES] + if conflicting: + raise ValueError( + f"Profile environment keys must not conflict with platform-reserved names: {sorted(conflicting)}" + ) + return self + + @model_validator(mode="after") + def validate_env_no_reserved_managed_names(self) -> JobExecutionProfileConfig: + conflicting = find_reserved_managed_job_environment_variable_names(self.env) + if conflicting: + raise ValueError( + f"Profile environment keys must not conflict with platform-reserved auth environment keys: {sorted(conflicting)}" + ) + return self + + _DEFAULT_TASK_IMAGE_NAME = "nmp-cpu-tasks" +JOB_RUNTIME_SERVICE_NAMES = ("auth", "jobs", "files", "models", "secrets") def resolve_task_image(container_image: str | None, default_task_image: str | None) -> str: @@ -74,6 +212,96 @@ def resolve_gpu_job_shm_size( return f"{max(1, num_gpus)}Gi" +def _split_url(url: str) -> SplitResult | None: + try: + parsed = urlsplit(url) + parsed.hostname + parsed.port + except ValueError: + return None + return parsed + + +def _is_loopback_url(url: str) -> bool: + parsed = _split_url(url) + return parsed is not None and parsed.hostname in LOOPBACK_ADDRESSES + + +def _format_netloc_with_hostname(parsed: SplitResult, hostname: str) -> str: + host = hostname[1:-1] if hostname.startswith("[") and hostname.endswith("]") else hostname + if ":" in host: + host = f"[{host}]" + + userinfo = "" + if parsed.username is not None: + userinfo = parsed.username + if parsed.password is not None: + userinfo = f"{userinfo}:{parsed.password}" + userinfo = f"{userinfo}@" + + netloc = userinfo + host + if parsed.port is not None: + netloc = f"{netloc}:{parsed.port}" + return netloc + + +def _replace_loopback_address(url: str, loopback_address: str | None) -> str: + if not loopback_address: + return url + + parsed = _split_url(url) + if parsed is None or parsed.hostname not in LOOPBACK_ADDRESSES: + return url + + netloc = _format_netloc_with_hostname(parsed, loopback_address) + return parsed._replace(netloc=netloc).geturl() + + +def _contains_loopback_address(url: str) -> bool: + return _is_loopback_url(url) + + +def _job_runtime_base_url(platform_config: PlatformConfig) -> str: + service_discovery_url = platform_config.service_discovery.get("platform") or platform_config.service_discovery.get( + "base" + ) + if service_discovery_url: + return service_discovery_url + + return platform_config.base_url + + +def get_job_runtime_shared_envvars( + platform_config: PlatformConfig, + *, + disable_warnings: bool = True, + loopback_address: str | None = None, +) -> dict[str, str]: + """Return platform URLs for job containers/pods. + + Jobs run outside the API server process. When services run in-process, + PlatformConfig.get_service_url() intentionally returns a local server URL, + but job runtimes need URLs reachable from their own network, usually the + configured service_discovery gateway. + """ + effective_loopback_address = loopback_address or platform_config.loopback_address or determine_loopback_override() + base_url = _job_runtime_base_url(platform_config) + base_url = _replace_loopback_address(base_url, effective_loopback_address) + + envvars = {"NMP_BASE_URL": base_url} + for service_name in JOB_RUNTIME_SERVICE_NAMES: + service_url = platform_config.service_discovery.get(service_name) or base_url + envvars[f"NMP_{service_name.upper()}_URL"] = _replace_loopback_address( + service_url, + effective_loopback_address, + ) + + if disable_warnings: + envvars[NMP_CONFIG_WARNINGS_DISABLED_ENV_VAR] = "1" + + return envvars + + class JobBackend(Generic[ExecutionProviderConfigT, ExecutionProfileConfigT], ABC): BACKEND_NAME: str = "generic_backend" @@ -286,7 +514,10 @@ def extract_provider_profile(step: PlatformJobStepWithContext) -> tuple[str, str def get_logs_endpoint_from_fileset( - platform_config: PlatformConfig, workspace: str, fileset_id: str, loopback_address: str | None = None + platform_config: PlatformConfig, + workspace: str, + fileset_id: str, + loopback_address: str | None = None, ) -> str: """Get the OTLP logs endpoint URL for a fileset. @@ -300,15 +531,14 @@ def get_logs_endpoint_from_fileset( Returns: Full OTLP logs endpoint URL with appropriate loopback address applied. """ - base_url = platform_config.get_service_url("files") + # Job telemetry is emitted from a separate process/container/pod. When Files + # runs in-process with the API server, local service URLs are not necessarily + # routable from job runtime networks, so fall back through the same + # workload-facing base URL used for job SDK env vars. + base_url = platform_config.service_discovery.get("files") or _job_runtime_base_url(platform_config) # Use configured loopback_address, or fall back to automatic detection effective_override = loopback_address or platform_config.loopback_address or determine_loopback_override() - - if effective_override: - for loopback in LOOPBACK_ADDRESSES: - if loopback in base_url: - base_url = base_url.replace(loopback, effective_override) - break + base_url = _replace_loopback_address(base_url, effective_override) return f"{base_url}/apis/files/v2/workspaces/{workspace}/filesets/{fileset_id}/otlp/v1/logs" diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py index 6c536dd8bd..b3eefe3640 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py @@ -14,7 +14,7 @@ from abc import abstractmethod from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass -from typing import Any, Generic, TypeVar +from typing import Any, Generic, Self, TypeVar import docker.types from docker.errors import APIError, ImageNotFound, NotFound @@ -89,12 +89,19 @@ GPUExecutionProvider, ) from nmp.core.jobs.controllers.backends.base import ( + JOB_LOGS_ENDPOINT_ENVVAR, + WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, + WORKLOAD_IDENTITY_TOKEN_FILE_PATH, + WORKLOAD_IDENTITY_VOLUME_PATH, JobBackend, JobUpdate, + get_job_runtime_shared_envvars, get_logs_endpoint_from_fileset, + is_workload_identity_token_exchange_enabled, resolve_gpu_job_shm_size, resolve_task_image, staleness_error_message, + validate_no_reserved_managed_job_environment_variable_names, ) from nmp.core.jobs.controllers.backends.exceptions import ( FailedToScheduleError, @@ -102,8 +109,16 @@ ResourceAllocationError, SchedulingDeferred, ) +from nmp.core.jobs.controllers.backends.workload_tokens import ( + DEFAULT_SUBJECT_TOKEN_REFRESH_MARGIN_SECONDS, + DEFAULT_SUBJECT_TOKEN_TTL_SECONDS, + OAuthPasswordGrantSubjectTokenIssuer, + SubjectTokenIssuer, + SubjectTokenRefreshLoop, + build_token_archive, +) from opentelemetry import trace -from pydantic import Field +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator import docker @@ -135,6 +150,9 @@ def k8s_shm_quantity_to_docker(quantity: str) -> str: # Default is 30 seconds which matches the Kubernetes default grace period for pod termination. DOCKER_STOP_TIMEOUT = int(os.getenv("NEMO_JOBS_DEFAULT_DOCKER_STOP_TIMEOUT", "30")) NMP_JOBS_DOCKER_OWNER_ID_ENVVAR = "NMP_JOBS_DOCKER_OWNER_ID" +DOCKER_WORKLOAD_IDENTITY_PASSWORD_ENV_VAR_DEFAULT = "AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD" +DOCKER_WORKLOAD_IDENTITY_TOKEN_FILE_LABEL = "nmp.nvidia.com/workload_identity_token_file" +DOCKER_WORKLOAD_IDENTITY_VOLUME_LABEL = "nmp.nvidia.com/workload_identity_volume" ProviderT = TypeVar("ProviderT", bound=ExecutionProviderT) @@ -170,6 +188,77 @@ class DockerJobNetworkConfig(PluginDockerJobNetworkConfig): ) +class DockerWorkloadIdentityConfig(BaseModel): + """Docker-only subject token issuer configuration for workload identity.""" + + model_config = ConfigDict(extra="forbid") + + enabled: bool | None = Field( + default=None, + description="Enable Docker workload identity token-file injection. Defaults to auth.oidc.workload_token_exchange_enabled.", + ) + token_endpoint: str | None = Field( + default=None, + description="OAuth token endpoint used by the Docker demo issuer. Defaults to auth.oidc.token_endpoint.", + ) + client_id: str | None = Field( + default=None, + description="OAuth client ID used by the Docker demo issuer. Defaults to auth.oidc.workload_client_id or auth.oidc.client_id.", + ) + client_secret: str | None = Field( + default=None, + description="OAuth client secret for the Docker demo issuer.", + json_schema_extra={"format": "password", "writeOnly": True}, + ) + username: str | None = Field(default=None, description="Username for the Docker demo issuer password grant.") + password_env_var: str = Field( + default=DOCKER_WORKLOAD_IDENTITY_PASSWORD_ENV_VAR_DEFAULT, + description="Controller environment variable that contains the Docker demo issuer password grant shared secret.", + ) + scope: str | None = Field(default=None, description="OAuth scope for the Docker demo issuer.") + subject_token_ttl_seconds: int = Field( + default_factory=lambda: int( + os.environ.get("NMP_WORKLOAD_IDENTITY_TOKEN_TTL_SECONDS", DEFAULT_SUBJECT_TOKEN_TTL_SECONDS) + ), + ge=1, + validate_default=True, + description="Fallback subject-token lifetime when the Docker demo issuer response omits expires_in.", + ) + refresh_margin_seconds: int = Field( + default=DEFAULT_SUBJECT_TOKEN_REFRESH_MARGIN_SECONDS, + ge=0, + validate_default=True, + description="Seconds before subject-token expiry when the Docker refresher issues a replacement token.", + ) + + @field_validator("password_env_var") + @classmethod + def validate_password_env_var(cls, value: str) -> str: + value = value.strip() + if not value: + raise ValueError("password_env_var must name a non-empty environment variable") + if value[0].isdigit() or not all(char == "_" or char.isalnum() for char in value): + raise ValueError("password_env_var must be a valid environment variable name") + return value + + @model_validator(mode="after") + def validate_refresh_margin_before_expiry(self) -> Self: + if self.refresh_margin_seconds >= self.subject_token_ttl_seconds: + raise ValueError("refresh_margin_seconds must be less than subject_token_ttl_seconds") + return self + + +def _resolve_docker_workload_identity_scope(workload_config: DockerWorkloadIdentityConfig, oidc_config: Any) -> str: + return workload_config.scope or getattr(oidc_config, "workload_scope", None) or "openid email groups" + + +def _resolve_docker_workload_identity_password(workload_config: DockerWorkloadIdentityConfig) -> str | None: + password = os.environ.get(workload_config.password_env_var) + if password: + return password + return None + + class DockerJobExecutionProfileConfig(PluginDockerJobExecutionProfileConfig): """Configuration for Docker Job execution profile.""" @@ -177,6 +266,10 @@ class DockerJobExecutionProfileConfig(PluginDockerJobExecutionProfileConfig): networking: DockerJobNetworkConfig = Field( default_factory=DockerJobNetworkConfig, description="Docker networking configuration" ) + workload_identity: DockerWorkloadIdentityConfig = Field( + default_factory=DockerWorkloadIdentityConfig, + description="Docker workload identity subject-token issuer configuration.", + ) class DockerJobExecutionProfile(PluginDockerJobExecutionProfile): @@ -196,6 +289,7 @@ def init(self) -> None: self._jobs_controller_instance_id = _resolve_jobs_controller_instance_id() self._container_start_admission = threading.BoundedSemaphore(DOCKER_CONTAINER_START_WORKERS) self._container_run_threadpool = ThreadPoolExecutor(max_workers=DOCKER_CONTAINER_START_WORKERS) + self._workload_identity_refreshers: dict[str, SubjectTokenRefreshLoop] = {} self._client = docker.from_env(timeout=180) if NEMO_JOBS_IMAGE_REGISTRY: logger.info( @@ -209,8 +303,217 @@ def init(self) -> None: def shutdown(self) -> None: self._container_run_threadpool.shutdown(wait=True) + self._stop_all_workload_identity_refreshers() self._client.close() + def _is_workload_identity_enabled(self) -> bool: + configured = self._execution_profile_config.workload_identity.enabled + if configured is not None: + return configured + return is_workload_identity_token_exchange_enabled() + + def _create_docker_subject_token_issuer(self) -> SubjectTokenIssuer: + workload_config = self._execution_profile_config.workload_identity + try: + from nmp.common.config import get_auth_config + + oidc_config = get_auth_config().oidc + except Exception: + oidc_config = None + + token_endpoint = workload_config.token_endpoint or getattr(oidc_config, "token_endpoint", None) + client_id = ( + workload_config.client_id + or getattr(oidc_config, "workload_client_id", None) + or getattr(oidc_config, "client_id", None) + ) + password = _resolve_docker_workload_identity_password(workload_config) + scope = _resolve_docker_workload_identity_scope(workload_config, oidc_config) + missing = [ + name + for name, value in { + "token_endpoint": token_endpoint, + "client_id": client_id, + "username": workload_config.username, + workload_config.password_env_var: password, + }.items() + if not value + ] + if missing: + raise JobStorageError( + "Docker workload identity token exchange is enabled, but the Docker profile is missing issuer " + f"configuration: {', '.join(missing)}" + ) + + assert token_endpoint is not None + assert client_id is not None + assert workload_config.username is not None + assert password is not None + return OAuthPasswordGrantSubjectTokenIssuer( + token_endpoint=token_endpoint, + client_id=client_id, + client_secret=workload_config.client_secret, + username=workload_config.username, + password=password, + scope=scope, + default_expires_in_seconds=workload_config.subject_token_ttl_seconds, + ) + + def _write_workload_identity_subject_token(self, volume_name: str, token: str) -> None: + storage_config = self._execution_profile_config.storage + permissions_image = ( + storage_config.volume_permissions_image if storage_config is not None else DEFAULT_VOLUME_PERMISSIONS_IMAGE + ) + token_volume_path = "/workload-identity-vol" + container_name = f"workload-token-write-{uuid.uuid4().hex[:8]}" + finalize_token_command = ( + f"mv {token_volume_path}/token.tmp {token_volume_path}/token && chmod 0444 {token_volume_path}/token" + ) + try: + container = self._client.containers.create( + name=container_name, + image=permissions_image, + command=[ + "sh", + "-c", + finalize_token_command, + ], + volumes={volume_name: {"bind": token_volume_path, "mode": "rw"}}, + labels={JOB_MANAGED_BY_LABEL: JOB_MANAGED_BY_JOBS_CONTROLLER}, + ) + except ImageNotFound: + self._client.images.pull(permissions_image) + container = self._client.containers.create( + name=container_name, + image=permissions_image, + command=[ + "sh", + "-c", + finalize_token_command, + ], + volumes={volume_name: {"bind": token_volume_path, "mode": "rw"}}, + labels={JOB_MANAGED_BY_LABEL: JOB_MANAGED_BY_JOBS_CONTROLLER}, + ) + except APIError as exc: + raise JobStorageError("Error creating workload identity token writer container") from exc + + try: + container.put_archive(path=token_volume_path, data=build_token_archive(token)) + container.start() + exit_status = container.wait() + if exit_status["StatusCode"] != 0: + raise JobStorageError( + f"Workload identity token writer exited with non-zero status {exit_status['StatusCode']}" + ) + finally: + try: + container.remove() + except Exception: + logger.debug("Failed to remove workload identity token writer container", exc_info=True) + + def _build_workload_identity_refresher(self, volume_name: str) -> SubjectTokenRefreshLoop: + issuer = self._create_docker_subject_token_issuer() + return SubjectTokenRefreshLoop( + issuer=issuer, + write_token=lambda token: self._write_workload_identity_subject_token(volume_name, token), + refresh_margin_seconds=self._execution_profile_config.workload_identity.refresh_margin_seconds, + ) + + def _start_workload_identity_refresher( + self, container_name: str, refresher: SubjectTokenRefreshLoop | None + ) -> None: + if refresher is None: + return + old_refresher = self._workload_identity_refreshers.get(container_name) + if old_refresher is not None: + old_refresher.stop() + self._workload_identity_refreshers.pop(container_name, None) + refresher.start() + self._workload_identity_refreshers[container_name] = refresher + + def _stop_workload_identity_refresher(self, container_name: str) -> None: + refresher = self._workload_identity_refreshers.get(container_name) + if refresher is not None: + refresher.stop() + self._workload_identity_refreshers.pop(container_name, None) + + def _stop_all_workload_identity_refreshers(self) -> None: + for container_name, refresher in list(self._workload_identity_refreshers.items()): + try: + refresher.stop() + except Exception: + logger.warning( + "Failed to stop workload identity subject token refresher for Docker container %s", + container_name, + exc_info=True, + ) + finally: + self._workload_identity_refreshers.pop(container_name, None) + + @staticmethod + def _get_workload_identity_volume_from_mounts(container: Container) -> str | None: + attrs = getattr(container, "attrs", None) or {} + mounts = attrs.get("Mounts", []) or [] + for mount in mounts: + if mount.get("Type") != "volume" or mount.get("Destination") != WORKLOAD_IDENTITY_VOLUME_PATH: + continue + volume_name = mount.get("Name") or mount.get("Source") + if isinstance(volume_name, str) and volume_name: + return volume_name + return None + + def _get_workload_identity_volume_for_container(self, container: Container) -> str | None: + labels = getattr(container, "labels", None) or {} + token_file = labels.get(DOCKER_WORKLOAD_IDENTITY_TOKEN_FILE_LABEL) + if token_file is not None and token_file != WORKLOAD_IDENTITY_TOKEN_FILE_PATH: + logger.warning( + "Skipping workload identity refresher restore for container with unexpected token file label", + extra={ + "container_name": getattr(container, "name", None), + "token_file": token_file, + }, + ) + return None + + volume_name = labels.get(DOCKER_WORKLOAD_IDENTITY_VOLUME_LABEL) + if volume_name: + return volume_name + return self._get_workload_identity_volume_from_mounts(container) + + def _restore_workload_identity_refresher_for_container(self, container: Container) -> None: + container_name = getattr(container, "name", None) + if not isinstance(container_name, str) or not container_name: + return + if container_name in self._workload_identity_refreshers: + return + if container.status != "running": + return + if not self._is_container_owned_by_this_controller(container): + return + + labels = getattr(container, "labels", None) or {} + if labels.get(JOB_TYPE_LABEL) != JOB_TYPE_JOB: + return + + volume_name = self._get_workload_identity_volume_for_container(container) + if volume_name is None: + return + + try: + self._start_workload_identity_refresher( + container_name, + self._build_workload_identity_refresher(volume_name), + ) + logger.info( + "Restored Docker workload identity refresher for running job container", + extra={"container_name": container_name, "volume_name": volume_name}, + ) + except Exception: + logger.exception( + "Failed to restore Docker workload identity refresher for running job container", + extra={"container_name": container_name, "volume_name": volume_name}, + ) + @staticmethod def get_label_from_container(container: Container, label: str) -> str: return container.labels[label] @@ -257,12 +560,17 @@ def task_config_volume_name(self, workspace: str, job: str, task: str) -> str: """Generate a unique volume name for task config space.""" return f"task-config-{workspace}-{job}-{task}" + def task_workload_identity_volume_name(self, workspace: str, job: str, task: str) -> str: + """Generate a unique volume name for workload identity token material.""" + return f"task-workload-identity-{workspace}-{job}-{task}" + def cleanup_task_storage_volumes(self, workspace: str, job: str, task: str) -> None: """Remove the task storage volume after the container is done.""" volumes_to_delete = [ self.task_storage_volume_name(workspace, job, task), self.task_config_volume_name(workspace, job, task), + self.task_workload_identity_volume_name(workspace, job, task), ] for volume_name in volumes_to_delete: try: @@ -408,6 +716,7 @@ def get_mounts( config_volume_path: str, task_volume_name: str, task_volume_path: str, + workload_identity_volume_name: str | None = None, additional_volume_mounts: list[DockerVolumeMount] | None = None, ) -> list[Mount]: """ @@ -433,6 +742,14 @@ def get_mounts( task_storage_mount, config_storage_mount, ] + if workload_identity_volume_name is not None: + workload_identity_mount = docker.types.Mount( + type="volume", + source=workload_identity_volume_name, + target=WORKLOAD_IDENTITY_VOLUME_PATH, + read_only=True, + ) + mounts.append(workload_identity_mount) if job_volume_path != "": job_storage_mount = docker.types.Mount( @@ -465,6 +782,7 @@ def ensure_job_storage( task: str, step_config_json: str, additional_volumes_mounts: list[DockerVolumeMount] | None = None, + workload_identity_volume_name: str | None = None, ) -> None: """ Ensure Docker volumes exist for the job and task, with proper permissions. @@ -492,6 +810,15 @@ def ensure_job_storage( except Exception as exc: raise JobStorageError(f"Error creating task config volume {config_volume_name}") from exc + if workload_identity_volume_name is not None: + try: + self._client.volumes.create(workload_identity_volume_name) + logger.debug("Created workload identity volume", extra={"volume_name": workload_identity_volume_name}) + except Exception as exc: + raise JobStorageError( + f"Error creating workload identity volume {workload_identity_volume_name}" + ) from exc + script = f"""#!/bin/sh set -ex chmod -R 777 {task_vol} @@ -505,6 +832,12 @@ def ensure_job_storage( task_volume_name: {"bind": task_vol, "mode": "rw"}, config_volume_name: {"bind": config_vol, "mode": "rw"}, } + if workload_identity_volume_name is not None: + volumes[workload_identity_volume_name] = {"bind": "/workload-identity-vol", "mode": "rw"} + script += """ +mkdir -p /workload-identity-vol +chmod -R 777 /workload-identity-vol +""" if job_storage_volume_name != "": job_vol = "/job-vol" @@ -599,6 +932,10 @@ def schedule_single_container( # Profile-level env vars first (e.g. HOME=/tmp); system, step, and shared env override these env = self._execution_profile_config.env.copy() + validate_no_reserved_managed_job_environment_variable_names( + (envvar.name for envvar in step.step_spec.environment or []), + source="Job step environment keys", + ) # identify the task using a uuid. In docker, there's only one task per step. # because parallelism and completions are not supported. @@ -614,18 +951,19 @@ def schedule_single_container( EPHEMERAL_TASK_STORAGE_PATH_ENVVAR: DEFAULT_TASK_STORAGE_PATH, CONFIG_TASK_STORAGE_PATH_ENVVAR: DEFAULT_CONFIG_STORAGE_PATH, NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR: DEFAULT_NEMO_JOB_STEP_CONFIG_FILE_PATH, - # Forward OTEL env vars for jobs-launcher to emit telemetry, particularly logs - "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT": get_logs_endpoint_from_fileset( + # Endpoint used by jobs-launcher to upload task stdout/stderr logs. + JOB_LOGS_ENDPOINT_ENVVAR: get_logs_endpoint_from_fileset( platform_config, step.workspace, step.fileset, ), - "OTEL_LOGS_EXPORTER": "otlp", - "OTEL_SERVICE_NAME": "nmp-job-task", # Inject secret environment variable mappings for the jobs-launcher to fetch NEMO_JOB_SECRETS_ENVVAR: self.get_secrets_environment_variable_for_injection(step), } ) + workload_identity_enabled = self._is_workload_identity_enabled() + if workload_identity_enabled: + env[WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR] = WORKLOAD_IDENTITY_TOKEN_FILE_PATH # Set auth context env var for job containers to make authenticated API calls if step.auth_context: @@ -635,8 +973,6 @@ def schedule_single_container( env_var_dict = principal.get_env_var() for name, value in env_var_dict.items(): env[name] = value - # Also set OTLP headers for telemetry (logs) to be authenticated - env["OTEL_EXPORTER_OTLP_LOGS_HEADERS"] = principal.get_otlp_headers_value() step_config_json = json.dumps(step.step_spec.config) @@ -660,7 +996,7 @@ def schedule_single_container( # Thread through shared platform envvars to the job # Note: address_override defaults to None, which triggers automatic loopback detection - env.update(platform_config.to_shared_envvars()) + env.update(get_job_runtime_shared_envvars(platform_config)) log_config = LogConfig( type=LogConfig.types.JSON, @@ -725,6 +1061,12 @@ def _prepare_container_args_for_start( job_volume_name = storage_config.volume_name if storage_config is not None else "" task_volume_name = self.task_storage_volume_name(workspace=step.workspace, job=step.job, task=task_id) config_volume_name = self.task_config_volume_name(workspace=step.workspace, job=step.job, task=task_id) + workload_identity_enabled = self._is_workload_identity_enabled() + workload_identity_volume_name = ( + self.task_workload_identity_volume_name(workspace=step.workspace, job=step.job, task=task_id) + if workload_identity_enabled + else None + ) additional_volume_mounts = storage_config.additional_volume_mounts if storage_config else None ensure_storage_started_at = time.monotonic() self.ensure_job_storage( @@ -738,7 +1080,16 @@ def _prepare_container_args_for_start( task=task_id, additional_volumes_mounts=additional_volume_mounts, step_config_json=step_config_json, + workload_identity_volume_name=workload_identity_volume_name, ) + workload_identity_refresher = None + if workload_identity_volume_name is not None: + try: + workload_identity_refresher = self._build_workload_identity_refresher(workload_identity_volume_name) + workload_identity_refresher.refresh_once() + except Exception: + self.cleanup_task_storage_volumes(step.workspace, step.job, task_id) + raise logger.debug( "Docker job storage ensured", extra={ @@ -767,6 +1118,9 @@ def _prepare_container_args_for_start( labels[JOB_USES_PERSISTENT_STORAGE_LABEL] = "true" else: labels[JOB_USES_PERSISTENT_STORAGE_LABEL] = "false" + if workload_identity_volume_name is not None: + labels[DOCKER_WORKLOAD_IDENTITY_TOKEN_FILE_LABEL] = WORKLOAD_IDENTITY_TOKEN_FILE_PATH + labels[DOCKER_WORKLOAD_IDENTITY_VOLUME_LABEL] = workload_identity_volume_name task_image = resolve_task_image( executor_config.container.image, self._execution_profile_config.default_task_image @@ -790,9 +1144,12 @@ def _prepare_container_args_for_start( config_volume_path=DEFAULT_CONFIG_STORAGE_PATH, task_volume_name=task_volume_name, task_volume_path=task_storage_mount, + workload_identity_volume_name=workload_identity_volume_name, additional_volume_mounts=additional_volume_mounts, ), } + if workload_identity_refresher is not None: + container_args["_nmp_workload_identity_refresher"] = workload_identity_refresher container_args["network"] = self._execution_profile_config.networking.job_container_network return self.configure_container(container_args, executor_config) @@ -886,6 +1243,7 @@ def run_container( def _run_container_in_thread(self, step: PlatformJobStepWithContext, container_args: dict): status_details = {} status = PlatformJobStatus.PENDING + workload_identity_refresher = container_args.pop("_nmp_workload_identity_refresher", None) # If a request to pause or cancel came in while we were waiting for scheduling loop, # cancel scheduling the container @@ -1112,6 +1470,7 @@ def _run_container_in_thread(self, step: PlatformJobStepWithContext, container_a max_attempts = 3 attempts = 0 start_started_at = time.monotonic() + self._start_workload_identity_refresher(container.name, workload_identity_refresher) while not started and attempts < max_attempts: attempts += 1 try: @@ -1139,6 +1498,7 @@ def _run_container_in_thread(self, step: PlatformJobStepWithContext, container_a }, ) except Exception as e: + self._stop_workload_identity_refresher(container.name) raise FailedToScheduleError( f"Failed to start container {container.name} for job step", error_details={"message": f"Failed to start container: {e}"}, @@ -1394,6 +1754,9 @@ def docker_state_debug_fields(self, container: Container) -> dict[str, Any]: } def create_step_update(self, step: PlatformJobStepWithContext, container: Container) -> JobUpdate: + if step.status in (PlatformJobStatus.ACTIVE, PlatformJobStatus.PENDING): + self._restore_workload_identity_refresher_for_container(container) + status, status_details, error_stack = self.map_docker_container_status_to_platform_status(step, container) task_id = self.get_label_from_container(container, JOB_TASK_ID_LABEL) error_details = {} @@ -1649,6 +2012,7 @@ def cleanup_single_container(self, container: Container) -> None: task = self.get_label_from_container(container, JOB_TASK_ID_LABEL) exit_code = container.attrs.get("State", {}).get("ExitCode", 0) + self._stop_workload_identity_refresher(container.name) self.cleanup_container(container) logger.debug( "Cleaned up container", diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/common.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/common.py index 7f6a8a574e..d6def96124 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/common.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/common.py @@ -78,12 +78,21 @@ ) from nmp.core.jobs.app.providers import ComputeResources, ContainerSpec from nmp.core.jobs.controllers.backends.base import ( + JOB_LOGS_ENDPOINT_ENVVAR, + WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, + WORKLOAD_IDENTITY_TOKEN_FILE_PATH, + WORKLOAD_IDENTITY_VOLUME_NAME, + WORKLOAD_IDENTITY_VOLUME_PATH, + get_job_runtime_shared_envvars, get_logs_endpoint_from_fileset, + get_workload_identity_token_audience, + is_workload_identity_token_exchange_enabled, resolve_gpu_job_shm_size, resolve_task_image, + validate_no_reserved_managed_job_environment_variable_names, ) from nmp.core.jobs.controllers.backends.exceptions import FailedToScheduleError, JobStorageError -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field, model_validator logger = logging.getLogger(__name__) @@ -606,8 +615,83 @@ def __init__(self, *args): # ``to_k8s()`` is available on nested volumes. +class KubernetesKeyToPath(BaseModel): + """Kubernetes volume key-to-path mapping.""" + + key: str = Field(description="Source key to project from the volume source") + path: str = Field(description="Relative file path to write the key to") + mode: int | None = Field(default=None, description="Optional file mode for this key") + + def to_k8s(self) -> client.V1KeyToPath: + """Convert to Kubernetes V1KeyToPath object.""" + return client.V1KeyToPath(key=self.key, path=self.path, mode=self.mode) + + +class KubernetesSecretVolume(BaseModel): + """Kubernetes Secret volume definition.""" + + secret_name: str = Field(description="Secret name to mount") + default_mode: int | None = Field(default=None, description="Optional default file mode") + optional: bool | None = Field(default=None, description="Whether the Secret is optional") + items: list[KubernetesKeyToPath] = Field(default_factory=list, description="Optional Secret keys to project") + + def to_k8s(self) -> client.V1SecretVolumeSource: + """Convert to Kubernetes V1SecretVolumeSource object.""" + return client.V1SecretVolumeSource( + secret_name=self.secret_name, + default_mode=self.default_mode, + optional=self.optional, + items=[item.to_k8s() for item in self.items] or None, + ) + + +class KubernetesConfigMapVolume(BaseModel): + """Kubernetes ConfigMap volume definition.""" + + name: str = Field(description="ConfigMap name to mount") + default_mode: int | None = Field(default=None, description="Optional default file mode") + optional: bool | None = Field(default=None, description="Whether the ConfigMap is optional") + items: list[KubernetesKeyToPath] = Field(default_factory=list, description="Optional ConfigMap keys to project") + + def to_k8s(self) -> client.V1ConfigMapVolumeSource: + """Convert to Kubernetes V1ConfigMapVolumeSource object.""" + return client.V1ConfigMapVolumeSource( + name=self.name, + default_mode=self.default_mode, + optional=self.optional, + items=[item.to_k8s() for item in self.items] or None, + ) + + class KubernetesVolume(PluginKubernetesVolume): - """Kubernetes Volume definition.""" + """Kubernetes Volume definition with secret and config_map support.""" + + model_config = ConfigDict( + json_schema_extra={ + "oneOf": [ + { + "required": ["persistent_volume_claim"], + "properties": {"persistent_volume_claim": {"not": {"type": "null"}}}, + }, + {"required": ["empty_dir"], "properties": {"empty_dir": {"not": {"type": "null"}}}}, + {"required": ["secret"], "properties": {"secret": {"not": {"type": "null"}}}}, + {"required": ["config_map"], "properties": {"config_map": {"not": {"type": "null"}}}}, + ] + } + ) + + secret: KubernetesSecretVolume | None = Field(default=None, description="Secret Volume configuration") + config_map: KubernetesConfigMapVolume | None = Field(default=None, description="ConfigMap Volume configuration") + + @model_validator(mode="after") + def validate_self(self): + """Ensure that exactly one volume source is specified.""" + sources = [self.persistent_volume_claim, self.empty_dir, self.secret, self.config_map] + if sum(source is not None for source in sources) != 1: + raise ValueError( + "Exactly one of 'persistent_volume_claim', 'empty_dir', 'secret', or 'config_map' must be specified." + ) + return self def to_k8s(self) -> client.V1Volume: """Convert to Kubernetes V1Volume object.""" @@ -622,6 +706,10 @@ def to_k8s(self) -> client.V1Volume: medium=self.empty_dir.medium, size_limit=self.empty_dir.size_limit, ) + if self.secret: + volume.secret = self.secret.to_k8s() + if self.config_map: + volume.config_map = self.config_map.to_k8s() return volume @@ -651,6 +739,17 @@ class KubernetesJobStorageConfig(PluginKubernetesJobStorageConfig): class BaseKubernetesExecutionProfileConfig(PluginBaseKubernetesExecutionProfileConfig): """Kubernetes execution config whose storage carries ``to_k8s()`` (server-side).""" + workload_identity_token_expiration_seconds: int = Field( + default=600, + ge=600, + description="Requested expirationSeconds for the projected service account token used as the workload identity subject token.", + ) + workload_identity_token_audience: str | None = Field( + default=None, + description="Audience for the projected service account token. Defaults to auth.oidc.workload_client_id, auth.oidc.client_id, then 'nemo-platform'.", + ) + + # Storage configurations for the job (re-typed to server subclass with to_k8s()) storage: KubernetesJobStorageConfig = Field( default_factory=KubernetesJobStorageConfig, description="Storage configuration for the Kubernetes job pods." ) @@ -879,9 +978,15 @@ def create_pod_template_spec( """ platform_config = get_platform_config() + workload_identity_enabled = is_workload_identity_token_exchange_enabled() + workload_identity_token_audience = config.workload_identity_token_audience or get_workload_identity_token_audience() # Profile-level env vars first (e.g. HOME=/tmp); system, step, and shared env override these env = [client.V1EnvVar(name=name, value=value) for name, value in config.env.items()] + validate_no_reserved_managed_job_environment_variable_names( + (envvar.name for envvar in step.step_spec.environment or []), + source="Job step environment keys", + ) env.extend( [ client.V1EnvVar(name=NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR, value=DEFAULT_NEMO_JOB_STEP_CONFIG_FILE_PATH), @@ -896,18 +1001,18 @@ def create_pod_template_spec( ), client.V1EnvVar(name=EPHEMERAL_TASK_STORAGE_PATH_ENVVAR, value=DEFAULT_TASK_STORAGE_PATH), client.V1EnvVar( - name="OTEL_EXPORTER_OTLP_LOGS_ENDPOINT", + name=JOB_LOGS_ENDPOINT_ENVVAR, value=get_logs_endpoint_from_fileset( platform_config, step.workspace, step.fileset, ), ), - client.V1EnvVar(name="OTEL_LOGS_EXPORTER", value="otlp"), - client.V1EnvVar(name="OTEL_SERVICE_NAME", value="nmp-job-task"), client.V1EnvVar(name=NEMO_JOB_SECRETS_ENVVAR, value=secret_env_var_str), ] ) + if workload_identity_enabled: + env.append(client.V1EnvVar(name=WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, value=WORKLOAD_IDENTITY_TOKEN_FILE_PATH)) # Set auth context env var for job containers to make authenticated API calls if step.auth_context: @@ -917,11 +1022,9 @@ def create_pod_template_spec( env_var_dict = principal.get_env_var() for name, value in env_var_dict.items(): env.append(client.V1EnvVar(name=name, value=value)) - # Also set OTLP headers for telemetry (logs) to be authenticated - env.append(client.V1EnvVar(name="OTEL_EXPORTER_OTLP_LOGS_HEADERS", value=principal.get_otlp_headers_value())) # Thread through shared platform envvars to the job - shared_envvars = platform_config.to_shared_envvars() + shared_envvars = get_job_runtime_shared_envvars(platform_config) env.extend([client.V1EnvVar(name=name, value=value) for name, value in shared_envvars.items()]) job_storage_mount = "" @@ -949,6 +1052,14 @@ def create_pod_template_spec( mount_path=DEFAULT_CONFIG_STORAGE_PATH, ), ] + if workload_identity_enabled: + volume_mounts.append( + client.V1VolumeMount( + name=WORKLOAD_IDENTITY_VOLUME_NAME, + mount_path=WORKLOAD_IDENTITY_VOLUME_PATH, + read_only=True, + ) + ) storage_config = config.storage # Persistent job storage (PVC mount) is only provisioned when the step @@ -973,6 +1084,23 @@ def create_pod_template_spec( client.V1Volume(name=LAUNCHER_VOLUME_NAME, empty_dir=client.V1EmptyDirVolumeSource()), client.V1Volume(name=STEP_CONFIG_VOLUME_NAME, config_map=client.V1ConfigMapVolumeSource(name=configmap_name)), ] + if workload_identity_enabled: + volumes.append( + client.V1Volume( + name=WORKLOAD_IDENTITY_VOLUME_NAME, + projected=client.V1ProjectedVolumeSource( + sources=[ + client.V1VolumeProjection( + service_account_token=client.V1ServiceAccountTokenProjection( + path="token", + expiration_seconds=config.workload_identity_token_expiration_seconds, + audience=workload_identity_token_audience, + ) + ) + ] + ), + ) + ) if storage_config.additional_volumes: volumes.extend(vol.to_k8s() for vol in storage_config.additional_volumes) diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/registry.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/registry.py index f5fd76accc..7d572ba1d9 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/registry.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/registry.py @@ -55,6 +55,7 @@ class BackendKey: BackendKey("gpu", "docker"): GPUDockerJobBackend, BackendKey("cpu", "kubernetes_job"): CPUKubernetesJobBackend, BackendKey("gpu", "kubernetes_job"): GPUKubernetesJobBackend, + BackendKey("gpu_distributed", "kubernetes_job"): GPUKubernetesJobBackend, BackendKey("gpu_distributed", "volcano_job"): VolcanoJobBackend, BackendKey("subprocess", "subprocess"): SubprocessJobBackend, BackendKey("cpu", "e2e"): TestE2ECPUJobBackend, diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess.py index 2fb747ab1e..6b71ad0d95 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess.py @@ -443,7 +443,10 @@ def _prepare_runtime(self, step: PlatformJobStepWithContext) -> tuple[dict[str, PERSISTENT_JOB_STORAGE_PATH_ENVVAR: str(persistent_dir), NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR: str(config_path), "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT": get_logs_endpoint_from_fileset( - platform_config, step.workspace, step.fileset, loopback_address="localhost" + platform_config, + step.workspace, + step.fileset, + loopback_address="localhost", ), "OTEL_LOGS_EXPORTER": "otlp", "OTEL_SERVICE_NAME": "nmp-job-task", diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/workload_tokens.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/workload_tokens.py new file mode 100644 index 0000000000..d726ae5fe8 --- /dev/null +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/workload_tokens.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Helpers for controller-managed workload identity subject tokens.""" + +from __future__ import annotations + +import io +import json +import logging +import tarfile +import threading +import time +from dataclasses import dataclass, field +from ipaddress import ip_address +from math import isfinite +from typing import Callable, Protocol +from urllib.parse import urlparse + +import httpx + +logger = logging.getLogger(__name__) + +DEFAULT_SUBJECT_TOKEN_TTL_SECONDS = 600 +DEFAULT_SUBJECT_TOKEN_REFRESH_MARGIN_SECONDS = 60 +DEFAULT_SUBJECT_TOKEN_FAILURE_BACKOFF_MAX_SECONDS = 60.0 +DEFAULT_SUBJECT_TOKEN_STOP_TIMEOUT_SECONDS = 5.0 + + +def _is_loopback_host(host: str | None) -> bool: + if host is None: + return False + if host.lower() == "localhost": + return True + try: + return ip_address(host).is_loopback + except ValueError: + return False + + +def _validate_token_endpoint(token_endpoint: str) -> None: + parsed = urlparse(token_endpoint) + scheme = parsed.scheme.lower() + if scheme == "https" and parsed.netloc: + return + if scheme == "http" and _is_loopback_host(parsed.hostname): + return + raise RuntimeError( + "Invalid Docker workload identity token_endpoint configuration: " + "token_endpoint must use https://, except http:// is allowed for loopback hosts only" + ) + + +@dataclass(frozen=True) +class SubjectToken: + """Issued workload identity subject token and its refresh deadline metadata.""" + + value: str + expires_at: float + + def seconds_until_refresh(self, margin_seconds: int) -> float: + return max(0.0, self.expires_at - time.time() - margin_seconds) + + +class SubjectTokenIssuer(Protocol): + """Issues a subject token suitable for SDK RFC 8693 token exchange.""" + + def issue(self) -> SubjectToken: ... + + +@dataclass(frozen=True) +class OAuthPasswordGrantSubjectTokenIssuer: + """Demo issuer that obtains a short-lived subject token with OAuth password grant.""" + + token_endpoint: str + client_id: str + username: str + password: str = field(repr=False) + client_secret: str | None = field(default=None, repr=False) + scope: str | None = None + default_expires_in_seconds: int = DEFAULT_SUBJECT_TOKEN_TTL_SECONDS + timeout: float = 30.0 + + def issue(self) -> SubjectToken: + _validate_token_endpoint(self.token_endpoint) + data = { + "grant_type": "password", + "client_id": self.client_id, + "username": self.username, + "password": self.password, + } + if self.client_secret: + data["client_secret"] = self.client_secret + if self.scope: + data["scope"] = self.scope + + response = httpx.post(self.token_endpoint, data=data, timeout=self.timeout) + if response.status_code != 200: + error_data: dict[str, object] = {} + if response.headers.get("content-type", "").startswith("application/json"): + try: + payload = response.json() + except (json.JSONDecodeError, ValueError): + payload = {} + if isinstance(payload, dict): + error_data = payload + error = error_data.get("error", "unknown_error") + description = error_data.get("error_description", response.text) + raise RuntimeError(f"Failed to issue workload subject token: {error} - {description}") + + try: + token_data = response.json() + except (json.JSONDecodeError, ValueError) as exc: + raise RuntimeError( + "Failed to issue workload subject token: invalid_response - " + "Token endpoint response was not a JSON object" + ) from exc + if not isinstance(token_data, dict): + raise RuntimeError( + "Failed to issue workload subject token: invalid_response - " + "Token endpoint response was not a JSON object" + ) + + access_token = token_data.get("access_token") + if not isinstance(access_token, str) or not access_token.strip(): + raise RuntimeError( + "Failed to issue workload subject token: invalid_response - " + "Token endpoint response did not include a non-empty access_token" + ) + + if "expires_in" in token_data: + expires_in = token_data["expires_in"] + if ( + isinstance(expires_in, bool) + or not isinstance(expires_in, int | float) + or not isfinite(expires_in) + or expires_in <= 0 + ): + raise RuntimeError( + "Failed to issue workload subject token: invalid_response - " + "Token endpoint response did not include a positive numeric expires_in" + ) + else: + expires_in = self.default_expires_in_seconds + return SubjectToken(value=access_token, expires_at=time.time() + expires_in) + + +def build_token_archive(token: str, *, name: str = "token.tmp") -> io.BytesIO: + """Build a tar archive containing one token file.""" + data = token.encode("utf-8") + archive = io.BytesIO() + with tarfile.open(fileobj=archive, mode="w") as tar: + info = tarfile.TarInfo(name=name) + info.size = len(data) + info.mode = 0o400 + tar.addfile(info, io.BytesIO(data)) + archive.seek(0) + return archive + + +class SubjectTokenRefreshLoop: + """Background refresher for a controller-owned workload subject token file.""" + + def __init__( + self, + *, + issuer: SubjectTokenIssuer, + write_token: Callable[[str], None], + refresh_margin_seconds: int = DEFAULT_SUBJECT_TOKEN_REFRESH_MARGIN_SECONDS, + min_sleep_seconds: float = 1.0, + max_failure_backoff_seconds: float = DEFAULT_SUBJECT_TOKEN_FAILURE_BACKOFF_MAX_SECONDS, + ) -> None: + self._issuer = issuer + self._write_token = write_token + self._refresh_margin_seconds = refresh_margin_seconds + self._min_sleep_seconds = min_sleep_seconds + self._max_failure_backoff_seconds = max(min_sleep_seconds, max_failure_backoff_seconds) + self._stop_timeout_seconds = DEFAULT_SUBJECT_TOKEN_STOP_TIMEOUT_SECONDS + self._stop = threading.Event() + self._thread: threading.Thread | None = None + + def start(self) -> None: + if self._thread is not None and self._thread.is_alive(): + return + self._stop.clear() + self._thread = threading.Thread(target=self._run, name="nmp-workload-token-refresh", daemon=True) + self._thread.start() + + def stop(self) -> None: + self._stop.set() + thread = self._thread + if thread is None: + return + + thread.join(timeout=self._stop_timeout_seconds) + if thread.is_alive(): + raise RuntimeError("Timed out stopping workload identity subject token refresher") + self._thread = None + + def refresh_once(self) -> SubjectToken: + token = self._issuer.issue() + self._write_token(token.value) + return token + + def _refresh_once_for_worker(self) -> SubjectToken | None: + token = self._issuer.issue() + if self._stop.is_set(): + return None + self._write_token(token.value) + return token + + def _run(self) -> None: + token: SubjectToken | None = None + failure_sleep_seconds = self._min_sleep_seconds + while not self._stop.is_set(): + try: + token = self._refresh_once_for_worker() + except Exception: + logger.exception("Failed to refresh workload identity subject token") + if self._stop.wait(failure_sleep_seconds): + return + failure_sleep_seconds = min(self._max_failure_backoff_seconds, failure_sleep_seconds * 2) + continue + if token is None: + return + + failure_sleep_seconds = self._min_sleep_seconds + sleep_seconds = max(self._min_sleep_seconds, token.seconds_until_refresh(self._refresh_margin_seconds)) + self._stop.wait(sleep_seconds) diff --git a/services/core/jobs/tests/controllers/test_base.py b/services/core/jobs/tests/controllers/test_base.py index 4a728249b9..d5ba6f2ace 100644 --- a/services/core/jobs/tests/controllers/test_base.py +++ b/services/core/jobs/tests/controllers/test_base.py @@ -7,18 +7,144 @@ from contextlib import contextmanager from types import SimpleNamespace from unittest.mock import MagicMock, patch +from urllib.parse import urlunsplit +import pytest from nmp.common.config import PlatformConfig +from nmp.common.jobs.constants import EPHEMERAL_TASK_STORAGE_PATH_ENVVAR, PERSISTENT_JOB_STORAGE_PATH_ENVVAR from nmp.common.jobs.schemas import PlatformJobStatus from nmp.core.jobs.api.v2.jobs.schemas import PlatformJobStepWithContext from nmp.core.jobs.app.providers import ContainerSpec, CPUExecutionProvider from nmp.core.jobs.app.schemas import PlatformJobStepSpec, StepLifecycle -from nmp.core.jobs.controllers.backends.base import get_logs_endpoint_from_fileset, resolve_task_image +from nmp.core.jobs.controllers.backends.base import ( + JOB_LOGS_ENDPOINT_ENVVAR, + WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, + JobExecutionProfileConfig, + _contains_loopback_address, + _replace_loopback_address, + find_reserved_managed_job_environment_variable_names, + get_job_runtime_shared_envvars, + get_logs_endpoint_from_fileset, + get_workload_identity_token_audience, + resolve_task_image, + validate_no_reserved_managed_job_environment_variable_names, +) from nmp.core.jobs.controllers.backends.test import MockKubernetesCPUJobBackend +from pydantic import ValidationError from services.core.jobs.tests.controllers.client_mocks import data_response +class TestManagedJobReservedEnvironment: + def test_finds_reserved_auth_names_and_prefixes(self): + assert find_reserved_managed_job_environment_variable_names( + [ + "NMP_ACCESS_TOKEN", + WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, + "NEMO_WORKFLOW_TOKEN", + "NMP_WORKFLOW_TOKEN", + "APP_SETTING", + ] + ) == [ + "NEMO_WORKFLOW_TOKEN", + "NMP_ACCESS_TOKEN", + "NMP_WORKFLOW_TOKEN", + WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, + ] + + def test_step_validation_allows_storage_path_contract_env_names(self): + validate_no_reserved_managed_job_environment_variable_names( + [PERSISTENT_JOB_STORAGE_PATH_ENVVAR, EPHEMERAL_TASK_STORAGE_PATH_ENVVAR], + source="Job step environment keys", + ) + + def test_step_validation_rejects_workload_identity_prefixes(self): + with pytest.raises(ValueError, match="NMP_WORKLOAD_TOKEN"): + validate_no_reserved_managed_job_environment_variable_names( + ["NMP_WORKLOAD_TOKEN"], + source="Job step environment keys", + ) + + def test_step_validation_rejects_platform_job_logs_endpoint(self): + with pytest.raises(ValueError, match=JOB_LOGS_ENDPOINT_ENVVAR): + validate_no_reserved_managed_job_environment_variable_names( + [JOB_LOGS_ENDPOINT_ENVVAR], + source="Job step environment keys", + ) + + def test_profile_env_rejects_access_token(self): + with pytest.raises(ValidationError, match="NMP_ACCESS_TOKEN"): + JobExecutionProfileConfig(env={"NMP_ACCESS_TOKEN": "token"}) + + def test_profile_env_rejects_auth_service_url(self): + with pytest.raises(ValidationError, match="NMP_AUTH_URL"): + JobExecutionProfileConfig(env={"NMP_AUTH_URL": "https://auth.example.com"}) + + def test_profile_env_rejects_workflow_typo_prefix(self): + with pytest.raises(ValidationError, match="NEMO_WORKFLOW_TOKEN"): + JobExecutionProfileConfig(env={"NEMO_WORKFLOW_TOKEN": "token"}) + + def test_profile_env_allows_user_otel_configuration(self): + config = JobExecutionProfileConfig( + env={ + "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT": "https://otel.example.com/v1/logs", + "OTEL_EXPORTER_OTLP_LOGS_HEADERS": "Authorization=Bearer%20external", + "OTEL_LOGS_EXPORTER": "otlp", + "OTEL_SERVICE_NAME": "custom-workload", + } + ) + + assert config.env["OTEL_EXPORTER_OTLP_LOGS_ENDPOINT"] == "https://otel.example.com/v1/logs" + assert config.env["OTEL_EXPORTER_OTLP_LOGS_HEADERS"] == "Authorization=Bearer%20external" + assert config.env["OTEL_LOGS_EXPORTER"] == "otlp" + assert config.env["OTEL_SERVICE_NAME"] == "custom-workload" + + def test_profile_env_rejects_platform_job_logs_endpoint(self): + with pytest.raises(ValidationError, match=JOB_LOGS_ENDPOINT_ENVVAR): + JobExecutionProfileConfig(env={JOB_LOGS_ENDPOINT_ENVVAR: "https://files.example.com/v1/logs"}) + + +class TestWorkloadIdentityAudience: + def test_uses_workload_client_id_for_projected_subject_tokens(self): + auth_config = SimpleNamespace( + oidc=SimpleNamespace( + workload_client_id="nemo-platform-workload", + client_id="nemo-platform-cli", + workload_audience="nemo-platform", + audience="nemo-platform", + ) + ) + + with patch("nmp.common.config.get_auth_config", return_value=auth_config): + assert get_workload_identity_token_audience() == "nemo-platform-workload" + + +class TestReplaceLoopbackAddress: + def test_contains_loopback_address_matches_hostname_exactly(self): + assert _contains_loopback_address("http://localhost:8080/v1/jobs") + assert not _contains_loopback_address("http://api.localhost.example:8080/v1/jobs/localhost") + + def test_preserves_url_when_loopback_text_is_not_the_hostname(self): + url = "http://api.localhost.example:8080/v1/proxy/localhost?next=http://127.0.0.1:8080" + + result = _replace_loopback_address(url, "host.docker.internal") + + assert result == url + + def test_replaces_ipv6_loopback_hostname_with_valid_bracket_formatting(self): + result = _replace_loopback_address("http://[::1]:8080/v1/jobs", "fd00::10") + + assert result == "http://[fd00::10]:8080/v1/jobs" + + def test_replaces_loopback_hostname_with_userinfo_from_parsed_url(self): + url = urlunsplit(("http", "user:pass@localhost:8080", "/v1/jobs", "", "")) + expected = urlunsplit(("http", "user:pass@host.docker.internal:8080", "/v1/jobs", "", "")) + + result = _replace_loopback_address(url, "host.docker.internal") + + assert result == expected + + class TestGetLogsEndpointFromFileset: """Tests for get_logs_endpoint_from_fileset.""" @@ -128,6 +254,28 @@ def test_uses_service_discovery_files_url_when_set(self): "http://files-service.svc.cluster.local:8080/apis/files/v2/workspaces/my-workspace/filesets/my-fileset-id/otlp/v1/logs" ) + def test_platform_service_discovery_platform_routes_logs_endpoint(self): + """Platform service_discovery can route job log export through a workload-facing gateway.""" + config = PlatformConfig( # type: ignore[abstract] + base_url="http://platform-api:8080", + service_discovery={"platform": "https://nemo-gateway:8080"}, + loopback_address=None, + ) + + with patch( + "nmp.core.jobs.controllers.backends.base.determine_loopback_override", + return_value=None, + ): + result = get_logs_endpoint_from_fileset( + config, + workspace="my-workspace", + fileset_id="my-fileset-id", + ) + + assert result == ( + "https://nemo-gateway:8080/apis/files/v2/workspaces/my-workspace/filesets/my-fileset-id/otlp/v1/logs" + ) + def test_service_discovery_takes_precedence_over_base_url(self): """When service_discovery['files'] is set, it is used instead of base_url for files.""" config = PlatformConfig( # type: ignore[abstract] @@ -145,6 +293,40 @@ def test_service_discovery_takes_precedence_over_base_url(self): # get_service_url("files") returns service_discovery["files"] when present assert result == ("http://discovered-files:9000/apis/files/v2/workspaces/ws1/filesets/fs-1/otlp/v1/logs") + def test_service_discovery_files_takes_precedence_over_local_service_url(self): + """Job log export uses routable service_discovery even when Files runs in-process.""" + config = PlatformConfig( # type: ignore[abstract] + base_url="http://platform:8080", + services="files", + service_discovery={"files": "http://files-direct:8080"}, + loopback_address=None, + ) + + with patch( + "nmp.core.jobs.controllers.backends.base.determine_loopback_override", + return_value=None, + ): + result = get_logs_endpoint_from_fileset(config, workspace="ws1", fileset_id="fs-1") + + assert result == ("http://files-direct:8080/apis/files/v2/workspaces/ws1/filesets/fs-1/otlp/v1/logs") + + def test_resolved_base_url_is_used_when_files_service_discovery_is_not_set(self): + """Job log export uses resolved runtime base URL, not local in-process service URL lookup.""" + config = PlatformConfig( # type: ignore[abstract] + base_url="http://127.0.0.1:59007", + services="files", + service_discovery={}, + loopback_address=None, + ) + + with patch( + "nmp.core.jobs.controllers.backends.base.determine_loopback_override", + return_value=None, + ): + result = get_logs_endpoint_from_fileset(config, workspace="ws1", fileset_id="fs-1") + + assert result == ("http://127.0.0.1:59007/apis/files/v2/workspaces/ws1/filesets/fs-1/otlp/v1/logs") + def test_service_discovery_files_with_loopback_replacement(self): """Loopback replacement works when files URL comes from service_discovery.""" config = PlatformConfig( # type: ignore[abstract] @@ -160,6 +342,134 @@ def test_service_discovery_files_with_loopback_replacement(self): ) +class TestGetJobRuntimeSharedEnvvars: + def test_uses_service_discovery_gateway_urls_for_job_runtimes(self): + config = PlatformConfig( # type: ignore[abstract] + base_url="http://127.0.0.1:8080", + services="jobs,files,models,secrets", + service_discovery={ + "platform": "https://nemo-gateway:8080", + "jobs": "https://nemo-gateway:8080", + "files": "https://nemo-gateway:8080", + "models": "https://nemo-gateway:8080", + "secrets": "https://nemo-gateway:8080", + }, + loopback_address="nemo-gateway", + ) + + envvars = get_job_runtime_shared_envvars(config) + + assert envvars == { + "NMP_BASE_URL": "https://nemo-gateway:8080", + "NMP_AUTH_URL": "https://nemo-gateway:8080", + "NMP_JOBS_URL": "https://nemo-gateway:8080", + "NMP_FILES_URL": "https://nemo-gateway:8080", + "NMP_MODELS_URL": "https://nemo-gateway:8080", + "NMP_SECRETS_URL": "https://nemo-gateway:8080", + "NMP_CONFIG_WARNINGS_DISABLED": "1", + } + + def test_gateway_base_url_is_used_for_job_runtimes_without_service_discovery(self): + config = PlatformConfig( # type: ignore[abstract] + base_url="https://nemo-gateway:8080", + services="jobs,files,models,secrets", + service_discovery={}, + loopback_address=None, + ) + + with patch( + "nmp.core.jobs.controllers.backends.base.determine_loopback_override", + return_value=None, + ): + envvars = get_job_runtime_shared_envvars(config) + + assert envvars == { + "NMP_BASE_URL": "https://nemo-gateway:8080", + "NMP_AUTH_URL": "https://nemo-gateway:8080", + "NMP_JOBS_URL": "https://nemo-gateway:8080", + "NMP_FILES_URL": "https://nemo-gateway:8080", + "NMP_MODELS_URL": "https://nemo-gateway:8080", + "NMP_SECRETS_URL": "https://nemo-gateway:8080", + "NMP_CONFIG_WARNINGS_DISABLED": "1", + } + + def test_resolved_base_url_is_used_for_job_runtimes_without_service_discovery(self): + config = PlatformConfig( # type: ignore[abstract] + base_url="http://127.0.0.1:59007", + services="jobs,files,models,secrets", + service_discovery={}, + loopback_address=None, + ) + + with patch( + "nmp.core.jobs.controllers.backends.base.determine_loopback_override", + return_value=None, + ): + envvars = get_job_runtime_shared_envvars(config) + + assert envvars == { + "NMP_BASE_URL": "http://127.0.0.1:59007", + "NMP_AUTH_URL": "http://127.0.0.1:59007", + "NMP_JOBS_URL": "http://127.0.0.1:59007", + "NMP_FILES_URL": "http://127.0.0.1:59007", + "NMP_MODELS_URL": "http://127.0.0.1:59007", + "NMP_SECRETS_URL": "http://127.0.0.1:59007", + "NMP_CONFIG_WARNINGS_DISABLED": "1", + } + + def test_auth_service_url_is_exported_for_job_runtime(self): + config = PlatformConfig( # type: ignore[abstract] + base_url="https://nemo-gateway:8080", + services="jobs,files,models,secrets", + service_discovery={"auth": "http://nemo-auth:8080"}, + loopback_address=None, + ) + + with patch( + "nmp.core.jobs.controllers.backends.base.determine_loopback_override", + return_value=None, + ): + envvars = get_job_runtime_shared_envvars(config) + + assert envvars["NMP_BASE_URL"] == "https://nemo-gateway:8080" + assert envvars["NMP_AUTH_URL"] == "http://nemo-auth:8080" + + def test_platform_service_discovery_platform_overrides_job_runtime_base_url(self): + config = PlatformConfig( # type: ignore[abstract] + base_url="http://nemo-platform-api:8080", + services="jobs,files,models,secrets", + service_discovery={ + "platform": "https://nemo-gateway:8080", + "auth": "https://nemo-auth:8080", + }, + loopback_address=None, + ) + + with patch( + "nmp.core.jobs.controllers.backends.base.determine_loopback_override", + return_value=None, + ): + envvars = get_job_runtime_shared_envvars(config) + + assert envvars["NMP_BASE_URL"] == "https://nemo-gateway:8080" + assert envvars["NMP_AUTH_URL"] == "https://nemo-auth:8080" + assert envvars["NMP_JOBS_URL"] == "https://nemo-gateway:8080" + assert envvars["NMP_FILES_URL"] == "https://nemo-gateway:8080" + + def test_loopback_replacement_applies_to_base_url_fallback(self): + config = PlatformConfig( # type: ignore[abstract] + base_url="http://127.0.0.1:8080", + service_discovery={}, + loopback_address="host.docker.internal", + ) + + envvars = get_job_runtime_shared_envvars(config) + + assert envvars["NMP_BASE_URL"] == "http://host.docker.internal:8080" + assert envvars["NMP_AUTH_URL"] == "http://host.docker.internal:8080" + assert envvars["NMP_FILES_URL"] == "http://host.docker.internal:8080" + + def _make_step( staleness_timeout: int = 0, created_at: datetime.datetime | None = None, diff --git a/services/core/jobs/tests/controllers/test_docker_backend.py b/services/core/jobs/tests/controllers/test_docker_backend.py index 1c828edbd9..08102731d7 100644 --- a/services/core/jobs/tests/controllers/test_docker_backend.py +++ b/services/core/jobs/tests/controllers/test_docker_backend.py @@ -3,7 +3,9 @@ import datetime import json +import time import uuid +from types import SimpleNamespace from typing import Iterator from unittest.mock import MagicMock, patch @@ -11,6 +13,7 @@ from docker.errors import APIError, NotFound from nemo_platform.types.shared import AuthContext as SdkAuthContext from nmp.common.auth import NMP_PRINCIPAL_ENVVAR, AuthContext, Principal +from nmp.common.config import PlatformConfig from nmp.common.docker.gpu_pool import DockerGPUPool from nmp.common.jobs.constants import ( EPHEMERAL_TASK_STORAGE_PATH_ENVVAR, @@ -49,20 +52,31 @@ PlatformJobSecretEnvironmentVariableRef, PlatformJobStepSpec, ) +from nmp.core.jobs.controllers.backends.base import ( + JOB_LOGS_ENDPOINT_ENVVAR, + WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, + WORKLOAD_IDENTITY_TOKEN_FILE_PATH, + WORKLOAD_IDENTITY_VOLUME_PATH, +) from nmp.core.jobs.controllers.backends.docker import ( DEFAULT_VOLUME_PERMISSIONS_IMAGE, DOCKER_CONTAINER_START_WORKERS, + DOCKER_WORKLOAD_IDENTITY_TOKEN_FILE_LABEL, + DOCKER_WORKLOAD_IDENTITY_VOLUME_LABEL, CPUDockerJobBackend, DockerJobExecutionProfileConfig, DockerJobStorageConfig, DockerVolumeMount, + DockerWorkloadIdentityConfig, GPUDockerJobBackend, ) from nmp.core.jobs.controllers.backends.exceptions import ( FailedToScheduleError, + JobStorageError, ResourceAllocationError, SchedulingDeferred, ) +from nmp.core.jobs.controllers.backends.workload_tokens import SubjectToken from pydantic import ValidationError from services.core.jobs.tests.controllers.client_mocks import data_response @@ -83,6 +97,18 @@ def owned_container_labels(labels: dict) -> dict: } +def assert_created_task_volumes_cleaned_up(docker_client_mock) -> None: + created_volume_names = [volume_call.args[0] for volume_call in docker_client_mock.volumes.create.call_args_list] + assert len(created_volume_names) == 3 + assert any(name.startswith("task-storage-") for name in created_volume_names) + assert any(name.startswith("task-config-") for name in created_volume_names) + assert any(name.startswith("task-workload-identity-") for name in created_volume_names) + + cleaned_volume_names = [volume_call.args[0] for volume_call in docker_client_mock.volumes.get.call_args_list] + assert sorted(cleaned_volume_names) == sorted(created_volume_names) + assert docker_client_mock.volumes.get.return_value.remove.call_count == 3 + + @pytest.fixture def docker_client_mock(monkeypatch): """Mock docker client for testing.""" @@ -708,6 +734,44 @@ def test_docker_job_profile_environment_applied(mock_nmp_client, docker_client_m assert env_vars.get("ENV_VAR") == "test_value" +def test_docker_job_uses_service_discovery_urls_for_job_runtime(mock_nmp_client, docker_client_mock, test_job_step): + """Job containers use routable service_discovery URLs instead of local in-process service URLs.""" + platform_config = PlatformConfig( # type: ignore[abstract] + base_url="http://127.0.0.1:8080", + services="jobs,files,models,secrets", + service_discovery={ + "platform": "https://nemo-gateway:8080", + "auth": "https://nemo-auth:8080", + }, + loopback_address="nemo-gateway", + ) + with patch("nmp.core.jobs.controllers.backends.docker.get_platform_config", return_value=platform_config): + backend = CPUDockerJobBackend( + mock_nmp_client, + DockerJobExecutionProfileConfig( + storage=DockerJobStorageConfig(volume_name="test_jobs_storage"), + ), + profile_name="default", + ) + backend._client = docker_client_mock + + backend.schedule(test_job_step.step_spec.executor, test_job_step) + backend._container_run_threadpool.shutdown(wait=True) + backend._container_run_threadpool = MagicMock() + + create_call_args = docker_client_mock.containers.create.call_args + container_args = create_call_args[1] if create_call_args[1] else create_call_args[0][0] + env_vars = container_args.get("environment", {}) + + assert env_vars["NMP_BASE_URL"] == "https://nemo-gateway:8080" + assert env_vars["NMP_AUTH_URL"] == "https://nemo-auth:8080" + assert env_vars["NMP_JOBS_URL"] == "https://nemo-gateway:8080" + assert env_vars["NMP_FILES_URL"] == "https://nemo-gateway:8080" + assert env_vars["NMP_MODELS_URL"] == "https://nemo-gateway:8080" + assert env_vars["NMP_SECRETS_URL"] == "https://nemo-gateway:8080" + assert env_vars[JOB_LOGS_ENDPOINT_ENVVAR].startswith("https://nemo-gateway:8080/apis/files/") + + def test_docker_job_execution_profile_config_rejects_reserved_env_vars(): """DockerJobExecutionProfileConfig raises when environment contains reserved names.""" with pytest.raises(ValidationError) as exc_info: @@ -719,6 +783,339 @@ def test_docker_job_execution_profile_config_rejects_reserved_env_vars(): assert "reserved" in str(exc_info.value).lower() +def test_docker_job_rejects_reserved_step_auth_env_vars(docker_job, test_job_step): + test_job_step.step_spec.environment.append( + PlatformJobEnvironmentVariable(name=WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, value="/tmp/token") + ) + + with pytest.raises(ValueError, match=WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR): + docker_job.schedule(test_job_step.step_spec.executor, test_job_step) + + +def test_docker_job_injects_workload_identity_volume_when_token_exchange_enabled( + docker_job, docker_client_mock, test_job_step +): + auth_config = SimpleNamespace(oidc=SimpleNamespace(workload_token_exchange_enabled=True)) + + class FakeIssuer: + def issue(self): + return SubjectToken(value="subject-token", expires_at=time.time() + 3600) + + with ( + patch("nmp.common.config.get_auth_config", return_value=auth_config), + patch.object(docker_job, "_create_docker_subject_token_issuer", return_value=FakeIssuer()), + ): + docker_job.schedule(test_job_step.step_spec.executor, test_job_step) + + docker_job._container_run_threadpool.shutdown(wait=True) + docker_job._container_run_threadpool = MagicMock() + + try: + job_create_call = next( + call + for call in docker_client_mock.containers.create.call_args_list + if call.kwargs.get("name") == "job-test-job-id-test-step" + ) + kwargs = job_create_call.kwargs + env = kwargs["environment"] + assert env[WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR] == WORKLOAD_IDENTITY_TOKEN_FILE_PATH + + mounts = kwargs["mounts"] + workload_identity_mount = next(m for m in mounts if m["Target"] == WORKLOAD_IDENTITY_VOLUME_PATH) + assert workload_identity_mount["Type"] == "volume" + assert workload_identity_mount["Source"].startswith( + f"task-workload-identity-{test_job_step.workspace}-{test_job_step.job}-" + ) + assert workload_identity_mount["ReadOnly"] is True + assert kwargs["labels"][DOCKER_WORKLOAD_IDENTITY_TOKEN_FILE_LABEL] == WORKLOAD_IDENTITY_TOKEN_FILE_PATH + assert kwargs["labels"][DOCKER_WORKLOAD_IDENTITY_VOLUME_LABEL] == workload_identity_mount["Source"] + workload_token_write_call = next( + call + for call in docker_client_mock.containers.create.call_args_list + if call.kwargs.get("name", "").startswith("workload-token-write-") + ) + assert workload_token_write_call.kwargs["command"] == [ + "sh", + "-c", + "mv /workload-identity-vol/token.tmp /workload-identity-vol/token && chmod 0444 /workload-identity-vol/token", + ] + finally: + for refresher in list(docker_job._workload_identity_refreshers.values()): + refresher.stop() + + +def test_docker_schedule_cleans_task_volumes_when_workload_identity_issuer_fails( + docker_job, docker_client_mock, test_job_step +): + docker_job._execution_profile_config.workload_identity.enabled = True + + with ( + patch.object( + docker_job, + "_create_docker_subject_token_issuer", + side_effect=JobStorageError("issuer failed"), + ), + pytest.raises(JobStorageError, match="issuer failed"), + ): + docker_job.schedule_single_container(test_job_step.step_spec.executor, test_job_step) + + assert_created_task_volumes_cleaned_up(docker_client_mock) + + +def test_docker_schedule_cleans_task_volumes_when_initial_workload_identity_refresh_fails( + docker_job, docker_client_mock, test_job_step +): + docker_job._execution_profile_config.workload_identity.enabled = True + refresher = MagicMock() + refresher.refresh_once.side_effect = RuntimeError("refresh failed") + + with ( + patch.object(docker_job, "_build_workload_identity_refresher", return_value=refresher), + pytest.raises(RuntimeError, match="refresh failed"), + ): + docker_job.schedule_single_container(test_job_step.step_spec.executor, test_job_step) + + refresher.refresh_once.assert_called_once() + assert_created_task_volumes_cleaned_up(docker_client_mock) + + +def test_docker_sync_restores_workload_identity_refresher_from_container_labels( + docker_job, docker_client_mock, test_job_step +): + volume_name = "task-workload-identity-default-job-test-job-id-task-restored" + container = MagicMock() + container.name = "job-test-job-id-test-step" + container.id = "workload-container-id" + container.status = "running" + container.attrs = {"State": {"Status": "running", "Running": True, "ExitCode": 0}, "HostConfig": {}} + container.labels = owned_container_labels( + { + JOB_WORKSPACE_ID_LABEL: test_job_step.workspace, + JOB_ID_LABEL: test_job_step.job, + JOB_STEP_NAME_LABEL: test_job_step.name, + JOB_TASK_ID_LABEL: "task-restored", + JOB_TYPE_LABEL: JOB_TYPE_JOB, + DOCKER_WORKLOAD_IDENTITY_TOKEN_FILE_LABEL: WORKLOAD_IDENTITY_TOKEN_FILE_PATH, + DOCKER_WORKLOAD_IDENTITY_VOLUME_LABEL: volume_name, + } + ) + + docker_client_mock.containers.get.side_effect = None + docker_client_mock.containers.get.return_value = container + refresher = MagicMock() + + test_job_step.status = PlatformJobStatus.ACTIVE + with patch.object(docker_job, "_build_workload_identity_refresher", return_value=refresher) as build_refresher: + update = docker_job.sync(test_job_step) + + assert update.status == PlatformJobStatus.ACTIVE + build_refresher.assert_called_once_with(volume_name) + refresher.start.assert_called_once() + assert docker_job._workload_identity_refreshers[container.name] is refresher + + +def test_docker_sync_restores_workload_identity_refresher_from_mounted_volume( + docker_job, docker_client_mock, test_job_step +): + volume_name = "task-workload-identity-default-job-test-job-id-task-mounted" + container = MagicMock() + container.name = "job-test-job-id-test-step" + container.id = "workload-container-id" + container.status = "running" + container.attrs = { + "State": {"Status": "running", "Running": True, "ExitCode": 0}, + "HostConfig": {}, + "Mounts": [{"Type": "volume", "Destination": WORKLOAD_IDENTITY_VOLUME_PATH, "Name": volume_name}], + } + container.labels = owned_container_labels( + { + JOB_WORKSPACE_ID_LABEL: test_job_step.workspace, + JOB_ID_LABEL: test_job_step.job, + JOB_STEP_NAME_LABEL: test_job_step.name, + JOB_TASK_ID_LABEL: "task-mounted", + JOB_TYPE_LABEL: JOB_TYPE_JOB, + } + ) + + docker_client_mock.containers.get.side_effect = None + docker_client_mock.containers.get.return_value = container + refresher = MagicMock() + + test_job_step.status = PlatformJobStatus.ACTIVE + with patch.object(docker_job, "_build_workload_identity_refresher", return_value=refresher) as build_refresher: + update = docker_job.sync(test_job_step) + + assert update.status == PlatformJobStatus.ACTIVE + build_refresher.assert_called_once_with(volume_name) + refresher.start.assert_called_once() + assert docker_job._workload_identity_refreshers[container.name] is refresher + + +def test_docker_stop_workload_identity_refresher_keeps_refresher_when_stop_fails(docker_job): + refresher = MagicMock() + refresher.stop.side_effect = RuntimeError("Timed out stopping workload identity subject token refresher") + docker_job._workload_identity_refreshers["job-container"] = refresher + + with pytest.raises(RuntimeError, match="Timed out stopping workload identity subject token refresher"): + docker_job._stop_workload_identity_refresher("job-container") + + assert docker_job._workload_identity_refreshers["job-container"] is refresher + + +def test_docker_shutdown_stops_workload_identity_refreshers(docker_job, docker_client_mock): + first_refresher = MagicMock() + second_refresher = MagicMock() + docker_job._workload_identity_refreshers = { + "job-container-one": first_refresher, + "job-container-two": second_refresher, + } + close_calls_before_shutdown = docker_client_mock.close.call_count + + docker_job.shutdown() + + first_refresher.stop.assert_called_once() + second_refresher.stop.assert_called_once() + assert docker_job._workload_identity_refreshers == {} + assert docker_client_mock.close.call_count == close_calls_before_shutdown + 1 + + +def test_docker_shutdown_continues_stopping_refreshers_when_one_stop_fails(docker_job, docker_client_mock): + failing_refresher = MagicMock() + failing_refresher.stop.side_effect = RuntimeError("Timed out stopping workload identity subject token refresher") + second_refresher = MagicMock() + docker_job._workload_identity_refreshers = { + "job-container-one": failing_refresher, + "job-container-two": second_refresher, + } + close_calls_before_shutdown = docker_client_mock.close.call_count + + docker_job.shutdown() + + failing_refresher.stop.assert_called_once() + second_refresher.stop.assert_called_once() + assert docker_job._workload_identity_refreshers == {} + assert docker_client_mock.close.call_count == close_calls_before_shutdown + 1 + + +def test_docker_subject_token_issuer_reads_password_from_configured_env_var(docker_job, monkeypatch): + monkeypatch.setenv("AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD", "shared-secret") + auth_config = SimpleNamespace( + oidc=SimpleNamespace( + token_endpoint="http://127.0.0.1:18080/application/o/token/", + workload_client_id="nemo-platform-workload", + client_id="nemo-platform-cli", + workload_scope="openid email groups", + ) + ) + docker_job._execution_profile_config.workload_identity = DockerWorkloadIdentityConfig( + username="svc-nemo", + password_env_var="AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD", + ) + + with patch("nmp.common.config.get_auth_config", return_value=auth_config): + issuer = docker_job._create_docker_subject_token_issuer() + + assert issuer.token_endpoint == "http://127.0.0.1:18080/application/o/token/" + assert issuer.client_id == "nemo-platform-workload" + assert issuer.username == "svc-nemo" + assert issuer.password == "shared-secret" + assert issuer.scope == "openid email groups" + + +def test_docker_subject_token_issuer_uses_internal_token_endpoint_override(docker_job, monkeypatch): + monkeypatch.setenv("AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD", "shared-secret") + auth_config = SimpleNamespace( + oidc=SimpleNamespace( + token_endpoint="http://127.0.0.1:18080/application/o/token/", + workload_client_id="nemo-platform-workload", + client_id="nemo-platform-cli", + workload_scope="openid email groups", + ) + ) + docker_job._execution_profile_config.workload_identity = DockerWorkloadIdentityConfig( + token_endpoint="https://nemo-gateway:8080/application/o/token/", + username="svc-nemo", + password_env_var="AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD", + ) + + with patch("nmp.common.config.get_auth_config", return_value=auth_config): + issuer = docker_job._create_docker_subject_token_issuer() + + assert issuer.token_endpoint == "https://nemo-gateway:8080/application/o/token/" + + +@pytest.mark.parametrize("env_value", [None, ""]) +def test_docker_subject_token_issuer_requires_non_empty_password_env_var(docker_job, monkeypatch, env_value): + if env_value is None: + monkeypatch.delenv("AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD", raising=False) + else: + monkeypatch.setenv("AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD", env_value) + auth_config = SimpleNamespace( + oidc=SimpleNamespace( + token_endpoint="http://127.0.0.1:18080/application/o/token/", + workload_client_id="nemo-platform-workload", + client_id="nemo-platform-cli", + workload_scope="openid email groups", + ) + ) + docker_job._execution_profile_config.workload_identity = DockerWorkloadIdentityConfig( + username="svc-nemo", + password_env_var="AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD", + ) + + with ( + patch("nmp.common.config.get_auth_config", return_value=auth_config), + pytest.raises(JobStorageError, match="AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD"), + ): + docker_job._create_docker_subject_token_issuer() + + +def test_docker_workload_identity_config_rejects_legacy_password_field(): + with pytest.raises(ValidationError) as exc_info: + DockerWorkloadIdentityConfig.model_validate( + { + "username": "svc-nemo", + "password": "inline-secret", + } + ) + + assert "password" in str(exc_info.value) + assert "extra" in str(exc_info.value).lower() + + +def test_docker_workload_identity_client_secret_schema_is_write_only(): + client_secret_schema = DockerWorkloadIdentityConfig.model_json_schema()["properties"]["client_secret"] + + assert client_secret_schema["format"] == "password" + assert client_secret_schema["writeOnly"] is True + + +def test_docker_workload_identity_token_timing_constraints(monkeypatch): + properties = DockerWorkloadIdentityConfig.model_json_schema()["properties"] + + assert properties["subject_token_ttl_seconds"]["minimum"] == 1 + assert properties["refresh_margin_seconds"]["minimum"] == 0 + assert ( + DockerWorkloadIdentityConfig(subject_token_ttl_seconds=1, refresh_margin_seconds=0).subject_token_ttl_seconds + == 1 + ) + assert DockerWorkloadIdentityConfig(refresh_margin_seconds=0).refresh_margin_seconds == 0 + with pytest.raises(ValidationError): + DockerWorkloadIdentityConfig.model_validate({"subject_token_ttl_seconds": 0}) + with pytest.raises(ValidationError): + DockerWorkloadIdentityConfig.model_validate({"refresh_margin_seconds": -1}) + for values in ( + {"subject_token_ttl_seconds": 1}, + {"subject_token_ttl_seconds": 30, "refresh_margin_seconds": 30}, + {"subject_token_ttl_seconds": 30, "refresh_margin_seconds": 31}, + ): + with pytest.raises(ValidationError, match="refresh_margin_seconds"): + DockerWorkloadIdentityConfig.model_validate(values) + monkeypatch.setenv("NMP_WORKLOAD_IDENTITY_TOKEN_TTL_SECONDS", "0") + with pytest.raises(ValidationError): + DockerWorkloadIdentityConfig() + + def test_schedule_docker_gpu(mock_nmp_client, docker_client_mock): """Test successful job scheduling.""" @@ -1562,13 +1959,16 @@ def test_cleanup_steps_by_ttl(docker_job, docker_client_mock, test_job_step, cle mock_container_running.remove.assert_not_called() # Verify task storage and config volumes were cleaned up for successful containers (2 volumes per task) - assert docker_client_mock.volumes.get.call_count == 6 + assert docker_client_mock.volumes.get.call_count == 9 docker_client_mock.volumes.get.assert_any_call("task-storage-default-test-job-id-task-success") docker_client_mock.volumes.get.assert_any_call("task-config-default-test-job-id-task-success") + docker_client_mock.volumes.get.assert_any_call("task-workload-identity-default-test-job-id-task-success") docker_client_mock.volumes.get.assert_any_call("task-storage-default-test-job-id-task-killed") docker_client_mock.volumes.get.assert_any_call("task-config-default-test-job-id-task-killed") + docker_client_mock.volumes.get.assert_any_call("task-workload-identity-default-test-job-id-task-killed") docker_client_mock.volumes.get.assert_any_call("task-storage-default-test-job-id-task-old-error") docker_client_mock.volumes.get.assert_any_call("task-config-default-test-job-id-task-old-error") + docker_client_mock.volumes.get.assert_any_call("task-workload-identity-default-test-job-id-task-old-error") else: # When cleanup_completed_jobs_immediately is False, no containers should be removed mock_container_success.remove.assert_not_called() @@ -1579,9 +1979,10 @@ def test_cleanup_steps_by_ttl(docker_job, docker_client_mock, test_job_step, cle assert mock_container_old_error.remove.call_count == 1 mock_container_old_error.remove.assert_called_with(force=True) - assert docker_client_mock.volumes.get.call_count == 2 + assert docker_client_mock.volumes.get.call_count == 3 docker_client_mock.volumes.get.assert_any_call("task-storage-default-test-job-id-task-old-error") docker_client_mock.volumes.get.assert_any_call("task-config-default-test-job-id-task-old-error") + docker_client_mock.volumes.get.assert_any_call("task-workload-identity-default-test-job-id-task-old-error") # Network cleanup should happen for all exited containers regardless of exit code assert mock_network.disconnect.call_count == 4 # success, killed, and error containers @@ -2070,11 +2471,12 @@ def test_job_step_with_auth_context(): def test_docker_job_schedule_with_auth_context(docker_job, docker_client_mock, test_job_step_with_auth_context): - """Test that scheduling sets NMP_PRINCIPAL and OTEL headers env vars when auth_context is present. + """Test that scheduling sets NMP_PRINCIPAL without injecting OTEL log headers. Verifies GitLab issue #3390 Gap 2: job tasks should run with the creating - user's auth context, propagated via the NMP_PRINCIPAL environment variable - and OTEL_EXPORTER_OTLP_LOGS_HEADERS for authenticated telemetry export. + user's auth context, propagated via the NMP_PRINCIPAL environment variable. + Job log upload auth is handled by jobs-launcher workload token exchange, + not by globally scoped OTEL header environment variables. """ step_spec = test_job_step_with_auth_context.step_spec executor_config = step_spec.executor @@ -2103,13 +2505,15 @@ def test_docker_job_schedule_with_auth_context(docker_job, docker_client_mock, t "groups": ["engineering", "ml-team"], } - # Verify OTEL headers env var is set for authenticated telemetry - assert "OTEL_EXPORTER_OTLP_LOGS_HEADERS" in env - otlp_headers = env["OTEL_EXPORTER_OTLP_LOGS_HEADERS"] - # URL-encoded: @ -> %40, , -> %2C - assert "X-NMP-Principal-Id=creator%40example.com" in otlp_headers - assert "X-NMP-Principal-Email=creator%40example.com" in otlp_headers - assert "X-NMP-Principal-Groups=engineering%2Cml-team" in otlp_headers + assert env[JOB_LOGS_ENDPOINT_ENVVAR].endswith( + "/apis/files/v2/workspaces/default/filesets/test-logs-fileset/otlp/v1/logs" + ) + assert "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT" not in env + assert "OTEL_EXPORTER_OTLP_PROTOCOL" not in env + assert "OTEL_EXPORTER_OTLP_LOGS_PROTOCOL" not in env + assert "OTEL_LOGS_EXPORTER" not in env + assert "OTEL_SERVICE_NAME" not in env + assert "OTEL_EXPORTER_OTLP_LOGS_HEADERS" not in env def test_docker_job_schedule_without_auth_context(docker_job, docker_client_mock, test_job_step): @@ -2240,7 +2644,7 @@ def test_cleanup_single_container_checks_job_terminal_before_persistent_storage_ # Verify container and task volumes were cleaned up assert mock_container.remove.call_count == 1 - assert docker_client_mock.volumes.get.call_count == 2 # task storage + config volumes + assert docker_client_mock.volumes.get.call_count == 3 # task storage + config + workload identity volumes # Verify persistent storage cleanup was NOT called docker_job.cleanup_job_persistent_storage.assert_not_called() @@ -2261,7 +2665,7 @@ def test_cleanup_single_container_checks_job_terminal_before_persistent_storage_ # Verify container and task volumes were cleaned up assert mock_container.remove.call_count == 1 - assert docker_client_mock.volumes.get.call_count == 2 # task storage + config volumes + assert docker_client_mock.volumes.get.call_count == 3 # task storage + config + workload identity volumes # Verify persistent storage cleanup WAS called docker_job.cleanup_job_persistent_storage.assert_called_once_with("default", "test-job-id") @@ -2299,7 +2703,7 @@ def test_cleanup_single_container_without_persistent_storage_label(docker_job, d # Verify container and task volumes were cleaned up assert mock_container.remove.call_count == 1 - assert docker_client_mock.volumes.get.call_count == 2 # task storage + config volumes + assert docker_client_mock.volumes.get.call_count == 3 # task storage + config + workload identity volumes # Verify job terminal check was NOT called (no persistent storage to cleanup) docker_job.check_job_is_terminal.assert_not_called() @@ -2348,7 +2752,7 @@ def test_cleanup_single_container_step_terminal_but_job_has_more_steps(docker_jo # Verify container and task volumes were cleaned up assert mock_container.remove.call_count == 1 - assert docker_client_mock.volumes.get.call_count == 2 # task storage + config volumes + assert docker_client_mock.volumes.get.call_count == 3 # task storage + config + workload identity volumes # Verify job terminal check WAS called (since container uses persistent storage) docker_job.check_job_is_terminal.assert_called_once_with(job="multi-step-job", workspace="default") @@ -2449,9 +2853,10 @@ def check_step_side_effect(job, step_name, workspace): mock_network.disconnect.assert_called_once_with(mock_container_step1) # Verify task storage volumes were cleaned up for step 1 (2 volumes per task) - assert docker_client_mock.volumes.get.call_count == 2 + assert docker_client_mock.volumes.get.call_count == 3 docker_client_mock.volumes.get.assert_any_call("task-storage-default-multi-step-job-task-step1") docker_client_mock.volumes.get.assert_any_call("task-config-default-multi-step-job-task-step1") + docker_client_mock.volumes.get.assert_any_call("task-workload-identity-default-multi-step-job-task-step1") # Verify job terminal check was called for step 1 docker_job.check_job_is_terminal.assert_called_once_with(job="multi-step-job", workspace="default") diff --git a/services/core/jobs/tests/controllers/test_kubernetes_backend.py b/services/core/jobs/tests/controllers/test_kubernetes_backend.py index 9fa58f34f0..728bb0354c 100644 --- a/services/core/jobs/tests/controllers/test_kubernetes_backend.py +++ b/services/core/jobs/tests/controllers/test_kubernetes_backend.py @@ -2,12 +2,15 @@ # SPDX-License-Identifier: Apache-2.0 import datetime +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest +from jsonschema import Draft202012Validator +from jsonschema.exceptions import ValidationError as JsonSchemaValidationError from kubernetes import client from kubernetes.client.rest import ApiException -from nmp.common.config import ImagePullSecret +from nmp.common.config import ImagePullSecret, PlatformConfig from nmp.common.jobs.constants import ( EPHEMERAL_TASK_STORAGE_PATH_ENVVAR, NEMO_JOB_FILESET_ENVVAR, @@ -34,6 +37,13 @@ PlatformJobSecretEnvironmentVariableRef, PlatformJobStepSpec, ) +from nmp.core.jobs.controllers.backends.base import ( + JOB_LOGS_ENDPOINT_ENVVAR, + WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, + WORKLOAD_IDENTITY_TOKEN_FILE_PATH, + WORKLOAD_IDENTITY_VOLUME_NAME, + WORKLOAD_IDENTITY_VOLUME_PATH, +) from nmp.core.jobs.controllers.backends.kubernetes import ( CPUKubernetesJobBackend, GPUKubernetesJobBackend, @@ -42,10 +52,13 @@ from nmp.core.jobs.controllers.backends.kubernetes.common import ( JOB_DSHM_VOLUME_NAME, JOB_STORAGE_VOLUME_NAME, + KubernetesConfigMapVolume, KubernetesEmptyDirVolume, KubernetesJobStorageConfig, + KubernetesKeyToPath, KubernetesObjectMetadata, KubernetesPersistentVolumeClaim, + KubernetesSecretVolume, KubernetesVolume, KubernetesVolumeMount, PodStatus, @@ -185,20 +198,96 @@ def test_kubernetes_volume_to_k8s_empty_dir(): assert k8s_volume.empty_dir.size_limit == "1Gi" +def test_kubernetes_volume_to_k8s_secret(): + """Test conversion of KubernetesVolume with Secret to V1Volume.""" + + secret_config = KubernetesSecretVolume( + secret_name="test-secret", + items=[KubernetesKeyToPath(key="ca.crt", path="ca.crt")], + ) + volume_config = KubernetesVolume(name="test-volume", secret=secret_config) + + k8s_volume = volume_config.to_k8s() + + assert k8s_volume is not None + assert k8s_volume.name == "test-volume" + assert k8s_volume.secret is not None + assert k8s_volume.secret.secret_name == "test-secret" + assert k8s_volume.secret.items[0].key == "ca.crt" + assert k8s_volume.secret.items[0].path == "ca.crt" + + +def test_kubernetes_volume_to_k8s_config_map(): + """Test conversion of KubernetesVolume with ConfigMap to V1Volume.""" + + config_map_config = KubernetesConfigMapVolume( + name="test-config-map", + items=[KubernetesKeyToPath(key="trust-bundle.pem", path="ca.crt")], + ) + volume_config = KubernetesVolume(name="test-volume", config_map=config_map_config) + + k8s_volume = volume_config.to_k8s() + + assert k8s_volume is not None + assert k8s_volume.name == "test-volume" + assert k8s_volume.config_map is not None + assert k8s_volume.config_map.name == "test-config-map" + assert k8s_volume.config_map.items[0].key == "trust-bundle.pem" + assert k8s_volume.config_map.items[0].path == "ca.crt" + + def test_kubernetes_volume_invalid_configuration(): """Test that invalid KubernetesVolume configuration raises ValueError.""" # Neither PVC nor EmptyDir specified - with pytest.raises(ValueError, match="Exactly one of 'persistent_volume_claim' or 'empty_dir' must be specified."): + with pytest.raises(ValueError, match="Exactly one of"): KubernetesVolume(name="invalid-volume") # Both PVC and EmptyDir specified pvc_config = KubernetesPersistentVolumeClaim(claim_name="test-pvc") empty_dir_config = KubernetesEmptyDirVolume() - with pytest.raises(ValueError, match="Exactly one of 'persistent_volume_claim' or 'empty_dir' must be specified."): + with pytest.raises(ValueError, match="Exactly one of"): KubernetesVolume(name="invalid-volume", persistent_volume_claim=pvc_config, empty_dir=empty_dir_config) +def test_kubernetes_volume_json_schema_requires_exactly_one_source(): + """Test OpenAPI schema requires exactly one non-null KubernetesVolume source.""" + + schema = KubernetesVolume.model_json_schema() + validator = Draft202012Validator(schema) + + assert "name" in schema["required"] + assert schema["oneOf"] == [ + { + "required": ["persistent_volume_claim"], + "properties": {"persistent_volume_claim": {"not": {"type": "null"}}}, + }, + {"required": ["empty_dir"], "properties": {"empty_dir": {"not": {"type": "null"}}}}, + {"required": ["secret"], "properties": {"secret": {"not": {"type": "null"}}}}, + {"required": ["config_map"], "properties": {"config_map": {"not": {"type": "null"}}}}, + ] + + validator.validate({"name": "pvc-volume", "persistent_volume_claim": {"claim_name": "test-pvc"}}) + validator.validate({"name": "empty-dir-volume", "empty_dir": {}}) + validator.validate({"name": "secret-volume", "secret": {"secret_name": "test-secret"}}) + validator.validate({"name": "config-map-volume", "config_map": {"name": "test-config-map"}}) + + with pytest.raises(JsonSchemaValidationError): + validator.validate({"name": "missing-source"}) + + with pytest.raises(JsonSchemaValidationError): + validator.validate( + { + "name": "multiple-sources", + "persistent_volume_claim": {"claim_name": "test-pvc"}, + "empty_dir": {}, + } + ) + + with pytest.raises(JsonSchemaValidationError): + validator.validate({"name": "null-source", "persistent_volume_claim": None}) + + def test_build_metadata(): """Test building Kubernetes metadata from configuration.""" @@ -749,6 +838,55 @@ def test_kubernetes_job_profile_environment_applied( assert env_vars.get("ENV_VAR") == "test_value" +def test_kubernetes_job_uses_service_discovery_urls_for_job_runtime( + mock_nmp_client, + kubernetes_client_mock, + kubernetes_execution_profile_config, + cpu_execution_provider, + test_step_pending, +): + """Job pods use routable service_discovery URLs instead of local in-process service URLs.""" + platform_config = PlatformConfig( # type: ignore[abstract] + base_url="http://127.0.0.1:8080", + services="jobs,files,models,secrets", + service_discovery={ + "platform": "https://nemo-gateway:8080", + "auth": "https://nemo-auth:8080", + }, + loopback_address="nemo-gateway", + ) + with ( + patch("nmp.core.jobs.controllers.backends.kubernetes.common.config.load_incluster_config"), + patch( + "nmp.core.jobs.controllers.backends.kubernetes.common.get_platform_config", + return_value=platform_config, + ), + ): + backend = CPUKubernetesJobBackend( + mock_nmp_client, + kubernetes_execution_profile_config, + profile_name="default", + ) + backend._batch_v1 = kubernetes_client_mock["batch_v1"] + backend._core_v1 = kubernetes_client_mock["core_v1"] + + backend._batch_v1.create_namespaced_job.return_value = MagicMock() + backend.schedule(cpu_execution_provider, test_step_pending) + + call_args = backend._batch_v1.create_namespaced_job.call_args + job_body = call_args.kwargs["body"] + main_container = job_body.spec.template.spec.containers[0] + env_vars = {env.name: env.value for env in main_container.env} + + assert env_vars["NMP_BASE_URL"] == "https://nemo-gateway:8080" + assert env_vars["NMP_AUTH_URL"] == "https://nemo-auth:8080" + assert env_vars["NMP_JOBS_URL"] == "https://nemo-gateway:8080" + assert env_vars["NMP_FILES_URL"] == "https://nemo-gateway:8080" + assert env_vars["NMP_MODELS_URL"] == "https://nemo-gateway:8080" + assert env_vars["NMP_SECRETS_URL"] == "https://nemo-gateway:8080" + assert env_vars[JOB_LOGS_ENDPOINT_ENVVAR].startswith("https://nemo-gateway:8080/apis/files/") + + def test_kubernetes_job_execution_profile_config_rejects_reserved_env_vars(): """KubernetesJobExecutionProfileConfig raises when environment contains reserved names.""" with pytest.raises(ValidationError) as exc_info: @@ -761,6 +899,46 @@ def test_kubernetes_job_execution_profile_config_rejects_reserved_env_vars(): assert "reserved" in str(exc_info.value).lower() +def test_kubernetes_job_rejects_reserved_step_auth_env_vars(kubernetes_job, cpu_execution_provider, test_step_pending): + test_step_pending.step_spec.environment.append( + PlatformJobEnvironmentVariable(name="NEMO_WORKFLOW_TOKEN", value="token") + ) + + with pytest.raises(ValueError, match="NEMO_WORKFLOW_TOKEN"): + kubernetes_job.schedule(cpu_execution_provider, test_step_pending) + + +def test_kubernetes_job_injects_projected_workload_identity_token_when_exchange_enabled( + kubernetes_job, cpu_execution_provider, test_step_pending +): + kubernetes_job._execution_profile_config.workload_identity_token_expiration_seconds = 600 + kubernetes_job._execution_profile_config.workload_identity_token_audience = "test-audience" + auth_config = SimpleNamespace(oidc=SimpleNamespace(workload_token_exchange_enabled=True)) + + with patch("nmp.common.config.get_auth_config", return_value=auth_config): + kubernetes_job.schedule(cpu_execution_provider, test_step_pending) + + call_args = kubernetes_job._batch_v1.create_namespaced_job.call_args + job_body = call_args.kwargs["body"] + pod_spec = job_body.spec.template.spec + main_container = pod_spec.containers[0] + + env_vars = {env.name: env.value for env in main_container.env} + assert env_vars[WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR] == WORKLOAD_IDENTITY_TOKEN_FILE_PATH + + workload_identity_volume = next( + volume for volume in pod_spec.volumes if volume.name == WORKLOAD_IDENTITY_VOLUME_NAME + ) + projection = workload_identity_volume.projected.sources[0].service_account_token + assert projection.path == "token" + assert projection.expiration_seconds == 600 + assert projection.audience == "test-audience" + + mount = next(vm for vm in main_container.volume_mounts if vm.name == WORKLOAD_IDENTITY_VOLUME_NAME) + assert mount.mount_path == WORKLOAD_IDENTITY_VOLUME_PATH + assert mount.read_only is True + + def test_schedule_job_with_args(kubernetes_job, cpu_execution_provider, test_step_pending): """Test job scheduling with custom args.""" @@ -1701,10 +1879,18 @@ def test_schedule_with_additional_volumes(kubernetes_job, cpu_execution_provider name="extra-volume-2", empty_dir=KubernetesEmptyDirVolume(medium="Memory"), ), + KubernetesVolume( + name="extra-volume-3", + secret=KubernetesSecretVolume( + secret_name="extra-secret", + items=[KubernetesKeyToPath(key="ca.crt", path="ca.crt")], + ), + ), ] kubernetes_job._execution_profile_config.storage.additional_volume_mounts = [ KubernetesVolumeMount(name="extra-volume-1", mount_path="/mnt/extra-1"), KubernetesVolumeMount(name="extra-volume-2", mount_path="/mnt/extra-2"), + KubernetesVolumeMount(name="extra-volume-3", mount_path="/mnt/extra-3", read_only=True), ] mock_create_job = kubernetes_job._batch_v1.create_namespaced_job @@ -1723,9 +1909,16 @@ def test_schedule_with_additional_volumes(kubernetes_job, cpu_execution_provider extra_volume_2 = next((v for v in volumes if v.name == "extra-volume-2"), None) assert extra_volume_2 is not None assert extra_volume_2.empty_dir.medium == "Memory" + extra_volume_3 = next((v for v in volumes if v.name == "extra-volume-3"), None) + assert extra_volume_3 is not None + assert extra_volume_3.secret.secret_name == "extra-secret" + assert extra_volume_3.secret.items[0].key == "ca.crt" # Check volume mounts main_container = job_body.spec.template.spec.containers[0] + volume_mounts = {mount.name: mount for mount in main_container.volume_mounts} + assert volume_mounts["extra-volume-3"].mount_path == "/mnt/extra-3" + assert volume_mounts["extra-volume-3"].read_only is True job_storage_mount = next((vm for vm in main_container.volume_mounts if vm.name == "extra-volume-1"), None) assert job_storage_mount is not None assert job_storage_mount.mount_path == "/mnt/extra-1" @@ -1952,11 +2145,12 @@ def test_step_pending_with_auth_context() -> PlatformJobStepWithContext: def test_kubernetes_job_schedule_with_auth_context( kubernetes_job, cpu_execution_provider, test_step_pending_with_auth_context ): - """Test that scheduling sets NMP_PRINCIPAL and OTEL headers env vars when auth_context is present. + """Test that scheduling sets NMP_PRINCIPAL without injecting OTEL log headers. Verifies GitLab issue #3390 Gap 2: job tasks should run with the creating - user's auth context, propagated via the NMP_PRINCIPAL environment variable - and OTEL_EXPORTER_OTLP_LOGS_HEADERS for authenticated telemetry export. + user's auth context, propagated via the NMP_PRINCIPAL environment variable. + Job log upload auth is handled by jobs-launcher workload token exchange, + not by globally scoped OTEL header environment variables. """ import json @@ -1977,6 +2171,7 @@ def test_kubernetes_job_schedule_with_auth_context( pod_spec = job_body.spec.template.spec main_container = pod_spec.containers[0] env_vars = {env.name: env.value for env in main_container.env if env.value is not None} + env_var_names = {env.name for env in main_container.env} # Verify NMP_PRINCIPAL env var is set assert NMP_PRINCIPAL_ENVVAR in env_vars @@ -1989,13 +2184,15 @@ def test_kubernetes_job_schedule_with_auth_context( "groups": ["engineering", "ml-team"], } - # Verify OTEL headers env var is set for authenticated telemetry - assert "OTEL_EXPORTER_OTLP_LOGS_HEADERS" in env_vars - otlp_headers = env_vars["OTEL_EXPORTER_OTLP_LOGS_HEADERS"] - # URL-encoded: @ -> %40, , -> %2C - assert "X-NMP-Principal-Id=creator%40example.com" in otlp_headers - assert "X-NMP-Principal-Email=creator%40example.com" in otlp_headers - assert "X-NMP-Principal-Groups=engineering%2Cml-team" in otlp_headers + assert env_vars[JOB_LOGS_ENDPOINT_ENVVAR].endswith( + "/apis/files/v2/workspaces/default/filesets/test-logs-fileset/otlp/v1/logs" + ) + assert "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT" not in env_var_names + assert "OTEL_EXPORTER_OTLP_PROTOCOL" not in env_var_names + assert "OTEL_EXPORTER_OTLP_LOGS_PROTOCOL" not in env_var_names + assert "OTEL_LOGS_EXPORTER" not in env_var_names + assert "OTEL_SERVICE_NAME" not in env_var_names + assert "OTEL_EXPORTER_OTLP_LOGS_HEADERS" not in env_var_names def test_kubernetes_job_schedule_without_auth_context(kubernetes_job, cpu_execution_provider, test_step_pending): diff --git a/services/core/jobs/tests/controllers/test_workload_tokens.py b/services/core/jobs/tests/controllers/test_workload_tokens.py new file mode 100644 index 0000000000..f91d5f0b3b --- /dev/null +++ b/services/core/jobs/tests/controllers/test_workload_tokens.py @@ -0,0 +1,395 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import tarfile +import threading +import time + +import httpx +import pytest +from nmp.core.jobs.controllers.backends.workload_tokens import ( + OAuthPasswordGrantSubjectTokenIssuer, + SubjectToken, + SubjectTokenRefreshLoop, + build_token_archive, +) + + +def test_build_token_archive_contains_read_only_token_file() -> None: + archive = build_token_archive("subject-token", name="token.tmp") + + with tarfile.open(fileobj=archive, mode="r") as tar: + member = tar.getmember("token.tmp") + extracted = tar.extractfile(member) + + assert member.mode == 0o400 + assert extracted is not None + assert extracted.read() == b"subject-token" + + +def test_oauth_password_grant_issuer_requests_subject_token(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict = {} + + def fake_post(url: str, *, data: dict, timeout: float) -> httpx.Response: + captured["url"] = url + captured["data"] = data + captured["timeout"] = timeout + return httpx.Response(200, json={"access_token": "subject-token", "expires_in": 120}) + + monkeypatch.setattr("nmp.core.jobs.controllers.backends.workload_tokens.httpx.post", fake_post) + + issuer = OAuthPasswordGrantSubjectTokenIssuer( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform", + client_secret="secret", + username="svc-nemo", + password="app-password", + scope="openid email groups", + timeout=5.0, + ) + before = time.time() + + token = issuer.issue() + + assert token.value == "subject-token" + assert token.expires_at >= before + 120 + assert captured == { + "url": "https://idp.example.com/token", + "data": { + "grant_type": "password", + "client_id": "nemo-platform", + "client_secret": "secret", + "username": "svc-nemo", + "password": "app-password", + "scope": "openid email groups", + }, + "timeout": 5.0, + } + + +def test_oauth_password_grant_issuer_reports_idp_errors(monkeypatch: pytest.MonkeyPatch) -> None: + def fake_post(url: str, *, data: dict, timeout: float) -> httpx.Response: + return httpx.Response( + 400, + json={"error": "invalid_grant", "error_description": "bad credentials"}, + headers={"content-type": "application/json"}, + ) + + monkeypatch.setattr("nmp.core.jobs.controllers.backends.workload_tokens.httpx.post", fake_post) + + issuer = OAuthPasswordGrantSubjectTokenIssuer( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform", + username="svc-nemo", + password="bad-password", + ) + + with pytest.raises(RuntimeError, match="invalid_grant - bad credentials"): + issuer.issue() + + +@pytest.mark.parametrize( + "response", + [ + httpx.Response(400, json=["invalid"], headers={"content-type": "application/json"}), + httpx.Response(400, json="invalid", headers={"content-type": "application/json"}), + httpx.Response(400, json=None, headers={"content-type": "application/json"}), + httpx.Response(400, content=b"not-json", headers={"content-type": "application/json"}), + ], +) +def test_oauth_password_grant_issuer_ignores_non_object_error_json( + monkeypatch: pytest.MonkeyPatch, response: httpx.Response +) -> None: + def fake_post(url: str, *, data: dict, timeout: float) -> httpx.Response: + return response + + monkeypatch.setattr("nmp.core.jobs.controllers.backends.workload_tokens.httpx.post", fake_post) + + issuer = OAuthPasswordGrantSubjectTokenIssuer( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform", + username="svc-nemo", + password="bad-password", + ) + + with pytest.raises(RuntimeError, match="unknown_error - "): + issuer.issue() + + +def test_oauth_password_grant_issuer_rejects_non_loopback_http_before_sending_credentials( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def fake_post(*args, **kwargs) -> httpx.Response: + raise AssertionError("token endpoint should not be called") + + monkeypatch.setattr("nmp.core.jobs.controllers.backends.workload_tokens.httpx.post", fake_post) + + issuer = OAuthPasswordGrantSubjectTokenIssuer( + token_endpoint="http://authentik-server:9000/application/o/token/", + client_id="nemo-platform", + username="svc-nemo", + password="app-password", + ) + + with pytest.raises(RuntimeError, match="token_endpoint must use https://"): + issuer.issue() + + +@pytest.mark.parametrize("token_endpoint", ["http://localhost:18080/token", "http://127.0.0.1:18080/token"]) +def test_oauth_password_grant_issuer_allows_loopback_http_for_local_development( + monkeypatch: pytest.MonkeyPatch, token_endpoint: str +) -> None: + captured: dict[str, object] = {} + + def fake_post(url: str, *, data: dict, timeout: float) -> httpx.Response: + captured["url"] = url + return httpx.Response(200, json={"access_token": "subject-token", "expires_in": 120}) + + monkeypatch.setattr("nmp.core.jobs.controllers.backends.workload_tokens.httpx.post", fake_post) + + issuer = OAuthPasswordGrantSubjectTokenIssuer( + token_endpoint=token_endpoint, + client_id="nemo-platform", + username="svc-nemo", + password="app-password", + ) + + assert issuer.issue().value == "subject-token" + assert captured["url"] == token_endpoint + + +def test_oauth_password_grant_issuer_uses_configured_default_expiry(monkeypatch: pytest.MonkeyPatch) -> None: + def fake_post(url: str, *, data: dict, timeout: float) -> httpx.Response: + return httpx.Response(200, json={"access_token": "subject-token"}) + + monkeypatch.setattr("nmp.core.jobs.controllers.backends.workload_tokens.httpx.post", fake_post) + + issuer = OAuthPasswordGrantSubjectTokenIssuer( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform", + username="svc-nemo", + password="app-password", + default_expires_in_seconds=45, + ) + before = time.time() + + token = issuer.issue() + + assert token.value == "subject-token" + assert token.expires_at >= before + 45 + + +@pytest.mark.parametrize( + ("payload", "message"), + [ + ([], "Token endpoint response was not a JSON object"), + ({}, "Token endpoint response did not include a non-empty access_token"), + ({"access_token": ""}, "Token endpoint response did not include a non-empty access_token"), + ({"access_token": None}, "Token endpoint response did not include a non-empty access_token"), + ( + {"access_token": "subject-token", "expires_in": None}, + "Token endpoint response did not include a positive numeric expires_in", + ), + ( + {"access_token": "subject-token", "expires_in": "120"}, + "Token endpoint response did not include a positive numeric expires_in", + ), + ( + {"access_token": "subject-token", "expires_in": 0}, + "Token endpoint response did not include a positive numeric expires_in", + ), + ( + {"access_token": "subject-token", "expires_in": -1}, + "Token endpoint response did not include a positive numeric expires_in", + ), + ( + b'{"access_token": "subject-token", "expires_in": NaN}', + "Token endpoint response did not include a positive numeric expires_in", + ), + ( + b'{"access_token": "subject-token", "expires_in": Infinity}', + "Token endpoint response did not include a positive numeric expires_in", + ), + ( + b'{"access_token": "subject-token", "expires_in": -Infinity}', + "Token endpoint response did not include a positive numeric expires_in", + ), + ( + {"access_token": "subject-token", "expires_in": True}, + "Token endpoint response did not include a positive numeric expires_in", + ), + ], +) +def test_oauth_password_grant_issuer_rejects_invalid_success_response( + monkeypatch: pytest.MonkeyPatch, payload: object, message: str +) -> None: + def fake_post(url: str, *, data: dict, timeout: float) -> httpx.Response: + if isinstance(payload, bytes): + return httpx.Response(200, content=payload) + return httpx.Response(200, json=payload) + + monkeypatch.setattr("nmp.core.jobs.controllers.backends.workload_tokens.httpx.post", fake_post) + + issuer = OAuthPasswordGrantSubjectTokenIssuer( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform", + username="svc-nemo", + password="app-password", + ) + + with pytest.raises(RuntimeError, match=f"invalid_response - {message}"): + issuer.issue() + + +def test_oauth_password_grant_issuer_repr_excludes_secrets() -> None: + issuer = OAuthPasswordGrantSubjectTokenIssuer( + token_endpoint="https://idp.example.com/token", + client_id="nemo-platform", + client_secret="super-sensitive-client-secret", + username="svc-nemo", + password="super-sensitive-password", + ) + + issuer_repr = repr(issuer) + + assert "client_secret=" not in issuer_repr + assert "password=" not in issuer_repr + assert "super-sensitive-client-secret" not in issuer_repr + assert "super-sensitive-password" not in issuer_repr + + +def test_refresh_loop_refresh_once_writes_issued_token() -> None: + class FakeIssuer: + def issue(self) -> SubjectToken: + return SubjectToken(value="subject-token", expires_at=time.time() + 120) + + writes: list[str] = [] + refresher = SubjectTokenRefreshLoop(issuer=FakeIssuer(), write_token=writes.append) + + token = refresher.refresh_once() + + assert token.value == "subject-token" + assert writes == ["subject-token"] + + +def test_refresh_loop_can_restart_after_stop() -> None: + class FakeIssuer: + def __init__(self) -> None: + self.issued = 0 + + def issue(self) -> SubjectToken: + self.issued += 1 + return SubjectToken(value=f"subject-token-{self.issued}", expires_at=time.time() + 120) + + writes: list[str] = [] + wrote = threading.Event() + + def write_token(token: str) -> None: + writes.append(token) + wrote.set() + + refresher = SubjectTokenRefreshLoop( + issuer=FakeIssuer(), + write_token=write_token, + min_sleep_seconds=0.01, + ) + + refresher.start() + assert wrote.wait(timeout=1) + refresher.stop() + + wrote.clear() + refresher.start() + assert wrote.wait(timeout=1) + refresher.stop() + + assert writes[:2] == ["subject-token-1", "subject-token-2"] + + +def test_refresh_loop_stop_timeout_preserves_thread_and_prevents_late_write() -> None: + issue_started = threading.Event() + release_issue = threading.Event() + writes: list[str] = [] + + class BlockingIssuer: + def issue(self) -> SubjectToken: + issue_started.set() + if not release_issue.wait(timeout=1): + raise RuntimeError("issuer was not released") + return SubjectToken(value="late-subject-token", expires_at=time.time() + 120) + + refresher = SubjectTokenRefreshLoop( + issuer=BlockingIssuer(), + write_token=writes.append, + min_sleep_seconds=0.01, + ) + refresher._stop_timeout_seconds = 0.01 + + refresher.start() + assert issue_started.wait(timeout=1) + + with pytest.raises(RuntimeError, match="Timed out stopping workload identity subject token refresher"): + refresher.stop() + + thread = refresher._thread + assert thread is not None + assert thread.is_alive() + + release_issue.set() + thread.join(timeout=1) + assert not thread.is_alive() + + refresher.stop() + + assert refresher._thread is None + assert writes == [] + + +def test_refresh_loop_backs_off_failures_and_resets_after_success() -> None: + class FakeIssuer: + def __init__(self) -> None: + self.calls = 0 + + def issue(self) -> SubjectToken: + self.calls += 1 + if self.calls in {1, 2, 3, 5}: + raise RuntimeError("idp unavailable") + return SubjectToken(value="subject-token", expires_at=time.time()) + + class FakeStop(threading.Event): + def __init__(self) -> None: + super().__init__() + self.waits: list[float] = [] + self.stopped = False + + def is_set(self) -> bool: + return self.stopped + + def wait(self, timeout: float | None = None) -> bool: + assert timeout is not None + self.waits.append(timeout) + if len(self.waits) == 5: + self.stopped = True + return True + return False + + stop = FakeStop() + writes: list[str] = [] + refresher = SubjectTokenRefreshLoop( + issuer=FakeIssuer(), + write_token=writes.append, + min_sleep_seconds=1.0, + max_failure_backoff_seconds=4.0, + ) + refresher._stop = stop + + refresher._run() + + assert writes == ["subject-token"] + assert stop.waits == [1.0, 2.0, 4.0, 1.0, 1.0] + + +def test_subject_token_seconds_until_refresh_never_negative() -> None: + token = SubjectToken(value="expired", expires_at=time.time() - 10) + + assert token.seconds_until_refresh(margin_seconds=60) == 0.0 diff --git a/services/core/jobs/tests/test_config.py b/services/core/jobs/tests/test_config.py index 2d9e5a0d21..1b67e0bbce 100644 --- a/services/core/jobs/tests/test_config.py +++ b/services/core/jobs/tests/test_config.py @@ -32,7 +32,7 @@ KubernetesJobStorageConfig, VolcanoJobExecutionProfileConfig, ) -from nmp.core.jobs.controllers.backends.registry import BackendKey, BackendRegistry +from nmp.core.jobs.controllers.backends.registry import BackendKey, BackendRegistry, backend_registry from nmp.core.jobs.controllers.backends.subprocess import ( SubprocessJobExecutionProfile, SubprocessJobExecutionProfileConfig, @@ -284,12 +284,63 @@ def test_kubernetes_job_service_account_name_from_executor_defaults(): assert config.executor_defaults.kubernetes_job.storage.pvc_name == "test-pvc" +def test_kubernetes_job_secret_volume_from_executor_defaults(): + """executor_defaults.kubernetes_job.storage supports Secret volume mounts.""" + global_settings = { + "jobs": { + "executor_defaults": { + "kubernetes_job": { + "storage": { + "additional_volumes": [ + { + "name": "trust-bundle", + "secret": { + "secret_name": "platform-ca", + "items": [{"key": "ca.crt", "path": "ca.crt"}], + }, + } + ], + "additional_volume_mounts": [ + { + "name": "trust-bundle", + "mount_path": "/etc/nmp/ca", + "read_only": True, + } + ], + }, + } + } + } + } + config = Configuration.global_settings_to_service_config(global_settings, JobsServiceConfig) + + storage = config.executor_defaults.kubernetes_job.storage + assert storage.additional_volumes[0].secret is not None + assert storage.additional_volumes[0].secret.secret_name == "platform-ca" + assert storage.additional_volumes[0].secret.items[0].key == "ca.crt" + assert storage.additional_volume_mounts[0].mount_path == "/etc/nmp/ca" + assert storage.additional_volume_mounts[0].read_only is True + + def test_volcano_job_service_account_name_default(): """VolcanoJobExecutionProfileConfig defaults service_account_name to 'default'.""" config = VolcanoJobExecutionProfileConfig() assert config.service_account_name == "default" +@pytest.mark.parametrize( + "config_cls", + [KubernetesJobExecutionProfileConfig, VolcanoJobExecutionProfileConfig], +) +def test_kubernetes_workload_identity_token_expiration_rejects_values_below_kubernetes_minimum(config_cls): + """Projected service account token expiration must honor Kubernetes' 600 second minimum.""" + with pytest.raises(ValidationError) as exc_info: + config_cls(workload_identity_token_expiration_seconds=599) + + assert "workload_identity_token_expiration_seconds" in str(exc_info.value) + assert "greater than or equal to 600" in str(exc_info.value) + + def test_job_execution_profile_config_rejects_reserved_env_vars(): """JobExecutionProfileConfig raises when environment contains reserved names.""" with pytest.raises(ValidationError) as exc_info: @@ -348,9 +399,35 @@ def test_default_profiles_exclude_subprocess_for_kubernetes_runtime(): profile_keys = [(p.provider, p.profile, p.backend) for p in profiles] + assert ("subprocess", "default", "subprocess") not in profile_keys + assert ("cpu", "default", "kubernetes_job") in profile_keys + assert ("gpu", "default", "kubernetes_job") in profile_keys assert ("cpu", "gpu", "kubernetes_job") in profile_keys assert ("gpu", "gpu", "kubernetes_job") in profile_keys - assert ("subprocess", "default", "subprocess") not in profile_keys + assert ("gpu_distributed", "default", "volcano_job") in profile_keys + + +def test_merge_executor_profiles_can_override_default_volcano_with_kubernetes_job(): + defaults = get_default_executor_profiles_for_runtime(Runtime.KUBERNETES, DefaultExecutionProfileConfig()) + custom = [ + KubernetesJobExecutionProfile( + provider="gpu_distributed", + profile="default", + backend="kubernetes_job", + config=KubernetesJobExecutionProfileConfig(namespace="authentik"), + ) + ] + + merged = merge_executor_profiles(custom, defaults) + + profile = next(p for p in merged if p.provider == "gpu_distributed" and p.profile == "default") + assert profile.backend == "kubernetes_job" + assert type(profile.config) is KubernetesJobExecutionProfileConfig + assert profile.config.namespace == "authentik" + + +def test_backend_registry_supports_gpu_distributed_kubernetes_job_override(): + assert BackendKey("gpu_distributed", "kubernetes_job") in backend_registry def test_merged_profiles(): diff --git a/services/core/models/tests/unit/controllers/test_deployment_reconciler.py b/services/core/models/tests/unit/controllers/test_deployment_reconciler.py index 4b73d60a75..a295a2d1bc 100644 --- a/services/core/models/tests/unit/controllers/test_deployment_reconciler.py +++ b/services/core/models/tests/unit/controllers/test_deployment_reconciler.py @@ -1959,6 +1959,7 @@ async def test_gc_not_found_on_status_update_handled(gc_reconciler, mock_backend ) async def test_gc_ttl_boundary_parametrized(mock_models_sdk, mock_backend_registry, ttl, age_seconds, should_gc): """Parametrized boundary tests for various TTL values and ages.""" + now = datetime.now(timezone.utc) config = ControllerConfig(error_deployment_ttl_seconds=ttl) reconciler = ModelDeploymentReconciler( models_sdk=mock_models_sdk, @@ -1975,10 +1976,12 @@ async def test_gc_ttl_boundary_parametrized(mock_models_sdk, mock_backend_regist reconciler._delete_model_provider = AsyncMock() dep = _make_error_deployment( - updated_at=datetime.now(timezone.utc) - timedelta(seconds=age_seconds), + updated_at=now - timedelta(seconds=age_seconds), ) - await reconciler.gc_error_deployments([dep]) + with patch("nmp.core.models.controllers.deployment_reconciler.datetime") as mock_datetime: + mock_datetime.now.return_value = now + await reconciler.gc_error_deployments([dep]) if should_gc: mock_backend.delete_model_deployment.assert_called_once() diff --git a/services/hello-world/src/nmp/hello_world/tasks/workload_workspace_get/run.py b/services/hello-world/src/nmp/hello_world/tasks/workload_workspace_get/run.py index 49977c403d..4984bff6aa 100644 --- a/services/hello-world/src/nmp/hello_world/tasks/workload_workspace_get/run.py +++ b/services/hello-world/src/nmp/hello_world/tasks/workload_workspace_get/run.py @@ -3,14 +3,10 @@ """Task that exercises workload-auth by reading a workspace through the public SDK.""" -import os - from nemo_platform import NeMoPlatform from nmp.common.jobs.config import get_task_config from pydantic import BaseModel -_WORKLOAD_TOKEN_ENV_VARS = ("NEMO_WORKLOAD_TOKEN", "NEMO_WORKLOAD_TOKEN_FILE") - class WorkloadWorkspaceGetConfig(BaseModel): """Configuration for the workload workspace read task.""" @@ -18,25 +14,12 @@ class WorkloadWorkspaceGetConfig(BaseModel): workspace: str -def _load_workload_token() -> str: - if token := os.environ.get("NEMO_WORKLOAD_TOKEN"): - return token - if token_file := os.environ.get("NEMO_WORKLOAD_TOKEN_FILE"): - with open(token_file, encoding="utf-8") as token_handle: - token = token_handle.read().strip() - if token: - return token - token_vars = " or ".join(_WORKLOAD_TOKEN_ENV_VARS) - raise RuntimeError(f"workload token not configured; set {token_vars}") - - def run(*, sdk: NeMoPlatform | None = None) -> int: - """Read the configured workspace using the public bearer-token SDK path.""" + """Read the configured workspace using the public SDK workload identity path.""" try: config = get_task_config(WorkloadWorkspaceGetConfig) if sdk is None: - token = _load_workload_token() - sdk = NeMoPlatform(default_headers={"Authorization": f"Bearer {token}"}) + sdk = NeMoPlatform() workspace = sdk.workspaces.retrieve(config.workspace) print(f"Successfully retrieved workspace: {workspace.name}") return 0 diff --git a/services/hello-world/tests/integration/tasks/test_workload_workspace_get_task.py b/services/hello-world/tests/integration/tasks/test_workload_workspace_get_task.py index dc195dc6aa..15a36a38c9 100644 --- a/services/hello-world/tests/integration/tasks/test_workload_workspace_get_task.py +++ b/services/hello-world/tests/integration/tasks/test_workload_workspace_get_task.py @@ -2,7 +2,13 @@ # SPDX-License-Identifier: Apache-2.0 from types import SimpleNamespace +from typing import cast +import httpx +import respx +from nemo_platform import NeMoPlatform +from nemo_platform.auth.helpers import NMPOIDCConfig +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR from nmp.common.jobs.constants import TASK_CONFIG_ENVVAR from nmp.hello_world.tasks.workload_workspace_get.run import run as task_run @@ -21,64 +27,78 @@ def __init__(self) -> None: self.workspaces = _StubWorkspaces() -def test_workload_workspace_get_reads_workspace_via_public_sdk(monkeypatch): - sdk = _StubSDK() - sdk_kwargs = {} - - def create_sdk(**kwargs): - sdk_kwargs.update(kwargs) - return sdk +@respx.mock +def test_workload_workspace_get_reads_workspace_via_public_sdk(monkeypatch, tmp_path): + subject_token_file = tmp_path / "workload-token" + subject_token_file.write_text("subject-token-from-file\n", encoding="utf-8") + config_file = tmp_path / "config.yaml" + config_file.write_text("{}\n", encoding="utf-8") + discovery_requests: list[str] = [] + exchange_requests: list[dict] = [] + + def discover_nmp_config(base_url: str) -> NMPOIDCConfig: + discovery_requests.append(base_url) + return NMPOIDCConfig( + auth_enabled=True, + workload_token_exchange_enabled=True, + workload_client_id="workload-client", + workload_token_endpoint="https://idp.example.test/oauth2/token", + workload_audience="nemo-platform", + workload_scope="openid email groups", + ) + + def token_exchange_grant(**kwargs): + exchange_requests.append(kwargs) + return {"access_token": "exchanged-access-token", "expires_in": 300} monkeypatch.setenv(TASK_CONFIG_ENVVAR, '{"workspace":"workload-read-target"}') - monkeypatch.setenv("NEMO_WORKLOAD_TOKEN", "workload-token-123") - monkeypatch.setattr("nmp.hello_world.tasks.workload_workspace_get.run.NeMoPlatform", create_sdk) - - exit_code = task_run() + monkeypatch.setenv("NMP_CONFIG_FILE", str(config_file)) + monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(subject_token_file)) + monkeypatch.delenv("NMP_ACCESS_TOKEN", raising=False) + monkeypatch.setattr("nemo_platform.client.factory.discover_nmp_config", discover_nmp_config) + monkeypatch.setattr("nemo_platform.auth.workload_exchange.token_exchange_grant", token_exchange_grant) + + workspace_route = respx.get("http://nmp.example.test/apis/entities/v2/workspaces/workload-read-target").mock( + return_value=httpx.Response( + 200, + json={ + "id": "workspace-id", + "name": "workload-read-target", + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-01T00:00:00Z", + }, + ) + ) + + sdk = NeMoPlatform(base_url="http://nmp.example.test") + try: + exit_code = task_run(sdk=sdk) + finally: + sdk.close() assert exit_code == 0 - assert sdk.workspaces.requested == ["workload-read-target"] - assert sdk_kwargs == {"default_headers": {"Authorization": "Bearer workload-token-123"}} - - -def test_workload_workspace_get_requires_workload_token_env(monkeypatch): - monkeypatch.setenv(TASK_CONFIG_ENVVAR, '{"workspace":"workload-read-target"}') - monkeypatch.delenv("NEMO_WORKLOAD_TOKEN", raising=False) - monkeypatch.delenv("NEMO_WORKLOAD_TOKEN_FILE", raising=False) - - exit_code = task_run() - - assert exit_code == 1 + assert workspace_route.called + assert workspace_route.calls[0].request.headers["Authorization"] == "Bearer exchanged-access-token" + assert [url.rstrip("/") for url in discovery_requests] == ["http://nmp.example.test"] + assert exchange_requests == [ + { + "token_endpoint": "https://idp.example.test/oauth2/token", + "client_id": "workload-client", + "subject_token": "subject-token-from-file", + "audience": "nemo-platform", + "scope": "openid email groups", + } + ] def test_workload_workspace_get_uses_injected_sdk_without_workload_token(monkeypatch): sdk = _StubSDK() monkeypatch.setenv(TASK_CONFIG_ENVVAR, '{"workspace":"workload-read-target"}') + monkeypatch.delenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, raising=False) monkeypatch.delenv("NEMO_WORKLOAD_TOKEN", raising=False) monkeypatch.delenv("NEMO_WORKLOAD_TOKEN_FILE", raising=False) - exit_code = task_run(sdk=sdk) - - assert exit_code == 0 - assert sdk.workspaces.requested == ["workload-read-target"] - - -def test_workload_workspace_get_accepts_workload_token_file_env(monkeypatch, tmp_path): - sdk = _StubSDK() - sdk_kwargs = {} - - def create_sdk(**kwargs): - sdk_kwargs.update(kwargs) - return sdk - - token_path = tmp_path / "workload.token" - token_path.write_text("workload-token-from-file\n", encoding="utf-8") - monkeypatch.setenv(TASK_CONFIG_ENVVAR, '{"workspace":"workload-read-target"}') - monkeypatch.delenv("NEMO_WORKLOAD_TOKEN", raising=False) - monkeypatch.setenv("NEMO_WORKLOAD_TOKEN_FILE", str(token_path)) - monkeypatch.setattr("nmp.hello_world.tasks.workload_workspace_get.run.NeMoPlatform", create_sdk) - - exit_code = task_run() + exit_code = task_run(sdk=cast(NeMoPlatform, sdk)) assert exit_code == 0 assert sdk.workspaces.requested == ["workload-read-target"] - assert sdk_kwargs == {"default_headers": {"Authorization": "Bearer workload-token-from-file"}} diff --git a/tests/auth_idp/authentik_live.py b/tests/auth_idp/authentik_live.py index 77ffa470dc..545a2f1a28 100644 --- a/tests/auth_idp/authentik_live.py +++ b/tests/auth_idp/authentik_live.py @@ -1,11 +1,158 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import os +import shutil +from datetime import datetime, timedelta, timezone +from ipaddress import ip_address +from pathlib import Path + import pytest +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from cryptography.x509.oid import ExtendedKeyUsageOID, NameOID +from nemo_platform_ext.client.tls import NMP_CLIENT_SSL_CERT_FILE_ENVVAR + +REPO_ROOT = Path(__file__).resolve().parents[2] +AUTHENTIK_ROOT = REPO_ROOT / "contrib/auth/authentik" +AUTHENTIK_COMPOSE_PROJECT_PREFIX = "authentik-e2e" +AUTHENTIK_WORKLOAD_NETWORK_NAME = f"{AUTHENTIK_COMPOSE_PROJECT_PREFIX}-${{gateway_port}}-workload" +AUTHENTIK_GATEWAY_TLS_VOLUME_NAME = f"{AUTHENTIK_COMPOSE_PROJECT_PREFIX}-${{gateway_port}}-gateway-tls" +AUTHENTIK_GATEWAY_BASE_URL = "${gateway_url}" +AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD = os.environ.get( + "AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD", "svc-nemo-token-secret-e2e" +) +_GATEWAY_TLS_OPENSSL_CONFIG = """[req] +prompt = no +distinguished_name = dn +x509_extensions = v3_req + +[dn] +CN = nemo-gateway + +[v3_req] +basicConstraints = critical, CA:TRUE +keyUsage = critical, digitalSignature, keyEncipherment, keyCertSign +extendedKeyUsage = serverAuth +subjectAltName = @alt_names + +[alt_names] +DNS.1 = localhost +DNS.2 = nemo-gateway +IP.1 = 127.0.0.1 +""" + + +def _authentik_relative_path(value: str, *, root: Path) -> Path: + path = Path(value) + if path.is_absolute(): + return path + return root / value.removeprefix("./") + + +def _blueprint_output_dir(*, root: Path) -> Path: + return _authentik_relative_path(os.environ.get("AUTHENTIK_BLUEPRINT_DIR", "./.generated/blueprints"), root=root) + + +def _gateway_tls_dir(*, root: Path) -> Path: + return _authentik_relative_path(os.environ.get("AUTHENTIK_GATEWAY_TLS_DIR", "./.generated/gateway-tls"), root=root) + + +def authentik_gateway_tls_ca_bundle(*, root: Path = AUTHENTIK_ROOT) -> Path: + return _gateway_tls_dir(root=root) / "tls.crt" + + +AUTHENTIK_GATEWAY_TLS_CA_BUNDLE = str(authentik_gateway_tls_ca_bundle()) + + +def _workload_token_private_key_file(*, root: Path) -> Path: + return root / ".generated/workload-token-private-key.pem" + + +def _ensure_workload_token_private_key(path: Path) -> None: + if path.exists(): + return + path.parent.mkdir(parents=True, exist_ok=True) + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + path.write_bytes( + key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.TraditionalOpenSSL, + encryption_algorithm=serialization.NoEncryption(), + ) + ) + path.chmod(0o600) + + +def _ensure_gateway_tls_certificate(tls_dir: Path) -> None: + cert_path = tls_dir / "tls.crt" + key_path = tls_dir / "tls.key" + if cert_path.exists() and key_path.exists(): + return + + tls_dir.mkdir(parents=True, exist_ok=True) + (tls_dir / "openssl.cnf").write_text(_GATEWAY_TLS_OPENSSL_CONFIG, encoding="utf-8") + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + subject = issuer = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "nemo-gateway")]) + now = datetime.now(timezone.utc) + cert = ( + x509.CertificateBuilder() + .subject_name(subject) + .issuer_name(issuer) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - timedelta(minutes=1)) + .not_valid_after(now + timedelta(days=365)) + .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True) + .add_extension( + x509.KeyUsage( + digital_signature=True, + key_encipherment=True, + key_cert_sign=True, + key_agreement=False, + content_commitment=False, + data_encipherment=False, + crl_sign=False, + encipher_only=False, + decipher_only=False, + ), + critical=True, + ) + .add_extension(x509.ExtendedKeyUsage([ExtendedKeyUsageOID.SERVER_AUTH]), critical=False) + .add_extension( + x509.SubjectAlternativeName( + [ + x509.DNSName("localhost"), + x509.DNSName("nemo-gateway"), + x509.IPAddress(ip_address("127.0.0.1")), + ] + ), + critical=False, + ) + .sign(key, hashes.SHA256()) + ) + key_path.write_bytes( + key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.TraditionalOpenSSL, + encryption_algorithm=serialization.NoEncryption(), + ) + ) + cert_path.write_bytes(cert.public_bytes(serialization.Encoding.PEM)) + key_path.chmod(0o600) + cert_path.chmod(0o644) + + +def prepare_authentik_compose_inputs(*, root: Path = AUTHENTIK_ROOT) -> None: + """Prepare generated files required before Docker Compose can start.""" + blueprint_source = root / "helm/files/blueprints/nemo.yaml" + blueprint_output = _blueprint_output_dir(root=root) / "nemo.yaml" + blueprint_output.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(blueprint_source, blueprint_output) + _ensure_workload_token_private_key(_workload_token_private_key_file(root=root)) + _ensure_gateway_tls_certificate(_gateway_tls_dir(root=root)) -AUTHENTIK_COMPOSE_PROJECT_NAME = "authentik-e2e" -AUTHENTIK_WORKLOAD_NETWORK_NAME = f"{AUTHENTIK_COMPOSE_PROJECT_NAME}_workload" -AUTHENTIK_NEMO_DIRECT_PORT = "38081" AUTHENTIK_DOCKER_E2E_CONFIG = pytest.mark.e2e_config( "contrib/auth/authentik/config/platform-compose-authentik.yaml", @@ -14,31 +161,79 @@ "oidc": { "additional_issuers": [ "http://authentik-server:9000/application/o/nemo/", - "http://127.0.0.1:38080/application/o/nemo-cli/", - "http://127.0.0.1:38080/application/o/nemo/", + "${gateway_url}/application/o/nemo-cli/", + "${gateway_url}/application/o/nemo/", + ], + "token_endpoint": "${gateway_url}/application/o/token/", + "device_authorization_endpoint": "${gateway_url}/application/o/device/", + "workload_subject_issuers": [ + "http://authentik-server:9000/application/o/nemo-workload/", + "https://nemo-gateway:8080/application/o/nemo-workload/", + "${gateway_url}/application/o/nemo-workload/", ], - "token_endpoint": "http://127.0.0.1:38080/application/o/token/", - "device_authorization_endpoint": "http://127.0.0.1:38080/application/o/device/", } }, }, + { + "jobs": { + "executors": [ + { + "provider": "cpu", + "profile": "workload", + "backend": "docker", + "config": { + "cleanup_completed_jobs_immediately": False, + "launcher_tool_path": "/tools/jobs-launcher", + "env": { + "SSL_CERT_FILE": "/etc/nmp/gateway-tls/tls.crt", + "REQUESTS_CA_BUNDLE": "/etc/nmp/gateway-tls/tls.crt", + }, + "storage": { + "additional_volume_mounts": [ + { + "volume_name": AUTHENTIK_GATEWAY_TLS_VOLUME_NAME, + "mount_path": "/etc/nmp/gateway-tls", + } + ] + }, + "workload_identity": { + "token_endpoint": "https://nemo-gateway:8080/application/o/token/", + "username": "svc-nemo", + "password_env_var": "AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD", + }, + }, + } + ] + } + }, harness={ "backend": "docker_compose", - "compose_file": "contrib/auth/authentik/docker-compose.yml", - "compose_project_name": AUTHENTIK_COMPOSE_PROJECT_NAME, - "service_url": "http://127.0.0.1:38080", - "auth_ready_url": f"http://127.0.0.1:{AUTHENTIK_NEMO_DIRECT_PORT}", - "wait_url": "http://127.0.0.1:38080/application/o/nemo/.well-known/openid-configuration", + "compose_file": "contrib/auth/authentik/compose/docker-compose.yml", + "compose_project_prefix": AUTHENTIK_COMPOSE_PROJECT_PREFIX, + "lifecycle": "fresh", + "dynamic_ports": { + "gateway": { + "host": "127.0.0.1", + "scheme": "https", + } + }, + "service_url": AUTHENTIK_GATEWAY_BASE_URL, + "auth_ready_url": f"{AUTHENTIK_GATEWAY_BASE_URL}/health/gateway/ready", "env": { - "AUTHENTIK_GATEWAY_PORT": "38080", - "NEMO_DIRECT_PORT": AUTHENTIK_NEMO_DIRECT_PORT, + "AUTHENTIK_GATEWAY_PORT": "${gateway_port}", + "AUTHENTIK_GATEWAY_TLS_VOLUME": AUTHENTIK_GATEWAY_TLS_VOLUME_NAME, "AUTHENTIK_WORKLOAD_NETWORK_NAME": AUTHENTIK_WORKLOAD_NETWORK_NAME, + "AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD": AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD, + NMP_CLIENT_SSL_CERT_FILE_ENVVAR: AUTHENTIK_GATEWAY_TLS_CA_BUNDLE, }, }, ) AUTHENTIK_DOCKER_PYTESTMARK = [ pytest.mark.auth_idp, + pytest.mark.auth_idp_docker, + pytest.mark.auth_idp_runtime("authentik-compose"), + pytest.mark.e2e, AUTHENTIK_DOCKER_E2E_CONFIG, pytest.mark.xdist_group("idp-live"), ] diff --git a/tests/auth_idp/common.py b/tests/auth_idp/common.py new file mode 100644 index 0000000000..9633faa41c --- /dev/null +++ b/tests/auth_idp/common.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import base64 +import json +import os + +import pytest + +from tests.auth_idp.runtime_contract import AuthIdpCase + + +def jwt_claims(token: str) -> dict[str, object]: + parts = token.split(".") + if len(parts) != 3: + return {} + payload = parts[1] + "=" * (-len(parts[1]) % 4) + decoded = json.loads(base64.urlsafe_b64decode(payload)) + if not isinstance(decoded, dict): + return {} + return decoded + + +def require_capability(case: AuthIdpCase, capability: str) -> None: + if capability not in case.capabilities: + pytest.skip(f"{case.id} does not declare auth-idp capability: {capability}") + + +def nmp_api_image() -> str: + registry = os.environ.get("IMAGE_REGISTRY", "my-registry") + tag = os.environ.get("BAKE_TAG", "local") + return f"{registry}/nmp-api:{tag}" diff --git a/tests/auth_idp/compose/test_authentik_cli_login.py b/tests/auth_idp/compose/test_authentik_cli_login.py new file mode 100644 index 0000000000..a323dfe884 --- /dev/null +++ b/tests/auth_idp/compose/test_authentik_cli_login.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import httpx +import pytest +from nemo_platform_ext.auth.helpers import discover_nmp_config +from nemo_platform_ext.client.tls import client_verify_from_env + +from tests.auth_idp.authentik_live import AUTHENTIK_DOCKER_E2E_CONFIG + +pytestmark = [ + pytest.mark.auth_idp, + pytest.mark.auth_idp_docker, + pytest.mark.e2e, + AUTHENTIK_DOCKER_E2E_CONFIG, + pytest.mark.xdist_group("idp-live"), +] + + +def test_authentik_discovery_exposes_gateway_reachable_device_flow(authentik_stack): + oidc = discover_nmp_config(authentik_stack.gateway_base_url) + + assert oidc.auth_enabled is True + assert oidc.client_id == "nemo-platform-cli" + assert oidc.token_endpoint == f"{authentik_stack.gateway_base_url}/application/o/token/" + assert oidc.device_authorization_endpoint == f"{authentik_stack.gateway_base_url}/application/o/device/" + assert oidc.default_scopes == "openid email offline_access groups" + + response = httpx.post( + oidc.device_authorization_endpoint, + data={ + "client_id": oidc.client_id, + "scope": oidc.default_scopes, + }, + timeout=30.0, + verify=client_verify_from_env(), + ) + response.raise_for_status() + body = response.json() + + assert body["verification_uri"] == f"{authentik_stack.gateway_base_url}/device" + assert body["verification_uri_complete"].startswith(f"{authentik_stack.gateway_base_url}/device?code=") + assert body["device_code"] + assert body["user_code"] + + +def test_authentik_cli_provider_rejects_unseeded_human_app_password(authentik_stack): + token_response = httpx.post( + authentik_stack.token_endpoint, + data={ + "grant_type": "password", + "client_id": "nemo-platform-cli", + "username": "nemo-user", + "password": "nemo-user-token-secret-dev", + "scope": "openid email offline_access groups", + }, + timeout=30.0, + verify=client_verify_from_env(), + ) + + assert token_response.status_code == 400 + assert token_response.json()["error"] == "invalid_grant" diff --git a/tests/auth_idp/conftest.py b/tests/auth_idp/conftest.py index 401992a93e..03751e53a7 100644 --- a/tests/auth_idp/conftest.py +++ b/tests/auth_idp/conftest.py @@ -1,21 +1,94 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import importlib +import os import time import uuid from collections.abc import Iterator from dataclasses import replace +from typing import Any, cast +from urllib.parse import urlparse import httpx import pytest from nemo_platform import NeMoPlatform +from nemo_platform_ext.client.tls import NMP_CLIENT_SSL_CERT_FILE_ENVVAR, client_verify_from_env +from e2e.services_pool import E2EHarnessConfig, E2EServicesPool, RunningServices +from tests.auth_idp.authentik_live import authentik_gateway_tls_ca_bundle, prepare_authentik_compose_inputs from tests.auth_idp.providers import ProviderConfig from tests.auth_idp.runtime import get_authentik_docker_test_runtime +from tests.auth_idp.runtime_contract import AuthIdpCase +from tests.auth_idp.runtime_factory import iter_auth_idp_cases, parametrize_cases, runtime_class_for_case pytest_plugins = ("e2e.conftest",) +def pytest_addoption(parser: pytest.Parser) -> None: + group = parser.getgroup("auth-idp") + group.addoption( + "--auth-idp-provider", + action="store", + default=None, + help="Run parametrized auth-idp contract tests only for the named provider.", + ) + group.addoption( + "--auth-idp-backend", + action="store", + choices=("compose", "kubernetes", "external"), + default=None, + help="Run parametrized auth-idp contract tests only for the selected backend.", + ) + group.addoption( + "--auth-idp-runtime", + action="store", + default=None, + help="Run parametrized auth-idp contract tests only for the selected runtime id.", + ) + + +def pytest_generate_tests(metafunc: pytest.Metafunc) -> None: + if "auth_idp_case" not in metafunc.fixturenames: + return + backend = metafunc.config.getoption("--auth-idp-backend") + if backend is None and metafunc.definition.get_closest_marker("auth_idp_k8s") is not None: + backend = "kubernetes" + if backend is None and metafunc.definition.get_closest_marker("auth_idp_docker") is not None: + backend = "compose" + cases = iter_auth_idp_cases( + backend=backend, + provider_name=metafunc.config.getoption("--auth-idp-provider"), + runtime_id=metafunc.config.getoption("--auth-idp-runtime"), + ) + metafunc.parametrize("auth_idp_case", parametrize_cases(cases), scope="session") + + +def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None: + runtime_id = config.getoption("--auth-idp-runtime") + selected: list[pytest.Item] = [] + deselected: list[pytest.Item] = [] + for item in items: + runtime_markers = list(item.iter_markers("auth_idp_runtime")) + if runtime_markers: + item.add_marker(pytest.mark.xdist_group("idp-live")) + if not runtime_markers: + selected.append(item) + continue + if not runtime_id: + selected.append(item) + continue + runtime_marker_args = {str(argument) for marker in runtime_markers for argument in marker.args} + if not runtime_marker_args or runtime_id in runtime_marker_args: + selected.append(item) + else: + deselected.append(item) + + if deselected: + config.hook.pytest_deselected(items=deselected) + items[:] = selected + + def _token_request_body(grant: dict[str, str]) -> dict[str, str]: grant_type = grant["grant_type"] body = { @@ -37,15 +110,53 @@ def _token_request_body(grant: dict[str, str]) -> dict[str, str]: raise ValueError(f"unsupported grant_type for auth_idp token exchange: {grant_type}") -def _exchange_token_with_retries(token_endpoint: str, grant: dict[str, str], timeout: float = 60.0) -> str: +def _compose_e2e_config_for_case( + auth_idp_case: AuthIdpCase, +) -> tuple[tuple[str | dict[str, Any], ...], E2EHarnessConfig]: + provider_module_name = auth_idp_case.provider.name.replace("-", "_") + marker_name = f"{provider_module_name.upper()}_DOCKER_E2E_CONFIG" + module = importlib.import_module(f"tests.auth_idp.{provider_module_name}_live") + marker_decorator = getattr(module, marker_name, None) + if marker_decorator is None: + raise ValueError( + f"compose auth-idp runtime {auth_idp_case.id!r} needs tests.auth_idp." + f"{provider_module_name}_live.{marker_name}" + ) + marker = marker_decorator.mark + config_layers = cast(tuple[str | dict[str, Any], ...], marker.args) + harness_config = cast(E2EHarnessConfig, dict(marker.kwargs.get("harness") or {})) + lifecycle = os.environ.get("NMP_E2E_COMPOSE_LIFECYCLE") + if lifecycle: + harness_config["lifecycle"] = cast(Any, lifecycle) + compose_project_name = os.environ.get("NMP_AUTHENTIK_COMPOSE_PROJECT_NAME") + if compose_project_name: + harness_config["compose_project_name"] = compose_project_name + gateway_port = os.environ.get("NMP_AUTHENTIK_COMPOSE_GATEWAY_PORT") + if gateway_port: + dynamic_ports = dict(harness_config.get("dynamic_ports") or {}) + gateway_config = dict(dynamic_ports.get("gateway") or {}) + gateway_config["port"] = gateway_port + dynamic_ports["gateway"] = gateway_config + harness_config["dynamic_ports"] = dynamic_ports + return config_layers, harness_config + + +def _exchange_token_with_retries( + token_endpoint: str, + grant: dict[str, str], + timeout: float = 60.0, + verify: str | bool | None = None, +) -> str: deadline = time.monotonic() + timeout last_error: Exception | None = None + request_verify = client_verify_from_env() if verify is None else verify while time.monotonic() < deadline: try: response = httpx.post( token_endpoint, data=_token_request_body(grant), timeout=30.0, + verify=request_verify, ) if response.status_code >= 500: last_error = httpx.HTTPStatusError( @@ -55,7 +166,14 @@ def _exchange_token_with_retries(token_endpoint: str, grant: dict[str, str], tim ) time.sleep(2) continue - response.raise_for_status() + try: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + raise httpx.HTTPStatusError( + f"{exc}; response body: {response.text}", + request=response.request, + response=response, + ) from exc return response.json()["access_token"] except httpx.RequestError as exc: last_error = exc @@ -65,23 +183,171 @@ def _exchange_token_with_retries(token_endpoint: str, grant: dict[str, str], tim raise TimeoutError(f"token endpoint did not become ready: {token_endpoint}") +def _url_port(url: str) -> str: + port = urlparse(url).port + return str(port) if port is not None else "default" + + +def _auth_idp_runtime_event_line( + event: str, + auth_idp_case: AuthIdpCase, + runtime: Any | None = None, + services: RunningServices | None = None, + status: str | None = None, +) -> str: + gateway_base_url = str(getattr(runtime, "gateway_base_url", "")) if runtime is not None else "" + parts = [ + f"Auth-idp runtime {event}:", + f"id={auth_idp_case.id}", + f"backend={auth_idp_case.backend}", + ] + if status is not None: + parts.append(f"status={status}") + if gateway_base_url: + parts.extend((f"url={gateway_base_url}", f"port={_url_port(gateway_base_url)}")) + if services is not None: + if services.compose_project_name is not None: + parts.append(f"compose_project={services.compose_project_name}") + if services.log_path is not None: + parts.append(f"log={services.log_path}") + cluster = getattr(runtime, "cluster", None) if runtime is not None else None + if cluster is not None: + parts.extend( + ( + f"cluster={cluster.name}", + f"context={cluster.context}", + f"runtime={cluster.runtime}", + ) + ) + kubeconfig = getattr(cluster, "kubeconfig", None) + if kubeconfig is not None: + parts.append(f"kubeconfig={kubeconfig}") + namespace = getattr(runtime, "namespace", None) if runtime is not None else None + if namespace is not None: + parts.append(f"namespace={namespace}") + helm_release = getattr(runtime, "helm_release", None) if runtime is not None else None + if helm_release is not None: + parts.append(f"helm_release={helm_release}") + return " ".join(parts) + + +def _write_terminal_line(request: pytest.FixtureRequest, message: str) -> None: + capture_manager = request.config.pluginmanager.get_plugin("capturemanager") + if capture_manager is None: + print(message, flush=True) + return + with capture_manager.global_and_fixture_disabled(): + print(message, flush=True) + + @pytest.fixture(scope="session") def idp_e2e_enabled(pytestconfig: pytest.Config) -> bool: - return bool(pytestconfig.getoption("--run-e2e")) + return bool(pytestconfig.getoption("--run-e2e") or pytestconfig.getoption("--auth-idp-runtime")) @pytest.fixture(scope="session") def require_idp_e2e(idp_e2e_enabled: bool) -> Iterator[None]: if not idp_e2e_enabled: - pytest.skip("set --run-e2e to execute provider stack validation") + pytest.skip("set --auth-idp-runtime or --run-e2e to execute provider stack validation") yield +@pytest.fixture(scope="session", autouse=True) +def _prepare_authentik_compose_inputs_for_e2e(idp_e2e_enabled: bool) -> None: + if idp_e2e_enabled: + prepare_authentik_compose_inputs() + os.environ.setdefault( + NMP_CLIENT_SSL_CERT_FILE_ENVVAR, + str(authentik_gateway_tls_ca_bundle()), + ) + + +@pytest.fixture(scope="session") +def auth_idp_runtime( + auth_idp_case: AuthIdpCase, + require_idp_e2e: None, + request: pytest.FixtureRequest, + _services_pool_manager: E2EServicesPool, +): + services: RunningServices | None = None + owner_id: str | None = None + runtime: Any | None = None + failures_before = request.session.testsfailed + _write_terminal_line(request, _auth_idp_runtime_event_line("starting", auth_idp_case)) + try: + runtime_class = runtime_class_for_case(auth_idp_case) + except ValueError as exc: + pytest.skip(str(exc)) + try: + if auth_idp_case.backend == "compose": + try: + config_layers, harness_config = _compose_e2e_config_for_case(auth_idp_case) + except ValueError as exc: + pytest.skip(str(exc)) + owner_id = f"auth-idp-runtime::{auth_idp_case.id}" + services = _services_pool_manager.acquire_for_config(owner_id, config_layers, harness_config) + if services.log_path is not None: + from e2e.conftest import _services_log_key + + request.session.stash[_services_log_key] = services.log_path + runtime = runtime_class( + auth_idp_case, + services.url, + cleanup=lambda: _services_pool_manager.release_for_config(owner_id), + ) + elif auth_idp_case.backend == "kubernetes": + runtime = runtime_class(auth_idp_case) + else: + runtime = runtime_class(auth_idp_case) + except Exception: + if owner_id is not None and runtime is None: + _services_pool_manager.release_for_config(owner_id) + _write_terminal_line(request, _auth_idp_runtime_event_line("startup_failed", auth_idp_case, runtime, services)) + raise + assert runtime is not None + _write_terminal_line(request, _auth_idp_runtime_event_line("available", auth_idp_case, runtime, services)) + try: + yield runtime + finally: + status = "passed" if request.session.testsfailed == failures_before else "failed" + _write_terminal_line(request, _auth_idp_runtime_event_line("result", auth_idp_case, runtime, services, status)) + _write_terminal_line( + request, _auth_idp_runtime_event_line("teardown_starting", auth_idp_case, runtime, services) + ) + try: + runtime.cleanup() + except Exception: + _write_terminal_line( + request, + _auth_idp_runtime_event_line("teardown_failed", auth_idp_case, runtime, services), + ) + raise + _write_terminal_line( + request, _auth_idp_runtime_event_line("teardown_complete", auth_idp_case, runtime, services) + ) + + +@pytest.fixture +def auth_idp_workspace(auth_idp_runtime) -> Iterator[str]: + workspace_name = f"auth-idp-ws-{uuid.uuid4().hex[:8]}" + sdk = auth_idp_runtime.e2e_setup_sdk() + sdk.workspaces.create( + name=workspace_name, + description="Workspace for auth-idp provider contract tests", + wait_role_propagation=True, + ) + try: + yield workspace_name + finally: + sdk.workspaces.delete(workspace_name) + + @pytest.fixture(scope="module") def authentik_provider(authentik_docker_runtime: ProviderConfig) -> ProviderConfig: provider = authentik_docker_runtime assert provider.token_endpoint is not None - assert provider.machine_grant is not None + assert provider.e2e_setup_password_grant is not None + assert provider.workload_provider_password_grant is not None return provider @@ -89,7 +355,8 @@ def authentik_provider(authentik_docker_runtime: ProviderConfig) -> ProviderConf def authentik_docker_runtime() -> ProviderConfig: provider = get_authentik_docker_test_runtime() assert provider.token_endpoint is not None - assert provider.machine_grant is not None + assert provider.e2e_setup_password_grant is not None + assert provider.workload_provider_password_grant is not None return provider @@ -108,43 +375,60 @@ def authentik_stack( @pytest.fixture(scope="module") -def machine_token(authentik_stack: ProviderConfig) -> str: - grant = authentik_stack.machine_grant +def e2e_setup_token(authentik_stack: ProviderConfig) -> str: + grant = authentik_stack.e2e_setup_password_grant + assert grant is not None + assert authentik_stack.token_endpoint is not None + return _exchange_token_with_retries(authentik_stack.token_endpoint, grant) + + +@pytest.fixture(scope="module") +def workload_provider_token(authentik_stack: ProviderConfig) -> str: + grant = authentik_stack.workload_provider_password_grant assert grant is not None assert authentik_stack.token_endpoint is not None return _exchange_token_with_retries(authentik_stack.token_endpoint, grant) @pytest.fixture(scope="module") -def human_token(authentik_stack: ProviderConfig) -> str: - grant = authentik_stack.human_grant +def interactive_user_token(authentik_stack: ProviderConfig) -> str: + grant = authentik_stack.interactive_user_password_grant assert grant is not None assert authentik_stack.token_endpoint is not None return _exchange_token_with_retries(authentik_stack.token_endpoint, grant) @pytest.fixture(scope="module") -def authentik_human_sdk(authentik_stack: ProviderConfig, human_token: str) -> NeMoPlatform: +def authentik_e2e_setup_sdk(authentik_stack: ProviderConfig, e2e_setup_token: str) -> NeMoPlatform: + return NeMoPlatform( + base_url=authentik_stack.gateway_base_url, + default_headers={"Authorization": f"Bearer {e2e_setup_token}"}, + max_retries=0, + ) + + +@pytest.fixture(scope="module") +def authentik_interactive_user_sdk(authentik_stack: ProviderConfig, interactive_user_token: str) -> NeMoPlatform: return NeMoPlatform( base_url=authentik_stack.gateway_base_url, - default_headers={"Authorization": f"Bearer {human_token}"}, + default_headers={"Authorization": f"Bearer {interactive_user_token}"}, max_retries=0, ) @pytest.fixture(scope="module") -def machine_sdk(authentik_stack: ProviderConfig, machine_token: str) -> NeMoPlatform: +def workload_provider_sdk(authentik_stack: ProviderConfig, workload_provider_token: str) -> NeMoPlatform: return NeMoPlatform( base_url=authentik_stack.gateway_base_url, - default_headers={"Authorization": f"Bearer {machine_token}"}, + default_headers={"Authorization": f"Bearer {workload_provider_token}"}, max_retries=0, ) @pytest.fixture -def authentik_workspace(authentik_human_sdk: NeMoPlatform) -> Iterator[str]: +def authentik_workspace(authentik_e2e_setup_sdk: NeMoPlatform) -> Iterator[str]: workspace_name = f"authentik-ws-{uuid.uuid4().hex[:8]}" - authentik_human_sdk.workspaces.create( + authentik_e2e_setup_sdk.workspaces.create( name=workspace_name, description="Workspace for Authentik live auth tests", wait_role_propagation=True, @@ -152,4 +436,4 @@ def authentik_workspace(authentik_human_sdk: NeMoPlatform) -> Iterator[str]: try: yield workspace_name finally: - authentik_human_sdk.workspaces.delete(workspace_name) + authentik_e2e_setup_sdk.workspaces.delete(workspace_name) diff --git a/tests/auth_idp/contracts/test_cli_refresh.py b/tests/auth_idp/contracts/test_cli_refresh.py new file mode 100644 index 0000000000..014badb944 --- /dev/null +++ b/tests/auth_idp/contracts/test_cli_refresh.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import time +from dataclasses import replace +from pathlib import Path + +import pytest +import yaml +from nemo_platform_ext.auth.helpers import decode_jwt_claims, discover_nmp_config, generate_unsigned_jwt +from nemo_platform_ext.cli.app import app +from nemo_platform_ext.client.tls import NMP_CLIENT_SSL_CERT_FILE_ENVVAR, client_verify_from_env +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR +from typer.testing import CliRunner + +from tests.auth_idp.common import require_capability +from tests.auth_idp.device_flow import authenticate_authentik_device_flow, with_url_origin + +pytestmark = [ + pytest.mark.auth_idp, + pytest.mark.auth_idp_runtime, + pytest.mark.e2e, + pytest.mark.xdist_group("idp-live"), +] + +CLI_REFRESH_CONTEXT_NAME = "auth-idp-refresh" + + +def _write_cli_config( + config_path: Path, + *, + base_url: str, + access_token: str, + refresh_token: str, +) -> None: + config = { + "current_context": CLI_REFRESH_CONTEXT_NAME, + "clusters": [ + { + "name": "auth-idp", + "base_url": base_url, + } + ], + "users": [ + { + "name": "interactive-user", + "type": "oauth", + "token": access_token, + "refresh_token": refresh_token, + } + ], + "contexts": [ + { + "name": CLI_REFRESH_CONTEXT_NAME, + "cluster": "auth-idp", + "user": "interactive-user", + "workspace": "default", + "preferences": { + "output_format": "json", + "timestamp_format": "iso8601", + "truncate": True, + "color_output": False, + }, + } + ], + } + config_path.write_text(yaml.safe_dump(config, sort_keys=False), encoding="utf-8") + + +def _read_config_user(config_path: Path) -> dict[str, object]: + config = yaml.safe_load(config_path.read_text(encoding="utf-8")) + assert isinstance(config, dict) + users = config["users"] + assert isinstance(users, list) + user = users[0] + assert isinstance(user, dict) + return user + + +def test_cli_api_command_auto_refreshes_expired_device_flow_token( + auth_idp_case, + auth_idp_runtime, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + require_capability(auth_idp_case, "device_flow") + require_capability(auth_idp_case, "gateway_authn") + + oidc = discover_nmp_config(auth_idp_runtime.gateway_base_url) + assert oidc.client_id + assert oidc.device_authorization_endpoint + assert oidc.token_endpoint + assert "offline_access" in oidc.default_scopes.split() + + verify = getattr(auth_idp_runtime, "verify", client_verify_from_env()) + runtime_device_authorization_endpoint = with_url_origin( + oidc.device_authorization_endpoint, + auth_idp_runtime.gateway_base_url, + ) + runtime_token_endpoint = with_url_origin(oidc.token_endpoint, auth_idp_runtime.gateway_base_url) + token_response = authenticate_authentik_device_flow( + gateway_base_url=auth_idp_runtime.gateway_base_url, + device_authorization_endpoint=runtime_device_authorization_endpoint, + token_endpoint=runtime_token_endpoint, + client_id=oidc.client_id, + scope=oidc.default_scopes, + username=auth_idp_case.provider.interactive_user_username, + password=auth_idp_case.provider.interactive_user_password, + verify=verify, + ) + refresh_token = token_response.get("refresh_token") + assert isinstance(refresh_token, str) + assert refresh_token + + expired_access_token = generate_unsigned_jwt( + auth_idp_case.provider.interactive_user_username, + email=auth_idp_case.provider.interactive_user_expected_email, + groups=auth_idp_case.provider.workload_expected_groups, + scopes=oidc.default_scopes.split(), + expires_in_seconds=-120, + ) + config_path = tmp_path / "config.yaml" + _write_cli_config( + config_path, + base_url=auth_idp_runtime.gateway_base_url, + access_token=expired_access_token, + refresh_token=refresh_token, + ) + + monkeypatch.setenv("NMP_CONFIG_FILE", str(config_path)) + monkeypatch.delenv("NMP_ACCESS_TOKEN", raising=False) + monkeypatch.delenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, raising=False) + runtime_oidc = replace( + oidc, + device_authorization_endpoint=runtime_device_authorization_endpoint, + token_endpoint=runtime_token_endpoint, + ) + monkeypatch.setattr("nemo_platform.client.factory.discover_nmp_config", lambda *_args, **_kwargs: runtime_oidc) + monkeypatch.setattr( + "nemo_platform_ext.client.factory.discover_nmp_config", + lambda *_args, **_kwargs: runtime_oidc, + ) + + cli_env = {"NMP_CONFIG_FILE": str(config_path)} + if isinstance(verify, str): + cli_env[NMP_CLIENT_SSL_CERT_FILE_ENVVAR] = verify + + result = CliRunner().invoke( + app, + [ + "--context", + CLI_REFRESH_CONTEXT_NAME, + "--no-auto-refresh", + "--output-format", + "json", + "workspaces", + "list", + ], + env=cli_env, + ) + + assert result.exit_code == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert "Jwt is expired" not in result.output + + saved_user = _read_config_user(config_path) + refreshed_access_token = saved_user["token"] + assert isinstance(refreshed_access_token, str) + assert refreshed_access_token != expired_access_token + + claims = decode_jwt_claims(refreshed_access_token) + assert claims.get("exp", 0) > time.time() + assert saved_user.get("refresh_token") diff --git a/tests/auth_idp/contracts/test_discovery.py b/tests/auth_idp/contracts/test_discovery.py new file mode 100644 index 0000000000..1cbb1ceb5d --- /dev/null +++ b/tests/auth_idp/contracts/test_discovery.py @@ -0,0 +1,126 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from urllib.parse import parse_qs, urlparse + +import httpx +import pytest +from nemo_platform_ext.auth.helpers import discover_nmp_config +from nemo_platform_ext.client.tls import client_verify_from_env + +from tests.auth_idp.common import require_capability +from tests.auth_idp.device_flow import ( + authenticate_authentik_device_flow, + url_origin, + with_url_origin, +) + +pytestmark = [ + pytest.mark.auth_idp, + pytest.mark.auth_idp_runtime, + pytest.mark.e2e, + pytest.mark.xdist_group("idp-live"), +] + + +def test_provider_gateway_serves_oidc_discovery(auth_idp_case, auth_idp_runtime): + require_capability(auth_idp_case, "gateway_discovery") + + verify = getattr(auth_idp_runtime, "verify", client_verify_from_env()) + response = httpx.get(auth_idp_runtime.discovery_url, timeout=10.0, verify=verify) + + response.raise_for_status() + discovery = response.json() + assert discovery["issuer"].endswith("/application/o/nemo/") + assert discovery["jwks_uri"] + + +def test_provider_discovery_exposes_device_flow_when_supported(auth_idp_case, auth_idp_runtime): + require_capability(auth_idp_case, "device_flow") + + oidc = discover_nmp_config(auth_idp_runtime.gateway_base_url) + + assert oidc.auth_enabled is True + assert oidc.client_id + assert oidc.token_endpoint + assert oidc.device_authorization_endpoint + assert oidc.default_scopes + + +def test_provider_device_authorization_endpoint_issues_user_code(auth_idp_case, auth_idp_runtime): + require_capability(auth_idp_case, "device_flow") + + oidc = discover_nmp_config(auth_idp_runtime.gateway_base_url) + verify = getattr(auth_idp_runtime, "verify", client_verify_from_env()) + device_authorization_endpoint = with_url_origin( + oidc.device_authorization_endpoint, + auth_idp_runtime.gateway_base_url, + ) + response = httpx.post( + device_authorization_endpoint, + data={ + "client_id": oidc.client_id, + "scope": oidc.default_scopes, + }, + timeout=30.0, + verify=verify, + ) + + response.raise_for_status() + body = response.json() + assert body["device_code"] + assert body["user_code"] + assert body["verification_uri"].startswith(url_origin(device_authorization_endpoint)) + + verification_complete = urlparse(body["verification_uri_complete"]) + verification_uri = urlparse(body["verification_uri"]) + assert verification_complete.scheme == verification_uri.scheme + assert verification_complete.netloc == verification_uri.netloc + assert verification_complete.path == verification_uri.path + assert parse_qs(verification_complete.query) == {"code": [body["user_code"]]} + + +def test_provider_device_flow_returns_refresh_token(auth_idp_case, auth_idp_runtime): + require_capability(auth_idp_case, "device_flow") + + oidc = discover_nmp_config(auth_idp_runtime.gateway_base_url) + assert oidc.token_endpoint + assert oidc.device_authorization_endpoint + assert "offline_access" in oidc.default_scopes.split() + + verify = getattr(auth_idp_runtime, "verify", client_verify_from_env()) + device_authorization_endpoint = with_url_origin( + oidc.device_authorization_endpoint, + auth_idp_runtime.gateway_base_url, + ) + token_endpoint = with_url_origin(oidc.token_endpoint, auth_idp_runtime.gateway_base_url) + token_response = authenticate_authentik_device_flow( + gateway_base_url=auth_idp_runtime.gateway_base_url, + device_authorization_endpoint=device_authorization_endpoint, + token_endpoint=token_endpoint, + client_id=oidc.client_id, + scope=oidc.default_scopes, + username=auth_idp_case.provider.interactive_user_username, + password=auth_idp_case.provider.interactive_user_password, + verify=verify, + ) + + refresh_token = token_response.get("refresh_token") + assert token_response.get("access_token") + assert isinstance(refresh_token, str) + assert refresh_token + + refresh_response = httpx.post( + token_endpoint, + data={ + "grant_type": "refresh_token", + "client_id": oidc.client_id, + "refresh_token": refresh_token, + "scope": oidc.default_scopes, + }, + timeout=30.0, + verify=verify, + ) + refresh_response.raise_for_status() + refreshed = refresh_response.json() + assert refreshed["access_token"] diff --git a/tests/auth_idp/contracts/test_gateway.py b/tests/auth_idp/contracts/test_gateway.py new file mode 100644 index 0000000000..cfbbfb10a7 --- /dev/null +++ b/tests/auth_idp/contracts/test_gateway.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import time +import uuid + +import httpx +import pytest +from nemo_platform_ext.client.tls import client_verify_from_env +from nmp.testing import grant_workspace_role + +from tests.auth_idp.common import jwt_claims, require_capability + +pytestmark = [ + pytest.mark.auth_idp, + pytest.mark.auth_idp_runtime, + pytest.mark.e2e, + pytest.mark.xdist_group("idp-live"), +] + +GATEWAY_REQUEST_TIMEOUT_SECONDS = 10.0 +GATEWAY_TRANSIENT_RETRY_TIMEOUT_SECONDS = 20.0 +GATEWAY_TRANSIENT_RETRY_SLEEP_SECONDS = 1.0 +GATEWAY_TRANSIENT_STATUS_CODES = {502, 503, 504} + + +def _runtime_verify(auth_idp_runtime) -> str | bool: + return getattr(auth_idp_runtime, "verify", client_verify_from_env()) + + +def _gateway_get_with_transient_retries( + url: str, + *, + headers: dict[str, str], + verify: str | bool, +) -> httpx.Response: + deadline = time.monotonic() + GATEWAY_TRANSIENT_RETRY_TIMEOUT_SECONDS + last_transient_response: httpx.Response | None = None + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + if last_transient_response is not None: + return last_transient_response + raise TimeoutError(f"gateway transient retry deadline reached before request: {url}") + + response = httpx.get( + url, + headers=headers, + timeout=min(GATEWAY_REQUEST_TIMEOUT_SECONDS, remaining), + verify=verify, + ) + if response.status_code not in GATEWAY_TRANSIENT_STATUS_CODES: + return response + + last_transient_response = response + remaining = deadline - time.monotonic() + if remaining <= 0: + return response + time.sleep(min(GATEWAY_TRANSIENT_RETRY_SLEEP_SECONDS, remaining)) + + +def test_provider_gateway_rejects_unauthenticated_requests(auth_idp_case, auth_idp_runtime): + require_capability(auth_idp_case, "gateway_authn") + + response = httpx.get( + f"{auth_idp_runtime.gateway_base_url}/apis/entities/v2/workspaces", + timeout=10.0, + verify=_runtime_verify(auth_idp_runtime), + ) + + assert response.status_code in {401, 403} + + +def test_provider_gateway_accepts_e2e_setup_token(auth_idp_case, auth_idp_runtime, auth_idp_workspace): + require_capability(auth_idp_case, "gateway_authn") + require_capability(auth_idp_case, "workspace_rbac") + + token = auth_idp_runtime.e2e_setup_token().access_token + response = _gateway_get_with_transient_retries( + f"{auth_idp_runtime.gateway_base_url}/apis/entities/v2/workspaces/{auth_idp_workspace}", + headers={"Authorization": f"Bearer {token}"}, + verify=_runtime_verify(auth_idp_runtime), + ) + + assert 200 <= response.status_code < 300, response.text + assert response.json()["name"] == auth_idp_workspace + + +def test_provider_gateway_rejects_spoofed_principal_headers(auth_idp_case, auth_idp_runtime): + require_capability(auth_idp_case, "spoofed_header_rejection") + require_capability(auth_idp_case, "workload_provider_token") + + workload_provider_token = auth_idp_runtime.workload_provider_token().access_token + workspace_name = f"spoof-check-{uuid.uuid4().hex[:8]}" + claims = jwt_claims(workload_provider_token) + authenticated_principal_id = str(claims["sub"]) + headers = { + "Authorization": f"Bearer {workload_provider_token}", + "X-NMP-Principal-Id": "service:bootstrap", + "X-NMP-Principal-Email": "attacker@example.com", + } + verify = _runtime_verify(auth_idp_runtime) + + try: + create_response = httpx.post( + f"{auth_idp_runtime.gateway_base_url}/apis/entities/v2/workspaces", + json={"name": workspace_name, "description": "Spoofed header check"}, + headers=headers, + timeout=GATEWAY_REQUEST_TIMEOUT_SECONDS, + verify=verify, + ) + create_response.raise_for_status() + assert create_response.json()["created_by"] == authenticated_principal_id + + members_response = httpx.get( + f"{auth_idp_runtime.gateway_base_url}/apis/entities/v2/workspaces/{workspace_name}/members", + headers=headers, + timeout=GATEWAY_REQUEST_TIMEOUT_SECONDS, + verify=verify, + ) + members_response.raise_for_status() + admin_member = next(member for member in members_response.json()["data"] if "Admin" in member["roles"]) + + assert admin_member["granted_by"] == authenticated_principal_id + assert admin_member["principal"] == authenticated_principal_id + assert admin_member["principal"] not in {"service:bootstrap", "attacker@example.com"} + finally: + httpx.delete( + f"{auth_idp_runtime.gateway_base_url}/apis/entities/v2/workspaces/{workspace_name}", + headers=headers, + timeout=GATEWAY_REQUEST_TIMEOUT_SECONDS, + verify=verify, + ) + + +def test_provider_gateway_forwards_workload_groups( + auth_idp_case, + auth_idp_runtime, + auth_idp_workspace, +): + require_capability(auth_idp_case, "workspace_rbac") + require_capability(auth_idp_case, "workload_token_exchange") + + workload_token = auth_idp_runtime.exchange_workload_token(auth_idp_runtime.workload_subject_token()).access_token + claims = jwt_claims(workload_token) + claim_groups = claims.get("groups") + assert isinstance(claim_groups, (list, str)) + token_groups = ( + set(claim_groups) + if isinstance(claim_groups, list) + else {group.strip() for group in claim_groups.split(",") if group.strip()} + ) + role_principals = auth_idp_runtime.workload_role_principals() + assert set(role_principals).intersection(token_groups) + + e2e_setup_sdk = auth_idp_runtime.e2e_setup_sdk() + for principal in role_principals: + grant_workspace_role(e2e_setup_sdk, workspace=auth_idp_workspace, principal=principal, roles=["Viewer"]) + + response = _gateway_get_with_transient_retries( + f"{auth_idp_runtime.gateway_base_url}/apis/entities/v2/workspaces/{auth_idp_workspace}", + headers={"Authorization": f"Bearer {workload_token}"}, + verify=_runtime_verify(auth_idp_runtime), + ) + + assert response.status_code == 200, response.text + assert response.json()["name"] == auth_idp_workspace diff --git a/tests/auth_idp/contracts/test_jobs.py b/tests/auth_idp/contracts/test_jobs.py new file mode 100644 index 0000000000..169579f1bd --- /dev/null +++ b/tests/auth_idp/contracts/test_jobs.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from nmp.testing import grant_workspace_role +from nmp.testing.e2e import wait_for_job_logs, wait_for_platform_job + +from tests.auth_idp.common import nmp_api_image, require_capability + +pytestmark = [ + pytest.mark.auth_idp, + pytest.mark.auth_idp_runtime, + pytest.mark.e2e, +] + + +def test_provider_workload_job_runs_via_workload_profile( + auth_idp_case, + auth_idp_runtime, + auth_idp_workspace, +): + require_capability(auth_idp_case, "workspace_rbac") + require_capability(auth_idp_case, "workload_job") + + e2e_setup_sdk = auth_idp_runtime.e2e_setup_sdk() + for principal in auth_idp_runtime.workload_role_principals(): + grant_workspace_role( + e2e_setup_sdk, + workspace=auth_idp_workspace, + principal=principal, + roles=["Viewer", "JobRunner"], + ) + + job = e2e_setup_sdk.jobs.create( + workspace=auth_idp_workspace, + source=f"{auth_idp_case.id}-workload-job", + spec={"test": "workload-job"}, + platform_spec={ + "steps": [ + { + "name": "workload-workspace-get", + "executor": { + "provider": "cpu", + "profile": "workload", + "container": { + "image": nmp_api_image(), + "entrypoint": ["nemo-platform"], + "command": [ + "run", + "task", + "--task", + "nmp.hello_world.tasks.workload_workspace_get", + ], + }, + }, + "config": { + "workspace": auth_idp_workspace, + }, + } + ] + }, + ) + + completed_job = wait_for_platform_job(e2e_setup_sdk, job.name, auth_idp_workspace, timeout=240) + assert completed_job.status == "completed" + + step_logs = wait_for_job_logs(e2e_setup_sdk, job.name, auth_idp_workspace, min_log_count=1, timeout=240) + assert step_logs.data + assert all(log.job == job.name for log in step_logs.data) + assert all(log.job_step == "workload-workspace-get" for log in step_logs.data) + assert all(log.job_task for log in step_logs.data) + assert all(log.message.strip() for log in step_logs.data) + assert any(f"Successfully retrieved workspace: {auth_idp_workspace}" in log.message for log in step_logs.data) diff --git a/tests/auth_idp/contracts/test_tokens.py b/tests/auth_idp/contracts/test_tokens.py new file mode 100644 index 0000000000..fd24d6f654 --- /dev/null +++ b/tests/auth_idp/contracts/test_tokens.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from tests.auth_idp.common import require_capability + +pytestmark = [ + pytest.mark.auth_idp, + pytest.mark.auth_idp_runtime, + pytest.mark.e2e, +] + + +def _claim_values(value: object) -> set[str]: + if isinstance(value, list): + return {str(item) for item in value} + if isinstance(value, str): + return {item.strip() for item in value.split(",") if item.strip()} + return set() + + +def test_provider_e2e_setup_token_is_real(auth_idp_case, auth_idp_runtime): + require_capability(auth_idp_case, "gateway_authn") + + token = auth_idp_runtime.e2e_setup_token() + grant = auth_idp_case.provider.e2e_setup_password_grant + assert grant is not None + + assert token.access_token + assert token.claims + assert token.claims["sub"] == grant["username"] + + +def test_provider_workload_provider_token_is_real(auth_idp_case, auth_idp_runtime): + require_capability(auth_idp_case, "workload_provider_token") + + token = auth_idp_runtime.workload_provider_token() + + assert token.access_token + assert token.claims + + +def test_provider_workload_provider_token_claims_match_manifest(auth_idp_case, auth_idp_runtime): + require_capability(auth_idp_case, "workload_provider_token") + + token = auth_idp_runtime.workload_provider_token() + provider = auth_idp_case.provider + grant = provider.workload_provider_password_grant + assert grant is not None + + assert token.claims["sub"] == provider.workload_principal_id + assert set(provider.workload_expected_groups).issubset( + _claim_values(token.claims.get(provider.workload_groups_claim)) + ) + assert grant["client_id"] in _claim_values(token.claims.get("aud")) + + +def test_provider_workload_subject_token_exchanges_for_access_token(auth_idp_case, auth_idp_runtime): + require_capability(auth_idp_case, "workload_subject_token") + require_capability(auth_idp_case, "workload_token_exchange") + + subject_token = auth_idp_runtime.workload_subject_token() + exchanged = auth_idp_runtime.exchange_workload_token(subject_token) + + assert exchanged.access_token + assert exchanged.claims diff --git a/tests/auth_idp/contracts/test_workspace.py b/tests/auth_idp/contracts/test_workspace.py new file mode 100644 index 0000000000..fd94c0d96f --- /dev/null +++ b/tests/auth_idp/contracts/test_workspace.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from nemo_platform import APIStatusError +from nmp.testing import grant_workspace_role + +from tests.auth_idp.common import require_capability + +pytestmark = [ + pytest.mark.auth_idp, + pytest.mark.auth_idp_runtime, + pytest.mark.e2e, +] + + +def test_provider_workload_identity_is_denied_before_binding( + auth_idp_case, + auth_idp_runtime, + auth_idp_workspace, +): + require_capability(auth_idp_case, "workspace_rbac") + + with pytest.raises(APIStatusError) as exc_info: + auth_idp_runtime.workload_provider_sdk().workspaces.retrieve(auth_idp_workspace) + + assert exc_info.value.status_code == 403 + + +def test_provider_workload_identity_is_allowed_after_binding( + auth_idp_case, + auth_idp_runtime, + auth_idp_workspace, +): + require_capability(auth_idp_case, "workspace_rbac") + + e2e_setup_sdk = auth_idp_runtime.e2e_setup_sdk() + for principal in auth_idp_runtime.workload_role_principals(): + grant_workspace_role(e2e_setup_sdk, workspace=auth_idp_workspace, principal=principal, roles=["Viewer"]) + + retrieved = auth_idp_runtime.workload_provider_sdk().workspaces.retrieve(auth_idp_workspace) + assert retrieved.name == auth_idp_workspace + + +def test_provider_workload_identity_returns_to_denied_after_revoke( + auth_idp_case, + auth_idp_runtime, + auth_idp_workspace, +): + require_capability(auth_idp_case, "workspace_rbac") + + e2e_setup_sdk = auth_idp_runtime.e2e_setup_sdk() + for principal in auth_idp_runtime.workload_role_principals(): + grant_workspace_role(e2e_setup_sdk, workspace=auth_idp_workspace, principal=principal, roles=["Viewer"]) + e2e_setup_sdk.workspaces.members.delete( + principal, + workspace=auth_idp_workspace, + wait_role_propagation=True, + ) + + with pytest.raises(APIStatusError) as exc_info: + auth_idp_runtime.workload_provider_sdk().workspaces.retrieve(auth_idp_workspace) + + assert exc_info.value.status_code == 403 diff --git a/tests/auth_idp/device_flow.py b/tests/auth_idp/device_flow.py new file mode 100644 index 0000000000..72b9bd9842 --- /dev/null +++ b/tests/auth_idp/device_flow.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import time +from json import JSONDecodeError +from urllib.parse import urlencode, urljoin, urlparse, urlunparse + +import httpx + +DEVICE_CODE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:device_code" +AUTHENTIK_DEFAULT_AUTHENTICATION_FLOW_SLUG = "default-authentication-flow" +DEVICE_TOKEN_POLL_ATTEMPTS = 3 +DEVICE_TOKEN_POLL_INTERVAL_SECONDS = 1.0 + + +def url_origin(url: str) -> str: + parsed = urlparse(url) + assert parsed.scheme + assert parsed.netloc + return f"{parsed.scheme}://{parsed.netloc}" + + +def with_url_origin(url: str, origin: str) -> str: + parsed_url = urlparse(url) + parsed_origin = urlparse(origin) + assert parsed_url.scheme + assert parsed_url.netloc + assert parsed_origin.scheme + assert parsed_origin.netloc + return urlunparse(parsed_url._replace(scheme=parsed_origin.scheme, netloc=parsed_origin.netloc)) + + +def authentik_flow_executor_url(gateway_base_url: str, location: str) -> str | None: + flow_url = urljoin(gateway_base_url, location) + parsed = urlparse(flow_url) + path_parts = [part for part in parsed.path.split("/") if part] + query = f"?{urlencode({'query': parsed.query})}" if parsed.query else "" + + if path_parts[:4] == ["api", "v3", "flows", "executor"] and len(path_parts) >= 5: + return flow_url + if path_parts[:2] == ["if", "flow"] and len(path_parts) >= 3: + return f"{gateway_base_url}/api/v3/flows/executor/{path_parts[2]}/{query}" + if path_parts == ["flows", "-", "default", "authentication"]: + return f"{gateway_base_url}/api/v3/flows/executor/{AUTHENTIK_DEFAULT_AUTHENTICATION_FLOW_SLUG}/{query}" + + return None + + +def next_authentik_challenge( + client: httpx.Client, + *, + gateway_base_url: str, + response: httpx.Response, +) -> tuple[dict[str, object], str]: + while response.status_code in {301, 302, 303, 307, 308}: + location = response.headers.get("location") + assert location + flow_executor_url = authentik_flow_executor_url(gateway_base_url, location) + response = client.get(flow_executor_url or urljoin(gateway_base_url, location), timeout=30.0) + + response.raise_for_status() + try: + challenge = response.json() + except JSONDecodeError as exc: + body = response.text[:500].replace("\n", " ") + raise AssertionError( + "Expected Authentik flow executor JSON challenge, got " + f"status={response.status_code} url={response.url} " + f"content_type={response.headers.get('content-type')!r} body={body!r}" + ) from exc + assert isinstance(challenge, dict) + return challenge, str(response.url) + + +def solve_authentik_device_flow( + *, + gateway_base_url: str, + verification_uri_complete: str, + user_code: str, + username: str, + password: str, + verify: str | bool, +) -> None: + with httpx.Client(verify=verify, follow_redirects=False) as client: + response = client.get(verification_uri_complete, timeout=30.0) + challenge, flow_url = next_authentik_challenge( + client, + gateway_base_url=gateway_base_url, + response=response, + ) + + for _ in range(10): + component = challenge.get("component") + if component == "xak-flow-redirect": + redirect_to = challenge.get("to") + assert isinstance(redirect_to, str) + flow_executor_url = authentik_flow_executor_url(gateway_base_url, redirect_to) + response = client.get(flow_executor_url or urljoin(gateway_base_url, redirect_to), timeout=30.0) + challenge, flow_url = next_authentik_challenge( + client, + gateway_base_url=gateway_base_url, + response=response, + ) + continue + if component == "ak-stage-access-denied": + raise AssertionError(f"Authentik device flow was denied: {challenge}") + + if component == "ak-stage-identification": + payload = {"component": component, "uid_field": username} + if challenge.get("password_fields"): + payload["password"] = password + elif component == "ak-stage-password": + payload = {"component": component, "password": password} + elif component == "ak-stage-user-login": + payload = {"component": component} + elif component == "ak-provider-oauth2-device-code": + payload = {"component": component, "code": user_code} + elif component == "ak-provider-oauth2-device-code-finish": + payload = {"component": component} + else: + raise AssertionError(f"Unexpected Authentik device flow component {component!r}: {challenge}") + + response = client.post(flow_url, json=payload, timeout=30.0) + challenge, flow_url = next_authentik_challenge( + client, + gateway_base_url=gateway_base_url, + response=response, + ) + if component == "ak-provider-oauth2-device-code-finish": + return + + raise AssertionError(f"Authentik device flow did not complete after 10 stages: {challenge}") + + +def poll_device_token( + *, + token_endpoint: str, + client_id: str, + device_code: str, + scope: str, + verify: str | bool, +) -> dict[str, object]: + last_response: httpx.Response | None = None + for _ in range(DEVICE_TOKEN_POLL_ATTEMPTS): + response = httpx.post( + token_endpoint, + data={ + "grant_type": DEVICE_CODE_GRANT_TYPE, + "client_id": client_id, + "device_code": device_code, + "scope": scope, + }, + timeout=30.0, + verify=verify, + ) + if response.status_code == 200: + token_response = response.json() + assert isinstance(token_response, dict) + return token_response + + last_response = response + error = response.json().get("error") + if error != "authorization_pending": + response.raise_for_status() + + time.sleep(DEVICE_TOKEN_POLL_INTERVAL_SECONDS) + + raise AssertionError( + "Device token endpoint did not return tokens after browser-side authorization completed: " + f"{last_response.text if last_response is not None else 'no response'}" + ) + + +def authenticate_authentik_device_flow( + *, + gateway_base_url: str, + device_authorization_endpoint: str, + token_endpoint: str, + client_id: str, + scope: str, + username: str, + password: str, + verify: str | bool, +) -> dict[str, object]: + device_response = httpx.post( + device_authorization_endpoint, + data={ + "client_id": client_id, + "scope": scope, + }, + timeout=30.0, + verify=verify, + ) + device_response.raise_for_status() + device_body = device_response.json() + + solve_authentik_device_flow( + gateway_base_url=gateway_base_url, + verification_uri_complete=with_url_origin(device_body["verification_uri_complete"], gateway_base_url), + user_code=device_body["user_code"], + username=username, + password=password, + verify=verify, + ) + + return poll_device_token( + token_endpoint=token_endpoint, + client_id=client_id, + device_code=device_body["device_code"], + scope=scope, + verify=verify, + ) diff --git a/tests/auth_idp/k8s/test_authentik_kubernetes_live.py b/tests/auth_idp/k8s/test_authentik_kubernetes_live.py new file mode 100644 index 0000000000..3cee2bc5ea --- /dev/null +++ b/tests/auth_idp/k8s/test_authentik_kubernetes_live.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from tests.auth_idp.runtime_factory import iter_auth_idp_cases +from tests.auth_idp.runtime_kubernetes import ( + HELM_UPGRADE_COMMAND_TIMEOUT_SECONDS, + PORT_FORWARD_READY_TIMEOUT_SECONDS, + _add_platform_helm_repositories, + _helm_upgrade_args, +) + +pytestmark = [ + pytest.mark.auth_idp, + pytest.mark.auth_idp_k8s, + pytest.mark.skipif( + not any(case.backend == "kubernetes" for case in iter_auth_idp_cases()), + reason="no Kubernetes auth-idp runtimes are declared", + ), +] + + +def test_authentik_kubernetes_runtime_exports_helm_contract_helpers() -> None: + assert _helm_upgrade_args + assert _add_platform_helm_repositories + assert HELM_UPGRADE_COMMAND_TIMEOUT_SECONDS == 900 + assert PORT_FORWARD_READY_TIMEOUT_SECONDS == 30 diff --git a/tests/auth_idp/providers.py b/tests/auth_idp/providers.py index 9b810015e6..2a691a543a 100644 --- a/tests/auth_idp/providers.py +++ b/tests/auth_idp/providers.py @@ -1,12 +1,21 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import os from dataclasses import dataclass from pathlib import Path import yaml +@dataclass(frozen=True) +class ProviderRuntimeConfig: + id: str + backend: str + capabilities: frozenset[str] + command: str | None = None + + @dataclass(frozen=True) class ProviderConfig: name: str @@ -16,6 +25,9 @@ class ProviderConfig: issuer_url: str discovery_url: str nemo_config: Path + interactive_user_username: str + interactive_user_password: str + interactive_user_expected_email: str workload_principal_id: str workload_expected_groups: list[str] workload_audience: str @@ -25,13 +37,21 @@ class ProviderConfig: workload_token_env_vars: list[str] workload_forwarded_headers: dict[str, str] token_endpoint: str | None - human_grant: dict[str, str] | None - machine_grant: dict[str, str] | None + e2e_setup_password_grant: dict[str, str] | None + interactive_user_password_grant: dict[str, str] | None + workload_provider_password_grant: dict[str, str] | None healthchecks: list[dict[str, str]] startup_timeouts: dict[str, int] + test_runtimes: tuple[ProviderRuntimeConfig, ...] = () compose_project_name: str | None = None +@dataclass(frozen=True) +class ProviderManifestSummary: + name: str + mode: str + + def load_provider_config(manifest_path: Path) -> ProviderConfig: data = yaml.safe_load(manifest_path.read_text()) return ProviderConfig( @@ -42,6 +62,9 @@ def load_provider_config(manifest_path: Path) -> ProviderConfig: issuer_url=data["issuer_url"], discovery_url=data["discovery_url"], nemo_config=manifest_path.parent / data["nemo_config"], + interactive_user_username=data["interactive_user_identity"]["username"], + interactive_user_password=data["interactive_user_identity"]["password"], + interactive_user_expected_email=data["interactive_user_identity"]["expected_email"], workload_principal_id=data["workload_identity"]["principal_id"], workload_expected_groups=list(data["workload_identity"]["expected_groups"]), workload_audience=data["workload_contract"]["audience"], @@ -51,13 +74,45 @@ def load_provider_config(manifest_path: Path) -> ProviderConfig: workload_token_env_vars=list(data["workload_contract"]["token_env_vars"]), workload_forwarded_headers=dict(data["workload_contract"]["forwarded_headers"]), token_endpoint=data.get("token_acquisition", {}).get("token_endpoint"), - human_grant=data.get("token_acquisition", {}).get("human_grant"), - machine_grant=data.get("token_acquisition", {}).get("machine_grant"), + e2e_setup_password_grant=_resolve_grant(data.get("token_acquisition", {}).get("e2e_setup_password_grant")), + interactive_user_password_grant=_resolve_grant( + data.get("token_acquisition", {}).get("interactive_user_password_grant") + ), + workload_provider_password_grant=_resolve_grant( + data.get("token_acquisition", {}).get("workload_provider_password_grant") + ), healthchecks=list(data.get("healthchecks", [])), startup_timeouts=dict(data.get("startup_timeouts", {})), + test_runtimes=_load_provider_runtimes(data), ) +def _load_provider_runtimes(data: dict) -> tuple[ProviderRuntimeConfig, ...]: + return tuple( + ProviderRuntimeConfig( + id=runtime["id"], + backend=runtime["backend"], + command=runtime.get("command"), + capabilities=frozenset(runtime["capabilities"]), + ) + for runtime in data.get("test_runtimes", []) + ) + + +def _resolve_grant(grant: dict[str, str] | None) -> dict[str, str] | None: + if grant is None: + return None + resolved = dict(grant) + password_env_var = resolved.pop("password_env_var", None) + if password_env_var and "password" not in resolved: + password = os.environ.get(password_env_var) + if not password: + resolved["password_env_var"] = password_env_var + return resolved + resolved["password"] = password + return resolved + + def load_provider_configs() -> list[ProviderConfig]: configs: list[ProviderConfig] = [] for manifest_path in sorted(Path("contrib/auth").glob("*/manifest.yaml")): @@ -65,13 +120,21 @@ def load_provider_configs() -> list[ProviderConfig]: return configs +def _load_provider_manifest_summaries() -> list[ProviderManifestSummary]: + summaries: list[ProviderManifestSummary] = [] + for manifest_path in sorted(Path("contrib/auth").glob("*/manifest.yaml")): + data = yaml.safe_load(manifest_path.read_text()) + summaries.append(ProviderManifestSummary(name=data["provider"], mode=data["mode"])) + return summaries + + def load_provider_configs_by_mode(mode: str) -> list[ProviderConfig]: return [provider for provider in load_provider_configs() if provider.mode == mode] def load_provider_names() -> list[str]: - return [provider.name for provider in load_provider_configs()] + return [provider.name for provider in _load_provider_manifest_summaries()] def load_provider_names_by_mode(mode: str) -> list[str]: - return [provider.name for provider in load_provider_configs_by_mode(mode)] + return [provider.name for provider in _load_provider_manifest_summaries() if provider.mode == mode] diff --git a/tests/auth_idp/runtime.py b/tests/auth_idp/runtime.py index 7d0385d7d6..7fd5554ae2 100644 --- a/tests/auth_idp/runtime.py +++ b/tests/auth_idp/runtime.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import os import uuid from dataclasses import replace from functools import lru_cache @@ -8,9 +9,13 @@ from tests.auth_idp.providers import ProviderConfig, load_provider_config +AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD_ENVVAR = "AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD" +AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD_DEFAULT = "svc-nemo-token-secret-e2e" + @lru_cache(maxsize=1) def get_authentik_docker_test_runtime() -> ProviderConfig: + os.environ.setdefault(AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD_ENVVAR, AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD_DEFAULT) provider = load_provider_config(Path("contrib/auth/authentik/manifest.yaml")) return replace( provider, diff --git a/tests/auth_idp/runtime_compose.py b/tests/auth_idp/runtime_compose.py new file mode 100644 index 0000000000..a3f27fb524 --- /dev/null +++ b/tests/auth_idp/runtime_compose.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +from dataclasses import replace +from typing import Callable + +import httpx +from nemo_platform import NeMoPlatform +from nemo_platform_ext.client.tls import client_verify_from_env + +from tests.auth_idp.common import jwt_claims +from tests.auth_idp.runtime_contract import AuthIdpCase, TokenSet + +AUTHENTIK_COMPOSE_WORKLOAD_IDENTITY_PASSWORD = "svc-nemo-token-secret-e2e" +AUTHENTIK_DEFAULT_PASSWORDS_BY_ENVVAR = { + "AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD": AUTHENTIK_COMPOSE_WORKLOAD_IDENTITY_PASSWORD, +} +TOKEN_EXCHANGE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange" +JWT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt" +ACCESS_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:access_token" +TOKEN_EXCHANGE_TIMEOUT_SECONDS = 30.0 + + +def _grant_with_password(grant: dict[str, str]) -> dict[str, str]: + resolved = dict(grant) + password_env_var = resolved.pop("password_env_var", None) + if password_env_var and "password" not in resolved: + password = os.environ.get(password_env_var) or AUTHENTIK_DEFAULT_PASSWORDS_BY_ENVVAR.get(password_env_var) + if password is None: + raise AssertionError(f"{password_env_var} must be set for token acquisition") + resolved["password"] = password + return resolved + + +class ComposeAuthIdpRuntime: + def __init__(self, case: AuthIdpCase, gateway_base_url: str, cleanup: Callable[[], None] | None = None): + self.case = case + self._cleanup = cleanup + self._cleaned_up = False + self.provider = replace( + case.provider, + gateway_base_url=gateway_base_url, + discovery_url=f"{gateway_base_url}/application/o/nemo/.well-known/openid-configuration", + token_endpoint=f"{gateway_base_url}/application/o/token/", + ) + self.gateway_base_url = self.provider.gateway_base_url + self.discovery_url = self.provider.discovery_url + self.token_endpoint = self.provider.token_endpoint + self.workload_token_endpoint = f"{gateway_base_url}/apis/auth/token" + + def e2e_setup_token(self) -> TokenSet: + assert self.provider.e2e_setup_password_grant is not None + assert self.token_endpoint is not None + token = self._exchange_token(self.token_endpoint, _grant_with_password(self.provider.e2e_setup_password_grant)) + return TokenSet(access_token=token, claims=jwt_claims(token)) + + def interactive_user_token(self) -> TokenSet: + assert self.provider.interactive_user_password_grant is not None + assert self.token_endpoint is not None + token = self._exchange_token(self.token_endpoint, self.provider.interactive_user_password_grant) + return TokenSet(access_token=token, claims=jwt_claims(token)) + + def workload_provider_token(self) -> TokenSet: + assert self.provider.workload_provider_password_grant is not None + assert self.token_endpoint is not None + token = self._exchange_token( + self.token_endpoint, + _grant_with_password(self.provider.workload_provider_password_grant), + ) + return TokenSet(access_token=token, claims=jwt_claims(token)) + + def workload_subject_token(self) -> str: + return self.workload_provider_token().access_token + + def exchange_workload_token(self, subject_token: str) -> TokenSet: + assert self.workload_token_endpoint is not None + assert self.provider.workload_provider_password_grant is not None + workload_grant = self.provider.workload_provider_password_grant + response = httpx.post( + self.workload_token_endpoint, + data={ + "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, + "client_id": workload_grant["client_id"], + "subject_token": subject_token, + "subject_token_type": JWT_TOKEN_TYPE, + "requested_token_type": ACCESS_TOKEN_TYPE, + "audience": self.provider.workload_audience, + "scope": workload_grant.get("scope", "openid email groups"), + }, + timeout=TOKEN_EXCHANGE_TIMEOUT_SECONDS, + verify=client_verify_from_env(), + ) + response.raise_for_status() + token_response = response.json() + access_token = token_response["access_token"] + assert token_response.get("token_type", "").lower() == "bearer" + return TokenSet(access_token=access_token, claims=jwt_claims(access_token)) + + def e2e_setup_sdk(self) -> NeMoPlatform: + token = self.e2e_setup_token().access_token + return NeMoPlatform( + base_url=self.gateway_base_url, + default_headers={"Authorization": f"Bearer {token}"}, + max_retries=0, + ) + + def interactive_user_sdk(self) -> NeMoPlatform: + token = self.interactive_user_token().access_token + return NeMoPlatform( + base_url=self.gateway_base_url, + default_headers={"Authorization": f"Bearer {token}"}, + max_retries=0, + ) + + def workload_provider_sdk(self) -> NeMoPlatform: + token = self.workload_provider_token().access_token + return NeMoPlatform( + base_url=self.gateway_base_url, + default_headers={"Authorization": f"Bearer {token}"}, + max_retries=0, + ) + + def workload_role_principals(self) -> list[str]: + return list(self.provider.workload_expected_groups) + + def cleanup(self) -> None: + if self._cleaned_up: + return + self._cleaned_up = True + if self._cleanup is not None: + self._cleanup() + + def _exchange_token(self, token_endpoint: str, grant: dict[str, str]) -> str: + from tests.auth_idp.conftest import _exchange_token_with_retries + + return _exchange_token_with_retries(token_endpoint, grant) diff --git a/tests/auth_idp/runtime_contract.py b/tests/auth_idp/runtime_contract.py new file mode 100644 index 0000000000..d7902281ac --- /dev/null +++ b/tests/auth_idp/runtime_contract.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Literal, Protocol + +from nemo_platform import NeMoPlatform + +from tests.auth_idp.providers import ProviderConfig + +AuthIdpBackend = Literal["compose", "kubernetes", "external"] + + +@dataclass(frozen=True) +class AuthIdpCase: + id: str + provider: ProviderConfig + backend: AuthIdpBackend + capabilities: frozenset[str] + marks: tuple[object, ...] = () + + +@dataclass(frozen=True) +class TokenSet: + access_token: str + claims: dict[str, object] + + +class AuthIdpRuntime(Protocol): + case: AuthIdpCase + gateway_base_url: str + discovery_url: str + token_endpoint: str | None + workload_token_endpoint: str | None + + def e2e_setup_token(self) -> TokenSet: + raise NotImplementedError + + def interactive_user_token(self) -> TokenSet: + raise NotImplementedError + + def workload_provider_token(self) -> TokenSet: + raise NotImplementedError + + def workload_subject_token(self) -> str: + raise NotImplementedError + + def exchange_workload_token(self, subject_token: str) -> TokenSet: + raise NotImplementedError + + def e2e_setup_sdk(self) -> NeMoPlatform: + raise NotImplementedError + + def interactive_user_sdk(self) -> NeMoPlatform: + raise NotImplementedError + + def workload_provider_sdk(self) -> NeMoPlatform: + raise NotImplementedError + + def workload_role_principals(self) -> list[str]: + raise NotImplementedError + + def cleanup(self) -> None: + raise NotImplementedError + + +class RuntimeManager(Protocol): + def start(self, case: AuthIdpCase) -> Iterator[AuthIdpRuntime]: + raise NotImplementedError diff --git a/tests/auth_idp/runtime_factory.py b/tests/auth_idp/runtime_factory.py new file mode 100644 index 0000000000..dfd343952f --- /dev/null +++ b/tests/auth_idp/runtime_factory.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Iterable +from typing import cast + +import pytest + +from tests.auth_idp.providers import load_provider_configs +from tests.auth_idp.runtime_contract import AuthIdpBackend, AuthIdpCase + + +def iter_auth_idp_cases( + *, + backend: AuthIdpBackend | None = None, + provider_name: str | None = None, + runtime_id: str | None = None, + required_capability: str | None = None, +) -> list[AuthIdpCase]: + cases: list[AuthIdpCase] = [] + for provider in load_provider_configs(): + if provider_name is not None and provider.name != provider_name: + continue + for runtime in provider.test_runtimes: + runtime_backend = cast(AuthIdpBackend, runtime.backend) + if runtime_id is not None and runtime.id != runtime_id: + continue + if backend is not None and runtime_backend != backend: + continue + if required_capability is not None and required_capability not in runtime.capabilities: + continue + cases.append( + AuthIdpCase( + id=runtime.id, + provider=provider, + backend=runtime_backend, + capabilities=runtime.capabilities, + marks=(), + ) + ) + return cases + + +def parametrize_cases(cases: Iterable[AuthIdpCase]) -> list[object]: + return [pytest.param(case, id=case.id, marks=case.marks) for case in cases] + + +def runtime_class_for_case(case: AuthIdpCase) -> type: + if case.backend == "compose": + from tests.auth_idp.runtime_compose import ComposeAuthIdpRuntime + + return ComposeAuthIdpRuntime + if case.backend == "kubernetes": + from tests.auth_idp.runtime_kubernetes import KubernetesAuthIdpRuntime + + return KubernetesAuthIdpRuntime + raise ValueError(f"unsupported auth-idp runtime backend: {case.backend}") diff --git a/tests/auth_idp/runtime_kubernetes.py b/tests/auth_idp/runtime_kubernetes.py new file mode 100644 index 0000000000..e16c387ac2 --- /dev/null +++ b/tests/auth_idp/runtime_kubernetes.py @@ -0,0 +1,917 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import base64 +import contextlib +import json +import os +import shutil +import socket +import subprocess +import tempfile +import time +import uuid +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import TypeVar + +import httpx +import pytest +from nemo_platform import DefaultHttpxClient, NeMoPlatform +from nemo_platform_ext.client.tls import NMP_CLIENT_SSL_CERT_FILE_ENVVAR + +from tests.auth_idp.common import jwt_claims +from tests.auth_idp.runtime_contract import AuthIdpCase, TokenSet + +REPO_ROOT = Path(__file__).resolve().parents[2] +NAMESPACE = os.environ.get("NMP_AUTHENTIK_K8S_NAMESPACE", "nemo-authentik") +HELM_RELEASE = os.environ.get("NMP_AUTHENTIK_K8S_HELM_RELEASE", "authentik-demo") +HELM_CHART = Path("contrib/auth/authentik/helm") +ENVOY_TLS_SECRET = "nemo-platform-envoy-tls" +WORKLOAD_AUDIENCE = "nemo-platform" +WORKLOAD_CLIENT_ID = "nemo-platform-workload" +AUTHENTIK_K8S_WORKLOAD_IDENTITY_PASSWORD = "workload-identity-dev-only" +WORKLOAD_TOKEN_PRIVATE_KEY_FILE_ENV = "NMP_AUTHENTIK_K8S_WORKLOAD_TOKEN_PRIVATE_KEY_FILE" +GATEWAY_PORT_ENV = "NMP_AUTHENTIK_K8S_GATEWAY_PORT" +DISCOVERY_PATH = "/application/o/nemo/.well-known/openid-configuration" +GATEWAY_READY_PATH = "/health/gateway/ready" +TOKEN_EXCHANGE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange" +JWT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt" +ACCESS_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:access_token" +T = TypeVar("T") + + +# Keep timeouts centralized so slow-cluster tuning is a single, visible edit. +PYTEST_TIMEOUT_SECONDS = 900 +DEFAULT_COMMAND_TIMEOUT_SECONDS = 120 +DIAGNOSTIC_COMMAND_TIMEOUT_SECONDS = 60 +POD_DISCOVERY_TIMEOUT_SECONDS = 60 +CLUSTER_CREATE_TIMEOUT_SECONDS = 360 +KIND_CREATE_WAIT_TIMEOUT = "180s" +IMAGE_LOAD_TIMEOUT_SECONDS = 300 +CLUSTER_DELETE_TIMEOUT_SECONDS = 180 +ROLLOUT_STATUS_TIMEOUT = "240s" +ROLLOUT_COMMAND_TIMEOUT_SECONDS = 300 +HELM_WAIT_TIMEOUT = "10m" +HELM_UPGRADE_COMMAND_TIMEOUT_SECONDS = 900 +HELM_REPO_TIMEOUT_SECONDS = 60 +HELM_DEPENDENCY_TIMEOUT_SECONDS = 300 +HTTP_RETRY_TIMEOUT_SECONDS = 180 +HTTP_REQUEST_TIMEOUT_SECONDS = 10.0 +RETRY_SLEEP_SECONDS = 2.0 +SECRET_TIMEOUT_SECONDS = 180 +SECRET_GET_TIMEOUT_SECONDS = 30 +PORT_FORWARD_READY_TIMEOUT_SECONDS = 30 +PORT_FORWARD_HTTP_TIMEOUT_SECONDS = 2.0 +PORT_FORWARD_RETRY_SLEEP_SECONDS = 0.5 +PORT_FORWARD_TERMINATE_TIMEOUT_SECONDS = 10 +SUBJECT_TOKEN_DURATION = "10m" +TOKEN_EXCHANGE_TIMEOUT_SECONDS = 30.0 + + +@dataclass(frozen=True) +class Cluster: + name: str + runtime: str + context: str + kubeconfig: Path | None = None + cleanup_kubeconfig: bool = False + + +def _require_tool(name: str) -> None: + if shutil.which(name) is None: + pytest.skip(f"{name} is required for the Authentik Kubernetes E2E test") + + +def _run(args: list[str], *, timeout: float = DEFAULT_COMMAND_TIMEOUT_SECONDS) -> subprocess.CompletedProcess[str]: + completed = subprocess.run( + args, + cwd=REPO_ROOT, + text=True, + capture_output=True, + timeout=timeout, + check=False, + ) + if completed.returncode != 0: + command = " ".join(args) + stdout = completed.stdout[-4000:] + stderr = completed.stderr[-4000:] + raise AssertionError( + f"command failed ({completed.returncode}): {command}\nstdout:\n{stdout}\nstderr:\n{stderr}" + ) + return completed + + +def _temporary_kubeconfig_path(cluster_name: str) -> Path: + temp_file = tempfile.NamedTemporaryFile( + prefix=f"nmp-authentik-{cluster_name}-", + suffix="-kubeconfig.yaml", + delete=False, + ) + temp_file.close() + return Path(temp_file.name) + + +def _kubectl_command(context: str, args: list[str], kubeconfig: Path | None = None) -> list[str]: + command = ["kubectl"] + if kubeconfig is not None: + command.extend(["--kubeconfig", str(kubeconfig)]) + command.extend(["--context", context, *args]) + return command + + +def _helm_command(context: str, args: list[str], kubeconfig: Path | None = None) -> list[str]: + command = ["helm"] + if kubeconfig is not None: + command.extend(["--kubeconfig", str(kubeconfig)]) + command.extend(["--kube-context", context, *args]) + return command + + +def _write_diagnostic_process( + log_dir: Path, + name: str, + args: list[str], + *, + timeout: float = DIAGNOSTIC_COMMAND_TIMEOUT_SECONDS, +) -> None: + completed = subprocess.run( + args, + cwd=REPO_ROOT, + text=True, + capture_output=True, + timeout=timeout, + check=False, + ) + (log_dir / name).write_text( + "\n".join( + ( + f"command: {' '.join(args)}", + f"exit_code: {completed.returncode}", + "stdout:", + completed.stdout, + "stderr:", + completed.stderr, + ) + ), + encoding="utf-8", + ) + + +def _diagnostic_output_text(output: str | bytes | None) -> str: + if output is None: + return "" + if isinstance(output, bytes): + return output.decode(errors="replace") + return output + + +def _write_diagnostic_timeout(log_dir: Path, name: str, args: list[str], exc: subprocess.TimeoutExpired) -> None: + (log_dir / name).write_text( + "\n".join( + ( + f"command: {' '.join(args)}", + f"timeout: {exc.timeout}", + "stdout:", + _diagnostic_output_text(exc.stdout), + "stderr:", + _diagnostic_output_text(exc.stderr), + ) + ), + encoding="utf-8", + ) + + +def _write_diagnostic_command( + context: str, + log_dir: Path, + name: str, + args: list[str], + *, + kubeconfig: Path | None = None, + timeout: float = DIAGNOSTIC_COMMAND_TIMEOUT_SECONDS, +) -> None: + _write_diagnostic_process(log_dir, name, _kubectl_command(context, args, kubeconfig), timeout=timeout) + + +def _collect_kubernetes_diagnostics(context: str, cluster_name: str, kubeconfig: Path | None = None) -> Path: + configured_dir = os.environ.get("NMP_AUTHENTIK_K8S_LOG_DIR") + log_root = Path(configured_dir) if configured_dir else REPO_ROOT / "docker" / "logs" + log_dir = log_root / f"k8s-authentik-{cluster_name}" + log_dir.mkdir(parents=True, exist_ok=True) + + for name, args in { + "helm-list.txt": _helm_command(context, ["-n", NAMESPACE, "list", "--all"], kubeconfig), + "helm-status.txt": _helm_command(context, ["-n", NAMESPACE, "status", HELM_RELEASE], kubeconfig), + }.items(): + with contextlib.suppress(Exception): + _write_diagnostic_process(log_dir, name, args) + + for name, args in { + "get-nodes.txt": ["get", "nodes", "-o", "wide"], + "get-all.txt": ["-n", NAMESPACE, "get", "all", "-o", "wide"], + "get-pods-json.txt": ["-n", NAMESPACE, "get", "pods", "-o", "json"], + "events.txt": ["-n", NAMESPACE, "get", "events", "--sort-by=.lastTimestamp"], + "describe-pods.txt": ["-n", NAMESPACE, "describe", "pods"], + "describe-blueprint-configmap.txt": ["-n", NAMESPACE, "describe", "configmap/authentik-nemo-blueprint"], + }.items(): + with contextlib.suppress(Exception): + _write_diagnostic_command(context, log_dir, name, args, kubeconfig=kubeconfig) + + pod_discovery_args = _kubectl_command( + context, + [ + "-n", + NAMESPACE, + "get", + "pods", + "-o", + "jsonpath={range .items[*]}{.metadata.name}{'\\n'}{end}", + ], + kubeconfig, + ) + try: + pods = subprocess.run( + pod_discovery_args, + cwd=REPO_ROOT, + text=True, + capture_output=True, + timeout=POD_DISCOVERY_TIMEOUT_SECONDS, + check=False, + ) + except subprocess.TimeoutExpired as exc: + with contextlib.suppress(Exception): + _write_diagnostic_timeout(log_dir, "get-pods-for-logs.txt", pod_discovery_args, exc) + return log_dir + + for pod in pods.stdout.splitlines(): + safe_pod = pod.replace("/", "_") + with contextlib.suppress(Exception): + _write_diagnostic_command( + context, + log_dir, + f"logs-{safe_pod}.txt", + ["-n", NAMESPACE, "logs", pod, "--all-containers", "--timestamps"], + kubeconfig=kubeconfig, + ) + with contextlib.suppress(Exception): + _write_diagnostic_command( + context, + log_dir, + f"logs-{safe_pod}-previous.txt", + ["-n", NAMESPACE, "logs", pod, "--all-containers", "--previous", "--timestamps"], + kubeconfig=kubeconfig, + ) + + return log_dir + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _configured_gateway_port() -> int | None: + configured = os.environ.get(GATEWAY_PORT_ENV) + if not configured: + return None + try: + port = int(configured) + except ValueError as exc: + raise AssertionError(f"{GATEWAY_PORT_ENV} must be an integer TCP port") from exc + if port < 1 or port > 65535: + raise AssertionError(f"{GATEWAY_PORT_ENV} must be between 1 and 65535") + return port + + +def _platform_image() -> str: + return os.environ.get( + "NMP_AUTHENTIK_K8S_PLATFORM_IMAGE", + f"{os.environ.get('IMAGE_REGISTRY', 'my-registry')}/nmp-api:{os.environ.get('BAKE_TAG', 'local')}", + ) + + +def _create_cluster(runtime: str, name: str) -> Cluster: + _require_tool("docker") + _require_tool("kubectl") + kubeconfig = _temporary_kubeconfig_path(name) + + try: + if runtime == "k3d": + _require_tool("k3d") + _run( + [ + "k3d", + "cluster", + "create", + name, + "--wait", + "--agents", + "0", + "--k3s-arg", + "--disable=traefik@server:0", + "--kubeconfig-update-default=false", + ], + timeout=CLUSTER_CREATE_TIMEOUT_SECONDS, + ) + kubeconfig.write_text(_run(["k3d", "kubeconfig", "get", name]).stdout, encoding="utf-8") + return Cluster( + name=name, + runtime=runtime, + context=f"k3d-{name}", + kubeconfig=kubeconfig, + cleanup_kubeconfig=True, + ) + + if runtime == "kind": + _require_tool("kind") + _run( + [ + "kind", + "create", + "cluster", + "--name", + name, + "--kubeconfig", + str(kubeconfig), + "--wait", + KIND_CREATE_WAIT_TIMEOUT, + ], + timeout=CLUSTER_CREATE_TIMEOUT_SECONDS, + ) + return Cluster( + name=name, + runtime=runtime, + context=f"kind-{name}", + kubeconfig=kubeconfig, + cleanup_kubeconfig=True, + ) + except Exception: + with contextlib.suppress(FileNotFoundError): + kubeconfig.unlink() + raise + + raise ValueError(f"unsupported NMP_AUTHENTIK_K8S_RUNTIME={runtime!r}; expected kind or k3d") + + +def _existing_cluster(runtime: str, name: str) -> Cluster | None: + _require_tool("kubectl") + kubeconfig = _temporary_kubeconfig_path(name) + try: + if runtime == "k3d": + _require_tool("k3d") + kubeconfig.write_text(_run(["k3d", "kubeconfig", "get", name]).stdout, encoding="utf-8") + return Cluster( + name=name, + runtime=runtime, + context=f"k3d-{name}", + kubeconfig=kubeconfig, + cleanup_kubeconfig=True, + ) + + if runtime == "kind": + _require_tool("kind") + kubeconfig.write_text(_run(["kind", "get", "kubeconfig", "--name", name]).stdout, encoding="utf-8") + return Cluster( + name=name, + runtime=runtime, + context=f"kind-{name}", + kubeconfig=kubeconfig, + cleanup_kubeconfig=True, + ) + except AssertionError: + with contextlib.suppress(FileNotFoundError): + kubeconfig.unlink() + return None + except Exception: + with contextlib.suppress(FileNotFoundError): + kubeconfig.unlink() + raise + + with contextlib.suppress(FileNotFoundError): + kubeconfig.unlink() + raise ValueError(f"unsupported NMP_AUTHENTIK_K8S_RUNTIME={runtime!r}; expected kind or k3d") + + +def _reuse_or_create_cluster(runtime: str, name: str) -> Cluster: + cluster = _existing_cluster(runtime, name) + if cluster is not None: + return cluster + return _create_cluster(runtime, name) + + +def _load_platform_image(runtime: str, name: str, image: str) -> None: + if os.environ.get("NMP_AUTHENTIK_K8S_SKIP_IMAGE_LOAD") == "1": + return + if runtime == "k3d": + _run(["k3d", "image", "import", image, "-c", name], timeout=IMAGE_LOAD_TIMEOUT_SECONDS) + return + if runtime == "kind": + _run(["kind", "load", "docker-image", image, "--name", name], timeout=IMAGE_LOAD_TIMEOUT_SECONDS) + return + raise ValueError(f"unsupported NMP_AUTHENTIK_K8S_RUNTIME={runtime!r}; expected kind or k3d") + + +def _delete_cluster(runtime: str, name: str, kubeconfig: Path | None = None) -> None: + with contextlib.suppress(Exception): + if runtime == "k3d": + _run(["k3d", "cluster", "delete", name], timeout=CLUSTER_DELETE_TIMEOUT_SECONDS) + elif runtime == "kind": + command = ["kind", "delete", "cluster", "--name", name] + if kubeconfig is not None: + command.extend(["--kubeconfig", str(kubeconfig)]) + _run(command, timeout=CLUSTER_DELETE_TIMEOUT_SECONDS) + + +def _reuse_context(runtime: str, name: str) -> str: + if runtime == "k3d": + return f"k3d-{name}" + if runtime == "kind": + return f"kind-{name}" + raise ValueError(f"unsupported NMP_AUTHENTIK_K8S_RUNTIME={runtime!r}; expected kind or k3d") + + +def _wait_for_authentik(context: str, kubeconfig: Path | None = None) -> None: + for deployment in ( + "authentik-server", + "authentik-worker", + "nemo-platform-api", + "nemo-platform-envoy", + ): + _run( + _kubectl_command( + context, + [ + "-n", + NAMESPACE, + "rollout", + "status", + f"deploy/{deployment}", + f"--timeout={ROLLOUT_STATUS_TIMEOUT}", + ], + kubeconfig, + ), + timeout=ROLLOUT_COMMAND_TIMEOUT_SECONDS, + ) + _run( + _kubectl_command( + context, + [ + "-n", + NAMESPACE, + "rollout", + "status", + "statefulset/shared-postgresql", + f"--timeout={ROLLOUT_STATUS_TIMEOUT}", + ], + kubeconfig, + ), + timeout=ROLLOUT_COMMAND_TIMEOUT_SECONDS, + ) + + +def _helm_upgrade_args(context: str, kubeconfig: Path | None = None) -> list[str]: + image = _platform_image() + registry, tag = image.rsplit("/nmp-api:", 1) + args = _helm_command( + context, + [ + "upgrade", + "--install", + HELM_RELEASE, + str(HELM_CHART), + "--namespace", + NAMESPACE, + "--create-namespace", + "--wait", + "--wait-for-jobs", + "--timeout", + HELM_WAIT_TIMEOUT, + "--set", + f"nemo-platform.api.image.repository={registry}/nmp-api", + "--set", + f"nemo-platform.api.image.tag={tag}", + "--set", + f"nemo-platform.core.image.repository={registry}/nmp-api", + "--set", + f"nemo-platform.core.image.tag={tag}", + "--set-string", + f"nemo-platform.platformConfig.platform.image_registry={registry}", + "--set-string", + f"nemo-platform.platformConfig.platform.image_tag={tag}", + ], + kubeconfig, + ) + ngc_existing_secret = os.environ.get("NMP_AUTHENTIK_K8S_NGC_EXISTING_SECRET") + if ngc_existing_secret: + args.extend( + [ + "--set-string", + f"nemo-platform.existingSecret={ngc_existing_secret}", + ] + ) + image_pull_secret = os.environ.get("NMP_AUTHENTIK_K8S_IMAGE_PULL_SECRET") + if image_pull_secret: + args.extend( + [ + "--set-string", + f"nemo-platform.imagePullSecrets[0].name={image_pull_secret}", + ] + ) + gateway_port = _configured_gateway_port() + if gateway_port is not None: + args.extend( + [ + "--set-string", + f"nemo-platform.authentikPublicGateway.port={gateway_port}", + ] + ) + workload_token_private_key = os.environ.get(WORKLOAD_TOKEN_PRIVATE_KEY_FILE_ENV) + if workload_token_private_key: + args.extend( + [ + "--set-file", + f"workloadTokenSigningKey.privateKeyPem={workload_token_private_key}", + ] + ) + return args + + +def _add_platform_helm_repositories() -> None: + _run( + ["helm", "repo", "add", "nvidia", "https://helm.ngc.nvidia.com/nvidia", "--force-update"], + timeout=HELM_REPO_TIMEOUT_SECONDS, + ) + _run( + ["helm", "repo", "add", "authentik", "https://charts.goauthentik.io", "--force-update"], + timeout=HELM_REPO_TIMEOUT_SECONDS, + ) + + +def _helm_install_authentik_demo(context: str, kubeconfig: Path | None = None) -> None: + _require_tool("helm") + _add_platform_helm_repositories() + _run(["helm", "dependency", "build", "k8s/helm"], timeout=HELM_DEPENDENCY_TIMEOUT_SECONDS) + _run(["helm", "dependency", "build", str(HELM_CHART)], timeout=HELM_DEPENDENCY_TIMEOUT_SECONDS) + + _run(_helm_upgrade_args(context, kubeconfig), timeout=HELM_UPGRADE_COMMAND_TIMEOUT_SECONDS) + + +def _retry_until_timeout( + attempt: Callable[[float], T | None], + *, + timeout: float, + sleep: float, + retry_exceptions: tuple[type[Exception], ...], + timeout_message: str, + raise_last_retry_error: bool = True, +) -> T: + deadline = time.monotonic() + timeout + last_error: Exception | None = None + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + try: + result = attempt(remaining) + if result is not None: + return result + except retry_exceptions as exc: + last_error = exc + remaining = deadline - time.monotonic() + if remaining <= 0: + break + time.sleep(min(sleep, remaining)) + if last_error is not None: + if raise_last_retry_error: + raise last_error + raise TimeoutError(timeout_message) from last_error + raise TimeoutError(timeout_message) + + +def _get_json_with_retries( + url: str, + *, + timeout: float = HTTP_RETRY_TIMEOUT_SECONDS, + verify: str | bool = True, +) -> dict: + def get_json(remaining: float) -> dict | None: + response = httpx.get(url, timeout=min(HTTP_REQUEST_TIMEOUT_SECONDS, remaining), verify=verify) + if response.status_code >= 500: + return None + response.raise_for_status() + return response.json() + + return _retry_until_timeout( + get_json, + timeout=timeout, + sleep=RETRY_SLEEP_SECONDS, + retry_exceptions=(httpx.HTTPError, ValueError), + timeout_message=f"timed out waiting for {url}", + ) + + +def _secret_data( + context: str, + secret_name: str, + key: str, + *, + kubeconfig: Path | None = None, + timeout: float = SECRET_TIMEOUT_SECONDS, +) -> bytes: + def get_secret(remaining: float) -> bytes: + secret = json.loads( + _run( + _kubectl_command( + context, + [ + "-n", + NAMESPACE, + "get", + "secret", + secret_name, + "-o", + "json", + ], + kubeconfig, + ), + timeout=min(SECRET_GET_TIMEOUT_SECONDS, remaining), + ).stdout + ) + return base64.b64decode(secret["data"][key]) + + return _retry_until_timeout( + get_secret, + timeout=timeout, + sleep=RETRY_SLEEP_SECONDS, + retry_exceptions=(AssertionError, KeyError, json.JSONDecodeError, ValueError), + timeout_message=f"timed out waiting for secret {NAMESPACE}/{secret_name} key {key}", + ) + + +def _start_port_forward_service( + context: str, + service: str, + ca_bundle: Path, + kubeconfig: Path | None = None, +) -> tuple[str, subprocess.Popen[str]]: + port = _configured_gateway_port() or _free_port() + process = subprocess.Popen( + _kubectl_command( + context, + [ + "-n", + NAMESPACE, + "port-forward", + f"svc/{service}", + f"{port}:8080", + ], + kubeconfig, + ), + cwd=REPO_ROOT, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + text=True, + ) + gateway_url = f"https://127.0.0.1:{port}" + + def wait_for_gateway_ready(remaining: float) -> httpx.Response: + if process.poll() is not None: + raise AssertionError(f"kubectl port-forward exited early with {process.returncode}") + return httpx.get( + gateway_url + GATEWAY_READY_PATH, + timeout=min(PORT_FORWARD_HTTP_TIMEOUT_SECONDS, remaining), + verify=str(ca_bundle), + ) + + try: + _retry_until_timeout( + wait_for_gateway_ready, + timeout=PORT_FORWARD_READY_TIMEOUT_SECONDS, + sleep=PORT_FORWARD_RETRY_SLEEP_SECONDS, + retry_exceptions=(httpx.RequestError,), + timeout_message=f"timed out waiting for port-forward readiness at {gateway_url + GATEWAY_READY_PATH}", + raise_last_retry_error=False, + ) + except Exception: + _terminate_process(process) + raise + return gateway_url, process + + +def _terminate_process(process: subprocess.Popen[str]) -> None: + process.terminate() + with contextlib.suppress(subprocess.TimeoutExpired): + process.wait(timeout=PORT_FORWARD_TERMINATE_TIMEOUT_SECONDS) + if process.poll() is None: + process.kill() + process.wait() + + +def _grant_with_password(grant: dict[str, str], *, default_password: str | None = None) -> dict[str, str]: + resolved = dict(grant) + password_env_var = resolved.pop("password_env_var", None) + if password_env_var and "password" not in resolved: + password = os.environ.get(password_env_var) or default_password + if password is None: + raise AssertionError(f"{password_env_var} must be set for token acquisition") + resolved["password"] = password + return resolved + + +class KubernetesAuthIdpRuntime: + def __init__(self, case: AuthIdpCase): + self.case = case + self.provider = case.provider + self.cluster: Cluster | None = None + self.namespace = NAMESPACE + self.helm_release = HELM_RELEASE + self.gateway_base_url = "" + self.discovery_url = "" + self.token_endpoint: str | None = None + self.workload_token_endpoint: str | None = None + self.ca_bundle: Path | None = None + self._ca_temp_file: tempfile._TemporaryFileWrapper[bytes] | None = None + self._port_forward_process: subprocess.Popen[str] | None = None + self._diagnostics_collected = False + self._reuse_cluster = False + self._keep_cluster = False + self._previous_client_ssl_cert_file = os.environ.get(NMP_CLIENT_SSL_CERT_FILE_ENVVAR) + self._start() + + @property + def verify(self) -> str: + assert self.ca_bundle is not None + return str(self.ca_bundle) + + def e2e_setup_token(self) -> TokenSet: + assert self.provider.e2e_setup_password_grant is not None + assert self.token_endpoint is not None + token = self._exchange_token(self.token_endpoint, self.provider.e2e_setup_password_grant) + return TokenSet(access_token=token, claims=jwt_claims(token)) + + def interactive_user_token(self) -> TokenSet: + assert self.provider.interactive_user_password_grant is not None + assert self.token_endpoint is not None + token = self._exchange_token(self.token_endpoint, self.provider.interactive_user_password_grant) + return TokenSet(access_token=token, claims=jwt_claims(token)) + + def workload_provider_token(self) -> TokenSet: + assert self.provider.workload_provider_password_grant is not None + assert self.token_endpoint is not None + grant = _grant_with_password( + self.provider.workload_provider_password_grant, + default_password=AUTHENTIK_K8S_WORKLOAD_IDENTITY_PASSWORD, + ) + token = self._exchange_token(self.token_endpoint, grant) + return TokenSet(access_token=token, claims=jwt_claims(token)) + + def workload_subject_token(self) -> str: + assert self.cluster is not None + return _run( + _kubectl_command( + self.cluster.context, + [ + "-n", + NAMESPACE, + "create", + "token", + "default", + "--audience", + WORKLOAD_CLIENT_ID, + "--duration", + SUBJECT_TOKEN_DURATION, + ], + self.cluster.kubeconfig, + ), + ).stdout.strip() + + def exchange_workload_token(self, subject_token: str) -> TokenSet: + assert self.workload_token_endpoint is not None + response = httpx.post( + self.workload_token_endpoint, + data={ + "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, + "client_id": WORKLOAD_CLIENT_ID, + "subject_token": subject_token, + "subject_token_type": JWT_TOKEN_TYPE, + "requested_token_type": ACCESS_TOKEN_TYPE, + "audience": WORKLOAD_AUDIENCE, + "scope": "openid email groups", + }, + timeout=TOKEN_EXCHANGE_TIMEOUT_SECONDS, + verify=self.verify, + ) + response.raise_for_status() + token_response = response.json() + access_token = token_response["access_token"] + assert token_response.get("token_type", "").lower() == "bearer" + return TokenSet(access_token=access_token, claims=jwt_claims(access_token)) + + def e2e_setup_sdk(self) -> NeMoPlatform: + return self._sdk_for_token(self.e2e_setup_token().access_token) + + def interactive_user_sdk(self) -> NeMoPlatform: + return self._sdk_for_token(self.interactive_user_token().access_token) + + def workload_provider_sdk(self) -> NeMoPlatform: + return self._sdk_for_token(self.exchange_workload_token(self.workload_subject_token()).access_token) + + def workload_role_principals(self) -> list[str]: + return [f"system:serviceaccounts:{NAMESPACE}"] + + def _collect_diagnostics_best_effort( + self, + context: str, + cluster_name: str, + kubeconfig: Path | None = None, + ) -> None: + if self._diagnostics_collected: + return + try: + with contextlib.suppress(Exception): + log_dir = _collect_kubernetes_diagnostics(context, cluster_name, kubeconfig) + print(f"Collected Authentik Kubernetes diagnostics: {log_dir}") + finally: + self._diagnostics_collected = True + + def cleanup(self) -> None: + if self._port_forward_process is not None: + _terminate_process(self._port_forward_process) + self._port_forward_process = None + if self._ca_temp_file is not None: + with contextlib.suppress(FileNotFoundError): + Path(self._ca_temp_file.name).unlink() + self._ca_temp_file = None + if self._previous_client_ssl_cert_file is None: + os.environ.pop(NMP_CLIENT_SSL_CERT_FILE_ENVVAR, None) + else: + os.environ[NMP_CLIENT_SSL_CERT_FILE_ENVVAR] = self._previous_client_ssl_cert_file + if self.cluster is not None: + cluster = self.cluster + try: + self._collect_diagnostics_best_effort(cluster.context, cluster.name, cluster.kubeconfig) + finally: + try: + if not self._reuse_cluster and not self._keep_cluster: + _delete_cluster(cluster.runtime, cluster.name, cluster.kubeconfig) + finally: + if cluster.cleanup_kubeconfig and not self._keep_cluster and cluster.kubeconfig is not None: + with contextlib.suppress(FileNotFoundError): + cluster.kubeconfig.unlink() + self.cluster = None + + def _start(self) -> None: + runtime = os.environ.get("NMP_AUTHENTIK_K8S_RUNTIME", "kind") + name = os.environ.get("NMP_AUTHENTIK_K8S_CLUSTER_NAME") or f"nmp-authentik-{uuid.uuid4().hex[:8]}" + self._reuse_cluster = os.environ.get("NMP_AUTHENTIK_K8S_REUSE_CLUSTER") == "1" + self._keep_cluster = os.environ.get("NMP_AUTHENTIK_K8S_KEEP_CLUSTER") == "1" + cluster = _reuse_or_create_cluster(runtime, name) if self._reuse_cluster else _create_cluster(runtime, name) + self.cluster = cluster + try: + _load_platform_image(runtime, name, _platform_image()) + _helm_install_authentik_demo(cluster.context, cluster.kubeconfig) + _wait_for_authentik(cluster.context, cluster.kubeconfig) + self._write_ca_bundle(cluster.context, cluster.kubeconfig) + assert self.ca_bundle is not None + self.gateway_base_url, self._port_forward_process = _start_port_forward_service( + cluster.context, + "nemo-platform-envoy", + self.ca_bundle, + cluster.kubeconfig, + ) + self.discovery_url = self.gateway_base_url + DISCOVERY_PATH + self.token_endpoint = self.gateway_base_url + "/application/o/token/" + self.workload_token_endpoint = self.gateway_base_url + "/apis/auth/token" + os.environ[NMP_CLIENT_SSL_CERT_FILE_ENVVAR] = self.verify + except Exception: + try: + self._collect_diagnostics_best_effort(cluster.context, cluster.name, cluster.kubeconfig) + finally: + self.cleanup() + raise + + def _write_ca_bundle(self, context: str, kubeconfig: Path | None = None) -> None: + temp_file = tempfile.NamedTemporaryFile(suffix="-nmp-ca.crt", delete=False) + temp_file.write(_secret_data(context, ENVOY_TLS_SECRET, "ca.crt", kubeconfig=kubeconfig)) + temp_file.flush() + temp_file.close() + self._ca_temp_file = temp_file + self.ca_bundle = Path(temp_file.name) + + def _exchange_token(self, token_endpoint: str, grant: dict[str, str]) -> str: + from tests.auth_idp.conftest import _exchange_token_with_retries + + return _exchange_token_with_retries(token_endpoint, grant, verify=self.verify) + + def _sdk_for_token(self, token: str) -> NeMoPlatform: + return NeMoPlatform( + base_url=self.gateway_base_url, + default_headers={"Authorization": f"Bearer {token}"}, + max_retries=0, + http_client=DefaultHttpxClient(verify=self.verify), + ) diff --git a/tests/auth_idp/static/test_authentik_blueprint.py b/tests/auth_idp/static/test_authentik_blueprint.py new file mode 100644 index 0000000000..b18eb1b855 --- /dev/null +++ b/tests/auth_idp/static/test_authentik_blueprint.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import pytest +import yaml +from yaml.nodes import MappingNode, Node, ScalarNode, SequenceNode + +pytestmark = [pytest.mark.auth_idp] + +BLUEPRINT = Path("contrib/auth/authentik/helm/files/blueprints/nemo.yaml") +LEGACY_BLUEPRINT = Path("contrib/auth/authentik/blueprints/nemo.yaml") + + +@dataclass(frozen=True) +class TaggedYamlValue: + tag: str + value: Any + + +class _BlueprintLoader(yaml.SafeLoader): + pass + + +def _construct_tagged_value(loader: _BlueprintLoader, tag_suffix: str, node: Node) -> TaggedYamlValue: + if isinstance(node, ScalarNode): + value = loader.construct_scalar(node) + elif isinstance(node, SequenceNode): + value = loader.construct_sequence(node, deep=True) + elif isinstance(node, MappingNode): + value = loader.construct_mapping(node, deep=True) + else: + raise AssertionError(f"Unsupported YAML node type for !{tag_suffix}: {type(node).__name__}") + + return TaggedYamlValue(f"!{tag_suffix}", value) + + +_BlueprintLoader.add_multi_constructor("!", _construct_tagged_value) + + +def _load_blueprint(path: Path = BLUEPRINT) -> Mapping[str, Any]: + blueprint = yaml.load(path.read_text(encoding="utf-8"), Loader=_BlueprintLoader) + assert isinstance(blueprint, dict) + return blueprint + + +def _entries(blueprint: Mapping[str, Any]) -> list[Mapping[str, Any]]: + entries = blueprint["entries"] + assert isinstance(entries, list) + assert all(isinstance(entry, dict) for entry in entries) + return entries + + +def _entry_by_id(blueprint: Mapping[str, Any], entry_id: str) -> Mapping[str, Any]: + matches = [entry for entry in _entries(blueprint) if entry.get("id") == entry_id] + assert len(matches) == 1 + return matches[0] + + +def _entry_by_identifier( + blueprint: Mapping[str, Any], model: str, identifier_name: str, identifier_value: str +) -> Mapping[str, Any]: + matches = [ + entry + for entry in _entries(blueprint) + if entry.get("model") == model + and isinstance(entry.get("identifiers"), dict) + and entry["identifiers"].get(identifier_name) == identifier_value + ] + assert len(matches) == 1 + return matches[0] + + +def _attrs(entry: Mapping[str, Any]) -> Mapping[str, Any]: + attrs = entry["attrs"] + assert isinstance(attrs, dict) + return attrs + + +def test_static_authentik_blueprint_declares_workload_provider_defaults() -> None: + blueprint = _load_blueprint() + + metadata = blueprint["metadata"] + assert isinstance(metadata, dict) + labels = metadata["labels"] + assert isinstance(labels, dict) + assert labels["blueprints.goauthentik.io/instantiate"] == "true" + + workload_provider = _entry_by_id(blueprint, "nemo-workload-provider") + assert workload_provider["model"] == "authentik_providers_oauth2.oauth2provider" + assert workload_provider["identifiers"] == {"name": "nemo-platform-workload"} + workload_provider_attrs = _attrs(workload_provider) + + assert workload_provider_attrs["name"] == "nemo-platform-workload" + assert workload_provider_attrs["client_type"] == "public" + assert workload_provider_attrs["client_id"] == "nemo-platform-workload" + assert workload_provider_attrs["access_token_validity"] == "minutes=5" + + cli_provider = _entry_by_id(blueprint, "nemo-cli-provider") + assert _attrs(cli_provider)["access_token_validity"] == "minutes=2" + assert 'Use a longer value such as "hours=1"' in BLUEPRINT.read_text(encoding="utf-8") + + workload_application = _entry_by_identifier(blueprint, "authentik_core.application", "slug", "nemo-workload") + assert workload_application["identifiers"] == {"slug": "nemo-workload"} + workload_application_attrs = _attrs(workload_application) + assert workload_application_attrs["name"] == "NeMo Platform Workload Identity" + assert workload_application_attrs["slug"] == "nemo-workload" + assert workload_application_attrs["provider"] == TaggedYamlValue("!KeyOf", "nemo-workload-provider") + + +def test_authentik_blueprint_keeps_human_and_workload_groups_separate() -> None: + blueprint = _load_blueprint() + + editors_group = _entry_by_id(blueprint, "group-nemo-editors") + workloads_group = _entry_by_id(blueprint, "group-nemo-workloads") + human_user = _entry_by_id(blueprint, "nemo-user") + workload_user = _entry_by_id(blueprint, "svc-nemo") + + assert _attrs(editors_group)["name"] == "nemo-editors" + assert _attrs(workloads_group)["name"] == "nemo-workloads" + assert _attrs(human_user)["groups"] == [TaggedYamlValue("!KeyOf", "group-nemo-editors")] + assert _attrs(workload_user)["groups"] == [TaggedYamlValue("!KeyOf", "group-nemo-workloads")] + + +def test_authentik_blueprint_reads_workload_identity_password_from_env() -> None: + blueprint = _load_blueprint() + legacy_secret = "svc-nemo" + "-token-secret-dev" + token_identifiers = [ + entry["identifiers"]["identifier"] + for entry in _entries(blueprint) + if entry.get("model") == "authentik_core.token" + and isinstance(entry.get("identifiers"), dict) + and "identifier" in entry["identifiers"] + ] + token_keys = [ + _attrs(entry).get("key") for entry in _entries(blueprint) if entry.get("model") == "authentik_core.token" + ] + + workload_token = _entry_by_identifier(blueprint, "authentik_core.token", "identifier", "svc-nemo-token") + workload_token_attrs = _attrs(workload_token) + + assert workload_token_attrs["intent"] == "app_password" + assert workload_token_attrs["user"] == TaggedYamlValue("!KeyOf", "svc-nemo") + assert workload_token_attrs["key"] == TaggedYamlValue("!Env", "AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD") + assert "nemo-user-token" not in token_identifiers + assert legacy_secret not in token_keys + + +def test_authentik_blueprint_declares_e2e_setup_identity_as_test_only() -> None: + blueprint_text = BLUEPRINT.read_text(encoding="utf-8") + blueprint = _load_blueprint() + + assert "E2E TEST HARNESS ONLY" in blueprint_text + + setup_user = _entry_by_id(blueprint, "nemo-setup") + setup_user_attrs = _attrs(setup_user) + assert setup_user["model"] == "authentik_core.user" + assert setup_user_attrs["type"] == "service_account" + assert setup_user_attrs["groups"] == [TaggedYamlValue("!KeyOf", "group-nemo-admins")] + + setup_token = _entry_by_identifier(blueprint, "authentik_core.token", "identifier", "nemo-setup-token") + setup_token_attrs = _attrs(setup_token) + assert setup_token_attrs["intent"] == "app_password" + assert setup_token_attrs["user"] == TaggedYamlValue("!KeyOf", "nemo-setup") + assert setup_token_attrs["key"] == "nemo-setup-token-secret-dev" + + +def test_authentik_blueprint_has_single_canonical_source() -> None: + assert BLUEPRINT.exists() + assert not LEGACY_BLUEPRINT.exists() diff --git a/tests/auth_idp/static/test_authentik_device_flow_helper.py b/tests/auth_idp/static/test_authentik_device_flow_helper.py new file mode 100644 index 0000000000..e4585f554f --- /dev/null +++ b/tests/auth_idp/static/test_authentik_device_flow_helper.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from tests.auth_idp.device_flow import authentik_flow_executor_url, url_origin, with_url_origin + +pytestmark = [pytest.mark.auth_idp] + + +def test_authentik_flow_executor_url_preserves_default_login_next_query() -> None: + assert authentik_flow_executor_url( + "https://127.0.0.1:38080", + "/flows/-/default/authentication/?next=/device%3Fcode%3D123456789", + ) == ( + "https://127.0.0.1:38080/api/v3/flows/executor/default-authentication-flow/" + "?query=next%3D%2Fdevice%253Fcode%253D123456789" + ) + + +def test_authentik_flow_executor_url_preserves_if_flow_query() -> None: + assert authentik_flow_executor_url( + "https://127.0.0.1:38080", + "/if/flow/default-provider-authorization-implicit-consent/?code=abc&state=xyz", + ) == ( + "https://127.0.0.1:38080/api/v3/flows/executor/default-provider-authorization-implicit-consent/" + "?query=code%3Dabc%26state%3Dxyz" + ) + + +def test_authentik_flow_executor_url_ignores_non_flow_browser_urls() -> None: + assert authentik_flow_executor_url("https://127.0.0.1:38080", "/device?code=123456789") is None + + +def test_url_origin_keeps_runtime_port_forward_separate_from_advertised_device_origin() -> None: + assert url_origin("https://127.0.0.1:18080/application/o/device/") == "https://127.0.0.1:18080" + + +def test_with_url_origin_preserves_advertised_oidc_path_for_runtime_port_forward() -> None: + assert ( + with_url_origin( + "https://127.0.0.1:18080/application/o/device/?x=1", + "https://127.0.0.1:65490", + ) + == "https://127.0.0.1:65490/application/o/device/?x=1" + ) diff --git a/tests/auth_idp/static/test_authentik_kubernetes_demo.py b/tests/auth_idp/static/test_authentik_kubernetes_demo.py new file mode 100644 index 0000000000..b1d0502e30 --- /dev/null +++ b/tests/auth_idp/static/test_authentik_kubernetes_demo.py @@ -0,0 +1,1193 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import ast +import importlib.util +import os +import shutil +import subprocess +import sys +from pathlib import Path +from types import ModuleType + +import pytest +import yaml +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR + +pytestmark = [pytest.mark.auth_idp] + +AUTHENTIK_DIR = Path("contrib/auth/authentik") +HELM_DIR = AUTHENTIK_DIR / "helm" +AUTHENTIK_SCRIPT_TIMEOUT_SECONDS = 30 +HELM_TEMPLATE_TIMEOUT_SECONDS = 60 +ENVOY_SERVICE_URL_TEMPLATE = ( + '{{ include "nemo-platform-authentik.serviceUrl" ' + '(dict "root" . "serviceName" "nemo-platform-envoy" ' + '"namespace" .Values.envoyProxy.serviceNamespace "scheme" "https" "port" 8080) }}' +) +ENVOY_CONTROLLER_ENV_URL = "https://nemo-platform-envoy.$(POD_NAMESPACE).svc.cluster.local:8080" +AUTHENTIK_SERVICE_URL_TEMPLATE = '{{ include "nemo-platform-authentik.serviceUrl" (dict "root" . "serviceName" "authentik-server" "scheme" "http") }}' +PUBLIC_GATEWAY_URL_TEMPLATE = '{{ include "nemo-platform-authentik.publicGatewayUrl" . }}' + + +def _load_yaml(path: Path) -> dict: + return yaml.safe_load(path.read_text(encoding="utf-8")) + + +def _literal_run_commands(path: Path) -> set[tuple[str, ...]]: + tree = ast.parse(path.read_text(encoding="utf-8")) + commands: set[tuple[str, ...]] = set() + + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + if not isinstance(node.func, ast.Name) or node.func.id != "_run": + continue + if not node.args: + continue + + try: + command = ast.literal_eval(node.args[0]) + except (SyntaxError, ValueError): + continue + + if isinstance(command, list) and all(isinstance(part, str) for part in command): + commands.add(tuple(command)) + + return commands + + +def _workflow_job_block(workflow: str, job_name: str) -> str: + lines = workflow.splitlines() + start = next(index for index, line in enumerate(lines) if line == f" {job_name}:") + end = len(lines) + for index in range(start + 1, len(lines)): + line = lines[index] + if line.startswith(" ") and not line.startswith(" ") and line.endswith(":"): + end = index + break + return "\n".join(lines[start:end]) + + +def _load_authentik_k8s_live_module() -> ModuleType: + module_path = Path("tests/auth_idp/runtime_kubernetes.py") + spec = importlib.util.spec_from_file_location("authentik_k8s_live_for_unit", module_path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _run_authentik_script(*args: str, env: dict[str, str] | None = None) -> str: + process_env = os.environ.copy() + if env: + process_env.update(env) + completed = subprocess.run( + [str(AUTHENTIK_DIR / "run.sh"), *args], + text=True, + capture_output=True, + check=False, + env=process_env, + timeout=AUTHENTIK_SCRIPT_TIMEOUT_SECONDS, + ) + assert completed.returncode == 0, completed.stdout + completed.stderr + return completed.stdout + + +def test_authentik_run_local_defaults_workload_identity_password_for_compose() -> None: + env = os.environ.copy() + env.pop("AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD", None) + completed = subprocess.run( + [str(AUTHENTIK_DIR / "run.sh"), "run-local", "--dry-run"], + text=True, + capture_output=True, + check=False, + env=env, + timeout=AUTHENTIK_SCRIPT_TIMEOUT_SECONDS, + ) + + assert completed.returncode == 0, completed.stdout + completed.stderr + assert "helm/files/blueprints/nemo.yaml" in completed.stdout + assert "contrib/auth/authentik/.generated/workload-token-private-key.pem" in completed.stdout + assert "contrib/auth/authentik/.generated/gateway-tls" in completed.stdout + assert "contrib/auth/authentik/compose" in completed.stdout + assert "AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD=" in completed.stdout + assert "docker compose up" in completed.stdout + + +def test_authentik_prepare_local_creates_shared_generated_inputs() -> None: + output = _run_authentik_script("prepare-local", "--dry-run") + + assert "helm/files/blueprints/nemo.yaml" in output + assert "contrib/auth/authentik/.generated/workload-token-private-key.pem" in output + assert "contrib/auth/authentik/.generated/gateway-tls" in output + assert "docker compose up" not in output + + +def test_authentik_user_startup_docs_use_manual_runtime_steps() -> None: + top_level_readme = (AUTHENTIK_DIR / "README.md").read_text(encoding="utf-8") + tutorial = (AUTHENTIK_DIR / "tutorial.md").read_text(encoding="utf-8") + compose_readme = (AUTHENTIK_DIR / "compose" / "README.md").read_text(encoding="utf-8") + compose_details = (AUTHENTIK_DIR / "compose" / "implementation-details.md").read_text(encoding="utf-8") + kubernetes_readme = (AUTHENTIK_DIR / "kubernetes" / "README.md").read_text(encoding="utf-8") + kubernetes_details = (AUTHENTIK_DIR / "kubernetes" / "implementation-details.md").read_text(encoding="utf-8") + + assert "run.sh --help" in top_level_readme + assert "(tutorial.md)" in top_level_readme + assert "(compose/implementation-details.md)" in top_level_readme + assert "(kubernetes/implementation-details.md)" in top_level_readme + assert "### Docker Compose" in tutorial + assert "### Kubernetes" in tutorial + assert "## Wait For The Gateway" in tutorial + assert "${AUTHENTIK_BASE_URL}/health/gateway/ready" in tutorial + assert "NeMo Platform and Authentik Ready" in tutorial + assert "uv run nemo auth login \\" in tutorial + assert '--context "$AUTHENTIK_CONTEXT"' in tutorial + assert '--base-url "$AUTHENTIK_BASE_URL"' in tutorial + assert '--principal "$AUTHENTIK_WORKLOAD_GROUP"' in tutorial + assert "contrib/auth/authentik/run.sh prepare-local" in tutorial + assert "cd contrib/auth/authentik/compose" not in tutorial + assert "docker compose -f contrib/auth/authentik/compose/docker-compose.yml up" in tutorial + assert "docker compose -f contrib/auth/authentik/compose/docker-compose.yml down -v" in tutorial + assert ( + "--set-file workloadTokenSigningKey.privateKeyPem=" + "contrib/auth/authentik/.generated/workload-token-private-key.pem" + ) in tutorial + assert "contrib/auth/authentik/run.sh run-local" not in tutorial + assert "contrib/auth/authentik/run.sh compose" not in tutorial + assert "contrib/auth/authentik/run.sh k8s" not in tutorial + assert "run.sh" not in compose_readme + assert "run.sh" not in kubernetes_readme + assert "docker compose up" in compose_readme + assert "nemo-platform-authentik" in compose_readme + assert "helm --kube-context" in kubernetes_readme + assert "(implementation-details.md)" in compose_readme + assert "(implementation-details.md)" in kubernetes_readme + assert "For the step-by-step test flow, see the" in compose_details + assert "[shared tutorial](../tutorial.md)" in compose_details + assert "COMPOSE_PROJECT_NAME" in compose_details + assert "For the step-by-step test flow, see the [shared tutorial](../tutorial.md)." in kubernetes_details + + +def test_authentik_tutorial_grants_workloads_job_log_permissions() -> None: + tutorial = (AUTHENTIK_DIR / "tutorial.md").read_text(encoding="utf-8") + + editor_group_block = """ +uv run nemo --context "$AUTHENTIK_CONTEXT" workspaces members create \\ + --workspace "$WORKSPACE" \\ + --principal nemo-editors \\ + --roles Viewer \\ + --wait-role-propagation +""" + service_account_group_block = """ +uv run nemo --context "$AUTHENTIK_CONTEXT" workspaces members create \\ + --workspace "$WORKSPACE" \\ + --principal "$AUTHENTIK_WORKLOAD_GROUP" \\ + --roles Viewer \\ + --roles JobRunner \\ + --wait-role-propagation +""" + + assert editor_group_block in tutorial + assert service_account_group_block in tutorial + assert "dedicated `nemo-workloads` Authentik group" in tutorial + assert "permission to upload workload" in tutorial + + +def test_authentik_e2e_ci_requires_published_nmp_api_image() -> None: + ci_workflow = Path(".github/workflows/ci.yaml").read_text(encoding="utf-8") + job = _workflow_job_block(ci_workflow, "python-auth-idp-e2e-test") + + assert "needs.policy-wasm.result == 'success'" in job + assert "needs.build-cpu-smoke-images.result == 'success'" in job + assert "needs.build-cpu-smoke-images.outputs.publish_images == 'true'" in job + assert "needs.python-auth-idp-static-test.result == 'success'" in job + assert '--image "${IMAGE_REGISTRY}/nmp-api:${BAKE_TAG}"' in job + + +def _helm_template_authentik_demo(template: str) -> str: + if shutil.which("helm") is None: + pytest.skip("helm is required to render the Authentik Kubernetes demo chart") + + completed = subprocess.run( + [ + "helm", + "template", + "authentik-demo", + str(AUTHENTIK_DIR / "helm"), + "-n", + "nemo-authentik", + "--show-only", + template, + ], + text=True, + capture_output=True, + check=False, + timeout=HELM_TEMPLATE_TIMEOUT_SECONDS, + ) + assert completed.returncode == 0, completed.stderr + return completed.stdout + + +def _load_rendered_authentik_envoy_config() -> dict: + rendered = _helm_template_authentik_demo("charts/nemo-platform/templates/proxy/envoy-configmap.yaml") + config_map = next( + document + for document in yaml.safe_load_all(rendered) + if document and document["kind"] == "ConfigMap" and document["metadata"]["name"] == "nemo-platform-envoy" + ) + return yaml.safe_load(config_map["data"]["envoy.yaml"]) + + +def test_authentik_umbrella_chart_declares_expected_dependencies() -> None: + chart = _load_yaml(HELM_DIR / "Chart.yaml") + dependencies = {dependency["name"]: dependency for dependency in chart["dependencies"]} + + assert chart["name"] == "nemo-platform-authentik" + assert "cert-manager" not in dependencies + assert dependencies["authentik"] == { + "name": "authentik", + "version": "2026.5.4", + "repository": "https://charts.goauthentik.io", + } + assert "postgresql" not in dependencies + assert dependencies["nemo-platform"] == { + "name": "nemo-platform", + "version": "0.1.0", + "repository": "file://../../../../k8s/helm", + } + + +def test_authentik_umbrella_values_use_latest_authentik_chart_without_image_tag_override() -> None: + values = _load_yaml(HELM_DIR / "values.yaml") + authentik_values = values["authentik"] + + assert authentik_values["fullnameOverride"] == "authentik" + assert authentik_values["postgresql"]["enabled"] is False + assert authentik_values["blueprints"]["configMaps"] == ["authentik-nemo-blueprint"] + assert { + "name": "AUTHENTIK_POSTGRESQL__PASSWORD", + "valueFrom": { + "secretKeyRef": { + "name": "shared-postgresql", + "key": "authentik-password", + } + }, + } in authentik_values["global"]["env"] + assert authentik_values["authentik"]["postgresql"] == { + "host": "shared-postgresql", + "name": "authentik", + "user": "authentik", + "port": 5432, + } + assert authentik_values.get("global", {}).get("image", {}).get("tag", "") == "" + + +def test_authentik_umbrella_values_define_one_shared_postgresql_instance() -> None: + values = _load_yaml(HELM_DIR / "values.yaml") + initdb_template = (HELM_DIR / "templates" / "shared-postgres-initdb-configmap.yaml").read_text(encoding="utf-8") + secret_template = (HELM_DIR / "templates" / "shared-postgres-secret.yaml").read_text(encoding="utf-8") + nemo_secret_template = (HELM_DIR / "templates" / "shared-postgres-nemo-secret.yaml").read_text(encoding="utf-8") + helpers_template = (HELM_DIR / "templates" / "_helpers.tpl").read_text(encoding="utf-8") + + assert values["sharedPostgresql"]["enabled"] is True + assert values["sharedPostgresql"]["serviceName"] == "shared-postgresql" + assert "password" not in values["sharedPostgresql"]["authentik"] + assert "cert-manager" not in values + assert "shared-postgresql" not in values + assert 'define "nemo-platform-authentik.sharedPostgresql.password"' in helpers_template + assert 'define "nemo-platform-authentik.existingSecretData"' in helpers_template + assert 'include "nemo-platform-authentik.existingSecretData"' in helpers_template + assert 'lookup "v1" "PersistentVolumeClaim" $root.Release.Namespace $pvcName' in helpers_template + assert "restore the Secret or rotate the PostgreSQL role before changing it" in helpers_template + assert secret_template.count('include "nemo-platform-authentik.sharedPostgresql.password"') == 3 + assert '"secretKey" "authentik-password" "value" .Values.sharedPostgresql.authentik.password "generate" true' in ( + secret_template + ) + assert "authentik-password: {{ $authentikPassword | quote }}" in secret_template + assert nemo_secret_template.startswith("{{- if .Values.sharedPostgresql.enabled }}\n{{- $nemoPassword :=") + assert nemo_secret_template.rstrip().endswith("{{- end }}") + assert 'include "nemo-platform-authentik.sharedPostgresql.password"' in nemo_secret_template + assert "--set=" not in initdb_template + assert "<<'EOSQL'" in initdb_template + assert "\\set authentik_password `printf '%s' \"$AUTHENTIK_PASSWORD\"`" in initdb_template + assert "\\set nemo_password `printf '%s' \"$NEMO_PASSWORD\"`" in initdb_template + assert "CREATE USER :\"authentik_username\" WITH PASSWORD :'authentik_password';" in initdb_template + assert "CREATE USER :\"nemo_username\" WITH PASSWORD :'nemo_password';" in initdb_template + + nemo_database = values["nemo-platform"]["externalDatabase"] + assert values["nemo-platform"]["postgresql"]["enabled"] is False + assert nemo_database == { + "host": "shared-postgresql", + "port": 5432, + "user": "nemo", + "database": "nemoplatform", + "existingSecret": "shared-postgresql-nemo", + "existingSecretPasswordKey": "password", + } + + +def test_authentik_kubernetes_helm_args_can_reuse_precreated_ngc_secret(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("NMP_AUTHENTIK_K8S_NGC_EXISTING_SECRET", "ngc-api") + + live_test = _load_authentik_k8s_live_module() + + args = live_test._helm_upgrade_args("kind-ci") + assert "nemo-platform.existingSecret=ngc-api" in args + + +def test_authentik_kubernetes_helm_args_can_override_public_gateway_port( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("NMP_AUTHENTIK_K8S_GATEWAY_PORT", "18082") + + live_test = _load_authentik_k8s_live_module() + + args = live_test._helm_upgrade_args("kind-ci") + assert "nemo-platform.authentikPublicGateway.port=18082" in args + + +def test_authentik_kubernetes_live_timeouts_are_named_constants() -> None: + live_test = _load_authentik_k8s_live_module() + + args = live_test._helm_upgrade_args("kind-ci") + assert args[args.index("--timeout") + 1] == live_test.HELM_WAIT_TIMEOUT + assert live_test.HELM_WAIT_TIMEOUT == "10m" + assert live_test.HELM_UPGRADE_COMMAND_TIMEOUT_SECONDS == 900 + assert live_test.PORT_FORWARD_READY_TIMEOUT_SECONDS == 30 + + +def test_authentik_kubernetes_reuse_context_validates_runtime() -> None: + live_test = _load_authentik_k8s_live_module() + + assert live_test._reuse_context("kind", "ci") == "kind-ci" + assert live_test._reuse_context("k3d", "ci") == "k3d-ci" + with pytest.raises(ValueError, match="unsupported NMP_AUTHENTIK_K8S_RUNTIME='minikube'; expected kind or k3d"): + live_test._reuse_context("minikube", "ci") + + +def test_authentik_kubernetes_reuse_or_create_uses_existing_cluster(monkeypatch: pytest.MonkeyPatch) -> None: + live_test = _load_authentik_k8s_live_module() + existing = live_test.Cluster(name="ci", runtime="kind", context="kind-ci") + + monkeypatch.setattr(live_test, "_existing_cluster", lambda runtime, name: existing) + monkeypatch.setattr( + live_test, + "_create_cluster", + lambda runtime, name: pytest.fail("reuse should not create an existing cluster"), + ) + + assert live_test._reuse_or_create_cluster("kind", "ci") == existing + + +def test_authentik_kubernetes_reuse_or_create_creates_missing_cluster(monkeypatch: pytest.MonkeyPatch) -> None: + live_test = _load_authentik_k8s_live_module() + created = live_test.Cluster(name="ci", runtime="kind", context="kind-ci") + + monkeypatch.setattr(live_test, "_existing_cluster", lambda runtime, name: None) + monkeypatch.setattr(live_test, "_create_cluster", lambda runtime, name: created) + + assert live_test._reuse_or_create_cluster("kind", "ci") == created + + +def test_authentik_kubernetes_kind_create_uses_isolated_kubeconfig( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + live_test = _load_authentik_k8s_live_module() + kubeconfig = tmp_path / "kind-kubeconfig.yaml" + commands: list[list[str]] = [] + + monkeypatch.setattr(live_test, "_temporary_kubeconfig_path", lambda cluster_name: kubeconfig) + monkeypatch.setattr(live_test, "_require_tool", lambda name: None) + + def record_command(args: list[str], **_kwargs: object) -> subprocess.CompletedProcess[str]: + commands.append(args) + return subprocess.CompletedProcess(args, 0, stdout="", stderr="") + + monkeypatch.setattr(live_test, "_run", record_command) + + cluster = live_test._create_cluster("kind", "ci") + + assert cluster == live_test.Cluster( + name="ci", + runtime="kind", + context="kind-ci", + kubeconfig=kubeconfig, + cleanup_kubeconfig=True, + ) + assert commands == [ + [ + "kind", + "create", + "cluster", + "--name", + "ci", + "--kubeconfig", + str(kubeconfig), + "--wait", + live_test.KIND_CREATE_WAIT_TIMEOUT, + ] + ] + + +def test_authentik_kubernetes_k3d_create_does_not_update_default_kubeconfig( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + live_test = _load_authentik_k8s_live_module() + kubeconfig = tmp_path / "k3d-kubeconfig.yaml" + commands: list[list[str]] = [] + + monkeypatch.setattr(live_test, "_temporary_kubeconfig_path", lambda cluster_name: kubeconfig) + monkeypatch.setattr(live_test, "_require_tool", lambda name: None) + + def record_command(args: list[str], **_kwargs: object) -> subprocess.CompletedProcess[str]: + commands.append(args) + if args == ["k3d", "kubeconfig", "get", "ci"]: + return subprocess.CompletedProcess(args, 0, stdout="apiVersion: v1\n", stderr="") + return subprocess.CompletedProcess(args, 0, stdout="", stderr="") + + monkeypatch.setattr(live_test, "_run", record_command) + + cluster = live_test._create_cluster("k3d", "ci") + + assert cluster == live_test.Cluster( + name="ci", + runtime="k3d", + context="k3d-ci", + kubeconfig=kubeconfig, + cleanup_kubeconfig=True, + ) + assert commands[0] == [ + "k3d", + "cluster", + "create", + "ci", + "--wait", + "--agents", + "0", + "--k3s-arg", + "--disable=traefik@server:0", + "--kubeconfig-update-default=false", + ] + assert commands[1] == ["k3d", "kubeconfig", "get", "ci"] + assert kubeconfig.read_text(encoding="utf-8") == "apiVersion: v1\n" + + +def test_authentik_kubernetes_commands_accept_isolated_kubeconfig() -> None: + live_test = _load_authentik_k8s_live_module() + kubeconfig = Path("/tmp/nmp-authentik-kubeconfig.yaml") + + assert live_test._kubectl_command("kind-ci", ["get", "pods"], kubeconfig) == [ + "kubectl", + "--kubeconfig", + str(kubeconfig), + "--context", + "kind-ci", + "get", + "pods", + ] + assert live_test._helm_upgrade_args("kind-ci", kubeconfig)[:5] == [ + "helm", + "--kubeconfig", + str(kubeconfig), + "--kube-context", + "kind-ci", + ] + + +def test_authentik_kubernetes_port_forward_times_out_without_readiness(monkeypatch: pytest.MonkeyPatch) -> None: + live_test = _load_authentik_k8s_live_module() + + class FakeProcess: + returncode = None + + def poll(self) -> None: + return None + + def terminate(self) -> None: + return None + + def wait(self, timeout: int | None = None) -> None: + return None + + def kill(self) -> None: + return None + + monotonic_values = iter([0.0, 0.0, 31.0]) + monkeypatch.setattr(live_test, "_free_port", lambda: 19001) + monkeypatch.setattr(live_test.subprocess, "Popen", lambda *args, **kwargs: FakeProcess()) + monkeypatch.setattr(live_test.time, "monotonic", lambda: next(monotonic_values)) + monkeypatch.setattr(live_test.time, "sleep", lambda seconds: None) + requested_urls = [] + + def fake_get(url: str, **_kwargs: object): + requested_urls.append(url) + raise live_test.httpx.ConnectError("not ready") + + monkeypatch.setattr(live_test.httpx, "get", fake_get) + + with pytest.raises(TimeoutError, match="timed out waiting for port-forward readiness"): + live_test._start_port_forward_service("kind-ci", "nemo-platform-envoy", Path("ca.crt")) + + assert requested_urls == ["https://127.0.0.1:19001/health/gateway/ready"] + + +def test_authentik_kubernetes_port_forward_waits_after_kill(monkeypatch: pytest.MonkeyPatch) -> None: + live_test = _load_authentik_k8s_live_module() + events: list[object] = [] + + class FakeProcess: + returncode = None + + def poll(self) -> int | None: + return self.returncode + + def terminate(self) -> None: + events.append("terminate") + + def wait(self, timeout: int | None = None) -> int: + events.append(("wait", timeout)) + if timeout is not None: + raise live_test.subprocess.TimeoutExpired(cmd="kubectl", timeout=timeout) + self.returncode = -9 + return self.returncode + + def kill(self) -> None: + events.append("kill") + + monotonic_values = iter([0.0, 0.0, 31.0]) + monkeypatch.setattr(live_test, "_free_port", lambda: 19001) + monkeypatch.setattr(live_test.subprocess, "Popen", lambda *args, **kwargs: FakeProcess()) + monkeypatch.setattr(live_test.time, "monotonic", lambda: next(monotonic_values)) + monkeypatch.setattr(live_test.time, "sleep", lambda seconds: None) + monkeypatch.setattr( + live_test.httpx, + "get", + lambda *args, **kwargs: (_ for _ in ()).throw(live_test.httpx.ConnectError("not ready")), + ) + + with pytest.raises(TimeoutError, match="timed out waiting for port-forward readiness"): + live_test._start_port_forward_service("kind-ci", "nemo-platform-envoy", Path("ca.crt")) + + assert events == [ + "terminate", + ("wait", live_test.PORT_FORWARD_TERMINATE_TIMEOUT_SECONDS), + "kill", + ("wait", None), + ] + + +def test_authentik_kubernetes_diagnostics_records_pod_discovery_timeout( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + live_test = _load_authentik_k8s_live_module() + monkeypatch.setenv("NMP_AUTHENTIK_K8S_LOG_DIR", str(tmp_path)) + monkeypatch.setattr(live_test, "_write_diagnostic_process", lambda *args, **kwargs: None) + + def raise_timeout(args: list[str], **kwargs: object) -> None: + raise live_test.subprocess.TimeoutExpired( + cmd=args, + timeout=kwargs["timeout"], + output="partial pod list\n", + stderr="kubectl timed out\n", + ) + + monkeypatch.setattr(live_test.subprocess, "run", raise_timeout) + + log_dir = live_test._collect_kubernetes_diagnostics("kind-ci", "ci") + + assert log_dir == tmp_path / "k8s-authentik-ci" + pod_discovery = log_dir / "get-pods-for-logs.txt" + assert pod_discovery.read_text(encoding="utf-8") == "\n".join( + ( + "command: kubectl --context kind-ci -n nemo-authentik get pods -o " + "jsonpath={range .items[*]}{.metadata.name}{'\\n'}{end}", + "timeout: 60", + "stdout:", + "partial pod list\n", + "stderr:", + "kubectl timed out\n", + ) + ) + assert not list(log_dir.glob("logs-*.txt")) + + +def test_authentik_kubernetes_cleanup_deletes_cluster_when_diagnostics_fail( + monkeypatch: pytest.MonkeyPatch, +) -> None: + live_test = _load_authentik_k8s_live_module() + runtime = live_test.KubernetesAuthIdpRuntime.__new__(live_test.KubernetesAuthIdpRuntime) + runtime.cluster = live_test.Cluster(name="ci", runtime="kind", context="kind-ci") + runtime._port_forward_process = None + runtime._ca_temp_file = None + runtime._previous_client_ssl_cert_file = os.environ.get(live_test.NMP_CLIENT_SSL_CERT_FILE_ENVVAR) + runtime._diagnostics_collected = False + runtime._reuse_cluster = False + runtime._keep_cluster = False + deleted: list[tuple[str, str]] = [] + + def raise_diagnostics(context: str, cluster_name: str, kubeconfig: Path | None = None) -> Path: + raise RuntimeError(f"diagnostics failed for {context}/{cluster_name}") + + def delete_cluster(runtime_name: str, cluster_name: str, kubeconfig: Path | None = None) -> None: + deleted.append((runtime_name, cluster_name)) + + monkeypatch.setattr(live_test, "_collect_kubernetes_diagnostics", raise_diagnostics) + monkeypatch.setattr(live_test, "_delete_cluster", delete_cluster) + + runtime.cleanup() + + assert deleted == [("kind", "ci")] + assert runtime.cluster is None + assert runtime._diagnostics_collected is True + + +def test_authentik_kubernetes_startup_preserves_original_error_when_diagnostics_fail( + monkeypatch: pytest.MonkeyPatch, +) -> None: + live_test = _load_authentik_k8s_live_module() + runtime = live_test.KubernetesAuthIdpRuntime.__new__(live_test.KubernetesAuthIdpRuntime) + runtime.cluster = None + runtime.ca_bundle = None + runtime._port_forward_process = None + runtime._ca_temp_file = None + runtime._previous_client_ssl_cert_file = os.environ.get(live_test.NMP_CLIENT_SSL_CERT_FILE_ENVVAR) + runtime._diagnostics_collected = False + runtime._reuse_cluster = False + runtime._keep_cluster = False + diagnostics_calls: list[tuple[str, str]] = [] + deleted: list[tuple[str, str]] = [] + + def raise_startup_error(runtime_name: str, cluster_name: str, image: str) -> None: + raise RuntimeError(f"startup failed for {runtime_name}/{cluster_name}/{image}") + + def raise_diagnostics(context: str, cluster_name: str, kubeconfig: Path | None = None) -> Path: + diagnostics_calls.append((context, cluster_name)) + raise RuntimeError(f"diagnostics failed for {context}/{cluster_name}") + + def delete_cluster(runtime_name: str, cluster_name: str, kubeconfig: Path | None = None) -> None: + deleted.append((runtime_name, cluster_name)) + + monkeypatch.setenv("NMP_AUTHENTIK_K8S_RUNTIME", "kind") + monkeypatch.setenv("NMP_AUTHENTIK_K8S_CLUSTER_NAME", "ci") + monkeypatch.delenv("NMP_AUTHENTIK_K8S_REUSE_CLUSTER", raising=False) + monkeypatch.delenv("NMP_AUTHENTIK_K8S_KEEP_CLUSTER", raising=False) + monkeypatch.setattr( + live_test, + "_create_cluster", + lambda runtime_name, cluster_name: live_test.Cluster( + name=cluster_name, + runtime=runtime_name, + context="kind-ci", + kubeconfig=Path("isolated-kubeconfig.yaml"), + cleanup_kubeconfig=True, + ), + ) + monkeypatch.setattr(live_test, "_platform_image", lambda: "nmp:test") + monkeypatch.setattr(live_test, "_load_platform_image", raise_startup_error) + monkeypatch.setattr(live_test, "_collect_kubernetes_diagnostics", raise_diagnostics) + monkeypatch.setattr(live_test, "_delete_cluster", delete_cluster) + + with pytest.raises(RuntimeError, match="startup failed"): + runtime._start() + + assert diagnostics_calls == [("kind-ci", "ci")] + assert deleted == [("kind", "ci")] + assert runtime.cluster is None + assert runtime._diagnostics_collected is True + + +@pytest.mark.auth_idp_k8s +def test_authentik_umbrella_values_configure_nemo_envoy_as_the_only_edge_proxy() -> None: + values = _load_yaml(HELM_DIR / "values.yaml") + helpers_template = (HELM_DIR / "templates" / "_helpers.tpl").read_text(encoding="utf-8") + nemo_values = values["nemo-platform"] + envoy = nemo_values["envoyProxy"] + envoy_config = _load_rendered_authentik_envoy_config() + http_manager = envoy_config["static_resources"]["listeners"][0]["filter_chains"][0]["filters"][0]["typed_config"] + routes = http_manager["route_config"]["virtual_hosts"][0]["routes"] + forwarded_proto_header = [ + { + "header": {"key": "x-forwarded-proto", "value": "https"}, + "append_action": "OVERWRITE_IF_EXISTS_OR_ADD", + } + ] + jwt_filter = next( + filter_config + for filter_config in http_manager["http_filters"] + if filter_config["name"] == "envoy.filters.http.jwt_authn" + ) + jwt_providers = jwt_filter["typed_config"]["providers"] + clusters = {cluster["name"]: cluster for cluster in envoy_config["static_resources"]["clusters"]} + + assert envoy["configOverride"] == '{{ include "nemo-platform-authentik.envoyConfig" . }}' + gateway_ready_route = next(route for route in routes if route["match"] == {"path": "/health/gateway/ready"}) + health_route = next(route for route in routes if route["match"] == {"prefix": "/health/"}) + assert routes.index(gateway_ready_route) < routes.index(health_route) + assert gateway_ready_route["direct_response"] == { + "status": 503, + "body": {"inline_string": '{"status":"not_ready"}'}, + } + for match in ( + {"prefix": "/.well-known/nemo-platform/"}, + {"prefix": "/apis/"}, + {"prefix": "/health/"}, + {"path": "/status"}, + {"prefix": "/studio/"}, + ): + route = next( + route for route in routes if route.get("route", {}).get("cluster") == "nemo" and route["match"] == match + ) + assert route["request_headers_to_add"] == forwarded_proto_header + lua_filter = next( + filter_config + for filter_config in http_manager["http_filters"] + if filter_config["name"] == "envoy.filters.http.lua" + ) + lua_code = lua_filter["typed_config"]["inline_code"] + assert 'headers:get(":path") ~= "/health/gateway/ready"' in lua_code + assert 'gateway_ready_http_call(request_handle, "nemo", "nemo-platform-api", "/health/ready")' in lua_code + assert ( + 'gateway_ready_http_call(request_handle, "authentik", "authentik-server", ' + '"/application/o/nemo/.well-known/openid-configuration")' + ) in lua_code + assert jwt_providers["authentik_workload"]["remote_jwks"]["http_uri"] == { + "uri": "https://nemo-platform-envoy:8080/application/o/nemo/jwks/", + "cluster": "nemo_envoy_https", + "timeout": "5s", + } + assert jwt_providers["workload_exchange"]["remote_jwks"]["http_uri"] == { + "uri": "https://nemo-platform-envoy:8080/apis/auth/jwks", + "cluster": "nemo_envoy_https", + "timeout": "5s", + } + envoy_jwks_cluster = clusters["nemo_envoy_https"] + assert envoy_jwks_cluster["transport_socket"]["typed_config"]["common_tls_context"]["validation_context"][ + "trusted_ca" + ] == {"filename": "/etc/nmp/workload-token-tls/ca.crt"} + + platform_config = nemo_values["platformConfig"].get("platform", {}) + assert "base_url" not in platform_config + assert "auth" not in platform_config.get("service_discovery", {}) + oidc = nemo_values["platformConfig"]["auth"]["oidc"] + assert oidc["issuer"] == f"{AUTHENTIK_SERVICE_URL_TEMPLATE}/application/o/nemo-cli/" + assert oidc["additional_issuers"][0] == f"{AUTHENTIK_SERVICE_URL_TEMPLATE}/application/o/nemo/" + assert oidc["additional_issuers"][1] == f"{PUBLIC_GATEWAY_URL_TEMPLATE}/application/o/nemo-cli/" + assert oidc["additional_issuers"][2] == f"{PUBLIC_GATEWAY_URL_TEMPLATE}/application/o/nemo/" + assert oidc["workload_token_issuer"] == f"{ENVOY_SERVICE_URL_TEMPLATE}/apis/auth" + assert oidc["workload_token_endpoint"] == f"{ENVOY_SERVICE_URL_TEMPLATE}/apis/auth/token" + assert oidc["token_endpoint"] == f"{PUBLIC_GATEWAY_URL_TEMPLATE}/application/o/token/" + assert oidc["device_authorization_endpoint"] == f"{PUBLIC_GATEWAY_URL_TEMPLATE}/application/o/device/" + assert nemo_values["authentikPublicGateway"] == { + "scheme": "https", + "host": "127.0.0.1", + "port": 18081, + } + assert "nemo-platform.authentikPublicGateway.host is required" in helpers_template + assert "nemo-platform.authentikPublicGateway.port is required" in helpers_template + assert 'default "127.0.0.1"' not in helpers_template + assert "default 18081" not in helpers_template + assert nemo_values["authentikEnvoy"] == {"serviceName": "authentik-server", "servicePort": 80} + assert nemo_values["envoyProxy"]["serviceNamespace"] == "" + assert values["integration"]["nemoPlatform"]["envoyServiceName"] == "nemo-platform-envoy" + assert nemo_values["rbac"]["volcanoEnabled"] is False + assert nemo_values["platformConfig"]["models"]["controller"]["backends"] == { + "none": {"enabled": True}, + "nim_operator": {"enabled": False}, + } + + +def test_nemo_platform_chart_does_not_use_volcano_disable_config() -> None: + values_template = Path("k8s/helm/values.yaml").read_text(encoding="utf-8") + + assert "enable_default_volcano_executor" not in values_template + + +def test_authentik_umbrella_values_mount_workload_token_signing_key_as_file() -> None: + values = _load_yaml(HELM_DIR / "values.yaml") + signing_key = values["workloadTokenSigningKey"] + nemo_values = values["nemo-platform"] + oidc = nemo_values["platformConfig"]["auth"]["oidc"] + + assert signing_key["secretName"] == "nemo-workload-token-signing-key" + assert signing_key["key"] == "private-key.pem" + assert signing_key["mountPath"] == "/etc/nmp/workload-token" + assert signing_key["privateKeyPem"] == "" + assert oidc["workload_token_private_key_file"] == "/etc/nmp/workload-token/private-key.pem" + + assert nemo_values["api"]["extraVolumes"] == [ + { + "name": "workload-token-signing-key", + "secret": {"secretName": "nemo-workload-token-signing-key"}, + } + ] + assert nemo_values["api"]["extraVolumeMounts"] == [ + { + "name": "workload-token-signing-key", + "mountPath": "/etc/nmp/workload-token", + "readOnly": True, + } + ] + + +def test_authentik_umbrella_chart_does_not_locally_template_subchart_workloads() -> None: + template_names = {path.name for path in (HELM_DIR / "templates").glob("*.yaml")} + + forbidden = { + "authentik-deployment.yaml", + "authentik-server-deployment.yaml", + "authentik-worker-deployment.yaml", + "postgres-deployment.yaml", + "postgres-statefulset.yaml", + "redis-deployment.yaml", + "nemo-deployment.yaml", + "nemo-service.yaml", + "gateway-configmap.yaml", + "gateway-deployment.yaml", + "gateway-service.yaml", + } + assert template_names.isdisjoint(forbidden) + + +def test_authentik_umbrella_chart_uses_single_canonical_authentik_blueprint() -> None: + packaged = (HELM_DIR / "files" / "blueprints" / "nemo.yaml").read_text(encoding="utf-8") + configmap_template = (HELM_DIR / "templates" / "blueprint-configmap.yaml").read_text(encoding="utf-8") + + assert not (AUTHENTIK_DIR / "blueprints" / "nemo.yaml").exists() + assert "grant_types:" in packaged + assert "- password" in packaged + assert "- urn:ietf:params:oauth:grant-type:device_code" in packaged + assert 'blueprints.goauthentik.io/instantiate: "true"' in packaged + assert '.Files.Get "files/blueprints/nemo.yaml"' in configmap_template + + values = _load_yaml(HELM_DIR / "values.yaml") + assert values["blueprintApplyJob"]["enabled"] is True + assert values["blueprintApplyJob"]["image"]["tag"] == "2026.5.4" + + +def test_authentik_umbrella_chart_applies_blueprint_with_waitable_helm_hook() -> None: + template = (HELM_DIR / "templates" / "blueprint-apply-job.yaml").read_text(encoding="utf-8") + + assert "kind: Job" in template + assert '"helm.sh/hook": post-install,post-upgrade' in template + assert '"helm.sh/hook-delete-policy": before-hook-creation,hook-succeeded' in template + assert "- ak" in template + assert "- apply_blueprint" in template + assert "- /blueprints/mounted/cm-authentik-nemo-blueprint/nemo.yaml" in template + assert "authentik-nemo-blueprint" in template + assert "automountServiceAccountToken: false" in template + assert "nmp.nvidia.com/blueprint-checksum" not in template + + +def test_authentik_kubernetes_runner_uses_helm_not_kustomize() -> None: + run_sh = (AUTHENTIK_DIR / "run.sh").read_text(encoding="utf-8") + live_test_path = Path("tests/auth_idp/k8s/test_authentik_kubernetes_live.py") + live_test = live_test_path.read_text(encoding="utf-8") + runtime_path = Path("tests/auth_idp/runtime_kubernetes.py") + runtime_impl = runtime_path.read_text(encoding="utf-8") + ci_workflow = Path(".github/workflows/ci.yaml").read_text(encoding="utf-8") + setup_kind_action = Path(".github/actions/setup-kind-cluster/action.yaml").read_text(encoding="utf-8") + run_commands = _literal_run_commands(runtime_path) + + assert "NMP_AUTHENTIK_K8S_HELM_RELEASE" in run_sh + assert 'K8S_RUNTIME="${NMP_AUTHENTIK_K8S_RUNTIME:-kind}"' in run_sh + assert "uv run --frozen pytest tests/auth_idp/contracts" in run_sh + assert "--auth-idp-runtime authentik-kubernetes" in run_sh + assert "-m auth_idp_runtime" in run_sh + assert "--run-e2e" not in run_sh + assert "uv run --frozen pytest tests/auth_idp/static -v" in ci_workflow + assert 'uv run --frozen pytest tests/auth_idp -v -m "auth_idp and not auth_idp_runtime"' not in ci_workflow + assert "--runtime RUNTIME" in run_sh + assert "validate_k8s_runtime" in run_sh + assert "--reuse" in run_sh + assert "--cluster-name" not in run_sh + assert "--reuse-cluster" not in run_sh + assert "--keep-cluster" not in run_sh + assert "--skip-image-load" in run_sh + expected_skip_image_load_line = ( + "NMP_AUTHENTIK_K8S_SKIP_IMAGE_LOAD: " + "${{ needs.build-cpu-smoke-images.outputs.publish_images == 'true' && '1' || '0' }}" + ) + assert expected_skip_image_load_line in ci_workflow + assert "NMP_AUTHENTIK_K8S_NAMESPACE: nemo-authentik" in ci_workflow + assert "kube-namespace: ${{ env.NMP_AUTHENTIK_K8S_NAMESPACE }}" in ci_workflow + assert "NMP_AUTHENTIK_K8S_IMAGE_PULL_SECRET: ghcr-pull" in ci_workflow + assert "NMP_AUTHENTIK_K8S_NGC_EXISTING_SECRET: ngc-api" in ci_workflow + assert 'K8S_IMAGE_PULL_SECRET="${NMP_AUTHENTIK_K8S_IMAGE_PULL_SECRET:-}"' in run_sh + assert "NMP_AUTHENTIK_K8S_IMAGE_PULL_SECRET=${K8S_IMAGE_PULL_SECRET}" in run_sh + assert "NMP_AUTHENTIK_K8S_IMAGE_PULL_SECRET" in runtime_impl + assert "nemo-platform.imagePullSecrets[0].name=" in runtime_impl + assert "NMP_AUTHENTIK_K8S_NGC_EXISTING_SECRET" in runtime_impl + assert "nemo-platform.existingSecret=" in runtime_impl + assert 'K8S_GATEWAY_PORT="${NMP_AUTHENTIK_K8S_GATEWAY_PORT:-18082}"' in run_sh + assert "NMP_AUTHENTIK_K8S_GATEWAY_PORT=${K8S_GATEWAY_PORT}" in run_sh + assert "NMP_AUTHENTIK_K8S_GATEWAY_PORT" in runtime_impl + assert "nemo-platform.authentikPublicGateway.port=" in runtime_impl + assert "GITHUB_TOKEN: ${{ inputs['kind-image-pull-token'] }}" in setup_kind_action + assert "CERT_MANAGER_CHART" not in runtime_impl + assert "_install_cert_manager" not in runtime_impl + assert ("helm", "repo", "add", "nvidia", "https://helm.ngc.nvidia.com/nvidia", "--force-update") in run_commands + assert ("helm", "repo", "add", "authentik", "https://charts.goauthentik.io", "--force-update") in run_commands + assert 'os.environ.get("NMP_AUTHENTIK_K8S_RUNTIME", "kind")' in runtime_impl + assert '"--no-hooks"' not in runtime_impl + assert "HELM_UPGRADE_COMMAND_TIMEOUT_SECONDS = 900" in runtime_impl + assert "_run(_helm_upgrade_args(context, kubeconfig), timeout=HELM_UPGRADE_COMMAND_TIMEOUT_SECONDS)" in runtime_impl + assert "PORT_FORWARD_READY_TIMEOUT_SECONDS = 30" in runtime_impl + assert "certificates.cert-manager.io" not in runtime_impl + assert "issuers.cert-manager.io" not in runtime_impl + assert "NMP_AUTHENTIK_K8S_KUSTOMIZATION" not in run_sh + assert "NMP_AUTHENTIK_K8S_KUSTOMIZATION" not in runtime_impl + assert 'kubectl", "--context", context, "apply"' not in runtime_impl + assert "auth_idp_k8s" in live_test + assert "nmp.nvidia.com/blueprint-checksum" not in (HELM_DIR / "values.yaml").read_text(encoding="utf-8") + + +def test_authentik_kubernetes_runner_builds_when_only_image_tag_env_is_set() -> None: + output = _run_authentik_script( + "k8s", + "--dry-run", + env={ + "IMAGE_REGISTRY": "registry.example.test/nemo", + "BAKE_TAG": "tag-from-env", + "NMP_AUTHENTIK_K8S_REUSE_CLUSTER": "0", + "NMP_AUTHENTIK_K8S_SKIP_IMAGE_LOAD": "0", + }, + ) + + assert "Building auth-idp test image for" in output + assert "make docker-load DOCKER_TARGET=nmp-api-docker" in output + assert "Using prebuilt auth-idp Kubernetes test image" not in output + + +def test_authentik_compose_runner_reuse_uses_stable_project_and_port() -> None: + output = _run_authentik_script("compose", "--dry-run", "--reuse") + + assert "workload-token-private-key.pem" in output + assert "NMP_E2E_COMPOSE_LIFECYCLE=reuse" in output + assert "NMP_AUTHENTIK_COMPOSE_PROJECT_NAME=authentik-e2e-reuse" in output + assert "NMP_AUTHENTIK_COMPOSE_GATEWAY_PORT=18083" in output + assert "--auth-idp-runtime authentik-compose" in output + + +def test_authentik_kubernetes_runner_reuse_uses_stable_cluster() -> None: + output = _run_authentik_script("k8s", "--dry-run", "--reuse") + + assert "NMP_AUTHENTIK_K8S_CLUSTER_NAME=nmp-authentik-reuse" in output + assert "NMP_AUTHENTIK_K8S_GATEWAY_PORT=18082" in output + assert "NMP_AUTHENTIK_K8S_REUSE_CLUSTER=1" in output + assert "NMP_AUTHENTIK_K8S_KEEP_CLUSTER=1" in output + assert "--auth-idp-runtime authentik-kubernetes" in output + + +def test_authentik_down_cleans_reused_compose_and_kubernetes_resources() -> None: + output = _run_authentik_script("down", "--dry-run") + + assert "docker compose down -v --remove-orphans" in output + assert "COMPOSE_PROJECT_NAME=authentik-e2e-reuse" in output + assert "AUTHENTIK_GATEWAY_PORT=18083" in output + assert "AUTHENTIK_GATEWAY_TLS_VOLUME=authentik-e2e-18083-gateway-tls" in output + assert "AUTHENTIK_WORKLOAD_NETWORK_NAME=authentik-e2e-18083-workload" in output + assert "kind delete cluster --name nmp-authentik-reuse" in output + + +def test_authentik_down_accepts_kubernetes_runtime_for_reuse_cleanup() -> None: + output = _run_authentik_script("down", "--dry-run", "--runtime", "k3d") + + assert "k3d cluster delete nmp-authentik-reuse" in output + + +def test_authentik_kubernetes_runner_skips_build_only_for_explicit_image() -> None: + output = _run_authentik_script("k8s", "--dry-run", "--image", "registry.example.test/nmp-api:prebuilt") + + assert "Using prebuilt auth-idp Kubernetes test image: registry.example.test/nmp-api:prebuilt" in output + assert "make docker-load DOCKER_TARGET=nmp-api-docker" not in output + assert "NMP_AUTHENTIK_K8S_WORKLOAD_TOKEN_PRIVATE_KEY_FILE=" in output + + +def test_authentik_kubernetes_runtime_uses_provisioned_signing_key_file() -> None: + run_sh = (AUTHENTIK_DIR / "run.sh").read_text(encoding="utf-8") + runtime_impl = Path("tests/auth_idp/runtime_kubernetes.py").read_text(encoding="utf-8") + helpers = (HELM_DIR / "templates" / "_helpers.tpl").read_text(encoding="utf-8") + + assert "ensure_workload_token_private_key" in run_sh + assert "NMP_AUTHENTIK_K8S_WORKLOAD_TOKEN_PRIVATE_KEY_FILE" in run_sh + assert "WORKLOAD_TOKEN_PRIVATE_KEY_FILE_ENV" in runtime_impl + assert '"--set-file"' in runtime_impl + assert "workloadTokenSigningKey.privateKeyPem=" in runtime_impl + assert "workloadTokenSigningKey.privateKeyPem" in helpers + assert 'genPrivateKey "rsa"' in helpers + + +def test_authentik_kubernetes_live_test_uses_workload_client_audience_for_subject_token() -> None: + runtime_impl = Path("tests/auth_idp/runtime_kubernetes.py").read_text(encoding="utf-8") + runtime_tree = ast.parse(runtime_impl) + + assert any( + isinstance(node, ast.List) + and any( + isinstance(argument_node, ast.Constant) + and argument_node.value == "--audience" + and isinstance(value_node, ast.Name) + and value_node.id == "WORKLOAD_CLIENT_ID" + for argument_node, value_node in zip(node.elts, node.elts[1:]) + ) + for node in ast.walk(runtime_tree) + ) + assert '"audience": WORKLOAD_AUDIENCE,' in runtime_impl + + +def test_authentik_runners_capture_ci_and_local_diagnostics() -> None: + run_sh = (AUTHENTIK_DIR / "run.sh").read_text(encoding="utf-8") + runtime_impl = Path("tests/auth_idp/runtime_kubernetes.py").read_text(encoding="utf-8") + + assert "diagnostics_dir()" in run_sh + assert "prepare_diagnostics_dir()" in run_sh + assert "write_diagnostics_metadata()" in run_sh + assert '>"${output}/run-metadata.txt"' in run_sh + assert "E2E_SERVICES_LOG_DIR=${diagnostics}" in run_sh + assert "NMP_AUTHENTIK_K8S_LOG_DIR=${k8s_diagnostics}" in run_sh + assert 'tee "${diagnostics}/pytest.log"' in run_sh + assert "Auth-idp Compose diagnostics:" in run_sh + assert "Auth-idp Kubernetes diagnostics:" in run_sh + + assert 'configured_dir = os.environ.get("NMP_AUTHENTIK_K8S_LOG_DIR")' in runtime_impl + assert '"helm-status.txt"' in runtime_impl + assert '"helm-list.txt"' in runtime_impl + assert '"get-nodes.txt"' in runtime_impl + assert '"get-pods-json.txt"' in runtime_impl + assert "self._diagnostics_collected = False" in runtime_impl + assert "def _collect_diagnostics_best_effort" in runtime_impl + assert "with contextlib.suppress(Exception):" in runtime_impl + assert "Collected Authentik Kubernetes diagnostics:" in runtime_impl + + +def test_authentik_compose_runner_uses_nemo_scoped_ca_bundle() -> None: + run_sh = (AUTHENTIK_DIR / "run.sh").read_text(encoding="utf-8") + + assert "NMP_CLIENT_SSL_CERT_FILE=$(gateway_tls_cert_file)" in run_sh + assert '"NMP_CLIENT_SSL_CERT_FILE=$(gateway_tls_cert_file)" \\' in run_sh + assert "uv run --frozen pytest tests/auth_idp/contracts" in run_sh + assert "--auth-idp-runtime authentik-compose" in run_sh + assert "-m auth_idp_runtime" in run_sh + stripped_lines = [line.strip() for line in run_sh.splitlines()] + assert not any(line.startswith('SSL_CERT_FILE="$(gateway_tls_cert_file)"') for line in stripped_lines) + assert not any(line.startswith('REQUESTS_CA_BUNDLE="$(gateway_tls_cert_file)"') for line in stripped_lines) + + +def test_authentik_umbrella_values_configure_workload_token_tls() -> None: + values = _load_yaml(HELM_DIR / "values.yaml") + tls_values = values["workloadTokenTls"] + nemo_values = values["nemo-platform"] + tls_template = (HELM_DIR / "templates" / "workload-token-tls.yaml").read_text(encoding="utf-8") + helpers_template = (HELM_DIR / "templates" / "_helpers.tpl").read_text(encoding="utf-8") + + assert tls_values["create"] is True + assert tls_values["secretName"] == "nemo-platform-envoy-tls" + assert tls_values["durationDays"] == 365 + assert tls_values["mountPath"] == "/etc/nmp/workload-token-tls" + assert tls_values["caBundleFile"] == "/etc/nmp/workload-token-ca/ca.crt" + assert tls_values["dnsNames"] == ["localhost"] + assert "127.0.0.1" in tls_values["ipAddresses"] + assert "selfSignedIssuerName" not in tls_values + assert "caIssuerName" not in tls_values + assert "caSecretName" not in tls_values + assert "type: kubernetes.io/tls" in tls_template + assert 'define "nemo-platform-authentik.serviceDnsNames"' in helpers_template + assert 'include "nemo-platform-authentik.serviceDnsNames"' in tls_template + assert ".Values.integration.nemoPlatform.envoyServiceName" in tls_template + assert "$nemoPlatformValues.envoyProxy.serviceNamespace" in tls_template + assert 'include "nemo-platform-authentik.existingSecretData"' in tls_template + assert "genSignedCert" in tls_template + assert "kind: Issuer" not in tls_template + assert "kind: Certificate" not in tls_template + assert "cert-manager.io/v1" not in tls_template + assert nemo_values["envoyProxy"]["extraVolumes"] == [ + {"name": "tmp", "emptyDir": {}}, + {"name": "workload-token-tls", "secret": {"secretName": "nemo-platform-envoy-tls"}}, + ] + assert nemo_values["envoyProxy"]["extraVolumeMounts"] == [ + {"name": "tmp", "mountPath": "/tmp"}, + {"name": "workload-token-tls", "mountPath": "/etc/nmp/workload-token-tls", "readOnly": True}, + ] + assert nemo_values["core"]["controller"]["env"] == { + "NMP_PLATFORM_URL": ENVOY_CONTROLLER_ENV_URL, + "NMP_AUTH_URL": ENVOY_CONTROLLER_ENV_URL, + } + assert nemo_values["platformConfig"]["jobs"]["executor_defaults"]["kubernetes_job"]["env"] == { + "SSL_CERT_FILE": "/etc/nmp/workload-token-ca/ca.crt", + "REQUESTS_CA_BUNDLE": "/etc/nmp/workload-token-ca/ca.crt", + } + assert "service_discovery" not in nemo_values["platformConfig"]["jobs"]["executor_defaults"]["kubernetes_job"] + assert nemo_values["platformConfig"]["jobs"]["executor_defaults"]["kubernetes_job"]["storage"] == { + "additional_volumes": [ + { + "name": "workload-token-tls-ca", + "secret": { + "secret_name": "nemo-platform-envoy-tls", + "items": [{"key": "ca.crt", "path": "ca.crt"}], + }, + } + ], + "additional_volume_mounts": [ + { + "name": "workload-token-tls-ca", + "mount_path": "/etc/nmp/workload-token-ca", + "read_only": True, + } + ], + } + workload_executor = next( + executor + for executor in nemo_values["platformConfig"]["jobs"]["executors"] + if executor["provider"] == "cpu" and executor["profile"] == "workload" + ) + workload_config = workload_executor["config"] + assert workload_config["launcher_image"] == '{{ include "nmp-core.image" . }}' + assert "default_task_image" not in workload_config + assert "service_discovery" not in workload_config + assert workload_config["env"] == { + "SSL_CERT_FILE": "/etc/nmp/workload-token-ca/ca.crt", + "REQUESTS_CA_BUNDLE": "/etc/nmp/workload-token-ca/ca.crt", + } + expected_storage = nemo_values["platformConfig"]["jobs"]["executor_defaults"]["kubernetes_job"]["storage"] + assert workload_config["storage"] == expected_storage + assert not any( + executor["provider"] == "gpu_distributed" and executor["profile"] == "default" + for executor in nemo_values["platformConfig"]["jobs"]["executors"] + ) + + +def test_nemo_platform_seed_hook_uses_internal_api_service_url() -> None: + template = Path("k8s/helm/templates/platform-seed-job.yaml").read_text(encoding="utf-8") + + assert "- name: NMP_BASE_URL" in template + assert 'include "nemo-platform.internalBaseUrl"' in template + + +def test_nemo_platform_controller_uses_internal_api_service_url_for_embedded_pdp() -> None: + template = Path("k8s/helm/templates/core/controller-deployment.yaml").read_text(encoding="utf-8") + + assert "- name: NMP_BASE_URL" in template + assert "- name: NMP_AUTH_POLICY_DECISION_POINT_BASE_URL" in template + assert ( + "name: NMP_AUTH_POLICY_DECISION_POINT_BASE_URL\n" + ' value: {{ include "nemo-platform.internalBaseUrl" . | quote }}' + ) in template + + +def test_authentik_helm_demo_does_not_inject_legacy_workload_token_envs() -> None: + template_files = sorted(path for path in (HELM_DIR / "templates").rglob("*") if path.is_file()) + text_files = [ + HELM_DIR / "values.yaml", + *template_files, + AUTHENTIK_DIR / "run.sh", + ] + combined = "\n".join(path.read_text(encoding="utf-8") for path in text_files) + + assert "NEMO_WORKLOAD_TOKEN" not in combined + assert "NEMO_WORKLOAD_TOKEN_FILE" not in combined + assert WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR not in combined diff --git a/tests/auth_idp/test_docs_links.py b/tests/auth_idp/static/test_docs_links.py similarity index 73% rename from tests/auth_idp/test_docs_links.py rename to tests/auth_idp/static/test_docs_links.py index e9804a9f73..ea228aaeaf 100644 --- a/tests/auth_idp/test_docs_links.py +++ b/tests/auth_idp/static/test_docs_links.py @@ -4,6 +4,7 @@ from pathlib import Path import pytest +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR pytestmark = [pytest.mark.auth_idp] @@ -11,4 +12,4 @@ def test_auth_docs_link_to_contrib_references(): content = Path("docs/auth/authentication/idp-integration.mdx").read_text() assert "contrib/auth/authentik" in content - assert "NEMO_WORKLOAD_TOKEN" in content + assert WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR in content diff --git a/tests/auth_idp/static/test_fixture_helpers.py b/tests/auth_idp/static/test_fixture_helpers.py new file mode 100644 index 0000000000..45dba588e0 --- /dev/null +++ b/tests/auth_idp/static/test_fixture_helpers.py @@ -0,0 +1,339 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for auth-idp pytest fixtures and fixture-only helper functions.""" + +import os +from pathlib import Path +from types import SimpleNamespace +from typing import Any, cast + +import pytest +from cryptography import x509 +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR + +from e2e.services_pool import RunningServices +from tests.auth_idp import conftest, runtime +from tests.auth_idp.authentik_live import authentik_gateway_tls_ca_bundle, prepare_authentik_compose_inputs +from tests.auth_idp.common import jwt_claims, require_capability +from tests.auth_idp.conftest import _token_request_body +from tests.auth_idp.providers import ProviderConfig, load_provider_configs +from tests.auth_idp.runtime_contract import AuthIdpCase +from tests.auth_idp.xdist import append_xdist_group_suffix + +pytestmark = [pytest.mark.auth_idp] + + +def test_authentik_stack_fixture_uses_pooled_gateway_metadata(): + provider = ProviderConfig( + name="authentik", + mode="compose-ci", + compose_file=Path("docker-compose.yml"), + gateway_base_url="https://127.0.0.1:18080", + issuer_url="http://authentik-server:9000/application/o/nemo/", + discovery_url="https://127.0.0.1:18080/application/o/nemo/.well-known/openid-configuration", + token_endpoint="https://127.0.0.1:18080/application/o/token/", + nemo_config=Path("config/platform-compose-authentik.yaml"), + interactive_user_username="nemo-user", + interactive_user_password="nemo-user-password-dev", + interactive_user_expected_email="nemo-user@example.com", + workload_principal_id="svc-nemo", + workload_expected_groups=["nemo-workloads"], + workload_audience="nemo-platform", + workload_principal_claim="sub", + workload_groups_claim="groups", + workload_groups_format="comma_string", + workload_token_env_vars=[WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR], + workload_forwarded_headers={ + "principal_id": "X-NMP-Principal-Id", + "principal_groups": "X-NMP-Principal-Groups", + }, + e2e_setup_password_grant={ + "grant_type": "password", + "client_id": "nemo-platform", + "username": "nemo-setup", + "password": "nemo-setup-token-secret-dev", + "scope": "openid email groups", + }, + interactive_user_password_grant=None, + workload_provider_password_grant={ + "grant_type": "password", + "username": "svc-nemo", + "password": "shared-secret", + }, + healthchecks=[], + startup_timeouts={}, + ) + fixture_fn = cast(Any, conftest.authentik_stack).__wrapped__ + stack = fixture_fn(None, provider, "https://127.0.0.1:28080") + + assert stack.gateway_base_url == "https://127.0.0.1:28080" + assert stack.discovery_url == "https://127.0.0.1:28080/application/o/nemo/.well-known/openid-configuration" + assert stack.token_endpoint == "https://127.0.0.1:28080/application/o/token/" + assert stack.nemo_config == provider.nemo_config + + +def test_auth_idp_runtime_event_line_includes_compose_instance_metadata(tmp_path): + case = AuthIdpCase( + id="authentik-compose", + provider=ProviderConfig( + name="authentik", + mode="compose-ci", + compose_file=Path("docker-compose.yml"), + gateway_base_url="https://127.0.0.1:18080", + issuer_url="http://authentik-server:9000/application/o/nemo/", + discovery_url="https://127.0.0.1:18080/application/o/nemo/.well-known/openid-configuration", + token_endpoint="https://127.0.0.1:18080/application/o/token/", + nemo_config=Path("config/platform-compose-authentik.yaml"), + interactive_user_username="nemo-user", + interactive_user_password="nemo-user-password-dev", + interactive_user_expected_email="nemo-user@example.com", + workload_principal_id="svc-nemo", + workload_expected_groups=["nemo-workloads"], + workload_audience="nemo-platform", + workload_principal_claim="sub", + workload_groups_claim="groups", + workload_groups_format="comma_string", + workload_token_env_vars=[WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR], + workload_forwarded_headers={}, + e2e_setup_password_grant=None, + interactive_user_password_grant=None, + workload_provider_password_grant=None, + healthchecks=[], + startup_timeouts={}, + ), + backend="compose", + capabilities=frozenset(), + ) + services = RunningServices( + url="https://127.0.0.1:49123", + compose_project_name="authentik-e2e-abc123-deadbeef", + log_path=tmp_path / "pytest.log", + proc=None, + config_path=None, + ) + runtime_obj = SimpleNamespace(gateway_base_url="https://127.0.0.1:49123") + + summary = conftest._auth_idp_runtime_event_line("available", case, runtime_obj, services) + + assert summary == ( + "Auth-idp runtime available: id=authentik-compose backend=compose " + "url=https://127.0.0.1:49123 port=49123 " + f"compose_project=authentik-e2e-abc123-deadbeef log={tmp_path / 'pytest.log'}" + ) + + +def test_auth_idp_runtime_event_line_includes_kubernetes_instance_metadata(): + case = AuthIdpCase( + id="authentik-kubernetes", + provider=ProviderConfig( + name="authentik", + mode="kubernetes-ci", + compose_file=None, + gateway_base_url="https://127.0.0.1:18081", + issuer_url="http://authentik-server:9000/application/o/nemo/", + discovery_url="https://127.0.0.1:18081/application/o/nemo/.well-known/openid-configuration", + token_endpoint="https://127.0.0.1:18081/application/o/token/", + nemo_config=Path("config/platform-compose-authentik.yaml"), + interactive_user_username="nemo-user", + interactive_user_password="nemo-user-password-dev", + interactive_user_expected_email="nemo-user@example.com", + workload_principal_id="svc-nemo", + workload_expected_groups=["nemo-workloads"], + workload_audience="nemo-platform", + workload_principal_claim="sub", + workload_groups_claim="groups", + workload_groups_format="comma_string", + workload_token_env_vars=[WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR], + workload_forwarded_headers={}, + e2e_setup_password_grant=None, + interactive_user_password_grant=None, + workload_provider_password_grant=None, + healthchecks=[], + startup_timeouts={}, + ), + backend="kubernetes", + capabilities=frozenset(), + ) + runtime_obj = SimpleNamespace( + gateway_base_url="https://127.0.0.1:39001", + cluster=SimpleNamespace( + name="ci", + context="kind-ci", + runtime="kind", + kubeconfig=Path("/tmp/nmp-authentik-kubeconfig.yaml"), + ), + namespace="nemo-authentik", + helm_release="authentik-demo", + ) + + summary = conftest._auth_idp_runtime_event_line("teardown_complete", case, runtime_obj) + + assert summary == ( + "Auth-idp runtime teardown_complete: id=authentik-kubernetes backend=kubernetes " + "url=https://127.0.0.1:39001 port=39001 " + "cluster=ci context=kind-ci runtime=kind kubeconfig=/tmp/nmp-authentik-kubeconfig.yaml " + "namespace=nemo-authentik helm_release=authentik-demo" + ) + + +def test_auth_idp_runtime_result_event_line_includes_status(): + case = AuthIdpCase( + id="authentik-compose", + provider=load_provider_configs()[0], + backend="compose", + capabilities=frozenset(), + ) + + summary = conftest._auth_idp_runtime_event_line("result", case, status="passed") + + assert summary == "Auth-idp runtime result: id=authentik-compose backend=compose status=passed" + + +def test_write_terminal_line_bypasses_pytest_capture(capsys): + events: list[str] = [] + + class FakeCaptureManager: + def global_and_fixture_disabled(self): + class CaptureDisabled: + def __enter__(self): + events.append("enter") + + def __exit__(self, exc_type, exc, traceback): + events.append("exit") + + return CaptureDisabled() + + class FakePluginManager: + def get_plugin(self, name: str): + assert name == "capturemanager" + return FakeCaptureManager() + + request = SimpleNamespace(config=SimpleNamespace(pluginmanager=FakePluginManager())) + + conftest._write_terminal_line(cast(Any, request), "Auth-idp runtime available: id=authentik-compose") + + assert events == ["enter", "exit"] + assert capsys.readouterr().out == "Auth-idp runtime available: id=authentik-compose\n" + + +def test_token_request_body_for_password_grant_includes_username_and_password(): + assert _token_request_body( + { + "grant_type": "password", + "client_id": "nemo-platform", + "client_secret": "secret", + "username": "akadmin", + "password": "akadmin-dev", + "scope": "openid profile email groups", + } + ) == { + "grant_type": "password", + "client_id": "nemo-platform", + "client_secret": "secret", + "username": "akadmin", + "password": "akadmin-dev", + "scope": "openid profile email groups", + } + + +def test_token_request_body_for_workload_password_grant_includes_username_and_password(): + assert _token_request_body( + { + "grant_type": "password", + "client_id": "nemo-platform", + "client_secret": "secret", + "username": "svc-nemo", + "password": "shared-secret", + "scope": "openid email groups", + } + ) == { + "grant_type": "password", + "client_id": "nemo-platform", + "client_secret": "secret", + "username": "svc-nemo", + "password": "shared-secret", + "scope": "openid email groups", + } + + +def test_jwt_claims_decodes_payload() -> None: + token = "eyJhbGciOiJub25lIn0.eyJzdWIiOiJwcm9qZWN0LXN1YmplY3QiLCJncm91cHMiOiJuZW1vLXdvcmtsb2FkcyJ9." + + claims = jwt_claims(token) + + assert claims["sub"] == "project-subject" + assert claims["groups"] == "nemo-workloads" + + +def test_require_capability_skips_when_missing() -> None: + case = AuthIdpCase( + id="provider-smoke", + provider=load_provider_configs()[0], + backend="external", + capabilities=frozenset({"gateway_discovery"}), + ) + + with pytest.raises(pytest.skip.Exception): + require_capability(case, "workspace_rbac") + + +def test_prepare_authentik_compose_inputs_creates_generated_assets(tmp_path, monkeypatch) -> None: + monkeypatch.delenv("AUTHENTIK_BLUEPRINT_DIR", raising=False) + monkeypatch.delenv("AUTHENTIK_GATEWAY_TLS_DIR", raising=False) + source_blueprint = tmp_path / "helm/files/blueprints/nemo.yaml" + source_blueprint.parent.mkdir(parents=True) + source_blueprint.write_text("version: 1\n", encoding="utf-8") + + prepare_authentik_compose_inputs(root=tmp_path) + + generated = tmp_path / ".generated" + assert (generated / "blueprints/nemo.yaml").read_text(encoding="utf-8") == "version: 1\n" + key_path = generated / "workload-token-private-key.pem" + assert key_path.read_text(encoding="utf-8").startswith("-----BEGIN RSA PRIVATE KEY-----") + assert key_path.stat().st_mode & 0o777 == 0o600 + + tls_cert = generated / "gateway-tls/tls.crt" + tls_key = generated / "gateway-tls/tls.key" + assert authentik_gateway_tls_ca_bundle(root=tmp_path) == tls_cert + assert tls_key.stat().st_mode & 0o777 == 0o600 + assert tls_cert.stat().st_mode & 0o777 == 0o644 + assert "DNS.2 = nemo-gateway" in (generated / "gateway-tls/openssl.cnf").read_text(encoding="utf-8") + cert = x509.load_pem_x509_certificate(tls_cert.read_bytes()) + san = cert.extensions.get_extension_for_class(x509.SubjectAlternativeName).value + assert san.get_values_for_type(x509.DNSName) == ["localhost", "nemo-gateway"] + assert [str(value) for value in san.get_values_for_type(x509.IPAddress)] == ["127.0.0.1"] + + +def test_authentik_gateway_tls_ca_bundle_honors_tls_dir_override(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("AUTHENTIK_GATEWAY_TLS_DIR", "./custom-gateway-tls") + + assert authentik_gateway_tls_ca_bundle(root=tmp_path) == tmp_path / "custom-gateway-tls/tls.crt" + + +def test_authentik_docker_runtime_defaults_workload_identity_password(monkeypatch): + monkeypatch.delenv(runtime.AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD_ENVVAR, raising=False) + runtime.get_authentik_docker_test_runtime.cache_clear() + try: + try: + provider = runtime.get_authentik_docker_test_runtime() + finally: + runtime.get_authentik_docker_test_runtime.cache_clear() + + assert os.environ[runtime.AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD_ENVVAR] == ( + runtime.AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD_DEFAULT + ) + assert provider.workload_provider_password_grant is not None + assert ( + provider.workload_provider_password_grant["password"] + == runtime.AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD_DEFAULT + ) + finally: + os.environ.pop(runtime.AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD_ENVVAR, None) + + +def test_append_xdist_group_suffix_only_appends_once_and_sorts_groups(): + nodeid = "tests/auth_idp/contracts/test_tokens.py::test_provider_workload_provider_token_is_real" + assert append_xdist_group_suffix(nodeid, {"idp-live"}) == f"{nodeid}@idp-live" + assert append_xdist_group_suffix(nodeid, {"b", "a"}) == f"{nodeid}@a_b" + assert append_xdist_group_suffix(f"{nodeid}@idp-live", {"idp-live"}) == f"{nodeid}@idp-live" diff --git a/tests/auth_idp/static/test_provider_layout.py b/tests/auth_idp/static/test_provider_layout.py new file mode 100644 index 0000000000..44c0814f09 --- /dev/null +++ b/tests/auth_idp/static/test_provider_layout.py @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +import pytest +import yaml + +from tests.auth_idp.providers import load_provider_configs_by_mode, load_provider_names_by_mode + +pytestmark = [pytest.mark.auth_idp] + + +def test_provider_name_discovery_does_not_require_grant_secret(monkeypatch): + monkeypatch.delenv("AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD", raising=False) + + assert "authentik" in load_provider_names_by_mode("compose-ci") + + +def test_compose_backed_providers_ship_required_assets(monkeypatch): + monkeypatch.setenv("AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD", "shared-secret") + for provider in load_provider_configs_by_mode("compose-ci"): + root = Path(f"contrib/auth/{provider.name}") + assert provider.compose_file is not None + assert provider.compose_file.exists() + assert (root / "gateway").exists() + assert (root / "README.md").exists() + assert (root / "manifest.yaml").exists() + + +def test_reference_only_providers_do_not_require_compose(monkeypatch): + monkeypatch.setenv("AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD", "shared-secret") + for provider in load_provider_names_by_mode("reference-only"): + root = Path(f"contrib/auth/{provider}") + assert (root / "README.md").exists() + assert not (root / "docker-compose.yml").exists() + + +def test_authentik_compose_disables_model_provider_seed_without_ngc_key(): + compose = yaml.safe_load(Path("contrib/auth/authentik/compose/docker-compose.yml").read_text()) + nemo_service = compose["services"]["nemo"] + nemo_env = nemo_service["environment"] + + assert nemo_env["NMP_SEED_ON_STARTUP"] == "true" + assert nemo_env["NMP_PLATFORM_SEED_MODEL_PROVIDER_ENABLED"] == "false" + assert "ports" not in nemo_service + + +def test_authentik_compose_defaults_support_direct_docker_compose_start(): + compose_path = Path("contrib/auth/authentik/compose/docker-compose.yml") + compose_text = compose_path.read_text(encoding="utf-8") + compose = yaml.safe_load(compose_text) + + password_default = "${AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD:-svc-nemo-token-secret-dev}" + blueprint_mount = "${AUTHENTIK_BLUEPRINT_DIR:-../helm/files/blueprints}:/blueprints/custom:ro" + + assert compose["name"] == "${COMPOSE_PROJECT_NAME:-nemo-platform-authentik}" + assert compose["x-authentik-env"]["AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD"] == password_default + assert compose["services"]["nemo"]["environment"]["AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD"] == password_default + assert ( + "../.generated/workload-token-private-key.pem:" + "/var/run/secrets/nemo-platform/workload-token-signing/private-key.pem:ro" + ) in compose["services"]["nemo"]["volumes"] + assert "../gateway/envoy.yaml:/etc/envoy/envoy.yaml:ro" in compose["services"]["gateway"]["volumes"] + assert ( + "${AUTHENTIK_GATEWAY_TLS_DIR:-../.generated/gateway-tls}:/source/tls:ro" + in compose["services"]["gateway-tls-init"]["volumes"] + ) + assert blueprint_mount in compose["services"]["authentik-server"]["volumes"] + assert blueprint_mount in compose["services"]["authentik-worker"]["volumes"] + assert "./.generated/blueprints" not in compose_text + + +def test_authentik_compose_uses_liveness_for_container_health_and_routes_status_through_gateway(): + compose = yaml.safe_load(Path("contrib/auth/authentik/compose/docker-compose.yml").read_text()) + envoy = yaml.safe_load(Path("contrib/auth/authentik/gateway/envoy.yaml").read_text()) + + nemo_healthcheck = compose["services"]["nemo"]["healthcheck"]["test"] + assert "http://127.0.0.1:8080/health/live" in nemo_healthcheck[-1] + assert "/health/ready" not in nemo_healthcheck[-1] + + http_manager = envoy["static_resources"]["listeners"][0]["filter_chains"][0]["filters"][0]["typed_config"] + routes = http_manager["route_config"]["virtual_hosts"][0]["routes"] + route_matches = [route["match"] for route in routes if route.get("route", {}).get("cluster") == "nemo"] + forwarded_proto_header = [ + { + "header": {"key": "x-forwarded-proto", "value": "https"}, + "append_action": "OVERWRITE_IF_EXISTS_OR_ADD", + } + ] + gateway_ready_route = next(route for route in routes if route["match"] == {"path": "/health/gateway/ready"}) + health_route = next(route for route in routes if route["match"] == {"prefix": "/health/"}) + assert routes.index(gateway_ready_route) < routes.index(health_route) + assert gateway_ready_route["direct_response"] == { + "status": 503, + "body": {"inline_string": '{"status":"not_ready"}'}, + } + assert {"prefix": "/health/"} in route_matches + assert {"path": "/status"} in route_matches + for match in ( + {"prefix": "/.well-known/nemo-platform/"}, + {"prefix": "/apis/"}, + {"prefix": "/health/"}, + {"path": "/status"}, + {"prefix": "/studio/"}, + ): + route = next( + route for route in routes if route.get("route", {}).get("cluster") == "nemo" and route["match"] == match + ) + assert route["request_headers_to_add"] == forwarded_proto_header + + lua_filter = next( + filter_config + for filter_config in http_manager["http_filters"] + if filter_config["name"] == "envoy.filters.http.lua" + ) + lua_code = lua_filter["typed_config"]["inline_code"] + assert 'headers:get(":path") ~= "/health/gateway/ready"' in lua_code + assert 'gateway_ready_http_call(request_handle, "nemo", "nemo", "/health/ready")' in lua_code + assert ( + 'gateway_ready_http_call(request_handle, "authentik", "authentik-server", ' + '"/application/o/nemo/.well-known/openid-configuration")' + ) in lua_code + + jwt_filter = next( + filter_config + for filter_config in http_manager["http_filters"] + if filter_config["name"] == "envoy.filters.http.jwt_authn" + ) + jwt_providers = jwt_filter["typed_config"]["providers"] + assert jwt_providers["authentik_workload"]["audiences"] == [ + "nemo-platform", + "nemo-platform-cli", + "nemo-platform-workload", + ] + assert jwt_providers["workload_exchange"]["audiences"] == ["nemo-platform"] + jwt_rules = jwt_filter["typed_config"]["rules"] + jwt_rule_matches = [rule["match"] for rule in jwt_rules] + assert {"prefix": "/health/"} in jwt_rule_matches + assert {"path": "/status"} in jwt_rule_matches + + +def test_authentik_compose_mounts_workload_token_signing_key(): + compose = yaml.safe_load(Path("contrib/auth/authentik/compose/docker-compose.yml").read_text()) + config = yaml.safe_load(Path("contrib/auth/authentik/config/platform-compose-authentik.yaml").read_text()) + + key_mount = ( + "../.generated/workload-token-private-key.pem:" + "/var/run/secrets/nemo-platform/workload-token-signing/private-key.pem:ro" + ) + + assert key_mount in compose["services"]["nemo"]["volumes"] + assert ( + config["auth"]["oidc"]["workload_token_private_key_file"] + == "/var/run/secrets/nemo-platform/workload-token-signing/private-key.pem" + ) + + +def test_authentik_compose_uses_https_gateway_for_workloads(): + compose = yaml.safe_load(Path("contrib/auth/authentik/compose/docker-compose.yml").read_text()) + config = yaml.safe_load(Path("contrib/auth/authentik/config/platform-compose-authentik.yaml").read_text()) + nemo = compose["services"]["nemo"] + gateway = compose["services"]["gateway"] + gateway_tls_init = compose["services"]["gateway-tls-init"] + + assert config["platform"]["base_url"] == "https://nemo-gateway:8080" + assert config["auth"]["policy_decision_point_base_url"] == "http://127.0.0.1:8080" + assert nemo["environment"]["NMP_AUTH_POLICY_DECISION_POINT_BASE_URL"] == "http://127.0.0.1:8080" + assert "loopback_address" not in config["platform"] + assert "service_discovery" not in config["platform"] + assert config["auth"]["oidc"]["token_endpoint"] == "https://127.0.0.1:18080/application/o/token/" + assert config["auth"]["oidc"]["workload_token_issuer"] == "https://nemo-gateway:8080/apis/auth" + assert config["auth"]["oidc"]["workload_token_endpoint"] == "https://nemo-gateway:8080/apis/auth/token" + assert ( + "https://nemo-gateway:8080/application/o/nemo-workload/" in config["auth"]["oidc"]["workload_subject_issuers"] + ) + workload_executor = next( + executor + for executor in config["jobs"]["executors"] + if executor["provider"] == "cpu" and executor["profile"] == "workload" + ) + assert workload_executor["config"]["workload_identity"]["token_endpoint"] == ( + "https://nemo-gateway:8080/application/o/token/" + ) + assert set(nemo["networks"]) == {"nemo-internal"} + assert "nemo-direct" not in yaml.safe_dump(nemo) + assert nemo["depends_on"]["gateway-tls-init"]["condition"] == "service_completed_successfully" + assert gateway["networks"]["nemo-internal"]["aliases"] == ["nemo-gateway"] + assert gateway["networks"]["workload"]["aliases"] == ["nemo-gateway"] + assert gateway["depends_on"]["gateway-tls-init"]["condition"] == "service_completed_successfully" + assert "gateway-tls:/etc/envoy/tls:ro" in gateway["volumes"] + assert gateway_tls_init["image"] == ( + "docker.io/library/busybox:1.37.0@sha256:9532d8c39891ca2ecde4d30d7710e01fb739c87a8b9299685c63704296b16028" + ) + assert gateway_tls_init["user"] == "0:0" + assert gateway_tls_init["entrypoint"][:2] == ["sh", "-euc"] + assert "chown 101:101 /target/tls/tls.crt /target/tls/tls.key" in gateway_tls_init["entrypoint"][2] + assert "chmod 600 /target/tls/tls.key" in gateway_tls_init["entrypoint"][2] + assert compose["volumes"]["gateway-tls"]["name"] == "${AUTHENTIK_GATEWAY_TLS_VOLUME:-authentik_gateway_tls}" + assert "gateway-tls:/etc/nmp/gateway-tls:ro" in nemo["volumes"] + assert nemo["environment"]["SSL_CERT_FILE"] == "/etc/nmp/gateway-tls/tls.crt" + assert nemo["environment"]["REQUESTS_CA_BUNDLE"] == "/etc/nmp/gateway-tls/tls.crt" + + +def test_authentik_compose_mounts_gateway_ca_into_docker_workloads(): + config = yaml.safe_load(Path("contrib/auth/authentik/config/platform-compose-authentik.yaml").read_text()) + gateway_tls_volume_name = "authentik_gateway_tls" + + workload_executor = next( + executor + for executor in config["jobs"]["executors"] + if executor["provider"] == "cpu" and executor["profile"] == "workload" + ) + docker_config = workload_executor["config"] + + assert docker_config["env"] == { + "SSL_CERT_FILE": "/etc/nmp/gateway-tls/tls.crt", + "REQUESTS_CA_BUNDLE": "/etc/nmp/gateway-tls/tls.crt", + } + assert docker_config["storage"]["additional_volume_mounts"] == [ + { + "volume_name": gateway_tls_volume_name, + "mount_path": "/etc/nmp/gateway-tls", + } + ] diff --git a/tests/auth_idp/static/test_provider_manifest.py b/tests/auth_idp/static/test_provider_manifest.py new file mode 100644 index 0000000000..059ee8d605 --- /dev/null +++ b/tests/auth_idp/static/test_provider_manifest.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +import pytest +import yaml +from jsonschema.exceptions import ValidationError +from jsonschema.validators import validator_for +from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR + +from tests.auth_idp.providers import load_provider_config, load_provider_configs + +pytestmark = [pytest.mark.auth_idp] + + +def _load_provider_manifest_schema() -> dict: + return yaml.safe_load(Path("contrib/auth/manifest.schema.yaml").read_text()) + + +def test_all_provider_manifests_share_the_same_contract(monkeypatch): + monkeypatch.setenv("AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD", "shared-secret") + schema = _load_provider_manifest_schema() + validator = validator_for(schema)(schema) + for provider in load_provider_configs(): + manifest = yaml.safe_load(Path(f"contrib/auth/{provider.name}/manifest.yaml").read_text()) + validator.validate(manifest) + assert manifest["provider"] == provider.name + + +@pytest.mark.parametrize( + ("grant_name", "base_grant", "extra_credential"), + [ + ( + "e2e_setup_password_grant", + None, + {"password_env_var": "AUTHENTIK_SETUP_PASSWORD"}, + ), + ( + "interactive_user_password_grant", + { + "grant_type": "password", + "client_id": "nemo-platform", + "username": "nemo-user", + "password": "shared-secret", + "scope": "openid email groups", + }, + {"password_env_var": "AUTHENTIK_USER_PASSWORD"}, + ), + ("workload_provider_password_grant", None, {"password": "shared-secret"}), + ], +) +def test_provider_manifest_rejects_multiple_grant_credential_sources(grant_name, base_grant, extra_credential): + schema = _load_provider_manifest_schema() + validator = validator_for(schema)(schema) + manifest = yaml.safe_load(Path("contrib/auth/authentik/manifest.yaml").read_text()) + if base_grant is not None: + manifest["token_acquisition"][grant_name] = base_grant + manifest["token_acquisition"][grant_name].update(extra_credential) + + with pytest.raises(ValidationError): + validator.validate(manifest) + + +def test_authentik_manifest_declares_real_token_acquisition_contract(): + manifest = yaml.safe_load(Path("contrib/auth/authentik/manifest.yaml").read_text()) + token_acquisition = manifest["token_acquisition"] + interactive_user_identity = manifest["interactive_user_identity"] + principal_contract = manifest["principal_contract"] + workload_identity = manifest["workload_identity"] + workload_contract = manifest["workload_contract"] + + assert interactive_user_identity["username"] == "nemo-user" + assert interactive_user_identity["password"] == "nemo-user-password-dev" + assert interactive_user_identity["expected_email"] == "nemo-user@example.com" + assert token_acquisition["token_endpoint"] + setup_grant = token_acquisition["e2e_setup_password_grant"] + assert setup_grant["grant_type"] == "password" + assert setup_grant["client_id"] == "nemo-platform" + assert setup_grant["username"] == "nemo-setup" + assert setup_grant["password"] == "nemo-setup-token-secret-dev" + assert "password_env_var" not in setup_grant + assert "interactive_user_password_grant" not in token_acquisition + workload_provider_grant = token_acquisition["workload_provider_password_grant"] + assert workload_provider_grant["grant_type"] == "password" + assert workload_provider_grant["client_id"] + assert workload_provider_grant["password_env_var"] == "AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD" + assert "password" not in workload_provider_grant + assert workload_identity["principal_id"] + assert not workload_identity["principal_id"].startswith(principal_contract["internal_service_prefix_reserved"]) + assert workload_identity["expected_groups"] == ["nemo-workloads"] + assert workload_contract["audience"] == "nemo-platform" + assert workload_contract["groups_format"] == "comma_string" + assert workload_contract["token_env_vars"] == [WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR] + assert workload_contract["forwarded_headers"]["principal_id"] == "X-NMP-Principal-Id" + assert workload_contract["forwarded_headers"]["principal_groups"] == "X-NMP-Principal-Groups" + + +def test_authentik_manifest_declares_provider_test_runtimes(): + manifest = yaml.safe_load(Path("contrib/auth/authentik/manifest.yaml").read_text()) + + runtime_ids = {runtime["id"] for runtime in manifest["test_runtimes"]} + + assert "authentik-compose" in runtime_ids + assert "authentik-kubernetes" in runtime_ids + + +def test_authentik_manifest_compose_and_kubernetes_runtime_capabilities_stay_in_parity(): + manifest = yaml.safe_load(Path("contrib/auth/authentik/manifest.yaml").read_text()) + runtimes = {runtime["id"]: set(runtime["capabilities"]) for runtime in manifest["test_runtimes"]} + + compose = runtimes["authentik-compose"] + kubernetes = runtimes["authentik-kubernetes"] + + assert compose - {"docker_subject_token_refresh"} == kubernetes - {"kubernetes_token_review"} + assert "interactive_user_token" not in compose + assert "interactive_user_token" not in kubernetes + assert "workload_provider_token" in compose + assert "workload_provider_token" in kubernetes + + +def test_authentik_common_contracts_do_not_require_removed_interactive_user_token_capability(): + contract_text = "\n".join( + path.read_text(encoding="utf-8") for path in Path("tests/auth_idp/contracts").glob("*.py") + ) + + assert '"interactive_user_token"' not in contract_text + + +def test_all_provider_test_runtimes_declare_capabilities(monkeypatch): + monkeypatch.setenv("AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD", "shared-secret") + + for provider in load_provider_configs(): + assert provider.test_runtimes + for runtime in provider.test_runtimes: + assert runtime.id + assert runtime.backend in {"compose", "kubernetes", "external"} + assert runtime.capabilities + + +def test_authentik_manifest_declares_extended_startup_timeouts_for_real_oidc(): + manifest = yaml.safe_load(Path("contrib/auth/authentik/manifest.yaml").read_text()) + startup_timeouts = manifest["startup_timeouts"] + + assert startup_timeouts["healthchecks_seconds"] >= 240 + assert startup_timeouts["gateway_seconds"] >= 30 + assert startup_timeouts["token_endpoint_seconds"] >= 60 + + +def test_authentik_provider_config_loads_token_acquisition_fields(monkeypatch): + monkeypatch.setenv("AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD", "shared-secret") + provider = next(config for config in load_provider_configs() if config.name == "authentik") + + assert provider.compose_file == Path("contrib/auth/authentik/compose/docker-compose.yml") + assert provider.nemo_config == Path("contrib/auth/authentik/config/platform-compose-authentik.yaml") + assert provider.gateway_base_url == "https://127.0.0.1:18080" + assert provider.discovery_url == "https://127.0.0.1:18080/application/o/nemo/.well-known/openid-configuration" + assert provider.token_endpoint == "https://127.0.0.1:18080/application/o/token/" + assert provider.interactive_user_username == "nemo-user" + assert provider.interactive_user_password == "nemo-user-password-dev" + assert provider.interactive_user_expected_email == "nemo-user@example.com" + assert provider.e2e_setup_password_grant is not None + assert provider.e2e_setup_password_grant["grant_type"] == "password" + assert provider.e2e_setup_password_grant["password"] == "nemo-setup-token-secret-dev" + assert provider.interactive_user_password_grant is None + assert provider.workload_provider_password_grant["grant_type"] == "password" + assert provider.workload_audience == "nemo-platform" + assert provider.workload_principal_claim == "sub" + assert provider.workload_groups_claim == "groups" + assert provider.workload_groups_format == "comma_string" + assert provider.workload_token_env_vars == [WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR] + assert provider.startup_timeouts == { + "healthchecks_seconds": 600, + "gateway_seconds": 30, + "token_endpoint_seconds": 180, + } + + +def test_authentik_provider_config_resolves_workload_provider_password_grant_env_var(monkeypatch): + monkeypatch.setenv("AUTHENTIK_WORKLOAD_IDENTITY_PASSWORD", "shared-secret") + + provider = load_provider_config(Path("contrib/auth/authentik/manifest.yaml")) + + assert provider.workload_provider_password_grant is not None + assert provider.workload_provider_password_grant["password"] == "shared-secret" + assert "password_env_var" not in provider.workload_provider_password_grant + assert provider.e2e_setup_password_grant is not None + assert provider.e2e_setup_password_grant["password"] == "nemo-setup-token-secret-dev" + assert "password_env_var" not in provider.e2e_setup_password_grant diff --git a/tests/auth_idp/static/test_runtime_compose.py b/tests/auth_idp/static/test_runtime_compose.py new file mode 100644 index 0000000000..b421de4161 --- /dev/null +++ b/tests/auth_idp/static/test_runtime_compose.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import base64 +import json +from pathlib import Path + +import httpx + +from tests.auth_idp.providers import ProviderConfig +from tests.auth_idp.runtime_compose import ( + ACCESS_TOKEN_TYPE, + JWT_TOKEN_TYPE, + TOKEN_EXCHANGE_GRANT_TYPE, + TOKEN_EXCHANGE_TIMEOUT_SECONDS, + ComposeAuthIdpRuntime, +) +from tests.auth_idp.runtime_contract import AuthIdpCase + + +def _jwt(claims: dict[str, object]) -> str: + def encode(value: dict[str, object]) -> str: + payload = json.dumps(value, separators=(",", ":")).encode("utf-8") + return base64.urlsafe_b64encode(payload).decode("ascii").rstrip("=") + + return f"{encode({'alg': 'none'})}.{encode(claims)}." + + +def test_compose_workload_exchange_posts_token_exchange_grant(monkeypatch) -> None: + provider = ProviderConfig( + name="authentik", + mode="compose-ci", + compose_file=None, + gateway_base_url="https://127.0.0.1:18080", + issuer_url="http://authentik-server:9000/application/o/nemo/", + discovery_url="https://127.0.0.1:18080/application/o/nemo/.well-known/openid-configuration", + nemo_config=Path("config.yaml"), + interactive_user_username="nemo-user", + interactive_user_password="nemo-user-password-dev", + interactive_user_expected_email="nemo-user@example.com", + workload_principal_id="svc-nemo", + workload_expected_groups=["nemo-workloads"], + workload_audience="nemo-platform", + workload_principal_claim="sub", + workload_groups_claim="groups", + workload_groups_format="comma_string", + workload_token_env_vars=["NMP_WORKLOAD_IDENTITY_TOKEN_FILE"], + workload_forwarded_headers={}, + token_endpoint="https://127.0.0.1:18080/application/o/token/", + e2e_setup_password_grant={ + "grant_type": "password", + "client_id": "nemo-platform", + "username": "nemo-setup", + "password": "nemo-setup-token-secret-dev", + "scope": "openid email groups", + }, + interactive_user_password_grant=None, + workload_provider_password_grant={ + "grant_type": "password", + "client_id": "nemo-platform-workload", + "username": "svc-nemo", + "password": "secret", + "scope": "openid email groups", + }, + healthchecks=[], + startup_timeouts={}, + ) + case = AuthIdpCase( + id="authentik-compose", + provider=provider, + backend="compose", + capabilities=frozenset({"workload_token_exchange"}), + ) + runtime = ComposeAuthIdpRuntime(case, "https://127.0.0.1:18080") + subject_token = _jwt({"sub": "svc-nemo"}) + exchanged_token = _jwt({"sub": "svc-nemo", "groups": "nemo-workloads"}) + captured: dict[str, object] = {} + + def fake_post(url: str, *, data: dict[str, str], timeout: float, verify: str | bool) -> httpx.Response: + captured.update({"url": url, "data": data, "timeout": timeout, "verify": verify}) + request = httpx.Request("POST", url) + return httpx.Response(200, json={"access_token": exchanged_token, "token_type": "Bearer"}, request=request) + + monkeypatch.setenv("NMP_CLIENT_SSL_CERT_FILE", "/tmp/nemo-ca.pem") + monkeypatch.setattr("tests.auth_idp.runtime_compose.httpx.post", fake_post) + + token = runtime.exchange_workload_token(subject_token) + + assert token.access_token == exchanged_token + assert token.claims == {"sub": "svc-nemo", "groups": "nemo-workloads"} + assert captured == { + "url": "https://127.0.0.1:18080/apis/auth/token", + "data": { + "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, + "client_id": "nemo-platform-workload", + "subject_token": subject_token, + "subject_token_type": JWT_TOKEN_TYPE, + "requested_token_type": ACCESS_TOKEN_TYPE, + "audience": "nemo-platform", + "scope": "openid email groups", + }, + "timeout": TOKEN_EXCHANGE_TIMEOUT_SECONDS, + "verify": "/tmp/nemo-ca.pem", + } diff --git a/tests/auth_idp/static/test_runtime_selection.py b/tests/auth_idp/static/test_runtime_selection.py new file mode 100644 index 0000000000..30e7fe042c --- /dev/null +++ b/tests/auth_idp/static/test_runtime_selection.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from tests.auth_idp.runtime_factory import iter_auth_idp_cases, runtime_class_for_case + +pytestmark = [pytest.mark.auth_idp] + + +def test_iter_auth_idp_cases_includes_authentik_compose() -> None: + case_ids = {case.id for case in iter_auth_idp_cases()} + + assert "authentik-compose" in case_ids + + +def test_iter_auth_idp_cases_can_filter_by_backend() -> None: + case_ids = {case.id for case in iter_auth_idp_cases(backend="kubernetes")} + + assert "authentik-kubernetes" in case_ids + assert "authentik-compose" not in case_ids + + +def test_iter_auth_idp_cases_can_filter_by_provider_name() -> None: + case_ids = {case.id for case in iter_auth_idp_cases(provider_name="authentik")} + + assert case_ids == {"authentik-compose", "authentik-kubernetes"} + + +def test_iter_auth_idp_cases_can_filter_by_runtime_id() -> None: + case_ids = {case.id for case in iter_auth_idp_cases(runtime_id="authentik-kubernetes")} + + assert case_ids == {"authentik-kubernetes"} + + +def test_runtime_case_has_pytest_id() -> None: + [case] = [case for case in iter_auth_idp_cases() if case.id == "authentik-compose"] + + assert pytest.param(case, id=case.id).id == "authentik-compose" + + +def test_compose_runtime_factory_selects_compose_runtime() -> None: + case = next(case for case in iter_auth_idp_cases() if case.id == "authentik-compose") + + runtime_class = runtime_class_for_case(case) + + assert runtime_class.__name__ == "ComposeAuthIdpRuntime" + + +def test_kubernetes_runtime_factory_selects_kubernetes_runtime() -> None: + case = next(case for case in iter_auth_idp_cases() if case.id == "authentik-kubernetes") + + runtime_class = runtime_class_for_case(case) + + assert runtime_class.__name__ == "KubernetesAuthIdpRuntime" diff --git a/tests/auth_idp/test_authentik_cli_login.py b/tests/auth_idp/test_authentik_cli_login.py deleted file mode 100644 index cb1e6552b3..0000000000 --- a/tests/auth_idp/test_authentik_cli_login.py +++ /dev/null @@ -1,89 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import base64 -import json -import uuid - -import httpx -from nemo_platform_ext.auth.helpers import discover_nmp_config - -from tests.auth_idp.authentik_live import AUTHENTIK_DOCKER_PYTESTMARK - -pytestmark = AUTHENTIK_DOCKER_PYTESTMARK - - -def _jwt_claims(token: str) -> dict[str, object]: - parts = token.split(".") - if len(parts) != 3: - return {} - payload = parts[1] + "=" * (-len(parts[1]) % 4) - return json.loads(base64.urlsafe_b64decode(payload)) - - -def _delete_workspace_for_cleanup(base_url: str, workspace_name: str, headers: dict[str, str]) -> None: - response = httpx.delete( - f"{base_url}/apis/entities/v2/workspaces/{workspace_name}", - headers=headers, - timeout=10.0, - ) - response.raise_for_status() - - -def test_authentik_discovery_exposes_gateway_reachable_device_flow(authentik_stack): - oidc = discover_nmp_config(authentik_stack.gateway_base_url) - - assert oidc.auth_enabled is True - assert oidc.client_id == "nemo-platform-cli" - assert oidc.token_endpoint == "http://127.0.0.1:38080/application/o/token/" - assert oidc.device_authorization_endpoint == "http://127.0.0.1:38080/application/o/device/" - assert oidc.default_scopes == "openid email offline_access groups" - - response = httpx.post( - oidc.device_authorization_endpoint, - data={ - "client_id": oidc.client_id, - "scope": oidc.default_scopes, - }, - timeout=30.0, - ) - response.raise_for_status() - body = response.json() - - assert body["verification_uri"] == "http://127.0.0.1:38080/device" - assert body["verification_uri_complete"].startswith("http://127.0.0.1:38080/device?code=") - assert body["device_code"] - assert body["user_code"] - - -def test_authentik_cli_provider_token_is_accepted_by_gateway(authentik_stack): - token_response = httpx.post( - authentik_stack.token_endpoint, - data={ - "grant_type": "password", - "client_id": "nemo-platform-cli", - "username": "nemo-user", - "password": "nemo-user-token-secret-dev", - "scope": "openid email offline_access groups", - }, - timeout=30.0, - ) - token_response.raise_for_status() - access_token = token_response.json()["access_token"] - claims = _jwt_claims(access_token) - workspace_name = f"cli-audience-check-{uuid.uuid4().hex[:8]}" - headers = {"Authorization": f"Bearer {access_token}"} - - assert claims["aud"] == "nemo-platform-cli" - - try: - create_response = httpx.post( - f"{authentik_stack.gateway_base_url}/apis/entities/v2/workspaces", - json={"name": workspace_name, "description": "CLI audience check"}, - headers=headers, - timeout=10.0, - ) - create_response.raise_for_status() - assert create_response.json()["created_by"] == "nemo-user" - finally: - _delete_workspace_for_cleanup(authentik_stack.gateway_base_url, workspace_name, headers) diff --git a/tests/auth_idp/test_authentik_gateway_live.py b/tests/auth_idp/test_authentik_gateway_live.py deleted file mode 100644 index 7961c3f6ac..0000000000 --- a/tests/auth_idp/test_authentik_gateway_live.py +++ /dev/null @@ -1,94 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import base64 -import json -import uuid - -import httpx -from nmp.testing import grant_workspace_role - -from tests.auth_idp.authentik_live import AUTHENTIK_DOCKER_PYTESTMARK - -pytestmark = AUTHENTIK_DOCKER_PYTESTMARK - - -def _jwt_claims(token: str) -> dict[str, object]: - parts = token.split(".") - if len(parts) != 3: - return {} - payload = parts[1] + "=" * (-len(parts[1]) % 4) - return json.loads(base64.urlsafe_b64decode(payload)) - - -def test_authentik_gateway_rejects_unauthenticated_requests(authentik_stack): - response = httpx.get(f"{authentik_stack.gateway_base_url}/apis/entities/v2/workspaces", timeout=10.0) - assert response.status_code in {401, 403} - - -def test_authentik_gateway_rejects_spoofed_principal_headers(authentik_stack, machine_token: str): - workspace_name = f"spoof-check-{uuid.uuid4().hex[:8]}" - claims = _jwt_claims(machine_token) - authenticated_principal_id = str(claims["sub"]) - expected_binding_principal = authenticated_principal_id - headers = { - "Authorization": f"Bearer {machine_token}", - "X-NMP-Principal-Id": "service:bootstrap", - "X-NMP-Principal-Email": "attacker@example.com", - } - - try: - create_response = httpx.post( - f"{authentik_stack.gateway_base_url}/apis/entities/v2/workspaces", - json={"name": workspace_name, "description": "Spoofed header check"}, - headers=headers, - timeout=10.0, - ) - create_response.raise_for_status() - assert create_response.json()["created_by"] == authenticated_principal_id - - members_response = httpx.get( - f"{authentik_stack.gateway_base_url}/apis/entities/v2/workspaces/{workspace_name}/members", - headers=headers, - timeout=10.0, - ) - members_response.raise_for_status() - admin_member = next(member for member in members_response.json()["data"] if "Admin" in member["roles"]) - - assert admin_member["granted_by"] == authenticated_principal_id - assert admin_member["principal"] == expected_binding_principal - assert admin_member["principal"] not in {"service:bootstrap", "attacker@example.com"} - finally: - httpx.delete( - f"{authentik_stack.gateway_base_url}/apis/entities/v2/workspaces/{workspace_name}", - headers=headers, - timeout=10.0, - ) - - -def test_authentik_gateway_forwards_workload_groups( - authentik_stack, - authentik_human_sdk, - authentik_workspace, - authentik_provider, - machine_token: str, -): - claims = _jwt_claims(machine_token) - claim_groups = claims.get("groups") - assert isinstance(claim_groups, str) - token_groups = {group.strip() for group in claim_groups.split(",") if group.strip()} - bound_group = authentik_provider.workload_expected_groups[0] - assert bound_group in token_groups - - grant_workspace_role(authentik_human_sdk, workspace=authentik_workspace, principal=bound_group, roles=["Viewer"]) - - headers = {"Authorization": f"Bearer {machine_token}"} - - response = httpx.get( - f"{authentik_stack.gateway_base_url}/apis/entities/v2/workspaces/{authentik_workspace}", - headers=headers, - timeout=10.0, - ) - - assert response.status_code == 200 - assert response.json()["name"] == authentik_workspace diff --git a/tests/auth_idp/test_authentik_real_oidc.py b/tests/auth_idp/test_authentik_real_oidc.py deleted file mode 100644 index 6ce4c5c2a2..0000000000 --- a/tests/auth_idp/test_authentik_real_oidc.py +++ /dev/null @@ -1,119 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import os - -import pytest -from nemo_platform import APIStatusError -from nmp.testing import grant_workspace_role -from nmp.testing.e2e import wait_for_job_logs, wait_for_platform_job - -from tests.auth_idp.authentik_live import AUTHENTIK_DOCKER_PYTESTMARK - -pytestmark = AUTHENTIK_DOCKER_PYTESTMARK - - -def _nmp_api_image() -> str: - registry = os.environ.get("IMAGE_REGISTRY", "my-registry") - tag = os.environ.get("BAKE_TAG", "local") - return f"{registry}/nmp-api:{tag}" - - -def test_authentik_workload_token_is_real(machine_token: str, authentik_provider): - assert machine_token - assert authentik_provider.token_endpoint - - -def test_authentik_workload_identity_is_denied_before_binding(machine_sdk, authentik_workspace): - with pytest.raises(APIStatusError) as exc_info: - machine_sdk.workspaces.retrieve(authentik_workspace) - assert exc_info.value.status_code == 403 - - -def test_authentik_workload_identity_is_allowed_after_binding( - authentik_human_sdk, - machine_sdk, - authentik_workspace, - authentik_provider, -): - for group in authentik_provider.workload_expected_groups: - grant_workspace_role(authentik_human_sdk, workspace=authentik_workspace, principal=group, roles=["Viewer"]) - - retrieved = machine_sdk.workspaces.retrieve(authentik_workspace) - assert retrieved.name == authentik_workspace - - -def test_authentik_workload_identity_returns_to_denied_after_revoke( - authentik_human_sdk, - machine_sdk, - authentik_workspace, - authentik_provider, -): - for group in authentik_provider.workload_expected_groups: - grant_workspace_role( - authentik_human_sdk, - workspace=authentik_workspace, - principal=group, - roles=["Viewer"], - ) - authentik_human_sdk.workspaces.members.delete( - group, - workspace=authentik_workspace, - wait_role_propagation=True, - ) - - with pytest.raises(APIStatusError) as exc_info: - machine_sdk.workspaces.retrieve(authentik_workspace) - assert exc_info.value.status_code == 403 - - -def test_authentik_workload_job_runs_via_docker_profile( - authentik_human_sdk, - authentik_workspace, - authentik_provider, - machine_token: str, -): - for group in authentik_provider.workload_expected_groups: - grant_workspace_role( - authentik_human_sdk, - workspace=authentik_workspace, - principal=group, - roles=["Viewer", "JobRunner"], - ) - - job = authentik_human_sdk.jobs.create( - workspace=authentik_workspace, - source="authentik-live-workload-job", - spec={"test": "workload-job"}, - platform_spec={ - "steps": [ - { - "name": "workload-workspace-get", - "executor": { - "provider": "cpu", - "profile": "workload", - "container": { - "image": _nmp_api_image(), - "entrypoint": ["sh", "-c"], - "command": ["nemo-platform run task --task nmp.hello_world.tasks.workload_workspace_get"], - }, - }, - "environment": [ - { - "name": "NEMO_WORKLOAD_TOKEN", - "value": machine_token, - } - ], - "config": { - "workspace": authentik_workspace, - }, - } - ] - }, - ) - - completed_job = wait_for_platform_job(authentik_human_sdk, job.name, authentik_workspace, timeout=240) - assert completed_job.status == "completed" - - step_logs = wait_for_job_logs(authentik_human_sdk, job.name, authentik_workspace, min_log_count=1, timeout=240) - assert any("Successfully retrieved workspace" in log.message for log in step_logs.data) diff --git a/tests/auth_idp/test_authentik_startup_smoke.py b/tests/auth_idp/test_authentik_startup_smoke.py deleted file mode 100644 index fd2aae0fc1..0000000000 --- a/tests/auth_idp/test_authentik_startup_smoke.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import httpx - -from tests.auth_idp.authentik_live import AUTHENTIK_DOCKER_PYTESTMARK - -pytestmark = AUTHENTIK_DOCKER_PYTESTMARK - - -def test_authentik_discovery_is_reachable(authentik_stack): - response = httpx.get(authentik_stack.discovery_url, timeout=10.0) - assert response.status_code == 200 - assert response.json()["issuer"].endswith("/application/o/nemo/") diff --git a/tests/auth_idp/test_fixture_helpers.py b/tests/auth_idp/test_fixture_helpers.py deleted file mode 100644 index 5bf76ed834..0000000000 --- a/tests/auth_idp/test_fixture_helpers.py +++ /dev/null @@ -1,97 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for auth-idp pytest fixtures and fixture-only helper functions.""" - -from pathlib import Path - -import pytest - -from tests.auth_idp import conftest -from tests.auth_idp.conftest import _token_request_body -from tests.auth_idp.providers import ProviderConfig -from tests.auth_idp.xdist import append_xdist_group_suffix - -pytestmark = [pytest.mark.auth_idp] - - -def test_authentik_stack_fixture_uses_pooled_gateway_metadata(): - provider = ProviderConfig( - name="authentik", - mode="compose-ci", - compose_file=Path("docker-compose.yml"), - gateway_base_url="http://127.0.0.1:18080", - issuer_url="http://authentik-server:9000/application/o/nemo/", - discovery_url="http://127.0.0.1:18080/application/o/nemo/.well-known/openid-configuration", - token_endpoint="http://127.0.0.1:18080/application/o/token/", - nemo_config=Path("config/platform-compose-authentik.yaml"), - workload_principal_id="svc-nemo", - workload_expected_groups=["nemo-editors"], - workload_audience="nemo-platform", - workload_principal_claim="sub", - workload_groups_claim="groups", - workload_groups_format="comma_string", - workload_token_env_vars=["NEMO_WORKLOAD_TOKEN", "NEMO_WORKLOAD_TOKEN_FILE"], - workload_forwarded_headers={ - "principal_id": "X-NMP-Principal-Id", - "principal_groups": "X-NMP-Principal-Groups", - }, - human_grant={"grant_type": "password"}, - machine_grant={"grant_type": "password", "username": "svc-nemo", "password": "svc-nemo-token-secret-dev"}, - healthchecks=[], - startup_timeouts={}, - ) - fixture_fn = conftest.authentik_stack.__wrapped__ - stack = fixture_fn(None, provider, "http://127.0.0.1:28080") - - assert stack.gateway_base_url == "http://127.0.0.1:28080" - assert stack.discovery_url == "http://127.0.0.1:28080/application/o/nemo/.well-known/openid-configuration" - assert stack.token_endpoint == "http://127.0.0.1:28080/application/o/token/" - assert stack.nemo_config == provider.nemo_config - - -def test_token_request_body_for_password_grant_includes_username_and_password(): - assert _token_request_body( - { - "grant_type": "password", - "client_id": "nemo-platform", - "client_secret": "secret", - "username": "akadmin", - "password": "akadmin-dev", - "scope": "openid profile email groups", - } - ) == { - "grant_type": "password", - "client_id": "nemo-platform", - "client_secret": "secret", - "username": "akadmin", - "password": "akadmin-dev", - "scope": "openid profile email groups", - } - - -def test_token_request_body_for_workload_password_grant_includes_username_and_password(): - assert _token_request_body( - { - "grant_type": "password", - "client_id": "nemo-platform", - "client_secret": "secret", - "username": "svc-nemo", - "password": "svc-nemo-token-secret-dev", - "scope": "openid email groups", - } - ) == { - "grant_type": "password", - "client_id": "nemo-platform", - "client_secret": "secret", - "username": "svc-nemo", - "password": "svc-nemo-token-secret-dev", - "scope": "openid email groups", - } - - -def test_append_xdist_group_suffix_only_appends_once_and_sorts_groups(): - nodeid = "tests/auth_idp/test_authentik_real_oidc.py::test_authentik_machine_token_is_real" - assert append_xdist_group_suffix(nodeid, {"idp-live"}) == f"{nodeid}@idp-live" - assert append_xdist_group_suffix(nodeid, {"b", "a"}) == f"{nodeid}@a_b" - assert append_xdist_group_suffix(f"{nodeid}@idp-live", {"idp-live"}) == f"{nodeid}@idp-live" diff --git a/tests/auth_idp/test_provider_layout.py b/tests/auth_idp/test_provider_layout.py deleted file mode 100644 index 573998d81d..0000000000 --- a/tests/auth_idp/test_provider_layout.py +++ /dev/null @@ -1,37 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from pathlib import Path - -import pytest -import yaml - -from tests.auth_idp.providers import load_provider_names_by_mode - -pytestmark = [pytest.mark.auth_idp] - - -def test_compose_backed_providers_ship_required_assets(): - for provider in load_provider_names_by_mode("compose-ci"): - root = Path(f"contrib/auth/{provider}") - assert (root / "docker-compose.yml").exists() - assert (root / "gateway").exists() - assert (root / "README.md").exists() - assert (root / "manifest.yaml").exists() - - -def test_reference_only_providers_do_not_require_compose(): - for provider in load_provider_names_by_mode("reference-only"): - root = Path(f"contrib/auth/{provider}") - assert (root / "README.md").exists() - assert not (root / "docker-compose.yml").exists() - - -def test_authentik_compose_disables_model_provider_seed_without_ngc_key(): - compose = yaml.safe_load(Path("contrib/auth/authentik/docker-compose.yml").read_text()) - nemo_service = compose["services"]["nemo"] - nemo_env = nemo_service["environment"] - - assert nemo_env["NMP_SEED_ON_STARTUP"] == "true" - assert nemo_env["NMP_PLATFORM_SEED_MODEL_PROVIDER_ENABLED"] == "false" - assert "${NEMO_DIRECT_PORT:-18081}:8080" in nemo_service["ports"] diff --git a/tests/auth_idp/test_provider_manifest.py b/tests/auth_idp/test_provider_manifest.py deleted file mode 100644 index f3fc1d50b2..0000000000 --- a/tests/auth_idp/test_provider_manifest.py +++ /dev/null @@ -1,76 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from pathlib import Path - -import pytest -import yaml -from jsonschema.validators import validator_for - -from tests.auth_idp.providers import load_provider_configs - -pytestmark = [pytest.mark.auth_idp] - - -def _load_provider_manifest_schema() -> dict: - return yaml.safe_load(Path("contrib/auth/manifest.schema.yaml").read_text()) - - -def test_all_provider_manifests_share_the_same_contract(): - schema = _load_provider_manifest_schema() - validator = validator_for(schema)(schema) - for provider in load_provider_configs(): - manifest = yaml.safe_load(Path(f"contrib/auth/{provider.name}/manifest.yaml").read_text()) - validator.validate(manifest) - assert manifest["provider"] == provider.name - - -def test_authentik_manifest_declares_real_token_acquisition_contract(): - manifest = yaml.safe_load(Path("contrib/auth/authentik/manifest.yaml").read_text()) - token_acquisition = manifest["token_acquisition"] - principal_contract = manifest["principal_contract"] - workload_identity = manifest["workload_identity"] - workload_contract = manifest["workload_contract"] - - assert token_acquisition["token_endpoint"] - assert token_acquisition["human_grant"]["grant_type"] == "password" - assert token_acquisition["machine_grant"]["grant_type"] == "password" - assert token_acquisition["human_grant"]["client_id"] - assert token_acquisition["machine_grant"]["client_id"] - assert token_acquisition["human_grant"]["password"] == "nemo-user-token-secret-dev" - assert "offline_access" in token_acquisition["human_grant"]["scope"].split() - assert workload_identity["principal_id"] - assert not workload_identity["principal_id"].startswith(principal_contract["internal_service_prefix_reserved"]) - assert workload_identity["expected_groups"] - assert workload_contract["audience"] == "nemo-platform" - assert workload_contract["groups_format"] == "comma_string" - assert workload_contract["forwarded_headers"]["principal_id"] == "X-NMP-Principal-Id" - assert workload_contract["forwarded_headers"]["principal_groups"] == "X-NMP-Principal-Groups" - - -def test_authentik_manifest_declares_extended_startup_timeouts_for_real_oidc(): - manifest = yaml.safe_load(Path("contrib/auth/authentik/manifest.yaml").read_text()) - startup_timeouts = manifest["startup_timeouts"] - - assert startup_timeouts["healthchecks_seconds"] >= 240 - assert startup_timeouts["gateway_seconds"] >= 30 - assert startup_timeouts["token_endpoint_seconds"] >= 60 - - -def test_authentik_provider_config_loads_token_acquisition_fields(): - provider = next(config for config in load_provider_configs() if config.name == "authentik") - - assert provider.nemo_config == Path("contrib/auth/authentik/config/platform-compose-authentik.yaml") - assert provider.token_endpoint == "http://127.0.0.1:18080/application/o/token/" - assert provider.human_grant["grant_type"] == "password" - assert provider.machine_grant["grant_type"] == "password" - assert provider.workload_audience == "nemo-platform" - assert provider.workload_principal_claim == "sub" - assert provider.workload_groups_claim == "groups" - assert provider.workload_groups_format == "comma_string" - assert provider.workload_token_env_vars == ["NEMO_WORKLOAD_TOKEN", "NEMO_WORKLOAD_TOKEN_FILE"] - assert provider.startup_timeouts == { - "healthchecks_seconds": 600, - "gateway_seconds": 30, - "token_endpoint_seconds": 180, - } diff --git a/tests/test_e2e_docker_compose_backend.py b/tests/test_e2e_docker_compose_backend.py index 55237fda9b..899fa023f5 100644 --- a/tests/test_e2e_docker_compose_backend.py +++ b/tests/test_e2e_docker_compose_backend.py @@ -7,10 +7,10 @@ import pytest -from e2e.backends.docker_compose import DockerComposeE2EBackend +from e2e.backends.docker_compose import DockerComposeE2EBackend, _compose_stack_readiness -def _compose_ps_json(*entries: dict[str, str]) -> str: +def _compose_ps_json(*entries: dict[str, object]) -> str: return json.dumps(list(entries)) @@ -61,6 +61,114 @@ class Response: assert first_env["NEMO_COMPOSE_CONFIG_PATH"] == str(config_path.resolve()) +def test_compose_stack_readiness_accepts_completed_init_services() -> None: + ready, not_ready = _compose_stack_readiness( + [ + {"Service": "gateway-tls-init", "State": "exited", "ExitCode": 0}, + {"Service": "gateway", "State": "running", "Health": "healthy"}, + ], + {"gateway-tls-init", "gateway"}, + ) + + assert ready is True + assert not_ready == [] + + +def test_compose_stack_readiness_rejects_failed_init_services() -> None: + ready, not_ready = _compose_stack_readiness( + [{"Service": "gateway-tls-init", "State": "exited", "ExitCode": 1}], + {"gateway-tls-init"}, + ) + + assert ready is False + assert not_ready == ["gateway-tls-init (state=exited)"] + + +def test_compose_stack_readiness_rejects_exited_non_init_services() -> None: + ready, not_ready = _compose_stack_readiness( + [{"Service": "gateway", "State": "exited", "ExitCode": 0}], + {"gateway"}, + ) + + assert ready is False + assert not_ready == ["gateway (state=exited)"] + + +def test_compose_backend_uses_env_ca_bundle_for_ready_probe(monkeypatch, tmp_path: Path) -> None: + compose_file = tmp_path / "docker-compose.yml" + compose_file.write_text("services: {}\n") + config_path = tmp_path / "platform.yaml" + config_path.write_text("platform: {}\n") + ca_bundle = tmp_path / "gateway-ca.crt" + ca_bundle.write_text("test ca\n") + verify_values = [] + + class Response: + status_code = 200 + + def fake_get(*args, **kwargs): + verify_values.append(kwargs.get("verify")) + return Response() + + monkeypatch.setattr("e2e.backends.docker_compose.httpx.get", fake_get) + + backend = DockerComposeE2EBackend( + compose_file=compose_file, + config_path=config_path, + project_name="authentik-e2e-test", + service_url="https://127.0.0.1:38080", + wait_url="https://127.0.0.1:38080/apis/auth/discovery", + env={"REQUESTS_CA_BUNDLE": str(ca_bundle)}, + ) + + backend._wait_ready() + + assert verify_values == [str(ca_bundle)] + + +def test_compose_backend_waits_for_all_ready_urls(monkeypatch, tmp_path: Path) -> None: + compose_file = tmp_path / "docker-compose.yml" + compose_file.write_text("services: {}\n") + config_path = tmp_path / "platform.yaml" + config_path.write_text("platform: {}\n") + calls: list[str] = [] + discovery_attempts = 0 + + class Response: + def __init__(self, status_code: int) -> None: + self.status_code = status_code + + def fake_get(url, **_kwargs): + nonlocal discovery_attempts + calls.append(url) + if url.endswith("/health/ready"): + return Response(200) + discovery_attempts += 1 + return Response(200 if discovery_attempts == 2 else 404) + + monkeypatch.setattr("e2e.backends.docker_compose.httpx.get", fake_get) + monkeypatch.setattr("e2e.backends.docker_compose.time.sleep", lambda _seconds: None) + + backend = DockerComposeE2EBackend( + compose_file=compose_file, + config_path=config_path, + project_name="authentik-e2e-test", + service_url="https://127.0.0.1:38080", + wait_urls=[ + "https://127.0.0.1:38080/health/ready", + "https://127.0.0.1:38080/application/o/nemo/.well-known/openid-configuration", + ], + ) + + backend._wait_ready() + + assert calls == [ + "https://127.0.0.1:38080/health/ready", + "https://127.0.0.1:38080/application/o/nemo/.well-known/openid-configuration", + "https://127.0.0.1:38080/application/o/nemo/.well-known/openid-configuration", + ] + + def test_compose_backend_stop_uses_same_project_and_env(monkeypatch, tmp_path: Path) -> None: calls: list[tuple[list[str], dict[str, str] | None]] = [] compose_file = tmp_path / "docker-compose.yml" @@ -103,7 +211,16 @@ def test_compose_backend_write_logs_uses_same_project_and_env(monkeypatch, tmp_p def fake_run(args, *, check, text=False, stdout=None, stderr=None, env=None, **_kwargs): calls.append((list(args), env)) assert stdout is not None - stdout.write("nemo log line\n") + if args[6:] == ["config", "--services"]: + stdout.write("nemo\ngateway\n") + elif args[6:] == ["ps", "--all", "--format", "json"]: + stdout.write('[{"Service":"nemo","State":"running"}]\n') + elif args[6:] == ["ps", "--all"]: + stdout.write("NAME STATE\nnemo running\n") + elif args[6:] == ["logs", "--no-color", "--timestamps"]: + stdout.write("nemo log line\n") + else: + raise AssertionError(f"unexpected compose diagnostics command: {args[6:]}") return subprocess.CompletedProcess(args=args, returncode=0, stdout="", stderr="") monkeypatch.setattr("e2e.backends.docker_compose.subprocess.run", fake_run) @@ -118,13 +235,64 @@ def fake_run(args, *, check, text=False, stdout=None, stderr=None, env=None, **_ backend.write_logs(log_path) - args, env = calls[0] - assert args[:6] == ["docker", "compose", "-f", str(compose_file), "-p", "authentik-e2e-test"] - assert args[6:] == ["logs", "--no-color", "--timestamps"] - assert env is not None - assert env["AUTHENTIK_GATEWAY_PORT"] == "38080" - assert env["NEMO_COMPOSE_CONFIG_PATH"] == str(config_path.resolve()) - assert log_path.read_text(encoding="utf-8") == "nemo log line\n" + assert [args[6:] for args, _env in calls] == [ + ["config", "--services"], + ["ps", "--all", "--format", "json"], + ["ps", "--all"], + ["logs", "--no-color", "--timestamps"], + ] + for args, env in calls: + assert args[:6] == ["docker", "compose", "-f", str(compose_file), "-p", "authentik-e2e-test"] + assert env is not None + assert env["AUTHENTIK_GATEWAY_PORT"] == "38080" + assert env["NEMO_COMPOSE_CONFIG_PATH"] == str(config_path.resolve()) + log_text = log_path.read_text(encoding="utf-8") + assert "===== docker compose config --services =====" in log_text + assert "===== docker compose ps --all --format json =====" in log_text + assert "===== docker compose ps --all =====" in log_text + assert "===== docker compose logs --no-color --timestamps =====" in log_text + assert "nemo log line" in log_text + + +def test_compose_backend_write_logs_continues_after_diagnostic_timeout(monkeypatch, tmp_path: Path) -> None: + calls: list[list[str]] = [] + compose_file = tmp_path / "docker-compose.yml" + compose_file.write_text("services: {}\n") + config_path = tmp_path / "platform.yaml" + config_path.write_text("platform: {}\n") + log_path = tmp_path / "compose.log" + + def fake_run(args, *, check, text=False, stdout=None, stderr=None, env=None, timeout=None, **_kwargs): + calls.append(list(args)) + assert timeout is not None + if args[6:] == ["config", "--services"]: + raise subprocess.TimeoutExpired(cmd=args, timeout=timeout) + assert stdout is not None + stdout.write("diagnostic output\n") + return subprocess.CompletedProcess(args=args, returncode=0, stdout="", stderr="") + + monkeypatch.setattr("e2e.backends.docker_compose.subprocess.run", fake_run) + + backend = DockerComposeE2EBackend( + compose_file=compose_file, + config_path=config_path, + project_name="authentik-e2e-test", + service_url="http://127.0.0.1:38080", + ) + + backend.write_logs(log_path) + + assert [args[6:] for args in calls] == [ + ["config", "--services"], + ["ps", "--all", "--format", "json"], + ["ps", "--all"], + ["logs", "--no-color", "--timestamps"], + ] + log_text = log_path.read_text(encoding="utf-8") + assert "===== docker compose config --services =====" in log_text + assert "[command timed out after 60s]" in log_text + assert "===== docker compose logs --no-color --timestamps =====" in log_text + assert "diagnostic output" in log_text def test_compose_backend_reuse_mode_reuses_healthy_stack_without_restart(monkeypatch, tmp_path: Path) -> None: diff --git a/tests/test_e2e_services_pool.py b/tests/test_e2e_services_pool.py index 1b2b8f019a..f3af52ace1 100644 --- a/tests/test_e2e_services_pool.py +++ b/tests/test_e2e_services_pool.py @@ -1,6 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from types import SimpleNamespace + +import pytest import e2e.services_pool as services_pool @@ -90,14 +93,16 @@ def test_render_e2e_config_for_docker_compose_preserves_container_paths(tmp_path assert rendered["files"]["default_storage_config"]["path"] == "/data/files" -def test_start_services_docker_compose_waits_for_auth_ready_when_enabled(tmp_path, monkeypatch) -> None: +def test_start_services_docker_compose_uses_auth_ready_url_override(tmp_path, monkeypatch) -> None: compose_file = tmp_path / "compose.yaml" compose_file.write_text("services: {}\n", encoding="utf-8") config_path = tmp_path / "config.yaml" config_path.write_text("{}\n", encoding="utf-8") + captured_kwargs = {} class FakeDockerComposeBackend: def __init__(self, **kwargs) -> None: + captured_kwargs.update(kwargs) self.service_url = kwargs["service_url"] def start(self) -> None: @@ -109,11 +114,8 @@ def stop(self) -> None: def write_logs(self, log_path) -> None: log_path.write_text("compose logs\n", encoding="utf-8") - wait_calls = [] - def fake_wait_for_auth_ready(url, proc) -> bool: - wait_calls.append((url, proc)) - return True + raise AssertionError("docker compose auth readiness should use DockerComposeE2EBackend.wait_url") monkeypatch.setattr(services_pool, "DockerComposeE2EBackend", FakeDockerComposeBackend) monkeypatch.setattr(services_pool, "_wait_for_auth_ready", fake_wait_for_auth_ready) @@ -124,28 +126,112 @@ def fake_wait_for_auth_ready(url, proc) -> bool: { "backend": "docker_compose", "compose_file": str(compose_file), - "service_url": "http://127.0.0.1:8080", + "service_url": "http://127.0.0.1:38080", + "auth_ready_url": "${service_url}/health/gateway/ready", + "wait_url": "http://127.0.0.1:38080/health/ready", "lifecycle": "fresh", }, "abc123", tmp_path / "services.log", ) - assert wait_calls == [("http://127.0.0.1:8080", None)] + assert captured_kwargs["wait_url"] == "http://127.0.0.1:38080/health/gateway/ready" + assert captured_kwargs["wait_urls"] == ["http://127.0.0.1:38080/health/gateway/ready"] assert services.auth_enabled is True - assert services.url == "http://127.0.0.1:8080" + assert services.url == "http://127.0.0.1:38080" + + +def test_auth_ready_url_override_uses_url_probe_and_skips_default_probe(monkeypatch) -> None: + calls = [] + def fake_wait_for_auth_ready_url(url, proc, *, env=None): + calls.append(("url", url, proc, env)) + return True + + def fake_wait_for_auth_ready(url, proc): + raise AssertionError("default auth probe should not run when auth_ready_url is configured") + + monkeypatch.setattr(services_pool, "_wait_for_auth_ready_url", fake_wait_for_auth_ready_url) + monkeypatch.setattr(services_pool, "_wait_for_auth_ready", fake_wait_for_auth_ready) -def test_start_services_docker_compose_uses_auth_ready_url_when_configured(tmp_path, monkeypatch) -> None: + assert services_pool._wait_for_configured_auth_ready( + "http://127.0.0.1:38080", + None, + { + "backend": "subprocess", + "auth_ready_url": "${service_url}/health/gateway/ready", + "env": {"NMP_CLIENT_SSL_CERT_FILE": "/tmp/ca.crt"}, + }, + ) + + assert calls == [ + ( + "url", + "http://127.0.0.1:38080/health/gateway/ready", + None, + {"NMP_CLIENT_SSL_CERT_FILE": "/tmp/ca.crt"}, + ) + ] + + +def test_start_services_docker_compose_exposes_log_path_and_captures_logs_on_close(tmp_path, monkeypatch) -> None: compose_file = tmp_path / "compose.yaml" compose_file.write_text("services: {}\n", encoding="utf-8") config_path = tmp_path / "config.yaml" config_path.write_text("{}\n", encoding="utf-8") + log_path = tmp_path / "services.log" + events = [] class FakeDockerComposeBackend: def __init__(self, **kwargs) -> None: self.service_url = kwargs["service_url"] + def start(self) -> None: + events.append("start") + + def stop(self) -> None: + events.append("stop") + + def write_logs(self, path) -> None: + events.append(("logs", path)) + path.write_text("compose logs\n", encoding="utf-8") + + monkeypatch.setattr(services_pool, "DockerComposeE2EBackend", FakeDockerComposeBackend) + + services = services_pool._start_services_docker_compose( + config_path, + {}, + { + "backend": "docker_compose", + "compose_file": str(compose_file), + "service_url": "http://127.0.0.1:38080", + "lifecycle": "fresh", + }, + "abc123", + log_path, + ) + + assert services.log_path == log_path + assert services.close is not None + + services.close() + + assert events == ["start", ("logs", log_path), "stop"] + assert log_path.read_text(encoding="utf-8") == "compose logs\n" + + +def test_start_services_docker_compose_resolves_dynamic_port_templates(tmp_path, monkeypatch) -> None: + compose_file = tmp_path / "compose.yaml" + compose_file.write_text("services: {}\n", encoding="utf-8") + config_path = tmp_path / "config.yaml" + config_path.write_text("stale: true\n", encoding="utf-8") + captured_kwargs = {} + + class FakeDockerComposeBackend: + def __init__(self, **kwargs) -> None: + captured_kwargs.update(kwargs) + self.service_url = kwargs["service_url"] + def start(self) -> None: return None @@ -155,75 +241,250 @@ def stop(self) -> None: def write_logs(self, log_path) -> None: log_path.write_text("compose logs\n", encoding="utf-8") - wait_calls = [] - - def fake_wait_for_auth_ready(url, proc) -> bool: - wait_calls.append((url, proc)) - return True - monkeypatch.setattr(services_pool, "DockerComposeE2EBackend", FakeDockerComposeBackend) - monkeypatch.setattr(services_pool, "_wait_for_auth_ready", fake_wait_for_auth_ready) + monkeypatch.setattr(services_pool, "_find_free_port", lambda: 49123) + monkeypatch.setattr(services_pool.uuid, "uuid4", lambda: SimpleNamespace(hex="deadbeefcafebabe")) services = services_pool._start_services_docker_compose( config_path, - {"auth": {"enabled": True}}, + { + "auth": { + "enabled": True, + "oidc": { + "token_endpoint": "${gateway_url}/application/o/token/", + }, + }, + "jobs": { + "executors": [ + { + "config": { + "storage": { + "additional_volume_mounts": [ + { + "volume_name": "authentik-gateway-tls-${gateway_port}", + "mount_path": "/etc/nmp/gateway-tls", + } + ] + } + } + } + ] + }, + }, { "backend": "docker_compose", "compose_file": str(compose_file), - "service_url": "http://127.0.0.1:38080", - "auth_ready_url": "http://127.0.0.1:38081", + "compose_project_prefix": "authentik-e2e", + "dynamic_ports": { + "gateway": { + "host": "127.0.0.1", + "scheme": "https", + } + }, + "service_url": "${gateway_url}", + "auth_ready_url": "${gateway_url}/health/gateway/ready", + "env": { + "AUTHENTIK_GATEWAY_PORT": "${gateway_port}", + "AUTHENTIK_GATEWAY_TLS_VOLUME": "authentik-gateway-tls-${gateway_port}", + }, "lifecycle": "fresh", }, "abc123", tmp_path / "services.log", ) - assert wait_calls == [("http://127.0.0.1:38081", None)] - assert services.auth_enabled is True - assert services.url == "http://127.0.0.1:38080" + assert captured_kwargs["project_name"] == "authentik-e2e-abc123-deadbeef" + assert captured_kwargs["service_url"] == "https://127.0.0.1:49123" + assert captured_kwargs["wait_url"] == "https://127.0.0.1:49123/health/gateway/ready" + assert captured_kwargs["wait_urls"] == ["https://127.0.0.1:49123/health/gateway/ready"] + assert "dynamic_ports" not in captured_kwargs + assert captured_kwargs["env"]["AUTHENTIK_GATEWAY_PORT"] == "49123" + assert captured_kwargs["env"]["AUTHENTIK_GATEWAY_TLS_VOLUME"] == "authentik-gateway-tls-49123" + rendered_config = services_pool.yaml.safe_load(config_path.read_text(encoding="utf-8")) + assert rendered_config["auth"]["oidc"]["token_endpoint"] == "https://127.0.0.1:49123/application/o/token/" + assert ( + rendered_config["jobs"]["executors"][0]["config"]["storage"]["additional_volume_mounts"][0]["volume_name"] + == "authentik-gateway-tls-49123" + ) + assert services.url == "https://127.0.0.1:49123" + assert services.compose_project_name == "authentik-e2e-abc123-deadbeef" -def test_start_services_docker_compose_exposes_log_path_and_captures_logs_on_close(tmp_path, monkeypatch) -> None: +def test_start_services_docker_compose_rejects_unfixed_dynamic_ports_with_reuse(tmp_path) -> None: compose_file = tmp_path / "compose.yaml" compose_file.write_text("services: {}\n", encoding="utf-8") config_path = tmp_path / "config.yaml" config_path.write_text("{}\n", encoding="utf-8") - log_path = tmp_path / "services.log" - events = [] + + with pytest.raises(pytest.UsageError, match="dynamic_ports require explicit ports"): + services_pool._start_services_docker_compose( + config_path, + {}, + { + "backend": "docker_compose", + "compose_file": str(compose_file), + "dynamic_ports": { + "gateway": { + "host": "127.0.0.1", + "scheme": "https", + } + }, + "service_url": "${gateway_url}", + "lifecycle": "reuse", + }, + "abc123", + tmp_path / "services.log", + ) + + +def test_start_services_docker_compose_allows_fixed_dynamic_ports_with_reuse(tmp_path, monkeypatch) -> None: + compose_file = tmp_path / "compose.yaml" + compose_file.write_text("services: {}\n", encoding="utf-8") + config_path = tmp_path / "config.yaml" + config_path.write_text("stale: true\n", encoding="utf-8") + captured_kwargs = {} class FakeDockerComposeBackend: def __init__(self, **kwargs) -> None: + captured_kwargs.update(kwargs) self.service_url = kwargs["service_url"] def start(self) -> None: - events.append("start") + return None def stop(self) -> None: - events.append("stop") + return None - def write_logs(self, path) -> None: - events.append(("logs", path)) - path.write_text("compose logs\n", encoding="utf-8") + def write_logs(self, log_path) -> None: + log_path.write_text("compose logs\n", encoding="utf-8") monkeypatch.setattr(services_pool, "DockerComposeE2EBackend", FakeDockerComposeBackend) services = services_pool._start_services_docker_compose( config_path, - {}, + { + "auth": { + "enabled": True, + "oidc": { + "token_endpoint": "${gateway_url}/application/o/token/", + }, + }, + }, { "backend": "docker_compose", "compose_file": str(compose_file), - "service_url": "http://127.0.0.1:38080", - "lifecycle": "fresh", + "compose_project_name": "authentik-e2e-reuse", + "dynamic_ports": { + "gateway": { + "host": "127.0.0.1", + "port": "18080", + "scheme": "https", + } + }, + "service_url": "${gateway_url}", + "auth_ready_url": "${gateway_url}/health/gateway/ready", + "env": { + "AUTHENTIK_GATEWAY_PORT": "${gateway_port}", + }, + "lifecycle": "reuse", }, "abc123", - log_path, + tmp_path / "services.log", ) - assert services.log_path == log_path - assert services.close is not None + assert captured_kwargs["project_name"] == "authentik-e2e-reuse" + assert captured_kwargs["service_url"] == "https://127.0.0.1:18080" + assert captured_kwargs["wait_url"] == "https://127.0.0.1:18080/health/gateway/ready" + assert captured_kwargs["wait_urls"] == ["https://127.0.0.1:18080/health/gateway/ready"] + assert captured_kwargs["env"]["AUTHENTIK_GATEWAY_PORT"] == "18080" + rendered_config = services_pool.yaml.safe_load(config_path.read_text(encoding="utf-8")) + assert rendered_config["auth"]["oidc"]["token_endpoint"] == "https://127.0.0.1:18080/application/o/token/" + assert services.compose_project_name == "authentik-e2e-reuse" + + +def test_release_for_module_clears_active_binding_while_shared_service_remains(tmp_path) -> None: + pool = services_pool.E2EServicesPool() + key = services_pool.ServicesPoolKey(config_hash="shared") + module_a = SimpleNamespace(nodeid="tests/test_a.py") + module_b = SimpleNamespace(nodeid="tests/test_b.py") + + for module in (module_a, module_b): + pool._module_states[module.nodeid] = services_pool.ModuleConfigState( + module_id=module.nodeid, + key=key, + config_path=None, + config_data={}, + harness_config={"backend": "subprocess"}, + config_layers=(), + auth_enabled=False, + ) + pool._active_service_key_by_module[module.nodeid] = key + + pool._remaining_modules_by_key[key] = {module_a.nodeid, module_b.nodeid} + pool._running_by_key[key] = services_pool.RunningServices( + url="http://127.0.0.1:8080", + log_path=tmp_path / "services.log", + proc=None, + config_path=None, + key=key, + ) - services.close() + pool.release_for_module(module_a) + + assert pool.describe_active_module_binding(module_a.nodeid) is None + assert pool.describe_active_module_binding(module_b.nodeid)["service_url"] == "http://127.0.0.1:8080" + + +def test_acquire_for_module_reregisters_released_owner_before_next_release(tmp_path, monkeypatch) -> None: + pool = services_pool.E2EServicesPool() + key = services_pool.ServicesPoolKey(config_hash="shared") + module_a = SimpleNamespace(nodeid="tests/test_a.py") + module_b = SimpleNamespace(nodeid="tests/test_b.py") + terminated = [] + + for module in (module_a, module_b): + pool._module_states[module.nodeid] = services_pool.ModuleConfigState( + module_id=module.nodeid, + key=key, + config_path=tmp_path / "platform.yaml", + config_data={}, + harness_config={"backend": "subprocess"}, + config_layers=(), + auth_enabled=False, + ) + + running_services = services_pool.RunningServices( + url="http://127.0.0.1:8080", + log_path=tmp_path / "services.log", + proc=None, + config_path=tmp_path / "platform.yaml", + key=key, + ) + pool._remaining_modules_by_key[key] = {module_a.nodeid, module_b.nodeid} + pool._running_by_key[key] = running_services + monkeypatch.setattr(pool, "_terminate_services", terminated.append) + + pool.release_for_module(module_a) + reacquired = pool.acquire_for_module(module_a) + pool.release_for_module(module_b) + + assert reacquired is running_services + assert terminated == [] + assert pool.describe_active_module_binding(module_a.nodeid)["service_url"] == "http://127.0.0.1:8080" + + +def test_describe_active_module_binding_ignores_stale_active_key() -> None: + pool = services_pool.E2EServicesPool() + key = services_pool.ServicesPoolKey(config_hash="shared") + module_id = "tests/test_released.py" + pool._active_service_key_by_module[module_id] = key + pool._remaining_modules_by_key[key] = {"tests/test_other.py"} + pool._running_by_key[key] = services_pool.RunningServices( + url="http://127.0.0.1:8080", + log_path=None, + proc=None, + config_path=None, + key=key, + ) - assert events == ["start", ("logs", log_path), "stop"] - assert log_path.read_text(encoding="utf-8") == "compose logs\n" + assert pool.describe_active_module_binding(module_id) is None diff --git a/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/core/openapi.py b/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/core/openapi.py index a16d618a00..8429d72255 100644 --- a/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/core/openapi.py +++ b/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/core/openapi.py @@ -12,6 +12,14 @@ # https://www.stainless.com/docs/reference/diagnostics/#endpoint shows the supported HTTP methods, this is created explicitly to avoid HEAD endpoints as they are not supported by Stainless. SUPPORTED_HTTP_METHODS = {"get", "post", "put", "patch", "delete", "query"} +SDK_EXCLUDED_PATHS = { + # Auth bootstrap and utility endpoints do not follow the generated SDK's + # versioned resource path convention. They are used directly by client auth + # code instead of through generated resource methods. + "/apis/auth/discovery", + "/apis/auth/jwks", + "/apis/auth/token", +} logger = logging.getLogger(__name__) @@ -128,6 +136,9 @@ def calculate_schema_to_endpoints(self) -> dict[str, list[OpenAPIEndpoint]]: # Analyze paths for path, path_item in self._spec.get("paths", {}).items(): + if _should_skip_sdk_endpoint(path): + continue + for method, spec in path_item.items(): method_lower = method.lower() if method_lower not in SUPPORTED_HTTP_METHODS: @@ -199,10 +210,11 @@ def extract_endpoints(self) -> Iterable[OpenAPIEndpoint]: continue if path.startswith("/v1/jobs"): continue - # Skip discovery endpoints — they don't follow the - # versioned /v1/ or /v2/ path convention and are accessed - # directly by the CLI, not through the generated SDK. - if path.startswith("/apis/auth/discovery"): + if _should_skip_sdk_endpoint(path): continue yield OpenAPIEndpoint(method_lower, path) + + +def _should_skip_sdk_endpoint(path: str) -> bool: + return path.lower() in SDK_EXCLUDED_PATHS diff --git a/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/vendor/vendor_package.py b/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/vendor/vendor_package.py index a1edd2945e..64eb931554 100644 --- a/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/vendor/vendor_package.py +++ b/tools/nemo-platform-sdk-tools/src/nemo_platform_sdk_tools/sdk/vendor/vendor_package.py @@ -31,6 +31,22 @@ GENERATED_BUNDLE_GROUP_COMMENT = "# Generated from [tool.bundle-package]; do not edit by hand." GENERATED_PROJECT_COMMENTS = {GENERATED_BUNDLE_GROUP_COMMENT, GENERATED_BUNDLE_TABLE_COMMENT} VALID_BUNDLE_INHERIT_VALUES = {"entry-points", "optional-dependencies", "scripts"} +GENERATED_EMPTY_INIT_HEADER = """# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" def _setup_logging() -> None: @@ -65,7 +81,11 @@ class ResourceReplacement(BaseModel): _CLIENT_CLASS_NAMES = ("NeMoPlatform", "AsyncNeMoPlatform") _CLIENT_METHOD_NAMES = ("__init__", "__getattr__") +_CLIENT_HELPER_FUNCTION_NAMES = ("_should_bootstrap_config",) _CLIENT_INIT_REQUIRED_IMPORTS: dict[str, tuple[str, ...]] = { + "nemo_platform._base_client": ("DefaultAsyncHttpxClient", "DefaultHttpxClient"), + "nemo_platform.client.tls": ("client_verify_from_env",), + "nemo_platform_plugin.client.constants": ("WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR",), "pathlib": ("Path",), } @@ -117,6 +137,49 @@ def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef return updated_node.with_changes(body=updated_node.body.with_changes(body=body_items)) +def _collect_client_helper_functions(source_module: cst.Module) -> dict[str, cst.FunctionDef]: + helpers: dict[str, cst.FunctionDef] = {} + for statement in source_module.body: + if isinstance(statement, cst.FunctionDef) and statement.name.value in _CLIENT_HELPER_FUNCTION_NAMES: + helpers[statement.name.value] = statement + return helpers + + +def _find_client_helper_insertion_index(module: cst.Module) -> int: + for index, statement in enumerate(module.body): + if not isinstance(statement, cst.SimpleStatementLine): + continue + for body_statement in statement.body: + if not isinstance(body_statement, cst.Assign): + continue + for target in body_statement.targets: + if isinstance(target.target, cst.Name) and target.target.value == "__all__": + return index + 1 + + return _find_import_insertion_index(module) + + +def _with_client_helper_functions( + target_module: cst.Module, + helper_functions: dict[str, cst.FunctionDef], +) -> cst.Module: + if not helper_functions: + return target_module + + helper_names = set(helper_functions) + body_without_stale_helpers = [ + statement + for statement in target_module.body + if not (isinstance(statement, cst.FunctionDef) and statement.name.value in helper_names) + ] + target_module = target_module.with_changes(body=body_without_stale_helpers) + + helpers = [helper_functions[name] for name in _CLIENT_HELPER_FUNCTION_NAMES if name in helper_functions] + insert_at = _find_client_helper_insertion_index(target_module) + body = list(target_module.body) + return target_module.with_changes(body=body[:insert_at] + helpers + body[insert_at:]) + + def _module_name_from_expr(module_expr: cst.BaseExpression | None) -> str | None: if module_expr is None: return None @@ -1379,6 +1442,13 @@ def _get_transitive_source_module_name(dependency_name: str) -> str: def _copy_included_paths(source_path: Path, destination_path: Path, included_paths: list[str]) -> None: """Copy included paths from source to destination, handling glob patterns.""" + + def copy_file(source_file: Path, dest_file: Path) -> None: + if source_file.name == "__init__.py" and source_file.read_text(encoding="utf-8") == "": + dest_file.write_text(GENERATED_EMPTY_INIT_HEADER, encoding="utf-8") + else: + shutil.copy(source_file, dest_file) + for pattern in included_paths: pattern = str(pattern) pattern_path = source_path / pattern @@ -1388,7 +1458,7 @@ def _copy_included_paths(source_path: Path, destination_path: Path, included_pat dest_file = destination_path / relative_file dest_file.parent.mkdir(parents=True, exist_ok=True) logger.debug(f"Copying {relative_file}") - shutil.copy(pattern_path, dest_file) + copy_file(pattern_path, dest_file) elif pattern_path.is_dir(): # If it's a directory, find all Python files in it for py_file in pattern_path.rglob("*.py"): @@ -1396,7 +1466,7 @@ def _copy_included_paths(source_path: Path, destination_path: Path, included_pat dest_file = destination_path / relative_file dest_file.parent.mkdir(parents=True, exist_ok=True) logger.debug(f"Copying {relative_file}") - shutil.copy(py_file, dest_file) + copy_file(py_file, dest_file) else: # Treat as glob pattern for matched_file in source_path.glob(pattern): @@ -1405,7 +1475,7 @@ def _copy_included_paths(source_path: Path, destination_path: Path, included_pat dest_file = destination_path / relative_file dest_file.parent.mkdir(parents=True, exist_ok=True) logger.debug(f"Copying {relative_file}") - shutil.copy(matched_file, dest_file) + copy_file(matched_file, dest_file) elif matched_file.is_dir(): # If glob matches a directory, find all Python files in it for py_file in matched_file.rglob("*.py"): @@ -1413,7 +1483,7 @@ def _copy_included_paths(source_path: Path, destination_path: Path, included_pat dest_file = destination_path / relative_file dest_file.parent.mkdir(parents=True, exist_ok=True) logger.debug(f"Copying {relative_file}") - shutil.copy(py_file, dest_file) + copy_file(py_file, dest_file) def _build_and_validate_package_path( @@ -1596,7 +1666,7 @@ def _create_init_files(target_path: Path) -> None: # Create __init__.py in the target directory itself init_file = target_path / "__init__.py" if not init_file.exists(): - init_file.touch() + init_file.write_text(GENERATED_EMPTY_INIT_HEADER, encoding="utf-8") logger.debug(f"Created {init_file.relative_to(target_path.parent.parent.parent)}") # Create __init__.py in all subdirectories @@ -1604,7 +1674,7 @@ def _create_init_files(target_path: Path) -> None: if subdir.is_dir(): init_file = subdir / "__init__.py" if not init_file.exists(): - init_file.touch() + init_file.write_text(GENERATED_EMPTY_INIT_HEADER, encoding="utf-8") logger.debug(f"Created {init_file.relative_to(target_path.parent.parent.parent)}") @@ -1998,6 +2068,7 @@ def _replace_client_methods( source_tree = cst.parse_module(source_content).visit(ImportRewriter(source_module, target_module)) collector = _ClientMethodCollector() source_tree.visit(collector) + helper_functions = _collect_client_helper_functions(source_tree) if not collector.class_methods: logger.warning(f"No NeMo client methods found in {source_path}, skipping.") @@ -2005,6 +2076,7 @@ def _replace_client_methods( target_tree = cst.parse_module(target_content) target_tree = _ensure_required_client_init_imports(target_tree) + target_tree = _with_client_helper_functions(target_tree, helper_functions) replacer = _ClientMethodReplacer(collector.class_methods) modified_tree = target_tree.visit(replacer) modified_content = modified_tree.code diff --git a/tools/nemo-platform-sdk-tools/tests/sdk/openapi_test.py b/tools/nemo-platform-sdk-tools/tests/sdk/openapi_test.py index a2adc12c57..467c1e78f7 100644 --- a/tools/nemo-platform-sdk-tools/tests/sdk/openapi_test.py +++ b/tools/nemo-platform-sdk-tools/tests/sdk/openapi_test.py @@ -22,6 +22,63 @@ def test_stainless_resource_path(): assert endpoint.approx_resource_path() == ["customization", "configs"] +def test_extract_endpoints_skips_auth_utility_routes(): + spec = { + "paths": { + "/apis/auth/discovery": {"get": {}}, + "/apis/auth/jwks": {"get": {}}, + "/apis/auth/token": {"post": {}}, + "/apis/auth/v2/iam/role-bindings": {"get": {}}, + "/apis/entities/v2/workspaces": {"get": {}}, + }, + "components": {"schemas": {"Workspace": {"type": "object"}}}, + } + + endpoints = list(OpenAPI(spec).extract_endpoints()) + + assert endpoints == [ + OpenAPIEndpoint(method="get", path="/apis/auth/v2/iam/role-bindings"), + OpenAPIEndpoint(method="get", path="/apis/entities/v2/workspaces"), + ] + + +def test_schema_usage_skips_auth_utility_routes(): + spec = { + "paths": { + "/apis/auth/token": { + "post": { + "responses": { + "200": { + "content": { + "application/json": {"schema": {"$ref": "#/components/schemas/TokenExchangeResponse"}} + } + } + } + } + }, + "/apis/entities/v2/workspaces": { + "get": { + "responses": { + "200": {"content": {"application/json": {"schema": {"$ref": "#/components/schemas/Workspace"}}}} + } + } + }, + }, + "components": { + "schemas": { + "TokenExchangeResponse": {"type": "object"}, + "Workspace": {"type": "object"}, + } + }, + } + + schema_usage = OpenAPI(spec).calculate_schema_to_endpoints() + + assert schema_usage == { + "Workspace": [OpenAPIEndpoint(method="get", path="/apis/entities/v2/workspaces")], + } + + def test_extract_schema_refs_with_cycle(caplog): """Test that _extract_schema_refs handles cycles correctly.""" # Create a mock OpenAPI spec with circular references diff --git a/tools/nemo-platform-sdk-tools/tests/sdk/vendor/test_vendor_package.py b/tools/nemo-platform-sdk-tools/tests/sdk/vendor/test_vendor_package.py index 1091792429..74b9832ea1 100644 --- a/tools/nemo-platform-sdk-tools/tests/sdk/vendor/test_vendor_package.py +++ b/tools/nemo-platform-sdk-tools/tests/sdk/vendor/test_vendor_package.py @@ -60,6 +60,17 @@ def test_build_and_validate_target_paths_supports_top_level_targets(tmp_path: Pa assert nested_path == sdk_path / "src/nemo_platform/services/runner" +def test_copy_included_paths_preserves_generated_header_for_empty_init(tmp_path: Path) -> None: + source_path = tmp_path / "source" + destination_path = tmp_path / "destination" + source_path.mkdir() + (source_path / "__init__.py").write_text("", encoding="utf-8") + + vendor_package._copy_included_paths(source_path, destination_path, ["**/*.py"]) + + assert (destination_path / "__init__.py").read_text(encoding="utf-8") == vendor_package.GENERATED_EMPTY_INIT_HEADER + + def test_update_dependencies_of_sdk_pyproject_merges_optional_dependency_groups(tmp_path: Path, monkeypatch) -> None: """SDK client extension deps are written to the SDK pyproject.""" sdk_path = tmp_path / "sdk/python/nemo-platform" @@ -744,6 +755,10 @@ def test_replace_client_methods_updates_init_and_getattr(tmp_path: Path) -> None from typing import Any +def _should_bootstrap_config(config_path: object | None = None) -> bool: + return False + + class NeMoPlatform: def __init__(self) -> None: self.value = 1 @@ -762,9 +777,14 @@ def __init__(self) -> None: from typing import Any +def _should_bootstrap_config(config_path: Path | None = None) -> bool: + return config_path is not None + + class NeMoPlatform: def __init__(self, config_path: Path | None = None) -> None: self.config_path = config_path + self.should_bootstrap = _should_bootstrap_config(config_path) def __getattr__(self, name: str) -> Any: return name @@ -773,6 +793,7 @@ def __getattr__(self, name: str) -> Any: class AsyncNeMoPlatform: def __init__(self, config_path: Path | None = None) -> None: self.config_path = config_path + self.should_bootstrap = _should_bootstrap_config(config_path) def __getattr__(self, name: str) -> Any: return name @@ -791,7 +812,14 @@ def __getattr__(self, name: str) -> Any: updated = client_path.read_text(encoding="utf-8") assert "from pathlib import Path" in updated + assert "from nemo_platform._base_client import DefaultAsyncHttpxClient, DefaultHttpxClient" in updated + assert "from nemo_platform.client.tls import client_verify_from_env" in updated + assert "from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR" in updated + assert "def _should_bootstrap_config(config_path: Path | None = None) -> bool:" in updated + assert "return config_path is not None" in updated + assert "return False" not in updated assert "def __init__(self, config_path: Path | None = None) -> None:" in updated + assert updated.count("self.should_bootstrap = _should_bootstrap_config(config_path)") == 2 assert updated.count("def __getattr__(self, name: str) -> Any:") == 2 assert "self.value = 1" not in updated assert "self.value = 2" not in updated