diff --git a/.github/workflows/staging-embedding-cleanup.yml b/.github/workflows/staging-embedding-cleanup.yml new file mode 100644 index 000000000..3c3cf88d0 --- /dev/null +++ b/.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 }}" diff --git a/libs/labelbox/tests/conftest.py b/libs/labelbox/tests/conftest.py index 87d59bdb1..94d756170 100644 --- a/libs/labelbox/tests/conftest.py +++ b/libs/labelbox/tests/conftest.py @@ -1,6 +1,8 @@ import json import os +import pathlib import re +import sys import time import uuid from datetime import datetime @@ -33,6 +35,18 @@ from labelbox.schema.project import Project from labelbox.schema.quality_mode import QualityMode +# CI invokes pytest as a console script, which (unlike `python -m pytest`) +# never puts the project directory on sys.path, so `tests.*` is not +# importable when this conftest loads. Insert it deterministically. +_PROJECT_DIR = str(pathlib.Path(__file__).resolve().parent.parent) +if _PROJECT_DIR not in sys.path: + sys.path.insert(0, _PROJECT_DIR) + +from tests.embedding_cleanup import ( # noqa: E402 (needs the sys.path insert above) + 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 # (e.g. picsum.photos) or a format the server may transcode (e.g. PNG) breaks them. @@ -1128,9 +1142,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() diff --git a/libs/labelbox/tests/embedding_cleanup.py b/libs/labelbox/tests/embedding_cleanup.py new file mode 100644 index 000000000..eed2a0062 --- /dev/null +++ b/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") diff --git a/libs/labelbox/tests/scripts/cleanup_staging_embeddings.py b/libs/labelbox/tests/scripts/cleanup_staging_embeddings.py new file mode 100644 index 000000000..7a3b295bb --- /dev/null +++ b/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()) diff --git a/libs/labelbox/tests/unit/test_embedding_cleanup.py b/libs/labelbox/tests/unit/test_embedding_cleanup.py new file mode 100644 index 000000000..882648d7a --- /dev/null +++ b/libs/labelbox/tests/unit/test_embedding_cleanup.py @@ -0,0 +1,259 @@ +from types import SimpleNamespace +from unittest.mock import Mock, call, patch + +import pytest +from lbox.exceptions import LabelboxError + +from tests.embedding_cleanup import ( + EMBEDDING_CAP_ERROR_SNIPPET, + EMBEDDING_NAME_PREFIX_V2, + EMBEDDING_STALE_TTL_SECONDS, + build_embedding_name, + create_embedding_with_heal, + is_embedding_cap_error, + parse_embedding_created_at, + select_stale_embeddings, +) +from tests.scripts.cleanup_staging_embeddings import ( + STAGING_GRAPHQL_ENDPOINT, + STAGING_REST_ENDPOINT, + _parse_boolean, + create_staging_client, + run_cleanup, +) + + +def _embedding( + embedding_id: str, + name: str, + *, + custom: bool = True, + delete: Mock = None, +): + return SimpleNamespace( + id=embedding_id, + name=name, + custom=custom, + delete=delete or Mock(), + ) + + +def _v2_name(created_at: int, suffix: str = "0123456789") -> str: + return f"{EMBEDDING_NAME_PREFIX_V2}{created_at}-{suffix}" + + +def _cap_error() -> LabelboxError: + return LabelboxError(f"{EMBEDDING_CAP_ERROR_SNIPPET}: 10") + + +def test_build_and_parse_embedding_name_round_trip(): + with patch( + "tests.embedding_cleanup.uuid.uuid4", + return_value=SimpleNamespace(hex="abcdef0123456789abcdef0123456789"), + ): + name = build_embedding_name(1_725_000_000.75) + + assert name == "sdk-int-ci-v2-1725000000-abcdef0123" + assert parse_embedding_created_at(name) == 1_725_000_000 + + +@pytest.mark.parametrize( + "name", + [ + "sdk-int-abcdef0123456789abcdef0123456789", + "sdk-int-ci-v2-1725000000-abcdef012", + "sdk-int-ci-v2-1725000000-abcdef01234", + "sdk-int-ci-v2-1725000000-abcdefghi0", + "foreign-ci-v2-1725000000-abcdef0123", + "sdk-int-ci-v2--1725000000-abcdef0123", + ], +) +def test_parse_embedding_created_at_rejects_non_v2_names(name): + assert parse_embedding_created_at(name) is None + + +def test_select_stale_embeddings_keeps_only_expired_custom_v2_names(): + now = 1_725_000_000 + stale = _embedding( + "stale", + _v2_name(now - EMBEDDING_STALE_TTL_SECONDS - 1), + ) + fresh = _embedding("fresh", _v2_name(now - 60)) + boundary = _embedding( + "boundary", _v2_name(now - EMBEDDING_STALE_TTL_SECONDS) + ) + legacy = _embedding("legacy", "sdk-int-abcdef0123456789abcdef0123456789") + non_custom = _embedding( + "non-custom", + _v2_name(now - EMBEDDING_STALE_TTL_SECONDS - 1), + custom=False, + ) + + assert select_stale_embeddings( + [fresh, stale, boundary, legacy, non_custom], now + ) == [stale] + + +def test_embedding_cap_error_matcher_is_fail_closed(): + assert is_embedding_cap_error(_cap_error()) + assert not is_embedding_cap_error(LabelboxError("permission denied")) + + +def test_workflow_boolean_input_is_parsed_defensively(): + assert _parse_boolean("true") is True + assert _parse_boolean("false") is False + + +def test_heal_retries_twice_then_creates_and_deletes_only_stale_v2(): + now_value = 1_725_000_000 + stale = _embedding( + "stale", + _v2_name(now_value - EMBEDDING_STALE_TTL_SECONDS - 1), + ) + fresh = _embedding("fresh", _v2_name(now_value - 1)) + legacy = _embedding("legacy", "sdk-int-abcdef0123456789abcdef0123456789") + non_custom = _embedding( + "non-custom", + _v2_name(now_value - EMBEDDING_STALE_TTL_SECONDS - 1), + custom=False, + ) + created = object() + create = Mock(side_effect=[_cap_error(), _cap_error(), created]) + list_embeddings = Mock(side_effect=[[stale, fresh, legacy, non_custom], []]) + delete = Mock() + sleep = Mock() + + result = create_embedding_with_heal( + create_embedding=create, + list_embeddings=list_embeddings, + delete_embedding=delete, + sleep=sleep, + now=Mock(return_value=now_value), + retry_delay=Mock(return_value=4), + print_fn=Mock(), + ) + + assert result is created + assert create.call_count == 3 + assert list_embeddings.call_count == 2 + assert delete.call_args_list == [call(stale)] + assert sleep.call_args_list == [call(4), call(4)] + + +def test_heal_third_cap_error_is_terminal_without_another_sweep_or_sleep(): + final_error = _cap_error() + create = Mock(side_effect=[_cap_error(), _cap_error(), final_error]) + list_embeddings = Mock(return_value=[]) + delete = Mock() + sleep = Mock() + print_fn = Mock() + + with pytest.raises(LabelboxError) as raised: + create_embedding_with_heal( + create_embedding=create, + list_embeddings=list_embeddings, + delete_embedding=delete, + sleep=sleep, + now=Mock(return_value=1_725_000_000), + retry_delay=Mock(return_value=4), + print_fn=print_fn, + ) + + assert raised.value is final_error + assert create.call_count == 3 + assert list_embeddings.call_count == 2 + delete.assert_not_called() + assert sleep.call_count == 2 + assert any( + "live fixtures and/or legacy/foreign names" in args[0] + for args, _ in print_fn.call_args_list + ) + + +def test_heal_non_cap_error_is_reraised_without_side_effects(): + non_cap_error = LabelboxError("permission denied") + create = Mock(side_effect=non_cap_error) + list_embeddings = Mock() + delete = Mock() + sleep = Mock() + + with pytest.raises(LabelboxError) as raised: + create_embedding_with_heal( + create_embedding=create, + list_embeddings=list_embeddings, + delete_embedding=delete, + sleep=sleep, + ) + + assert raised.value is non_cap_error + create.assert_called_once_with() + list_embeddings.assert_not_called() + delete.assert_not_called() + sleep.assert_not_called() + + +def test_bootstrap_client_uses_staging_endpoints_and_checks_effective_host(): + staging_client = SimpleNamespace(rest_endpoint=STAGING_REST_ENDPOINT) + client_factory = Mock(return_value=staging_client) + + assert ( + create_staging_client( + api_key="staging-key", client_factory=client_factory + ) + is staging_client + ) + client_factory.assert_called_once_with( + api_key="staging-key", + endpoint=STAGING_GRAPHQL_ENDPOINT, + rest_endpoint=STAGING_REST_ENDPOINT, + ) + + client_factory.return_value = SimpleNamespace( + rest_endpoint="https://api.labelbox.com/api/v1" + ) + with pytest.raises(RuntimeError, match="effective REST endpoint"): + create_staging_client( + api_key="staging-key", client_factory=client_factory + ) + + +def test_bootstrap_host_check_fails_before_listing(): + client = SimpleNamespace( + rest_endpoint="https://api.labelbox.com/api/v1", + get_embeddings=Mock(), + ) + + with pytest.raises(RuntimeError, match="effective REST endpoint"): + run_cleanup(client, dry_run=True, now=1_725_000_000) + + client.get_embeddings.assert_not_called() + + +def test_cleanup_dry_run_lists_without_deleting(capsys): + legacy = _embedding("legacy", "sdk-int-abcdef0123456789abcdef0123456789") + client = SimpleNamespace( + rest_endpoint=STAGING_REST_ENDPOINT, + get_embeddings=Mock(return_value=[legacy]), + ) + + assert run_cleanup(client, dry_run=True, now=1_725_000_000) == 0 + legacy.delete.assert_not_called() + assert "candidate id=legacy" in capsys.readouterr().out + + +def test_cleanup_real_run_attempts_every_candidate_and_reports_failures(): + failed_delete = Mock(side_effect=LabelboxError("delete failed")) + first = _embedding( + "first", + "sdk-int-abcdef0123456789abcdef0123456789", + delete=failed_delete, + ) + second = _embedding("second", "sdk-int-0123456789abcdef0123456789abcdef") + client = SimpleNamespace( + rest_endpoint=STAGING_REST_ENDPOINT, + get_embeddings=Mock(return_value=[first, second]), + ) + + assert run_cleanup(client, dry_run=False, now=1_725_000_000) == 1 + first.delete.assert_called_once_with() + second.delete.assert_called_once_with()