feat: enrich assign_server_event with outcome and configuration

src/jupyter/assignments.ts

@@ -39,7 +39,7 @@ import {
 } from '../colab/headers';
 import { log } from '../common/logging';
 import { telemetry } from '../telemetry';
-import { CommandSource } from '../telemetry/api';
+import { AssignmentOutcome, CommandSource } from '../telemetry/api';
 import { ProxiedJupyterClient } from './client';
 import { colabProxyWebSocket } from './colab-proxy-websocket';
 import {
@@ -313,60 +313,76 @@ export class AssignmentManager implements Disposable {
     this.guardDisposed();
     const id = randomUUID();
     const { label, variant, accelerator, shape, version } = descriptor;
-    let assignment: Assignment;
+    let outcome = AssignmentOutcome.ASSIGNMENT_OUTCOME_UNSPECIFIED;
+    let hadFallback = false;
     try {
-      if (isColabServerDescriptorWithAccelerator(descriptor)) {
-        assignment = await this.assignWithFallback(
-          id,
-          descriptor,
-          /* fallbacks= */ undefined,
-          signal,
-        );
-      } else {
-        ({ assignment } = await this.client.assign(
-          id,
-          { variant, accelerator, shape, version },
-          signal,
-        ));
-      }
-    } catch (error) {
-      log.trace(`Failed assigning server ${id}`, error);
-      if (error instanceof AllAcceleratorsUnavailableError) {
-        void this.notifyAllAcceleratorsUnavailable(error);
-      }
-      // TODO: Consider listing assignments to check if there are too many
-      // before the user goes through the assignment flow. This handling logic
-      // would still be needed for the rare race condition where an assignment
-      // is made (e.g. in Colab web) during the extension assignment flow.
-      if (error instanceof TooManyAssignmentsError) {
-        void this.notifyMaxAssignmentsExceeded();
-      }
-      if (error instanceof InsufficientQuotaError) {
-        void this.notifyInsufficientQuota(error);
-      }
-      if (error instanceof DenylistedError) {
-        this.notifyBanned(error);
+      let assignment: Assignment;
+      try {
+        if (isColabServerDescriptorWithAccelerator(descriptor)) {
+          assignment = await this.assignWithFallback(
+            id,
+            descriptor,
+            /* fallbacks= */ undefined,
+            signal,
+          );
+          hadFallback = assignment.accelerator !== descriptor.accelerator;
+        } else {
+          ({ assignment } = await this.client.assign(
+            id,
+            { variant, accelerator, shape, version },
+            signal,
+          ));
+        }
+      } catch (error) {
+        log.trace(`Failed assigning server ${id}`, error);
+        outcome = errorToAssignmentOutcome(error);
+        if (error instanceof AllAcceleratorsUnavailableError) {
+          hadFallback = error.attempted.length > 1;
+          void this.notifyAllAcceleratorsUnavailable(error);
+        }
+        // TODO: Consider listing assignments to check if there are too many
+        // before the user goes through the assignment flow. This handling logic
+        // would still be needed for the rare race condition where an assignment
+        // is made (e.g. in Colab web) during the extension assignment flow.
+        if (error instanceof TooManyAssignmentsError) {
+          void this.notifyMaxAssignmentsExceeded();
+        }
+        if (error instanceof InsufficientQuotaError) {
+          void this.notifyInsufficientQuota(error);
+        }
+        if (error instanceof DenylistedError) {
+          this.notifyBanned(error);
+        }
+        throw error;
       }
-      throw error;
+      const server = this.toAssignedServer(
+        {
+          id,
+          label,
+          variant: assignment.variant,
+          accelerator: assignment.accelerator,
+        },
+        assignment.endpoint,
+        assignment.runtimeProxyInfo,
+        new Date(),
+      );
+      await this.storage.store([server]);
+      this.assignmentChange.fire({
+        added: [server],
+        removed: [],
+        changed: [],
+      });
+      outcome = AssignmentOutcome.ASSIGNMENT_OUTCOME_SUCCEEDED;
+      return server;
+    } finally {
+      telemetry.logAssignServer(outcome, {
+        variant,
+        accelerator: accelerator ?? '',
+        shape: shape !== undefined ? Shape[shape] : '',
+        version: version ?? '',
+        hadFallback,
+      });
     }
-    const server = this.toAssignedServer(
-      {
-        id,
-        label,
-        variant: assignment.variant,
-        accelerator: assignment.accelerator,
-      },
-      assignment.endpoint,
-      assignment.runtimeProxyInfo,
-      new Date(),
-    );
-    await this.storage.store([server]);
-    this.assignmentChange.fire({
-      added: [server],
-      removed: [],
-      changed: [],
-    });
-    return server;
   }
 
   /**
@@ -902,3 +918,22 @@ function isColabServerDescriptorWithAccelerator(
 ): descriptor is ColabServerDescriptorWithAccelerator {
   return !!descriptor.accelerator;
 }
+
+function errorToAssignmentOutcome(error: unknown): AssignmentOutcome {
+  if (error instanceof AllAcceleratorsUnavailableError) {
+    return AssignmentOutcome.ASSIGNMENT_OUTCOME_ALL_ACCELERATORS_UNAVAILABLE;
+  }
+  if (error instanceof AcceleratorUnavailableError) {
+    return AssignmentOutcome.ASSIGNMENT_OUTCOME_ACCELERATOR_UNAVAILABLE;
+  }
+  if (error instanceof TooManyAssignmentsError) {
+    return AssignmentOutcome.ASSIGNMENT_OUTCOME_TOO_MANY_ASSIGNMENTS;
+  }
+  if (error instanceof InsufficientQuotaError) {
+    return AssignmentOutcome.ASSIGNMENT_OUTCOME_INSUFFICIENT_QUOTA;
+  }
+  if (error instanceof DenylistedError) {
+    return AssignmentOutcome.ASSIGNMENT_OUTCOME_DENYLISTED;
+  }
+  return AssignmentOutcome.ASSIGNMENT_OUTCOME_OTHER_FAILURE;
+}

src/jupyter/assignments.unit.test.ts

@@ -7,7 +7,11 @@
 import { randomUUID } from 'crypto';
 import { assert, expect } from 'chai';
 import fetch, { Headers, Request, Response } from 'node-fetch';
-import sinon, { SinonFakeTimers, SinonStubbedInstance } from 'sinon';
+import sinon, {
+  SinonFakeTimers,
+  SinonStubbedFunction,
+  SinonStubbedInstance,
+} from 'sinon';
 import { MessageItem, Uri } from 'vscode';
 import {
   Assignment,
@@ -32,7 +36,8 @@ import {
   COLAB_CLIENT_AGENT_HEADER,
   COLAB_RUNTIME_PROXY_TOKEN_HEADER,
 } from '../colab/headers';
-import { CommandSource } from '../telemetry/api';
+import { telemetry } from '../telemetry';
+import { AssignmentOutcome, CommandSource } from '../telemetry/api';
 import { TestEventEmitter } from '../test/helpers/events';
 import {
   createJupyterClientStub,
@@ -1318,6 +1323,196 @@ describe('AssignmentManager', () => {
         );
       });
     });
+
+    describe('telemetry', () => {
+      let logStub: SinonStubbedFunction<typeof telemetry.logAssignServer>;
+
+      beforeEach(() => {
+        logStub = sinon.stub(telemetry, 'logAssignServer');
+      });
+
+      afterEach(() => {
+        logStub.restore();
+      });
+
+      it('logs OUTCOME_SUCCEEDED with the requested configuration', async () => {
+        colabClientStub.assign.resolves({
+          assignment: defaultAssignment,
+          isNew: false,
+        });
+
+        await assignmentManager.assignServer(defaultAssignmentDescriptor);
+
+        sinon.assert.calledOnceWithExactly(
+          logStub,
+          AssignmentOutcome.ASSIGNMENT_OUTCOME_SUCCEEDED,
+          {
+            variant: Variant.GPU,
+            accelerator: 'A100',
+            shape: '',
+            version: '',
+            hadFallback: false,
+          },
+        );
+      });
+
+      it('logs hadFallback=true when a fallback succeeds', async () => {
+        colabClientStub.getUserInfo.resolves({
+          subscriptionTier: SubscriptionTier.PRO,
+          paidComputeUnitsBalance: 1,
+          eligibleAccelerators: [
+            { variant: Variant.GPU, models: ['T4', 'A100'] },
+          ],
+          ineligibleAccelerators: [],
+        });
+        colabClientStub.assign
+          .withArgs(sinon.match(isUUID), {
+            variant: Variant.GPU,
+            accelerator: 'A100',
+            shape: undefined,
+            version: undefined,
+          })
+          .rejects(new AcceleratorUnavailableError('A100'))
+          .withArgs(sinon.match(isUUID), {
+            variant: Variant.GPU,
+            accelerator: 'T4',
+            shape: undefined,
+            version: undefined,
+          })
+          .resolves({
+            assignment: { ...defaultAssignment, accelerator: 'T4' },
+            isNew: false,
+          });
+
+        await assignmentManager.assignServer(defaultAssignmentDescriptor);
+
+        sinon.assert.calledOnceWithExactly(
+          logStub,
+          AssignmentOutcome.ASSIGNMENT_OUTCOME_SUCCEEDED,
+          {
+            variant: Variant.GPU,
+            accelerator: 'A100',
+            shape: '',
+            version: '',
+            hadFallback: true,
+          },
+        );
+      });
+
+      it('logs OUTCOME_ALL_ACCELERATORS_UNAVAILABLE when fallbacks are exhausted', async () => {
+        colabClientStub.getUserInfo.resolves({
+          subscriptionTier: SubscriptionTier.PRO,
+          paidComputeUnitsBalance: 1,
+          eligibleAccelerators: [
+            { variant: Variant.GPU, models: ['T4', 'A100'] },
+          ],
+          ineligibleAccelerators: [],
+        });
+        colabClientStub.assign.rejects(new AcceleratorUnavailableError('any'));
+
+        await expect(
+          assignmentManager.assignServer(defaultAssignmentDescriptor),
+        ).to.be.rejected;
+
+        sinon.assert.calledOnceWithExactly(
+          logStub,
+          AssignmentOutcome.ASSIGNMENT_OUTCOME_ALL_ACCELERATORS_UNAVAILABLE,
+          {
+            variant: Variant.GPU,
+            accelerator: 'A100',
+            shape: '',
+            version: '',
+            hadFallback: true,
+          },
+        );
+      });
+
+      const errorOutcomeCases = [
+        {
+          label: 'OUTCOME_TOO_MANY_ASSIGNMENTS',
+          error: new TooManyAssignmentsError(),
+          outcome: AssignmentOutcome.ASSIGNMENT_OUTCOME_TOO_MANY_ASSIGNMENTS,
+        },
+        {
+          label: 'OUTCOME_INSUFFICIENT_QUOTA',
+          error: new InsufficientQuotaError('💰🐖'),
+          outcome: AssignmentOutcome.ASSIGNMENT_OUTCOME_INSUFFICIENT_QUOTA,
+        },
+        {
+          label: 'OUTCOME_DENYLISTED',
+          error: new DenylistedError('👨‍⚖️'),
+          outcome: AssignmentOutcome.ASSIGNMENT_OUTCOME_DENYLISTED,
+        },
+        {
+          label: 'OUTCOME_OTHER_FAILURE for unexpected errors',
+          error: new Error('boom'),
+          outcome: AssignmentOutcome.ASSIGNMENT_OUTCOME_OTHER_FAILURE,
+        },
+      ];
+      for (const { label, error, outcome } of errorOutcomeCases) {
+        it(`logs ${label}`, async () => {
+          colabClientStub.assign.rejects(error);
+
+          await expect(
+            assignmentManager.assignServer(defaultAssignmentDescriptor),
+          ).to.be.rejected;
+
+          sinon.assert.calledOnceWithExactly(logStub, outcome, {
+            variant: Variant.GPU,
+            accelerator: 'A100',
+            shape: '',
+            version: '',
+            hadFallback: false,
+          });
+        });
+      }
+
+      it('logs the requested shape and version when present', async () => {
+        colabClientStub.assign.resolves({
+          assignment: defaultAssignment,
+          isNew: false,
+        });
+
+        await assignmentManager.assignServer({
+          ...defaultAssignmentDescriptor,
+          shape: Shape.HIGHMEM,
+          version: 'v1',
+        });
+
+        sinon.assert.calledOnceWithExactly(
+          logStub,
+          AssignmentOutcome.ASSIGNMENT_OUTCOME_SUCCEEDED,
+          {
+            variant: Variant.GPU,
+            accelerator: 'A100',
+            shape: 'HIGHMEM',
+            version: 'v1',
+            hadFallback: false,
+          },
+        );
+      });
+
+      it('logs an empty accelerator for the default CPU descriptor', async () => {
+        colabClientStub.assign.resolves({
+          assignment: { ...defaultAssignment, accelerator: 'NONE' },
+          isNew: false,
+        });
+
+        await assignmentManager.assignServer(DEFAULT_CPU_SERVER);
+
+        sinon.assert.calledOnceWithExactly(
+          logStub,
+          AssignmentOutcome.ASSIGNMENT_OUTCOME_SUCCEEDED,
+          {
+            variant: Variant.DEFAULT,
+            accelerator: '',
+            shape: '',
+            version: '',
+            hadFallback: false,
+          },
+        );
+      });
+    });
   });
 
   describe('unassignServer', () => {

src/jupyter/provider.ts

@@ -230,7 +230,6 @@ export class ColabJupyterServerProvider
           telemetry.logAutoConnect();
           return await this.assignmentManager.latestOrAutoAssignServer();
         case NEW_SERVER.label:
-          telemetry.logAssignServer();
           return await this.assignServer();
         case OPEN_COLAB_WEB.label:
           openColabWeb(this.vs, CommandSource.COMMAND_SOURCE_SERVER_PROVIDER);