Refactor Azure job submission (#2687)

Refactoring Azure job submission so that this PR is easier to read:
#2686


There's a slight change in order of the telemetry events. We were
sending `submitJob` twice when the Copilot tools were invoked. This is
the only behavior change I'm expecting.

Co-authored-by: Mine Starks <>

Author: minestarks

source/vscode/src/azure/commands.ts

@@ -1,9 +1,28 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+import { log, TargetProfile } from "qsharp-lang";
 import * as vscode from "vscode";
-import { log } from "qsharp-lang";
-
+import { qsharpExtensionId } from "../common";
+import { FullProgramConfig, getActiveProgram } from "../programConfig";
+import { getQirForProgram, QirGenerationError } from "../qirGeneration";
+import {
+  EventType,
+  getActiveDocumentType,
+  QsharpDocumentType,
+  sendTelemetryEvent,
+  UserFlowStatus,
+  UserTaskInvocationType,
+} from "../telemetry";
+import { getRandomGuid } from "../utils";
+import { sendMessageToPanel } from "../webviewPanel";
+import { getTokenForWorkspace } from "./auth";
+import { QuantumUris } from "./networkRequests";
+import {
+  getPreferredTargetProfile,
+  targetSupportQir,
+} from "./providerProperties";
+import { startRefreshCycle } from "./treeRefresher";
 import {
   Job,
   Target,
@@ -18,16 +37,6 @@ import {
   queryWorkspaces,
   submitJob,
 } from "./workspaceActions";
-import { QuantumUris } from "./networkRequests";
-import { getQirForActiveWindow } from "../qirGeneration";
-import {
-  getPreferredTargetProfile,
-  targetSupportQir,
-} from "./providerProperties";
-import { startRefreshCycle } from "./treeRefresher";
-import { getTokenForWorkspace } from "./auth";
-import { qsharpExtensionId } from "../common";
-import { sendMessageToPanel } from "../webviewPanel";
 
 const workspacesSecret = `${qsharpExtensionId}.workspaces`;
 
@@ -112,42 +121,33 @@ export async function initAzureWorkspaces(context: vscode.ExtensionContext) {
 
         const target = treeItem.itemData as Target;
 
-        let qir = "";
         try {
-          qir = await getQirForActiveWindow(
-            getPreferredTargetProfile(target.id),
-          );
-        } catch (e: any) {
-          if (e?.name === "QirGenerationError") {
-            vscode.window.showErrorMessage(e.message, { modal: true });
-            return;
-          }
-        }
-        if (!qir) return;
+          const preferredTargetProfile = getPreferredTargetProfile(target.id);
+          const program = await getActiveProgram({
+            showModalError: true,
+            targetProfileFallback: preferredTargetProfile,
+          });
 
-        const quantumUris = new QuantumUris(
-          treeItem.workspace.endpointUri,
-          treeItem.workspace.id,
-        );
+          if (!program.success) {
+            throw new QirGenerationError(program.errorMsg);
+          }
 
-        try {
-          const token = await getTokenForWorkspace(treeItem.workspace);
-          if (!token) return;
-
-          const jobId = await submitJob(
-            token,
-            quantumUris,
-            qir,
-            target.providerId,
-            target.id,
+          await compileAndSubmit(
+            program.programConfig,
+            preferredTargetProfile,
+            getActiveDocumentType(),
+            UserTaskInvocationType.Command,
+            workspaceTreeProvider,
+            treeItem.workspace,
+            target,
           );
-          if (jobId) {
-            // The job submitted fine. Refresh the workspace until it shows up
-            // and all jobs are finished. Don't await on this, just let it run
-            startRefreshCycle(workspaceTreeProvider, treeItem.workspace, jobId);
+        } catch (e: unknown) {
+          log.warn("Failed to submit job. ", e);
+
+          if (e instanceof QirGenerationError) {
+            vscode.window.showErrorMessage(e.message, { modal: true });
+            return;
           }
-        } catch (e: any) {
-          log.error("Failed to submit job. ", e);
 
           vscode.window.showErrorMessage("Failed to submit the job to Azure.", {
             modal: true,
@@ -401,3 +401,104 @@ function getHistogramBucketsFromData(
   }
   return undefined;
 }
+
+export async function compileAndSubmit(
+  program: FullProgramConfig,
+  targetProfile: TargetProfile,
+  telemetryDocumentType: QsharpDocumentType,
+  telemetryInvocationType: UserTaskInvocationType,
+  workspaceTreeProvider: WorkspaceTreeProvider,
+  workspace: WorkspaceConnection,
+  target: Target,
+  parameters: { jobName: string; shots: number } | undefined = undefined,
+) {
+  const associationId = getRandomGuid();
+  const start = performance.now();
+  sendTelemetryEvent(
+    EventType.SubmitToAzureStart,
+    { associationId, invocationType: telemetryInvocationType },
+    {},
+  );
+
+  const qir = await getQirForProgram(
+    program,
+    targetProfile,
+    telemetryDocumentType,
+  );
+
+  if (!parameters) {
+    const result = await promptForJobParameters();
+    if (!result) {
+      sendTelemetryEvent(
+        EventType.SubmitToAzureEnd,
+        {
+          associationId,
+          reason: "user cancelled parameter input",
+          flowStatus: UserFlowStatus.Aborted,
+        },
+        { timeToCompleteMs: performance.now() - start },
+      );
+      return;
+    }
+    parameters = { jobName: result.jobName, shots: result.numberOfShots };
+  }
+
+  const { jobId } = await submitJob(
+    workspace,
+    associationId,
+    qir,
+    target.providerId,
+    target.id,
+    parameters.jobName,
+    parameters.shots,
+  );
+
+  sendTelemetryEvent(
+    EventType.SubmitToAzureEnd,
+    {
+      associationId,
+      reason: "job submitted",
+      flowStatus: UserFlowStatus.Succeeded,
+    },
+    { timeToCompleteMs: performance.now() - start },
+  );
+
+  // The job submitted fine. Refresh the workspace until it shows up
+  // and all jobs are finished. Don't await on this, just let it run
+  startRefreshCycle(workspaceTreeProvider, workspace, jobId);
+
+  return jobId;
+}
+
+async function promptForJobParameters(): Promise<
+  { jobName: string; numberOfShots: number } | undefined
+> {
+  const jobName = await vscode.window.showInputBox({
+    prompt: "Job name",
+    value: new Date().toISOString(),
+  });
+  if (!jobName) return;
+
+  // validator for the user-provided number of shots input
+  const validateShotsInput = (input: string) => {
+    const result = parseFloat(input);
+    if (isNaN(result) || Math.floor(result) !== result) {
+      return "Number of shots must be an integer";
+    }
+  };
+
+  // prompt the user for the number of shots
+  const numberOfShotsPrompted = await vscode.window.showInputBox({
+    value: "100",
+    prompt: "Number of shots",
+    validateInput: validateShotsInput,
+  });
+
+  // abort if the user hits <Esc> during shots entry
+  if (numberOfShotsPrompted === undefined) {
+    return;
+  }
+
+  const numberOfShots = parseInt(numberOfShotsPrompted);
+  return { jobName, numberOfShots };
+}

source/vscode/src/azure/networkRequests.ts

@@ -3,6 +3,7 @@
 
 import { log } from "qsharp-lang";
 import { EventType, getUserAgent, sendTelemetryEvent } from "../telemetry";
+import * as vscode from "vscode";
 
 export async function azureRequest(
   uri: string,
@@ -240,30 +241,35 @@ export class QuantumUris {
 }
 
 export class StorageUris {
-  // Requests use a Shared Access Signature. See https://learn.microsoft.com/en-us/rest/api/storageservices/service-sas-examples
+  private storageAccount: string;
+  private sasTokenRaw: string;
 
-  // x-ms-date header should be present in format: Sun, 06 Nov 1994 08:49:37 GMT
-  // See https://learn.microsoft.com/en-us/rest/api/storageservices/representation-of-date-time-values-in-headers
+  constructor(
+    private sasUri: string,
+    private containerName: string,
+  ) {
+    // Parse the Uri to get the storage account and sasToken
+    const sasUriObj = vscode.Uri.parse(sasUri);
+    this.storageAccount = sasUriObj.scheme + "://" + sasUriObj.authority;
 
-  readonly apiVersion = "2023-01-03"; // Pass as x-ms-version header (see https://learn.microsoft.com/en-us/rest/api/storageservices/versioning-for-the-azure-storage-services#authorize-requests-by-using-azure-ad-shared-key-or-shared-key-lite)
+    // Get the raw value to append to other query strings
+    this.sasTokenRaw = sasUri.substring(sasUri.indexOf("?") + 1);
+  }
 
-  // Same for PUT, with a status code of 201 if successful
-  getContainer(storageAccount: string, container: string, sas: string) {
-    return `https://${storageAccount}.blob.core.windows.net/${container}?restype=container&${sas}`;
+  blob(blobName: string) {
+    return `${this.storageAccount}/${this.containerName}/${blobName}`;
   }
 
-  // Same for DELETE, with a status code of 202 if successful
-  // Also same URI for PUT, but must include the following headers:
-  // - x-ms-blob-type: BlockBlob
-  // - Content-Length: <n>
-  // It will return 201 if created.
-  getBlob(
-    storageAccount: string,
-    container: string,
-    blob: string,
-    sas: string,
-  ) {
-    return `https://${storageAccount}.blob.core.windows.net/${container}/${blob}?${sas}`;
+  blobWithSasToken(blobName: string) {
+    return `${this.storageAccount}/${this.containerName}/${blobName}?${this.sasTokenRaw}`;
+  }
+
+  containerWithSasToken() {
+    return this.sasUri;
+  }
+
+  containerPutWithSasToken() {
+    return `${this.storageAccount}/${this.containerName}?restype=container&${this.sasTokenRaw}`;
   }
 }