gen ai test

Author: liustve

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/exporter/otlp/aws/common/constants.py

@@ -1,4 +1,4 @@
-BASE_LOG_BUFFER_BYTE_SIZE = 450000
+BASE_LOG_BUFFER_BYTE_SIZE = 2000
 MAX_LOG_REQUEST_BYTE_SIZE = (
     1048576  # https://docs.aws.amazon.com/AmazonCloudWatch/latest/monitoring/CloudWatch-OTLPEndpoint.html
 )

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/exporter/otlp/aws/logs/otlp_aws_logs_exporter.py

@@ -12,6 +12,7 @@
 
 from amazon.opentelemetry.distro.exporter.otlp.aws.common.aws_auth_session import AwsAuthSession
 from amazon.opentelemetry.distro.exporter.otlp.aws.common.constants import BASE_LOG_BUFFER_BYTE_SIZE
+from opentelemetry.exporter.otlp.proto.common._log_encoder import encode_logs
 from opentelemetry.exporter.otlp.proto.http import Compression
 from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter, _create_exp_backoff_generator
 from opentelemetry.sdk._logs import (
@@ -26,7 +27,7 @@
 
 class OTLPAwsLogExporter(OTLPLogExporter):
     _LARGE_LOG_HEADER = {"x-aws-log-semantics": "otel"}
-    _RETRY_AFTER_HEADER = "Retry-After"
+    _RETRY_AFTER_HEADER = "Retry-After"  # https://opentelemetry.io/docs/specs/otlp/#otlphttp-throttling
 
     def __init__(
         self,
@@ -55,31 +56,22 @@ def __init__(
             session=AwsAuthSession(aws_region=self._aws_region, service="logs"),
         )
 
+    # Code based off of:
+    # https://github.com/open-telemetry/opentelemetry-python/blob/main/exporter/opentelemetry-exporter-otlp-proto-http/src/opentelemetry/exporter/otlp/proto/http/_log_exporter/__init__.py#L167
     def export(self, batch: Sequence[LogData]) -> LogExportResult:
+
         print(f"Exporting batch of {len(batch)} logs")
         print("TOTAL DATA SIZE " + str(sum(self._get_size_of_log(logz) for logz in batch)))
         print("GEN_AI_FLAG " + str(self._gen_ai_flag))
 
-        return super().export(batch)
-
-    def set_gen_ai_flag(self):
-        self._gen_ai_flag = True
-
-    @staticmethod
-    def _retryable(resp: requests.Response) -> bool:
-        if resp.status_code == 429 or resp.status_code == 503:
-            return True
-
-        return OTLPLogExporter._retryable(resp)
-
-    def _export(self, serialized_data: bytes) -> requests.Response:
         """
-        Exports the given serialized OTLP log data. Behaviors of how this export will work.
+        Exports the given batch of OTLP log data.
+        Behaviors of how this export will work -
 
         1. Always compresses the serialized data into gzip before sending.
 
-        2. If self._gen_ai_flag is enabled, the log data is > 1 MB and we assume that the log contains normalized gen.ai attributes.
-            - in this case we inject the 'x-aws-log-semantics' flag into the header.
+        2. If self._gen_ai_flag is enabled, the log data is > 1 MB and the assumption is that the log is a normalized gen.ai LogEvent.
+            - inject the 'x-aws-log-semantics' flag into the header.
 
         3. Retry behavior is now the following:
             - if the response contains a status code that is retryable and the response contains Retry-After in its headers,
@@ -88,69 +80,109 @@ def _export(self, serialized_data: bytes) -> requests.Response:
             - if the response does not contain that Retry-After header, default back to the current iteration of the
               exponential backoff delay
         """
+
+        if self._shutdown:
+            _logger.warning("Exporter already shutdown, ignoring batch")
+            return LogExportResult.FAILURE
+
+        serialized_data = encode_logs(batch).SerializeToString()
+
         gzip_data = BytesIO()
         with gzip.GzipFile(fileobj=gzip_data, mode="w") as gzip_stream:
             gzip_stream.write(serialized_data)
 
         data = gzip_data.getvalue()
 
-        def send():
-            try:
-                return self._session.post(
-                    url=self._endpoint,
-                    headers=self._LARGE_LOG_HEADER if self._gen_ai_flag else None,
-                    data=data,
-                    verify=self._certificate_file,
-                    timeout=self._timeout,
-                    cert=self._client_cert,
-                )
-            except ConnectionError:
-                return self._session.post(
-                    url=self._endpoint,
-                    headers=self._LARGE_LOG_HEADER if self._gen_ai_flag else None,
-                    data=data,
-                    verify=self._certificate_file,
-                    timeout=self._timeout,
-                    cert=self._client_cert,
-                )
-
         backoff = _create_exp_backoff_generator(self._MAX_RETRY_TIMEOUT)
 
         while True:
-            resp = send()
+            resp = self._send(data)
+
+            if resp.ok:
+                return LogExportResult.SUCCESS
 
             if not self._retryable(resp):
+                _logger.error(
+                    "Failed to export logs batch code: %s, reason: %s",
+                    resp.status_code,
+                    resp.text,
+                )
                 self._gen_ai_flag = False
-                return resp
-
-            _logger.warning(
-                "Transient error %s encountered while exporting logs batch, retrying in %ss.",
-                resp.reason,
-                delay,
-            )
+                return LogExportResult.FAILURE
 
+            # https://opentelemetry.io/docs/specs/otlp/#otlphttp-throttling
             retry_after = resp.headers.get(self._RETRY_AFTER_HEADER, None)
 
             # Set the next retry delay to the value of the Retry-After response in the headers.
             # If Retry-After is not present in the headers, default to the next iteration of the exponential backoff strategy.
             delay = next(backoff, -1) if retry_after == None else self._parse_retryable_header(retry_after)
 
             if delay == -1:
+                _logger.error(
+                    "Transient error %s encountered while exporting logs batch. "
+                    "No Retry-After header found and all backoff retries exhausted. "
+                    "Logs will not be exported.",
+                    resp.reason,
+                )
                 self._gen_ai_flag = False
-                return resp
+                return LogExportResult.FAILURE
+
+            _logger.warning(
+                "Transient error %s encountered while exporting logs batch, retrying in %ss.",
+                resp.reason,
+                delay,
+            )
 
             sleep(delay)
-            continue
+    
+    def set_gen_ai_flag(self):
+        """
+        Sets the gen_ai flag to true to signal injecting the LLO flag to the headers of the export request.
+        """
+        self._gen_ai_flag = True
+
+    def _send(self, serialized_data: bytes):
+        try:
+            return self._session.post(
+                url=self._endpoint,
+                headers=self._LARGE_LOG_HEADER if self._gen_ai_flag else None,
+                data=serialized_data,
+                verify=self._certificate_file,
+                timeout=self._timeout,
+                cert=self._client_cert,
+            )
+        except ConnectionError:
+            return self._session.post(
+                url=self._endpoint,
+                headers=self._LARGE_LOG_HEADER if self._gen_ai_flag else None,
+                data=serialized_data,
+                verify=self._certificate_file,
+                timeout=self._timeout,
+                cert=self._client_cert,
+            )
+
+    @staticmethod
+    def _retryable(resp: requests.Response) -> bool:
+        """
+        Is it a retryable response?
+        """
+        if resp.status_code == 429 or resp.status_code == 503:
+            return True
+
+        return OTLPLogExporter._retryable(resp)
 
     def _parse_retryable_header(self, retry_header: str) -> float:
-        "Converts the given retryable header into a delay in seconds, returns -1 if there's an error with the parsing"
+        """
+        Converts the given retryable header into a delay in seconds, returns -1 if there's an error with the parsing
+        """
         try:
             return float(retry_header)
         except ValueError:
             return -1
 
     def _get_size_of_log(self, log_data: LogData):
-        # Rough estimate of the size of the LogData based on size of the content body + a buffer to account for other information in logs.
+        # Rough estimate of the size of the LogData based on size of
+        # the content body + a buffer to account for other information in logs.
         size = BASE_LOG_BUFFER_BYTE_SIZE
         body = log_data.log_record.body
 

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/exporter/otlp/aws/traces/otlp_aws_span_exporter.py

@@ -5,7 +5,6 @@
 
 from amazon.opentelemetry.distro._utils import is_agent_observability_enabled
 from amazon.opentelemetry.distro.exporter.otlp.aws.common.aws_auth_session import AwsAuthSession
-from amazon.opentelemetry.distro.exporter.otlp.aws.logs.otlp_aws_logs_exporter import OTLPAwsLogExporter
 from amazon.opentelemetry.distro.llo_handler import LLOHandler
 from opentelemetry.exporter.otlp.proto.http import Compression
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter