From 2ca96fdcd2a5e7dcacba8c432408cf94d5fa19d2 Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Mon, 4 Nov 2024 17:13:04 -0500 Subject: [PATCH 1/3] fix: allow passing promise of TelemetryConfig to enableTelemetry --- js/core/src/tracing.ts | 16 ++++++++++++---- js/core/src/tracing/instrumentation.ts | 4 ++-- js/plugins/google-cloud/src/index.ts | 12 +++++++----- 3 files changed, 21 insertions(+), 11 deletions(-) diff --git a/js/core/src/tracing.ts b/js/core/src/tracing.ts index dab2ffdf48..5670c35b04 100644 --- a/js/core/src/tracing.ts +++ b/js/core/src/tracing.ts @@ -35,18 +35,26 @@ let nodeOtelConfig: TelemetryConfig | null = null; const instrumentationKey = '__GENKIT_TELEMETRY_INSTRUMENTED'; -export function ensureBasicTelemetryInstrumentation() { +export async function ensureBasicTelemetryInstrumentation() { if (global[instrumentationKey]) { return; } - enableTelemetry({}); - global[instrumentationKey] = true; + await enableTelemetry({}); } /** * Enables tracing and metrics open telemetry configuration. */ -export function enableTelemetry(telemetryConfig: TelemetryConfig) { +export async function enableTelemetry( + telemetryConfig: TelemetryConfig | Promise +) { + global[instrumentationKey] = true; + + telemetryConfig = + telemetryConfig instanceof Promise + ? await telemetryConfig + : telemetryConfig; + nodeOtelConfig = telemetryConfig || {}; const processors: SpanProcessor[] = [createTelemetryServerProcessor()]; diff --git a/js/core/src/tracing/instrumentation.ts b/js/core/src/tracing/instrumentation.ts index f11cd675ac..84f063f701 100644 --- a/js/core/src/tracing/instrumentation.ts +++ b/js/core/src/tracing/instrumentation.ts @@ -45,7 +45,7 @@ export async function newTrace( }, fn: (metadata: SpanMetadata, rootSpan: ApiSpan) => Promise ) { - ensureBasicTelemetryInstrumentation(); + await ensureBasicTelemetryInstrumentation(); const traceMetadata: TraceMetadata = traceMetadataAls.getStore() || { paths: new Set(), timestamp: performance.now(), @@ -78,7 +78,7 @@ export async function runInNewSpan( }, fn: (metadata: SpanMetadata, otSpan: ApiSpan, isRoot: boolean) => Promise ): Promise { - ensureBasicTelemetryInstrumentation(); + await ensureBasicTelemetryInstrumentation(); const tracer = trace.getTracer(TRACER_NAME, TRACER_VERSION); const parentStep = spanMetadataAls.getStore(); diff --git a/js/plugins/google-cloud/src/index.ts b/js/plugins/google-cloud/src/index.ts index 8112d6ff45..4c67643c20 100644 --- a/js/plugins/google-cloud/src/index.ts +++ b/js/plugins/google-cloud/src/index.ts @@ -23,13 +23,15 @@ import { GcpOpenTelemetry } from './gcpOpenTelemetry.js'; import { TelemetryConfigs } from './telemetry/defaults.js'; import { GcpTelemetryConfig, GcpTelemetryConfigOptions } from './types.js'; -export async function enableGoogleCloudTelemetry( +export function enableGoogleCloudTelemetry( options?: GcpTelemetryConfigOptions ) { - const pluginConfig = await configureGcpPlugin(options); - - enableTelemetry(await new GcpOpenTelemetry(pluginConfig).getConfig()); - logger.init(await new GcpLogger(pluginConfig).getLogger(getCurrentEnv())); + return enableTelemetry( + configureGcpPlugin(options).then(async (pluginConfig) => { + logger.init(await new GcpLogger(pluginConfig).getLogger(getCurrentEnv())); + return new GcpOpenTelemetry(pluginConfig).getConfig(); + }) + ); } /** From e45ecc53ed7704fedcf9a732d289fe7e7de20165 Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Mon, 4 Nov 2024 21:12:16 -0500 Subject: [PATCH 2/3] await --- js/core/src/tracing.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/js/core/src/tracing.ts b/js/core/src/tracing.ts index 5670c35b04..5b981bbe12 100644 --- a/js/core/src/tracing.ts +++ b/js/core/src/tracing.ts @@ -37,7 +37,7 @@ const instrumentationKey = '__GENKIT_TELEMETRY_INSTRUMENTED'; export async function ensureBasicTelemetryInstrumentation() { if (global[instrumentationKey]) { - return; + return await global[instrumentationKey]; } await enableTelemetry({}); } @@ -48,7 +48,8 @@ export async function ensureBasicTelemetryInstrumentation() { export async function enableTelemetry( telemetryConfig: TelemetryConfig | Promise ) { - global[instrumentationKey] = true; + global[instrumentationKey] = + telemetryConfig instanceof Promise ? telemetryConfig : Promise.resolve(); telemetryConfig = telemetryConfig instanceof Promise From b025bf84bebd7489a104182de69244bc3939403d Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Tue, 5 Nov 2024 09:33:22 -0500 Subject: [PATCH 3/3] fix tests --- js/plugins/google-cloud/tests/traces_test.ts | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/js/plugins/google-cloud/tests/traces_test.ts b/js/plugins/google-cloud/tests/traces_test.ts index 470cb1251e..654ce5f5fa 100644 --- a/js/plugins/google-cloud/tests/traces_test.ts +++ b/js/plugins/google-cloud/tests/traces_test.ts @@ -188,7 +188,7 @@ describe('GoogleCloudTracing', () => { ); const testFlow = createFlow(ai, 'modelFlow', async () => { return run('runFlow', async () => { - ai.generate({ + await ai.generate({ model: echoModel, prompt: 'Testing model telemetry', }); @@ -199,10 +199,11 @@ describe('GoogleCloudTracing', () => { const spans = await getExportedSpans(); - assert.equal(spans[0].name, 'runFlow'); - assert.equal(spans[1].name, 'modelFlow'); - assert.equal(spans[2].name, 'echoModel'); - assert.equal(spans[2].attributes['genkit/model'], 'echoModel'); + assert.equal(spans[0].name, 'echoModel'); + assert.equal(spans[0].attributes['genkit/model'], 'echoModel'); + assert.equal(spans[1].name, 'generate'); + assert.equal(spans[2].name, 'runFlow'); + assert.equal(spans[3].name, 'modelFlow'); }); it('attaches additional span', async () => {