def _capture_telemetry()

in src/sagemaker/serve/utils/telemetry_logger.py [0:0]

• 84 lines of code
• 23 McCabe index (conditional complexity)

def _capture_telemetry(func_name: str):
    """Placeholder docstring"""

    def decorator(func):
        def wrapper(self, *args, **kwargs):
            # Call the original function
            logger.info(TELEMETRY_OPT_OUT_MESSAGING)
            response = None
            caught_ex = None
            status = "1"
            failure_reason = None
            failure_type = None
            extra = f"{func_name}"

            start_timer = perf_counter()
            try:
                response = func(self, *args, **kwargs)
            except (
                ModelBuilderException,
                exceptions.CapacityError,
                exceptions.UnexpectedStatusException,
                exceptions.AsyncInferenceError,
            ) as e:
                status = "0"
                caught_ex = e
                failure_reason = str(e)
                failure_type = e.__class__.__name__
            except Exception as e:  # pylint: disable=W0703
                raise e

            stop_timer = perf_counter()
            elapsed = stop_timer - start_timer

            if self.model_server:
                extra += f"&x-modelServer={MODEL_SERVER_TO_CODE[str(self.model_server)]}"

            if self.image_uri:
                image_uri_option = _get_image_uri_option(
                    self.image_uri, getattr(self, "_is_custom_image_uri", False)
                )
                split_image_uri = self.image_uri.split("/")
                if len(split_image_uri) > 1:
                    extra += f"&x-imageTag={split_image_uri[1]}"

            extra += f"&x-sdkVersion={SDK_VERSION}"

            if self.image_uri:
                extra += f"&x-defaultImageUsage={image_uri_option}"

            if self.model_server == ModelServer.DJL_SERVING or self.model_server == ModelServer.TGI:
                extra += f"&x-modelName={self.model}"

            if self.sagemaker_session and self.sagemaker_session.endpoint_arn:
                extra += f"&x-endpointArn={self.sagemaker_session.endpoint_arn}"

            if getattr(self, "_is_mlflow_model", False):
                mlflow_model_path = self.model_metadata[MLFLOW_MODEL_PATH]
                mlflow_model_path_type = _get_mlflow_model_path_type(mlflow_model_path)
                extra += f"&x-mlflowModelPathType={MLFLOW_MODEL_PATH_CODE[mlflow_model_path_type]}"
                mlflow_model_tracking_server_arn = self.model_metadata.get(MLFLOW_TRACKING_ARN)
                if mlflow_model_tracking_server_arn is not None:
                    extra += f"&x-mlflowTrackingServerArn={mlflow_model_tracking_server_arn}"

            if getattr(self, "model_hub", False):
                extra += f"&x-modelHub={MODEL_HUB_TO_CODE[str(self.model_hub)]}"

            if getattr(self, "is_fine_tuned", False):
                extra += "&x-fineTuned=1"

            if getattr(self, "is_compiled", False):
                extra += "&x-compiled=1"
            if getattr(self, "is_quantized", False):
                extra += "&x-quantized=1"
            if getattr(self, "speculative_decoding_draft_model_source", False):
                model_provider_enum = (
                    SpeculativeDecodingDraftModelSource.SAGEMAKER
                    if self.speculative_decoding_draft_model_source == "sagemaker"
                    else SpeculativeDecodingDraftModelSource.CUSTOM
                )
                model_provider_value = SD_DRAFT_MODEL_SOURCE_TO_CODE[str(model_provider_enum)]
                extra += f"&x-sdDraftModelSource={model_provider_value}"

            if getattr(self, "deployment_config_name", False):
                config_name_code = self.deployment_config_name.lower()
                extra += f"&x-configName={config_name_code}"

            extra += f"&x-latency={round(elapsed, 2)}"

            if hasattr(self, "serve_settings") and not self.serve_settings.telemetry_opt_out:
                _send_telemetry(
                    status,
                    MODE_TO_CODE[str(self.mode)],
                    self.sagemaker_session,
                    failure_reason,
                    failure_type,
                    extra,
                )

            if caught_ex:
                raise caught_ex

            return response

        return wrapper

    return decorator