Contextual Request ID for flyte requests (#531)

Signed-off-by: Ketan Umare <kumare3@users.noreply.github.com>
Co-authored-by: Ketan Umare <kumare3@users.noreply.github.com>

Author: kumare3

src/flyte/_internal/controllers/remote/_core.py

@@ -382,7 +382,7 @@ async def _bg_launch(self, action: Action):
                             e.code().name, f"Precondition failed: {e.details()}"
                         ) from e
                     # For all other errors, we will retry with backoff
-                    logger.exception(
+                    logger.error(
                         f"Failed to launch action: {action.name}, Code: {e.code()}, "
                         f"Details {e.details()} backing off..."
                     )

src/flyte/_internal/controllers/remote/_informer.py

@@ -248,11 +248,11 @@ async def watch(self):
                 logger.info(f"Watch cancelled: {self.name}")
                 return
             except asyncio.TimeoutError as e:
-                logger.error(f"Watch timeout: {self.name}", exc_info=e)
+                logger.error(f"Watch timeout: {self.name}, {e}", exc_info=False)
                 last_exc = e
                 retries += 1
             except grpc.aio.AioRpcError as e:
-                logger.exception(f"RPC error: {self.name}", exc_info=e)
+                logger.error(f"RPC error: {self.name}, {e}", exc_info=False)
                 last_exc = e
                 retries += 1
             except Exception as e:

src/flyte/models.py

@@ -77,6 +77,20 @@ def new_sub_action_from(self, task_call_seq: int, task_hash: str, input_hash: st
         new_name = base36_encode(bytes_digest)
         return self.new_sub_action(new_name)
 
+    def unique_id_str(self, salt: str | None = None) -> str:
+        """
+        Generate a unique ID string for this action in the format:
+        {project}-{domain}-{run_name}-{action_name}
+
+        This is optimized for performance assuming all fields are available.
+
+        :return: A unique ID string
+        """
+        v = f"{self.project}-{self.domain}-{self.run_name}-{self.name}"
+        if salt is not None:
+            return f"{v}-{salt}"
+        return v
+
 
 @rich.repr.auto
 @dataclass

src/flyte/remote/_client/auth/_channel.py

@@ -159,7 +159,7 @@ async def create_channel(
         DefaultMetadataUnaryUnaryInterceptor,
     )
 
-    # Add all types of default metadata interceptors
+    # Add all types of default metadata interceptors (includes x-request-id)
     interceptors: typing.List[grpc.aio.ClientInterceptor] = [
         DefaultMetadataUnaryUnaryInterceptor(),
         DefaultMetadataUnaryStreamInterceptor(),

src/flyte/remote/_client/auth/_grpc_utils/default_metadata_interceptor.py

@@ -1,9 +1,36 @@
+import random
+import string
 import typing
+from uuid import uuid4
 
 import grpc.aio
 from grpc.aio import ClientCallDetails, Metadata
 
+import flyte
+
 _default_metadata = Metadata(("accept", "application/grpc"))
+ALPHABET = string.ascii_lowercase + string.digits
+
+
+def fast_short_id():
+    return "".join(random.choices(ALPHABET, k=4))
+
+
+def _generate_request_id() -> str:
+    """
+    Generate a request ID based on the current Flyte context.
+
+    If running within a Flyte task context, creates a request ID using the action's unique_id_str method.
+    Otherwise, falls back to a UUID4.
+
+    :return: The generated request ID string
+    """
+    ctx = flyte.ctx()
+    if ctx and ctx.action:
+        return ctx.action.unique_id_str(salt=fast_short_id())
+
+    # Fall back to UUID4 if context is not available
+    return str(uuid4())
 
 
 def with_metadata(call_details: ClientCallDetails, new_metadata: Metadata) -> ClientCallDetails:
@@ -27,20 +54,27 @@ def with_metadata(call_details: ClientCallDetails, new_metadata: Metadata) -> Cl
 class _BaseDefaultMetadataInterceptor:
     """
     Base class for all default metadata interceptors that provides common functionality.
+    Injects both default metadata (accept header) and x-request-id.
     """
 
     async def _inject_default_metadata(self, call_details: grpc.aio.ClientCallDetails):
         """
-        Injects default metadata into the client call details.
+        Injects default metadata and request ID into the client call details.
 
-        This method adds all key-value pairs from the default metadata dictionary to the
-        client call details metadata. If the client call details don't have metadata,
-        a new Metadata object is created.
+        This method adds:
+        - Default metadata (accept: application/grpc)
+        - x-request-id header with context-based or UUID4 value
 
         :param call_details: The client call details to inject metadata into
         :return: A new ClientCallDetails object with the injected metadata
         """
-        return with_metadata(call_details, _default_metadata)
+        # Generate request ID and combine with default metadata
+        request_id = _generate_request_id()
+        combined_metadata = Metadata(
+            ("accept", "application/grpc"),
+            ("x-request-id", request_id),
+        )
+        return with_metadata(call_details, combined_metadata)
 
 
 class DefaultMetadataUnaryUnaryInterceptor(_BaseDefaultMetadataInterceptor, grpc.aio.UnaryUnaryClientInterceptor):