From 3242b1b002b7107fe07a3dab0e1fa164a2904398 Mon Sep 17 00:00:00 2001 From: Joel Dixon Date: Fri, 5 Sep 2025 15:07:40 -0500 Subject: [PATCH 01/12] Add a GrpcServiceClientBase and use it in pinmap.v1.client --- .../pinmap/v1/client/_client.py | 92 +++++-------------- .../pinmap/v1/client/_client_base.py | 59 ++++++++++++ 2 files changed, 80 insertions(+), 71 deletions(-) create mode 100644 packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py diff --git a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py index d3d3b0f5..d0347fae 100644 --- a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py +++ b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py @@ -2,93 +2,43 @@ from __future__ import annotations -import logging import pathlib -import threading import grpc import ni.measurementlink.pinmap.v1.pin_map_service_pb2 as pin_map_service_pb2 -import ni.measurementlink.pinmap.v1.pin_map_service_pb2_grpc as pin_map_service_pb2_grpc -from ni.measurementlink.discovery.v1.client import DiscoveryClient as DiscoveryClient +from ni.measurementlink.discovery.v1.client import DiscoveryClient +from ni.measurementlink.pinmap.v1 import pin_map_service_pb2_grpc from ni_grpc_extensions.channelpool import GrpcChannelPool -_logger = logging.getLogger(__name__) +from ni.measurementlink.pinmap.v1.client._client_base import GrpcServiceClientBase -GRPC_SERVICE_INTERFACE_NAME = "ni.measurementlink.pinmap.v1.PinMapService" -GRPC_SERVICE_CLASS = "ni.measurementlink.pinmap.v1.PinMapService" - -class PinMapClient: - """Client for accessing the NI Pin Map Service.""" +class PinMapClient(GrpcServiceClientBase): + """Client for accessing the NI Pin Map Service via gRPC.""" def __init__( self, *, - discovery_client: DiscoveryClient | None = None, - grpc_channel: grpc.Channel | None = None, - grpc_channel_pool: GrpcChannelPool | None = None, + discovery_client: "DiscoveryClient | None" = None, + grpc_channel: "grpc.Channel | None" = None, + grpc_channel_pool: "GrpcChannelPool | None" = None, ) -> None: - """Initialize the pin map client. - - Args: - discovery_client: An optional discovery client (recommended). - - grpc_channel: An optional pin map gRPC channel. - - grpc_channel_pool: An optional gRPC channel pool (recommended). - """ - self._initialization_lock = threading.Lock() - self._discovery_client = discovery_client - self._grpc_channel_pool = grpc_channel_pool - self._stub: pin_map_service_pb2_grpc.PinMapServiceStub | None = None - - if grpc_channel is not None: - self._stub = pin_map_service_pb2_grpc.PinMapServiceStub(grpc_channel) - - def _get_stub(self) -> pin_map_service_pb2_grpc.PinMapServiceStub: - if self._stub is None: - with self._initialization_lock: - if self._grpc_channel_pool is None: - _logger.debug("Creating unshared GrpcChannelPool.") - self._grpc_channel_pool = GrpcChannelPool() - if self._discovery_client is None: - _logger.debug("Creating unshared DiscoveryClient.") - self._discovery_client = DiscoveryClient( - grpc_channel_pool=self._grpc_channel_pool - ) - if self._stub is None: - compute_nodes = self._discovery_client.enumerate_compute_nodes() - remote_compute_nodes = [node for node in compute_nodes if not node.is_local] - # Use remote node URL as deployment target if only one remote node is found. - # If more than one remote node exists, use empty string for deployment target. - first_remote_node_url = ( - remote_compute_nodes[0].url if len(remote_compute_nodes) == 1 else "" - ) - service_location = self._discovery_client.resolve_service( - provided_interface=GRPC_SERVICE_INTERFACE_NAME, - deployment_target=first_remote_node_url, - service_class=GRPC_SERVICE_CLASS, - ) - channel = self._grpc_channel_pool.get_channel(service_location.insecure_address) - self._stub = pin_map_service_pb2_grpc.PinMapServiceStub(channel) - return self._stub + """Initialize a PinMapClient instance.""" + super().__init__( + discovery_client=discovery_client, + grpc_channel=grpc_channel, + grpc_channel_pool=grpc_channel_pool, + service_interface_name="ni.measurementlink.pinmap.v1.PinMapService", + service_class="ni.measurementlink.pinmap.v1.PinMapService", + stub_class=pin_map_service_pb2_grpc.PinMapServiceStub, + ) def update_pin_map(self, pin_map_path: str | pathlib.Path) -> str: - """Update registered pin map contents. - - Create and register a pin map if a pin map resource for the specified pin map id is not - found. - - Args: - pin_map_path: The file path of the pin map to register as a pin map resource. - - Returns: - The resource id of the pin map that is registered to the pin map service. - """ - # By convention, the pin map id is the .pinmap file path. + """Update the pin map from an XML file and return the pin map ID.""" request = pin_map_service_pb2.UpdatePinMapFromXmlRequest( pin_map_id=str(pin_map_path), pin_map_xml=pathlib.Path(pin_map_path).read_text(encoding="utf-8-sig"), ) - response = self._get_stub().UpdatePinMapFromXml(request) - return response.pin_map_id + stub: pin_map_service_pb2_grpc.PinMapServiceStub = self._get_stub() # type: ignore + response = stub.UpdatePinMapFromXml(request) + return str(response.pin_map_id) diff --git a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py new file mode 100644 index 00000000..f4d313fc --- /dev/null +++ b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py @@ -0,0 +1,59 @@ +import logging +import threading + +import grpc +from ni.measurementlink.discovery.v1.client import DiscoveryClient +from ni_grpc_extensions.channelpool import GrpcChannelPool + +_logger = logging.getLogger(__name__) + + +class GrpcServiceClientBase: + """Base class for NI gRPC service clients.""" + + def __init__( + self, + *, + discovery_client: DiscoveryClient | None = None, + grpc_channel: grpc.Channel | None = None, + grpc_channel_pool: GrpcChannelPool | None = None, + service_interface_name: str, + service_class: str, + stub_class: type, + ) -> None: + """Initialize a GrpcServiceClientBase instance.""" + self._initialization_lock = threading.Lock() + self._discovery_client = discovery_client + self._grpc_channel_pool = grpc_channel_pool + self._stub = stub_class(grpc_channel) if grpc_channel is not None else None + self._service_interface_name = service_interface_name + self._service_class = service_class + self._stub_class = stub_class + + def _get_stub(self) -> object: + if self._stub is None: + with self._initialization_lock: + if self._grpc_channel_pool is None: + _logger.debug("Creating unshared GrpcChannelPool.") + self._grpc_channel_pool = GrpcChannelPool() + + if self._discovery_client is None: + _logger.debug("Creating unshared DiscoveryClient.") + self._discovery_client = DiscoveryClient( + grpc_channel_pool=self._grpc_channel_pool + ) + + compute_nodes = self._discovery_client.enumerate_compute_nodes() + remote_nodes = [node for node in compute_nodes if not node.is_local] + target_url = remote_nodes[0].url if len(remote_nodes) == 1 else "" + + service_location = self._discovery_client.resolve_service( + provided_interface=self._service_interface_name, + deployment_target=target_url, + service_class=self._service_class, + ) + + channel = self._grpc_channel_pool.get_channel(service_location.insecure_address) + self._stub = self._stub_class(channel) + + return self._stub From 4ba6ac691404eec59abfa5229683d9e314314e48 Mon Sep 17 00:00:00 2001 From: Joel Dixon Date: Fri, 5 Sep 2025 15:14:21 -0500 Subject: [PATCH 02/12] Add a GrpcServiceClientBase and use it in sessionmanagement.v1.client --- .../sessionmanagement/v1/client/_client.py | 41 ++++--------- .../v1/client/_client_base.py | 59 +++++++++++++++++++ 2 files changed, 70 insertions(+), 30 deletions(-) create mode 100644 packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py diff --git a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py index 83351a5d..f8ee85f5 100644 --- a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py +++ b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py @@ -3,12 +3,11 @@ from __future__ import annotations import logging -import threading import warnings from collections.abc import Iterable, Mapping import google.protobuf.internal.containers -import grpc +from ni.measurementlink.pinmap.v1.client._client_base import GrpcServiceClientBase import ni.measurementlink.sessionmanagement.v1.session_management_service_pb2 as session_management_service_pb2 import ni.measurementlink.sessionmanagement.v1.session_management_service_pb2_grpc as session_management_service_pb2_grpc from ni.measurementlink.discovery.v1.client import DiscoveryClient @@ -31,35 +30,17 @@ _logger = logging.getLogger(__name__) +class SessionManagementClient(GrpcServiceClientBase): + def __init__(self, *, discovery_client=None, grpc_channel=None, grpc_channel_pool=None): + super().__init__( + discovery_client=discovery_client, + grpc_channel=grpc_channel, + grpc_channel_pool=grpc_channel_pool, + service_interface_name=GRPC_SERVICE_INTERFACE_NAME, + service_class=GRPC_SERVICE_CLASS, + stub_class=session_management_service_pb2_grpc.SessionManagementServiceStub, + ) -class SessionManagementClient: - """Client for accessing the measurement plug-in session management service.""" - - def __init__( - self, - *, - discovery_client: DiscoveryClient | None = None, - grpc_channel: grpc.Channel | None = None, - grpc_channel_pool: GrpcChannelPool | None = None, - ) -> None: - """Initialize session management client. - - Args: - discovery_client: An optional discovery client (recommended). - - grpc_channel: An optional session management gRPC channel. - - grpc_channel_pool: An optional gRPC channel pool (recommended). - """ - self._initialization_lock = threading.Lock() - self._discovery_client = discovery_client - self._grpc_channel_pool = grpc_channel_pool - self._stub: session_management_service_pb2_grpc.SessionManagementServiceStub | None = None - - if grpc_channel is not None: - self._stub = session_management_service_pb2_grpc.SessionManagementServiceStub( - grpc_channel - ) def _get_stub(self) -> session_management_service_pb2_grpc.SessionManagementServiceStub: if self._stub is None: diff --git a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py new file mode 100644 index 00000000..f4d313fc --- /dev/null +++ b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py @@ -0,0 +1,59 @@ +import logging +import threading + +import grpc +from ni.measurementlink.discovery.v1.client import DiscoveryClient +from ni_grpc_extensions.channelpool import GrpcChannelPool + +_logger = logging.getLogger(__name__) + + +class GrpcServiceClientBase: + """Base class for NI gRPC service clients.""" + + def __init__( + self, + *, + discovery_client: DiscoveryClient | None = None, + grpc_channel: grpc.Channel | None = None, + grpc_channel_pool: GrpcChannelPool | None = None, + service_interface_name: str, + service_class: str, + stub_class: type, + ) -> None: + """Initialize a GrpcServiceClientBase instance.""" + self._initialization_lock = threading.Lock() + self._discovery_client = discovery_client + self._grpc_channel_pool = grpc_channel_pool + self._stub = stub_class(grpc_channel) if grpc_channel is not None else None + self._service_interface_name = service_interface_name + self._service_class = service_class + self._stub_class = stub_class + + def _get_stub(self) -> object: + if self._stub is None: + with self._initialization_lock: + if self._grpc_channel_pool is None: + _logger.debug("Creating unshared GrpcChannelPool.") + self._grpc_channel_pool = GrpcChannelPool() + + if self._discovery_client is None: + _logger.debug("Creating unshared DiscoveryClient.") + self._discovery_client = DiscoveryClient( + grpc_channel_pool=self._grpc_channel_pool + ) + + compute_nodes = self._discovery_client.enumerate_compute_nodes() + remote_nodes = [node for node in compute_nodes if not node.is_local] + target_url = remote_nodes[0].url if len(remote_nodes) == 1 else "" + + service_location = self._discovery_client.resolve_service( + provided_interface=self._service_interface_name, + deployment_target=target_url, + service_class=self._service_class, + ) + + channel = self._grpc_channel_pool.get_channel(service_location.insecure_address) + self._stub = self._stub_class(channel) + + return self._stub From 3911b2f9cc0e4582778adc30c95b510782f015e2 Mon Sep 17 00:00:00 2001 From: Joel Dixon Date: Fri, 5 Sep 2025 15:21:09 -0500 Subject: [PATCH 03/12] Fix some linter errors --- .../measurementlink/sessionmanagement/v1/client/_client.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py index f8ee85f5..a7e867c3 100644 --- a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py +++ b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py @@ -7,10 +7,10 @@ from collections.abc import Iterable, Mapping import google.protobuf.internal.containers -from ni.measurementlink.pinmap.v1.client._client_base import GrpcServiceClientBase import ni.measurementlink.sessionmanagement.v1.session_management_service_pb2 as session_management_service_pb2 import ni.measurementlink.sessionmanagement.v1.session_management_service_pb2_grpc as session_management_service_pb2_grpc from ni.measurementlink.discovery.v1.client import DiscoveryClient +from ni.measurementlink.pinmap.v1.client._client_base import GrpcServiceClientBase from ni_grpc_extensions.channelpool import GrpcChannelPool from ni.measurementlink.sessionmanagement.v1.client._constants import ( @@ -30,8 +30,12 @@ _logger = logging.getLogger(__name__) + class SessionManagementClient(GrpcServiceClientBase): + """Client for accessing the NI Session Management Service via gRPC.""" + def __init__(self, *, discovery_client=None, grpc_channel=None, grpc_channel_pool=None): + """Initialize a SessionManagementClient instance.""" super().__init__( discovery_client=discovery_client, grpc_channel=grpc_channel, @@ -41,7 +45,6 @@ def __init__(self, *, discovery_client=None, grpc_channel=None, grpc_channel_poo stub_class=session_management_service_pb2_grpc.SessionManagementServiceStub, ) - def _get_stub(self) -> session_management_service_pb2_grpc.SessionManagementServiceStub: if self._stub is None: with self._initialization_lock: From abb4212fb01af0c6de86403642d2f13c41ef8519 Mon Sep 17 00:00:00 2001 From: Joel Dixon Date: Fri, 5 Sep 2025 15:34:06 -0500 Subject: [PATCH 04/12] Fix some linter errors --- .../sessionmanagement/v1/client/_client.py | 16 ++++++++++++++-- .../sessionmanagement/v1/client/_reservation.py | 10 ++++++---- 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py index a7e867c3..cacb6fbc 100644 --- a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py +++ b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py @@ -7,12 +7,15 @@ from collections.abc import Iterable, Mapping import google.protobuf.internal.containers +import grpc import ni.measurementlink.sessionmanagement.v1.session_management_service_pb2 as session_management_service_pb2 import ni.measurementlink.sessionmanagement.v1.session_management_service_pb2_grpc as session_management_service_pb2_grpc from ni.measurementlink.discovery.v1.client import DiscoveryClient -from ni.measurementlink.pinmap.v1.client._client_base import GrpcServiceClientBase from ni_grpc_extensions.channelpool import GrpcChannelPool +from ni.measurementlink.sessionmanagement.v1.client._client_base import ( + GrpcServiceClientBase, +) from ni.measurementlink.sessionmanagement.v1.client._constants import ( GRPC_SERVICE_CLASS, GRPC_SERVICE_INTERFACE_NAME, @@ -34,7 +37,13 @@ class SessionManagementClient(GrpcServiceClientBase): """Client for accessing the NI Session Management Service via gRPC.""" - def __init__(self, *, discovery_client=None, grpc_channel=None, grpc_channel_pool=None): + def __init__( + self, + *, + discovery_client: DiscoveryClient | None = None, + grpc_channel: grpc.Channel | None = None, + grpc_channel_pool: GrpcChannelPool | None = None, + ) -> None: """Initialize a SessionManagementClient instance.""" super().__init__( discovery_client=discovery_client, @@ -44,6 +53,9 @@ def __init__(self, *, discovery_client=None, grpc_channel=None, grpc_channel_poo service_class=GRPC_SERVICE_CLASS, stub_class=session_management_service_pb2_grpc.SessionManagementServiceStub, ) + self._discovery_client: DiscoveryClient | None = discovery_client + self._grpc_channel_pool: GrpcChannelPool | None = grpc_channel_pool + self._stub: session_management_service_pb2_grpc.SessionManagementServiceStub | None = None def _get_stub(self) -> session_management_service_pb2_grpc.SessionManagementServiceStub: if self._stub is None: diff --git a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_reservation.py b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_reservation.py index 6a1c5b47..54cb6147 100644 --- a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_reservation.py +++ b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_reservation.py @@ -172,15 +172,17 @@ def __init__( @property def _discovery_client(self) -> DiscoveryClient: - if not self._session_management_client._discovery_client: + client = self._session_management_client._discovery_client + if client is None: raise ValueError("This method requires a discovery client.") - return self._session_management_client._discovery_client + return client @property def _grpc_channel_pool(self) -> GrpcChannelPool: - if not self._session_management_client._grpc_channel_pool: + pool = self._session_management_client._grpc_channel_pool + if pool is None: raise ValueError("This method requires a gRPC channel pool.") - return self._session_management_client._grpc_channel_pool + return pool class MultiplexerSessionContainer(_BaseSessionContainer): From adb15116f0f062a0c19b6c236238c8cda8fe3365 Mon Sep 17 00:00:00 2001 From: Joel Dixon Date: Fri, 5 Sep 2025 15:47:16 -0500 Subject: [PATCH 05/12] Fix some unit test failures only found in Python 3.9 --- .../ni/measurementlink/pinmap/v1/client/_client_base.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py index f4d313fc..8bb5b3b2 100644 --- a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py +++ b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py @@ -1,5 +1,6 @@ import logging import threading +from typing import Optional import grpc from ni.measurementlink.discovery.v1.client import DiscoveryClient @@ -14,9 +15,9 @@ class GrpcServiceClientBase: def __init__( self, *, - discovery_client: DiscoveryClient | None = None, - grpc_channel: grpc.Channel | None = None, - grpc_channel_pool: GrpcChannelPool | None = None, + discovery_client: Optional[DiscoveryClient] = None, + grpc_channel: Optional[grpc.Channel] = None, + grpc_channel_pool: Optional[GrpcChannelPool] = None, service_interface_name: str, service_class: str, stub_class: type, From 0c5c618e94ba1ee91c9b43d436eae0dd2d361c9c Mon Sep 17 00:00:00 2001 From: Joel Dixon Date: Fri, 5 Sep 2025 15:55:44 -0500 Subject: [PATCH 06/12] Make the _client_base.py files identical --- .../sessionmanagement/v1/client/_client_base.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py index f4d313fc..8bb5b3b2 100644 --- a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py +++ b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py @@ -1,5 +1,6 @@ import logging import threading +from typing import Optional import grpc from ni.measurementlink.discovery.v1.client import DiscoveryClient @@ -14,9 +15,9 @@ class GrpcServiceClientBase: def __init__( self, *, - discovery_client: DiscoveryClient | None = None, - grpc_channel: grpc.Channel | None = None, - grpc_channel_pool: GrpcChannelPool | None = None, + discovery_client: Optional[DiscoveryClient] = None, + grpc_channel: Optional[grpc.Channel] = None, + grpc_channel_pool: Optional[GrpcChannelPool] = None, service_interface_name: str, service_class: str, stub_class: type, From e73ab06ed8b8e387c09db3f2887a123ab3ac4b94 Mon Sep 17 00:00:00 2001 From: Joel Dixon Date: Tue, 9 Sep 2025 11:56:55 -0500 Subject: [PATCH 07/12] Fix string type hints to use from __future__ import annotations --- .../src/ni/measurementlink/pinmap/v1/client/_client.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py index d0347fae..5e905619 100644 --- a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py +++ b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py @@ -19,9 +19,9 @@ class PinMapClient(GrpcServiceClientBase): def __init__( self, *, - discovery_client: "DiscoveryClient | None" = None, - grpc_channel: "grpc.Channel | None" = None, - grpc_channel_pool: "GrpcChannelPool | None" = None, + discovery_client: DiscoveryClient | None = None, + grpc_channel: grpc.Channel | None = None, + grpc_channel_pool: GrpcChannelPool | None = None, ) -> None: """Initialize a PinMapClient instance.""" super().__init__( From f48b045fdc063144fe8a2e919128535dbae2d2d2 Mon Sep 17 00:00:00 2001 From: Joel Dixon Date: Tue, 9 Sep 2025 12:01:36 -0500 Subject: [PATCH 08/12] Fix _client.py comments --- .../pinmap/v1/client/_client.py | 24 ++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py index 5e905619..02f2ec78 100644 --- a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py +++ b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py @@ -23,7 +23,15 @@ def __init__( grpc_channel: grpc.Channel | None = None, grpc_channel_pool: GrpcChannelPool | None = None, ) -> None: - """Initialize a PinMapClient instance.""" + """Initialize the pin map client. + + Args: + discovery_client: An optional discovery client (recommended). + + grpc_channel: An optional pin map gRPC channel. + + grpc_channel_pool: An optional gRPC channel pool (recommended). + """ super().__init__( discovery_client=discovery_client, grpc_channel=grpc_channel, @@ -34,11 +42,21 @@ def __init__( ) def update_pin_map(self, pin_map_path: str | pathlib.Path) -> str: - """Update the pin map from an XML file and return the pin map ID.""" + """Update registered pin map contents. + + Create and register a pin map if a pin map resource for the specified pin map id is not + found. + + Args: + pin_map_path: The file path of the pin map to register as a pin map resource. + + Returns: + The resource id of the pin map that is registered to the pin map service. + """ request = pin_map_service_pb2.UpdatePinMapFromXmlRequest( pin_map_id=str(pin_map_path), pin_map_xml=pathlib.Path(pin_map_path).read_text(encoding="utf-8-sig"), ) stub: pin_map_service_pb2_grpc.PinMapServiceStub = self._get_stub() # type: ignore response = stub.UpdatePinMapFromXml(request) - return str(response.pin_map_id) + return response.pin_map_id From d8015fa81ba8c717886a97315a68defb04ba2368 Mon Sep 17 00:00:00 2001 From: Joel Dixon Date: Tue, 9 Sep 2025 13:48:51 -0500 Subject: [PATCH 09/12] Fix a few PR comments, add a type param to GrpcServiceClientBase --- .../pinmap/v1/client/_client.py | 4 +- .../pinmap/v1/client/_client_base.py | 42 +++++++++++++++---- .../sessionmanagement/v1/client/_client.py | 4 +- .../v1/client/_client_base.py | 42 +++++++++++++++---- .../v1/client/_reservation.py | 6 +-- 5 files changed, 75 insertions(+), 23 deletions(-) diff --git a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py index 02f2ec78..d9a32b34 100644 --- a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py +++ b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py @@ -13,7 +13,7 @@ from ni.measurementlink.pinmap.v1.client._client_base import GrpcServiceClientBase -class PinMapClient(GrpcServiceClientBase): +class PinMapClient(GrpcServiceClientBase[pin_map_service_pb2_grpc.PinMapServiceStub]): """Client for accessing the NI Pin Map Service via gRPC.""" def __init__( @@ -57,6 +57,6 @@ def update_pin_map(self, pin_map_path: str | pathlib.Path) -> str: pin_map_id=str(pin_map_path), pin_map_xml=pathlib.Path(pin_map_path).read_text(encoding="utf-8-sig"), ) - stub: pin_map_service_pb2_grpc.PinMapServiceStub = self._get_stub() # type: ignore + stub: pin_map_service_pb2_grpc.PinMapServiceStub = self._get_stub() response = stub.UpdatePinMapFromXml(request) return response.pin_map_id diff --git a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py index 8bb5b3b2..84c489b9 100644 --- a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py +++ b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py @@ -1,6 +1,8 @@ +from __future__ import annotations + import logging import threading -from typing import Optional +from typing import Generic, Protocol, TypeVar import grpc from ni.measurementlink.discovery.v1.client import DiscoveryClient @@ -9,20 +11,44 @@ _logger = logging.getLogger(__name__) -class GrpcServiceClientBase: +class StubProtocol(Protocol): + """Protocol for gRPC stub classes.""" + + def __init__(self, channel: grpc.Channel) -> None: + """Initialize the gRPC client.""" + + +TStub = TypeVar("TStub", bound=StubProtocol) + + +class GrpcServiceClientBase(Generic[TStub]): """Base class for NI gRPC service clients.""" def __init__( self, *, - discovery_client: Optional[DiscoveryClient] = None, - grpc_channel: Optional[grpc.Channel] = None, - grpc_channel_pool: Optional[GrpcChannelPool] = None, + discovery_client: DiscoveryClient | None = None, + grpc_channel: grpc.Channel | None = None, + grpc_channel_pool: GrpcChannelPool | None = None, service_interface_name: str, service_class: str, - stub_class: type, + stub_class: type[TStub], ) -> None: - """Initialize a GrpcServiceClientBase instance.""" + """Initialize the gRPC client. + + Args: + discovery_client: An optional discovery client (recommended). + + grpc_channel: An optional pin map gRPC channel. + + grpc_channel_pool: An optional gRPC channel pool (recommended). + + service_interface_name: The fully qualified name of the service interface. + + service_class: The name of the service class. + + stub_class: The gRPC stub class for the service. + """ self._initialization_lock = threading.Lock() self._discovery_client = discovery_client self._grpc_channel_pool = grpc_channel_pool @@ -31,7 +57,7 @@ def __init__( self._service_class = service_class self._stub_class = stub_class - def _get_stub(self) -> object: + def _get_stub(self) -> TStub: if self._stub is None: with self._initialization_lock: if self._grpc_channel_pool is None: diff --git a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py index cacb6fbc..3e2fba05 100644 --- a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py +++ b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py @@ -34,7 +34,9 @@ _logger = logging.getLogger(__name__) -class SessionManagementClient(GrpcServiceClientBase): +class SessionManagementClient( + GrpcServiceClientBase[session_management_service_pb2_grpc.SessionManagementServiceStub] +): """Client for accessing the NI Session Management Service via gRPC.""" def __init__( diff --git a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py index 8bb5b3b2..84c489b9 100644 --- a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py +++ b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py @@ -1,6 +1,8 @@ +from __future__ import annotations + import logging import threading -from typing import Optional +from typing import Generic, Protocol, TypeVar import grpc from ni.measurementlink.discovery.v1.client import DiscoveryClient @@ -9,20 +11,44 @@ _logger = logging.getLogger(__name__) -class GrpcServiceClientBase: +class StubProtocol(Protocol): + """Protocol for gRPC stub classes.""" + + def __init__(self, channel: grpc.Channel) -> None: + """Initialize the gRPC client.""" + + +TStub = TypeVar("TStub", bound=StubProtocol) + + +class GrpcServiceClientBase(Generic[TStub]): """Base class for NI gRPC service clients.""" def __init__( self, *, - discovery_client: Optional[DiscoveryClient] = None, - grpc_channel: Optional[grpc.Channel] = None, - grpc_channel_pool: Optional[GrpcChannelPool] = None, + discovery_client: DiscoveryClient | None = None, + grpc_channel: grpc.Channel | None = None, + grpc_channel_pool: GrpcChannelPool | None = None, service_interface_name: str, service_class: str, - stub_class: type, + stub_class: type[TStub], ) -> None: - """Initialize a GrpcServiceClientBase instance.""" + """Initialize the gRPC client. + + Args: + discovery_client: An optional discovery client (recommended). + + grpc_channel: An optional pin map gRPC channel. + + grpc_channel_pool: An optional gRPC channel pool (recommended). + + service_interface_name: The fully qualified name of the service interface. + + service_class: The name of the service class. + + stub_class: The gRPC stub class for the service. + """ self._initialization_lock = threading.Lock() self._discovery_client = discovery_client self._grpc_channel_pool = grpc_channel_pool @@ -31,7 +57,7 @@ def __init__( self._service_class = service_class self._stub_class = stub_class - def _get_stub(self) -> object: + def _get_stub(self) -> TStub: if self._stub is None: with self._initialization_lock: if self._grpc_channel_pool is None: diff --git a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_reservation.py b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_reservation.py index 54cb6147..d2967a4d 100644 --- a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_reservation.py +++ b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_reservation.py @@ -276,9 +276,7 @@ def _validate_and_get_matching_multiplexer_session_infos( return multiplexer_session_infos @contextlib.contextmanager - def _cache_multiplexer_session( - self, session_name: str, session: TMultiplexerSession - ) -> Generator[None]: # pyright: ignore[reportInvalidTypeVarUse] + def _cache_multiplexer_session(self, session_name: str, session: object) -> Generator[None]: if session_name in self._multiplexer_session_cache: raise RuntimeError(f"Multiplexer session '{session_name}' already exists.") self._multiplexer_session_cache[session_name] = session @@ -666,7 +664,7 @@ def unreserve(self) -> None: self._session_management_client._unreserve_sessions(self._grpc_session_info) @contextlib.contextmanager - def _cache_session(self, session_name: str, session: TSession) -> Generator[None]: + def _cache_session(self, session_name: str, session: object) -> Generator[None]: if session_name in self._session_cache: raise RuntimeError(f"Session '{session_name}' already exists.") self._session_cache[session_name] = session From d98ec51ba82cb98d8375bfb295b14a7f449e2ac2 Mon Sep 17 00:00:00 2001 From: Joel Dixon Date: Tue, 9 Sep 2025 13:52:04 -0500 Subject: [PATCH 10/12] Remove _get_stub from _client.py in sessionmanagement.client package --- .../sessionmanagement/v1/client/_client.py | 30 ------------------- 1 file changed, 30 deletions(-) diff --git a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py index 3e2fba05..4d9b8b7d 100644 --- a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py +++ b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py @@ -59,36 +59,6 @@ def __init__( self._grpc_channel_pool: GrpcChannelPool | None = grpc_channel_pool self._stub: session_management_service_pb2_grpc.SessionManagementServiceStub | None = None - def _get_stub(self) -> session_management_service_pb2_grpc.SessionManagementServiceStub: - if self._stub is None: - with self._initialization_lock: - if self._grpc_channel_pool is None: - _logger.debug("Creating unshared GrpcChannelPool.") - self._grpc_channel_pool = GrpcChannelPool() - if self._discovery_client is None: - _logger.debug("Creating unshared DiscoveryClient.") - self._discovery_client = DiscoveryClient( - grpc_channel_pool=self._grpc_channel_pool - ) - if self._stub is None: - compute_nodes = self._discovery_client.enumerate_compute_nodes() - remote_compute_nodes = [node for node in compute_nodes if not node.is_local] - # Use remote node URL as deployment target if only one remote node is found. - # If more than one remote node exists, use empty string for deployment target. - first_remote_node_url = ( - remote_compute_nodes[0].url if len(remote_compute_nodes) == 1 else "" - ) - service_location = self._discovery_client.resolve_service( - provided_interface=GRPC_SERVICE_INTERFACE_NAME, - deployment_target=first_remote_node_url, - service_class=GRPC_SERVICE_CLASS, - ) - channel = self._grpc_channel_pool.get_channel(service_location.insecure_address) - self._stub = session_management_service_pb2_grpc.SessionManagementServiceStub( - channel - ) - return self._stub - def reserve_session( self, context: PinMapContext, From bc6ae096291192b05b5c2e56137afe1cd113866f Mon Sep 17 00:00:00 2001 From: Joel Dixon Date: Tue, 9 Sep 2025 14:10:10 -0500 Subject: [PATCH 11/12] Fix type hints and use slots --- .../pinmap/v1/client/_client_base.py | 23 +++++++++++++++---- .../sessionmanagement/v1/client/_client.py | 5 ++-- .../v1/client/_client_base.py | 23 +++++++++++++++---- 3 files changed, 38 insertions(+), 13 deletions(-) diff --git a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py index 84c489b9..ec527546 100644 --- a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py +++ b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py @@ -24,6 +24,24 @@ def __init__(self, channel: grpc.Channel) -> None: class GrpcServiceClientBase(Generic[TStub]): """Base class for NI gRPC service clients.""" + __slots__ = ( + "_initialization_lock", + "_discovery_client", + "_grpc_channel_pool", + "_stub", + "_service_interface_name", + "_service_class", + "_stub_class", + ) + + _initialization_lock: threading.Lock + _discovery_client: DiscoveryClient | None + _grpc_channel_pool: GrpcChannelPool | None + _stub: TStub | None + _service_interface_name: str + _service_class: str + _stub_class: type[TStub] + def __init__( self, *, @@ -38,15 +56,10 @@ def __init__( Args: discovery_client: An optional discovery client (recommended). - grpc_channel: An optional pin map gRPC channel. - grpc_channel_pool: An optional gRPC channel pool (recommended). - service_interface_name: The fully qualified name of the service interface. - service_class: The name of the service class. - stub_class: The gRPC stub class for the service. """ self._initialization_lock = threading.Lock() diff --git a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py index 4d9b8b7d..9a58c865 100644 --- a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py +++ b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client.py @@ -39,6 +39,8 @@ class SessionManagementClient( ): """Client for accessing the NI Session Management Service via gRPC.""" + __slots__ = () + def __init__( self, *, @@ -55,9 +57,6 @@ def __init__( service_class=GRPC_SERVICE_CLASS, stub_class=session_management_service_pb2_grpc.SessionManagementServiceStub, ) - self._discovery_client: DiscoveryClient | None = discovery_client - self._grpc_channel_pool: GrpcChannelPool | None = grpc_channel_pool - self._stub: session_management_service_pb2_grpc.SessionManagementServiceStub | None = None def reserve_session( self, diff --git a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py index 84c489b9..ec527546 100644 --- a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py +++ b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py @@ -24,6 +24,24 @@ def __init__(self, channel: grpc.Channel) -> None: class GrpcServiceClientBase(Generic[TStub]): """Base class for NI gRPC service clients.""" + __slots__ = ( + "_initialization_lock", + "_discovery_client", + "_grpc_channel_pool", + "_stub", + "_service_interface_name", + "_service_class", + "_stub_class", + ) + + _initialization_lock: threading.Lock + _discovery_client: DiscoveryClient | None + _grpc_channel_pool: GrpcChannelPool | None + _stub: TStub | None + _service_interface_name: str + _service_class: str + _stub_class: type[TStub] + def __init__( self, *, @@ -38,15 +56,10 @@ def __init__( Args: discovery_client: An optional discovery client (recommended). - grpc_channel: An optional pin map gRPC channel. - grpc_channel_pool: An optional gRPC channel pool (recommended). - service_interface_name: The fully qualified name of the service interface. - service_class: The name of the service class. - stub_class: The gRPC stub class for the service. """ self._initialization_lock = threading.Lock() From 570ec5e1ea6eb06585b3f543763ef4a266faa2b5 Mon Sep 17 00:00:00 2001 From: Joel Dixon Date: Wed, 10 Sep 2025 09:29:55 -0500 Subject: [PATCH 12/12] Various PR feedback --- .../pinmap/v1/client/_client.py | 6 ++-- .../pinmap/v1/client/_client_base.py | 33 ++++++++++--------- .../v1/client/_client_base.py | 33 ++++++++++--------- 3 files changed, 38 insertions(+), 34 deletions(-) diff --git a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py index d9a32b34..54e84129 100644 --- a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py +++ b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client.py @@ -16,6 +16,8 @@ class PinMapClient(GrpcServiceClientBase[pin_map_service_pb2_grpc.PinMapServiceStub]): """Client for accessing the NI Pin Map Service via gRPC.""" + __slots__ = () + def __init__( self, *, @@ -53,10 +55,10 @@ def update_pin_map(self, pin_map_path: str | pathlib.Path) -> str: Returns: The resource id of the pin map that is registered to the pin map service. """ + # By convention, the pin map id is the .pinmap file path. request = pin_map_service_pb2.UpdatePinMapFromXmlRequest( pin_map_id=str(pin_map_path), pin_map_xml=pathlib.Path(pin_map_path).read_text(encoding="utf-8-sig"), ) - stub: pin_map_service_pb2_grpc.PinMapServiceStub = self._get_stub() - response = stub.UpdatePinMapFromXml(request) + response = self._get_stub().UpdatePinMapFromXml(request) return response.pin_map_id diff --git a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py index ec527546..59c62f7d 100644 --- a/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py +++ b/packages/ni.measurementlink.pinmap.v1.client/src/ni/measurementlink/pinmap/v1/client/_client_base.py @@ -44,23 +44,23 @@ class GrpcServiceClientBase(Generic[TStub]): def __init__( self, + service_interface_name: str, + service_class: str, + stub_class: type[TStub], *, discovery_client: DiscoveryClient | None = None, grpc_channel: grpc.Channel | None = None, grpc_channel_pool: GrpcChannelPool | None = None, - service_interface_name: str, - service_class: str, - stub_class: type[TStub], ) -> None: """Initialize the gRPC client. Args: - discovery_client: An optional discovery client (recommended). - grpc_channel: An optional pin map gRPC channel. - grpc_channel_pool: An optional gRPC channel pool (recommended). service_interface_name: The fully qualified name of the service interface. service_class: The name of the service class. stub_class: The gRPC stub class for the service. + discovery_client: An optional discovery client (recommended). + grpc_channel: An optional pin map gRPC channel. + grpc_channel_pool: An optional gRPC channel pool (recommended). """ self._initialization_lock = threading.Lock() self._discovery_client = discovery_client @@ -83,17 +83,18 @@ def _get_stub(self) -> TStub: grpc_channel_pool=self._grpc_channel_pool ) - compute_nodes = self._discovery_client.enumerate_compute_nodes() - remote_nodes = [node for node in compute_nodes if not node.is_local] - target_url = remote_nodes[0].url if len(remote_nodes) == 1 else "" + if self._stub is None: + compute_nodes = self._discovery_client.enumerate_compute_nodes() + remote_nodes = [node for node in compute_nodes if not node.is_local] + target_url = remote_nodes[0].url if len(remote_nodes) == 1 else "" - service_location = self._discovery_client.resolve_service( - provided_interface=self._service_interface_name, - deployment_target=target_url, - service_class=self._service_class, - ) + service_location = self._discovery_client.resolve_service( + provided_interface=self._service_interface_name, + deployment_target=target_url, + service_class=self._service_class, + ) - channel = self._grpc_channel_pool.get_channel(service_location.insecure_address) - self._stub = self._stub_class(channel) + channel = self._grpc_channel_pool.get_channel(service_location.insecure_address) + self._stub = self._stub_class(channel) return self._stub diff --git a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py index ec527546..59c62f7d 100644 --- a/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py +++ b/packages/ni.measurementlink.sessionmanagement.v1.client/src/ni/measurementlink/sessionmanagement/v1/client/_client_base.py @@ -44,23 +44,23 @@ class GrpcServiceClientBase(Generic[TStub]): def __init__( self, + service_interface_name: str, + service_class: str, + stub_class: type[TStub], *, discovery_client: DiscoveryClient | None = None, grpc_channel: grpc.Channel | None = None, grpc_channel_pool: GrpcChannelPool | None = None, - service_interface_name: str, - service_class: str, - stub_class: type[TStub], ) -> None: """Initialize the gRPC client. Args: - discovery_client: An optional discovery client (recommended). - grpc_channel: An optional pin map gRPC channel. - grpc_channel_pool: An optional gRPC channel pool (recommended). service_interface_name: The fully qualified name of the service interface. service_class: The name of the service class. stub_class: The gRPC stub class for the service. + discovery_client: An optional discovery client (recommended). + grpc_channel: An optional pin map gRPC channel. + grpc_channel_pool: An optional gRPC channel pool (recommended). """ self._initialization_lock = threading.Lock() self._discovery_client = discovery_client @@ -83,17 +83,18 @@ def _get_stub(self) -> TStub: grpc_channel_pool=self._grpc_channel_pool ) - compute_nodes = self._discovery_client.enumerate_compute_nodes() - remote_nodes = [node for node in compute_nodes if not node.is_local] - target_url = remote_nodes[0].url if len(remote_nodes) == 1 else "" + if self._stub is None: + compute_nodes = self._discovery_client.enumerate_compute_nodes() + remote_nodes = [node for node in compute_nodes if not node.is_local] + target_url = remote_nodes[0].url if len(remote_nodes) == 1 else "" - service_location = self._discovery_client.resolve_service( - provided_interface=self._service_interface_name, - deployment_target=target_url, - service_class=self._service_class, - ) + service_location = self._discovery_client.resolve_service( + provided_interface=self._service_interface_name, + deployment_target=target_url, + service_class=self._service_class, + ) - channel = self._grpc_channel_pool.get_channel(service_location.insecure_address) - self._stub = self._stub_class(channel) + channel = self._grpc_channel_pool.get_channel(service_location.insecure_address) + self._stub = self._stub_class(channel) return self._stub