Fix for keyring errors when initializing Flyte for_sandbox config client

.github/workflows/monodocs_build.yml

@@ -57,4 +57,5 @@ jobs:
           DOCSEARCH_API_KEY: fake_docsearch_api_key  # must be set to get doc build to succeed
         run: |
           conda activate monodocs-env
+          pip install grpcio-health-checking==1.49.0
           make -C docs clean html SPHINXOPTS="-W -vvv"

dev-requirements.in

@@ -60,3 +60,7 @@ ipykernel
 
 orjson
 kubernetes>=12.0.1
+
+grpcio-health-checking
+grpcio-tools
+grpcio

dev-requirements.txt

@@ -564,3 +564,5 @@ zipp==3.19.1
 
 # The following packages are considered to be unsafe in a requirements file:
 # setuptools
+
+grpcio-health-checking==1.49.0

flytekit/clients/auth_helper.py

@@ -122,9 +122,12 @@ def upgrade_channel_to_proxy_authenticated(cfg: PlatformConfig, in_channel: grpc
     :param in_channel: grpc.Channel Precreated channel
     :return: grpc.Channel. New composite channel
     """
+
+    def authenticator_factory():
+        return get_proxy_authenticator(cfg)
+
     if cfg.proxy_command:
-        proxy_authenticator = get_proxy_authenticator(cfg)
-        return grpc.intercept_channel(in_channel, AuthUnaryInterceptor(proxy_authenticator))
+        return grpc.intercept_channel(in_channel, AuthUnaryInterceptor(authenticator_factory))
     else:
         return in_channel
 
@@ -137,8 +140,11 @@ def upgrade_channel_to_authenticated(cfg: PlatformConfig, in_channel: grpc.Chann
     :param in_channel: grpc.Channel Precreated channel
     :return: grpc.Channel. New composite channel
     """
-    authenticator = get_authenticator(cfg, RemoteClientConfigStore(in_channel))
-    return grpc.intercept_channel(in_channel, AuthUnaryInterceptor(authenticator))
+
+    def authenticator_factory():
+        return get_authenticator(cfg, RemoteClientConfigStore(in_channel))
+
+    return grpc.intercept_channel(in_channel, AuthUnaryInterceptor(authenticator_factory))
 
 
 def get_authenticated_channel(cfg: PlatformConfig) -> grpc.Channel:

flytekit/clients/grpc_utils/auth_interceptor.py

@@ -25,15 +25,22 @@
     is needed.
     """
 
-    def __init__(self, authenticator: Authenticator):
-        self._authenticator = authenticator
+    def __init__(self, get_authenticator: typing.Callable[[], Authenticator]):
+        self._get_authenticator = get_authenticator
+        self._authenticator = None
+
+    @property
+    def authenticator(self) -> Authenticator:
+        if self._authenticator is None:
+            self._authenticator = self._get_authenticator()
+        return self._authenticator
 
     def _call_details_with_auth_metadata(self, client_call_details: grpc.ClientCallDetails) -> grpc.ClientCallDetails:
         """
         Returns new ClientCallDetails with metadata added.
         """
         metadata = client_call_details.metadata
-        auth_metadata = self._authenticator.fetch_grpc_call_auth_metadata()
+        auth_metadata = self.authenticator.fetch_grpc_call_auth_metadata()
         if auth_metadata:
             metadata = []
             if client_call_details.metadata:
@@ -65,6 +72,7 @@
                 raise e
             if e.code() == grpc.StatusCode.UNAUTHENTICATED or e.code() == grpc.StatusCode.UNKNOWN:
                 self._authenticator.refresh_credentials()
+                self.authenticator.refresh_credentials()
                 updated_call_details = self._call_details_with_auth_metadata(client_call_details)
                 return continuation(updated_call_details, request)
         return fut
@@ -77,6 +85,7 @@
         c: grpc.Call = continuation(updated_call_details, request)
         if c.code() == grpc.StatusCode.UNAUTHENTICATED:
             self._authenticator.refresh_credentials()
+            self.authenticator.refresh_credentials()
             updated_call_details = self._call_details_with_auth_metadata(client_call_details)
             return continuation(updated_call_details, request)
         return c

flytekit/clients/raw.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 import typing
 
 import grpc
@@ -18,6 +19,12 @@
     wrap_exceptions_channel,
 )
 from flytekit.configuration import PlatformConfig
+from flytekit.exceptions.system import FlyteSystemUnavailableException
+from flytekit.exceptions.user import (
+    FlyteEntityAlreadyExistsException,
+    FlyteEntityNotExistException,
+    FlyteInvalidInputException,
+)
 from flytekit.loggers import logger
 
 
@@ -51,12 +58,14 @@
         # 32KB for error messages, 20MB for actual messages.
         options = (("grpc.max_metadata_size", 32 * 1024), ("grpc.max_receive_message_length", 20 * 1024 * 1024))
         self._cfg = cfg
+
         self._channel = wrap_exceptions_channel(
             cfg,
             upgrade_channel_to_authenticated(
                 cfg, upgrade_channel_to_proxy_authenticated(cfg, get_channel(cfg, options=options))
             ),
         )
+
         self._stub = _admin_service.AdminServiceStub(self._channel)
         self._signal = signal_service.SignalServiceStub(self._channel)
         self._dataproxy_stub = dataproxy_service.DataProxyServiceStub(self._channel)
@@ -67,6 +76,29 @@
         # metadata will hold the value of the token to send to the various endpoints.
         self._metadata = None
 
+    @staticmethod
+    def check_grpc_health_with_authentication(in_channel):
+        from grpc_health.v1 import health_pb2, health_pb2_grpc
+
+        health_stub = health_pb2_grpc.HealthStub(in_channel)
+        request = health_pb2.HealthCheckRequest()
+        try:
+            response = health_stub.Check(request)
+            if response.status == health_pb2.HealthCheckResponse.SERVING:
+                logging.info("Service is healthy and ready to serve.")
+                return True
+        except grpc.RpcError as e:
+            if e.code() == grpc.StatusCode.UNAUTHENTICATED:
+                return False
+            elif e.code() == grpc.StatusCode.ALREADY_EXISTS:
+                raise FlyteEntityAlreadyExistsException() from e
+            elif e.code() == grpc.StatusCode.NOT_FOUND:
+                raise FlyteEntityNotExistException() from e
+            elif e.code() == grpc.StatusCode.INVALID_ARGUMENT:
+                raise FlyteInvalidInputException(request) from e
+            elif e.code() == grpc.StatusCode.UNAVAILABLE:
+                raise FlyteSystemUnavailableException() from e
+
     @classmethod
     def with_root_certificate(cls, cfg: PlatformConfig, root_cert_file: str) -> RawSynchronousFlyteClient:
         b = None

tests/flytekit/unit/clients/test_auth_helper.py

@@ -172,7 +172,7 @@ def test_upgrade_channel_to_proxy_auth():
         ch,
     )
     assert isinstance(out_ch._interceptor, AuthUnaryInterceptor)
-    assert isinstance(out_ch._interceptor._authenticator, CommandAuthenticator)
+    assert isinstance(out_ch._interceptor.authenticator, CommandAuthenticator)
 
 
 def test_get_proxy_authenticated_session():

tests/flytekit/unit/clients/test_friendly.py

@@ -8,26 +8,29 @@
 from flytekit.clients.friendly import SynchronousFlyteClient as _SynchronousFlyteClient
 from flytekit.configuration import PlatformConfig
 from flytekit.models.project import Project as _Project
-
+from grpc_health.v1 import health_pb2
 
 @mock.patch("flytekit.clients.friendly._RawSynchronousFlyteClient.update_project")
-def test_update_project(mock_raw_update_project):
+@mock.patch("flytekit.clients.raw.RawSynchronousFlyteClient.check_grpc_health_with_authentication", return_value=health_pb2.HealthCheckResponse.SERVING)
+def test_update_project(mock_check_health, mock_raw_update_project):
     client = _SynchronousFlyteClient(PlatformConfig.for_endpoint("a.b.com", True))
     project = _Project("foo", "name", "description", state=_Project.ProjectState.ACTIVE)
     client.update_project(project)
     mock_raw_update_project.assert_called_with(project.to_flyte_idl())
 
 
 @mock.patch("flytekit.clients.friendly._RawSynchronousFlyteClient.list_projects")
-def test_list_projects_paginated(mock_raw_list_projects):
+@mock.patch("flytekit.clients.raw.RawSynchronousFlyteClient.check_grpc_health_with_authentication", return_value=health_pb2.HealthCheckResponse.SERVING)
+def test_list_projects_paginated(mock_check_health, mock_raw_list_projects):
     client = _SynchronousFlyteClient(PlatformConfig.for_endpoint("a.b.com", True))
     client.list_projects_paginated(limit=100, token="")
     project_list_request = _project_pb2.ProjectListRequest(limit=100, token="", filters=None, sort_by=None)
     mock_raw_list_projects.assert_called_with(project_list_request=project_list_request)
 
 
 @mock.patch("flytekit.clients.friendly._RawSynchronousFlyteClient.create_upload_location")
-def test_create_upload_location(mock_raw_create_upload_location):
+@mock.patch("flytekit.clients.raw.RawSynchronousFlyteClient.check_grpc_health_with_authentication", return_value=health_pb2.HealthCheckResponse.SERVING)
+def test_create_upload_location(mock_check_health, mock_raw_create_upload_location):
     client = _SynchronousFlyteClient(PlatformConfig.for_endpoint("a.b.com", True))
     client.get_upload_signed_url("foo", "bar", bytes(), "baz.qux", timedelta(minutes=42), add_content_md5_metadata=True)
     duration_pb = Duration()

tests/flytekit/unit/clients/test_raw.py

@@ -4,11 +4,13 @@
 
 from flytekit.clients.raw import RawSynchronousFlyteClient
 from flytekit.configuration import PlatformConfig
-
+from grpc_health.v1 import health_pb2
 
 @mock.patch("flytekit.clients.raw._admin_service")
 @mock.patch("flytekit.clients.raw.grpc.insecure_channel")
-def test_update_project(mock_channel, mock_admin):
+@mock.patch.object(RawSynchronousFlyteClient, "check_grpc_health_with_authentication", return_value=True)
+def test_update_project(mock_check_health, mock_channel, mock_admin):
+    mock_health_stub = mock.Mock()
     client = RawSynchronousFlyteClient(PlatformConfig(endpoint="a.b.com", insecure=True))
     project = _project_pb2.Project(id="foo", name="name", description="description", state=_project_pb2.Project.ACTIVE)
     client.update_project(project)
@@ -17,7 +19,8 @@ def test_update_project(mock_channel, mock_admin):
 
 @mock.patch("flytekit.clients.raw._admin_service")
 @mock.patch("flytekit.clients.raw.grpc.insecure_channel")
-def test_list_projects_paginated(mock_channel, mock_admin):
+@mock.patch("flytekit.clients.raw.RawSynchronousFlyteClient.check_grpc_health_with_authentication", return_value=health_pb2.HealthCheckResponse.SERVING)
+def test_list_projects_paginated(mock_check_health, mock_channel, mock_admin):
     client = RawSynchronousFlyteClient(PlatformConfig(endpoint="a.b.com", insecure=True))
     project_list_request = _project_pb2.ProjectListRequest(limit=100, token="", filters=None, sort_by=None)
     client.list_projects(project_list_request)