fix(Observe): missing traces in Mlflow

Signed-off-by: Milan Gallas <[email protected]>

package.json

@@ -84,6 +84,7 @@
   },
   "devDependencies": {
     "@mikro-orm/cli": "6.2.9",
+    "@opentelemetry/exporter-trace-otlp-proto": "^0.57.0",
     "@release-it/conventional-changelog": "^9.0.2",
     "@types/jsonwebtoken": "^9.0.6",
     "@types/node": "^20.14.1",

src/opentelemetry.ts

@@ -18,38 +18,46 @@ import 'dotenv/config';
 
 import { setTimeout } from 'node:timers/promises';
 
-import { NodeSDK, resources, metrics } from '@opentelemetry/sdk-node';
+import { NodeSDK, resources, metrics, tracing } from '@opentelemetry/sdk-node';
 import { getNodeAutoInstrumentations } from '@opentelemetry/auto-instrumentations-node';
 import { ATTR_SERVICE_NAME } from '@opentelemetry/semantic-conventions';
 import { OTLPMetricExporter } from '@opentelemetry/exporter-metrics-otlp-http';
 import { ATTR_DEPLOYMENT_ENVIRONMENT_NAME } from '@opentelemetry/semantic-conventions/incubating';
+import { OTLPTraceExporter } from '@opentelemetry/exporter-trace-otlp-proto';
 
 const ENV = process.env.ENVIRONMENT;
 
+export const isOpenTelemetryEnabled = Boolean(process.env.NODE_ENV === 'production');
+export const batchSpanProcessor = new tracing.BatchSpanProcessor(new OTLPTraceExporter({}));
+
 export const opentelemetrySDK = new NodeSDK({
   resource: new resources.Resource({
     [ATTR_SERVICE_NAME]: `bee-api`,
     [ATTR_DEPLOYMENT_ENVIRONMENT_NAME]: ENV
   }),
+  spanProcessors: [batchSpanProcessor],
   metricReader: new metrics.PeriodicExportingMetricReader({ exporter: new OTLPMetricExporter() }),
   instrumentations: [...getNodeAutoInstrumentations()]
 });
-opentelemetrySDK.start();
-
-let isShuttingDown = false;
-const { promise, resolve } = Promise.withResolvers<void>();
-
-for (const event of ['beforeExit', 'SIGINT', 'SIGTERM']) {
-  process.once(event, () => {
-    if (!isShuttingDown) {
-      isShuttingDown = true;
-      Promise.race([opentelemetrySDK.shutdown(), setTimeout(5_000, null, { ref: false })])
-        .catch((err) => {
-          // eslint-disable-next-line no-console
-          console.error(`Failed to execute shutdown hook`, err);
-        })
-        .finally(() => resolve());
-    }
-    return promise;
-  });
+
+if (isOpenTelemetryEnabled) {
+  opentelemetrySDK.start();
+
+  let isShuttingDown = false;
+  const { promise, resolve } = Promise.withResolvers<void>();
+
+  for (const event of ['beforeExit', 'SIGINT', 'SIGTERM']) {
+    process.once(event, () => {
+      if (!isShuttingDown) {
+        isShuttingDown = true;
+        Promise.race([opentelemetrySDK.shutdown(), setTimeout(5_000, null, { ref: false })])
+          .catch((err) => {
+            // eslint-disable-next-line no-console
+            console.error(`Failed to execute shutdown hook`, err);
+          })
+          .finally(() => resolve());
+      }
+      return promise;
+    });
+  }
 }

src/runs/execution/execute.ts

@@ -46,6 +46,7 @@ import { LoadedRun } from '@/runs/execution/types.js';
 import { UserResource } from '@/tools/entities/tool-resources/user-resource.entity.js';
 import { SystemResource } from '@/tools/entities/tool-resources/system-resource.entity.js';
 import { Attachment } from '@/messages/attachment.entity';
+import { batchSpanProcessor, isOpenTelemetryEnabled } from '@/opentelemetry.js';
 
 const agentExecutionTime = new Summary({
   name: 'agent_execution_time_seconds',
@@ -144,6 +145,7 @@ export async function executeRun(run: LoadedRun) {
       );
 
       await agentRun;
+      if (isOpenTelemetryEnabled) await batchSpanProcessor.forceFlush();
 
       endAgentExecutionTimer();
       run.complete();

src/server.ts

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+import './opentelemetry.js';
 import '@/ui/auth-server.js';
 
 import { JsonSchemaToTsProvider } from '@fastify/type-provider-json-schema-to-ts';