diff --git a/src/AzureClient/AzureClient.cs b/src/AzureClient/AzureClient.cs index ed059d99f9..749e674d22 100644 --- a/src/AzureClient/AzureClient.cs +++ b/src/AzureClient/AzureClient.cs @@ -7,6 +7,7 @@ using System.Collections.Generic; using System.Linq; using System.Net; +using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Quantum; using Microsoft.Azure.Quantum.Client.Models; @@ -110,7 +111,11 @@ public async Task GetConnectionStatusAsync(IChannel channel) return ValidExecutionTargets.ToExecutionResult(); } - private async Task SubmitOrExecuteJobAsync(IChannel channel, AzureSubmissionContext submissionContext, bool execute) + private async Task SubmitOrExecuteJobAsync( + IChannel channel, + AzureSubmissionContext submissionContext, + bool execute, + CancellationToken cancellationToken) { if (ActiveWorkspace == null) { @@ -178,38 +183,45 @@ private async Task SubmitOrExecuteJobAsync(IChannel channel, Az return AzureClientError.JobSubmissionFailed.ToExecutionResult(); } + // If the command was not %azure.execute, simply return the job status. if (!execute) { return await GetJobStatusAsync(channel, MostRecentJobId); } + // If the command was %azure.execute, wait for the job to complete and return the job output. channel.Stdout($"Waiting up to {submissionContext.ExecutionTimeout} seconds for Azure Quantum job to complete..."); - using (var cts = new System.Threading.CancellationTokenSource(TimeSpan.FromSeconds(submissionContext.ExecutionTimeout))) + using var executionTimeoutTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(submissionContext.ExecutionTimeout)); + using var executionCancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(executionTimeoutTokenSource.Token, cancellationToken); { - CloudJob? cloudJob = null; - do + try { - // TODO: Once jupyter-core supports interrupt requests (https://github.com/microsoft/jupyter-core/issues/55), - // handle Jupyter kernel interrupt here and break out of this loop - await Task.Delay(TimeSpan.FromSeconds(submissionContext.ExecutionPollingInterval)); - if (cts.IsCancellationRequested) break; - cloudJob = await ActiveWorkspace.GetJobAsync(MostRecentJobId); - channel.Stdout($"[{DateTime.Now.ToLongTimeString()}] Current job status: {cloudJob?.Status ?? "Unknown"}"); + CloudJob? cloudJob = null; + while (cloudJob == null || cloudJob.InProgress) + { + executionCancellationTokenSource.Token.ThrowIfCancellationRequested(); + await Task.Delay(TimeSpan.FromSeconds(submissionContext.ExecutionPollingInterval), executionCancellationTokenSource.Token); + cloudJob = await ActiveWorkspace.GetJobAsync(MostRecentJobId); + channel.Stdout($"[{DateTime.Now.ToLongTimeString()}] Current job status: {cloudJob?.Status ?? "Unknown"}"); + } + } + catch (Exception e) when (e is TaskCanceledException || e is OperationCanceledException) + { + Logger?.LogInformation($"Operation canceled while waiting for job execution to complete: {e.Message}"); } - while (cloudJob == null || cloudJob.InProgress); } return await GetJobResultAsync(channel, MostRecentJobId); } /// - public async Task SubmitJobAsync(IChannel channel, AzureSubmissionContext submissionContext) => - await SubmitOrExecuteJobAsync(channel, submissionContext, execute: false); + public async Task SubmitJobAsync(IChannel channel, AzureSubmissionContext submissionContext, CancellationToken? cancellationToken = null) => + await SubmitOrExecuteJobAsync(channel, submissionContext, execute: false, cancellationToken ?? CancellationToken.None); /// - public async Task ExecuteJobAsync(IChannel channel, AzureSubmissionContext submissionContext) => - await SubmitOrExecuteJobAsync(channel, submissionContext, execute: true); + public async Task ExecuteJobAsync(IChannel channel, AzureSubmissionContext submissionContext, CancellationToken? cancellationToken = null) => + await SubmitOrExecuteJobAsync(channel, submissionContext, execute: true, cancellationToken ?? CancellationToken.None); /// public async Task GetActiveTargetAsync(IChannel channel) diff --git a/src/AzureClient/IAzureClient.cs b/src/AzureClient/IAzureClient.cs index 441abc7bff..c820564c3e 100644 --- a/src/AzureClient/IAzureClient.cs +++ b/src/AzureClient/IAzureClient.cs @@ -5,6 +5,7 @@ using System.Collections.Generic; using System.ComponentModel; +using System.Threading; using System.Threading.Tasks; using Microsoft.Jupyter.Core; @@ -44,7 +45,7 @@ public Task ConnectAsync(IChannel channel, /// /// Details of the submitted job, or an error if submission failed. /// - public Task SubmitJobAsync(IChannel channel, AzureSubmissionContext submissionContext); + public Task SubmitJobAsync(IChannel channel, AzureSubmissionContext submissionContext, CancellationToken? token); /// /// Executes the specified Q# operation as a job to the currently active target @@ -53,7 +54,7 @@ public Task ConnectAsync(IChannel channel, /// /// The result of the executed job, or an error if execution failed. /// - public Task ExecuteJobAsync(IChannel channel, AzureSubmissionContext submissionContext); + public Task ExecuteJobAsync(IChannel channel, AzureSubmissionContext submissionContext, CancellationToken? token); /// /// Sets the specified target for job submission. diff --git a/src/AzureClient/Magic/AzureClientMagicBase.cs b/src/AzureClient/Magic/AzureClientMagicBase.cs index 823554c48f..da17815756 100644 --- a/src/AzureClient/Magic/AzureClientMagicBase.cs +++ b/src/AzureClient/Magic/AzureClientMagicBase.cs @@ -4,6 +4,7 @@ #nullable enable using System; +using System.Threading; using System.Threading.Tasks; using Microsoft.Jupyter.Core; using Microsoft.Quantum.IQSharp.Jupyter; @@ -35,11 +36,15 @@ public AzureClientMagicBase(IAzureClient azureClient, string keyword, Documentat /// public override ExecutionResult Run(string input, IChannel channel) => - RunAsync(input, channel).GetAwaiter().GetResult(); + RunCancellable(input, channel, CancellationToken.None); + + /// + public override ExecutionResult RunCancellable(string input, IChannel channel, CancellationToken cancellationToken) => + RunAsync(input, channel, cancellationToken).GetAwaiter().GetResult(); /// /// Executes the magic command functionality for the given input. /// - public abstract Task RunAsync(string input, IChannel channel); + public abstract Task RunAsync(string input, IChannel channel, CancellationToken cancellationToken); } } diff --git a/src/AzureClient/Magic/ConnectMagic.cs b/src/AzureClient/Magic/ConnectMagic.cs index ffdceff925..867815d2cd 100644 --- a/src/AzureClient/Magic/ConnectMagic.cs +++ b/src/AzureClient/Magic/ConnectMagic.cs @@ -6,6 +6,7 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Threading; using System.Threading.Tasks; using Microsoft.Jupyter.Core; using Microsoft.Quantum.IQSharp.Jupyter; @@ -90,7 +91,7 @@ credentials when connecting to Azure. /// Connects to an Azure workspace given a subscription ID, resource group name, /// workspace name, and connection string as a JSON-encoded object. /// - public override async Task RunAsync(string input, IChannel channel) + public override async Task RunAsync(string input, IChannel channel, CancellationToken cancellationToken) { var inputParameters = ParseInputParameters(input); diff --git a/src/AzureClient/Magic/ExecuteMagic.cs b/src/AzureClient/Magic/ExecuteMagic.cs index 122882abe5..6a675e9f3e 100644 --- a/src/AzureClient/Magic/ExecuteMagic.cs +++ b/src/AzureClient/Magic/ExecuteMagic.cs @@ -5,6 +5,7 @@ using System; using System.Collections.Generic; +using System.Threading; using System.Threading.Tasks; using Microsoft.Jupyter.Core; using Microsoft.Quantum.IQSharp.Jupyter; @@ -56,9 +57,9 @@ The Azure Quantum workspace must previously have been initialized /// name that is present in the current Q# Jupyter workspace, and /// waits for the job to complete before returning. /// - public override async Task RunAsync(string input, IChannel channel) + public override async Task RunAsync(string input, IChannel channel, CancellationToken cancellationToken) { - return await AzureClient.ExecuteJobAsync(channel, AzureSubmissionContext.Parse(input)); + return await AzureClient.ExecuteJobAsync(channel, AzureSubmissionContext.Parse(input), cancellationToken); } } } \ No newline at end of file diff --git a/src/AzureClient/Magic/JobsMagic.cs b/src/AzureClient/Magic/JobsMagic.cs index 590e77b418..1ac082350c 100644 --- a/src/AzureClient/Magic/JobsMagic.cs +++ b/src/AzureClient/Magic/JobsMagic.cs @@ -6,6 +6,7 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Threading; using System.Threading.Tasks; using Microsoft.Jupyter.Core; using Microsoft.Quantum.IQSharp.Jupyter; @@ -52,7 +53,7 @@ The Azure Quantum workspace must previously have been initialized /// /// Lists all jobs in the active workspace. /// - public override async Task RunAsync(string input, IChannel channel) => + public override async Task RunAsync(string input, IChannel channel, CancellationToken cancellationToken) => await AzureClient.GetJobListAsync(channel); } } \ No newline at end of file diff --git a/src/AzureClient/Magic/OutputMagic.cs b/src/AzureClient/Magic/OutputMagic.cs index dd542d7329..624b4a5a89 100644 --- a/src/AzureClient/Magic/OutputMagic.cs +++ b/src/AzureClient/Magic/OutputMagic.cs @@ -6,6 +6,7 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Threading; using System.Threading.Tasks; using Microsoft.Jupyter.Core; using Microsoft.Quantum.IQSharp.Jupyter; @@ -68,7 +69,7 @@ The Azure Quantum workspace must previously have been initialized /// Displays the output of a given completed job ID, if provided, /// or all jobs submitted in the current session. /// - public override async Task RunAsync(string input, IChannel channel) + public override async Task RunAsync(string input, IChannel channel, CancellationToken cancellationToken) { var inputParameters = ParseInputParameters(input, firstParameterInferredName: ParameterNameJobId); string jobId = inputParameters.DecodeParameter(ParameterNameJobId); diff --git a/src/AzureClient/Magic/StatusMagic.cs b/src/AzureClient/Magic/StatusMagic.cs index 80a2a85733..2ba8bcd0be 100644 --- a/src/AzureClient/Magic/StatusMagic.cs +++ b/src/AzureClient/Magic/StatusMagic.cs @@ -6,6 +6,7 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Threading; using System.Threading.Tasks; using Microsoft.Jupyter.Core; using Microsoft.Quantum.IQSharp.Jupyter; @@ -66,7 +67,7 @@ The Azure Quantum workspace must previously have been initialized /// Displays the status corresponding to a given job ID, if provided, /// or the most recently-submitted job in the current session. /// - public override async Task RunAsync(string input, IChannel channel) + public override async Task RunAsync(string input, IChannel channel, CancellationToken cancellationToken) { var inputParameters = ParseInputParameters(input, firstParameterInferredName: ParameterNameJobId); string jobId = inputParameters.DecodeParameter(ParameterNameJobId); diff --git a/src/AzureClient/Magic/SubmitMagic.cs b/src/AzureClient/Magic/SubmitMagic.cs index 3c7bcfccdd..2a46a9c03b 100644 --- a/src/AzureClient/Magic/SubmitMagic.cs +++ b/src/AzureClient/Magic/SubmitMagic.cs @@ -4,6 +4,7 @@ #nullable enable using System.Collections.Generic; +using System.Threading; using System.Threading.Tasks; using Microsoft.Jupyter.Core; using Microsoft.Quantum.IQSharp.Jupyter; @@ -52,9 +53,9 @@ The Azure Quantum workspace must previously have been initialized /// Submits a new job to an Azure Quantum workspace given a Q# operation /// name that is present in the current Q# Jupyter workspace. /// - public override async Task RunAsync(string input, IChannel channel) + public override async Task RunAsync(string input, IChannel channel, CancellationToken cancellationToken) { - return await AzureClient.SubmitJobAsync(channel, AzureSubmissionContext.Parse(input)); + return await AzureClient.SubmitJobAsync(channel, AzureSubmissionContext.Parse(input), cancellationToken); } } } \ No newline at end of file diff --git a/src/AzureClient/Magic/TargetMagic.cs b/src/AzureClient/Magic/TargetMagic.cs index d0c93be7e1..d5c4f8b801 100644 --- a/src/AzureClient/Magic/TargetMagic.cs +++ b/src/AzureClient/Magic/TargetMagic.cs @@ -6,6 +6,7 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Threading; using System.Threading.Tasks; using Microsoft.Jupyter.Core; using Microsoft.Quantum.IQSharp.Jupyter; @@ -63,7 +64,7 @@ available in the workspace. /// /// Sets or views the target for job submission to the current Azure Quantum workspace. /// - public override async Task RunAsync(string input, IChannel channel) + public override async Task RunAsync(string input, IChannel channel, CancellationToken cancellationToken) { var inputParameters = ParseInputParameters(input, firstParameterInferredName: ParameterNameTargetId); if (inputParameters.ContainsKey(ParameterNameTargetId)) diff --git a/src/Jupyter/Jupyter.csproj b/src/Jupyter/Jupyter.csproj index 67a9b01310..b092621bf2 100644 --- a/src/Jupyter/Jupyter.csproj +++ b/src/Jupyter/Jupyter.csproj @@ -35,7 +35,7 @@ - + diff --git a/src/Jupyter/Magic/AbstractMagic.cs b/src/Jupyter/Magic/AbstractMagic.cs index 64a1326266..9ddcd9e81c 100644 --- a/src/Jupyter/Magic/AbstractMagic.cs +++ b/src/Jupyter/Magic/AbstractMagic.cs @@ -6,18 +6,18 @@ using System.IO; using System.Linq; using System.Text.RegularExpressions; +using System.Threading; using System.Threading.Tasks; using Microsoft.Jupyter.Core; using Microsoft.Quantum.IQSharp.Common; using Microsoft.Quantum.QsCompiler.Serialization; -using Newtonsoft.Json.Linq; namespace Microsoft.Quantum.IQSharp.Jupyter { /// /// Abstract base class for IQ# magic symbols. /// - public abstract class AbstractMagic : MagicSymbol + public abstract class AbstractMagic : CancellableMagicSymbol { /// /// Constructs a new magic symbol given its name and documentation. @@ -28,7 +28,7 @@ public AbstractMagic(string keyword, Documentation docs) this.Documentation = docs; this.Kind = SymbolKind.Magic; - this.Execute = SafeExecute(this.Run); + this.ExecuteCancellable = this.SafeExecute(this.RunCancellable); } /// @@ -38,31 +38,32 @@ public AbstractMagic(string keyword, Documentation docs) /// returned execution function displays the given exceptions to its /// display channel. /// - public Func> SafeExecute(Func magic) => - async (input, channel) => - { - channel = channel.WithNewLines(); - - try - { - return magic(input, channel); - } - catch (InvalidWorkspaceException ws) + public Func> SafeExecute( + Func magic) => + async (input, channel, cancellationToken) => { - foreach (var m in ws.Errors) channel.Stderr(m); - return ExecuteStatus.Error.ToExecutionResult(); - } - catch (AggregateException agg) - { - foreach (var e in agg.InnerExceptions) channel.Stderr(e?.Message); - return ExecuteStatus.Error.ToExecutionResult(); - } - catch (Exception e) - { - channel.Stderr(e.Message); - return ExecuteStatus.Error.ToExecutionResult(); - } - }; + channel = channel.WithNewLines(); + + try + { + return magic(input, channel, cancellationToken); + } + catch (InvalidWorkspaceException ws) + { + foreach (var m in ws.Errors) channel.Stderr(m); + return ExecuteStatus.Error.ToExecutionResult(); + } + catch (AggregateException agg) + { + foreach (var e in agg.InnerExceptions) channel.Stderr(e?.Message); + return ExecuteStatus.Error.ToExecutionResult(); + } + catch (Exception e) + { + channel.Stderr(e.Message); + return ExecuteStatus.Error.ToExecutionResult(); + } + }; /// /// Parses the input to a magic command, interpreting the input as @@ -146,5 +147,17 @@ public static Dictionary ParseInputParameters(string input, stri /// A method to be run when the magic command is executed. /// public abstract ExecutionResult Run(string input, IChannel channel); + + /// + /// A method to be run when the magic command is executed, including a cancellation + /// token to use for requesting cancellation. + /// + /// + /// The default implementation in ignores the cancellation token. + /// Derived classes should override this method and monitor the cancellation token if they + /// wish to support cancellation. + /// + public virtual ExecutionResult RunCancellable(string input, IChannel channel, CancellationToken cancellationToken) => + Run(input, channel); } } diff --git a/src/Kernel/Kernel.csproj b/src/Kernel/Kernel.csproj index 242398b43b..b42ee4c7be 100644 --- a/src/Kernel/Kernel.csproj +++ b/src/Kernel/Kernel.csproj @@ -22,7 +22,6 @@ - diff --git a/src/Tests/AzureClientMagicTests.cs b/src/Tests/AzureClientMagicTests.cs index 51763fdbac..19ebf98982 100644 --- a/src/Tests/AzureClientMagicTests.cs +++ b/src/Tests/AzureClientMagicTests.cs @@ -6,6 +6,7 @@ using Microsoft.VisualStudio.TestTools.UnitTesting; using System.Collections.Generic; using System.Linq; +using System.Threading; using System.Threading.Tasks; using Microsoft.Jupyter.Core; using Microsoft.Quantum.IQSharp; @@ -223,14 +224,14 @@ public async Task GetActiveTargetAsync(IChannel channel) return ActiveTargetId.ToExecutionResult(); } - public async Task SubmitJobAsync(IChannel channel, AzureSubmissionContext submissionContext) + public async Task SubmitJobAsync(IChannel channel, AzureSubmissionContext submissionContext, CancellationToken? token) { LastAction = AzureClientAction.SubmitJob; SubmittedJobs.Add(submissionContext.OperationName); return ExecuteStatus.Ok.ToExecutionResult(); } - public async Task ExecuteJobAsync(IChannel channel, AzureSubmissionContext submissionContext) + public async Task ExecuteJobAsync(IChannel channel, AzureSubmissionContext submissionContext, CancellationToken? token) { LastAction = AzureClientAction.ExecuteJob; ExecutedJobs.Add(submissionContext.OperationName); diff --git a/src/Tests/AzureClientTests.cs b/src/Tests/AzureClientTests.cs index 099c251268..773d7a6d8a 100644 --- a/src/Tests/AzureClientTests.cs +++ b/src/Tests/AzureClientTests.cs @@ -6,6 +6,7 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Quantum; using Microsoft.Azure.Quantum.Client.Models; @@ -202,14 +203,14 @@ public void TestJobSubmission() var submissionContext = new AzureSubmissionContext(); // not yet connected - ExpectError(AzureClientError.NotConnected, azureClient.SubmitJobAsync(new MockChannel(), submissionContext)); + ExpectError(AzureClientError.NotConnected, azureClient.SubmitJobAsync(new MockChannel(), submissionContext, CancellationToken.None)); // connect var targets = ExpectSuccess>(ConnectToWorkspaceAsync(azureClient)); Assert.IsFalse(targets.Any()); // no target yet - ExpectError(AzureClientError.NoTarget, azureClient.SubmitJobAsync(new MockChannel(), submissionContext)); + ExpectError(AzureClientError.NoTarget, azureClient.SubmitJobAsync(new MockChannel(), submissionContext, CancellationToken.None)); // add a target var azureWorkspace = azureClient.ActiveWorkspace as MockAzureWorkspace; @@ -221,15 +222,15 @@ public void TestJobSubmission() Assert.AreEqual("ionq.simulator", target.Id); // no operation name specified - ExpectError(AzureClientError.NoOperationName, azureClient.SubmitJobAsync(new MockChannel(), submissionContext)); + ExpectError(AzureClientError.NoOperationName, azureClient.SubmitJobAsync(new MockChannel(), submissionContext, CancellationToken.None)); // specify an operation name, but have missing parameters submissionContext.OperationName = "Tests.qss.HelloAgain"; - ExpectError(AzureClientError.JobSubmissionFailed, azureClient.SubmitJobAsync(new MockChannel(), submissionContext)); + ExpectError(AzureClientError.JobSubmissionFailed, azureClient.SubmitJobAsync(new MockChannel(), submissionContext, CancellationToken.None)); // specify input parameters and verify that the job was submitted submissionContext.InputParameters = new Dictionary() { ["count"] = "3", ["name"] = "testing" }; - var job = ExpectSuccess(azureClient.SubmitJobAsync(new MockChannel(), submissionContext)); + var job = ExpectSuccess(azureClient.SubmitJobAsync(new MockChannel(), submissionContext, CancellationToken.None)); var retrievedJob = ExpectSuccess(azureClient.GetJobStatusAsync(new MockChannel(), job.Id)); Assert.AreEqual(job.Id, retrievedJob.Id); } @@ -261,7 +262,7 @@ public void TestJobExecution() ExecutionTimeout = 5, ExecutionPollingInterval = 1, }; - var histogram = ExpectSuccess(azureClient.ExecuteJobAsync(new MockChannel(), submissionContext)); + var histogram = ExpectSuccess(azureClient.ExecuteJobAsync(new MockChannel(), submissionContext, CancellationToken.None)); Assert.IsNotNull(histogram); } }