feature: context manager for reservation arns #949
Conversation
mbeach-aws
force-pushed
the
reservation-context
branch
from
ff8c616 to
0ca5c2f
Compare
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #949 +/- ##
=========================================
Coverage 100.00% 100.00%
=========================================
Files 134 136 +2
Lines 8911 8932 +21
Branches 2007 2010 +3
=========================================
+ Hits 8911 8932 +21 ☔ View full report in Codecov by Sentry. |
|
|
||
| Args: | ||
| device (Device): The Braket device for which you have a reservation ARN. | ||
| reservation_arn (str | None): The Braket Direct reservation ARN to be applied to all |
There was a problem hiding this comment.
What would be the use case of this with no reservation arn?
There was a problem hiding this comment.
So that a customer doesn't have to modify their code if they don't have an ARN. Example
def run_tasks(device, circuit, shots, reservation_arn=None):
with Reservation(device, reservation_arn=reservation_arn) as reservation:
task1 = device.run(circuit, shots)
task2 = device.run(circuit, shots)
return task1, task2
vs
def run_tasks(device, circuit, shots, reservation_arn=None):
if reservation_arn is None:
task1 = device.run(circuit, shots)
task2 = device.run(circuit, shots)
return task1, task2
elif reservation_arn:
with Reservation(device, reservation_arn=reservation_arn) as reservation:
task1 = device.run(circuit, shots)
task2 = device.run(circuit, shots)
return task1, task2
There was a problem hiding this comment.
updated to make reservation_arn required now, but its still allowed to be None.
Co-authored-by: Abe Coull <85974725+math411@users.noreply.github.com>
|
|
||
|
|
||
| @contextmanager | ||
| def Reservation(device: Device, reservation_arn: str | None = None) -> Device | None: |
There was a problem hiding this comment.
Lower case, since this isn't a class; also just take an AwsDevice instead of Device
| if not isinstance(device, Device): | ||
| raise ValueError("The provided device is not a valid Braket device.") | ||
|
|
||
| # Local simulators do not accept reservation ARNs | ||
| if isinstance(device, LocalSimulator) or device.type == AwsDeviceType.SIMULATOR: | ||
| yield device | ||
|
|
||
| elif isinstance(device, AwsDevice): |
There was a problem hiding this comment.
I'd rearrange as
if isinstance(device, AwsDevice):
...
elif isinstance(device, Device):
...
raise TypeError(...)There was a problem hiding this comment.
Does the API allow reservation arns for simulators? I'm not sure now
Co-authored-by: Cody Wang <speller26@gmail.com>
| modification. | ||
|
|
||
| Examples: | ||
| >>> with Reservation(device, reservation_arn="<my_reservation_arn>"): |
There was a problem hiding this comment.
Minor nit: can you also add an example of this used as a decorator?
There was a problem hiding this comment.
interesting, I did not know the context could also be a decorator. Why don't we also have Tracker available as a decorator?
There was a problem hiding this comment.
I would have to ask around for why that was not the case. Could be that the functionality in contextmanager was not known at the time.
|
|
||
|
|
||
| @contextmanager | ||
| def reservation(device: Device, reservation_arn: str | None) -> Device | None: |
There was a problem hiding this comment.
When customers use reservations, do they only do it for specific portions of their code or would they typically have the entire script using reservation arns? How would that affect the UX we expose here?
There was a problem hiding this comment.
do they only do it for specific portions of their code or would they typically have the entire script using reservation arns?
They are only required to pass reservation arns to their quantum tasks. Hence this wrapper makes it easier to add the require kw args to their tasks without manually modifying every single task.
There was a problem hiding this comment.
Is there a way to make this less intrusive and more agnostic?
| try: | ||
| yield device | ||
| finally: | ||
| device.run = original_run |
There was a problem hiding this comment.
do I need to also modify "run_batch? or is just "run" sufficient?
| @contextmanager | ||
| def reservation(device: Device, reservation_arn: str | None) -> Device | None: | ||
| """ |
There was a problem hiding this comment.
given this takes a Device, I'm not sure how to make it work with a qiskit-braket BraketAwsBackend, or even pennylane qml.device.
There was a problem hiding this comment.
This is usually an indicator that there might be an opportunity to work at a different abstraction. Would it help if we moved this functionality to a different abstraction?
|
closing to make sure qiskit/PL is supported |
feature: context manager for reservation arns
Issue #, if available:
Description of changes:
Adds a context manager that temporarily modifies the
runmethod of adeviceobjectto include a specified
reservation_arnin all calls within the context. This is usefulfor ensuring that all quantum task run within the context are associated with a given
reservation ARN (Amazon Resource Name).
Example usage
Testing done:
Merge Checklist
Put an
xin the boxes that apply. You can also fill these out after creating the PR. If you're unsure about any of them, don't hesitate to ask. We're here to help! This is simply a reminder of what we are going to look for before merging your pull request.General
Tests
By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.