[PLT-0] Add safe staging embedding cleanup

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: gramos-lb

.github/workflows/staging-embedding-cleanup.yml

@@ -0,0 +1,35 @@
+name: Staging Embedding Cleanup
+
+on:
+  workflow_dispatch:
+    inputs:
+      dry_run:
+        description: >-
+          Dry-run first. Set false only after confirming both Labelbox Python
+          SDK Staging and LBox Develop are quiet.
+        required: true
+        type: boolean
+        default: true
+
+permissions:
+  contents: read
+
+jobs:
+  cleanup:
+    runs-on: ubuntu-latest
+    env:
+      LABELBOX_TEST_API_KEY: ${{ secrets.STAGING_API_KEY_ORG_CMOI3PQ7801GM070D4TPG7FMH }}
+    steps:
+      - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5
+        with:
+          python-version: "3.11"
+          cache: pip
+      - name: Install SDK from this repository
+        working-directory: libs/labelbox
+        run: python -m pip install .
+      - name: List or clean leaked staging embeddings
+        working-directory: libs/labelbox
+        run: >-
+          python -m tests.scripts.cleanup_staging_embeddings
+          --dry-run "${{ inputs.dry_run }}"

libs/labelbox/tests/conftest.py

@@ -32,6 +32,10 @@
 from labelbox.schema.ontology import Ontology
 from labelbox.schema.project import Project
 from labelbox.schema.quality_mode import QualityMode
+from tests.embedding_cleanup import (
+    build_embedding_name,
+    create_embedding_with_heal,
+)
 
 # Must be a stable, deterministic JPEG: several tests assert byte-equality
 # between the source and the server-rehosted copy, so a random image service
