telemetry_rust/middleware/aws/instrumentation/fluent_builder/
sagemaker_runtime.rs

1/// AWS SageMaker Runtime fluent builder instrumentation implementations
2use super::{utils::*, *};
3
4// InvokeEndpoint
5impl<'a> AwsBuilderInstrument<'a> for InvokeEndpointFluentBuilder {
6    fn build_aws_span(&self) -> AwsSpanBuilder<'a> {
7        let endpoint_name = self.get_endpoint_name().clone().unwrap_or_default();
8        let attributes = attributes![
9            self.get_target_model()
10                .as_attribute("aws.sagemaker.target_model"),
11            self.get_target_variant()
12                .as_attribute("aws.sagemaker.target_variant"),
13            self.get_inference_component_name()
14                .as_attribute("aws.sagemaker.inference_component_name"),
15        ];
16        SageMakerRuntimeSpanBuilder::invoke_endpoint(endpoint_name).attributes(attributes)
17    }
18}
19impl InstrumentedFluentBuilderOutput for InvokeEndpointOutput {
20    fn extract_attributes(&self) -> impl IntoIterator<Item = KeyValue> {
21        attributes![
22            self.invoked_production_variant()
23                .as_attribute("aws.sagemaker.invoked_production_variant"),
24        ]
25    }
26}
27instrument_aws_operation!(aws_sdk_sagemakerruntime::operation::invoke_endpoint);
28
29// InvokeEndpointAsync
30impl<'a> AwsBuilderInstrument<'a> for InvokeEndpointAsyncFluentBuilder {
31    fn build_aws_span(&self) -> AwsSpanBuilder<'a> {
32        let endpoint_name = self.get_endpoint_name().clone().unwrap_or_default();
33        let attributes = attributes![
34            self.get_inference_id()
35                .as_attribute("aws.sagemaker.inference_id"),
36            self.get_input_location()
37                .as_attribute("aws.sagemaker.input_location"),
38        ];
39        SageMakerRuntimeSpanBuilder::invoke_endpoint_async(endpoint_name)
40            .attributes(attributes)
41    }
42}
43impl InstrumentedFluentBuilderOutput for InvokeEndpointAsyncOutput {
44    fn extract_attributes(&self) -> impl IntoIterator<Item = KeyValue> {
45        attributes![
46            self.inference_id()
47                .as_attribute("aws.sagemaker.inference_id"),
48        ]
49    }
50}
51instrument_aws_operation!(aws_sdk_sagemakerruntime::operation::invoke_endpoint_async);
52
53// InvokeEndpointWithResponseStream
54impl<'a> AwsBuilderInstrument<'a> for InvokeEndpointWithResponseStreamFluentBuilder {
55    fn build_aws_span(&self) -> AwsSpanBuilder<'a> {
56        let endpoint_name = self.get_endpoint_name().clone().unwrap_or_default();
57        let attributes = attributes![
58            self.get_target_variant()
59                .as_attribute("aws.sagemaker.target_variant"),
60            self.get_inference_component_name()
61                .as_attribute("aws.sagemaker.inference_component_name"),
62        ];
63        SageMakerRuntimeSpanBuilder::invoke_endpoint_with_response_stream(endpoint_name)
64            .attributes(attributes)
65    }
66}
67impl InstrumentedFluentBuilderOutput for InvokeEndpointWithResponseStreamOutput {
68    fn extract_attributes(&self) -> impl IntoIterator<Item = KeyValue> {
69        attributes![
70            self.invoked_production_variant()
71                .as_attribute("aws.sagemaker.invoked_production_variant"),
72        ]
73    }
74}
75instrument_aws_operation!(
76    aws_sdk_sagemakerruntime::operation::invoke_endpoint_with_response_stream
77);
telemetry_rust/middleware/aws/instrumentation/fluent_builder/sagemaker_runtime.rs

telemetry_rust/middleware/aws/instrumentation/fluent_builder/
sagemaker_runtime.rs
