diff --git a/python/simpler/__init__.py b/python/simpler/__init__.py index 553daa242..bf67a2b72 100644 --- a/python/simpler/__init__.py +++ b/python/simpler/__init__.py @@ -41,12 +41,13 @@ "TIMING", "get_current_config", "get_logger", + "comm_endpoints", "task_interface", ] # name -> (module, attribute). Resolved by __getattr__ on first access. _LAZY_ATTRS = {"Worker": (f"{__name__}.worker", "Worker")} -_LAZY_SUBMODULES = ("task_interface",) +_LAZY_SUBMODULES = ("comm_endpoints", "task_interface") def __getattr__(name: str) -> Any: diff --git a/python/simpler/comm_endpoints.py b/python/simpler/comm_endpoints.py new file mode 100644 index 000000000..86a64484e --- /dev/null +++ b/python/simpler/comm_endpoints.py @@ -0,0 +1,1023 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Endpoint selectors, registry resolution, and W2 backend planning.""" + +from __future__ import annotations + +import ipaddress +import re +from collections.abc import Sequence +from dataclasses import dataclass +from enum import Enum +from typing import Protocol + +from _task_interface import BackendKind # pyright: ignore[reportMissingImports] + +EndpointId = int +NodeScopeId = int + + +class EndpointDeployment(str, Enum): + HOST_CPU = "HOST_CPU" + DEVICE_AICORE = "DEVICE_AICORE" + DEVICE_AICPU = "DEVICE_AICPU" + + +HOST_CPU = EndpointDeployment.HOST_CPU +DEVICE_AICORE = EndpointDeployment.DEVICE_AICORE +DEVICE_AICPU = EndpointDeployment.DEVICE_AICPU + + +class EndpointSelectorKind(str, Enum): + AT = "AT" + UNDER = "UNDER" + + +@dataclass(frozen=True) +class EndpointSelector: + kind: EndpointSelectorKind + path: str + deployment: EndpointDeployment + + def __post_init__(self) -> None: + kind, path, deployment = _normalize_selector_values(self.kind, self.path, self.deployment) + object.__setattr__(self, "kind", kind) + object.__setattr__(self, "path", path) + object.__setattr__(self, "deployment", deployment) + + +def at(path: str, deployment: EndpointDeployment | str) -> EndpointSelector: + kind, normalized_path, normalized_deployment = _normalize_selector_values( + EndpointSelectorKind.AT, + path, + deployment, + ) + return EndpointSelector(kind, normalized_path, normalized_deployment) + + +def under(path: str, deployment: EndpointDeployment | str) -> EndpointSelector: + kind, normalized_path, normalized_deployment = _normalize_selector_values( + EndpointSelectorKind.UNDER, + path, + deployment, + ) + return EndpointSelector(kind, normalized_path, normalized_deployment) + + +def _normalize_selector_values( + kind: EndpointSelectorKind | str, path: str, deployment: EndpointDeployment | str +) -> tuple[EndpointSelectorKind, str, EndpointDeployment]: + try: + kind_value = EndpointSelectorKind(kind) + except ValueError as exc: + raise ValueError(f"invalid endpoint selector kind {kind!r}") from exc + if not isinstance(path, str) or not path: + raise ValueError("endpoint path must be a non-empty string") + if any(segment == "" for segment in path.split("/")): + raise ValueError(f"endpoint path must not contain empty segments: {path!r}") + try: + deployment_value = EndpointDeployment(deployment) + except ValueError as exc: + raise ValueError(f"invalid endpoint deployment {deployment!r}") from exc + return kind_value, path, deployment_value + + +@dataclass(frozen=True) +class EndpointPathSegment: + level: int + index: int | None = None + + +@dataclass(frozen=True) +class ParsedEndpointPath: + segments: tuple[EndpointPathSegment, ...] + text: str + + @property + def sort_key(self) -> tuple[tuple[int, int], ...]: + return tuple((segment.level, -1 if segment.index is None else segment.index) for segment in self.segments) + + +_PATH_SEGMENT_RE = re.compile(r"^L(?P[0-9]+)(?:\[(?P[0-9]+)\])?$") + + +def parse_endpoint_path(path: str, *, root_level: int) -> ParsedEndpointPath: + if not isinstance(path, str) or not path: + raise ValueError("endpoint path must be a non-empty string") + raw_segments = path.split("/") + if any(segment == "" for segment in raw_segments): + raise ValueError(f"invalid endpoint path {path!r}: empty segment") + segments: list[EndpointPathSegment] = [] + for i, raw in enumerate(raw_segments): + match = _PATH_SEGMENT_RE.match(raw) + if match is None: + raise ValueError(f"invalid endpoint path segment {raw!r} in {path!r}") + level = int(match.group("level")) + index_s = match.group("index") + index = None if index_s is None else int(index_s) + if i == 0: + if level != root_level or index is not None: + raise ValueError(f"endpoint path {path!r} must start with root L{root_level}") + elif index is None: + raise ValueError(f"child endpoint path segment {raw!r} in {path!r} must include an index") + segments.append(EndpointPathSegment(level=level, index=index)) + return ParsedEndpointPath(segments=tuple(segments), text=path) + + +def _format_worker_path(level: int, *, parent_path: str | None = None, index: int | None = None) -> str: + segment = f"L{int(level)}" if index is None else f"L{int(level)}[{int(index)}]" + return segment if parent_path is None else f"{parent_path}/{segment}" + + +def _normalize_node_identity(host: str) -> str: + normalized = str(host).strip().lower() + if normalized == "localhost": + return "local" + try: + address = ipaddress.ip_address(normalized) + except ValueError: + return normalized + if address.is_loopback: + return "local" + return str(address).lower() + + +@dataclass(frozen=True) +class _EndpointTopologyEntry: + path: str + deployment: EndpointDeployment + node_identity: str + + +@dataclass(frozen=True) +class _EndpointTopologySnapshot: + root_level: int + session_instance_id: bytes + entries: tuple[_EndpointTopologyEntry, ...] + + +@dataclass(frozen=True) +class SingleOwner: + provider: EndpointSelector | None = None + + +@dataclass(frozen=True) +class EndpointIdentity: + session_instance_id: bytes + registry_epoch: int + endpoint_id: EndpointId + + +@dataclass(frozen=True) +class EndpointRecord: + identity: EndpointIdentity + path: str + deployment: EndpointDeployment + node_scope_id: NodeScopeId + + @property + def endpoint_id(self) -> EndpointId: + return self.identity.endpoint_id + + +@dataclass(frozen=True) +class ResolvedSingleOwner: + provider_endpoint: EndpointIdentity | None + + +@dataclass(frozen=True) +class ResolvedRegionSpec: + members: tuple[EndpointRecord, ...] + topology: ResolvedSingleOwner + + +class EndpointResolveReason(str, Enum): + PATH_NOT_FOUND = "PATH_NOT_FOUND" + ENDPOINT_NOT_REGISTERED = "ENDPOINT_NOT_REGISTERED" + EMPTY_UNDER_SELECTOR = "EMPTY_UNDER_SELECTOR" + DUPLICATE_ENDPOINT = "DUPLICATE_ENDPOINT" + INVALID_PATH = "INVALID_PATH" + PROVIDER_NOT_SINGLE_ENDPOINT = "PROVIDER_NOT_SINGLE_ENDPOINT" + + +class EndpointResolveError(ValueError): + def __init__(self, reason: EndpointResolveReason, message: str, offending: Sequence[str] = ()) -> None: + self.reason = reason + self.message = message + self.offending = tuple(offending) + super().__init__(message) + + +class EndpointRegistry: + def __init__( + self, + *, + root_level: int, + session_instance_id: bytes, + registry_epoch: int, + records: Sequence[EndpointRecord], + ) -> None: + self.root_level = int(root_level) + self.session_instance_id = bytes(session_instance_id) + self.registry_epoch = int(registry_epoch) + self._records = tuple(records) + self._by_id = {record.endpoint_id: record for record in self._records} + self._by_identity = {record.identity: record for record in self._records} + self._by_key = {(record.path, record.deployment): record for record in self._records} + self._known_paths = {record.path for record in self._records} + self._parsed_paths = {path: self._parse(path) for path in self._known_paths} + + @classmethod + def from_snapshot( + cls, + snapshot: _EndpointTopologySnapshot, + *, + registry_epoch: int, + ) -> EndpointRegistry: + entries = tuple(snapshot.entries) + session_instance_id = bytes(snapshot.session_instance_id) + node_scopes = {"local": 0} + next_node_scope_id = 1 + records: list[EndpointRecord] = [] + seen: set[tuple[str, EndpointDeployment]] = set() + for endpoint_id, entry in enumerate(entries): + deployment = EndpointDeployment(entry.deployment) + key = (entry.path, deployment) + if key in seen: + raise ValueError(f"duplicate endpoint topology entry: {entry.path} {deployment.value}") + seen.add(key) + node_identity = _normalize_node_identity(entry.node_identity) + node_scope_id = node_scopes.get(node_identity) + if node_scope_id is None: + node_scope_id = next_node_scope_id + node_scopes[node_identity] = node_scope_id + next_node_scope_id += 1 + identity = EndpointIdentity( + session_instance_id=session_instance_id, + registry_epoch=int(registry_epoch), + endpoint_id=endpoint_id, + ) + records.append( + EndpointRecord( + identity=identity, + path=entry.path, + deployment=deployment, + node_scope_id=node_scope_id, + ) + ) + return cls( + root_level=int(snapshot.root_level), + session_instance_id=session_instance_id, + registry_epoch=int(registry_epoch), + records=records, + ) + + @property + def records(self) -> tuple[EndpointRecord, ...]: + return self._records + + def resolve_members(self, selectors: Sequence[EndpointSelector]) -> tuple[EndpointRecord, ...]: + resolved: list[EndpointRecord] = [] + seen: dict[EndpointIdentity, EndpointRecord] = {} + for selector in selectors: + for record in self._resolve_selector(selector): + duplicate = seen.get(record.identity) + if duplicate is not None: + label = _endpoint_label(record) + raise EndpointResolveError( + EndpointResolveReason.DUPLICATE_ENDPOINT, + f"duplicate endpoint in region members: {label}", + (label,), + ) + seen[record.identity] = record + resolved.append(record) + return tuple(resolved) + + def resolve_region_spec(self, members: Sequence[EndpointSelector], topology: SingleOwner) -> ResolvedRegionSpec: + resolved_members = self.resolve_members(members) + if not isinstance(topology, SingleOwner): + raise TypeError("W2 region planning supports SingleOwner topology only") + provider_endpoint = None + if topology.provider is not None: + provider_records = self._resolve_provider(topology.provider) + provider_endpoint = provider_records[0].identity + return ResolvedRegionSpec( + members=resolved_members, + topology=ResolvedSingleOwner(provider_endpoint=provider_endpoint), + ) + + def same_node( + self, a: EndpointRecord | EndpointIdentity | EndpointId, b: EndpointRecord | EndpointIdentity | EndpointId + ) -> bool: + return self._record_for(a).node_scope_id == self._record_for(b).node_scope_id + + def all_same_node(self, members: Sequence[EndpointRecord]) -> bool: + if len(members) <= 1: + return True + first = members[0].node_scope_id + return all(member.node_scope_id == first for member in members) + + def _resolve_provider(self, selector: EndpointSelector) -> tuple[EndpointRecord, ...]: + try: + records = self._resolve_selector(selector) + except EndpointResolveError as exc: + raise self._provider_error(selector) from exc + if len(records) != 1: + raise self._provider_error(selector, records) + return records + + def _provider_error( + self, selector: EndpointSelector, records: Sequence[EndpointRecord] = () + ) -> EndpointResolveError: + labels = tuple(_endpoint_label(record) for record in records) or (_selector_label(selector),) + return EndpointResolveError( + EndpointResolveReason.PROVIDER_NOT_SINGLE_ENDPOINT, + f"SingleOwner provider must resolve to exactly one endpoint: {_selector_label(selector)}", + labels, + ) + + def _resolve_selector(self, selector: EndpointSelector) -> tuple[EndpointRecord, ...]: + if not isinstance(selector, EndpointSelector): + raise TypeError("region members must be EndpointSelector values") + parsed = self._parse_or_error(selector.path) + if selector.path not in self._known_paths: + raise EndpointResolveError( + EndpointResolveReason.PATH_NOT_FOUND, + f"endpoint path not found: {_selector_label(selector)}", + (_selector_label(selector),), + ) + if selector.kind is EndpointSelectorKind.AT: + record = self._by_key.get((selector.path, selector.deployment)) + if record is None: + raise EndpointResolveError( + EndpointResolveReason.ENDPOINT_NOT_REGISTERED, + f"endpoint is not registered: {_selector_label(selector)}", + (_selector_label(selector),), + ) + return (record,) + if selector.kind is EndpointSelectorKind.UNDER: + records = [ + record + for record in self._records + if record.deployment is selector.deployment and self._is_descendant(record.path, parsed) + ] + records.sort(key=lambda record: self._parsed_paths[record.path].sort_key) + if not records: + raise EndpointResolveError( + EndpointResolveReason.EMPTY_UNDER_SELECTOR, + f"endpoint selector expanded to no endpoints: {_selector_label(selector)}", + (_selector_label(selector),), + ) + return tuple(records) + raise TypeError(f"unsupported endpoint selector kind: {selector.kind!r}") + + def _is_descendant(self, path: str, parent: ParsedEndpointPath) -> bool: + parsed = self._parsed_paths[path] + parent_len = len(parent.segments) + return len(parsed.segments) > parent_len and parsed.segments[:parent_len] == parent.segments + + def _record_for(self, endpoint: EndpointRecord | EndpointIdentity | EndpointId) -> EndpointRecord: + if isinstance(endpoint, EndpointRecord): + return endpoint + if isinstance(endpoint, EndpointIdentity): + record = self._by_identity.get(endpoint) + if record is None: + raise ValueError( + "unknown EndpointIdentity in registry " + f"session={self.session_instance_id!r} epoch={self.registry_epoch}" + ) + return record + try: + return self._by_id[int(endpoint)] + except KeyError as exc: + raise ValueError(f"unknown endpoint_id {endpoint!r} in registry epoch {self.registry_epoch}") from exc + + def _parse(self, path: str) -> ParsedEndpointPath: + return parse_endpoint_path(path, root_level=self.root_level) + + def _parse_or_error(self, path: str) -> ParsedEndpointPath: + try: + return self._parse(path) + except ValueError as exc: + raise EndpointResolveError( + EndpointResolveReason.INVALID_PATH, + f"invalid endpoint path for root L{self.root_level}: {path!r}", + (path,), + ) from exc + + +class RegionPartKind(str, Enum): + PAYLOAD = "PAYLOAD" + COUNTER = "COUNTER" + + +class AttachmentRole(str, Enum): + PROVIDER = "PROVIDER" + CONSUMER = "CONSUMER" + + +class AdapterKind(str, Enum): + DIRECT_MAP = "DIRECT_MAP" + DEVICE_PEER = "DEVICE_PEER" + OWNER_DELEGATED_COPY = "OWNER_DELEGATED_COPY" + EXPLICIT_TRANSFER = "EXPLICIT_TRANSFER" + COLLECTIVE = "COLLECTIVE" + + +class AdapterProfile(str, Enum): + HOST_SVM_MAP = "HOST_SVM_MAP" + HOST_VMM_COPY = "HOST_VMM_COPY" + DEVICE_VMM_PEER_IMPORT = "DEVICE_VMM_PEER_IMPORT" + DEVICE_FABRIC_V2_PEER_IMPORT = "DEVICE_FABRIC_V2_PEER_IMPORT" + HOST_SHM_MAP = "HOST_SHM_MAP" + REMOTE_COPY = "REMOTE_COPY" + + +class RegionAccessReasonCode(str, Enum): + SUPPORTED = "SUPPORTED" + UNSUPPORTED_BACKEND_KIND = "UNSUPPORTED_BACKEND_KIND" + UNSUPPORTED_ENDPOINT_RELATION = "UNSUPPORTED_ENDPOINT_RELATION" + NO_IMPLEMENTED_DIRECT_MAP_PROBE = "NO_IMPLEMENTED_DIRECT_MAP_PROBE" + NO_COPY_BACKEND = "NO_COPY_BACKEND" + STATIC_UNSUPPORTED = "STATIC_UNSUPPORTED" + + +@dataclass(frozen=True) +class RegionLayoutSpec: + payload_bytes: int + counter_bytes: int + + +@dataclass(frozen=True) +class RegionAccessQuery: + topology: str + part: RegionPartKind + backend_kind: BackendKind + provider: EndpointRecord + consumer: EndpointRecord + layout: RegionLayoutSpec + same_node: bool + platform: str | None = None + runtime: str | None = None + + +@dataclass(frozen=True) +class RegionAccessDiagnostics: + reason_code: RegionAccessReasonCode + message: str + provider_label: str | None = None + consumer_label: str | None = None + platform: str | None = None + runtime: str | None = None + backend_kind: BackendKind | None = None + part: RegionPartKind | None = None + adapter_kind: AdapterKind | None = None + adapter_profile: AdapterProfile | None = None + + +@dataclass(frozen=True) +class RegionAccessDecision: + supported: bool + reason: str | None = None + diagnostics: RegionAccessDiagnostics | None = None + + +@dataclass(frozen=True) +class MemberAttachmentPlan: + member: EndpointIdentity + role: AttachmentRole + adapter_kind: AdapterKind | None + adapter_profile: AdapterProfile | None + + +@dataclass(frozen=True) +class RegionPartPlan: + part: RegionPartKind + backend_kind: BackendKind + attachments: tuple[MemberAttachmentPlan, ...] + + +@dataclass(frozen=True) +class SingleOwnerPlan: + provider_endpoint: EndpointIdentity + + +@dataclass(frozen=True) +class BackendPlan: + ordered_members: tuple[EndpointIdentity, ...] + payload: RegionPartPlan + counter: RegionPartPlan + topology_plan: SingleOwnerPlan + + +class RegionAccessService(Protocol): + def evaluate_region_access( + self, + query: RegionAccessQuery, + candidate: _AdapterCandidate, + ) -> RegionAccessDecision: ... + + +class BackendUnsupportedReason(str, Enum): + PROVIDER_NOT_IN_MEMBERS = "PROVIDER_NOT_IN_MEMBERS" + NO_DEFAULT_PROVIDER = "NO_DEFAULT_PROVIDER" + ADAPTER_UNSUPPORTED = "ADAPTER_UNSUPPORTED" + UNSUPPORTED_DEPLOYMENT_COMBINATION = "UNSUPPORTED_DEPLOYMENT_COMBINATION" + + +@dataclass(frozen=True) +class AdapterAttempt: + part: RegionPartKind + member: EndpointRecord + backend_kind: BackendKind + adapter_kind: AdapterKind + adapter_profile: AdapterProfile + reason: str | None = None + + +@dataclass(frozen=True) +class UnsupportedRegionPlan: + reason: BackendUnsupportedReason + message: str + offending_endpoints: tuple[EndpointRecord, ...] = () + attempted_adapters: tuple[AdapterAttempt, ...] = () + + +@dataclass(frozen=True) +class _AdapterCandidate: + kind: AdapterKind + profile: AdapterProfile + unsupported_reason: str | None = None + + +class DefaultRegionAccessService: + def evaluate_region_access( + self, + query: RegionAccessQuery, + candidate: _AdapterCandidate, + ) -> RegionAccessDecision: + if candidate.profile is AdapterProfile.HOST_VMM_COPY: + return self._evaluate_host_vmm_copy(query, candidate) + if candidate.profile is AdapterProfile.HOST_SVM_MAP: + return _region_access_unsupported( + RegionAccessReasonCode.NO_IMPLEMENTED_DIRECT_MAP_PROBE, + "direct map probe is not implemented for this region access profile", + query, + candidate, + ) + if candidate.unsupported_reason is not None: + return _region_access_unsupported( + RegionAccessReasonCode.NO_COPY_BACKEND, + candidate.unsupported_reason, + query, + candidate, + ) + return _region_access_unsupported( + RegionAccessReasonCode.STATIC_UNSUPPORTED, + "region access profile is not supported by the default service", + query, + candidate, + ) + + def _evaluate_host_vmm_copy( + self, + query: RegionAccessQuery, + candidate: _AdapterCandidate, + ) -> RegionAccessDecision: + if query.backend_kind is not BackendKind.VMM_WINDOW: + return _region_access_unsupported( + RegionAccessReasonCode.UNSUPPORTED_BACKEND_KIND, + "host VMM copy requires a VMM window backend", + query, + candidate, + ) + if query.provider.deployment not in (DEVICE_AICORE, DEVICE_AICPU): + return _region_access_unsupported( + RegionAccessReasonCode.UNSUPPORTED_ENDPOINT_RELATION, + "host VMM copy requires a device provider", + query, + candidate, + ) + if query.consumer.deployment is not HOST_CPU or not query.same_node: + return _region_access_unsupported( + RegionAccessReasonCode.UNSUPPORTED_ENDPOINT_RELATION, + "host VMM copy requires a same-node host consumer", + query, + candidate, + ) + return _region_access_supported(query, candidate) + + +class StaticRegionAccessService: + def __init__( + self, + decisions: dict[tuple[BackendKind, RegionPartKind, AdapterKind, AdapterProfile], RegionAccessDecision | bool] + | None = None, + ) -> None: + self._decisions: dict[ + tuple[BackendKind, RegionPartKind, AdapterKind, AdapterProfile], RegionAccessDecision + ] = {} + for key, decision in (decisions or {}).items(): + normalized = _normalize_region_access_key(key) + if isinstance(decision, RegionAccessDecision): + self._decisions[normalized] = decision + else: + self._decisions[normalized] = RegionAccessDecision(bool(decision)) + + def evaluate_region_access( + self, + query: RegionAccessQuery, + candidate: _AdapterCandidate, + ) -> RegionAccessDecision: + key = (query.backend_kind, query.part, candidate.kind, candidate.profile) + decision = self._decisions.get(key) + if decision is not None: + if decision.supported and decision.diagnostics is None: + return _region_access_supported(query, candidate) + if not decision.supported and decision.diagnostics is None: + return _region_access_unsupported( + RegionAccessReasonCode.STATIC_UNSUPPORTED, + decision.reason or "region access is not supported by the static service", + query, + candidate, + ) + return decision + return _region_access_unsupported( + RegionAccessReasonCode.STATIC_UNSUPPORTED, + "region access is not supported by the static service", + query, + candidate, + ) + + +class BackendResolver: + def __init__(self, registry: EndpointRegistry, region_access: RegionAccessService) -> None: + self._registry = registry + self._region_access = region_access + + def plan(self, resolved: ResolvedRegionSpec, layout: RegionLayoutSpec) -> BackendPlan | UnsupportedRegionPlan: + self._validate_layout(layout) + members = resolved.members + provider = self._choose_provider(resolved) + if isinstance(provider, UnsupportedRegionPlan): + return provider + backend_kind = _backend_kind_for_provider(provider) + payload = self._plan_part(RegionPartKind.PAYLOAD, backend_kind, provider, members, layout) + if isinstance(payload, UnsupportedRegionPlan): + return payload + counter = self._plan_part(RegionPartKind.COUNTER, backend_kind, provider, members, layout) + if isinstance(counter, UnsupportedRegionPlan): + return counter + return BackendPlan( + ordered_members=tuple(member.identity for member in members), + payload=payload, + counter=counter, + # Only SingleOwnerPlan is defined here; topology-specific plans extend this field. + topology_plan=SingleOwnerPlan(provider_endpoint=provider.identity), + ) + + def _validate_layout(self, layout: RegionLayoutSpec) -> None: + if not isinstance(layout, RegionLayoutSpec): + raise TypeError("BackendResolver.plan expects RegionLayoutSpec") + if int(layout.payload_bytes) < 0 or int(layout.counter_bytes) < 0: + raise ValueError("RegionLayoutSpec byte sizes must be non-negative") + + def _choose_provider(self, resolved: ResolvedRegionSpec) -> EndpointRecord | UnsupportedRegionPlan: + members = resolved.members + by_identity = {member.identity: member for member in members} + provider_identity = resolved.topology.provider_endpoint + if provider_identity is not None: + provider = by_identity.get(provider_identity) + if provider is None: + try: + provider_record = self._registry._record_for(provider_identity) + except ValueError: + provider_record = None + return _unsupported( + BackendUnsupportedReason.PROVIDER_NOT_IN_MEMBERS, + "SingleOwner provider is not included in region members", + () if provider_record is None else (provider_record,), + ) + return provider + for deployment in (DEVICE_AICORE, DEVICE_AICPU, HOST_CPU): + for member in members: + if member.deployment is deployment: + return member + return _unsupported( + BackendUnsupportedReason.NO_DEFAULT_PROVIDER, + "no default SingleOwner provider is available", + ) + + def _plan_part( + self, + part: RegionPartKind, + backend_kind: BackendKind, + provider: EndpointRecord, + members: Sequence[EndpointRecord], + layout: RegionLayoutSpec, + ) -> RegionPartPlan | UnsupportedRegionPlan: + attachments: list[MemberAttachmentPlan] = [] + for member in members: + if member.identity == provider.identity: + attachments.append( + MemberAttachmentPlan( + member=member.identity, + role=AttachmentRole.PROVIDER, + adapter_kind=None, + adapter_profile=None, + ) + ) + continue + attachment = self._consumer_attachment(part, backend_kind, provider, member, layout) + if isinstance(attachment, UnsupportedRegionPlan): + return attachment + attachments.append(attachment) + return RegionPartPlan(part=part, backend_kind=backend_kind, attachments=tuple(attachments)) + + def _consumer_attachment( + self, + part: RegionPartKind, + backend_kind: BackendKind, + provider: EndpointRecord, + member: EndpointRecord, + layout: RegionLayoutSpec, + ) -> MemberAttachmentPlan | UnsupportedRegionPlan: + attempts: list[AdapterAttempt] = [] + same_node = self._registry.same_node(provider, member) + query = RegionAccessQuery( + topology="SingleOwner", + part=part, + backend_kind=backend_kind, + provider=provider, + consumer=member, + layout=layout, + same_node=same_node, + ) + for candidate in self._adapter_candidates(part, backend_kind, provider, member): + decision = self._region_access.evaluate_region_access(query, candidate) + if decision.supported: + return MemberAttachmentPlan( + member=member.identity, + role=AttachmentRole.CONSUMER, + adapter_kind=candidate.kind, + adapter_profile=candidate.profile, + ) + attempts.append(_attempt(part, member, backend_kind, candidate, decision.reason)) + if attempts: + return _unsupported( + BackendUnsupportedReason.ADAPTER_UNSUPPORTED, + "no adapter can attach region member to provider", + (provider, member), + attempted_adapters=attempts, + ) + return _unsupported( + BackendUnsupportedReason.UNSUPPORTED_DEPLOYMENT_COMBINATION, + "unsupported endpoint deployment combination for SingleOwner region", + (provider, member), + ) + + def _adapter_candidates( + self, part: RegionPartKind, backend_kind: BackendKind, provider: EndpointRecord, member: EndpointRecord + ) -> tuple[_AdapterCandidate, ...]: + same_node = self._registry.same_node(provider, member) + if backend_kind is BackendKind.VMM_WINDOW: + if member.deployment is HOST_CPU: + if same_node: + if part is RegionPartKind.PAYLOAD: + return ( + _AdapterCandidate( + AdapterKind.OWNER_DELEGATED_COPY, + AdapterProfile.HOST_VMM_COPY, + ), + ) + return ( + _AdapterCandidate( + AdapterKind.DIRECT_MAP, + AdapterProfile.HOST_SVM_MAP, + ), + _AdapterCandidate( + AdapterKind.OWNER_DELEGATED_COPY, + AdapterProfile.HOST_VMM_COPY, + ), + ) + return _remote_copy_candidates() + if member.deployment in (DEVICE_AICORE, DEVICE_AICPU): + if same_node: + return ( + _AdapterCandidate( + AdapterKind.DEVICE_PEER, + AdapterProfile.DEVICE_VMM_PEER_IMPORT, + ), + ) + return ( + _AdapterCandidate( + AdapterKind.DEVICE_PEER, + AdapterProfile.DEVICE_FABRIC_V2_PEER_IMPORT, + unsupported_reason="device fabric peer import is not available for this endpoint", + ), + *_remote_copy_candidates(), + ) + return () + if backend_kind is BackendKind.POSIX_SHM: + if member.deployment is HOST_CPU: + if same_node: + return ( + _AdapterCandidate( + AdapterKind.DIRECT_MAP, + AdapterProfile.HOST_SHM_MAP, + ), + ) + return _remote_copy_candidates() + if member.deployment in (DEVICE_AICORE, DEVICE_AICPU): + return ( + _AdapterCandidate( + AdapterKind.EXPLICIT_TRANSFER, + AdapterProfile.REMOTE_COPY, + unsupported_reason="explicit transfer materializer is not implemented yet", + ), + ) + return () + return () + + +def _remote_copy_candidates() -> tuple[_AdapterCandidate, ...]: + return ( + _AdapterCandidate( + AdapterKind.OWNER_DELEGATED_COPY, + AdapterProfile.REMOTE_COPY, + unsupported_reason="remote copy materializer is not implemented yet", + ), + _AdapterCandidate( + AdapterKind.EXPLICIT_TRANSFER, + AdapterProfile.REMOTE_COPY, + unsupported_reason="explicit transfer materializer is not implemented yet", + ), + ) + + +def _backend_kind_for_provider(provider: EndpointRecord) -> BackendKind: + if provider.deployment in (DEVICE_AICORE, DEVICE_AICPU): + return BackendKind.VMM_WINDOW + if provider.deployment is HOST_CPU: + return BackendKind.POSIX_SHM + raise ValueError(f"unsupported provider deployment: {_endpoint_label(provider)}") + + +def _normalize_region_access_key( + key: tuple[BackendKind | str, RegionPartKind | str, AdapterKind | str, AdapterProfile | str], +) -> tuple[BackendKind, RegionPartKind, AdapterKind, AdapterProfile]: + backend_kind, part, adapter_kind, adapter_profile = key + return ( + BackendKind(backend_kind), + RegionPartKind(part), + AdapterKind(adapter_kind), + AdapterProfile(adapter_profile), + ) + + +def _region_access_supported(query: RegionAccessQuery, candidate: _AdapterCandidate) -> RegionAccessDecision: + diagnostics = _region_access_diagnostics( + RegionAccessReasonCode.SUPPORTED, + "region access is supported", + query, + candidate, + ) + return RegionAccessDecision(True, diagnostics=diagnostics) + + +def _region_access_unsupported( + reason_code: RegionAccessReasonCode, + message: str, + query: RegionAccessQuery, + candidate: _AdapterCandidate, +) -> RegionAccessDecision: + diagnostics = _region_access_diagnostics(reason_code, message, query, candidate) + return RegionAccessDecision(False, message, diagnostics) + + +def _region_access_diagnostics( + reason_code: RegionAccessReasonCode, + message: str, + query: RegionAccessQuery, + candidate: _AdapterCandidate, +) -> RegionAccessDiagnostics: + return RegionAccessDiagnostics( + reason_code=reason_code, + message=message, + provider_label=_endpoint_label(query.provider), + consumer_label=_endpoint_label(query.consumer), + platform=query.platform, + runtime=query.runtime, + backend_kind=query.backend_kind, + part=query.part, + adapter_kind=candidate.kind, + adapter_profile=candidate.profile, + ) + + +def _attempt( + part: RegionPartKind, + member: EndpointRecord, + backend_kind: BackendKind, + candidate: _AdapterCandidate, + reason: str | None, +) -> AdapterAttempt: + return AdapterAttempt( + part=part, + member=member, + backend_kind=backend_kind, + adapter_kind=candidate.kind, + adapter_profile=candidate.profile, + reason=reason, + ) + + +def _unsupported( + reason: BackendUnsupportedReason, + message: str, + offending: Sequence[EndpointRecord] = (), + *, + attempted_adapters: Sequence[AdapterAttempt] = (), +) -> UnsupportedRegionPlan: + unique_offending = _unique_endpoints(offending) + attempts = tuple(attempted_adapters) + if reason is not BackendUnsupportedReason.ADAPTER_UNSUPPORTED and attempts: + raise ValueError("attempted_adapters are only valid for ADAPTER_UNSUPPORTED") + if unique_offending: + message = f"{message}: {', '.join(_endpoint_label(endpoint) for endpoint in unique_offending)}" + return UnsupportedRegionPlan( + reason=reason, + message=message, + offending_endpoints=unique_offending, + attempted_adapters=attempts, + ) + + +def _unique_endpoints(endpoints: Sequence[EndpointRecord]) -> tuple[EndpointRecord, ...]: + unique: list[EndpointRecord] = [] + seen: set[EndpointIdentity] = set() + for endpoint in endpoints: + if endpoint.identity in seen: + continue + seen.add(endpoint.identity) + unique.append(endpoint) + return tuple(unique) + + +def _endpoint_label(record: EndpointRecord) -> str: + return f"{record.path} {record.deployment.value}" + + +def _selector_label(selector: EndpointSelector) -> str: + return f"{selector.path} {selector.deployment.value}" + + +__all__ = [ + "AdapterAttempt", + "AdapterKind", + "AdapterProfile", + "AttachmentRole", + "BackendKind", + "BackendPlan", + "BackendResolver", + "BackendUnsupportedReason", + "DEVICE_AICORE", + "DEVICE_AICPU", + "EndpointDeployment", + "EndpointId", + "EndpointIdentity", + "EndpointPathSegment", + "EndpointRecord", + "EndpointRegistry", + "EndpointResolveError", + "EndpointResolveReason", + "EndpointSelector", + "EndpointSelectorKind", + "HOST_CPU", + "MemberAttachmentPlan", + "NodeScopeId", + "ParsedEndpointPath", + "RegionAccessDecision", + "RegionAccessDiagnostics", + "RegionAccessQuery", + "RegionAccessReasonCode", + "RegionAccessService", + "RegionLayoutSpec", + "RegionPartKind", + "RegionPartPlan", + "ResolvedRegionSpec", + "ResolvedSingleOwner", + "SingleOwner", + "SingleOwnerPlan", + "StaticRegionAccessService", + "UnsupportedRegionPlan", + "at", + "parse_endpoint_path", + "under", +] diff --git a/python/simpler/worker.py b/python/simpler/worker.py index c432965b3..ea31f4226 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -118,6 +118,24 @@ def my_l4_orch(orch, args, config): parse_python_callable_payload, parse_python_import_target, ) +from .comm_endpoints import ( + DEVICE_AICORE, + DEVICE_AICPU, + HOST_CPU, + BackendPlan, + BackendResolver, + DefaultRegionAccessService, + EndpointRegistry, + RegionAccessService, + RegionLayoutSpec, + SingleOwner, + UnsupportedRegionPlan, + _EndpointTopologyEntry, + _EndpointTopologySnapshot, + _format_worker_path, + _normalize_node_identity, + parse_endpoint_path, +) from .orchestrator import Orchestrator, _callback_run, direct_control from .remote_l3_protocol import HOST_TCP_TRANSPORT_PROFILE from .task_interface import ( @@ -3839,6 +3857,7 @@ def __init__( self._live_handles: dict[int, bytes] = {} self._next_handle_id: int = 0 self._owner_id = uuid.uuid4().hex + self._owner_instance_id = uuid.uuid4().bytes self._shm_token: str = "" self._shm_tree_tokens: set[str] = set() self._uncertain_hashids: set[bytes] = set() @@ -4029,6 +4048,10 @@ def __init__( # starts the C++ scheduler; no comm work happens there. self._comm_base_ready: bool = False + self._endpoint_registry: EndpointRegistry | None = None + self._endpoint_registry_epoch: int = 0 + self._region_access_service: RegionAccessService | None = None + self._live_worker_chip_regions: list[Any] = [] self._worker_chip_orch_comm_host_buffers: dict[int, int] = {} @@ -5324,6 +5347,102 @@ def _operation_lease(self, api: str): self._lease_depth[tid] = depth self._hierarchical_start_cv.notify_all() + def _invalidate_endpoint_registry(self) -> None: + self._endpoint_registry = None + self._region_access_service = None + self._endpoint_registry_epoch += 1 + + def _require_ready_for_region_planning(self, api: str = "region planning") -> None: + if self.level < 3: + raise RuntimeError(f"Worker.{api}: region planning requires a level >= 3 Worker") + + def _endpoint_topology_snapshot(self) -> _EndpointTopologySnapshot: + self._require_ready_for_region_planning("_endpoint_topology_snapshot") + root_level = int(self.level) + root_path = _format_worker_path(root_level) + entries: list[_EndpointTopologyEntry] = [] + self._append_endpoint_topology(entries, self, root_path, "local", include_self=True) + entries.sort(key=lambda entry: self._endpoint_topology_sort_key(entry, root_level)) + return _EndpointTopologySnapshot( + root_level=root_level, + session_instance_id=self._owner_instance_id, + entries=tuple(entries), + ) + + def _append_endpoint_topology( + self, + entries: list[_EndpointTopologyEntry], + worker: Worker, + path: str, + node_identity: str, + *, + include_self: bool, + ) -> None: + if include_self: + entries.append(_EndpointTopologyEntry(path, HOST_CPU, node_identity)) + if int(worker.level) == 3: + self._append_device_endpoint_topology(entries, path, worker._config.get("device_ids", ()), node_identity) + for child_index, child in zip(worker._next_level_worker_ids, worker._next_level_workers): + child_path = _format_worker_path(int(child.level), parent_path=path, index=int(child_index)) + self._append_endpoint_topology(entries, child, child_path, node_identity, include_self=True) + for child_index, spec in zip(worker._remote_worker_ids, worker._remote_worker_specs): + remote_path = _format_worker_path(3, parent_path=path, index=int(child_index)) + remote_node_identity = self._node_identity_from_remote_endpoint(spec.endpoint) + entries.append(_EndpointTopologyEntry(remote_path, HOST_CPU, remote_node_identity)) + self._append_device_endpoint_topology(entries, remote_path, spec.device_ids, remote_node_identity) + + def _append_device_endpoint_topology( + self, + entries: list[_EndpointTopologyEntry], + path_to_l3: str, + device_ids, + node_identity: str, + ) -> None: + for child_index, _device_id in enumerate(tuple(device_ids)): + device_path = _format_worker_path(2, parent_path=path_to_l3, index=child_index) + entries.append(_EndpointTopologyEntry(device_path, DEVICE_AICORE, node_identity)) + entries.append(_EndpointTopologyEntry(device_path, DEVICE_AICPU, node_identity)) + + def _node_identity_from_remote_endpoint(self, endpoint: str) -> str: + host, _port = self._parse_remote_endpoint(endpoint) + return _normalize_node_identity(host) + + def _endpoint_topology_sort_key(self, entry: _EndpointTopologyEntry, root_level: int): + deployment_order = {HOST_CPU: 0, DEVICE_AICORE: 1, DEVICE_AICPU: 2} + return (parse_endpoint_path(entry.path, root_level=root_level).sort_key, deployment_order[entry.deployment]) + + def _get_endpoint_registry(self) -> EndpointRegistry: + self._require_ready_for_region_planning("_get_endpoint_registry") + registry = self._endpoint_registry + if registry is None: + registry = EndpointRegistry.from_snapshot( + self._endpoint_topology_snapshot(), registry_epoch=self._endpoint_registry_epoch + ) + self._endpoint_registry = registry + return registry + + def _get_region_access_service(self) -> RegionAccessService: + service = self._region_access_service + if service is None: + service = DefaultRegionAccessService() + self._region_access_service = service + return service + + def _resolve_region_spec(self, members, topology: SingleOwner): + self._require_ready_for_region_planning("_resolve_region_spec") + with self._operation_lease("_resolve_region_spec"): + return self._get_endpoint_registry().resolve_region_spec(members, topology) + + def _plan_region( + self, members, topology: SingleOwner, layout_summary: RegionLayoutSpec + ) -> BackendPlan | UnsupportedRegionPlan: + self._require_ready_for_region_planning("_plan_region") + with self._operation_lease("_plan_region"): + registry = self._get_endpoint_registry() + resolved = registry.resolve_region_spec(members, topology) + resolver = BackendResolver(registry, self._get_region_access_service()) + return resolver.plan(resolved, layout_summary) + def _register_into_snapshot_or_wait(self, reg: _CallableRegistration) -> CallableHandle | None: """Linearize a level>=3 register against the startup epoch. @@ -8343,7 +8462,7 @@ def _create_buffer_locked(self, nbytes: int) -> Buffer: nbytes, owner_instance_id=self._owner_instance_id, buffer_id=buffer_id, - owner_worker_path=f"L{self.level}", + owner_worker_path=_format_worker_path(int(self.level)), generation=1, ) with self._registry_lock: @@ -9111,6 +9230,7 @@ def close(self) -> None: # noqa: PLR0912, PLR0915 -- lifecycle linearization: r # Claim: publish CLOSED (permanent admission fence) and install a # fresh teardown attempt. self._lifecycle = _Lifecycle.CLOSED + self._invalidate_endpoint_registry() attempt = _CloseAttempt() self._close_completion = attempt self._hierarchical_start_cv.notify_all() diff --git a/tests/ut/py/test_package_surface.py b/tests/ut/py/test_package_surface.py index a8c28dd61..94924fa9a 100644 --- a/tests/ut/py/test_package_surface.py +++ b/tests/ut/py/test_package_surface.py @@ -16,11 +16,13 @@ import simpler -def test_worker_and_task_interface_are_advertised(): +def test_worker_task_interface_and_comm_endpoints_are_advertised(): assert "Worker" in simpler.__all__ assert "task_interface" in simpler.__all__ + assert "comm_endpoints" in simpler.__all__ assert "Worker" in dir(simpler) assert "task_interface" in dir(simpler) + assert "comm_endpoints" in dir(simpler) def test_logging_helpers_remain_exported(): @@ -29,6 +31,31 @@ def test_logging_helpers_remain_exported(): assert hasattr(simpler, name) +def test_comm_endpoints_region_access_surface_is_module_scoped(): + ce = importlib.import_module("simpler.comm_endpoints") + for name in ( + "RegionAccessQuery", + "RegionAccessDecision", + "RegionAccessDiagnostics", + "RegionAccessReasonCode", + "RegionAccessService", + "StaticRegionAccessService", + ): + assert name in ce.__all__ + assert hasattr(ce, name) + assert name not in simpler.__all__ + assert not hasattr(simpler, name) + + for removed in ( + "PlatformCapability", + "CapabilityResult", + "PlatformCapabilityCache", + "StaticPlatformCapabilityCache", + "_AdapterCandidate", + ): + assert removed not in ce.__all__ + + def test_unknown_attribute_raises_attribute_error(): with pytest.raises(AttributeError, match="has no attribute"): getattr(simpler, "definitely_not_an_attribute") # noqa: B009 -- exercises __getattr__ diff --git a/tests/ut/py/test_worker/test_comm_endpoints.py b/tests/ut/py/test_worker/test_comm_endpoints.py new file mode 100644 index 000000000..7d3e1ec64 --- /dev/null +++ b/tests/ut/py/test_worker/test_comm_endpoints.py @@ -0,0 +1,537 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Unit tests for W2 endpoint selectors, registry resolution, and planning.""" + +import dataclasses + +import pytest +from simpler import comm_endpoints as ce +from simpler.buffer import BackendKind as BufferBackendKind +from simpler.worker import RemoteWorkerSpec, Worker, _Lifecycle + + +def _ready(worker: Worker) -> Worker: + worker._lifecycle = _Lifecycle.READY + return worker + + +def _l3(device_ids=(), *, num_sub_workers: int = 0) -> Worker: + return _ready(Worker(level=3, device_ids=list(device_ids), num_sub_workers=num_sub_workers)) + + +def _l4_with_local_l3(device_ids=()) -> Worker: + l3 = Worker(level=3, device_ids=list(device_ids), num_sub_workers=0) + w4 = Worker(level=4, num_sub_workers=0) + w4.add_worker(l3) + return _ready(w4) + + +def _l4_with_remote(*specs: RemoteWorkerSpec) -> Worker: + worker = Worker(level=4, num_sub_workers=0) + for spec in specs: + worker.add_remote_worker(spec) + return _ready(worker) + + +def _record(worker: Worker, path: str, deployment: ce.EndpointDeployment) -> ce.EndpointRecord: + return worker._resolve_region_spec([ce.at(path, deployment)], ce.SingleOwner()).members[0] + + +def _access_key( + backend_kind: ce.BackendKind, + part: ce.RegionPartKind, + adapter_kind: ce.AdapterKind, + adapter_profile: ce.AdapterProfile, +): + return (backend_kind, part, adapter_kind, adapter_profile) + + +def _supported_parts( + backend_kind: ce.BackendKind, + adapter_kind: ce.AdapterKind, + adapter_profile: ce.AdapterProfile, +): + return { + _access_key(backend_kind, part, adapter_kind, adapter_profile): True + for part in (ce.RegionPartKind.PAYLOAD, ce.RegionPartKind.COUNTER) + } + + +def _plan(worker: Worker, members, topology=None, access=None): + if access is not None: + worker._region_access_service = ce.StaticRegionAccessService(access) + return worker._plan_region( + members, + topology or ce.SingleOwner(), + ce.RegionLayoutSpec(payload_bytes=64, counter_bytes=8), + ) + + +def _attachments_by_member(part: ce.RegionPartPlan) -> dict[ce.EndpointIdentity, ce.MemberAttachmentPlan]: + return {attachment.member: attachment for attachment in part.attachments} + + +def test_selector_constructors_validate_shape_and_preserve_hashability(): + selector = ce.at("L3/L2[0]", "DEVICE_AICORE") + assert selector == ce.EndpointSelector(ce.EndpointSelectorKind.AT, "L3/L2[0]", ce.DEVICE_AICORE) + assert hash(selector) == hash(ce.at("L3/L2[0]", ce.DEVICE_AICORE)) + assert ce.under("L3", ce.DEVICE_AICPU).kind is ce.EndpointSelectorKind.UNDER + with pytest.raises(ValueError, match="invalid endpoint deployment"): + ce.at("L3", "DEVICE") + with pytest.raises(ValueError, match="non-empty"): + ce.at("", ce.HOST_CPU) + with pytest.raises(ValueError, match="empty segments"): + ce.EndpointSelector(ce.EndpointSelectorKind.AT, "L3//L2[0]", ce.DEVICE_AICORE) + + +def test_snapshot_is_frozen_flat_and_uses_worker_owner_instance_id(): + worker = _l3(device_ids=[8, 9]) + snapshot = worker._endpoint_topology_snapshot() + assert dataclasses.is_dataclass(snapshot) + assert snapshot.root_level == 3 + assert snapshot.session_instance_id == worker._owner_instance_id + assert isinstance(snapshot.entries, tuple) + assert [entry.path for entry in snapshot.entries] == [ + "L3", + "L3/L2[0]", + "L3/L2[0]", + "L3/L2[1]", + "L3/L2[1]", + ] + with pytest.raises(dataclasses.FrozenInstanceError): + snapshot.entries[0].path = "L3/changed" + + +def test_registry_only_builds_from_snapshot_and_identity_binds_epoch_and_session(): + worker = _l3(device_ids=[0]) + snapshot = worker._endpoint_topology_snapshot() + assert not hasattr(ce.EndpointRegistry, "from_worker") + registry0 = ce.EndpointRegistry.from_snapshot(snapshot, registry_epoch=0) + registry1 = ce.EndpointRegistry.from_snapshot(snapshot, registry_epoch=1) + record0 = registry0.resolve_members([ce.at("L3/L2[0]", ce.DEVICE_AICORE)])[0] + record1 = registry1.resolve_members([ce.at("L3/L2[0]", ce.DEVICE_AICORE)])[0] + other_session = dataclasses.replace(snapshot, session_instance_id=b"other-session") + record_other = ce.EndpointRegistry.from_snapshot(other_session, registry_epoch=0).resolve_members( + [ce.at("L3/L2[0]", ce.DEVICE_AICORE)] + )[0] + root = registry0.resolve_members([ce.at("L3", ce.HOST_CPU)])[0] + assert record0.endpoint_id == record0.identity.endpoint_id + assert record0.identity != record1.identity + assert record0.identity != record_other.identity + assert record0.identity != root.identity + + +def test_backend_kind_reuses_buffer_abi_vocabulary(): + assert ce.BackendKind is BufferBackendKind + assert ce.BackendKind.VMM_WINDOW is BufferBackendKind.VMM_WINDOW + + +def test_path_grammar_requires_root_and_uses_numeric_canonical_sort(): + worker = _l3(device_ids=range(11)) + with pytest.raises(ce.EndpointResolveError) as excinfo: + worker._resolve_region_spec([ce.at("L3[0]/L2[0]", ce.DEVICE_AICORE)], ce.SingleOwner()) + assert excinfo.value.reason is ce.EndpointResolveReason.INVALID_PATH + + members = worker._resolve_region_spec([ce.under("L3", ce.DEVICE_AICORE)], ce.SingleOwner()).members + paths = [record.path for record in members] + assert paths[2] == "L3/L2[2]" + assert paths[10] == "L3/L2[10]" + + +def test_l3_registry_registers_host_and_device_views_but_not_l2_host_cpu(): + worker = _l3(device_ids=[8, 9]) + assert _record(worker, "L3", ce.HOST_CPU).path == "L3" + assert _record(worker, "L3/L2[0]", ce.DEVICE_AICORE).deployment is ce.DEVICE_AICORE + assert _record(worker, "L3/L2[1]", ce.DEVICE_AICPU).deployment is ce.DEVICE_AICPU + with pytest.raises(ce.EndpointResolveError) as excinfo: + worker._resolve_region_spec([ce.at("L3/L2[0]", ce.HOST_CPU)], ce.SingleOwner()) + assert excinfo.value.reason is ce.EndpointResolveReason.ENDPOINT_NOT_REGISTERED + + +def test_l4_local_registry_registers_child_l3_devices_on_same_node(): + worker = _l4_with_local_l3(device_ids=[4, 7]) + root = _record(worker, "L4", ce.HOST_CPU) + child_host = _record(worker, "L4/L3[0]", ce.HOST_CPU) + child_aicore = _record(worker, "L4/L3[0]/L2[1]", ce.DEVICE_AICORE) + child_aicpu = _record(worker, "L4/L3[0]/L2[1]", ce.DEVICE_AICPU) + registry = worker._get_endpoint_registry() + assert registry.same_node(root, child_host) + assert registry.same_node(child_host, child_aicore) + assert registry.same_node(child_aicore, child_aicpu) + + +def test_remote_registry_normalizes_node_identity_by_host_not_remote_status(): + worker = _l4_with_remote( + RemoteWorkerSpec(endpoint="127.0.0.1:1234", platform="a2a3", device_ids=(6,)), + RemoteWorkerSpec(endpoint="10.0.0.7:1234", platform="a2a3", device_ids=(7,)), + RemoteWorkerSpec(endpoint="10.0.0.7:2345", platform="a2a3", device_ids=(8,)), + RemoteWorkerSpec(endpoint="10.0.0.8:1234", platform="a2a3", device_ids=(9,)), + ) + registry = worker._get_endpoint_registry() + root = _record(worker, "L4", ce.HOST_CPU) + loopback = _record(worker, "L4/L3[0]", ce.HOST_CPU) + remote_a = _record(worker, "L4/L3[1]", ce.HOST_CPU) + remote_b = _record(worker, "L4/L3[2]", ce.HOST_CPU) + remote_c = _record(worker, "L4/L3[3]", ce.HOST_CPU) + assert registry.same_node(root, loopback) + assert registry.same_node(remote_a, remote_b) + assert not registry.same_node(root, remote_a) + assert not registry.same_node(remote_a, remote_c) + + +def test_at_missing_path_reports_path_not_found(): + worker = _l3(device_ids=[0]) + with pytest.raises(ce.EndpointResolveError) as excinfo: + worker._resolve_region_spec([ce.at("L3/L2[9]", ce.DEVICE_AICORE)], ce.SingleOwner()) + assert excinfo.value.reason is ce.EndpointResolveReason.PATH_NOT_FOUND + assert "L3/L2[9] DEVICE_AICORE" in excinfo.value.message + + +def test_under_excludes_self_and_empty_expansion_is_error(): + worker = _l4_with_local_l3(device_ids=[]) + members = worker._resolve_region_spec([ce.under("L4", ce.HOST_CPU)], ce.SingleOwner()).members + assert [member.path for member in members] == ["L4/L3[0]"] + + with pytest.raises(ce.EndpointResolveError) as excinfo: + worker._resolve_region_spec([ce.under("L4/L3[0]", ce.HOST_CPU)], ce.SingleOwner()) + assert excinfo.value.reason is ce.EndpointResolveReason.EMPTY_UNDER_SELECTOR + + +def test_duplicate_member_expansion_is_rejected(): + worker = _l3(device_ids=[0, 1]) + with pytest.raises(ce.EndpointResolveError) as excinfo: + worker._resolve_region_spec( + [ce.under("L3", ce.DEVICE_AICORE), ce.at("L3/L2[0]", ce.DEVICE_AICORE)], ce.SingleOwner() + ) + assert excinfo.value.reason is ce.EndpointResolveReason.DUPLICATE_ENDPOINT + + +def test_provider_resolve_is_registry_api_boundary(): + worker = _l3(device_ids=[0, 1]) + resolved = worker._resolve_region_spec( + [ce.at("L3/L2[0]", ce.DEVICE_AICORE)], + ce.SingleOwner(provider=ce.at("L3/L2[0]", ce.DEVICE_AICORE)), + ) + assert resolved.topology.provider_endpoint == resolved.members[0].identity + + with pytest.raises(ce.EndpointResolveError) as excinfo: + worker._resolve_region_spec( + [ce.at("L3/L2[0]", ce.DEVICE_AICORE)], + ce.SingleOwner(provider=ce.under("L3", ce.DEVICE_AICORE)), + ) + assert excinfo.value.reason is ce.EndpointResolveReason.PROVIDER_NOT_SINGLE_ENDPOINT + + +def test_provider_not_in_members_is_backend_resolver_error(): + worker = _l3(device_ids=[0]) + plan = _plan( + worker, + [ce.at("L3", ce.HOST_CPU)], + ce.SingleOwner(provider=ce.at("L3/L2[0]", ce.DEVICE_AICORE)), + ) + assert isinstance(plan, ce.UnsupportedRegionPlan) + assert plan.reason is ce.BackendUnsupportedReason.PROVIDER_NOT_IN_MEMBERS + assert not plan.attempted_adapters + + +def test_default_provider_order_aicore_then_aicpu_then_host_cpu(): + worker = _l3(device_ids=[0]) + plan = _plan( + worker, + [ce.at("L3/L2[0]", ce.DEVICE_AICPU), ce.at("L3/L2[0]", ce.DEVICE_AICORE)], + access=_supported_parts( + ce.BackendKind.VMM_WINDOW, + ce.AdapterKind.DEVICE_PEER, + ce.AdapterProfile.DEVICE_VMM_PEER_IMPORT, + ), + ) + assert isinstance(plan, ce.BackendPlan) + assert _record(worker, "L3/L2[0]", ce.DEVICE_AICORE).identity == plan.topology_plan.provider_endpoint + + worker = _l3(device_ids=[0]) + plan = _plan( + worker, + [ce.at("L3", ce.HOST_CPU), ce.at("L3/L2[0]", ce.DEVICE_AICPU)], + ) + assert isinstance(plan, ce.BackendPlan) + assert _record(worker, "L3/L2[0]", ce.DEVICE_AICPU).identity == plan.topology_plan.provider_endpoint + + worker = _l4_with_local_l3() + plan = _plan( + worker, + [ce.at("L4", ce.HOST_CPU), ce.at("L4/L3[0]", ce.HOST_CPU)], + access=_supported_parts( + ce.BackendKind.POSIX_SHM, + ce.AdapterKind.DIRECT_MAP, + ce.AdapterProfile.HOST_SHM_MAP, + ), + ) + assert isinstance(plan, ce.BackendPlan) + assert _record(worker, "L4", ce.HOST_CPU).identity == plan.topology_plan.provider_endpoint + + +def test_default_provider_missing_for_empty_members(): + worker = _l3() + plan = _plan(worker, []) + assert isinstance(plan, ce.UnsupportedRegionPlan) + assert plan.reason is ce.BackendUnsupportedReason.NO_DEFAULT_PROVIDER + assert not plan.attempted_adapters + + +def test_device_backend_default_host_consumer_uses_copy_for_payload_and_counter(): + worker = _l3(device_ids=[0]) + members = [ce.at("L3", ce.HOST_CPU), ce.at("L3/L2[0]", ce.DEVICE_AICORE)] + plan = _plan(worker, members) + assert isinstance(plan, ce.BackendPlan) + assert not hasattr(plan, "required_capabilities") + assert not hasattr(ce, "BackingKind") + assert not hasattr(ce, "MaterializationMode") + assert not hasattr(ce, "PlatformCapability") + assert not hasattr(ce, "CapabilityResult") + assert not hasattr(ce, "PlatformCapabilityCache") + assert not hasattr(ce, "StaticPlatformCapabilityCache") + resolved = worker._resolve_region_spec(members, ce.SingleOwner()).members + assert plan.ordered_members == tuple(record.identity for record in resolved) + assert plan.payload.part is ce.RegionPartKind.PAYLOAD + assert plan.counter.part is ce.RegionPartKind.COUNTER + assert plan.payload.backend_kind is ce.BackendKind.VMM_WINDOW + assert plan.counter.backend_kind is ce.BackendKind.VMM_WINDOW + + host = _record(worker, "L3", ce.HOST_CPU) + provider = _record(worker, "L3/L2[0]", ce.DEVICE_AICORE) + payload_by_member = _attachments_by_member(plan.payload) + assert payload_by_member[provider.identity].role is ce.AttachmentRole.PROVIDER + assert payload_by_member[provider.identity].adapter_kind is None + assert payload_by_member[host.identity].role is ce.AttachmentRole.CONSUMER + assert payload_by_member[host.identity].adapter_kind is ce.AdapterKind.OWNER_DELEGATED_COPY + assert payload_by_member[host.identity].adapter_profile is ce.AdapterProfile.HOST_VMM_COPY + + counter_host = _attachments_by_member(plan.counter)[host.identity] + assert counter_host.adapter_kind is ce.AdapterKind.OWNER_DELEGATED_COPY + assert counter_host.adapter_profile is ce.AdapterProfile.HOST_VMM_COPY + + +def test_static_service_can_select_counter_direct_map_without_payload_direct_attempt(): + worker = _l3(device_ids=[0]) + members = [ce.at("L3", ce.HOST_CPU), ce.at("L3/L2[0]", ce.DEVICE_AICORE)] + access = { + _access_key( + ce.BackendKind.VMM_WINDOW, + ce.RegionPartKind.PAYLOAD, + ce.AdapterKind.OWNER_DELEGATED_COPY, + ce.AdapterProfile.HOST_VMM_COPY, + ): True, + _access_key( + ce.BackendKind.VMM_WINDOW, + ce.RegionPartKind.COUNTER, + ce.AdapterKind.DIRECT_MAP, + ce.AdapterProfile.HOST_SVM_MAP, + ): True, + } + plan = _plan(worker, members, access=access) + assert isinstance(plan, ce.BackendPlan) + + host = _record(worker, "L3", ce.HOST_CPU) + payload_host = _attachments_by_member(plan.payload)[host.identity] + counter_host = _attachments_by_member(plan.counter)[host.identity] + assert payload_host.adapter_kind is ce.AdapterKind.OWNER_DELEGATED_COPY + assert payload_host.adapter_profile is ce.AdapterProfile.HOST_VMM_COPY + assert counter_host.adapter_kind is ce.AdapterKind.DIRECT_MAP + assert counter_host.adapter_profile is ce.AdapterProfile.HOST_SVM_MAP + + +def test_device_backend_attempts_are_recorded_when_direct_and_copy_are_absent(): + worker = _l3(device_ids=[0]) + plan = _plan(worker, [ce.at("L3", ce.HOST_CPU), ce.at("L3/L2[0]", ce.DEVICE_AICORE)], access={}) + assert isinstance(plan, ce.UnsupportedRegionPlan) + assert plan.reason is ce.BackendUnsupportedReason.ADAPTER_UNSUPPORTED + assert [attempt.adapter_profile for attempt in plan.attempted_adapters] == [ + ce.AdapterProfile.HOST_VMM_COPY, + ] + assert all(attempt.part is ce.RegionPartKind.PAYLOAD for attempt in plan.attempted_adapters) + assert all(attempt.backend_kind is ce.BackendKind.VMM_WINDOW for attempt in plan.attempted_adapters) + assert all(attempt.member.path == "L3" for attempt in plan.attempted_adapters) + assert plan.message.count("L3 HOST_CPU") == 1 + + +def test_default_counter_direct_map_evaluator_reports_probe_not_implemented(): + worker = _l3(device_ids=[0]) + registry = worker._get_endpoint_registry() + provider = _record(worker, "L3/L2[0]", ce.DEVICE_AICORE) + host = _record(worker, "L3", ce.HOST_CPU) + layout = ce.RegionLayoutSpec(payload_bytes=64, counter_bytes=8) + resolver = ce.BackendResolver(registry, worker._get_region_access_service()) + query = ce.RegionAccessQuery( + topology="SingleOwner", + part=ce.RegionPartKind.COUNTER, + backend_kind=ce.BackendKind.VMM_WINDOW, + provider=provider, + consumer=host, + layout=layout, + same_node=True, + platform=worker._config.get("platform"), + runtime=worker._config.get("runtime"), + ) + direct_candidate = resolver._adapter_candidates( # pyright: ignore[reportPrivateUsage] + ce.RegionPartKind.COUNTER, + ce.BackendKind.VMM_WINDOW, + provider, + host, + )[0] + decision = worker._get_region_access_service().evaluate_region_access(query, direct_candidate) + assert not decision.supported + assert decision.diagnostics is not None + assert decision.diagnostics.reason_code is ce.RegionAccessReasonCode.NO_IMPLEMENTED_DIRECT_MAP_PROBE + + +def test_consumer_attachment_passes_part_to_candidate_selection(): + worker = _l3(device_ids=[0]) + registry = worker._get_endpoint_registry() + provider = _record(worker, "L3/L2[0]", ce.DEVICE_AICORE) + host = _record(worker, "L3", ce.HOST_CPU) + resolver = ce.BackendResolver( + registry, + ce.StaticRegionAccessService( + { + _access_key( + ce.BackendKind.VMM_WINDOW, + ce.RegionPartKind.COUNTER, + ce.AdapterKind.DIRECT_MAP, + ce.AdapterProfile.HOST_SVM_MAP, + ): True + } + ), + ) + candidate = resolver._adapter_candidates( # pyright: ignore[reportPrivateUsage] + ce.RegionPartKind.COUNTER, + ce.BackendKind.VMM_WINDOW, + provider, + host, + )[0] + seen_parts = [] + + def candidate_order(part, backend_kind, provider_record, member_record): + seen_parts.append(part) + assert backend_kind is ce.BackendKind.VMM_WINDOW + assert provider_record == provider + assert member_record == host + return (candidate,) + + resolver._adapter_candidates = candidate_order # pyright: ignore[reportPrivateUsage,reportAttributeAccessIssue] + attachment = resolver._consumer_attachment( # pyright: ignore[reportPrivateUsage] + ce.RegionPartKind.COUNTER, + ce.BackendKind.VMM_WINDOW, + provider, + host, + ce.RegionLayoutSpec(payload_bytes=64, counter_bytes=8), + ) + assert isinstance(attachment, ce.MemberAttachmentPlan) + assert seen_parts == [ce.RegionPartKind.COUNTER] + + +def test_host_shm_plan_and_host_provider_device_member_attempts(): + worker = _l4_with_local_l3() + plan = _plan( + worker, + [ce.at("L4", ce.HOST_CPU), ce.at("L4/L3[0]", ce.HOST_CPU)], + access=_supported_parts( + ce.BackendKind.POSIX_SHM, + ce.AdapterKind.DIRECT_MAP, + ce.AdapterProfile.HOST_SHM_MAP, + ), + ) + assert isinstance(plan, ce.BackendPlan) + child_host = _record(worker, "L4/L3[0]", ce.HOST_CPU) + attachment = _attachments_by_member(plan.payload)[child_host.identity] + assert attachment.adapter_kind is ce.AdapterKind.DIRECT_MAP + assert attachment.adapter_profile is ce.AdapterProfile.HOST_SHM_MAP + assert plan.payload.backend_kind is ce.BackendKind.POSIX_SHM + + worker = _l3(device_ids=[0]) + plan = _plan( + worker, + [ce.at("L3", ce.HOST_CPU), ce.at("L3/L2[0]", ce.DEVICE_AICORE)], + ce.SingleOwner(provider=ce.at("L3", ce.HOST_CPU)), + ) + assert isinstance(plan, ce.UnsupportedRegionPlan) + assert plan.reason is ce.BackendUnsupportedReason.ADAPTER_UNSUPPORTED + assert [(attempt.adapter_kind, attempt.adapter_profile, attempt.reason) for attempt in plan.attempted_adapters] == [ + ( + ce.AdapterKind.EXPLICIT_TRANSFER, + ce.AdapterProfile.REMOTE_COPY, + "explicit transfer materializer is not implemented yet", + ) + ] + + +def test_cross_node_resolves_then_reports_adapter_attempts_without_hard_reject(): + worker = _l4_with_remote(RemoteWorkerSpec(endpoint="10.0.0.7:1234", platform="a2a3", device_ids=(6,))) + resolved = worker._resolve_region_spec( + [ce.at("L4", ce.HOST_CPU), ce.at("L4/L3[0]/L2[0]", ce.DEVICE_AICORE)], + ce.SingleOwner(provider=ce.at("L4/L3[0]/L2[0]", ce.DEVICE_AICORE)), + ) + assert [record.path for record in resolved.members] == ["L4", "L4/L3[0]/L2[0]"] + + plan = _plan( + worker, + [ce.at("L4", ce.HOST_CPU), ce.at("L4/L3[0]/L2[0]", ce.DEVICE_AICORE)], + ce.SingleOwner(provider=ce.at("L4/L3[0]/L2[0]", ce.DEVICE_AICORE)), + ) + assert isinstance(plan, ce.UnsupportedRegionPlan) + assert plan.reason is ce.BackendUnsupportedReason.ADAPTER_UNSUPPORTED + assert not hasattr(ce.BackendUnsupportedReason, "CROSS_NODE_UNSUPPORTED") + assert [(attempt.adapter_kind, attempt.adapter_profile, attempt.reason) for attempt in plan.attempted_adapters] == [ + ( + ce.AdapterKind.OWNER_DELEGATED_COPY, + ce.AdapterProfile.REMOTE_COPY, + "remote copy materializer is not implemented yet", + ), + ( + ce.AdapterKind.EXPLICIT_TRANSFER, + ce.AdapterProfile.REMOTE_COPY, + "explicit transfer materializer is not implemented yet", + ), + ] + + +def test_cross_node_device_member_attempt_order_includes_fabric_before_remote_copy(): + worker = _l4_with_remote( + RemoteWorkerSpec(endpoint="10.0.0.7:1234", platform="a2a3", device_ids=(6,)), + RemoteWorkerSpec(endpoint="10.0.0.8:1234", platform="a2a3", device_ids=(7,)), + ) + plan = _plan( + worker, + [ce.at("L4/L3[0]/L2[0]", ce.DEVICE_AICORE), ce.at("L4/L3[1]/L2[0]", ce.DEVICE_AICORE)], + ce.SingleOwner(provider=ce.at("L4/L3[0]/L2[0]", ce.DEVICE_AICORE)), + ) + assert isinstance(plan, ce.UnsupportedRegionPlan) + assert [attempt.adapter_profile for attempt in plan.attempted_adapters] == [ + ce.AdapterProfile.DEVICE_FABRIC_V2_PEER_IMPORT, + ce.AdapterProfile.REMOTE_COPY, + ce.AdapterProfile.REMOTE_COPY, + ] + assert plan.attempted_adapters[0].reason == "device fabric peer import is not available for this endpoint" + + +def test_worker_region_planning_uses_lease_admission_and_close_invalidates_epoch(): + worker = Worker(level=3, device_ids=[0]) + with pytest.raises(RuntimeError, match="READY"): + worker._resolve_region_spec([ce.at("L3", ce.HOST_CPU)], ce.SingleOwner()) + + _ready(worker) + registry = worker._get_endpoint_registry() + record = _record(worker, "L3", ce.HOST_CPU) + assert registry.registry_epoch == 0 + assert record.identity.registry_epoch == 0 + worker.close() + assert worker._endpoint_registry is None + assert worker._region_access_service is None + assert worker._endpoint_registry_epoch == 1 + with pytest.raises(RuntimeError, match="READY"): + worker._resolve_region_spec([ce.at("L3", ce.HOST_CPU)], ce.SingleOwner())