chore: add safe-synthesizer test for kubernetes

conftest.py

@@ -284,6 +284,20 @@ def pytest_addoption(parser):
         default=True,
         help="Run integration tests (enabled by default)",
     )
+    parser.addoption(
+        "--feature",
+        action="append",
+        default=[],
+        metavar="NAME",
+        help="Enable optional e2e feature sets (repeatable), e.g. --feature gpu",
+    )
+
+
+def _e2e_features_enabled(config: pytest.Config) -> set[str]:
+    features = {feature.lower() for feature in config.getoption("--feature") or []}
+    if os.environ.get("RUN_NSS_K8S_E2E") == "1":
+        features.add("gpu")
+    return features
 
 
 def pytest_runtest_setup(item):
@@ -302,6 +316,9 @@ def pytest_runtest_setup(item):
     if "container_only" in [marker.name for marker in item.iter_markers()]:
         if not os.environ.get("NMP_BASE_URL"):
             skip_test("Skipping container-only test (requires NMP_BASE_URL)")
+    if "requires_gpu" in [marker.name for marker in item.iter_markers()]:
+        if "gpu" not in _e2e_features_enabled(item.config):
+            skip_test("Skipping GPU container e2e (pass --feature gpu)")
 
 
 from xdist.scheduler.loadgroup import LoadGroupScheduling  # noqa: E402

docker/Dockerfile.safe-synthesizer-tasks

@@ -35,7 +35,11 @@ ARG USERNAME=nemo
 ARG USER_UID=1000
 ARG USER_GID=1000
 
-ENV UV_HTTP_TIMEOUT=120
+# Raise uv's HTTP retry budget above the default of 3. uv retries transient
+# failures (connection errors and 429/5xx status codes such as the intermittent
+# 503s from download.pytorch.org) with exponential backoff between attempts.
+ENV UV_HTTP_TIMEOUT=120 \
+    UV_HTTP_RETRIES=8
 
 COPY --from=uv /uv /uvx /usr/local/bin/
 

e2e/k8s/scripts/local_build_and_upgrade.sh

@@ -8,8 +8,13 @@
 #   MINIKUBE_PROFILE  - minikube profile name (default: minikube)
 #   NMP_REGISTRY      - image registry (default: docker.io/my-registry)
 #   IMAGE_TAG         - image tag (default: local-<epoch>)
-#   BUILD_ARCH        - target platform (default: auto-detected from host)
-#   HELM_VALUES       - values file (default: e2e/k8s/values/minikube.yaml)
+#   BUILD_ARCH            - target platform (default: auto-detected from host)
+#   MINIKUBE_GPU          - when 1, start minikube with GPU passthrough
+#   BUILD_SAFE_SYNTHESIZER - when 1, also build safe-synthesizer-tasks (amd64;
+#                            set BUILD_ARCH=linux/amd64)
+#   BUILD_GPU             - deprecated convenience alias: sets MINIKUBE_GPU=1 and
+#                            BUILD_SAFE_SYNTHESIZER=1
+#   HELM_VALUES           - values file (default: e2e/k8s/values/minikube.yaml)
 
 set -e
 
@@ -18,10 +23,21 @@ REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
 
 MINIKUBE_PROFILE="${MINIKUBE_PROFILE:-minikube}"
 
+if [ "${BUILD_GPU:-0}" = "1" ]; then
+  MINIKUBE_GPU="${MINIKUBE_GPU:-1}"
+  BUILD_SAFE_SYNTHESIZER="${BUILD_SAFE_SYNTHESIZER:-1}"
+fi
+MINIKUBE_GPU="${MINIKUBE_GPU:-0}"
+BUILD_SAFE_SYNTHESIZER="${BUILD_SAFE_SYNTHESIZER:-0}"
+
 # Check if minikube is running
 if ! minikube status -p "${MINIKUBE_PROFILE}" &>/dev/null; then
   echo "Minikube profile ${MINIKUBE_PROFILE} is not running. Starting..."
