Add tags support in job submission, fix workspace's cancel job operation

azure-quantum/azure/quantum/job/base_job.py

@@ -86,6 +86,7 @@ def from_input_data(
         input_params: Dict[str, Any] = None,
         session_id: Optional[str] = None,
         priority: Optional[str] = None,
+        tags: Optional[list[str]] = None,
         **kwargs
     ) -> "BaseJob":
         """Create a new Azure Quantum job based on a raw input_data payload.
@@ -120,6 +121,8 @@ def from_input_data(
         :type input_params: Dict[str, Any]
         :param priority: Priority of job.
         :type priority: str
+        :param tags: Tags for the job.
+        :type tags: list[str]
         :return: Azure Quantum Job
         :rtype: Job
         """
@@ -157,6 +160,7 @@ def from_input_data(
             input_params=input_params,
             session_id=session_id,
             priority=priority,
+            tags=tags,
             **kwargs
         )
 
@@ -176,6 +180,7 @@ def from_storage_uri(
         submit_job: bool = True,
         session_id: Optional[str] = None,
         priority: Optional[str] = None,
+        tags: Optional[list[str]] = None,
         **kwargs
     ) -> "BaseJob":
         """Create new Job from URI if input data is already uploaded
@@ -205,6 +210,8 @@ def from_storage_uri(
         :type submit_job: bool
         :param priority: Priority of job.
         :type priority: str
+        :param tags: Tags for the job.
+        :type tags: list[str]
         :return: Job instance
         :rtype: Job
         """
@@ -231,6 +238,7 @@ def from_storage_uri(
             input_params=input_params,
             session_id=session_id,
             priority=priority,
+            tags=tags,
             **kwargs
         )
         job = cls(workspace, details, **kwargs)

azure-quantum/azure/quantum/job/session.py

@@ -51,6 +51,9 @@ class Session(WorkspaceItem):
     :param priority: Priority of the session.
     :type priority: Optional[str]
 
+    :param tags: Tags for the session.
+    :type tags: Optional[list[str]]
+
     :raises ValueError: if details is passed along individual parameters,
                         or if required parameters are missing.
     """
@@ -65,6 +68,7 @@ def __init__(
             name: Optional[str] = None,
             job_failure_policy: Union[str, SessionJobFailurePolicy, None] = None,
             priority: Optional[str] = None,
+            tags: Optional[list[str]] = None,
             **kwargs):
         from azure.quantum.target import Target
         target_name = target.name if isinstance(target, Target) else target
@@ -97,6 +101,7 @@ def __init__(
                                      target=target_name,
                                      job_failure_policy=job_failure_policy,
                                      priority=priority,
+                                     tags=tags,
                                      **kwargs)
 
         super().__init__(
@@ -265,6 +270,7 @@ def open_session(
         name: Optional[str] = None,
         job_failure_policy: Union[str, SessionJobFailurePolicy, None] = None,
         priority: Optional[str] = None,
+        tags: Optional[list[str]] = None,
         **kwargs
     ) -> Session:
         """Opens a session and associates all future job submissions to that