@@ -1128,9 +1132,14 @@ def configured_project_with_complex_ontology(
 
 @pytest.fixture
 def embedding(client: Client, environ):
-    uuid_str = uuid.uuid4().hex
     time.sleep(randint(1, 5))
-    embedding = client.create_embedding(f"sdk-int-{uuid_str}", 8)
+    embedding = create_embedding_with_heal(
+        create_embedding=lambda: client.create_embedding(
+            build_embedding_name(time.time()), 8
+        ),
+        list_embeddings=client.get_embeddings,
+        delete_embedding=lambda stale_embedding: stale_embedding.delete(),
+    )
     yield embedding
 
     embedding.delete()

libs/labelbox/tests/embedding_cleanup.py

@@ -0,0 +1,97 @@
+import os
+import re
+import time
+import uuid
+from random import uniform
+from typing import Any, Callable, Iterable, List, Optional
+
+from lbox.exceptions import LabelboxError
+
+EMBEDDING_NAME_PREFIX_V2 = "sdk-int-ci-v2-"
+EMBEDDING_CAP_ERROR_SNIPPET = "Max limit of custom embeddings"
+EMBEDDING_STALE_TTL_SECONDS = 12 * 3600
+LEGACY_NAME_RE = re.compile(r"^sdk-int-[0-9a-f]{32}$")
+
+_V2_NAME_RE = re.compile(
+    rf"^{re.escape(EMBEDDING_NAME_PREFIX_V2)}(\d+)-[0-9a-f]{{10}}$"
+)
+_MAX_CREATE_ATTEMPTS = 3
+
+
+def build_embedding_name(now: float) -> str:
+    return f"{EMBEDDING_NAME_PREFIX_V2}{int(now)}-{uuid.uuid4().hex[:10]}"
+
+
+def parse_embedding_created_at(name: str) -> Optional[int]:
+    match = _V2_NAME_RE.fullmatch(name)
+    return int(match.group(1)) if match is not None else None
+
+
+def select_stale_embeddings(embeddings: Iterable[Any], now: float) -> List[Any]:
+    stale_embeddings = []
+    for embedding in embeddings:
+        created_at = parse_embedding_created_at(embedding.name)
+        if (
+            embedding.custom
+            and created_at is not None
+            and now - created_at > EMBEDDING_STALE_TTL_SECONDS
+        ):
+            stale_embeddings.append(embedding)
+    return stale_embeddings
+
+
+def is_embedding_cap_error(error: LabelboxError) -> bool:
+    return EMBEDDING_CAP_ERROR_SNIPPET in str(error)
+
+
+def create_embedding_with_heal(
+    *,
+    create_embedding: Callable[[], Any],
+    list_embeddings: Callable[[], Iterable[Any]],
+    delete_embedding: Callable[[Any], None],
+    sleep: Callable[[float], None] = time.sleep,
+    now: Callable[[], float] = time.time,
+    retry_delay: Callable[[float, float], float] = uniform,
+    print_fn: Callable[[str], None] = print,
+) -> Any:
+    last_swept_count = None
+
+    for attempt in range(1, _MAX_CREATE_ATTEMPTS + 1):
+        try:
+            return create_embedding()
+        except LabelboxError as error:
+            if not is_embedding_cap_error(error):
+                raise
+
+            if attempt == _MAX_CREATE_ATTEMPTS:
+                if last_swept_count == 0:
+                    print_fn(
+                        "[embedding-fixture-heal] no stale embeddings were "
+                        "swept; the cap appears held by live fixtures and/or "
+                        "legacy/foreign names that automated healing "
+                        "deliberately does not touch"
+                    )
+                raise
+
+            stale_embeddings = select_stale_embeddings(list_embeddings(), now())
+            for embedding in stale_embeddings:
+                try:
+                    delete_embedding(embedding)
+                except LabelboxError:
+                    # Another worker may have deleted the same stale embedding.
+                    pass
+
+            last_swept_count = len(stale_embeddings)
+            sample = ",".join(
+                embedding.id for embedding in stale_embeddings[:3]
+            )
+            print_fn(
+                "[embedding-fixture-heal] "
+                f"run={os.getenv('GITHUB_RUN_ID', '-')} "
+                f"worker={os.getenv('PYTEST_XDIST_WORKER', '-')} "
+                f"attempt={attempt} cap_hit=1 "
+                f"swept={last_swept_count} sample={sample or '-'}"
+            )
+            sleep(retry_delay(2, 8))
+
+    raise AssertionError("unreachable")

libs/labelbox/tests/scripts/cleanup_staging_embeddings.py

@@ -0,0 +1,143 @@
+import argparse
+import os
+from typing import Any, Callable, Iterable, List, Optional, Sequence
+from urllib.parse import urlparse
+
+from lbox.exceptions import LabelboxError
+
+from labelbox import Client
+from tests.embedding_cleanup import (
+    LEGACY_NAME_RE,
+    select_stale_embeddings,
+)
+
+STAGING_GRAPHQL_ENDPOINT = "https://api.lb-stage.xyz/graphql"
+STAGING_REST_ENDPOINT = "https://api.lb-stage.xyz/api/v1"
+STAGING_REST_HOST = "api.lb-stage.xyz"
+
+
+def _parse_boolean(value: str) -> bool:
+    normalized = value.strip().lower()
+    if normalized == "true":
+        return True
+    if normalized == "false":
+        return False
+    raise argparse.ArgumentTypeError("expected 'true' or 'false'")
+
+
+def assert_staging_rest_endpoint(client: Any) -> None:
+    effective_host = urlparse(client.rest_endpoint).hostname
+    if effective_host != STAGING_REST_HOST:
+        raise RuntimeError(
+            "Refusing to inspect embeddings: the effective REST endpoint "
+            f"host is {effective_host!r}, expected {STAGING_REST_HOST!r}"
+        )
+
+
+def create_staging_client(
+    *,
+    api_key: Optional[str] = None,
+    client_factory: Callable[..., Any] = Client,
+) -> Any:
+    effective_api_key = api_key or os.environ.get("LABELBOX_TEST_API_KEY")
+    if not effective_api_key:
+        raise RuntimeError("LABELBOX_TEST_API_KEY is required")
+
+    client = client_factory(
+        api_key=effective_api_key,
+        endpoint=STAGING_GRAPHQL_ENDPOINT,
+        rest_endpoint=STAGING_REST_ENDPOINT,
+    )
+    assert_staging_rest_endpoint(client)
+    return client
+
+
+def select_cleanup_candidates(
+    embeddings: Iterable[Any], now: float
+) -> List[Any]:
+    embeddings = list(embeddings)
+    stale_v2_ids = {
+        embedding.id for embedding in select_stale_embeddings(embeddings, now)
+    }
+    return [
+        embedding
+        for embedding in embeddings
+        if embedding.custom
+        and (
+            LEGACY_NAME_RE.fullmatch(embedding.name) is not None
+            or embedding.id in stale_v2_ids
+        )
+    ]
+
+
+def run_cleanup(client: Any, *, dry_run: bool, now: float) -> int:
+    # Validate the effective endpoint immediately before the first API read.
+    assert_staging_rest_endpoint(client)
+    print(
+        "WARNING: this cleanup cannot detect active owners. Confirm both "
+        "Labelbox Python SDK Staging and LBox Develop are quiet, and run "
+        "dry-run first."
+    )
+    candidates = select_cleanup_candidates(client.get_embeddings(), now)
+
+    print(f"Embedding cleanup candidates ({len(candidates)}):")
+    for embedding in candidates:
+        print(f"candidate id={embedding.id} name={embedding.name}")
+
+    if dry_run:
+        print(
+            "[embedding-cleanup] "
+            f"run={os.getenv('GITHUB_RUN_ID', '-')} dry_run=1 "
+            f"candidates={len(candidates)} deleted=0 failed=0"
+        )
+        return 0
+
+    deleted = 0
+    failed = []
+    for embedding in candidates:
+        try:
+            embedding.delete()
+            deleted += 1
+            print(f"deleted id={embedding.id} name={embedding.name}")
+        except LabelboxError as error:
+            failed.append(embedding.id)
+            print(
+                f"failed id={embedding.id} name={embedding.name} error={error}"
+            )
+
+    print(
+        "[embedding-cleanup] "
+        f"run={os.getenv('GITHUB_RUN_ID', '-')} dry_run=0 "
+        f"candidates={len(candidates)} deleted={deleted} "
+        f"failed={len(failed)}"
+    )
+    return 1 if failed else 0
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Clean leaked custom embeddings from the shared staging org. "
+            "Confirm Labelbox Python SDK Staging and LBox Develop are quiet "
+            "and run dry-run first."
+        )
+    )
+    parser.add_argument(
+        "--dry-run",
+        type=_parse_boolean,
+        default=True,
+        help="true (default) lists only; false deletes every candidate",
+    )
+    args = parser.parse_args(argv)
+
+    import time
+
+    return run_cleanup(
+        create_staging_client(),
+        dry_run=args.dry_run,
+        now=time.time(),
+    )
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())