# Source code for tensorplay.distributed.elastic.rendezvous.p10d_rendezvous_backend Source: https://www.tensorplay.cn/docs/_modules/tensorplay/distributed/elastic/rendezvous/p10d_rendezvous_backend.html ``` from __future__ import annotations import base64 import binascii import logging import os import tempfile from datetime import timedelta from typing import Any from tensorplay.distributed import FileStore, Store, TCPStore from tensorplay.distributed.elastic.events import ( NodeState, construct_and_record_rdzv_event, ) from .api import ( RendezvousConnectionError, RendezvousError, RendezvousParameters, RendezvousStateError, ) from .core_rendezvous_backend import RendezvousBackend, Token from .utils import _matches_machine_hostname, parse_rendezvous_endpoint __all__ = ["P10dRendezvousBackend", "create_backend"] logger = logging.getLogger(__name__) DEFAULT_PORT = 29400 [docs] class P10dRendezvousBackend(RendezvousBackend): """Rendezvous state backend backed by a TP key/value Store.""" _NULL_SENTINEL = b"tp-rendezvous-null" def __init__(self, store: Store, run_id: str) -> None: if not run_id: raise ValueError("The run id must be a non-empty string.") self._store = store self._key = "tp.rendezvous." + run_id try: self._store.compare_and_swap(self._key, b"", self._NULL_SENTINEL) except (ValueError, RuntimeError, TimeoutError) as exc: raise RendezvousConnectionError( "The connection to the rendezvous store has failed. " "See the inner exception for details." ) from exc @property def name(self) -> str: return "p10d" def get_state(self) -> tuple[bytes, Token] | None: raw = self._call_store("get", self._key) return self._decode_state(raw) def set_state( self, state: bytes, token: Token | None = None ) -> tuple[bytes, Token] | None: if token is None: token = 0 if not isinstance(token, int): current = self.get_state() return current current_raw = self._call_store("get", self._key) current = self._decode_state(current_raw) if current is None: if token != 0: return None expected = self._NULL_SENTINEL else: _, current_token = current if current_token != token: return current expected = current_raw new_token = token + 1 encoded = base64.b64encode(state).decode("ascii") new_raw = f"{new_token}:{encoded}".encode("ascii") swapped, observed = self._call_store( "compare_and_swap", self._key, expected, new_raw ) if swapped: return state, new_token return self._decode_state(observed) def _call_store(self, store_op: str, *args: Any, **kwargs: Any) -> Any: try: return getattr(self._store, store_op)(*args, **kwargs) except (ValueError, RuntimeError, TimeoutError) as exc: raise RendezvousConnectionError( "The connection to the rendezvous store has failed. " "See the inner exception for details." ) from exc def _decode_state(self, raw: bytes) -> tuple[bytes, Token] | None: if raw == self._NULL_SENTINEL: return None try: token_text, separator, payload = raw.decode("ascii").partition(":") if not separator: raise ValueError("missing state token") token = int(token_text) state = base64.b64decode(payload, validate=True) except (binascii.Error, UnicodeDecodeError, ValueError, TypeError) as exc: raise RendezvousStateError( "The rendezvous state object is corrupt. " "See the inner exception for details." ) from exc return state, token def _create_tcp_store(params: RendezvousParameters) -> TCPStore: host, port = parse_rendezvous_endpoint(params.endpoint, default_port=DEFAULT_PORT) configured_host = params.get_as_bool("is_host") if configured_host is not None: is_host = configured_host else: is_host = _matches_machine_hostname(host) read_timeout = params.get_as_int("read_timeout", 60) if read_timeout is None or read_timeout <= 0: raise ValueError("The read timeout must be a positive integer.") store: TCPStore | None = None for is_server in (is_host, False): try: store = TCPStore( host, port, is_master=is_server, timeout=timedelta(seconds=read_timeout).total_seconds(), wait_for_workers=False, ) if is_server: message = ( f"Process {os.getpid()} hosts the rendezvous TCP store." ) construct_and_record_rdzv_event( run_id=params.run_id, message=message, node_state=NodeState.INITIALIZED, ) logger.info(message) break except (ValueError, RuntimeError, TimeoutError, OSError) as exc: if not is_server or configured_host is not None: raise RendezvousConnectionError( "The connection to the rendezvous store has failed. " "See the inner exception for details." ) from exc if store is None: raise RendezvousConnectionError( "The connection to the rendezvous store has failed." ) return store def _create_file_store(params: RendezvousParameters) -> FileStore: path = params.endpoint if not path: try: _, path = tempfile.mkstemp(prefix="tp_rendezvous_") except OSError as exc: raise RendezvousError( "The rendezvous store file could not be created. " "See the inner exception for details." ) from exc try: return FileStore(path) except (ValueError, RuntimeError, OSError) as exc: raise RendezvousConnectionError( "The connection to the rendezvous store has failed. " "See the inner exception for details." ) from exc def create_backend( params: RendezvousParameters, ) -> tuple[P10dRendezvousBackend, Store]: """Create a rendezvous backend and its shared Store.""" store_type = str(params.get("store_type", "tcp")).strip().lower() try: if store_type == "file": store = _create_file_store(params) elif store_type == "tcp": store = _create_tcp_store(params) else: raise ValueError( "Invalid store type. Supported values are 'file' and 'tcp'." ) backend = P10dRendezvousBackend(store, params.run_id) except Exception as exc: construct_and_record_rdzv_event( message=f"{type(exc).__name__}: {exc}", run_id=params.run_id, node_state=NodeState.FAILED, ) raise return backend, store ```