diff --git a/src/jupyter/assignments.ts b/src/jupyter/assignments.ts index 8b08dded..8f78dd51 100644 --- a/src/jupyter/assignments.ts +++ b/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; +} diff --git a/src/jupyter/assignments.unit.test.ts b/src/jupyter/assignments.unit.test.ts index 22f59ad0..36568770 100644 --- a/src/jupyter/assignments.unit.test.ts +++ b/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; + + 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', () => { diff --git a/src/jupyter/provider.ts b/src/jupyter/provider.ts index 9e2106f3..7f6661d9 100644 --- a/src/jupyter/provider.ts +++ b/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); diff --git a/src/telemetry/api.ts b/src/telemetry/api.ts index 95c8af28..46bcbed9 100644 --- a/src/telemetry/api.ts +++ b/src/telemetry/api.ts @@ -188,8 +188,54 @@ export enum AuthFlow { /** An event representing extension activation. */ type ActivationEvent = Record; -/** An event representing a server assignment */ -type AssignServerEvent = Record; +/** The final outcome of a server assignment attempt. */ +export enum AssignmentOutcome { + ASSIGNMENT_OUTCOME_UNSPECIFIED = 0, + ASSIGNMENT_OUTCOME_SUCCEEDED = 1, + /** + * The requested accelerator was unavailable and no fallback was attempted + * (e.g., the user explicitly requested a CPU server). + */ + ASSIGNMENT_OUTCOME_ACCELERATOR_UNAVAILABLE = 2, + /** + * The requested accelerator was unavailable and the fallback chain was + * exhausted. + */ + ASSIGNMENT_OUTCOME_ALL_ACCELERATORS_UNAVAILABLE = 3, + /** The user already has the maximum number of assignments. */ + ASSIGNMENT_OUTCOME_TOO_MANY_ASSIGNMENTS = 4, + /** + * The user does not have enough quota to be assigned the requested + * configuration. + */ + ASSIGNMENT_OUTCOME_INSUFFICIENT_QUOTA = 5, + /** The user is denylisted from being assigned a server. */ + ASSIGNMENT_OUTCOME_DENYLISTED = 6, + /** Catch-all for unexpected failures. */ + ASSIGNMENT_OUTCOME_OTHER_FAILURE = 7, +} + +/** An event representing a server assignment attempt. */ +interface AssignServerEvent { + /** The final outcome of the assignment attempt. */ + outcome: AssignmentOutcome; + /** The variant of the requested machine type (e.g. "DEFAULT", "GPU", "TPU"). */ + variant: string; + /** The requested accelerator (e.g. "T4", "L4"). Empty when none. */ + accelerator: string; + /** + * The requested machine shape ("STANDARD", "HIGHMEM"). Empty when not + * applicable. + */ + shape: string; + /** The version of the requested runtime image. Empty when not specified. */ + version: string; + /** + * Whether one or more fallback accelerators were attempted before reaching + * the final outcome. + */ + had_fallback: boolean; +} /** An event representing a server auto connection */ type AutoConnectEvent = Record; diff --git a/src/telemetry/index.ts b/src/telemetry/index.ts index 9dfb0b04..ffe29351 100644 --- a/src/telemetry/index.ts +++ b/src/telemetry/index.ts @@ -14,6 +14,7 @@ import { ColabLogEventBase, ColabEvent, CommandSource, + AssignmentOutcome, AuthFlow, ContentBrowserOperation, ContentBrowserTarget, @@ -76,8 +77,26 @@ export const telemetry = { logAutoConnect: () => { log({ auto_connect_event: {} }); }, - logAssignServer: () => { - log({ assign_server_event: {} }); + logAssignServer: ( + outcome: AssignmentOutcome, + config: { + variant: string; + accelerator: string; + shape: string; + version: string; + hadFallback: boolean; + }, + ) => { + log({ + assign_server_event: { + outcome, + variant: config.variant, + accelerator: config.accelerator, + shape: config.shape, + version: config.version, + had_fallback: config.hadFallback, + }, + }); }, logColabToolbar: () => { log({ colab_toolbar_event: {} }); diff --git a/src/telemetry/telemetry.unit.test.ts b/src/telemetry/telemetry.unit.test.ts index ffb37de7..4f3b5ed0 100644 --- a/src/telemetry/telemetry.unit.test.ts +++ b/src/telemetry/telemetry.unit.test.ts @@ -15,6 +15,7 @@ import { newVsCodeStub, VsCodeStub } from '../test/helpers/vscode'; import { ColabLogEventBase, CommandSource, + AssignmentOutcome, AuthFlow, ContentBrowserOperation, ContentBrowserTarget, @@ -219,11 +220,27 @@ describe('Telemetry Module', () => { }); it('logs on server assignment', () => { - telemetry.logAssignServer(); + telemetry.logAssignServer( + AssignmentOutcome.ASSIGNMENT_OUTCOME_SUCCEEDED, + { + variant: 'GPU', + accelerator: 'T4', + shape: 'STANDARD', + version: '', + hadFallback: false, + }, + ); sinon.assert.calledOnceWithExactly(logStub, { ...baseLog, - assign_server_event: {}, + assign_server_event: { + outcome: AssignmentOutcome.ASSIGNMENT_OUTCOME_SUCCEEDED, + variant: 'GPU', + accelerator: 'T4', + shape: 'STANDARD', + version: '', + had_fallback: false, + }, }); });