diff --git a/.gitignore b/.gitignore index 2583a455..be21934d 100644 --- a/.gitignore +++ b/.gitignore @@ -408,3 +408,6 @@ azure-quantum/build !azure-quantum/requirements*.txt [.][v]env/ version.py + +# Local scratch/manual test file (not part of the package) +azure-quantum/test_targets.py diff --git a/azure-quantum/azure/quantum/__init__.py b/azure-quantum/azure/quantum/__init__.py index 213f6287..19bb6d04 100644 --- a/azure-quantum/azure/quantum/__init__.py +++ b/azure-quantum/azure/quantum/__init__.py @@ -13,7 +13,7 @@ from .job.session import * from .workspace import * -from ._client.models._enums import JobStatus, SessionStatus, SessionJobFailurePolicy, ItemType +from ._client.models._enums import JobStatus, SessionStatus, SessionJobFailurePolicy, ItemType, Priority logger = logging.getLogger(__name__) logger.info(f"version: {__version__}") diff --git a/azure-quantum/azure/quantum/job/job.py b/azure-quantum/azure/quantum/job/job.py index 9ad0b957..600bc574 100644 --- a/azure-quantum/azure/quantum/job/job.py +++ b/azure-quantum/azure/quantum/job/job.py @@ -8,7 +8,7 @@ import time import json -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, List, Optional, Union from azure.quantum._client.models import JobDetails from azure.quantum.job.job_failed_with_results_error import JobFailedWithResultsError @@ -21,6 +21,7 @@ if TYPE_CHECKING: from azure.quantum.workspace import Workspace + from azure.quantum._client.models import Priority _log = logging.getLogger(__name__) @@ -60,6 +61,36 @@ def delete(self): """Delete the given job.""" self.workspace.delete_job(self) + def update( + self, + *, + name: Optional[str] = None, + priority: Optional[Union[str, "Priority"]] = None, + tags: Optional[List[str]] = None, + ) -> "Job": + """Update the job's name, priority and/or tags after submission. + + Only the arguments that are explicitly provided are updated; + any argument left as ``None`` is left unchanged on the service. + + :param name: The new name of the job. + :param priority: The new priority of the job + (one of :class:`~azure.quantum.Priority`, ``"Standard"`` or ``"High"``). + :param tags: The new list of user-supplied tags associated with the job. + This replaces the existing tags. + + :return: This job, with refreshed details. + :rtype: Job + """ + updated = self.workspace.update_job( + self, + name=name, + priority=priority, + tags=tags, + ) + self.details = updated.details + return self + def has_completed(self) -> bool: """Check if the job has completed.""" return ( diff --git a/azure-quantum/azure/quantum/workspace.py b/azure-quantum/azure/quantum/workspace.py index f7b37f6e..0ad43032 100644 --- a/azure-quantum/azure/quantum/workspace.py +++ b/azure-quantum/azure/quantum/workspace.py @@ -36,6 +36,8 @@ from azure.quantum._client.models import ( BlobDetails, JobStatus, + JobUpdateOptions, + Priority, TargetStatus, ) from azure.quantum import Job, Session @@ -464,6 +466,67 @@ def cancel_job(self, job: Job) -> Job: job.id) return Job(self, details) + def update_job( + self, + job: Job, + *, + name: Optional[str] = None, + priority: Optional[Union[str, Priority]] = None, + tags: Optional[List[str]] = None, + ) -> Job: + """ + Updates the name, priority and/or tags of a job after it has + been submitted. + + Only the arguments that are explicitly provided are updated; + any argument left as ``None`` is left unchanged on the service. + + :param job: + Job to update. + + :param name: + The new name of the job. + + :param priority: + The new priority of the job. + One of :class:`~azure.quantum.Priority` (``"Standard"`` or ``"High"``). + + :param tags: + The new list of user-supplied tags associated with the job. + This replaces the existing tags. + + :return: Azure Quantum Job with updated details. + :rtype: Job + """ + if name is None and priority is None and tags is None: + raise ValueError( + "At least one of 'name', 'priority' or 'tags' must be specified.") + + client = self._get_jobs_client() + job_id = job.id + + update_options = JobUpdateOptions() + if name is not None: + update_options.name = name + if priority is not None: + update_options.priority = priority + if tags is not None: + update_options.tags = tags + + client.update( + self.subscription_id, + self.resource_group, + self.name, + job_id, + update_options) + + details = client.get( + self.subscription_id, + self.resource_group, + self.name, + job_id) + return Job(self, details) + def delete_job(self, job: Job) -> None: """Deletes a job. :param job: diff --git a/azure-quantum/tests/mock_client.py b/azure-quantum/tests/mock_client.py index 6a815d57..328ea8ac 100644 --- a/azure-quantum/tests/mock_client.py +++ b/azure-quantum/tests/mock_client.py @@ -19,6 +19,7 @@ from azure.quantum._client.models import ( ItemDetails, JobDetails, + JobUpdateOptions, ProviderStatus, SessionDetails, TargetStatus, @@ -193,6 +194,40 @@ def get( return jd raise KeyError(job_id) + def update( + self, + subscription_id: str, + resource_group_name: str, + workspace_name: str, + job_id: str, + resource, + ): + def _get(field): + if isinstance(resource, dict): + return resource.get(field) + return getattr(resource, field, None) + + for jd in self._store: + if jd.id == job_id: + name = _get("name") + priority = _get("priority") + tags = _get("tags") + if name is not None: + jd.name = name + if priority is not None: + jd.priority = priority + if tags is not None: + jd.tags = tags + # Match the generated client's contract: the update response is a + # JobUpdateOptions that includes the required job id. + return JobUpdateOptions({ + "id": jd.id, + "name": jd.name, + "priority": jd.priority, + "tags": jd.tags, + }) + raise KeyError(job_id) + # Cancel/delete for older API; mark job as cancelled def delete( self, diff --git a/azure-quantum/tests/test_workspace.py b/azure-quantum/tests/test_workspace.py index 412be557..4fc0151f 100644 --- a/azure-quantum/tests/test_workspace.py +++ b/azure-quantum/tests/test_workspace.py @@ -8,6 +8,7 @@ from unittest import mock from azure.quantum.job.job import Job from azure.quantum._client.models import JobDetails +from azure.quantum import Priority from azure.quantum._constants import EnvironmentVariables, ConnectionConstants from azure.core.credentials import AzureKeyCredential from azure.core.pipeline.policies import AzureKeyCredentialPolicy @@ -439,6 +440,149 @@ def test_job_delete_success(): ws.get_job(job_id) +def test_workspace_update_job_success(): + ws = WorkspaceMock( + subscription_id=SUBSCRIPTION_ID, + resource_group=RESOURCE_GROUP, + name=WORKSPACE, + ) + + job_id = "test-update-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", + priority="Standard", + tags=["old-tag"], + ) + ws._client.services.jobs._store.append(details) + + job = Job(ws, details) + result = ws.update_job( + job, + name="new-name", + priority=Priority.HIGH, + tags=["tag-a", "tag-b"], + ) + + assert result.id == job_id + assert result.details.name == "new-name" + assert result.details.priority == "High" + assert result.details.tags == ["tag-a", "tag-b"] + + +def test_workspace_update_job_partial_leaves_other_fields_unchanged(): + ws = WorkspaceMock( + subscription_id=SUBSCRIPTION_ID, + resource_group=RESOURCE_GROUP, + name=WORKSPACE, + ) + + job_id = "test-update-partial" + details = JobDetails( + id=job_id, + name="original-name", + container_uri="https://example.com/container", + input_data_format="microsoft.resource-estimate.v2", + provider_id="ionq", + target="ionq.simulator", + status="Executing", + priority="Standard", + tags=["keep-me"], + ) + ws._client.services.jobs._store.append(details) + + job = Job(ws, details) + result = ws.update_job(job, name="renamed-only") + + assert result.details.name == "renamed-only" + # untouched fields remain unchanged + assert result.details.priority == "Standard" + assert result.details.tags == ["keep-me"] + + +def test_job_update_success(): + ws = WorkspaceMock( + subscription_id=SUBSCRIPTION_ID, + resource_group=RESOURCE_GROUP, + name=WORKSPACE, + ) + + job_id = "test-job-update-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", + priority="Standard", + tags=["old"], + ) + ws._client.services.jobs._store.append(details) + + job = Job(ws, details) + returned = job.update(name="updated", priority="High", tags=["new"]) + + # update mutates the job in place and returns itself + assert returned is job + assert job.details.name == "updated" + assert job.details.priority == "High" + assert job.details.tags == ["new"] + + +def test_workspace_update_job_not_found_raises(): + ws = WorkspaceMock( + subscription_id=SUBSCRIPTION_ID, + resource_group=RESOURCE_GROUP, + name=WORKSPACE, + ) + + details = JobDetails( + id="missing-job", + name="missing-job", + container_uri="https://example.com/container", + input_data_format="microsoft.resource-estimate.v2", + provider_id="ionq", + target="ionq.simulator", + status="Executing", + ) + job = Job(ws, details) + + with pytest.raises(KeyError): + ws.update_job(job, name="nope") + + +def test_workspace_update_job_requires_at_least_one_field(): + ws = WorkspaceMock( + subscription_id=SUBSCRIPTION_ID, + resource_group=RESOURCE_GROUP, + name=WORKSPACE, + ) + + job_id = "test-update-no-fields" + details = JobDetails( + id=job_id, + name=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) + + with pytest.raises(ValueError): + ws.update_job(job) + + def test_workspace_user_agent_appid(): app_id = "MyEnvVarAppId" user_agent = "MyUserAgent"