Class ManagerClient

axlearn/ft/manager_client.py:23–360

FT trainer manager gRPC client with simplified, consistent interface.

class ManagerClient:
    """FT trainer manager gRPC client with simplified, consistent interface."""

    def __init__(self, timeout: float = 10.0, port: int = DEFAULT_MANAGER_PORT):
        """Initialize manager client.

        Args:
            timeout: gRPC request timeout in seconds
            port: Default port for manager connections
        """
        self.timeout = timeout
        self.port = port
        self._channels: dict[tuple[str, int], grpc.Channel] = {}
        self._channels_lock = threading.Lock()
        logging.debug("ManagerClient initialized: timeout=%.1fs, port=%d", timeout, port)

    def _get_stub(self, hostname: str, port: int) -> manager_pb2_grpc.ManagerServiceStub:
        """Return a cached gRPC stub for the target host, creating the channel if needed."""
        key = (hostname, port)
        with self._channels_lock:
            if key not in self._channels:
                self._channels[key] = grpc.insecure_channel(f"{hostname}:{port}")
            return manager_pb2_grpc.ManagerServiceStub(self._channels[key])

    def _grpc_call(
        self,
        hostname: str,
        port: int,
        method_name: str,
        request,
        log_prefix: Optional[str] = None,
    ) -> Any:
        """Execute a gRPC call with uniform debug logging.

        Args:
            hostname: Target hostname.
            port: Target port.
            method_name: Name of the gRPC stub method to invoke (e.g. "ReportStatus",
                "RestartReplicaTraining"). Must match the RPC name in the proto service definition.
            request: Protobuf request message.
            log_prefix: Short string used in log messages; defaults to method_name.

        Returns:
            The response message. Raises exceptions on failure.
        """
        prefix = log_prefix or method_name
        logging.debug("%s: sending to %s:%d", prefix, hostname, port)
        stub = self._get_stub(hostname, port)
        response = getattr(stub, method_name)(request, timeout=self.timeout)
        logging.debug(
            "%s: response from %s:%d acknowledged=%s",
            prefix,
            hostname,
            port,
            response.acknowledged,
        )
        return response