From 63957e7a2e17a394889e795bf5742c42d3c08f0d Mon Sep 17 00:00:00 2001 From: Matt Kornfield Date: Tue, 14 Jul 2026 22:14:02 +0000 Subject: [PATCH 1/3] chore: update e2e safe-synthesizer tests Signed-off-by: Matt Kornfield --- e2e/test_safe_synthesizer.py | 623 +++++++++++++++++++++++++++++------ 1 file changed, 526 insertions(+), 97 deletions(-) diff --git a/e2e/test_safe_synthesizer.py b/e2e/test_safe_synthesizer.py index 33b2bf8fc3..a802fa92e1 100644 --- a/e2e/test_safe_synthesizer.py +++ b/e2e/test_safe_synthesizer.py @@ -1,74 +1,354 @@ -"""Opt-in container E2E for Safe Synthesizer GPU jobs (Docker or Kubernetes). +"""E2E coverage for the Safe Synthesizer plugin. -These tests exercise the full platform path: plugin job API -> Jobs controller -> -GPU container step -> safe-synthesizer-tasks image -> Files results. +Smoke coverage runs in the local subprocess harness and verifies the API, +Files, and job-entity surfaces without requiring the task image to complete. +Full workflow coverage is opt-in and targets a Kubernetes deployment such as +minikube at ``NMP_BASE_URL=http://localhost:30080``. -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: +Examples: + + uv run --frozen pytest e2e/test_safe_synthesizer.py -v --run-e2e - # 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 + uv run --frozen pytest e2e/test_safe_synthesizer.py -v \ + --run-e2e --run-slow --feature gpu """ from __future__ import annotations +import csv +import io +import json import os -import random import subprocess -from datetime import date +import sys +import time +import uuid +from collections.abc import Iterator +from contextlib import suppress from pathlib import Path +from typing import Any -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", -] +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.files.client import FilesClient +from nemo_platform_plugin.files.types import CreateFilesetRequest pytestmark = [ - pytest.mark.e2e, - pytest.mark.container_only, - pytest.mark.requires_gpu, - pytest.mark.slow, - pytest.mark.timeout(7200), + pytest.mark.timeout(600), + pytest.mark.e2e_config( + "e2e/configs/local-subprocess.yaml", + {"safe_synthesizer": {"runtime_python": sys.executable}}, + ), ] +TERMINAL_STATUSES = {"completed", "error", "cancelled"} +STARTED_STATUSES = {"pending", "active", "completed", "error", "cancelled"} + +INPUT_REMOTE_PATH = "inputs/safe-synthesizer-e2e.csv" +SMOKE_JOB_TIMEOUT_SECONDS = 180.0 +K8S_JOB_TIMEOUT_SECONDS = float(os.environ.get("NSS_E2E_JOB_TIMEOUT_SECONDS", "5400")) +POLL_INTERVAL_SECONDS = float(os.environ.get("NSS_E2E_POLL_INTERVAL_SECONDS", "10")) +RESULT_DOWNLOAD_TIMEOUT_SECONDS = float(os.environ.get("NSS_E2E_RESULT_DOWNLOAD_TIMEOUT_SECONDS", "600")) +DEFAULT_INPUT_ROWS = int(os.environ.get("NSS_E2E_INPUT_ROWS", "250")) +DEFAULT_NUM_RECORDS = int(os.environ.get("NSS_E2E_NUM_RECORDS", "250")) + + +def _unique_name(prefix: str) -> str: + return f"{prefix}-{uuid.uuid4().hex[:8]}" + + +def _string_headers(sdk: NeMoPlatform) -> dict[str, str]: + return {key: value for key, value in sdk.default_headers.items() if isinstance(value, str)} + + +def _nss_url(sdk: NeMoPlatform, workspace: str, path: str) -> str: + return f"{str(sdk.base_url).rstrip('/')}/apis/safe-synthesizer/v2/workspaces/{workspace}/{path.lstrip('/')}" + + +def _files_client(sdk: NeMoPlatform) -> FilesClient: + return client_from_platform(sdk, FilesClient) + + +def _create_fileset(sdk: NeMoPlatform, workspace: str, name: str) -> None: + _files_client(sdk).create_fileset( + workspace=workspace, + body=CreateFilesetRequest( + name=name, + description="Safe Synthesizer E2E fileset", + ), + ) + + +def _delete_fileset(sdk: NeMoPlatform, workspace: str, name: str) -> None: + with suppress(Exception): + _files_client(sdk).delete_fileset(name=name, workspace=workspace) + + +def _dataset_csv(rows: int = DEFAULT_INPUT_ROWS) -> str: + output = io.StringIO() + writer = csv.DictWriter( + output, + fieldnames=[ + "record_id", + "name", + "email", + "phone_number", + "city", + "signup_date", + "favorite_ice_cream_flavor", + "review", + "rating", + ], + ) + writer.writeheader() + cities = ["New York", "Los Angeles", "Chicago", "Houston", "Phoenix"] + flavors = ["Vanilla", "Chocolate", "Strawberry", "Mint Chip", "Coffee"] + for index in range(1, rows + 1): + writer.writerow( + { + "record_id": str(index), + "name": f"Customer {index}", + "email": f"customer{index}@example.com", + "phone_number": f"415-555-{index % 10000:04d}", + "city": cities[index % len(cities)], + "signup_date": f"2024-{(index % 12) + 1:02d}-{(index % 27) + 1:02d}", + "favorite_ice_cream_flavor": flavors[index % len(flavors)], + "review": f"Customer {index} asked support to call 415-555-{index % 10000:04d}.", + "rating": str((index % 5) + 1), + } + ) + return output.getvalue() + + +def _upload_dataset(sdk: NeMoPlatform, workspace: str, *, rows: int = DEFAULT_INPUT_ROWS) -> tuple[str, str]: + fileset = _unique_name("nss-inputs") + _create_fileset(sdk, workspace, fileset) + sdk.files.upload_content( + fileset=fileset, + workspace=workspace, + remote_path=INPUT_REMOTE_PATH, + content=_dataset_csv(rows), + ) + return fileset, f"{workspace}/{fileset}#{INPUT_REMOTE_PATH}" + + +def _job_payload( + name: str, + data_source: str, + config: dict[str, Any], + *, + description: str | None = None, +) -> dict[str, Any]: + payload: dict[str, Any] = { + "name": name, + "spec": { + "data_source": data_source, + "config": config, + }, + } + if description is not None: + payload["description"] = description + return payload + + +def _create_nss_job( + sdk: NeMoPlatform, + workspace: str, + *, + name: str, + data_source: str, + config: dict[str, Any], +) -> dict[str, Any]: + response = sdk._client.post( + _nss_url(sdk, workspace, "jobs"), + json=_job_payload(name, data_source, config), + headers=_string_headers(sdk), + timeout=60.0, + ) + assert response.status_code == 201, f"Failed to create Safe Synthesizer job: {response.text}" + return response.json() + + +def _list_nss_jobs(sdk: NeMoPlatform, workspace: str) -> dict[str, Any]: + response = sdk._client.get( + _nss_url(sdk, workspace, "jobs"), + headers=_string_headers(sdk), + timeout=60.0, + ) + assert response.status_code == 200, f"Failed to list Safe Synthesizer jobs: {response.text}" + return response.json() + + +def _retrieve_nss_job(sdk: NeMoPlatform, workspace: str, name: str) -> dict[str, Any]: + response = sdk._client.get( + _nss_url(sdk, workspace, f"jobs/{name}"), + headers=_string_headers(sdk), + timeout=60.0, + ) + assert response.status_code == 200, f"Failed to retrieve Safe Synthesizer job {name}: {response.text}" + return response.json() + + +def _delete_nss_job(sdk: NeMoPlatform, workspace: str, name: str) -> None: + response = sdk._client.delete( + _nss_url(sdk, workspace, f"jobs/{name}"), + headers=_string_headers(sdk), + timeout=60.0, + ) + if response.status_code not in {200, 202, 204, 404}: + response.raise_for_status() + + +def _cancel_nss_job(sdk: NeMoPlatform, workspace: str, name: str) -> dict[str, Any] | None: + response = sdk._client.post( + _nss_url(sdk, workspace, f"jobs/{name}/cancel"), + headers=_string_headers(sdk), + timeout=60.0, + ) + if response.status_code == 404: + return None + assert response.status_code == 200, f"Failed to cancel Safe Synthesizer job {name}: {response.text}" + return response.json() + + +def _list_nss_results(sdk: NeMoPlatform, workspace: str, job_name: str) -> dict[str, Any]: + response = sdk._client.get( + _nss_url(sdk, workspace, f"jobs/{job_name}/results"), + headers=_string_headers(sdk), + timeout=RESULT_DOWNLOAD_TIMEOUT_SECONDS, + ) + assert response.status_code == 200, f"Failed to list Safe Synthesizer results for {job_name}: {response.text}" + return response.json() + + +def _download_nss_result(sdk: NeMoPlatform, workspace: str, job_name: str, result_name: str) -> bytes: + response = sdk._client.get( + _nss_url(sdk, workspace, f"jobs/{job_name}/results/{result_name}/download"), + headers=_string_headers(sdk), + timeout=RESULT_DOWNLOAD_TIMEOUT_SECONDS, + ) + assert response.status_code == 200, ( + f"Failed to download Safe Synthesizer result {result_name!r} for {job_name}: {response.text}" + ) + return response.content + + +def _result_names(results: dict[str, Any]) -> set[str]: + return {str(result["name"]) for result in results.get("data", [])} + + +def _status_details(sdk: NeMoPlatform, workspace: str, job_name: str) -> str: + details = [f"Safe Synthesizer job {job_name} did not complete successfully."] + with suppress(Exception): + job = sdk.jobs.retrieve(job_name, workspace=workspace) + details.append(f"Job: {job.model_dump_json(indent=2)}") + with suppress(Exception): + status = sdk.jobs.get_status(job_name, workspace=workspace) + details.append(f"Status: {status.model_dump_json(indent=2)}") + with suppress(Exception): + logs = sdk.jobs.get_logs(job_name, workspace=workspace) + tail = logs.data[-30:] if logs.data else [] + details.append("Recent logs:") + details.extend(f"[{entry.job_step}] {entry.message}" for entry in tail) + return "\n".join(details) + + +def _wait_for_status( + sdk: NeMoPlatform, + workspace: str, + job_name: str, + *, + target_statuses: set[str] | None = None, + timeout_seconds: float, + poll_interval_seconds: float = POLL_INTERVAL_SECONDS, +) -> tuple[str, list[str]]: + target_statuses = target_statuses or TERMINAL_STATUSES + deadline = time.monotonic() + timeout_seconds + history: list[str] = [] + last_error: BaseException | None = None + + while time.monotonic() < deadline: + try: + status_info = sdk.jobs.get_status(job_name, workspace=workspace) + status = str(status_info.status) + if not history or history[-1] != status: + history.append(status) + if status in target_statuses: + return status, history + except Exception as exc: + last_error = exc + time.sleep(poll_interval_seconds) + + detail = _status_details(sdk, workspace, job_name) + if last_error is not None: + detail = f"{detail}\nLast polling error: {last_error!r}" + raise TimeoutError( + f"Timed out waiting for {job_name} to reach {sorted(target_statuses)}; history={history}\n{detail}" + ) + + +def _assert_job_completed(sdk: NeMoPlatform, workspace: str, job_name: str) -> list[str]: + status, history = _wait_for_status( + sdk, + workspace, + job_name, + timeout_seconds=K8S_JOB_TIMEOUT_SECONDS, + ) + assert status == "completed", _status_details(sdk, workspace, job_name) + assert any(seen in STARTED_STATUSES for seen in history), f"Unexpected job status history: {history}" + return history + + +def _assert_csv_rows(content: bytes, *, expected_rows: int | None = None, min_rows: int = 1) -> list[dict[str, str]]: + text = content.decode("utf-8") + rows = list(csv.DictReader(io.StringIO(text))) + assert len(rows) >= min_rows, f"Expected at least {min_rows} CSV rows, got {len(rows)}. Content: {text[:500]}" + if expected_rows is not None: + assert len(rows) == expected_rows, f"Expected {expected_rows} CSV rows, got {len(rows)}" + return rows + + +def _platform_root() -> Path: + candidates: list[Path] = [] + if os.environ.get("NMP_PLATFORM_ROOT"): + candidates.append(Path(os.environ["NMP_PLATFORM_ROOT"])) + candidates.extend( + [ + Path(__file__).resolve().parents[1], + Path.cwd() / "platform", + Path.cwd(), + ] + ) + for candidate in candidates: + script = candidate / "plugins/nemo-safe-synthesizer/scripts/setup_model_filesets.py" + if script.is_file(): + return candidate + raise FileNotFoundError("Could not locate plugins/nemo-safe-synthesizer/scripts/setup_model_filesets.py") + @pytest.fixture(scope="module") -def nss_model_filesets(sdk: NeMoPlatform, _services: str) -> None: - """Register HuggingFace-backed model filesets required by Safe Synthesizer tasks.""" +def nss_model_filesets(sdk: NeMoPlatform) -> None: + if os.environ.get("NSS_E2E_SKIP_MODEL_FILESETS") == "1": + return + + platform_root = _platform_root() + script = platform_root / "plugins/nemo-safe-synthesizer/scripts/setup_model_filesets.py" result = subprocess.run( [ "uv", "run", + "--project", + str(platform_root), "python", - str(_SETUP_MODEL_FILESETS), + str(script), "--files-api-url", - _services, + str(sdk.base_url).rstrip("/"), + "--workspace", + "default", ], - cwd=_REPO_ROOT, + cwd=platform_root, check=False, capture_output=True, text=True, @@ -79,65 +359,214 @@ def nss_model_filesets(sdk: NeMoPlatform, _services: str) -> None: ) -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( +@pytest.fixture +def nss_dataset(sdk: NeMoPlatform, workspace: str) -> Iterator[tuple[str, str]]: + fileset, data_source = _upload_dataset(sdk, workspace) + try: + yield fileset, data_source + finally: + _delete_fileset(sdk, workspace, fileset) + + +def test_safe_synthesizer_api_health(sdk: NeMoPlatform, workspace: str) -> None: + response = sdk._client.get( + f"{str(sdk.base_url).rstrip('/')}/status", + headers=_string_headers(sdk), + timeout=60.0, + ) + response.raise_for_status() + ready_services = response.json().get("services", {}).get("ready", []) + assert "safe-synthesizer" in ready_services + + jobs = _list_nss_jobs(sdk, workspace) + assert isinstance(jobs.get("data"), list) + + +def test_safe_synthesizer_fileset_upload_download_round_trips( + sdk: NeMoPlatform, + workspace: str, + nss_dataset: tuple[str, str], +) -> None: + fileset, _ = nss_dataset + downloaded = sdk.files.download_content( + fileset=fileset, + workspace=workspace, + remote_path=INPUT_REMOTE_PATH, + ) + + assert downloaded.decode("utf-8") == _dataset_csv() + + +def test_safe_synthesizer_job_create_list_retrieve_cancel_delete( sdk: NeMoPlatform, workspace: str, + nss_dataset: tuple[str, str], +) -> None: + _, data_source = nss_dataset + job_name = _unique_name("nss-smoke") + + job = _create_nss_job( + sdk, + workspace, + name=job_name, + data_source=data_source, + config={ + "enable_synthesis": False, + "enable_replace_pii": False, + }, + ) + try: + assert job["name"] == job_name + + jobs = _list_nss_jobs(sdk, workspace) + assert job_name in {entry["name"] for entry in jobs["data"]} + + retrieved = _retrieve_nss_job(sdk, workspace, job_name) + assert retrieved["name"] == job_name + assert retrieved["spec"]["data_source"] == data_source + + _cancel_nss_job(sdk, workspace, job_name) + status, _ = _wait_for_status( + sdk, + workspace, + job_name, + timeout_seconds=SMOKE_JOB_TIMEOUT_SECONDS, + poll_interval_seconds=2.0, + ) + assert status in TERMINAL_STATUSES + finally: + _cancel_nss_job(sdk, workspace, job_name) + _delete_nss_job(sdk, workspace, job_name) + + +@pytest.mark.container_only +@pytest.mark.requires_gpu +@pytest.mark.slow +@pytest.mark.timeout(7200) +def test_safe_synthesizer_k8s_job_cancel_transitions( + sdk: NeMoPlatform, + workspace: str, + nss_dataset: tuple[str, str], + nss_model_filesets: None, +) -> None: + _, data_source = nss_dataset + job_name = _unique_name("nss-cancel") + _create_nss_job( + sdk, + workspace, + name=job_name, + data_source=data_source, + config={ + "enable_synthesis": True, + "enable_replace_pii": False, + "generation": {"num_records": DEFAULT_NUM_RECORDS}, + "evaluation": {"enabled": False}, + "privacy": {"dp_enabled": False}, + }, + ) + try: + _, history = _wait_for_status( + sdk, + workspace, + job_name, + target_statuses=STARTED_STATUSES, + timeout_seconds=300, + ) + cancel_response = _cancel_nss_job(sdk, workspace, job_name) + assert cancel_response is not None + + final_status, final_history = _wait_for_status( + sdk, + workspace, + job_name, + timeout_seconds=600, + ) + assert final_status in {"cancelled", "completed", "error"} + assert history + final_history + finally: + _cancel_nss_job(sdk, workspace, job_name) + _delete_nss_job(sdk, workspace, job_name) + + +@pytest.mark.container_only +@pytest.mark.requires_gpu +@pytest.mark.slow +@pytest.mark.timeout(7200) +def test_safe_synthesizer_pii_replacement_job_completes( + sdk: NeMoPlatform, + workspace: str, + nss_dataset: tuple[str, 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}") + _, data_source = nss_dataset + job_name = _unique_name("nss-pii") + _create_nss_job( + sdk, + workspace, + name=job_name, + data_source=data_source, + config={ + "enable_synthesis": False, + "enable_replace_pii": True, + }, + ) + try: + _assert_job_completed(sdk, workspace, job_name) + results = _list_nss_results(sdk, workspace, job_name) + assert {"summary", "synthetic-data"}.issubset(_result_names(results)) + _assert_csv_rows( + _download_nss_result(sdk, workspace, job_name, "synthetic-data"), + expected_rows=DEFAULT_INPUT_ROWS, + ) + summary = json.loads(_download_nss_result(sdk, workspace, job_name, "summary")) + assert summary.get("timing") is not None + finally: + _delete_nss_job(sdk, workspace, job_name) + - job = ( - SafeSynthesizerJobBuilder(sdk, workspace=workspace) - .with_data_source(_synthesis_dataset()) - .synthesize() - .with_generate(num_records=num_records) - .with_evaluate(enabled=True) - .create_job() +@pytest.mark.container_only +@pytest.mark.requires_gpu +@pytest.mark.slow +@pytest.mark.timeout(7200) +def test_safe_synthesizer_full_workflow_downloads_artifacts( + sdk: NeMoPlatform, + workspace: str, + nss_dataset: tuple[str, str], + nss_model_filesets: None, +) -> None: + _, data_source = nss_dataset + job_name = _unique_name("nss-full") + _create_nss_job( + sdk, + workspace, + name=job_name, + data_source=data_source, + config={ + "enable_synthesis": True, + "enable_replace_pii": True, + "generation": {"num_records": DEFAULT_NUM_RECORDS}, + "evaluation": {"enabled": True}, + "privacy": {"dp_enabled": False}, + }, ) + try: + _assert_job_completed(sdk, workspace, job_name) + results = _list_nss_results(sdk, workspace, job_name) + result_names = _result_names(results) + assert {"summary", "synthetic-data", "evaluation-report", "adapter"}.issubset(result_names) - nss_job = SafeSynthesizerJob(job.job_name, sdk, workspace=workspace) - nss_job.wait_for_completion(poll_interval=15, verbose=True) + synthetic_rows = _assert_csv_rows( + _download_nss_result(sdk, workspace, job_name, "synthetic-data"), + expected_rows=DEFAULT_NUM_RECORDS, + ) + assert set(synthetic_rows[0]) >= {"name", "email", "favorite_ice_cream_flavor"} - summary = nss_job.fetch_summary() - assert summary.timing.training_time_sec is not None - assert summary.timing.generation_time_sec is not None + summary = json.loads(_download_nss_result(sdk, workspace, job_name, "summary")) + timing = summary["timing"] + assert timing["training_time_sec"] is not None + assert timing["generation_time_sec"] is not None - synthetic = nss_job.fetch_data() - assert len(synthetic) == num_records + report = _download_nss_result(sdk, workspace, job_name, "evaluation-report") + assert b" Date: Tue, 14 Jul 2026 22:35:42 +0000 Subject: [PATCH 2/3] chore: address 409 Signed-off-by: Matt Kornfield --- e2e/test_safe_synthesizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/e2e/test_safe_synthesizer.py b/e2e/test_safe_synthesizer.py index a802fa92e1..53f1695783 100644 --- a/e2e/test_safe_synthesizer.py +++ b/e2e/test_safe_synthesizer.py @@ -207,7 +207,7 @@ def _cancel_nss_job(sdk: NeMoPlatform, workspace: str, name: str) -> dict[str, A headers=_string_headers(sdk), timeout=60.0, ) - if response.status_code == 404: + if response.status_code in {404, 409}: return None assert response.status_code == 200, f"Failed to cancel Safe Synthesizer job {name}: {response.text}" return response.json() From 2894ca115b1fee3c086e3282b42b3387707c046e Mon Sep 17 00:00:00 2001 From: Matt Kornfield Date: Wed, 15 Jul 2026 17:45:22 +0000 Subject: [PATCH 3/3] chore: address safe synthesizer review comments Signed-off-by: Matt Kornfield --- e2e/test_safe_synthesizer.py | 304 +++++++++++++++++++++-------------- 1 file changed, 181 insertions(+), 123 deletions(-) diff --git a/e2e/test_safe_synthesizer.py b/e2e/test_safe_synthesizer.py index 53f1695783..d16eb04969 100644 --- a/e2e/test_safe_synthesizer.py +++ b/e2e/test_safe_synthesizer.py @@ -24,7 +24,7 @@ import sys import time import uuid -from collections.abc import Iterator +from collections.abc import Callable, Iterator from contextlib import suppress from pathlib import Path from typing import Any @@ -51,9 +51,13 @@ K8S_JOB_TIMEOUT_SECONDS = float(os.environ.get("NSS_E2E_JOB_TIMEOUT_SECONDS", "5400")) POLL_INTERVAL_SECONDS = float(os.environ.get("NSS_E2E_POLL_INTERVAL_SECONDS", "10")) RESULT_DOWNLOAD_TIMEOUT_SECONDS = float(os.environ.get("NSS_E2E_RESULT_DOWNLOAD_TIMEOUT_SECONDS", "600")) +MODEL_FILESETS_TIMEOUT_SECONDS = float(os.environ.get("NSS_E2E_MODEL_FILESETS_TIMEOUT_SECONDS", "300")) +DELETE_VERIFY_TIMEOUT_SECONDS = float(os.environ.get("NSS_E2E_DELETE_VERIFY_TIMEOUT_SECONDS", "60")) DEFAULT_INPUT_ROWS = int(os.environ.get("NSS_E2E_INPUT_ROWS", "250")) DEFAULT_NUM_RECORDS = int(os.environ.get("NSS_E2E_NUM_RECORDS", "250")) +NssJobFactory = Callable[[str, str, dict[str, Any]], dict[str, Any]] + def _unique_name(prefix: str) -> str: return f"{prefix}-{uuid.uuid4().hex[:8]}" @@ -181,6 +185,10 @@ def _list_nss_jobs(sdk: NeMoPlatform, workspace: str) -> dict[str, Any]: return response.json() +def _job_names(jobs: dict[str, Any]) -> set[str]: + return {str(entry["name"]) for entry in jobs.get("data", [])} + + def _retrieve_nss_job(sdk: NeMoPlatform, workspace: str, name: str) -> dict[str, Any]: response = sdk._client.get( _nss_url(sdk, workspace, f"jobs/{name}"), @@ -191,7 +199,25 @@ def _retrieve_nss_job(sdk: NeMoPlatform, workspace: str, name: str) -> dict[str, return response.json() -def _delete_nss_job(sdk: NeMoPlatform, workspace: str, name: str) -> None: +def _wait_for_job_absent( + sdk: NeMoPlatform, + workspace: str, + name: str, + *, + timeout_seconds: float = DELETE_VERIFY_TIMEOUT_SECONDS, + poll_interval_seconds: float = 2.0, +) -> None: + deadline = time.monotonic() + timeout_seconds + last_names: set[str] = set() + while time.monotonic() < deadline: + last_names = _job_names(_list_nss_jobs(sdk, workspace)) + if name not in last_names: + return + time.sleep(poll_interval_seconds) + pytest.fail(f"Safe Synthesizer job {name!r} still exists after delete; visible jobs: {sorted(last_names)}") + + +def _delete_nss_job(sdk: NeMoPlatform, workspace: str, name: str, *, verify: bool = True) -> None: response = sdk._client.delete( _nss_url(sdk, workspace, f"jobs/{name}"), headers=_string_headers(sdk), @@ -199,6 +225,8 @@ def _delete_nss_job(sdk: NeMoPlatform, workspace: str, name: str) -> None: ) if response.status_code not in {200, 202, 204, 404}: response.raise_for_status() + if verify: + _wait_for_job_absent(sdk, workspace, name) def _cancel_nss_job(sdk: NeMoPlatform, workspace: str, name: str) -> dict[str, Any] | None: @@ -310,6 +338,20 @@ def _assert_csv_rows(content: bytes, *, expected_rows: int | None = None, min_ro return rows +def _assert_known_pii_replaced(content: bytes) -> None: + text = content.decode("utf-8") + for source_value in ("customer1@example.com", "customer250@example.com", "415-555-0001", "415-555-0250"): + assert source_value not in text + + +def _process_output_text(output: str | bytes | None) -> str: + if output is None: + return "" + if isinstance(output, bytes): + return output.decode("utf-8", errors="replace") + return output + + def _platform_root() -> Path: candidates: list[Path] = [] if os.environ.get("NMP_PLATFORM_ROOT"): @@ -335,24 +377,32 @@ def nss_model_filesets(sdk: NeMoPlatform) -> None: platform_root = _platform_root() script = platform_root / "plugins/nemo-safe-synthesizer/scripts/setup_model_filesets.py" - result = subprocess.run( - [ - "uv", - "run", - "--project", - str(platform_root), - "python", - str(script), - "--files-api-url", - str(sdk.base_url).rstrip("/"), - "--workspace", - "default", - ], - cwd=platform_root, - check=False, - capture_output=True, - text=True, - ) + try: + result = subprocess.run( + [ + "uv", + "run", + "--project", + str(platform_root), + "python", + str(script), + "--files-api-url", + str(sdk.base_url).rstrip("/"), + "--workspace", + "default", + ], + cwd=platform_root, + timeout=MODEL_FILESETS_TIMEOUT_SECONDS, + check=False, + capture_output=True, + text=True, + ) + except subprocess.TimeoutExpired as exc: + pytest.fail( + f"Timed out after {MODEL_FILESETS_TIMEOUT_SECONDS:g}s registering Safe Synthesizer model filesets\n" + f"stdout:\n{_process_output_text(exc.stdout)}\n" + f"stderr:\n{_process_output_text(exc.stderr)}" + ) if result.returncode != 0: pytest.fail( f"Failed to register Safe Synthesizer model filesets\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}" @@ -368,6 +418,24 @@ def nss_dataset(sdk: NeMoPlatform, workspace: str) -> Iterator[tuple[str, str]]: _delete_fileset(sdk, workspace, fileset) +@pytest.fixture +def nss_job(sdk: NeMoPlatform, workspace: str) -> Iterator[NssJobFactory]: + job_names: list[str] = [] + + def create(prefix: str, data_source: str, config: dict[str, Any]) -> dict[str, Any]: + job_name = _unique_name(prefix) + job = _create_nss_job(sdk, workspace, name=job_name, data_source=data_source, config=config) + job_names.append(job_name) + return job + + try: + yield create + finally: + for job_name in reversed(job_names): + _cancel_nss_job(sdk, workspace, job_name) + _delete_nss_job(sdk, workspace, job_name) + + def test_safe_synthesizer_api_health(sdk: NeMoPlatform, workspace: str) -> None: response = sdk._client.get( f"{str(sdk.base_url).rstrip('/')}/status", @@ -401,42 +469,39 @@ def test_safe_synthesizer_job_create_list_retrieve_cancel_delete( sdk: NeMoPlatform, workspace: str, nss_dataset: tuple[str, str], + nss_job: NssJobFactory, ) -> None: _, data_source = nss_dataset - job_name = _unique_name("nss-smoke") - job = _create_nss_job( - sdk, - workspace, - name=job_name, - data_source=data_source, - config={ + job = nss_job( + "nss-smoke", + data_source, + { "enable_synthesis": False, "enable_replace_pii": False, }, ) - try: - assert job["name"] == job_name - - jobs = _list_nss_jobs(sdk, workspace) - assert job_name in {entry["name"] for entry in jobs["data"]} - - retrieved = _retrieve_nss_job(sdk, workspace, job_name) - assert retrieved["name"] == job_name - assert retrieved["spec"]["data_source"] == data_source - - _cancel_nss_job(sdk, workspace, job_name) - status, _ = _wait_for_status( - sdk, - workspace, - job_name, - timeout_seconds=SMOKE_JOB_TIMEOUT_SECONDS, - poll_interval_seconds=2.0, - ) - assert status in TERMINAL_STATUSES - finally: - _cancel_nss_job(sdk, workspace, job_name) - _delete_nss_job(sdk, workspace, job_name) + job_name = str(job["name"]) + assert job["name"] == job_name + + jobs = _list_nss_jobs(sdk, workspace) + assert job_name in _job_names(jobs) + + retrieved = _retrieve_nss_job(sdk, workspace, job_name) + assert retrieved["name"] == job_name + assert retrieved["spec"]["data_source"] == data_source + + cancel_response = _cancel_nss_job(sdk, workspace, job_name) + status, _ = _wait_for_status( + sdk, + workspace, + job_name, + timeout_seconds=SMOKE_JOB_TIMEOUT_SECONDS, + poll_interval_seconds=2.0, + ) + assert status != "error" + if cancel_response is not None: + assert status == "cancelled" @pytest.mark.container_only @@ -447,16 +512,14 @@ def test_safe_synthesizer_k8s_job_cancel_transitions( sdk: NeMoPlatform, workspace: str, nss_dataset: tuple[str, str], + nss_job: NssJobFactory, nss_model_filesets: None, ) -> None: _, data_source = nss_dataset - job_name = _unique_name("nss-cancel") - _create_nss_job( - sdk, - workspace, - name=job_name, - data_source=data_source, - config={ + job = nss_job( + "nss-cancel", + data_source, + { "enable_synthesis": True, "enable_replace_pii": False, "generation": {"num_records": DEFAULT_NUM_RECORDS}, @@ -464,28 +527,27 @@ def test_safe_synthesizer_k8s_job_cancel_transitions( "privacy": {"dp_enabled": False}, }, ) - try: - _, history = _wait_for_status( - sdk, - workspace, - job_name, - target_statuses=STARTED_STATUSES, - timeout_seconds=300, - ) - cancel_response = _cancel_nss_job(sdk, workspace, job_name) - assert cancel_response is not None - - final_status, final_history = _wait_for_status( - sdk, - workspace, - job_name, - timeout_seconds=600, - ) - assert final_status in {"cancelled", "completed", "error"} - assert history + final_history - finally: - _cancel_nss_job(sdk, workspace, job_name) - _delete_nss_job(sdk, workspace, job_name) + job_name = str(job["name"]) + + _, history = _wait_for_status( + sdk, + workspace, + job_name, + target_statuses=STARTED_STATUSES, + timeout_seconds=300, + ) + assert "error" not in history + cancel_response = _cancel_nss_job(sdk, workspace, job_name) + assert cancel_response is not None + + final_status, final_history = _wait_for_status( + sdk, + workspace, + job_name, + timeout_seconds=600, + ) + assert final_status == "cancelled" + assert "cancelled" in final_history @pytest.mark.container_only @@ -496,32 +558,31 @@ def test_safe_synthesizer_pii_replacement_job_completes( sdk: NeMoPlatform, workspace: str, nss_dataset: tuple[str, str], + nss_job: NssJobFactory, nss_model_filesets: None, ) -> None: _, data_source = nss_dataset - job_name = _unique_name("nss-pii") - _create_nss_job( - sdk, - workspace, - name=job_name, - data_source=data_source, - config={ + job = nss_job( + "nss-pii", + data_source, + { "enable_synthesis": False, "enable_replace_pii": True, }, ) - try: - _assert_job_completed(sdk, workspace, job_name) - results = _list_nss_results(sdk, workspace, job_name) - assert {"summary", "synthetic-data"}.issubset(_result_names(results)) - _assert_csv_rows( - _download_nss_result(sdk, workspace, job_name, "synthetic-data"), - expected_rows=DEFAULT_INPUT_ROWS, - ) - summary = json.loads(_download_nss_result(sdk, workspace, job_name, "summary")) - assert summary.get("timing") is not None - finally: - _delete_nss_job(sdk, workspace, job_name) + job_name = str(job["name"]) + + _assert_job_completed(sdk, workspace, job_name) + results = _list_nss_results(sdk, workspace, job_name) + assert {"summary", "synthetic-data"}.issubset(_result_names(results)) + synthetic_content = _download_nss_result(sdk, workspace, job_name, "synthetic-data") + _assert_csv_rows( + synthetic_content, + expected_rows=DEFAULT_INPUT_ROWS, + ) + _assert_known_pii_replaced(synthetic_content) + summary = json.loads(_download_nss_result(sdk, workspace, job_name, "summary")) + assert summary.get("timing") is not None @pytest.mark.container_only @@ -532,16 +593,14 @@ def test_safe_synthesizer_full_workflow_downloads_artifacts( sdk: NeMoPlatform, workspace: str, nss_dataset: tuple[str, str], + nss_job: NssJobFactory, nss_model_filesets: None, ) -> None: _, data_source = nss_dataset - job_name = _unique_name("nss-full") - _create_nss_job( - sdk, - workspace, - name=job_name, - data_source=data_source, - config={ + job = nss_job( + "nss-full", + data_source, + { "enable_synthesis": True, "enable_replace_pii": True, "generation": {"num_records": DEFAULT_NUM_RECORDS}, @@ -549,24 +608,23 @@ def test_safe_synthesizer_full_workflow_downloads_artifacts( "privacy": {"dp_enabled": False}, }, ) - try: - _assert_job_completed(sdk, workspace, job_name) - results = _list_nss_results(sdk, workspace, job_name) - result_names = _result_names(results) - assert {"summary", "synthetic-data", "evaluation-report", "adapter"}.issubset(result_names) - - synthetic_rows = _assert_csv_rows( - _download_nss_result(sdk, workspace, job_name, "synthetic-data"), - expected_rows=DEFAULT_NUM_RECORDS, - ) - assert set(synthetic_rows[0]) >= {"name", "email", "favorite_ice_cream_flavor"} + job_name = str(job["name"]) - summary = json.loads(_download_nss_result(sdk, workspace, job_name, "summary")) - timing = summary["timing"] - assert timing["training_time_sec"] is not None - assert timing["generation_time_sec"] is not None + _assert_job_completed(sdk, workspace, job_name) + results = _list_nss_results(sdk, workspace, job_name) + result_names = _result_names(results) + assert {"summary", "synthetic-data", "evaluation-report", "adapter"}.issubset(result_names) - report = _download_nss_result(sdk, workspace, job_name, "evaluation-report") - assert b"= {"name", "email", "favorite_ice_cream_flavor"} + + summary = json.loads(_download_nss_result(sdk, workspace, job_name, "summary")) + timing = summary["timing"] + assert timing["training_time_sec"] is not None + assert timing["generation_time_sec"] is not None + + report = _download_nss_result(sdk, workspace, job_name, "evaluation-report") + assert b"