@@ -306,6 +312,9 @@ def open_session(
         :param priority: Priority of the session.
         :type priority: Optional[str]
 
+        :param tags: Tags for the session.
+        :type tags: Optional[list[str]]
+
         :return: The session object with updated details after its opening.
         :rtype: Session
         """
@@ -320,6 +329,7 @@ def open_session(
                           target=self._get_azure_target_id(),
                           provider_id=self._get_azure_provider_id(),
                           priority=priority,
+                          tags=tags,
                           **kwargs)
         self.latest_session = session
         return session.open()

azure-quantum/azure/quantum/target/target.py

@@ -211,6 +211,7 @@ def submit(
         shots: int = None,
         input_params: Union[Dict[str, Any], InputParams, None] = None,
         priority: Optional[str] = None,
+        tags: Optional[list[str]] = None,
         **kwargs
     ) -> Job:
         """Submit input data and return Job.
@@ -228,6 +229,8 @@ def submit(
         :type input_params: Dict[str, Any]
         :param priority: Priority of job.
         :type priority: str
+        :param tags: Tags for the job.
+        :type tags: list[str]
         :return: Azure Quantum job
         :rtype: azure.quantum.job.Job
         """
@@ -324,6 +327,7 @@ def _get_entrypoint(input_data):
             input_params=input_params,
             session_id=self.get_latest_session_id(),
             priority=priority,
+            tags=tags,
             **kwargs
         )
 

azure-quantum/azure/quantum/workspace.py

@@ -22,6 +22,7 @@
     Union,
 )
 from typing_extensions import Self
+from azure.core.exceptions import HttpResponseError
 from azure.core.paging import ItemPaged
 from azure.quantum._client import WorkspaceClient
 from azure.quantum._client.models import JobDetails, ItemDetails, SessionDetails
@@ -443,11 +444,19 @@ def cancel_job(self, job: Job) -> Job:
         :rtype: Job
         """
         client = self._get_jobs_client()
-        client.delete(
-            self.subscription_id,
-            self.resource_group,
-            self.name,
-            job.details.id)
+
+        try:
+            client.cancel(
+                self.subscription_id,
+                self.resource_group,
+                self.name,
+                job.details.id)
+        except HttpResponseError as e:
+            # because of historical behavior of the service, the 204 No Content response is returned when cancellation request is succeeded.
+            # while backend is not updated to return 200 according to a guideline, let's handle that here to align with typespecs
+            if e.status_code != 204:
+                raise
+
         details = client.get(
             self.subscription_id,
             self.resource_group,

azure-quantum/tests/mock_client.py

@@ -11,6 +11,7 @@
 from typing import List, Optional
 from datetime import datetime, UTC, timedelta
 
+from azure.core.exceptions import HttpResponseError
 from azure.core.paging import ItemPaged
 from azure.quantum.workspace import Workspace
 from types import SimpleNamespace
@@ -200,6 +201,21 @@ def delete(
                 return None
         raise KeyError(job_id)
 
+    def cancel(
+        self,
+        subscription_id: str,
+        resource_group_name: str,
+        workspace_name: str,
+        job_id: str,
+    ) -> None:
+        for jd in self._store:
+            if jd.id == job_id:
+                jd.status = "Cancelled"
+                error = HttpResponseError("204 No Content")
+                error.status_code = 204
+                raise error
+        raise KeyError(job_id)
+
     def list(
         self,
         subscription_id: str,

azure-quantum/tests/test_workspace.py

@@ -5,6 +5,8 @@
 
 import os
 from unittest import mock
+from azure.quantum.job.job import Job
+from azure.quantum._client.models import JobDetails
 from azure.quantum._constants import EnvironmentVariables, ConnectionConstants
 from azure.core.credentials import AzureKeyCredential
 from azure.core.pipeline.policies import AzureKeyCredentialPolicy
@@ -357,6 +359,32 @@ def assert_value_error(exception: Exception):
             assert "Invalid resource id" in e.args[0]
 
 
+def test_workspace_cancel_job_success():
+    ws = WorkspaceMock(
+        subscription_id=SUBSCRIPTION_ID,
+        resource_group=RESOURCE_GROUP,
+        name=WORKSPACE,
+    )
+
+    job_id = "test-cancel-success"
+    details = JobDetails(
+        id=job_id,
+        name=f"job-{job_id}",
+        container_uri="https://example.com/container",
+        input_data_format="microsoft.resource-estimate.v2",
+        provider_id="ionq",
+        target="ionq.simulator",
+        status="Executing",
+    )
+    ws._client.services.jobs._store.append(details)
+
+    job = Job(ws, details)
+    result = ws.cancel_job(job)
+
+    assert result.details.status == "Cancelled"
+    assert result.id == job_id
+
+
 def test_workspace_user_agent_appid():
     app_id = "MyEnvVarAppId"
     user_agent = "MyUserAgent"