-  MINIKUBE_PROFILE="${MINIKUBE_PROFILE}" "$SCRIPT_DIR/setup_local_minikube_cpu.sh"
+  if [ "${MINIKUBE_GPU}" = "1" ]; then
+    MINIKUBE_PROFILE="${MINIKUBE_PROFILE}" "$SCRIPT_DIR/setup_local_minikube_gpu.sh"
+  else
+    MINIKUBE_PROFILE="${MINIKUBE_PROFILE}" "$SCRIPT_DIR/setup_local_minikube_cpu.sh"
+  fi
 fi
 
 # Wait for minikube to be ready
@@ -48,10 +64,22 @@ eval "$(minikube -p "${MINIKUBE_PROFILE}" docker-env)"
     IMAGE_REGISTRY="${NMP_REGISTRY}" \
     BUILD_ARCH="$BUILD_ARCH" \
     docker buildx bake docker-cpu --set "*.platform=$BUILD_ARCH"
+
+  if [ "${BUILD_SAFE_SYNTHESIZER}" = "1" ]; then
+    echo "Building safe-synthesizer-tasks (BUILD_SAFE_SYNTHESIZER=1)..."
+    CI_COMMIT_SHA="$GIT_SHA" \
+      BAKE_TAG="$IMAGE_TAG" \
+      IMAGE_REGISTRY="${NMP_REGISTRY}" \
+      BUILD_ARCH="$BUILD_ARCH" \
+      docker buildx bake safe-synthesizer-tasks-docker --set "safe-synthesizer-tasks-docker.platform=$BUILD_ARCH"
+  fi
 )
 
 echo "----------------------------------------"
 echo "Images built with tag: $IMAGE_TAG"
+if [ "${BUILD_SAFE_SYNTHESIZER}" = "1" ]; then
+  echo "Also built: safe-synthesizer-tasks"
+fi
 echo "----------------------------------------"
 
 # Delegate helm install to install_helm_e2e.sh

e2e/test_safe_synthesizer.py

