add accept grpc (#1841)

* add accept grpc

Signed-off-by: Yee Hing Tong <[email protected]>
Signed-off-by: Jeev B <[email protected]>

* unpin setup.py grpc

Signed-off-by: Yee Hing Tong <[email protected]>
Signed-off-by: Jeev B <[email protected]>

* Revert "add accept grpc"

This reverts commit 2294592.

Signed-off-by: Jeev B <[email protected]>

* default headers interceptor

Signed-off-by: Jeev B <[email protected]>

* setup.py

Signed-off-by: Jeev B <[email protected]>

* fixes

Signed-off-by: Jeev B <[email protected]>

* fmt

Signed-off-by: Jeev B <[email protected]>

* move prometheus-client import

Signed-off-by: Jeev B <[email protected]>

---------

Signed-off-by: Yee Hing Tong <[email protected]>
Signed-off-by: Jeev B <[email protected]>
Co-authored-by: Jeev B <[email protected]>

Author: wild-endeavor

flytekit/clients/auth_helper.py

@@ -16,6 +16,7 @@
     PKCEAuthenticator,
 )
 from flytekit.clients.grpc_utils.auth_interceptor import AuthUnaryInterceptor
+from flytekit.clients.grpc_utils.default_metadata_interceptor import DefaultMetadataInterceptor
 from flytekit.clients.grpc_utils.wrap_exception_interceptor import RetryExceptionWrapperInterceptor
 from flytekit.configuration import AuthType, PlatformConfig
 
@@ -171,7 +172,7 @@ def get_channel(cfg: PlatformConfig, **kwargs) -> grpc.Channel:
     :return: grpc.Channel (secure / insecure)
     """
     if cfg.insecure:
-        return grpc.insecure_channel(cfg.endpoint, **kwargs)
+        return grpc.intercept_channel(grpc.insecure_channel(cfg.endpoint, **kwargs), DefaultMetadataInterceptor())
 
     credentials = None
     if "credentials" not in kwargs:
@@ -189,11 +190,14 @@ def get_channel(cfg: PlatformConfig, **kwargs) -> grpc.Channel:
             )
     else:
         credentials = kwargs["credentials"]
-    return grpc.secure_channel(
-        target=cfg.endpoint,
-        credentials=credentials,
-        options=kwargs.get("options", None),
-        compression=kwargs.get("compression", None),
+    return grpc.intercept_channel(
+        grpc.secure_channel(
+            target=cfg.endpoint,
+            credentials=credentials,
+            options=kwargs.get("options", None),
+            compression=kwargs.get("compression", None),
+        ),
+        DefaultMetadataInterceptor(),
     )
 
 

flytekit/clients/grpc_utils/auth_interceptor.py

@@ -32,7 +32,7 @@ def _call_details_with_auth_metadata(self, client_call_details: grpc.ClientCallD
         """
         Returns new ClientCallDetails with metadata added.
         """
-        metadata = None
+        metadata = client_call_details.metadata
         auth_metadata = self._authenticator.fetch_grpc_call_auth_metadata()
         if auth_metadata:
             metadata = []

flytekit/clients/grpc_utils/default_metadata_interceptor.py

@@ -0,0 +1,43 @@
+import typing
+
+import grpc
+
+from flytekit.clients.grpc_utils.auth_interceptor import _ClientCallDetails
+
+
+class DefaultMetadataInterceptor(grpc.UnaryUnaryClientInterceptor, grpc.UnaryStreamClientInterceptor):
+    def _inject_default_metadata(self, call_details: grpc.ClientCallDetails):
+        metadata = [("accept", "application/grpc")]
+        if call_details.metadata:
+            metadata.extend(list(call_details.metadata))
+        new_details = _ClientCallDetails(
+            call_details.method,
+            call_details.timeout,
+            metadata,
+            call_details.credentials,
+        )
+        return new_details
+
+    def intercept_unary_unary(
+        self,
+        continuation: typing.Callable,
+        client_call_details: grpc.ClientCallDetails,
+        request: typing.Any,
+    ):
+        """
+        Intercepts unary calls and inject default metadata
+        """
+        updated_call_details = self._inject_default_metadata(client_call_details)
+        return continuation(updated_call_details, request)
+
+    def intercept_unary_stream(
+        self,
+        continuation: typing.Callable,
+        client_call_details: grpc.ClientCallDetails,
+        request: typing.Any,
+    ):
+        """
+        Handles a stream call and inject default metadata
+        """
+        updated_call_details = self._inject_default_metadata(client_call_details)
+        return continuation(updated_call_details, request)

flytekit/clis/sdk_in_container/serve.py

@@ -4,8 +4,6 @@
 from flyteidl.service.agent_pb2_grpc import add_AsyncAgentServiceServicer_to_server
 from grpc import aio
 
-from flytekit.extend.backend.agent_service import AsyncAgentService
-
 _serve_help = """Start a grpc server for the agent service."""
 
 
@@ -47,6 +45,8 @@ async def _start_grpc_server(port: int, worker: int, timeout: int):
     try:
         from prometheus_client import start_http_server
 
+        from flytekit.extend.backend.agent_service import AsyncAgentService
+
         start_http_server(9090)
     except ImportError as e:
         click.secho(f"Failed to start the prometheus server with error {e}", fg="red")

setup.py

@@ -38,10 +38,8 @@
         "deprecated>=1.0,<2.0",
         "docker>=4.0.0,<7.0.0",
         "python-dateutil>=2.1",
-        # Restrict grpcio and grpcio-status.  Version 1.50.0 pulls in a version of protobuf that is not compatible
-        # with the old protobuf library (as described in https://developers.google.com/protocol-buffers/docs/news/2022-05-06)
-        "grpcio>=1.50.0,!=1.55.0,<1.53.1,<2.0",
-        "grpcio-status>=1.50.0,!=1.55.0,<1.53.1,<2.0",
+        "grpcio",
+        "grpcio-status",
         "importlib-metadata",
         "fsspec>=2023.3.0",
         "adlfs",