diff --git a/durabletask-azuremanaged/durabletask/azuremanaged/preview/sandboxes/profile_builder.py b/durabletask-azuremanaged/durabletask/azuremanaged/preview/sandboxes/profile_builder.py index ca455eed..6fe9cc2a 100644 --- a/durabletask-azuremanaged/durabletask/azuremanaged/preview/sandboxes/profile_builder.py +++ b/durabletask-azuremanaged/durabletask/azuremanaged/preview/sandboxes/profile_builder.py @@ -7,7 +7,6 @@ from durabletask.azuremanaged.internal import sandbox_service_pb2 as pb from durabletask.azuremanaged.preview.sandboxes.helpers import ( SandboxActivity, - activities_overlap, format_activity, normalize_required, resolve_activities, @@ -91,22 +90,109 @@ def _build_sandbox_worker_profile( return worker_profile +class _ActivityOwnerSlot: + """Earliest worker profiles recorded for a single activity overlap bucket. + + Only two owners ever need to be retained to answer "which worker profile + first claimed an activity in this bucket, ignoring `worker_profile_id`": + the very first owner, plus the first owner that differs from it. + """ + + __slots__ = ("_first", "_first_other") + + def __init__(self) -> None: + self._first: Optional[tuple[int, str]] = None + self._first_other: Optional[tuple[int, str]] = None + + def add(self, order: int, worker_profile_id: str) -> None: + if self._first is None: + self._first = (order, worker_profile_id) + elif self._first_other is None and worker_profile_id != self._first[1]: + self._first_other = (order, worker_profile_id) + + def first_owner_other_than(self, worker_profile_id: str) -> Optional[tuple[int, str]]: + if self._first is not None and self._first[1] != worker_profile_id: + return self._first + return self._first_other + + +class _ActivityOwnerIndex: + """Indexes activity ownership so overlap checks cost O(1) per activity. + + Reproduces the semantics of + :func:`durabletask.azuremanaged.preview.sandboxes.helpers.activities_overlap` + exactly: activity names are compared case insensitively, an unversioned + activity overlaps every version of the same name, and two activities with + the same name overlap when their explicit versions are equal. Registration + order is tracked so the reported conflict matches the first overlapping + activity, exactly as a linear scan would report it. + """ + + def __init__(self) -> None: + self._registration_count = 0 + self._by_name: dict[str, _ActivityOwnerSlot] = {} + self._by_name_and_version: dict[tuple[str, Optional[str]], _ActivityOwnerSlot] = {} + + def find_conflicting_profile( + self, + activity: SandboxActivity, + worker_profile_id: str) -> Optional[str]: + """Return the profile that first claimed an overlapping activity, if any.""" + name_key = activity.name.casefold() + if activity.version is None: + # An unversioned activity overlaps every version of the same name. + owner = _first_owner_other_than(self._by_name.get(name_key), worker_profile_id) + else: + # A versioned activity overlaps the same version plus any + # unversioned registration of the same name. + owner = _earlier_owner( + _first_owner_other_than( + self._by_name_and_version.get((name_key, None)), worker_profile_id), + _first_owner_other_than( + self._by_name_and_version.get((name_key, activity.version)), worker_profile_id)) + return None if owner is None else owner[1] + + def add(self, activity: SandboxActivity, worker_profile_id: str) -> None: + """Record `worker_profile_id` as an owner of `activity`.""" + name_key = activity.name.casefold() + order = self._registration_count + self._registration_count += 1 + self._by_name.setdefault(name_key, _ActivityOwnerSlot()).add(order, worker_profile_id) + self._by_name_and_version.setdefault( + (name_key, activity.version), _ActivityOwnerSlot()).add(order, worker_profile_id) + + +def _first_owner_other_than( + slot: Optional[_ActivityOwnerSlot], + worker_profile_id: str) -> Optional[tuple[int, str]]: + return None if slot is None else slot.first_owner_other_than(worker_profile_id) + + +def _earlier_owner( + left: Optional[tuple[int, str]], + right: Optional[tuple[int, str]]) -> Optional[tuple[int, str]]: + if left is None: + return right + if right is None: + return left + return left if left[0] <= right[0] else right + + def build_sandbox_worker_profiles() -> list[pb.SandboxWorkerProfile]: """Build sandbox worker_profiles from worker profile configuration.""" worker_profiles: list[pb.SandboxWorkerProfile] = [] - activity_owners: list[tuple[SandboxActivity, str]] = [] + activity_owners = _ActivityOwnerIndex() for profile in registered_sandbox_worker_profiles(): activities = resolve_activities(profile.activities) for activity in activities: - existing_profile = next((owner_profile for owner_activity, owner_profile in activity_owners - if activities_overlap(owner_activity, activity) - and owner_profile != profile.worker_profile_id), None) + existing_profile = activity_owners.find_conflicting_profile( + activity, profile.worker_profile_id) if existing_profile: raise ValueError( f"Sandbox activity '{format_activity(activity)}' is assigned to both worker profile " f"'{existing_profile}' and '{profile.worker_profile_id}'.") - activity_owners.append((activity, profile.worker_profile_id)) + activity_owners.add(activity, profile.worker_profile_id) worker_profiles.append(_build_sandbox_worker_profile( activities=activities, diff --git a/tests/durabletask-azuremanaged/test_sandboxes_extension.py b/tests/durabletask-azuremanaged/test_sandboxes_extension.py index 15991eeb..9d482bf7 100644 --- a/tests/durabletask-azuremanaged/test_sandboxes_extension.py +++ b/tests/durabletask-azuremanaged/test_sandboxes_extension.py @@ -2,6 +2,7 @@ # Licensed under the MIT License. import inspect +import random import threading import grpc @@ -20,6 +21,7 @@ SandboxWorkerProfileImageOptions, ) from durabletask.azuremanaged.preview.sandboxes.profile_builder import ( + _ActivityOwnerIndex, _build_sandbox_worker_profile, build_sandbox_worker_profiles, ) @@ -28,6 +30,7 @@ build_sandbox_worker_start, ) from durabletask.azuremanaged.preview.sandboxes.helpers import resolve_activities +from durabletask.azuremanaged.preview.sandboxes.helpers import activities_overlap from durabletask.azuremanaged.preview.sandboxes.helpers import SandboxActivity from durabletask.azuremanaged.internal import sandbox_service_pb2 as pb from durabletask.azuremanaged.internal import sandbox_service_pb2_grpc as stubs @@ -206,6 +209,9 @@ def configure(self, options: SandboxWorkerProfileOptions) -> None: try: build_sandbox_worker_profiles() except ValueError as ex: + assert str(ex) == ( + "Sandbox activity 'pytestoverlapremotehello' is assigned to both worker profile " + "'pytest-overlap-profile-a' and 'pytest-overlap-profile-b'.") assert "pytestoverlapremotehello" in str(ex) assert "pytest-overlap-profile-a" in str(ex) assert "pytest-overlap-profile-b" in str(ex) @@ -252,6 +258,128 @@ def configure(self, options: SandboxWorkerProfileOptions) -> None: sandbox_worker_profiles._worker_profiles.pop("pytest-version-profile-b", None) +def test_build_sandbox_worker_profiles_rejects_versioned_activity_overlapping_unversioned_owner() -> None: + @sandbox_worker_profile("pytest-unversioned-owner-profile-a") + class PytestUnversionedOwnerProfileA(SandboxWorkerProfile): + def configure(self, options: SandboxWorkerProfileOptions) -> None: + options.image.image_ref = "example.azurecr.io/python-worker-a:v1" + options.image.managed_identity_client_id = "image-pull-client-id" + options.scheduler_managed_identity_client_id = "scheduler-client-id" + options.add_activity("PytestUnversionedOwner", version=None) + + @sandbox_worker_profile("pytest-unversioned-owner-profile-b") + class PytestUnversionedOwnerProfileB(SandboxWorkerProfile): + def configure(self, options: SandboxWorkerProfileOptions) -> None: + options.image.image_ref = "example.azurecr.io/python-worker-b:v1" + options.image.managed_identity_client_id = "image-pull-client-id" + options.scheduler_managed_identity_client_id = "scheduler-client-id" + options.add_activity("pytestunversionedowner", version="v9") + + try: + try: + build_sandbox_worker_profiles() + except ValueError as ex: + assert str(ex) == ( + "Sandbox activity 'pytestunversionedowner@v9' is assigned to both worker profile " + "'pytest-unversioned-owner-profile-a' and 'pytest-unversioned-owner-profile-b'.") + else: + raise AssertionError( + "Expected an unversioned sandbox activity to overlap every version of the same name.") + finally: + sandbox_worker_profiles._worker_profiles.pop("pytest-unversioned-owner-profile-a", None) + sandbox_worker_profiles._worker_profiles.pop("pytest-unversioned-owner-profile-b", None) + + +def test_build_sandbox_worker_profiles_rejects_unversioned_activity_overlapping_versioned_owner() -> None: + @sandbox_worker_profile("pytest-versioned-owner-profile-a") + class PytestVersionedOwnerProfileA(SandboxWorkerProfile): + def configure(self, options: SandboxWorkerProfileOptions) -> None: + options.image.image_ref = "example.azurecr.io/python-worker-a:v1" + options.image.managed_identity_client_id = "image-pull-client-id" + options.scheduler_managed_identity_client_id = "scheduler-client-id" + options.add_activity("PytestVersionedOwner", version="v9") + + @sandbox_worker_profile("pytest-versioned-owner-profile-b") + class PytestVersionedOwnerProfileB(SandboxWorkerProfile): + def configure(self, options: SandboxWorkerProfileOptions) -> None: + options.image.image_ref = "example.azurecr.io/python-worker-b:v1" + options.image.managed_identity_client_id = "image-pull-client-id" + options.scheduler_managed_identity_client_id = "scheduler-client-id" + options.add_activity("pytestversionedowner", version=None) + + try: + try: + build_sandbox_worker_profiles() + except ValueError as ex: + assert str(ex) == ( + "Sandbox activity 'pytestversionedowner' is assigned to both worker profile " + "'pytest-versioned-owner-profile-a' and 'pytest-versioned-owner-profile-b'.") + else: + raise AssertionError( + "Expected an unversioned sandbox activity to overlap an existing versioned owner.") + finally: + sandbox_worker_profiles._worker_profiles.pop("pytest-versioned-owner-profile-a", None) + sandbox_worker_profiles._worker_profiles.pop("pytest-versioned-owner-profile-b", None) + + +def test_build_sandbox_worker_profiles_rejects_identical_activity_versions() -> None: + @sandbox_worker_profile("pytest-same-version-profile-a") + class PytestSameVersionProfileA(SandboxWorkerProfile): + def configure(self, options: SandboxWorkerProfileOptions) -> None: + options.image.image_ref = "example.azurecr.io/python-worker-a:v1" + options.image.managed_identity_client_id = "image-pull-client-id" + options.scheduler_managed_identity_client_id = "scheduler-client-id" + options.add_activity("PytestSameVersionActivity", version="v3") + + @sandbox_worker_profile("pytest-same-version-profile-b") + class PytestSameVersionProfileB(SandboxWorkerProfile): + def configure(self, options: SandboxWorkerProfileOptions) -> None: + options.image.image_ref = "example.azurecr.io/python-worker-b:v1" + options.image.managed_identity_client_id = "image-pull-client-id" + options.scheduler_managed_identity_client_id = "scheduler-client-id" + options.add_activity("pytestsameversionactivity", version="v3") + + try: + try: + build_sandbox_worker_profiles() + except ValueError as ex: + assert str(ex) == ( + "Sandbox activity 'pytestsameversionactivity@v3' is assigned to both worker profile " + "'pytest-same-version-profile-a' and 'pytest-same-version-profile-b'.") + else: + raise AssertionError("Expected identical sandbox activity versions to overlap.") + finally: + sandbox_worker_profiles._worker_profiles.pop("pytest-same-version-profile-a", None) + sandbox_worker_profiles._worker_profiles.pop("pytest-same-version-profile-b", None) + + +def test_activity_owner_index_matches_linear_overlap_scan() -> None: + registrations = [ + (SandboxActivity(name, version), worker_profile_id) + for name in ("Alpha", "alpha", "BETA", "Gamma") + for version in (None, "v1", "v2", "V1") + for worker_profile_id in ("profile-a", "profile-b", "profile-c") + ] + + for seed in range(25): + shuffled = list(registrations) + random.Random(seed).shuffle(shuffled) + + index = _ActivityOwnerIndex() + owners: list[tuple[SandboxActivity, str]] = [] + for activity, worker_profile_id in shuffled: + expected = next( + (owner_profile for owner_activity, owner_profile in owners + if activities_overlap(owner_activity, activity) + and owner_profile != worker_profile_id), + None) + + assert index.find_conflicting_profile(activity, worker_profile_id) == expected + + owners.append((activity, worker_profile_id)) + index.add(activity, worker_profile_id) + + def test_profile_options_add_activity_accepts_callable() -> None: def pytest_callable_remote_hello(_ctx, value): return value