@@ -0,0 +1,143 @@
+"""Opt-in container E2E for Safe Synthesizer GPU jobs (Docker or Kubernetes).
+
+These tests exercise the full platform path: plugin job API -> Jobs controller ->
+GPU container step -> safe-synthesizer-tasks image -> Files results.
+
+Excluded from default kind-cpu CI (no GPU, no safe-synthesizer-tasks image).
+Run manually against minikube GPU, dev-blue, or a GPU-enabled Docker backend:
+
+    # After nss-k8s-deploy.sh (or MINIKUBE_GPU=1 BUILD_SAFE_SYNTHESIZER=1 local_build_and_upgrade.sh)
+    NMP_BASE_URL=http://localhost:30080 \
+      uv run --frozen pytest e2e/test_safe_synthesizer.py -v --run-e2e --run-slow --feature gpu
+"""
+
+from __future__ import annotations
+
+import os
+import random
+import subprocess
+from datetime import date
+from pathlib import Path
+
+import pandas as pd
+import pytest
+from nemo_platform import NeMoPlatform
+from nemo_safe_synthesizer_plugin.sdk.job import SafeSynthesizerJob
+from nemo_safe_synthesizer_plugin.sdk.job_builder import SafeSynthesizerJobBuilder
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+_SETUP_MODEL_FILESETS = _REPO_ROOT / "plugins/nemo-safe-synthesizer/scripts/setup_model_filesets.py"
+
+_MIN_INPUT_ROWS = 200
+_DEFAULT_INPUT_ROWS = 250
+_DEFAULT_NUM_RECORDS = _DEFAULT_INPUT_ROWS
+
+_ICE_CREAM_FLAVORS = [
+    "Vanilla",
+    "Chocolate",
+    "Strawberry",
+    "Mint Chocolate Chip",
+    "Cookies and Cream",
+    "Pistachio",
+    "Rocky Road",
+    "Butter Pecan",
+    "Coffee",
+    "Mango Sorbet",
+    "Salted Caramel",
+    "Cookie Dough",
+]
+
+pytestmark = [
+    pytest.mark.e2e,
+    pytest.mark.container_only,
+    pytest.mark.requires_gpu,
+    pytest.mark.slow,
+    pytest.mark.timeout(7200),
+]
+
+
+@pytest.fixture(scope="module")
+def nss_model_filesets(sdk: NeMoPlatform, _services: str) -> None:
+    """Register HuggingFace-backed model filesets required by Safe Synthesizer tasks."""
+    result = subprocess.run(
+        [
+            "uv",
+            "run",
+            "python",
+            str(_SETUP_MODEL_FILESETS),
+            "--files-api-url",
+            _services,
+        ],
+        cwd=_REPO_ROOT,
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0:
+        pytest.fail(
+            f"Failed to register Safe Synthesizer model filesets\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+        )
+
+
+def _synthesis_dataset(rows: int | None = None) -> pd.DataFrame:
+    """Build a tabular dataset suitable for Safe Synthesizer training (>= 200 rows).
+
+    Schema matches plugins/nemo-safe-synthesizer/tests/e2e/test_local_synthesis.py:
+    names, dates, and a categorical column with realistic variation.
+    """
+    if rows is None:
+        rows = int(os.environ.get("NSS_E2E_INPUT_ROWS", str(_DEFAULT_INPUT_ROWS)))
+    if rows < _MIN_INPUT_ROWS:
+        raise ValueError(f"Safe Synthesizer container E2E requires at least {_MIN_INPUT_ROWS} input rows, got {rows}")
+
+    faker_mod = pytest.importorskip("faker")
+    fake = faker_mod.Faker()
+    faker_mod.Faker.seed(42)
+    random.seed(42)
+
+    records = [
+        {
+            "name": fake.name(),
+            "signup_date": fake.date_between_dates(
+                date_start=date(2020, 1, 1),
+                date_end=date(2026, 5, 4),
+            ).isoformat(),
+            "birthdate": fake.date_between_dates(
+                date_start=date(1945, 1, 1),
+                date_end=date(2006, 12, 31),
+            ).isoformat(),
+            "favorite_ice_cream_flavor": random.choice(_ICE_CREAM_FLAVORS),
+        }
+        for _ in range(rows)
+    ]
+    return pd.DataFrame.from_records(records)
+
+
+def test_safe_synthesizer_container_job_completes(
+    sdk: NeMoPlatform,
+    workspace: str,
+    nss_model_filesets: None,
+) -> None:
+    """Submit a GPU container job and verify synthetic data is produced."""
+    num_records = int(os.environ.get("NSS_E2E_NUM_RECORDS", str(_DEFAULT_NUM_RECORDS)))
+    if num_records < _MIN_INPUT_ROWS:
+        raise ValueError(f"NSS_E2E_NUM_RECORDS must be at least {_MIN_INPUT_ROWS}, got {num_records}")
+
+    job = (
+        SafeSynthesizerJobBuilder(sdk, workspace=workspace)
+        .with_data_source(_synthesis_dataset())
+        .synthesize()
+        .with_generate(num_records=num_records)
+        .with_evaluate(enabled=True)
+        .create_job()
+    )
+
+    nss_job = SafeSynthesizerJob(job.job_name, sdk, workspace=workspace)
+    nss_job.wait_for_completion(poll_interval=15, verbose=True)
+
+    summary = nss_job.fetch_summary()
+    assert summary.timing.training_time_sec is not None
+    assert summary.timing.generation_time_sec is not None
+
+    synthetic = nss_job.fetch_data()
+    assert len(synthetic) == num_records

pytest.ini

@@ -66,6 +66,7 @@ markers =
     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
+    requires_gpu: Container e2e requiring GPU job scheduling; skipped unless --feature gpu is passed
     regression: Regression tests - test individual functional microservices for baseline functionality
     infrastructure: Infrastructure tests - ensure services are compatible with customer infrastructure
     canary: Canary tests - test deployed integration environments like top of tree