Skip to content

Commit

Permalink
use onStartTracing and onEndTracing
Browse files Browse the repository at this point in the history
  • Loading branch information
@howieleung
howieleung committed Sep 6, 2024
1 parent a142bb7 commit 3aee693
Show file tree
Hide file tree
Showing 5 changed files with 101 additions and 103 deletions.
51 changes: 24 additions & 27 deletions sdk/ai/ai-inference-rest/src/trace.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
// Licensed under the MIT license.

import { PathUncheckedResponse, RequestParameters, StreamableMethod } from "@azure-rest/core-client";
import { createTracingClient, OperationTracingOptions, SpanStatus } from "@azure/core-tracing";
import { createTracingClient, OperationTracingOptions, SpanStatus, TracingSpan } from "@azure/core-tracing";
import { GetChatCompletionsBodyParam, GetEmbeddingsBodyParam, GetImageEmbeddingsBodyParam } from "./parameters.js";

const traceClient = createTracingClient({ namespace: "Micirsoft.CognitiveServices", packageName: "ai-inference-rest", packageVersion: "1.0.0" });
Expand Down Expand Up @@ -56,52 +56,49 @@ export function traceInference(
}
}

const requestAttributeMapper = (request: RequestParameterWithBodyType) => {
const map = new Map<string, unknown>();
function onStartTracing(span: TracingSpan, request: RequestParameters) {
const body = (request as RequestParameterWithBodyType).body;
const urlObj = new URL(url);
const port = urlObj.port || (urlObj.protocol === 'https:' ? 443 : 80);

map.set(TracingAttributesEnum.Server_Address, urlObj.hostname);
map.set(TracingAttributesEnum.Server_Port, port);
map.set(TracingAttributesEnum.Operation_Name, getOperationName(routePath));
map.set(TracingAttributesEnum.System, "az.ai_inference");
map.set(TracingAttributesEnum.Request_Model, model);
map.set(TracingAttributesEnum.Request_Frequency_Penalty, body?.frequency_penalty);
map.set(TracingAttributesEnum.Request_Max_Tokens, body?.max_tokens);
map.set(TracingAttributesEnum.Request_Presence_Penalty, body?.presence_penalty);
map.set(TracingAttributesEnum.Request_Stop_Sequences, body?.stop);
map.set(TracingAttributesEnum.Request_Temperature, body?.temperature);
map.set(TracingAttributesEnum.Request_Top_P, body?.top_p);
return map;
span.setAttribute(TracingAttributesEnum.Server_Address, urlObj.hostname);
span.setAttribute(TracingAttributesEnum.Server_Port, port);
span.setAttribute(TracingAttributesEnum.Operation_Name, getOperationName(routePath));
span.setAttribute(TracingAttributesEnum.System, "az.ai_inference");
span.setAttribute(TracingAttributesEnum.Request_Model, model);
span.setAttribute(TracingAttributesEnum.Request_Frequency_Penalty, body?.frequency_penalty);
span.setAttribute(TracingAttributesEnum.Request_Max_Tokens, body?.max_tokens);
span.setAttribute(TracingAttributesEnum.Request_Presence_Penalty, body?.presence_penalty);
span.setAttribute(TracingAttributesEnum.Request_Stop_Sequences, body?.stop);
span.setAttribute(TracingAttributesEnum.Request_Temperature, body?.temperature);
span.setAttribute(TracingAttributesEnum.Request_Top_P, body?.top_p);
}

const responseAttributeMapper: (args: RequestParameters, rt?: PathUncheckedResponse, error?: unknown) => Map<string, unknown> | [Map<string, unknown>, SpanStatus] = (_: RequestParameters, response?: PathUncheckedResponse, error?: unknown) => {
const map = new Map<string, unknown>();
let status: SpanStatus | undefined;
function onEndTracing(span: TracingSpan, _request: RequestParameters, response?: PathUncheckedResponse, error?: unknown) {
let status: SpanStatus = { status: "unset" };
if (error) {
if (error instanceof Error) {
map.set(TracingAttributesEnum.Error_Type, error.message);
span.setAttribute(TracingAttributesEnum.Error_Type, error.message);
status = { status: "error", error: error.message };
} else {
map.set(TracingAttributesEnum.Error_Type, error);
span.setAttribute(TracingAttributesEnum.Error_Type, error);
status = { status: "error", error: error.toString() };
}

}
if (response) {
let body = response.body;
map.set(TracingAttributesEnum.Response_Id, body.id);
map.set(TracingAttributesEnum.Response_Model, body.model);
span.setAttribute(TracingAttributesEnum.Response_Id, body.id);
span.setAttribute(TracingAttributesEnum.Response_Model, body.model);
if (body.usage) {
map.set(TracingAttributesEnum.Usage_Input_Tokens, body.usage.prompt_tokens);
map.set(TracingAttributesEnum.Usage_Output_Tokens, body.usage.completion_tokens);
span.setAttribute(TracingAttributesEnum.Usage_Input_Tokens, body.usage.prompt_tokens);
span.setAttribute(TracingAttributesEnum.Usage_Output_Tokens, body.usage.completion_tokens);
}
if (body.error) {
map.set(TracingAttributesEnum.Error_Type, body.error.code);
span.setAttribute(TracingAttributesEnum.Error_Type, body.error.code);
}
}
return status ? [map, status] : map;
span.setStatus(status);
}

const request = args as RequestParameterWithBodyType;
Expand All @@ -111,5 +108,5 @@ export function traceInference(
return methodToTrace();
}
const name = `${getOperationName(routePath)} ${model ?? ""}`.trim();
return traceClient.traceAsync(name, request, methodToTrace, requestAttributeMapper, responseAttributeMapper, options);
return traceClient.traceAsync(name, request, methodToTrace, onStartTracing, onEndTracing, options);
}
9 changes: 6 additions & 3 deletions sdk/core/core-tracing/review/core-tracing.api.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,8 +44,10 @@ export type Resolved<T> = T extends {
then(onfulfilled: infer F): any;
} ? F extends (value: infer V) => any ? Resolved<V> : never : T;

// Warning: (ae-forgotten-export) The symbol "SpanStatusUnset" needs to be exported by the entry point index.d.ts
//
// @public
export type SpanStatus = SpanStatusSuccess | SpanStatusError;
export type SpanStatus = SpanStatusSuccess | SpanStatusError | SpanStatusUnset;

// @public
export type SpanStatusError = {
Expand All @@ -68,8 +70,8 @@ export interface TracingClient {
span: TracingSpan;
updatedOptions: OptionsWithTracingContext<Options>;
};
trace<Arguments, Return>(name: string, args: Arguments, methodToTrace: () => Return, paramAttributeMapper?: (args: Arguments) => Map<string, unknown>, returnAttributeMapper?: (args: Arguments, rt?: Return, error?: unknown) => Map<string, unknown> | [Map<string, unknown>, SpanStatus], options?: OperationTracingOptions): Return;
traceAsync<Arguments, ResolvedReturn, PromiseReturn extends Promise<ResolvedReturn> | PromiseLike<ResolvedReturn>>(name: string, args: Arguments, methodToTrace: () => PromiseReturn, paramAttributeMapper?: (args: Arguments) => Map<string, unknown>, returnAttributeMapper?: (args: Arguments, rt?: ResolvedReturn, error?: unknown) => Map<string, unknown> | [Map<string, unknown>, SpanStatus], options?: OperationTracingOptions): PromiseReturn;
trace<Arguments, Return>(name: string, args: Arguments, methodToTrace: () => Return, onStartTracing?: (span: TracingSpan, args: Arguments) => void, onEndTracing?: (span: TracingSpan, args: Arguments, rt?: Return, error?: unknown) => void, options?: OperationTracingOptions): Return;
traceAsync<Arguments, ResolvedReturn, PromiseReturn extends Promise<ResolvedReturn> | PromiseLike<ResolvedReturn>>(name: string, args: Arguments, methodToTrace: () => PromiseReturn, onStartTracing?: (span: TracingSpan, args: Arguments) => void, onEndTracing?: (span: TracingSpan, args: Arguments, rt?: ResolvedReturn, error?: unknown) => void, options?: OperationTracingOptions): PromiseReturn;
withContext<CallbackArgs extends unknown[], Callback extends (...args: CallbackArgs) => ReturnType<Callback>>(context: TracingContext, callback: Callback, ...callbackArgs: CallbackArgs): ReturnType<Callback>;
withSpan<Options extends {
tracingOptions?: OperationTracingOptions;
Expand All @@ -92,6 +94,7 @@ export interface TracingContext {

// @public
export interface TracingSpan {
addEvent(name: string, attributesOrStartTime?: unknown, startTime?: unknown): void;
end(): void;
isRecording(): boolean;
recordException(exception: Error | string): void;
Expand Down
3 changes: 3 additions & 0 deletions sdk/core/core-tracing/src/instrumenter.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,9 @@ export function createDefaultTracingSpan(): TracingSpan {
setStatus: () => {
// noop
},
addEvent: () => {
// noop
},
};
}

Expand Down
41 changes: 28 additions & 13 deletions sdk/core/core-tracing/src/interfaces.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -102,20 +102,20 @@ export interface TracingClient {
createRequestHeaders(tracingContext?: TracingContext): Record<string, string>;

/**
* Capture the arguments and return of a function and create a span.
* @param name - name of the span.
* @param args - arguments of the method to be traced. Generally, you should pass in `arguments` reserve word.
* @param methodToTrace - function pointer of the implementation.
* @param paramAttributeMapper - mapping function to map the arguments to span's attributes.
* @param returnAttributeMapper - mapping function to map the return object to span's attributes.
* @returns - return back the return from methodToTrace.
*/
* This method will create a span, call the methodToTrace, and end the span.
* @param name - name of the span.
* @param args - arguments of the method to be traced. Generally, you should pass in `arguments` reserve word.
* @param methodToTrace - function pointer of the implementation.
* @param onStartTracing - callback function to set attributes and events before calling methodTotrace.
* @param onEndTracing - callback function to set attributes, events, and status before ending the span.
* @returns - return back the return from methodToTrace.
*/
trace<Arguments, Return>(
name: string,
args: Arguments,
methodToTrace: () => Return,
paramAttributeMapper?: (args: Arguments) => Map<string, unknown>,
returnAttributeMapper?: (args: Arguments, rt?: Return, error?: unknown) => Map<string, unknown> | [Map<string, unknown>, SpanStatus],
onStartTracing?: (span: TracingSpan, args: Arguments) => void,
onEndTracing?: (span: TracingSpan, args: Arguments, rt?: Return, error?: unknown) => void,
options?: OperationTracingOptions): Return;

/**
Expand All @@ -131,8 +131,8 @@ export interface TracingClient {
name: string,
args: Arguments,
methodToTrace: () => PromiseReturn,
paramAttributeMapper?: (args: Arguments) => Map<string, unknown>,
returnAttributeMapper?: (args: Arguments, rt?: ResolvedReturn, error?: unknown) => Map<string, unknown> | [Map<string, unknown>, SpanStatus],
onStartTracing?: (span: TracingSpan, args: Arguments) => void,
onEndTracing?: (span: TracingSpan, args: Arguments, rt?: ResolvedReturn, error?: unknown) => void,
options?: OperationTracingOptions): PromiseReturn;
}

Expand Down Expand Up @@ -225,6 +225,11 @@ export interface InstrumenterSpanOptions extends TracingSpanOptions {
tracingContext?: TracingContext;
}

/**
* Status representing an unknow operation status that can be sent to {@link TracingSpan.setStatus}
*/
export type SpanStatusUnset = { status: "unset" };

/**
* Status representing a successful operation that can be sent to {@link TracingSpan.setStatus}
*/
Expand All @@ -240,7 +245,7 @@ export type SpanStatusError = { status: "error"; error?: Error | string };
*
* By default, all spans will be created with status "unset".
*/
export type SpanStatus = SpanStatusSuccess | SpanStatusError;
export type SpanStatus = SpanStatusSuccess | SpanStatusError | SpanStatusUnset;

/**
* Represents an implementation agnostic tracing span.
Expand Down Expand Up @@ -282,6 +287,16 @@ export interface TracingSpan {
* Depending on the span implementation, this may return false if the span is not being sampled.
*/
isRecording(): boolean;

/**
* Adds an event to the Span.
*
* @param name the name of the event.
* @param [attributesOrStartTime] the attributes that will be added; these are
* associated with this event. Can be also a start time
* @param [startTime] start time of the event.
*/
addEvent(name: string, attributesOrStartTime?: unknown, startTime?: unknown): void;
}

/** An immutable context bag of tracing values for the current operation. */
Expand Down
100 changes: 40 additions & 60 deletions sdk/core/core-tracing/src/tracingClient.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ import {
OperationTracingOptions,
OptionsWithTracingContext,
Resolved,
SpanStatus,
TracingClient,
TracingClientOptions,
TracingContext,
Expand Down Expand Up @@ -117,94 +116,83 @@ export function createTracingClient(options: TracingClientOptions): TracingClien
}

/**
* Capture the arguments and return of a function and create a span.
* @param name - name of the span.
* @param args - arguments of the method to be traced. Generally, you should pass in `arguments` reserve word.
* @param methodToTrace - function pointer of the implementation.
* @param paramAttributeMapper - mapping function to map the arguments to span's attributes.
* @param returnAttributeMapper - mapping function to map the return object to span's attributes.
* @returns - return back the return from methodToTrace.
*/
* This method will create a span, call the methodToTrace, and end the span.
* @param name - name of the span.
* @param args - arguments of the method to be traced. Generally, you should pass in `arguments` reserve word.
* @param methodToTrace - function pointer of the implementation.
* @param onStartTracing - callback function to set attributes and events before calling methodTotrace.
* @param onEndTracing - callback function to set attributes, events, and status before ending the span.
* @returns - return back the return from methodToTrace.
*/
function trace<Arguments, Return>(
name: string,
args: Arguments,
methodToTrace: () => Return,
paramAttributeMapper?: (args: Arguments) => Map<string, unknown>,
returnAttributeMapper?: (args: Arguments, rt?: Return, error?: unknown) => Map<string, unknown> & [Map<string, unknown>, SpanStatus],
tracingOptions?: OperationTracingOptions): Return {
onStartTracing?: (span: TracingSpan, args: Arguments) => void,
onEndTracing?: (span: TracingSpan, args: Arguments, rt?: Return, error?: unknown) => void,
options?: OperationTracingOptions): Return {

let spanAttributesInObject = {};
const { span, tracingContext } = tryCreateSpan(name, {}, options) ?? {};

if (paramAttributeMapper) {
const spanAttributesInMap = paramAttributeMapper(args);
spanAttributesInObject = mapToObject(spanAttributesInMap);
}
const { span, tracingContext } = tryCreateSpan(name, spanAttributesInObject, tracingOptions) ?? {};

if (!span || !tracingContext) {
return methodToTrace();
}

if (onStartTracing) {
onStartTracing(span, args);
}

try {
const returnObj = withContext(tracingContext, methodToTrace)
tryProcessReturn(span, args, returnObj, returnAttributeMapper);
tryProcessReturn(span, args, returnObj, onEndTracing);

return returnObj;
} catch (err: any) {
tryProcessReturn(span, args, undefined, err, returnAttributeMapper);
tryProcessReturn(span, args, undefined, err, onEndTracing);
throw err;
}
}

/**
* Capture the arguments and return of a function and create a span.
* @param name - name of the span.
* @param args - arguments of the method to be traced. Generally, you should pass in `arguments` reserve word.
* @param methodToTrace - function pointer of the implementation.
* @param paramAttributeMapper - mapping function to map the arguments to span's attributes.
* @param returnAttributeMapper - mapping function to map the return object to span's attributes.
* @returns - return back the return from methodToTrace.
*/
* This method will create a span, call the methodToTrace, and end the span.
* @param name - name of the span.
* @param args - arguments of the method to be traced. Generally, you should pass in `arguments` reserve word.
* @param methodToTrace - function pointer of the implementation.
* @param onStartTracing - callback function to set attributes and events before calling methodTotrace.
* @param onEndTracing - callback function to set attributes, events, and status before ending the span.
* @returns - return back the return from methodToTrace.
*/
function traceAsync<Arguments, ResolvedReturn, PromiseReturn extends Promise<ResolvedReturn> | PromiseLike<ResolvedReturn>>(
name: string,
args: Arguments,
methodToTrace: () => PromiseReturn,
paramAttributeMapper?: (args: Arguments) => Map<string, unknown>,
returnAttributeMapper?: (args: Arguments, rt?: ResolvedReturn, error?: unknown) => Map<string, unknown> & [Map<string, unknown>, SpanStatus],
tracingOptions?: OperationTracingOptions): PromiseReturn {

let spanAttributesInObject = {};
onStartTracing?: (span: TracingSpan, args: Arguments) => void,
onEndTracing?: (span: TracingSpan, args: Arguments, rt?: ResolvedReturn, error?: unknown) => void,
options?: OperationTracingOptions): PromiseReturn {

if (paramAttributeMapper) {
const spanAttributesInMap = paramAttributeMapper(args);
spanAttributesInObject = mapToObject(spanAttributesInMap);
}
const { span, tracingContext } = tryCreateSpan(name, spanAttributesInObject, tracingOptions) ?? {};
const { span, tracingContext } = tryCreateSpan(name, {}, options) ?? {};

if (!span || !tracingContext) {
return methodToTrace();
}

if (onStartTracing) {
onStartTracing(span, args);
}

try {
return withContext(tracingContext, methodToTrace).
then((response) => {
tryProcessReturn(span, args, response, undefined, returnAttributeMapper);
tryProcessReturn(span, args, response, undefined, onEndTracing);
return response;
}) as PromiseReturn;
} catch (err) {
tryProcessReturn(span, args, undefined, err, returnAttributeMapper);
tryProcessReturn(span, args, undefined, err, onEndTracing);
throw err;
}
}

function mapToObject(map: Map<string, any>): any {
const object: any = {};
map.forEach((value, key) => {
object[key] = value;
});
return object;
}

function tryCreateSpan(
spanName: string,
spanAttributes: Record<string, unknown>,
Expand Down Expand Up @@ -238,21 +226,13 @@ export function createTracingClient(options: TracingClientOptions): TracingClien
function tryProcessReturn<Arguments, Return>(
span: TracingSpan,
args: Arguments,
returnObj?: Return,
rt?: Return,
error?: unknown,
returnAttributeMapper?: (args: Arguments, rt?: Return, error?: unknown) => Map<string, unknown> & [Map<string, unknown>, SpanStatus]): void {
onEndTracing?: (span: TracingSpan, args: Arguments, rt?: Return, error?: unknown) => void) {
try {
if (returnAttributeMapper) {
const rt = returnAttributeMapper(args, returnObj, error);
const returnAttributes = Array.isArray(rt) ? rt[0] : rt;
returnAttributes.forEach((value, key) => {
span.setAttribute(key, value);
});
if (Array.isArray(rt)) {
span.setStatus(rt[1]);
}
if (onEndTracing) {
onEndTracing(span, args, rt, error);
}

span.end();
} catch (e: any) {
logger.warning(`Skipping tracing span processing due to an error: ${getErrorMessage(e)}`);
Expand Down

0 comments on commit 3aee693

Please sign in to comment.